shift-sdk 0.3.4

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

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

• api_key: your plain project key (for example aura_...), not the SHA256 hash
• base_url is optional. If omitted, SDK uses: explicit constructor value -> env URL -> SDK default gateway URL.

Environment shortcuts are supported:

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

Quick start

import asyncio
import os
from switch_sdk import SwitchClient, ChatMessage

API_KEY = os.environ["SHIFT_API_KEY"]

async def main() -> None:
    async with SwitchClient(api_key=API_KEY) as client:
        completion = await client.chat(
            model="auto",
            messages=[ChatMessage(role="user", content="Reply with: SDK_OK")],
            residency="US",
            sla="realtime",
            capability_flags={"auto_model": True},
        )

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


asyncio.run(main())

Set env vars first:

export SHIFT_API_KEY=aura_your_plain_project_key

Optional URL override:

export SHIFT_BASE_URL=http://localhost:8000

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"))

Model garden endpoints

Use these when you want to read or toggle which routed targets are enabled for the current org/project.

garden = await client.get_model_garden()
for target in garden.targets:
    print(target.id, target.model, target.region, target.is_enabled)

# Disable one target
updated = await client.set_model_garden_target_enabled(
    target_id="target-uuid",
    is_enabled=False,
)
print(updated.id, updated.is_enabled)

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

Model garden script

python /Users/proguy/Documents/projects/switch/switch-sdk/examples/test_model_garden.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.