JDSS/mem0
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Enam Biswas
2025-08-01 20:53:34 +02:00
committed by GitHub
parent 0e03d69ed1
commit 907328aafe
4 changed files with 88 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/components/vectordbs/dbs/pinecone.mdx
+6
View File
@@ -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
+1
View File
@@ -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
+17 -8
View File
@@ -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
+64 -3
View File
@@ -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()