GriptapeCloudVectorStoreDriver

CHANGELOG.md

@@ -8,6 +8,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Added
 - Native function calling support to `OpenAiChatPromptDriver`, `AzureOpenAiChatPromptDriver`, `AnthropicPromptDriver`, `AmazonBedrockPromptDriver`, `GooglePromptDriver`, and `CoherePromptDriver`.
 - `OllamaEmbeddingDriver` for generating embeddings with Ollama.
+- `GriptapeCloudKnowledgeBaseVectorStoreDriver` to query Griptape Cloud Knowledge Bases.
 
 ### Changed
 

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

@@ -49,6 +49,30 @@ print("\n\n".join(values))
 
 ```
 
+### Griptape Cloud Knowledge Base
+
+The [GriptapeCloudKnowledgeBaseVectorStoreDriver](../../reference/griptape/drivers/vector/griptape_cloud_knowledge_base_vector_store_driver.md) can be used to query data from a Griptape Cloud Knowledge Base. Loading into Knowledge Bases is not support, only querying. Here is a complete example of how the driver can be used to query an existing Knowledge Base:
+
+```python
+import os 
+from griptape.artifacts import BaseArtifact
+from griptape.drivers import GriptapeCloudKnowledgeBaseVectorStoreDriver
+
+
+# Initialize environment variables
+gt_cloud_api_key = os.environ["GT_CLOUD_API_KEY"]
+gt_cloud_knowledge_base_id = os.environ["GT_CLOUD_KNOWLEDGE_BASE_ID"]
+
+vector_store_driver = GriptapeCloudKnowledgeBaseVectorStoreDriver(api_key=gt_cloud_api_key, knowledge_base_id=gt_cloud_knowledge_base_id)
+
+result = vector_store_driver.query(query="What is griptape?")
+
+values = [r.to_artifact().value for r in results]
+
+print("\n\n".join(values))
+
+```
+
 ### Pinecone
 
 !!! info
@@ -127,7 +151,7 @@ prompt_driver = OpenAiChatPromptDriver(model="gpt-3.5-turbo")
 namespace = 'griptape-ai'
 
 # Initialize the vector store driver
-vector_store = MarqoVectorStoreDriver(
+vector_store_driver = MarqoVectorStoreDriver(
     api_key=os.environ["MARQO_API_KEY"],
     url=os.environ["MARQO_URL"],
     index=os.environ["MARQO_INDEX_NAME"],
@@ -138,13 +162,13 @@ vector_store = MarqoVectorStoreDriver(
 artifacts = WebLoader(max_tokens=200).load("https://www.griptape.ai")
 
 # Upsert the artifacts into the vector store
-vector_store.upsert_text_artifacts(
+vector_store_driver.upsert_text_artifacts(
     {
         "griptape": artifacts,
     }
 )
 
-result = vector_store.query(query="What is griptape?")
+result = vector_store_driver.query(query="What is griptape?")
 print(result)
 ```
 
@@ -174,7 +198,7 @@ index_name = os.environ["MONGODB_INDEX_NAME"]
 vector_path = os.environ["MONGODB_VECTOR_PATH"]
 
 # Initialize the vector store driver
-vector_store = MongoDbAtlasVectorStoreDriver(
+vector_store_driver = MongoDbAtlasVectorStoreDriver(
     connection_string=f"mongodb+srv://{username}:{password}@{host}/{database_name}",
     database_name=database_name,
     collection_name=collection_name,
@@ -187,13 +211,13 @@ vector_store = MongoDbAtlasVectorStoreDriver(
 artifacts = WebLoader(max_tokens=200).load("https://www.griptape.ai")
 
 # Upsert the artifacts into the vector store
-vector_store.upsert_text_artifacts(
+vector_store_driver.upsert_text_artifacts(
     {
         "griptape": artifacts,
     }
 )
 
-result = vector_store.query(query="What is griptape?")
+result = vector_store_driver.query(query="What is griptape?")
 print(result)
 ```
 
@@ -242,7 +266,7 @@ index_name = os.environ["AZURE_MONGODB_INDEX_NAME"]
 vector_path = os.environ["AZURE_MONGODB_VECTOR_PATH"]
 
 # Initialize the vector store driver
-vector_store = AzureMongoDbVectorStoreDriver(
+vector_store_driver = AzureMongoDbVectorStoreDriver(
     connection_string=f"mongodb+srv://{username}:{password}@{azure_host}/{database_name}?tls=true&authMechanism=SCRAM-SHA-256&retrywrites=false&maxIdleTimeMS=120000",
     database_name=database_name,
     collection_name=collection_name,
@@ -255,13 +279,13 @@ vector_store = AzureMongoDbVectorStoreDriver(
 artifacts = WebLoader(max_tokens=200).load("https://www.griptape.ai")
 
 # Upsert the artifacts into the vector store
-vector_store.upsert_text_artifacts(
+vector_store_driver.upsert_text_artifacts(
     {
         "griptape": artifacts,
     }
 )
 
-result = vector_store.query(query="What is griptape?")
+result = vector_store_driver.query(query="What is griptape?")
 print(result)
 ```
 
@@ -403,7 +427,7 @@ vector_store_driver.upsert_text_artifacts(
     }
 )
 
-result = vector_store_driver.query("What is griptape?")
+result = vector_store_driver.query(query="What is griptape?")
 print(result)
 ```
 

griptape/drivers/__init__.py

@@ -41,6 +41,7 @@
 from .vector.azure_mongodb_vector_store_driver import AzureMongoDbVectorStoreDriver
 from .vector.dummy_vector_store_driver import DummyVectorStoreDriver
 from .vector.qdrant_vector_store_driver import QdrantVectorStoreDriver
+from .vector.griptape_cloud_knowledge_base_vector_store_driver import GriptapeCloudKnowledgeBaseVectorStoreDriver
 
 from .sql.base_sql_driver import BaseSqlDriver
 from .sql.amazon_redshift_sql_driver import AmazonRedshiftSqlDriver
@@ -149,6 +150,7 @@
     "PgVectorVectorStoreDriver",
     "QdrantVectorStoreDriver",
     "DummyVectorStoreDriver",
+    "GriptapeCloudKnowledgeBaseVectorStoreDriver",
     "BaseSqlDriver",
     "AmazonRedshiftSqlDriver",
     "SnowflakeSqlDriver",

griptape/drivers/vector/griptape_cloud_knowledge_base_vector_store_driver.py

@@ -0,0 +1,107 @@
+from urllib.parse import urljoin
+import requests
+from typing import Optional, Any
+from attrs import Factory, define, field
+from griptape.artifacts import TextArtifact, ListArtifact
+from griptape.drivers import BaseEmbeddingDriver, BaseVectorStoreDriver, DummyEmbeddingDriver
+
+
+@define
+class GriptapeCloudKnowledgeBaseVectorStoreDriver(BaseVectorStoreDriver):
+    """A vector store driver for Griptape Cloud Knowledge Bases.
+
+    Attributes:
+        api_key: API Key for Griptape Cloud.
+        knowledge_base_id: Knowledge Base ID for Griptape Cloud.
+        base_url: Base URL for Griptape Cloud.
+        headers: Headers for Griptape Cloud.
+    """
+
+    api_key: str = field(kw_only=True, metadata={"serializable": True})
+    knowledge_base_id: str = field(kw_only=True, metadata={"serializable": True})
+    base_url: str = field(default="https://cloud.griptape.ai", kw_only=True)
+    headers: dict = field(
+        default=Factory(lambda self: {"Authorization": f"Bearer {self.api_key}"}, takes_self=True), kw_only=True
+    )
+    embedding_driver: BaseEmbeddingDriver = field(
+        default=Factory(lambda: DummyEmbeddingDriver()), metadata={"serializable": True}, kw_only=True, init=False
+    )
+
+    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:
+        """Retrieves a specific vector entry from the collection based on its identifier and optional namespace."""
+        raise NotImplementedError(f"{self.__class__.__name__} does not support entry loading.")
+
+    def load_entries(self, namespace: Optional[str] = None) -> list[BaseVectorStoreDriver.Entry]:
+        """Retrieves all vector entries from the collection, optionally filtering to only
+        those that match the provided namespace.
+        """
+        raise NotImplementedError(f"{self.__class__.__name__} does not support entry loading.")
+
+    def load_artifacts(self, namespace: Optional[str] = None) -> ListArtifact:
+        """Retrieves all Artifacts from the collection, optionally filtering to only
+        those that match the provided namespace.
+        """
+        raise NotImplementedError(f"{self.__class__.__name__} does not support Artifact loading.")
+
+    def query(
+        self,
+        query: str,
+        count: Optional[int] = None,
+        namespace: Optional[str] = None,
+        include_vectors: Optional[bool] = None,
+        distance_metric: Optional[str] = None,
+        # GriptapeCloudKnowledgeBaseVectorStoreDriver-specific params:
+        filter: Optional[dict] = None,
+        **kwargs,
+    ) -> list[BaseVectorStoreDriver.Entry]:
+        """Performs a search on the Knowledge Base to find vectors similar to the provided input vector,
+        optionally filtering to only those that match the provided filter(s).
+        """
+        url = urljoin(self.base_url.strip("/"), f"/api/knowledge-bases/{self.knowledge_base_id}/query")
+
+        request: dict[str, Any] = {
+            "query": query,
+            "count": count,
+            "distance_metric": distance_metric,
+            "filter": filter,
+            "include_vectors": include_vectors,
+        }
+        request = {k: v for k, v in request.items() if v is not None}
+
+        response = requests.post(url, json=request, headers=self.headers).json()
+        entries = response.get("entries", [])
+        entry_list = [BaseVectorStoreDriver.Entry.from_dict(entry) for entry in entries]
+        return entry_list
+
+    def delete_vector(self, vector_id: str):
+        raise NotImplementedError(f"{self.__class__.__name__} does not support deletion.")

tests/unit/drivers/vector/test_griptape_cloud_knowledge_base_vector_store_driver.py

@@ -0,0 +1,79 @@
+import pytest
+import uuid
+
+
+class MockResponse:
+    def __init__(self, json_data, status_code):
+        self.json_data = json_data
+        self.status_code = status_code
+
+    def json(self):
+        return self.json_data
+
+
+class TestGriptapeCloudKnowledgeBaseVectorStoreDriver:
+    test_ids = [str(uuid.uuid4()), str(uuid.uuid4())]
+    test_vecs = [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]]
+    test_namespaces = [str(uuid.uuid4()), str(uuid.uuid4())]
+    test_metas = [{"key": "value1"}, {"key": "value2"}]
+    test_scores = [0.7, 0.8]
+
+    @pytest.fixture
+    def driver(self, mocker):
+        from griptape.drivers import GriptapeCloudKnowledgeBaseVectorStoreDriver
+
+        test_entries = {
+            "entries": [
+                {
+                    "id": self.test_ids[0],
+                    "vector": self.test_vecs[0],
+                    "namespace": self.test_namespaces[0],
+                    "meta": self.test_metas[0],
+                    "score": self.test_scores[0],
+                },
+                {
+                    "id": self.test_ids[1],
+                    "vector": self.test_vecs[1],
+                    "namespace": self.test_namespaces[1],
+                    "meta": self.test_metas[1],
+                    "score": self.test_scores[1],
+                },
+            ]
+        }
+
+        mock_response = mocker.Mock()
+        mock_response.status_code = 200
+        mock_response.json.return_value = test_entries
+        mocker.patch("requests.post", return_value=mock_response)
+
+        return GriptapeCloudKnowledgeBaseVectorStoreDriver(api_key="foo bar", knowledge_base_id="1")
+
+    def test_query(self, driver):
+        result = driver.query(
+            "some query", count=10, namespace="foo", include_vectors=True, distance_metric="bar", filter={"foo": "bar"}
+        )
+
+        assert result[0].id == self.test_ids[0]
+        assert result[1].id == self.test_ids[1]
+        assert result[0].vector == self.test_vecs[0]
+        assert result[1].vector == self.test_vecs[1]
+        assert result[0].namespace == self.test_namespaces[0]
+        assert result[1].namespace == self.test_namespaces[1]
+        assert result[0].meta == self.test_metas[0]
+        assert result[1].meta == self.test_metas[1]
+        assert result[0].score == self.test_scores[0]
+        assert result[1].score == self.test_scores[1]
+
+    def test_query_defaults(self, driver):
+        result = driver.query("some query")
+
+        assert result[0].id == self.test_ids[0]
+        assert result[1].id == self.test_ids[1]
+        assert result[0].vector == self.test_vecs[0]
+        assert result[1].vector == self.test_vecs[1]
+        assert result[0].namespace == self.test_namespaces[0]
+        assert result[1].namespace == self.test_namespaces[1]
+        assert result[0].meta == self.test_metas[0]
+        assert result[1].meta == self.test_metas[1]
+        assert result[0].score == self.test_scores[0]
+        assert result[1].score == self.test_scores[1]