Speed up client initialization with lazy loading

.github/workflows/testing-integration.yaml

@@ -3,6 +3,31 @@ name: "Integration Tests"
   workflow_call: {}
 
 jobs:
+  reorg:
+    name: Reorg tests
+    runs-on: ubuntu-latest
+    env:
+      PINECONE_DEBUG_CURL: 'true'
+      PINECONE_API_KEY: '${{ secrets.PINECONE_API_KEY }}'
+      PINECONE_ADDITIONAL_HEADERS: '{"sdk-test-suite": "pinecone-python-client"}'
+    strategy:
+      matrix:
+        python_version: [3.9, 3.12]
+    steps:
+      - uses: actions/checkout@v4
+      - name: 'Set up Python ${{ matrix.python_version }}'
+        uses: actions/setup-python@v5
+        with:
+          python-version: '${{ matrix.python_version }}'
+      - name: Setup Poetry
+        uses: ./.github/actions/setup-poetry
+        with:
+          include_asyncio: true
+      - name: 'Run index tests'
+        run: poetry run pytest tests/integration/control/index --retries 5 --retry-delay 35 -s -vv --log-cli-level=DEBUG
+      - name: 'Run collection tests'
+        run: poetry run pytest tests/integration/control/collections --retries 5 --retry-delay 35 -s -vv --log-cli-level=DEBUG
+
 
   inference:
     name: Inference tests

pinecone/__init__.py

@@ -2,22 +2,143 @@
 .. 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.dataclasses", "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.models", "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",
+    ),
+    "Index": ("pinecone.db_data.import_error", "Index"),
+    "Inference": ("pinecone.db_data.import_error", "Inference"),
+}
+
+_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",
+    # All lazy-loaded types
+    *list(_LAZY_IMPORTS.keys()),
+    # Exception classes
+    "PineconeException",
+    "PineconeApiException",
+    "PineconeConfigurationError",
+    "PineconeProtocolError",
+    "PineconeApiAttributeError",
+    "PineconeApiTypeError",
+    "PineconeApiValueError",
+    "PineconeApiKeyError",
+    "NotFoundException",
+    "UnauthorizedException",
+    "ForbiddenException",
+    "ServiceException",
+    "ListConversionException",
+    "VectorDictionaryMissingKeysError",
+    "VectorDictionaryExcessKeysError",
+    "VectorTupleLengthError",
+    "SparseValuesTypeError",
+    "SparseValuesMissingKeysError",
+    "SparseValuesDictionaryExpectedError",
+]

pinecone/__init__.pyi

@@ -0,0 +1,118 @@
+from pinecone.config import Config
+from pinecone.config import ConfigBuilder
+from pinecone.config import PineconeConfig
+from pinecone.inference import RerankModel
+from pinecone.inference import EmbedModel
+from pinecone.db_data.dataclasses import (
+    Vector,
+    SparseValues,
+    SearchQuery,
+    SearchQueryVector,
+    SearchRerank,
+)
+from pinecone.db_data.models import (
+    FetchResponse,
+    DeleteRequest,
+    DescribeIndexStatsRequest,
+    IndexDescription as DescribeIndexStatsResponse,
+    RpcStatus,
+    ScoredVector,
+    SingleQueryResults,
+    QueryRequest,
+    QueryResponse,
+    UpsertResponse,
+    UpdateRequest,
+)
+from pinecone.core.openapi.db_data.models import ImportErrorMode
+from pinecone.db_data.errors import (
+    VectorDictionaryMissingKeysError,
+    VectorDictionaryExcessKeysError,
+    VectorTupleLengthError,
+    SparseValuesTypeError,
+    SparseValuesMissingKeysError,
+    SparseValuesDictionaryExpectedError,
+)
+from pinecone.db_control.enums import (
+    CloudProvider,
+    AwsRegion,
+    GcpRegion,
+    AzureRegion,
+    PodIndexEnvironment,
+    Metric,
+    VectorType,
+    DeletionProtection,
+    PodType,
+)
+from pinecone.db_control.models import (
+    CollectionDescription,
+    CollectionList,
+    IndexList,
+    IndexModel,
+    IndexEmbed,
+    ServerlessSpec,
+    ServerlessSpecDefinition,
+    PodSpec,
+    PodSpecDefinition,
+)
+from pinecone.pinecone import Pinecone
+from pinecone.pinecone_asyncio import PineconeAsyncio
+
+# Re-export all the types
+__all__ = [
+    # Primary client classes
+    "Pinecone",
+    "PineconeAsyncio",
+    # Config classes
+    "Config",
+    "ConfigBuilder",
+    "PineconeConfig",
+    # Inference classes
+    "RerankModel",
+    "EmbedModel",
+    # Data classes
+    "Vector",
+    "SparseValues",
+    "SearchQuery",
+    "SearchQueryVector",
+    "SearchRerank",
+    # Model classes
+    "FetchResponse",
+    "DeleteRequest",
+    "DescribeIndexStatsRequest",
+    "DescribeIndexStatsResponse",
+    "RpcStatus",
+    "ScoredVector",
+    "SingleQueryResults",
+    "QueryRequest",
+    "QueryResponse",
+    "UpsertResponse",
+    "UpdateRequest",
+    "ImportErrorMode",
+    # Error classes
+    "VectorDictionaryMissingKeysError",
+    "VectorDictionaryExcessKeysError",
+    "VectorTupleLengthError",
+    "SparseValuesTypeError",
+    "SparseValuesMissingKeysError",
+    "SparseValuesDictionaryExpectedError",
+    # Control plane enums
+    "CloudProvider",
+    "AwsRegion",
+    "GcpRegion",
+    "AzureRegion",
+    "PodIndexEnvironment",
+    "Metric",
+    "VectorType",
+    "DeletionProtection",
+    "PodType",
+    # Control plane models
+    "CollectionDescription",
+    "CollectionList",
+    "IndexList",
+    "IndexModel",
+    "IndexEmbed",
+    "ServerlessSpec",
+    "ServerlessSpecDefinition",
+    "PodSpec",
+    "PodSpecDefinition",
+]

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