upgraded astra packages (#173)

* upgraded astra packages

* fix tests

Author: epinzur

docs/examples/code-generation.ipynb

@@ -616,7 +616,11 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {},
+   "metadata": {
+    "tags": [
+     "skip-execution"
+    ]
+   },
    "outputs": [],
    "source": [
     "import os\n",
@@ -644,7 +648,9 @@
    "execution_count": null,
    "metadata": {
     "tags": [
-     "keep_output"
+     "skip_execution",
+     "keep_output",
+     "raises-exception"
     ]
    },
    "outputs": [
@@ -699,6 +705,7 @@
    "execution_count": null,
    "metadata": {
     "tags": [
+     "skip_execution",
      "keep_output"
     ]
    },
@@ -768,6 +775,7 @@
    "execution_count": null,
    "metadata": {
     "tags": [
+     "skip_execution",
      "keep_output"
     ]
    },

packages/graph-rag-example-helpers/src/graph_rag_example_helpers/datasets/wikimultihop/load.py

@@ -128,15 +128,16 @@ async def aload_2wikimultihop(
         max_tries=MAX_RETRIES,
     )
     async def add_docs(batch_docs, offset) -> None:
-        from astrapy.exceptions import InsertManyException
+        from astrapy.exceptions import CollectionInsertManyException
 
         try:
             await store.aadd_documents(batch_docs)
             persistence.ack(offset)
-        except InsertManyException as err:
-            for err_desc in err.error_descriptors:
-                if err_desc.error_code != "DOCUMENT_ALREADY_EXISTS":
-                    print(err_desc)  # noqa: T201
+        except CollectionInsertManyException as err:
+            for exp in err.exceptions:
+                exp_desc = str(exp)
+                if "DOCUMENT_ALREADY_EXISTS" not in exp_desc:
+                    print(exp_desc)  # noqa: T201
             raise
 
     # We can't use asyncio.TaskGroup in 3.10. This would be simpler with that.

packages/langchain-graph-retriever/pyproject.toml

@@ -95,9 +95,9 @@ graph-retriever = { workspace = true }
 
 [project.optional-dependencies]
 astra = [
-    "astrapy>=1.5.2",
+    "astrapy>=2.0.0",
     "httpx>=0.28.1",
-    "langchain-astradb>=0.5.3",
+    "langchain-astradb>=0.6.0",
 ]
 html = [
     "beautifulsoup4>=4.12.3",

packages/langchain-graph-retriever/tests/adapters/test_astra.py

@@ -5,7 +5,6 @@
 from typing import Any
 
 import pytest
-from astrapy.authentication import StaticTokenProvider
 from graph_retriever.testing.adapter_tests import (
     AdapterComplianceCase,
     AdapterComplianceSuite,
@@ -16,7 +15,9 @@
 from langchain_graph_retriever.adapters.astra import AstraAdapter, _queries
 from typing_extensions import override
 
-TEST_CODEC = _DefaultVSDocumentCodec("page_content", ignore_invalid_documents=True)
+TEST_CODEC = _DefaultVSDocumentCodec(
+    "page_content", ignore_invalid_documents=True, has_lexical=False
+)
 
 
 def create_queries(
@@ -172,7 +173,7 @@ def test_create_metadata_query_user() -> None:
 
 @dataclasses.dataclass
 class _AstraConfig:
-    token: StaticTokenProvider
+    token: str
     keyspace: str
     api_endpoint: str
 
@@ -183,16 +184,17 @@ def astra_config(enabled_stores: set[str]) -> Iterator[_AstraConfig | None]:
         pytest.skip("Pass --stores=astra to test Astra")
         return
 
-    from astrapy import AstraDBDatabaseAdmin
+    from astrapy import DataAPIClient
     from dotenv import load_dotenv
 
     load_dotenv()
 
-    token = StaticTokenProvider(os.environ["ASTRA_DB_APPLICATION_TOKEN"])
+    token = os.environ["ASTRA_DB_APPLICATION_TOKEN"]
     keyspace = os.environ.get("ASTRA_DB_KEYSPACE", "default_keyspace")
     api_endpoint = os.environ["ASTRA_DB_API_ENDPOINT"]
 
-    admin = AstraDBDatabaseAdmin(api_endpoint=api_endpoint, token=token)
+    my_client = DataAPIClient(token=token)
+    admin = my_client.get_admin().get_database_admin(api_endpoint)
     admin.create_keyspace(keyspace)
 
     # Sometimes the creation of the store fails because the keyspace isn't
@@ -281,17 +283,17 @@ def adapter(
         animal_docs: list[Document],
         astra_config: _AstraConfig,
     ) -> Iterator["AstraAdapter"]:
-        from astrapy.info import CollectionVectorServiceOptions
+        from astrapy.info import VectorServiceOptions
         from langchain_astradb import AstraDBVectorStore
 
-        service = CollectionVectorServiceOptions(
+        service_options = VectorServiceOptions(
             provider="nvidia",
             model_name="NV-Embed-QA",
         )
 
         store = AstraDBVectorStore(
             collection_name="animals_vectorize",
-            collection_vector_service_options=service,
+            collection_vector_service_options=service_options,
             namespace=astra_config.keyspace,
             token=astra_config.token,
             api_endpoint=astra_config.api_endpoint,

scripts/drop-astra-keyspace.py

@@ -1,18 +1,18 @@
 import os
 import sys
 
-from astrapy import AstraDBDatabaseAdmin
-from astrapy.authentication import StaticTokenProvider
+from astrapy import DataAPIClient
 
-token = StaticTokenProvider(os.environ["ASTRA_DB_APPLICATION_TOKEN"])
+token = os.environ["ASTRA_DB_APPLICATION_TOKEN"]
 keyspace = os.environ.get("ASTRA_DB_KEYSPACE")
 if keyspace is None:
     print("No keyspace to drop.")  # noqa: T201
     sys.exit()
 
 api_endpoint = os.environ["ASTRA_DB_API_ENDPOINT"]
 
-admin = AstraDBDatabaseAdmin(api_endpoint=api_endpoint, token=token)
+my_client = DataAPIClient(token=token)
+admin = my_client.get_admin().get_database_admin(api_endpoint)
 keyspaces = admin.list_keyspaces()
 if keyspace in keyspaces:
     print(f"Dropping keyspace '{keyspace}'")  # noqa: T201