Wire up plugin interface

pinecone/control/pinecone.py

@@ -1,13 +1,16 @@
 import time
 import warnings
+import logging
 from typing import Optional, Dict, Any, Union, List, cast, NamedTuple
 
 from .index_host_store import IndexHostStore
 
 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,

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

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

pinecone/utils/setup_openapi_client.py

@@ -1,14 +1,33 @@
-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
     )
     api_client.user_agent = get_user_agent(config)
     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

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

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"

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"