Skip to main content

Python SDK for Shift managed AI routing, telemetry, and local-first execution

Project description

shift-sdk

Python SDK for the Shift (Switch gateway) managed API.

Install

pip install shift-sdk

Import path remains:

from switch_sdk import SwitchClient

For local development:

pip install -e .[dev]

For local ExecuTorch runtime work:

pip install -e .[dev,local]

For packaging and publishing:

pip install -e .[publish]

Note: ExecuTorch wheels are not available for Python 3.14 yet. Use Python 3.10-3.13 (3.11 works well).

Required values

  • base_url: your gateway URL, for example http://localhost:8000
  • api_key: your plain project key (for example aura_...), not the SHA256 hash

Environment shortcuts are supported:

  • SHIFT_BASE_URL (fallback: SWITCH_BASE_URL)
  • SHIFT_API_KEY (fallbacks: SWITCH_API_KEY, API_KEY)

Quick start

import asyncio
from switch_sdk import SwitchClient, ChatMessage


async def main() -> None:
    async with SwitchClient.from_env() as client:
        completion = await client.chat(
            model="gpt-5",
            messages=[ChatMessage(role="user", content="Reply with: SDK_OK")],
            residency="US",
            sla="realtime",
            capability_flags={"force_cloud": True, "preferred_region": "eastus"},
        )

        print(completion.choices[0].message.content)
        print(completion.switch_meta["route"]["target"]["region"])


asyncio.run(main())

Set env vars first:

export SHIFT_BASE_URL=http://localhost:8000
export SHIFT_API_KEY=aura_your_plain_project_key

BYOK mode (prompt privacy)

Use chat_byok() when you want Shift to do routing/observability while your prompt is sent directly to Azure OpenAI with your own API key.

  • Shift sees: route request + telemetry metadata.
  • Shift does not see: prompt/messages payload.
import asyncio
from switch_sdk import AzureBYOKConfig, AzureRegionCredential, ChatMessage, SwitchClient


async def main() -> None:
    byok = AzureBYOKConfig(
        api_version="2025-01-01-preview",
        regions={
            "eastus": AzureRegionCredential(
                endpoint="https://shift-eastus.openai.azure.com",
                api_key="AZURE_EASTUS_KEY",
            ),
            "westus": AzureRegionCredential(
                endpoint="https://shift-westus.openai.azure.com",
                api_key="AZURE_WESTUS_KEY",
            ),
            "centralus": AzureRegionCredential(
                endpoint="https://shift-centralus.openai.azure.com",
                api_key="AZURE_CENTRALUS_KEY",
            ),
        },
    )

    async with SwitchClient.from_env(byok_azure=byok) as client:
        completion = await client.chat_byok(
            model="auto",
            messages=[ChatMessage(role="user", content="Reply exactly: BYOK_OK")],
            residency="US",
            capability_flags={"auto_model": True},
        )
        print(completion.choices[0].message.content)
        print(completion.switch_meta["route"]["target"]["region"])
        print(completion.switch_meta["resolved_model"])


asyncio.run(main())

Environment-based BYOK config is also supported:

export SHIFT_BYOK_AZURE_EASTUS_ENDPOINT=https://shift-eastus.openai.azure.com
export SHIFT_BYOK_AZURE_EASTUS_API_KEY=...
export SHIFT_BYOK_AZURE_WESTUS_ENDPOINT=https://shift-westus.openai.azure.com
export SHIFT_BYOK_AZURE_WESTUS_API_KEY=...
export SHIFT_BYOK_AZURE_CENTRALUS_ENDPOINT=https://shift-centralus.openai.azure.com
export SHIFT_BYOK_AZURE_CENTRALUS_API_KEY=...
export SHIFT_BYOK_AZURE_API_VERSION=2025-01-01-preview
async with SwitchClient.from_env(load_byok_azure_from_env=True) as client:
    completion = await client.chat_byok(
        model="auto",
        messages=[ChatMessage(role="user", content="Reply exactly: PRIVACY_OK")],
        capability_flags={"auto_model": True},
    )

Hybrid local-first mode (ExecuTorch-ready)

chat_hybrid() tries local execution first, then falls back to cloud when needed. Local models are cached on disk and downloaded only once per model version.

import asyncio
from switch_sdk import ChatMessage, LocalModelManager, SwitchClient


manifest = [
    {
        "model_id": "smollm2-135m",
        "task": "chat",
        "download_url": "https://your-model-host/smollm2-135m.pte",
        "sha256": "replace_with_sha256",
        "size_mb": 550,
        "min_ram_gb": 4,
        "max_prompt_chars": 280,
        "rank": 10,
    },
]


async def main() -> None:
    manager = LocalModelManager(cache_dir="~/.shift/models", manifest=manifest)
    # Optional: real ExecuTorch adapter (requires deps below)
    from switch_sdk import build_executorch_text_runtime
    local_runtime = build_executorch_text_runtime(
        tokenizer_source="HuggingFaceTB/SmolLM2-135M-Instruct",
        max_new_tokens=96,
        prefer_optimum=True,
    )

    async with SwitchClient(
        base_url="http://localhost:8000",
        api_key="aura_your_plain_project_key",
        local_model_manager=manager,
        local_runtime=local_runtime,
    ) as client:
        completion = await client.chat_hybrid(
            model="auto",
            messages=[ChatMessage(role="user", content="Reply exactly: LOCAL_OK")],
            capability_flags={"auto_model": True},
        )
        print(completion.model)
        print(completion.choices[0].message.content)
        print(completion.switch_meta)


asyncio.run(main())

Notes:

  • Default local runtime is a stub (for wiring/tests).
  • build_executorch_text_runtime(...) provides a real adapter that prefers Optimum ExecuTorch and falls back to raw executorch.runtime.
  • Cache path format: ~/.shift/models/<model_id>/<version>/model.pte
  • LRU eviction is applied when cache exceeds max_cache_gb.

Install local runtime dependencies:

pip install -e .[local]

Ready-made demo manifest:

  • /Users/proguy/Documents/projects/switch/switch-sdk/examples/local_manifest_smollm2_135m.json

Runtime callable contract:

from switch_sdk import ChatMessage, LocalModelHandle

async def my_executorch_runtime(messages: list[ChatMessage], handle: LocalModelHandle) -> str:
    # Load/use handle.path (.pte) with your ExecuTorch integration.
    # Return assistant text.
    return "LOCAL_EXECUTORCH_OK"

Routing-only call

decision = await client.route(
    model="gpt-5",
    residency="US",
    sla="realtime",
    capability_flags={"force_cloud": True},
)

print(decision.target.region)
print(decision.scores)
print(decision.candidate_breakdown)

Dashboard + carbon endpoints

summary = await client.get_dashboard_summary()
feed = await client.get_dashboard_feed(limit=20)
carbon = await client.get_live_carbon()

print(summary.summary.total_requests)
print(len(feed.items))
print(carbon.provider, carbon.regions.get("eastus"))

Custom telemetry event

from switch_sdk import TelemetryEvent

await client.track_event(
    TelemetryEvent(
        event_type="sdk_custom",
        request_id="custom-123",
        model="gpt-5",
        metadata={"feature": "my_feature"},
    )
)

await client.flush_telemetry()

Error handling

from switch_sdk import SwitchAPIError, SwitchNetworkError, SwitchTimeoutError

try:
    await client.route(model="gpt-5")
except SwitchAPIError as exc:
    print(exc.status_code, exc.detail)
except SwitchTimeoutError:
    print("Request timed out")
except SwitchNetworkError as exc:
    print(f"Network issue: {exc}")

Notes

  • The SDK is async-first.
  • Use async with SwitchClient(...) so telemetry flushes cleanly on exit.
  • Retries/backoff are built in for transient failures.
  • Telemetry is best-effort and never blocks successful chat/route calls.

Live switching checks

Automatic east/west region-switch verification script:

cd switch-sdk
.venv/bin/python examples/test_region_switching.py \
  --base-url http://localhost:8000 \
  --api-key aura_your_plain_project_key \
  --east-region eastus \
  --west-region westus \
  --central-region centralus \
  --check-chat

From-env example script

export SHIFT_BASE_URL=http://localhost:8000
export SHIFT_API_KEY=aura_your_plain_project_key
python /Users/proguy/Documents/projects/switch/switch-sdk/examples/test_from_env.py

Full user-journey script

python /Users/proguy/Documents/projects/switch/switch-sdk/examples/test_user_journey.py \
  --base-url http://localhost:8000 \
  --api-key aura_your_plain_project_key

Local ExecuTorch sanity check

Force local execution and fail if local runtime does not work:

cd /Users/proguy/Documents/projects/switch/switch-sdk
.venv311/bin/python examples/test_hybrid_local.py \
  --base-url http://localhost:8000 \
  --api-key dummy_local_only \
  --manifest-path examples/local_manifest_smollm2_135m.json \
  --executorch \
  --prefer-runtime \
  --tokenizer-source HuggingFaceTB/SmolLM2-135M-Instruct \
  --no-download \
  --no-cloud-fallback

Expected: output JSON includes "source": "sdk-local" in switch_meta.

Release

See /Users/proguy/Documents/projects/switch/switch-sdk/RELEASING.md for TestPyPI and PyPI release steps.

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

shift_sdk-0.3.2.tar.gz (32.7 kB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

shift_sdk-0.3.2-py3-none-any.whl (31.2 kB view details)

Uploaded Python 3

File details

Details for the file shift_sdk-0.3.2.tar.gz.

File metadata

  • Download URL: shift_sdk-0.3.2.tar.gz
  • Upload date:
  • Size: 32.7 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.2.0 CPython/3.11.13

File hashes

Hashes for shift_sdk-0.3.2.tar.gz
Algorithm Hash digest
SHA256 dc67cf1022e09a3de6f6dc88eec033cde938787668f248dc75468295c9d5b040
MD5 a0d62c91971b39ba2891930c9bc4d2cc
BLAKE2b-256 e6434ae213acc592341a7baecc5a9d533efc605ac2c0d28c0783edb1f8d6f142

See more details on using hashes here.

File details

Details for the file shift_sdk-0.3.2-py3-none-any.whl.

File metadata

  • Download URL: shift_sdk-0.3.2-py3-none-any.whl
  • Upload date:
  • Size: 31.2 kB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.2.0 CPython/3.11.13

File hashes

Hashes for shift_sdk-0.3.2-py3-none-any.whl
Algorithm Hash digest
SHA256 6c2a27670c0affbf7199743b5a01afc11aaf5163e8bdb0d918db479e0c16ab69
MD5 965adc6f765b24e0d8154bd910d99264
BLAKE2b-256 7516abdd21ca352bbed98c58b5cfbdaf188716b6b01dbc56ab74ba9f5bfde754

See more details on using hashes here.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page