Inference class scans for plugins

Author: jhamon

pinecone/data/features/inference/inference.py

@@ -1,17 +1,22 @@
+import logging
 from typing import Optional, Dict, List, Union, Any
 
 from pinecone.openapi_support import ApiClient
 from pinecone.core.openapi.inference.apis import InferenceApi
 from pinecone.core.openapi.inference.models import EmbeddingsList, RerankResult
 from pinecone.core.openapi.inference import API_VERSION
-from pinecone.utils import setup_openapi_client
+from pinecone.utils import setup_openapi_client, build_plugin_setup_client
+
+from pinecone_plugin_interface import load_and_install as install_plugins
 
 from .inference_request_builder import (
     InferenceRequestBuilder,
     EmbedModel as EmbedModelEnum,
     RerankModel as RerankModelEnum,
 )
 
+logger = logging.getLogger(__name__)
+
 
 class Inference:
     """
@@ -27,6 +32,8 @@ class Inference:
 
     def __init__(self, config, openapi_config, **kwargs):
         self.config = config
+        self.openapi_config = openapi_config
+        self.pool_threads = kwargs.get("pool_threads", 1)
 
         self.__inference_api = setup_openapi_client(
             api_client_klass=ApiClient,
@@ -37,6 +44,23 @@ def __init__(self, config, openapi_config, **kwargs):
             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: {e}")
+
     def embed(
         self,
         model: Union[EmbedModel, str],