wehub-resource-sync
692 lines
30 KiB
Python
692 lines
30 KiB
Python
# Copyright (c) Microsoft. All rights reserved.
|
|
|
|
import ast
|
|
import logging
|
|
import sys
|
|
from collections.abc import Sequence
|
|
from inspect import isawaitable
|
|
from typing import Any, ClassVar, Final, Generic, TypeVar
|
|
|
|
from pinecone import FetchResponse, IndexModel, Metric, PineconeAsyncio, ServerlessSpec, Vector
|
|
from pinecone.db_data.index_asyncio import _IndexAsyncio as IndexAsyncio
|
|
from pinecone.grpc import GRPCIndex, GRPCVector, PineconeGRPC
|
|
from pydantic import SecretStr, ValidationError
|
|
|
|
from semantic_kernel.connectors.ai.embedding_generator_base import EmbeddingGeneratorBase
|
|
from semantic_kernel.data.vector import (
|
|
DistanceFunction,
|
|
GetFilteredRecordOptions,
|
|
KernelSearchResults,
|
|
SearchType,
|
|
TModel,
|
|
VectorSearch,
|
|
VectorSearchOptions,
|
|
VectorSearchResult,
|
|
VectorStore,
|
|
VectorStoreCollection,
|
|
VectorStoreCollectionDefinition,
|
|
VectorStoreField,
|
|
_get_collection_name_from_model,
|
|
)
|
|
from semantic_kernel.exceptions.vector_store_exceptions import (
|
|
VectorStoreInitializationException,
|
|
VectorStoreModelException,
|
|
VectorStoreOperationException,
|
|
)
|
|
from semantic_kernel.kernel_pydantic import KernelBaseSettings
|
|
from semantic_kernel.kernel_types import OneOrMany
|
|
from semantic_kernel.utils.feature_stage_decorator import release_candidate
|
|
|
|
if sys.version_info >= (3, 12):
|
|
from typing import override # pragma: no cover
|
|
else:
|
|
from typing_extensions import override # pragma: no cover
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
TKey = TypeVar("TKey", bound=str)
|
|
|
|
DISTANCE_METRIC_MAP: Final[dict[DistanceFunction, Metric]] = {
|
|
DistanceFunction.COSINE_SIMILARITY: Metric.COSINE,
|
|
DistanceFunction.EUCLIDEAN_DISTANCE: Metric.EUCLIDEAN,
|
|
DistanceFunction.DOT_PROD: Metric.DOTPRODUCT,
|
|
DistanceFunction.DEFAULT: Metric.COSINE,
|
|
}
|
|
|
|
|
|
class PineconeSettings(KernelBaseSettings):
|
|
"""Pinecone model settings.
|
|
|
|
Args:
|
|
- api_key: SecretStr - Pinecone API key
|
|
(Env var PINECONE_API_KEY)
|
|
- namespace: str - Pinecone namespace (optional, default is "")
|
|
- embed_model: str - Embedding model (optional, default is None)
|
|
(Env var PINECONE_EMBED_MODEL)
|
|
"""
|
|
|
|
env_prefix: ClassVar[str] = "PINECONE_"
|
|
|
|
api_key: SecretStr
|
|
namespace: str = ""
|
|
embed_model: str | None = None
|
|
|
|
|
|
@release_candidate
|
|
class PineconeCollection(
|
|
VectorStoreCollection[TKey, TModel],
|
|
VectorSearch[TKey, TModel],
|
|
Generic[TKey, TModel],
|
|
):
|
|
"""Interact with a Pinecone Index."""
|
|
|
|
client: PineconeGRPC | PineconeAsyncio
|
|
namespace: str = ""
|
|
index: IndexModel | None = None
|
|
index_client: IndexAsyncio | GRPCIndex | None = None
|
|
supported_key_types: ClassVar[set[str] | None] = {"str"}
|
|
supported_search_types: ClassVar[set[SearchType]] = {SearchType.VECTOR}
|
|
embed_settings: dict[str, Any] | None = None
|
|
|
|
def __init__(
|
|
self,
|
|
record_type: type[TModel],
|
|
definition: VectorStoreCollectionDefinition | None = None,
|
|
collection_name: str | None = None,
|
|
client: PineconeGRPC | PineconeAsyncio | None = None,
|
|
embedding_generator: EmbeddingGeneratorBase | None = None,
|
|
embed_model: str | None = None,
|
|
embed_settings: dict[str, Any] | None = None,
|
|
use_grpc: bool = False,
|
|
api_key: str | None = None,
|
|
namespace: str | None = None,
|
|
env_file_path: str | None = None,
|
|
env_file_encoding: str | None = None,
|
|
**kwargs: Any,
|
|
) -> None:
|
|
"""Initialize the Pinecone collection.
|
|
|
|
Args:
|
|
record_type: The type of the data model.
|
|
definition: The definition of the data model.
|
|
collection_name: The name of the Pinecone collection.
|
|
client: The Pinecone client to use. If not provided, a new client will be created.
|
|
use_grpc: Whether to use the GRPC client or not. Default is False.
|
|
embedding_generator: The embedding generator to use. If not provided, it will be read from the environment.
|
|
embed_model: The settings for the embedding model. If not provided, it will be read from the environment.
|
|
This cannot be combined with a GRPC client.
|
|
embed_settings: The settings for the embedding model. If not provided, the model can be read
|
|
from the environment.
|
|
The other settings are created based on the data model.
|
|
See the pinecone documentation for more details.
|
|
This cannot be combined with a GRPC client.
|
|
api_key: The Pinecone API key. If not provided, it will be read from the environment.
|
|
namespace: The namespace to use. Default is "".
|
|
env_file_path: The path to the environment file. If not provided, it will be read from the default location.
|
|
env_file_encoding: The encoding of the environment file.
|
|
kwargs: Additional arguments to pass to the Pinecone client.
|
|
"""
|
|
if not collection_name:
|
|
collection_name = _get_collection_name_from_model(record_type, definition)
|
|
managed_client = not client
|
|
try:
|
|
settings = PineconeSettings(
|
|
api_key=api_key,
|
|
embed_model=embed_model,
|
|
namespace=namespace,
|
|
env_file_path=env_file_path,
|
|
env_file_encoding=env_file_encoding,
|
|
)
|
|
except ValidationError as exc:
|
|
raise VectorStoreInitializationException(f"Failed to create Pinecone settings: {exc}") from exc
|
|
|
|
if embed_settings:
|
|
if "model" not in embed_settings:
|
|
embed_settings["model"] = settings.embed_model
|
|
if settings.embed_model and embed_settings["model"] != settings.embed_model:
|
|
logger.warning(
|
|
"The model in the embed_settings is different from the one in "
|
|
"the settings. The one in the settings will be used."
|
|
)
|
|
elif settings.embed_model:
|
|
embed_settings = {
|
|
"model": settings.embed_model,
|
|
}
|
|
if not client:
|
|
if use_grpc:
|
|
client = PineconeGRPC(
|
|
api_key=settings.api_key.get_secret_value(),
|
|
**kwargs,
|
|
)
|
|
else:
|
|
client = PineconeAsyncio(
|
|
api_key=settings.api_key.get_secret_value(),
|
|
**kwargs,
|
|
)
|
|
|
|
super().__init__(
|
|
collection_name=collection_name,
|
|
record_type=record_type,
|
|
definition=definition,
|
|
client=client,
|
|
embed_settings=embed_settings,
|
|
namespace=settings.namespace,
|
|
managed_client=managed_client,
|
|
embedding_generator=embedding_generator,
|
|
**kwargs,
|
|
)
|
|
|
|
def _validate_data_model(self):
|
|
"""Check if there is exactly one vector."""
|
|
super()._validate_data_model()
|
|
if len(self.definition.vector_field_names) > 1:
|
|
raise VectorStoreInitializationException(
|
|
"Pinecone only supports one (or zero when using the integrated inference) vector field. "
|
|
"Please use a different data model or "
|
|
f"remove {len(self.definition.vector_field_names) - 1} vector fields."
|
|
)
|
|
|
|
@override
|
|
async def ensure_collection_exists(self, **kwargs: Any) -> None:
|
|
"""Create the Pinecone collection.
|
|
|
|
Args:
|
|
kwargs: Additional arguments to pass to the Pinecone collection creation.
|
|
- embed: if you want to support vectorizable text search,
|
|
you need to set this to a dict with the parameters
|
|
see https://docs.pinecone.io/guides/inference/understanding-inference
|
|
for more details.
|
|
Optionally, the `metric` and `field_map` will be filled based on the data model.
|
|
This can not be used with the GRPC client.
|
|
- cloud: The cloud provider to use. Default is "aws".
|
|
- region: The region to use. Default is "us-east-1".
|
|
"""
|
|
vector_field = self.definition.vector_fields[0] if self.definition.vector_fields else None
|
|
await (
|
|
self._create_index_with_integrated_embeddings(vector_field, **kwargs)
|
|
if self.embed_settings is not None or "embed" in kwargs
|
|
else self._create_regular_index(vector_field, **kwargs)
|
|
)
|
|
|
|
async def _create_index_with_integrated_embeddings(
|
|
self, vector_field: VectorStoreField | None, **kwargs: Any
|
|
) -> None:
|
|
"""Create the Pinecone index with the embed parameter."""
|
|
if isinstance(self.client, PineconeGRPC):
|
|
raise VectorStoreOperationException(
|
|
"Pinecone GRPC client does not support integrated embeddings. Please use the Pinecone Asyncio client."
|
|
)
|
|
if self.embed_settings:
|
|
embed = self.embed_settings.copy()
|
|
embed.update(kwargs.pop("embed", {}))
|
|
else:
|
|
embed = kwargs.pop("embed", {})
|
|
cloud = kwargs.pop("cloud", "aws")
|
|
region = kwargs.pop("region", "us-east-1")
|
|
if "metric" not in embed and vector_field:
|
|
if vector_field.distance_function not in DISTANCE_METRIC_MAP:
|
|
raise VectorStoreOperationException(
|
|
f"Distance function {vector_field.distance_function} is not supported by Pinecone."
|
|
)
|
|
embed["metric"] = DISTANCE_METRIC_MAP[vector_field.distance_function]
|
|
if "field_map" not in embed:
|
|
for field in self.definition.vector_fields:
|
|
if not field.embedding_generator and not self.embedding_generator:
|
|
embed["field_map"] = {"text": field.storage_name or field.name}
|
|
break
|
|
index_creation_args = {
|
|
"name": self.collection_name,
|
|
"cloud": cloud,
|
|
"region": region,
|
|
"embed": embed,
|
|
}
|
|
index_creation_args.update(kwargs)
|
|
self.index = await self.client.create_index_for_model(**index_creation_args)
|
|
await self._load_index_client()
|
|
|
|
async def _create_regular_index(self, vector_field: VectorStoreField | None, **kwargs: Any) -> None:
|
|
"""Create the Pinecone index with the embed parameter."""
|
|
if not vector_field:
|
|
raise VectorStoreOperationException(
|
|
"Pinecone collection needs a vector field, when not using the integrated embeddings."
|
|
)
|
|
if vector_field.distance_function not in DISTANCE_METRIC_MAP:
|
|
raise VectorStoreOperationException(
|
|
f"Distance function {vector_field.distance_function} is not supported by Pinecone."
|
|
)
|
|
cloud = kwargs.pop("cloud", "aws")
|
|
region = kwargs.pop("region", "us-east-1")
|
|
spec = kwargs.pop("spec", ServerlessSpec(cloud=cloud, region=region))
|
|
index_creation_args = {
|
|
"name": self.collection_name,
|
|
"spec": spec,
|
|
"dimension": vector_field.dimensions,
|
|
"metric": DISTANCE_METRIC_MAP[vector_field.distance_function],
|
|
"vector_type": "dense",
|
|
}
|
|
index_creation_args.update(kwargs)
|
|
index = self.client.create_index(**index_creation_args)
|
|
if isawaitable(index):
|
|
index = await index # type: ignore
|
|
self.index = index
|
|
await self._load_index_client()
|
|
|
|
async def _load_index_client(self) -> None:
|
|
if not self.index:
|
|
index = self.client.describe_index(self.collection_name)
|
|
if isawaitable(index):
|
|
index = await index # type: ignore
|
|
self.index = index
|
|
if self.index.embed is not None:
|
|
if isinstance(self.client, PineconeGRPC):
|
|
raise VectorStoreOperationException(
|
|
"Pinecone GRPC client does not support integrated embeddings. "
|
|
"Please use the Pinecone Asyncio client."
|
|
)
|
|
self.embed_settings = self.index.embed
|
|
if not self.index_client:
|
|
self.index_client = (
|
|
self.client.IndexAsyncio(host=self.index.host)
|
|
if isinstance(self.client, PineconeAsyncio)
|
|
else self.client.Index(host=self.index.host)
|
|
)
|
|
|
|
@override
|
|
async def collection_exists(self, **kwargs) -> bool:
|
|
"""Check if the Pinecone collection exists."""
|
|
exists = (
|
|
await self.client.has_index(self.collection_name)
|
|
if isinstance(self.client, PineconeAsyncio)
|
|
else self.client.has_index(self.collection_name)
|
|
)
|
|
if exists:
|
|
await self._load_index_client()
|
|
return exists
|
|
|
|
@override
|
|
async def ensure_collection_deleted(self, **kwargs: Any) -> None:
|
|
"""Delete the Pinecone collection."""
|
|
if not await self.collection_exists():
|
|
if self.index or self.index_client:
|
|
self.index = None
|
|
self.index_client = None
|
|
return
|
|
await self.client.delete_index(self.collection_name) if isinstance(
|
|
self.client, PineconeAsyncio
|
|
) else self.client.delete_index(self.collection_name)
|
|
self.index = None
|
|
if self.index_client:
|
|
await self.index_client.close() if isinstance(
|
|
self.index_client, IndexAsyncio
|
|
) else self.index_client.close()
|
|
self.index_client = None
|
|
|
|
def _record_to_pinecone_vector(self, record: dict[str, Any]) -> Vector | GRPCVector | dict[str, Any]:
|
|
"""Convert a record to a Pinecone vector."""
|
|
metadata_fields = self.definition.get_storage_names(include_key_field=False, include_vector_fields=False)
|
|
vector_field = self.definition.vector_fields[0]
|
|
if isinstance(self.client, PineconeGRPC):
|
|
return GRPCVector(
|
|
id=record[self._key_field_storage_name],
|
|
values=record.get(vector_field.storage_name or vector_field.name, None),
|
|
metadata={key: value for key, value in record.items() if key in metadata_fields},
|
|
)
|
|
if self.embed_settings is not None:
|
|
record.pop(vector_field.storage_name or vector_field.name, None)
|
|
record["_id"] = record.pop(self._key_field_name)
|
|
return record
|
|
return Vector(
|
|
id=record[self._key_field_storage_name],
|
|
values=record.get(vector_field.storage_name or vector_field.name, None) or list(),
|
|
metadata={key: value for key, value in record.items() if key in metadata_fields},
|
|
)
|
|
|
|
def _pinecone_vector_to_record(self, record: Vector | dict[str, Any]) -> dict[str, Any]:
|
|
"""Convert a Pinecone vector to a record."""
|
|
if isinstance(record, dict):
|
|
record[self._key_field_storage_name] = record.pop("_id")
|
|
return record
|
|
vector_field = self.definition.vector_fields[0]
|
|
ret_record: dict[str, Any] = {
|
|
self._key_field_storage_name: record.id,
|
|
vector_field.storage_name or vector_field.name: record.values,
|
|
}
|
|
if record.metadata:
|
|
ret_record.update(record.metadata)
|
|
return ret_record
|
|
|
|
@override
|
|
def _serialize_dicts_to_store_models(
|
|
self, records: Sequence[dict[str, Any]], **kwargs: Any
|
|
) -> Sequence[Vector | GRPCVector | dict[str, Any]]:
|
|
return [self._record_to_pinecone_vector(record) for record in records]
|
|
|
|
@override
|
|
def _deserialize_store_models_to_dicts(self, records: Sequence[Vector], **kwargs: Any) -> Sequence[dict[str, Any]]:
|
|
return [self._pinecone_vector_to_record(record) for record in records]
|
|
|
|
@override
|
|
async def _inner_upsert(
|
|
self,
|
|
records: Sequence[Any],
|
|
**kwargs: Any,
|
|
) -> Sequence[TKey]:
|
|
"""Upsert the records to the Pinecone collection."""
|
|
if not self.index_client:
|
|
await self._load_index_client()
|
|
if not self.index_client:
|
|
raise VectorStoreOperationException("Pinecone collection is not initialized.")
|
|
if "namespace" not in kwargs:
|
|
kwargs["namespace"] = self.namespace
|
|
if self.embed_settings is not None:
|
|
if isinstance(self.index_client, GRPCIndex):
|
|
raise VectorStoreOperationException(
|
|
"Pinecone GRPC client does not support integrated embeddings. "
|
|
"Please use the Pinecone Asyncio client."
|
|
)
|
|
await self.index_client.upsert_records(records=records, **kwargs) # type: ignore
|
|
return [record["_id"] for record in records]
|
|
if isinstance(self.index_client, GRPCIndex):
|
|
self.index_client.upsert(records, **kwargs) # type: ignore
|
|
else:
|
|
await self.index_client.upsert(records, **kwargs) # type: ignore
|
|
return [record.id for record in records]
|
|
|
|
@override
|
|
async def _inner_get(
|
|
self, keys: Sequence[TKey] | None = None, options: GetFilteredRecordOptions | None = None, **kwargs: Any
|
|
) -> OneOrMany[Any] | None:
|
|
"""Get the records from the Pinecone collection."""
|
|
if not keys:
|
|
if options is not None:
|
|
raise NotImplementedError("Get without keys is not yet implemented.")
|
|
return None
|
|
if not self.index_client:
|
|
await self._load_index_client()
|
|
if not self.index_client:
|
|
raise VectorStoreOperationException("Pinecone collection is not initialized.")
|
|
namespace = kwargs.get("namespace", self.namespace)
|
|
if isinstance(self.index_client, GRPCIndex):
|
|
response: FetchResponse = self.index_client.fetch(ids=keys, namespace=namespace) # type: ignore
|
|
else:
|
|
response = await self.index_client.fetch(ids=keys, namespace=namespace) # type: ignore
|
|
return list(response.vectors.values())
|
|
|
|
@override
|
|
async def _inner_delete(self, keys: Sequence[TKey], **kwargs: Any) -> None:
|
|
"""Delete the records from the Pinecone collection."""
|
|
if not self.index_client:
|
|
await self._load_index_client()
|
|
if not self.index_client:
|
|
raise VectorStoreOperationException("Pinecone collection is not initialized.")
|
|
if "namespace" not in kwargs:
|
|
kwargs["namespace"] = self.namespace
|
|
if isinstance(self.index_client, GRPCIndex):
|
|
self.index_client.delete(ids=keys, **kwargs) # type: ignore
|
|
else:
|
|
await self.index_client.delete(ids=keys, **kwargs) # type: ignore
|
|
|
|
@override
|
|
async def _inner_search(
|
|
self,
|
|
search_type: SearchType,
|
|
options: VectorSearchOptions,
|
|
values: Any | None = None,
|
|
vector: Sequence[float | int] | None = None,
|
|
**kwargs: Any,
|
|
) -> KernelSearchResults[VectorSearchResult[TModel]]:
|
|
"""Search the records in the Pinecone collection."""
|
|
if not self.index_client:
|
|
await self._load_index_client()
|
|
if not self.index_client:
|
|
raise VectorStoreOperationException("Pinecone collection is not initialized.")
|
|
if search_type != SearchType.VECTOR:
|
|
raise VectorStoreOperationException(f"Search type {search_type} is not supported by Pinecone.")
|
|
if "namespace" not in kwargs:
|
|
kwargs["namespace"] = self.namespace
|
|
vector_field = self.definition.try_get_vector_field(options.vector_property_name)
|
|
if not vector_field:
|
|
raise VectorStoreModelException(
|
|
f"Vector field '{options.vector_property_name}' not found in the data model definition."
|
|
)
|
|
filter = self._build_filter(options.filter)
|
|
# is embedded mode
|
|
if self.embed_settings is not None:
|
|
if not self.index_client or isinstance(self.index_client, GRPCIndex):
|
|
raise VectorStoreOperationException(
|
|
"Pinecone GRPC client does not support integrated embeddings. "
|
|
"Please use the Pinecone Asyncio client."
|
|
)
|
|
search_args = {
|
|
"query": {
|
|
"inputs": {"text": values},
|
|
"top_k": options.top,
|
|
},
|
|
"namespace": kwargs.get("namespace", self.namespace),
|
|
}
|
|
if filter:
|
|
search_args["query"]["filter"] = {"$and": filter} if isinstance(filter, list) else filter
|
|
search_results = await self.index_client.search_records(**search_args)
|
|
return KernelSearchResults(
|
|
results=self._get_vector_search_results_from_results(search_results.result.hits, options),
|
|
total_count=len(search_results.result.hits),
|
|
)
|
|
if not vector:
|
|
vector = await self._generate_vector_from_values(values, options)
|
|
if not vector:
|
|
raise VectorStoreOperationException("No vector found for the given values.")
|
|
search_args = {
|
|
"vector": vector,
|
|
"top_k": options.top,
|
|
"include_metadata": True,
|
|
"include_values": options.include_vectors,
|
|
"namespace": kwargs.get("namespace", self.namespace),
|
|
}
|
|
if filter:
|
|
search_args["filter"] = {"$and": filter} if isinstance(filter, list) else filter
|
|
query_results = self.index_client.query(**search_args)
|
|
if isawaitable(query_results):
|
|
query_results = await query_results # type: ignore
|
|
return KernelSearchResults(
|
|
results=self._get_vector_search_results_from_results(query_results.matches, options), # type: ignore
|
|
total_count=len(query_results.matches), # type: ignore
|
|
)
|
|
|
|
@override
|
|
def _lambda_parser(self, node: ast.AST) -> Any:
|
|
# Comparison operations
|
|
match node:
|
|
case ast.Compare():
|
|
if len(node.ops) > 1:
|
|
# Chain comparisons (e.g., 1 < x < 3) become $and of each comparison
|
|
values = []
|
|
for idx in range(len(node.ops)):
|
|
left = node.left if idx == 0 else node.comparators[idx - 1]
|
|
right = node.comparators[idx]
|
|
op = node.ops[idx]
|
|
values.append(self._lambda_parser(ast.Compare(left=left, ops=[op], comparators=[right])))
|
|
return {"$and": values}
|
|
left = self._lambda_parser(node.left)
|
|
right = self._lambda_parser(node.comparators[0])
|
|
op = node.ops[0]
|
|
match op:
|
|
case ast.In():
|
|
return {left: {"$in": right}}
|
|
case ast.NotIn():
|
|
return {left: {"$nin": right}}
|
|
case ast.Eq():
|
|
# Pinecone allows short form: {field: value}
|
|
return {left: right}
|
|
case ast.NotEq():
|
|
return {left: {"$ne": right}}
|
|
case ast.Gt():
|
|
return {left: {"$gt": right}}
|
|
case ast.GtE():
|
|
return {left: {"$gte": right}}
|
|
case ast.Lt():
|
|
return {left: {"$lt": right}}
|
|
case ast.LtE():
|
|
return {left: {"$lte": right}}
|
|
raise NotImplementedError(f"Unsupported operator: {type(op)}")
|
|
case ast.BoolOp():
|
|
op = node.op # type: ignore
|
|
values = [self._lambda_parser(v) for v in node.values]
|
|
if isinstance(op, ast.And):
|
|
return {"$and": values}
|
|
if isinstance(op, ast.Or):
|
|
return {"$or": values}
|
|
raise NotImplementedError(f"Unsupported BoolOp: {type(op)}")
|
|
case ast.UnaryOp():
|
|
match node.op:
|
|
case ast.Not():
|
|
operand = self._lambda_parser(node.operand)
|
|
# Pinecone only supports $not over $in (becomes $nin)
|
|
if (
|
|
isinstance(operand, dict)
|
|
and len(operand) == 1
|
|
and isinstance(next(operand.values()), dict) # type: ignore
|
|
and "$in" in next(operand.values()) # type: ignore
|
|
):
|
|
field = next(operand.keys()) # type: ignore
|
|
values = next(operand.values())["$in"] # type: ignore
|
|
return {field: {"$nin": values}}
|
|
raise NotImplementedError(
|
|
"$not is only supported over $in (i.e., for ![...].contains(field)). "
|
|
"Other NOT expressions are not supported by Pinecone."
|
|
)
|
|
case ast.UAdd() | ast.USub() | ast.Invert():
|
|
raise NotImplementedError("Unary +, -, ~ are not supported in Pinecone filters.")
|
|
case ast.Attribute():
|
|
# Only allow attributes that are in the data model
|
|
if node.attr not in self.definition.storage_names:
|
|
raise VectorStoreOperationException(
|
|
f"Field '{node.attr}' not in data model (storage property names are used)."
|
|
)
|
|
return node.attr
|
|
case ast.Name():
|
|
# Only allow names that are in the data model
|
|
if node.id not in self.definition.storage_names:
|
|
raise VectorStoreOperationException(
|
|
f"Field '{node.id}' not in data model (storage property names are used)."
|
|
)
|
|
return node.id
|
|
case ast.Constant():
|
|
if node.value is None:
|
|
raise NotImplementedError("Pinecone does not support null checks in vector search pre-filters.")
|
|
return node.value
|
|
raise NotImplementedError(f"Unsupported AST node: {type(node)}")
|
|
|
|
@override
|
|
def _get_record_from_result(self, result: dict[str, Any]) -> dict[str, Any]:
|
|
if self.embed_settings is not None:
|
|
return {"_id": result["_id"], **result["fields"]}
|
|
return result
|
|
|
|
@override
|
|
def _get_score_from_result(self, result: Any) -> float | None:
|
|
if self.embed_settings is not None:
|
|
return result._score
|
|
return result.score
|
|
|
|
@override
|
|
async def __aexit__(self, exc_type, exc_value, traceback) -> None:
|
|
"""Exit the context manager."""
|
|
if self.index_client:
|
|
if isinstance(self.index_client, GRPCIndex):
|
|
self.index_client.close()
|
|
else:
|
|
await self.index_client.close()
|
|
self.index_client = None
|
|
if isinstance(self.client, PineconeAsyncio) and self.managed_client:
|
|
await self.client.close()
|
|
|
|
|
|
@release_candidate
|
|
class PineconeStore(VectorStore):
|
|
"""Pinecone Vector Store, for interacting with Pinecone collections."""
|
|
|
|
client: PineconeGRPC | PineconeAsyncio
|
|
|
|
def __init__(
|
|
self,
|
|
client: PineconeGRPC | PineconeAsyncio | None = None,
|
|
api_key: str | None = None,
|
|
embedding_generator: EmbeddingGeneratorBase | None = None,
|
|
env_file_path: str | None = None,
|
|
env_file_encoding: str | None = None,
|
|
use_grpc: bool = False,
|
|
**kwargs: Any,
|
|
) -> None:
|
|
"""Initialize the Pinecone store.
|
|
|
|
Args:
|
|
client: The Pinecone client to use. If not provided, a new client will be created.
|
|
api_key: The Pinecone API key. If not provided, it will be read from the environment.
|
|
embedding_generator: The embedding generator to use. If not provided, it will be read from the environment.
|
|
env_file_path: The path to the environment file. If not provided, it will be read from the default location.
|
|
env_file_encoding: The encoding of the environment file.
|
|
use_grpc: Whether to use the GRPC client or not. Default is False.
|
|
kwargs: Additional arguments to pass to the Pinecone client.
|
|
|
|
"""
|
|
managed_client = not client
|
|
if not client:
|
|
try:
|
|
settings = PineconeSettings(
|
|
api_key=api_key,
|
|
env_file_path=env_file_path,
|
|
env_file_encoding=env_file_encoding,
|
|
)
|
|
except ValidationError as exc:
|
|
raise VectorStoreInitializationException(f"Failed to create Pinecone settings: {exc}") from exc
|
|
|
|
if not client:
|
|
if use_grpc:
|
|
client = PineconeGRPC(
|
|
api_key=settings.api_key.get_secret_value(),
|
|
**kwargs,
|
|
)
|
|
else:
|
|
client = PineconeAsyncio(
|
|
api_key=settings.api_key.get_secret_value(),
|
|
**kwargs,
|
|
)
|
|
super().__init__(
|
|
client=client,
|
|
managed_client=managed_client,
|
|
embedding_generator=embedding_generator,
|
|
**kwargs,
|
|
)
|
|
|
|
@override
|
|
async def list_collection_names(self, **kwargs) -> Sequence[str]:
|
|
"""List the Pinecone collection names."""
|
|
if isinstance(self.client, PineconeGRPC):
|
|
return [idx.name for idx in self.client.list_indexes()]
|
|
return [idx.name for idx in await self.client.list_indexes()]
|
|
|
|
@override
|
|
def get_collection(
|
|
self,
|
|
record_type: type[TModel],
|
|
*,
|
|
definition: VectorStoreCollectionDefinition | None = None,
|
|
collection_name: str | None = None,
|
|
embedding_generator: EmbeddingGeneratorBase | None = None,
|
|
**kwargs: Any,
|
|
) -> PineconeCollection:
|
|
return PineconeCollection(
|
|
collection_name=collection_name,
|
|
record_type=record_type,
|
|
definition=definition,
|
|
client=self.client,
|
|
embedding_generator=embedding_generator or self.embedding_generator,
|
|
**kwargs,
|
|
)
|
|
|
|
@override
|
|
async def __aexit__(self, exc_type, exc_value, traceback) -> None:
|
|
"""Exit the context manager."""
|
|
if isinstance(self.client, PineconeAsyncio) and self.managed_client:
|
|
await self.client.close()
|