Metadata-Version: 2.4
Name: splx-proxy-sdk
Version: 1.6.3
Summary: SDK for building proxy servers for integration between user applications and the SPLX Platform
License: MIT
License-File: LICENSE
Author: Luka Simac
Author-email: lsimac@zscaler.com
Requires-Python: >=3.11,<4.0
Classifier: License :: OSI Approved :: MIT License
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Programming Language :: Python :: 3.14
Requires-Dist: certifi (<2026.5.20)
Requires-Dist: click (<8.4.0)
Requires-Dist: fastapi (>=0.136.1,<0.137.0)
Requires-Dist: loguru (>=0.7.2,<0.8.0)
Requires-Dist: pydantic (>=2.11.7,<3.0.0)
Requires-Dist: typing-extensions (>=4.15.0,<5.0.0)
Requires-Dist: virtualenv (<21.3.3)
Description-Content-Type: text/markdown

# SPLX Proxy SDK

A Python SDK for building proxy servers that connect your applications to the [Zscaler AI Security Platform](https://ai-console.zscaler.com/). It handles authentication, request routing, structured logging, and error handling so you can focus on your integration logic.

Built on [FastAPI](https://fastapi.tiangolo.com/) — your proxy server is a standard ASGI app with full access to FastAPI's ecosystem.

## Installation

```bash
# poetry
poetry add splx-proxy-sdk

# pip
pip install splx-proxy-sdk
```

Requires Python 3.11+.

## Quick Start

```python
from uuid import uuid4

from fastapi import Request

from splx_proxy_sdk import (
    CloseSessionRequest,
    CloseSessionResponse,
    OpenSessionRequest,
    OpenSessionResponse,
    SendMessageRequest,
    SendMessageResponse,
    Server,
)


class MyProxy(Server):
    async def open_session(
        self, request: OpenSessionRequest, raw: Request
    ) -> OpenSessionResponse:
        session_id = request.session_id or str(uuid4())
        return OpenSessionResponse(session_id=session_id)

    async def close_session(
        self, request: CloseSessionRequest, raw: Request
    ) -> CloseSessionResponse:
        return CloseSessionResponse()

    async def send_message(
        self, request: SendMessageRequest, raw: Request
    ) -> SendMessageResponse:
        return SendMessageResponse(
            session_id=request.session_id,
            message=f"Echo: {request.message}",
        )


# Set SPLX_PROXY_API_KEY env variable before starting
app = MyProxy()
```

Run it:

```bash
export SPLX_PROXY_API_KEY="your-secret-key"
uvicorn main:app --reload
```

The AI Security Platform must be able to reach the host and port you expose.

## Testing App

The SDK ships with a lightweight testing app mounted at `/` so you can test your proxy endpoints without leaving the browser. Launch your server locally, then open `http://localhost:8000/` to see the UI. The settings panel lets you set the proxy URL and API key, and attach custom HTTP headers (as a JSON object) to every request sent to your proxy.

![Testing app overview](examples/echo_multimodal/testing-app.png)

## How It Works

The SDK defines three endpoints that the AI Security Platform calls on your proxy:

| Endpoint         | Method | Purpose                                          |
| ---------------- | ------ | ------------------------------------------------ |
| `/open-session`  | POST   | Initialize a new session. Return a `session_id`. |
| `/send-message`  | POST   | Handle a user message. Return the response.      |
| `/close-session` | POST   | Clean up when a session ends.                    |

You subclass `Server` and implement these three methods. The SDK handles authentication, request parsing, logging, and error serialization.

Since `Server` extends `FastAPI`, you can add your own routes, middleware, startup/shutdown events, and anything else FastAPI supports.

## Features

### Authentication

Every request from the AI Security Platform includes an `x-api-key` header. The SDK validates it automatically against the `SPLX_PROXY_API_KEY` environment variable. Invalid keys get a `401 Unauthorized` response.

This variable must be set before the server starts — the SDK raises a `RuntimeError` on startup if it's missing.

### Extra Arguments

You can pass custom, typed arguments to any endpoint using generics. Define a Pydantic model and parametrize your `Server`:

```python
from pydantic import BaseModel

class ExtraArgs(BaseModel):
    tone: str = "friendly"
    max_tokens: int = 1000

class MyProxy(Server[ExtraArgs]):
    async def send_message(
        self, request: SendMessageRequest[ExtraArgs], raw: Request
    ) -> SendMessageResponse:
        if request.extra_args:
            print(f"Tone: {request.extra_args.tone}")
            print(f"Max tokens: {request.extra_args.max_tokens}")
        # ...
```

The `extra_args` field is optional and defaults to `None`. If omitted, `Server` works without any type parameter — `Server` is equivalent to `Server[None]`. Unknown fields in `extra_args` are ignored rather than rejected.

### Multimodal Content

The SDK supports sending and receiving images, audio, and documents. Multimodal content is sent as base64-encoded [Data URIs](https://developer.mozilla.org/en-US/docs/Web/URI/Schemes/data).

```python
from splx_proxy_sdk import MultiModal, MultiModalType

async def send_message(self, request, raw):
    if request.multimodal:
        for name, modal in request.multimodal.items():
            # modal.type is one of: MultiModalType.IMAGE, AUDIO, DOCUMENT
            # modal.content is the base64 Data URI string
            mime_type, base64_data = MultiModal.parse_content(modal.content)
            # mime_type: e.g. "image/png"
            # base64_data: the raw base64 string
```

### Logging

All requests and responses are logged as structured JSON via [Loguru](https://github.com/Delgan/loguru), with a timestamp on every log line. Sensitive headers and fields are redacted by default.

Configure via `LoggingConfig` in code:

```python
from splx_proxy_sdk import LoggingConfig

app = MyProxy(logging_config=LoggingConfig(
    enabled=True,
    log_request_body=True,
    log_response_body=False,
    max_body_length=2048,
    redact_headers=("x-api-key", "authorization"),
    redact_fields=("password", "secret", "token"),
))
```

Or via environment variables:

| Variable                         | Default                                                                  | Description                           |
| -------------------------------- | ------------------------------------------------------------------------ | ------------------------------------- |
| `SPLX_PROXY_LOG_ENABLED`         | `True`                                                                   | Enable/disable logging                |
| `SPLX_PROXY_LOG_REQUEST_BODY`    | `True`                                                                   | Log request bodies                    |
| `SPLX_PROXY_LOG_RESPONSE_BODY`   | `True`                                                                   | Log response bodies                   |
| `SPLX_PROXY_LOG_MAX_BODY_LENGTH` | `4096`                                                                   | Truncate bodies beyond this length    |
| `SPLX_PROXY_LOG_REDACT_HEADERS`  | `x-api-key`                                                              | Comma-separated headers to redact     |
| `SPLX_PROXY_LOG_REDACT_FIELDS`   | `api_key,access_token,refresh_token,token,secret,password,authorization` | Comma-separated JSON fields to redact |

### Error Handling

The SDK provides typed exception classes that produce structured JSON error responses understood by the AI Security Platform:

| Exception                      | Status Code | Default Code            |
| ------------------------------ | ----------- | ----------------------- |
| `BadRequestException`          | 400         | `BAD_REQUEST`           |
| `UnauthorizedException`        | 401         | `UNAUTHORIZED`          |
| `ForbiddenException`           | 403         | `FORBIDDEN`             |
| `NotFoundException`            | 404         | `NOT_FOUND`             |
| `TooManyRequestsException`     | 429         | `TOO_MANY_REQUESTS`     |
| `SessionClosedException`       | 452         | `SESSION_CLOSED`        |
| `InternalServerErrorException` | 500         | `INTERNAL_SERVER_ERROR` |

Basic usage:

```python
from splx_proxy_sdk import NotFoundException, SessionClosedException

# Session doesn't exist
raise NotFoundException(details="Session not found.")

# Session has been closed
raise SessionClosedException(details="This session has ended.")
```

#### Exception Configuration

Use `ProxyExceptionConfig` to attach metadata that the AI Security Platform (Probe) understands:

```python
from datetime import timedelta
from splx_proxy_sdk import ProxyExceptionConfig, TooManyRequestsException

raise TooManyRequestsException(
    details="Rate limit exceeded.",
    config=ProxyExceptionConfig(
        session_closed=False,        # Session is still open
        retry_after=timedelta(seconds=30),  # Retry-After header
        headers={"X-Custom": "value"},      # Additional headers
    ),
)
```

- `session_closed` — tells Probe whether the session was closed as a result of this error (sets `X-Splx-Session-Status` header)
- `retry_after` — tells Probe when to retry (sets `Retry-After` header)
- `headers` — any additional headers to include in the response

## Examples

The [examples/](https://github.com/splx-ai/splx-proxy-sdk-examples/tree/main/) directory contains six proxy implementations, each demonstrating a different integration pattern. Every example is a standalone project you can run locally.

### [echo](https://github.com/splx-ai/splx-proxy-sdk-examples/tree/main/echo/) — Standalone Stateful Proxy

The simplest starting point. A self-contained proxy with no external dependencies that manages sessions in memory and echoes messages back with a counter.

**What it demonstrates:**

- In-memory session management (open, close, lookup)
- Typed `extra_args` with defaults (`repeat: int = 1`)
- Conversation history tracking
- Error handling with `NotFoundException` and `SessionClosedException`

**Run it:**

```bash
cd examples/echo
poetry install
cp proxy/.env.example proxy/.env
uvicorn proxy.main:app --reload
```

---

### [echo_rest](https://github.com/splx-ai/splx-proxy-sdk-examples/tree/main/echo_rest/) — REST API Integration

A proxy that bridges the AI Security Platform protocol to a standard REST API with Bearer token authentication. Includes both the proxy and a sample backend API.

**What it demonstrates:**

- Calling an external REST API with `httpx`
- Bearer token authentication with the downstream service
- Mapping AI Security Platform session IDs to API-specific thread IDs
- Translating between AI Security Platform and REST request/response formats

**Architecture:**

```
AI Security Platform --> Proxy (:8000) --> REST API (:8001)
```

**Run it:**

```bash
cd examples/echo_rest
poetry install
cp proxy/.env.example proxy/.env
# Terminal 1: Start the backend API
uvicorn app.main:app --port 8001 --reload
# Terminal 2: Start the proxy
uvicorn proxy.main:app --port 8000 --reload
```

---

### [echo_sse](https://github.com/splx-ai/splx-proxy-sdk-examples/tree/main/echo_sse/) — Server-Sent Events Integration

A proxy that consumes SSE (Server-Sent Events) from a downstream app and assembles the streamed chunks into a single response for the AI Security Platform protocol.

**What it demonstrates:**

- Consuming SSE streams with `httpx` streaming
- JWT-based authentication with periodic token refresh
- A background token-refresh task managed with a FastAPI lifespan
- Assembling streamed chunks into a single synchronous response

**Architecture:**

```
AI Security Platform --> Proxy (:8000) --> SSE App (:8001)
                    |
            Background auth loop
            refreshes JWT every 8min
```

**Run it:**

```bash
cd examples/echo_sse
poetry install
cp proxy/.env.example proxy/.env
# Terminal 1: Start the SSE app
uvicorn app.main:app --port 8001 --reload
# Terminal 2: Start the proxy
uvicorn proxy.main:app --port 8000 --reload
```

---

### [echo_ws](https://github.com/splx-ai/splx-proxy-sdk-examples/tree/main/echo_ws/) — WebSocket Integration

A proxy that maintains persistent WebSocket connections to a backend app, with one connection per session.

**What it demonstrates:**

- Persistent WebSocket connections per session using `websockets`
- REST call to create a backend conversation, then upgrading to WebSocket
- Consuming chunked WebSocket responses (JSON frames with a "done" signal)
- Graceful connection cleanup on session close and server shutdown

**Architecture:**

```
AI Security Platform --> Proxy (:8000) --> REST /create-conversation
                                    --> WS   /chat (one connection per session)
```

**Run it:**

```bash
cd examples/echo_ws
poetry install
cp proxy/.env.example proxy/.env
# Terminal 1: Start the WebSocket app
uvicorn app.main:app --port 8001 --reload
# Terminal 2: Start the proxy
uvicorn proxy.main:app --port 8000 --reload
```

---

### [echo_multimodal](https://github.com/splx-ai/splx-proxy-sdk-examples/tree/main/echo_multimodal/) — Multimodal File Handling

A proxy that handles multimodal content (images, audio, documents) by decoding base64 Data URIs and forwarding them as multipart file uploads to a downstream API.

**What it demonstrates:**

- Decoding base64 Data URIs using `MultiModal.parse_content()`
- Converting Data URIs to multipart file uploads for a downstream REST API
- MIME type detection and file extension mapping
- Handling image, audio, and document attachments

**Architecture:**

```
AI Security Platform --> Proxy (:8000) --> Multimodal App (:8001)
              (base64 Data URI)   (multipart file upload)
```

**Run it:**

```bash
cd examples/echo_multimodal
poetry install
cp proxy/.env.example proxy/.env
# Terminal 1: Start the multimodal app
uvicorn app.main:app --port 8001 --reload
# Terminal 2: Start the proxy
uvicorn proxy.main:app --port 8000 --reload
```

---

### [broker](https://github.com/splx-ai/splx-proxy-sdk-examples/tree/main/broker/) — Embedded AIRT Broker (One Container)

A proxy with the AIRT broker embedded in the same process, so your AI app and its tunnel ship as one artifact. The broker opens an outbound reverse tunnel (WSS) to the AIRT gateway, letting the platform reach your app for red-teaming without any inbound firewall changes.

**What it demonstrates:**

- Embedding `airt_broker` in the proxy process via a FastAPI lifespan (`start()`/`stop()`)
- Fail-fast startup — a misconfigured broker aborts boot with `BrokerStartError`
- Shutting the process down on terminal tunnel death so the orchestrator restarts it
- Docker packaging with runtime-mounted CA bundles for TLS-inspecting egress

**Run it:** see the [broker README](https://github.com/splx-ai/splx-proxy-sdk-examples/blob/main/broker/README.md) — unlike the other examples it is pip/requirements-based, and requires Python 3.13 plus the broker wheel and credentials from your AIRT onboarding.

## Development

```bash
poetry install
make check          # isort, black, pylint, mypy
poetry run pytest   # run the test suite
```

## Links

- [PyPI](https://pypi.org/project/splx-proxy-sdk/)
- [Zscaler AI Security docs](https://help.zscaler.com/secure-ai-apps-infra/ai-red-teaming)

## License

MIT

