From 6119e5cfe82104a935d99feab70fe0ec980440b7 Mon Sep 17 00:00:00 2001 From: Eric Hare Date: Mon, 26 Feb 2024 06:43:41 -0700 Subject: [PATCH] Fix example code for Astra DB pipeline (#481) --- integrations/astra/examples/example.py | 8 +++----- integrations/astra/examples/pipeline_example.py | 10 ++++------ integrations/astra/examples/requirements.txt | 3 ++- 3 files changed, 9 insertions(+), 12 deletions(-) diff --git a/integrations/astra/examples/example.py b/integrations/astra/examples/example.py index eda62834a..00da589bb 100644 --- a/integrations/astra/examples/example.py +++ b/integrations/astra/examples/example.py @@ -22,16 +22,14 @@ logger.info(file_paths) collection_name = os.getenv("COLLECTION_NAME", "haystack_vector_search") -keyspace_name = os.getenv("KEYSPACE_NAME", "recommender_demo") -# Make sure ASTRA_API_ENDPOINT and ASTRA_TOKEN environment variables are set before proceeding +# Make sure ASTRA_DB_API_ENDPOINT and ASTRA_DB_APPLICATION_TOKEN environment variables are set before proceeding # We support many different databases. Here, we load a simple and lightweight in-memory database. document_store = AstraDocumentStore( - astra_collection=collection_name, - astra_keyspace=keyspace_name, + collection_name=collection_name, duplicates_policy=DuplicatePolicy.OVERWRITE, - embedding_dim=384, + embedding_dimension=384, ) # Create components and an indexing pipeline that converts txt files to documents, diff --git a/integrations/astra/examples/pipeline_example.py b/integrations/astra/examples/pipeline_example.py index 731bffd54..826d85956 100644 --- a/integrations/astra/examples/pipeline_example.py +++ b/integrations/astra/examples/pipeline_example.py @@ -30,16 +30,14 @@ """ collection_name = os.getenv("COLLECTION_NAME", "haystack_vector_search") -keyspace_name = os.getenv("KEYSPACE_NAME", "recommender_demo") -# Make sure ASTRA_API_ENDPOINT and ASTRA_TOKEN environment variables are set before proceeding +# Make sure ASTRA_DB_API_ENDPOINT and ASTRA_DB_APPLICATION_TOKEN environment variables are set before proceeding # We support many different databases. Here, we load a simple and lightweight in-memory database. document_store = AstraDocumentStore( - astra_collection=collection_name, - astra_keyspace=keyspace_name, + collection_name=collection_name, duplicates_policy=DuplicatePolicy.SKIP, - embedding_dim=384, + embedding_dimension=384, ) @@ -74,7 +72,7 @@ ) rag_pipeline.add_component(instance=AstraEmbeddingRetriever(document_store=document_store), name="retriever") rag_pipeline.add_component(instance=PromptBuilder(template=prompt_template), name="prompt_builder") -rag_pipeline.add_component(instance=OpenAIGenerator(api_key=os.environ.get("OPENAI_API_KEY")), name="llm") +rag_pipeline.add_component(instance=OpenAIGenerator(), name="llm") rag_pipeline.add_component(instance=AnswerBuilder(), name="answer_builder") rag_pipeline.connect("embedder", "retriever") rag_pipeline.connect("retriever", "prompt_builder.documents") diff --git a/integrations/astra/examples/requirements.txt b/integrations/astra/examples/requirements.txt index be27e7427..284fd46fe 100644 --- a/integrations/astra/examples/requirements.txt +++ b/integrations/astra/examples/requirements.txt @@ -1,3 +1,4 @@ haystack-ai==2.0.0b4 sentence_transformers==2.2.2 -openai==1.6.1 \ No newline at end of file +openai==1.6.1 +astrapy>=0.7.0 \ No newline at end of file