Skip to content

Commit

Permalink
Wire up plugin interface (#353)
Browse files Browse the repository at this point in the history
## 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.
  • Loading branch information
jhamon authored Jun 5, 2024
1 parent cab72a1 commit 701f6b6
Show file tree
Hide file tree
Showing 8 changed files with 191 additions and 13 deletions.
34 changes: 32 additions & 2 deletions pinecone/control/pinecone.py
Original file line number Diff line number Diff line change
@@ -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,
Expand All @@ -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__(
Expand Down Expand Up @@ -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,
Expand Down
8 changes: 7 additions & 1 deletion pinecone/data/index.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
2 changes: 1 addition & 1 deletion pinecone/utils/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
25 changes: 22 additions & 3 deletions pinecone/utils/setup_openapi_client.py
Original file line number Diff line number Diff line change
@@ -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
13 changes: 12 additions & 1 deletion poetry.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

1 change: 1 addition & 0 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
13 changes: 13 additions & 0 deletions tests/unit/test_control.py
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -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"
Expand Down
108 changes: 103 additions & 5 deletions tests/unit/utils/test_setup_openapi_client.py
Original file line number Diff line number Diff line change
@@ -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"

0 comments on commit 701f6b6

Please sign in to comment.