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

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.

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.3.tar.gz (34.3 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.3-py3-none-any.whl (32.2 kB view details)

Uploaded Python 3

File details

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

File metadata

  • Download URL: shift_sdk-0.3.3.tar.gz
  • Upload date:
  • Size: 34.3 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.3.tar.gz
Algorithm Hash digest
SHA256 b869b305a649a9c5a5f622766489f5217ebfcd3cacc63ae26e703456120fdc0e
MD5 44158c55e5fcee82ec13db322bb0cb72
BLAKE2b-256 851011915e362e4d2a36ec4fd2033f698c602b9a9f6e15b7080c69aa01ccda4c

See more details on using hashes here.

File details

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

File metadata

  • Download URL: shift_sdk-0.3.3-py3-none-any.whl
  • Upload date:
  • Size: 32.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.3-py3-none-any.whl
Algorithm Hash digest
SHA256 0478c1f311cf88c3760a410977a9b5f6139c2cb0c7695dd8ace542978caf0fe9
MD5 25aa12c77fb1574b7f46fb07e7a1c323
BLAKE2b-256 bba9df8223413f90b9846c144775ea41d081365b5fac9295edaffe46c11c1687

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