Fix some ruff preview rules (#153)

Author: cbornet

libs/astradb/langchain_astradb/cache.py

@@ -503,7 +503,12 @@ def lookup_with_id(
     ) -> tuple[str, RETURN_VAL_TYPE] | None:
         """Look up based on prompt and llm_string.
 
-        If there are hits, return (document_id, cached_entry) for the top hit
+        Args:
+            prompt: the prompt string to look up
+            llm_string: the str representation of the model parameters
+
+        Returns:
+            If there are hits, (document_id, cached_entry) for the top hit
         """
         self.astra_env.ensure_db_setup()
         prompt_embedding: list[float] = self._get_embedding(text=prompt)
@@ -533,7 +538,12 @@ async def alookup_with_id(
     ) -> tuple[str, RETURN_VAL_TYPE] | None:
         """Look up based on prompt and llm_string.
 
-        If there are hits, return (document_id, cached_entry) for the top hit
+        Args:
+            prompt: the prompt string to look up
+            llm_string: the str representation of the model parameters
+
+        Returns:
+            If there are hits, (document_id, cached_entry) for the top hit
         """
         await self.astra_env.aensure_db_setup()
         prompt_embedding: list[float] = await self._aget_embedding(text=prompt)
@@ -563,7 +573,13 @@ def lookup_with_id_through_llm(
     ) -> tuple[str, RETURN_VAL_TYPE] | None:
         """Look up based on prompt and LLM.
 
-        If there are hits, return (document_id, cached_entry) for the top hit
+        Args:
+            prompt: the prompt string to look up
+            llm: the LLM instance whose parameters are used in the lookup
+            stop: optional list of stop words passed to the LLM calls
+
+        Returns:
+            If there are hits, (document_id, cached_entry) for the top hit.
         """
         llm_string = get_prompts(
             {**llm.dict(), "stop": stop},
@@ -576,7 +592,13 @@ async def alookup_with_id_through_llm(
     ) -> tuple[str, RETURN_VAL_TYPE] | None:
         """Look up based on prompt and LLM.
 
-        If there are hits, return (document_id, cached_entry) for the top hit
+        Args:
+            prompt: the prompt string to look up
+            llm: the LLM instance whose parameters are used in the lookup
+            stop: optional list of stop words passed to the LLM calls
+
+        Returns:
+            If there are hits, (document_id, cached_entry) for the top hit.
         """
         llm_string = (
             await aget_prompts(

libs/astradb/langchain_astradb/document_loaders.py

@@ -145,8 +145,8 @@ def lazy_load(self) -> Iterator[Document]:
         ):
             yield self._to_langchain_doc(doc)
 
+    @override
     async def aload(self) -> list[Document]:
-        """Load data into Document objects."""
         return [doc async for doc in self.alazy_load()]
 
     @override

libs/astradb/langchain_astradb/storage.py

@@ -38,7 +38,12 @@
 
 class AstraDBBaseStore(BaseStore[str, V], Generic[V]):
     def __init__(self, *args: Any, **kwargs: Any) -> None:
-        """Base class for the DataStax Astra DB data store."""
+        """Base class for the DataStax Astra DB data store.
+
+        Raises:
+            ValueError: if 'requested_indexing_policy' or
+                'default_indexing_policy' are passed in kwargs.
+        """
         if "requested_indexing_policy" in kwargs:
             msg = "Do not pass 'requested_indexing_policy' to AstraDBBaseStore init"
             raise ValueError(msg)

libs/astradb/langchain_astradb/utils/astradb.py

@@ -505,6 +505,10 @@ def copy(
                 feature and a required secret is not stored with the database.
                 In order to suppress the API Key in the copy, explicitly pass
                 ``astrapy.authentication.RerankingAPIKeyHeaderProvider(None)``.
+
+        Returns:
+            A shallow copy of this environment, with the possibly updated
+            attributes.
         """
         return _AstraDBCollectionEnvironment(
             collection_name=self.collection_name,

libs/astradb/langchain_astradb/utils/vector_store_codecs.py

@@ -216,7 +216,14 @@ def encode_filter(self, filter_dict: dict[str, Any]) -> dict[str, Any]:
 
     @abstractmethod
     def metadata_key_to_field_identifier(self, md_key: str) -> str:
-        """Express an 'abstract' metadata key as a full Data API field identifier."""
+        """Express an 'abstract' metadata key as a full Data API field identifier.
+
+        Args:
+            md_key: the metadata key.
+
+        Returns:
+            the full Data API field identifier.
+        """
 
     @abstractmethod
     def encode_vectorize_sort(self, query_text: str) -> dict[str, Any]:

libs/astradb/langchain_astradb/vectorstores.py

@@ -245,9 +245,9 @@ def _decide_hybrid_search_setting(
 
 
 def _make_hybrid_limits(
-    hlf: None | float | dict[str, float],
+    hlf: float | dict[str, float] | None,
     k: int,
-) -> None | int | dict[str, int]:
+) -> int | dict[str, int] | None:
     if hlf is None:
         return None
     if isinstance(hlf, float):
@@ -258,9 +258,9 @@ def _make_hybrid_limits(
 
 def _normalize_hybrid_limit_factor(
     hybrid_limit_factor: float
-    | None
     | dict[str, float]
-    | HybridLimitFactorPrescription,
+    | HybridLimitFactorPrescription
+    | None,
 ) -> float | dict[str, float] | None:
     """Bring `hybrid_limit_factor` to a normal form."""
     if hybrid_limit_factor is None:
@@ -740,9 +740,9 @@ def __init__(
         | None = None,
         hybrid_search: HybridSearchMode | None = None,
         hybrid_limit_factor: float
-        | None
         | dict[str, float]
-        | HybridLimitFactorPrescription = None,
+        | HybridLimitFactorPrescription
+        | None = None,
     ) -> None:
         """A vector store which uses DataStax Astra DB as backend.
 
@@ -891,6 +891,9 @@ def __init__(
                 and the lexical subsearches. Alternatively, a simple dictionary
                 with keys "$lexical" and "$vector" achieves the same effect.
 
+        Raises:
+            ValueError: if the parameters are inconsistent or invalid.
+
         Note:
             For concurrency in synchronous :meth:`~add_texts`:, as a rule of thumb,
             on a typical client machine it is suggested to keep the quantity
@@ -941,7 +944,7 @@ def __init__(
         self.has_lexical: bool
         self.has_hybrid: bool
         self.hybrid_search: bool  # affecting the actual behaviour when running searches
-        self.hybrid_limit_factor: None | float | dict[str, float]
+        self.hybrid_limit_factor: float | dict[str, float] | None
         self.collection_reranking_api_key = collection_reranking_api_key
 
         if not self.autodetect_collection:
@@ -1178,7 +1181,7 @@ def copy(
         In those cases, one should create a new instance of ``AstraDBVectorStore``
         from scratch.
 
-        Attributes:
+        Args:
             token: API token for Astra DB usage, either in the form of a string
                 or a subclass of ``astrapy.authentication.TokenProvider``.
                 In order to suppress token usage in the copy, explicitly pass
@@ -1200,6 +1203,10 @@ def copy(
                 This is useful when the service is configured for the collection,
                 but no corresponding secret is stored within
                 Astra's key management system.
+
+        Returns:
+            a shallow copy of this vector store, possibly with some changed
+            attributes.
         """
         copy = AstraDBVectorStore(
             collection_name="moot",
@@ -1310,6 +1317,9 @@ def delete(
 
         Returns:
             True if deletion is (entirely) successful, False otherwise.
+
+        Raises:
+            ValueError: if no ids are provided.
         """
         if kwargs:
             warnings.warn(
@@ -1350,6 +1360,9 @@ async def adelete(
 
         Returns:
             True if deletion is (entirely) successful, False otherwise.
+
+        Raises:
+            ValueError: if no ids are provided.
         """
         if kwargs:
             warnings.warn(
@@ -1383,6 +1396,9 @@ def delete_by_metadata_filter(
 
         Returns:
             A number expressing the amount of deleted documents.
+
+        Raises:
+            ValueError: if the provided filter is empty.
         """
         if not filter:
             msg = (
@@ -1412,6 +1428,9 @@ async def adelete_by_metadata_filter(
 
         Returns:
             A number expressing the amount of deleted documents.
+
+        Raises:
+            ValueError: if the provided filter is empty.
         """
         if not filter:
             msg = (
@@ -1524,6 +1543,9 @@ def add_texts(
 
         Returns:
             The list of ids of the added texts.
+
+        Raises:
+            AstraDBVectorStoreError: if not all documents could be inserted.
         """
         if kwargs:
             warnings.warn(
@@ -1665,6 +1687,9 @@ async def aadd_texts(
 
         Returns:
             The list of ids of the added texts.
+
+        Raises:
+            AstraDBVectorStoreError: if not all documents could be inserted.
         """
         if kwargs:
             warnings.warn(
@@ -2473,6 +2498,9 @@ def metadata_search(
         Args:
             filter: the metadata to query for.
             n: the maximum number of documents to return.
+
+        Returns:
+            The documents found.
         """
         docs_ite = self.run_query(n=n, filter=filter)
         return [doc for doc, _, _, _ in docs_ite]
@@ -2487,6 +2515,9 @@ async def ametadata_search(
         Args:
             filter: the metadata to query for.
             n: the maximum number of documents to return.
+
+        Returns:
+            The documents found.
         """
         docs_ite = await self.arun_query(n=n, filter=filter)
         return [doc async for doc, _, _, _ in docs_ite]
@@ -2933,6 +2964,9 @@ def similarity_search_with_score_id_by_vector(
 
         Returns:
             The list of (Document, score, id), the most similar to the query vector.
+
+        Raises:
+            ValueError: if the vector store uses server-side embeddings.
         """
         if self.document_codec.server_side_embeddings:
             msg = (
@@ -3232,6 +3266,9 @@ async def asimilarity_search_with_score_id_by_vector(
 
         Returns:
             The list of (Document, score, id), the most similar to the query vector.
+
+        Raises:
+            ValueError: If the vector store uses server-side embeddings.
         """
         if self.document_codec.server_side_embeddings:
             msg = (
@@ -3559,8 +3596,8 @@ async def _arun_mmr_find_by_sort(
             prefetch_hit_pairs=prefetch_hit_pairs,
         )
 
+    @staticmethod
     def _get_mmr_hits(
-        self,
         embedding: list[float],
         k: int,
         lambda_mult: float,

libs/astradb/tests/integration_tests/test_chat_message_histories.py

@@ -55,7 +55,7 @@ def history2(
 
 
 @pytest.fixture
-async def async_history1(
+def async_history1(
     astra_db_credentials: AstraDBCredentials,
     history1: AstraDBChatMessageHistory,
 ) -> AstraDBChatMessageHistory:
@@ -71,7 +71,7 @@ async def async_history1(
 
 
 @pytest.fixture
-async def async_history2(
+def async_history2(
     astra_db_credentials: AstraDBCredentials,
     history1: AstraDBChatMessageHistory,
 ) -> AstraDBChatMessageHistory:

libs/astradb/tests/integration_tests/test_document_loaders.py

@@ -32,7 +32,7 @@ def document_loader_collection(
 
 
 @pytest.fixture
-async def async_document_loader_collection(
+def async_document_loader_collection(
     collection_idxall: Collection,
 ) -> AsyncCollection:
     return collection_idxall.to_async()
@@ -257,7 +257,7 @@ async def test_page_content_mapper_async(
             page_content_mapper=itemgetter("foo"),
             filter_criteria={"foo": "bar"},
         )
-        doc = await loader.alazy_load().__anext__()
+        doc = await anext(loader.alazy_load())
         assert doc.page_content == "bar"
 
     async def test_metadata_mapper_async(
@@ -275,5 +275,5 @@ async def test_metadata_mapper_async(
             metadata_mapper=lambda x: {"a": x["foo"]},
             filter_criteria={"foo": "bar"},
         )
-        doc = await loader.alazy_load().__anext__()
+        doc = await anext(loader.alazy_load())
         assert doc.metadata == {"a": "bar"}

libs/astradb/tests/integration_tests/test_vectorstore_hybrid.py

@@ -235,7 +235,7 @@ def test_astradb_vectorstore_explicit_hybrid_lifecycle_vectorize_sync(
             assert rscore < 100
             assert isinstance(rid, str)
             # run other search methods
-            with pytest.raises(ValueError, match="not allowed."):
+            with pytest.raises(ValueError, match="not allowed"):
                 store2_ad.similarity_search_by_vector([1, 2, 3])
             mmr_hits_docs = store2_ad.max_marginal_relevance_search(
                 query=QUERY_TEXT,
@@ -349,7 +349,7 @@ def test_astradb_vectorstore_explicit_hybrid_lifecycle_vectorize_sync(
             assert rscore < 100
             assert isinstance(rid, str)
             # run other search methods
-            with pytest.raises(ValueError, match="not allowed."):
+            with pytest.raises(ValueError, match="not allowed"):
                 store4.similarity_search_by_vector([1, 2, 3])
             mmr_hits_docs = store4.max_marginal_relevance_search(
                 query=QUERY_TEXT,
@@ -615,7 +615,7 @@ async def test_astradb_vectorstore_explicit_hybrid_lifecycle_vectorize_async(
             assert rscore < 100
             assert isinstance(rid, str)
             # run other search methods
-            with pytest.raises(ValueError, match="not allowed."):
+            with pytest.raises(ValueError, match="not allowed"):
                 await store4.asimilarity_search_by_vector([1, 2, 3])
             mmr_hits_docs = await store4.amax_marginal_relevance_search(
                 query=QUERY_TEXT,

libs/astradb/tests/unit_tests/test_astra_db_environment.py

@@ -59,7 +59,7 @@ def test_initialization(self) -> None:
 
             # just tokenn, no endpoint
             with pytest.raises(
-                ValueError, match="API endpoint for Data API not provided."
+                ValueError, match="API endpoint for Data API not provided"
             ):
                 _AstraDBEnvironment(
                     token=FAKE_TOKEN,