import logging import os from typing import Any, Dict, List, Optional, Union from pydantic import BaseModel try: from pinecone import Pinecone, PodSpec, ServerlessSpec, Vector except ImportError: raise ImportError( "Pinecone requires extra dependencies. Install with `pip install pinecone pinecone-text`" ) from None from mem0.vector_stores.base import VectorStoreBase logger = logging.getLogger(__name__) class OutputData(BaseModel): id: Optional[str] # memory id score: Optional[float] # distance payload: Optional[Dict] # metadata class PineconeDB(VectorStoreBase): def __init__( self, collection_name: str, embedding_model_dims: int, client: Optional["Pinecone"], api_key: Optional[str], environment: Optional[str], serverless_config: Optional[Dict[str, Any]], pod_config: Optional[Dict[str, Any]], hybrid_search: bool, metric: str, batch_size: int, extra_params: Optional[Dict[str, Any]], namespace: Optional[str] = None, ): """ Initialize the Pinecone vector store. Args: collection_name (str): Name of the index/collection. embedding_model_dims (int): Dimensions of the embedding model. client (Pinecone, optional): Existing Pinecone client instance. Defaults to None. api_key (str, optional): API key for Pinecone. Defaults to None. environment (str, optional): Pinecone environment. Defaults to None. serverless_config (Dict, optional): Configuration for serverless deployment. Defaults to None. pod_config (Dict, optional): Configuration for pod-based deployment. Defaults to None. hybrid_search (bool, optional): Whether to enable hybrid search. Defaults to False. 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 else: api_key = api_key or os.environ.get("PINECONE_API_KEY") if not api_key: raise ValueError( "Pinecone API key must be provided either as a parameter or as an environment variable" ) params = extra_params or {} self.client = Pinecone(api_key=api_key, **params) self.collection_name = collection_name self.embedding_model_dims = embedding_model_dims self.environment = environment self.serverless_config = serverless_config self.pod_config = pod_config self.hybrid_search = hybrid_search self.metric = metric self.batch_size = batch_size self.namespace = namespace self.sparse_encoder = None if self.hybrid_search: try: from pinecone_text.sparse import BM25Encoder logger.info("Initializing BM25Encoder for sparse vectors...") self.sparse_encoder = BM25Encoder.default() except ImportError: logger.warning("pinecone-text not installed. Hybrid search will be disabled.") self.hybrid_search = False self.create_col(embedding_model_dims, metric) def create_col(self, vector_size: int, metric: str = "cosine"): """ Create a new index/collection. Args: vector_size (int): Size of the vectors to be stored. metric (str, optional): Distance metric for vector similarity. Defaults to "cosine". """ existing_indexes = self.list_cols().names() if self.collection_name in existing_indexes: logger.debug(f"Index {self.collection_name} already exists. Skipping creation.") self.index = self.client.Index(self.collection_name) return if self.serverless_config: spec = ServerlessSpec(**self.serverless_config) elif self.pod_config: spec = PodSpec(**self.pod_config) else: spec = ServerlessSpec(cloud="aws", region="us-west-2") self.client.create_index( name=self.collection_name, dimension=vector_size, metric=metric, spec=spec, ) self.index = self.client.Index(self.collection_name) def insert( self, vectors: List[List[float]], payloads: Optional[List[Dict]] = None, ids: Optional[List[Union[str, int]]] = None, ): """ Insert vectors into an index. Args: vectors (list): List of vectors to insert. payloads (list, optional): List of payloads corresponding to vectors. Defaults to None. ids (list, optional): List of IDs corresponding to vectors. Defaults to None. """ logger.info(f"Inserting {len(vectors)} vectors into index {self.collection_name}") items = [] for idx, vector in enumerate(vectors): item_id = str(ids[idx]) if ids is not None else str(idx) payload = payloads[idx] if payloads else {} vector_record = {"id": item_id, "values": vector, "metadata": payload} if self.hybrid_search and self.sparse_encoder and "text" in payload: sparse_vector = self.sparse_encoder.encode_documents(payload["text"]) vector_record["sparse_values"] = sparse_vector items.append(vector_record) if len(items) >= self.batch_size: self.index.upsert(vectors=items, namespace=self.namespace) items = [] if items: self.index.upsert(vectors=items, namespace=self.namespace) def _parse_output(self, data: Dict) -> List[OutputData]: """ Parse the output data from Pinecone search results. Args: data (Dict): Output data from Pinecone query. Returns: List[OutputData]: Parsed output data. """ if isinstance(data, Vector): result = OutputData( id=data.id, score=0.0, payload=data.metadata, ) return result else: result = [] for match in data: entry = OutputData( id=match.get("id"), score=match.get("score"), payload=match.get("metadata"), ) result.append(entry) return result OPERATOR_MAP = { "eq": "$eq", "ne": "$ne", "gt": "$gt", "gte": "$gte", "lt": "$lt", "lte": "$lte", "in": "$in", "nin": "$nin", } def _create_filter(self, filters: Optional[Dict]) -> Dict: """ Create a filter dictionary from the provided filters. """ if not filters: return {} pinecone_filter = {} for key, value in filters.items(): if isinstance(value, dict): condition = {} for op, operand in value.items(): pc_op = self.OPERATOR_MAP.get(op) if pc_op: condition[pc_op] = operand else: condition[f"${op}"] = operand pinecone_filter[key] = condition else: pinecone_filter[key] = {"$eq": value} return pinecone_filter def search( self, query: str, vectors: List[float], top_k: int = 5, filters: Optional[Dict] = None ) -> List[OutputData]: """ Search for similar vectors. Args: query (str): Query. vectors (list): List of vectors to search. top_k (int, optional): Number of results to return. Defaults to 5. filters (dict, optional): Filters to apply to the search. Defaults to None. Returns: list: Search results. """ filter_dict = self._create_filter(filters) if filters else None query_params = { "vector": vectors, "top_k": top_k, "include_metadata": True, "include_values": False, } if filter_dict: query_params["filter"] = filter_dict if self.hybrid_search and self.sparse_encoder and filters and "text" in filters: query_text = filters.get("text") if query_text: sparse_vector = self.sparse_encoder.encode_queries(query_text) query_params["sparse_vector"] = sparse_vector response = self.index.query(**query_params, namespace=self.namespace) results = self._parse_output(response.matches) return results def keyword_search(self, query, top_k=5, filters=None): """ Search using BM25 sparse vectors for keyword-based retrieval. Args: query (str): The search query text. top_k (int, optional): Number of results to return. Defaults to 5. filters (dict, optional): Filters to apply to the search. Defaults to None. Returns: List[OutputData]: Search results, or None if hybrid search is not configured. """ if not self.hybrid_search or self.sparse_encoder is None: return None try: filter_dict = self._create_filter(filters) if filters else None sparse_vector = self.sparse_encoder.encode_queries(query) query_params = { "sparse_vector": sparse_vector, "top_k": top_k, "include_metadata": True, "include_values": False, } if filter_dict: query_params["filter"] = filter_dict response = self.index.query(**query_params, namespace=self.namespace) return self._parse_output(response.matches) except Exception: return None def delete(self, vector_id: Union[str, int]): """ Delete a vector by ID. Args: vector_id (Union[str, int]): ID of the vector to delete. """ 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): """ Update a vector and its payload. Args: vector_id (Union[str, int]): ID of the vector to update. vector (list, optional): Updated vector. Defaults to None. payload (dict, optional): Updated payload. Defaults to None. """ item = { "id": str(vector_id), } if vector is not None: item["values"] = vector if payload is not None: item["metadata"] = payload if self.hybrid_search and self.sparse_encoder and "text" in payload: sparse_vector = self.sparse_encoder.encode_documents(payload["text"]) item["sparse_values"] = sparse_vector self.index.upsert(vectors=[item], namespace=self.namespace) def get(self, vector_id: Union[str, int]) -> OutputData: """ Retrieve a vector by ID. Args: vector_id (Union[str, int]): ID of the vector to retrieve. Returns: dict: Retrieved vector or None if not found. """ try: 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 except Exception as e: logger.error(f"Error retrieving vector {vector_id}: {e}") return None def list_cols(self): """ List all indexes/collections. Returns: list: List of index information. """ return self.client.list_indexes() def delete_col(self): """Delete an index/collection.""" try: self.client.delete_index(self.collection_name) logger.info(f"Index {self.collection_name} deleted successfully") except Exception as e: logger.error(f"Error deleting index {self.collection_name}: {e}") def col_info(self) -> Dict: """ Get information about an index/collection. Returns: dict: Index information. """ return self.client.describe_index(self.collection_name) def list(self, filters: Optional[Dict] = None, top_k: int = 100) -> List[OutputData]: """ List vectors in an index with optional filtering. Args: filters (dict, optional): Filters to apply to the list. Defaults to None. top_k (int, optional): Number of vectors to return. Defaults to 100. Returns: dict: List of vectors with their metadata. """ filter_dict = self._create_filter(filters) if filters else None stats = self.index.describe_index_stats() dimension = stats.dimension zero_vector = [0.0] * dimension query_params = { "vector": zero_vector, "top_k": top_k, "include_metadata": True, "include_values": True, } if filter_dict: query_params["filter"] = filter_dict try: response = self.index.query(**query_params, namespace=self.namespace) response = response.to_dict() results = self._parse_output(response["matches"]) return [results] except Exception as e: logger.error(f"Error listing vectors: {e}") return [[]] def count(self) -> int: """ Count number of vectors in the index. Returns: int: Total number of vectors. """ stats = self.index.describe_index_stats() 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): """ Reset the index by deleting and recreating it. """ logger.warning(f"Resetting index {self.collection_name}...") self.delete_col() self.create_col(self.embedding_model_dims, self.metric)