diff --git a/pinecone/data/index.py b/pinecone/data/index.py index cc1bae6c..3f63d046 100644 --- a/pinecone/data/index.py +++ b/pinecone/data/index.py @@ -1,11 +1,11 @@ from tqdm.autonotebook import tqdm +import logging from typing import Union, List, Optional, Dict, Any from pinecone.config import ConfigBuilder from pinecone.core.openapi.shared import API_VERSION -from pinecone.core.openapi.data.models import SparseValues from pinecone.core.openapi.data import ApiClient from pinecone.core.openapi.data.models import ( FetchResponse, @@ -22,12 +22,22 @@ UpdateRequest, DescribeIndexStatsRequest, ListResponse, + SparseValues, ) -from .features.bulk_import import ImportFeatureMixin from pinecone.core.openapi.data.api.data_plane_api import DataPlaneApi -from ..utils import setup_openapi_client, parse_non_empty_args +from ..utils import ( + setup_openapi_client, + parse_non_empty_args, + build_plugin_setup_client, + validate_and_convert_errors, +) +from .features.bulk_import import ImportFeatureMixin from .vector_factory import VectorFactory +from pinecone_plugin_interface import load_and_install as install_plugins + +logger = logging.getLogger(__name__) + __all__ = [ "Index", "FetchResponse", @@ -47,8 +57,6 @@ "SparseValues", ] -from ..utils.error_handling import validate_and_convert_errors - _OPENAPI_ENDPOINT_PARAMS = ( "_return_http_data_only", "_preload_content", @@ -89,20 +97,38 @@ def __init__( **kwargs, ) - self._config = ConfigBuilder.build( + self.config = ConfigBuilder.build( api_key=api_key, host=host, additional_headers=additional_headers, **kwargs ) - openapi_config = ConfigBuilder.build_openapi_config(self._config, openapi_config) + self._openapi_config = ConfigBuilder.build_openapi_config(self.config, openapi_config) + self._pool_threads = pool_threads self._vector_api = setup_openapi_client( api_client_klass=ApiClient, api_klass=DataPlaneApi, - config=self._config, - openapi_config=openapi_config, - pool_threads=pool_threads, + config=self.config, + openapi_config=self._openapi_config, + pool_threads=self._pool_threads, api_version=API_VERSION, ) + self._load_plugins() + + def _load_plugins(self): + """@private""" + try: + # I don't expect this to ever throw, but wrapping this in a + # try block just in case to make sure a bad plugin doesn't + # halt client initialization. + openapi_client_builder = build_plugin_setup_client( + config=self.config, + openapi_config=self._openapi_config, + pool_threads=self._pool_threads, + ) + install_plugins(self, openapi_client_builder) + except Exception as e: + logger.error(f"Error loading plugins in Index: {e}") + def __enter__(self): return self diff --git a/pinecone/grpc/base.py b/pinecone/grpc/base.py index 17580d7e..cc4ca2c6 100644 --- a/pinecone/grpc/base.py +++ b/pinecone/grpc/base.py @@ -1,6 +1,7 @@ from abc import ABC, abstractmethod from typing import Optional +import logging import grpc from grpc._channel import Channel @@ -10,6 +11,10 @@ from .config import GRPCClientConfig from .grpc_runner import GrpcRunner +from pinecone_plugin_interface import load_and_install as install_plugins + +_logger = logging.getLogger(__name__) + class GRPCIndexBase(ABC): """ @@ -40,6 +45,19 @@ def __init__( self._channel = channel or self._gen_channel() self.stub = self.stub_class(self._channel) + self._load_plugins() + + def _load_plugins(self): + """@private""" + try: + + def stub_openapi_client_builder(*args, **kwargs): + pass + + install_plugins(self, stub_openapi_client_builder) + except Exception as e: + _logger.error(f"Error loading plugins in GRPCIndex: {e}") + @property @abstractmethod def stub_class(self): diff --git a/pinecone/utils/__init__.py b/pinecone/utils/__init__.py index 999daabf..34e1d5aa 100644 --- a/pinecone/utils/__init__.py +++ b/pinecone/utils/__init__.py @@ -9,6 +9,7 @@ from .parse_args import parse_non_empty_args from .docslinks import docslinks from .repr_overrides import install_json_repr_override +from .error_handling import validate_and_convert_errors __all__ = [ "check_kwargs", @@ -23,4 +24,5 @@ "parse_non_empty_args", "docslinks", "install_json_repr_override", + "validate_and_convert_errors", ]