Speed up client initialization with lazy loading

pinecone/__init__.py

@@ -2,22 +2,183 @@
 .. include:: ../pdoc/README.md
 """
 
-from .deprecated_plugins import check_for_deprecated_plugins
+from .deprecated_plugins import check_for_deprecated_plugins as _check_for_deprecated_plugins
 from .deprecation_warnings import *
-from .config import *
+from .pinecone import Pinecone
+from .pinecone_asyncio import PineconeAsyncio
 from .exceptions import *
-from .control import *
-from .data import *
-from .models import *
-from .enums import *
 
 from .utils import __version__
 
 import logging
 
+# Set up lazy import handling
+from .utils.lazy_imports import setup_lazy_imports as _setup_lazy_imports
+
+_inference_lazy_imports = {
+    "RerankModel": ("pinecone.inference", "RerankModel"),
+    "EmbedModel": ("pinecone.inference", "EmbedModel"),
+}
+
+_db_data_lazy_imports = {
+    "Vector": ("pinecone.db_data.dataclasses", "Vector"),
+    "SparseValues": ("pinecone.db_data.dataclasses", "SparseValues"),
+    "SearchQuery": ("pinecone.db_data.dataclasses", "SearchQuery"),
+    "SearchQueryVector": ("pinecone.db_data.dataclasses", "SearchQueryVector"),
+    "SearchRerank": ("pinecone.db_data.dataclasses", "SearchRerank"),
+    "FetchResponse": ("pinecone.db_data.models", "FetchResponse"),
+    "DeleteRequest": ("pinecone.db_data.models", "DeleteRequest"),
+    "DescribeIndexStatsRequest": ("pinecone.db_data.models", "DescribeIndexStatsRequest"),
+    "DescribeIndexStatsResponse": ("pinecone.db_data.models", "IndexDescription"),
+    "RpcStatus": ("pinecone.db_data.models", "RpcStatus"),
+    "ScoredVector": ("pinecone.db_data.models", "ScoredVector"),
+    "SingleQueryResults": ("pinecone.db_data.models", "SingleQueryResults"),
+    "QueryRequest": ("pinecone.db_data.models", "QueryRequest"),
+    "QueryResponse": ("pinecone.db_data.models", "QueryResponse"),
+    "UpsertResponse": ("pinecone.db_data.models", "UpsertResponse"),
+    "UpdateRequest": ("pinecone.db_data.models", "UpdateRequest"),
+    "ImportErrorMode": ("pinecone.core.openapi.db_data.model", "ImportErrorMode"),
+    "VectorDictionaryMissingKeysError": (
+        "pinecone.db_data.errors",
+        "VectorDictionaryMissingKeysError",
+    ),
+    "VectorDictionaryExcessKeysError": (
+        "pinecone.db_data.errors",
+        "VectorDictionaryExcessKeysError",
+    ),
+    "VectorTupleLengthError": ("pinecone.db_data.errors", "VectorTupleLengthError"),
+    "SparseValuesTypeError": ("pinecone.db_data.errors", "SparseValuesTypeError"),
+    "SparseValuesMissingKeysError": ("pinecone.db_data.errors", "SparseValuesMissingKeysError"),
+    "SparseValuesDictionaryExpectedError": (
+        "pinecone.db_data.errors",
+        "SparseValuesDictionaryExpectedError",
+    ),
+}
+
+_db_control_lazy_imports = {
+    "CloudProvider": ("pinecone.db_control.enums", "CloudProvider"),
+    "AwsRegion": ("pinecone.db_control.enums", "AwsRegion"),
+    "GcpRegion": ("pinecone.db_control.enums", "GcpRegion"),
+    "AzureRegion": ("pinecone.db_control.enums", "AzureRegion"),
+    "PodIndexEnvironment": ("pinecone.db_control.enums", "PodIndexEnvironment"),
+    "Metric": ("pinecone.db_control.enums", "Metric"),
+    "VectorType": ("pinecone.db_control.enums", "VectorType"),
+    "DeletionProtection": ("pinecone.db_control.enums", "DeletionProtection"),
+    "CollectionDescription": ("pinecone.db_control.models", "CollectionDescription"),
+    "CollectionList": ("pinecone.db_control.models", "CollectionList"),
+    "IndexList": ("pinecone.db_control.models", "IndexList"),
+    "IndexModel": ("pinecone.db_control.models", "IndexModel"),
+    "IndexEmbed": ("pinecone.db_control.models", "IndexEmbed"),
+    "ServerlessSpec": ("pinecone.db_control.models", "ServerlessSpec"),
+    "ServerlessSpecDefinition": ("pinecone.db_control.models", "ServerlessSpecDefinition"),
+    "PodSpec": ("pinecone.db_control.models", "PodSpec"),
+    "PodSpecDefinition": ("pinecone.db_control.models", "PodSpecDefinition"),
+    "PodType": ("pinecone.db_control.enums", "PodType"),
+}
+
+_config_lazy_imports = {
+    "Config": ("pinecone.config", "Config"),
+    "ConfigBuilder": ("pinecone.config", "ConfigBuilder"),
+    "PineconeConfig": ("pinecone.config", "PineconeConfig"),
+}
+
+# Define imports to be lazily loaded
+_LAZY_IMPORTS = {
+    **_inference_lazy_imports,
+    **_db_data_lazy_imports,
+    **_db_control_lazy_imports,
+    **_config_lazy_imports,
+}
+
+# Set up the lazy import handler
+_setup_lazy_imports(_LAZY_IMPORTS)
+
 # Raise an exception if the user is attempting to use the SDK with
 # deprecated plugins installed in their project.
-check_for_deprecated_plugins()
+_check_for_deprecated_plugins()
 
 # Silence annoying log messages from the plugin interface
 logging.getLogger("pinecone_plugin_interface").setLevel(logging.CRITICAL)
+
+__all__ = [
+    "__version__",
+    # Deprecated top-levelfunctions
+    "init",
+    "create_index",
+    "delete_index",
+    "list_indexes",
+    "describe_index",
+    "configure_index",
+    "scale_index",
+    "create_collection",
+    "delete_collection",
+    "describe_collection",
+    "list_collections",
+    # Primary client classes
+    "Pinecone",
+    "PineconeAsyncio",
+    # Config classes
+    "Config",
+    "ConfigBuilder",
+    "PineconeConfig",
+    # OpenAPI classes
+    "CloudProvider",
+    "AwsRegion",
+    "GcpRegion",
+    "AzureRegion",
+    "PodIndexEnvironment",
+    "Metric",
+    "VectorType",
+    "DeletionProtection",
+    "CollectionDescription",
+    "CollectionList",
+    "IndexList",
+    "IndexModel",
+    "IndexEmbed",
+    "ImportErrorMode",
+    "ServerlessSpec",
+    "ServerlessSpecDefinition",
+    "PodSpec",
+    "PodSpecDefinition",
+    "PodType",
+    "Vector",
+    "FetchResponse",
+    "DeleteRequest",
+    "DescribeIndexStatsRequest",
+    "DescribeIndexStatsResponse",
+    "RpcStatus",
+    "ScoredVector",
+    "SingleQueryResults",
+    "QueryRequest",
+    "QueryResponse",
+    "SearchQuery",
+    "SearchQueryVector",
+    "SearchRerank",
+    "UpsertResponse",
+    "UpdateRequest",
+    "SparseValues",
+    # Inference classes
+    "RerankModel",
+    "EmbedModel",
+    # Exception classes
+    "PineconeException",
+    "PineconeApiException",
+    "PineconeConfigurationError",
+    "PineconeProtocolError",
+    "PineconeApiAttributeError",
+    "PineconeApiTypeError",
+    "PineconeApiValueError",
+    "PineconeApiKeyError",
+    "PineconeApiException",
+    "NotFoundException",
+    "UnauthorizedException",
+    "ForbiddenException",
+    "ServiceException",
+    "ListConversionException",
+    "VectorDictionaryMissingKeysError",
+    "VectorDictionaryExcessKeysError",
+    "VectorTupleLengthError",
+    "SparseValuesTypeError",
+    "SparseValuesMissingKeysError",
+    "SparseValuesDictionaryExpectedError",
+]

pinecone/config/__init__.py

@@ -2,6 +2,7 @@
 import os
 
 from .config import ConfigBuilder, Config
+from .openapi_configuration import Configuration as OpenApiConfiguration
 from .pinecone_config import PineconeConfig
 
 if os.getenv("PINECONE_DEBUG") is not None:

pinecone/config/config.py

@@ -1,9 +1,11 @@
-from typing import NamedTuple, Optional, Dict
+from typing import NamedTuple, Optional, Dict, TYPE_CHECKING
 import os
 
-from pinecone.exceptions.exceptions import PineconeConfigurationError
-from pinecone.config.openapi import OpenApiConfigFactory
-from pinecone.openapi_support.configuration import Configuration as OpenApiConfiguration
+from pinecone.exceptions import PineconeConfigurationError
+from pinecone.config.openapi_config_factory import OpenApiConfigFactory
+
+if TYPE_CHECKING:
+    from pinecone.config.openapi_configuration import Configuration as OpenApiConfiguration
 
 
 # Duplicated this util to help resolve circular imports
@@ -81,8 +83,8 @@ def build(
 
     @staticmethod
     def build_openapi_config(
-        config: Config, openapi_config: Optional[OpenApiConfiguration] = None, **kwargs
-    ) -> OpenApiConfiguration:
+        config: Config, openapi_config: Optional["OpenApiConfiguration"] = None, **kwargs
+    ) -> "OpenApiConfiguration":
         if openapi_config:
             openapi_config = OpenApiConfigFactory.copy(
                 openapi_config=openapi_config, api_key=config.api_key, host=config.host

pinecone/config/openapi.py → pinecone/config/openapi_config_factory.py

@@ -1,13 +1,11 @@
 import sys
-from typing import List, Optional
+from typing import List, Optional, Tuple
 
 import certifi
 import socket
 import copy
 
-from urllib3.connection import HTTPConnection
-
-from pinecone.openapi_support.configuration import Configuration as OpenApiConfiguration
+from pinecone.config.openapi_configuration import Configuration as OpenApiConfiguration
 
 TCP_KEEPINTVL = 60  # Sec
 TCP_KEEPIDLE = 300  # Sec
@@ -58,7 +56,7 @@ def _get_socket_options(
         keep_alive_idle_sec: int = TCP_KEEPIDLE,
         keep_alive_interval_sec: int = TCP_KEEPINTVL,
         keep_alive_tries: int = TCP_KEEPCNT,
-    ) -> List[tuple]:
+    ) -> List[Tuple[int, int, int]]:
         """
         Returns the socket options to pass to OpenAPI's Rest client
         Args:
@@ -72,7 +70,8 @@ def _get_socket_options(
         """
         # Source: https://www.finbourne.com/blog/the-mysterious-hanging-client-tcp-keep-alives
 
-        socket_params = HTTPConnection.default_socket_options
+        # urllib3.connection.HTTPConnection.default_socket_options
+        socket_params = [(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)]
         if not do_keep_alive:
             return socket_params