Skip to content

Commit

Permalink
Detect plugins for Index and IndexGRPC classes (#402)
Browse files Browse the repository at this point in the history
## 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]>
  • Loading branch information
jhamon and austin-denoble authored Oct 28, 2024
1 parent efc751d commit 36373a1
Show file tree
Hide file tree
Showing 3 changed files with 56 additions and 10 deletions.
46 changes: 36 additions & 10 deletions pinecone/data/index.py
Original file line number Diff line number Diff line change
@@ -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,
Expand All @@ -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",
Expand All @@ -47,8 +57,6 @@
"SparseValues",
]

from ..utils.error_handling import validate_and_convert_errors

_OPENAPI_ENDPOINT_PARAMS = (
"_return_http_data_only",
"_preload_content",
Expand Down Expand Up @@ -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

Expand Down
18 changes: 18 additions & 0 deletions pinecone/grpc/base.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
from abc import ABC, abstractmethod
from typing import Optional

import logging
import grpc
from grpc._channel import Channel

Expand All @@ -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):
"""
Expand Down Expand Up @@ -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):
Expand Down
2 changes: 2 additions & 0 deletions pinecone/utils/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand All @@ -23,4 +24,5 @@
"parse_non_empty_args",
"docslinks",
"install_json_repr_override",
"validate_and_convert_errors",
]

0 comments on commit 36373a1

Please sign in to comment.