feat (pinecone): Add namespace support and improve type safety (#3216)

docs/components/vectordbs/dbs/pinecone.mdx

@@ -1,5 +1,7 @@
 [Pinecone](https://www.pinecone.io/) is a fully managed vector database designed for machine learning applications, offering high performance vector search with low latency at scale. It's particularly well-suited for semantic search, recommendation systems, and other AI-powered applications.
 
+> **New**: Pinecone integration now supports custom namespaces! Use the `namespace` parameter to logically separate data within the same index. This is especially useful for multi-tenant or multi-user applications.
+
 > **Note**: Before configuring Pinecone, you need to select an embedding model (e.g., OpenAI, Cohere, or custom models) and ensure the `embedding_model_dims` in your config matches your chosen model's dimensions. For example, OpenAI's text-embedding-3-small uses 1536 dimensions.
 
 ### Usage
@@ -18,6 +20,7 @@ config = {
         "config": {
             "collection_name": "testing",
             "embedding_model_dims": 1536,  # Matches OpenAI's text-embedding-3-small
+            "namespace": "my-namespace", # Optional: specify a namespace for multi-tenancy
             "serverless_config": {
                 "cloud": "aws",  # Choose between 'aws' or 'gcp' or 'azure'
                 "region": "us-east-1"
@@ -53,6 +56,7 @@ Here are the parameters available for configuring Pinecone:
 | `hybrid_search` | Whether to enable hybrid search | `False` |
 | `metric` | Distance metric for vector similarity | `"cosine"` |
 | `batch_size` | Batch size for operations | `100` |
+| `namespace` | Namespace for the collection, useful for multi-tenancy. | `None` |
 
 > **Important**: You must choose either `serverless_config` or `pod_config` for your deployment, but not both.
 
@@ -64,6 +68,7 @@ config = {
         "config": {
             "collection_name": "memory_index",
             "embedding_model_dims": 1536,  # For OpenAI's text-embedding-3-small
+            "namespace": "my-namespace",  # Optional: custom namespace
             "serverless_config": {
                 "cloud": "aws",  # or "gcp" or "azure"
                 "region": "us-east-1"  # Choose appropriate region
@@ -81,6 +86,7 @@ config = {
         "config": {
             "collection_name": "memory_index",
             "embedding_model_dims": 1536,  # For OpenAI's text-embedding-ada-002
+            "namespace": "my-namespace",  # Optional: custom namespace
             "pod_config": {
                 "environment": "gcp-starter",
                 "replicas": 1,

mem0/configs/vector_stores/pinecone.py

@@ -18,6 +18,7 @@ class PineconeConfig(BaseModel):
     metric: str = Field("cosine", description="Distance metric for vector similarity")
     batch_size: int = Field(100, description="Batch size for operations")
     extra_params: Optional[Dict[str, Any]] = Field(None, description="Additional parameters for Pinecone client")
+    namespace: Optional[str] = Field(None, description="Namespace for the collection")
 
     @model_validator(mode="before")
     @classmethod

mem0/vector_stores/pinecone.py

@@ -36,6 +36,7 @@ class PineconeDB(VectorStoreBase):
         metric: str,
         batch_size: int,
         extra_params: Optional[Dict[str, Any]],
+        namespace: Optional[str] = None,
     ):
         """
         Initialize the Pinecone vector store.
@@ -52,6 +53,7 @@ class PineconeDB(VectorStoreBase):
             metric (str, optional): Distance metric for vector similarity. Defaults to "cosine".
             batch_size (int, optional): Batch size for operations. Defaults to 100.
             extra_params (Dict, optional): Additional parameters for Pinecone client. Defaults to None.
+            namespace (str, optional): Namespace for the collection. Defaults to None.
         """
         if client:
             self.client = client
@@ -73,6 +75,7 @@ class PineconeDB(VectorStoreBase):
         self.hybrid_search = hybrid_search
         self.metric = metric
         self.batch_size = batch_size
+        self.namespace = namespace
 
         self.sparse_encoder = None
         if self.hybrid_search:
@@ -148,11 +151,11 @@ class PineconeDB(VectorStoreBase):
             items.append(vector_record)
 
             if len(items) >= self.batch_size:
-                self.index.upsert(vectors=items)
+                self.index.upsert(vectors=items, namespace=self.namespace)
                 items = []
 
         if items:
-            self.index.upsert(vectors=items)
+            self.index.upsert(vectors=items, namespace=self.namespace)
 
     def _parse_output(self, data: Dict) -> List[OutputData]:
         """
@@ -233,7 +236,7 @@ class PineconeDB(VectorStoreBase):
                 sparse_vector = self.sparse_encoder.encode_queries(query_text)
                 query_params["sparse_vector"] = sparse_vector
 
-        response = self.index.query(**query_params)
+        response = self.index.query(**query_params, namespace=self.namespace)
 
         results = self._parse_output(response.matches)
         return results
@@ -245,7 +248,7 @@ class PineconeDB(VectorStoreBase):
         Args:
             vector_id (Union[str, int]): ID of the vector to delete.
         """
-        self.index.delete(ids=[str(vector_id)])
+        self.index.delete(ids=[str(vector_id)], namespace=self.namespace)
 
     def update(self, vector_id: Union[str, int], vector: Optional[List[float]] = None, payload: Optional[Dict] = None):
         """
@@ -270,7 +273,7 @@ class PineconeDB(VectorStoreBase):
                 sparse_vector = self.sparse_encoder.encode_documents(payload["text"])
                 item["sparse_values"] = sparse_vector
 
-        self.index.upsert(vectors=[item])
+        self.index.upsert(vectors=[item], namespace=self.namespace)
 
     def get(self, vector_id: Union[str, int]) -> OutputData:
         """
@@ -283,7 +286,7 @@ class PineconeDB(VectorStoreBase):
             dict: Retrieved vector or None if not found.
         """
         try:
-            response = self.index.fetch(ids=[str(vector_id)])
+            response = self.index.fetch(ids=[str(vector_id)], namespace=self.namespace)
             if str(vector_id) in response.vectors:
                 return self._parse_output(response.vectors[str(vector_id)])
             return None
@@ -346,7 +349,7 @@ class PineconeDB(VectorStoreBase):
             query_params["filter"] = filter_dict
 
         try:
-            response = self.index.query(**query_params)
+            response = self.index.query(**query_params, namespace=self.namespace)
             response = response.to_dict()
             results = self._parse_output(response["matches"])
             return [results]
@@ -362,7 +365,13 @@ class PineconeDB(VectorStoreBase):
             int: Total number of vectors.
         """
         stats = self.index.describe_index_stats()
-        return stats.total_vector_count
+        if self.namespace:
+            # Safely get the namespace stats and return vector_count, defaulting to 0 if not found
+            namespace_summary = (stats.namespaces or {}).get(self.namespace)
+            if namespace_summary:
+                return namespace_summary.vector_count or 0
+            return 0
+        return stats.total_vector_count or 0
 
     def reset(self):
         """

tests/vector_stores/test_pinecone.py

@@ -27,6 +27,7 @@ def pinecone_db(mock_pinecone_client):
         metric="cosine",
         batch_size=100,
         extra_params=None,
+        namespace="test_namespace",
     )
 
 
@@ -46,6 +47,7 @@ def test_create_col_existing_index(mock_pinecone_client):
         metric="cosine",
         batch_size=100,
         extra_params=None,
+        namespace="test_namespace",
     )
 
     # Reset the mock to verify it wasn't called during the test
@@ -67,12 +69,25 @@ def test_insert_vectors(pinecone_db):
     payloads = [{"name": "vector1"}, {"name": "vector2"}]
     ids = ["id1", "id2"]
     pinecone_db.insert(vectors, payloads, ids)
-    pinecone_db.index.upsert.assert_called()
+    pinecone_db.index.upsert.assert_called_with(
+        vectors=[
+            {"id": "id1", "values": [0.1] * 128, "metadata": {"name": "vector1"}},
+            {"id": "id2", "values": [0.2] * 128, "metadata": {"name": "vector2"}},
+        ],
+        namespace="test_namespace",
+    )
 
 
 def test_search_vectors(pinecone_db):
     pinecone_db.index.query.return_value.matches = [{"id": "id1", "score": 0.9, "metadata": {"name": "vector1"}}]
     results = pinecone_db.search("test query", [0.1] * 128, limit=1)
+    pinecone_db.index.query.assert_called_with(
+        vector=[0.1] * 128,
+        top_k=1,
+        include_metadata=True,
+        include_values=False,
+        namespace="test_namespace",
+    )
     assert len(results) == 1
     assert results[0].id == "id1"
     assert results[0].score == 0.9
@@ -80,7 +95,10 @@ def test_search_vectors(pinecone_db):
 
 def test_update_vector(pinecone_db):
     pinecone_db.update("id1", vector=[0.5] * 128, payload={"name": "updated"})
-    pinecone_db.index.upsert.assert_called()
+    pinecone_db.index.upsert.assert_called_with(
+        vectors=[{"id": "id1", "values": [0.5] * 128, "metadata": {"name": "updated"}}],
+        namespace="test_namespace",
+    )
 
 
 def test_get_vector_found(pinecone_db):
@@ -98,6 +116,7 @@ def test_get_vector_found(pinecone_db):
     pinecone_db.index.fetch.return_value = mock_response
 
     result = pinecone_db.get("id1")
+    pinecone_db.index.fetch.assert_called_with(ids=["id1"], namespace="test_namespace")
     assert result is not None
     assert result.id == "id1"
     assert result.payload == {"name": "vector1"}
@@ -105,12 +124,13 @@ def test_get_vector_found(pinecone_db):
 
 def test_delete_vector(pinecone_db):
     pinecone_db.delete("id1")
-    pinecone_db.index.delete.assert_called_with(ids=["id1"])
+    pinecone_db.index.delete.assert_called_with(ids=["id1"], namespace="test_namespace")
 
 
 def test_get_vector_not_found(pinecone_db):
     pinecone_db.index.fetch.return_value.vectors = {}
     result = pinecone_db.get("id1")
+    pinecone_db.index.fetch.assert_called_with(ids=["id1"], namespace="test_namespace")
     assert result is None
 
 
@@ -127,3 +147,44 @@ def test_delete_col(pinecone_db):
 def test_col_info(pinecone_db):
     pinecone_db.col_info()
     pinecone_db.client.describe_index.assert_called_with("test_index")
+
+
+def test_count_with_namespace(pinecone_db):
+    stats_mock = MagicMock()
+    stats_mock.namespaces = {"test_namespace": MagicMock(vector_count=10)}
+    pinecone_db.index.describe_index_stats.return_value = stats_mock
+
+    count = pinecone_db.count()
+    assert count == 10
+    pinecone_db.index.describe_index_stats.assert_called_once()
+
+
+def test_count_without_namespace(pinecone_db):
+    pinecone_db.namespace = None
+    stats_mock = MagicMock()
+    stats_mock.total_vector_count = 20
+    pinecone_db.index.describe_index_stats.return_value = stats_mock
+
+    count = pinecone_db.count()
+    assert count == 20
+    pinecone_db.index.describe_index_stats.assert_called_once()
+
+
+def test_count_with_non_existent_namespace(pinecone_db):
+    stats_mock = MagicMock()
+    stats_mock.namespaces = {"another_namespace": MagicMock(vector_count=5)}
+    pinecone_db.index.describe_index_stats.return_value = stats_mock
+
+    count = pinecone_db.count()
+    assert count == 0
+    pinecone_db.index.describe_index_stats.assert_called_once()
+
+
+def test_count_with_none_vector_count(pinecone_db):
+    stats_mock = MagicMock()
+    stats_mock.namespaces = {"test_namespace": MagicMock(vector_count=None)}
+    pinecone_db.index.describe_index_stats.return_value = stats_mock
+
+    count = pinecone_db.count()
+    assert count == 0
+    pinecone_db.index.describe_index_stats.assert_called_once()