scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-21 00:21:49 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 73

Merge remote-tracking branch 'origin/main' into litellm_max_parallel_requests_queue_size

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

# Conflicts:
#	tests/test_litellm/test_utils.py
This commit is contained in:
yassin 2026-09-18 08:57:53 +00:00
commit 61ce1b46d9
752 changed files with 54515 additions and 30560 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
.circleci/config.yml
View file

@ -3009,7 +3009,7 @@ workflows:
name: integration-<< matrix.suite >>
matrix:
parameters:
suite: [management, accounting, database, providers, extensions, browser]
suite: [management, accounting, database, providers, extensions, sdk, browser]
filters:
branches:
only:

78
.github/scripts/auto_merge_price_sync.py vendored
View file

@ -1,9 +1,9 @@
"""Auto-merge the provider-info-sync bot's cost-map pull requests.
Evaluates every gate (author allowlist, cost-map-only diff, required and
non-required checks, Greptile confidence, Bugbot review, human reviews) and
merges with a merge commit when all of them hold. Every hold reason is
logged; the process exits 0 on hold and 1 only on API or programming errors.
non-required checks, human reviews) and merges with a merge commit when
all of them hold. Every hold reason is logged; the process exits 0 on hold
and 1 only on API or programming errors.
``DRY_RUN=1`` prints the verdict without calling the merge endpoint.
"""
@ -11,7 +11,6 @@ from __future__ import annotations
import json
import os
import re
import subprocess
import sys
import time
@ -27,12 +26,6 @@ CLASSIFY_SCRIPT: Final = os.path.join(REPO_ROOT, ".circleci", "scripts", "classi
API_ROOT: Final = "https://api.github.com"
CHANGED_FILE_CEILING: Final = 3000
OK_CHECK_CONCLUSIONS: Final = frozenset({"success", "skipped", "neutral"})
GREPTILE_LOGIN: Final = "greptile-apps[bot]"
BUGBOT_LOGIN: Final = "cursor[bot]"
GREPTILE_SCORE_RE: Final = re.compile(r"Confidence Score:\s*(\d)/5")
BUGBOT_REVIEW_MARKER: Final = "<!-- BUGBOT_REVIEW -->"
BUGBOT_STALE_MARKER: Final = "<!-- BUGBOT_REVIEW_STALE -->"
BUGBOT_CLEAN: Final = "found no new issues"
@dataclass(frozen=True, slots=True)
@ -60,13 +53,6 @@ class CommitStatus:
state: str
@dataclass(frozen=True, slots=True)
class IssueComment:
author_login: str
body: str
updated_at: datetime
@dataclass(frozen=True, slots=True)
class Review:
author_login: str
@ -89,9 +75,7 @@ class EvaluationInputs:
required_contexts: frozenset[str]
check_runs: tuple[CheckRun, ...]
statuses: tuple[CommitStatus, ...]
comments: tuple[IssueComment, ...]
reviews: tuple[Review, ...]
head_commit_date: datetime
self_check_name: str
author_allowlist: frozenset[str]
@ -155,37 +139,6 @@ def evaluate(
if status.state != "success":
reasons.append(f"commit status {status.context!r} is {status.state}")
greptile: Final = tuple(
comment
for comment in inputs.comments
if comment.author_login == GREPTILE_LOGIN and GREPTILE_SCORE_RE.search(comment.body)
)
if not greptile:
reasons.append("greptile score not available")
else:
latest: Final = max(greptile, key=lambda comment: comment.updated_at)
match: Final = GREPTILE_SCORE_RE.search(latest.body)
score: Final = int(match.group(1)) if match else 0
if latest.updated_at < inputs.head_commit_date:
reasons.append("greptile score older than head commit")
elif score != 5:
reasons.append(f"greptile score {score}/5 below 5")
bugbot: Final = tuple(
review
for review in inputs.reviews
if review.author_login == BUGBOT_LOGIN
and BUGBOT_REVIEW_MARKER in review.body
and BUGBOT_STALE_MARKER not in review.body
and review.commit_id == pr.head_sha
)
if not bugbot:
reasons.append("bugbot review not available")
else:
latest_review: Final = max(bugbot, key=lambda review: review.submitted_at)
if BUGBOT_CLEAN not in latest_review.body:
reasons.append("bugbot reported issues")
latest_state_by_reviewer: Final[dict[str, str]] = {}
for review in sorted(inputs.reviews, key=lambda review: review.submitted_at):
if _is_bot_login(review.author_login):
@ -350,19 +303,6 @@ def _statuses(token: str, repo: str, sha: str) -> tuple[CommitStatus, ...]:
)
def _comments(token: str, repo: str, number: int) -> tuple[IssueComment, ...]:
comments: Final = _paginate(token, f"/repos/{repo}/issues/{number}/comments")
return tuple(
IssueComment(
author_login=_text(_nested(item, "user", "login")),
body=_text(item.get("body")),
updated_at=_parse_time(item.get("updated_at")),
)
for item in comments
if isinstance(item, Mapping)
)
def _reviews(token: str, repo: str, number: int) -> tuple[Review, ...]:
reviews: Final = _paginate(token, f"/repos/{repo}/pulls/{number}/reviews")
return tuple(
@ -378,16 +318,6 @@ def _reviews(token: str, repo: str, number: int) -> tuple[Review, ...]:
)
def _head_commit_date(token: str, repo: str, number: int) -> datetime:
commits: Final = _paginate(token, f"/repos/{repo}/pulls/{number}/commits")
if not commits:
return datetime.min.replace(tzinfo=timezone.utc)
last: Final = commits[-1]
if not isinstance(last, Mapping):
return datetime.min.replace(tzinfo=timezone.utc)
return _parse_time(_nested(last, "commit", "committer", "date"))
def _mergeable_or_refetch(token: str, repo: str, pr: PullRequest) -> PullRequest:
if pr.mergeable is not None:
return pr
@ -410,9 +340,7 @@ def _gather_inputs(
required_contexts=_required_contexts(token, repo, base),
check_runs=_check_runs(token, repo, pr.head_sha),
statuses=_statuses(token, repo, pr.head_sha),
comments=_comments(token, repo, number),
reviews=_reviews(token, repo, number),
head_commit_date=_head_commit_date(token, repo, number),
self_check_name=self_check_name,
author_allowlist=allowlist,
)

5
.github/workflows/test-unit-proxy-db.yml vendored
View file

@ -94,7 +94,6 @@ jobs:
tests/proxy_unit_tests/test_jwt_key_mapping.py
tests/proxy_unit_tests/test_proxy_custom_auth.py
tests/proxy_unit_tests/test_key_generate_dynamodb.py
tests/proxy_unit_tests/test_deployed_proxy_keygen.py
workers: 4
dist: loadscope
timeout: 15
@ -110,8 +109,6 @@ jobs:
- test-group: proxy-server-core
test-path: >-
tests/proxy_unit_tests/test_proxy_server.py
tests/proxy_unit_tests/test_proxy_server_keys.py
tests/proxy_unit_tests/test_proxy_server_spend.py
tests/proxy_unit_tests/test_aproxy_startup.py
workers: 4
dist: loadscope
@ -120,7 +117,6 @@ jobs:
test-path: >-
tests/proxy_unit_tests/test_proxy_config_unit_test.py
tests/proxy_unit_tests/test_proxy_routes.py
tests/proxy_unit_tests/test_proxy_gunicorn.py
tests/proxy_unit_tests/test_server_root_path.py
tests/proxy_unit_tests/test_proxy_pass_user_config.py
tests/proxy_unit_tests/test_proxy_token_counter.py
@ -198,7 +194,6 @@ jobs:
tests/proxy_unit_tests/test_realtime_cache.py
tests/proxy_unit_tests/test_proxy_exception_mapping.py
tests/proxy_unit_tests/test_custom_tokenizer_bug.py
tests/proxy_unit_tests/test_model_response_typing
workers: 4
dist: loadscope
timeout: 15

3
.github/workflows/test-unit.yml vendored
View file

@ -100,6 +100,7 @@ jobs:
tests/test_litellm/secret_managers
tests/test_litellm/a2a_protocol
tests/test_litellm/anthropic_interface
tests/test_litellm/chat_completions
tests/test_litellm/completion_extras
tests/test_litellm/compression
tests/test_litellm/containers
@ -109,6 +110,7 @@ jobs:
tests/test_litellm/repositories
tests/test_litellm/images
tests/test_litellm/interactions
tests/test_litellm/messages
tests/test_litellm/ocr
tests/test_litellm/passthrough
tests/test_litellm/rag
@ -211,7 +213,6 @@ jobs:
test-path: >-
tests/local_testing/test_cache_preset_key.py
tests/local_testing/test_caching_handler.py
tests/local_testing/test_prompt_caching.py
tests/local_testing/test_responses_stream_cache_keys.py
tests/local_testing/test_unit_test_caching.py
workers: 2

6
cookbook/litellm_proxy_server/cli_token_usage.py
View file

@ -3,7 +3,7 @@
Example: Using CLI token with LiteLLM SDK
This example shows how to use the CLI authentication token
in your Python scripts after running `litellm-proxy login`.
in your Python scripts after running `lite login`.
"""
from textwrap import indent
@ -22,7 +22,7 @@ def main():
api_key = litellm.get_litellm_gateway_api_key()
if not api_key:
print("❌ No CLI token found. Please run 'litellm-proxy login' first.")
print("❌ No CLI token found. Please run 'lite login' first.")
return
print("✅ Found CLI token.")
@ -58,6 +58,6 @@ if __name__ == "__main__":
main()
print("\n💡 Tips:")
print("1. Run 'litellm-proxy login' to authenticate first")
print("1. Run 'lite login' to authenticate first")
print("2. Replace 'https://your-proxy.com' with your actual proxy URL")
print("3. The token is stored in your OS keychain, or in ~/.litellm/token.json when there is none")

614
cookbook/litellm_proxy_server/grafana_dashboard/dashboard_1/grafana_dashboard.json
View file

@ -1,614 +0,0 @@
{
"annotations": {
"list": [
{
"builtIn": 1,
"datasource": {
"type": "grafana",
"uid": "-- Grafana --"
},
"enable": true,
"hide": true,
"iconColor": "rgba(0, 211, 255, 1)",
"name": "Annotations & Alerts",
"target": {
"limit": 100,
"matchAny": false,
"tags": [],
"type": "dashboard"
},
"type": "dashboard"
}
]
},
"description": "",
"editable": true,
"fiscalYearStartMonth": 0,
"graphTooltip": 0,
"id": 2039,
"links": [],
"liveNow": false,
"panels": [
{
"datasource": {
"type": "prometheus",
"uid": "${DS_PROMETHEUS}"
},
"fieldConfig": {
"defaults": {
"color": {
"mode": "palette-classic"
},
"custom": {
"axisCenteredZero": false,
"axisColorMode": "text",
"axisLabel": "",
"axisPlacement": "auto",
"barAlignment": 0,
"drawStyle": "line",
"fillOpacity": 0,
"gradientMode": "none",
"hideFrom": {
"legend": false,
"tooltip": false,
"viz": false
},
"lineInterpolation": "linear",
"lineWidth": 1,
"pointSize": 5,
"scaleDistribution": {
"type": "linear"
},
"showPoints": "auto",
"spanNulls": false,
"stacking": {
"group": "A",
"mode": "none"
},
"thresholdsStyle": {
"mode": "off"
}
},
"mappings": [],
"thresholds": {
"mode": "absolute",
"steps": [
{
"color": "green",
"value": null
},
{
"color": "red",
"value": 80
}
]
},
"unit": "s"
},
"overrides": []
},
"gridPos": {
"h": 8,
"w": 12,
"x": 0,
"y": 0
},
"id": 10,
"options": {
"legend": {
"calcs": [],
"displayMode": "list",
"placement": "bottom",
"showLegend": true
},
"tooltip": {
"mode": "single",
"sort": "none"
}
},
"targets": [
{
"datasource": {
"type": "prometheus",
"uid": "${DS_PROMETHEUS}"
},
"editorMode": "code",
"expr": "histogram_quantile(0.99, sum(rate(litellm_self_latency_bucket{self=\"self\"}[1m])) by (le))",
"legendFormat": "Time to first token",
"range": true,
"refId": "A"
}
],
"title": "Time to first token (latency)",
"type": "timeseries"
},
{
"datasource": {
"type": "prometheus",
"uid": "${DS_PROMETHEUS}"
},
"fieldConfig": {
"defaults": {
"color": {
"mode": "palette-classic"
},
"custom": {
"axisCenteredZero": false,
"axisColorMode": "text",
"axisLabel": "",
"axisPlacement": "auto",
"barAlignment": 0,
"drawStyle": "line",
"fillOpacity": 0,
"gradientMode": "none",
"hideFrom": {
"legend": false,
"tooltip": false,
"viz": false
},
"lineInterpolation": "linear",
"lineWidth": 1,
"pointSize": 5,
"scaleDistribution": {
"type": "linear"
},
"showPoints": "auto",
"spanNulls": false,
"stacking": {
"group": "A",
"mode": "none"
},
"thresholdsStyle": {
"mode": "off"
}
},
"mappings": [],
"thresholds": {
"mode": "absolute",
"steps": [
{
"color": "green",
"value": null
},
{
"color": "red",
"value": 80
}
]
},
"unit": "currencyUSD"
},
"overrides": [
{
"matcher": {
"id": "byName",
"options": "7e4b0627fd32efdd2313c846325575808aadcf2839f0fde90723aab9ab73c78f"
},
"properties": [
{
"id": "displayName",
"value": "Translata"
}
]
}
]
},
"gridPos": {
"h": 8,
"w": 12,
"x": 0,
"y": 8
},
"id": 11,
"options": {
"legend": {
"calcs": [],
"displayMode": "list",
"placement": "bottom",
"showLegend": true
},
"tooltip": {
"mode": "single",
"sort": "none"
}
},
"targets": [
{
"datasource": {
"type": "prometheus",
"uid": "${DS_PROMETHEUS}"
},
"editorMode": "code",
"expr": "sum(increase(litellm_spend_metric_total[30d])) by (hashed_api_key)",
"legendFormat": "{{team}}",
"range": true,
"refId": "A"
}
],
"title": "Spend by team",
"transformations": [],
"type": "timeseries"
},
{
"datasource": {
"type": "prometheus",
"uid": "${DS_PROMETHEUS}"
},
"fieldConfig": {
"defaults": {
"color": {
"mode": "palette-classic"
},
"custom": {
"axisCenteredZero": false,
"axisColorMode": "text",
"axisLabel": "",
"axisPlacement": "auto",
"barAlignment": 0,
"drawStyle": "line",
"fillOpacity": 0,
"gradientMode": "none",
"hideFrom": {
"legend": false,
"tooltip": false,
"viz": false
},
"lineInterpolation": "linear",
"lineWidth": 1,
"pointSize": 5,
"scaleDistribution": {
"type": "linear"
},
"showPoints": "auto",
"spanNulls": false,
"stacking": {
"group": "A",
"mode": "none"
},
"thresholdsStyle": {
"mode": "off"
}
},
"mappings": [],
"thresholds": {
"mode": "absolute",
"steps": [
{
"color": "green",
"value": null
},
{
"color": "red",
"value": 80
}
]
}
},
"overrides": []
},
"gridPos": {
"h": 9,
"w": 12,
"x": 0,
"y": 16
},
"id": 2,
"options": {
"legend": {
"calcs": [],
"displayMode": "list",
"placement": "bottom",
"showLegend": true
},
"tooltip": {
"mode": "single",
"sort": "none"
}
},
"targets": [
{
"datasource": {
"type": "prometheus",
"uid": "${DS_PROMETHEUS}"
},
"editorMode": "code",
"expr": "sum by (model) (increase(litellm_requests_metric_total[5m]))",
"legendFormat": "{{model}}",
"range": true,
"refId": "A"
}
],
"title": "Requests by model",
"type": "timeseries"
},
{
"datasource": {
"type": "prometheus",
"uid": "${DS_PROMETHEUS}"
},
"fieldConfig": {
"defaults": {
"color": {
"mode": "thresholds"
},
"mappings": [],
"noValue": "0",
"thresholds": {
"mode": "absolute",
"steps": [
{
"color": "green",
"value": null
},
{
"color": "red",
"value": 80
}
]
}
},
"overrides": []
},
"gridPos": {
"h": 7,
"w": 3,
"x": 0,
"y": 25
},
"id": 8,
"options": {
"colorMode": "value",
"graphMode": "area",
"justifyMode": "auto",
"orientation": "auto",
"reduceOptions": {
"calcs": [
"lastNotNull"
],
"fields": "",
"values": false
},
"textMode": "auto"
},
"pluginVersion": "9.4.17",
"targets": [
{
"datasource": {
"type": "prometheus",
"uid": "${DS_PROMETHEUS}"
},
"editorMode": "code",
"expr": "sum(increase(litellm_llm_api_failed_requests_metric_total[1h]))",
"legendFormat": "__auto",
"range": true,
"refId": "A"
}
],
"title": "Faild Requests",
"type": "stat"
},
{
"datasource": {
"type": "prometheus",
"uid": "${DS_PROMETHEUS}"
},
"fieldConfig": {
"defaults": {
"color": {
"mode": "palette-classic"
},
"custom": {
"axisCenteredZero": false,
"axisColorMode": "text",
"axisLabel": "",
"axisPlacement": "auto",
"barAlignment": 0,
"drawStyle": "line",
"fillOpacity": 0,
"gradientMode": "none",
"hideFrom": {
"legend": false,
"tooltip": false,
"viz": false
},
"lineInterpolation": "linear",
"lineWidth": 1,
"pointSize": 5,
"scaleDistribution": {
"type": "linear"
},
"showPoints": "auto",
"spanNulls": false,
"stacking": {
"group": "A",
"mode": "none"
},
"thresholdsStyle": {
"mode": "off"
}
},
"mappings": [],
"thresholds": {
"mode": "absolute",
"steps": [
{
"color": "green",
"value": null
},
{
"color": "red",
"value": 80
}
]
},
"unit": "currencyUSD"
},
"overrides": []
},
"gridPos": {
"h": 7,
"w": 3,
"x": 3,
"y": 25
},
"id": 6,
"options": {
"legend": {
"calcs": [],
"displayMode": "list",
"placement": "bottom",
"showLegend": true
},
"tooltip": {
"mode": "single",
"sort": "none"
}
},
"targets": [
{
"datasource": {
"type": "prometheus",
"uid": "${DS_PROMETHEUS}"
},
"editorMode": "code",
"expr": "sum(increase(litellm_spend_metric_total[30d])) by (model)",
"legendFormat": "{{model}}",
"range": true,
"refId": "A"
}
],
"title": "Spend",
"type": "timeseries"
},
{
"datasource": {
"type": "prometheus",
"uid": "${DS_PROMETHEUS}"
},
"fieldConfig": {
"defaults": {
"color": {
"mode": "palette-classic"
},
"custom": {
"axisCenteredZero": false,
"axisColorMode": "text",
"axisLabel": "",
"axisPlacement": "auto",
"barAlignment": 0,
"drawStyle": "line",
"fillOpacity": 0,
"gradientMode": "none",
"hideFrom": {
"legend": false,
"tooltip": false,
"viz": false
},
"lineInterpolation": "linear",
"lineWidth": 1,
"pointSize": 5,
"scaleDistribution": {
"type": "linear"
},
"showPoints": "auto",
"spanNulls": false,
"stacking": {
"group": "A",
"mode": "none"
},
"thresholdsStyle": {
"mode": "off"
}
},
"mappings": [],
"thresholds": {
"mode": "absolute",
"steps": [
{
"color": "green",
"value": null
},
{
"color": "red",
"value": 80
}
]
}
},
"overrides": []
},
"gridPos": {
"h": 7,
"w": 6,
"x": 6,
"y": 25
},
"id": 4,
"options": {
"legend": {
"calcs": [],
"displayMode": "list",
"placement": "bottom",
"showLegend": true
},
"tooltip": {
"mode": "single",
"sort": "none"
}
},
"targets": [
{
"datasource": {
"type": "prometheus",
"uid": "${DS_PROMETHEUS}"
},
"editorMode": "code",
"expr": "sum(increase(litellm_total_tokens_total[5m])) by (model)",
"legendFormat": "__auto",
"range": true,
"refId": "A"
}
],
"title": "Tokens",
"type": "timeseries"
}
],
"refresh": "1m",
"revision": 1,
"schemaVersion": 38,
"style": "dark",
"tags": [],
"templating": {
"list": [
{
"current": {
"selected": false,
"text": "prometheus",
"value": "edx8memhpd9tsa"
},
"hide": 0,
"includeAll": false,
"label": "datasource",
"multi": false,
"name": "DS_PROMETHEUS",
"options": [],
"query": "prometheus",
"queryValue": "",
"refresh": 1,
"regex": "",
"skipUrlSync": false,
"type": "datasource"
}
]
},
"time": {
"from": "now-1h",
"to": "now"
},
"timepicker": {},
"timezone": "",
"title": "LLM Proxy",
"uid": "rgRrHxESz",
"version": 15,
"weekStart": ""
}

6
cookbook/litellm_proxy_server/grafana_dashboard/dashboard_1/readme.md
View file

@ -1,6 +0,0 @@
## This folder contains the `json` for creating the following Grafana Dashboard
### Pre-Requisites
- Setup LiteLLM Proxy Prometheus Metrics https://docs.litellm.ai/docs/proxy/prometheus
![1716623265684](https://github.com/BerriAI/litellm/assets/29436595/0e12c57e-4a2d-4850-bd4f-e4294f87a814)

6312
cookbook/litellm_proxy_server/grafana_dashboard/dashboard_all_metrics/grafana_dashboard.json Normal file
View file

File diff suppressed because it is too large Load diff

11
cookbook/litellm_proxy_server/grafana_dashboard/dashboard_all_metrics/readme.md Normal file
View file

@ -0,0 +1,11 @@
# LiteLLM All Prometheus Metrics dashboard
Every `litellm_*` metric family the proxy can expose on `/metrics` (134 families across 95 panels), grouped into rows: proxy traffic, latency, spend and tokens, cache, LLM API deployments, key and team rate limits, budgets, guardrails, MCP, managed files and batches, users and teams, the Redis circuit breaker, the spend log cleanup job, and the `prometheus_system` service callback metrics (per-service latency, request and failure rates, spend update queue sizes). Panel titles are the metric names so you can grep the JSON for the metric you care about
Import `grafana_dashboard.json` from **Dashboards > New > Import** and pick your Prometheus data source when prompted (the `DS_PROMETHEUS` variable). Counters are plotted as `rate()` over `$__rate_interval`, histograms as p50 / p95 / p99, gauges as the raw value grouped by the most useful label. Every query names the metric exactly as the proxy emits it (counters carry the `_total` suffix the Prometheus client adds), and `tests/test_litellm/integrations/test_prometheus_metric_name_consistency.py` fails if a metric is renamed without updating this dashboard
The first eleven rows need only `callbacks: ["prometheus"]`. The last three rows and the `litellm_admission_*` panels are emitted by other subsystems and stay empty until those are on: the service callback row needs `service_callback: ["prometheus_system"]` in `litellm_settings`, the circuit breaker row needs a Redis cache, the cleanup row needs spend log retention, and admission control needs its middleware enabled. Within the base rows, many panels only fill in once the matching feature is in use: budgets need keys, teams, users or orgs with `max_budget` set, cache panels need caching on, guardrail and MCP panels need those features configured, deployment health needs the router with more than one deployment or a failure to record, and `litellm_in_flight_requests` needs traffic at scrape time. An empty panel for a feature you do not use is expected
## Pre-requisites
Prometheus metrics on the proxy: https://docs.litellm.ai/docs/proxy/prometheus

4
cookbook/litellm_proxy_server/grafana_dashboard/dashboard_v2/grafana_dashboard.json
View file

@ -476,7 +476,7 @@
"uid": "${DS_PROMETHEUS}"
},
"editorMode": "code",
"expr": "topk(5, sort(litellm_remaining_requests))",
"expr": "topk(5, sort(litellm_remaining_requests_metric))",
"legendFormat": "__auto",
"range": true,
"refId": "A"
@ -573,7 +573,7 @@
"uid": "${DS_PROMETHEUS}"
},
"editorMode": "code",
"expr": "topk(5, sort(litellm_remaining_tokens))",
"expr": "topk(5, sort(litellm_remaining_tokens_metric))",
"legendFormat": "__auto",
"range": true,
"refId": "A"

6
cookbook/litellm_proxy_server/grafana_dashboard/readme.md
View file

@ -6,8 +6,14 @@ This folder contains the `json` for creating Grafana Dashboards
Charts the `gen_ai.*` metrics from the OpenTelemetry v2 integration: spend, tokens, request rate, and latency percentiles by model. Separate from the dashboards below, which chart the `litellm_*` Prometheus metrics.
## [LiteLLM All Prometheus Metrics dashboard](./dashboard_all_metrics)
Every `litellm_*` Prometheus metric family the proxy can emit (134 families, 95 panels) grouped by theme: traffic, latency, spend and tokens, cache, deployments, rate limits, budgets, guardrails, MCP, managed files and batches, users and teams, plus the Redis circuit breaker, spend log cleanup and `prometheus_system` service metrics. Start here if you want everything on one screen; see its [readme](./dashboard_all_metrics/readme.md) for import steps and which panels need a feature enabled before they show data
## [LiteLLM v2 Dashboard](./dashboard_v2)
A compact view of proxy request rate, failures, latency and the top remaining-request / remaining-token gauges per model group
<img width="1316" alt="grafana_1" src="https://github.com/user-attachments/assets/d0df802d-0cb9-4906-a679-941c547789ab">
<img width="1289" alt="grafana_2" src="https://github.com/user-attachments/assets/b11f755f-e113-42ab-b21d-83f91f451a28">
<img width="1323" alt="grafana_3" src="https://github.com/user-attachments/assets/cb29ffdb-477d-4be1-a5cd-c3f7f2cb21c5">

5
db_scripts/partition_spend_logs.sql
View file

@ -53,6 +53,8 @@ ALTER INDEX IF EXISTS "LiteLLM_SpendLogs_end_user_idx"
RENAME TO "LiteLLM_SpendLogs_legacy_end_user_idx";
ALTER INDEX IF EXISTS "LiteLLM_SpendLogs_session_id_idx"
RENAME TO "LiteLLM_SpendLogs_legacy_session_id_idx";
ALTER INDEX IF EXISTS "LiteLLM_SpendLogs_api_key_startTime_idx"
RENAME TO "LiteLLM_SpendLogs_legacy_api_key_startTime_idx";
CREATE TABLE "LiteLLM_SpendLogs" (
LIKE "LiteLLM_SpendLogs_legacy" INCLUDING DEFAULTS INCLUDING GENERATED
@ -78,6 +80,9 @@ CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_end_user_idx"
CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_session_id_idx"
ON "LiteLLM_SpendLogs" ("session_id");
CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_api_key_startTime_idx"
ON "LiteLLM_SpendLogs" ("api_key", "startTime");
-- Safety net: any row whose startTime has no explicit partition lands here so
-- writes never fail. The cleanup job never drops the DEFAULT partition.
CREATE TABLE IF NOT EXISTS "LiteLLM_SpendLogs_pdefault"

5
db_scripts/unpartition_spend_logs.sql
View file

@ -40,6 +40,8 @@ ALTER INDEX IF EXISTS "LiteLLM_SpendLogs_end_user_idx"
RENAME TO "LiteLLM_SpendLogs_partitioned_end_user_idx";
ALTER INDEX IF EXISTS "LiteLLM_SpendLogs_session_id_idx"
RENAME TO "LiteLLM_SpendLogs_partitioned_session_id_idx";
ALTER INDEX IF EXISTS "LiteLLM_SpendLogs_api_key_startTime_idx"
RENAME TO "LiteLLM_SpendLogs_partitioned_api_key_startTime_idx";
CREATE TABLE "LiteLLM_SpendLogs" (
LIKE "LiteLLM_SpendLogs_partitioned" INCLUDING DEFAULTS INCLUDING GENERATED
@ -60,6 +62,9 @@ CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_end_user_idx"
CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_session_id_idx"
ON "LiteLLM_SpendLogs" ("session_id");
CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_api_key_startTime_idx"
ON "LiteLLM_SpendLogs" ("api_key", "startTime");
INSERT INTO "LiteLLM_SpendLogs"
SELECT * FROM "LiteLLM_SpendLogs_partitioned"
ON CONFLICT ("request_id") DO NOTHING;

1
gateway/routes/allowlist.py
View file

@ -96,6 +96,7 @@ GATEWAY_PATH_PREFIXES: tuple[str, ...] = (
"/langfuse/",
"/vllm/",
"/mistral/",
"/typesafe/",
"/nvidia_nim/",
"/groq/",
"/voyage/",

2
litellm-proxy-extras/litellm_proxy_extras/migrations/20260823000000_add_spend_logs_api_key_starttime_index/migration.sql Normal file
View file

@ -0,0 +1,2 @@
-- CreateIndex
CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_api_key_startTime_idx" ON "LiteLLM_SpendLogs"("api_key", "startTime");

1
litellm-proxy-extras/litellm_proxy_extras/migrations/20260917055603_add_policy_attachment_priority/migration.sql Normal file
View file

@ -0,0 +1 @@
ALTER TABLE "LiteLLM_PolicyAttachmentTable" ADD COLUMN IF NOT EXISTS "priority" INTEGER;

2
litellm-proxy-extras/litellm_proxy_extras/schema.prisma
View file

@ -678,6 +678,7 @@ model LiteLLM_SpendLogs {
@@index([end_user])
@@index([session_id])
@@index([litellm_call_id])
@@index([api_key, startTime])
}
model LiteLLM_BudgetWindowSpend {
@ -1378,6 +1379,7 @@ model LiteLLM_PolicyAttachmentTable {
keys String[] @default([]) // Key aliases or patterns
models String[] @default([]) // Model names or patterns
tags String[] @default([]) // Tag patterns (e.g., ["healthcare", "prod-*"])
priority Int? // Explicit execution order
created_at DateTime @default(now())
created_by String?
updated_at DateTime @default(now()) @updatedAt

4
litellm-proxy-extras/pyproject.toml
View file

@ -1,6 +1,6 @@
[project]
name = "litellm-proxy-extras"
version = "0.4.98"
version = "0.4.99"
description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package."
readme = "README.md"
requires-python = ">=3.9"
@ -26,7 +26,7 @@ required-version = ">=0.10.9"
module-root = ""
[tool.commitizen]
version = "0.4.98"
version = "0.4.99"
version_files = [
"pyproject.toml:^version",
"../pyproject.toml:litellm-proxy-extras==",

227
litellm-rust/Cargo.lock generated
View file

@ -948,8 +948,18 @@ version = "0.20.11"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "fc7f46116c46ff9ab3eb1597a45688b6715c6e628b5c133e288e709a29bcb4ee"
dependencies = [
"darling_core",
"darling_macro",
"darling_core 0.20.11",
"darling_macro 0.20.11",
]
[[package]]
name = "darling"
version = "0.21.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9cdf337090841a411e2a7f3deb9187445851f91b309c0c0a29e05f74a00a48c0"
dependencies = [
"darling_core 0.21.3",
"darling_macro 0.21.3",
]
[[package]]
@ -966,13 +976,38 @@ dependencies = [
"syn 2.0.119",
]
[[package]]
name = "darling_core"
version = "0.21.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1247195ecd7e3c85f83c8d2a366e4210d588e802133e1e355180a9870b517ea4"
dependencies = [
"fnv",
"ident_case",
"proc-macro2",
"quote",
"strsim",
"syn 2.0.119",
]
[[package]]
name = "darling_macro"
version = "0.20.11"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "fc34b93ccb385b40dc71c6fceac4b2ad23662c7eeb248cf10d529b7e055b6ead"
dependencies = [
"darling_core",
"darling_core 0.20.11",
"quote",
"syn 2.0.119",
]
[[package]]
name = "darling_macro"
version = "0.21.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d38308df82d1080de0afee5d069fa14b0326a88c14f15c5ccda35b4a6c414c81"
dependencies = [
"darling_core 0.21.3",
"quote",
"syn 2.0.119",
]
@ -1022,7 +1057,7 @@ version = "0.20.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "2d5bcf7b024d6835cfb3d473887cd966994907effbe9227e8c8219824d06c4e8"
dependencies = [
"darling",
"darling 0.20.11",
"proc-macro2",
"quote",
"syn 2.0.119",
@ -1363,7 +1398,7 @@ dependencies = [
"futures-sink",
"futures-util",
"http 0.2.12",
"indexmap",
"indexmap 2.14.0",
"slab",
"tokio",
"tokio-util",
@ -1382,7 +1417,7 @@ dependencies = [
"futures-core",
"futures-sink",
"http 1.4.2",
"indexmap",
"indexmap 2.14.0",
"slab",
"tokio",
"tokio-util",
@ -1400,6 +1435,12 @@ dependencies = [
"zerocopy",
]
[[package]]
name = "hashbrown"
version = "0.12.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "8a9ee70c43aaf417c914396645a0fa852624801b24ebb7ae78fe8272889ac888"
[[package]]
name = "hashbrown"
version = "0.17.1"
@ -1736,6 +1777,17 @@ dependencies = [
"icu_properties",
]
[[package]]
name = "indexmap"
version = "1.9.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "bd070e393353796e801d209ad339e89596eb4c8d430d18ede6a1cced8fafbd99"
dependencies = [
"autocfg",
"hashbrown 0.12.3",
"serde",
]
[[package]]
name = "indexmap"
version = "2.14.0"
@ -1743,7 +1795,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d466e9454f08e4a911e14806c24e16fba1b4c121d1ea474396f396069cf949d9"
dependencies = [
"equivalent",
"hashbrown",
"hashbrown 0.17.1",
"serde",
"serde_core",
]
@ -1949,10 +2001,32 @@ dependencies = [
"tokio",
]
[[package]]
name = "litellm-callbacks"
version = "0.1.0"
dependencies = [
"rstest",
"serde_json",
"tokio",
]
[[package]]
name = "litellm-callbacks-legacy"
version = "0.1.0"
dependencies = [
"litellm-callbacks",
"litellm-host-python",
"pyo3",
"rstest",
"serde_json",
]
[[package]]
name = "litellm-core"
version = "0.1.0"
dependencies = [
"aws-smithy-eventstream",
"aws-smithy-types",
"base64 0.22.1",
"bytes",
"data-url",
@ -1961,20 +2035,26 @@ dependencies = [
"litellm-auth-aws",
"litellm-auth-azure",
"litellm-auth-gcp",
"litellm-callbacks",
"litellm-framing",
"litellm-providers",
"mime_guess",
"moka",
"rand 0.8.7",
"reqwest 0.12.28",
"rstest",
"rstest_reuse",
"rustls 0.23.42",
"rustls-native-certs",
"serde",
"serde_json",
"serde_path_to_error",
"serde_with",
"sha2 0.10.9",
"strum",
"subtle",
"thiserror 2.0.19",
"time",
"tokio",
"tokio-tungstenite",
"url",
@ -1995,6 +2075,33 @@ dependencies = [
"tokio",
]
[[package]]
name = "litellm-host-python"
version = "0.1.0"
dependencies = [
"futures-util",
"litellm-callbacks",
"pyo3",
"pyo3-async-runtimes",
"pythonize",
"rstest",
"serde",
"serde_json",
"tokio",
]
[[package]]
name = "litellm-providers"
version = "0.1.0"
dependencies = [
"litellm-auth",
"litellm-auth-aws",
"rstest",
"serde",
"serde_json",
"thiserror 2.0.19",
]
[[package]]
name = "litellm-python-bridge"
version = "0.1.0"
@ -2003,36 +2110,25 @@ dependencies = [
"criterion",
"futures-util",
"litellm-auth",
"litellm-callbacks-legacy",
"litellm-core",
"litellm-python-interop",
"litellm-host-python",
"litellm-token-counter",
"pyo3",
"pyo3-async-runtimes",
"rstest",
"serde",
"serde_json",
"tokio",
"tokio-tungstenite",
]
[[package]]
name = "litellm-python-interop"
version = "0.1.0"
dependencies = [
"pyo3",
"pythonize",
"rstest",
"serde",
"serde_json",
]
[[package]]
name = "litellm-token-counter"
version = "0.1.0"
dependencies = [
"base64 0.22.1",
"criterion",
"indexmap",
"indexmap 2.14.0",
"itoa",
"rand 0.8.7",
"rstest",
@ -2753,6 +2849,26 @@ dependencies = [
"bitflags",
]
[[package]]
name = "ref-cast"
version = "1.0.27"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "7e440fb4e4b4147295338efb76001ab9e4efc0e5839df2c47fc5ac2381d365c3"
dependencies = [
"ref-cast-impl",
]
[[package]]
name = "ref-cast-impl"
version = "1.0.27"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "92ecd8964f8453721699a1ed72037b0db49ce2f5a5138486ee89bed6f67cdf3a"
dependencies = [
"proc-macro2",
"quote",
"syn 3.0.0",
]
[[package]]
name = "regex"
version = "1.13.1"
@ -2919,6 +3035,17 @@ dependencies = [
"unicode-ident",
]
[[package]]
name = "rstest_reuse"
version = "0.7.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b3a8fb4672e840a587a66fc577a5491375df51ddb88f2a2c2a792598c326fe14"
dependencies = [
"quote",
"rand 0.8.7",
"syn 2.0.119",
]
[[package]]
name = "rustc-hash"
version = "2.1.3"
@ -3075,6 +3202,30 @@ dependencies = [
"windows-sys 0.61.2",
]
[[package]]
name = "schemars"
version = "0.9.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "4cd191f9397d57d581cddd31014772520aa448f65ef991055d7f61582c65165f"
dependencies = [
"dyn-clone",
"ref-cast",
"serde",
"serde_json",
]
[[package]]
name = "schemars"
version = "1.2.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "687274d293b6cdc6e73e0fee520bf2049650090d7164f87672d212a3c530cf4a"
dependencies = [
"dyn-clone",
"ref-cast",
"serde",
"serde_json",
]
[[package]]
name = "scopeguard"
version = "1.2.0"
@ -3156,6 +3307,7 @@ version = "1.0.150"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e8014e44b4736ed0538adeecded0fce2a272f22dc9578a7eb6b2d9993c74cfb9"
dependencies = [
"indexmap 2.14.0",
"itoa",
"memchr",
"serde",
@ -3186,6 +3338,37 @@ dependencies = [
"serde",
]
[[package]]
name = "serde_with"
version = "3.16.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "4fa237f2807440d238e0364a218270b98f767a00d3dada77b1c53ae88940e2e7"
dependencies = [
"base64 0.22.1",
"chrono",
"hex",
"indexmap 1.9.3",
"indexmap 2.14.0",
"schemars 0.9.0",
"schemars 1.2.2",
"serde_core",
"serde_json",
"serde_with_macros",
"time",
]
[[package]]
name = "serde_with_macros"
version = "3.16.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "52a8e3ca0ca629121f70ab50f95249e5a6f925cc0f6ffe8256c45b728875706c"
dependencies = [
"darling 0.21.3",
"proc-macro2",
"quote",
"syn 2.0.119",
]
[[package]]
name = "sha1"
version = "0.10.7"
@ -3661,7 +3844,7 @@ version = "0.25.13+spec-1.1.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "6975367e4d2ef766d86af01ffad14b622fecc8d4357a998fbc4deb6e9bacaf9b"
dependencies = [
"indexmap",
"indexmap 2.14.0",
"toml_datetime",
"toml_parser",
"winnow",

12
litellm-rust/Cargo.toml
View file

@ -9,26 +9,33 @@ license = "MIT"
repository = "https://github.com/BerriAI/litellm"
[workspace.dependencies]
bytes = "1"
litellm-core = { path = "crates/core" }
litellm-callbacks = { path = "crates/callbacks" }
litellm-callbacks-legacy = { path = "crates/callbacks-legacy" }
litellm-framing = { path = "crates/framer" }
litellm-auth = { path = "crates/auth" }
litellm-auth-aws = { path = "crates/auth-aws" }
litellm-auth-azure = { path = "crates/auth-azure" }
litellm-auth-gcp = { path = "crates/auth-gcp" }
litellm-providers = { path = "crates/providers" }
litellm-cache = { path = "crates/cache" }
litellm-cache-memory = { path = "crates/cache-memory" }
litellm-token-counter = { path = "crates/token-counter" }
litellm-python-interop = { path = "crates/python-interop" }
litellm-host-python = { path = "crates/host-python" }
bytes = "1"
pyo3 = "0.29.2"
pyo3-async-runtimes = { version = "0.29.0", features = ["tokio-runtime"] }
pythonize = "0.29.0"
rand = "0.8"
reqwest = { version = "0.12", default-features = false, features = ["blocking", "json", "multipart", "rustls-tls", "http2", "stream"] }
rstest = "0.26.1"
rstest_reuse = "0.7.0"
rustls = { version = "0.23", default-features = false, features = ["ring", "std", "tls12"] }
rustls-native-certs = "0.8"
serde = { version = "1.0", features = ["derive"] }
serde_json = { version = "1.0", features = ["float_roundtrip"] }
serde_with = { version = "=3.16.1", default-features = false, features = ["std", "macros"] }
sha2 = "0.10"
subtle = "2"
thiserror = "2.0"
@ -39,6 +46,7 @@ base64 = "0.22"
moka = { version = "0.12.16", features = ["future"] }
strum = { version = "0.28.0", features = ["derive"] }
url = "2.5.8"
time = { version = "0.3.53", features = ["parsing"] }
criterion = "0.8.2"
veil = "0.3.0"

45
litellm-rust/crates/auth-azure/src/resolve.rs
View file

@ -657,4 +657,49 @@ mod tests {
assert!(matches!(error, Error::CredentialChain(errors) if errors.len() == 2));
}
#[derive(Debug)]
struct CallerToken(&'static str);
impl litellm_auth::TokenProvider for CallerToken {
fn acquire(&self) -> litellm_auth::TokenFuture<'_> {
Box::pin(async move {
Ok(ResolvedCredential::AccessToken {
token: SecretValue::new(self.0),
expires_on: None,
})
})
}
}
fn caller_inputs(token: &'static str) -> AzureAuthInputs {
let params = json!({"azure_ad_token": "static-token"});
AzureAuthInputs {
azure_ad_token_provider: Some(litellm_auth::TokenProviderHandle::new(Arc::new(
CallerToken(token),
))),
..AzureAuthInputs::from_optional_params(params.as_object().unwrap()).unwrap()
}
}
#[tokio::test]
async fn caller_token_is_chosen_over_supplied_static_token() {
let credential = AzureAuthService::default()
.get_azure_ad_token(&caller_inputs("caller-token"), &|_| None)
.await
.unwrap()
.unwrap();
assert_eq!(credential.value().secret().expose(), "caller-token");
}
#[tokio::test]
async fn empty_caller_token_is_rejected() {
let error = AzureAuthService::default()
.get_azure_ad_token(&caller_inputs(""), &|_| None)
.await
.unwrap_err();
assert!(matches!(error, Error::EmptyAzureToken));
}
}

15
litellm-rust/crates/auth/src/credential.rs
View file

@ -9,21 +9,6 @@ use crate::Error;
use super::{ResolvedCredential, SecretValue, TokenProviderHandle};
pub fn credential_index(requested: &str, names: &[String]) -> Option<usize> {
names.iter().position(|name| name == requested)
}
pub fn credential_default_fields<'a>(
supplied: &[String],
credential_fields: &'a [String],
) -> Vec<&'a str> {
credential_fields
.iter()
.filter(|name| !supplied.contains(name))
.map(String::as_str)
.collect()
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum CredentialFileRef {
Path(PathBuf),

1
litellm-rust/crates/auth/src/lib.rs
View file

@ -47,7 +47,6 @@ impl<T> Sourced<T> {
pub use credential::{
CredentialFileRef, CredentialLookup, CredentialLookupFuture, CredentialPlan,
CredentialPlanResolution, CredentialRef, CredentialResolver, CredentialResolverHandle,
credential_default_fields, credential_index,
};
pub use error::Error;
pub use http::{CredentialPlacement, RequestAuth};

17
litellm-rust/crates/callbacks-legacy/AGENTS.md Normal file
View file

@ -0,0 +1,17 @@
- Target invariants, not completion claims
- Keep this crate the legacy `@client` wrapper as the native call sees it, and nothing else: the `Logging` contract (`function_setup`, the deployment hooks, `pre_call`/`post_call`, the sync and async success and failure fan-out, the deferred proxy release, the argument sharing those callbacks rely on) plus the kwargs rewrites the wrapper makes on the way in (credential-name inheritance, the budget and retry-count limits)
- The driver in `litellm-host-python`, the routes and core see one `CallbackAdapter`; they never learn which Python objects consume a call
- `PublicCall` is the caller's call as `Logging` sees it: the positional arguments, the keyword view as the legacy path rewrites it (setup, deployment hook, prepare) and the bound request object whose attributes back keywords the caller omitted; routes hand it over through `run_legacy_call` and keep no copy
- `setup` decides once who owns the `Logging` instance and returns it as `CallSetup.bridge_owned`; `PythonLogger` carries it and nothing on the instance records it
- A logger the caller passed as `litellm_logging_obj` is caller-owned and observed in full, because the caller reads it after the call; the proxy is the live case
- A logger `function_setup` built for this call is bridge-owned, so each fan-out phase is skipped when `callbacks_needed` finds no registry, dynamic callback, `logger_fn` or debug switch for it; cost, timing and response metadata still run
- Callbacks receive the caller's own objects and may mutate them; this crate alone carries that obligation
- Retain complete boundary arguments, opaque unknown values, aliases, omitted/default distinctions and deliberate copies; preserve the established deployment-hook kwargs view
- Re-alias every `passthrough_fields` body key to the caller's object before `pre_call`; a keyword wins over the request attribute even when it is an explicit `None`
- Retain independently captured body/header roots from `pre_call` to `post_call`; in-place mutation reaches the wire, envelope field replacement is visible to later callbacks only
- A later kind of callback host (WASM, in-process Rust) has none of these obligations, so they stay out of `litellm-callbacks`, `litellm-host-python` and the bridge; the only facts that cross from the route are the prepared keyword view and `RequestContext.passthrough_fields`
- Success and failure handlers receive the exact selected public response or exception; logging projections, redaction and snapshots keep their own copy contracts
- Ordinary failure-handler errors cannot suppress the other eligible family or replace the mapped provider error; a cancellation ends the call with no further dispatch
- Dispatch errors never replay provider work or trigger the opposite outcome; the proxy's acceptance or rejection releases deferred success at most once
- Delivery follows the registry, not the callable's type: direct, awaited, executor-submitted, logging-worker and deferred paths stay distinct
- Traverse every retained Python edge; `close` is idempotent and restores the correlation context once

16
litellm-rust/crates/callbacks-legacy/Cargo.toml Normal file
View file

@ -0,0 +1,16 @@
[package]
name = "litellm-callbacks-legacy"
version = "0.1.0"
edition.workspace = true
license.workspace = true
repository.workspace = true
autotests = false
[dependencies]
litellm-callbacks.workspace = true
litellm-host-python.workspace = true
pyo3.workspace = true
[dev-dependencies]
rstest.workspace = true
serde_json.workspace = true

385
litellm-rust/crates/callbacks-legacy/src/adapter.rs Normal file
View file

@ -0,0 +1,385 @@
//! The legacy `Logging` contract as one adapter: every event and interception the driver
//! raises is answered with the same `Logging` calls, in the same order, as the Python
//! `@client` path makes them.
use litellm_callbacks::event::{CallEvent, FailureOrigin, RequestContext, Timing, WireRequest};
use litellm_host_python::{
AdapterStep, CallbackAdapter, PublicValue, from_py, missing_state, to_py,
};
use pyo3::{
exceptions::{PyBaseException, PyException},
gc::{PyTraverseError, PyVisit},
prelude::*,
types::PyDict,
};
use crate::{
DeploymentHooks, LegacyCallbacks, PublicCall, PythonLogger,
deferred::{PendingLogging, PendingSuccess},
finalize, is_internal_call, prepare, setup,
};
/// What the legacy contract needs to know about the route it is logging.
#[derive(Clone, Copy, Debug)]
pub struct LegacySurface {
pub call_type: &'static str,
/// What `Logging.pre_call` is told the input was.
pub input_description: &'static str,
}
enum Pending {
DeploymentPreCall,
DeploymentPostCall,
DeploymentFailure,
AsyncFailure,
}
pub struct LegacyLogging {
surface: LegacySurface,
call: PublicCall,
logger: Option<PythonLogger>,
start: Py<PyAny>,
end: Option<Py<PyAny>>,
response: Option<Py<PyAny>>,
error: Option<Py<PyBaseException>>,
body: Option<Py<PyDict>>,
headers: Option<Py<PyDict>>,
asynchronous: bool,
internal: bool,
pending: Option<Pending>,
}
fn datetime(py: Python<'_>, epoch_seconds: f64) -> PyResult<Py<PyAny>> {
py.import("datetime")?
.getattr("datetime")?
.call_method1("fromtimestamp", (epoch_seconds,))
.map(Bound::unbind)
}
fn is_cancellation(py: Python<'_>, error: &PyErr) -> bool {
!error.is_instance_of::<PyException>(py)
}
impl LegacyLogging {
pub fn new(
py: Python<'_>,
surface: LegacySurface,
call: PublicCall,
asynchronous: bool,
) -> Self {
Self {
surface,
call,
logger: None,
start: py.None(),
end: None,
response: None,
error: None,
body: None,
headers: None,
asynchronous,
internal: false,
pending: None,
}
}
/// Deployment hooks are awaited, and Python's synchronous `@client` wrapper never
/// runs them.
fn deployment_hooks(&self, py: Python<'_>) -> PyResult<bool> {
Ok(self.asynchronous && DeploymentHooks::needed(py)?)
}
fn logger(&self) -> PyResult<&PythonLogger> {
self.logger.as_ref().ok_or_else(|| {
pyo3::exceptions::PyRuntimeError::new_err("call logging is not initialized")
})
}
fn prepare(&mut self, py: Python<'_>) -> PyResult<AdapterStep> {
let prepared = prepare(py, self.call.kwargs().bind(py), self.logger()?)?.unbind();
self.call.set_kwargs(prepared);
Ok(AdapterStep::Arguments(self.call.kwargs().clone_ref(py)))
}
fn finalize(&mut self, py: Python<'_>) -> PyResult<AdapterStep> {
finalize(
py,
&self.response,
self.logger()?,
self.call.kwargs(),
&self.start,
&self.end,
)?;
self.response
.as_ref()
.map(|response| AdapterStep::Response(response.clone_ref(py)))
.ok_or_else(missing_state)
}
fn dispatch_success(&self, py: Python<'_>) -> PyResult<()> {
match self.try_dispatch_success(py) {
Err(error) if error.is_instance_of::<PyException>(py) => {
error.write_unraisable(py, self.logger.as_ref().map(|logger| logger.object(py)));
Ok(())
}
result => result,
}
}
fn try_dispatch_success(&self, py: Python<'_>) -> PyResult<()> {
let logger = self.logger()?;
let pending = || PendingSuccess {
logger: logger.clone_ref(py),
response: self.response.as_ref().map(|value| value.clone_ref(py)),
start: self.start.clone_ref(py),
end: self.end.as_ref().map(|value| value.clone_ref(py)),
};
if !self.asynchronous {
return pending().sync(py);
}
if !self.internal
&& self
.call
.kwargs()
.bind(py)
.get_item("fallbacks")?
.is_none_or(|value| value.is_none())
{
if !logger.callbacks_needed(py, "async_success")? {
logger.success_bookkeeping(py, &self.response, &self.start, &self.end, true)?;
} else if logger.defers_async_logging(py) {
let pending = Py::new(
py,
PendingLogging {
pending: Some(pending()),
},
)?;
logger.defer_success(py, pending.bind(py).as_any())?;
} else {
pending().asynchronous(py)?;
}
}
logger.sync_success_for_async_call(py, &self.response, &self.start, &self.end)
}
/// The sync failure handler, then the async one for async calls. Ordinary handler
/// errors never replace the selected failure or suppress the other family; a
/// cancellation does end the call.
fn dispatch_failure(&mut self, py: Python<'_>) -> PyResult<AdapterStep> {
let (Some(logger), Some(error)) = (&self.logger, &self.error) else {
return Ok(AdapterStep::Done);
};
if self.asynchronous && self.internal {
return Ok(AdapterStep::Done);
}
if let Err(failure) = logger.failure(py, error, &self.start, &self.end, false)
&& is_cancellation(py, &failure)
{
return Err(failure);
}
if !self.asynchronous {
return Ok(AdapterStep::Done);
}
match logger.failure(py, error, &self.start, &self.end, true) {
Ok(Some(awaitable)) => {
self.pending = Some(Pending::AsyncFailure);
Ok(AdapterStep::Await(awaitable))
}
Ok(None) => Ok(AdapterStep::Done),
Err(failure) if is_cancellation(py, &failure) => Err(failure),
Err(_) => Ok(AdapterStep::Done),
}
}
}
impl CallbackAdapter for LegacyLogging {
fn begin(
&mut self,
py: Python<'_>,
arguments: Py<PyDict>,
started_at: f64,
) -> PyResult<AdapterStep> {
self.call.set_kwargs(arguments);
self.start = datetime(py, started_at)?;
self.internal = is_internal_call(py)?;
let result = setup(
py,
self.surface.call_type,
self.call.args(),
self.call.kwargs(),
&self.start,
self.asynchronous,
)?;
self.logger = Some(result.logger()?);
self.call.set_kwargs(result.kwargs()?);
if self.deployment_hooks(py)? {
self.pending = Some(Pending::DeploymentPreCall);
return Ok(AdapterStep::Await(DeploymentHooks::before_call(
py,
self.call.kwargs(),
self.surface.call_type,
)?));
}
self.prepare(py)
}
fn before_send(
&mut self,
py: Python<'_>,
wire: Box<WireRequest>,
context: &RequestContext,
) -> PyResult<AdapterStep> {
let logger = self.logger()?;
logger.update_from_kwargs(py, self.call.kwargs(), &wire, context)?;
if !logger.callbacks_needed(py, "payload")? {
logger.record_api_call_start(py)?;
return Ok(AdapterStep::Wire(wire));
}
let body = to_py(py, &wire.body)?
.into_bound(py)
.cast_into::<PyDict>()?;
for name in context.passthrough_fields.iter() {
if let Some(value) = self.call.lookup(py, name)? {
body.set_item(name, value)?;
}
}
let headers = PyDict::new(py);
for (name, value) in &wire.headers {
headers.set_item(name, value)?;
}
self.body = Some(body.clone().unbind());
self.headers = Some(headers.clone().unbind());
let api_key = self.call.lookup(py, "api_key")?;
self.logger()?.pre_call(
py,
self.surface.input_description,
api_key.as_ref(),
&body,
&headers,
&wire.url,
)?;
let headers = headers
.iter()
.map(|(name, value)| Ok((name.extract::<String>()?, value.extract::<String>()?)))
.collect::<PyResult<Vec<_>>>()?;
Ok(AdapterStep::Wire(Box::new(WireRequest {
body: from_py(&body)?,
headers,
..*wire
})))
}
fn after_success(
&mut self,
py: Python<'_>,
response: Py<PyAny>,
timing: Timing,
) -> PyResult<AdapterStep> {
self.end = Some(datetime(py, timing.end_time)?);
self.response = Some(response);
if self.deployment_hooks(py)? {
self.pending = Some(Pending::DeploymentPostCall);
return Ok(AdapterStep::Await(DeploymentHooks::after_success(
py,
self.call.kwargs(),
&self.response,
self.surface.call_type,
)?));
}
self.finalize(py)
}
fn emit(
&mut self,
py: Python<'_>,
event: &CallEvent,
public: Option<PublicValue<'_>>,
) -> PyResult<AdapterStep> {
match (event, public) {
(CallEvent::ResponseReceived { raw }, _) => {
let logger = self.logger()?;
if logger.callbacks_needed(py, "payload")? {
logger.post_call(py, &raw.body, self.body.as_ref(), self.headers.as_ref())?;
}
Ok(AdapterStep::Done)
}
(CallEvent::Succeeded { timing }, Some(PublicValue::Response(response))) => {
self.end = Some(datetime(py, timing.end_time)?);
self.response = Some(response.clone_ref(py));
self.dispatch_success(py)?;
Ok(AdapterStep::Done)
}
(CallEvent::Failed { timing, origin }, Some(PublicValue::Error(error))) => {
self.end = Some(datetime(py, timing.end_time)?);
self.error = Some(error.clone_ref(py).into_value(py));
if *origin == FailureOrigin::Call
&& self.logger.is_some()
&& self.deployment_hooks(py)?
{
let error = self.error.as_ref().ok_or_else(missing_state)?;
self.pending = Some(Pending::DeploymentFailure);
return Ok(AdapterStep::Await(DeploymentHooks::after_failure(
py,
self.call.kwargs(),
error,
self.surface.call_type,
)?));
}
self.dispatch_failure(py)
}
_ => Err(missing_state()),
}
}
fn resume(&mut self, py: Python<'_>, result: PyResult<Py<PyAny>>) -> PyResult<AdapterStep> {
match self.pending.take().ok_or_else(missing_state)? {
Pending::DeploymentPreCall => {
self.call
.set_kwargs(result?.into_bound(py).cast_into::<PyDict>()?.unbind());
self.prepare(py)
}
Pending::DeploymentPostCall => {
self.response = Some(result?);
self.finalize(py)
}
Pending::DeploymentFailure => self.dispatch_failure(py),
Pending::AsyncFailure => match result {
Err(failure) if is_cancellation(py, &failure) => Err(failure),
_ => Ok(AdapterStep::Done),
},
}
}
fn close(&mut self, py: Python<'_>) {
if let Some(logger) = self.logger.take()
&& let Err(error) = logger.restore_context(py)
{
error.write_unraisable(py, None);
}
self.body = None;
self.headers = None;
}
fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
self.call.traverse(visit)?;
if let Some(logger) = &self.logger {
logger.traverse(visit)?;
}
visit.call(&self.start)?;
visit.call(&self.end)?;
visit.call(&self.response)?;
visit.call(&self.error)?;
visit.call(&self.body)?;
visit.call(&self.headers)
}
}
#[cfg(test)]
#[path = "../tests/deployment_hooks.rs"]
mod deployment_hooks_tests;
#[cfg(test)]
#[path = "../tests/payload.rs"]
mod payload_tests;
#[cfg(test)]
#[path = "../tests/terminal.rs"]
mod terminal_tests;

179
litellm-rust/crates/callbacks-legacy/src/call.rs Normal file
View file

@ -0,0 +1,179 @@
//! The caller's public call as the legacy `Logging` contract sees it. Legacy callbacks
//! receive these exact objects and may mutate them, so the call keeps them for its whole
//! lifetime. No other callback host has that obligation, which is why nothing outside
//! this crate holds them.
use litellm_callbacks::{machine::Machine, route::Route};
use litellm_host_python::{RouteHost, run_call};
use pyo3::{
gc::{PyTraverseError, PyVisit},
prelude::*,
types::{PyDict, PyTuple},
};
use crate::{LegacyLogging, LegacySurface};
pub struct PublicCall {
args: Py<PyTuple>,
kwargs: Py<PyDict>,
request: Py<PyAny>,
}
impl PublicCall {
/// Copies the keyword arguments once, so the legacy path's rewrites never reach the
/// caller's own dict while every value keeps its identity.
pub fn capture(
request: &Bound<'_, PyAny>,
args: &Bound<'_, PyTuple>,
kwargs: &Bound<'_, PyDict>,
) -> PyResult<Self> {
Ok(Self {
args: args.clone().unbind(),
kwargs: kwargs.copy()?.unbind(),
request: request.clone().unbind(),
})
}
pub(crate) fn args(&self) -> &Py<PyTuple> {
&self.args
}
/// The keyword view the legacy path currently reads: the caller's copy until
/// `function_setup`, then each rewrite (setup, deployment hook, prepare) in turn.
pub(crate) fn kwargs(&self) -> &Py<PyDict> {
&self.kwargs
}
pub(crate) fn set_kwargs(&mut self, kwargs: Py<PyDict>) {
self.kwargs = kwargs;
}
pub(crate) fn lookup<'py>(
&self,
py: Python<'py>,
name: &str,
) -> PyResult<Option<Bound<'py, PyAny>>> {
lookup(self.kwargs.bind(py), self.request.bind(py), name)
}
pub(crate) fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
visit.call(&self.args)?;
visit.call(&self.kwargs)?;
visit.call(&self.request)
}
}
/// The caller's own object for a public argument, as every legacy reader resolves it: the
/// keyword if given, even an explicit `None`, else the bound request's attribute. A route
/// host projecting from the prepared keyword view uses the same rule, so the callbacks
/// and the provider see one object per argument.
pub fn lookup<'py>(
kwargs: &Bound<'py, PyDict>,
request: &Bound<'py, PyAny>,
name: &str,
) -> PyResult<Option<Bound<'py, PyAny>>> {
if let Some(value) = kwargs.get_item(name)? {
return Ok(Some(value));
}
request.getattr_opt(name)
}
/// Runs one native call under the legacy `Logging` contract: the route host projects from
/// the keyword view the contract prepares, and the contract observes the call.
pub fn run_legacy_call<H, M>(
py: Python<'_>,
surface: LegacySurface,
call: PublicCall,
machine: M,
route: H,
asynchronous: bool,
) -> PyResult<Py<PyAny>>
where
H: RouteHost + 'static,
M: Machine<Route = H::Route, Complete = <H::Route as Route>::Response> + 'static,
{
let arguments = call.kwargs.clone_ref(py);
run_call(
py,
machine,
route,
Box::new(LegacyLogging::new(py, surface, call, asynchronous)),
arguments,
asynchronous,
)
}
#[cfg(test)]
mod tests {
use super::*;
fn capture<'py>(py: Python<'py>, source: &std::ffi::CStr) -> (PublicCall, Bound<'py, PyDict>) {
let locals = PyDict::new(py);
py.run(source, Some(&locals), Some(&locals)).unwrap();
let request = locals.get_item("request").unwrap().unwrap();
let kwargs = locals
.get_item("kwargs")
.unwrap()
.unwrap()
.cast_into::<PyDict>()
.unwrap();
let call = PublicCall::capture(&request, &PyTuple::empty(py), &kwargs).unwrap();
(call, locals)
}
#[test]
fn lookup_prefers_the_keyword_even_when_none_and_falls_back_to_the_request() {
Python::initialize();
Python::attach(|py| {
let (call, locals) = capture(
py,
c"
key = object()
document = {'type': 'document_url'}
class Request:
api_key = 'from-request'
api_base = 'from-request'
document = document
request = Request()
kwargs = {'api_key': key, 'api_base': None}
",
);
let key = locals.get_item("key").unwrap().unwrap();
let document = locals.get_item("document").unwrap().unwrap();
assert!(call.lookup(py, "api_key").unwrap().unwrap().is(&key));
assert!(call.lookup(py, "api_base").unwrap().unwrap().is_none());
assert!(call.lookup(py, "document").unwrap().unwrap().is(&document));
assert!(call.lookup(py, "model").unwrap().is_none());
});
}
#[test]
fn capture_copies_the_keyword_dict_without_copying_its_values() {
Python::initialize();
Python::attach(|py| {
let (call, locals) = capture(
py,
c"
pages = [0]
class Request:
pass
request = Request()
kwargs = {'pages': pages}
",
);
let caller = locals
.get_item("kwargs")
.unwrap()
.unwrap()
.cast_into::<PyDict>()
.unwrap();
call.kwargs()
.bind(py)
.set_item("litellm_call_id", "call")
.unwrap();
assert!(!caller.contains("litellm_call_id").unwrap());
let pages = locals.get_item("pages").unwrap().unwrap();
assert!(call.lookup(py, "pages").unwrap().unwrap().is(&pages));
});
}
}

404
litellm-rust/crates/callbacks-legacy/src/callbacks.rs Normal file
View file

@ -0,0 +1,404 @@
//! Callback fan-out over litellm's `Logging` object: which callbacks are registered,
//! the deferred and worker-submitted success paths, and the sync-callbacks-for-async-calls
//! duplication. All of it expires with the legacy callback contract.
use litellm_callbacks::event::{RequestContext, WireRequest};
use litellm_host_python::to_py;
use pyo3::{exceptions::PyBaseException, prelude::*, types::PyDict};
use crate::logger::PythonLogger;
pub trait LegacyCallbacks {
fn callbacks_needed(&self, py: Python<'_>, phase: &str) -> PyResult<bool>;
/// `Logging.update_from_kwargs`: what the logger is told about the request it is
/// about to see, with consumed credentials redacted.
fn update_from_kwargs(
&self,
py: Python<'_>,
kwargs: &Py<PyDict>,
wire: &WireRequest,
context: &RequestContext,
) -> PyResult<()>;
fn record_api_call_start(&self, py: Python<'_>) -> PyResult<()>;
/// `Logging.pre_call`, or its payload-free shortcut when no input callback listens.
fn pre_call(
&self,
py: Python<'_>,
input: &str,
api_key: Option<&Bound<'_, PyAny>>,
body: &Bound<'_, PyDict>,
headers: &Bound<'_, PyDict>,
url: &str,
) -> PyResult<()>;
/// `Logging.post_call`, or its payload-free shortcut when no input callback listens.
fn post_call(
&self,
py: Python<'_>,
original_response: &str,
body: Option<&Py<PyDict>>,
headers: Option<&Py<PyDict>>,
) -> PyResult<()>;
fn defers_async_logging(&self, py: Python<'_>) -> bool;
fn defer_success(&self, py: Python<'_>, pending: &Bound<'_, PyAny>) -> PyResult<()>;
fn sync_success_for_async_call(
&self,
py: Python<'_>,
response: &Option<Py<PyAny>>,
start: &Py<PyAny>,
end: &Option<Py<PyAny>>,
) -> PyResult<()>;
fn failure(
&self,
py: Python<'_>,
error: &Py<PyBaseException>,
start: &Py<PyAny>,
end: &Option<Py<PyAny>>,
asynchronous: bool,
) -> PyResult<Option<Py<PyAny>>>;
fn submit_success(
&self,
py: Python<'_>,
response: &Option<Py<PyAny>>,
start: &Py<PyAny>,
end: &Option<Py<PyAny>>,
) -> PyResult<()>;
fn enqueue_success(
&self,
py: Python<'_>,
response: &Option<Py<PyAny>>,
start: &Py<PyAny>,
end: &Option<Py<PyAny>>,
) -> PyResult<()>;
}
impl LegacyCallbacks for PythonLogger {
fn callbacks_needed(&self, py: Python<'_>, phase: &str) -> PyResult<bool> {
if !self.bridge_owned() {
return Ok(true);
}
py.import("litellm.rust_bridge.legacy_callbacks")?
.getattr("callbacks_needed")?
.call1((self.object(py), phase))?
.extract()
}
fn update_from_kwargs(
&self,
py: Python<'_>,
kwargs: &Py<PyDict>,
wire: &WireRequest,
context: &RequestContext,
) -> PyResult<()> {
let secret_fields: Vec<&str> = context.secret_fields.iter().map(String::as_str).collect();
let update = PyDict::new(py);
update.set_item("kwargs", redact(py, kwargs.bind(py), &secret_fields)?)?;
update.set_item("model", &context.model)?;
update.set_item(
"optional_params",
redact(
py,
&to_py(py, &context.optional_params)?
.into_bound(py)
.cast_into::<PyDict>()?,
&secret_fields,
)?,
)?;
let params = PyDict::new(py);
params.set_item(
"litellm_call_id",
kwargs.bind(py).get_item("litellm_call_id")?,
)?;
params.set_item("api_base", &wire.url)?;
for name in ["logger_fn", "litellm_request_debug"] {
if let Some(value) = kwargs.bind(py).get_item(name)? {
params.set_item(name, value)?;
}
}
for name in custom_pricing_fields(py)? {
if let Some(value) = kwargs.bind(py).get_item(&name)?
&& !value.is_none()
{
params.set_item(name, value)?;
}
}
update.set_item("litellm_params", params)?;
update.set_item("custom_llm_provider", &context.custom_llm_provider)?;
self.object(py)
.call_method("update_from_kwargs", (), Some(&update))?;
Ok(())
}
fn record_api_call_start(&self, py: Python<'_>) -> PyResult<()> {
self.object(py).call_method0("record_api_call_start_time")?;
Ok(())
}
fn pre_call(
&self,
py: Python<'_>,
input: &str,
api_key: Option<&Bound<'_, PyAny>>,
body: &Bound<'_, PyDict>,
headers: &Bound<'_, PyDict>,
url: &str,
) -> PyResult<()> {
let additional = PyDict::new(py);
additional.set_item("complete_input_dict", body)?;
additional.set_item("headers", headers)?;
additional.set_item("api_base", url)?;
let kwargs = PyDict::new(py);
kwargs.set_item("input", input)?;
kwargs.set_item("api_key", api_key)?;
kwargs.set_item("additional_args", &additional)?;
if self.callbacks_needed(py, "input")? {
self.object(py).call_method("pre_call", (), Some(&kwargs))?;
} else {
self.object(py)
.call_method("_pre_call", (), Some(&kwargs))?;
self.record_api_call_start(py)?;
}
Ok(())
}
fn post_call(
&self,
py: Python<'_>,
original_response: &str,
body: Option<&Py<PyDict>>,
headers: Option<&Py<PyDict>>,
) -> PyResult<()> {
let additional = PyDict::new(py);
additional.set_item("complete_input_dict", body)?;
additional.set_item("headers", headers)?;
if self.callbacks_needed(py, "input")? {
let kwargs = PyDict::new(py);
kwargs.set_item("original_response", original_response)?;
kwargs.set_item("additional_args", &additional)?;
self.object(py)
.call_method("post_call", (), Some(&kwargs))?;
} else {
let response = py
.import("json")?
.call_method1("dumps", (original_response,))?;
self.object(py).call_method1(
"record_post_call",
(response, py.None(), py.None(), additional),
)?;
}
Ok(())
}
fn defers_async_logging(&self, py: Python<'_>) -> bool {
self.object(py)
.getattr("_defer_async_logging")
.is_ok_and(|value| value.is_truthy().unwrap_or(false))
}
fn defer_success(&self, py: Python<'_>, pending: &Bound<'_, PyAny>) -> PyResult<()> {
self.object(py).setattr("_native_pending_logging", pending)
}
fn sync_success_for_async_call(
&self,
py: Python<'_>,
response: &Option<Py<PyAny>>,
start: &Py<PyAny>,
end: &Option<Py<PyAny>>,
) -> PyResult<()> {
if !self.callbacks_needed(py, "sync_success_async")? {
return Ok(());
}
self.object(py).call_method1(
"handle_sync_success_callbacks_for_async_calls",
(response, start, end),
)?;
Ok(())
}
fn failure(
&self,
py: Python<'_>,
error: &Py<PyBaseException>,
start: &Py<PyAny>,
end: &Option<Py<PyAny>>,
asynchronous: bool,
) -> PyResult<Option<Py<PyAny>>> {
if !self.callbacks_needed(
py,
if asynchronous {
"async_failure"
} else {
"sync_failure"
},
)? {
py.import("litellm.rust_bridge.legacy_callbacks")?
.getattr("failure_bookkeeping")?
.call1((self.object(py), error, start, end, asynchronous))?;
return Ok(None);
}
let trace = py
.import("traceback")?
.getattr("format_exception")?
.call1((error,))?;
let trace = pyo3::types::PyString::new(py, "").call_method1("join", (trace,))?;
let value = self.object(py).call_method1(
if asynchronous {
"async_failure_handler"
} else {
"failure_handler"
},
(error, trace, start, end),
)?;
Ok(asynchronous.then(|| value.unbind()))
}
fn submit_success(
&self,
py: Python<'_>,
response: &Option<Py<PyAny>>,
start: &Py<PyAny>,
end: &Option<Py<PyAny>>,
) -> PyResult<()> {
if !self.callbacks_needed(py, "sync_success")? {
return self.success_bookkeeping(py, response, start, end, false);
}
let context = py.import("contextvars")?.call_method0("copy_context")?;
py.import("litellm.litellm_core_utils.litellm_logging")?
.getattr("executor")?
.call_method1(
"submit",
(
context.getattr("run")?,
self.object(py).getattr("success_handler")?,
response,
start,
end,
),
)?;
Ok(())
}
fn enqueue_success(
&self,
py: Python<'_>,
response: &Option<Py<PyAny>>,
start: &Py<PyAny>,
end: &Option<Py<PyAny>>,
) -> PyResult<()> {
if !self.callbacks_needed(py, "async_success")? {
return self.success_bookkeeping(py, response, start, end, true);
}
let context = py.import("contextvars")?.call_method0("copy_context")?;
let worker = py
.import("litellm.litellm_core_utils.logging_worker")?
.getattr("GLOBAL_LOGGING_WORKER")?
.getattr("ensure_initialized_and_enqueue")?;
let coroutine = self
.object(py)
.call_method1("async_success_handler", (response, start, end))?;
let enqueue = context.call_method1("run", (worker, &coroutine));
if enqueue.is_err()
&& let Err(error) = coroutine.call_method0("close")
{
error.write_unraisable(py, Some(&coroutine));
}
enqueue.map(|_| ())
}
}
fn custom_pricing_fields(py: Python<'_>) -> PyResult<Vec<String>> {
py.import("litellm.types.utils")?
.getattr("CustomPricingLiteLLMParams")?
.getattr("model_fields")?
.cast_into::<PyDict>()?
.keys()
.iter()
.map(|name| name.extract::<String>())
.collect()
}
fn redact(
py: Python<'_>,
params: &Bound<'_, PyDict>,
secret_fields: &[&str],
) -> PyResult<Py<PyDict>> {
let redacted = PyDict::new(py);
for (name, value) in params {
let name = name.extract::<String>()?;
if name == "proxy_server_request" {
continue;
}
if secret_fields.contains(&name.as_str()) {
redacted.set_item(name, "****")?;
} else {
redacted.set_item(name, value)?;
}
}
Ok(redacted.unbind())
}
/// Proxy-internal calls skip the legacy success fan-out.
pub fn is_internal_call(py: Python<'_>) -> PyResult<bool> {
py.import("litellm._internal_context")?
.getattr("is_internal_call")?
.call_method0("get")?
.extract()
}
#[cfg(test)]
mod tests {
use pyo3::types::PyDict;
use super::*;
fn logger_whose_registries_need_no_input(py: Python<'_>, bridge_owned: bool) -> PythonLogger {
let locals = PyDict::new(py);
py.run(
c"
import sys
import types
for name in ('litellm', 'litellm.rust_bridge', 'litellm.rust_bridge.legacy_callbacks'):
sys.modules.setdefault(name, types.ModuleType(name))
legacy = sys.modules['litellm.rust_bridge.legacy_callbacks']
legacy.callbacks_needed = lambda logger, phase: logger.needed.get(phase, True)
class Logger:
needed = {'input': False}
logger = Logger()
",
Some(&locals),
Some(&locals),
)
.unwrap();
PythonLogger::new(
locals.get_item("logger").unwrap().unwrap().unbind(),
bridge_owned,
)
}
#[test]
fn a_caller_owned_logger_is_observed_in_full() {
Python::initialize();
Python::attach(|py| {
let logger = logger_whose_registries_need_no_input(py, false);
assert!(logger.callbacks_needed(py, "input").unwrap());
});
}
#[test]
fn a_bridge_owned_logger_is_elided_where_no_registry_needs_it() {
Python::initialize();
Python::attach(|py| {
let logger = logger_whose_registries_need_no_input(py, true);
assert!(!logger.callbacks_needed(py, "input").unwrap());
assert!(logger.callbacks_needed(py, "payload").unwrap());
});
}
}

67
litellm-rust/crates/callbacks-legacy/src/deferred.rs Normal file
View file

@ -0,0 +1,67 @@
//! The proxy's deferred success release: the async success handler is queued only once
//! the proxy accepts the response, and at most once.
use pyo3::{exceptions::PyException, prelude::*};
use crate::{LegacyCallbacks, PythonLogger};
pub(crate) struct PendingSuccess {
pub(crate) logger: PythonLogger,
pub(crate) response: Option<Py<PyAny>>,
pub(crate) start: Py<PyAny>,
pub(crate) end: Option<Py<PyAny>>,
}
impl PendingSuccess {
pub(crate) fn sync(&self, py: Python<'_>) -> PyResult<()> {
self.logger
.submit_success(py, &self.response, &self.start, &self.end)
}
pub(crate) fn asynchronous(&self, py: Python<'_>) -> PyResult<()> {
self.logger
.enqueue_success(py, &self.response, &self.start, &self.end)
}
}
#[pyclass]
pub(crate) struct PendingLogging {
pub(crate) pending: Option<PendingSuccess>,
}
#[pymethods]
impl PendingLogging {
fn release(slf: &Bound<'_, Self>, py: Python<'_>, success: bool) -> PyResult<()> {
let pending = slf.borrow_mut().pending.take();
if let Some(pending) = pending
&& success
{
match pending.asynchronous(py) {
Err(error) if error.is_instance_of::<PyException>(py) => {
error.write_unraisable(py, Some(pending.logger.object(py)));
}
result => return result,
}
}
Ok(())
}
fn __traverse__(&self, visit: pyo3::gc::PyVisit<'_>) -> Result<(), pyo3::gc::PyTraverseError> {
if let Some(pending) = &self.pending {
pending.logger.traverse(&visit)?;
visit.call(&pending.response)?;
visit.call(&pending.start)?;
visit.call(&pending.end)?;
}
Ok(())
}
fn __clear__(slf: &Bound<'_, Self>) {
let pending = slf.borrow_mut().pending.take();
drop(pending);
}
}
#[cfg(test)]
#[path = "../tests/deferred.rs"]
mod tests;

27
litellm-rust/crates/callbacks-legacy/src/lib.rs Normal file
View file

@ -0,0 +1,27 @@
//! The legacy `@client` wrapper as the native call sees it: litellm's `Logging` object, the
//! sync and async callback registries it fans out to, the deployment hooks, the deferred
//! proxy release, and the kwargs rewrites the wrapper makes on the way in (credential-name
//! inheritance, budget and retry-count limits). All of it sits behind one
//! [`CallbackAdapter`](litellm_host_python::CallbackAdapter), so the driver, the routes and
//! core never learn which Python object is on the other end.
//!
//! Legacy callbacks receive the caller's own objects and may mutate them. [`PublicCall`]
//! is where those objects live, and [`run_legacy_call`] is how a route hands them over
//! without keeping a copy.
mod adapter;
mod call;
mod callbacks;
mod deferred;
mod logger;
mod preparation;
#[cfg(test)]
#[path = "../tests/support.rs"]
mod test_support;
pub(crate) use adapter::LegacyLogging;
pub use adapter::LegacySurface;
pub use call::{PublicCall, lookup, run_legacy_call};
pub(crate) use callbacks::{LegacyCallbacks, is_internal_call};
pub(crate) use logger::{DeploymentHooks, PythonLogger, finalize, setup};
pub(crate) use preparation::prepare;

236
litellm-rust/crates/callbacks-legacy/src/logger.rs Normal file
View file

@ -0,0 +1,236 @@
use pyo3::{
exceptions::PyBaseException,
gc::{PyTraverseError, PyVisit},
prelude::*,
types::{PyDict, PyTuple},
};
/// The `Logging` instance one call fans out through, and who owns it. A logger the caller
/// handed in is observed in full, because the caller reads it after the call; one this
/// crate built through `function_setup` is elided wherever no registry needs it.
pub struct PythonLogger {
object: Py<PyAny>,
bridge_owned: bool,
}
impl PythonLogger {
pub(crate) fn new(object: Py<PyAny>, bridge_owned: bool) -> Self {
Self {
object,
bridge_owned,
}
}
pub(crate) fn object<'py>(&self, py: Python<'py>) -> &Bound<'py, PyAny> {
self.object.bind(py)
}
pub(crate) fn bridge_owned(&self) -> bool {
self.bridge_owned
}
pub fn clone_ref(&self, py: Python<'_>) -> Self {
Self {
object: self.object.clone_ref(py),
bridge_owned: self.bridge_owned,
}
}
pub fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
visit.call(&self.object)
}
pub fn success_bookkeeping(
&self,
py: Python<'_>,
response: &Option<Py<PyAny>>,
start: &Py<PyAny>,
end: &Option<Py<PyAny>>,
asynchronous: bool,
) -> PyResult<()> {
py.import("litellm.rust_bridge.legacy_callbacks")?
.getattr("success_bookkeeping")?
.call1((self.object(py), response, start, end, asynchronous))?;
Ok(())
}
pub fn restore_context(&self, py: Python<'_>) -> PyResult<()> {
py.import("litellm.utils")?
.getattr("_restore_correlation_context_if_supported")?
.call1((self.object(py),))?;
Ok(())
}
}
/// A bare Python object was not obtained from `setup`, so it is caller-owned.
impl FromPyObject<'_, '_> for PythonLogger {
type Error = PyErr;
fn extract(object: Borrowed<'_, '_, PyAny>) -> PyResult<Self> {
Ok(Self::new(object.to_owned().unbind(), false))
}
}
pub struct SetupResult<'py>(Bound<'py, PyAny>);
impl SetupResult<'_> {
pub fn logger(&self) -> PyResult<PythonLogger> {
let object = self.0.getattr("logger")?.unbind();
let bridge_owned = self.0.getattr("bridge_owned")?.extract()?;
Ok(PythonLogger::new(object, bridge_owned))
}
pub fn kwargs(&self) -> PyResult<Py<PyDict>> {
Ok(self.0.getattr("kwargs")?.extract()?)
}
}
pub fn setup<'py>(
py: Python<'py>,
call_type: &str,
args: &Py<PyTuple>,
kwargs: &Py<PyDict>,
start: &Py<PyAny>,
asynchronous: bool,
) -> PyResult<SetupResult<'py>> {
py.import("litellm.rust_bridge.legacy_callbacks")?
.getattr("setup")?
.call1((call_type, args, kwargs, start, asynchronous))
.map(SetupResult)
}
pub fn finalize(
py: Python<'_>,
response: &Option<Py<PyAny>>,
logger: &PythonLogger,
kwargs: &Py<PyDict>,
start: &Py<PyAny>,
end: &Option<Py<PyAny>>,
) -> PyResult<()> {
py.import("litellm.rust_bridge.legacy_callbacks")?
.getattr("finalize")?
.call1((response, logger.object(py), kwargs, start, end))?;
Ok(())
}
pub struct DeploymentHooks;
impl DeploymentHooks {
pub fn needed(py: Python<'_>) -> PyResult<bool> {
py.import("litellm.rust_bridge.legacy_callbacks")?
.getattr("deployment_callbacks_needed")?
.call0()?
.extract()
}
pub fn before_call(
py: Python<'_>,
kwargs: &Py<PyDict>,
call_type: &str,
) -> PyResult<Py<PyAny>> {
py.import("litellm.utils")?
.getattr("async_pre_call_deployment_hook")?
.call1((kwargs, call_type))
.map(Bound::unbind)
}
pub fn after_success(
py: Python<'_>,
kwargs: &Py<PyDict>,
response: &Option<Py<PyAny>>,
call_type: &str,
) -> PyResult<Py<PyAny>> {
py.import("litellm.utils")?
.getattr("async_post_call_success_deployment_hook")?
.call1((kwargs, response, call_type))
.map(Bound::unbind)
}
pub fn after_failure(
py: Python<'_>,
kwargs: &Py<PyDict>,
error: &Py<PyBaseException>,
call_type: &str,
) -> PyResult<Py<PyAny>> {
py.import("litellm.utils")?
.getattr("async_post_call_failure_deployment_hook")?
.call1((kwargs, error, call_type))
.map(Bound::unbind)
}
}
#[cfg(test)]
mod tests {
use pyo3::exceptions::PyTypeError;
use super::*;
#[test]
fn setup_fields_are_checked_lazily() {
Python::initialize();
Python::attach(|py| {
let locals = PyDict::new(py);
py.run(
pyo3::ffi::c_str!(
r#"
reads = []
class Logger:
def __getattribute__(self, name):
reads.append(name)
raise AssertionError('logger methods must remain lazy')
logger = Logger()
class Setup:
@property
def logger(self):
reads.append('logger')
return logger
@property
def bridge_owned(self):
reads.append('bridge_owned')
return True
@property
def kwargs(self):
reads.append('kwargs')
return []
result = Setup()
"#
),
Some(&locals),
Some(&locals),
)
.unwrap();
let result = SetupResult(locals.get_item("result").unwrap().unwrap());
let logger = result.logger().unwrap();
assert!(
logger
.object(py)
.is(locals.get_item("logger").unwrap().unwrap())
);
assert!(logger.bridge_owned());
assert!(
result
.kwargs()
.unwrap_err()
.is_instance_of::<PyTypeError>(py)
);
assert_eq!(
locals
.get_item("reads")
.unwrap()
.unwrap()
.extract::<Vec<String>>()
.unwrap(),
["logger", "bridge_owned", "kwargs"]
);
});
}
#[test]
fn a_logger_extracted_from_a_bare_object_is_caller_owned() {
Python::initialize();
Python::attach(|py| {
let logger: PythonLogger = py.None().into_bound(py).extract().unwrap();
assert!(!logger.bridge_owned());
});
}
}

21
litellm-rust/crates/python-bridge/src/lifecycle/preparation.rs → litellm-rust/crates/callbacks-legacy/src/preparation.rs
View file

@ -1,6 +1,7 @@
use litellm_auth::{credential_default_fields, credential_index};
use pyo3::prelude::*;
use pyo3::types::{PyDict, PyList};
use pyo3::{
prelude::*,
types::{PyDict, PyList},
};
struct CredentialEntry<'py>(Bound<'py, PyAny>);
@ -14,16 +15,16 @@ impl<'py> CredentialEntry<'py> {
}
}
pub(super) fn prepare<'py>(
pub fn prepare<'py>(
py: Python<'py>,
kwargs: &Bound<'py, PyDict>,
logger: &super::PythonLogger,
logger: &crate::PythonLogger,
) -> PyResult<Bound<'py, PyDict>> {
let arguments = kwargs.copy()?;
arguments.set_item("litellm_logging_obj", logger.object(py))?;
let litellm = py.import("litellm")?;
inherit_credentials(py, &litellm, &arguments)?;
py.import("litellm.rust_bridge.lifecycle")?
py.import("litellm.rust_bridge.legacy_callbacks")?
.getattr("check_limits")?
.call1((&arguments,))?;
Ok(arguments)
@ -49,7 +50,7 @@ fn inherit_credentials(
.iter()
.map(|credential| CredentialEntry(credential).name())
.collect::<PyResult<Vec<_>>>()?;
let Some(index) = credential_index(&requested, &names) else {
let Some(index) = names.iter().position(|name| *name == requested) else {
py.import("litellm._logging")?.getattr("verbose_logger")?.call_method1(
"warning",
("litellm_credential_name=%s matched none of the %d loaded credentials; the request runs without it", requested, names.len()),
@ -60,9 +61,9 @@ fn inherit_credentials(
let values = selected.values()?;
let supplied: Vec<String> = arguments.keys().extract()?;
let fields: Vec<String> = values.keys().extract()?;
for name in credential_default_fields(&supplied, &fields) {
if let Some(value) = values.get_item(name)? {
arguments.set_item(name, value)?;
for name in fields.iter().filter(|name| !supplied.contains(name)) {
if let Some(value) = values.get_item(name.as_str())? {
arguments.set_item(name.as_str(), value)?;
}
}
Ok(())

162
litellm-rust/crates/callbacks-legacy/tests/deferred.rs Normal file
View file

@ -0,0 +1,162 @@
use std::ffi::CStr;
use pyo3::prelude::*;
use pyo3::types::PyDict;
use rstest::rstest;
use super::{PendingLogging, PendingSuccess};
use crate::PythonLogger;
use crate::test_support::{local, namespace, run};
/// A deferred success for the namespace's `logger` and `response`, bound as `pending`.
fn defer<'py>(py: Python<'py>, script: &CStr) -> Bound<'py, PyDict> {
let locals = namespace(py, c"response = object()");
run(py, &locals, script);
let pending = Py::new(
py,
PendingLogging {
pending: Some(PendingSuccess {
logger: PythonLogger::new(local(&locals, "logger").unbind(), true),
response: Some(local(&locals, "response").unbind()),
start: py.None(),
end: Some(py.None()),
}),
},
)
.unwrap();
locals.set_item("pending", pending).unwrap();
locals
}
#[test]
fn release_enqueues_the_success_once_in_the_releasing_context() {
Python::initialize();
Python::attach(|py| {
let locals = defer(
py,
c"
from contextvars import ContextVar
marker = ContextVar('marker', default='unset')
observed = []
def on_enqueue(coroutine):
observed.append(marker.get())
pending.release(True)
logger.on_enqueue = on_enqueue
",
);
run(
py,
&locals,
c"
marker.set('release')
pending.release(True)
pending.release(True)
assert observed == ['release'], observed
assert logger.names() == ['async_success_handler', 'enqueued'], logger.calls
assert logger.calls[0][1] is response
",
);
});
}
#[test]
fn a_blocked_release_drops_the_success_for_good() {
Python::initialize();
Python::attach(|py| {
let locals = defer(py, c"");
run(
py,
&locals,
c"
pending.release(False)
pending.release(True)
assert logger.calls == [], logger.calls
",
);
});
}
#[test]
fn a_release_after_the_async_callbacks_went_away_only_keeps_the_books() {
Python::initialize();
Python::attach(|py| {
let locals = defer(py, c"logger.needed = {'async_success': False}");
run(
py,
&locals,
c"
pending.release(True)
assert logger.calls == [('success_bookkeeping', True)], logger.calls
",
);
});
}
#[rstest]
#[case::ordinary_error(c"RuntimeError('queue full')", false)]
#[case::cancellation(c"asyncio.CancelledError()", true)]
fn a_failed_enqueue_closes_the_coroutine_and_is_never_replayed(
#[case] failure: &CStr,
#[case] propagates: bool,
) {
Python::initialize();
Python::attach(|py| {
let locals = defer(
py,
c"
import asyncio
def on_enqueue(coroutine):
raise failure
logger.on_enqueue = on_enqueue
",
);
locals
.set_item("failure", py.eval(failure, None, Some(&locals)).unwrap())
.unwrap();
let released = local(&locals, "pending").call_method1("release", (true,));
match released {
Ok(_) => assert!(!propagates),
Err(error) => {
assert!(propagates);
assert!(error.value(py).is(local(&locals, "failure")));
}
}
locals.set_item("propagates", propagates).unwrap();
run(
py,
&locals,
c"
pending.release(True)
assert logger.names() == ['async_success_handler', 'enqueued', 'closed'], logger.calls
assert unraisable_from(logger) == ([] if propagates else [failure])
",
);
});
}
#[test]
fn an_unreleased_success_does_not_keep_its_logger_alive() {
Python::initialize();
Python::attach(|py| {
let locals = defer(py, c"");
run(
py,
&locals,
c"
import gc
import weakref
logger.pending = pending
reference = weakref.ref(logger)
del logger, pending
gc.collect()
assert reference() is None
",
);
});
}

246
litellm-rust/crates/callbacks-legacy/tests/deployment_hooks.rs Normal file
View file

@ -0,0 +1,246 @@
use std::ffi::CStr;
use litellm_callbacks::event::{CallEvent, FailureOrigin, Timing};
use litellm_host_python::{AdapterStep, CallbackAdapter, PublicValue};
use pyo3::exceptions::asyncio::CancelledError;
use pyo3::prelude::*;
use pyo3::types::PyDict;
use rstest::rstest;
use super::LegacyLogging;
use crate::test_support::{legacy_call, local, namespace, run};
const CALL: &CStr = c"
document = {'type': 'document_url', 'document_url': 'data:application/pdf;base64,YWJj'}
kwargs = {'logger': logger, 'document': document}
";
const TIMING: Timing = Timing {
start_time: 0.0,
end_time: 1.0,
};
fn begin<'py>(
py: Python<'py>,
locals: &Bound<'py, PyDict>,
asynchronous: bool,
) -> (LegacyLogging, AdapterStep) {
let mut logging = legacy_call(py, locals, asynchronous);
let kwargs = local(locals, "kwargs")
.cast_into::<PyDict>()
.unwrap()
.unbind();
let step = logging.begin(py, kwargs, 0.0).unwrap();
(logging, step)
}
fn arguments<'py>(py: Python<'py>, step: AdapterStep) -> Bound<'py, PyDict> {
let AdapterStep::Arguments(arguments) = step else {
panic!("expected the prepared arguments");
};
arguments.into_bound(py)
}
fn awaits_deployment_hook(step: &AdapterStep) -> bool {
matches!(step, AdapterStep::Await(_))
}
#[rstest]
#[case::synchronous(false)]
#[case::asynchronous(true)]
fn deployment_pre_call_hook_runs_only_for_asynchronous_calls(#[case] asynchronous: bool) {
Python::initialize();
Python::attach(|py| {
let locals = namespace(py, CALL);
let (_, step) = begin(py, &locals, asynchronous);
assert_eq!(awaits_deployment_hook(&step), asynchronous);
let names: Vec<String> = local(&locals, "logger")
.call_method0("names")
.unwrap()
.extract()
.unwrap();
assert_eq!(names.contains(&"pre_hook".to_string()), asynchronous);
});
}
#[test]
fn kwargs_returned_by_the_pre_call_hook_are_what_the_call_prepares() {
Python::initialize();
Python::attach(|py| {
let locals = namespace(
py,
c"
document = {'type': 'document_url', 'document_url': 'data:application/pdf;base64,YWJj'}
replacement = {'type': 'document_url', 'document_url': 'data:application/pdf;base64,ZWRpdGVk'}
kwargs = {'logger': logger, 'document': document}
replaced_kwargs = {'logger': logger, 'document': replacement, 'pages': [0]}
",
);
let (mut logging, step) = begin(py, &locals, true);
assert!(awaits_deployment_hook(&step));
let step = logging
.resume(py, Ok(local(&locals, "replaced_kwargs").unbind()))
.unwrap();
locals.set_item("prepared", arguments(py, step)).unwrap();
run(
py,
&locals,
c"
assert prepared['document'] is replacement
assert prepared['pages'] is replaced_kwargs['pages']
assert prepared['litellm_logging_obj'] is logger
assert 'litellm_logging_obj' not in replaced_kwargs
[checked] = [value for name, value in logger.calls if name == 'check_limits']
assert checked is prepared
",
);
});
}
#[test]
fn response_returned_by_the_post_call_hook_is_finalized_and_returned() {
Python::initialize();
Python::attach(|py| {
let locals = namespace(
py,
c"
kwargs = {'logger': logger}
response = object()
replacement = object()
logger.hooks = {'pre': lambda kwargs: kwargs}
",
);
let (mut logging, _) = begin(py, &locals, true);
logging
.resume(py, Ok(local(&locals, "kwargs").unbind()))
.unwrap();
let step = logging
.after_success(py, local(&locals, "response").unbind(), TIMING)
.unwrap();
assert!(awaits_deployment_hook(&step));
let step = logging
.resume(py, Ok(local(&locals, "replacement").unbind()))
.unwrap();
let AdapterStep::Response(returned) = step else {
panic!("expected the finalized response");
};
assert!(returned.bind(py).is(local(&locals, "replacement")));
run(
py,
&locals,
c"
[finalized] = [value for name, value in logger.calls if name == 'finalize']
assert finalized is replacement
",
);
});
}
#[rstest]
#[case::pre_call(false)]
#[case::post_call(true)]
fn cancelling_a_deployment_hook_ends_the_call_with_that_cancellation(#[case] post_call: bool) {
Python::initialize();
Python::attach(|py| {
let locals = namespace(py, c"kwargs = {'logger': logger}\nresponse = object()");
let (mut logging, _) = begin(py, &locals, true);
if post_call {
logging
.resume(py, Ok(local(&locals, "kwargs").unbind()))
.unwrap();
logging
.after_success(py, local(&locals, "response").unbind(), TIMING)
.unwrap();
}
let cancellation = CancelledError::new_err("cancelled");
let cancelled = cancellation.value(py).clone();
let error = logging.resume(py, Err(cancellation)).err().unwrap();
assert!(error.value(py).is(&cancelled));
let names: Vec<String> = local(&locals, "logger")
.call_method0("names")
.unwrap()
.extract()
.unwrap();
assert!(!names.iter().any(|name| name.contains("handler")));
});
}
#[rstest]
#[case::hook_completed(false)]
#[case::hook_cancelled(true)]
fn failure_callbacks_run_after_the_failure_hook_however_it_ends(#[case] cancelled: bool) {
Python::initialize();
Python::attach(|py| {
let locals = namespace(
py,
c"kwargs = {'logger': logger}\nfailure = ValueError('provider')",
);
let (mut logging, _) = begin(py, &locals, true);
logging
.resume(py, Ok(local(&locals, "kwargs").unbind()))
.unwrap();
let failure = PyErr::from_value(local(&locals, "failure"));
let failed = CallEvent::Failed {
timing: TIMING,
origin: FailureOrigin::Call,
};
let step = logging
.emit(py, &failed, Some(PublicValue::Error(&failure)))
.unwrap();
assert!(awaits_deployment_hook(&step));
let hook_result = if cancelled {
Err(CancelledError::new_err("cancelled"))
} else {
Ok(py.None())
};
assert!(matches!(
logging.resume(py, hook_result).unwrap(),
AdapterStep::Await(_)
));
run(
py,
&locals,
c"
assert logger.names()[-3:] == ['failure_hook', 'failure_handler', 'async_failure_handler'], logger.calls
assert all(value is failure for name, value in logger.calls if name.endswith('_handler'))
",
);
});
}
#[rstest]
#[case::synchronous(false)]
#[case::asynchronous(true)]
fn a_limit_rejected_before_the_call_surfaces_as_the_callers_error(#[case] asynchronous: bool) {
Python::initialize();
Python::attach(|py| {
let locals = namespace(
py,
c"
class BudgetExceeded(Exception):
pass
rejection = BudgetExceeded('over budget')
class LimitedLogger(StubLogger):
def check_limits(self, arguments):
raise rejection
logger = LimitedLogger()
logger.hooks = {'pre': lambda kwargs: kwargs}
kwargs = {'logger': logger}
",
);
let mut logging = legacy_call(py, &locals, asynchronous);
let kwargs = local(&locals, "kwargs")
.cast_into::<PyDict>()
.unwrap()
.unbind();
let result = logging.begin(py, kwargs, 0.0).and_then(|step| match step {
AdapterStep::Await(_) => logging.resume(py, Ok(local(&locals, "kwargs").unbind())),
step => Ok(step),
});
let error = result.err().unwrap();
assert!(error.value(py).is(local(&locals, "rejection")));
});
}

365
litellm-rust/crates/callbacks-legacy/tests/payload.rs Normal file
View file

@ -0,0 +1,365 @@
use std::ffi::CStr;
use litellm_callbacks::event::{CallEvent, Passthrough, RawResponse, RequestContext, WireRequest};
use litellm_host_python::{AdapterStep, CallbackAdapter};
use pyo3::prelude::*;
use rstest::rstest;
use serde_json::{Value, json};
use super::LegacyLogging;
use crate::PythonLogger;
use crate::test_support::{legacy_call, local, namespace, run};
/// The payload phases of `Logging` on top of `StubLogger`, with `pre_call` handing the
/// payload to the case's `on_pre_call`.
const PAYLOAD_LOGGER: &CStr = c"
class Request:
pass
class PayloadLogger(StubLogger):
def update_from_kwargs(self, **update):
self.update = update
def pre_call(self, input, api_key, additional_args):
self.record('pre_call', None)
self.pre = additional_args
on_pre_call(additional_args)
def _pre_call(self, input, api_key, additional_args):
self.record('_pre_call', None)
def record_api_call_start_time(self):
self.record('record_api_call_start_time', None)
def post_call(self, original_response, additional_args):
self.record('post_call', None)
self.post = (original_response, additional_args)
def record_post_call(self, response, *rest):
self.record('record_post_call', response)
request = Request()
kwargs = {}
logger = PayloadLogger()
on_pre_call = lambda additional_args: None
check = lambda: None
";
const DOCUMENT: &str = "data:application/pdf;base64,YWJj";
const EDITED: &str = "data:application/pdf;base64,ZWRpdGVk";
fn document(source: &str) -> Value {
json!({"type": "document_url", "document_url": source})
}
fn before_send(script: &CStr, caller: Value, body: Value) -> WireRequest {
before_send_with_secrets(script, caller, body, &[])
}
/// Runs `before_send` over `body` for a caller whose route-side view is `caller`, with the
/// Python objects `script` binds, then delivers the provider's raw response the way the
/// driver does and runs the script's `check()`.
fn before_send_with_secrets(
script: &CStr,
caller: Value,
body: Value,
secret_fields: &[&str],
) -> WireRequest {
Python::initialize();
Python::attach(|py| {
let locals = namespace(py, PAYLOAD_LOGGER);
run(py, &locals, script);
let mut logging = LegacyLogging {
logger: Some(PythonLogger::new(local(&locals, "logger").unbind(), true)),
..legacy_call(py, &locals, false)
};
let context = RequestContext {
model: "model".into(),
custom_llm_provider: "provider".into(),
optional_params: caller.clone(),
passthrough_fields: Passthrough::unchanged(caller.as_object().unwrap(), &body),
secret_fields: secret_fields.iter().map(|name| name.to_string()).collect(),
};
let wire = WireRequest {
url: "https://provider.invalid/ocr".into(),
headers: vec![("x-route".into(), "route".into())],
body,
};
let step = logging.before_send(py, Box::new(wire), &context).unwrap();
let raw = CallEvent::ResponseReceived {
raw: RawResponse {
body: "raw response".into(),
},
};
assert!(matches!(
logging.emit(py, &raw, None).unwrap(),
AdapterStep::Done
));
run(py, &locals, c"check()");
let AdapterStep::Wire(wire) = step else {
panic!("before_send did not hand back the wire request");
};
*wire
})
}
#[rstest]
#[case::caller_keyword(c"
document = {'type': 'document_url', 'document_url': 'data:application/pdf;base64,YWJj'}
pages = [0]
kwargs = {'document': document, 'pages': pages}
observed = []
on_pre_call = lambda args: observed.append(
(args['complete_input_dict']['document'] is document, args['complete_input_dict']['pages'] is pages)
)
def check():
assert observed == [(True, True)], observed
")]
#[case::request_attribute_behind_an_omitted_keyword(c"
document = {'type': 'document_url', 'document_url': 'data:application/pdf;base64,YWJj'}
pages = [0]
request.document = document
kwargs = {'pages': pages}
observed = []
on_pre_call = lambda args: observed.append(
(args['complete_input_dict']['document'] is document, args['complete_input_dict']['pages'] is pages)
)
def check():
assert observed == [(True, True)], observed
")]
fn passthrough_keys_reach_pre_call_as_the_callers_own_objects(#[case] script: &CStr) {
let body = json!({"model": "model", "document": document(DOCUMENT), "pages": [0]});
let wire = before_send(
script,
json!({"document": document(DOCUMENT), "pages": [0]}),
body.clone(),
);
assert_eq!(wire.body, body);
}
#[test]
fn pre_call_edit_of_a_passthrough_object_reaches_the_caller_and_the_wire() {
let wire = before_send(
c"
document = {'type': 'document_url', 'document_url': 'data:application/pdf;base64,YWJj'}
kwargs = {'document': document}
def on_pre_call(args):
args['complete_input_dict']['document']['document_url'] = 'data:application/pdf;base64,ZWRpdGVk'
def check():
assert document['document_url'] == 'data:application/pdf;base64,ZWRpdGVk'
",
json!({"document": document(DOCUMENT)}),
json!({"document": document(DOCUMENT)}),
);
assert_eq!(wire.body["document"], document(EDITED));
}
#[test]
fn a_body_key_the_route_rewrote_is_not_the_callers_object() {
let wire = before_send(
c"
document = {'type': 'document_url', 'document_url': 'https://example.invalid/scan.pdf'}
kwargs = {'document': document}
observed = []
def on_pre_call(args):
observed.append(args['complete_input_dict']['document'] is document)
args['complete_input_dict']['document']['document_name'] = 'edited.pdf'
def check():
assert observed == [False], observed
assert document == {'type': 'document_url', 'document_url': 'https://example.invalid/scan.pdf'}
",
json!({"document": document("https://example.invalid/scan.pdf")}),
json!({"document": document(DOCUMENT)}),
);
assert_eq!(
wire.body["document"],
json!({"type": "document_url", "document_url": DOCUMENT, "document_name": "edited.pdf"})
);
}
#[rstest]
#[case::body(
c"
def on_pre_call(args):
args['complete_input_dict'] = {'replacement': True}
"
)]
#[case::headers(
c"
def on_pre_call(args):
args['headers'] = {'x-replacement': 'yes'}
"
)]
fn rebinding_the_payload_envelope_does_not_reach_the_wire(#[case] script: &CStr) {
let body = json!({"document": document(DOCUMENT)});
let wire = before_send(script, json!({}), body.clone());
assert_eq!(wire.body, body);
assert_eq!(wire.headers, [("x-route".to_string(), "route".to_string())]);
}
#[test]
fn pre_call_header_edit_reaches_the_wire() {
let wire = before_send(
c"
def on_pre_call(args):
args['headers']['x-callback'] = 'edited'
",
json!({}),
json!({}),
);
assert_eq!(
wire.headers,
[
("x-route".to_string(), "route".to_string()),
("x-callback".to_string(), "edited".to_string()),
]
);
}
#[test]
fn pre_call_receives_the_wire_request_and_the_logger_its_redacted_request() {
let body = json!({"model": "model", "document": document(DOCUMENT)});
before_send_with_secrets(
c"
logger_fn = lambda *args: None
kwargs = {
'litellm_call_id': 'call-1',
'client_secret': 'shh',
'proxy_server_request': {'body': {}},
'logger_fn': logger_fn,
'litellm_request_debug': True,
'ocr_cost_per_page': 0.05,
}
observed = []
on_pre_call = observed.append
def check():
[args] = observed
assert args['api_base'] == 'https://provider.invalid/ocr', args
assert args['complete_input_dict'] == {
'model': 'model',
'document': {'type': 'document_url', 'document_url': 'data:application/pdf;base64,YWJj'},
}, args
update = logger.update
assert update['model'] == 'model' and update['custom_llm_provider'] == 'provider', update
assert update['litellm_params']['litellm_call_id'] == 'call-1', update
assert update['litellm_params']['api_base'] == 'https://provider.invalid/ocr', update
assert update['litellm_params']['logger_fn'] is logger_fn, update
assert update['litellm_params']['litellm_request_debug'] is True, update
assert update['litellm_params']['ocr_cost_per_page'] == 0.05, update
assert update['kwargs']['client_secret'] == '****', update
assert 'proxy_server_request' not in update['kwargs'], update
assert update['optional_params']['client_secret'] == '****', update
",
json!({"client_secret": "shh"}),
body,
&["client_secret"],
);
}
#[rstest]
#[case::added_key(
c"
def on_pre_call(args):
args['complete_input_dict']['include_image_base64'] = True
",
json!({"document": document(DOCUMENT), "include_image_base64": true})
)]
#[case::replaced_document(
c"
document = {'type': 'document_url', 'document_url': 'data:application/pdf;base64,YWJj'}
kwargs = {'document': document}
def on_pre_call(args):
args['complete_input_dict']['document'] = {
'type': 'document_url', 'document_url': 'data:application/pdf;base64,ZWRpdGVk'
}
def check():
assert document['document_url'] == 'data:application/pdf;base64,YWJj', document
",
json!({"document": document(EDITED)})
)]
#[case::retained_body_edited_after_rebinding(
c"
def on_pre_call(args):
retained = args['complete_input_dict']
args['complete_input_dict'] = {'rebound': True}
retained['include_image_base64'] = True
",
json!({"document": document(DOCUMENT), "include_image_base64": true})
)]
fn pre_call_body_edits_reach_the_wire(#[case] script: &CStr, #[case] expected: Value) {
let body = json!({"document": document(DOCUMENT)});
let wire = before_send(script, json!({"document": document(DOCUMENT)}), body);
assert_eq!(wire.body, expected);
}
#[test]
fn retained_headers_edited_after_rebinding_reach_the_wire() {
let wire = before_send(
c"
def on_pre_call(args):
retained = args['headers']
args['headers'] = {'x-rebound': 'rebound'}
retained['x-retained'] = 'sent'
",
json!({}),
json!({}),
);
assert_eq!(
wire.headers,
[
("x-route".to_string(), "route".to_string()),
("x-retained".to_string(), "sent".to_string()),
]
);
}
#[test]
fn post_call_receives_the_raw_response_and_the_payload_dicts_pre_call_saw() {
before_send(
c"
def check():
original_response, additional_args = logger.post
assert original_response == 'raw response', original_response
assert additional_args['complete_input_dict'] is logger.pre['complete_input_dict']
assert additional_args['headers'] is logger.pre['headers']
",
json!({}),
json!({"document": document(DOCUMENT)}),
);
}
#[rstest]
#[case::every_phase_listens(c"{}", &["pre_call", "post_call"])]
#[case::no_input_callback(
c"{'input': False}",
&["_pre_call", "record_api_call_start_time", "record_post_call"]
)]
#[case::no_payload_consumer(c"{'payload': False}", &["record_api_call_start_time"])]
fn payload_callbacks_run_only_for_the_phases_someone_listens_to(
#[case] needed: &CStr,
#[case] expected_calls: &[&str],
) {
let script = std::ffi::CString::new(format!(
"
logger.needed = {needed}
def on_pre_call(args):
args['complete_input_dict']['include_image_base64'] = True
def check():
assert logger.names() == {expected_calls:?}, logger.calls
",
needed = needed.to_str().unwrap(),
expected_calls = expected_calls,
))
.unwrap();
let body = json!({"document": document(DOCUMENT)});
let wire = before_send(&script, json!({}), body.clone());
let edited = json!({"document": document(DOCUMENT), "include_image_base64": true});
assert_eq!(
wire.body,
if expected_calls.contains(&"pre_call") {
edited
} else {
body
}
);
}

188
litellm-rust/crates/callbacks-legacy/tests/support.rs Normal file
View file

@ -0,0 +1,188 @@
use std::ffi::CStr;
use pyo3::prelude::*;
use pyo3::types::{PyDict, PyTuple};
use crate::{LegacyLogging, LegacySurface, PublicCall};
/// Stand-ins for every litellm function the legacy contract calls. Tests share one
/// interpreter and run concurrently, so each stub is installed idempotently and forwards to
/// the per-test `StubLogger` it is handed (directly, or as `kwargs['logger']`).
const STUBS: &CStr = c"
import contextvars
import sys
import types
for name in (
'litellm',
'litellm.utils',
'litellm.types',
'litellm.types.utils',
'litellm._internal_context',
'litellm.litellm_core_utils',
'litellm.litellm_core_utils.logging_worker',
'litellm.litellm_core_utils.litellm_logging',
'litellm.rust_bridge',
'litellm.rust_bridge.legacy_callbacks',
):
sys.modules.setdefault(name, types.ModuleType(name))
legacy = sys.modules['litellm.rust_bridge.legacy_callbacks']
legacy.setup = lambda call_type, args, kwargs, start, asynchronous: types.SimpleNamespace(
logger=kwargs['logger_factory'](kwargs) if 'logger_factory' in kwargs else kwargs['logger'],
kwargs=kwargs,
bridge_owned=True,
)
legacy.deployment_callbacks_needed = lambda: True
legacy.check_limits = lambda arguments: arguments['logger'].check_limits(arguments)
legacy.callbacks_needed = lambda logger, phase: logger.needed.get(phase, True)
legacy.success_bookkeeping = lambda logger, response, start, end, asynchronous: logger.record(
'success_bookkeeping', asynchronous
)
legacy.failure_bookkeeping = lambda logger, error, start, end, asynchronous: logger.record(
'failure_bookkeeping', asynchronous
)
legacy.finalize = lambda response, logger, kwargs, start, end: logger.record('finalize', response)
utils = sys.modules['litellm.utils']
utils.async_pre_call_deployment_hook = lambda kwargs, call_type: kwargs['logger'].hook(
'pre', kwargs, call_type
)
utils.async_post_call_success_deployment_hook = lambda kwargs, response, call_type: kwargs[
'logger'
].hook('success', response, call_type)
utils.async_post_call_failure_deployment_hook = lambda kwargs, error, call_type: kwargs[
'logger'
].hook('failure', error, call_type)
utils._restore_correlation_context_if_supported = lambda logger: logger.record('restore', None)
internal = sys.modules['litellm._internal_context']
if not hasattr(internal, 'is_internal_call'):
internal.is_internal_call = contextvars.ContextVar('is_internal_call', default=False)
sys.modules['litellm.types.utils'].CustomPricingLiteLLMParams = type(
'CustomPricingLiteLLMParams', (), {'model_fields': {'ocr_cost_per_page': None}}
)
unraisable = sys.modules.setdefault(
'litellm_test_unraisable', types.ModuleType('litellm_test_unraisable')
)
if not hasattr(unraisable, 'events'):
unraisable.events = []
sys.unraisablehook = lambda event: unraisable.events.append((event.object, event.exc_value))
def unraisable_from(owner):
return [error for source, error in unraisable.events if source is owner]
class Worker:
def ensure_initialized_and_enqueue(self, coroutine):
return coroutine.enqueue()
class Executor:
def submit(self, run, handler, *args):
handler.__self__.record('submit', args)
sys.modules['litellm.litellm_core_utils.logging_worker'].GLOBAL_LOGGING_WORKER = Worker()
sys.modules['litellm.litellm_core_utils.litellm_logging'].executor = Executor()
class StubCoroutine:
def __init__(self, logger):
self.logger = logger
def enqueue(self):
self.logger.record('enqueued', None)
self.logger.on_enqueue(self)
def close(self):
self.logger.record('closed', None)
class StubLogger:
def __init__(self):
self.calls = []
self.needed = {}
self.hooks = {}
self.on_enqueue = lambda coroutine: None
def record(self, name, value):
self.calls.append((name, value))
def names(self):
return [name for name, _ in self.calls]
def hook(self, phase, value, call_type):
self.record(phase + '_hook', call_type)
return self.hooks.get(phase, lambda value: 'awaitable')(value)
def check_limits(self, arguments):
self.record('check_limits', arguments)
def failure_handler(self, error, trace, start, end):
self.record('failure_handler', error)
def async_failure_handler(self, error, trace, start, end):
self.record('async_failure_handler', error)
return 'awaitable'
def success_handler(self, response, start, end):
self.record('success_handler', response)
def async_success_handler(self, response, start, end):
self.record('async_success_handler', response)
return StubCoroutine(self)
def handle_sync_success_callbacks_for_async_calls(self, response, start, end):
self.record('sync_success_for_async_call', response)
logger = StubLogger()
";
/// A namespace with the stubs, `StubLogger` and a fresh `logger`, after `script` ran in it.
pub(crate) fn namespace<'py>(py: Python<'py>, script: &CStr) -> Bound<'py, PyDict> {
let locals = PyDict::new(py);
py.run(STUBS, Some(&locals), Some(&locals)).unwrap();
py.run(script, Some(&locals), Some(&locals)).unwrap();
locals
}
pub(crate) fn run(py: Python<'_>, locals: &Bound<'_, PyDict>, code: &CStr) {
py.run(code, Some(locals), Some(locals)).unwrap();
}
pub(crate) fn local<'py>(locals: &Bound<'py, PyDict>, name: &str) -> Bound<'py, PyAny> {
locals.get_item(name).unwrap().unwrap()
}
/// A legacy call over the namespace's `kwargs` (or none) and `request` (or `None`).
pub(crate) fn legacy_call(
py: Python<'_>,
locals: &Bound<'_, PyDict>,
asynchronous: bool,
) -> LegacyLogging {
let request = locals
.get_item("request")
.unwrap()
.unwrap_or_else(|| py.None().into_bound(py));
let kwargs = locals
.get_item("kwargs")
.unwrap()
.map(|kwargs| kwargs.cast_into::<PyDict>().unwrap())
.unwrap_or_else(|| PyDict::new(py));
let call = PublicCall::capture(&request, &PyTuple::empty(py), &kwargs).unwrap();
LegacyLogging::new(
py,
LegacySurface {
call_type: "test",
input_description: "test input",
},
call,
asynchronous,
)
}

291
litellm-rust/crates/callbacks-legacy/tests/terminal.rs Normal file
View file

@ -0,0 +1,291 @@
use std::ffi::CStr;
use litellm_callbacks::event::{CallEvent, FailureOrigin, Timing};
use litellm_host_python::{AdapterStep, CallbackAdapter, PublicValue};
use pyo3::exceptions::PyRuntimeError;
use pyo3::exceptions::asyncio::CancelledError;
use pyo3::prelude::*;
use pyo3::types::PyDict;
use rstest::rstest;
use super::LegacyLogging;
use crate::PythonLogger;
use crate::test_support::{legacy_call, local, namespace, run};
const TIMING: Timing = Timing {
start_time: 0.0,
end_time: 1.0,
};
fn logged(py: Python<'_>, locals: &Bound<'_, PyDict>, asynchronous: bool) -> LegacyLogging {
LegacyLogging {
logger: Some(PythonLogger::new(local(locals, "logger").unbind(), true)),
..legacy_call(py, locals, asynchronous)
}
}
fn succeed(py: Python<'_>, locals: &Bound<'_, PyDict>, logging: &mut LegacyLogging) -> AdapterStep {
let response = local(locals, "response").unbind();
logging
.emit(
py,
&CallEvent::Succeeded { timing: TIMING },
Some(PublicValue::Response(&response)),
)
.unwrap()
}
fn fail(py: Python<'_>, locals: &Bound<'_, PyDict>, logging: &mut LegacyLogging) -> AdapterStep {
let failure = PyErr::from_value(local(locals, "failure"));
logging
.emit(
py,
&CallEvent::Failed {
timing: TIMING,
origin: FailureOrigin::Host,
},
Some(PublicValue::Error(&failure)),
)
.unwrap()
}
#[rstest]
#[case::sync_listened(false, c"", &["submit"])]
#[case::sync_unlistened(false, c"logger.needed = {'sync_success': False}", &["success_bookkeeping"])]
#[case::async_listened(
true,
c"",
&["async_success_handler", "enqueued", "sync_success_for_async_call"]
)]
#[case::async_unlistened(
true,
c"logger.needed = {'async_success': False, 'sync_success_async': False}",
&["success_bookkeeping"]
)]
#[case::async_deferred(true, c"logger._defer_async_logging = True", &["sync_success_for_async_call"])]
#[case::async_with_fallbacks(true, c"kwargs = {'fallbacks': ['other']}", &["sync_success_for_async_call"])]
fn success_reaches_only_the_callbacks_that_listen(
#[case] asynchronous: bool,
#[case] script: &CStr,
#[case] expected: &[&str],
) {
Python::initialize();
Python::attach(|py| {
let locals = namespace(py, c"response = object()");
run(py, &locals, script);
let mut logging = logged(py, &locals, asynchronous);
assert!(matches!(
succeed(py, &locals, &mut logging),
AdapterStep::Done
));
let names: Vec<String> = local(&locals, "logger")
.call_method0("names")
.unwrap()
.extract()
.unwrap();
assert_eq!(names, expected);
run(
py,
&locals,
c"
assert all(value is response for name, value in logger.calls if name.endswith('_handler'))
assert hasattr(logger, '_native_pending_logging') == getattr(logger, '_defer_async_logging', False)
",
);
});
}
#[rstest]
#[case::synchronous(false, &["failure_handler"])]
#[case::asynchronous(true, &[])]
fn internal_calls_skip_failure_callbacks_only_when_asynchronous(
#[case] asynchronous: bool,
#[case] expected: &[&str],
) {
Python::initialize();
Python::attach(|py| {
let locals = namespace(py, c"failure = ValueError('provider')");
let mut logging = LegacyLogging {
internal: true,
..logged(py, &locals, asynchronous)
};
assert!(matches!(fail(py, &locals, &mut logging), AdapterStep::Done));
let names: Vec<String> = local(&locals, "logger")
.call_method0("names")
.unwrap()
.extract()
.unwrap();
assert_eq!(names, expected);
});
}
#[test]
fn internal_async_calls_skip_the_async_success_fan_out() {
Python::initialize();
Python::attach(|py| {
let locals = namespace(py, c"response = object()");
let mut logging = LegacyLogging {
internal: true,
..logged(py, &locals, true)
};
succeed(py, &locals, &mut logging);
run(
py,
&locals,
c"assert logger.names() == ['sync_success_for_async_call'], logger.calls",
);
});
}
#[test]
fn a_failing_success_callback_is_reported_without_replacing_the_response() {
Python::initialize();
Python::attach(|py| {
let locals = namespace(
py,
c"
response = object()
failure = ValueError('terminal diagnostic')
class FailingLogger(StubLogger):
def handle_sync_success_callbacks_for_async_calls(self, *args):
raise failure
logger = FailingLogger()
",
);
let mut logging = logged(py, &locals, true);
assert!(matches!(
succeed(py, &locals, &mut logging),
AdapterStep::Done
));
assert!(
logging
.response
.as_ref()
.unwrap()
.bind(py)
.is(local(&locals, "response"))
);
run(py, &locals, c"assert unraisable_from(logger) == [failure]");
});
}
#[rstest]
#[case::sync_listened(false, c"", &["failure_handler"])]
#[case::sync_unlistened(false, c"logger.needed = {'sync_failure': False}", &["failure_bookkeeping"])]
#[case::async_listened(true, c"", &["failure_handler", "async_failure_handler"])]
#[case::async_unlistened(
true,
c"logger.needed = {'sync_failure': False, 'async_failure': False}",
&["failure_bookkeeping", "failure_bookkeeping"]
)]
fn failure_reaches_only_the_callbacks_that_listen(
#[case] asynchronous: bool,
#[case] script: &CStr,
#[case] expected: &[&str],
) {
Python::initialize();
Python::attach(|py| {
let locals = namespace(py, c"failure = ValueError('provider')");
run(py, &locals, script);
let mut logging = logged(py, &locals, asynchronous);
let step = fail(py, &locals, &mut logging);
let awaits_async_handler = expected.contains(&"async_failure_handler");
assert_eq!(matches!(step, AdapterStep::Await(_)), awaits_async_handler);
let names: Vec<String> = local(&locals, "logger")
.call_method0("names")
.unwrap()
.extract()
.unwrap();
assert_eq!(names, expected);
run(
py,
&locals,
c"assert all(value is failure for name, value in logger.calls if name.endswith('_handler'))",
);
});
}
#[test]
fn a_failing_sync_failure_callback_keeps_the_error_and_still_runs_the_async_family() {
Python::initialize();
Python::attach(|py| {
let locals = namespace(
py,
c"
failure = ValueError('selected')
class FailingLogger(StubLogger):
def failure_handler(self, error, trace, start, end):
self.record('failure_handler', error)
raise RuntimeError('handler failed')
logger = FailingLogger()
",
);
let mut logging = logged(py, &locals, true);
assert!(matches!(
fail(py, &locals, &mut logging),
AdapterStep::Await(_)
));
assert!(
logging
.error
.as_ref()
.unwrap()
.bind(py)
.is(local(&locals, "failure"))
);
run(
py,
&locals,
c"assert logger.names() == ['failure_handler', 'async_failure_handler'], logger.calls",
);
});
}
#[rstest]
#[case::completed(None, true)]
#[case::handler_error(Some(false), true)]
#[case::cancelled(Some(true), false)]
fn the_async_failure_handler_ends_the_call_unless_it_was_cancelled(
#[case] error: Option<bool>,
#[case] done: bool,
) {
Python::initialize();
Python::attach(|py| {
let locals = namespace(py, c"failure = ValueError('provider')");
let mut logging = logged(py, &locals, true);
fail(py, &locals, &mut logging);
let result = match error {
None => Ok(py.None()),
Some(false) => Err(PyRuntimeError::new_err("handler failed")),
Some(true) => Err(CancelledError::new_err("cancelled")),
};
let expected = result.as_ref().err().map(|error| error.value(py).clone());
match logging.resume(py, result) {
Ok(step) => assert!(done && matches!(step, AdapterStep::Done)),
Err(propagated) => {
assert!(!done);
assert!(propagated.value(py).is(expected.unwrap()));
}
}
});
}
#[test]
fn closing_restores_the_correlation_context_once() {
Python::initialize();
Python::attach(|py| {
let locals = namespace(py, c"");
let mut logging = logged(py, &locals, true);
logging.close(py);
logging.close(py);
run(
py,
&locals,
c"assert logger.names() == ['restore'], logger.calls",
);
});
}

8
litellm-rust/crates/python-interop/Cargo.toml → litellm-rust/crates/callbacks/Cargo.toml
View file

@ -1,15 +1,13 @@
[package]
name = "litellm-python-interop"
name = "litellm-callbacks"
version = "0.1.0"
edition.workspace = true
license.workspace = true
repository.workspace = true
[dependencies]
pyo3.workspace = true
pythonize.workspace = true
serde.workspace = true
serde_json.workspace = true
[dev-dependencies]
rstest.workspace = true
serde_json.workspace = true
tokio = { workspace = true, features = ["macros"] }

135
litellm-rust/crates/callbacks/src/event.rs Normal file
View file

@ -0,0 +1,135 @@
use std::time::{SystemTime, UNIX_EPOCH};
use serde_json::{Map, Value};
/// Seconds since the Unix epoch, on one clock for every host.
pub fn epoch_seconds() -> f64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|duration| duration.as_secs_f64())
.unwrap_or(0.0)
}
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct Timing {
pub start_time: f64,
pub end_time: f64,
}
/// The provider request as it is about to leave, offered to the host for rewriting.
#[derive(Clone, Debug, PartialEq)]
pub struct WireRequest {
pub url: String,
pub headers: Vec<(String, String)>,
pub body: Value,
}
/// What the route knows about the request it is sending, for a host that logs it. The
/// route owns these facts; a host reads them beside the wire request and never rewrites
/// them.
#[derive(Clone, Debug, PartialEq)]
pub struct RequestContext {
pub model: String,
pub custom_llm_provider: String,
/// The route's parameters before the provider transformation.
pub optional_params: Value,
pub passthrough_fields: Passthrough,
/// Optional-param names that carry credentials and must be redacted when logged.
pub secret_fields: Vec<String>,
}
/// Body keys whose values are the caller's inputs, unchanged by the route. The only way to
/// build one is to compare the two, so a route cannot name a key it rewrote.
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct Passthrough(Vec<String>);
impl Passthrough {
pub fn unchanged(caller: &Map<String, Value>, body: &Value) -> Self {
Self(
caller
.iter()
.filter(|(name, value)| body.get(name.as_str()) == Some(*value))
.map(|(name, _)| name.clone())
.collect(),
)
}
pub fn iter(&self) -> impl Iterator<Item = &str> {
self.0.iter().map(String::as_str)
}
pub fn contains(&self, name: &str) -> bool {
self.0.iter().any(|field| field == name)
}
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct RawResponse {
pub body: String,
}
/// Whether a failure surfaced inside the call, including a host op the call asked for,
/// or in a host step around it (preparing the arguments, finalizing the response).
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum FailureOrigin {
Call,
Host,
}
#[derive(Clone, Debug, PartialEq)]
pub enum CallEvent {
ResponseReceived {
raw: RawResponse,
},
Succeeded {
timing: Timing,
},
Failed {
timing: Timing,
origin: FailureOrigin,
},
}
#[cfg(test)]
mod tests {
use rstest::rstest;
use serde_json::json;
use super::*;
#[rstest]
#[case::unchanged_scalar(json!({"pages": [0]}), json!({"pages": [0]}), &["pages"])]
#[case::unchanged_explicit_null(json!({"pages": null}), json!({"pages": null}), &["pages"])]
#[case::unchanged_nested_object(
json!({"document": {"type": "document_url", "document_url": "https://a/b.pdf"}}),
json!({"document": {"type": "document_url", "document_url": "https://a/b.pdf"}, "model": "m"}),
&["document"]
)]
#[case::rewritten_value(
json!({"document": {"type": "document_url", "document_url": "https://a/b.pdf"}}),
json!({"document": {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"}}),
&[]
)]
#[case::dropped_nested_field(
json!({"document": {"type": "image_url", "image_url": "https://a/b.png", "document_name": "b.png"}}),
json!({"document": {"type": "image_url", "image_url": "https://a/b.png"}}),
&[]
)]
#[case::added_nested_field(
json!({"document": {"type": "image_url", "image_url": "https://a/b.png"}}),
json!({"document": {"type": "image_url", "image_url": "https://a/b.png", "detail": "high"}}),
&[]
)]
#[case::reordered_array(json!({"pages": [0, 1]}), json!({"pages": [1, 0]}), &[])]
#[case::consumed_by_the_route(json!({"api_key": "k", "pages": [0]}), json!({"pages": [0]}), &["pages"])]
#[case::added_by_the_route(json!({}), json!({"model": "m"}), &[])]
#[case::non_object_body(json!({"pages": [0]}), json!([{"pages": [0]}]), &[])]
fn passthrough_is_exactly_the_callers_unchanged_keys(
#[case] caller: Value,
#[case] body: Value,
#[case] expected: &[&str],
) {
let passthrough = Passthrough::unchanged(caller.as_object().unwrap(), &body);
assert_eq!(passthrough.iter().collect::<Vec<_>>(), expected);
}
}

45
litellm-rust/crates/callbacks/src/host.rs Normal file
View file

@ -0,0 +1,45 @@
use std::future::Future;
use crate::event::{CallEvent, RequestContext, WireRequest};
use crate::route::Route;
/// One suspension point of a native call, performed by the host.
pub enum HostOp<R: Route> {
Route(R::Op),
BeforeSend {
wire: Box<WireRequest>,
context: Box<RequestContext>,
},
Emit(CallEvent),
}
pub enum HostResult<R: Route> {
Route(R::OpResult),
BeforeSend(Box<WireRequest>),
Emitted,
}
/// A host answer that is either available now or arrives once the host's own
/// suspension (a Python awaitable, for example) resolves.
pub enum HostStep<V, S> {
Ready(V),
Suspend(S),
}
/// An in-process host: answers route operations and observes the call without leaving
/// the Rust runtime. Language hosts implement their own driver instead.
pub trait Host<R: Route>: Send + Sync {
fn route(&self, op: R::Op) -> impl Future<Output = Result<R::OpResult, R::Error>> + Send;
fn before_send(
&self,
wire: WireRequest,
_context: &RequestContext,
) -> impl Future<Output = Result<WireRequest, R::Error>> + Send {
async move { Ok(wire) }
}
fn emit(&self, _event: &CallEvent) -> impl Future<Output = Result<(), R::Error>> + Send {
async { Ok(()) }
}
}

12
litellm-rust/crates/callbacks/src/lib.rs Normal file
View file

@ -0,0 +1,12 @@
//! The contract between a native call and the host runtime that drives it.
//!
//! A host is whatever sits on the far side of the language boundary: CPython today,
//! another runtime later. Core implements [`machine::Machine`] per route and never learns
//! which host is on the other end. The machine yields [`host::HostOp`]s; a driver answers
//! them, observes [`event::CallEvent`]s and may rewrite the wire request before it is sent.
pub mod event;
pub mod host;
pub mod machine;
pub mod route;
pub mod run;

63
litellm-rust/crates/callbacks/src/machine.rs Normal file
View file

@ -0,0 +1,63 @@
use std::future::Future;
use std::pin::Pin;
use crate::host::{HostOp, HostResult};
use crate::route::Route;
pub enum MachineStep<R: Route, C> {
Host(HostOp<R>),
Complete(C),
}
pub type Step<'a, M> = Pin<
Box<
dyn Future<
Output = Result<
MachineStep<<M as Machine>::Route, <M as Machine>::Complete>,
<<M as Machine>::Route as Route>::Error,
>,
> + Send
+ 'a,
>,
>;
pub type Interrupted<'a, M> = Pin<
Box<
dyn Future<
Output = Result<<M as Machine>::Complete, <<M as Machine>::Route as Route>::Error>,
> + Send
+ 'a,
>,
>;
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum HostFailure<E> {
Error(E),
Cancelled(E),
}
impl<E> HostFailure<E> {
pub fn into_error(self) -> E {
match self {
Self::Error(error) | Self::Cancelled(error) => error,
}
}
}
/// A resumable call. Core implements it per route; a host drives it. Every suspension
/// point is an op the host performs and answers with a result.
pub trait Machine: Send {
type Route: Route;
type Complete: Send + 'static;
/// `None` on the first call and whenever the previous step completed without
/// yielding an op; otherwise the result of the op last yielded.
fn resume(&mut self, result: Option<HostResult<Self::Route>>) -> Step<'_, Self>;
/// The host failed to perform the pending op, or the caller cancelled. The call
/// yields no further ops.
fn interrupt(
&mut self,
failure: HostFailure<<Self::Route as Route>::Error>,
) -> Interrupted<'_, Self>;
}

9
litellm-rust/crates/callbacks/src/route.rs Normal file
View file

@ -0,0 +1,9 @@
/// One public call surface: what a completed call produces, how it fails, and the
/// route-specific operations only its host can perform (request projection, file reads,
/// token acquisition).
pub trait Route: Send + Sync + 'static {
type Response: Send + 'static;
type Error: Clone + Send + Sync + 'static;
type Op: Send + 'static;
type OpResult: Send + 'static;
}

149
litellm-rust/crates/callbacks/src/run.rs Normal file
View file

@ -0,0 +1,149 @@
use crate::event::{CallEvent, FailureOrigin, Timing, epoch_seconds};
use crate::host::{Host, HostOp, HostResult};
use crate::machine::{HostFailure, Machine, MachineStep};
use crate::route::Route;
/// Drives a machine to completion against an in-process host and emits exactly one
/// terminal event.
pub async fn run<M, H>(mut machine: M, host: &H) -> Result<M::Complete, <M::Route as Route>::Error>
where
M: Machine,
H: Host<M::Route>,
{
let start_time = epoch_seconds();
let mut result = None;
let outcome = loop {
let step = match machine.resume(result.take()).await {
Ok(MachineStep::Complete(complete)) => break Ok(complete),
Ok(MachineStep::Host(op)) => op,
Err(error) => break Err(error),
};
let answer = match step {
HostOp::Route(op) => host.route(op).await.map(HostResult::Route),
HostOp::BeforeSend { wire, context } => host
.before_send(*wire, &context)
.await
.map(|wire| HostResult::BeforeSend(Box::new(wire))),
HostOp::Emit(event) => host.emit(&event).await.map(|()| HostResult::Emitted),
};
match answer {
Ok(answer) => result = Some(answer),
Err(error) => break machine.interrupt(HostFailure::Error(error)).await,
}
};
let timing = Timing {
start_time,
end_time: epoch_seconds(),
};
let terminal = match &outcome {
Ok(_) => CallEvent::Succeeded { timing },
Err(_) => CallEvent::Failed {
timing,
origin: FailureOrigin::Call,
},
};
let _ = host.emit(&terminal).await;
outcome
}
#[cfg(test)]
mod tests {
use std::sync::Mutex;
use super::*;
use crate::machine::{Interrupted, Step};
struct Unit;
impl Route for Unit {
type Response = ();
type Error = &'static str;
type Op = &'static str;
type OpResult = ();
}
struct Scripted {
ops: Vec<&'static str>,
outcome: Result<(), &'static str>,
}
impl Machine for Scripted {
type Route = Unit;
type Complete = ();
fn resume(&mut self, _: Option<HostResult<Unit>>) -> Step<'_, Self> {
Box::pin(async move {
if !self.ops.is_empty() {
return Ok(MachineStep::Host(HostOp::Route(self.ops.remove(0))));
}
self.outcome.map(MachineStep::Complete)
})
}
fn interrupt(&mut self, failure: HostFailure<&'static str>) -> Interrupted<'_, Self> {
Box::pin(async move { Err(failure.into_error()) })
}
}
#[derive(Default)]
struct Recording {
seen: Mutex<Vec<String>>,
fail: Option<&'static str>,
}
impl Host<Unit> for Recording {
async fn route(&self, op: &'static str) -> Result<(), &'static str> {
self.seen.lock().unwrap().push(format!("route:{op}"));
match self.fail {
Some(failing) if failing == op => Err("host failed"),
_ => Ok(()),
}
}
async fn emit(&self, event: &CallEvent) -> Result<(), &'static str> {
self.seen.lock().unwrap().push(match event {
CallEvent::Succeeded { .. } => "succeeded".into(),
CallEvent::Failed { .. } => "failed".into(),
other => format!("{other:?}"),
});
Ok(())
}
}
fn scripted(ops: &[&'static str], outcome: Result<(), &'static str>) -> Scripted {
Scripted {
ops: ops.to_vec(),
outcome,
}
}
#[tokio::test]
async fn forwards_every_op_then_emits_one_succeeded() {
let host = Recording::default();
let outcome = run(scripted(&["project", "send"], Ok(())), &host).await;
assert_eq!(outcome, Ok(()));
assert_eq!(
*host.seen.lock().unwrap(),
["route:project", "route:send", "succeeded"]
);
}
#[tokio::test]
async fn errors_and_host_failures_each_emit_failed_once() {
let host = Recording::default();
let outcome = run(scripted(&[], Err("boom")), &host).await;
assert_eq!(outcome, Err("boom"));
assert_eq!(*host.seen.lock().unwrap(), ["failed"]);
let host = Recording {
fail: Some("send"),
..Recording::default()
};
let outcome = run(scripted(&["project", "send", "never"], Ok(())), &host).await;
assert_eq!(outcome, Err("host failed"));
assert_eq!(
*host.seen.lock().unwrap(),
["route:project", "route:send", "failed"]
);
}
}

24
litellm-rust/crates/core/AGENTS.md
View file

@ -1,7 +1,27 @@
litellm-core is the LiteLLM SDK in Rust — it makes the LLM call. Each top-level call is a module under `src/<route>/` exposing a public entrypoint named after the route (`messages::messages()`, the Rust equivalent of `litellm.messages()`): you call it and get a typed non-streaming response back.
A route module owns everything the call needs: types, the provider template trait, provider transforms (under `providers/`), provider/auth/URL resolution, and the handler that performs the HTTP call. Handlers belong here, never in a host crate.
A route module owns the call entrypoint, runtime types, provider/auth/URL resolution, and the handler that performs the HTTP call. Provider code and base config traits live under `src/llms/`, mirroring their Python source paths. This applies to every API surface: shared orchestration stays in its route module (`ocr/`, `chat_completions/`, `messages/`, `audio_transcription/`, or `responses/`), while provider transformations live under the corresponding Python-mirrored `llms/<provider>/` path. Import implementations directly from their canonical paths; do not add a `src/providers/` layer or compatibility re-exports. Shared provider resolution lives under `src/litellm_core_utils/get_llm_provider_logic.rs`. Handlers belong in core, never in a host crate
Not here: serving HTTP (axum routes, extractors), config file reading, rollout state, databases, or host-specific callback execution. Core owns lifecycle sequencing and callback payload construction; hosts execute the selected integrations. Env reads are limited to credential fallback in a route's `prepare.rs`.
Not here: serving HTTP (axum routes, extractors), config file reading, rollout state, databases, or callback execution of any kind. Core runs each route as a machine that yields host operations and call events; which integrations consume those events is the host's business. Env reads are limited to credential fallback in a route's `prepare.rs`.
Routes (messages, ocr, realtime) and providers (anthropic, mistral, openai) are modules, not crates.
## Python/Rust transformation pairs
Use the base OCR and Mistral OCR pairs as the reference when aligning transformations. Derive `src/<relative_path>.rs` from `litellm/<relative_path>.py`, preserving meaningful basenames such as `messages_transformation`
Keep corresponding operation names and parameter names when their responsibilities match. Rust types retain the Python semantic name with Rust acronym casing (`BaseOCRConfig` / `BaseOcrConfig`, `MistralOCRConfig` / `MistralOcrConfig`). Private Python helpers can drop their leading underscore. Give Rust adapter helpers distinct responsibility names rather than duplicating trait method names
Order OCR config methods as supported parameters, credential metadata and connection resolution, health-check input, parameter mapping, environment validation, URL construction, request transformation, async request transformation, response transformation, async response transformation, and error conversion. Put constants and data types before the config, private helpers after it in operation order, and tests last. Rust-only trait hooks follow the corresponding Python methods
Use trait defaults for unchanged inherited behavior and explicit delegation for shared provider behavior. Keep typed inputs, ownership, `Result`, and async I/O idiomatic. A matching path or symbol identifies the counterpart, not a claim of full behavioral parity
Use named `#[rstest]` cases for independent input/output scenarios instead of loops or repeated calls in one test. Inject reusable setup with `#[fixture]` arguments and use `#[with(...)]` for fixture overrides. Keep assertions about the same result together
For base OCR, Python response models correspond to `src/ocr/types.rs`; Rust context/environment types support the runtime. `BaseOcrConfig::prepare_request` corresponds to Python's HTTP-handler preparation rather than a `BaseOCRConfig` method, and `validate_request_body` is a Rust-only hook
For Mistral, `async_transform_ocr_request` uses the base default in both languages. `resolve_headers` and `build_ocr_url` implement the respective environment and URL operations, and `normalize_response` implements the typed part of response transformation. Existing auth key/header handling and top-level response-extra preservation differ between languages; layout refactors must preserve those behaviors and verify them with the existing tests
For non-OCR pairs, order corresponding methods as parameter support/mapping, environment validation, URL construction, request transformation, and response transformation, followed by Rust-only runtime hooks. Auth resolution remains split between configs and route preparation. Chat `supported_openai_param_mappings` describes accepted OpenAI/provider name pairs, unlike Python's `get_supported_openai_params` name list. Audio `map_transcription_params` remains a Rust filtering helper
Azure Messages maps to `llms/azure_ai/anthropic/messages_transformation.py`; Bedrock Converse maps to `llms/bedrock/chat/converse_transformation.py`. `AnthropicConfig`, `AmazonConverseConfig`, and the non-OCR base traits are partial ports. `OpenAiResponsesApiConfig` currently implements only the WebSocket surface. Preserve their acceptance gates, passthrough behavior, and host fallback contracts when aligning layout

10
litellm-rust/crates/core/Cargo.toml
View file

@ -7,6 +7,7 @@ repository.workspace = true
autotests = false
[dependencies]
litellm-callbacks.workspace = true
bytes.workspace = true
futures-util.workspace = true
base64.workspace = true
@ -15,6 +16,8 @@ litellm-auth.workspace = true
litellm-auth-aws.workspace = true
litellm-auth-azure.workspace = true
litellm-auth-gcp.workspace = true
litellm-providers.workspace = true
litellm-framing.workspace = true
moka.workspace = true
mime_guess = "2.0.5"
rand.workspace = true
@ -22,16 +25,21 @@ reqwest.workspace = true
rustls.workspace = true
rustls-native-certs.workspace = true
serde.workspace = true
serde_json.workspace = true
serde_json = { workspace = true, features = ["preserve_order"] }
serde_with.workspace = true
serde_path_to_error = "0.1"
strum.workspace = true
subtle.workspace = true
tokio = { workspace = true, features = ["sync"] }
tokio-tungstenite.workspace = true
thiserror.workspace = true
time.workspace = true
sha2.workspace = true
url.workspace = true
veil.workspace = true
[dev-dependencies]
aws-smithy-eventstream = "=0.61.1"
aws-smithy-types = "1.6.1"
rstest.workspace = true
rstest_reuse.workspace = true

3
litellm-rust/crates/core/src/audio_transcription/client.rs
View file

@ -1,5 +1,4 @@
use std::sync::OnceLock;
use std::time::Duration;
use std::{sync::OnceLock, time::Duration};
use crate::constants::AUDIO_TRANSCRIPTION_TIMEOUT_SECS;

20
litellm-rust/crates/core/src/audio_transcription/error.rs
View file

@ -24,3 +24,23 @@ pub enum Error {
#[error(transparent)]
Aws(#[from] litellm_auth_aws::Error),
}
impl From<litellm_providers::audio_transcription::Error> for Error {
fn from(error: litellm_providers::audio_transcription::Error) -> Self {
match error {
litellm_providers::audio_transcription::Error::InvalidType { expected, actual } => {
Self::InvalidType { expected, actual }
}
litellm_providers::audio_transcription::Error::MissingField(field) => {
Self::MissingField(field)
}
litellm_providers::audio_transcription::Error::InvalidRequest(message) => {
Self::InvalidRequest(message)
}
litellm_providers::audio_transcription::Error::InvalidResponse(message) => {
Self::InvalidResponse(message)
}
litellm_providers::audio_transcription::Error::Auth(error) => Self::Auth(error),
}
}
}

15
litellm-rust/crates/core/src/audio_transcription/handler.rs
View file

@ -1,11 +1,8 @@
use serde_json::Value;
use super::Error;
use super::{Error, client::http_client, types::ProviderAudioTranscriptionRequest};
use crate::http_utils::{http_request, truncate_error_body};
use super::client::http_client;
use super::types::ProviderAudioTranscriptionRequest;
pub async fn execute_audio_transcription_provider_call(
request: ProviderAudioTranscriptionRequest,
) -> Result<Value, Error> {
@ -37,7 +34,7 @@ pub async fn execute_audio_transcription_provider_call(
.map_err(|error| Error::InvalidResponse(format!("invalid audio response JSON: {error}")))?;
Ok(request
.config
.transform_transcription_response(&request.model, response_json)?
.transform_audio_transcription_response(&request.model, response_json)?
.into_json())
}
@ -45,12 +42,10 @@ async fn signed_headers(
request: &ProviderAudioTranscriptionRequest,
body: &[u8],
) -> Result<Vec<(String, String)>, Error> {
use std::collections::BTreeMap;
use std::time::SystemTime;
use std::{collections::BTreeMap, time::SystemTime};
use crate::audio_transcription::transformation::AudioTranscriptionAuth;
use crate::providers::bedrock::audio_transcription::aws_auth_config;
use crate::providers::bedrock::aws_base::{resolve_credentials, sign_bedrock_post};
use litellm_auth_aws::{aws_auth_config, resolve_credentials, sign_bedrock_post};
use litellm_providers::base_llm::audio_transcription::transformation::AudioTranscriptionAuth;
let AudioTranscriptionAuth::AwsSigV4 { region, .. } = &request.auth else {
return Ok(request.upstream_headers.clone());

7
litellm-rust/crates/core/src/audio_transcription/mod.rs
View file

@ -3,13 +3,10 @@ pub use error::Error;
mod client;
mod handler;
mod prepare;
pub mod transformation;
pub mod types;
use serde_json::Value;
pub use handler::execute_audio_transcription_provider_call;
pub use litellm_providers::audio_transcription::types;
pub use prepare::prepare_audio_transcription_provider_call;
use serde_json::Value;
pub use types::{AudioTranscriptionRequest, ProviderAudioTranscriptionRequest};
pub async fn audio_transcription(request: AudioTranscriptionRequest<'_>) -> Result<Value, Error> {

26
litellm-rust/crates/core/src/audio_transcription/prepare.rs
View file

@ -1,12 +1,20 @@
use super::Error;
use crate::http_utils::{has_header, string_headers};
use crate::providers::bedrock::audio_transcription::BEDROCK_AUDIO_TRANSCRIPTION_CONFIG;
use crate::providers::custom_llm_provider::{CustomLlmProvider, get_custom_llm_provider};
use litellm_providers::{
base_llm::audio_transcription::transformation::{
AudioTranscriptionAuth, BaseAudioTranscriptionConfig,
},
bedrock::audio_transcription::BEDROCK_AUDIO_TRANSCRIPTION_CONFIG,
};
use super::transformation::{AudioTranscriptionAuth, AudioTranscriptionProviderConfig};
use super::types::{AudioTranscriptionRequest, ProviderAudioTranscriptionRequest};
use super::{
Error,
types::{AudioTranscriptionRequest, ProviderAudioTranscriptionRequest},
};
use crate::{
http_utils::{has_header, string_headers},
litellm_core_utils::get_llm_provider_logic::{CustomLlmProvider, get_custom_llm_provider},
};
fn provider_config(provider: &str) -> Option<&'static dyn AudioTranscriptionProviderConfig> {
fn provider_config(provider: &str) -> Option<&'static dyn BaseAudioTranscriptionConfig> {
if provider == "bedrock" {
return Some(&BEDROCK_AUDIO_TRANSCRIPTION_CONFIG);
}
@ -46,7 +54,7 @@ pub fn prepare_audio_transcription_provider_call(
if !has_header(&headers, "content-type") {
headers.push(("Content-Type".to_string(), "application/json".to_string()));
}
let url = config.complete_url(
let url = config.get_complete_url(
request.api_base,
&model,
&request.optional_params,
@ -54,7 +62,7 @@ pub fn prepare_audio_transcription_provider_call(
)?;
let filtered_params = config.map_transcription_params(&request.optional_params);
let transformed =
config.transform_transcription_request(&model, request.audio, filtered_params)?;
config.transform_audio_transcription_request(&model, request.audio, filtered_params)?;
Ok(ProviderAudioTranscriptionRequest {
model,
custom_llm_provider: provider_info.custom_llm_provider.to_string(),

11
litellm-rust/crates/core/src/audio_transcription/tests.rs
View file

@ -1,11 +1,12 @@
use std::io::{Read, Write};
use std::net::TcpListener;
use std::thread;
use std::{
io::{Read, Write},
net::TcpListener,
thread,
};
use serde_json::{Map, json};
use super::audio_transcription;
use super::types::AudioTranscriptionRequest;
use super::{audio_transcription, types::AudioTranscriptionRequest};
#[tokio::test]
async fn bedrock_request_is_signed_and_contains_audio() {

181
litellm-rust/crates/core/src/call_arguments.rs Normal file
View file

@ -0,0 +1,181 @@
use std::ops::Deref;
use serde::{Deserialize, Serialize, de::DeserializeOwned};
use serde_json::{Map, Value};
#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
#[serde(transparent)]
pub struct CallArguments(Map<String, Value>);
impl CallArguments {
pub(crate) fn select(&self, names: &[&str]) -> Map<String, Value> {
self.iter()
.filter(|(name, _)| names.contains(&name.as_str()))
.map(|(name, value)| (name.clone(), value.clone()))
.collect()
}
}
#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)]
#[error("invalid argument: {path}")]
pub struct ArgumentError {
pub path: String,
}
pub fn parse_options<T: DeserializeOwned>(arguments: &CallArguments) -> Result<T, ArgumentError> {
let deserializer = serde::de::value::MapDeserializer::new(
arguments.iter().map(|(name, value)| (name.as_str(), value)),
);
serde_path_to_error::deserialize(deserializer).map_err(|error| ArgumentError {
path: error.path().to_string(),
})
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct ArgumentSpec {
pub name: &'static str,
pub secret: bool,
}
pub fn compose_body<B: Serialize>(
arguments: &CallArguments,
body: &B,
consumed: &[&str],
) -> Result<Value, crate::params::Error> {
let Value::Object(fields) =
serde_json::to_value(body).map_err(|_| crate::params::Error::Body)?
else {
return Err(crate::params::Error::Body);
};
let overrides = match arguments.get("extra_body") {
None | Some(Value::Null) => None,
Some(Value::Object(fields)) => Some(fields),
Some(_) => return Err(crate::params::Error::ExtraBody),
};
let extensions = arguments
.iter()
.filter(|(name, _)| !consumed.contains(&name.as_str()));
Ok(Value::Object(
fields
.into_iter()
.chain(
extensions
.chain(overrides.into_iter().flatten())
.filter(|(name, _)| {
name.as_str() != "model"
&& name.as_str() != "extra_body"
&& !crate::params::is_control_param(name)
})
.map(|(name, value)| (name.clone(), value.clone())),
)
.collect(),
))
}
impl Deref for CallArguments {
type Target = Map<String, Value>;
fn deref(&self) -> &Self::Target {
&self.0
}
}
impl From<Map<String, Value>> for CallArguments {
fn from(values: Map<String, Value>) -> Self {
Self(values)
}
}
impl From<CallArguments> for Map<String, Value> {
fn from(arguments: CallArguments) -> Self {
arguments.0
}
}
impl FromIterator<(String, Value)> for CallArguments {
fn from_iter<T: IntoIterator<Item = (String, Value)>>(iter: T) -> Self {
Self(iter.into_iter().collect())
}
}
impl IntoIterator for CallArguments {
type Item = (String, Value);
type IntoIter = serde_json::map::IntoIter;
fn into_iter(self) -> Self::IntoIter {
self.0.into_iter()
}
}
#[cfg(test)]
mod tests {
use serde_json::json;
use super::*;
#[test]
fn composition_preserves_extensions_and_applies_shallow_explicit_overrides() {
let original = json!({
"known": false, "future": {"old": 1}, "null": null, "zero": 0,
"metadata": {"host": true}, "timeout": 30, "api_key": "secret",
"extra_body": {
"known": null, "future": {"new": [false, 0, null]},
"metadata": {"provider": true}, "model": "ignored", "api_key": "ignored"
}
});
let arguments = serde_json::from_value(original.clone()).unwrap();
let body = compose_body(
&arguments,
&json!({"model":"resolved", "known":false}),
&["known"],
)
.unwrap();
assert_eq!(
body,
json!({
"model":"resolved", "known":null, "future":{"new":[false,0,null]},
"null":null, "zero":0, "metadata":{"provider":true}
})
);
assert_eq!(serde_json::to_value(arguments).unwrap(), original);
}
#[test]
fn invalid_extra_body_is_rejected_without_coercing_it_to_empty() {
for value in [json!(false), json!(0), json!([]), json!("")] {
let arguments = serde_json::from_value(json!({"extra_body":value})).unwrap();
assert_eq!(
compose_body(&arguments, &json!({}), &[]),
Err(crate::params::Error::ExtraBody)
);
}
let arguments = serde_json::from_value(json!({"extra_body":null})).unwrap();
assert_eq!(
compose_body(&arguments, &json!({}), &[]).unwrap(),
json!({})
);
}
#[test]
fn typed_views_preserve_missing_and_explicit_null_in_the_source() {
#[derive(Deserialize)]
struct Options {
enabled: Option<bool>,
}
let arguments: CallArguments =
serde_json::from_value(json!({"enabled":null,"future":0})).unwrap();
assert!(
parse_options::<Options>(&arguments)
.unwrap()
.enabled
.is_none()
);
assert_eq!(arguments.get("enabled"), Some(&Value::Null));
assert_eq!(arguments.get("missing"), None);
let invalid = serde_json::from_value(json!({"enabled":0})).unwrap();
assert_eq!(
parse_options::<Options>(&invalid).err().unwrap().path,
"enabled"
);
}
}

122
litellm-rust/crates/core/src/call_lifecycle/host.rs
View file

@ -1,122 +0,0 @@
use std::future::Future;
use std::pin::Pin;
pub enum HostCallStep<O, C> {
Host(O),
Complete(C),
}
pub type HostCallFuture<'a, O, C, E> =
Pin<Box<dyn Future<Output = Result<HostCallStep<O, C>, E>> + Send + 'a>>;
pub trait HostCall: Send + Sync {
type Error: Send + Sync + 'static;
type Operation: Send + 'static;
type Result: Send + 'static;
type Complete: Send + 'static;
fn resume(
&mut self,
result: Option<Self::Result>,
) -> HostCallFuture<'_, Self::Operation, Self::Complete, Self::Error>;
fn interrupt(
&mut self,
failure: HostFailure<Self::Error>,
) -> HostCallFuture<'_, Self::Operation, Self::Complete, Self::Error>;
}
pub enum HostStep<V, S> {
Ready(V),
Suspend(S),
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum HostPhase {
Setup,
DeploymentPreCall,
Prepare,
Execute,
ConstructResponse,
DeploymentPostCall,
Finalize,
Success,
MapFailure,
DeploymentFailure,
Failure,
AsyncFailure,
Complete,
}
#[derive(Clone, Debug)]
pub enum HostFailure<E> {
Error(E),
Cancelled(E),
}
pub struct HostLifecycle {
phase: HostPhase,
asynchronous: bool,
}
impl HostLifecycle {
pub fn new(asynchronous: bool) -> Self {
Self {
phase: HostPhase::Setup,
asynchronous,
}
}
pub fn phase(&self) -> HostPhase {
self.phase
}
pub fn accept<E>(&mut self, result: Result<(), HostFailure<E>>) -> Option<E> {
if let Err(failure) = result {
if self.phase == HostPhase::DeploymentFailure {
self.phase = HostPhase::Failure;
return None;
}
let error = match failure {
HostFailure::Cancelled(error) => {
self.phase = HostPhase::Complete;
return Some(error);
}
HostFailure::Error(error) => error,
};
match self.phase {
HostPhase::Failure | HostPhase::AsyncFailure => {
self.advance();
return None;
}
HostPhase::Success => self.phase = HostPhase::Complete,
HostPhase::Execute | HostPhase::ConstructResponse => {
self.phase = HostPhase::MapFailure;
}
_ => self.phase = HostPhase::Failure,
}
return Some(error);
}
self.advance();
None
}
fn advance(&mut self) {
self.phase = match self.phase {
HostPhase::Setup if self.asynchronous => HostPhase::DeploymentPreCall,
HostPhase::Setup | HostPhase::DeploymentPreCall => HostPhase::Prepare,
HostPhase::Prepare => HostPhase::Execute,
HostPhase::Execute => HostPhase::ConstructResponse,
HostPhase::ConstructResponse if self.asynchronous => HostPhase::DeploymentPostCall,
HostPhase::ConstructResponse | HostPhase::DeploymentPostCall => HostPhase::Finalize,
HostPhase::Finalize => HostPhase::Success,
HostPhase::MapFailure if self.asynchronous => HostPhase::DeploymentFailure,
HostPhase::MapFailure | HostPhase::DeploymentFailure => HostPhase::Failure,
HostPhase::Failure if self.asynchronous => HostPhase::AsyncFailure,
HostPhase::Failure
| HostPhase::AsyncFailure
| HostPhase::Success
| HostPhase::Complete => HostPhase::Complete,
};
}
}

426
litellm-rust/crates/core/src/call_lifecycle/mod.rs
View file

@ -1,426 +0,0 @@
use std::future::Future;
use std::time::{Instant, SystemTime, UNIX_EPOCH};
pub mod host;
#[cfg(test)]
#[path = "../../tests/host_lifecycle.rs"]
mod host_tests;
pub mod types;
pub use types::{
CallLifecycleContext, CallLifecyclePhase, CallLifecyclePhaseTiming, CallLifecycleRequest,
CallLifecycleTiming,
};
pub trait CallLifecycleHooks<InitialReq, ProviderReq, Resp>: Send + Sync {
type Error: Send + Sync;
type PreCallFuture<'a>: Future<Output = Result<InitialReq, Self::Error>> + Send + 'a
where
Self: 'a,
InitialReq: 'a,
ProviderReq: 'a,
Resp: 'a;
type DuringCallFuture<'a>: Future<Output = Result<ProviderReq, Self::Error>> + Send + 'a
where
Self: 'a,
InitialReq: 'a,
ProviderReq: 'a,
Resp: 'a;
type SuccessFuture<'a>: Future<Output = ()> + Send + 'a
where
Self: 'a,
Resp: 'a;
type FailureFuture<'a>: Future<Output = ()> + Send + 'a
where
Self: 'a;
fn async_pre_call_hook<'a>(
&'a self,
context: &'a CallLifecycleContext,
request: InitialReq,
) -> Self::PreCallFuture<'a>;
fn async_during_call_hook<'a>(
&'a self,
context: &'a CallLifecycleContext,
request: InitialReq,
) -> Self::DuringCallFuture<'a>;
fn async_log_success_event<'a>(
&'a self,
context: &'a CallLifecycleContext,
response: &'a Resp,
timing: &'a CallLifecycleTiming,
) -> Self::SuccessFuture<'a>;
fn async_log_failure_event<'a>(
&'a self,
context: &'a CallLifecycleContext,
error: &'a Self::Error,
timing: &'a CallLifecycleTiming,
) -> Self::FailureFuture<'a>;
}
pub trait CallLifecycleObserver: Send + Sync {
fn on_phase_start(&self, _context: &CallLifecycleContext, _phase: CallLifecyclePhase) {}
fn on_phase_end(&self, _context: &CallLifecycleContext, _timing: &CallLifecyclePhaseTiming) {}
}
#[derive(Default)]
pub struct NoopCallLifecycleObserver;
impl CallLifecycleObserver for NoopCallLifecycleObserver {}
pub struct CallLifecycle<'a> {
observer: &'a dyn CallLifecycleObserver,
}
impl<'a> CallLifecycle<'a> {
pub fn new(observer: &'a dyn CallLifecycleObserver) -> Self {
Self { observer }
}
pub async fn run_request<InitialReq, ProviderReq, Resp, Hooks, ProviderCall, ProviderFuture>(
&self,
request: InitialReq,
hooks: &Hooks,
provider_call: ProviderCall,
) -> Result<Resp, Hooks::Error>
where
InitialReq: CallLifecycleRequest,
Hooks: CallLifecycleHooks<InitialReq, ProviderReq, Resp>,
ProviderCall: FnOnce(ProviderReq) -> ProviderFuture,
ProviderFuture: Future<Output = Result<Resp, Hooks::Error>>,
{
let context = request.lifecycle_context();
self.run(context, request, hooks, provider_call).await
}
pub async fn run<InitialReq, ProviderReq, Resp, Hooks, ProviderCall, ProviderFuture>(
&self,
context: CallLifecycleContext,
request: InitialReq,
hooks: &Hooks,
provider_call: ProviderCall,
) -> Result<Resp, Hooks::Error>
where
Hooks: CallLifecycleHooks<InitialReq, ProviderReq, Resp>,
ProviderCall: FnOnce(ProviderReq) -> ProviderFuture,
ProviderFuture: Future<Output = Result<Resp, Hooks::Error>>,
{
let call_start = epoch_seconds();
let mut phases = Vec::new();
let pre_call = self.start_phase(&context, CallLifecyclePhase::PreCall);
let request = match hooks.async_pre_call_hook(&context, request).await {
Ok(request) => {
phases.push(self.finish_phase(&context, pre_call));
request
}
Err(error) => {
phases.push(self.finish_phase(&context, pre_call));
self.log_failure(&context, hooks, &error, call_start, &mut phases)
.await;
return Err(error);
}
};
let during_call = self.start_phase(&context, CallLifecyclePhase::DuringCall);
let provider_request = match hooks.async_during_call_hook(&context, request).await {
Ok(request) => {
phases.push(self.finish_phase(&context, during_call));
request
}
Err(error) => {
phases.push(self.finish_phase(&context, during_call));
self.log_failure(&context, hooks, &error, call_start, &mut phases)
.await;
return Err(error);
}
};
let provider_phase = self.start_phase(&context, CallLifecyclePhase::ProviderCall);
let result = provider_call(provider_request).await;
phases.push(self.finish_phase(&context, provider_phase));
match &result {
Ok(response) => {
let success_phase = self.start_phase(&context, CallLifecyclePhase::SuccessCallback);
let timing = CallLifecycleTiming::new(call_start, epoch_seconds(), phases.clone());
hooks
.async_log_success_event(&context, response, &timing)
.await;
phases.push(self.finish_phase(&context, success_phase));
}
Err(error) => {
self.log_failure(&context, hooks, error, call_start, &mut phases)
.await;
}
}
result
}
async fn log_failure<InitialReq, ProviderReq, Resp, Hooks>(
&self,
context: &CallLifecycleContext,
hooks: &Hooks,
error: &Hooks::Error,
call_start: f64,
phases: &mut Vec<CallLifecyclePhaseTiming>,
) where
Hooks: CallLifecycleHooks<InitialReq, ProviderReq, Resp>,
{
let failure_phase = self.start_phase(context, CallLifecyclePhase::FailureCallback);
let timing = CallLifecycleTiming::new(call_start, epoch_seconds(), phases.clone());
hooks.async_log_failure_event(context, error, &timing).await;
phases.push(self.finish_phase(context, failure_phase));
}
fn start_phase(&self, context: &CallLifecycleContext, phase: CallLifecyclePhase) -> PhaseStart {
self.observer.on_phase_start(context, phase);
PhaseStart {
phase,
start_time: epoch_seconds(),
started_at: Instant::now(),
}
}
fn finish_phase(
&self,
context: &CallLifecycleContext,
phase_start: PhaseStart,
) -> CallLifecyclePhaseTiming {
let timing = CallLifecyclePhaseTiming {
phase: phase_start.phase,
start_time: phase_start.start_time,
end_time: epoch_seconds(),
duration: phase_start.started_at.elapsed(),
};
self.observer.on_phase_end(context, &timing);
timing
}
}
impl Default for CallLifecycle<'static> {
fn default() -> Self {
static OBSERVER: NoopCallLifecycleObserver = NoopCallLifecycleObserver;
Self::new(&OBSERVER)
}
}
struct PhaseStart {
phase: CallLifecyclePhase,
start_time: f64,
started_at: Instant,
}
fn epoch_seconds() -> f64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|duration| duration.as_secs_f64())
.unwrap_or(0.0)
}
#[cfg(test)]
mod tests {
use super::*;
use std::pin::Pin;
use std::sync::Mutex;
type BoxFuture<'a, T> = Pin<Box<dyn Future<Output = T> + Send + 'a>>;
#[derive(Default)]
struct RecordingHooks {
events: Mutex<Vec<&'static str>>,
}
struct RecordingRequest(String);
impl CallLifecycleRequest for RecordingRequest {
fn lifecycle_context(&self) -> CallLifecycleContext {
CallLifecycleContext::new("ocr", "mistral-ocr-latest", "mistral", "call_1")
}
}
impl RecordingHooks {
fn events(&self) -> Vec<&'static str> {
self.events.lock().unwrap().clone()
}
}
impl CallLifecycleHooks<String, String, String> for RecordingHooks {
type Error = crate::messages::Error;
type PreCallFuture<'a> = BoxFuture<'a, Result<String, crate::messages::Error>>;
type DuringCallFuture<'a> = BoxFuture<'a, Result<String, crate::messages::Error>>;
type SuccessFuture<'a> = BoxFuture<'a, ()>;
type FailureFuture<'a> = BoxFuture<'a, ()>;
fn async_pre_call_hook<'a>(
&'a self,
_context: &'a CallLifecycleContext,
request: String,
) -> Self::PreCallFuture<'a> {
Box::pin(async move {
self.events.lock().unwrap().push("pre_call");
Ok(format!("{request}:pre"))
})
}
fn async_during_call_hook<'a>(
&'a self,
_context: &'a CallLifecycleContext,
request: String,
) -> Self::DuringCallFuture<'a> {
Box::pin(async move {
self.events.lock().unwrap().push("during_call");
Ok(format!("{request}:during"))
})
}
fn async_log_success_event<'a>(
&'a self,
_context: &'a CallLifecycleContext,
_response: &'a String,
timing: &'a CallLifecycleTiming,
) -> Self::SuccessFuture<'a> {
Box::pin(async move {
assert!(timing.end_time >= timing.start_time);
assert_eq!(timing.phases.len(), 3);
self.events.lock().unwrap().push("success");
})
}
fn async_log_failure_event<'a>(
&'a self,
_context: &'a CallLifecycleContext,
_error: &'a crate::messages::Error,
_timing: &'a CallLifecycleTiming,
) -> Self::FailureFuture<'a> {
Box::pin(async move {
self.events.lock().unwrap().push("failure");
})
}
}
impl CallLifecycleHooks<RecordingRequest, String, String> for RecordingHooks {
type Error = crate::messages::Error;
type PreCallFuture<'a> = BoxFuture<'a, Result<RecordingRequest, crate::messages::Error>>;
type DuringCallFuture<'a> = BoxFuture<'a, Result<String, crate::messages::Error>>;
type SuccessFuture<'a> = BoxFuture<'a, ()>;
type FailureFuture<'a> = BoxFuture<'a, ()>;
fn async_pre_call_hook<'a>(
&'a self,
_context: &'a CallLifecycleContext,
request: RecordingRequest,
) -> Self::PreCallFuture<'a> {
Box::pin(async move {
self.events.lock().unwrap().push("pre_call");
Ok(RecordingRequest(format!("{}:pre", request.0)))
})
}
fn async_during_call_hook<'a>(
&'a self,
_context: &'a CallLifecycleContext,
request: RecordingRequest,
) -> Self::DuringCallFuture<'a> {
Box::pin(async move {
self.events.lock().unwrap().push("during_call");
Ok(format!("{}:during", request.0))
})
}
fn async_log_success_event<'a>(
&'a self,
_context: &'a CallLifecycleContext,
_response: &'a String,
_timing: &'a CallLifecycleTiming,
) -> Self::SuccessFuture<'a> {
Box::pin(async move {
self.events.lock().unwrap().push("success");
})
}
fn async_log_failure_event<'a>(
&'a self,
_context: &'a CallLifecycleContext,
_error: &'a crate::messages::Error,
_timing: &'a CallLifecycleTiming,
) -> Self::FailureFuture<'a> {
Box::pin(async move {
self.events.lock().unwrap().push("failure");
})
}
}
#[tokio::test]
async fn lifecycle_runs_hooks_around_provider_call() {
let hooks = RecordingHooks::default();
let response = CallLifecycle::default()
.run(
CallLifecycleContext::new("ocr", "mistral-ocr-latest", "mistral", "call_1"),
"request".to_string(),
&hooks,
|request| async move {
assert_eq!(request, "request:pre:during");
Ok("response".to_string())
},
)
.await
.expect("call succeeds");
assert_eq!(response, "response");
assert_eq!(hooks.events(), vec!["pre_call", "during_call", "success"]);
}
#[tokio::test]
async fn lifecycle_logs_failure_when_provider_fails() {
let hooks = RecordingHooks::default();
let error = CallLifecycle::default()
.run(
CallLifecycleContext::new("ocr", "mistral-ocr-latest", "mistral", "call_1"),
"request".to_string(),
&hooks,
|_request| async move {
Err::<String, crate::messages::Error>(crate::messages::Error::Transport(
crate::transport::Error::Network("provider down".to_string()),
))
},
)
.await
.expect_err("call fails");
assert_eq!(
error,
crate::messages::Error::Transport(crate::transport::Error::Network(
"provider down".to_string()
))
);
assert_eq!(hooks.events(), vec!["pre_call", "during_call", "failure"]);
}
#[tokio::test]
async fn lifecycle_can_run_any_request_with_embedded_context() {
let hooks = RecordingHooks::default();
let response = CallLifecycle::default()
.run_request(
RecordingRequest("request".to_string()),
&hooks,
|request| async move {
assert_eq!(request, "request:pre:during");
Ok("response".to_string())
},
)
.await
.expect("call succeeds");
assert_eq!(response, "response");
assert_eq!(hooks.events(), vec!["pre_call", "during_call", "success"]);
}
}

75
litellm-rust/crates/core/src/call_lifecycle/types.rs
View file

@ -1,75 +0,0 @@
use std::time::Duration;
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct CallLifecycleContext {
pub call_type: String,
pub model: String,
pub custom_llm_provider: String,
pub litellm_call_id: String,
}
impl CallLifecycleContext {
pub fn new(
call_type: impl Into<String>,
model: impl Into<String>,
custom_llm_provider: impl Into<String>,
litellm_call_id: impl Into<String>,
) -> Self {
Self {
call_type: call_type.into(),
model: model.into(),
custom_llm_provider: custom_llm_provider.into(),
litellm_call_id: litellm_call_id.into(),
}
}
}
pub trait CallLifecycleRequest {
fn lifecycle_context(&self) -> CallLifecycleContext;
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum CallLifecyclePhase {
PreCall,
DuringCall,
ProviderCall,
SuccessCallback,
FailureCallback,
}
impl CallLifecyclePhase {
pub fn as_str(self) -> &'static str {
match self {
Self::PreCall => "pre_call",
Self::DuringCall => "during_call",
Self::ProviderCall => "provider_call",
Self::SuccessCallback => "success_callback",
Self::FailureCallback => "failure_callback",
}
}
}
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct CallLifecyclePhaseTiming {
pub phase: CallLifecyclePhase,
pub start_time: f64,
pub end_time: f64,
pub duration: Duration,
}
#[derive(Clone, Debug, PartialEq)]
pub struct CallLifecycleTiming {
pub start_time: f64,
pub end_time: f64,
pub phases: Vec<CallLifecyclePhaseTiming>,
}
impl CallLifecycleTiming {
pub fn new(start_time: f64, end_time: f64, phases: Vec<CallLifecyclePhaseTiming>) -> Self {
Self {
start_time,
end_time,
phases,
}
}
}

3
litellm-rust/crates/core/src/chat_completions/client.rs
View file

@ -1,5 +1,4 @@
use std::sync::OnceLock;
use std::time::Duration;
use std::{sync::OnceLock, time::Duration};
use crate::constants::{CHAT_COMPLETIONS_CONNECT_TIMEOUT_SECS, CHAT_COMPLETIONS_TIMEOUT_SECS};

16
litellm-rust/crates/core/src/chat_completions/common_utils.rs
View file

@ -1,19 +1,19 @@
use super::Error;
use crate::http_utils::string_headers as shared_string_headers;
use crate::providers::anthropic::chat_completions::transformation::ANTHROPIC_CHAT_COMPLETIONS_CONFIG;
use litellm_providers::{
anthropic::chat::transformation::ANTHROPIC_CHAT_COMPLETIONS_CONFIG,
base_llm::chat::transformation::BaseConfig,
};
use serde_json::{Map, Value};
use super::transformation::ChatCompletionsProviderConfig;
use super::Error;
use crate::http_utils::string_headers as shared_string_headers;
const HEADER_CONTEXT: &str = "chat completions";
pub(super) fn chat_completions_provider_config(
provider: &str,
) -> Option<&'static dyn ChatCompletionsProviderConfig> {
pub(super) fn chat_completions_provider_config(provider: &str) -> Option<&'static dyn BaseConfig> {
match provider {
"anthropic" => Some(&ANTHROPIC_CHAT_COMPLETIONS_CONFIG),
"bedrock" => Some(
&crate::providers::bedrock::chat_completions::transformation::BEDROCK_CHAT_COMPLETIONS_CONFIG,
&litellm_providers::bedrock::chat::converse_transformation::BEDROCK_CHAT_COMPLETIONS_CONFIG,
),
_ => None,
}

16
litellm-rust/crates/core/src/chat_completions/error.rs
View file

@ -24,3 +24,19 @@ pub enum Error {
#[error(transparent)]
Aws(#[from] litellm_auth_aws::Error),
}
impl From<litellm_providers::chat::Error> for Error {
fn from(error: litellm_providers::chat::Error) -> Self {
match error {
litellm_providers::chat::Error::MissingField(field) => Self::MissingField(field),
litellm_providers::chat::Error::InvalidRequest(message) => {
Self::InvalidRequest(message)
}
litellm_providers::chat::Error::InvalidResponse(message) => {
Self::InvalidResponse(message)
}
litellm_providers::chat::Error::Unsupported(reason) => Self::Unsupported(reason),
litellm_providers::chat::Error::Auth(error) => Self::Auth(error),
}
}
}

25
litellm-rust/crates/core/src/chat_completions/handler.rs
View file

@ -1,15 +1,16 @@
use litellm_providers::base_llm::chat::transformation::ChatCompletionsAuth;
use serde_json::Value;
use super::Error;
use crate::http_utils::{http_request, truncate_error_body};
use super::client::http_client;
use super::prepare::prepare_provider_request;
use super::transformation::ChatCompletionsAuth;
use super::types::{
ChatCompletionsResponse, ProviderChatCompletionsRequest, ProviderChatResponseData,
ResolvedChatCompletionsRequest,
use super::{
Error,
client::http_client,
prepare::prepare_provider_request,
types::{
ChatCompletionsResponse, ProviderChatCompletionsRequest, ProviderChatResponseData,
ResolvedChatCompletionsRequest,
},
};
use crate::http_utils::{http_request, truncate_error_body};
pub(super) async fn execute_chat_completions_provider_call(
request: ResolvedChatCompletionsRequest<'_>,
@ -60,6 +61,7 @@ pub(super) async fn execute_chat_completions_provider_call(
request
.config
.transform_response(&request.model, ProviderChatResponseData { body })
.map_err(Error::from)
.map_err(as_response_error)
}
@ -84,10 +86,9 @@ pub(super) async fn signed_headers(
request: &ProviderChatCompletionsRequest,
body: &[u8],
) -> Result<Vec<(String, String)>, Error> {
use std::collections::BTreeMap;
use std::time::SystemTime;
use std::{collections::BTreeMap, time::SystemTime};
use crate::providers::bedrock::aws_base::{
use litellm_auth_aws::{
aws_auth_config, aws_signature_headers, host_supplied_credentials,
is_sigv4_computed_header, resolve_credentials, sign_bedrock_post,
};

11
litellm-rust/crates/core/src/chat_completions/mod.rs
View file

@ -10,17 +10,14 @@ mod error;
pub use error::Error;
mod client;
mod common_utils;
pub mod conversation;
pub use litellm_providers::chat::{conversation, response_utils};
pub(crate) mod handler;
mod prepare;
pub mod response_utils;
pub mod transformation;
pub mod types;
use serde_json::{Map, Value};
pub mod streaming;
use handler::execute_chat_completions_provider_call;
pub use litellm_providers::chat::types;
use prepare::{parse_messages, resolve_provider_config, resolve_request};
use serde_json::{Map, Value};
use types::{ChatCompletionsRequest, ChatCompletionsResponse};
pub async fn chat_completions(

27
litellm-rust/crates/core/src/chat_completions/prepare.rs
View file

@ -1,20 +1,23 @@
use litellm_providers::base_llm::chat::transformation::{BaseConfig, ChatCompletionsAuth};
use serde_json::Value;
use super::Error;
use crate::http_utils::has_header;
use crate::providers::custom_llm_provider::{CustomLlmProvider, get_custom_llm_provider};
use super::common_utils::{chat_completions_provider_config, string_headers};
use super::transformation::{ChatCompletionsAuth, ChatCompletionsProviderConfig};
use super::types::{
ChatCompletionsRequest, ChatMessage, ProviderChatCompletionsRequest,
ResolvedChatCompletionsRequest,
use super::{
Error,
common_utils::{chat_completions_provider_config, string_headers},
types::{
ChatCompletionsRequest, ChatMessage, ProviderChatCompletionsRequest,
ResolvedChatCompletionsRequest,
},
};
use crate::{
http_utils::has_header,
litellm_core_utils::get_llm_provider_logic::{CustomLlmProvider, get_custom_llm_provider},
};
pub(super) fn resolve_provider_config<'a>(
model: &'a str,
custom_llm_provider: Option<&'a str>,
) -> Result<(String, &'static dyn ChatCompletionsProviderConfig), Error> {
) -> Result<(String, &'static dyn BaseConfig), Error> {
let provider_info = get_custom_llm_provider(model, custom_llm_provider)
.or_else(|| {
custom_llm_provider.map(|provider| CustomLlmProvider {
@ -65,7 +68,7 @@ pub(super) fn resolve_request(
fn validate_environment(
request: &ResolvedChatCompletionsRequest<'_>,
model: &str,
config: &dyn ChatCompletionsProviderConfig,
config: &dyn BaseConfig,
) -> Result<(Vec<(String, String)>, ChatCompletionsAuth), Error> {
let env_lookup = |key: &str| std::env::var(key).ok();
let mut headers = string_headers(request.extra_headers.clone())?;
@ -122,7 +125,7 @@ pub(super) fn prepare_provider_request(
let model = request.model;
let config = request.config;
let env_lookup = |key: &str| std::env::var(key).ok();
let url = config.complete_url(
let url = config.get_complete_url(
request.api_base,
&model,
&request.optional_params,

9
litellm-rust/crates/core/src/chat_completions/streaming.rs Normal file
View file

@ -0,0 +1,9 @@
pub trait StreamTransformer {
type Input;
type Output;
type Error;
fn transform(&mut self, input: Self::Input) -> Result<Vec<Self::Output>, Self::Error>;
fn finish(&mut self) -> Result<Vec<Self::Output>, Self::Error>;
}

19
litellm-rust/crates/core/src/chat_completions/tests.rs
View file

@ -1,10 +1,11 @@
use litellm_providers::base_llm::chat::transformation::ChatCompletionsAuth;
use serde_json::{Map, Value, json};
use super::Error;
use super::prepare::{prepare_provider_request, resolve_request};
use super::transformation::ChatCompletionsAuth;
use super::types::{ChatCompletionsRequest, ProviderChatCompletionsRequest};
use super::{
Error,
prepare::{prepare_provider_request, resolve_request},
types::{ChatCompletionsRequest, ProviderChatCompletionsRequest},
};
fn prepare_chat_completions_call(
request: ChatCompletionsRequest<'_>,
@ -588,10 +589,12 @@ fn the_gate_agrees_with_prepare_on_every_case_it_accepts() {
}
mod round_trip {
use super::*;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream};
use tokio::{
io::{AsyncReadExt, AsyncWriteExt},
net::{TcpListener, TcpStream},
};
use super::*;
use crate::chat_completions::chat_completions;
async fn read_http_request(socket: &mut TcpStream) -> String {

122
litellm-rust/crates/core/src/chat_completions/types.rs
View file

@ -1,122 +0,0 @@
use std::time::Duration;
use serde::{Deserialize, Serialize};
use serde_json::{Map, Value};
use super::transformation::{ChatCompletionsAuth, ChatCompletionsProviderConfig};
/// A `/chat/completions` call as it crosses into the core.
///
/// `optional_params` arrives already mapped to the provider's own parameter
/// names by the host, exactly as the messages route receives an already
/// Anthropic-shaped body. The core owns the conversation translation, the
/// provider call, and the response normalization.
pub struct ChatCompletionsRequest<'a> {
pub model: &'a str,
pub messages: Value,
pub optional_params: Map<String, Value>,
pub api_key: Option<&'a str>,
pub api_base: Option<&'a str>,
pub custom_llm_provider: Option<&'a str>,
pub extra_headers: Option<Map<String, Value>>,
pub timeout: Option<Duration>,
}
pub(super) struct ResolvedChatCompletionsRequest<'a> {
pub(super) model: String,
pub(super) config: &'static dyn ChatCompletionsProviderConfig,
pub(super) messages: Vec<ChatMessage>,
pub(super) optional_params: Map<String, Value>,
pub(super) api_key: Option<&'a str>,
pub(super) api_base: Option<&'a str>,
pub(super) extra_headers: Option<Map<String, Value>>,
pub(super) timeout: Option<Duration>,
}
pub(super) struct ProviderChatCompletionsRequest {
pub(super) model: String,
pub(super) config: &'static dyn ChatCompletionsProviderConfig,
pub(super) url: String,
pub(super) body: Value,
pub(super) upstream_headers: Vec<(String, String)>,
pub(super) auth: ChatCompletionsAuth,
pub(super) optional_params: Map<String, Value>,
pub(super) timeout: Option<Duration>,
}
/// The provider-shaped request body a config produces. Named rather than a bare
/// `Value` so the transform contract stays a typed one, mirroring
/// [`crate::audio_transcription::types::AudioTranscriptionRequestData`].
pub struct ProviderChatRequestData {
pub body: Value,
}
/// The raw provider response body handed back to a config for normalization.
pub struct ProviderChatResponseData {
pub body: Value,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[serde(untagged)]
pub enum ChatMessageContent {
Text(String),
Parts(Vec<Value>),
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct ChatMessage {
pub role: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub content: Option<ChatMessageContent>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub name: Option<String>,
#[serde(flatten)]
pub extra: Map<String, Value>,
}
/// OpenAI `usage`, including the `prompt_tokens_details` split LiteLLM's Python
/// path reports so cost tracking sees the same numbers on either path.
#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
pub struct PromptTokensDetails {
pub cached_tokens: u64,
pub cache_creation_tokens: u64,
pub text_tokens: u64,
}
#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
pub struct ChatCompletionsUsage {
pub prompt_tokens: u64,
pub completion_tokens: u64,
pub total_tokens: u64,
pub prompt_tokens_details: PromptTokensDetails,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct ChatCompletionsChoiceMessage {
pub role: String,
// Whether an empty turn is `None` or `""` is the provider's choice, not a
// shared invariant: Anthropic's transform ends on `merged_text or None`
// while Converse assigns the joined string unconditionally. Each config
// mirrors its own, so keep this optional and serialize it even when None.
pub content: Option<String>,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct ChatCompletionsChoice {
pub index: u64,
pub message: ChatCompletionsChoiceMessage,
pub finish_reason: String,
}
/// The normalized response handed back to the host.
///
/// There is deliberately no `id`: Python mints the `chatcmpl-…` id on the
/// `ModelResponse` it already created, and echoing the provider's own id here
/// would change it. Pinned by `response_carries_no_id` in `tests.rs`.
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct ChatCompletionsResponse {
pub created: u64,
pub model: String,
pub choices: Vec<ChatCompletionsChoice>,
pub usage: ChatCompletionsUsage,
}

3
litellm-rust/crates/core/src/http_utils.rs
View file

@ -131,9 +131,10 @@ pub fn json_type_name(value: &serde_json::Value) -> &'static str {
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
use super::*;
#[rstest::rstest]
#[case(HeaderPolicy::All, true, true)]
#[case(HeaderPolicy::Only(&["authorization"]), true, false)]

8
litellm-rust/crates/core/src/lib.rs
View file

@ -1,14 +1,18 @@
pub mod audio_transcription;
pub mod call_lifecycle;
pub mod call_arguments;
pub mod chat_completions;
pub mod constants;
pub mod error;
pub mod http_utils;
pub mod litellm_core_utils;
pub mod llms;
pub mod machine;
mod media;
pub mod messages;
pub mod ocr;
pub mod providers;
pub mod params;
pub mod responses;
mod serde_compat;
pub mod transport;
mod url_utils;

34
litellm-rust/crates/core/src/providers/custom_llm_provider.rs → litellm-rust/crates/core/src/litellm_core_utils/get_llm_provider_logic.rs
View file

@ -1,36 +1,4 @@
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct CustomLlmProvider<'a> {
pub model: &'a str,
pub custom_llm_provider: &'a str,
}
pub fn get_custom_llm_provider<'a>(
model: &'a str,
custom_llm_provider: Option<&'a str>,
) -> Option<CustomLlmProvider<'a>> {
if let Some(custom_llm_provider) = custom_llm_provider.filter(|provider| !provider.is_empty()) {
return Some(CustomLlmProvider {
model: strip_custom_llm_provider_prefix(model, custom_llm_provider),
custom_llm_provider,
});
}
let (custom_llm_provider, model) = model.split_once('/')?;
if custom_llm_provider.is_empty() || model.is_empty() {
return None;
}
Some(CustomLlmProvider {
model,
custom_llm_provider,
})
}
fn strip_custom_llm_provider_prefix<'a>(model: &'a str, custom_llm_provider: &str) -> &'a str {
model
.strip_prefix(custom_llm_provider)
.and_then(|model| model.strip_prefix('/'))
.unwrap_or(model)
}
pub use litellm_providers::provider_resolution::{CustomLlmProvider, get_custom_llm_provider};
#[cfg(test)]
mod tests {

1
litellm-rust/crates/core/src/litellm_core_utils/mod.rs Normal file
View file

@ -0,0 +1 @@
pub mod get_llm_provider_logic;

1
litellm-rust/crates/core/src/llms/anthropic/chat/mod.rs Normal file
View file

@ -0,0 +1 @@
pub mod streaming;

166
litellm-rust/crates/core/src/llms/anthropic/chat/streaming.rs Normal file
View file

@ -0,0 +1,166 @@
use std::collections::HashMap;
use serde_json::Value;
use super::super::experimental_pass_through::messages::streaming::{
AnthropicContentBlock, AnthropicContentBlockDelta, AnthropicMessagesStreamEvent,
AnthropicStreamUsage,
};
use crate::chat_completions::{
Error,
streaming::StreamTransformer,
types::{
ChatCompletionChunk, ChatCompletionThinkingBlock, ChatCompletionToolCallChunk,
ChatCompletionsUsage,
},
};
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum AnthropicJsonChunkType {
ValidJson,
AccumulatedJson,
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum AnthropicContentBlockType {
Text,
ToolUse,
ServerToolUse,
Thinking,
RedactedThinking,
Compaction,
ToolResult(String),
Other(String),
}
#[derive(Clone, Debug, PartialEq)]
pub struct AnthropicContentBlockDeltaEvent {
pub index: u64,
pub delta: AnthropicContentBlockDelta,
}
pub struct AnthropicChatCompletionsStreamTransformer {
pub content_blocks: Vec<AnthropicContentBlockDeltaEvent>,
pub tool_index: i64,
pub json_mode: bool,
pub speed: Option<String>,
pub tool_name_reverse_map: HashMap<String, String>,
pub response_id: String,
pub served_model: Option<String>,
pub is_response_format_tool: bool,
pub converted_response_format_tool: bool,
pub accumulated_json: String,
pub chunk_type: AnthropicJsonChunkType,
pub current_content_block_type: Option<AnthropicContentBlockType>,
pub web_search_results: Vec<Value>,
pub web_search_calls: HashMap<String, Value>,
pub compaction_blocks: Vec<Value>,
pub reasoning_content_chunks: Vec<String>,
pub server_tool_inputs: HashMap<String, Value>,
pub tool_results: Vec<Value>,
pub current_server_tool_id: Option<String>,
pub container_id: Option<String>,
}
impl AnthropicChatCompletionsStreamTransformer {
pub fn new(
_json_mode: bool,
_speed: Option<String>,
_tool_name_reverse_map: HashMap<String, String>,
) -> Self {
todo!()
}
pub fn check_empty_tool_call_args(&self) -> bool {
todo!()
}
pub fn handle_usage(&mut self, _usage: AnthropicStreamUsage) -> ChatCompletionsUsage {
todo!()
}
pub fn handle_content_block_delta(
&mut self,
_index: u64,
_delta: AnthropicContentBlockDelta,
) -> (
String,
Option<ChatCompletionToolCallChunk>,
Vec<ChatCompletionThinkingBlock>,
Option<Value>,
Option<String>,
) {
todo!()
}
pub fn handle_content_block_start(
&mut self,
_index: u64,
_content_block: AnthropicContentBlock,
) -> Result<ChatCompletionChunk, Error> {
todo!()
}
pub fn handle_json_mode_chunk(
&mut self,
_text: String,
_tool_use: Option<ChatCompletionToolCallChunk>,
) -> (String, Option<ChatCompletionToolCallChunk>) {
todo!()
}
pub fn handle_accumulated_json_chunk(
&mut self,
_data: &str,
_is_final: bool,
) -> Result<Option<ChatCompletionChunk>, Error> {
todo!()
}
pub fn handle_redacted_thinking_content(
&mut self,
_content_block: &AnthropicContentBlock,
) -> Vec<ChatCompletionThinkingBlock> {
todo!()
}
pub fn web_search_call_snapshot(&self) -> HashMap<String, Value> {
todo!()
}
pub fn complete_web_search_call(&mut self, _result: Value) {
todo!()
}
pub fn build_code_interpreter_results(&self) -> Vec<Value> {
todo!()
}
pub fn handle_message_delta(
&mut self,
_event: AnthropicMessagesStreamEvent,
) -> (Option<String>, Option<ChatCompletionsUsage>, Option<Value>) {
todo!()
}
pub fn chunk_parser(
&mut self,
_event: AnthropicMessagesStreamEvent,
) -> Result<ChatCompletionChunk, Error> {
todo!()
}
}
impl StreamTransformer for AnthropicChatCompletionsStreamTransformer {
type Input = AnthropicMessagesStreamEvent;
type Output = ChatCompletionChunk;
type Error = Error;
fn transform(&mut self, _input: Self::Input) -> Result<Vec<Self::Output>, Self::Error> {
todo!()
}
fn finish(&mut self) -> Result<Vec<Self::Output>, Self::Error> {
todo!()
}
}

337
litellm-rust/crates/core/src/llms/anthropic/experimental_pass_through/messages/batches.rs Normal file
View file

@ -0,0 +1,337 @@
use litellm_providers::anthropic::experimental_pass_through::messages::transformation::resolve_anthropic_api_base;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use time::OffsetDateTime;
use url::Url;
use crate::messages::{Error, types::AnthropicMessagesResponse};
const BATCHES_PATH_SUFFIX: &str = "/v1/messages/batches";
#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct AnthropicBatchRequestCounts {
#[serde(default)]
pub processing: u64,
#[serde(default)]
pub succeeded: u64,
#[serde(default)]
pub errored: u64,
#[serde(default)]
pub canceled: u64,
#[serde(default)]
pub expired: u64,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct AnthropicMessageBatch {
#[serde(default)]
pub id: String,
#[serde(default = "default_processing_status")]
pub processing_status: String,
pub created_at: Option<String>,
pub ended_at: Option<String>,
pub expires_at: Option<String>,
pub cancel_initiated_at: Option<String>,
pub archived_at: Option<String>,
#[serde(default)]
pub request_counts: AnthropicBatchRequestCounts,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum BatchStatus {
InProgress,
Cancelling,
Completed,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct BatchRequestCounts {
pub total: u64,
pub completed: u64,
pub failed: u64,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct LiteLlmMessageBatch {
pub id: String,
pub object: String,
pub endpoint: String,
pub input_file_id: String,
pub completion_window: String,
pub status: BatchStatus,
pub output_file_id: String,
pub created_at: i64,
pub in_progress_at: Option<i64>,
pub expires_at: Option<i64>,
pub completed_at: Option<i64>,
pub expired_at: Option<i64>,
pub cancelling_at: Option<i64>,
pub cancelled_at: Option<i64>,
pub request_counts: BatchRequestCounts,
}
pub trait AnthropicBatchesConfig {
fn create_batch_url(
&self,
api_base: Option<&str>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<String, Error>;
fn transform_create_batch_request(&self) -> Result<Value, Error>;
fn transform_create_batch_response(
&self,
response: AnthropicMessageBatch,
now: i64,
) -> Result<LiteLlmMessageBatch, Error>;
fn retrieve_batch_url(
&self,
api_base: Option<&str>,
batch_id: &str,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<String, Error>;
fn transform_retrieve_batch_request(&self) -> Value;
fn transform_retrieve_batch_response(
&self,
response: AnthropicMessageBatch,
now: i64,
) -> LiteLlmMessageBatch;
fn transform_batch_results(&self, body: &str) -> Result<Vec<AnthropicMessagesResponse>, Error>;
}
pub struct AnthropicBatchesTransformation;
pub const ANTHROPIC_BATCHES_TRANSFORMATION: AnthropicBatchesTransformation =
AnthropicBatchesTransformation;
fn default_processing_status() -> String {
"in_progress".into()
}
fn timestamp(value: Option<&str>) -> Option<i64> {
value
.and_then(|value| {
OffsetDateTime::parse(value, &time::format_description::well_known::Rfc3339).ok()
})
.map(OffsetDateTime::unix_timestamp)
}
fn batches_base_url(
api_base: Option<&str>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<Url, Error> {
let api_base = resolve_anthropic_api_base(api_base, env_lookup);
let api_base = api_base.trim_end_matches('/');
let complete_url = if api_base.ends_with(BATCHES_PATH_SUFFIX) {
api_base.to_string()
} else if let Some(base) = api_base.strip_suffix("/v1/messages") {
format!("{base}{BATCHES_PATH_SUFFIX}")
} else {
format!("{api_base}{BATCHES_PATH_SUFFIX}")
};
Url::parse(&complete_url)
.map_err(|error| Error::InvalidRequest(format!("invalid Anthropic API base: {error}")))
}
impl AnthropicBatchesConfig for AnthropicBatchesTransformation {
fn create_batch_url(
&self,
api_base: Option<&str>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<String, Error> {
Ok(batches_base_url(api_base, env_lookup)?.into())
}
fn transform_create_batch_request(&self) -> Result<Value, Error> {
Err(Error::Unsupported("Anthropic message batch creation"))
}
fn transform_create_batch_response(
&self,
_response: AnthropicMessageBatch,
_now: i64,
) -> Result<LiteLlmMessageBatch, Error> {
Err(Error::Unsupported("Anthropic message batch creation"))
}
fn retrieve_batch_url(
&self,
api_base: Option<&str>,
batch_id: &str,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<String, Error> {
if batch_id.is_empty() {
return Err(Error::MissingField("batch_id"));
}
let mut url = batches_base_url(api_base, env_lookup)?;
url.path_segments_mut()
.map_err(|_| Error::InvalidRequest("Anthropic API base cannot be a base URL".into()))?
.push(batch_id);
Ok(url.into())
}
fn transform_retrieve_batch_request(&self) -> Value {
Value::Object(Default::default())
}
fn transform_retrieve_batch_response(
&self,
response: AnthropicMessageBatch,
now: i64,
) -> LiteLlmMessageBatch {
let created_at = timestamp(response.created_at.as_deref());
let ended_at = timestamp(response.ended_at.as_deref());
let expires_at = timestamp(response.expires_at.as_deref());
let cancel_initiated_at = timestamp(response.cancel_initiated_at.as_deref());
let archived_at = timestamp(response.archived_at.as_deref());
let status = match response.processing_status.as_str() {
"canceling" => BatchStatus::Cancelling,
"ended" => BatchStatus::Completed,
_ => BatchStatus::InProgress,
};
let request_counts = BatchRequestCounts {
total: response.request_counts.processing
+ response.request_counts.succeeded
+ response.request_counts.errored
+ response.request_counts.canceled
+ response.request_counts.expired,
completed: response.request_counts.succeeded,
failed: response.request_counts.errored,
};
LiteLlmMessageBatch {
id: response.id.clone(),
object: "batch".into(),
endpoint: "/v1/messages".into(),
input_file_id: "None".into(),
completion_window: "24h".into(),
status,
output_file_id: response.id,
created_at: created_at.unwrap_or(now),
in_progress_at: (response.processing_status == "in_progress")
.then_some(created_at)
.flatten(),
expires_at,
completed_at: (response.processing_status == "ended")
.then_some(ended_at)
.flatten(),
expired_at: archived_at,
cancelling_at: (response.processing_status == "canceling")
.then_some(cancel_initiated_at)
.flatten(),
cancelled_at: (response.processing_status == "canceling")
.then_some(ended_at)
.flatten(),
request_counts,
}
}
fn transform_batch_results(&self, body: &str) -> Result<Vec<AnthropicMessagesResponse>, Error> {
body.lines()
.filter(|line| !line.trim().is_empty())
.filter_map(|line| serde_json::from_str::<Value>(line.trim()).ok())
.map(|record| {
serde_json::from_value(record["result"]["message"].clone()).map_err(|error| {
Error::InvalidResponse(format!("invalid Anthropic batch result: {error}"))
})
})
.collect()
}
}
#[cfg(test)]
mod tests {
use serde_json::json;
use super::*;
#[test]
fn builds_and_encodes_message_batch_urls() {
assert_eq!(
ANTHROPIC_BATCHES_TRANSFORMATION
.create_batch_url(None, &|_| None)
.unwrap(),
"https://api.anthropic.com/v1/messages/batches"
);
assert_eq!(
ANTHROPIC_BATCHES_TRANSFORMATION
.create_batch_url(Some("https://proxy.test/v1/messages/batches"), &|_| None)
.unwrap(),
"https://proxy.test/v1/messages/batches"
);
assert_eq!(
ANTHROPIC_BATCHES_TRANSFORMATION
.retrieve_batch_url(Some("https://proxy.test"), "batch/id ?", &|_| None)
.unwrap(),
"https://proxy.test/v1/messages/batches/batch%2Fid%20%3F"
);
assert_eq!(
ANTHROPIC_BATCHES_TRANSFORMATION.transform_retrieve_batch_request(),
json!({})
);
}
#[test]
fn maps_retrieved_batch_status_counts_and_timestamps_like_python() {
let response: AnthropicMessageBatch = serde_json::from_value(json!({
"id": "msgbatch_1",
"processing_status": "ended",
"created_at": "2025-01-01T00:00:00Z",
"ended_at": "2025-01-01T00:01:00Z",
"expires_at": "not-a-timestamp",
"request_counts": {
"processing": 1,
"succeeded": 2,
"errored": 3,
"canceled": 4,
"expired": 5
}
}))
.unwrap();
let batch = ANTHROPIC_BATCHES_TRANSFORMATION.transform_retrieve_batch_response(response, 7);
assert_eq!(batch.status, BatchStatus::Completed);
assert_eq!(batch.created_at, 1_735_689_600);
assert_eq!(batch.completed_at, Some(1_735_689_660));
assert_eq!(batch.expires_at, None);
assert_eq!(
batch.request_counts,
BatchRequestCounts {
total: 15,
completed: 2,
failed: 3
}
);
}
#[test]
fn extracts_message_responses_from_ndjson_and_skips_non_json_lines() {
let body = r#"not-json
{"result":{"message":{"id":"msg_1","type":"message","role":"assistant","model":"claude-test","content":[],"stop_reason":"end_turn","stop_sequence":null}}}
"#;
let messages = ANTHROPIC_BATCHES_TRANSFORMATION
.transform_batch_results(body)
.unwrap();
assert_eq!(messages.len(), 1);
assert_eq!(messages[0].id, "msg_1");
}
#[test]
fn preserves_python_placeholder_for_batch_creation() {
assert!(matches!(
ANTHROPIC_BATCHES_TRANSFORMATION.transform_create_batch_request(),
Err(Error::Unsupported("Anthropic message batch creation"))
));
let response: AnthropicMessageBatch = serde_json::from_value(json!({})).unwrap();
assert!(matches!(
ANTHROPIC_BATCHES_TRANSFORMATION.transform_create_batch_response(response, 0),
Err(Error::Unsupported("Anthropic message batch creation"))
));
}
}

172
litellm-rust/crates/core/src/llms/anthropic/experimental_pass_through/messages/count_tokens.rs Normal file
View file

@ -0,0 +1,172 @@
use serde::{Deserialize, Serialize};
use serde_json::Value;
use crate::{
constants::ANTHROPIC_OAUTH_TOKEN_PREFIX,
messages::{
Error,
types::{AnthropicMessage, SystemPrompt},
},
};
const COUNT_TOKENS_ENDPOINT: &str = "https://api.anthropic.com/v1/messages/count_tokens";
const TOKEN_COUNTING_BETA: &str = "token-counting-2024-11-01";
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct AnthropicCountTokensRequest {
pub model: String,
pub messages: Vec<AnthropicMessage>,
#[serde(skip_serializing_if = "Option::is_none")]
pub tools: Option<Vec<Value>>,
#[serde(skip_serializing_if = "Option::is_none")]
pub system: Option<SystemPrompt>,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct AnthropicCountTokensResponse {
pub input_tokens: u64,
}
pub trait AnthropicCountTokensConfig {
fn endpoint(&self) -> &'static str;
fn validate_request(&self, model: &str, messages: &[AnthropicMessage]) -> Result<(), Error>;
fn transform_request(
&self,
model: &str,
messages: Vec<AnthropicMessage>,
tools: Option<Vec<Value>>,
system: Option<SystemPrompt>,
) -> Result<AnthropicCountTokensRequest, Error>;
fn required_headers(&self, api_key: &str) -> Vec<(&'static str, String)>;
}
pub struct AnthropicCountTokensTransformation;
pub const ANTHROPIC_COUNT_TOKENS_TRANSFORMATION: AnthropicCountTokensTransformation =
AnthropicCountTokensTransformation;
impl AnthropicCountTokensConfig for AnthropicCountTokensTransformation {
fn endpoint(&self) -> &'static str {
COUNT_TOKENS_ENDPOINT
}
fn transform_request(
&self,
model: &str,
messages: Vec<AnthropicMessage>,
tools: Option<Vec<Value>>,
system: Option<SystemPrompt>,
) -> Result<AnthropicCountTokensRequest, Error> {
self.validate_request(model, &messages)?;
Ok(AnthropicCountTokensRequest {
model: model.to_string(),
messages,
tools,
system,
})
}
fn validate_request(&self, model: &str, messages: &[AnthropicMessage]) -> Result<(), Error> {
if model.is_empty() {
return Err(Error::MissingField("model"));
}
if messages.is_empty() {
return Err(Error::MissingField("messages"));
}
Ok(())
}
fn required_headers(&self, api_key: &str) -> Vec<(&'static str, String)> {
let auth = if api_key.starts_with(ANTHROPIC_OAUTH_TOKEN_PREFIX) {
("authorization", format!("Bearer {api_key}"))
} else {
("x-api-key", api_key.to_string())
};
vec![
("content-type", "application/json".to_string()),
auth,
("anthropic-version", "2023-06-01".to_string()),
("anthropic-beta", TOKEN_COUNTING_BETA.to_string()),
]
}
}
#[cfg(test)]
mod tests {
use serde_json::{Map, json};
use super::*;
use crate::messages::types::MessageContent;
fn message() -> AnthropicMessage {
AnthropicMessage {
role: "user".into(),
content: MessageContent::Text("hello".into()),
extra: Map::new(),
}
}
#[test]
fn maps_the_python_count_tokens_contract() {
let request = ANTHROPIC_COUNT_TOKENS_TRANSFORMATION
.transform_request(
"claude-test",
vec![message()],
Some(vec![json!({"name": "lookup"})]),
Some(SystemPrompt::Text("system".into())),
)
.unwrap();
assert_eq!(
serde_json::to_value(request).unwrap(),
json!({
"model": "claude-test",
"messages": [{"role": "user", "content": "hello"}],
"tools": [{"name": "lookup"}],
"system": "system"
})
);
assert_eq!(
ANTHROPIC_COUNT_TOKENS_TRANSFORMATION.endpoint(),
COUNT_TOKENS_ENDPOINT
);
}
#[test]
fn rejects_the_invalid_requests_python_rejects() {
assert!(matches!(
ANTHROPIC_COUNT_TOKENS_TRANSFORMATION.transform_request(
"",
vec![message()],
None,
None
),
Err(Error::MissingField("model"))
));
assert!(matches!(
ANTHROPIC_COUNT_TOKENS_TRANSFORMATION.transform_request(
"claude-test",
vec![],
None,
None
),
Err(Error::MissingField("messages"))
));
}
#[test]
fn uses_api_key_or_oauth_headers_without_combining_credentials() {
let api_key = ANTHROPIC_COUNT_TOKENS_TRANSFORMATION.required_headers("sk-ant-api");
assert!(api_key.contains(&("x-api-key", "sk-ant-api".into())));
assert!(!api_key.iter().any(|(name, _)| *name == "authorization"));
let oauth = ANTHROPIC_COUNT_TOKENS_TRANSFORMATION.required_headers("sk-ant-oat-test");
assert!(oauth.contains(&("authorization", "Bearer sk-ant-oat-test".into())));
assert!(!oauth.iter().any(|(name, _)| *name == "x-api-key"));
assert!(oauth.contains(&("anthropic-beta", TOKEN_COUNTING_BETA.into())));
}
}

3
litellm-rust/crates/core/src/llms/anthropic/experimental_pass_through/messages/mod.rs Normal file
View file

@ -0,0 +1,3 @@
pub mod batches;
pub mod count_tokens;
pub mod streaming;

284
litellm-rust/crates/core/src/llms/anthropic/experimental_pass_through/messages/streaming.rs Normal file
View file

@ -0,0 +1,284 @@
use base64::Engine;
use bytes::Buf;
use futures_util::{Stream, StreamExt};
use litellm_framing::{
Framer,
aws_event_stream::{AwsEventStreamFrame, AwsEventStreamFramer},
sse::{SseFrame, SseFramer},
};
use serde::{Deserialize, Serialize};
use serde_json::{Map, Value};
use crate::messages::Error;
#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
pub struct AnthropicStreamUsage {
#[serde(default)]
pub input_tokens: u64,
#[serde(default)]
pub output_tokens: u64,
#[serde(default)]
pub cache_creation_input_tokens: u64,
#[serde(default)]
pub cache_read_input_tokens: u64,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub server_tool_use: Option<Value>,
#[serde(flatten)]
pub extra: Map<String, Value>,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct AnthropicStreamMessage {
pub id: String,
#[serde(rename = "type")]
pub message_type: String,
pub role: String,
pub model: String,
pub content: Vec<Value>,
pub stop_reason: Option<String>,
pub stop_sequence: Option<String>,
pub usage: AnthropicStreamUsage,
#[serde(flatten)]
pub extra: Map<String, Value>,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum AnthropicContentBlockDelta {
TextDelta {
text: String,
},
InputJsonDelta {
partial_json: String,
},
#[serde(rename = "citations_delta")]
Citations {
citation: Value,
},
ThinkingDelta {
thinking: String,
},
SignatureDelta {
signature: String,
},
CompactionDelta {
content: String,
},
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct AnthropicContentBlock {
#[serde(rename = "type")]
pub block_type: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub id: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub name: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub text: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub input: Option<Value>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub thinking: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub signature: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub data: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub content: Option<Value>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub caller: Option<Value>,
#[serde(flatten)]
pub extra: Map<String, Value>,
}
#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
pub struct AnthropicMessageDelta {
#[serde(default, skip_serializing_if = "Option::is_none")]
pub stop_reason: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub stop_sequence: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub stop_details: Option<Value>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub container: Option<Value>,
#[serde(flatten)]
pub extra: Map<String, Value>,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct AnthropicStreamError {
#[serde(rename = "type")]
pub error_type: String,
pub message: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub details: Option<Value>,
#[serde(flatten)]
pub extra: Map<String, Value>,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum AnthropicMessagesStreamEvent {
MessageStart {
message: AnthropicStreamMessage,
},
ContentBlockStart {
index: u64,
content_block: AnthropicContentBlock,
},
ContentBlockDelta {
index: u64,
delta: AnthropicContentBlockDelta,
},
ContentBlockStop {
index: u64,
},
MessageDelta {
delta: AnthropicMessageDelta,
#[serde(default, skip_serializing_if = "Option::is_none")]
usage: Option<AnthropicStreamUsage>,
#[serde(default, skip_serializing_if = "Option::is_none")]
context_management: Option<Value>,
},
MessageStop,
Ping,
Error {
error: AnthropicStreamError,
},
}
#[derive(Deserialize)]
struct BedrockChunkPayload {
bytes: String,
}
pub fn decode_anthropic_sse_frame(frame: SseFrame) -> Result<AnthropicMessagesStreamEvent, Error> {
let data = frame.data.ok_or(Error::MissingStreamData)?;
serde_json::from_str(&data).map_err(|error| Error::InvalidStreamEvent(error.to_string()))
}
pub fn decode_bedrock_anthropic_frame(
frame: AwsEventStreamFrame,
) -> Result<AnthropicMessagesStreamEvent, Error> {
let payload: BedrockChunkPayload = serde_json::from_slice(&frame.payload)
.map_err(|error| Error::InvalidBedrockPayload(error.to_string()))?;
let event = base64::engine::general_purpose::STANDARD
.decode(payload.bytes)
.map_err(|error| Error::InvalidBedrockBase64(error.to_string()))?;
serde_json::from_slice(&event).map_err(|error| Error::InvalidStreamEvent(error.to_string()))
}
pub fn direct_anthropic_event_stream<S, B, E>(
input: S,
) -> impl Stream<Item = Result<AnthropicMessagesStreamEvent, Error>> + Send
where
S: Stream<Item = Result<B, E>> + Send,
B: Buf + Send,
E: std::error::Error + Send + Sync + 'static,
{
SseFramer.frame(input).map(|frame| {
let frame = frame.map_err(|error| Error::StreamFraming(error.to_string()))?;
decode_anthropic_sse_frame(frame)
})
}
pub fn bedrock_anthropic_event_stream<S, B, E>(
input: S,
) -> impl Stream<Item = Result<AnthropicMessagesStreamEvent, Error>> + Send
where
S: Stream<Item = Result<B, E>> + Send,
B: Buf + Send,
E: std::error::Error + Send + Sync + 'static,
{
AwsEventStreamFramer.frame(input).map(|frame| {
let frame = frame.map_err(|error| Error::StreamFraming(error.to_string()))?;
decode_bedrock_anthropic_frame(frame)
})
}
#[cfg(test)]
mod tests {
use std::io;
use aws_smithy_eventstream::frame::write_message_to;
use aws_smithy_types::event_stream::{Header, HeaderValue, Message};
use base64::engine::general_purpose::STANDARD;
use bytes::Bytes;
use futures_util::TryStreamExt;
use super::*;
const TEXT_DELTA: &str =
r#"{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"hello"}}"#;
#[tokio::test]
async fn direct_anthropic_sse_frames_into_typed_events() {
let wire = format!("event: content_block_delta\ndata: {TEXT_DELTA}\n\n");
let events = direct_anthropic_event_stream(futures_util::stream::iter(
wire.as_bytes().chunks(3).map(Ok::<_, io::Error>),
))
.try_collect::<Vec<_>>()
.await
.unwrap();
assert_eq!(
events,
vec![AnthropicMessagesStreamEvent::ContentBlockDelta {
index: 0,
delta: AnthropicContentBlockDelta::TextDelta {
text: "hello".into(),
},
}]
);
}
#[test]
fn decodes_citations_delta_events() {
let event = decode_anthropic_sse_frame(SseFrame {
event: Some("content_block_delta".into()),
data: Some(
r#"{"type":"content_block_delta","index":0,"delta":{"type":"citations_delta","citation":{"type":"char_location"}}}"#
.into(),
),
id: None,
retry: None,
})
.unwrap();
assert!(matches!(
event,
AnthropicMessagesStreamEvent::ContentBlockDelta {
delta: AnthropicContentBlockDelta::Citations { .. },
..
}
));
}
#[tokio::test]
async fn bedrock_aws_frames_into_the_same_typed_events() {
let payload = serde_json::json!({"bytes": STANDARD.encode(TEXT_DELTA)});
let message = Message::new(Bytes::from(serde_json::to_vec(&payload).unwrap())).add_header(
Header::new(":event-type", HeaderValue::String("chunk".into())),
);
let mut wire = Vec::new();
write_message_to(&message, &mut wire).unwrap();
let events = bedrock_anthropic_event_stream(futures_util::stream::iter(
wire.chunks(3).map(Ok::<_, io::Error>),
))
.try_collect::<Vec<_>>()
.await
.unwrap();
assert_eq!(
events,
vec![AnthropicMessagesStreamEvent::ContentBlockDelta {
index: 0,
delta: AnthropicContentBlockDelta::TextDelta {
text: "hello".into(),
},
}]
);
}
}

0
litellm-rust/crates/core/src/providers/azure_ai/mod.rs → litellm-rust/crates/core/src/llms/anthropic/experimental_pass_through/mod.rs
View file

2
litellm-rust/crates/core/src/llms/anthropic/mod.rs Normal file
View file

@ -0,0 +1,2 @@
pub mod chat;
pub mod experimental_pass_through;

1
litellm-rust/crates/core/src/llms/azure_ai/mod.rs Normal file
View file

@ -0,0 +1 @@
pub(crate) mod ocr;

175
litellm-rust/crates/core/src/llms/azure_ai/ocr/cohere_parse_transformation.rs Normal file
View file

@ -0,0 +1,175 @@
use serde_json::Value;
use crate::{
call_arguments::CallArguments,
llms::{
base_llm::ocr::transformation::{BaseOcrConfig, OcrRequestContext},
cohere::ocr::{
CohereOptions,
transformation::{CohereParseConfig, CohereRequest},
validate_document,
},
},
ocr::{
OcrClient,
document::{inline_remote_document, validate_inline_document},
types::{LiteLLMOcrResponse, OcrDocument, PreparedOcrRequest},
},
url_utils::ApiUrl,
};
#[derive(Default)]
pub(crate) struct AzureAICohereParseConfig;
impl BaseOcrConfig for AzureAICohereParseConfig {
type OcrParams = CohereOptions;
type ProviderRequest = CohereRequest;
type Environment = Vec<(String, String)>;
fn get_api_key_env_var(&self) -> Option<&'static str> {
super::transformation::AzureAiOcrConfig.get_api_key_env_var()
}
fn get_health_check_document(&self) -> OcrDocument {
CohereParseConfig.get_health_check_document()
}
async fn validate_environment(
&self,
request: &PreparedOcrRequest,
client: &OcrClient,
) -> Result<Self::Environment, crate::ocr::Error> {
BaseOcrConfig::validate_environment(
&super::transformation::AzureAiOcrConfig,
request,
client,
)
.await
}
fn get_complete_url(
&self,
request: &PreparedOcrRequest,
_params: &Self::OcrParams,
_environment: &Self::Environment,
) -> Result<String, crate::ocr::Error> {
let base = super::transformation::AzureAiOcrConfig::resolve_api_base(
request.connection.api_base.as_deref(),
&crate::ocr::prepare::credential_env,
)?;
self.get_complete_url(&base)
}
fn transform_ocr_request(
&self,
model: &str,
document: OcrDocument,
params: &CohereOptions,
headers: &[(String, String)],
) -> Result<CohereRequest, crate::ocr::Error> {
CohereParseConfig.transform_ocr_request(model, document, params, headers)
}
fn get_supported_ocr_params(&self, model: &str) -> &'static [&'static str] {
CohereParseConfig.get_supported_ocr_params(model)
}
fn map_ocr_params(
&self,
arguments: &CallArguments,
model: &str,
) -> Result<CohereOptions, crate::ocr::Error> {
CohereParseConfig.map_ocr_params(arguments, model)
}
async fn async_transform_ocr_request(
&self,
model: &str,
document: OcrDocument,
optional_params: &CohereOptions,
headers: &[(String, String)],
context: OcrRequestContext<'_>,
) -> Result<CohereRequest, crate::ocr::Error> {
validate_document(&document)?;
let document = inline_remote_document(
context.client.document_fetcher(),
document,
context.connection,
)
.await?;
self.transform_ocr_request(model, document, optional_params, headers)
}
fn transform_ocr_response(
&self,
model: &str,
raw_response: &[u8],
request_format: crate::ocr::types::OcrResponseFormat,
) -> Result<LiteLLMOcrResponse, crate::ocr::Error> {
CohereParseConfig.transform_ocr_response(model, raw_response, request_format)
}
fn validate_request_body(&self, body: &Value) -> Result<(), crate::ocr::Error> {
let document = crate::ocr::prepare::body_document(body)?;
validate_document(&document)?;
validate_inline_document(&document)
}
}
impl AzureAICohereParseConfig {
fn get_complete_url(&self, base: &str) -> Result<String, crate::ocr::Error> {
let mut url = reqwest::Url::parse(base).map_err(|_| invalid_api_base())?;
if !matches!(url.scheme(), "http" | "https") {
return Err(invalid_api_base());
}
let path = url.path().trim_end_matches('/').to_string();
if path.ends_with("/v2/parse") {
url.set_path(&path);
return Ok(url.into());
}
url.set_path(path.strip_suffix("/models").unwrap_or(&path));
ApiUrl::parse(url.as_str())
.and_then(|url| url.complete_path(&["providers", "cohere", "v2", "parse"]))
.map(|url| url.into_string())
.map_err(|_| invalid_api_base())
}
}
fn invalid_api_base() -> crate::ocr::Error {
crate::ocr::Error::RequestField {
path: "api_base".into(),
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn completes_foundry_urls_without_duplicate_paths_and_preserves_queries() {
for suffix in [
"",
"/models",
"/providers/cohere/v2",
"/providers/cohere/v2/parse",
] {
assert_eq!(
AzureAICohereParseConfig
.get_complete_url(&format!("https://example.com{suffix}?tenant=a"))
.unwrap(),
"https://example.com/providers/cohere/v2/parse?tenant=a"
);
}
assert_eq!(
AzureAICohereParseConfig
.get_complete_url("https://example.com/v2/parse?tenant=a")
.unwrap(),
"https://example.com/v2/parse?tenant=a"
);
assert!(
AzureAICohereParseConfig
.get_complete_url("relative/path")
.is_err()
);
}
}

25
litellm-rust/crates/core/src/ocr/adapters/azure/mod.rs → litellm-rust/crates/core/src/llms/azure_ai/ocr/common_utils.rs
View file

@ -1,25 +1,14 @@
mod cohere;
mod document_intelligence;
mod mistral;
use std::sync::OnceLock;
use crate::ocr::Error;
use crate::ocr::error::OcrError;
use crate::ocr::types::OcrConnection;
use litellm_auth::{InputSource, Sourced};
use litellm_auth_azure::{AzureAuthInputs, AzureAuthService};
pub(crate) use cohere::AzureCohereAdapter;
pub(crate) use document_intelligence::AzureDocumentIntelligenceAdapter;
pub(crate) use mistral::AzureMistralAdapter;
pub(super) use mistral::validate_environment as validate_ai_environment;
use crate::ocr::types::OcrConnection;
async fn resolve_entra(
pub(super) async fn resolve_entra(
config: &AzureAuthInputs,
env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
) -> Result<Option<Sourced<String>>, Error> {
) -> Result<Option<Sourced<String>>, crate::ocr::Error> {
static SERVICE: OnceLock<AzureAuthService> = OnceLock::new();
SERVICE
.get_or_init(AzureAuthService::default)
@ -36,18 +25,18 @@ async fn resolve_entra(
Sourced::new(value, source)
})
})
.map_err(Error::from)
.map_err(crate::ocr::Error::from)
}
fn validate_destination(
pub(super) fn validate_destination(
connection: &OcrConnection,
credential_source: InputSource,
) -> Result<(), OcrError> {
) -> Result<(), crate::ocr::Error> {
if connection.api_base.is_some()
&& connection.api_base_source == InputSource::Request
&& credential_source != InputSource::Request
{
return Err(Error::from(litellm_auth::Error::RequestAzureCredentialDestination).into());
return Err(litellm_auth::Error::RequestAzureCredentialDestination.into());
}
Ok(())
}

1
litellm-rust/crates/core/src/llms/azure_ai/ocr/document_intelligence/mod.rs Normal file
View file

@ -0,0 +1 @@
pub(crate) mod transformation;

1216
litellm-rust/crates/core/src/llms/azure_ai/ocr/document_intelligence/transformation.rs Normal file
View file

File diff suppressed because it is too large Load diff

4
litellm-rust/crates/core/src/llms/azure_ai/ocr/mod.rs Normal file
View file

@ -0,0 +1,4 @@
pub(crate) mod cohere_parse_transformation;
pub(crate) mod common_utils;
pub(crate) mod document_intelligence;
pub(crate) mod transformation;

615
litellm-rust/crates/core/src/llms/azure_ai/ocr/transformation.rs Normal file
View file

@ -0,0 +1,615 @@
use litellm_auth::{InputSource, Sourced};
use litellm_auth_azure::AzureAuthInputs;
use serde_json::Value;
use crate::call_arguments::CallArguments;
use crate::constants::AZURE_AI_OCR_PATH;
use crate::llms::base_llm::ocr::transformation::{BaseOcrConfig, OcrRequestContext};
use crate::llms::mistral::ocr::transformation::{MistralOcrConfig, MistralOcrRequest};
use crate::ocr::OcrClient;
use crate::ocr::document::{inline_remote_document, validate_inline_document};
use crate::ocr::prepare::credential_env;
use crate::ocr::types::{LiteLLMOcrResponse, OcrConnection, OcrDocument, PreparedOcrRequest};
use crate::params::OpaqueParams;
use crate::url_utils::ApiUrl;
const AZURE_AI_API_KEY_ENV: &str = "AZURE_AI_API_KEY";
const AZURE_AI_API_BASE_ENV: &str = "AZURE_AI_API_BASE";
#[derive(Clone, Debug, Default)]
pub(crate) struct AzureAiOcrConfig;
impl BaseOcrConfig for AzureAiOcrConfig {
type OcrParams = OpaqueParams;
type ProviderRequest = MistralOcrRequest;
type Environment = Vec<(String, String)>;
fn get_supported_ocr_params(&self, model: &str) -> &'static [&'static str] {
MistralOcrConfig.get_supported_ocr_params(model)
}
fn get_api_key_env_var(&self) -> Option<&'static str> {
Some(AZURE_AI_API_KEY_ENV)
}
fn map_ocr_params(
&self,
non_default_params: &CallArguments,
model: &str,
) -> Result<OpaqueParams, crate::ocr::Error> {
MistralOcrConfig.map_ocr_params(non_default_params, model)
}
async fn validate_environment(
&self,
request: &PreparedOcrRequest,
_client: &OcrClient,
) -> Result<Self::Environment, crate::ocr::Error> {
let config = AzureAuthInputs {
azure_ad_token_provider: request.azure_ad_token_provider.clone(),
..AzureAuthInputs::from_sourced_optional_params(
&request.optional_params,
&request.input_sources,
)?
};
self.resolve_headers(&request.connection, &config, &credential_env)
.await
}
fn get_complete_url(
&self,
request: &PreparedOcrRequest,
_optional_params: &Self::OcrParams,
_environment: &Self::Environment,
) -> Result<String, crate::ocr::Error> {
self.build_ocr_url(request.connection.api_base.as_deref(), &credential_env)
}
fn transform_ocr_request(
&self,
model: &str,
document: OcrDocument,
optional_params: &OpaqueParams,
headers: &[(String, String)],
) -> Result<MistralOcrRequest, crate::ocr::Error> {
MistralOcrConfig.transform_ocr_request(model, document, optional_params, headers)
}
async fn async_transform_ocr_request(
&self,
model: &str,
document: OcrDocument,
optional_params: &OpaqueParams,
headers: &[(String, String)],
context: OcrRequestContext<'_>,
) -> Result<MistralOcrRequest, crate::ocr::Error> {
let document = inline_remote_document(
context.client.document_fetcher(),
document,
context.connection,
)
.await?;
self.transform_ocr_request(model, document, optional_params, headers)
}
fn transform_ocr_response(
&self,
model: &str,
raw_response: &[u8],
request_format: crate::ocr::types::OcrResponseFormat,
) -> Result<LiteLLMOcrResponse, crate::ocr::Error> {
MistralOcrConfig.transform_ocr_response(model, raw_response, request_format)
}
fn validate_request_body(&self, body: &Value) -> Result<(), crate::ocr::Error> {
validate_inline_document(&crate::ocr::prepare::body_document(body)?)
}
}
impl AzureAiOcrConfig {
/// Python `AzureAIOCRConfig.validate_environment` requires the endpoint
/// before it resolves credentials; keep that order so a missing base is
/// reported without invoking any token provider.
pub(super) fn resolve_api_base(
api_base: Option<&str>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<String, crate::ocr::Error> {
nonblank(api_base.map(str::to_string))
.or_else(|| nonblank(env_lookup(AZURE_AI_API_BASE_ENV)))
.ok_or(crate::ocr::Error::Auth(
litellm_auth::Error::MissingApiBase {
provider: "Azure AI",
environment_variable: AZURE_AI_API_BASE_ENV,
},
))
}
async fn resolve_headers(
&self,
connection: &OcrConnection,
config: &AzureAuthInputs,
env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
) -> Result<Vec<(String, String)>, crate::ocr::Error> {
Self::resolve_api_base(connection.api_base.as_deref(), env_lookup)?;
if crate::http_utils::has_header(&connection.extra_headers, "authorization") {
if config.azure_ad_token_provider.is_some() {
super::common_utils::resolve_entra(config, env_lookup).await?;
}
super::common_utils::validate_destination(connection, connection.extra_headers_source)?;
return Ok(connection.extra_headers.clone());
}
let key = nonblank(connection.api_key.clone())
.map(|value| Sourced::new(value, connection.api_key_source))
.or_else(|| {
nonblank(self.get_api_key_env_var().and_then(env_lookup))
.map(|value| Sourced::new(value, InputSource::Environment))
});
if let Some(key) = key {
super::common_utils::validate_destination(connection, key.source())?;
return Ok(bearer_headers(connection, key.value()));
}
let key = super::common_utils::resolve_entra(config, env_lookup)
.await?
.ok_or(crate::ocr::Error::MissingAzureAiCredentials)?;
super::common_utils::validate_destination(connection, key.source())?;
Ok(bearer_headers(connection, key.value()))
}
fn build_ocr_url(
&self,
api_base: Option<&str>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<String, crate::ocr::Error> {
let base = Self::resolve_api_base(api_base, env_lookup)?;
let path: Vec<&str> = AZURE_AI_OCR_PATH.trim_matches('/').split('/').collect();
ApiUrl::parse(&base)
.and_then(|url| url.complete_path(&path))
.map(|url| url.into_string())
.map_err(|_| crate::ocr::Error::RequestField {
path: "api_base".into(),
})
}
}
fn bearer_headers(connection: &OcrConnection, key: &str) -> Vec<(String, String)> {
std::iter::once(("Authorization".into(), format!("Bearer {key}")))
.chain(connection.extra_headers.clone())
.collect()
}
fn nonblank(value: Option<String>) -> Option<String> {
value
.map(|value| value.trim().to_string())
.filter(|value| !value.is_empty())
}
#[cfg(test)]
mod tests {
use rstest::{fixture, rstest};
use super::*;
#[fixture]
fn connection() -> OcrConnection {
OcrConnection {
api_key: Some("request-key".into()),
api_base: Some("https://example.com".into()),
..Default::default()
}
}
#[rstest]
#[case::base_with_query(
"https://example.com/?tenant=a",
"https://example.com/providers/mistral/azure/ocr?tenant=a"
)]
#[case::complete_endpoint(
"https://example.com/providers/mistral/azure/ocr",
"https://example.com/providers/mistral/azure/ocr"
)]
fn completes_azure_path_and_preserves_query(#[case] api_base: &str, #[case] expected: &str) {
assert_eq!(
AzureAiOcrConfig
.build_ocr_url(Some(api_base), &|_| None)
.unwrap(),
expected
);
}
#[test]
fn missing_api_base_is_structured() {
assert!(matches!(
AzureAiOcrConfig::resolve_api_base(None, &|_| None),
Err(crate::ocr::Error::Auth(
litellm_auth::Error::MissingApiBase {
provider: "Azure AI",
environment_variable: AZURE_AI_API_BASE_ENV,
}
))
));
}
#[rstest]
#[tokio::test]
async fn supplied_authorization_precedes_keys(connection: OcrConnection) {
let connection = OcrConnection {
extra_headers: vec![("authorization".into(), "Bearer prepared".into())],
..connection
};
assert_eq!(
AzureAiOcrConfig
.resolve_headers(&connection, &Default::default(), &|_| {
Some("environment-key".into())
})
.await
.unwrap(),
connection.extra_headers
);
}
#[rstest]
#[tokio::test]
async fn request_key_precedes_environment_key(connection: OcrConnection) {
assert_eq!(
AzureAiOcrConfig
.resolve_headers(&connection, &Default::default(), &|_| {
Some("environment-key".into())
})
.await
.unwrap()[0],
("Authorization".into(), "Bearer request-key".into())
);
}
#[tokio::test]
async fn request_endpoint_cannot_receive_environment_key() {
let connection = OcrConnection {
api_base: Some("https://request.example".into()),
api_base_source: InputSource::Request,
..Default::default()
};
let error = AzureAiOcrConfig
.resolve_headers(&connection, &Default::default(), &|name| {
(name == AZURE_AI_API_KEY_ENV).then(|| "environment-key".into())
})
.await
.unwrap_err();
assert!(
error
.to_string()
.contains("request-controlled Azure endpoint")
);
}
#[tokio::test]
async fn request_endpoint_accepts_request_owned_key() {
let connection = OcrConnection {
api_key: Some("request-key".into()),
api_key_source: InputSource::Request,
api_base: Some("https://request.example".into()),
api_base_source: InputSource::Request,
..Default::default()
};
let headers = AzureAiOcrConfig
.resolve_headers(&connection, &Default::default(), &|_| None)
.await
.unwrap();
assert_eq!(
headers[0],
("Authorization".into(), "Bearer request-key".into())
);
}
use serde_json::json;
use crate::ocr::LocalOcrHost;
use crate::ocr::test_support::{
MockResponse, mock_server, perform_ocr, perform_ocr_with, wire_request,
};
#[tokio::test]
async fn facade_executes_azure_mistral_with_prepared_auth() {
let (base, seen, server) = mock_server(vec![MockResponse::json(json!({
"pages":[{"index":0,"markdown":"hello"}],
"usage_info":{"pages_processed":1}
}))])
.await;
let mut request = wire_request(
"azure_ai/model",
&base,
json!({"include_image_base64":true}),
);
request.credentials.api_key = None;
request.transport.extra_headers = vec![(
"Authorization".into(),
"Bearer python-prepared-token".into(),
)];
let result = perform_ocr(request).await.unwrap();
server.await.unwrap();
assert_eq!(result.pages[0].markdown, "hello");
let requests = seen.lock().unwrap();
assert_eq!(requests.len(), 1);
assert!(requests[0].starts_with("POST /providers/mistral/azure/ocr "));
assert!(
requests[0]
.to_ascii_lowercase()
.contains("authorization: bearer python-prepared-token\r\n")
);
let body: Value =
serde_json::from_str(requests[0].split_once("\r\n\r\n").unwrap().1).unwrap();
assert_eq!(
body,
json!({
"model":"model",
"document":{"type":"document_url","document_url":"data:application/pdf;base64,YWJj"},
"include_image_base64":true
})
);
}
#[tokio::test]
async fn facade_acquires_supplied_entra_token_for_final_request() {
let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await;
let mut request = wire_request(
"azure_ai/model",
&base,
json!({"azure_ad_token":"rust-owned-token"}),
);
request.credentials.api_key = None;
perform_ocr(request).await.unwrap();
server.await.unwrap();
let requests = seen.lock().unwrap();
assert_eq!(requests.len(), 1);
assert!(
requests[0]
.to_ascii_lowercase()
.contains("authorization: bearer rust-owned-token\r\n")
);
}
#[tokio::test]
async fn rejects_non_inline_body_after_guardrails() {
let request = wire_request("azure_ai/model", "http://127.0.0.1:1", json!({}));
let host = LocalOcrHost::new(request).with_before_send(|mut wire, _| {
wire.body["document"] = json!({
"type":"document_url",
"document_url":"https://example.com/not-inline.pdf"
});
Ok(wire)
});
let error = perform_ocr_with(host).await.unwrap_err();
assert!(error.to_string().contains("data URI"));
}
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
use litellm_auth::{
ResolvedCredential, SecretValue, TokenFuture, TokenProvider, TokenProviderHandle,
};
use crate::ocr::LiteLLMOcrRequest;
use crate::ocr::test_support::header;
use crate::ocr::wire::decode_request;
#[derive(Debug)]
struct CountingToken {
token: fn(usize) -> String,
calls: AtomicUsize,
}
impl CountingToken {
fn new(token: fn(usize) -> String) -> Arc<Self> {
Arc::new(Self {
token,
calls: AtomicUsize::new(0),
})
}
fn calls(&self) -> usize {
self.calls.load(Ordering::SeqCst)
}
}
impl TokenProvider for CountingToken {
fn acquire(&self) -> TokenFuture<'_> {
let call = self.calls.fetch_add(1, Ordering::SeqCst) + 1;
let token = SecretValue::new((self.token)(call));
Box::pin(async move {
Ok(ResolvedCredential::AccessToken {
token,
expires_on: None,
})
})
}
}
fn numbered_token(call: usize) -> String {
format!("callback-{call}")
}
fn azure_request(
provider: &Arc<CountingToken>,
api_base: Option<&str>,
api_key: Option<&str>,
extra_headers: Value,
optional_params: Value,
) -> LiteLLMOcrRequest {
let wire = serde_json::from_value(json!({
"model": "azure_ai/mistral-ocr-latest",
"document": {"type":"document_url","document_url":"data:application/pdf;base64,YWJj"},
"api_key": api_key,
"api_base": api_base,
"custom_llm_provider": null,
"extra_headers": extra_headers,
"optional_params": optional_params,
"timeout_seconds": 2.0
}))
.unwrap();
LiteLLMOcrRequest {
azure_ad_token_provider: Some(TokenProviderHandle::new(provider.clone())),
..decode_request(wire).unwrap()
}
}
fn ocr_page() -> MockResponse {
MockResponse::json(json!({"pages":[{"index":0,"markdown":"hello"}]}))
}
#[tokio::test]
async fn token_provider_result_is_the_bearer_and_is_acquired_for_each_request() {
let provider = CountingToken::new(numbered_token);
let (base, seen, server) = mock_server(vec![ocr_page(), ocr_page()]).await;
for _ in 0..2 {
perform_ocr(azure_request(
&provider,
Some(&base),
None,
Value::Null,
json!({}),
))
.await
.unwrap();
}
server.await.unwrap();
assert_eq!(provider.calls(), 2);
let requests = seen.lock().unwrap();
assert_eq!(
requests
.iter()
.map(|request| header(request, "authorization"))
.collect::<Vec<_>>(),
[Some("Bearer callback-1"), Some("Bearer callback-2")]
);
}
#[rstest]
#[case::api_key_skips_provider(Some("resource-key"), Value::Null, json!({}), "Bearer resource-key", 0)]
#[case::provider_beats_static_token(
None,
Value::Null,
json!({"azure_ad_token":"static-token"}),
"Bearer callback-1",
1
)]
#[case::header_wins_on_the_wire_but_provider_still_runs(
None,
json!({"Authorization":"Bearer override"}),
json!({}),
"Bearer override",
1
)]
#[tokio::test]
async fn credential_precedence(
#[case] api_key: Option<&str>,
#[case] extra_headers: Value,
#[case] optional_params: Value,
#[case] expected_authorization: &str,
#[case] expected_calls: usize,
) {
let provider = CountingToken::new(numbered_token);
let (base, seen, server) = mock_server(vec![ocr_page()]).await;
perform_ocr(azure_request(
&provider,
Some(&base),
api_key,
extra_headers,
optional_params,
))
.await
.unwrap();
server.await.unwrap();
assert_eq!(provider.calls(), expected_calls);
let requests = seen.lock().unwrap();
assert_eq!(requests.len(), 1);
assert_eq!(
header(&requests[0], "authorization"),
Some(expected_authorization)
);
}
#[rstest]
#[case::missing_api_base(
false,
json!({}),
numbered_token,
|error: &crate::ocr::Error| matches!(error, crate::ocr::Error::Auth(litellm_auth::Error::MissingApiBase {
provider: "Azure AI",
environment_variable: AZURE_AI_API_BASE_ENV,
})),
0
)]
#[case::unsupported_oidc_reference(
true,
json!({"azure_ad_token":"oidc/assertion","client_id":"client","tenant_id":"tenant"}),
numbered_token,
|error: &crate::ocr::Error| matches!(error, crate::ocr::Error::Auth(litellm_auth::Error::UnsupportedOidcReference)),
0
)]
#[case::empty_provider_token_ignores_static_token(
true,
json!({"azure_ad_token":"static-token"}),
|_| String::new(),
|error: &crate::ocr::Error| matches!(error, crate::ocr::Error::MissingAzureAiCredentials),
1
)]
#[tokio::test]
async fn credential_failures_send_no_provider_request(
#[case] with_api_base: bool,
#[case] optional_params: Value,
#[case] token: fn(usize) -> String,
#[case] expected: fn(&crate::ocr::Error) -> bool,
#[case] expected_calls: usize,
) {
let provider = CountingToken::new(token);
let (base, seen, server) = mock_server(vec![ocr_page()]).await;
let error = perform_ocr(azure_request(
&provider,
with_api_base.then_some(base.as_str()),
None,
Value::Null,
optional_params,
))
.await
.unwrap_err();
server.abort();
assert!(expected(&error), "unexpected error: {error:?}");
assert_eq!(provider.calls(), expected_calls);
assert!(seen.lock().unwrap().is_empty());
}
#[tokio::test]
async fn environment_supplies_api_base_and_bearer_key() {
let env = |name: &str| match name {
AZURE_AI_API_BASE_ENV => Some("https://env.example".to_string()),
AZURE_AI_API_KEY_ENV => Some("env-key".to_string()),
_ => None,
};
let connection = OcrConnection::default();
let headers = AzureAiOcrConfig
.resolve_headers(&connection, &Default::default(), &env)
.await
.unwrap();
let url = AzureAiOcrConfig.build_ocr_url(None, &env).unwrap();
assert_eq!(
headers,
[("Authorization".to_string(), "Bearer env-key".to_string())]
);
assert_eq!(url, "https://env.example/providers/mistral/azure/ocr");
}
}

1
litellm-rust/crates/core/src/llms/base_llm/mod.rs Normal file
View file

@ -0,0 +1 @@
pub(crate) mod ocr;

1
litellm-rust/crates/core/src/llms/base_llm/ocr/mod.rs Normal file
View file

@ -0,0 +1 @@
pub(crate) mod transformation;

213
litellm-rust/crates/core/src/llms/base_llm/ocr/transformation.rs Normal file
View file

@ -0,0 +1,213 @@
use std::future::Future;
use serde::{Serialize, de::DeserializeOwned};
use serde_json::Value;
use crate::{
call_arguments::CallArguments,
ocr::{
OcrClient,
route::OcrHost,
types::{
LiteLLMOcrResponse, OcrConnection, OcrCredentialInputs, OcrDocument, OcrResponseFormat,
PreparedOcrRequest, ResolvedOcrCredentials,
},
},
};
const HEALTH_CHECK_PDF_DATA_URI: &str = "data:application/pdf;base64,JVBERi0xLjQKJeLjz9MKMyAwIG9iago8PC9UeXBlIC9QYWdlCi9QYXJlbnQgMSAwIFIKL01lZGlhQm94IFswIDAgNjEyIDc5Ml0KL0NvbnRlbnRzIDQgMCBSCi9SZXNvdXJjZXMgPDwvRm9udCA8PC9GMSAyIDAgUj4+Pj4+PgplbmRvYmoKNCAwIG9iago8PC9MZW5ndGggNDQ+PgpzdHJlYW0KQlQKL0YxIDI0IFRmCjEwMCA3MDAgVGQKKHRlc3QpIFRqCkVUCmVuZHN0cmVhbQplbmRvYmoKMiAwIG9iago8PC9UeXBlIC9Gb250Ci9TdWJ0eXBlIC9UeXBlMQovQmFzZUZvbnQgL0hlbHZldGljYT4+CmVuZG9iagoxIDAgb2JqCjw8L1R5cGUgL1BhZ2VzCi9LaWRzIFszIDAgUl0KL0NvdW50IDE+PgplbmRvYmoKNSAwIG9iago8PC9UeXBlIC9DYXRhbG9nCi9QYWdlcyAxIDAgUj4+CmVuZG9iagp0cmFpbGVyCjw8L1NpemUgNgovUm9vdCA1IDAgUj4+CnN0YXJ0eHJlZgozMjQKJSVFT0Y=";
/// Output of `validate_environment`: whatever a provider resolves up front
/// (headers at minimum; Vertex also carries the project id).
pub(crate) trait OcrEnvironment: Send + Sync {
fn headers(&self) -> &[(String, String)];
}
impl OcrEnvironment for Vec<(String, String)> {
fn headers(&self) -> &[(String, String)] {
self
}
}
#[derive(Clone, Copy)]
pub(crate) struct OcrRequestContext<'a> {
pub client: &'a OcrClient,
pub connection: &'a OcrConnection,
}
#[derive(Clone, Copy)]
pub(crate) struct OcrResponseContext<'a> {
pub client: &'a OcrClient,
pub connection: &'a OcrConnection,
pub host: &'a OcrHost,
pub request_format: OcrResponseFormat,
pub url: &'a str,
pub headers: &'a [(String, String)],
}
pub(crate) trait BaseOcrConfig: Send + Sync + Sized + 'static {
type OcrParams: Send + Sync;
type ProviderRequest: Serialize + Send;
type Environment: OcrEnvironment;
fn get_supported_ocr_params(&self, _model: &str) -> &'static [&'static str] {
&[]
}
fn get_api_key_env_var(&self) -> Option<&'static str> {
None
}
fn resolve_connection_params(&self, inputs: OcrCredentialInputs) -> ResolvedOcrCredentials {
ResolvedOcrCredentials {
api_key: inputs
.dynamic_api_key
.filter(|value| !value.value().is_empty())
.or(inputs.api_key),
api_base: inputs
.dynamic_api_base
.filter(|value| !value.value().is_empty())
.or(inputs.api_base),
}
}
fn get_health_check_document(&self) -> OcrDocument {
OcrDocument::DocumentUrl {
document_url: HEALTH_CHECK_PDF_DATA_URI.into(),
extra_fields: Default::default(),
}
}
fn map_ocr_params(
&self,
non_default_params: &CallArguments,
model: &str,
) -> Result<Self::OcrParams, crate::ocr::Error>;
fn validate_environment(
&self,
request: &PreparedOcrRequest,
client: &OcrClient,
) -> impl Future<Output = Result<Self::Environment, crate::ocr::Error>> + Send;
fn get_complete_url(
&self,
request: &PreparedOcrRequest,
optional_params: &Self::OcrParams,
environment: &Self::Environment,
) -> Result<String, crate::ocr::Error>;
fn transform_ocr_request(
&self,
model: &str,
document: OcrDocument,
optional_params: &Self::OcrParams,
headers: &[(String, String)],
) -> Result<Self::ProviderRequest, crate::ocr::Error>;
fn async_transform_ocr_request(
&self,
model: &str,
document: OcrDocument,
optional_params: &Self::OcrParams,
headers: &[(String, String)],
_context: OcrRequestContext<'_>,
) -> impl Future<Output = Result<Self::ProviderRequest, crate::ocr::Error>> + Send {
async move { self.transform_ocr_request(model, document, optional_params, headers) }
}
fn transform_ocr_response(
&self,
model: &str,
raw_response: &[u8],
request_format: OcrResponseFormat,
) -> Result<LiteLLMOcrResponse, crate::ocr::Error>;
fn async_transform_ocr_response(
&self,
model: &str,
raw_response: reqwest::Response,
context: OcrResponseContext<'_>,
) -> impl Future<Output = Result<LiteLLMOcrResponse, crate::ocr::Error>> + Send {
async move {
let bytes = crate::ocr::client::read_response_bytes(
raw_response,
context.connection.max_response_bytes,
)
.await?;
crate::ocr::handler::emit_response_received(context.host, &bytes).await?;
self.transform_ocr_response(model, &bytes, context.request_format)
}
}
fn get_error_class(
&self,
error_message: String,
status_code: u16,
headers: Vec<(String, String)>,
) -> crate::ocr::Error {
crate::ocr::Error::Provider {
status: status_code,
body: error_message,
headers,
}
}
/// Provider-specific check applied to the composed body, both before and
/// after guardrail hooks. Defaults to accepting any body.
fn validate_request_body(&self, _body: &Value) -> Result<(), crate::ocr::Error> {
Ok(())
}
/// Rust counterpart of `BaseLLMHTTPHandler._async_prepare_ocr_request`:
/// map params, validate environment, build URL, transform, compose body.
fn prepare_request(
&self,
request: &PreparedOcrRequest,
client: &OcrClient,
) -> impl Future<Output = Result<reqwest::Request, crate::ocr::Error>> + Send {
async move {
let params = self.map_ocr_params(&request.optional_params, &request.model)?;
let environment = self.validate_environment(request, client).await?;
let url = self.get_complete_url(request, &params, &environment)?;
let headers = environment.headers();
let body = self
.async_transform_ocr_request(
&request.model,
request.document.clone(),
&params,
headers,
OcrRequestContext {
client,
connection: &request.connection,
},
)
.await?;
crate::ocr::prepare::transform_request_body(
client,
request,
&url,
headers,
body,
|body| self.validate_request_body(body),
)
.await
}
}
}
pub(crate) fn decode_and_normalize_response<T: DeserializeOwned>(
model: &str,
raw_response: &[u8],
request_format: OcrResponseFormat,
normalize: impl FnOnce(&str, T) -> Result<LiteLLMOcrResponse, crate::ocr::Error>,
) -> Result<LiteLLMOcrResponse, crate::ocr::Error> {
let decoded = crate::ocr::json::decode_response(
raw_response,
request_format == OcrResponseFormat::Native,
)?;
Ok(LiteLLMOcrResponse {
provider_native_response: decoded.native,
..normalize(model, decoded.data)?
})
}

1
litellm-rust/crates/core/src/llms/cohere/mod.rs Normal file
View file

@ -0,0 +1 @@
pub(crate) mod ocr;

3
litellm-rust/crates/core/src/llms/cohere/ocr/mod.rs Normal file
View file

@ -0,0 +1,3 @@
pub(crate) mod transformation;
pub(crate) use transformation::{CohereOptions, validate_document};

870
litellm-rust/crates/core/src/llms/cohere/ocr/transformation.rs Normal file
View file

@ -0,0 +1,870 @@
use serde::{Deserialize, Serialize};
use serde_json::{Map, Value};
use serde_with::serde_as;
use crate::{
call_arguments::{CallArguments, parse_options},
constants::{COHERE_API_KEY_ENV, COHERE_PARSE_API_BASE},
llms::base_llm::ocr::transformation::{BaseOcrConfig, decode_and_normalize_response},
ocr::{
OcrClient,
document::InlineDocument,
prepare::credential_env,
types::{
LiteLLMOcrResponse, OcrConnection, OcrDocument, OcrPage, OcrPageImage,
OcrResponseFormat, OcrUsageInfo, PreparedOcrRequest,
},
},
serde_compat::LaxI64,
url_utils::ApiUrl,
};
const COHERE_PARSE_HEALTH_CHECK_IMAGE_DATA_URI: &str = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAIAAACQd1PeAAAADElEQVR4nGP4//8/AAX+Av4N70a4AAAAAElFTkSuQmCC";
#[derive(Clone, Copy, Debug, Default, Deserialize, Serialize)]
#[serde(rename_all = "lowercase")]
pub(crate) enum OutputFormat {
#[default]
Markdown,
Blocks,
}
#[derive(Default, Deserialize, Serialize)]
pub(crate) struct CohereOptions {
#[serde(skip_serializing_if = "Option::is_none")]
pub output_format: Option<OutputFormat>,
}
#[derive(Deserialize, Serialize)]
pub(crate) struct CohereRequest {
pub model: String,
pub document: CohereParseDocument,
pub output_format: String,
}
#[derive(Deserialize, Serialize)]
#[serde(tag = "type")]
pub(crate) enum CohereParseDocument {
#[serde(rename = "image_url")]
ImageUrl { image_url: String },
}
#[derive(Deserialize)]
pub(crate) struct CohereResponse {
#[serde(default)]
pages: Vec<CoherePage>,
meta: Option<CohereMeta>,
}
#[serde_as]
#[derive(Deserialize)]
struct CoherePage {
#[serde_as(deserialize_as = "Option<LaxI64>")]
index: Option<i64>,
markdown: Option<CohereMarkdown>,
blocks: Option<Vec<Map<String, Value>>>,
}
#[derive(Deserialize, Serialize)]
struct CohereMarkdown {
#[serde(default)]
content: String,
images: Option<Vec<Map<String, Value>>>,
}
#[derive(Deserialize)]
struct CohereMeta {
billed_units: Option<CohereBilledUnits>,
}
#[serde_as]
#[derive(Deserialize)]
struct CohereBilledUnits {
#[serde_as(deserialize_as = "Option<LaxI64>")]
pages: Option<i64>,
}
#[derive(Default)]
pub(crate) struct CohereParseConfig;
impl BaseOcrConfig for CohereParseConfig {
type OcrParams = CohereOptions;
type ProviderRequest = CohereRequest;
type Environment = Vec<(String, String)>;
fn get_supported_ocr_params(&self, _model: &str) -> &'static [&'static str] {
&["output_format", "req_format"]
}
fn get_api_key_env_var(&self) -> Option<&'static str> {
Some(COHERE_API_KEY_ENV)
}
fn get_health_check_document(&self) -> OcrDocument {
OcrDocument::ImageUrl {
image_url: COHERE_PARSE_HEALTH_CHECK_IMAGE_DATA_URI.into(),
extra_fields: Default::default(),
}
}
fn map_ocr_params(
&self,
non_default_params: &CallArguments,
_model: &str,
) -> Result<CohereOptions, crate::ocr::Error> {
Ok(parse_options(non_default_params)?)
}
async fn validate_environment(
&self,
request: &PreparedOcrRequest,
_client: &OcrClient,
) -> Result<Self::Environment, crate::ocr::Error> {
self.resolve_headers(&request.connection, &credential_env)
}
fn get_complete_url(
&self,
request: &PreparedOcrRequest,
_optional_params: &Self::OcrParams,
_environment: &Self::Environment,
) -> Result<String, crate::ocr::Error> {
self.build_ocr_url(
request
.connection
.api_base
.as_deref()
.unwrap_or(COHERE_PARSE_API_BASE),
)
}
fn transform_ocr_request(
&self,
model: &str,
document: OcrDocument,
optional_params: &CohereOptions,
_headers: &[(String, String)],
) -> Result<CohereRequest, crate::ocr::Error> {
let image_url = image_url(document)?;
Ok(build_request(model, image_url, optional_params))
}
fn transform_ocr_response(
&self,
model: &str,
raw_response: &[u8],
request_format: OcrResponseFormat,
) -> Result<LiteLLMOcrResponse, crate::ocr::Error> {
decode_and_normalize_response(model, raw_response, request_format, normalize_response)
}
fn validate_request_body(&self, body: &Value) -> Result<(), crate::ocr::Error> {
validate_document(&crate::ocr::prepare::body_document(body)?)
}
}
impl CohereParseConfig {
fn resolve_headers(
&self,
connection: &OcrConnection,
env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
) -> Result<Vec<(String, String)>, crate::ocr::Error> {
if crate::http_utils::has_header(&connection.extra_headers, "authorization") {
return Ok(connection.extra_headers.clone());
}
let key = connection
.api_key
.as_deref()
.map(str::trim)
.filter(|key| !key.is_empty())
.map(str::to_string)
.or_else(|| {
self.get_api_key_env_var()
.and_then(env_lookup)
.filter(|key| !key.trim().is_empty())
})
.ok_or_else(|| {
crate::ocr::Error::Auth(litellm_auth::Error::ProviderAuthentication(
"Missing COHERE_API_KEY - set it in the environment or pass api_key".into(),
))
})?;
Ok(
std::iter::once(("Authorization".into(), format!("Bearer {key}")))
.chain(connection.extra_headers.clone())
.collect(),
)
}
fn build_ocr_url(&self, api_base: &str) -> Result<String, crate::ocr::Error> {
let parsed = reqwest::Url::parse(api_base).map_err(|_| invalid_api_base())?;
if !matches!(parsed.scheme(), "http" | "https") {
return Err(invalid_api_base());
}
ApiUrl::parse(api_base)
.and_then(|url| url.complete_path(&["v2", "parse"]))
.map(|url| url.into_string())
.map_err(|_| invalid_api_base())
}
}
pub(crate) fn validate_document(document: &OcrDocument) -> Result<(), crate::ocr::Error> {
let OcrDocument::ImageUrl { image_url, .. } = document else {
return Err(crate::ocr::Error::CohereImageOnly);
};
if image_url.is_empty() {
return Err(crate::ocr::Error::CohereImageOnly);
}
if let Some(inline) = InlineDocument::parse(image_url)? {
if !inline.mime_type().type_.eq_ignore_ascii_case("image") {
return Err(crate::ocr::Error::CohereImageOnly);
}
inline.decode(crate::constants::OCR_INLINE_MAX_BYTES)?;
}
Ok(())
}
pub(crate) fn normalize_response(
model: &str,
response: CohereResponse,
) -> Result<LiteLLMOcrResponse, crate::ocr::Error> {
let pages_processed = billed_pages(&response).map(Ok).unwrap_or_else(|| {
i64::try_from(response.pages.len()).map_err(|_| crate::ocr::Error::NumericRange("pages"))
})?;
let pages = response
.pages
.into_iter()
.enumerate()
.map(|(position, page)| normalize_page(page, position))
.collect::<Result<Vec<_>, crate::ocr::Error>>()?;
Ok(LiteLLMOcrResponse {
usage_info: Some(OcrUsageInfo {
pages_processed: Some(pages_processed),
..Default::default()
}),
..LiteLLMOcrResponse::new(model, pages)
})
}
fn image_url(document: OcrDocument) -> Result<String, crate::ocr::Error> {
validate_document(&document)?;
let OcrDocument::ImageUrl { image_url, .. } = document else {
return Err(crate::ocr::Error::CohereImageOnly);
};
Ok(image_url)
}
fn build_request(model: &str, image_url: String, params: &CohereOptions) -> CohereRequest {
CohereRequest {
model: model.into(),
document: CohereParseDocument::ImageUrl { image_url },
output_format: match params.output_format.unwrap_or_default() {
OutputFormat::Markdown => "markdown",
OutputFormat::Blocks => "blocks",
}
.into(),
}
}
fn page_image(
mut image: Map<String, Value>,
path: &str,
) -> Result<OcrPageImage, crate::ocr::Error> {
if let Some(Value::Object(bbox)) = image.get("bounding_box") {
image.insert("bbox".into(), Value::Object(bbox.clone()));
}
crate::ocr::json::decode_response_value(Value::Object(image), path)
}
fn normalize_page(page: CoherePage, position: usize) -> Result<OcrPage, crate::ocr::Error> {
let index = page.index.map(Ok).unwrap_or_else(|| {
i64::try_from(position).map_err(|_| crate::ocr::Error::NumericRange("page index"))
})?;
let (markdown, images) = match page.markdown {
Some(markdown) => {
let images = markdown
.images
.filter(|images| !images.is_empty())
.map(|images| {
images
.into_iter()
.enumerate()
.map(|(image_index, image)| {
page_image(
image,
&format!("pages[{position}].markdown.images[{image_index}]"),
)
})
.collect::<Result<Vec<_>, _>>()
})
.transpose()?;
(markdown.content, images)
}
None => (String::new(), None),
};
let extra_fields = page
.blocks
.map(|blocks| {
(
"blocks".into(),
Value::Array(blocks.into_iter().map(Value::Object).collect()),
)
})
.into_iter()
.collect();
Ok(OcrPage {
index,
markdown,
images,
extra_fields,
..Default::default()
})
}
fn billed_pages(response: &CohereResponse) -> Option<i64> {
response.meta.as_ref()?.billed_units.as_ref()?.pages
}
fn invalid_api_base() -> crate::ocr::Error {
crate::ocr::Error::RequestField {
path: "api_base".into(),
}
}
#[cfg(test)]
mod tests {
use rstest::rstest;
use serde_json::json;
use super::*;
#[tokio::test]
async fn composed_body_preserves_native_document_fields_and_untyped_overrides() {
let request = crate::ocr::test_support::wire_request(
"cohere/parse",
"https://example.com",
json!({
"output_format":"markdown", "timeout":30,
"extra_body":{
"output_format": {"future":true},
"document":{"type":"image_url","image_url":"https://example.com/a.png",
"provider_options":{"nested":[false,0,null]}}
}
}),
);
let request = request.with_document(
serde_json::from_value(json!({
"type":"image_url","image_url":"https://example.com/original.png"
}))
.unwrap(),
);
let request = crate::ocr::prepare::prepare_request_for_test(request);
let http = CohereParseConfig
.prepare_request(&request, &crate::ocr::test_support::ocr_client())
.await
.unwrap();
let body: Value = serde_json::from_slice(http.body().unwrap().as_bytes().unwrap()).unwrap();
assert_eq!(
body,
json!({
"model":"parse", "output_format":{"future":true},
"document":{"type":"image_url","image_url":"https://example.com/a.png",
"provider_options":{"nested":[false,0,null]}}
})
);
}
#[rstest]
#[case::cohere(false)]
#[case::azure(true)]
fn options_read_known_fields_without_changing_arguments(#[case] azure: bool) {
let arguments = serde_json::from_value(json!({
"output_format":"blocks", "req_format":"native", "extension":false
}))
.unwrap();
let mapped = if azure {
crate::llms::azure_ai::ocr::cohere_parse_transformation::AzureAICohereParseConfig
.map_ocr_params(&arguments, "parse")
} else {
CohereParseConfig.map_ocr_params(&arguments, "parse")
}
.unwrap();
assert_eq!(
serde_json::to_value(mapped).unwrap(),
json!({"output_format":"blocks"})
);
assert_eq!(arguments["req_format"], "native");
assert_eq!(arguments["extension"], false);
}
#[test]
fn options_reject_invalid_output_format() {
let invalid = serde_json::from_value(json!({"output_format":"html"})).unwrap();
assert!(matches!(
CohereParseConfig.map_ocr_params(&invalid, "parse"),
Err(crate::ocr::Error::RequestField { path })
if path == "optional_params.output_format"
));
}
#[test]
fn billed_pages_accept_integral_doubles() {
let response = serde_json::from_str::<CohereResponse>(
r#"{"pages":[],"meta":{"billed_units":{"pages":1.0}}}"#,
)
.unwrap();
let normalized = normalize_response("parse", response).unwrap();
assert_eq!(normalized.usage_info.unwrap().pages_processed, Some(1));
}
#[test]
fn billed_pages_reject_fractional_counts() {
assert!(
serde_json::from_str::<CohereResponse>(
r#"{"pages":[],"meta":{"billed_units":{"pages":1.5}}}"#,
)
.is_err()
);
}
#[test]
fn response_preserves_python_mapping_shapes_and_extensions() {
let blocks = json!([
{"type":"text", "text":"Total Due: $4.00"},
{"type":"future", "payload":{"nested":[null,false,0]}}
]);
let response = serde_json::from_value(json!({
"pages":[{
"index":"2",
"markdown":{"content":"receipt", "images":[
{"bounding_box":{"x":1}, "bbox":"replaced", "category":"future", "extension":null},
{"image_base64":"encoded"}
]},
"blocks":blocks
}],
"meta":{"billed_units":{"pages":0}}
})).unwrap();
let response = normalize_response("parse", response).unwrap();
assert_eq!(response.pages[0].index, 2);
assert_eq!(response.usage_info.unwrap().pages_processed, Some(0));
assert_eq!(response.pages[0].extra_fields["blocks"], blocks);
let images = response.pages[0].images.as_ref().unwrap();
assert_eq!(images[0].bbox.as_ref().unwrap()["x"], 1);
assert_eq!(images[0].extra_fields["category"], "future");
assert_eq!(images[0].extra_fields.get("extension"), Some(&Value::Null));
assert_eq!(images[1].image_base64.as_deref(), Some("encoded"));
assert!(images[1].bbox.is_none());
}
#[test]
fn malformed_normalized_image_fields_report_the_original_path() {
let response = serde_json::from_value(json!({
"pages":[{"markdown":{"images":[{"image_base64":42}]}}]
}))
.unwrap();
assert!(matches!(
normalize_response("parse", response).unwrap_err(),
crate::ocr::Error::ResponseField { path }
if path == "pages[0].markdown.images[0].image_base64"
));
}
#[rstest]
fn provider_options_exclude_response_controls_and_extensions(
#[values("markdown", "blocks")] output_format: &str,
#[values("https://example.com/a.png", "data:image/png;base64,YWJj")] source: &str,
) {
let arguments = serde_json::from_value(
json!({"output_format":output_format,"req_format":"native","unknown":true}),
)
.unwrap();
let params = CohereParseConfig
.map_ocr_params(&arguments, "parse")
.unwrap();
assert_eq!(
serde_json::to_value(&params).unwrap(),
json!({"output_format":output_format})
);
let document = serde_json::from_value(
json!({"type":"image_url","image_url":source,"ignored":"field"}),
)
.unwrap();
let body = CohereParseConfig
.transform_ocr_request("parse", document, &params, &[])
.unwrap();
assert_eq!(
serde_json::to_value(body).unwrap(),
json!({
"model":"parse", "document":{"type":"image_url","image_url":source}, "output_format":output_format
})
);
}
#[tokio::test]
async fn explicit_null_options_use_defaults_before_http() {
let request = crate::ocr::test_support::wire_request(
"cohere/parse",
"https://example.com",
json!({"output_format":null,"req_format":null}),
);
let request = request.with_document(
serde_json::from_value(
json!({"type":"image_url","image_url":"https://example.com/a.png"}),
)
.unwrap(),
);
assert_eq!(
request.response_format().unwrap(),
crate::ocr::types::OcrResponseFormat::Litellm
);
let request = crate::ocr::prepare::prepare_request_for_test(request);
let http = CohereParseConfig
.prepare_request(&request, &crate::ocr::test_support::ocr_client())
.await
.unwrap();
let body: Value = serde_json::from_slice(http.body().unwrap().as_bytes().unwrap()).unwrap();
assert_eq!(body["output_format"], "markdown");
assert!(body.get("req_format").is_none());
}
#[rstest]
fn response_normalizes_markdown_images_blocks_and_billed_pages() {
let payload = json!({
"pages": [
{
"type":"markdown",
"index":4,
"markdown":{
"content":"receipt",
"images":[{
"id":"image",
"bounding_box":{
"top_left_x":1,
"top_left_y":2,
"bottom_right_x":48,
"bottom_right_y":49
},
"bounding_box_normalized":{
"top_left_x":0.04,
"top_left_y":0.05,
"bottom_right_x":0.15,
"bottom_right_y":0.16
},
"description":"scan",
"category":"logo",
"provider_extension":"preserved"
}]
}
},
{"type":"blocks","blocks":[{"type":"text","text":{"content":"total"}}]}
],
"meta":{"api_version":{"version":"2"},"billed_units":{"pages":3}}
});
let response = serde_json::from_value(payload.clone()).unwrap();
let normalized = normalize_response("parse-v5.0", response).unwrap();
assert_eq!(normalized.pages[0].index, 4);
assert_eq!(normalized.pages[0].markdown, "receipt");
let image = &normalized.pages[0].images.as_ref().unwrap()[0];
let original_image = &payload["pages"][0]["markdown"]["images"][0];
assert_eq!(
serde_json::to_value(&image.bbox).unwrap(),
original_image["bounding_box"]
);
assert_eq!(
image.extra_fields["bounding_box_normalized"],
original_image["bounding_box_normalized"]
);
assert_eq!(image.extra_fields["id"], original_image["id"]);
assert_eq!(image.extra_fields["description"], "scan");
assert_eq!(image.extra_fields["category"], "logo");
assert_eq!(image.extra_fields["provider_extension"], "preserved");
assert_eq!(normalized.pages[1].index, 1);
assert_eq!(normalized.pages[1].markdown, "");
assert_eq!(
normalized.pages[1].extra_fields["blocks"][0]["text"]["content"],
"total"
);
assert_eq!(normalized.usage_info.unwrap().pages_processed, Some(3));
}
#[rstest]
#[case::empty(json!({}))]
#[case::null_meta(json!({"meta":null}))]
#[case::null_billed_units(json!({"pages":[],"meta":{"billed_units":null}}))]
fn response_defaults(#[case] value: Value) {
let normalized =
normalize_response("parse", serde_json::from_value(value).unwrap()).unwrap();
assert!(normalized.pages.is_empty());
assert_eq!(normalized.usage_info.unwrap().pages_processed, Some(0));
}
#[rstest]
#[case::null_pages(json!({"pages":null}))]
#[case::invalid_markdown(json!({"pages":[{"markdown":"text"}]}))]
#[case::invalid_index(json!({"pages":[{"index":"bad"}]}))]
fn response_rejects_invalid_fields(#[case] value: Value) {
assert!(serde_json::from_value::<CohereResponse>(value).is_err());
}
#[test]
fn null_markdown_uses_page_defaults() {
let normalized = normalize_response(
"parse",
serde_json::from_value(json!({"pages":[{"markdown":null}]})).unwrap(),
)
.unwrap();
assert_eq!(normalized.usage_info.unwrap().pages_processed, Some(1));
assert!(normalized.pages[0].images.is_none());
}
#[rstest]
fn response_types_documented_block_variants(
#[values(
crate::ocr::types::OcrResponseFormat::Litellm,
crate::ocr::types::OcrResponseFormat::Native
)]
response_format: crate::ocr::types::OcrResponseFormat,
) {
let payload = json!({
"pages": [{
"type": "blocks",
"index": 0,
"blocks": [
{"type": "text", "text": {"content": "hello"}},
{
"type": "image",
"image": {
"id": "img-0",
"description": "logo",
"category": "logo",
"bounding_box": {
"top_left_x": 1,
"top_left_y": 2,
"bottom_right_x": 3,
"bottom_right_y": 4
},
"bounding_box_normalized": {
"top_left_x": 0.1,
"top_left_y": 0.2,
"bottom_right_x": 0.3,
"bottom_right_y": 0.4
}
}
},
{
"type": "table",
"table": {
"type": "html",
"html": "<table></table>",
"bounding_box": {
"top_left_x": 5,
"top_left_y": 6,
"bottom_right_x": 7,
"bottom_right_y": 8
},
"bounding_box_normalized": {
"top_left_x": 0.5,
"top_left_y": 0.6,
"bottom_right_x": 0.7,
"bottom_right_y": 0.8
},
"title": "Totals",
"description": "Invoice totals"
}
}
]
}]
});
let normalized = CohereParseConfig
.transform_ocr_response(
"parse-v5.0",
&serde_json::to_vec(&payload).unwrap(),
response_format,
)
.unwrap();
assert_eq!(
normalized.pages[0].extra_fields["blocks"],
payload["pages"][0]["blocks"]
);
assert_eq!(normalized.pages[0].markdown, "");
assert_eq!(normalized.pages[0].index, 0);
assert_eq!(
normalized.usage_info.as_ref().unwrap().pages_processed,
Some(1)
);
match response_format {
crate::ocr::types::OcrResponseFormat::Litellm => {
assert!(normalized.provider_native_response.is_none());
}
crate::ocr::types::OcrResponseFormat::Native => {
assert_eq!(
normalized.provider_native_response.as_ref(),
payload.as_object()
);
}
}
assert_eq!(
normalized.into_json()["pages"][0]["blocks"],
payload["pages"][0]["blocks"]
);
}
#[rstest]
#[case::document_url(json!({"type":"document_url","document_url":"https://example.com/a.pdf"}))]
#[case::empty_image_url(json!({"type":"image_url","image_url":""}))]
#[case::pdf_data_uri(json!({"type":"image_url","image_url":"data:application/pdf;base64,YQ=="}))]
fn request_requires_image(#[case] value: Value) {
assert!(matches!(
validate_document(&serde_json::from_value(value).unwrap()),
Err(crate::ocr::Error::CohereImageOnly)
));
}
#[rstest]
#[case::markdown("markdown", true)]
#[case::blocks("blocks", true)]
#[case::unsupported("html", false)]
fn request_requires_supported_output_format(#[case] format: &str, #[case] valid: bool) {
assert_eq!(
serde_json::from_value::<CohereOptions>(json!({"output_format":format})).is_ok(),
valid
);
}
#[test]
fn request_defaults_to_markdown() {
let request = CohereParseConfig
.transform_ocr_request(
"parse-v5.0",
serde_json::from_value(json!({
"type":"image_url",
"image_url":"https://example.com/image.png"
}))
.unwrap(),
&serde_json::from_value(json!({})).unwrap(),
&[],
)
.unwrap();
assert_eq!(
serde_json::to_value(request).unwrap()["output_format"],
"markdown"
);
}
#[rstest]
#[case::base("", "/v2/parse")]
#[case::version("/v2", "/v2/parse")]
#[case::complete("/v2/parse", "/v2/parse")]
#[case::proxy_prefix("/cohere/", "/cohere/v2/parse")]
fn completes_provider_urls_without_duplicate_paths_and_preserves_queries(
#[case] suffix: &str,
#[case] path: &str,
) {
assert_eq!(
CohereParseConfig
.build_ocr_url(&format!("https://example.com{suffix}?tenant=a"))
.unwrap(),
format!("https://example.com{path}?tenant=a")
);
}
#[rstest]
#[case::relative("relative/path")]
#[case::unsupported_scheme("ftp://example.com")]
fn rejects_invalid_urls(#[case] api_base: &str) {
assert!(CohereParseConfig.build_ocr_url(api_base).is_err());
}
#[test]
fn rejects_blank_keys() {
assert!(matches!(
CohereParseConfig.resolve_headers(
&OcrConnection {
api_key: Some(" ".into()),
..Default::default()
},
&|_| None,
),
Err(crate::ocr::Error::Auth(_))
));
}
#[test]
fn environment_key_becomes_the_bearer() {
let headers = CohereParseConfig
.resolve_headers(&OcrConnection::default(), &|name| {
(name == COHERE_API_KEY_ENV).then(|| "env-key".to_string())
})
.unwrap();
assert_eq!(
headers,
[("Authorization".to_string(), "Bearer env-key".to_string())]
);
}
#[test]
fn missing_key_names_the_environment_variable() {
let error = CohereParseConfig
.resolve_headers(&OcrConnection::default(), &|_| None)
.unwrap_err();
assert!(error.to_string().contains(COHERE_API_KEY_ENV), "{error}");
}
#[rstest]
#[case::cohere("cohere/parse-v5.0", "POST /v2/parse ")]
#[case::azure_ai("azure_ai/Cohere-parse-v5.0", "POST /providers/cohere/v2/parse ")]
#[tokio::test]
async fn route_sends_image_to_its_parse_endpoint_with_the_bearer_key(
#[case] model: &str,
#[case] request_line: &str,
) {
use crate::ocr::test_support::{MockResponse, header, mock_server, perform_ocr};
let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await;
let request = crate::ocr::test_support::wire_request(model, &base, json!({}))
.with_document(
serde_json::from_value::<OcrDocument>(
json!({"type":"image_url","image_url":"data:image/png;base64,YWJj"}),
)
.unwrap()
.into(),
);
perform_ocr(request).await.unwrap();
server.await.unwrap();
let requests = seen.lock().unwrap();
assert_eq!(requests.len(), 1);
assert!(requests[0].starts_with(request_line), "{}", requests[0]);
assert_eq!(
header(&requests[0], "authorization"),
Some("Bearer test-key")
);
}
#[rstest]
#[tokio::test]
async fn route_rejects_non_image_document_without_a_request(
#[values("cohere/parse-v5.0", "azure_ai/Cohere-parse-v5.0")] model: &str,
) {
use crate::ocr::test_support::{MockResponse, mock_server, perform_ocr};
let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await;
let error = perform_ocr(crate::ocr::test_support::wire_request(
model,
&base,
json!({}),
))
.await
.unwrap_err();
server.abort();
assert!(
matches!(error, crate::ocr::Error::CohereImageOnly),
"{error:?}"
);
assert!(seen.lock().unwrap().is_empty());
}
}

1
litellm-rust/crates/core/src/llms/mistral/mod.rs Normal file
View file

@ -0,0 +1 @@
pub(crate) mod ocr;

1
litellm-rust/crates/core/src/llms/mistral/ocr/mod.rs Normal file
View file

@ -0,0 +1 @@
pub(crate) mod transformation;

653
litellm-rust/crates/core/src/llms/mistral/ocr/transformation.rs Normal file
View file

@ -0,0 +1,653 @@
use serde::{Deserialize, Serialize};
use serde_json::Value;
use crate::{
call_arguments::CallArguments,
constants::MISTRAL_OCR_API_BASE,
llms::base_llm::ocr::transformation::{BaseOcrConfig, decode_and_normalize_response},
ocr::{
OcrClient,
prepare::credential_env,
types::{
LiteLLMOcrResponse, OcrConnection, OcrDocument, OcrPage, OcrResponseFormat,
OcrUsageInfo, PreparedOcrRequest,
},
},
params::OpaqueParams,
url_utils::ApiUrl,
};
const MISTRAL_OCR_API_KEY_ENV_VAR: &str = "MISTRAL_API_KEY";
#[derive(Clone, Debug, Serialize, Deserialize)]
pub(crate) struct MistralOcrRequest {
pub model: String,
pub document: OcrDocument,
#[serde(flatten)]
pub params: OpaqueParams,
}
#[derive(Clone, Debug, Default, Deserialize)]
pub(crate) struct MistralOcrResponse {
#[serde(default)]
pub pages: Vec<OcrPage>,
#[serde(
default,
deserialize_with = "serde_with::rust::double_option::deserialize"
)]
pub model: Option<Option<String>>,
pub document_annotation: Option<Value>,
pub usage_info: Option<OcrUsageInfo>,
#[serde(flatten)]
pub extra_fields: serde_json::Map<String, Value>,
}
#[derive(Clone, Debug, Default)]
pub(crate) struct MistralOcrConfig;
impl BaseOcrConfig for MistralOcrConfig {
type OcrParams = OpaqueParams;
type ProviderRequest = MistralOcrRequest;
type Environment = Vec<(String, String)>;
fn get_supported_ocr_params(&self, _model: &str) -> &'static [&'static str] {
&[
"pages",
"include_image_base64",
"image_limit",
"image_min_size",
"bbox_annotation_format",
"document_annotation_format",
"document_annotation_prompt",
"extract_header",
"extract_footer",
"table_format",
"confidence_scores_granularity",
"include_blocks",
"id",
]
}
fn get_api_key_env_var(&self) -> Option<&'static str> {
Some(MISTRAL_OCR_API_KEY_ENV_VAR)
}
fn map_ocr_params(
&self,
non_default_params: &CallArguments,
model: &str,
) -> Result<OpaqueParams, crate::ocr::Error> {
Ok(non_default_params
.select(self.get_supported_ocr_params(model))
.into())
}
async fn validate_environment(
&self,
request: &PreparedOcrRequest,
_client: &OcrClient,
) -> Result<Self::Environment, crate::ocr::Error> {
self.resolve_headers(&request.connection, &credential_env)
}
fn get_complete_url(
&self,
request: &PreparedOcrRequest,
_optional_params: &Self::OcrParams,
_environment: &Self::Environment,
) -> Result<String, crate::ocr::Error> {
self.build_ocr_url(request.connection.api_base.as_deref())
}
fn transform_ocr_request(
&self,
model: &str,
document: OcrDocument,
optional_params: &OpaqueParams,
_headers: &[(String, String)],
) -> Result<MistralOcrRequest, crate::ocr::Error> {
Ok(MistralOcrRequest {
model: model.to_string(),
document,
params: optional_params.clone(),
})
}
fn transform_ocr_response(
&self,
model: &str,
raw_response: &[u8],
request_format: OcrResponseFormat,
) -> Result<LiteLLMOcrResponse, crate::ocr::Error> {
decode_and_normalize_response(model, raw_response, request_format, normalize_response)
}
}
impl MistralOcrConfig {
fn resolve_headers(
&self,
connection: &OcrConnection,
env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
) -> Result<Vec<(String, String)>, crate::ocr::Error> {
if crate::http_utils::has_header(&connection.extra_headers, "authorization") {
return Ok(connection.extra_headers.clone());
}
let api_key = connection
.api_key
.as_deref()
.map(str::trim)
.filter(|key| !key.is_empty())
.map(str::to_string)
.or_else(|| {
self.get_api_key_env_var()
.and_then(env_lookup)
.filter(|key| !key.trim().is_empty())
})
.ok_or(litellm_auth::Error::MissingApiKey {
provider: "Mistral",
environment_variable: MISTRAL_OCR_API_KEY_ENV_VAR,
})?;
Ok(
std::iter::once(("Authorization".into(), format!("Bearer {api_key}")))
.chain(connection.extra_headers.clone())
.collect(),
)
}
fn build_ocr_url(&self, api_base: Option<&str>) -> Result<String, crate::ocr::Error> {
let base = api_base
.map(str::trim)
.filter(|base| !base.is_empty())
.unwrap_or(MISTRAL_OCR_API_BASE);
ApiUrl::parse(base)
.and_then(|url| url.complete_path(&["v1", "ocr"]))
.map(|url| url.into_string())
.map_err(|_| crate::ocr::Error::RequestField {
path: "api_base".into(),
})
}
}
pub(crate) fn normalize_response(
model: &str,
response: MistralOcrResponse,
) -> Result<LiteLLMOcrResponse, crate::ocr::Error> {
let model = match response.model {
Some(Some(model)) => model,
Some(None) => {
return Err(crate::ocr::Error::ResponseField {
path: "model".into(),
});
}
None => model.to_string(),
};
Ok(LiteLLMOcrResponse {
extra_fields: response.extra_fields,
document_annotation: response.document_annotation,
usage_info: response.usage_info,
..LiteLLMOcrResponse::new(model, response.pages)
})
}
#[cfg(test)]
mod tests {
use rstest::{fixture, rstest};
use serde_json::{Value, json};
use super::*;
#[fixture]
fn document() -> OcrDocument {
serde_json::from_value(
json!({"type":"document_url","document_url":"https://example.com/a.pdf"}),
)
.unwrap()
}
#[fixture]
fn connection(
#[default(None)] api_key: Option<&str>,
#[default(vec![])] extra_headers: Vec<(String, String)>,
) -> OcrConnection {
OcrConnection {
api_key: api_key.map(str::to_string),
extra_headers,
..OcrConnection::default()
}
}
#[test]
fn explicit_null_model_does_not_use_the_missing_model_default() {
let response = serde_json::from_value(json!({"model":null})).unwrap();
assert!(matches!(
normalize_response("fallback", response).unwrap_err(),
crate::ocr::Error::ResponseField { path } if path == "model"
));
}
#[rstest]
#[case::non_object_page(json!({"pages":[42]}), "pages[0]")]
#[case::missing_markdown(json!({"pages":[{"index":0}]}), "pages[0]")]
#[case::non_string_markdown(
json!({"pages":[{"index":0,"markdown":42}]}),
"pages[0].markdown"
)]
#[case::non_object_image(
json!({"pages":[{"index":0,"markdown":"","images":[42]}]}),
"pages[0].images[0]"
)]
#[case::fractional_width(
json!({"pages":[{"index":0,"markdown":"","dimensions":{"width":1.5}}]}),
"pages[0].dimensions.width"
)]
#[case::invalid_page_count(
json!({"usage_info":{"pages_processed":"bad"}}),
"usage_info.pages_processed"
)]
fn response_validates_normalized_shapes_at_the_provider_boundary(
#[case] payload: Value,
#[case] path: &str,
) {
let error = crate::ocr::json::decode_response::<MistralOcrResponse>(
&serde_json::to_vec(&payload).unwrap(),
false,
)
.unwrap_err();
assert!(matches!(
error,
crate::ocr::Error::ResponseField { path: actual } if actual == path
));
}
#[test]
fn response_normalizes_python_numeric_inputs_and_shared_defaults() {
let response = serde_json::from_value(json!({
"pages":[{"index":"2","markdown":"text","dimensions":{"width":1.0},"extension":false}],
"usage_info":{"pages_processed":true,"credits":"1.5","custom":0},
"extra":"ignored"
}))
.unwrap();
let response = normalize_response("model", response).unwrap();
assert_eq!(response.pages[0].index, 2);
assert_eq!(
response.pages[0].dimensions.as_ref().unwrap().width,
Some(1)
);
assert_eq!(
response.usage_info.as_ref().unwrap().pages_processed,
Some(1)
);
assert_eq!(response.usage_info.as_ref().unwrap().credits, Some(1.5));
let serialized = response.into_json();
assert_eq!(serialized["pages"][0]["extension"], false);
assert!(serialized["pages"][0]["images"].is_null());
assert!(serialized["usage_info"]["doc_size_bytes"].is_null());
assert_eq!(serialized["usage_info"]["custom"], 0);
assert!(serialized["content"].is_null());
assert_eq!(serialized["extra"], "ignored");
}
#[test]
fn map_ocr_params_selects_known_fields_without_changing_arguments() {
let input =
serde_json::from_value(json!({"pages":null,"extract_header":false,"unknown":true}))
.unwrap();
let params = MistralOcrConfig.map_ocr_params(&input, "model").unwrap();
assert_eq!(
serde_json::to_value(params).unwrap(),
json!({"pages":null,"extract_header":false})
);
assert_eq!(input["unknown"], true);
assert_eq!(input.get("pages"), Some(&Value::Null));
}
#[rstest]
fn request_transform_uses_already_mapped_params_without_filtering_again(document: OcrDocument) {
let params = serde_json::from_value(json!({"extension":{"nested":null}})).unwrap();
let body = MistralOcrConfig
.transform_ocr_request("model", document, &params, &[])
.unwrap();
assert_eq!(
serde_json::to_value(body).unwrap()["extension"],
json!({"nested":null})
);
}
#[test]
fn raw_response_transform_keeps_native_payload_separate_from_typed_normalization() {
let raw = br#"{"pages":[{"index":"2","markdown":"text"}],"provider_extension":false}"#;
let response = MistralOcrConfig
.transform_ocr_response("model", raw, crate::ocr::types::OcrResponseFormat::Native)
.unwrap();
assert_eq!(response.pages[0].index, 2);
let native = response.provider_native_response.unwrap();
assert_eq!(native["pages"][0]["index"], "2");
assert_eq!(native["provider_extension"], false);
assert_eq!(response.extra_fields["provider_extension"], false);
}
#[rstest]
fn raw_response_transform_rejects_invalid_page(
#[values(OcrResponseFormat::Litellm, OcrResponseFormat::Native)]
request_format: OcrResponseFormat,
) {
assert!(
MistralOcrConfig
.transform_ocr_response("model", br#"{"pages":[{"index":0}]}"#, request_format)
.is_err()
);
}
fn mapped_params(value: Value) -> Value {
let params = serde_json::from_value(value).unwrap();
serde_json::to_value(MistralOcrConfig.map_ocr_params(&params, "model").unwrap()).unwrap()
}
#[rstest]
fn extract_header_is_a_supported_ocr_param() {
assert_eq!(
mapped_params(json!({"extract_header":true}))["extract_header"],
true
);
}
#[rstest]
fn extract_footer_is_a_supported_ocr_param() {
assert_eq!(
mapped_params(json!({"extract_footer":false}))["extract_footer"],
false
);
}
#[rstest]
fn existing_ocr_params_remain_supported() {
let mapped = mapped_params(json!({
"pages":[0,2],
"include_image_base64":true,
"image_limit":2,
"image_min_size":100,
"bbox_annotation_format":{"type":"json_schema"},
"document_annotation_format":{"type":"json_schema"}
}));
assert_eq!(mapped["pages"], json!([0, 2]));
assert_eq!(mapped["include_image_base64"], true);
assert_eq!(mapped["image_limit"], 2);
assert_eq!(mapped["image_min_size"], 100);
assert_eq!(mapped["bbox_annotation_format"]["type"], "json_schema");
assert_eq!(mapped["document_annotation_format"]["type"], "json_schema");
}
#[rstest]
fn map_ocr_params_forwards_extract_header() {
assert_eq!(
mapped_params(json!({"extract_header":true}))["extract_header"],
true
);
}
#[rstest]
fn map_ocr_params_forwards_extract_footer() {
assert_eq!(
mapped_params(json!({"extract_footer":true}))["extract_footer"],
true
);
}
#[rstest]
fn map_ocr_params_forwards_extract_header_and_footer() {
let mapped = mapped_params(json!({"extract_header":true,"extract_footer":false}));
assert_eq!(mapped["extract_header"], true);
assert_eq!(mapped["extract_footer"], false);
}
#[rstest]
fn map_ocr_params_excludes_extensions_from_the_provider_options() {
let mapped = mapped_params(json!({"extract_header":true,"unsupported_param":"value"}));
assert_eq!(mapped["extract_header"], true);
assert!(mapped.get("unsupported_param").is_none());
}
#[rstest]
fn map_ocr_params_preserves_unvalidated_values_and_explicit_null() {
let mapped = mapped_params(json!({
"pages":{"future":"shape"},
"include_image_base64":null
}));
assert_eq!(mapped["pages"], json!({"future":"shape"}));
assert!(mapped.get("include_image_base64").unwrap().is_null());
}
#[rstest]
#[case("table_format", json!("html"))]
#[case("confidence_scores_granularity", json!("word"))]
#[case("confidence_scores_granularity", json!("block"))]
#[case("document_annotation_prompt", json!("extract"))]
#[case("include_blocks", json!(true))]
#[case("id", json!("req-123"))]
fn new_ocr_params_are_supported(#[case] name: &str, #[case] value: Value) {
assert_eq!(mapped_params(json!({name:value.clone()}))[name], value);
}
#[rstest]
#[case("table_format", json!("html"))]
#[case("table_format", json!("markdown"))]
#[case("confidence_scores_granularity", json!("word"))]
#[case("confidence_scores_granularity", json!("page"))]
#[case("document_annotation_prompt", json!("extract"))]
#[case("include_blocks", json!(true))]
#[case("id", json!("req-123"))]
fn map_ocr_params_forwards_new_ocr_params(#[case] name: &str, #[case] value: Value) {
assert_eq!(mapped_params(json!({name:value.clone()}))[name], value);
}
#[rstest]
#[case("pages", json!([0, 2]))]
#[case("pages", json!("0,2-4"))]
#[case("pages", Value::Null)]
#[case("include_image_base64", json!(true))]
#[case("include_image_base64", json!(false))]
#[case("image_limit", json!(2))]
#[case("image_min_size", json!(100))]
#[case("bbox_annotation_format", json!({"type":"json_schema"}))]
#[case("document_annotation_format", json!({"type":"json_schema"}))]
#[case("document_annotation_prompt", json!("extract"))]
#[case("extract_header", json!(true))]
#[case("extract_footer", json!(false))]
#[case("table_format", json!("html"))]
#[case("table_format", json!("markdown"))]
#[case("confidence_scores_granularity", json!("word"))]
#[case("confidence_scores_granularity", json!("page"))]
#[case("confidence_scores_granularity", json!("block"))]
#[case("include_blocks", json!(true))]
#[case("include_blocks", json!(false))]
#[case("id", json!("req-123"))]
fn request_mapping_preserves_supplied_options(
document: OcrDocument,
#[case] name: &str,
#[case] value: Value,
) {
let arguments = serde_json::from_value(json!({name: value.clone()})).unwrap();
let params = MistralOcrConfig
.map_ocr_params(&arguments, "model")
.unwrap();
let result = serde_json::to_value(
MistralOcrConfig
.transform_ocr_request("model", document.clone(), &params, &[])
.unwrap(),
)
.unwrap();
assert_eq!(
result,
json!({"model":"model", "document":document, name:value})
);
}
#[rstest]
#[case("table_format", json!("html"))]
#[case("confidence_scores_granularity", json!("word"))]
#[case("document_annotation_prompt", json!("extract"))]
#[case("id", json!("req-123"))]
#[case("extract_header", json!(true))]
#[case("include_blocks", json!(true))]
#[case("pages", json!([0,1]))]
fn transform_ocr_request_includes_each_optional_param(
document: OcrDocument,
#[case] name: &str,
#[case] value: Value,
) {
let params: OpaqueParams = serde_json::from_value(json!({name:value.clone()})).unwrap();
let result = serde_json::to_value(
MistralOcrConfig
.transform_ocr_request("mistral-ocr-latest", document, &params, &[])
.unwrap(),
)
.unwrap();
assert_eq!(result[name], value);
assert_eq!(result["model"], "mistral-ocr-latest");
}
#[rstest]
fn transform_ocr_request_includes_multiple_new_params(document: OcrDocument) {
let params: OpaqueParams = serde_json::from_value(json!({
"table_format":"html",
"confidence_scores_granularity":"page",
"extract_header":true
}))
.unwrap();
let result = serde_json::to_value(
MistralOcrConfig
.transform_ocr_request("mistral-ocr-latest", document, &params, &[])
.unwrap(),
)
.unwrap();
assert_eq!(result["table_format"], "html");
assert_eq!(result["confidence_scores_granularity"], "page");
assert_eq!(result["extract_header"], true);
}
#[rstest]
fn transform_ocr_response_preserves_blocks_and_confidence_scores() {
let payload = json!({
"pages":[{
"index":0,
"markdown":"hello",
"images":[{"id":"img-0","image_base64":"data:image/png;base64,AA=="}],
"dimensions":{"width":612,"height":792,"dpi":72},
"blocks":[{"type":"title","bbox":{"x":1},"confidence_scores":{"mean":0.98}}],
"confidence_scores":{"average_page_confidence_score":0.99,"minimum_page_confidence_score":0.97}
}],
"model":"returned-model",
"document_annotation":"{\"language\":\"en\"}",
"usage_info":{"pages_processed":1}
});
let response: MistralOcrResponse = serde_json::from_value(payload.clone()).unwrap();
let result = normalize_response("model", response).unwrap().into_json();
assert_eq!(result["pages"][0]["blocks"], payload["pages"][0]["blocks"]);
assert_eq!(
result["pages"][0]["confidence_scores"],
payload["pages"][0]["confidence_scores"]
);
assert_eq!(result["pages"][0]["images"][0]["id"], "img-0");
assert_eq!(result["pages"][0]["dimensions"]["dpi"], 72);
assert_eq!(result["model"], "returned-model");
assert_eq!(result["document_annotation"], "{\"language\":\"en\"}");
assert_eq!(result["usage_info"]["pages_processed"], 1);
}
#[rstest]
fn transform_ocr_response_preserves_ocr4_page_fields() {
let page = json!({
"index":0,
"markdown":"table page",
"tables":[{"rows":2,"cols":3}],
"hyperlinks":["https://example.com"],
"header":"header",
"footer":"footer"
});
let response: MistralOcrResponse =
serde_json::from_value(json!({"pages":[page.clone()]})).unwrap();
let result = normalize_response("model", response).unwrap().into_json();
assert_eq!(result["pages"][0]["tables"], page["tables"]);
assert_eq!(result["pages"][0]["hyperlinks"], page["hyperlinks"]);
assert_eq!(result["pages"][0]["header"], page["header"]);
assert_eq!(result["pages"][0]["footer"], page["footer"]);
assert!(result["pages"][0]["images"].is_null());
assert!(result["pages"][0]["dimensions"].is_null());
}
#[rstest]
#[case::default_base(None, "https://api.mistral.ai/v1/ocr")]
#[case::versioned_base(
Some("https://example.com/v1?tenant=a"),
"https://example.com/v1/ocr?tenant=a"
)]
#[case::complete_endpoint(
Some("https://example.com/v1/ocr?tenant=a"),
"https://example.com/v1/ocr?tenant=a"
)]
fn complete_url_defaults_and_dedupes_v1(
#[case] api_base: Option<&str>,
#[case] expected: &str,
) {
assert_eq!(MistralOcrConfig.build_ocr_url(api_base).unwrap(), expected);
}
#[rstest]
#[case::explicit_key(Some("explicit"), "Bearer explicit")]
#[case::environment_fallback(None, "Bearer environment")]
fn environment_prefers_explicit_key_then_environment(
#[case] _api_key: Option<&str>,
#[case] expected: &str,
#[with(_api_key)] connection: OcrConnection,
) {
assert_eq!(
MistralOcrConfig
.resolve_headers(&connection, &|_| Some("environment".into()))
.unwrap()[0],
("Authorization".into(), expected.into())
);
}
#[rstest]
fn environment_preserves_forwarded_authorization(
#[with(None, vec![("authorization".into(), "Bearer forwarded".into())])]
connection: OcrConnection,
) {
assert_eq!(
MistralOcrConfig
.resolve_headers(&connection, &|_| None)
.unwrap(),
connection.extra_headers
);
}
#[rstest]
fn environment_keeps_extra_headers_after_the_bearer_key(
#[with(Some("explicit"), vec![("X-Trace".into(), "trace-1".into())])]
connection: OcrConnection,
) {
assert_eq!(
MistralOcrConfig
.resolve_headers(&connection, &|_| None)
.unwrap(),
[
("Authorization".to_string(), "Bearer explicit".to_string()),
("X-Trace".to_string(), "trace-1".to_string()),
]
);
}
#[rstest]
fn environment_rejects_missing_key(connection: OcrConnection) {
assert!(matches!(
MistralOcrConfig.resolve_headers(&connection, &|_| None),
Err(crate::ocr::Error::Auth(
litellm_auth::Error::MissingApiKey {
provider: "Mistral",
environment_variable: MISTRAL_OCR_API_KEY_ENV_VAR,
}
))
));
}
}

8
litellm-rust/crates/core/src/llms/mod.rs Normal file
View file

@ -0,0 +1,8 @@
pub mod anthropic;
pub mod azure_ai;
pub mod base_llm;
pub(crate) mod cohere;
pub(crate) mod mistral;
pub mod openai;
pub(crate) mod reducto;
pub(crate) mod vertex_ai;

0
litellm-rust/crates/core/src/providers/openai/mod.rs → litellm-rust/crates/core/src/llms/openai/mod.rs
View file

0
litellm-rust/crates/core/src/providers/anthropic/chat_completions/mod.rs → litellm-rust/crates/core/src/llms/openai/responses/mod.rs
View file

14
litellm-rust/crates/core/src/providers/openai/responses/transformation.rs → litellm-rust/crates/core/src/llms/openai/responses/transformation.rs
View file

@ -1,12 +1,14 @@
use crate::responses::Error;
use crate::responses::types::{ResponsesWsEvent, ResponsesWsTransformResult};
use crate::responses::websocket::{ResponsesWebSocketProviderConfig, enforce_model};
use crate::responses::{
Error,
types::{ResponsesWsEvent, ResponsesWsTransformResult},
websocket::{ResponsesWebSocketProviderConfig, enforce_model},
};
pub struct OpenAIResponsesWsConfig;
pub struct OpenAiResponsesApiConfig;
pub const OPENAI_RESPONSES_WS_CONFIG: OpenAIResponsesWsConfig = OpenAIResponsesWsConfig;
pub const OPENAI_RESPONSES_WS_CONFIG: OpenAiResponsesApiConfig = OpenAiResponsesApiConfig;
impl ResponsesWebSocketProviderConfig for OpenAIResponsesWsConfig {
impl ResponsesWebSocketProviderConfig for OpenAiResponsesApiConfig {
fn supports_native_websocket(&self) -> bool {
true
}

1
litellm-rust/crates/core/src/llms/reducto/mod.rs Normal file
View file

@ -0,0 +1 @@
pub(crate) mod ocr;

1
litellm-rust/crates/core/src/llms/reducto/ocr/mod.rs Normal file
View file

@ -0,0 +1 @@
pub(crate) mod transformation;

979
litellm-rust/crates/core/src/llms/reducto/ocr/transformation.rs Normal file
View file

@ -0,0 +1,979 @@
use std::collections::BTreeMap;
use serde::{Deserialize, Deserializer, Serialize};
use serde_json::{Map, Value, json};
use crate::{
call_arguments::{CallArguments, compose_body},
constants::{REDUCTO_API_BASE, REDUCTO_API_KEY_ENV, REDUCTO_ID_PREFIX},
llms::base_llm::ocr::transformation::{
BaseOcrConfig, OcrRequestContext, decode_and_normalize_response,
},
ocr::{
OcrClient,
document::InlineDocument,
prepare::{build_http_request, credential_env, guardrail_document},
types::{
LiteLLMOcrResponse, OcrConnection, OcrDocument, OcrPage, OcrResponseFormat,
OcrUsageInfo, PreparedOcrRequest,
},
},
params::OpaqueParams,
url_utils::ApiUrl,
};
#[derive(Clone, Debug, Serialize, Deserialize)]
#[serde(transparent)]
pub(crate) struct ReductoFileId(String);
pub(crate) type ReductoV3Params = OpaqueParams;
pub(crate) type ReductoLegacyParams = OpaqueParams;
#[derive(Clone, Debug, Serialize, Deserialize)]
pub(crate) struct ReductoV3Request {
pub input: ReductoFileId,
#[serde(flatten)]
pub params: ReductoV3Params,
}
#[derive(Clone, Debug, Serialize, Deserialize)]
pub(crate) struct ReductoLegacyRequest {
pub document_url: ReductoFileId,
#[serde(skip_serializing_if = "Option::is_none")]
pub options: Option<ReductoLegacyOptions>,
}
#[derive(Clone, Debug, Serialize, Deserialize)]
pub(crate) struct ReductoLegacyOptions {
pub enhance: Value,
}
#[derive(Deserialize)]
struct ReductoUploadResponse {
pub file_id: Option<String>,
}
#[derive(Clone, Debug, Deserialize)]
pub(crate) struct ReductoResponse {
#[serde(default, deserialize_with = "present_nullable")]
result: Option<Option<ReductoResult>>,
usage: Option<ReductoUsage>,
#[serde(default)]
chunks: Option<Vec<ReductoChunk>>,
}
#[derive(Clone, Debug, Default, Deserialize)]
struct ReductoResult {
pub chunks: Option<Vec<ReductoChunk>>,
}
#[serde_with::serde_as]
#[derive(Clone, Debug, Default, Deserialize)]
struct ReductoUsage {
#[serde_as(deserialize_as = "Option<crate::serde_compat::LaxI64>")]
pub num_pages: Option<i64>,
#[serde_as(deserialize_as = "Option<crate::serde_compat::FiniteF64>")]
pub credits: Option<f64>,
}
#[derive(Clone, Debug, Deserialize)]
struct ReductoChunk {
pub content: Option<String>,
pub blocks: Option<Vec<Map<String, Value>>>,
}
#[derive(Clone, Debug)]
pub(crate) struct ReductoParseV3Config;
impl BaseOcrConfig for ReductoParseV3Config {
type OcrParams = ReductoV3Params;
type ProviderRequest = ReductoV3Request;
type Environment = Vec<(String, String)>;
fn get_supported_ocr_params(&self, _model: &str) -> &'static [&'static str] {
&["formatting", "retrieval", "settings"]
}
fn map_ocr_params(
&self,
non_default_params: &CallArguments,
model: &str,
) -> Result<ReductoV3Params, crate::ocr::Error> {
Ok(non_default_params
.select(self.get_supported_ocr_params(model))
.into())
}
async fn validate_environment(
&self,
request: &PreparedOcrRequest,
_client: &OcrClient,
) -> Result<Self::Environment, crate::ocr::Error> {
resolve_headers(&request.connection, &credential_env)
}
fn get_complete_url(
&self,
request: &PreparedOcrRequest,
_optional_params: &Self::OcrParams,
_environment: &Self::Environment,
) -> Result<String, crate::ocr::Error> {
build_ocr_url(request.connection.api_base.as_deref())
}
fn transform_ocr_request(
&self,
_model: &str,
document: OcrDocument,
optional_params: &Self::OcrParams,
_headers: &[(String, String)],
) -> Result<Self::ProviderRequest, crate::ocr::Error> {
Ok(ReductoV3Request {
input: uploaded_file_id(document)?,
params: optional_params.clone(),
})
}
async fn async_transform_ocr_request(
&self,
_model: &str,
document: OcrDocument,
optional_params: &ReductoV3Params,
headers: &[(String, String)],
context: OcrRequestContext<'_>,
) -> Result<ReductoV3Request, crate::ocr::Error> {
let file_id = ensure_file_id_async(document, headers, context).await?;
Ok(ReductoV3Request {
input: file_id,
params: optional_params.clone(),
})
}
fn transform_ocr_response(
&self,
model: &str,
raw_response: &[u8],
request_format: OcrResponseFormat,
) -> Result<LiteLLMOcrResponse, crate::ocr::Error> {
decode_and_normalize_response(model, raw_response, request_format, normalize_response)
}
async fn prepare_request(
&self,
request: &PreparedOcrRequest,
client: &OcrClient,
) -> Result<reqwest::Request, crate::ocr::Error> {
prepare_upload_request(self, request, client).await
}
}
#[derive(Clone, Debug)]
pub(crate) struct ReductoParseLegacyConfig;
impl BaseOcrConfig for ReductoParseLegacyConfig {
type OcrParams = ReductoLegacyParams;
type ProviderRequest = ReductoLegacyRequest;
type Environment = Vec<(String, String)>;
fn get_supported_ocr_params(&self, _model: &str) -> &'static [&'static str] {
&["enhance"]
}
fn map_ocr_params(
&self,
non_default_params: &CallArguments,
model: &str,
) -> Result<ReductoLegacyParams, crate::ocr::Error> {
Ok(non_default_params
.select(self.get_supported_ocr_params(model))
.into())
}
async fn validate_environment(
&self,
request: &PreparedOcrRequest,
client: &OcrClient,
) -> Result<Self::Environment, crate::ocr::Error> {
ReductoParseV3Config
.validate_environment(request, client)
.await
}
fn get_complete_url(
&self,
request: &PreparedOcrRequest,
optional_params: &Self::OcrParams,
environment: &Self::Environment,
) -> Result<String, crate::ocr::Error> {
ReductoParseV3Config.get_complete_url(request, optional_params, environment)
}
fn transform_ocr_request(
&self,
_model: &str,
document: OcrDocument,
optional_params: &Self::OcrParams,
_headers: &[(String, String)],
) -> Result<Self::ProviderRequest, crate::ocr::Error> {
Ok(build_legacy_body(
uploaded_file_id(document)?,
optional_params,
))
}
async fn async_transform_ocr_request(
&self,
_model: &str,
document: OcrDocument,
optional_params: &ReductoLegacyParams,
headers: &[(String, String)],
context: OcrRequestContext<'_>,
) -> Result<ReductoLegacyRequest, crate::ocr::Error> {
let file_id = ensure_file_id_async(document, headers, context).await?;
Ok(build_legacy_body(file_id, optional_params))
}
fn transform_ocr_response(
&self,
model: &str,
raw_response: &[u8],
request_format: OcrResponseFormat,
) -> Result<LiteLLMOcrResponse, crate::ocr::Error> {
ReductoParseV3Config.transform_ocr_response(model, raw_response, request_format)
}
async fn prepare_request(
&self,
request: &PreparedOcrRequest,
client: &OcrClient,
) -> Result<reqwest::Request, crate::ocr::Error> {
prepare_upload_request(self, request, client).await
}
}
/// Reducto differs from the shared `BaseOcrConfig::prepare_request` flow:
/// guardrails see the *source* document before it is uploaded, because the
/// final body only carries the opaque Reducto file id.
async fn prepare_upload_request<C: BaseOcrConfig<Environment = Vec<(String, String)>>>(
config: &C,
request: &PreparedOcrRequest,
client: &OcrClient,
) -> Result<reqwest::Request, crate::ocr::Error> {
let params = config.map_ocr_params(&request.optional_params, &request.model)?;
let headers = config.validate_environment(request, client).await?;
let url = config.get_complete_url(request, &params, &headers)?;
let (document, headers) = guardrail_document(request, &url, &headers).await?;
let body = config
.async_transform_ocr_request(
&request.model,
document,
&params,
&headers,
OcrRequestContext {
client,
connection: &request.connection,
},
)
.await?;
let body = compose_body(
&request.optional_params,
&body,
config.get_supported_ocr_params(&request.model),
)?;
build_http_request(client, request, &url, &headers, &body)
}
fn uploaded_file_id(document: OcrDocument) -> Result<ReductoFileId, crate::ocr::Error> {
if !document.source().starts_with(REDUCTO_ID_PREFIX) {
return Err(crate::ocr::Error::ReductoSource);
}
if document.source()[REDUCTO_ID_PREFIX.len()..]
.trim()
.is_empty()
{
return Err(crate::ocr::Error::RequestField {
path: "document file id".into(),
});
}
Ok(ReductoFileId(document.source().into()))
}
fn present_nullable<'de, D: Deserializer<'de>, T: Deserialize<'de>>(
deserializer: D,
) -> Result<Option<Option<T>>, D::Error> {
Option::<T>::deserialize(deserializer).map(Some)
}
fn block_page_number(value: &Value) -> Option<i64> {
match value {
Value::Number(number) => number
.as_i64()
.or_else(|| number.as_f64().and_then(checked_truncated_i64)),
Value::String(value) => value.trim().parse::<i64>().ok(),
Value::Bool(value) => Some(i64::from(*value)),
_ => None,
}
}
fn checked_truncated_i64(value: f64) -> Option<i64> {
(value.is_finite() && value >= i64::MIN as f64 && value <= i64::MAX as f64)
.then(|| value.trunc() as i64)
}
pub(crate) fn normalize_response(
model: &str,
response: ReductoResponse,
) -> Result<LiteLLMOcrResponse, crate::ocr::Error> {
let result = match response.result {
Some(result) => result.unwrap_or_default(),
None => ReductoResult {
chunks: response.chunks,
},
};
let usage = response.usage.unwrap_or_default();
Ok(LiteLLMOcrResponse {
usage_info: Some(OcrUsageInfo {
pages_processed: usage.num_pages,
credits: usage.credits,
..Default::default()
}),
..LiteLLMOcrResponse::new(
model,
build_pages_from_reducto(result.chunks.unwrap_or_default())?,
)
})
}
fn build_pages_from_reducto(chunks: Vec<ReductoChunk>) -> Result<Vec<OcrPage>, crate::ocr::Error> {
let blocks_by_page = chunks
.iter()
.flat_map(|chunk| chunk.blocks.iter().flatten())
.filter_map(|block| {
block_page_number(block.get("bbox")?.get("page")?).map(|page| (page, block))
})
.fold(
BTreeMap::<i64, Vec<&Map<String, Value>>>::new(),
|mut pages, (page, block)| {
pages.entry(page).or_default().push(block);
pages
},
);
if blocks_by_page.is_empty() {
let markdown = join_content(chunks.iter().map(|chunk| chunk.content.as_deref()));
return Ok(if markdown.is_empty() {
Vec::new()
} else {
vec![page(0, markdown, None)]
});
}
blocks_by_page
.into_iter()
.map(|(index, blocks)| {
let content = blocks
.iter()
.map(|block| match block.get("content") {
None | Some(Value::Null) => Ok(None),
Some(Value::String(content)) => Ok(Some(content.as_str())),
Some(_) => Err(crate::ocr::Error::ResponseField {
path: "result.chunks.blocks.content".into(),
}),
})
.collect::<Result<Vec<_>, _>>()?;
let markdown = join_content(content.into_iter());
Ok(page(
index.saturating_sub(1).max(0),
markdown,
Some(json!(blocks)),
))
})
.collect()
}
fn join_content<'a>(content: impl Iterator<Item = Option<&'a str>>) -> String {
content
.flatten()
.filter(|text| !text.is_empty())
.collect::<Vec<_>>()
.join("\n\n")
}
fn page(index: i64, markdown: String, blocks: Option<Value>) -> OcrPage {
OcrPage {
index,
markdown,
extra_fields: blocks
.map(|blocks| ("blocks".into(), blocks))
.into_iter()
.collect(),
..Default::default()
}
}
fn build_ocr_url(api_base: Option<&str>) -> Result<String, crate::ocr::Error> {
complete_endpoint_url(api_base, "parse")
}
fn complete_endpoint_url(api_base: Option<&str>, path: &str) -> Result<String, crate::ocr::Error> {
let base = api_base
.map(str::trim)
.filter(|base| !base.is_empty())
.unwrap_or(REDUCTO_API_BASE);
ApiUrl::parse(base)
.and_then(|url| url.complete_path(&[path]))
.map(|url| url.into_string())
.map_err(|_| crate::ocr::Error::RequestField {
path: "api_base".into(),
})
}
fn resolve_headers(
connection: &OcrConnection,
env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
) -> Result<Vec<(String, String)>, crate::ocr::Error> {
if crate::http_utils::has_header(&connection.extra_headers, "authorization") {
return Ok(connection.extra_headers.clone());
}
let api_key = connection
.api_key
.as_deref()
.map(str::trim)
.filter(|key| !key.is_empty())
.map(str::to_string)
.or_else(|| {
env_lookup(REDUCTO_API_KEY_ENV)
.map(|key| key.trim().to_string())
.filter(|key| !key.is_empty())
})
.ok_or(crate::ocr::Error::MissingReductoApiKey)?;
Ok(
std::iter::once(("Authorization".into(), format!("Bearer {api_key}")))
.chain(connection.extra_headers.clone())
.collect(),
)
}
fn build_legacy_body(
file_id: ReductoFileId,
optional_params: &ReductoLegacyParams,
) -> ReductoLegacyRequest {
ReductoLegacyRequest {
document_url: file_id,
options: optional_params
.get("enhance")
.filter(|value| !value.is_null())
.map(|enhance| ReductoLegacyOptions {
enhance: enhance.clone(),
}),
}
}
async fn ensure_file_id_async(
document: OcrDocument,
headers: &[(String, String)],
context: OcrRequestContext<'_>,
) -> Result<ReductoFileId, crate::ocr::Error> {
if document.source().starts_with(REDUCTO_ID_PREFIX) {
if document.source()[REDUCTO_ID_PREFIX.len()..]
.trim()
.is_empty()
{
return Err(crate::ocr::Error::RequestField {
path: "document file id".into(),
});
}
return Ok(ReductoFileId(document.source().to_string()));
}
let inline =
InlineDocument::parse(document.source())?.ok_or(crate::ocr::Error::ReductoSource)?;
let mime = inline.mime_type().to_string();
let bytes = inline.decode(crate::constants::OCR_INLINE_MAX_BYTES)?;
upload_bytes_async(bytes, &mime, headers, context).await
}
async fn upload_bytes_async(
bytes: Vec<u8>,
mime: &str,
headers: &[(String, String)],
context: OcrRequestContext<'_>,
) -> Result<ReductoFileId, crate::ocr::Error> {
let OcrRequestContext { client, connection } = context;
let part = reqwest::multipart::Part::bytes(bytes)
.file_name("document")
.mime_str(mime)
.map_err(|_| crate::ocr::Error::InvalidDataUri)?;
let builder = client
.provider_http()
.post(complete_endpoint_url(
connection.api_base.as_deref(),
"upload",
)?)
.multipart(reqwest::multipart::Form::new().part("file", part))
.timeout(connection.timeout);
let builder = crate::http_utils::with_headers(
builder,
headers,
crate::http_utils::HeaderPolicy::Except(&["content-type", "content-length"]),
);
let response = crate::http_utils::http_request(builder)
.await
.map_err(crate::transport::Error::from)?;
let uploaded = crate::ocr::client::read_json_response::<ReductoUploadResponse>(
response,
false,
connection.max_response_bytes,
)
.await?
.data;
let file_id = uploaded
.file_id
.as_deref()
.map(str::trim)
.filter(|id| !id.is_empty());
let Some(file_id) = file_id else {
return Err(crate::ocr::Error::ResponseField {
path: "file_id".into(),
});
};
Ok(ReductoFileId(file_id.to_string()))
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn options_preserve_null_and_select_the_provider_fields() {
let overrides = serde_json::from_value(json!({
"formatting":null, "enhance":null, "ignored":true
}))
.unwrap();
let v3 = ReductoParseV3Config
.map_ocr_params(&overrides, "parse-v3")
.unwrap();
assert_eq!(
serde_json::to_value(v3).unwrap(),
json!({
"formatting":null
})
);
let legacy = ReductoParseLegacyConfig
.map_ocr_params(&overrides, "parse-legacy")
.unwrap();
assert_eq!(
serde_json::to_value(legacy).unwrap(),
json!({
"enhance":null
})
);
}
#[test]
fn usage_uses_shared_validation_while_block_page_numbers_are_best_effort() {
for usage in [
json!({"num_pages":1.5}),
json!({"num_pages":[]}),
json!({"credits":{}}),
] {
assert!(serde_json::from_value::<ReductoResponse>(json!({"usage":usage})).is_err());
}
let response = serde_json::from_value(json!({"result":{"chunks":[{"blocks":[
{"content":"ignored", "bbox":{"page":"invalid"}},
{"content":"kept", "bbox":{"page":2.5}, "extra":null}
]}]}, "usage":{"num_pages":2.0, "credits":true}}))
.unwrap();
let normalized = normalize_response("model", response).unwrap();
assert_eq!(normalized.pages[0].index, 1);
assert_eq!(normalized.pages[0].markdown, "kept");
assert_eq!(
normalized.pages[0].extra_fields["blocks"][0]["bbox"]["page"],
2.5
);
assert_eq!(normalized.usage_info.unwrap().credits, Some(1.0));
}
#[tokio::test]
async fn v3_options_preserve_explicit_null() {
let overrides =
serde_json::from_value(json!({"formatting":null,"settings":{},"unknown":true}))
.unwrap();
let params = ReductoParseV3Config
.map_ocr_params(&overrides, "parse-v3")
.unwrap();
let client = crate::ocr::test_support::ocr_client();
let connection = OcrConnection::default();
let document = serde_json::from_value(
json!({"type":"document_url","document_url":"reducto://ready.pdf"}),
)
.unwrap();
let body = ReductoParseV3Config
.async_transform_ocr_request(
"parse-v3",
document,
&params,
&[],
OcrRequestContext {
client: &client,
connection: &connection,
},
)
.await
.unwrap();
assert_eq!(
serde_json::to_value(body).unwrap(),
json!({
"input":"reducto://ready.pdf", "formatting":null, "settings":{}
})
);
let absent = ReductoParseV3Config
.map_ocr_params(&crate::call_arguments::CallArguments::default(), "parse-v3")
.unwrap();
assert_eq!(serde_json::to_value(absent).unwrap(), json!({}));
}
#[test]
fn legacy_body_omits_null_enhance_and_wraps_mapped_options() {
for (value, expected) in [
(json!(null), json!({"document_url":"reducto://ready.pdf"})),
(
json!({}),
json!({"document_url":"reducto://ready.pdf","options":{"enhance":{}}}),
),
] {
let overrides =
serde_json::from_value(json!({"enhance":value,"unknown":true})).unwrap();
let params = ReductoParseLegacyConfig
.map_ocr_params(&overrides, "parse-legacy")
.unwrap();
assert_eq!(
serde_json::to_value(build_legacy_body(
ReductoFileId("reducto://ready.pdf".into()),
&params
))
.unwrap(),
expected
);
}
}
#[test]
fn explicit_key_precedes_environment_key() {
let connection = OcrConnection {
api_key: Some("passed-key".into()),
..Default::default()
};
let headers = resolve_headers(&connection, &|_| Some("env-key".into())).unwrap();
assert_eq!(headers[0].1, "Bearer passed-key");
}
#[test]
fn blank_explicit_key_uses_environment_key() {
let connection = OcrConnection {
api_key: Some(" ".into()),
..Default::default()
};
let headers = resolve_headers(&connection, &|_| Some(" env-key ".into())).unwrap();
assert_eq!(headers[0].1, "Bearer env-key");
}
#[test]
fn existing_authorization_skips_key_lookup() {
let connection = OcrConnection {
extra_headers: vec![("authorization".into(), "Bearer existing".into())],
..Default::default()
};
assert_eq!(
resolve_headers(&connection, &|_| None).unwrap(),
connection.extra_headers
);
}
use litellm_callbacks::event::{CallEvent, WireRequest};
use rstest::rstest;
use crate::ocr::{
LocalOcrHost,
test_support::{MockResponse, mock_server, perform_ocr, perform_ocr_with, wire_request},
};
fn request_body(request: &str) -> Value {
serde_json::from_str(request.split_once("\r\n\r\n").unwrap().1).unwrap()
}
#[rstest]
#[case(
"reducto/parse-v3",
json!({
"formatting":{"table_output_format":"html"},
"retrieval":{"chunk_mode":"section"},
"settings":{"ocr_system":"standard"},
"future_ocr_option":true,
"extra_body":{"provider_option":"value"}
}),
"reducto://already.pdf",
json!({
"input":"reducto://already.pdf",
"formatting":{"table_output_format":"html"},
"retrieval":{"chunk_mode":"section"},
"settings":{"ocr_system":"standard"},
"future_ocr_option":true,
"provider_option":"value"
})
)]
#[case(
"reducto/parse-legacy",
json!({
"enhance":{"agentic":[{"type":"table"}]},
"future_ocr_option":true,
"extra_body":{"provider_option":"value"}
}),
"reducto://legacy.pdf",
json!({
"document_url":"reducto://legacy.pdf",
"options":{"enhance":{"agentic":[{"type":"table"}]}},
"future_ocr_option":true,
"provider_option":"value"
})
)]
#[tokio::test]
async fn request_mapping_matches_python(
#[case] model: &str,
#[case] options: Value,
#[case] source: &str,
#[case] expected: Value,
) {
let (base, seen, server) = mock_server(vec![MockResponse::json(json!({
"result":{"chunks":[]}
}))])
.await;
let request =
crate::ocr::test_support::with_source(wire_request(model, &base, options), source);
perform_ocr(request).await.unwrap();
server.await.unwrap();
let requests = seen.lock().unwrap();
assert_eq!(requests.len(), 1);
assert!(requests[0].starts_with("POST /parse "));
assert_eq!(request_body(&requests[0]), expected);
}
#[rstest]
#[case("parse-v3")]
#[case("parse-legacy")]
#[tokio::test]
async fn data_uri_upload_preserves_multipart_headers(#[case] model: &str) {
let (base, seen, server) = mock_server(vec![
MockResponse::json(json!({"file_id":"reducto://uploaded.pdf"})),
MockResponse::json(json!({"result":{"chunks":[{"content":"hello"}]}})),
])
.await;
let mut request = wire_request(&format!("reducto/{model}"), &base, json!({}));
request.transport.extra_headers = vec![
("Content-Type".into(), "application/json".into()),
("X-Trace".into(), "upload-test".into()),
];
let response = perform_ocr(request).await.unwrap();
server.await.unwrap();
assert_eq!(response.pages[0].markdown, "hello");
let requests = seen.lock().unwrap();
assert_eq!(requests.len(), 2);
assert!(requests[0].starts_with("POST /upload "));
assert!(
requests[0]
.to_ascii_lowercase()
.contains("content-type: multipart/form-data; boundary=")
);
assert!(requests[0].contains("x-trace: upload-test"));
assert!(requests[0].contains("application/pdf"));
assert!(requests[0].contains("abc"));
assert!(requests[1].starts_with("POST /parse "));
}
#[tokio::test]
async fn response_received_stays_after_reducto_upload_and_parse() {
let (base, seen, server) = mock_server(vec![
MockResponse::json(json!({"file_id":"reducto://uploaded.pdf"})),
MockResponse::json(json!({"result":{"chunks":[]}})),
])
.await;
let request_count = seen.clone();
let host = LocalOcrHost::new(wire_request("reducto/parse-v3", &base, json!({})))
.with_observer(move |event| {
if let CallEvent::ResponseReceived { raw } = event {
assert_eq!(request_count.lock().unwrap().len(), 2);
assert_eq!(raw.body, r#"{"result":{"chunks":[]}}"#);
}
});
perform_ocr_with(host).await.unwrap();
server.await.unwrap();
assert_eq!(seen.lock().unwrap().len(), 2);
}
#[rstest]
#[case(json!({"file_id":""}))]
#[case(json!({}))]
#[case(json!({"file_id":null}))]
#[tokio::test]
async fn invalid_upload_ids_stop_before_parse(#[case] response: Value) {
let (base, seen, server) = mock_server(vec![MockResponse::json(response)]).await;
let error = perform_ocr(wire_request("reducto/parse-v3", &base, json!({})))
.await
.unwrap_err();
server.await.unwrap();
assert!(error.to_string().contains("file_id"));
assert_eq!(seen.lock().unwrap().len(), 1);
}
#[tokio::test]
async fn upload_failure_stops_before_parse() {
let (base, seen, server) = mock_server(vec![MockResponse {
status: 503,
headers: vec![],
body: json!({"error":"unavailable"}),
}])
.await;
assert!(
perform_ocr(wire_request("reducto/parse-v3", &base, json!({})))
.await
.is_err()
);
server.await.unwrap();
assert_eq!(seen.lock().unwrap().len(), 1);
}
#[rstest]
#[case("https://example.com/a.pdf")]
#[case("reducto://")]
#[case("data:application/pdf;base64")]
#[case("data:application/pdf;base64,INVALID!")]
#[tokio::test]
async fn rejects_invalid_document_sources_before_network(#[case] source: &str) {
let request = crate::ocr::test_support::with_source(
wire_request("reducto/parse-v3", "http://127.0.0.1:1", json!({})),
source,
);
assert!(perform_ocr(request).await.is_err());
}
#[test]
fn response_normalization_groups_blocks_and_distinguishes_null_result() {
use crate::llms::reducto::ocr::transformation::{ReductoResponse, normalize_response};
let raw = json!({"usage":{"num_pages":"2","credits":"3"},"result":{"type":"full","chunks":[
{"blocks":[{
"type":"Table",
"content":"B",
"bbox":{"left":0.1,"top":0.2,"width":0.8,"height":0.3,"page":2,"original_page":4},
"confidence":"high",
"granular_confidence":{"parse_confidence":0.95,"extract_confidence":null},
"image_url":null
}]},
{"blocks":[{"content":"A","bbox":{"page":1},"type":"Text"},{"content":"C","bbox":{"page":1}}]}
]}});
let response: ReductoResponse = serde_json::from_value(raw).unwrap();
let normalized = normalize_response("parse-v3", response)
.unwrap()
.into_json();
assert_eq!(normalized["pages"][0]["markdown"], "A\n\nC");
assert_eq!(normalized["pages"][1]["markdown"], "B");
assert_eq!(normalized["pages"][1]["blocks"][0]["type"], "Table");
assert_eq!(
normalized["pages"][1]["blocks"][0]["bbox"],
json!({"left":0.1,"top":0.2,"width":0.8,"height":0.3,"page":2,"original_page":4})
);
assert_eq!(normalized["pages"][1]["blocks"][0]["confidence"], "high");
assert_eq!(
normalized["pages"][1]["blocks"][0]["granular_confidence"]["parse_confidence"],
0.95
);
assert!(normalized["pages"][1]["blocks"][0]["image_url"].is_null());
assert_eq!(normalized["usage_info"]["pages_processed"], 2);
assert_eq!(normalized["usage_info"]["credits"], 3.0);
let missing: ReductoResponse =
serde_json::from_value(json!({"chunks":[{"content":"text"}]})).unwrap();
let missing = normalize_response("parse-v3", missing).unwrap();
assert_eq!(missing.pages[0].markdown, "text");
let null: ReductoResponse = serde_json::from_value(
json!({"result":null,"chunks":[{"content":"ignored"}],"usage":null}),
)
.unwrap();
let null = normalize_response("parse-v3", null).unwrap();
assert!(null.pages.is_empty());
}
#[tokio::test]
async fn facade_omits_native_response_by_default_and_preserves_auth_priority() {
let raw = json!({"job_id":"job-1","result":{"chunks":[]}});
let (base, seen, server) = mock_server(vec![MockResponse::json(raw)]).await;
let mut request = crate::ocr::test_support::with_source(
wire_request("reducto/parse-v3", &base, json!({})),
"reducto://ready.pdf",
);
request.transport.extra_headers = vec![("authorization".into(), "Bearer existing".into())];
let response = perform_ocr(request).await.unwrap();
server.await.unwrap();
assert_eq!(response.provider_native_response, None);
assert!(
seen.lock().unwrap()[0]
.to_ascii_lowercase()
.contains("authorization: bearer existing")
);
}
#[rstest]
#[case("reducto/parse-v3")]
#[case("reducto/parse-legacy")]
#[tokio::test]
async fn guardrail_headers_reach_upload_and_parse(#[case] model: &str) {
let (base, seen, server) = mock_server(vec![
MockResponse::json(json!({"file_id":"reducto://uploaded.pdf"})),
MockResponse::json(json!({"result":{"chunks":[]}})),
])
.await;
let mut request = wire_request(model, &base, json!({}));
request.transport.extra_headers = vec![("authorization".into(), "Bearer original".into())];
let host = LocalOcrHost::new(request).with_before_send(|wire, _| {
Ok(WireRequest {
headers: vec![("authorization".into(), "Bearer guarded".into())],
..wire
})
});
perform_ocr_with(host).await.unwrap();
server.await.unwrap();
let requests = seen.lock().unwrap();
assert_eq!(requests.len(), 2);
assert!(requests[0].starts_with("POST /upload "));
assert!(requests[1].starts_with("POST /parse "));
for request in requests.iter() {
assert!(request.contains("authorization: Bearer guarded"));
assert!(!request.contains("Bearer original"));
}
}
#[tokio::test]
async fn guardrail_rewrites_document_before_upload() {
let (base, seen, server) =
mock_server(vec![MockResponse::json(json!({"result":{"chunks":[]}}))]).await;
let host = LocalOcrHost::new(wire_request("reducto/parse-v3", &base, json!({})))
.with_before_send(|wire, _| {
assert_eq!(
wire.body["document_url"],
"data:application/pdf;base64,YWJj"
);
Ok(WireRequest {
body: json!({"type":"document_url","document_url":"reducto://guarded.pdf"}),
..wire
})
});
perform_ocr_with(host).await.unwrap();
server.await.unwrap();
let requests = seen.lock().unwrap();
assert_eq!(requests.len(), 1);
assert!(requests[0].starts_with("POST /parse "));
assert!(requests[0].contains("reducto://guarded.pdf"));
}
}

Some files were not shown because too many files have changed in this diff Show more