Wire up plugin interface (#353)

## Problem

We need a way to develop, test, demo, version and release features when
they are in a preview or pre-release state independent from the stable
core sdk. New features should only enter the sdk core when rapid
refinement and iteration has tapered off and we're ready to stabilize
the feature. This will help us innovate and ship faster.

## Solution

We have implemented a plugin interface that allows us to extend the
`Pinecone` class with experimental features defined elsewhere. For now,
this plugin interface is only intended for internal use at Pinecone.

There are three pieces involved here: the changes in this diff, a small
separate package called `pinecone_plugin_interface`, and the plugin
implementations themself. Currently, there are no publicly available
plugins.

The `Pinecone` class here imports a `load_and_install` function from the
`pinecone_plugin_interface` package. When invoked, this function scans
the user's environment looking for packages that have been implemented
within a namespace package called `pinecone_plugins`. If it finds any,
it will use the `client_builder` function supplied by the `Pinecone`
class to pass user configuration into the plugin for consistent handling
of configurations such as api keys, proxy settings, etc.

Although the amount of code involved in `pinecone_plugin_interface` is
quite meager, externalizing it should give some flexibility to share
code between the sdk and plugin implementations without creating a
circular dependency should we ever wish to include a plugin as a
dependency of the sdk (after functionality has stabilized).

I have been defensive here by wrapping the plugin installation in a
try/catch. I want to be confident that the mere presence of an invalid
or broken plugin cannot break the rest of the sdk.

## Type of Change

- [x] New feature (non-breaking change which adds functionality)

## Test Plan

Tests should be green. I added some unit tests covering the new builder
function stuff. As far as integration testing, I can't really add those
here until we have publicly available plugins. But I've done manual
testing to gain confidence.

Author: jhamon

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

poetry.lock

@@ -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

pyproject.toml

@@ -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/test_control.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"