diff --git a/pinecone/control/pinecone.py b/pinecone/control/pinecone.py index 2de6cd4e..83b8384e 100644 --- a/pinecone/control/pinecone.py +++ b/pinecone/control/pinecone.py @@ -1,5 +1,6 @@ import time import warnings +import logging from typing import Optional, Dict, Any, Union, List, cast, NamedTuple from .index_host_store import IndexHostStore @@ -7,7 +8,9 @@ from pinecone.config import PineconeConfig, Config, ConfigBuilder from pinecone.core.client.api.manage_indexes_api import ManageIndexesApi -from pinecone.utils import normalize_host, setup_openapi_client +from pinecone.core.client.api_client import ApiClient + +from pinecone.utils import normalize_host, setup_openapi_client, build_plugin_setup_client from pinecone.core.client.models import ( CreateCollectionRequest, CreateIndexRequest, @@ -20,6 +23,10 @@ from pinecone.data import Index +from pinecone_plugin_interface import load_and_install as install_plugins + +logger = logging.getLogger(__name__) + class Pinecone: def __init__( @@ -203,11 +210,34 @@ def __init__( if index_api: self.index_api = index_api else: - self.index_api = setup_openapi_client(ManageIndexesApi, self.config, self.openapi_config, pool_threads) + self.index_api = setup_openapi_client( + api_client_klass=ApiClient, + api_klass=ManageIndexesApi, + config=self.config, + openapi_config=self.openapi_config, + pool_threads=pool_threads + ) self.index_host_store = IndexHostStore() """ @private """ + 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: {e}") + def create_index( self, name: str, diff --git a/pinecone/data/index.py b/pinecone/data/index.py index 8596c7eb..d613996e 100644 --- a/pinecone/data/index.py +++ b/pinecone/data/index.py @@ -86,7 +86,13 @@ def __init__( ) openapi_config = ConfigBuilder.build_openapi_config(self._config, openapi_config) - self._vector_api = setup_openapi_client(DataPlaneApi, self._config, openapi_config, 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 + ) def __enter__(self): return self diff --git a/pinecone/utils/__init__.py b/pinecone/utils/__init__.py index 56599293..e72df335 100644 --- a/pinecone/utils/__init__.py +++ b/pinecone/utils/__init__.py @@ -5,5 +5,5 @@ from .fix_tuple_length import fix_tuple_length from .convert_to_list import convert_to_list from .normalize_host import normalize_host -from .setup_openapi_client import setup_openapi_client +from .setup_openapi_client import setup_openapi_client, build_plugin_setup_client from .docslinks import docslinks \ No newline at end of file diff --git a/pinecone/utils/setup_openapi_client.py b/pinecone/utils/setup_openapi_client.py index 02e0b14f..f1d33424 100644 --- a/pinecone/utils/setup_openapi_client.py +++ b/pinecone/utils/setup_openapi_client.py @@ -1,8 +1,19 @@ -from pinecone.core.client.api_client import ApiClient from .user_agent import get_user_agent +import copy -def setup_openapi_client(api_klass, config, openapi_config, pool_threads): - api_client = ApiClient( +def setup_openapi_client(api_client_klass, api_klass, config, openapi_config, pool_threads, api_version=None, **kwargs): + # It is important that we allow the user to pass in a reference to api_client_klass + # instead of creating a direct dependency on ApiClient because plugins have their + # own ApiClient implementations. Even if those implementations seem like they should + # be functionally identical, they are not the same class and have references to + # different copies of the ModelNormal class. Therefore cannot be used interchangeably. + # without breaking the generated client code anywhere it is relying on isinstance to make + # a decision about something. + if kwargs.get("host"): + openapi_config = copy.deepcopy(openapi_config) + openapi_config._base_path = kwargs['host'] + + api_client = api_client_klass( configuration=openapi_config, pool_threads=pool_threads ) @@ -10,5 +21,13 @@ def setup_openapi_client(api_klass, config, openapi_config, pool_threads): extra_headers = config.additional_headers or {} for key, value in extra_headers.items(): api_client.set_default_header(key, value) + + if api_version: + api_client.set_default_header("X-Pinecone-API-Version", api_version) client = api_klass(api_client) return client + +def build_plugin_setup_client(config, openapi_config, pool_threads): + def setup_plugin_client(api_client_klass, api_klass, api_version, **kwargs): + return setup_openapi_client(api_client_klass, api_klass, config, openapi_config, pool_threads, api_version, **kwargs) + return setup_plugin_client diff --git a/poetry.lock b/poetry.lock index ed05c50f..925445ee 100644 --- a/poetry.lock +++ b/poetry.lock @@ -1,4 +1,4 @@ -# This file is automatically @generated by Poetry 1.8.2 and should not be changed by hand. +# This file is automatically @generated by Poetry 1.6.1 and should not be changed by hand. [[package]] name = "astunparse" @@ -734,6 +734,17 @@ pygments = ">=2.12.0" [package.extras] dev = ["black", "hypothesis", "mypy", "pygments (>=2.14.0)", "pytest", "pytest-cov", "pytest-timeout", "ruff", "tox", "types-pygments"] +[[package]] +name = "pinecone-plugin-interface" +version = "0.0.7" +description = "Plugin interface for the Pinecone python client" +optional = false +python-versions = "<4.0,>=3.8" +files = [ + {file = "pinecone_plugin_interface-0.0.7-py3-none-any.whl", hash = "sha256:875857ad9c9fc8bbc074dbe780d187a2afd21f5bfe0f3b08601924a61ef1bba8"}, + {file = "pinecone_plugin_interface-0.0.7.tar.gz", hash = "sha256:b8e6675e41847333aa13923cc44daa3f85676d7157324682dc1640588a982846"}, +] + [[package]] name = "pluggy" version = "1.3.0" @@ -1170,4 +1181,4 @@ grpc = ["googleapis-common-protos", "grpcio", "grpcio", "lz4", "protobuf", "prot [metadata] lock-version = "2.0" python-versions = "^3.8" -content-hash = "4dd1a293f07b49e250679ca58f6ab3c5a090f2d4807d721dab64d5d574a38c10" +content-hash = "3f06d23e45560281fcd773e7bde9a5d3be62db0de27456d8b8ca332187e8e031" diff --git a/pyproject.toml b/pyproject.toml index 2071996e..ee663154 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -67,6 +67,7 @@ googleapis-common-protos = { version = ">=1.53.0", optional = true } lz4 = { version = ">=3.1.3", optional = true } protobuf = { version = "^4.25", optional = true } protoc-gen-openapiv2 = {version = "^0.0.1", optional = true } +pinecone-plugin-interface = "^0.0.7" [tool.poetry.group.types] optional = true diff --git a/tests/unit/test_control.py b/tests/unit/test_control.py index 7e4bd6e1..d7186cdf 100644 --- a/tests/unit/test_control.py +++ b/tests/unit/test_control.py @@ -1,5 +1,6 @@ import pytest import re +from unittest.mock import patch from pinecone import ConfigBuilder, Pinecone, PodSpec, ServerlessSpec from pinecone.core.client.models import IndexList, IndexModel from pinecone.core.client.api.manage_indexes_api import ManageIndexesApi @@ -17,6 +18,18 @@ def index_list_response(): ]) class TestControl: + def test_plugins_are_installed(self): + with patch('pinecone.control.pinecone.install_plugins') as mock_install_plugins: + p = Pinecone(api_key='asdf') + mock_install_plugins.assert_called_once() + + def test_bad_plugin_doesnt_break_sdk(self): + with patch('pinecone.control.pinecone.install_plugins', side_effect=Exception("bad plugin")): + try: + p = Pinecone(api_key='asdf') + except Exception as e: + assert False, f"Unexpected exception: {e}" + def test_default_host(self): p = Pinecone(api_key="123-456-789") assert p.index_api.api_client.configuration.host == "https://api.pinecone.io" diff --git a/tests/unit/utils/test_setup_openapi_client.py b/tests/unit/utils/test_setup_openapi_client.py index 0c44688a..9c4f355a 100644 --- a/tests/unit/utils/test_setup_openapi_client.py +++ b/tests/unit/utils/test_setup_openapi_client.py @@ -1,11 +1,109 @@ +import pytest import re from pinecone.config import ConfigBuilder from pinecone.core.client.api.manage_indexes_api import ManageIndexesApi -from pinecone.utils.setup_openapi_client import setup_openapi_client +from pinecone.core.client.api_client import ApiClient +from pinecone.utils.setup_openapi_client import setup_openapi_client, build_plugin_setup_client class TestSetupOpenAPIClient(): def test_setup_openapi_client(self): - "" - # config = ConfigBuilder.build(api_key="my-api-key", host="https://my-controller-host") - # api_client = setup_openapi_client(ManageIndexesApi, config=config, pool_threads=2) - # # assert api_client.user_agent == "pinecone-python-client/0.0.1" + config = ConfigBuilder.build( + api_key="my-api-key", + host="https://my-controller-host" + ) + openapi_config = ConfigBuilder.build_openapi_config(config) + assert openapi_config.host == "https://my-controller-host" + + control_plane_client = setup_openapi_client(ApiClient, ManageIndexesApi, config=config, openapi_config=openapi_config, pool_threads=2) + user_agent_regex = re.compile(r"python-client-\d+\.\d+\.\d+ \(urllib3\:\d+\.\d+\.\d+\)") + assert re.match(user_agent_regex, control_plane_client.api_client.user_agent) + assert re.match(user_agent_regex, control_plane_client.api_client.default_headers['User-Agent']) + + def test_setup_openapi_client_with_api_version(self): + config = ConfigBuilder.build( + api_key="my-api-key", + host="https://my-controller-host", + ) + openapi_config = ConfigBuilder.build_openapi_config(config) + assert openapi_config.host == "https://my-controller-host" + + control_plane_client = setup_openapi_client(ApiClient, ManageIndexesApi, config=config, openapi_config=openapi_config, pool_threads=2, api_version="2024-04") + user_agent_regex = re.compile(r"python-client-\d+\.\d+\.\d+ \(urllib3\:\d+\.\d+\.\d+\)") + assert re.match(user_agent_regex, control_plane_client.api_client.user_agent) + assert re.match(user_agent_regex, control_plane_client.api_client.default_headers['User-Agent']) + assert control_plane_client.api_client.default_headers['X-Pinecone-API-Version'] == "2024-04" + + +class TestBuildPluginSetupClient(): + @pytest.mark.parametrize("plugin_api_version,plugin_host", [ + (None, None), + ("2024-07", "https://my-plugin-host") + ]) + def test_setup_openapi_client_with_host_override(self, plugin_api_version, plugin_host): + # These configurations represent the configurations that the core sdk + # (e.g. Pinecone class) will have built prior to invoking the plugin setup. + # In real usage, this takes place during the Pinecone class initialization + # and pulls together configuration from all sources (kwargs and env vars). + # It reflects a merging of the user's configuration and the defaults set + # by the sdk. + config = ConfigBuilder.build( + api_key="my-api-key", + host="https://api.pinecone.io", + source_tag="my_source_tag", + proxy_url="http://my-proxy.com", + ssl_ca_certs="path/to/bundle.pem" + ) + openapi_config = ConfigBuilder.build_openapi_config(config) + + # The core sdk (e.g. Pinecone class) will be responsible for invoking the + # build_plugin_setup_client method before passing the result to the plugin + # install method. This is + # somewhat like currying the openapi setup function, because we want some + # information to be controled by the core sdk (e.g. the user-agent string, + # proxy settings, etc) while allowing the plugin to pass the parts of the + # configuration that are relevant to it such as api version, base url if + # served from somewhere besides api.pinecone.io, etc. + client_builder = build_plugin_setup_client(config=config, openapi_config=openapi_config, pool_threads=2) + + # The plugin machinery in pinecone_plugin_interface will be the one to call + # this client_builder function using classes and other config it discovers inside the + # pinecone_plugin namespace package. Putting plugin configuration and references + # to the implementation classes into a spot where the pinecone_plugin_interface + # can find them is the responsibility of the plugin developer. + # + # Passing ManagedIndexesApi and ApiClient here are just a standin for testing + # purposes; in a real plugin, the class would be something else related + # to a new feature, but to test that this setup works I just need a FooApi + # class generated off the openapi spec. + plugin_api=ManageIndexesApi + plugin_client = client_builder( + api_client_klass=ApiClient, + api_klass=plugin_api, + api_version=plugin_api_version, + host=plugin_host + ) + + # Returned client is an instance of the input class + assert isinstance(plugin_client, plugin_api) + + # We want requests from plugins to have a user-agent matching the host SDK. + user_agent_regex = re.compile(r"python-client-\d+\.\d+\.\d+ \(urllib3\:\d+\.\d+\.\d+\)") + assert re.match(user_agent_regex, plugin_client.api_client.user_agent) + assert re.match(user_agent_regex, plugin_client.api_client.default_headers['User-Agent']) + + # User agent still contains the source tag that was set in the sdk config + assert 'my_source_tag' in plugin_client.api_client.default_headers['User-Agent'] + + # Proxy settings should be passed from the core sdk to the plugin client + assert plugin_client.api_client.configuration.proxy == "http://my-proxy.com" + assert plugin_client.api_client.configuration.ssl_ca_cert == "path/to/bundle.pem" + + # Plugins need to be able to pass their own API version (optionally) + assert plugin_client.api_client.default_headers.get('X-Pinecone-API-Version') == plugin_api_version + + # Plugins need to be able to override the host (optionally) + if plugin_host: + assert plugin_client.api_client.configuration._base_path == plugin_host + else: + # When plugin does not set a host, it should default to the host set in the core sdk + assert plugin_client.api_client.configuration._base_path == "https://api.pinecone.io"