cmem-plugin-pgvector 0.7.0__tar.gz → 0.8.0__tar.gz

{cmem_plugin_pgvector-0.7.0 → cmem_plugin_pgvector-0.8.0}/PKG-INFO

@@ -1,8 +1,9 @@
-Metadata-Version: 2.3
+Metadata-Version: 2.4
 Name: cmem-plugin-pgvector
-Version: 0.7.0
+Version: 0.8.0
 Summary: Store and search for embedding vectors in a Postgres vector store.
 License: Apache-2.0
+License-File: LICENSE
 Keywords: eccenca Corporate Memory,plugin
 Author: eccenca GmbH
 Author-email: cmempy-developer@eccenca.com
@@ -16,6 +17,7 @@ Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3.13
+Classifier: Programming Language :: Python :: 3.14
 Requires-Dist: cmem-plugin-base (>=4.12.1,<5.0.0)
 Requires-Dist: langchain-core (>=0.3.71,<0.4.0)
 Requires-Dist: langchain-postgres (>=0.0.15,<0.0.16)

{cmem_plugin_pgvector-0.7.0 → cmem_plugin_pgvector-0.8.0}/cmem_plugin_pgvector/commons.py

@@ -15,6 +15,37 @@ from cmem_plugin_base.dataintegration.types import (
 )
 
 
+class DatabaseConnectionError(Exception):
+    """Custom exception for database connection issues"""
+
+
+def check_database_connection(dbname: str, user: str, password: str, host: str, port: int) -> str:
+    """Test database connection and return success message or raise exception on failure"""
+    try:
+        with (
+            psycopg.connect(
+                dbname=dbname,
+                user=user,
+                password=password,
+                host=host,
+                port=port,
+                connect_timeout=10,
+            ) as conn,
+            conn.cursor() as cursor,
+        ):
+            # Test basic connectivity
+            cursor.execute("SELECT version();")
+            version = cursor.fetchone()[0]  # type: ignore[index]
+            return f"Connection successful. PostgreSQL version: {version[:50]}..."
+
+    except psycopg.OperationalError as e:
+        raise DatabaseConnectionError(f"Connection failed: {e!s}") from e
+    except psycopg.Error as e:
+        raise DatabaseConnectionError(f"Database error: {e!s}") from e
+    except Exception as e:
+        raise DatabaseConnectionError(f"Unexpected error: {e!s}") from e
+
+
 def get_collection_names(
     dbname: str, user: str, password: str, host: str = "localhost", port: int = 5432
 ) -> list[str]:
@@ -59,24 +90,19 @@ class PGVectorCollection(StringParameterType):
         user = depend_on_parameter_values[3]
         password = depend_on_parameter_values[4]
         password = password if isinstance(password, str) else password.decrypt()
-        result = []
+        result: list[Autocompletion] = []
         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
+        except psycopg.Error:
+            return result
+        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.7.0 → cmem_plugin_pgvector-0.8.0}/cmem_plugin_pgvector/search_task.py

@@ -5,7 +5,7 @@ from ast import literal_eval
 from collections.abc import Generator, Sequence
 
 from cmem_plugin_base.dataintegration.context import ExecutionContext, ExecutionReport
-from cmem_plugin_base.dataintegration.description import Icon, Plugin, PluginParameter
+from cmem_plugin_base.dataintegration.description import Icon, Plugin, PluginAction, PluginParameter
 from cmem_plugin_base.dataintegration.entity import Entities, Entity, EntityPath, EntitySchema
 from cmem_plugin_base.dataintegration.parameter.password import Password
 from cmem_plugin_base.dataintegration.plugins import WorkflowPlugin
@@ -19,7 +19,11 @@ from langchain_core.embeddings import Embeddings
 from langchain_postgres import PGVector
 from langchain_postgres.vectorstores import DistanceStrategy
 
-from cmem_plugin_pgvector.commons import DatabaseParams
+from cmem_plugin_pgvector.commons import (
+    DatabaseConnectionError,
+    DatabaseParams,
+    check_database_connection,
+)
 
 
 class DummyEmbeddings(Embeddings):
@@ -66,6 +70,13 @@ The results in this output are structured like this:
 """,
     icon=Icon(package=__package__, file_name="postgresql.svg"),
     plugin_id="cmem_plugin_pgvector-Search",
+    actions=[
+        PluginAction(
+            name="test_connection",
+            label="Test Connection",
+            description="Test database connectivity",
+        )
+    ],
     parameters=[
         *DatabaseParams().as_list(),
         PluginParameter(
@@ -147,17 +158,21 @@ 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,
-            distance_strategy=distance_strategy,
-        )
         self._setup_ports()
 
+    def test_connection(self) -> str:
+        """Plugin Action to test database connection"""
+        try:
+            return check_database_connection(
+                dbname=self.database,
+                user=self.user,
+                password=self.password,
+                host=self.host,
+                port=self.port,
+            )
+        except DatabaseConnectionError as e:
+            raise ValueError(f"Connection test failed: {e!s}") from e
+
     def _setup_ports(self) -> None:
         """Configure input and output ports depending on the configuration"""
         input_paths = [EntityPath(path=self.embedding_query_path)]
@@ -203,7 +218,7 @@ class PGVectorSearchPlugin(WorkflowPlugin):
         for doc_tuple in docs:
             json_entity = {
                 "id": doc_tuple[0].id,
-                "metadata": str(doc_tuple[0].metadata),
+                "metadata": doc_tuple[0].metadata,
                 "_embedding_source": doc_tuple[0].page_content,
                 "distance": str(doc_tuple[1]),
             }
@@ -237,6 +252,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.7.0 → cmem_plugin_pgvector-0.8.0}/cmem_plugin_pgvector/store_task.py

@@ -5,7 +5,7 @@ from collections.abc import Sequence
 from typing import Any
 
 from cmem_plugin_base.dataintegration.context import ExecutionContext, ExecutionReport
-from cmem_plugin_base.dataintegration.description import Icon, Plugin, PluginParameter
+from cmem_plugin_base.dataintegration.description import Icon, Plugin, PluginAction, PluginParameter
 from cmem_plugin_base.dataintegration.entity import Entities, Entity, EntityPath
 from cmem_plugin_base.dataintegration.parameter.password import Password
 from cmem_plugin_base.dataintegration.plugins import WorkflowPlugin
@@ -15,7 +15,11 @@ from cmem_plugin_base.dataintegration.ports import (
 )
 from langchain_postgres import PGVector
 
-from cmem_plugin_pgvector.commons import DatabaseParams
+from cmem_plugin_pgvector.commons import (
+    DatabaseConnectionError,
+    DatabaseParams,
+    check_database_connection,
+)
 
 
 class DataContainer:
@@ -56,6 +60,13 @@ metadata.
 """,
     icon=Icon(package=__package__, file_name="postgresql.svg"),
     plugin_id="cmem_plugin_pgvector-Store",
+    actions=[
+        PluginAction(
+            name="test_connection",
+            label="Test Connection",
+            description="Test database connectivity",
+        )
+    ],
     parameters=[
         *DatabaseParams().as_list(),
         PluginParameter(
@@ -137,6 +148,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,13 +161,18 @@ 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 test_connection(self) -> str:
+        """Plugin Action to test database connection"""
+        try:
+            return check_database_connection(
+                dbname=self.database,
+                user=self.user,
+                password=self.password,
+                host=self.host,
+                port=self.port,
+            )
+        except DatabaseConnectionError as e:
+            raise ValueError(f"Connection test failed: {e!s}") from e
 
     def _update_report(self, count: int) -> None:
         self.report.entity_count = count
@@ -228,6 +245,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.7.0 → cmem_plugin_pgvector-0.8.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "cmem-plugin-pgvector"
-version = "0.7.0"
+version = "0.8.0"
 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"]