pgai knowledge base vector store driver init (#1889)

* pgai knowledge base vector store driver init

* cleanup

* more format

* fix docs for now

* fix unit test

* remove debug log

* lint

* pgai extra

* address comments

* format

* camelcase

Author: cjkindel

docs/griptape-framework/drivers/src/vector_store_drivers_12.py

@@ -0,0 +1,18 @@
+import os
+
+from griptape.drivers.vector.pgai import PgAiKnowledgeBaseVectorStoreDriver
+
+# PG.AI connection parameters
+connection_string = os.environ["PGAI_CONNECTION_STRING"]
+knowledge_base_name = os.environ["PGAI_KNOWLEDGE_BASE_NAME"]
+
+vector_store_driver = PgAiKnowledgeBaseVectorStoreDriver(
+    connection_string=connection_string,
+    knowledge_base_name=knowledge_base_name,  # optional
+)
+
+results = vector_store_driver.query(query="What is griptape?")
+
+values = [r.to_artifact().value for r in results]
+
+print("\n\n".join(values))

docs/griptape-framework/drivers/vector-store-drivers.md

@@ -294,3 +294,25 @@ The following example shows how to store vector entries and query the informatio
     ```text
     --8<-- "docs/griptape-framework/drivers/logs/vector_store_drivers_11.txt"
     ```
+
+### PG.AI
+
+!!! info
+
+    This Driver requires the `drivers-vector-pgai` [extra](../index.md#extras).
+
+The PgAiKnowledgeBaseVectorStoreDriver integrates with PG.AI, a managed postgres platform from [EnterpriseDB](https://www.enterprisedb.com/).
+
+Here is an example of how the Driver can be used to load and query information in a PG.AI Knowledge Base:
+
+=== "Code"
+
+    ```python
+    --8<-- "docs/griptape-framework/drivers/src/vector_store_drivers_12.py"
+    ```
+
+=== "Logs"
+
+    ```text
+    --8<-- "docs/griptape-framework/drivers/logs/vector_store_drivers_12.txt"
+    ```

griptape/drivers/__init__.py

@@ -49,6 +49,7 @@
 from .vector.qdrant import QdrantVectorStoreDriver
 from .vector.astradb import AstraDbVectorStoreDriver
 from .vector.griptape_cloud import GriptapeCloudVectorStoreDriver
+from .vector.pgai import PgAiKnowledgeBaseVectorStoreDriver
 
 from .sql import BaseSqlDriver
 from .sql.sql_driver import SqlDriver
@@ -226,6 +227,7 @@
     "OpenTelemetryObservabilityDriver",
     "PerplexityPromptDriver",
     "PerplexityWebSearchDriver",
+    "PgAiKnowledgeBaseVectorStoreDriver",
     "PgVectorVectorStoreDriver",
     "PineconeVectorStoreDriver",
     "ProxyWebScraperDriver",

griptape/drivers/vector/pgai/__init__.py

@@ -0,0 +1,3 @@
+from griptape.drivers.vector.pgai_knowledge_base_vector_store_driver import PgAiKnowledgeBaseVectorStoreDriver
+
+__all__ = ["PgAiKnowledgeBaseVectorStoreDriver"]

griptape/drivers/vector/pgai_knowledge_base_vector_store_driver.py

@@ -0,0 +1,108 @@
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING, NoReturn, Optional
+
+from attrs import Factory, define, field
+
+from griptape.artifacts import ImageArtifact, ListArtifact, TextArtifact
+from griptape.drivers.embedding.dummy import DummyEmbeddingDriver
+from griptape.drivers.vector import BaseVectorStoreDriver
+from griptape.utils import import_optional_dependency
+from griptape.utils.decorators import lazy_property
+
+if TYPE_CHECKING:
+    import sqlalchemy
+
+    from griptape.drivers.embedding import BaseEmbeddingDriver
+
+logger = logging.getLogger(__name__)
+
+
+@define
+class PgAiKnowledgeBaseVectorStoreDriver(BaseVectorStoreDriver):
+    connection_string: str = field(kw_only=True, metadata={"serializable": True})
+    knowledge_base_name: str = field(kw_only=True, metadata={"serializable": True})
+    embedding_driver: BaseEmbeddingDriver = field(
+        default=Factory(lambda: DummyEmbeddingDriver()),
+        metadata={"serializable": True},
+        kw_only=True,
+        init=False,
+    )
+    _engine: sqlalchemy.Engine = field(default=None, kw_only=True, alias="engine", metadata={"serializable": False})
+
+    @lazy_property()
+    def engine(self) -> sqlalchemy.Engine:
+        return import_optional_dependency("sqlalchemy").create_engine(self.connection_string)
+
+    def query(
+        self,
+        query: str | TextArtifact | ImageArtifact,
+        *,
+        count: Optional[int] = BaseVectorStoreDriver.DEFAULT_QUERY_COUNT,
+        **kwargs,
+    ) -> list[BaseVectorStoreDriver.Entry]:
+        if isinstance(query, ImageArtifact):
+            raise ValueError(f"{self.__class__.__name__} does not support querying with Image Artifacts.")
+
+        sqlalchemy = import_optional_dependency("sqlalchemy")
+
+        with sqlalchemy.orm.Session(self.engine) as session:
+            rows = session.query(sqlalchemy.func.aidb.retrieve_text(self.knowledge_base_name, query, count)).all()
+
+        entries = []
+        for (row,) in rows:
+            # Remove the first and last parentheses from the row and list by commas
+            # Example: '(foo,bar)' -> ['foo', 'bar']
+            row_list = "".join(row.replace("(", "", 1).rsplit(")", 1)).split(",")
+            entries.append(
+                BaseVectorStoreDriver.Entry(
+                    id=row_list[0],
+                    score=float(row_list[2]),
+                    meta={"artifact": TextArtifact(row_list[1]).to_json()},
+                )
+            )
+
+        return entries
+
+    def upsert_vector(
+        self,
+        vector: list[float],
+        vector_id: Optional[str] = None,
+        namespace: Optional[str] = None,
+        meta: Optional[dict] = None,
+        **kwargs,
+    ) -> str:
+        raise NotImplementedError(f"{self.__class__.__name__} does not support vector upsert.")
+
+    def upsert_text_artifact(
+        self,
+        artifact: TextArtifact,
+        namespace: Optional[str] = None,
+        meta: Optional[dict] = None,
+        vector_id: Optional[str] = None,
+        **kwargs,
+    ) -> str:
+        raise NotImplementedError(f"{self.__class__.__name__} does not support text artifact upsert.")
+
+    def upsert_text(
+        self,
+        string: str,
+        vector_id: Optional[str] = None,
+        namespace: Optional[str] = None,
+        meta: Optional[dict] = None,
+        **kwargs,
+    ) -> str:
+        raise NotImplementedError(f"{self.__class__.__name__} does not support text upsert.")
+
+    def load_entry(self, vector_id: str, *, namespace: Optional[str] = None) -> BaseVectorStoreDriver.Entry:
+        raise NotImplementedError(f"{self.__class__.__name__} does not support entry loading.")
+
+    def load_entries(self, *, namespace: Optional[str] = None) -> list[BaseVectorStoreDriver.Entry]:
+        raise NotImplementedError(f"{self.__class__.__name__} does not support entry loading.")
+
+    def load_artifacts(self, *, namespace: Optional[str] = None) -> ListArtifact:
+        raise NotImplementedError(f"{self.__class__.__name__} does not support Artifact loading.")
+
+    def delete_vector(self, vector_id: str) -> NoReturn:
+        raise NotImplementedError(f"{self.__class__.__name__} does not support deletion.")

pyproject.toml

@@ -60,6 +60,9 @@ drivers-vector-pgvector = [
 ]
 drivers-vector-qdrant = ["qdrant-client>=1.10.1"]
 drivers-vector-astra-db = ["astrapy>=2.0"]
+drivers-vector-pgai = [
+  "sqlalchemy>=2.0.31"
+]
 drivers-embedding-amazon-bedrock = ["boto3>=1.34.119"]
 drivers-embedding-amazon-sagemaker = ["boto3>=1.34.119"]
 drivers-embedding-huggingface = [

tests/integration/test_code_blocks.py

@@ -23,6 +23,7 @@
     "docs/griptape-framework/structures/src/observability_2.py",
     "docs/griptape-framework/data/src/loaders_9.py",
     "docs/recipes/src/talk_to_an_audio_2.py",
+    "docs/griptape-framework/drivers/src/vector_store_drivers_12.py",
 ]
 
 

tests/unit/drivers/vector/test_pgai_knowledge_base_vector_store_driver.py

@@ -0,0 +1,55 @@
+import json
+from unittest.mock import MagicMock
+
+import pytest
+
+from griptape.drivers.vector.pgai import PgAiKnowledgeBaseVectorStoreDriver
+
+
+class TestPgAiKnowledgeBaseVectorStoreDriver:
+    connection_string = "postgresql://postgres:postgres@localhost:5432/postgres"
+    knowledge_base_name = "example_knowledge_base"
+
+    @pytest.fixture()
+    def mock_engine(self):
+        return MagicMock()
+
+    @pytest.fixture()
+    def mock_session(self, mocker):
+        session = MagicMock()
+        mock_session_manager = MagicMock()
+        mock_session_manager.__enter__.return_value = session
+        mocker.patch("sqlalchemy.orm.Session", return_value=mock_session_manager)
+
+        return session
+
+    def test_initialize(self):
+        PgAiKnowledgeBaseVectorStoreDriver(
+            connection_string=self.connection_string, knowledge_base_name=self.knowledge_base_name
+        )
+
+    def test_query(self, mock_engine, mock_session):
+        test_ids = [17, 23]
+        test_values = ['"foo"', "bar"]
+        test_scores = [0.4, 0.6]
+        mock_query = MagicMock()
+        mock_query.all.return_value = [
+            (f"{test_ids[0]},{test_values[0]},{test_scores[0]}",),
+            (f"{test_ids[1]},{test_values[1]},{test_scores[1]}",),
+        ]
+        mock_session.query.return_value = mock_query
+
+        driver = PgAiKnowledgeBaseVectorStoreDriver(
+            engine=mock_engine, connection_string=self.connection_string, knowledge_base_name=self.knowledge_base_name
+        )
+
+        result = driver.query("some query")
+
+        assert result[0].id == str(test_ids[0])
+        assert result[1].id == str(test_ids[1])
+        assert result[0].meta
+        assert json.loads(result[0].meta["artifact"])["value"] == test_values[0]
+        assert result[1].meta
+        assert json.loads(result[1].meta["artifact"])["value"] == test_values[1]
+        assert result[0].score == test_scores[0]
+        assert result[1].score == test_scores[1]