diff --git a/docs/components/vectordbs/dbs/pinecone.mdx b/docs/components/vectordbs/dbs/pinecone.mdx index 84673f18d..8633ab256 100644 --- a/docs/components/vectordbs/dbs/pinecone.mdx +++ b/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, diff --git a/mem0/configs/vector_stores/pinecone.py b/mem0/configs/vector_stores/pinecone.py index bca8dfd82..0474b9ac6 100644 --- a/mem0/configs/vector_stores/pinecone.py +++ b/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 diff --git a/mem0/vector_stores/pinecone.py b/mem0/vector_stores/pinecone.py index c551638ef..08ccf8bc6 100644 --- a/mem0/vector_stores/pinecone.py +++ b/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): """ diff --git a/tests/vector_stores/test_pinecone.py b/tests/vector_stores/test_pinecone.py index 0a7fad054..fb7043398 100644 --- a/tests/vector_stores/test_pinecone.py +++ b/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()