mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-06 02:48:13 +00:00
docs(rust-ocr): AGENTS.md explaining the python<->rust bridge
This commit is contained in:
Ishaan Jaffer
parent
27d6fa9303
commit
3005e6db11
1 changed files with 160 additions and 0 deletions
160
litellm-rust/AGENTS.md
Normal file
160
litellm-rust/AGENTS.md
Normal file
|
|
@ -0,0 +1,160 @@
|
|||
# AGENTS.md — how the Python ↔ Rust OCR bridge works
|
||||
|
||||
This explains the wiring between `litellm.ocr()` / `litellm.aocr()` and the Rust
|
||||
core, so you can extend it without re-deriving the design. For the *rules* of
|
||||
Rust work here (the pure-transform boundary, the production bar), see
|
||||
[`CLAUDE.md`](./CLAUDE.md). This file is the *map*.
|
||||
|
||||
## The one-paragraph version
|
||||
|
||||
`litellm.ocr()` and `litellm.aocr()` are thin Python wrappers. They resolve
|
||||
credentials and pre-process the document, then hand a typed call to the compiled
|
||||
`litellm_python_bridge` extension. The bridge parses the Python arguments into a
|
||||
typed `OcrRequest` (GIL held), runs the whole HTTP call on an async Tokio runtime
|
||||
with the GIL released, and converts the typed `OcrResponse` back to a dict.
|
||||
`aocr` returns a Python awaitable; `ocr` blocks on the same future. Everything
|
||||
provider-specific (URL, headers, request/response transform) is typed Rust — no
|
||||
`serde_json::Value` blobs except where Python itself is `Any`.
|
||||
|
||||
## Crate layout
|
||||
|
||||
```
|
||||
litellm-rust/crates/
|
||||
core/ # typed domain model + errors. No I/O, no deps on the others.
|
||||
src/ocr/types.rs OcrRequest, OcrResponse, OcrProvider, OcrDocument, OcrParams, …
|
||||
src/error.rs CoreError (the typed error enum)
|
||||
providers/ # the async entry point + pure per-provider transforms.
|
||||
src/ocr.rs `pub async fn ocr(OcrRequest) -> CoreResult<OcrResponse>` ← re-exported as `litellm_providers::ocr`
|
||||
src/mistral/ocr/transformation.rs complete_url / resolve_api_key / request_body / parse_response (pure)
|
||||
python-bridge/ # the PyO3 boundary. The ONLY crate that touches Python objects.
|
||||
src/lib.rs `ocr`, `aocr`, `gil_stats` pyfunctions
|
||||
src/gil.rs GIL-release accounting
|
||||
```
|
||||
|
||||
Boundary (per `CLAUDE.md`): `mistral/ocr/transformation.rs` is **pure** — no
|
||||
network, no env, no logging. `ocr.rs` is the transport host: it owns the shared
|
||||
`reqwest` client, the HTTP call, and dispatch. The bridge owns Python interop.
|
||||
|
||||
## Naming — it's `ocr` all the way down
|
||||
|
||||
| Python | Rust |
|
||||
|--------|------|
|
||||
| `litellm.ocr()` | bridge `ocr()` → `litellm_providers::ocr(...)` (`block_on`) |
|
||||
| `litellm.aocr()` | bridge `aocr()` → `litellm_providers::ocr(...)` (awaited) |
|
||||
| `MistralOCRConfig.transform_ocr_request/response` | `mistral::request_body` / `mistral::parse_response` |
|
||||
|
||||
There is **one** async core function, `litellm_providers::ocr`. The sync and
|
||||
async bridge entry points are two ways to drive the same future.
|
||||
|
||||
## Async model (why there's no `run_in_executor`)
|
||||
|
||||
The proxy calls `aocr()`. The old path ran the blocking call in a thread-pool
|
||||
worker — one OS thread pinned per in-flight OCR for up to 600s, capped at
|
||||
`min(32, cpu+4)` threads. That bottlenecks under load.
|
||||
|
||||
Instead:
|
||||
|
||||
- `providers::ocr` is `async` and uses the **async** `reqwest` client.
|
||||
- `python-bridge`'s `aocr` wraps it with `pyo3_async_runtimes::tokio::future_into_py`,
|
||||
returning a Python awaitable. `litellm.aocr()` just `await`s it. The HTTP wait
|
||||
happens on a small fixed Tokio worker pool — no thread-per-request.
|
||||
- `python-bridge`'s `ocr` calls `get_runtime().block_on(...)` inside
|
||||
`gil::release_gil` (`py.allow_threads`), so sync SDK callers stay correct and
|
||||
other Python threads keep running during the wait.
|
||||
|
||||
On the Python side, `litellm.aocr()`'s rust branch returns the *coroutine*
|
||||
`_arun_rust_ocr(...)`; the existing `aocr` wrapper awaits it on the event loop, so
|
||||
the executor thread is released before the HTTP wait — not held through it.
|
||||
|
||||
## Request lifecycle (one call)
|
||||
|
||||
1. **Python shell** (`litellm/ocr/main.py`): detect provider, convert a
|
||||
`type:"file"` document to a base64 data-URI, map OCR params, resolve the API
|
||||
key via secret managers, run `pre_call` logging. (These stay in Python — see
|
||||
"What stays in Python".)
|
||||
2. **Bridge `ocr`/`aocr`** (`python-bridge/src/lib.rs`): `build_request(...)`
|
||||
parses the loose Python args into a typed `OcrRequest` while the GIL is held
|
||||
(Python `json` round-trip → serde). Unknown params are dropped.
|
||||
3. **Core `providers::ocr`** (`providers/src/ocr.rs`): `match request.provider`.
|
||||
For Mistral: `resolve_api_key` → `complete_url` → `request_body` (typed,
|
||||
serde-serialized) → async POST on the shared client → status check → `parse_response`.
|
||||
4. **Back across the bridge**: the typed `OcrResponse` is serde-serialized and
|
||||
`json.loads`'d into a Python dict; Python wraps it in `OCRResponse.model_validate`.
|
||||
|
||||
## The typed model (and the only `Value`s)
|
||||
|
||||
`core/src/ocr/types.rs` mirrors the Python `OCRResponse` pydantic models 1:1 as
|
||||
real Rust types. Serde derives produce the exact wire JSON — no hand-built JSON.
|
||||
`serde_json::Value` appears in exactly three leaves, only where Python's own type
|
||||
is `Any`/`Dict[str, Any]`:
|
||||
|
||||
- `AnnotationFormat` — a user-supplied JSON Schema we forward verbatim.
|
||||
- `OcrResponse::document_annotation` — Python `Optional[Any]`.
|
||||
- `OcrPageImage::bbox` — Python `Optional[Dict[str, Any]]`.
|
||||
|
||||
Strict by default: serde ignores unknown upstream fields (vs Python's
|
||||
`extra="allow"` passthrough).
|
||||
|
||||
## Errors
|
||||
|
||||
`CoreError` (typed) → `core_error_to_pyerr` in the bridge:
|
||||
|
||||
- `Auth` / `UnsupportedProvider` / `InvalidType` / `MissingField` → `ValueError`.
|
||||
- `Http { status, body }` / `Network` / `InvalidResponse` → `RuntimeError`, with
|
||||
the HTTP **status preserved in the message**. The Python `exception_type` layer
|
||||
keys off that to raise the right `litellm` exception (RateLimitError on 429,
|
||||
AuthenticationError on 401, …). Upstream bodies are truncated to 256 chars so
|
||||
document contents/secrets never cross the boundary.
|
||||
|
||||
## GIL accounting
|
||||
|
||||
`gil_stats()` returns `{"releases": N}` — the count of calls whose HTTP work ran
|
||||
off the GIL (sync `block_on` releases + async `aocr` offloads). `aocr` work is
|
||||
off-GIL by construction (Tokio threads); `ocr` uses `py.allow_threads`.
|
||||
|
||||
## Adding a provider
|
||||
|
||||
1. **core**: the `OcrProvider` enum already lists the variants. Add typed fields
|
||||
to `OcrParams` / response types only if the provider needs ones Python models.
|
||||
2. **providers**: add `providers/src/<provider>/ocr/transformation.rs` with pure
|
||||
`complete_url`, key resolution, a typed `request_body` (`#[derive(Serialize)]`),
|
||||
and `parse_response` (`#[derive(Deserialize)]` → `OcrResponse`). Unit-test
|
||||
param filtering, request shape, response normalization, null/missing, bad input.
|
||||
3. **providers/src/ocr.rs**: add a `match` arm calling your provider's transforms.
|
||||
Non-Bearer auth goes through the `extra_headers` map; pollers (e.g. Azure DocAI)
|
||||
use `tokio::time::sleep` in the async path.
|
||||
4. **python shell** (`litellm/ocr/rust_bridge.py`): add the provider to
|
||||
`RUST_SUPPORTED_PROVIDERS`, and its API-key env to `_RUST_PROVIDER_API_KEY_ENV`
|
||||
in `main.py`. The shell pre-resolves GCP/Azure tokens and `url→base64` and
|
||||
passes them in (Rust does not mint credentials — see `CLAUDE.md`).
|
||||
5. **tests**: Rust unit tests + Python parity tests (disabled / enabled /
|
||||
bridge-unavailable fallback).
|
||||
|
||||
## What stays in Python (by design)
|
||||
|
||||
`pre_call`/`post_call`, cost & spend tracking, success/failure callbacks, retries,
|
||||
fallbacks, caching, guardrails, secret-manager and GCP/Azure token minting, and
|
||||
`exception_type`. These are cross-cutting and FFI/GIL-bound; the Rust core owns
|
||||
transport + transform + typed errors only.
|
||||
|
||||
## Build & test
|
||||
|
||||
```bash
|
||||
cd litellm-rust
|
||||
cargo fmt --check
|
||||
cargo clippy --workspace --all-targets -- -D warnings
|
||||
cargo test --workspace # core + providers unit tests
|
||||
```
|
||||
|
||||
Building the Python extension locally (macOS needs the dynamic-lookup flags so
|
||||
the cdylib can resolve Python symbols at load time):
|
||||
|
||||
```bash
|
||||
RUSTFLAGS="-C link-arg=-undefined -C link-arg=dynamic_lookup" \
|
||||
cargo build -p litellm-python-bridge --release
|
||||
cp target/release/liblitellm_python_bridge.dylib /tmp/litellm_python_bridge.so
|
||||
# then: PYTHONPATH=/tmp python -c "import litellm, litellm_python_bridge; litellm.use_litellm_rust()"
|
||||
```
|
||||
|
||||
Python-side tests inject a fake bridge (`use_litellm_rust(True, bridge=...)`), so
|
||||
they run without the compiled wheel: `tests/test_litellm/ocr/test_rust_bridge.py`.
|
||||
Loading…
Add table
Reference in a new issue