lint

Author: averikitsch

langchain_postgres/v2/async_vectorstore.py

@@ -111,9 +111,7 @@ def __init__(
         self.schema_name = schema_name
         self.content_column = content_column
         self.embedding_column = embedding_column
-        self.metadata_columns = (
-            metadata_columns if metadata_columns is not None else []
-        )
+        self.metadata_columns = metadata_columns if metadata_columns is not None else []
         self.id_column = id_column
         self.metadata_json_column = metadata_json_column
         self.distance_strategy = distance_strategy
@@ -189,27 +187,21 @@ async def create(
         if id_column not in columns:
             raise ValueError(f"Id column, {id_column}, does not exist.")
         if content_column not in columns:
-            raise ValueError(
-                f"Content column, {content_column}, does not exist."
-            )
+            raise ValueError(f"Content column, {content_column}, does not exist.")
         content_type = columns[content_column]
         if content_type != "text" and "char" not in content_type:
             raise ValueError(
                 f"Content column, {content_column}, is type, {content_type}. It must be a type of character string."
             )
         if embedding_column not in columns:
-            raise ValueError(
-                f"Embedding column, {embedding_column}, does not exist."
-            )
+            raise ValueError(f"Embedding column, {embedding_column}, does not exist.")
         if columns[embedding_column] != "USER-DEFINED":
             raise ValueError(
                 f"Embedding column, {embedding_column}, is not type Vector."
             )
 
         metadata_json_column = (
-            None
-            if metadata_json_column not in columns
-            else metadata_json_column
+            None if metadata_json_column not in columns else metadata_json_column
         )
 
         # If using metadata_columns check to make sure column exists
@@ -272,14 +264,10 @@ async def aadd_embeddings(
             metadatas = [{} for _ in texts]
 
         # Check for inline embedding capability
-        inline_embed_func = getattr(
-            self.embedding_service, "embed_query_inline", None
-        )
+        inline_embed_func = getattr(self.embedding_service, "embed_query_inline", None)
         can_inline_embed = callable(inline_embed_func)
         # Insert embeddings
-        for id, content, embedding, metadata in zip(
-            ids, texts, embeddings, metadatas
-        ):
+        for id, content, embedding, metadata in zip(ids, texts, embeddings, metadatas):
             metadata_col_names = (
                 ", " + ", ".join(f'"{col}"' for col in self.metadata_columns)
                 if len(self.metadata_columns) > 0
@@ -348,15 +336,11 @@ async def aadd_texts(
             :class:`InvalidTextRepresentationError <asyncpg.exceptions.InvalidTextRepresentationError>`: if the `ids` data type does not match that of the `id_column`.
         """
         # Check for inline embedding query
-        inline_embed_func = getattr(
-            self.embedding_service, "embed_query_inline", None
-        )
+        inline_embed_func = getattr(self.embedding_service, "embed_query_inline", None)
         if callable(inline_embed_func):
             embeddings: list[list[float]] = [[] for _ in list(texts)]
         else:
-            embeddings = await self.embedding_service.aembed_documents(
-                list(texts)
-            )
+            embeddings = await self.embedding_service.aembed_documents(list(texts))
 
         ids = await self.aadd_embeddings(
             texts, embeddings, metadatas=metadatas, ids=ids, **kwargs
@@ -378,9 +362,7 @@ async def aadd_documents(
         metadatas = [doc.metadata for doc in documents]
         if not ids:
             ids = [doc.id for doc in documents]
-        ids = await self.aadd_texts(
-            texts, metadatas=metadatas, ids=ids, **kwargs
-        )
+        ids = await self.aadd_texts(texts, metadatas=metadatas, ids=ids, **kwargs)
         return ids
 
     async def adelete(
@@ -576,9 +558,7 @@ async def __query_collection(
         if filter and isinstance(filter, dict):
             safe_filter, filter_dict = self._create_filter_clause(filter)
         filter = f"WHERE {safe_filter}" if safe_filter else ""
-        inline_embed_func = getattr(
-            self.embedding_service, "embed_query_inline", None
-        )
+        inline_embed_func = getattr(self.embedding_service, "embed_query_inline", None)
         if not embedding and callable(inline_embed_func) and "query" in kwargs:
             query_embedding = self.embedding_service.embed_query_inline(kwargs["query"])  # type: ignore
         else:
@@ -613,9 +593,7 @@ async def asimilarity_search(
         **kwargs: Any,
     ) -> list[Document]:
         """Return docs selected by similarity search on query."""
-        inline_embed_func = getattr(
-            self.embedding_service, "embed_query_inline", None
-        )
+        inline_embed_func = getattr(self.embedding_service, "embed_query_inline", None)
         embedding = (
             []
             if callable(inline_embed_func)
@@ -646,9 +624,7 @@ async def asimilarity_search_with_score(
         **kwargs: Any,
     ) -> list[tuple[Document, float]]:
         """Return docs and distance scores selected by similarity search on query."""
-        inline_embed_func = getattr(
-            self.embedding_service, "embed_query_inline", None
-        )
+        inline_embed_func = getattr(self.embedding_service, "embed_query_inline", None)
         embedding = (
             []
             if callable(inline_embed_func)
@@ -770,9 +746,7 @@ async def amax_marginal_relevance_search_with_score_by_vector(
         k = k if k else self.k
         fetch_k = fetch_k if fetch_k else self.fetch_k
         lambda_mult = lambda_mult if lambda_mult else self.lambda_mult
-        embedding_list = [
-            json.loads(row[self.embedding_column]) for row in results
-        ]
+        embedding_list = [json.loads(row[self.embedding_column]) for row in results]
         mmr_selected = utils.maximal_marginal_relevance(
             np.array(embedding, dtype=np.float32),
             embedding_list,
@@ -800,9 +774,7 @@ async def amax_marginal_relevance_search_with_score_by_vector(
                 )
             )
 
-        return [
-            r for i, r in enumerate(documents_with_scores) if i in mmr_selected
-        ]
+        return [r for i, r in enumerate(documents_with_scores) if i in mmr_selected]
 
     async def aapply_vector_index(
         self,
@@ -820,16 +792,12 @@ async def aapply_vector_index(
         if index.extension_name:
             async with self.engine.connect() as conn:
                 await conn.execute(
-                    text(
-                        f"CREATE EXTENSION IF NOT EXISTS {index.extension_name}"
-                    )
+                    text(f"CREATE EXTENSION IF NOT EXISTS {index.extension_name}")
                 )
                 await conn.commit()
         function = index.get_index_function()
 
-        filter = (
-            f"WHERE ({index.partial_indexes})" if index.partial_indexes else ""
-        )
+        filter = f"WHERE ({index.partial_indexes})" if index.partial_indexes else ""
         params = "WITH " + index.index_options()
         if name is None:
             if index.name == None:
@@ -993,9 +961,7 @@ def _handle_field_filter(
             #     filter_value = f"'{filter_value}'"
             native = COMPARISONS_TO_NATIVE[operator]
             id = str(uuid.uuid4()).split("-")[0]
-            return f"{field} {native} :{field}_{id}", {
-                f"{field}_{id}": filter_value
-            }
+            return f"{field} {native} :{field}_{id}", {f"{field}_{id}": filter_value}
         elif operator == "$between":
             # Use AND with two comparisons
             low, high = filter_value
@@ -1019,17 +985,11 @@ def _handle_field_filter(
                         )
 
             if operator in {"$in"}:
-                return f"{field} = ANY(:{field}_in)", {
-                    f"{field}_in": filter_value
-                }
+                return f"{field} = ANY(:{field}_in)", {f"{field}_in": filter_value}
             elif operator in {"$nin"}:
-                return f"{field} <> ALL (:{field}_nin)", {
-                    f"{field}_nin": filter_value
-                }
+                return f"{field} <> ALL (:{field}_nin)", {f"{field}_nin": filter_value}
             elif operator in {"$like"}:
-                return f"({field} LIKE :{field}_like)", {
-                    f"{field}_like": filter_value
-                }
+                return f"({field} LIKE :{field}_like)", {f"{field}_like": filter_value}
             elif operator in {"$ilike"}:
                 return f"({field} ILIKE :{field}_ilike)", {
                     f"{field}_ilike": filter_value
@@ -1108,9 +1068,7 @@ def _create_filter_clause(self, filters: Any) -> tuple[str, dict]:
                     params = {}
                     for clause in not_conditions:
                         params.update(clause[1])
-                    not_stmts = [
-                        f"NOT {condition}" for condition in all_clauses
-                    ]
+                    not_stmts = [f"NOT {condition}" for condition in all_clauses]
                     return f"({' AND '.join(not_stmts)})", params
                 elif isinstance(value, dict):
                     not_, params = self._create_filter_clause(value)
@@ -1134,8 +1092,7 @@ def _create_filter_clause(self, filters: Any) -> tuple[str, dict]:
                     )
             # These should all be fields and combined using an $and operator
             and_ = [
-                self._handle_field_filter(field=k, value=v)
-                for k, v in filters.items()
+                self._handle_field_filter(field=k, value=v) for k, v in filters.items()
             ]
             if len(and_) > 1:
                 all_clauses = [clause[0] for clause in and_]

tests/unit_tests/v2/test_async_pg_vectorstore_search.py

@@ -33,15 +33,13 @@
 ids = [str(uuid.uuid4()) for i in range(len(texts))]
 metadatas = [{"page": str(i), "source": "postgres"} for i in range(len(texts))]
 docs = [
-    Document(page_content=texts[i], metadata=metadatas[i])
-    for i in range(len(texts))
+    Document(page_content=texts[i], metadata=metadatas[i]) for i in range(len(texts))
 ]
 
 embeddings = [embeddings_service.embed_query("foo") for i in range(len(texts))]
 
 filter_docs = [
-    Document(page_content=texts[i], metadata=METADATAS[i])
-    for i in range(len(texts))
+    Document(page_content=texts[i], metadata=METADATAS[i]) for i in range(len(texts))
 ]
 
 
@@ -87,9 +85,7 @@ async def vs(self, engine: PGEngine) -> AsyncIterator[AsyncPGVectorStore]:
         yield vs
 
     @pytest_asyncio.fixture(scope="class")
-    async def vs_custom(
-        self, engine: PGEngine
-    ) -> AsyncIterator[AsyncPGVectorStore]:
+    async def vs_custom(self, engine: PGEngine) -> AsyncIterator[AsyncPGVectorStore]:
         await engine._ainit_vectorstore_table(
             CUSTOM_TABLE,
             VECTOR_SIZE,
@@ -153,24 +149,18 @@ async def vs_custom_filter(
         await vs_custom_filter.aadd_documents(filter_docs, ids=ids)
         yield vs_custom_filter
 
-    async def test_asimilarity_search_score(
-        self, vs: AsyncPGVectorStore
-    ) -> None:
+    async def test_asimilarity_search_score(self, vs: AsyncPGVectorStore) -> None:
         results = await vs.asimilarity_search_with_score("foo")
         assert len(results) == 4
         assert results[0][0] == Document(page_content="foo", id=ids[0])
         assert results[0][1] == 0
 
-    async def test_asimilarity_search_by_vector(
-        self, vs: AsyncPGVectorStore
-    ) -> None:
+    async def test_asimilarity_search_by_vector(self, vs: AsyncPGVectorStore) -> None:
         embedding = embeddings_service.embed_query("foo")
         results = await vs.asimilarity_search_by_vector(embedding)
         assert len(results) == 4
         assert results[0] == Document(page_content="foo", id=ids[0])
-        result = await vs.asimilarity_search_with_score_by_vector(
-            embedding=embedding
-        )
+        result = await vs.asimilarity_search_with_score_by_vector(embedding=embedding)
         assert result[0][0] == Document(page_content="foo", id=ids[0])
         assert result[0][1] == 0
 
@@ -244,9 +234,7 @@ async def test_amax_marginal_relevance_search_vector_score(
         )
         assert results[0][0] == Document(page_content="bar", id=ids[1])
 
-    async def test_similarity_search_score(
-        self, vs_custom: AsyncPGVectorStore
-    ) -> None:
+    async def test_similarity_search_score(self, vs_custom: AsyncPGVectorStore) -> None:
         results = await vs_custom.asimilarity_search_with_score("foo")
         assert len(results) == 4
         assert results[0][0] == Document(page_content="foo", id=ids[0])
@@ -269,26 +257,20 @@ async def test_max_marginal_relevance_search_vector(
         self, vs_custom: AsyncPGVectorStore
     ) -> None:
         embedding = embeddings_service.embed_query("bar")
-        results = await vs_custom.amax_marginal_relevance_search_by_vector(
-            embedding
-        )
+        results = await vs_custom.amax_marginal_relevance_search_by_vector(embedding)
         assert results[0] == Document(page_content="bar", id=ids[1])
 
     async def test_max_marginal_relevance_search_vector_score(
         self, vs_custom: AsyncPGVectorStore
     ) -> None:
         embedding = embeddings_service.embed_query("bar")
-        results = (
-            await vs_custom.amax_marginal_relevance_search_with_score_by_vector(
-                embedding
-            )
+        results = await vs_custom.amax_marginal_relevance_search_with_score_by_vector(
+            embedding
         )
         assert results[0][0] == Document(page_content="bar", id=ids[1])
 
-        results = (
-            await vs_custom.amax_marginal_relevance_search_with_score_by_vector(
-                embedding, lambda_mult=0.75, fetch_k=10
-            )
+        results = await vs_custom.amax_marginal_relevance_search_with_score_by_vector(
+            embedding, lambda_mult=0.75, fetch_k=10
         )
         assert results[0][0] == Document(page_content="bar", id=ids[1])
 
@@ -298,9 +280,7 @@ async def test_aget_by_ids(self, vs: AsyncPGVectorStore) -> None:
 
         assert results[0] == Document(page_content="foo", id=ids[0])
 
-    async def test_aget_by_ids_custom_vs(
-        self, vs_custom: AsyncPGVectorStore
-    ) -> None:
+    async def test_aget_by_ids_custom_vs(self, vs_custom: AsyncPGVectorStore) -> None:
         test_ids = [ids[0]]
         results = await vs_custom.aget_by_ids(ids=test_ids)
 
@@ -322,6 +302,4 @@ async def test_vectorstore_with_metadata_filters(
         docs = await vs_custom_filter.asimilarity_search(
             "meow", k=5, filter=test_filter
         )
-        assert [
-            doc.metadata["code"] for doc in docs
-        ] == expected_ids, test_filter
+        assert [doc.metadata["code"] for doc in docs] == expected_ids, test_filter