lint

Author: epinzur

libs/community/langchain_community/graph_vectorstores/__init__.py

@@ -152,5 +152,5 @@
     "Node",
     "Link",
     "CassandraGraphVectorStore",
-    "MmrHelper"
+    "MmrHelper",
 ]

libs/community/langchain_community/graph_vectorstores/base.py

@@ -632,7 +632,7 @@ async def asearch(
         elif search_type == "traversal":
             return [doc async for doc in self.atraversal_search(query, **kwargs)]
         elif search_type == "mmr_traversal":
-            return list(self.ammr_traversal_search(query, **kwargs))
+            return [doc async for doc in self.ammr_traversal_search(query, **kwargs)]
         else:
             raise ValueError(
                 f"search_type of {search_type} not allowed. Expected "

libs/community/langchain_community/graph_vectorstores/cassandra.py

@@ -447,7 +447,7 @@ def metadata_search(
         return [
             self._restore_links(doc)
             for doc in self.vector_store.metadata_search(
-                filter=filter,
+                filter=filter or {},
                 n=n,
             )
         ]
@@ -466,7 +466,7 @@ async def ametadata_search(
         return [
             self._restore_links(doc)
             for doc in await self.vector_store.ametadata_search(
-                filter=filter,
+                filter=filter or {},
                 n=n,
             )
         ]
@@ -594,7 +594,7 @@ async def fetch_initial_candidates() -> None:
                     retrieved_docs[doc_id] = doc
 
                 if doc_id not in outgoing_links_map:
-                    node = self._doc_to_node(doc)
+                    node = _doc_to_node(doc)
                     outgoing_links_map[doc_id] = _outgoing_links(node=node)
                     candidates[doc_id] = embedding
             helper.add_candidates(candidates)
@@ -952,7 +952,7 @@ async def _get_outgoing_links(self, source_ids: Iterable[str]) -> set[Link]:
 
         for doc in docs:
             if doc is not None:
-                node = self._doc_to_node(doc=doc)
+                node = _doc_to_node(doc=doc)
                 links.update(_outgoing_links(node=node))
 
         return links

libs/community/langchain_community/vectorstores/cassandra.py

@@ -500,100 +500,6 @@ def _row_to_document(row: Dict[str, Any]) -> Document:
             metadata=row["metadata"],
         )
 
-    def get_by_document_id(self, document_id: str) -> Document | None:
-        """Get by document ID.
-
-        Args:
-            document_id: the document ID to get.
-        """
-        row = self.table.get(row_id=document_id)
-        if row is None:
-            return None
-        return self._row_to_document(row=row)
-
-    async def aget_by_document_id(self, document_id: str) -> Document | None:
-        """Get by document ID.
-
-        Args:
-            document_id: the document ID to get.
-        """
-        row = await self.table.aget(row_id=document_id)
-        if row is None:
-            return None
-        return self._row_to_document(row=row)
-
-    def metadata_search(
-        self,
-        metadata: dict[str, Any] = {},  # noqa: B006
-        n: int = 5,
-    ) -> Iterable[Document]:
-        """Get documents via a metadata search.
-
-        Args:
-            metadata: the metadata to query for.
-        """
-        rows = self.table.find_entries(metadata=metadata, n=n)
-        return [self._row_to_document(row=row) for row in rows if row]
-
-    async def ametadata_search(
-        self,
-        metadata: dict[str, Any] = {},  # noqa: B006
-        n: int = 5,
-    ) -> Iterable[Document]:
-        """Get documents via a metadata search.
-
-        Args:
-            metadata: the metadata to query for.
-        """
-        rows = await self.table.afind_entries(metadata=metadata, n=n)
-        return [self._row_to_document(row=row) for row in rows]
-
-    async def asimilarity_search_with_embedding_id_by_vector(
-        self,
-        embedding: List[float],
-        k: int = 4,
-        filter: Optional[Dict[str, str]] = None,
-        body_search: Optional[Union[str, List[str]]] = None,
-    ) -> List[Tuple[Document, List[float], str]]:
-        """Return docs most similar to embedding vector.
-
-        Args:
-            embedding: Embedding to look up documents similar to.
-            k: Number of Documents to return. Defaults to 4.
-            filter: Filter on the metadata to apply.
-            body_search: Document textual search terms to apply.
-                Only supported by Astra DB at the moment.
-        Returns:
-            List of (Document, embedding, id), the most similar to the query vector.
-        """
-        kwargs: Dict[str, Any] = {}
-        if filter is not None:
-            kwargs["metadata"] = filter
-        if body_search is not None:
-            kwargs["body_search"] = body_search
-
-        hits = await self.table.aann_search(
-            vector=embedding,
-            n=k,
-            **kwargs,
-        )
-        return [
-            (
-                self._row_to_document(row=hit),
-                hit["vector"],
-                hit["row_id"],
-            )
-            for hit in hits
-        ]
-
-    @staticmethod
-    def _row_to_document(row: Dict[str, Any]) -> Document:
-        return Document(
-            id=row["row_id"],
-            page_content=row["body_blob"],
-            metadata=row["metadata"],
-        )
-
     def get_by_document_id(self, document_id: str) -> Document | None:
         """Retrieve a single document from the store, given its document ID.
 
@@ -624,7 +530,7 @@ async def aget_by_document_id(self, document_id: str) -> Document | None:
 
     def metadata_search(
         self,
-        filter: Optional[Dict[str, str]] = None,
+        filter: dict[str, Any] = {},  # noqa: B006
         n: int = 5,
     ) -> Iterable[Document]:
         """Get documents via a metadata search.
@@ -638,7 +544,7 @@ def metadata_search(
 
     async def ametadata_search(
         self,
-        filter: Optional[Dict[str, str]] = None,
+        filter: dict[str, Any] = {},  # noqa: B006
         n: int = 5,
     ) -> Iterable[Document]:
         """Get documents via a metadata search.
@@ -681,11 +587,7 @@ async def asimilarity_search_with_embedding_id_by_vector(
         )
         return [
             (
-                Document(
-                    page_content=hit["body_blob"],
-                    metadata=hit["metadata"],
-                    id=hit["row_id"],
-                ),
+                self._row_to_document(row=hit),
                 hit["vector"],
                 hit["row_id"],
             )