cmem-plugin-pgvector 0.6.3__tar.gz → 0.7.1__tar.gz

{cmem_plugin_pgvector-0.6.3 → cmem_plugin_pgvector-0.7.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: cmem-plugin-pgvector
-Version: 0.6.3
+Version: 0.7.1
 Summary: Store and search for embedding vectors in a Postgres vector store.
 License: Apache-2.0
 Keywords: eccenca Corporate Memory,plugin

{cmem_plugin_pgvector-0.6.3 → cmem_plugin_pgvector-0.7.1}/cmem_plugin_pgvector/commons.py

@@ -60,23 +60,15 @@ class PGVectorCollection(StringParameterType):
         password = depend_on_parameter_values[4]
         password = password if isinstance(password, str) else password.decrypt()
         result = []
-        try:
-            collections = get_collection_names(
-                host=host, port=port, dbname=dbname, user=user, password=password
-            )
-            filtered_models = set()
-            if query_terms:
-                for term in query_terms:
-                    for collection in collections:
-                        if term in collection:
-                            filtered_models.add(collection)
-            else:
-                filtered_models = set(collections)
-            result = [Autocompletion(value=f"{_}", label=f"{_}") for _ in filtered_models]
-        except Exception as error:
-            raise ValueError(
-                "Failed to authenticate with OpenAI API, Please check URL and API key."
-            ) from error
+        collections = get_collection_names(
+            host=host, port=port, dbname=dbname, user=user, password=password
+        )
+        filtered_collections = set()
+        for term in query_terms:
+            for collection in collections:
+                if term in collection:
+                    filtered_collections.add(collection)
+        result = [Autocompletion(value=f"{_}", label=f"{_}") for _ in filtered_collections]
         result.sort(key=lambda x: x.label)
         return result
 

{cmem_plugin_pgvector-0.6.3 → cmem_plugin_pgvector-0.7.1}/cmem_plugin_pgvector/search_task.py

@@ -13,10 +13,11 @@ from cmem_plugin_base.dataintegration.ports import (
     FixedNumberOfInputs,
     FixedSchemaPort,
 )
-from cmem_plugin_base.dataintegration.types import IntParameterType
+from cmem_plugin_base.dataintegration.types import EnumParameterType, IntParameterType
 from langchain_core.documents import Document
 from langchain_core.embeddings import Embeddings
 from langchain_postgres import PGVector
+from langchain_postgres.vectorstores import DistanceStrategy
 
 from cmem_plugin_pgvector.commons import DatabaseParams
 
@@ -86,6 +87,14 @@ The results in this output are structured like this:
             default_value=10,
             param_type=IntParameterType(),
         ),
+        PluginParameter(
+            name="distance_strategy",
+            label="Distance Strategy",
+            description="The distance strategy to use. (default: COSINE)",
+            param_type=EnumParameterType(enum_type=DistanceStrategy),
+            default_value=DistanceStrategy.COSINE,
+            advanced=True,
+        ),
     ],
 )
 class PGVectorSearchPlugin(WorkflowPlugin):
@@ -105,6 +114,7 @@ class PGVectorSearchPlugin(WorkflowPlugin):
     report: ExecutionReport
     search_result_path: str
     top_k: int
+    distance_strategy: DistanceStrategy
 
     def __init__(  # noqa: PLR0913
         self,
@@ -117,6 +127,7 @@ class PGVectorSearchPlugin(WorkflowPlugin):
         search_result_path: str = "_search_result",
         embedding_query_path: str = "_embedding",
         top_k: int = 10,
+        distance_strategy: DistanceStrategy = DistanceStrategy.COSINE,
     ) -> None:
         self.collection_name = collection_name
         self.user = user
@@ -126,6 +137,7 @@ class PGVectorSearchPlugin(WorkflowPlugin):
         self.embedding_query_path = embedding_query_path
         self.search_result_path = search_result_path
         self.top_k = top_k
+        self.distance_strategy = distance_strategy
 
         str_password = self.password = password if isinstance(password, str) else password.decrypt()
         self.connection_string = (
@@ -135,14 +147,6 @@ class PGVectorSearchPlugin(WorkflowPlugin):
         self.report = ExecutionReport()
         self.report.operation = "search"
         self.report.operation_desc = "searches"
-
-        self.db = PGVector(
-            collection_name=self.collection_name,
-            connection=self.connection_string,
-            embeddings=DummyEmbeddings(),
-            use_jsonb=True,
-            pre_delete_collection=False,
-        )
         self._setup_ports()
 
     def _setup_ports(self) -> None:
@@ -224,6 +228,14 @@ class PGVectorSearchPlugin(WorkflowPlugin):
     ) -> Entities:
         """Run the workflow operator."""
         self.log.info("Start searching collection.")
+        self.db = PGVector(
+            collection_name=self.collection_name,
+            connection=self.connection_string,
+            embeddings=DummyEmbeddings(),
+            use_jsonb=True,
+            pre_delete_collection=False,
+            distance_strategy=self.distance_strategy,
+        )
         self.inputs = inputs
         self.execution_context = context
         try:

{cmem_plugin_pgvector-0.6.3 → cmem_plugin_pgvector-0.7.1}/cmem_plugin_pgvector/store_task.py

@@ -137,6 +137,7 @@ class PGVectorStorePlugin(WorkflowPlugin):
         self.embedding_path = embedding_path
         self.metadata_paths = metadata_paths
         self.source_path = source_path
+        self.pre_delete_collection = pre_delete_collection
 
         self.output_port = None
         self.input_ports = FixedNumberOfInputs([UnknownSchemaPort()])
@@ -149,14 +150,6 @@ class PGVectorStorePlugin(WorkflowPlugin):
         self.report.operation = "store"
         self.report.operation_desc = "vectors stored"
 
-        self.db = PGVector(
-            collection_name=self.collection_name,
-            connection=self.connection_string,
-            embeddings=None,  # type: ignore  # noqa: PGH003
-            use_jsonb=True,
-            pre_delete_collection=pre_delete_collection,
-        )
-
     def _update_report(self, count: int) -> None:
         self.report.entity_count = count
         self.execution_context.report.update(self.report)
@@ -228,6 +221,13 @@ class PGVectorStorePlugin(WorkflowPlugin):
     ) -> None:
         """Run the workflow operator."""
         self.log.info("Start storing vectors.")
+        self.db = PGVector(
+            collection_name=self.collection_name,
+            connection=self.connection_string,
+            embeddings=None,  # type: ignore  # noqa: PGH003
+            use_jsonb=True,
+            pre_delete_collection=self.pre_delete_collection,
+        )
         self.inputs = inputs
         self.execution_context = context
         try:

{cmem_plugin_pgvector-0.6.3 → cmem_plugin_pgvector-0.7.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "cmem-plugin-pgvector"
-version = "0.6.3"
+version = "0.7.1"
 license = "Apache-2.0"
 description = "Store and search for embedding vectors in a Postgres vector store."
 authors = ["eccenca GmbH <cmempy-developer@eccenca.com>"]
@@ -40,6 +40,7 @@ pytest-memray = { version = "^1.7.0",  markers = "platform_system != 'Windows'"
 ruff = "^0.12.0"
 safety = "^1.10.3"
 aiohttp = "^3.10.11"
+testcontainers = "^4.12.0"
 
 [build-system]
 requires = ["poetry-core>=1.0.0", "poetry-dynamic-versioning"]