diff --git a/lumen/ai/vector_store.py b/lumen/ai/vector_store.py index 94912ffb5..7237ff383 100644 --- a/lumen/ai/vector_store.py +++ b/lumen/ai/vector_store.py @@ -1,9 +1,11 @@ import asyncio +import hashlib import importlib import io import json import os import typing as t +import uuid from abc import abstractmethod from collections.abc import Callable @@ -1626,3 +1628,382 @@ def close(self) -> None: if self.connection: self.connection.close() self.connection = None + + +class PineconeVectorStore(VectorStore): + """ + Vector store implementation using Pinecone for persistent storage. + + :Example: + + .. code-block:: python + + from lumen.ai.vector_store import PineconeVectorStore + + vector_store = PineconeVectorStore( + api_key="YOUR_API_KEY", + index_name="lumen-demo", + ) + await vector_store.add([{"text": "Hello!", "metadata": {"source": "greeting"}}]) + await vector_store.query("Hello", threshold=0.1) + """ + + api_key = param.String( + default=os.environ.get("PINECONE_API_KEY", ""), + doc="Pinecone API key. Defaults to the PINECONE_API_KEY environment variable.", + ) + + index_name = param.String( + default="lumen", + doc="Name of the Pinecone index to use.", + ) + + namespace = param.String( + default="", + doc="Pinecone namespace to use for all operations.", + ) + + metric = param.String( + default="cosine", + doc="Pinecone index metric to use when creating an index.", + ) + + cloud = param.String( + default="aws", + doc="Cloud provider for serverless Pinecone indexes.", + ) + + region = param.String( + default="us-east-1", + doc="Region for serverless Pinecone indexes.", + ) + + create_index = param.Boolean( + default=True, + doc="Whether to create the index if it does not exist.", + ) + + dimension = param.Integer( + default=None, + allow_None=True, + doc="Embedding dimension. Required to create a new index.", + ) + + batch_size = param.Integer( + default=100, + bounds=(1, None), + doc="Batch size for Pinecone upserts.", + ) + + max_filter_results = param.Integer( + default=1000, + bounds=(1, None), + doc="Maximum number of results returned from filter_by when limit is None.", + ) + + def __init__(self, **params): + super().__init__(**params) + self._pc = None + self._index = None + self._dimension = None + self._text_key = "_lumen_text" + self._init_client() + + def _init_client(self) -> None: + try: + from pinecone import ( # type: ignore[import-not-found] + Pinecone, ServerlessSpec, + ) + except Exception as e: + raise ImportError( + "Pinecone client not available. Install with `pip install pinecone`." + ) from e + + api_key = self.api_key or os.environ.get("PINECONE_API_KEY", "") + if not api_key: + raise ValueError( + "Pinecone API key not provided. Set api_key or PINECONE_API_KEY." + ) + + self._pc = Pinecone(api_key=api_key) + self._serverless_spec = ServerlessSpec + self._ensure_index() + + def _list_index_names(self) -> list[str]: + try: + indexes = self._pc.list_indexes() + except Exception: + return [] + + if isinstance(indexes, (list, tuple, set)): + return list(indexes) + if hasattr(indexes, "names") and callable(indexes.names): + return list(indexes.names()) + if hasattr(indexes, "names"): + return list(indexes.names) + if hasattr(indexes, "indexes"): + return [ + idx.get("name") if isinstance(idx, dict) else getattr(idx, "name", idx) + for idx in indexes.indexes + ] + return [] + + def _get_index_dimension(self) -> int | None: + if self.dimension is not None: + return self.dimension + if self._dimension is not None: + return self._dimension + if self._index is None: + return None + stats = self._describe_index_stats() + if isinstance(stats, dict): + dimension = stats.get("dimension") + if dimension: + self._dimension = dimension + return dimension + return None + + def _describe_index_stats(self) -> dict: + if self._index is None: + return {} + try: + return self._index.describe_index_stats() + except Exception: + return {} + + def _ensure_index(self, dimension: int | None = None) -> None: + if self._pc is None: + return + index_names = self._list_index_names() + if self.index_name not in index_names: + if not self.create_index: + raise ValueError( + f"Pinecone index '{self.index_name}' does not exist." + ) + dim = dimension or self.dimension + if dim is None: + return + self._pc.create_index( + name=self.index_name, + dimension=dim, + metric=self.metric, + spec=self._serverless_spec(cloud=self.cloud, region=self.region), + ) + self._dimension = dim + self._index = self._pc.Index(self.index_name) + + def _normalize_metadata(self, metadata: dict[str, Any]) -> dict[str, Any]: + normalized = {} + for key, value in metadata.items(): + if isinstance(value, (str, int, float, bool)) or value is None: + normalized[key] = value + elif isinstance(value, (list, tuple)): + normalized[key] = [ + v if isinstance(v, (str, int, float, bool)) or v is None + else self._format_metadata_value(v) + for v in value + ] + else: + normalized[key] = self._format_metadata_value(value) + return normalized + + def _build_vector_payload( + self, + ids: list[int], + texts: list[str], + metadata: list[dict], + embeddings: np.ndarray, + ) -> list[dict]: + vectors = [] + for idx, text in enumerate(texts): + metadata_copy = self._normalize_metadata(metadata[idx]) + metadata_copy[self._text_key] = text + vectors.append( + { + "id": str(ids[idx]), + "values": embeddings[idx].tolist(), + "metadata": metadata_copy, + } + ) + return vectors + + def _generate_random_ids(self, count: int) -> list[int]: + return [uuid.uuid4().int for _ in range(count)] + + def _generate_deterministic_id(self, text: str, metadata: dict) -> int: + payload = json.dumps( + {"text": text, "metadata": self._normalize_metadata(metadata)}, + sort_keys=True, + ) + digest = hashlib.md5(payload.encode("utf-8")).digest() + return int.from_bytes(digest[:8], byteorder="big", signed=False) + + async def _add_items( + self, + texts: list[str], + metadata: list[dict], + embeddings: np.ndarray, + force_ids: list[int] | None = None, + ) -> list[int]: + if not texts: + return [] + + dimension = embeddings.shape[1] if embeddings.ndim == 2 else None + self._ensure_index(dimension=dimension) + if self._index is None: + raise ValueError( + "Pinecone index is not initialized. Provide dimension or create the index first." + ) + + ids = force_ids if force_ids is not None else self._generate_random_ids(len(texts)) + vectors = self._build_vector_payload(ids, texts, metadata, embeddings) + + for start in range(0, len(vectors), self.batch_size): + batch = vectors[start : start + self.batch_size] + await asyncio.to_thread( + self._index.upsert, + vectors=batch, + namespace=self.namespace or None, + ) + return ids + + async def upsert(self, items: list[dict], situate: bool | None = None) -> list[int]: + if not items: + return [] + all_texts, all_metadata, text_and_metadata_list = await self._prepare_items_for_embedding( + items, situate + ) + if not all_texts: + return [] + + embeddings = np.array( + await self.embeddings.embed(text_and_metadata_list), dtype=np.float32 + ) + ids = [ + self._generate_deterministic_id(text, meta) + for text, meta in zip(all_texts, all_metadata, strict=False) + ] + await self._add_items(all_texts, all_metadata, embeddings, force_ids=ids) + return ids + + async def query( + self, + text: str, + top_k: int = 5, + filters: dict | None = None, + threshold: float = -1.0, + ) -> list[dict]: + if self._index is None: + self._ensure_index() + if self._index is None: + return [] + + query_embedding = np.array( + (await self.embeddings.embed([text]))[0], dtype=np.float32 + ).tolist() + + response = await asyncio.to_thread( + self._index.query, + vector=query_embedding, + top_k=top_k, + include_metadata=True, + filter=filters or None, + namespace=self.namespace or None, + ) + + matches = response.get("matches", []) if isinstance(response, dict) else response.matches + results = [] + for match in matches: + match_metadata = match.get("metadata", {}) if isinstance(match, dict) else getattr(match, "metadata", {}) + match_id = match.get("id") if isinstance(match, dict) else getattr(match, "id", None) + score = match.get("score") if isinstance(match, dict) else getattr(match, "score", None) + + metadata_copy = dict(match_metadata) if match_metadata else {} + text_value = metadata_copy.pop(self._text_key, "") + try: + match_id = int(match_id) + except (TypeError, ValueError): + pass + similarity = float(score) if score is not None else 0.0 + if similarity < threshold: + continue + results.append( + { + "id": match_id, + "text": text_value, + "metadata": metadata_copy, + "similarity": similarity, + } + ) + return results + + def filter_by( + self, filters: dict, limit: int | None = None, offset: int = 0 + ) -> list[dict]: + if self._index is None: + self._ensure_index() + if self._index is None: + return [] + dimension = self._get_index_dimension() + if dimension is None: + return [] + + top_k = limit or self.max_filter_results + zero_vector = [0.0] * dimension + response = self._index.query( + vector=zero_vector, + top_k=top_k, + include_metadata=True, + filter=filters or None, + namespace=self.namespace or None, + ) + matches = response.get("matches", []) if isinstance(response, dict) else response.matches + + results = [] + for match in matches: + match_metadata = match.get("metadata", {}) if isinstance(match, dict) else getattr(match, "metadata", {}) + match_id = match.get("id") if isinstance(match, dict) else getattr(match, "id", None) + metadata_copy = dict(match_metadata) if match_metadata else {} + text_value = metadata_copy.pop(self._text_key, "") + try: + match_id = int(match_id) + except (TypeError, ValueError): + pass + results.append( + {"id": match_id, "text": text_value, "metadata": metadata_copy} + ) + + if offset: + results = results[offset:] + if limit is not None: + results = results[:limit] + return results + + def delete(self, ids: list[int]) -> None: + if not ids or self._index is None: + return + str_ids = [str(item_id) for item_id in ids] + self._index.delete(ids=str_ids, namespace=self.namespace or None) + + def clear(self) -> None: + if self._index is None: + self._ensure_index() + if self._index is None: + return + self._index.delete(delete_all=True, namespace=self.namespace or None) + + def __len__(self) -> int: + if self._index is None: + self._ensure_index() + stats = self._describe_index_stats() + if not isinstance(stats, dict): + return 0 + if self.namespace: + namespaces = stats.get("namespaces", {}) + namespace_stats = namespaces.get(self.namespace, {}) + return namespace_stats.get("vector_count", 0) + return stats.get("total_vector_count", 0) + + def close(self) -> None: + """Close the Pinecone client (no-op for Pinecone indexes).""" + self._index = None