Skip to content

Commit 36373a1

Browse files
Detect plugins for Index and IndexGRPC classes (#402)
## Problem Plugins not detected for Index/IndexGRPC classes ## Solution - Copy and modify installation method in pinecone.py used for `Pinecone` and `PineconeGRPC` classes. - For IndexGRPC, pass in a stub function. ## Type of Change - [ ] Bug fix (non-breaking change which fixes an issue) - [x] New feature (non-breaking change which adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to not work as expected) - [ ] This change requires a documentation update - [ ] Infrastructure change (CI configs, etc) - [ ] Non-code change (docs, etc) - [ ] None of the above: (explain here) ## Test Plan To really validate this will probably require making a dev build release and installing alongside some actual plugins. --------- Co-authored-by: Austin DeNoble <[email protected]>
1 parent efc751d commit 36373a1

File tree

3 files changed

+56
-10
lines changed

3 files changed

+56
-10
lines changed

pinecone/data/index.py

Lines changed: 36 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,11 @@
11
from tqdm.autonotebook import tqdm
22

3+
import logging
34
from typing import Union, List, Optional, Dict, Any
45

56
from pinecone.config import ConfigBuilder
67

78
from pinecone.core.openapi.shared import API_VERSION
8-
from pinecone.core.openapi.data.models import SparseValues
99
from pinecone.core.openapi.data import ApiClient
1010
from pinecone.core.openapi.data.models import (
1111
FetchResponse,
@@ -22,12 +22,22 @@
2222
UpdateRequest,
2323
DescribeIndexStatsRequest,
2424
ListResponse,
25+
SparseValues,
2526
)
26-
from .features.bulk_import import ImportFeatureMixin
2727
from pinecone.core.openapi.data.api.data_plane_api import DataPlaneApi
28-
from ..utils import setup_openapi_client, parse_non_empty_args
28+
from ..utils import (
29+
setup_openapi_client,
30+
parse_non_empty_args,
31+
build_plugin_setup_client,
32+
validate_and_convert_errors,
33+
)
34+
from .features.bulk_import import ImportFeatureMixin
2935
from .vector_factory import VectorFactory
3036

37+
from pinecone_plugin_interface import load_and_install as install_plugins
38+
39+
logger = logging.getLogger(__name__)
40+
3141
__all__ = [
3242
"Index",
3343
"FetchResponse",
@@ -47,8 +57,6 @@
4757
"SparseValues",
4858
]
4959

50-
from ..utils.error_handling import validate_and_convert_errors
51-
5260
_OPENAPI_ENDPOINT_PARAMS = (
5361
"_return_http_data_only",
5462
"_preload_content",
@@ -89,20 +97,38 @@ def __init__(
8997
**kwargs,
9098
)
9199

92-
self._config = ConfigBuilder.build(
100+
self.config = ConfigBuilder.build(
93101
api_key=api_key, host=host, additional_headers=additional_headers, **kwargs
94102
)
95-
openapi_config = ConfigBuilder.build_openapi_config(self._config, openapi_config)
103+
self._openapi_config = ConfigBuilder.build_openapi_config(self.config, openapi_config)
104+
self._pool_threads = pool_threads
96105

97106
self._vector_api = setup_openapi_client(
98107
api_client_klass=ApiClient,
99108
api_klass=DataPlaneApi,
100-
config=self._config,
101-
openapi_config=openapi_config,
102-
pool_threads=pool_threads,
109+
config=self.config,
110+
openapi_config=self._openapi_config,
111+
pool_threads=self._pool_threads,
103112
api_version=API_VERSION,
104113
)
105114

115+
self._load_plugins()
116+
117+
def _load_plugins(self):
118+
"""@private"""
119+
try:
120+
# I don't expect this to ever throw, but wrapping this in a
121+
# try block just in case to make sure a bad plugin doesn't
122+
# halt client initialization.
123+
openapi_client_builder = build_plugin_setup_client(
124+
config=self.config,
125+
openapi_config=self._openapi_config,
126+
pool_threads=self._pool_threads,
127+
)
128+
install_plugins(self, openapi_client_builder)
129+
except Exception as e:
130+
logger.error(f"Error loading plugins in Index: {e}")
131+
106132
def __enter__(self):
107133
return self
108134

pinecone/grpc/base.py

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
from abc import ABC, abstractmethod
22
from typing import Optional
33

4+
import logging
45
import grpc
56
from grpc._channel import Channel
67

@@ -10,6 +11,10 @@
1011
from .config import GRPCClientConfig
1112
from .grpc_runner import GrpcRunner
1213

14+
from pinecone_plugin_interface import load_and_install as install_plugins
15+
16+
_logger = logging.getLogger(__name__)
17+
1318

1419
class GRPCIndexBase(ABC):
1520
"""
@@ -40,6 +45,19 @@ def __init__(
4045
self._channel = channel or self._gen_channel()
4146
self.stub = self.stub_class(self._channel)
4247

48+
self._load_plugins()
49+
50+
def _load_plugins(self):
51+
"""@private"""
52+
try:
53+
54+
def stub_openapi_client_builder(*args, **kwargs):
55+
pass
56+
57+
install_plugins(self, stub_openapi_client_builder)
58+
except Exception as e:
59+
_logger.error(f"Error loading plugins in GRPCIndex: {e}")
60+
4361
@property
4462
@abstractmethod
4563
def stub_class(self):

pinecone/utils/__init__.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
from .parse_args import parse_non_empty_args
1010
from .docslinks import docslinks
1111
from .repr_overrides import install_json_repr_override
12+
from .error_handling import validate_and_convert_errors
1213

1314
__all__ = [
1415
"check_kwargs",
@@ -23,4 +24,5 @@
2324
"parse_non_empty_args",
2425
"docslinks",
2526
"install_json_repr_override",
27+
"validate_and_convert_errors",
2628
]

0 commit comments

Comments
 (0)