fix vector store

datastax · Feb 23, 2024 · 31cb1c2 · 31cb1c2
1 parent 71de102
commit 31cb1c2
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/ragstack-e2e-tests/e2e_tests/langchain/test_unstructured.py b/ragstack-e2e-tests/e2e_tests/langchain/test_unstructured.py
@@ -6,6 +6,7 @@
 from langchain.prompts import PromptTemplate
 from langchain_core.runnables import RunnablePassthrough
 from langchain.chat_models import ChatOpenAI
+from langchain.embeddings import OpenAIEmbeddings
 
 from e2e_tests.langchain.rag_application import BASIC_QA_PROMPT
 from e2e_tests.test_utils import get_local_resource_path
@@ -47,9 +48,8 @@ def test_unstructured(vector_store, unstructured_mode, request):
     )
 
     vector_store_context: VectorStoreTestContext = request.getfixturevalue(vector_store)
-    vector_store = vector_store_context.new_langchain_vector_store(
-        embedding_dimension=1536
-    )
+    embedding = OpenAIEmbeddings(openai_api_key=get_required_env("OPEN_AI_KEY"))
+    vector_store = vector_store_context.new_langchain_vector_store(embedding=embedding)
 
     loader = UnstructuredAPIFileLoader(
         file_path=get_local_resource_path("tree.pdf"),