shift-sdk 0.3.4
Python SDK for Shift managed AI routing, telemetry, and local-first execution
Requires Python: >=3.10
Keywords: ai, llm, routing, telemetry, sustainability, executorch

# shift-sdk

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

## Install

```bash
pip install shift-sdk
```

Import path remains:

```python
from switch_sdk import SwitchClient
```

For local development:

```bash
pip install -e .[dev]
```

For local ExecuTorch runtime work:

```bash
pip install -e .[dev,local]
```

For packaging and publishing:

```bash
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

```python
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:

```bash
export SHIFT_API_KEY=aura_your_plain_project_key
```

Optional URL override:

```bash
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.

```python
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:

```bash
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
```

```python
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.

```python
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:

```bash
pip install -e .[local]
```

Ready-made demo manifest:

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

Runtime callable contract:

```python
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

```python
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

```python
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.

```python
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

```python
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

```python
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:

```bash
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

```bash
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

```bash
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

```bash
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:

```bash
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.