Typing improvements (#444)

## Problem

The more type hints we have in our package, the better people will be
able to understand how to use it.

## Solution

Much of the untyped code in the package is derived from old generated
code that we have extracted. So there's a fair amount of refactoring in
this difif to try to break out some smaller classes and functions with
seams that you can start to analyze and type; when everything is just
big mutable state blobs it is quite hard to reason about.

Along the way I uncovered that bulk import features were in a broken
state because some of those operation ids was modified since the last
release and these functions are not well-covered with automated tests.
This sort of thing really highlights why we need better type coverage in
the package.

## Type of Change

- [x] Bug fix (non-breaking change which fixes an issue)
- [x] None of the above: Refactoring to improve type safety

Author: jhamon

codegen/apis

@@ -34,7 +34,7 @@ update_templates_repo() {
 	echo "Updating templates repo"
 	pushd codegen/python-oas-templates
 		git fetch
-		git checkout main
+		git checkout jhamon/core-sdk
 		git pull
 	popd
 }
@@ -119,19 +119,6 @@ remove_shared_classes() {
 			rm "${source_directory}/${file}.py"
 		done
 	done
-
-	# Adjust import paths in every file
-	find "${destination}" -name "*.py" | while IFS= read -r file; do
-		sed -i '' "s/from \.\.model_utils/from pinecone\.openapi_support\.model_utils/g" "$file"
-
-		for module in "${modules[@]}"; do
-			sed -i '' "s/from pinecone\.$py_module_name\.openapi\.$module import rest/from pinecone\.openapi_support import rest/g" "$file"
-
-			for sharedFile in "${sharedFiles[@]}"; do
-				sed -i '' "s/from pinecone\.$py_module_name\.openapi\.$module\.$sharedFile/from pinecone\.openapi_support/g" "$file"
-			done
-		done
-	done
 }
 
 # Generated Python code attempts to internally map OpenAPI fields that begin
@@ -203,3 +190,5 @@ remove_shared_classes
 
 # Format generated files
 poetry run ruff format "${destination}"
+
+rm -rf "$build_dir"

codegen/build-oas.sh

@@ -0,0 +1,15 @@
+[mypy]
+; pretty = True
+; disallow_untyped_calls = True
+; check_untyped_defs = True
+; disallow_untyped_defs = True
+; warn_return_any = True
+; warn_unused_configs = True
+
+# Per-module options:
+
+; [mypy-mycode.foo.*]
+; disallow_untyped_defs = True
+
+[mypy-google.api.*]
+ignore_missing_imports = True

codegen/python-oas-templates

@@ -4,8 +4,17 @@
 from pinecone.exceptions.exceptions import PineconeConfigurationError
 from pinecone.config.openapi import OpenApiConfigFactory
 from pinecone.openapi_support.configuration import Configuration as OpenApiConfiguration
-from pinecone.utils import normalize_host
-from pinecone.utils.constants import SOURCE_TAG
+
+
+# Duplicated this util to help resolve circular imports
+def normalize_host(host: Optional[str]) -> str:
+    if host is None:
+        return ""
+    if host.startswith("https://"):
+        return host
+    if host.startswith("http://"):
+        return host
+    return "https://" + host
 
 
 class Config(NamedTuple):
@@ -50,7 +59,7 @@ def build(
         api_key = api_key or kwargs.pop("api_key", None) or os.getenv("PINECONE_API_KEY")
         host = host or kwargs.pop("host", None)
         host = normalize_host(host)
-        source_tag = kwargs.pop(SOURCE_TAG, None)
+        source_tag = kwargs.pop("source_tag", None)
 
         if not api_key:
             raise PineconeConfigurationError("You haven't specified an Api-Key.")

mypy.ini

@@ -18,7 +18,9 @@ def __call__(cls, *args, **kwargs):
 
 
 class IndexHostStore(metaclass=SingletonMeta):
-    def __init__(self):
+    _indexHosts: Dict[str, str]
+
+    def __init__(self) -> None:
         self._indexHosts = {}
 
     def _key(self, config: Config, index_name: str) -> str:

pinecone/config/config.py

@@ -1,6 +1,7 @@
 import time
 import logging
 from typing import Optional, Dict, Union
+from multiprocessing import cpu_count
 
 from .index_host_store import IndexHostStore
 from .pinecone_interface import PineconeDBControlInterface
@@ -60,7 +61,7 @@ def __init__(
         ssl_verify: Optional[bool] = None,
         config: Optional[Config] = None,
         additional_headers: Optional[Dict[str, str]] = {},
-        pool_threads: Optional[int] = 1,
+        pool_threads: Optional[int] = None,
         **kwargs,
     ):
         if config:
@@ -86,7 +87,11 @@ def __init__(
             )
 
         self.openapi_config = ConfigBuilder.build_openapi_config(self.config, **kwargs)
-        self.pool_threads = pool_threads
+
+        if pool_threads is None:
+            self.pool_threads = 5 * cpu_count()
+        else:
+            self.pool_threads = pool_threads
 
         self._inference = None  # Lazy initialization
 
@@ -102,7 +107,9 @@ def __init__(
         self.index_host_store = IndexHostStore()
         """ @private """
 
-        self.load_plugins()
+        self.load_plugins(
+            config=self.config, openapi_config=self.openapi_config, pool_threads=self.pool_threads
+        )
 
     @property
     def inference(self):
@@ -164,7 +171,7 @@ def create_index_for_model(
     def __poll_describe_index_until_ready(self, name: str, timeout: Optional[int] = None):
         description = None
 
-        def is_ready():
+        def is_ready() -> bool:
             nonlocal description
             description = self.describe_index(name=name)
             return description.status.ready
@@ -203,17 +210,14 @@ def delete_index(self, name: str, timeout: Optional[int] = None):
         self.index_api.delete_index(name)
         self.index_host_store.delete_host(self.config, name)
 
-        def get_remaining():
-            return name in self.list_indexes().names()
-
         if timeout == -1:
             return
 
         if timeout is None:
-            while get_remaining():
+            while self.has_index(name):
                 time.sleep(5)
         else:
-            while get_remaining() and timeout >= 0:
+            while self.has_index(name) and timeout >= 0:
                 time.sleep(5)
                 timeout -= 5
         if timeout and timeout < 0:

pinecone/control/index_host_store.py

@@ -2,14 +2,12 @@
 import asyncio
 from typing import Optional, Dict, Union
 
-from .index_host_store import IndexHostStore
-
-from pinecone.config import PineconeConfig, Config, ConfigBuilder
+from pinecone.config import PineconeConfig, ConfigBuilder
 
 from pinecone.core.openapi.db_control.api.manage_indexes_api import AsyncioManageIndexesApi
 from pinecone.openapi_support import AsyncioApiClient
 
-from pinecone.utils import normalize_host, setup_openapi_client
+from pinecone.utils import normalize_host, setup_async_openapi_client
 from pinecone.core.openapi.db_control import API_VERSION
 from pinecone.models import (
     ServerlessSpec,
@@ -56,52 +54,31 @@ def __init__(
         proxy_headers: Optional[Dict[str, str]] = None,
         ssl_ca_certs: Optional[str] = None,
         ssl_verify: Optional[bool] = None,
-        config: Optional[Config] = None,
         additional_headers: Optional[Dict[str, str]] = {},
         **kwargs,
     ):
-        if config:
-            if not isinstance(config, Config):
-                raise TypeError("config must be of type pinecone.config.Config")
-            else:
-                self.config = config
-        else:
-            self.config = PineconeConfig.build(
-                api_key=api_key,
-                host=host,
-                additional_headers=additional_headers,
-                proxy_url=proxy_url,
-                proxy_headers=proxy_headers,
-                ssl_ca_certs=ssl_ca_certs,
-                ssl_verify=ssl_verify,
-                **kwargs,
-            )
-
-        if kwargs.get("openapi_config", None):
-            raise Exception(
-                "Passing openapi_config is no longer supported. Please pass settings such as proxy_url, proxy_headers, ssl_ca_certs, and ssl_verify directly to the Pinecone constructor as keyword arguments. See the README at https://github.com/pinecone-io/pinecone-python-client for examples."
-            )
-
+        self.config = PineconeConfig.build(
+            api_key=api_key,
+            host=host,
+            additional_headers=additional_headers,
+            proxy_url=proxy_url,
+            proxy_headers=proxy_headers,
+            ssl_ca_certs=ssl_ca_certs,
+            ssl_verify=ssl_verify,
+            **kwargs,
+        )
         self.openapi_config = ConfigBuilder.build_openapi_config(self.config, **kwargs)
 
-        # I am pretty sure these aren't used, but need to pass a
-        # value for now to satisfy the method signatures
-        self.pool_threads = 1
-
         self._inference = None  # Lazy initialization
 
-        self.index_api = setup_openapi_client(
+        self.index_api = setup_async_openapi_client(
             api_client_klass=AsyncioApiClient,
             api_klass=AsyncioManageIndexesApi,
             config=self.config,
             openapi_config=self.openapi_config,
-            pool_threads=self.pool_threads,
             api_version=API_VERSION,
         )
 
-        self.index_host_store = IndexHostStore()
-        """ @private """
-
     async def __aenter__(self):
         return self
 
@@ -171,7 +148,7 @@ async def create_index_for_model(
     async def __poll_describe_index_until_ready(self, name: str, timeout: Optional[int] = None):
         description = None
 
-        async def is_ready():
+        async def is_ready() -> bool:
             nonlocal description
             description = await self.describe_index(name=name)
             return description.status.ready
@@ -208,20 +185,15 @@ async def is_ready():
 
     async def delete_index(self, name: str, timeout: Optional[int] = None):
         await self.index_api.delete_index(name)
-        self.index_host_store.delete_host(self.config, name)
-
-        async def get_remaining():
-            available_indexes = await self.list_indexes()
-            return name in available_indexes.names()
 
         if timeout == -1:
             return
 
         if timeout is None:
-            while await get_remaining():
+            while await self.has_index(name):
                 await asyncio.sleep(5)
         else:
-            while await get_remaining() and timeout >= 0:
+            while await self.has_index(name) and timeout >= 0:
                 await asyncio.sleep(5)
                 timeout -= 5
         if timeout and timeout < 0:
@@ -239,9 +211,6 @@ async def list_indexes(self) -> IndexList:
 
     async def describe_index(self, name: str) -> IndexModel:
         description = await self.index_api.describe_index(name)
-        host = description.host
-        self.index_host_store.set_host(self.config, name, host)
-
         return IndexModel(description)
 
     async def has_index(self, name: str) -> bool:
@@ -284,25 +253,17 @@ async def delete_collection(self, name: str):
     async def describe_collection(self, name: str):
         return await self.index_api.describe_collection(name).to_dict()
 
-    def Index(self, name: str = "", host: str = "", **kwargs):
-        if name == "" and host == "":
-            raise ValueError("Either name or host must be specified")
-
-        pt = kwargs.pop("pool_threads", None) or self.pool_threads
+    def Index(self, host: str, **kwargs):
         api_key = self.config.api_key
         openapi_config = self.openapi_config
 
-        if host != "":
-            # Use host url if it is provided
-            index_host = normalize_host(host)
-        else:
-            # Otherwise, get host url from describe_index using the index name
-            index_host = self.index_host_store.get_host(self.index_api, self.config, name)
+        index_host = normalize_host(host)
+        if index_host == "":
+            raise ValueError("host must be specified")
 
         return _AsyncioIndex(
             host=index_host,
             api_key=api_key,
-            pool_threads=pt,
             openapi_config=openapi_config,
             source_tag=self.config.source_tag,
             **kwargs,

pinecone/control/pinecone.py

@@ -3,8 +3,6 @@
 from typing import Optional, Dict, Union
 
 
-from pinecone.config import Config
-
 from pinecone.core.openapi.db_control.api.manage_indexes_api import ManageIndexesApi
 
 
@@ -39,7 +37,6 @@ def __init__(
         proxy_headers: Optional[Dict[str, str]] = None,
         ssl_ca_certs: Optional[str] = None,
         ssl_verify: Optional[bool] = None,
-        config: Optional[Config] = None,
         additional_headers: Optional[Dict[str, str]] = {},
         pool_threads: Optional[int] = 1,
         index_api: Optional[ManageIndexesApi] = None,
@@ -61,8 +58,6 @@ def __init__(
         :type ssl_ca_certs: str, optional
         :param ssl_verify: SSL verification is performed by default, but can be disabled using the boolean flag. Default: `True`
         :type ssl_verify: bool, optional
-        :param config: A `pinecone.config.Config` object. If passed, the `api_key` and `host` parameters will be ignored.
-        :type config: pinecone.config.Config, optional
         :param additional_headers: Additional headers to pass to the API. Default: `{}`
         :type additional_headers: Dict[str, str], optional
         :param pool_threads: The number of threads to use for the connection pool. Default: `1`

pinecone/control/pinecone_asyncio.py

@@ -529,7 +529,7 @@ async def describe_collection(self, name: str):
         pass
 
     @abstractmethod
-    async def Index(self, name: str = "", host: str = "", **kwargs):
+    async def Index(self, host: str, **kwargs):
         """
         Target an index for data operations.
 

pinecone/control/pinecone_interface.py

@@ -10,7 +10,7 @@ class VectorServiceStub(object):
     This service could also be called a `gRPC` service or a `REST`-like api.
     """
 
-    def __init__(self, channel):
+    def __init__(self, channel) -> None:
         """Constructor.
 
         Args:

pinecone/control/pinecone_interface_asyncio.py

@@ -5,18 +5,19 @@
 
 Pinecone is a vector database that makes it easy to search and retrieve billions of high-dimensional vectors.  # noqa: E501
 
+This file is @generated using OpenAPI.
+
 The version of the OpenAPI document: 2025-01
 Contact: [email protected]
-Generated by: https://openapi-generator.tech
 """
 
 __version__ = "1.0.0"
 
 # import ApiClient
-from pinecone.openapi_support import ApiClient
+from pinecone.openapi_support.api_client import ApiClient
 
 # import Configuration
-from pinecone.openapi_support import Configuration
+from pinecone.openapi_support.configuration import Configuration
 
 # import exceptions
 from pinecone.openapi_support.exceptions import PineconeException