Use langchain-astradb 0.1.0 (#312)

cbornet · web-flow · commit fd7df61d458d · 2024-03-14T13:41:42.000+01:00
diff --git a/docs/modules/ROOT/pages/dev-environment.adoc b/docs/modules/ROOT/pages/dev-environment.adoc
@@ -148,7 +148,7 @@ To install the `load_dotenv` package, run `pip install python-dotenv`.
 ----
 import os
 from dotenv import load_dotenv
-from langchain_community.vectorstores import AstraDB
+from langchain_astradb import AstraDBVectorStore
 from langchain_openai import OpenAIEmbeddings
 
 load_dotenv()
@@ -159,7 +159,7 @@ OPEN_AI_API_KEY = os.environ.get("OPENAI_API_KEY")
 ASTRA_DB_COLLECTION = os.environ.get("ASTRA_DB_COLLECTION")
 
 embedding = OpenAIEmbeddings()
-vstore = AstraDB(
+vstore = AstraDBVectorStore(
     embedding=embedding,
     collection_name="test",
     token=os.environ["ASTRA_DB_APPLICATION_TOKEN"],
diff --git a/docs/modules/ROOT/pages/migration.adoc b/docs/modules/ROOT/pages/migration.adoc
@@ -40,7 +40,7 @@ import os
 from datasets import load_dataset
 from dotenv import load_dotenv
 from langchain_community.document_loaders import PyPDFDirectoryLoader
-from langchain_community.vectorstores import AstraDB
+from langchain_astradb import AstraDBVectorStore
 from langchain_openai import OpenAIEmbeddings
 from langchain_core import Document
 
@@ -52,7 +52,7 @@ OPEN_AI_API_KEY = os.environ.get("OPENAI_API_KEY")
 ASTRA_DB_COLLECTION = os.environ.get("ASTRA_DB_COLLECTION")
 
 embedding = OpenAIEmbeddings()
-vstore = AstraDB(
+vstore = AstraDBVectorStore(
     embedding=embedding,
     collection_name="test3",
     token=os.environ["ASTRA_DB_APPLICATION_TOKEN"],
diff --git a/docs/modules/ROOT/pages/quickstart.adoc b/docs/modules/ROOT/pages/quickstart.adoc
@@ -118,7 +118,7 @@ With your environment set up, you're ready to create a RAG workflow in Python.
 import os
 from dotenv import load_dotenv
 from datasets import load_dataset
-from langchain_community.vectorstores import AstraDB
+from langchain_astradb import AstraDBVectorStore
 from langchain_openai import OpenAIEmbeddings
 from langchain_core import Document
 from langchain.prompts import ChatPromptTemplate
@@ -145,7 +145,7 @@ ASTRA_DB_COLLECTION = os.environ.get("ASTRA_DB_COLLECTION")
 [source,python]
 ----
 embedding = OpenAIEmbeddings()
-vstore = AstraDB(
+vstore = AstraDBVectorStore(
     embedding=embedding,
     collection_name="test",
     token=os.environ["ASTRA_DB_APPLICATION_TOKEN"],
@@ -244,7 +244,7 @@ Python::
 import os
 from dotenv import load_dotenv
 from datasets import load_dataset
-from langchain_community.vectorstores import AstraDB
+from langchain_astradb import AstraDBVectorStore
 from langchain_openai import OpenAIEmbeddings
 from langchain_core import Document
 from langchain_core.prompts import ChatPromptTemplate
@@ -260,7 +260,7 @@ OPEN_AI_API_KEY = os.environ.get("OPENAI_API_KEY")
 ASTRA_DB_COLLECTION = os.environ.get("ASTRA_DB_COLLECTION")
 
 embedding = OpenAIEmbeddings()
-vstore = AstraDB(
+vstore = AstraDBVectorStore(
     embedding=embedding,
     collection_name="test",
     token=os.environ["ASTRA_DB_APPLICATION_TOKEN"],
diff --git a/docs/modules/default-architecture/pages/retrieval.adoc b/docs/modules/default-architecture/pages/retrieval.adoc
@@ -16,7 +16,7 @@ Python::
 ----
 import os
 from dotenv import load_dotenv
-from langchain_community.vectorstores import AstraDB
+from langchain_astradb import AstraDBVectorStore
 from langchain_openai import OpenAIEmbeddings
 from langchain_core.prompts import ChatPromptTemplate
 from langchain_core.runnables import RunnablePassthrough
@@ -27,7 +27,7 @@ load_dotenv()
 
 OPEN_AI_API_KEY = os.environ["OPENAI_API_KEY"]
 
-vstore = AstraDB(
+vstore = AstraDBVectorStore(
     embedding=OpenAIEmbeddings(openai_api_key=OPEN_AI_API_KEY),
     collection_name="test",
     token=os.environ["ASTRA_DB_APPLICATION_TOKEN"],
diff --git a/docs/modules/default-architecture/pages/storing.adoc b/docs/modules/default-architecture/pages/storing.adoc
@@ -13,15 +13,15 @@ This code embeds the loaded `Documents` from the xref:splitting.adoc[] example a
 ----
 import os
 from dotenv import load_dotenv
-from langchain_community.vectorstores import AstraDB
+from langchain_astradb import AstraDBVectorStore
 from langchain_openai import OpenAIEmbeddings
 
 load_dotenv()
 
 ASTRA_DB_COLLECTION = os.environ.get("ASTRA_DB_COLLECTION")
 
 embedding = OpenAIEmbeddings()
-vstore = AstraDB(
+vstore = AstraDBVectorStore(
     embedding=embedding,
     collection_name="test",
     token=os.environ["ASTRA_DB_APPLICATION_TOKEN"],
diff --git a/docs/modules/examples/pages/advanced-rag.adoc b/docs/modules/examples/pages/advanced-rag.adoc
@@ -39,7 +39,7 @@ from langchain_openai import OpenAIEmbeddings
 from langchain_community.retrievers import ParentDocumentRetriever
 from langchain_community.storage import InMemoryStore
 from langchain.text_splitter import TokenTextSplitter
-from langchain_community.vectorstores import AstraDB
+from langchain_astradb import AstraDBVectorStore
 ----
 +
 . The two advanced RAG techniques require different dependencies.
@@ -156,7 +156,7 @@ model = ChatOpenAI(openai_api_key=openai_api_key, model_name="gpt-3.5-turbo", te
 embedding = OpenAIEmbeddings(openai_api_key=openai_api_key)
 
 # Initialize a vector store for storing the child chunks
-vstore = AstraDB(
+vstore = AstraDBVectorStore(
     collection_name=collection_name,
     embedding=embedding,
     token=astra_token,
diff --git a/docs/modules/examples/pages/flare.adoc b/docs/modules/examples/pages/flare.adoc
@@ -36,7 +36,7 @@ pip install ragstack-ai
 ----
 import os
 from dotenv import load_dotenv
-from langchain_community.vectorstores import AstraDB
+from langchain_astradb import AstraDBVectorStore
 from langchain_openai import OpenAIEmbeddings
 from langchain_community.document_loaders import TextLoader
 from langchain.globals import set_verbose
@@ -53,7 +53,7 @@ from langchain.chains.flare.base import QuestionGeneratorChain
 [source,python]
 ----
 embedding = OpenAIEmbeddings()
-vstore = AstraDB(
+vstore = AstraDBVectorStore(
         collection_name=collection,
         embedding=embedding,
         token=astra_token,
diff --git a/docs/modules/examples/pages/langchain-evaluation.adoc b/docs/modules/examples/pages/langchain-evaluation.adoc
@@ -38,7 +38,7 @@ pip install ragstack-ai langchain[openai]
 import os
 from dotenv import load_dotenv
 from langchain.embeddings.openai import OpenAIEmbeddings
-from langchain_community.vectorstores import AstraDB
+from langchain_astradb import AstraDBVectorStore
 from langchain_community.document_loaders import TextLoader
 from langchain_community.document_loaders import PyPDFLoader
 from langchain.evaluation import EvaluatorType
@@ -55,7 +55,7 @@ from langchain.chains import RetrievalQA
 [source,python]
 ----
 embedding = OpenAIEmbeddings(openai_api_key=openai_api_key)
-vstore = AstraDB(
+vstore = AstraDBVectorStore(
     collection_name=collection,
     embedding=embedding,
     token=astra_token,
diff --git a/docs/modules/examples/pages/langchain-unstructured-astra.adoc b/docs/modules/examples/pages/langchain-unstructured-astra.adoc
@@ -57,7 +57,7 @@ import os
 import requests
 
 from dotenv import load_dotenv
-from langchain_community.vectorstores import AstraDB
+from langchain_astradb import AstraDBVectorStore
 from langchain_core.documents import Document
 from langchain_core.output_parsers import StrOutputParser
 from langchain_core.prompts import PromptTemplate
@@ -144,7 +144,7 @@ See the Colab notebook linked at the top of this page for a more detailed invest
 +
 [source,python]
 ----
-astra_db_store = AstraDB(
+astra_db_store = AstraDBVectorStore(
     collection_name="langchain_unstructured",
     embedding=OpenAIEmbeddings(),
     token=os.getenv("ASTRA_DB_APPLICATION_TOKEN"),
@@ -236,7 +236,7 @@ import requests
 
 from dotenv import load_dotenv
 from langchain_community.document_loaders import unstructured
-from langchain_community.vectorstores import AstraDB
+from langchain_astradb import AstraDBVectorStore
 from langchain_core.documents import Document
 from langchain_core.output_parsers import StrOutputParser
 from langchain_core.prompts import PromptTemplate
@@ -269,7 +269,7 @@ elements = unstructured.get_elements_from_api(
     pdf_infer_table_structure=True,
 )
 
-astra_db_store = AstraDB(
+astra_db_store = AstraDBVectorStore(
     collection_name="langchain_unstructured",
     embedding=OpenAIEmbeddings(),
     token=os.getenv("ASTRA_DB_APPLICATION_TOKEN"),
diff --git a/docs/modules/examples/pages/nvidia_embeddings.adoc b/docs/modules/examples/pages/nvidia_embeddings.adoc
@@ -69,9 +69,9 @@ embedding = NVIDIAEmbeddings(
 [source,python]
 ----
 import os
-from langchain.vectorstores.astradb import AstraDB
+from langchain_astradb import AstraDBVectorStore
 
-vstore = AstraDB(
+vstore = AstraDBVectorStore(
     collection_name=collection,
     embedding=embedding,
     token=os.getenv("ASTRA_DB_APPLICATION_TOKEN"),
@@ -194,7 +194,7 @@ vstore.delete_collection()
 ----
 from datasets import load_dataset
 from langchain_nvidia_ai_endpoints import NVIDIAEmbeddings, ChatNVIDIA
-from langchain.vectorstores.astradb import AstraDB
+from langchain_astradb import AstraDBVectorStore
 from langchain.schema import Document
 from langchain.prompts import ChatPromptTemplate
 from langchain.chat_models import ChatOpenAI
@@ -210,7 +210,7 @@ embedding = NVIDIAEmbeddings(nvidia_api_key=nvidia_api_key, model="nvolveqa_40k"
 collection_name = "test"
 astra_token = os.getenv("ASTRA_DB_APPLICATION_TOKEN")
 astra_api_endpoint = os.getenv("ASTRA_DB_API_ENDPOINT")
-vstore = AstraDB(collection_name=collection_name, embedding=embedding, 
+vstore = AstraDBVectorStore(collection_name=collection_name, embedding=embedding,
                  token=astra_token, api_endpoint=astra_api_endpoint)
 print("Astra vector store configured")
 
diff --git a/evaluation/tru_shared.py b/evaluation/tru_shared.py
@@ -17,7 +17,7 @@
 
 from langchain_community.chat_models import AzureChatOpenAI
 from langchain_community.embeddings import AzureOpenAIEmbeddings
-from langchain.vectorstores.astradb import AstraDB
+from langchain_astradb import AstraDBVectorStore as LangChainAstraDBVectorStore
 
 # this code assumes the following env vars exist in a .env file:
 #
@@ -171,7 +171,7 @@ def get_azure_embeddings_model(framework: Framework):
 
 def get_astra_vector_store(framework: Framework, collection_name: str):
     if framework == Framework.LANG_CHAIN:
-        return AstraDB(
+        return LangChainAstraDBVectorStore(
             collection_name=collection_name,
             embedding=get_azure_embeddings_model(framework),
             token=os.getenv("ASTRA_DB_TOKEN"),
diff --git a/pyproject.toml b/pyproject.toml
@@ -19,7 +19,7 @@ llama-parse = { version = "0.1.4" }
 langchain = "0.1.12"
 langchain-core = "0.1.31"
 langchain-community = "0.0.28"
-langchain-astradb = "0.0.1"
+langchain-astradb = "0.1.0"
 langchain-openai = "0.0.8"
 langchain-google-genai = { version = "0.0.9", optional = true }
 langchain-google-vertexai = { version = "0.1.0", optional = true }
@@ -38,5 +38,5 @@ tox = "^4"
 requires = ["poetry-core"]
 build-backend = "poetry.core.masonry.api"
 
-[tool.poetry.dev-dependencies]
+[tool.poetry.group.dev.dependencies]
 yamllint = "^1.34.0"
diff --git a/ragstack-e2e-tests/e2e_tests/langchain/test_astra.py b/ragstack-e2e-tests/e2e_tests/langchain/test_astra.py
@@ -7,7 +7,7 @@
 from httpx import ConnectError, HTTPStatusError
 
 from langchain.schema.embeddings import Embeddings
-from langchain.vectorstores import AstraDB
+from langchain_astradb import AstraDBVectorStore
 from e2e_tests.conftest import (
     is_astra,
 )
@@ -20,14 +20,14 @@
 from e2e_tests.test_utils.vector_store_handler import VectorStoreImplementation
 
 
-def test_basic_vector_search(vectorstore: AstraDB):
+def test_basic_vector_search(vectorstore: AstraDBVectorStore):
     print("Running test_basic_vector_search")
     vectorstore.add_texts(["RAGStack is a framework to run LangChain in production"])
     retriever = vectorstore.as_retriever()
     assert len(retriever.get_relevant_documents("RAGStack")) > 0
 
 
-def test_ingest_errors(vectorstore: AstraDB):
+def test_ingest_errors(vectorstore: AstraDBVectorStore):
     print("Running test_ingestion")
 
     empty_text = ""
@@ -59,9 +59,9 @@ def test_ingest_errors(vectorstore: AstraDB):
             )
 
 
-def test_wrong_connection_parameters(vectorstore: AstraDB):
+def test_wrong_connection_parameters(vectorstore: AstraDBVectorStore):
     try:
-        AstraDB(
+        AstraDBVectorStore(
             collection_name="something",
             embedding=MockEmbeddings(),
             token="xxxxx",
@@ -77,7 +77,7 @@ def test_wrong_connection_parameters(vectorstore: AstraDB):
     api_endpoint = vectorstore.api_endpoint
     try:
         print("api_endpoint:", api_endpoint)
-        AstraDB(
+        AstraDBVectorStore(
             collection_name="something",
             embedding=MockEmbeddings(),
             token="this-is-a-wrong-token",
@@ -92,7 +92,7 @@ def test_wrong_connection_parameters(vectorstore: AstraDB):
             )
 
 
-def test_basic_metadata_filtering_no_vector(vectorstore: AstraDB):
+def test_basic_metadata_filtering_no_vector(vectorstore: AstraDBVectorStore):
     collection = vectorstore.collection
     vectorstore.add_texts(
         texts=["RAGStack is a framework to run LangChain in production"],
@@ -423,7 +423,7 @@ def embed_query(self, text: str) -> List[float]:
 
 
 @pytest.fixture
-def vectorstore() -> AstraDB:
+def vectorstore() -> AstraDBVectorStore:
     if not is_astra:
         skip_test_due_to_implementation_not_supported("astradb")
     handler = AstraDBVectorStoreHandler(VectorStoreImplementation.ASTRADB)
diff --git a/ragstack-e2e-tests/e2e_tests/langchain_llamaindex/test_astra.py b/ragstack-e2e-tests/e2e_tests/langchain_llamaindex/test_astra.py
@@ -6,7 +6,7 @@
 from langchain.chains import ConversationalRetrievalChain
 from langchain.chat_models import ChatOpenAI
 from langchain.embeddings import OpenAIEmbeddings
-from langchain.vectorstores.astradb import AstraDB
+from langchain_astradb import AstraDBVectorStore as LangChainAstraDBVectorStore
 
 try:
     # llamaindex 0.9.x
@@ -54,7 +54,7 @@ def test_ingest_llama_retrieve_langchain(astra_ref: AstraRef):
 
     langchain_embeddings = OpenAIEmbeddings(openai_api_key=openai_key)
 
-    langchain_vector_db = AstraDB(
+    langchain_vector_db = LangChainAstraDBVectorStore(
         collection_name=collection,
         embedding=langchain_embeddings,
         token=token,
@@ -157,7 +157,7 @@ def test_ingest_langchain_retrieve_llama_index(astra_ref: AstraRef):
 
     langchain_embeddings = OpenAIEmbeddings(openai_api_key=openai_key)
 
-    vector_db = AstraDB(
+    vector_db = LangChainAstraDBVectorStore(
         collection_name=collection,
         embedding=langchain_embeddings,
         token=token,
diff --git a/ragstack-e2e-tests/e2e_tests/test_utils/astradb_vector_store_handler.py b/ragstack-e2e-tests/e2e_tests/test_utils/astradb_vector_store_handler.py
@@ -8,7 +8,7 @@
 
 import cassio
 from langchain_community.chat_message_histories import AstraDBChatMessageHistory
-from langchain_community.vectorstores.astradb import AstraDB
+from langchain_astradb import AstraDBVectorStore as LangChainVectorStore
 from langchain_core.chat_history import BaseChatMessageHistory
 
 try:
@@ -30,7 +30,7 @@
     CassandraChatMessageHistory,
 )
 
-from astrapy.db import AstraDB as AstraPyClient
+from astrapy.db import AstraDB
 
 
 @dataclass
@@ -94,7 +94,9 @@ def shutdown(self, wait=True):
         self.executor.shutdown(wait=wait)
 
 
-class EnhancedAstraDBLangChainVectorStore(EnhancedLangChainVectorStore, AstraDB):
+class EnhancedAstraDBLangChainVectorStore(
+    EnhancedLangChainVectorStore, LangChainVectorStore
+):
     def put_document(
         self, doc_id: str, document: str, metadata: dict, vector: List[float]
     ) -> None:
@@ -121,7 +123,6 @@ def search_documents(self, vector: List[float], limit: int) -> List[str]:
 class EnhancedAstraDBLlamaIndexVectorStore(
     AstraDBVectorStore, EnhancedLlamaIndexVectorStore
 ):
-
     def put_document(
         self, doc_id: str, document: str, metadata: dict, vector: List[float]
     ) -> None:
@@ -244,7 +245,7 @@ def initialize(cls):
             cls.api_endpoint = get_required_env("ASTRA_DB_ENDPOINT")
             cls.env = os.environ.get("ASTRA_DB_ENV", "prod").lower()
             cls.database_id = get_required_env("ASTRA_DB_ID")
-            cls.default_astra_client = AstraPyClient(
+            cls.default_astra_client = AstraDB(
                 api_endpoint=cls.api_endpoint, token=cls.token
             )
             cls.delete_collection_handler = DeleteCollectionHandler(
diff --git a/ragstack-e2e-tests/pyproject.llamaindex.toml b/ragstack-e2e-tests/pyproject.llamaindex.toml
diff --git a/ragstack-e2e-tests/pyproject.ragstack-ai.toml b/ragstack-e2e-tests/pyproject.ragstack-ai.toml
