xinyin025/LangBot
1
0
Fork 0
mirror of https://github.com/langbot-app/LangBot.git synced 2026-10-22 08:25:21 +08:00
Code Issues Packages Projects Releases Wiki Activity

Merge origin/master into feat/cloud-v2-subscriptions

Browse Source 
Keep i18n'd requester error messages while adopting master's structured
RequesterError metadata.

Co-Authored-By: Claude Opus 5.5 <[email protected]>
This commit is contained in:
RockChinQ
2026-10-17 01:37:14 +08:00
co-authored by Claude Opus 5.5
parent 16cfb6f8bd 1c5ab3227c
commit 653f9e731d
32 changed files with 4563 additions and 410 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/LLM_LATENCY_BENCHMARK.md
+100
View File
@@ -0,0 +1,100 @@
# LLM request-layer benchmark (2026-10-15)
This controlled test did not reproduce a large LangBot slowdown. Warm request
preparation was sub-millisecond; most measured latency occurred upstream.
It did identify two local transport problems, now fixed in this PR.
## Method
One Windows host, Python 3.12.7, LiteLLM 1.88.1 and OpenCode 2.0.16 used the
same user-provided NewAPI gateway and requested `claude-haiku-4.5`. Each client
made one warmup and six measured streaming calls, alternating all six client
orders. The synthetic system prefix was about 6,000 tokens. Parameters were
`temperature=0`, `max_tokens=256`, `stream=true`, `include_usage=true`, with an
ephemeral cache marker on the system text. No reasoning budget was requested.
The measurement relay deliberately forwards identical serialized requests and
headers for all three clients. It replaces OpenCode's added agent/environment
instructions and normalizes SDK request differences. The canonical SHA-256 is
`7834a66a188e70b4eb1b6234b3a14be4c66b8c72dd8a2c06395fe2505c1c371b`.
OpenCode runs as an isolated resident server, using its
[versioned session API](https://github.com/anomalyco/opencode/blob/v2.0.16/packages/protocol/openapi.json).
Process startup and session creation are outside the warm dispatch measurement.
All clients share the relay's upstream connection pool. Therefore this experiment
isolates local request dispatch and upstream response timing; it does **not**
compare their independent direct-to-provider TLS/proxy behavior. It also excludes
the complete bot pipeline, Runner/plugin RPC and OpenCode UI rendering. It cannot
establish which application is faster in every production deployment.
## Results
Medians, milliseconds; six measured samples per client:
| Client | Dispatch to relay | First content | Complete upstream response |
| --- | ---: | ---: | ---: |
| Direct HTTPX | 3.952 | 1119.541 | 2199.413 |
| LangBot Core requester | 6.802 | 1089.092 | 2135.935 |
| OpenCode | 41.158 | 1104.020 | 2148.265 |
Dispatch includes client preparation and local HTTP dispatch. First content and
completion use the same relay boundary for all clients, measured from dispatch;
they are not UI latency. LangBot's separately instrumented preparation median
was **0.316 ms**, maximum **0.497 ms**. These small samples and upstream variation
do not justify interpreting the small TTFT differences as a speed ranking.
Every measured response reported 5,990 input tokens, 5,987 cache-read tokens and
109 output tokens. This demonstrates equal warm cache accounting under the
normalized upstream request, not native cache-marker forwarding by all clients.
No cache-miss latency comparison or independent provider-route verification was
performed. Missing cache accounting is represented as null, never assumed zero.
The gateway may perform its own routing, retries and cache translation.
Sanitized individual samples and summaries are in
[llm-latency-2026-10-15.json](llm-latency-2026-10-15.json).
## Bottlenecks fixed
1. The OpenAI SDK stops consuming SSE at `[DONE]` before HTTP framing reaches
EOF. The healthy connection was then discarded. In the baseline relay run,
LangBot reused zero of six measured client connections; after bounded terminal
draining it reused six of six. A real TCP regression test also reproduces this
with split terminal markers and late usage. OpenRouter's generic adapter
already consumed EOF. Early close and cancellation still close immediately.
2. On this host, constructing default HTTPX clients took 1,761 / 1,443 / 1,435 ms.
Loading one TLS context per requester pool and sharing it across that client's
proxy mounts reduced measured pool setup to 666 / 430 / 431 ms. Certificate
loading now runs off the event loop. These are local setup measurements, not
a claim of equivalent percentage improvement in LLM latency. Runtime lazy
initialization and process startup remain separate cold costs.
3. High-resolution monotonic timing replaces the low-resolution clock used on
Python 3.12 Windows. Preparation no longer misleadingly rounds to zero.
`transport_setup_ms` exposes pool initialization separately.
The relay socket observations concern client-to-relay connections. HTTPX and
OpenCode have different keepalive policies, so those counts alone are not a
cross-client pooling ranking. The shared relay also prevents quantifying the
provider-side handshake savings from the SSE fix in this run.
## Reproduce
From a development environment with LangBot dependencies installed, set
`LLM_BENCHMARK_API_KEY` securely in the environment, then run:
```sh
python scripts/llm_latency_probe.py --opencode /path/to/opencode \
--opencode-version 2.0.16 --base-url https://gateway.example/v1 \
--model claude-haiku-4.5 --long-prefix --cache-control --rounds 6 \
--output /path/outside/repository/latency.json
```
This makes 21 billable requests by default. `--mock` uses synthetic local SSE
instead; it still needs OpenCode. `--baseline-no-drain` disables terminal draining
for a transport regression comparison. The harness records only hashes and
message sizes, not prompt bodies or credentials, and terminates its own isolated
OpenCode process. Keep raw results outside the repository because they include
the configured gateway URL. Other OpenCode versions may need API adjustments.
The next diagnostic for a reported slow production session should capture the
actual pipeline/Runner timing, streaming mode, prompt size, reasoning settings,
cache usage and gateway routing alongside these requester-level measurements.
docs/LLM_REQUEST_E2E.md
+108
View File
@@ -0,0 +1,108 @@
# Request-chain acceptance (2026-10-15)
The acceptance suite starts an isolated Core process, a real SDK Plugin Runtime
process, and a packaged Runner plugin. It installs the plugin and creates the
provider, model and Agent through the same authenticated HTTP APIs used by the
application. Agent Debug requests then traverse:
HTTP -> Core Agent orchestration -> Plugin Runtime -> plugin SDK -> Runtime ->
Core model manager/request policy -> LiteLLM/provider SDK -> HTTP/SSE provider.
Assertions cover plugin return values, captured provider requests, TCP disconnects
and persisted SQLite monitoring records. No requester, action transport, model
manager or monitoring implementation is replaced with a mock.
The remote provider is a deterministic local HTTP server. It injects failures and
delays reproducibly; these are real process/network integration tests, **not** a
claim that every external provider has been tested. See the separate
[latency benchmark](LLM_LATENCY_BENCHMARK.md) for the live NewAPI experiment.
## Matrix
Measured results on Windows / Python 3.12.7 / LiteLLM 1.88.1:
- Core with SDK companion commit `c25ff9b`: **31 passed**, strict stream closure,
including all four historical API cases (106.30 seconds).
- Core with the existing SDK 0.7.11rc5 environment: **27 passed, 4 XFAIL** for
the reproduced SDK cancellation defect (107.61 seconds).
- Published stable SDK 0.7.11 also reproduced the cancellation defect in the
preceding 23-case run. The old and stable wrappers were verified identical.
- Focused Core regressions: **371 passed** on Python 3.13.7, ResourceWarnings fatal.
- SDK API and stream flow-control regressions: **507 passed**.
- Existing real Runner discovery/orchestration/event-processor E2E: **3 passed**.
| Area | Cases and evidence |
| --- | --- |
| Public APIs | Runner and invocation-bound plugin APIs: message, message with usage, chunk stream and event stream; unchanged SDK entities and nested usage |
| Historical plugin APIs | Verbatim SDK 0.5.8 proxy and provider-message entities, all four APIs, ordinary calls without `run_id` |
| Request configuration | Persisted defaults, call overrides, OpenAI cache key/retention, OpenRouter numeric reasoning budget and Anthropic cache marker |
| Headers | Explicit application attribution and scoped conversation ID reach the provider |
| Multimodal/tools | Image data URL and function schema survive plugin serialization and provider conversion |
| Retry/failure | 503 then success, `num_retries=0`, `max_retries=0`, 401 without retry, first-event timeout, total timeout, idle timeout and connection loss after content |
| Recovery | Each fault followed by a successful request through the same running processes |
| Cancellation | Early close of Runner and plugin APIs, both chunk and event streams; provider socket closes and the next invocation succeeds |
| Reasoning | Reasoning content shown/hidden with the existing `remove_think` option; unsupported numeric budget rejected before HTTP |
| Monitoring | Full cache/reasoning usage, transport timing and structured errors persisted |
Historical API files come from SDK tag `0.5.8`, commit
`41e320b68e899a7b5e74cb48c5246f6ccbf9d73f`.
They use their own message serializers and the current Runtime transport. This
checks old plugin call/response compatibility; it is not a test of running the
entire old Runtime against the current Core.
## Defects exposed by the full chain
- The plugin bridge flattened model and call policy layers, allowing persisted
`request_policy.retries` to override a call's legacy zero-retry option. LLM
policy layers are now resolved separately; unrelated argument merging remains.
- Ordinary plugin calls without a Runner/query returned content but lost usage.
Core now creates a trusted, invocation-local query for accounting, using the
existing plugin-action query builder. No public payload/signature changes.
- The shipped OpenRouter requester uses OpenAI metadata. Numeric budgets selected
the wrong mapping, and LiteLLM's OpenAI chat adapter stripped cache markers.
OpenRouter chat now uses its own protocol adapter. Embedding/rerank routing is
unchanged, and a regression test reads the actual shipped requester YAML.
- Published SDK 0.7.11 leaves inner streaming generators unowned in several public
proxy wrappers. Core's direct TCP close test passes, but the complete plugin
chain does not promptly cancel its upstream request. The companion
[SDK PR #177](https://github.com/langbot-app/langbot-plugin-sdk/pull/177) adds
explicit generator ownership without changing signatures, permissions, action
payloads or error translation.
**Immediate remote cancellation requires SDK 0.7.12 or later.** Core now pins
published SDK 0.7.12, which contains the companion fix. Separately deployed
Runtime and plugin development environments also need the fixed SDK. The E2E
suite reports the known close failure as XFAIL only for SDK <=0.7.11. Set
`LLM_E2E_REQUIRE_STREAM_CLOSE=1` to require strict passing with the patched SDK;
newer SDK versions must pass without this override. Historical API checks are
explicitly skipped unless their released source files are supplied.
## Reproduction
```sh
python -m pytest tests/e2e/test_llm_request_chain.py -q --tb=short
```
For the historical API matrix, export the two unmodified files below from SDK tag
`0.5.8` into a temporary directory, named `legacy_api.py` and `legacy_message.py`,
then set `LLM_E2E_LEGACY_API_DIR` to its absolute path:
- `src/langbot_plugin/api/proxies/langbot_api.py`
- `src/langbot_plugin/api/entities/builtin/provider/message.py`
For strict cancellation acceptance, install the companion SDK into an isolated
test environment and set `LLM_E2E_REQUIRE_STREAM_CLOSE=1` before running the suite.
The existing CI E2E job discovers this test module automatically. Tests use random
localhost ports, test-only credentials and pytest temporary data directories;
they leave application data and running user instances untouched.
## Limits
The suite does not exercise IM delivery, browser rendering, external provider
authentication/billing, native Anthropic/Gemini network services, Codex OAuth, or
complete third-party Runner business logic. It validates multimodal/tool request
transport rather than image understanding or actual tool execution. Native Codex,
retry, reasoning, transport and plugin regressions also run in the focused Core
test suite. A Python 3.13 Windows full-app attempt timed out during startup before
any LLM request; full process acceptance therefore uses Python 3.12.7, while the
focused Core regressions run on Python 3.13.7.
docs/LLM_REQUEST_POLICY.md
+256
View File
@@ -0,0 +1,256 @@
# LLM request policy and diagnostics
The Core owns chat-call deadlines and retry decisions. The existing SDK
`Message`, `MessageChunk`, invoke-with-usage and stream-usage shapes are unchanged.
No SDK or Runner upgrade is required. This policy does not add automatic tool
replay, model fallback, key rotation, or embedding/rerank retries.
The public API remains compatible with existing plugins. Full-chain testing
identified a pre-existing SDK <=0.7.11 early-close defect: prompt remote
cancellation through public plugin/Runner wrappers requires SDK 0.7.12 or later,
which includes the [SDK fix](https://github.com/langbot-app/langbot-plugin-sdk/pull/177).
Core now pins 0.7.12; separately deployed Runtime and plugin development
environments must also install a version containing the fix.
See [E2E acceptance results](LLM_REQUEST_E2E.md) for tested versions and scope.
## Configuration
Instance defaults live in `data/config.yaml` under `llm_request_policy`; restart
Core after editing them. Persist model-specific overrides in the existing model
advanced parameters (`extra_args`). Runners can override them through their
existing call-level `extra_args` argument:
```json
{
"request_policy": {
"retries": 1,
"total_timeout": 180,
"first_event_timeout": 60,
"stream_idle_timeout": 45,
"connect_timeout": 10,
"read_timeout": 120,
"write_timeout": 30,
"pool_timeout": 10,
"backoff_base": 0.5,
"backoff_max": 8,
"attribution": true,
"conversation_header": null
}
}
```
Precedence is instance/requester defaults, model extra arguments, then call extra
arguments. Within each layer, `request_policy.retries` overrides legacy
`num_retries` / `max_retries`. Zero disables retries. Legacy numeric `timeout`
sets read, first-event and idle timeouts; it does not extend the total deadline.
Timeouts are positive finite seconds; retries must be an integer from 0 to 5.
Unknown policy keys are rejected. Policy keys are removed before calling upstream.
Default retries are **zero**, total timeout is 300 seconds, and first-event and
idle timeouts are 120 seconds. Retry backoff uses exponential full jitter;
`Retry-After` seconds or HTTP dates are respected only within the remaining total
budget. Quota, authentication and invalid-request failures are not retried.
Core disables both LiteLLM and provider-SDK retries for managed calls. A gateway
may still retry internally; configure its budget separately.
Retries stop after any raw upstream event (including reasoning/control/usage)
or, for legacy adapters without raw instrumentation, the first yielded chunk.
Cancellation is propagated and the owned stream is closed. Total deadlines also
include downstream consumption time, but Core does not cancel a consumer while
the generator is suspended at a yield; expiration is enforced on the next pull.
Transport connect/read/write/pool limits apply to LiteLLM and native Codex HTTP.
Legacy third-party requesters receive Host-level deadlines but remain responsible
for their internal transport configuration and internal retries.
## Diagnostics and usage
HTTP monitoring and `get_monitoring_records(kind="llm_calls")` return additive
`details` containing `timing`, complete provider `usage`, and structured `error`.
Older records have null details. Existing top-level token counters remain intact.
Migration `0038_llm_call_details` adds a nullable JSON column and preserves rows.
Timing offsets use a high-resolution monotonic clock, in milliseconds from invocation start:
| Field | Meaning |
| --- | --- |
| `request_prepared_ms` | First attempt finished building its request |
| `request_started_ms` | First adapter call into the network SDK; not a socket-level timestamp |
| `response_open_ms` | SDK returned a response/stream object |
| `first_event_ms` | First parsed upstream stream event, including reasoning/control |
| `first_reasoning_ms` | First upstream reasoning delta when supported |
| `first_content_ms` | First upstream content delta; inline think tags may still be content |
| `first_visible_chunk_ms` | First nonempty normalized chunk; may include displayed reasoning |
| `response_completed_ms` | Host finished consuming the response |
| `total_ms` | Invocation duration, excluding the subsequent monitoring write |
| `attempts`, `retry_wait_ms` | Host attempts and total backoff wait |
| `preparation_ms` | Cumulative LiteLLM request-building time across attempts |
| `transport_setup_ms` | Cumulative managed HTTP pool acquisition, including first-use TLS setup |
| `preparation_cache_hits` | Message preparation reused within a retrying invocation |
Missing fields mean not observed or not instrumented, not zero latency. The Core
does not label these measurements as network-only TTFT. Usage preserves nested
cache and reasoning counts, costs and provider-specific counters when supplied;
missing upstream accounting cannot be reconstructed. Partial usage received
before cancellation/failure is retained. No fabricated usage event is sent early.
`RequesterError` retains its legacy string representation and adds `code`,
`status_code`, `retryable`, `retry_after`, `request_id`, `phase`, and `as_dict()`.
Native Codex errors retain their existing ValueError/HTTP error contract.
Legacy SDK error transport remains unchanged; structured diagnostics are available
in Core monitoring. The metadata never copies upstream error bodies or headers
wholesale. Requester-specific raw error strings are not included by the LiteLLM
error translator.
## Application and conversation headers
Direct OpenRouter requests automatically receive:
```text
HTTP-Referer: https://langbot.app
X-OpenRouter-Title: LangBot
x-session-id: lb-<opaque scoped hash>
```
Existing headers in model/call `extra_headers` win, case-insensitively. Legacy
`X-Title` is respected. Set `attribution: false` to disable automatic application
attribution (explicit headers remain untouched). Detection uses the exact
`openrouter.ai` hostname, or an OpenRouter requester without an overridden base
URL; arbitrary gateway hosts do not inherit OpenRouter identity headers.
The conversation hash includes Workspace, bot, pipeline and session identity;
it is stable across turns and differs across those scopes. It is pseudonymous,
not a claim of cryptographic anonymization. Calls without conversation context
omit generated conversation headers. No prompt, raw user ID or group ID is sent.
`x-langbot-conversation-id` is a custom convention, **not a documented universal
provider standard**. OpenRouter documents `x-session-id` (or body `session_id`,
which takes precedence). LiteLLM Proxy documents `x-litellm-session-id` and
vendor session headers. Configure `conversation_header` explicitly for gateways:
```json
{
"request_policy": {"conversation_header": "x-litellm-session-id"},
"extra_headers": {
"HTTP-Referer": "https://langbot.app",
"X-OpenRouter-Title": "LangBot"
}
}
```
Supported generated header names are `x-session-id`, `x-litellm-session-id` and
`x-langbot-conversation-id`. The latter is intended for a gateway configured to
consume that custom header; do not expect automatic routing or caching benefits.
New API may supply its own OpenRouter attribution and may filter/override client
headers. Configure explicit channel forwarding/override rules and verify the
outgoing request. Sending a header to New API is not proof OpenRouter received it.
The local `langbot-new-api` snapshot inspected during implementation overwrites
attribution with New API; upstream main uses defaults when its outgoing headers
are empty. This change does not modify or deploy either gateway.
Sources checked on 2026-10-15:
- https://openrouter.ai/docs/app-attribution
- https://openrouter.ai/docs/guides/best-practices/prompt-caching
- https://docs.litellm.ai/docs/proxy/request_headers
- https://docs.litellm.ai/docs/routing
- https://github.com/QuantumNous/new-api/blob/main/relay/channel/openai/adaptor.go
- https://github.com/QuantumNous/new-api/blob/main/relay/channel/api_request.go
## Connection and preparation reuse
Initialized OpenAI-compatible, OpenRouter and native Codex requesters own a lazy
HTTP pool per runtime Provider. Keys, headers and timeout budgets stay request-local.
The pool holds at most 100 connections, with 20 idle keepalive connections and a
30-second idle expiry. Provider retirement rejects new leases and closes its pool
after existing streams finish or are cancelled. Native LiteLLM providers retain
their own transport handling; custom LiteLLM transport/TLS settings, explicit caller clients and redirected model routes
are preserved. Bare OpenAI names on OpenRouter retain LiteLLM's alternate routing;
use OpenRouter's fully qualified names such as `openai/gpt-4o`.
Each pool loads its TLS context once in a worker thread, preserving HTTPX
certificate and proxy environment settings. Concurrent initialization is serialized;
cancellation or retirement during certificate loading creates no HTTP client.
OpenAI-compatible streams drain HTTP framing after the SSE `[DONE]` record
(at most 50 ms / 64 KiB) so the SDK can return healthy connections to the pool.
Early close and cancellation never trigger this drain; late usage preceding
`[DONE]` remains available. OpenRouter already consumes EOF in its generic adapter.
Codex briefly drains the HTTP tail after a completion event (at most 50 ms) to
release a healthy connection for reuse. Incomplete/cancelled streams are closed.
Capability lookups use a bounded requester-local cache (128 entries, five minutes).
Message/tool conversion is reused only within an invocation that permits retries,
with independent copies for each attempt. No conversation content is cached across
invocations. Token counting runs in the application's executor instead of blocking
the event loop. These reduce local overhead; they do not guarantee faster upstream
inference or a cache hit.
## Unified thinking budgets
Existing effort levels and Runner/SDK signatures remain unchanged. Model
`reasoning_config` accepts an optional positive integer `budget_tokens`:
```json
{"reasoning_config": {"level": "enabled", "budget_tokens": 2048}}
```
Old SDKs can send the same object inside their existing call `extra_args`.
Precedence is persisted model configuration, model `extra_args.reasoning_config`,
call `extra_args.reasoning_config`, then the existing Runner override. Each object
replaces the previous one; a Runner effort override clears a lower-priority budget.
Omitting `level` with a budget means `enabled`. Budgets cannot be combined with an
effort level, `disabled`, or native thinking arguments. Unsupported budgets fail
before the network call; they are never silently converted into an effort level.
Numeric budgets map to Anthropic `thinking.budget_tokens`, Gemini 2.5 thinking
configuration, supported Qwen `thinking_budget`, or OpenRouter
`reasoning.max_tokens`. Direct OpenAI/Codex and adaptive-only Claude models keep
effort-based control. Query `reasoning_capabilities.budget_tokens` for the Core's
accepted bounds; upstream models/gateways may impose additional limits. OpenRouter
may translate numeric budgets to effort for models without native numeric support.
Direct Anthropic calls require an explicit `max_tokens` greater than the budget;
Core never silently raises the output limit. Existing UI effort selectors are
unchanged; numeric budgets can be configured through the model API or extra args.
## Prompt caching
Default behavior remains provider-controlled. `extra_args.prompt_caching` is a
Core-only object and is stripped before sending the request. Call-level objects
replace model-level objects. There is no local generated-response cache.
```json
{"prompt_caching": {"mode": "openai", "retention": "in_memory"}}
```
OpenAI mode adds an opaque conversation-and-model-scoped `prompt_cache_key` when
conversation context exists. `retention` is optional (`in_memory` or `24h`) and
must be supported by the destination model/gateway. Native explicit
`prompt_cache_key`, `prompt_cache_retention`, and newer `prompt_cache_options`
take precedence. Routing keys are hints, not cache-hit guarantees; current models
may route caching automatically or use newer native options instead. This mode
supports OpenAI-compatible/OpenRouter OpenAI models; forwarding through gateways
must be verified separately.
```json
{"prompt_caching": {"mode": "anthropic", "ttl": "5m"}}
```
Anthropic mode supports native Claude and OpenRouter Claude. It adds one ephemeral
breakpoint to the last text block in the leading system/developer prefix, without
reordering tools/history or marking user messages. TTL is `5m` (default) or `1h`.
If any native cache-control marker exists, Core preserves it without adding more.
With no eligible prefix, no marker is inserted. Provider minimum prefix lengths,
TTL pricing and eligibility still apply. Explicit `provider_default` adds no fields.
Codex's restricted Responses body does not accept this caching object; its existing
server-managed caching behavior is preserved. Gemini explicit cache resource
creation/lifecycle and general Responses migration are outside this batch.
References checked on 2026-10-15:
- https://developers.openai.com/api/docs/guides/prompt-caching
- https://platform.claude.com/docs/en/build-with-claude/prompt-caching
- https://platform.claude.com/docs/en/build-with-claude/extended-thinking
- https://ai.google.dev/gemini-api/docs/generate-content/thinking
- https://openrouter.ai/docs/guides/best-practices/reasoning-tokens
See [the latency benchmark](LLM_LATENCY_BENCHMARK.md) for measured results and limitations.
docs/llm-latency-2026-10-15.json
+411
View File
@@ -0,0 +1,411 @@
{
"model": "claude-haiku-4.5",
"litellm": "1.88.1",
"opencode": "2.0.16",
"canonical_sha256": "7834a66a188e70b4eb1b6234b3a14be4c66b8c72dd8a2c06395fe2505c1c371b",
"canonical_bytes": 31785,
"summary": {
"direct": {
"samples": 6,
"dispatch_to_wire_ms": {
"median": 3.952,
"max": 5.44
},
"upstream_ttft_ms": {
"median": 1115.744,
"max": 1158.609
},
"ttft_ms": {
"median": 1119.541,
"max": 1164.048
},
"total_ms": {
"median": 2199.413,
"max": 2219.7
},
"cache_read_tokens": [
5987,
5987,
5987,
5987,
5987,
5987
]
},
"langbot": {
"samples": 6,
"dispatch_to_wire_ms": {
"median": 6.802,
"max": 10.028
},
"upstream_ttft_ms": {
"median": 1079.658,
"max": 1113.394
},
"ttft_ms": {
"median": 1089.092,
"max": 1120.475
},
"total_ms": {
"median": 2135.935,
"max": 2175.882
},
"preparation_ms": {
"median": 0.316,
"max": 0.497
},
"cache_read_tokens": [
5987,
5987,
5987,
5987,
5987,
5987
]
},
"opencode": {
"samples": 6,
"dispatch_to_wire_ms": {
"median": 41.158,
"max": 50.283
},
"upstream_ttft_ms": {
"median": 1063.875,
"max": 1158.768
},
"ttft_ms": {
"median": 1104.02,
"max": 1201.731
},
"total_ms": {
"median": 2148.265,
"max": 2244.956
},
"cache_read_tokens": [
5987,
5987,
5987,
5987,
5987,
5987
]
}
},
"records": [
{
"client": "direct",
"round": -1,
"warmup": true,
"dispatch_to_wire_ms": 3.500899998471141,
"upstream_ttft_ms": 2094.1667999722995,
"ttft_ms": 2097.6676999707706,
"total_ms": 3145.2215999597684,
"client_socket_reused": false,
"consumer_first_content_ms": 2098.3635999727994,
"consumer_total_ms": 3146.0422999807633,
"input_tokens": 5990,
"output_tokens": 109,
"cache_read_tokens": 5987
},
{
"client": "langbot",
"round": -1,
"warmup": true,
"dispatch_to_wire_ms": 1452.2201999789104,
"upstream_ttft_ms": 1122.500299999956,
"ttft_ms": 2574.7204999788664,
"total_ms": 3640.2128000045195,
"preparation_ms": 0.373,
"transport_setup_ms": 441.243,
"client_socket_reused": false,
"consumer_first_content_ms": 2576.894700003322,
"consumer_total_ms": 3646.833000006154,
"input_tokens": 5990,
"output_tokens": 109,
"cache_read_tokens": 5987
},
{
"client": "opencode",
"round": -1,
"warmup": true,
"dispatch_to_wire_ms": 846.3009999832138,
"upstream_ttft_ms": 1082.7391999773681,
"ttft_ms": 1929.040199960582,
"total_ms": 3569.3477999884635,
"client_socket_reused": false,
"input_tokens": 5990,
"output_tokens": 109,
"cache_read_tokens": 5987
},
{
"client": "direct",
"round": 0,
"warmup": false,
"dispatch_to_wire_ms": 3.5812999703921378,
"upstream_ttft_ms": 1012.8218000172637,
"ttft_ms": 1016.4030999876559,
"total_ms": 2067.0370999723673,
"client_socket_reused": false,
"consumer_first_content_ms": 1017.0787000097334,
"consumer_total_ms": 2067.508200008888,
"input_tokens": 5990,
"output_tokens": 109,
"cache_read_tokens": 5987
},
{
"client": "langbot",
"round": 0,
"warmup": false,
"dispatch_to_wire_ms": 6.439700024202466,
"upstream_ttft_ms": 1067.0202999608591,
"ttft_ms": 1073.4599999850616,
"total_ms": 2151.4558999915607,
"preparation_ms": 0.281,
"transport_setup_ms": 0.017,
"client_socket_reused": true,
"consumer_first_content_ms": 1074.6371999848634,
"consumer_total_ms": 2156.441500002984,
"input_tokens": 5990,
"output_tokens": 109,
"cache_read_tokens": 5987
},
{
"client": "opencode",
"round": 0,
"warmup": false,
"dispatch_to_wire_ms": 50.28290004702285,
"upstream_ttft_ms": 1151.4483999926597,
"ttft_ms": 1201.7313000396825,
"total_ms": 2244.956100010313,
"client_socket_reused": false,
"input_tokens": 5990,
"output_tokens": 109,
"cache_read_tokens": 5987
},
{
"client": "direct",
"round": 1,
"warmup": false,
"dispatch_to_wire_ms": 3.269899985753,
"upstream_ttft_ms": 1149.7626000200398,
"ttft_ms": 1153.0325000057928,
"total_ms": 2198.293499997817,
"client_socket_reused": true,
"consumer_first_content_ms": 1153.3279999857768,
"consumer_total_ms": 2198.7193999812007,
"input_tokens": 5990,
"output_tokens": 109,
"cache_read_tokens": 5987
},
{
"client": "opencode",
"round": 1,
"warmup": false,
"dispatch_to_wire_ms": 40.65959999570623,
"upstream_ttft_ms": 1158.767800021451,
"ttft_ms": 1199.4274000171572,
"total_ms": 2242.273700016085,
"client_socket_reused": false,
"input_tokens": 5990,
"output_tokens": 109,
"cache_read_tokens": 5987
},
{
"client": "langbot",
"round": 1,
"warmup": false,
"dispatch_to_wire_ms": 6.523299962282181,
"upstream_ttft_ms": 1070.9955000202172,
"ttft_ms": 1077.5187999824993,
"total_ms": 2109.3437999952585,
"preparation_ms": 0.497,
"transport_setup_ms": 0.015,
"client_socket_reused": true,
"consumer_first_content_ms": 1081.5072999685071,
"consumer_total_ms": 2113.767499977257,
"input_tokens": 5990,
"output_tokens": 109,
"cache_read_tokens": 5987
},
{
"client": "langbot",
"round": 2,
"warmup": false,
"dispatch_to_wire_ms": 7.08060001488775,
"upstream_ttft_ms": 1113.394399988465,
"ttft_ms": 1120.4750000033528,
"total_ms": 2175.882300012745,
"preparation_ms": 0.221,
"transport_setup_ms": 0.009,
"client_socket_reused": true,
"consumer_first_content_ms": 1122.9100999771617,
"consumer_total_ms": 2187.183100031689,
"input_tokens": 5990,
"output_tokens": 109,
"cache_read_tokens": 5987
},
{
"client": "direct",
"round": 2,
"warmup": false,
"dispatch_to_wire_ms": 5.439500033389777,
"upstream_ttft_ms": 1158.6087999748997,
"ttft_ms": 1164.0483000082895,
"total_ms": 2219.7004000190645,
"client_socket_reused": false,
"consumer_first_content_ms": 1164.5996000152081,
"consumer_total_ms": 2220.197000016924,
"input_tokens": 5990,
"output_tokens": 109,
"cache_read_tokens": 5987
},
{
"client": "opencode",
"round": 2,
"warmup": false,
"dispatch_to_wire_ms": 41.65710002416745,
"upstream_ttft_ms": 1063.7803999707103,
"ttft_ms": 1105.4374999948777,
"total_ms": 2145.6969000282697,
"client_socket_reused": false,
"input_tokens": 5990,
"output_tokens": 109,
"cache_read_tokens": 5987
},
{
"client": "langbot",
"round": 3,
"warmup": false,
"dispatch_to_wire_ms": 10.027599986642599,
"upstream_ttft_ms": 1080.2179000456817,
"ttft_ms": 1090.2455000323243,
"total_ms": 2105.4020000156015,
"preparation_ms": 0.406,
"transport_setup_ms": 0.012,
"client_socket_reused": true,
"consumer_first_content_ms": 1092.152299999725,
"consumer_total_ms": 2109.886699996423,
"input_tokens": 5990,
"output_tokens": 109,
"cache_read_tokens": 5987
},
{
"client": "opencode",
"round": 3,
"warmup": false,
"dispatch_to_wire_ms": 38.63289998844266,
"upstream_ttft_ms": 1063.9694000128657,
"ttft_ms": 1102.6023000013083,
"total_ms": 2139.4490000093356,
"client_socket_reused": false,
"input_tokens": 5990,
"output_tokens": 109,
"cache_read_tokens": 5987
},
{
"client": "direct",
"round": 3,
"warmup": false,
"dispatch_to_wire_ms": 4.3889000080525875,
"upstream_ttft_ms": 1157.878800004255,
"ttft_ms": 1162.2677000123076,
"total_ms": 2214.4845000002533,
"client_socket_reused": false,
"consumer_first_content_ms": 1162.9892000346445,
"consumer_total_ms": 2214.923900028225,
"input_tokens": 5990,
"output_tokens": 109,
"cache_read_tokens": 5987
},
{
"client": "opencode",
"round": 4,
"warmup": false,
"dispatch_to_wire_ms": 34.960800025146455,
"upstream_ttft_ms": 1052.52319999272,
"ttft_ms": 1087.4840000178665,
"total_ms": 2137.3740999842994,
"client_socket_reused": true,
"input_tokens": 5990,
"output_tokens": 109,
"cache_read_tokens": 5987
},
{
"client": "direct",
"round": 4,
"warmup": false,
"dispatch_to_wire_ms": 2.0186000037938356,
"upstream_ttft_ms": 987.3963000136428,
"ttft_ms": 989.4149000174366,
"total_ms": 2200.532699993346,
"client_socket_reused": true,
"consumer_first_content_ms": 989.9038000148721,
"consumer_total_ms": 2201.0204999824055,
"input_tokens": 5990,
"output_tokens": 109,
"cache_read_tokens": 5987
},
{
"client": "langbot",
"round": 4,
"warmup": false,
"dispatch_to_wire_ms": 6.363800028339028,
"upstream_ttft_ms": 1112.4532999820076,
"ttft_ms": 1118.8171000103466,
"total_ms": 2150.362300046254,
"preparation_ms": 0.325,
"transport_setup_ms": 0.011,
"client_socket_reused": true,
"consumer_first_content_ms": 1120.592100021895,
"consumer_total_ms": 2154.264000011608,
"input_tokens": 5990,
"output_tokens": 109,
"cache_read_tokens": 5987
},
{
"client": "opencode",
"round": 5,
"warmup": false,
"dispatch_to_wire_ms": 41.722599999047816,
"upstream_ttft_ms": 1059.4015000388026,
"ttft_ms": 1101.1241000378504,
"total_ms": 2150.8331000222825,
"client_socket_reused": false,
"input_tokens": 5990,
"output_tokens": 109,
"cache_read_tokens": 5987
},
{
"client": "langbot",
"round": 5,
"warmup": false,
"dispatch_to_wire_ms": 8.840300026349723,
"upstream_ttft_ms": 1079.098400019575,
"ttft_ms": 1087.9387000459246,
"total_ms": 2121.5078000095673,
"preparation_ms": 0.307,
"transport_setup_ms": 0.021,
"client_socket_reused": true,
"consumer_first_content_ms": 1091.6245000553317,
"consumer_total_ms": 2126.561200013384,
"input_tokens": 5990,
"output_tokens": 109,
"cache_read_tokens": 5987
},
{
"client": "direct",
"round": 5,
"warmup": false,
"dispatch_to_wire_ms": 4.323699977248907,
"upstream_ttft_ms": 1081.7263000062667,
"ttft_ms": 1086.0499999835156,
"total_ms": 2126.9287999602966,
"client_socket_reused": false,
"consumer_first_content_ms": 1086.7821999709122,
"consumer_total_ms": 2127.6911000022665,
"input_tokens": 5990,
"output_tokens": 109,
"cache_read_tokens": 5987
}
]
}
pyproject.toml
+1 -1
View File
@@ -71,7 +71,7 @@ dependencies = [
"langchain-text-splitters>=1.1.2",
"chromadb>=1.0.0,<2.0.0",
"qdrant-client (>=1.15.1,<2.0.0)",
"langbot-plugin==0.7.11",
"langbot-plugin==0.7.12",
"asyncpg>=0.30.0",
"line-bot-sdk>=3.19.0",
"matrix-nio>=0.25.2",
scripts/llm_latency_probe.py
+418
View File
@@ -0,0 +1,418 @@
"""Compare a direct stream, Core requester and isolated OpenCode 2.x server.
All three clients use a local measurement relay. The relay sends the same
synthetic request bytes and headers upstream, deliberately removing agent prompt
differences. This measures the request layer, not a production bot pipeline/UI.
Credentials are read from an environment variable or an untracked local file.
"""
import argparse
import asyncio
import contextlib
import hashlib
import itertools
import json
import os
from pathlib import Path
import socket
import statistics
import secrets
import subprocess
import tempfile
import time
from types import SimpleNamespace
import httpx
from aiohttp import web
def digest(value):
return hashlib.sha256(json.dumps(value, sort_keys=True, separators=(',', ':')).encode()).hexdigest()
class SSEMeter:
"""Incrementally parse SSE without assuming network chunk boundaries."""
def __init__(self):
self.buffer = b''
self.first_event = None
self.first_content = None
self.usage = None
self.content_chars = 0
self.done = False
def feed(self, data, now):
self.buffer += data
while b'\n' in self.buffer:
line, self.buffer = self.buffer.split(b'\n', 1)
if not line.startswith(b'data:'):
continue
raw = line[5:].strip()
if raw == b'[DONE]':
self.done = True
continue
item = json.loads(raw)
if self.first_event is None:
self.first_event = now
if item.get('usage'):
self.usage = item['usage']
for choice in item.get('choices', []):
content = (choice.get('delta') or {}).get('content')
if content:
if self.first_content is None:
self.first_content = now
self.content_chars += len(content)
def cache_tokens(usage):
if not isinstance(usage, dict):
return None
detail = usage.get('prompt_tokens_details') or usage.get('input_tokens_details') or {}
return detail.get('cached_tokens', usage.get('cache_read_input_tokens'))
def summary(records):
result = {}
for name in sorted({row['client'] for row in records}):
rows = [r for r in records if r['client'] == name and not r.get('warmup')]
stats = {'samples': len(rows)}
for key in ('dispatch_to_wire_ms', 'upstream_ttft_ms', 'ttft_ms', 'total_ms', 'preparation_ms'):
values = [r[key] for r in rows if r.get(key) is not None]
if values:
stats[key] = {'median': round(statistics.median(values), 3), 'max': round(max(values), 3)}
stats['cache_read_tokens'] = [cache_tokens(r.get('usage')) for r in rows]
result[name] = stats
return result
def free_port():
with socket.socket() as sock:
sock.bind(('127.0.0.1', 0))
return sock.getsockname()[1]
async def run(options):
# Import/initialization cost is reported separately from resident calls.
started = time.perf_counter()
from langbot.pkg.provider.modelmgr.requesters.litellmchat import LiteLLMRequester
from langbot.pkg.provider.modelmgr.request_policy import CallState, RequestPolicy, current_call
from langbot_plugin.api.entities.builtin.provider.message import Message
import importlib.metadata
import_ms = (time.perf_counter() - started) * 1000
credentials = json.loads(Path(options.credentials_file).read_text()) if options.credentials_file else {}
base_url = options.base_url or credentials.get('base_url')
key = os.environ.get(options.api_key_env) or credentials.get('api_key')
if not options.mock and not (base_url and key):
raise ValueError('A base URL and API key are required for live mode')
prefix = 'You are a concise technical assistant. Follow only the final question.\n'
if options.long_prefix:
prefix += ''.join(
f'Reference entry {i}: HTTP connection pooling reuses established transports. '
'Keep authentication and timeouts request scoped. Measure first content separately from reasoning.\n'
for i in range(180)
)
canonical = {
'model': options.model,
'stream': True,
'stream_options': {'include_usage': True},
'temperature': 0,
'max_tokens': 256,
'messages': [
{
'role': 'system',
'content': [
{
'type': 'text',
'text': prefix,
**({'cache_control': {'type': 'ephemeral'}} if options.cache_control else {}),
}
],
},
{'role': 'user', 'content': 'Explain HTTP connection pooling in exactly five short numbered sentences.'},
],
}
payload = json.dumps(canonical, separators=(',', ':')).encode()
records = []
active = None
upstream = httpx.AsyncClient(timeout=90)
relay_sockets = {}
async def relay(request):
nonlocal active
if active is None or active.get('wire_started') is not None:
return web.json_response({'error': 'Unexpected extra request'}, status=409)
row = active
finished = row['relay_finished']
row['wire_started'] = time.perf_counter()
incoming = await request.json()
row['received_body_hash'] = digest(incoming)
row['received_fields'] = sorted(incoming)
row['received_message_chars'] = len(json.dumps(incoming.get('messages', [])))
row['upstream_body_hash'] = digest(canonical)
row['dispatch_to_wire_ms'] = (row['wire_started'] - row['start']) * 1000
connection = id(request.transport)
row['client_socket_reused'] = connection in relay_sockets
relay_sockets[connection] = request.transport
meter = SSEMeter()
response = web.StreamResponse(headers={'Content-Type': 'text/event-stream'})
try:
if options.mock:
await response.prepare(request)
await asyncio.sleep(0.05)
for i in range(20):
event = {
'id': 'synthetic',
'object': 'chat.completion.chunk',
'created': 1,
'model': options.model,
'choices': [{'index': 0, 'delta': {'content': 'word '}, 'finish_reason': None}],
}
data = ('data: ' + json.dumps(event) + '\n\n').encode()
meter.feed(data, time.perf_counter())
await response.write(data)
await asyncio.sleep(0.03)
event['choices'] = [{'index': 0, 'delta': {}, 'finish_reason': 'stop'}]
event['usage'] = {'prompt_tokens': 50, 'completion_tokens': 20, 'total_tokens': 70}
data = ('data: ' + json.dumps(event) + '\n\ndata: [DONE]\n\n').encode()
meter.feed(data, time.perf_counter())
await response.write(data)
else:
async with upstream.stream(
'POST',
base_url.rstrip('/') + '/chat/completions',
content=payload,
headers={'Authorization': 'Bearer ' + key, 'Content-Type': 'application/json'},
) as remote:
row['status'] = remote.status_code
if remote.status_code != 200:
await remote.aread()
# Do not log upstream response bodies, headers or credentials.
return web.json_response(
{'error': f'Upstream HTTP {remote.status_code}'}, status=remote.status_code
)
row['upstream_headers_ms'] = (time.perf_counter() - row['wire_started']) * 1000
await response.prepare(request)
async for chunk in remote.aiter_bytes():
meter.feed(chunk, time.perf_counter())
try:
await response.write(chunk)
except ConnectionResetError:
row['client_closed_before_http_eof'] = True
row['upstream_ttft_ms'] = (
(meter.first_content - row['wire_started']) * 1000 if meter.first_content is not None else None
)
row['ttft_ms'] = (meter.first_content - row['start']) * 1000 if meter.first_content is not None else None
row['usage'] = meter.usage
row['content_chars'] = meter.content_chars
row['upstream_complete_ms'] = (time.perf_counter() - row['wire_started']) * 1000
with contextlib.suppress(ConnectionResetError):
await response.write_eof()
return response
finally:
finished.set()
app = web.Application()
app.router.add_post('/v1/chat/completions', relay)
runner = web.AppRunner(app)
await runner.setup()
site = web.TCPSite(runner, '127.0.0.1', 0)
await site.start()
local = f'http://127.0.0.1:{site._server.sockets[0].getsockname()[1]}/v1'
adapter = LiteLLMRequester(
SimpleNamespace(),
{'custom_llm_provider': 'openai', 'base_url': local, 'requester_name': 'new-api-chat-completions'},
)
await adapter.initialize()
if options.baseline_no_drain:
adapter._http_pool.drain_sse = False
model = SimpleNamespace(
model_entity=SimpleNamespace(name=options.model, extra_args={}, reasoning_config=None, abilities=[]),
reasoning_config_override=None,
provider=SimpleNamespace(
token_mgr=SimpleNamespace(get_token=lambda: 'local-benchmark'),
provider_entity=SimpleNamespace(requester='new-api-chat-completions'),
),
)
messages = [Message(role=x['role'], content=x['content']) for x in canonical['messages']]
process = None
with tempfile.TemporaryDirectory(prefix='langbot-llm-benchmark-') as directory:
root = Path(directory)
config = {
'model': f'bench/{options.model}',
'small_model': f'bench/{options.model}',
'autoupdate': False,
'provider': {
'bench': {
'npm': '@ai-sdk/openai-compatible',
'name': 'Benchmark',
'options': {'baseURL': local, 'apiKey': 'local-benchmark'},
'models': {options.model: {'name': 'Benchmark', 'limit': {'context': 200000, 'output': 256}}},
}
},
'agent': {'bench': {'mode': 'primary', 'prompt': prefix, 'tools': {'*': False}, 'steps': 2}},
'permission': {'*': 'deny'},
}
config_path = root / 'opencode.json'
config_path.write_text(json.dumps(config))
env = os.environ.copy()
server_password = secrets.token_urlsafe(24)
env.update(
{
'OPENCODE_CONFIG': str(config_path),
'XDG_CONFIG_HOME': str(root / 'config'),
'XDG_DATA_HOME': str(root / 'data'),
'XDG_CACHE_HOME': str(root / 'cache'),
'NO_PROXY': os.environ.get('NO_PROXY', ''),
'OPENCODE_SERVER_PASSWORD': server_password,
}
)
port = free_port()
server = f'http://127.0.0.1:{port}'
log = (root / 'server.log').open('w')
try:
start_server = time.perf_counter()
process = await asyncio.create_subprocess_exec(
options.opencode,
'serve',
'--hostname',
'127.0.0.1',
'--port',
str(port),
cwd=root,
env=env,
stdout=log,
stderr=log,
creationflags=subprocess.CREATE_NO_WINDOW if os.name == 'nt' else 0,
)
async with httpx.AsyncClient(trust_env=False, timeout=100, auth=('opencode', server_password)) as client:
for _ in range(120):
try:
response = await client.get(server + '/api/info')
if response.is_success:
break
except httpx.ConnectError:
pass
await asyncio.sleep(0.25)
else:
raise RuntimeError('Isolated OpenCode server did not start')
startup_ms = (time.perf_counter() - start_server) * 1000
orders = list(itertools.permutations(('direct', 'langbot', 'opencode')))
for round_index in range(-1, options.rounds):
for name in orders[max(0, round_index) % len(orders)]:
session_id = None
if name == 'opencode':
response = await client.post(
server + '/api/session',
json={
'title': 'LLM benchmark',
'agent': 'bench',
'model': {'id': options.model, 'providerID': 'bench'},
'location': {'directory': str(root)},
},
)
response.raise_for_status()
session_id = response.json()['data']['id']
row = {
'client': name,
'round': round_index,
'warmup': round_index < 0,
'start': time.perf_counter(),
'relay_finished': asyncio.Event(),
}
active = row
if name == 'direct':
meter = SSEMeter()
async with client.stream('POST', local + '/chat/completions', json=canonical) as response:
response.raise_for_status()
async for chunk in response.aiter_bytes():
meter.feed(chunk, time.perf_counter())
row['consumer_first_content_ms'] = (meter.first_content - row['start']) * 1000
elif name == 'langbot':
state = CallState(RequestPolicy(retries=0, total_timeout=90))
token = current_call.set(state)
try:
async with contextlib.aclosing(
adapter.invoke_llm_stream(
None, model, messages, extra_args={'temperature': 0, 'max_tokens': 256}
)
) as stream:
async for chunk in stream:
if chunk.content and 'consumer_first_content_ms' not in row:
row['consumer_first_content_ms'] = (
time.perf_counter() - row['start']
) * 1000
row['preparation_ms'] = state.metrics.get('preparation_ms')
row['transport_setup_ms'] = state.metrics.get('transport_setup_ms')
row['client_usage'] = state.usage
finally:
current_call.reset(token)
else:
response = await client.post(
server + f'/api/session/{session_id}/prompt',
json={'text': canonical['messages'][-1]['content']},
)
response.raise_for_status()
await asyncio.wait_for(row['relay_finished'].wait(), 100)
if name != 'opencode':
row['consumer_total_ms'] = (time.perf_counter() - row['start']) * 1000
await asyncio.wait_for(row['relay_finished'].wait(), 100)
if 'upstream_complete_ms' not in row:
raise RuntimeError(f'Incomplete benchmark response (HTTP {row.get("status", "unknown")})')
# Use one completion boundary for all clients. This
# relay cannot measure OpenCode's UI rendering latency.
row['total_ms'] = row['dispatch_to_wire_ms'] + row['upstream_complete_ms']
row.pop('relay_finished')
row.pop('start')
row.pop('wire_started', None)
records.append(row)
active = None
print(json.dumps(row), flush=True)
await asyncio.sleep(0.1)
report = {
'model': options.model,
'base_url': base_url,
'mock': options.mock,
'long_prefix': options.long_prefix,
'cache_control': options.cache_control,
'baseline_no_drain': options.baseline_no_drain,
'canonical_sha256': digest(canonical),
'canonical_bytes': len(payload),
'python_import_ms': import_ms,
'opencode_startup_ms': startup_ms,
'litellm': importlib.metadata.version('litellm'),
'opencode': options.opencode_version,
'records': records,
'summary': summary(records),
}
Path(options.output).write_text(json.dumps(report, indent=2), encoding='utf-8')
print('SUMMARY', json.dumps(report['summary']), flush=True)
finally:
if process and process.returncode is None:
process.terminate()
await process.wait()
log.close()
await adapter.aclose()
await upstream.aclose()
await runner.cleanup()
if __name__ == '__main__':
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument('--opencode', required=True)
parser.add_argument('--opencode-version', default='unknown')
parser.add_argument('--model', default='claude-haiku-4.5')
parser.add_argument('--base-url')
parser.add_argument('--credentials-file')
parser.add_argument('--api-key-env', default='LLM_BENCHMARK_API_KEY')
parser.add_argument('--rounds', type=int, default=6)
parser.add_argument('--mock', action='store_true')
parser.add_argument('--long-prefix', action='store_true')
parser.add_argument('--cache-control', action='store_true')
parser.add_argument('--baseline-no-drain', action='store_true', help='Reproduce the pre-fix stream transport')
parser.add_argument('--output', required=True)
args = parser.parse_args()
if not 1 <= args.rounds <= 12:
parser.error('--rounds must be between 1 and 12 (plus one warmup round)')
os.environ['NO_PROXY'] = ','.join(filter(None, [os.environ.get('NO_PROXY'), '127.0.0.1', 'localhost']))
asyncio.run(run(args))
skills/skills/langbot-mcp-ops/SKILL.md
+23
View File
@@ -242,3 +242,26 @@ Use `reset_session_context(bot_id, session_id)` only when a user asks to start a
## Cloud plan platform restrictions
Bot creation and enabling use the same Workspace entitlement policy as HTTP. A disabled platform is rejected even through MCP or plugin calls; do not retry using an Omni or legacy alias to bypass the restriction. The Workspace-scoped HTTP adapter catalog omits disabled platforms.
## LLM request diagnostics
`get_monitoring_records(kind="llm_calls")` includes nullable `details` on each call.
Read `details.timing`, `details.usage`, and `details.error` to distinguish provider
latency, retry waits, cached/reasoning tokens, and failures. Timing values are
milliseconds since invocation start; `first_event_ms` includes reasoning/control
frames, while `first_content_ms` measures the first upstream content delta.
`first_visible_chunk_ms` may contain displayed reasoning, not just answer text.
Old records have no details. Do not equate total duration with provider TTFT.
See `docs/LLM_REQUEST_POLICY.md` for policy precedence and gateway header behavior.
`get_llm_model` also exposes optional `reasoning_capabilities.budget_tokens`
min/max bounds. Numeric thinking budgets use `reasoning_config.budget_tokens`
with level `enabled`; do not combine budgets with effort levels or native thinking
arguments. Existing Runner reasoning overrides take precedence. Prompt caching
is opt-in through model/call `extra_args.prompt_caching`; use provider-specific
modes documented in `docs/LLM_REQUEST_POLICY.md`. Cache keys are routing hints,
not evidence of hits; read the actual usage counters. `preparation_ms` measures
request construction and `preparation_cache_hits` counts retry preparation reuse.
`transport_setup_ms` measures managed pool acquisition, including first-use TLS
initialization; it is not network TTFT. See `docs/LLM_LATENCY_BENCHMARK.md` for
controlled request-layer measurements and their scope.
src/langbot/pkg/api/http/service/monitoring.py
+3 -1
View File
@@ -608,6 +608,7 @@ class MonitoringService(ExecutionDetailsMixin):
cost: float | None = None,
error_message: str | None = None,
message_id: str | None = None,
details: dict | None = None,
) -> str:
"""Record an LLM call"""
workspace_uuid = self._require_write_context(context)
@@ -619,7 +620,7 @@ class MonitoringService(ExecutionDetailsMixin):
'model_name': model_name,
'input_tokens': input_tokens,
'output_tokens': output_tokens,
'total_tokens': input_tokens + output_tokens,
'total_tokens': ((details or {}).get('usage') or {}).get('total_tokens') or (input_tokens + output_tokens),
'duration': duration,
'cost': cost,
'status': status,
@@ -630,6 +631,7 @@ class MonitoringService(ExecutionDetailsMixin):
'session_id': session_id,
'error_message': error_message,
'message_id': message_id,
'details': details,
}
await self.ap.persistence_mgr.execute_async(
src/langbot/pkg/api/mcp/server.py
+8 -2
View File
@@ -182,7 +182,7 @@ class LangBotMCPServer:
return _dump(ap.box_service.get_recent_errors(context))
@mcp.tool(
description='Read bounded Workspace runtime records. Filters follow the monitoring service: bot_ids, pipeline_ids, session_ids, start_time, end_time, mode (all/real/debug), execution_statuses (owning execution status groups), knowledge_base_id, user_query, is_active as supported by the record kind.'
description='Read bounded Workspace runtime records. LLM calls include optional details (timing in ms, complete provider usage, structured error); older records have null details. Filters follow the monitoring service: bot_ids, pipeline_ids, session_ids, start_time, end_time, mode (all/real/debug), execution_statuses (owning execution status groups), knowledge_base_id, user_query, is_active as supported by the record kind.'
)
async def get_monitoring_records(
kind: typing.Literal['messages', 'llm_calls', 'tool_calls', 'embedding_calls', 'sessions', 'errors'],
@@ -530,7 +530,13 @@ class LangBotMCPServer:
context = _authorized(Permission.RESOURCE_VIEW)
return _dump(await ap.llm_model_service.get_llm_models(context, include_secret=False))
@mcp.tool(description='Get a single LLM model by UUID.')
@mcp.tool(
description=(
'Get a single LLM model by UUID, including reasoning capabilities. '
'Optional budget_tokens min/max describes numeric thinking-budget support; '
'reasoning_config and extra_args retain the existing model API shape.'
)
)
async def get_llm_model(model_uuid: str) -> str:
context = _authorized(Permission.RESOURCE_VIEW)
return _dump(await ap.llm_model_service.get_llm_model(context, model_uuid, include_secret=False))
src/langbot/pkg/entity/persistence/monitoring.py
+2
View File
@@ -70,6 +70,8 @@ class MonitoringLLMCall(Base):
error_message = sqlalchemy.Column(sqlalchemy.Text, nullable=True)
message_id = sqlalchemy.Column(sqlalchemy.String(255), nullable=True, index=True) # Associated message ID
details = sqlalchemy.Column(sqlalchemy.JSON, nullable=True)
__table_args__ = (
sqlalchemy.Index('ix_monitoring_llm_calls_workspace_timestamp', 'workspace_uuid', 'timestamp'),
sqlalchemy.Index('ix_monitoring_llm_calls_workspace_session', 'workspace_uuid', 'session_id'),
src/langbot/pkg/persistence/alembic/versions/0038_llm_call_details.py
+25
View File
@@ -0,0 +1,25 @@
"""Persist additive LLM timing, usage and structured failure diagnostics."""
from alembic import op
import sqlalchemy as sa
revision = '0038_llm_call_details'
down_revision = '0037_metadata_text'
branch_labels = None
depends_on = None
def upgrade():
inspector = sa.inspect(op.get_bind())
if inspector.has_table('monitoring_llm_calls'):
columns = {column['name'] for column in inspector.get_columns('monitoring_llm_calls')}
if 'details' not in columns:
op.add_column('monitoring_llm_calls', sa.Column('details', sa.JSON(), nullable=True))
def downgrade():
inspector = sa.inspect(op.get_bind())
if inspector.has_table('monitoring_llm_calls'):
columns = {column['name'] for column in inspector.get_columns('monitoring_llm_calls')}
if 'details' in columns:
op.drop_column('monitoring_llm_calls', 'details')
src/langbot/pkg/pipeline/monitoring_helper.py
+2
View File
@@ -380,6 +380,7 @@ class MonitoringHelper:
cost: float | None = None,
error_message: str | None = None,
message_id: str | None = None,
details: dict | None = None,
):
"""Record LLM call"""
try:
@@ -400,6 +401,7 @@ class MonitoringHelper:
cost=cost,
error_message=error_message,
message_id=message_id,
details=details,
)
except Exception as e:
ap.logger.error(f'Failed to record LLM call: {e}')
src/langbot/pkg/plugin/handler.py
+13 -3
View File
@@ -443,13 +443,19 @@ def _resolve_remove_think(data: dict[str, Any], query: Any | None) -> bool:
return False
def _merge_model_extra_args(model: Any, call_extra_args: Any) -> dict[str, Any]:
def _merge_model_extra_args(model: Any, call_extra_args: Any, *, llm_policy: bool = False) -> dict[str, Any]:
"""Merge persisted model extra_args with action-level overrides."""
merged: dict[str, Any] = {}
model_extra_args = getattr(getattr(model, 'model_entity', None), 'extra_args', None)
if isinstance(model_extra_args, dict):
merged.update(model_extra_args)
if llm_policy:
# RuntimeProvider resolves model and call policy layers independently.
# Flattening them here makes a model request_policy.retries override a
# caller's legacy num_retries=0 (and similarly hides timeout overrides).
for key in ('request_policy', 'num_retries', 'max_retries', 'timeout'):
merged.pop(key, None)
if isinstance(call_extra_args, dict):
merged.update(call_extra_args)
@@ -1484,7 +1490,9 @@ class RuntimeConnectionHandler(handler.Handler):
funcs_obj = [resource_tool.LLMTool.model_validate({**func, 'func': _placeholder_func}) for func in funcs]
query = _resolve_action_query(data, session, self.ap, action_context)
effective_extra_args = _merge_model_extra_args(llm_model, extra_args)
if query is None:
query = _build_plugin_action_query(action_context)
effective_extra_args = _merge_model_extra_args(llm_model, extra_args, llm_policy=True)
remove_think = _resolve_remove_think(data, query)
effective_funcs = funcs_obj if 'func_call' in (llm_model.model_entity.abilities or []) else []
@@ -1580,7 +1588,9 @@ class RuntimeConnectionHandler(handler.Handler):
funcs_obj = [resource_tool.LLMTool.model_validate({**func, 'func': _placeholder_func}) for func in funcs]
query = _resolve_action_query(data, session, self.ap, action_context)
effective_extra_args = _merge_model_extra_args(llm_model, extra_args)
if query is None:
query = _build_plugin_action_query(action_context)
effective_extra_args = _merge_model_extra_args(llm_model, extra_args, llm_policy=True)
remove_think = _resolve_remove_think(data, query)
effective_funcs = funcs_obj if 'func_call' in (llm_model.model_entity.abilities or []) else []
src/langbot/pkg/provider/modelmgr/codex_errors.py
+9 -1
View File
@@ -1,12 +1,20 @@
"""Explicitly safe Codex failures; never construct messages from upstream bodies."""
from ...i18n import t
from .errors import RequesterError
class CodexProviderError(ValueError):
"""A known provider failure safe to expose at the HTTP boundary."""
def __init__(self, message: str, status_code: int = 502, error_code: str = 'codex_upstream_failure'):
super().__init__(t(message))
ValueError.__init__(self, t(message))
self.code = error_code
self.retryable = status_code in (500, 502, 503, 504)
self.retry_after = None
self.request_id = None
self.phase = None
self.status_code = status_code
self.error_code = error_code
as_dict = RequesterError.as_dict
src/langbot/pkg/provider/modelmgr/errors.py
+24 -1
View File
@@ -4,5 +4,28 @@ from ...i18n import t
class RequesterError(Exception):
"""Base class for all Requester errors."""
def __init__(self, message: str):
def __init__(
self,
message: str,
*,
code: str = 'provider_error',
status_code: int | None = None,
retryable: bool = False,
retry_after: float | None = None,
request_id: str | None = None,
phase: str | None = None,
):
super().__init__(t('Model request failed') + ': ' + message)
self.code = code
self.status_code = status_code
self.retryable = retryable
self.retry_after = retry_after
self.request_id = request_id
self.phase = phase
def as_dict(self) -> dict:
"""Machine-readable diagnostics without upstream bodies or credentials."""
return {
key: getattr(self, key)
for key in ('code', 'status_code', 'retryable', 'retry_after', 'request_id', 'phase')
}
src/langbot/pkg/provider/modelmgr/http_pool.py
+114
View File
@@ -0,0 +1,114 @@
"""Requester-local connections, retired only after their active leases finish."""
from contextlib import asynccontextmanager
import asyncio
import re
import httpx
from .errors import RequesterError
class _CompletedSSEStream(httpx.AsyncByteStream):
"""Finish bounded HTTP framing before SDKs stop at the SSE DONE record."""
_done = re.compile(rb'\ndata: ?\[DONE\]\r?\n\r?\n')
def __init__(self, stream):
self.stream = stream
async def __aiter__(self):
iterator = self.stream.__aiter__()
tail = b'\n'
async for chunk in iterator:
probe = tail + chunk
if self._done.search(probe):
# SDK parsers stop at [DONE] before httpcore observes EOF. This
# otherwise discards the connection on every successful stream.
# Do not drain on cancellation, truncation, or merely a finish
# reason: providers can still send usage before the DONE record.
try:
async with asyncio.timeout(0.05):
drained = 0
async for extra in iterator:
drained += len(extra)
if drained >= 65536:
break
except (TimeoutError, httpx.HTTPError):
pass
yield chunk
return
tail = probe[-32:]
yield chunk
async def aclose(self):
await self.stream.aclose()
async def _finish_sse_body(response):
if response.headers.get('content-type', '').split(';')[0].strip() == 'text/event-stream':
response.stream = _CompletedSSEStream(response.stream)
class RequesterHTTPPool:
def __init__(self, *, follow_redirects=False, drain_sse=False):
self.client = None
self.active = 0
self.retired = False
self.follow_redirects = follow_redirects
self.drain_sse = drain_sse
self._creation_lock = asyncio.Lock()
@asynccontextmanager
async def lease(self):
if self.retired:
raise RequesterError('Provider has been retired', code='provider_retired')
if self.client is None:
async with self._creation_lock:
if self.retired:
raise RequesterError('Provider has been retired', code='provider_retired')
if self.client is None:
# HTTPX otherwise reloads the same CA bundle for its direct
# transport and each environment proxy mount. Load it once,
# off the event loop, preserving HTTPX's TLS/env semantics.
# Cancellation here leaves no open HTTP client or sockets.
ssl_context = await asyncio.to_thread(httpx.create_ssl_context)
if self.retired:
raise RequesterError('Provider has been retired', code='provider_retired')
# Auth and deadlines remain request-local. No pool or TLS
# context is shared across runtime providers.
self.client = httpx.AsyncClient(
verify=ssl_context,
follow_redirects=self.follow_redirects,
limits=httpx.Limits(max_connections=100, max_keepalive_connections=20, keepalive_expiry=30),
event_hooks={'response': [_finish_sse_body]} if self.drain_sse else None,
)
self.active += 1
try:
yield self.client
finally:
self.active -= 1
if self.retired and not self.active:
await self._close()
async def _close(self):
client, self.client = self.client, None
if client is not None:
await client.aclose()
async def aclose(self):
self.retired = True
if not self.active:
await self._close()
@asynccontextmanager
async def requester_http_client(owner):
pool = getattr(owner, '_http_pool', None)
if pool is not None:
async with pool.lease() as client:
yield client
else:
# Preserve direct/uninitialized requester use by older extensions/tests.
async with httpx.AsyncClient(follow_redirects=False) as client:
yield client
src/langbot/pkg/provider/modelmgr/modelmgr.py
+1
View File
@@ -743,6 +743,7 @@ class ModelManager:
config = {
'base_url': provider_entity.base_url,
'requester_name': provider_entity.requester,
'request_policy': self.ap.instance_config.data.get('llm_request_policy', {}),
}
if provider_entity.requester == 'openai-codex':
src/langbot/pkg/provider/modelmgr/prompt_caching.py
+82
View File
@@ -0,0 +1,82 @@
"""Opt-in upstream prompt caching; never cache generated answers locally."""
import copy
import hashlib
from .errors import RequesterError
def _has_breakpoint(value):
if isinstance(value, dict):
return 'cache_control' in value or any(_has_breakpoint(item) for item in value.values())
if isinstance(value, list):
return any(_has_breakpoint(item) for item in value)
return False
def apply_prompt_caching(args, state, *, provider, family):
config = args.pop('prompt_caching', None)
if config is None:
return
if not isinstance(config, dict) or set(config) - {'mode', 'retention', 'ttl'}:
raise RequesterError('Invalid prompt_caching fields', code='invalid_policy')
mode = config.get('mode', 'provider_default')
if mode not in ('provider_default', 'openai', 'anthropic'):
raise RequesterError('Unsupported prompt_caching mode', code='invalid_policy')
if mode == 'provider_default':
if set(config) - {'mode'}:
raise RequesterError('Cache settings require an explicit mode', code='invalid_policy')
return
if mode == 'openai':
if provider not in ('openai', 'openrouter') or family not in ('openai', 'openrouter'):
raise RequesterError('OpenAI caching requires an OpenAI-compatible model', code='invalid_policy')
if 'ttl' in config or config.get('retention') not in (None, 'in_memory', '24h'):
raise RequesterError('Invalid OpenAI cache retention', code='invalid_policy')
body = copy.deepcopy(args.get('extra_body') or {})
if not isinstance(body, dict):
raise RequesterError('extra_body must be an object', code='invalid_policy')
# Bypass LiteLLM's model-specific parameter filter for explicit opt-in
# caching fields as well as generated hints. The SDK's extra_body wins
# when the same native field appears in both locations.
for key in ('prompt_cache_key', 'prompt_cache_retention', 'prompt_cache_options'):
if key in args:
body.setdefault(key, args.pop(key))
# Explicit provider-native settings win. A key is a routing hint only;
# it neither enables caching nor guarantees a hit.
if (
state
and state.conversation_id
and not any(key in body for key in ('prompt_cache_key', 'prompt_cache_options'))
):
scope = f'{state.conversation_id}:{args["model"]}'
body['prompt_cache_key'] = 'lb-' + hashlib.sha256(scope.encode()).hexdigest()
if config.get('retention') and not any(
key in location for key in ('prompt_cache_retention', 'prompt_cache_options') for location in (args, body)
):
body['prompt_cache_retention'] = config['retention']
if body:
args['extra_body'] = body
return
if family != 'anthropic':
raise RequesterError('Anthropic caching requires an Anthropic model', code='invalid_policy')
if 'retention' in config or config.get('ttl', '5m') not in ('5m', '1h'):
raise RequesterError('Invalid Anthropic cache TTL', code='invalid_policy')
# Existing manual or automatic breakpoints are authoritative. Do not exceed
# a provider's breakpoint limit or change its TTL ordering.
if _has_breakpoint(args):
return
messages = copy.deepcopy(args['messages'])
boundary = None
for message in messages:
if message.get('role') not in ('system', 'developer'):
break
content = message.get('content')
if isinstance(content, str) and content:
message['content'] = [{'type': 'text', 'text': content}]
if isinstance(message.get('content'), list):
for part in message['content']:
if isinstance(part, dict) and part.get('type') == 'text' and part.get('text'):
boundary = part
if boundary is not None:
boundary['cache_control'] = {'type': 'ephemeral', 'ttl': config.get('ttl', '5m')}
args['messages'] = messages
src/langbot/pkg/provider/modelmgr/reasoning.py
+32 -5
View File
@@ -48,20 +48,27 @@ _CONFLICTING_EXTRA_BODY_ARGS = {
}
def normalize_reasoning_config(value: typing.Any) -> dict[str, str]:
def normalize_reasoning_config(value: typing.Any) -> dict[str, typing.Any]:
"""Return the canonical model reasoning configuration."""
if value is None:
return dict(DEFAULT_REASONING_CONFIG)
if not isinstance(value, dict):
raise ValueError('reasoning_config must be an object')
unknown_fields = set(value) - {'level'}
unknown_fields = set(value) - {'level', 'budget_tokens'}
if unknown_fields:
raise ValueError(f'Unsupported reasoning_config fields: {", ".join(sorted(unknown_fields))}')
level = value.get('level', 'provider_default')
budget = value.get('budget_tokens')
level = value.get('level', 'enabled' if budget is not None else 'provider_default')
if level not in REASONING_LEVELS:
raise ValueError(f'Unsupported reasoning level: {level}')
if budget is not None:
if type(budget) is not int or budget <= 0:
raise ValueError('budget_tokens must be a positive integer')
if level != 'enabled':
raise ValueError('budget_tokens requires level enabled; do not combine a token budget with an effort level')
return {'level': level, 'budget_tokens': budget}
return {'level': typing.cast(str, level)}
@@ -69,7 +76,7 @@ def validate_reasoning_config(
value: typing.Any,
abilities: typing.Iterable[str] | None,
extra_args: typing.Any,
) -> dict[str, str]:
) -> dict[str, typing.Any]:
"""Validate a model-facing reasoning config and conflicting raw arguments."""
config = normalize_reasoning_config(value)
if config['level'] == 'provider_default':
@@ -101,7 +108,13 @@ def validate_reasoning_capabilities(
model_name: str,
) -> None:
"""Ensure an explicit reasoning level can be honored by the requester."""
level = normalize_reasoning_config(config)['level']
config = normalize_reasoning_config(config)
level = config['level']
if 'budget_tokens' in config:
limits = capabilities.get('budget_tokens')
if not isinstance(limits, dict) or not limits['min'] <= config['budget_tokens'] <= limits['max']:
raise ValueError(f'Unsupported reasoning token budget for model {model_name}: {config["budget_tokens"]}')
return
if level == 'provider_default':
return
@@ -154,3 +167,17 @@ def model_with_reasoning_level(
scoped_model = copy.copy(model)
scoped_model.reasoning_config_override = copy.deepcopy(config)
return scoped_model
def resolve_reasoning_config(model, extra_args=None):
"""Runner policy wins; a call object replaces, rather than merges, a model policy."""
value = getattr(model.model_entity, 'reasoning_config', None)
if value is not None and not isinstance(value, (dict, str, int, float, bool, list, tuple)):
value = None
for args in (model.model_entity.extra_args, extra_args):
if isinstance(args, dict) and 'reasoning_config' in args:
value = args['reasoning_config']
override = getattr(model, 'reasoning_config_override', None)
if override is not None:
value = override
return normalize_reasoning_config(value)
src/langbot/pkg/provider/modelmgr/request_policy.py
+173
View File
@@ -0,0 +1,173 @@
"""Host-owned LLM deadlines, retries and additive diagnostics.
No SDK wire shape changes: request_policy is a reserved model/call extra argument.
Context-local state lets adapters report raw events without sharing mutable state.
"""
from __future__ import annotations
import asyncio
from contextvars import ContextVar
from dataclasses import dataclass, fields
from email.utils import parsedate_to_datetime
import hashlib
import json
import math
import random
import time
from .errors import RequesterError
from .codex_errors import CodexProviderError
STRUCTURED_ERRORS = (RequesterError, CodexProviderError)
@dataclass
class RequestPolicy:
retries: int = 0
total_timeout: float = 300
first_event_timeout: float = 120
stream_idle_timeout: float = 120
connect_timeout: float = 10
read_timeout: float = 120
write_timeout: float = 30
pool_timeout: float = 10
backoff_base: float = 0.5
backoff_max: float = 8
attribution: bool = True
conversation_header: str | None = None
@classmethod
def resolve(cls, *layers):
values = {}
names = {field.name for field in fields(cls)}
for layer in layers:
if not isinstance(layer, dict):
continue
# Preserve the legacy per-model/call retry knob, including zero.
if 'max_retries' in layer:
values['retries'] = layer['max_retries']
if 'num_retries' in layer:
values['retries'] = layer['num_retries']
if isinstance(layer.get('timeout'), (int, float)):
values['read_timeout'] = layer['timeout']
values['first_event_timeout'] = layer['timeout']
values['stream_idle_timeout'] = layer['timeout']
policy = layer.get('request_policy', {})
if not isinstance(policy, dict) or set(policy) - names:
raise RequesterError('Invalid request_policy fields', code='invalid_policy')
values.update(policy)
policy = cls(**values)
for field in fields(cls):
value = getattr(policy, field.name)
if field.name == 'conversation_header':
if value not in (None, 'x-session-id', 'x-litellm-session-id', 'x-langbot-conversation-id'):
raise RequesterError('Unsupported conversation header', code='invalid_policy')
elif field.name == 'attribution':
if not isinstance(value, bool):
raise RequesterError('attribution must be boolean', code='invalid_policy')
elif field.name == 'retries':
if type(value) is not int or not 0 <= value <= 5:
raise RequesterError('retries must be an integer from 0 to 5', code='invalid_policy')
elif (
isinstance(value, bool) or not isinstance(value, (int, float)) or not math.isfinite(value) or value <= 0
):
raise RequesterError(f'{field.name} must be finite and positive', code='invalid_policy')
return policy
class CallState:
def __init__(self, policy, conversation_id=None):
self.policy = policy
self.conversation_id = conversation_id
self.start = time.perf_counter()
self.attempt_start = self.start
self.metrics = {'attempts': 0, 'retry_wait_ms': 0}
self.raw_started = False
self.error = None
self.usage = None
self.prepared = {}
def mark(self, name):
self.metrics.setdefault(name, round((time.perf_counter() - self.start) * 1000, 3))
async def wait(self, operation, phase, limit):
remaining = self.policy.total_timeout - (time.perf_counter() - self.start)
if phase == 'first_event':
limit = min(limit, self.policy.first_event_timeout - (time.perf_counter() - self.attempt_start))
token = current_call.set(self)
try:
# timeout(0) only cancels at a scheduling point. Buffered SDK
# responses can complete synchronously, so enforce expired budgets
# before starting another operation as well.
if min(limit, remaining) <= 0:
raise TimeoutError
async with asyncio.timeout(max(0, min(limit, remaining))):
return await operation()
except TimeoutError as exc:
raise RequesterError(
'LLM request timed out', code='timeout', retryable=True, phase='total' if remaining <= limit else phase
) from exc
finally:
current_call.reset(token)
async def retry(self, error):
self.error = error.as_dict() if isinstance(error, STRUCTURED_ERRORS) else {'code': 'provider_error'}
self.metrics.setdefault('attempt_errors', []).append(dict(self.error))
if not isinstance(error, STRUCTURED_ERRORS) or not error.retryable or self.raw_started:
return False
if self.metrics['attempts'] > self.policy.retries:
return False
delay = random.uniform(
0, min(self.policy.backoff_max, self.policy.backoff_base * 2 ** (self.metrics['attempts'] - 1))
)
delay = max(delay, error.retry_after or 0)
remaining = self.policy.total_timeout - (time.perf_counter() - self.start)
if delay >= remaining:
return False
started = time.perf_counter()
await asyncio.sleep(delay)
self.metrics['retry_wait_ms'] += round((time.perf_counter() - started) * 1000, 3)
return True
current_call: ContextVar[CallState | None] = ContextVar('llm_call', default=None)
def opaque_conversation_id(query, context):
if query is None:
return None
session = getattr(query, 'session', None)
session_id = getattr(session, 'session_id', None)
if not isinstance(session_id, (str, int)):
launcher = getattr(query, 'launcher_id', None)
if not isinstance(launcher, (str, int)):
return None
launcher_type = getattr(query, 'launcher_type', '')
session_id = [str(getattr(launcher_type, 'value', launcher_type)), str(launcher)]
scope = [context.workspace_uuid, query.bot_uuid, query.pipeline_uuid, session_id]
return 'lb-' + hashlib.sha256(json.dumps(scope, ensure_ascii=True).encode()).hexdigest()
def error_metadata(exc):
status = getattr(exc, 'status_code', None)
status = status if isinstance(status, int) else None
response = getattr(exc, 'response', None)
headers = getattr(response, 'headers', {})
headers = headers if hasattr(headers, 'get') else {}
retry_after = None
raw = headers.get('retry-after')
if raw:
try:
retry_after = float(raw)
except (ValueError, TypeError):
try:
retry_after = parsedate_to_datetime(raw).timestamp() - time.time()
except (ValueError, TypeError, OverflowError):
pass
if retry_after is not None:
retry_after = max(0, retry_after) if math.isfinite(retry_after) else None
request_id = headers.get('x-request-id') or headers.get('request-id')
if not isinstance(request_id, str):
request_id = None
return {'status_code': status, 'retry_after': retry_after, 'request_id': request_id[:256] if request_id else None}
src/langbot/pkg/provider/modelmgr/requester.py
+135 -144
View File
@@ -14,6 +14,7 @@ from ...workspace.errors import WorkspaceInvariantError
import langbot_plugin.api.entities.builtin.resource.tool as resource_tool
from . import token
from . import reasoning
from . import request_policy
import langbot_plugin.api.entities.builtin.pipeline.query as pipeline_query
import langbot_plugin.api.entities.builtin.provider.message as provider_message
@@ -101,6 +102,48 @@ class RuntimeProvider:
raise WorkspaceInvariantError('LLM invocation requires an ExecutionContext when query is absent')
return execution_context
def _call_state(self, query, model, extra_args, context):
config = getattr(self.requester, 'requester_cfg', {})
policy = request_policy.RequestPolicy.resolve(config, model.model_entity.extra_args, extra_args)
if query is not None:
if query.variables is None:
query.variables = {}
query.variables.pop(LLM_USAGE_QUERY_VARIABLE, None)
query.variables.pop(STREAM_USAGE_QUERY_VARIABLE, None)
return request_policy.CallState(policy, request_policy.opaque_conversation_id(query, context))
async def _record_call(self, query, model, state, status, error):
state.mark('total_ms')
if query is None:
return
stream_usage = query.variables.pop(STREAM_USAGE_QUERY_VARIABLE, None)
usage = state.usage if state.usage is not None else stream_usage
if usage is not None:
_store_llm_usage(query, usage)
details = {'timing': dict(state.metrics), 'usage': usage, 'error': error}
query.variables['_llm_diagnostics'] = details
try:
from ...pipeline import monitoring_helper
await monitoring_helper.MonitoringHelper.record_llm_call(
ap=self.requester.ap,
query=query,
bot_id=query.bot_uuid or 'unknown',
bot_name=query.variables.get('_monitoring_bot_name', 'Unknown'),
pipeline_id=query.pipeline_uuid or 'unknown',
pipeline_name=query.variables.get('_monitoring_pipeline_name', 'Unknown'),
model_name=model.model_entity.name,
input_tokens=(usage or {}).get('prompt_tokens', 0) or 0,
output_tokens=(usage or {}).get('completion_tokens', 0) or 0,
duration_ms=int(state.metrics['total_ms']),
status=status,
error_message=error.get('message') if error else None,
message_id=query.variables.get('_monitoring_message_id'),
details=details,
)
except Exception as exc:
self.requester.ap.logger.error(f'[Monitoring] Failed to record LLM call: {exc}')
async def invoke_llm(
self,
query: pipeline_query.Query | None,
@@ -111,79 +154,46 @@ class RuntimeProvider:
remove_think: bool = False,
execution_context: ExecutionContext | None = None,
) -> provider_message.Message:
"""Bridge method for invoking LLM with monitoring"""
invocation_context = self._resolve_llm_execution_context(query, execution_context)
self._validate_invocation(model, invocation_context)
# Start timing for monitoring
start_time = time.time()
input_tokens = 0
output_tokens = 0
status = 'success'
error_message = None
"""Preserve the legacy Message return while owning retries and deadlines."""
context = self._resolve_llm_execution_context(query, execution_context)
self._validate_invocation(model, context)
state = self._call_state(query, model, extra_args, context)
status, error = 'success', None
try:
# Call the underlying requester
result = await self.requester.invoke_llm(
query=query,
model=model,
messages=messages,
funcs=funcs,
extra_args=extra_args,
remove_think=remove_think,
)
# Try to extract token usage if the requester returns it
# For requesters that return tuple (message, usage_info)
if isinstance(result, tuple):
msg, usage_info = result
if usage_info:
_store_llm_usage(query, usage_info)
input_tokens = usage_info.get('prompt_tokens', 0)
output_tokens = usage_info.get('completion_tokens', 0)
return msg
else:
return result
except Exception as e:
while True:
state.metrics['attempts'] += 1
state.attempt_start = time.perf_counter()
try:
result = await state.wait(
lambda: self.requester.invoke_llm(
query=query,
model=model,
messages=messages,
funcs=funcs,
extra_args=extra_args,
remove_think=remove_think,
),
'response',
state.policy.total_timeout,
)
state.mark('response_completed_ms')
if isinstance(result, tuple):
message, state.usage = result
return message
return result
except Exception as exc:
if not await state.retry(exc):
raise
except (asyncio.CancelledError, GeneratorExit):
status, error = 'error', {'code': 'cancelled', 'message': 'LLM invocation cancelled'}
raise
except Exception as exc:
status = 'error'
error_message = str(e)
error = exc.as_dict() if isinstance(exc, request_policy.STRUCTURED_ERRORS) else {'code': 'provider_error'}
error['message'] = str(exc)
raise
finally:
# Record LLM call monitoring data (only if query is provided)
if query is not None:
duration_ms = int((time.time() - start_time) * 1000)
# Import monitoring helper
try:
from ...pipeline import monitoring_helper
# Get monitoring metadata from query variables
if query.variables:
bot_name = query.variables.get('_monitoring_bot_name', 'Unknown')
pipeline_name = query.variables.get('_monitoring_pipeline_name', 'Unknown')
message_id = query.variables.get('_monitoring_message_id')
else:
bot_name = 'Unknown'
pipeline_name = 'Unknown'
message_id = None
await monitoring_helper.MonitoringHelper.record_llm_call(
ap=self.requester.ap,
query=query,
bot_id=query.bot_uuid or 'unknown',
bot_name=bot_name,
pipeline_id=query.pipeline_uuid or 'unknown',
pipeline_name=pipeline_name,
model_name=model.model_entity.name,
input_tokens=input_tokens,
output_tokens=output_tokens,
duration_ms=duration_ms,
status=status,
error_message=error_message,
message_id=message_id,
)
except Exception as monitor_err:
self.requester.ap.logger.error(f'[Monitoring] Failed to record LLM call: {monitor_err}')
await self._record_call(query, model, state, status, error)
async def invoke_llm_stream(
self,
@@ -194,85 +204,66 @@ class RuntimeProvider:
extra_args: dict[str, typing.Any] = {},
remove_think: bool = False,
execution_context: ExecutionContext | None = None,
) -> provider_message.MessageChunk:
"""Bridge method for invoking LLM stream with monitoring"""
invocation_context = self._resolve_llm_execution_context(query, execution_context)
self._validate_invocation(model, invocation_context)
# Start timing for monitoring
start_time = time.time()
status = 'success'
error_message = None
input_tokens = 0
output_tokens = 0
) -> typing.AsyncIterator[provider_message.MessageChunk]:
"""Never retry after a raw upstream event or a legacy adapter chunk."""
context = self._resolve_llm_execution_context(query, execution_context)
self._validate_invocation(model, context)
state = self._call_state(query, model, extra_args, context)
status, error = 'success', None
try:
# Stream the response
async with aclosing(
self.requester.invoke_llm_stream(
query=query,
model=model,
messages=messages,
funcs=funcs,
extra_args=extra_args,
remove_think=remove_think,
)
) as owned_stream:
async for chunk in owned_stream:
yield chunk
# Extract usage from stream if available (stored by LiteLLM requester)
if query:
if query.variables is None:
query.variables = {}
if STREAM_USAGE_QUERY_VARIABLE in query.variables:
usage_info = query.variables[STREAM_USAGE_QUERY_VARIABLE]
_store_llm_usage(query, usage_info)
input_tokens = usage_info.get('prompt_tokens', 0)
output_tokens = usage_info.get('completion_tokens', 0)
del query.variables[STREAM_USAGE_QUERY_VARIABLE]
while True:
state.metrics['attempts'] += 1
state.attempt_start = time.perf_counter()
try:
async with aclosing(
self.requester.invoke_llm_stream(
query=query,
model=model,
messages=messages,
funcs=funcs,
extra_args=extra_args,
remove_think=remove_think,
)
) as stream:
while True:
try:
chunk = await state.wait(
lambda: anext(stream),
'total'
if getattr(self.requester, 'reports_raw_timing', False) is True
else ('stream_idle' if state.raw_started else 'first_event'),
state.policy.total_timeout
if getattr(self.requester, 'reports_raw_timing', False) is True
else (
state.policy.stream_idle_timeout
if state.raw_started
else state.policy.first_event_timeout
),
)
except StopAsyncIteration:
break
if chunk is None:
continue
state.raw_started = True
state.mark('first_event_ms')
if chunk.content:
state.mark('first_visible_chunk_ms')
yield chunk
state.mark('response_completed_ms')
return
except Exception as exc:
if not await state.retry(exc):
raise
except (asyncio.CancelledError, GeneratorExit):
status = 'error'
error_message = 'Streaming generation cancelled by consumer'
status, error = 'error', {'code': 'cancelled', 'message': 'LLM stream cancelled'}
raise
except Exception as e:
except Exception as exc:
status = 'error'
error_message = str(e)
error = exc.as_dict() if isinstance(exc, request_policy.STRUCTURED_ERRORS) else {'code': 'provider_error'}
error['message'] = str(exc)
raise
finally:
# Record LLM call monitoring data (only if query is provided)
if query is not None:
duration_ms = int((time.time() - start_time) * 1000)
# Import monitoring helper
try:
from ...pipeline import monitoring_helper
# Get monitoring metadata from query variables
if query.variables:
bot_name = query.variables.get('_monitoring_bot_name', 'Unknown')
pipeline_name = query.variables.get('_monitoring_pipeline_name', 'Unknown')
message_id = query.variables.get('_monitoring_message_id')
else:
bot_name = 'Unknown'
pipeline_name = 'Unknown'
message_id = None
await monitoring_helper.MonitoringHelper.record_llm_call(
ap=self.requester.ap,
query=query,
bot_id=query.bot_uuid or 'unknown',
bot_name=bot_name,
pipeline_id=query.pipeline_uuid or 'unknown',
pipeline_name=pipeline_name,
model_name=model.model_entity.name,
input_tokens=input_tokens,
output_tokens=output_tokens,
duration_ms=duration_ms,
status=status,
error_message=error_message,
message_id=message_id,
)
except Exception as monitor_err:
self.requester.ap.logger.error(f'[Monitoring] Failed to record LLM stream call: {monitor_err}')
await self._record_call(query, model, state, status, error)
async def invoke_embedding(
self,
@@ -388,7 +379,7 @@ class RuntimeLLMModel:
provider: RuntimeProvider
"""提供商实例"""
reasoning_config_override: dict[str, str] | None
reasoning_config_override: dict[str, typing.Any] | None
"""Request-scoped reasoning policy supplied by the active pipeline."""
def __init__(
@@ -396,7 +387,7 @@ class RuntimeLLMModel:
execution_context: ExecutionContext,
model_entity: persistence_model.LLMModel,
provider: RuntimeProvider,
reasoning_config_override: dict[str, str] | None = None,
reasoning_config_override: dict[str, typing.Any] | None = None,
):
_ensure_same_execution_scope(provider.execution_context, execution_context, resource='LLM model')
if model_entity.workspace_uuid != execution_context.workspace_uuid:
src/langbot/pkg/provider/modelmgr/requesters/codex.py
+158 -76
View File
@@ -7,12 +7,16 @@ import json
import secrets
import time
from collections import OrderedDict
from contextlib import AsyncExitStack, aclosing
import httpx
import langbot
import langbot_plugin.api.entities.builtin.provider.message as pm
from .. import requester, reasoning
from ..request_policy import current_call
from ..http_pool import RequesterHTTPPool, requester_http_client
from ..errors import RequesterError
from ..codex_auth import BASE_URL, CodexAuth, LOGIN_REQUIRED
from ..codex_errors import CodexProviderError
@@ -21,7 +25,8 @@ async def sse_events(response):
"""Decode SSE records, including CRLF, comments, and multiline data."""
data = []
size = 0
async for line in response.aiter_lines():
lines = response.aiter_lines()
async for line in lines:
if not line:
if data:
text = '\n'.join(data)
@@ -33,6 +38,18 @@ async def sse_events(response):
raise ValueError
except ValueError:
raise ValueError('Codex returned an invalid stream event') from None
if event.get('type') == 'response.completed':
# Finish the HTTP body so a healthy connection can return
# to the pool. Never wait indefinitely for a broken SSE
# server after it has already sent the terminal result.
try:
async with asyncio.timeout(0.05):
async for _ in lines:
pass
except (TimeoutError, httpx.HTTPError):
pass
yield event
return
yield event
data, size = [], 0
elif line.startswith('data:'):
@@ -85,6 +102,7 @@ def _tool(item):
def _usage(response):
usage = response.get('usage') or {}
return {
**usage,
'prompt_tokens': usage.get('input_tokens', 0),
'completion_tokens': usage.get('output_tokens', 0),
'total_tokens': usage.get('total_tokens', usage.get('input_tokens', 0) + usage.get('output_tokens', 0)),
@@ -94,7 +112,10 @@ def _usage(response):
class CodexRequester(requester.ProviderAPIRequester):
reports_raw_timing = True
async def initialize(self):
self._http_pool = RequesterHTTPPool()
self.auth = CodexAuth(self.ap)
self.workspace = self.requester_cfg['workspace_uuid']
self.provider = self.requester_cfg['provider_uuid']
@@ -104,6 +125,8 @@ class CodexRequester(requester.ProviderAPIRequester):
async def aclose(self):
self._replay.clear()
if getattr(self, '_http_pool', None) is not None:
await self._http_pool.aclose()
def get_reasoning_capabilities(self, model):
return {
@@ -184,6 +207,8 @@ class CodexRequester(requester.ProviderAPIRequester):
def _body(self, query, model, messages, funcs, extra_args, tokens):
args = {**(model.model_entity.extra_args or {}), **(extra_args or {})}
args.pop('request_policy', None)
args.pop('reasoning_config', None)
# Never permit credentials, transport overrides, store/history or arbitrary
# SDK kwargs to be smuggled through model advanced parameters.
allowed = {'reasoning', 'text', 'parallel_tool_calls', 'tool_choice'}
@@ -235,13 +260,11 @@ class CodexRequester(requester.ProviderAPIRequester):
'stream': True,
'include': ['reasoning.encrypted_content'],
}
reasoning_config = getattr(model, 'reasoning_config_override', None)
if reasoning_config is None:
reasoning_config = getattr(model.model_entity, 'reasoning_config', None)
level = reasoning.normalize_reasoning_config(reasoning_config)['level']
reasoning_config = reasoning.resolve_reasoning_config(model, extra_args)
level = reasoning_config['level']
if level != 'provider_default':
reasoning.validate_reasoning_capabilities(
{'level': level}, self.get_reasoning_capabilities(model), model.model_entity.name
reasoning_config, self.get_reasoning_capabilities(model), model.model_entity.name
)
body['reasoning'] = {'effort': level, 'summary': 'auto'}
if funcs:
@@ -259,13 +282,39 @@ class CodexRequester(requester.ProviderAPIRequester):
async def _events(self, query, model, messages, funcs, extra_args):
tokens = await self.auth.access(self.workspace, self.provider)
state = current_call.get()
timeout = 120
if state:
p = state.policy
timeout = httpx.Timeout(
connect=p.connect_timeout, read=p.read_timeout, write=p.write_timeout, pool=p.pool_timeout
)
try:
async with asyncio.timeout(300), httpx.AsyncClient(timeout=120, follow_redirects=False) as client:
async with (
asyncio.timeout(None if state else 300),
requester_http_client(self) as client,
):
for attempt in range(2):
body = self._body(query, model, messages, funcs, extra_args, tokens)
async with client.stream(
'POST', BASE_URL + '/responses', json=body, headers=self._headers(tokens, stream=True)
) as response:
if state:
state.mark('request_prepared_ms')
state.mark('request_started_ms')
async with AsyncExitStack() as stack:
response_context = client.stream(
'POST',
BASE_URL + '/responses',
json=body,
headers=self._headers(tokens, stream=True),
timeout=timeout,
)
if state:
response = await state.wait(
lambda: stack.enter_async_context(response_context),
'first_event',
state.policy.first_event_timeout,
)
else:
response = await stack.enter_async_context(response_context)
if response.status_code == 401 and attempt == 0:
tokens = await self.auth.access(
self.workspace, self.provider, rejected_token=tokens['access_token']
@@ -273,85 +322,117 @@ class CodexRequester(requester.ProviderAPIRequester):
continue
if response.status_code != 200:
raise await self._response_error(response)
async for event in sse_events(response):
if state:
state.mark('response_open_ms')
events = sse_events(response)
while True:
try:
if state:
event = await state.wait(
lambda: anext(events),
'stream_idle' if state.raw_started else 'first_event',
state.policy.stream_idle_timeout
if state.raw_started
else state.policy.first_event_timeout,
)
else:
event = await anext(events)
except StopAsyncIteration:
break
if state:
state.raw_started = True
state.mark('first_event_ms')
kind = event.get('type', '')
if kind == 'response.output_text.delta' and event.get('delta'):
state.mark('first_content_ms')
elif 'reasoning' in kind and event.get('delta'):
state.mark('first_reasoning_ms')
yield event, tokens
return
except (httpx.HTTPError, TimeoutError):
raise ValueError('ChatGPT Codex network error or timeout. Please retry.') from None
except (httpx.HTTPError, TimeoutError) as exc:
raise RequesterError(
'ChatGPT Codex network error or timeout. Please retry.',
code='timeout' if isinstance(exc, (httpx.TimeoutException, TimeoutError)) else 'connection',
retryable=True,
) from None
async def _chunks(self, query, model, messages, funcs, extra_args, remove_think, usage_out):
text = ''
seen_calls = set()
output_items = {}
response_id = None
async for event, tokens in self._events(query, model, messages, funcs, extra_args):
kind = event.get('type')
response = event.get('response') or {}
response_id = response.get('id') or response_id
if kind in ('error', 'response.failed', 'response.incomplete'):
raise CodexProviderError('ChatGPT Codex response failed or was incomplete. Please retry.')
if kind == 'response.output_text.delta':
delta = event.get('delta', '')
text += delta
yield pm.MessageChunk(role='assistant', content=delta, resp_message_id=response_id)
elif kind in ('response.reasoning_summary_text.delta', 'response.reasoning_text.delta'):
if not remove_think:
yield pm.MessageChunk(
role='assistant',
content='',
provider_specific_fields={'reasoning_content': event.get('delta', '')},
)
elif kind == 'response.output_item.done':
item = event.get('item') or {}
output_items[event.get('output_index', len(output_items))] = item
if item.get('type') == 'function_call' and item.get('call_id') not in seen_calls:
seen_calls.add(item.get('call_id'))
yield pm.MessageChunk(role='assistant', content='', tool_calls=[_tool(item)])
elif kind in ('response.completed', 'response.done'):
if response.get('status') not in (None, 'completed'):
raise CodexProviderError('ChatGPT Codex response was not completed')
output = response.get('output') or [output_items[k] for k in sorted(output_items)]
for item in output:
async with aclosing(self._events(query, model, messages, funcs, extra_args)) as events:
async for event, tokens in events:
kind = event.get('type')
response = event.get('response') or {}
response_id = response.get('id') or response_id
if kind in ('error', 'response.failed', 'response.incomplete'):
raise CodexProviderError('ChatGPT Codex response failed or was incomplete. Please retry.')
if kind == 'response.output_text.delta':
delta = event.get('delta', '')
text += delta
yield pm.MessageChunk(role='assistant', content=delta, resp_message_id=response_id)
elif kind in ('response.reasoning_summary_text.delta', 'response.reasoning_text.delta'):
if not remove_think:
yield pm.MessageChunk(
role='assistant',
content='',
provider_specific_fields={'reasoning_content': event.get('delta', '')},
)
elif kind == 'response.output_item.done':
item = event.get('item') or {}
output_items[event.get('output_index', len(output_items))] = item
if item.get('type') == 'function_call' and item.get('call_id') not in seen_calls:
seen_calls.add(item.get('call_id'))
yield pm.MessageChunk(role='assistant', content='', tool_calls=[_tool(item)])
# Some servers send only the terminal output, without text deltas.
final_text = ''.join(
p.get('text', '')
for item in output
if item.get('type') == 'message'
for p in item.get('content', [])
if p.get('type') == 'output_text'
)
if not text and final_text:
text = final_text
yield pm.MessageChunk(role='assistant', content=text, resp_message_id=response_id)
usage_out.update(_usage(response))
if query is not None:
if query.variables is None:
query.variables = {}
query.variables[requester.STREAM_USAGE_QUERY_VARIABLE] = dict(usage_out)
fields = None
if query is not None and output:
handle = secrets.token_urlsafe(24)
self._replay[handle] = (self._scope(query, model, tokens), time.time() + 3600, output)
while len(self._replay) > 64:
self._replay.popitem(last=False)
fields = {'codex_replay_id': handle}
yield pm.MessageChunk(
role='assistant',
content='',
all_content=text,
is_final=True,
resp_message_id=response_id,
provider_specific_fields=fields,
)
return
elif kind in ('response.completed', 'response.done'):
if response.get('status') not in (None, 'completed'):
raise CodexProviderError('ChatGPT Codex response was not completed')
output = response.get('output') or [output_items[k] for k in sorted(output_items)]
for item in output:
if item.get('type') == 'function_call' and item.get('call_id') not in seen_calls:
seen_calls.add(item.get('call_id'))
yield pm.MessageChunk(role='assistant', content='', tool_calls=[_tool(item)])
# Some servers send only the terminal output, without text deltas.
final_text = ''.join(
p.get('text', '')
for item in output
if item.get('type') == 'message'
for p in item.get('content', [])
if p.get('type') == 'output_text'
)
if not text and final_text:
text = final_text
yield pm.MessageChunk(role='assistant', content=text, resp_message_id=response_id)
usage_out.update(_usage(response))
if current_call.get():
current_call.get().usage = dict(usage_out)
if query is not None:
if query.variables is None:
query.variables = {}
query.variables[requester.STREAM_USAGE_QUERY_VARIABLE] = dict(usage_out)
fields = None
if query is not None and output:
handle = secrets.token_urlsafe(24)
self._replay[handle] = (self._scope(query, model, tokens), time.time() + 3600, output)
while len(self._replay) > 64:
self._replay.popitem(last=False)
fields = {'codex_replay_id': handle}
yield pm.MessageChunk(
role='assistant',
content='',
all_content=text,
is_final=True,
resp_message_id=response_id,
provider_specific_fields=fields,
)
return
raise CodexProviderError('ChatGPT Codex stream ended before completion. Please retry.')
async def invoke_llm_stream(self, query, model, messages, funcs=None, extra_args=None, remove_think=False):
async for chunk in self._chunks(query, model, messages, funcs, extra_args, remove_think, {}):
yield chunk
async with aclosing(self._chunks(query, model, messages, funcs, extra_args, remove_think, {})) as chunks:
async for chunk in chunks:
yield chunk
async def invoke_llm(self, query, model, messages, funcs=None, extra_args=None, remove_think=False):
usage = {}
@@ -376,11 +457,12 @@ class CodexRequester(requester.ProviderAPIRequester):
async def scan_models(self, api_key=None):
tokens = await self.auth.access(self.workspace, self.provider)
try:
async with asyncio.timeout(90), httpx.AsyncClient(timeout=30, follow_redirects=False) as client:
async with asyncio.timeout(90), requester_http_client(self) as client:
for attempt in range(2):
response = await client.get(
BASE_URL + '/models',
params={'client_version': langbot.__version__},
timeout=30,
headers=self._headers(tokens),
)
if response.status_code == 401 and attempt == 0:
src/langbot/pkg/provider/modelmgr/requesters/litellmchat.py
+463 -172
View File
@@ -5,10 +5,23 @@ from __future__ import annotations
from ....i18n import t
import typing
import copy
import time
import os
from collections import OrderedDict
from contextlib import asynccontextmanager
from openai import AsyncOpenAI
from ..http_pool import RequesterHTTPPool
from ..prompt_caching import apply_prompt_caching
import asyncio
import httpx
from urllib.parse import urlsplit
from ..request_policy import current_call, RequestPolicy, error_metadata
import litellm
from litellm import acompletion, aembedding, arerank
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
from litellm.secret_managers.main import get_secret, get_secret_bool
from .. import errors, reasoning, requester
from ....utils import httpclient
@@ -17,6 +30,37 @@ import langbot_plugin.api.entities.builtin.pipeline.query as pipeline_query
import langbot_plugin.api.entities.builtin.provider.message as provider_message
class _LeasedHTTPHandler(AsyncHTTPHandler):
"""Adapt a leased transport without LiteLLM's hidden connection-error retry."""
def __init__(self, client, timeout):
self.client = client
self.timeout = timeout
self.event_hooks = None
self.client_alias = 'langbot'
async def post(self, url, *, timeout=None, stream=False, logging_obj=None, **kwargs):
if isinstance(kwargs.get('data'), (str, bytes)):
kwargs['content'] = kwargs.pop('data')
request = self.client.build_request('POST', url, timeout=timeout or self.timeout, **kwargs)
response = await self.client.send(request, stream=stream)
try:
response.raise_for_status()
except BaseException:
await response.aclose()
raise
return response
async def close(self):
pass # The requester's lease owns the transport.
async def __aexit__(self, *args):
pass
def __del__(self):
pass
class _ThinkStripState:
"""Drop explicit think blocks, buffering only possible split tag prefixes.
@@ -133,6 +177,8 @@ class _ThinkStripState:
class LiteLLMRequester(requester.ProviderAPIRequester):
"""LiteLLM unified API requester supporting chat, embedding, and rerank."""
reports_raw_timing = True
_EMBEDDING_MODEL_HINTS = ('embedding', 'embed', 'bge-', 'e5-', 'm3e', 'gte-', 'text-embedding')
_RERANK_MODEL_HINTS = ('rerank', 're-rank', 're_rank')
_QWEN_DEDICATED_THINKING_MODELS = frozenset(
@@ -181,9 +227,83 @@ class LiteLLMRequester(requester.ProviderAPIRequester):
async def initialize(self):
"""Initialize LiteLLM client settings."""
# LiteLLM doesn't require explicit client initialization
# Configuration is passed per-request via litellm params
pass
self._http_pool = RequesterHTTPPool(
follow_redirects=True, drain_sse=self._get_custom_llm_provider() == 'openai'
)
async def aclose(self):
if getattr(self, '_http_pool', None) is not None:
await self._http_pool.aclose()
if hasattr(self, '_metadata_cache'):
self._metadata_cache.clear()
@asynccontextmanager
async def _client_args(self, args):
pool = getattr(self, '_http_pool', None)
provider = args.get('custom_llm_provider') or self._get_custom_llm_provider()
if pool is None or provider not in ('openai', 'openrouter') or args.get('client') is not None:
yield args
return
# Preserve explicitly customized LiteLLM transports and TLS settings.
if (
any(args.get(key) is not None for key in ('shared_session', 'http_client', 'ssl_verify', 'ssl_certificate'))
or getattr(litellm, 'aclient_session', None) is not None
or getattr(litellm, 'ssl_verify', True) is not True
or getattr(litellm, 'ssl_certificate', None)
or getattr(litellm, 'force_ipv4', False)
or any(name in os.environ for name in ('SSL_VERIFY', 'SSL_CERTIFICATE'))
):
yield args
return
if not args['model'].startswith(f'{provider}/'):
yield args # A caller redirected the model; let LiteLLM resolve its transport.
return
model_name = args['model'].removeprefix(f'{provider}/')
if provider == 'openrouter' and model_name in litellm.open_ai_chat_completion_models:
yield args # Preserve LiteLLM's alternate bare-model routing.
return
transport_started = time.perf_counter()
async with pool.lease() as client:
state = current_call.get()
if state:
state.metrics['transport_setup_ms'] = round(
state.metrics.get('transport_setup_ms', 0) + (time.perf_counter() - transport_started) * 1000, 3
)
if provider == 'openrouter' or get_secret_bool('EXPERIMENTAL_OPENAI_BASE_LLM_HTTP_HANDLER'):
yield {**args, 'client': _LeasedHTTPHandler(client, args.get('timeout', 120))}
return
# A fresh lightweight SDK wrapper keeps rotating keys, organization
# and deadlines request-local; only its HTTP transport is reused.
sdk = AsyncOpenAI(
api_key=args.get('api_key') or litellm.api_key or litellm.openai_key or get_secret('OPENAI_API_KEY'),
base_url=(
args.get('api_base')
or litellm.api_base
or get_secret('OPENAI_BASE_URL')
or get_secret('OPENAI_API_BASE')
or 'https://api.openai.com/v1'
),
http_client=client,
organization=args.get('organization') or litellm.organization,
timeout=args.get('timeout', 120),
max_retries=0,
)
yield {**args, 'client': sdk}
def _cached_metadata(self, key, compute):
if not hasattr(self, '_metadata_cache'):
self._metadata_cache = OrderedDict()
cache = self._metadata_cache
cached = cache.get(key)
if cached is not None and time.monotonic() - cached[0] < 300:
cache.move_to_end(key)
return copy.deepcopy(cached[1])
value = compute()
cache[key] = (time.monotonic(), copy.deepcopy(value))
cache.move_to_end(key)
while len(cache) > 128:
cache.popitem(last=False)
return value
def _build_litellm_model_name(self, model_name: str, custom_llm_provider: str | None = None) -> str:
"""Build LiteLLM model name with provider prefix if needed."""
@@ -199,8 +319,21 @@ class LiteLLMRequester(requester.ProviderAPIRequester):
def _get_custom_llm_provider(self) -> str | None:
return self.requester_cfg.get('custom_llm_provider') or None
def _uses_openrouter_protocol(self) -> bool:
# OpenRouter's shipped metadata also serves OpenAI-compatible non-chat
# operations. Select gateway chat semantics independently of it.
return (
self._get_custom_llm_provider() == 'openrouter'
or self.requester_cfg.get('requester_name') == 'openrouter-chat-completions'
)
def _safe_litellm_bool_helper(self, helper_name: str, model_name: str) -> bool:
"""Call a LiteLLM boolean capability helper without letting metadata gaps fail requests."""
return self._cached_metadata(
('bool', helper_name, model_name), lambda: self._uncached_litellm_bool_helper(helper_name, model_name)
)
def _uncached_litellm_bool_helper(self, helper_name, model_name):
helper = getattr(litellm, helper_name, None)
if not callable(helper):
return False
@@ -299,6 +432,9 @@ class LiteLLMRequester(requester.ProviderAPIRequester):
return None
def _safe_context_length(self, model_name: str) -> int | None:
return self._cached_metadata(('model', model_name), lambda: self._uncached_model_info(model_name))
def _uncached_model_info(self, model_name):
helper = getattr(litellm, 'get_model_info', None)
if not callable(helper):
return self._known_context_length_fallback(model_name)
@@ -563,6 +699,20 @@ class LiteLLMRequester(requester.ProviderAPIRequester):
}
if family == 'qwen' and 'disabled' in capabilities['levels'] and 'enabled' not in capabilities['levels']:
capabilities['legacy_levels'] = ['enabled']
if self._uses_openrouter_protocol():
capabilities['budget_tokens'] = {'min': 1, 'max': 32768}
elif family == 'anthropic' and not self._is_anthropic_adaptive_model(normalized_name):
if normalized_name.rsplit('/', 1)[-1].startswith(
('claude-3-7-', 'claude-sonnet-4', 'claude-opus-4', 'claude-haiku-4')
):
capabilities['budget_tokens'] = {'min': 1024, 'max': 32768}
elif family == 'qwen' and self._supports_qwen_thinking_budget(model_name):
capabilities['budget_tokens'] = {'min': 1, 'max': 32768}
elif family == 'gemini' and 'gemini-2.5' in normalized_name:
capabilities['budget_tokens'] = {
'min': 128 if 'pro' in normalized_name else (512 if 'flash-lite' in normalized_name else 1),
'max': 32768 if 'pro' in normalized_name else 24576,
}
return capabilities
def _build_reasoning_args(self, model: requester.RuntimeLLMModel) -> dict[str, typing.Any]:
@@ -570,14 +720,21 @@ class LiteLLMRequester(requester.ProviderAPIRequester):
if level == 'provider_default':
return {}
config = {'level': level}
config = reasoning.normalize_reasoning_config(self._reasoning_config_value(model))
capabilities = self.get_reasoning_capabilities(model)
try:
reasoning.validate_reasoning_capabilities(config, capabilities, model.model_entity.name)
except ValueError as exc:
raise errors.RequesterError(str(exc)) from exc
raise errors.RequesterError(str(exc), code='invalid_policy') from exc
family = self._reasoning_family(model.model_entity.name, model)
budget = config.get('budget_tokens')
if budget is not None:
if self._uses_openrouter_protocol():
return {'extra_body': {'reasoning': {'max_tokens': budget}}}
if family == 'qwen':
return {'extra_body': {'enable_thinking': True, 'thinking_budget': budget}}
return {'thinking': {'type': 'enabled', 'budget_tokens': budget}}
if level == 'disabled':
if family in {'deepseek', 'kimi', 'mimo', 'doubao'}:
return {'extra_body': {'thinking': {'type': 'disabled'}}}
@@ -1029,7 +1186,7 @@ class LiteLLMRequester(requester.ProviderAPIRequester):
if include_retry_params:
if self.requester_cfg.get('drop_params'):
args['drop_params'] = self.requester_cfg['drop_params']
if self.requester_cfg.get('num_retries'):
if self.requester_cfg.get('num_retries') is not None:
args['num_retries'] = self.requester_cfg['num_retries']
if self.requester_cfg.get('api_version'):
args['api_version'] = self.requester_cfg['api_version']
@@ -1037,24 +1194,102 @@ class LiteLLMRequester(requester.ProviderAPIRequester):
def _handle_litellm_error(self, e: Exception) -> None:
"""Convert LiteLLM exceptions to RequesterError. Never returns, always raises."""
# Check more specific exceptions first (they inherit from base exceptions)
if isinstance(e, litellm.ContextWindowExceededError):
raise errors.RequesterError(t('Context length exceeded') + f': {e}')
if isinstance(e, litellm.BadRequestError):
raise errors.RequesterError(t('Invalid request parameters') + f': {e}')
if isinstance(e, litellm.AuthenticationError):
raise errors.RequesterError(t('Invalid API key') + f': {e}')
if isinstance(e, litellm.NotFoundError):
raise errors.RequesterError(t('Invalid model or endpoint') + f': {e}')
if isinstance(e, litellm.RateLimitError):
raise errors.RequesterError(t('Rate limit exceeded or insufficient balance') + f': {e}')
if isinstance(e, litellm.Timeout):
raise errors.RequesterError(t('Request timed out') + f': {e}')
if isinstance(e, litellm.APIConnectionError):
raise errors.RequesterError(t('Connection failed') + f': {e}')
if isinstance(e, litellm.APIError):
raise errors.RequesterError(t('API error') + f': {e}')
raise errors.RequesterError(t('Unknown error') + f': {e}')
if isinstance(e, errors.RequesterError):
raise e
metadata = error_metadata(e)
mappings = (
(litellm.ContextWindowExceededError, 'context_length', t('Context length exceeded'), False),
(litellm.AuthenticationError, 'authentication', t('Invalid API key'), False),
(litellm.NotFoundError, 'not_found', t('Invalid model or endpoint'), False),
(litellm.RateLimitError, 'rate_limit', t('Rate limit exceeded or insufficient balance'), True),
(litellm.Timeout, 'timeout', t('Request timed out'), True),
(litellm.APIConnectionError, 'connection', t('Connection failed'), True),
(litellm.BadRequestError, 'bad_request', t('Invalid request parameters'), False),
)
for exception_type, code, message, retryable in mappings:
if isinstance(e, exception_type):
body = getattr(e, 'body', None)
body = body if isinstance(body, dict) else {}
upstream_error = body.get('error', body)
upstream_code = upstream_error.get('code') if isinstance(upstream_error, dict) else None
if upstream_code in ('insufficient_quota', 'quota_exceeded', 'billing_hard_limit_reached'):
code, retryable = 'quota_exceeded', False
raise errors.RequesterError(message, code=code, retryable=retryable, **metadata) from e
status = metadata['status_code']
raise errors.RequesterError(
t('API error') if isinstance(e, litellm.APIError) else t('Unknown error'),
code='upstream_error',
retryable=status in (408, 409, 429, 500, 502, 503, 504),
**metadata,
) from e
def _apply_transport(self, args, model, extra_args):
state = current_call.get()
policy = (
state.policy
if state
else RequestPolicy.resolve(self.requester_cfg, model.model_entity.extra_args, extra_args)
)
args.pop('request_policy', None)
if state:
# The Host owns retries; do not multiply attempts in nested SDKs.
args['num_retries'] = 0
args['max_retries'] = 0
args['timeout'] = httpx.Timeout(
connect=policy.connect_timeout,
read=policy.read_timeout,
write=policy.write_timeout,
pool=policy.pool_timeout,
)
state.mark('request_prepared_ms')
base = args.get('api_base', '')
is_openrouter = urlsplit(base).hostname == 'openrouter.ai' or (
not base and self._get_custom_llm_provider() == 'openrouter'
)
headers = dict(args.get('extra_headers') or {})
existing = {key.lower() for key in headers}
if policy.attribution and is_openrouter:
if 'http-referer' not in existing:
headers['HTTP-Referer'] = 'https://langbot.app'
if not {'x-title', 'x-openrouter-title'} & existing:
headers['X-OpenRouter-Title'] = 'LangBot'
conversation_header = policy.conversation_header or ('x-session-id' if is_openrouter else None)
if state and state.conversation_id and conversation_header and conversation_header not in existing:
headers[conversation_header] = state.conversation_id
if headers:
args['extra_headers'] = headers
async def _completion(self, args):
state = current_call.get()
if state:
state.mark('request_started_ms')
result = await state.wait(
lambda: acompletion(**args),
'first_event' if args.get('stream') else 'response',
state.policy.first_event_timeout if args.get('stream') else state.policy.total_timeout,
)
state.mark('response_open_ms')
return result
return await acompletion(**args)
async def _raw_chunks(self, response):
state = current_call.get()
iterator = response.__aiter__()
while True:
try:
if state:
chunk = await state.wait(
lambda: anext(iterator),
'stream_idle' if state.raw_started else 'first_event',
state.policy.stream_idle_timeout if state.raw_started else state.policy.first_event_timeout,
)
state.raw_started = True
state.mark('first_event_ms')
else:
chunk = await anext(iterator)
except StopAsyncIteration:
return
yield chunk
async def _build_completion_args(
self,
@@ -1065,14 +1300,31 @@ class LiteLLMRequester(requester.ProviderAPIRequester):
stream: bool = False,
) -> dict:
"""Build common completion arguments for invoke_llm and invoke_llm_stream."""
preparation_start = time.perf_counter()
state = current_call.get()
cache_preparation = state is not None and state.policy.retries > 0
scoped_model = copy.copy(model)
try:
scoped_model.reasoning_config_override = reasoning.resolve_reasoning_config(model, extra_args)
except ValueError as exc:
raise errors.RequesterError(str(exc), code='invalid_policy') from exc
model = scoped_model
reasoning_family = self._reasoning_family(model.model_entity.name, model)
reasoning_level = self._reasoning_level(model)
req_messages = self._convert_messages(
messages,
reasoning_family=reasoning_family,
include_reasoning_context=reasoning_level != 'disabled',
)
model_name = self._build_litellm_model_name(model.model_entity.name)
message_key = ('messages', id(messages), reasoning_family, reasoning_level != 'disabled')
if cache_preparation and message_key in state.prepared:
req_messages = copy.deepcopy(state.prepared[message_key])
state.metrics['preparation_cache_hits'] = state.metrics.get('preparation_cache_hits', 0) + 1
else:
req_messages = self._convert_messages(
messages, reasoning_family=reasoning_family, include_reasoning_context=reasoning_level != 'disabled'
)
if cache_preparation:
state.prepared[message_key] = copy.deepcopy(req_messages)
# OpenRouter shares requester metadata with OpenAI-compatible non-chat
# operations, but chat needs its own adapter to preserve cache markers.
chat_provider = 'openrouter' if self._uses_openrouter_protocol() else self._get_custom_llm_provider()
model_name = self._build_litellm_model_name(model.model_entity.name, chat_provider)
api_key = model.provider.token_mgr.get_token()
args = {
@@ -1085,10 +1337,14 @@ class LiteLLMRequester(requester.ProviderAPIRequester):
args['stream_options'] = {'include_usage': True}
self._build_common_args(args)
if chat_provider == 'openrouter':
args['custom_llm_provider'] = chat_provider
# Apply model-level extra_args first, then call-level extra_args
if model.model_entity.extra_args:
args.update(model.model_entity.extra_args)
args.update(extra_args)
args.pop('reasoning_config', None)
reasoning_args = self._build_reasoning_args(model)
if reasoning_args:
@@ -1113,8 +1369,23 @@ class LiteLLMRequester(requester.ProviderAPIRequester):
raise errors.RequesterError('allowed_openai_params must be an array')
args['allowed_openai_params'] = list(dict.fromkeys([*allowed_openai_params, 'reasoning_effort']))
budget = model.reasoning_config_override.get('budget_tokens')
if budget is not None and reasoning_family == 'anthropic' and not self._uses_openrouter_protocol():
output_limit = args.get('max_completion_tokens', args.get('max_tokens'))
if type(output_limit) is not int or output_limit <= budget:
raise errors.RequesterError(
'Anthropic budget_tokens requires explicit max_tokens greater than the budget',
code='invalid_policy',
)
if funcs:
tools = await self.ap.tool_mgr.generate_tools_for_openai(funcs)
tool_key = ('tools', id(funcs))
if cache_preparation and tool_key in state.prepared:
tools = copy.deepcopy(state.prepared[tool_key])
else:
tools = copy.deepcopy(await self.ap.tool_mgr.generate_tools_for_openai(funcs))
if cache_preparation:
state.prepared[tool_key] = copy.deepcopy(tools)
if tools:
args['tools'] = tools
args.setdefault('tool_choice', 'auto')
@@ -1123,6 +1394,15 @@ class LiteLLMRequester(requester.ProviderAPIRequester):
# and imports proxy-only dependencies such as FastAPI.
args['_skip_mcp_handler'] = True
cache_family = reasoning_family
if cache_family == 'openrouter':
cache_family = self._infer_reasoning_family_from_model_name(model.model_entity.name) or cache_family
apply_prompt_caching(args, state, provider=self._get_custom_llm_provider(), family=cache_family)
self._apply_transport(args, model, extra_args)
if state:
state.metrics['preparation_ms'] = round(
state.metrics.get('preparation_ms', 0) + (time.perf_counter() - preparation_start) * 1000, 3
)
return args
async def count_tokens(
@@ -1144,7 +1424,7 @@ class LiteLLMRequester(requester.ProviderAPIRequester):
count_args['tool_choice'] = args['tool_choice']
try:
tokens = litellm.token_counter(**count_args)
tokens = await asyncio.to_thread(litellm.token_counter, **count_args)
except Exception as e:
self._handle_litellm_error(e)
@@ -1164,38 +1444,41 @@ class LiteLLMRequester(requester.ProviderAPIRequester):
"""Invoke LLM and return message with usage info."""
args = await self._build_completion_args(model, messages, funcs, extra_args, stream=False)
try:
response = await acompletion(**args)
async with self._client_args(args) as args:
try:
response = await self._completion(args)
message_data = response.choices[0].message.model_dump()
if 'role' not in message_data or message_data['role'] is None:
message_data['role'] = 'assistant'
message_data = response.choices[0].message.model_dump()
if 'role' not in message_data or message_data['role'] is None:
message_data['role'] = 'assistant'
content = message_data.get('content', '')
reasoning_content = message_data.get('reasoning_content', None)
thinking_blocks = message_data.get('thinking_blocks')
if reasoning_content or thinking_blocks:
provider_fields = dict(message_data.get('provider_specific_fields') or {})
if reasoning_content:
provider_fields['reasoning_content'] = reasoning_content
if thinking_blocks:
provider_fields['thinking_blocks'] = thinking_blocks
message_data['provider_specific_fields'] = provider_fields
display_reasoning = reasoning_content or self._thinking_blocks_text(thinking_blocks) or None
message_data['content'] = self._process_thinking_content(content, display_reasoning, remove_think)
content = message_data.get('content', '')
reasoning_content = message_data.get('reasoning_content', None)
thinking_blocks = message_data.get('thinking_blocks')
if reasoning_content or thinking_blocks:
provider_fields = dict(message_data.get('provider_specific_fields') or {})
if reasoning_content:
provider_fields['reasoning_content'] = reasoning_content
if thinking_blocks:
provider_fields['thinking_blocks'] = thinking_blocks
message_data['provider_specific_fields'] = provider_fields
display_reasoning = reasoning_content or self._thinking_blocks_text(thinking_blocks) or None
message_data['content'] = self._process_thinking_content(content, display_reasoning, remove_think)
if 'reasoning_content' in message_data:
del message_data['reasoning_content']
if 'thinking_blocks' in message_data:
del message_data['thinking_blocks']
if 'reasoning_content' in message_data:
del message_data['reasoning_content']
if 'thinking_blocks' in message_data:
del message_data['thinking_blocks']
message = provider_message.Message(**message_data)
usage_info = self._extract_usage(response)
message = provider_message.Message(**message_data)
usage_info = self._extract_usage(response)
if current_call.get():
current_call.get().usage = usage_info
return message, usage_info
return message, usage_info
except Exception as e:
self._handle_litellm_error(e)
except Exception as e:
self._handle_litellm_error(e)
async def invoke_llm_stream(
self,
@@ -1218,141 +1501,149 @@ class LiteLLMRequester(requester.ProviderAPIRequester):
thinking_blocks_state: list[dict[str, typing.Any]] = []
response = None
try:
response = await acompletion(**args)
async for chunk in response:
# Capture usage whenever a chunk carries it.
#
# Important: many OpenAI-compatible gateways (e.g. new-api) and
# providers send the final usage payload in a chunk that STILL
# contains a (empty-delta) choice, not an empty `choices` list.
# The previous implementation only captured usage when `choices`
# was empty, so streamed calls always recorded 0 tokens.
# We therefore capture usage independently of `choices`, and then
# fall through to also process any content this chunk may carry.
if getattr(chunk, 'usage', None):
usage_info = self._normalize_usage(chunk.usage)
if query is not None:
if query.variables is None:
query.variables = {}
query.variables[requester.STREAM_USAGE_QUERY_VARIABLE] = usage_info
async with self._client_args(args) as args:
try:
response = await self._completion(args)
async for chunk in self._raw_chunks(response):
# Capture usage whenever a chunk carries it.
#
# Important: many OpenAI-compatible gateways (e.g. new-api) and
# providers send the final usage payload in a chunk that STILL
# contains a (empty-delta) choice, not an empty `choices` list.
# The previous implementation only captured usage when `choices`
# was empty, so streamed calls always recorded 0 tokens.
# We therefore capture usage independently of `choices`, and then
# fall through to also process any content this chunk may carry.
if getattr(chunk, 'usage', None):
if current_call.get():
current_call.get().usage = self._extract_usage(chunk)
usage_info = self._normalize_usage(chunk.usage)
if query is not None:
if query.variables is None:
query.variables = {}
query.variables[requester.STREAM_USAGE_QUERY_VARIABLE] = usage_info
if not hasattr(chunk, 'choices') or not chunk.choices:
continue
if not hasattr(chunk, 'choices') or not chunk.choices:
continue
choice = chunk.choices[0]
delta = choice.delta.model_dump() if hasattr(choice, 'delta') else {}
finish_reason = getattr(choice, 'finish_reason', None)
choice = chunk.choices[0]
delta = choice.delta.model_dump() if hasattr(choice, 'delta') else {}
finish_reason = getattr(choice, 'finish_reason', None)
if 'role' in delta and delta['role']:
role = delta['role']
if 'role' in delta and delta['role']:
role = delta['role']
delta_content = delta.get('content') or ''
reasoning_content = delta.get('reasoning_content') or ''
provider_fields = dict(delta.get('provider_specific_fields') or {})
if finish_reason:
# Preserve an explicit provider stop even when the final delta is empty.
provider_fields['finish_reason'] = finish_reason
raw_thinking_blocks = delta.get('thinking_blocks')
if raw_thinking_blocks:
thinking_blocks_state = self._merge_thinking_blocks(thinking_blocks_state, raw_thinking_blocks)
provider_fields['thinking_blocks'] = thinking_blocks_state
thinking_blocks_text = self._thinking_blocks_text(raw_thinking_blocks)
display_reasoning_content = reasoning_content or thinking_blocks_text
state = current_call.get()
if state and delta.get('content'):
state.mark('first_content_ms')
if state and (delta.get('reasoning_content') or delta.get('thinking_blocks')):
state.mark('first_reasoning_ms')
delta_content = delta.get('content') or ''
reasoning_content = delta.get('reasoning_content') or ''
provider_fields = dict(delta.get('provider_specific_fields') or {})
if finish_reason:
# Preserve an explicit provider stop even when the final delta is empty.
provider_fields['finish_reason'] = finish_reason
raw_thinking_blocks = delta.get('thinking_blocks')
if raw_thinking_blocks:
thinking_blocks_state = self._merge_thinking_blocks(thinking_blocks_state, raw_thinking_blocks)
provider_fields['thinking_blocks'] = thinking_blocks_state
thinking_blocks_text = self._thinking_blocks_text(raw_thinking_blocks)
display_reasoning_content = reasoning_content or thinking_blocks_text
# Handle reasoning_content based on remove_think flag
if reasoning_content:
provider_fields['reasoning_content'] = reasoning_content
if remove_think:
delta_content = delta_content or None
else:
# Stream explicit markers so downstream adapters and
# the debug page see the same format as non-streaming
# responses.
if not reasoning_started:
delta_content = '<think>\n'
reasoning_started = True
# Handle reasoning_content based on remove_think flag
if reasoning_content:
provider_fields['reasoning_content'] = reasoning_content
if remove_think:
delta_content = delta_content or None
else:
delta_content = ''
delta_content += display_reasoning_content
if delta.get('content'):
delta_content += f'\n</think>\n{delta.get("content")}'
reasoning_closed = True
# Stream explicit markers so downstream adapters and
# the debug page see the same format as non-streaming
# responses.
if not reasoning_started:
delta_content = '<think>\n'
reasoning_started = True
else:
delta_content = ''
delta_content += display_reasoning_content
if delta.get('content'):
delta_content += f'\n</think>\n{delta.get("content")}'
reasoning_closed = True
elif display_reasoning_content:
if remove_think:
delta_content = delta_content or None
else:
if not reasoning_started:
delta_content = '<think>\n'
reasoning_started = True
elif display_reasoning_content:
if remove_think:
delta_content = delta_content or None
else:
delta_content = ''
delta_content += display_reasoning_content
if delta.get('content'):
delta_content += f'\n</think>\n{delta.get("content")}'
reasoning_closed = True
if not reasoning_started:
delta_content = '<think>\n'
reasoning_started = True
else:
delta_content = ''
delta_content += display_reasoning_content
if delta.get('content'):
delta_content += f'\n</think>\n{delta.get("content")}'
reasoning_closed = True
elif delta_content and not remove_think and reasoning_started and not reasoning_closed:
delta_content = f'\n</think>\n{delta_content}'
reasoning_closed = True
elif delta_content and not remove_think and reasoning_started and not reasoning_closed:
delta_content = f'\n</think>\n{delta_content}'
reasoning_closed = True
if finish_reason and not remove_think and reasoning_started and not reasoning_closed:
delta_content = f'{delta_content}\n</think>\n'
reasoning_closed = True
if finish_reason and not remove_think and reasoning_started and not reasoning_closed:
delta_content = f'{delta_content}\n</think>\n'
reasoning_closed = True
if think_state is not None and delta_content:
delta_content = think_state.feed(delta_content)
if think_state is not None and delta_content:
delta_content = think_state.feed(delta_content)
tool_calls = self._normalize_stream_tool_calls(delta.get('tool_calls'), tool_call_state)
tool_calls = self._normalize_stream_tool_calls(delta.get('tool_calls'), tool_call_state)
if not delta_content and not tool_calls and not provider_fields and not finish_reason:
if not delta_content and not tool_calls and not provider_fields and not finish_reason:
chunk_idx += 1
continue
chunk_data: dict[str, typing.Any] = {
'role': role,
'content': delta_content if delta_content else None,
'tool_calls': tool_calls,
'is_final': bool(finish_reason),
}
# Preserve provider_specific_fields from delta (e.g., Gemini thought_signatures)
if provider_fields:
chunk_data['provider_specific_fields'] = provider_fields
chunk_data = {k: v for k, v in chunk_data.items() if v is not None}
yield provider_message.MessageChunk(**chunk_data)
chunk_idx += 1
continue
chunk_data: dict[str, typing.Any] = {
'role': role,
'content': delta_content if delta_content else None,
'tool_calls': tool_calls,
'is_final': bool(finish_reason),
}
# Preserve provider_specific_fields from delta (e.g., Gemini thought_signatures)
if provider_fields:
chunk_data['provider_specific_fields'] = provider_fields
chunk_data = {k: v for k, v in chunk_data.items() if v is not None}
yield provider_message.MessageChunk(**chunk_data)
chunk_idx += 1
if reasoning_started and not reasoning_closed:
yield provider_message.MessageChunk(
role=role,
content='\n</think>\n',
is_final=True,
)
if think_state is not None:
pending_content = think_state.flush()
if pending_content:
if reasoning_started and not reasoning_closed:
yield provider_message.MessageChunk(
role=role,
content=pending_content,
content='\n</think>\n',
is_final=True,
)
except Exception as e:
self._handle_litellm_error(e)
if think_state is not None:
pending_content = think_state.flush()
if pending_content:
yield provider_message.MessageChunk(
role=role,
content=pending_content,
is_final=True,
)
finally:
# Closing a downstream async generator must release the provider's
# HTTP stream immediately, including cancellation while at yield.
close = getattr(response, 'aclose', None)
if callable(close):
try:
await asyncio.wait_for(close(), timeout=5.0)
except Exception as exc:
self.ap.logger.debug(f'Failed to close LLM stream: {type(exc).__name__}')
except Exception as e:
self._handle_litellm_error(e)
finally:
# Closing a downstream async generator must release the provider's
# HTTP stream immediately, including cancellation while at yield.
close = getattr(response, 'aclose', None)
if callable(close):
try:
await asyncio.wait_for(close(), timeout=5.0)
except Exception as exc:
self.ap.logger.debug(f'Failed to close LLM stream: {type(exc).__name__}')
async def invoke_embedding(
self,
src/langbot/templates/config.yaml
+17
View File
@@ -458,3 +458,20 @@ space:
# Complete execution records are batched as {"records": [...]}; records are
# removed from the buffer only once the sender accepted them.
telemetry_flush_seconds: 180
# Host-owned LLM policy. Models/calls may override through extra_args.request_policy.
llm_request_policy:
retries: 0
total_timeout: 300
first_event_timeout: 120
stream_idle_timeout: 120
connect_timeout: 10
read_timeout: 120
write_timeout: 30
pool_timeout: 10
backoff_base: 0.5
backoff_max: 8
attribution: true
# Optional: x-session-id, x-litellm-session-id, x-langbot-conversation-id.
# OpenRouter uses x-session-id automatically; other providers default to none.
conversation_header: null
tests/e2e/test_llm_request_chain.py
+442
View File
@@ -0,0 +1,442 @@
"""Real HTTP -> Core -> Plugin Runtime -> Runner -> Core -> HTTP provider tests.
Only the remote provider is deterministic. No requester, SDK action transport,
model manager, persistence or monitoring implementation is mocked.
"""
import json
import importlib.metadata
import os
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path
import shutil
import select
import socket
import sqlite3
import threading
import time
import pytest
import yaml
import httpx
from packaging.version import Version
from tests.e2e import test_runner_plugin_runtime as runtime
# Reuse the real process lifecycle; override the package/config and port below.
runner_runtime_ports = runtime.runner_runtime_ports
runner_runtime_process = runtime.runner_runtime_process
runner_langbot_process = runtime.runner_langbot_process
runner_client = runtime.runner_client
pytestmark = pytest.mark.e2e
RUNNER = runtime.QA_RUNNER_ID
USAGE = {
'prompt_tokens': 24,
'completion_tokens': 4,
'total_tokens': 28,
'prompt_tokens_details': {'cached_tokens': 16},
'completion_tokens_details': {'reasoning_tokens': 2},
}
@pytest.fixture(scope='session')
def runner_e2e_tmpdir(tmp_path_factory):
# Retain logs/SQLite in pytest's bounded temp history for failure diagnosis.
return tmp_path_factory.mktemp('llm-chain')
@pytest.fixture(scope='session')
def runner_e2e_port():
return runtime._free_port()
@pytest.fixture(scope='session')
def runner_e2e_config_path(runner_e2e_tmpdir, runner_e2e_port, runner_runtime_ports):
path = runtime.runner_e2e_config_path.__wrapped__(runner_e2e_tmpdir, runner_e2e_port, runner_runtime_ports)
root = runner_e2e_tmpdir / 'agent-runner-qa-package'
manifest = root / 'components/runner/default.yaml'
data = yaml.safe_load(manifest.read_text(encoding='utf-8'))
data['spec']['permissions'] = {'models': ['invoke', 'stream', 'count_tokens']}
data['spec']['config'] = [
{'name': 'model', 'label': {'en_US': 'Model'}, 'type': 'llm-model-selector', 'required': True},
]
manifest.write_text(yaml.safe_dump(data), encoding='utf-8')
shutil.copyfile(Path(__file__).parent / 'utils/llm_probe_runner.py', root / 'components/runner/default.py')
if legacy_root := os.environ.get('LLM_E2E_LEGACY_API_DIR'):
for filename in ('legacy_api.py', 'legacy_message.py'):
shutil.copyfile(Path(legacy_root) / filename, root / filename)
shutil.make_archive(str(runner_e2e_tmpdir / 'agent-runner-qa'), 'zip', root_dir=root)
return path
@pytest.fixture(scope='module')
def provider_wire():
records = []
class Handler(BaseHTTPRequestHandler):
protocol_version = 'HTTP/1.1'
def log_message(self, *args):
pass
def do_POST(self):
body = json.loads(self.rfile.read(int(self.headers['Content-Length'])))
scenario = self.headers.get('x-e2e-case', 'normal')
row = {'body': body, 'headers': dict(self.headers), 'case': scenario, 'closed': False}
records.append(row)
attempt = sum(x['case'] == scenario for x in records)
status = 503 if scenario.startswith('retry') and attempt == 1 else 200
if scenario == 'auth':
status = 401
if status != 200:
data = json.dumps({'error': {'message': 'injected failure', 'type': 'api_error'}}).encode()
self.send_response(status)
self.send_header('Content-Type', 'application/json')
self.send_header('Content-Length', str(len(data)))
self.send_header('Retry-After', '0')
self.end_headers()
self.wfile.write(data)
return
if scenario in ('first-timeout', 'total-timeout'):
time.sleep(1)
try:
if not body.get('stream'):
data = json.dumps(
{
'id': 'e2e',
'object': 'chat.completion',
'created': 1,
'model': body['model'],
'choices': [
{
'index': 0,
'message': {'role': 'assistant', 'content': 'hello world'},
'finish_reason': 'stop',
}
],
'usage': USAGE,
}
).encode()
self.send_response(200)
self.send_header('Content-Type', 'application/json')
self.send_header('Content-Length', str(len(data)))
self.end_headers()
self.wfile.write(data)
return
self.send_response(200)
self.send_header('Content-Type', 'text/event-stream')
self.send_header('Transfer-Encoding', 'chunked')
self.end_headers()
def send(event):
data = ('data: ' + (event if isinstance(event, str) else json.dumps(event)) + '\n\n').encode()
self.wfile.write(f'{len(data):X}\r\n'.encode() + data + b'\r\n')
self.wfile.flush()
base = {'id': 'e2e', 'object': 'chat.completion.chunk', 'created': 1, 'model': body['model']}
if scenario.startswith('reasoning-'):
send(
{
**base,
'choices': [
{'index': 0, 'delta': {'reasoning_content': 'private thought'}, 'finish_reason': None}
],
}
)
for index, content in enumerate(['hello ', 'world']):
send({**base, 'choices': [{'index': 0, 'delta': {'content': content}, 'finish_reason': None}]})
if index == 0 and (scenario in ('idle-timeout', 'disconnect') or scenario.startswith('cancel-')):
if scenario == 'disconnect':
self.close_connection = True
return
deadline = time.monotonic() + 1
while time.monotonic() < deadline:
if select.select([self.connection], [], [], 0.02)[0]:
if not self.connection.recv(1, socket.MSG_PEEK):
row['closed'] = True
return
send({**base, 'choices': [{'index': 0, 'delta': {}, 'finish_reason': 'stop'}]})
send({**base, 'choices': [], 'usage': USAGE})
send('[DONE]')
self.wfile.write(b'0\r\n\r\n')
self.wfile.flush()
except (BrokenPipeError, ConnectionResetError, ConnectionAbortedError):
row['closed'] = True
server = ThreadingHTTPServer(('127.0.0.1', 0), Handler)
thread = threading.Thread(target=server.serve_forever, daemon=True)
thread.start()
yield f'http://127.0.0.1:{server.server_port}/v1', records
server.shutdown()
server.server_close()
thread.join(timeout=5)
@pytest.fixture(scope='module')
def chain(request, provider_wire):
request.getfixturevalue('runner_langbot_process')
port = request.getfixturevalue('runner_e2e_port')
client = httpx.Client(base_url=f'http://127.0.0.1:{port}', timeout=90, trust_env=False)
request.addfinalizer(client.close)
tmpdir = request.getfixturevalue('runner_e2e_tmpdir')
token = runtime._init_and_auth(client)
runtime._ensure_qa_plugin(client, token, tmpdir / 'agent-runner-qa.zip')
client.headers['Authorization'] = f'Bearer {token}'
def post(path, body):
response = client.post('/api/v1/' + path, json=body)
assert response.status_code == 200, response.text
payload = response.json()
assert payload['code'] == 0, payload
return payload['data']
provider = post(
'provider/providers',
{
'name': 'LLM chain E2E',
'requester': 'openai-chat-completions',
'base_url': provider_wire[0],
'api_keys': ['e2e-not-a-secret'],
},
)['uuid']
model = post(
'provider/models/llm',
{
'name': 'gpt-4o',
'provider_uuid': provider,
'abilities': ['func_call', 'vision'],
'extra_args': {
'temperature': 0.1,
'request_policy': {'retries': 1, 'backoff_base': 0.01, 'backoff_max': 0.01},
},
},
)['uuid']
agent = post(
'agents',
{
'kind': 'agent',
'name': 'LLM chain E2E',
'config': {
'runner': {'id': RUNNER},
'runner_config': {RUNNER: {'model': model}},
},
},
)['uuid']
def invoke(mode, case='normal', extra=None, *, configuration=None, options=None):
selected_model, selected_agent = model, agent
if configuration:
configured_provider = post(
'provider/providers',
{
'name': case,
'requester': configuration.get('requester', 'openai-chat-completions'),
'base_url': provider_wire[0],
'api_keys': ['e2e-not-a-secret'],
},
)['uuid']
selected_model = post(
'provider/models/llm',
{
'name': configuration.get('name', 'gpt-4o'),
'provider_uuid': configured_provider,
'abilities': ['func_call', 'vision'],
'extra_args': configuration.get('extra_args', {}),
},
)['uuid']
selected_agent = post(
'agents',
{
'kind': 'agent',
'name': case,
'config': {
'runner': {'id': RUNNER},
'runner_config': {RUNNER: {'model': selected_model}},
},
},
)['uuid']
result = post(
f'agents/{selected_agent}/debug',
{
'text': json.dumps(
{'mode': mode, 'model': selected_model, 'case': case, 'extra': extra or {}, **(options or {})}
),
'conversation_id': 'e2e-' + case,
'event_type': 'message.received',
},
)
assert result.get('final_text'), result
return json.loads(result['final_text'])
return invoke, provider_wire[1], tmpdir
def test_multimodal_and_tool_definitions_survive_plugin_transport(chain):
invoke, records, _ = chain
image = 'data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8/x8AAwMCAO+jRZkAAAAASUVORK5CYII='
messages = [
{
'role': 'user',
'content': [
{'type': 'text', 'text': 'Describe this image'},
{'type': 'image_url', 'image_url': {'url': image}},
],
}
]
result = invoke('runner-events', 'multimodal', options={'messages': messages, 'tool': True})
assert result['text'] == 'hello world', result
wire = records[-1]['body']
assert wire['messages'][0]['content'][1]['image_url']['url'] == image
assert wire['tools'][0]['function']['name'] == 'e2e_lookup'
def test_openrouter_budget_cache_and_explicit_attribution(chain):
invoke, records, _ = chain
result = invoke(
'runner-events',
'openrouter-budget-cache',
configuration={
'requester': 'openrouter-chat-completions',
'name': 'anthropic/claude-sonnet-4',
'extra_args': {
'reasoning_config': {'budget_tokens': 2048},
'prompt_caching': {'mode': 'anthropic'},
'max_tokens': 4096,
},
},
extra={
'extra_headers': {'HTTP-Referer': 'https://langbot.app', 'X-OpenRouter-Title': 'LangBot'},
'request_policy': {'conversation_header': 'x-langbot-conversation-id'},
},
options={'messages': [{'role': 'system', 'content': 'Stable prefix'}, {'role': 'user', 'content': 'Hello'}]},
)
assert result['text'] == 'hello world', result
wire = records[-1]
assert wire['body']['reasoning']['max_tokens'] == 2048
assert wire['body']['messages'][0]['content'][0]['cache_control']['type'] == 'ephemeral'
headers = {k.lower(): v for k, v in wire['headers'].items()}
assert headers['http-referer'] == 'https://langbot.app'
assert headers['x-openrouter-title'] == 'LangBot'
assert headers['x-langbot-conversation-id'].startswith('lb-')
assert 'prompt_caching' not in wire['body'] and 'reasoning_config' not in wire['body']
def test_openai_cache_options_and_call_override(chain):
invoke, records, _ = chain
result = invoke(
'runner-usage',
'openai-cache',
configuration={
'extra_args': {'temperature': 0.1, 'prompt_caching': {'mode': 'openai', 'retention': 'in_memory'}},
},
extra={'temperature': 0.7, 'prompt_cache_key': 'explicit-test-cache'},
)
assert result['text'] == 'hello world', result
assert records[-1]['body']['temperature'] == 0.7
assert records[-1]['body']['prompt_cache_key'] == 'explicit-test-cache'
assert records[-1]['body']['prompt_cache_retention'] == 'in_memory'
@pytest.mark.parametrize(
'mode',
[
'runner-message',
'runner-usage',
'runner-stream',
'runner-events',
'plugin-message',
'plugin-usage',
'plugin-stream',
'plugin-events',
],
)
def test_public_api_shapes_and_usage(chain, mode):
invoke, records, _ = chain
result = invoke(mode, mode)
assert result['text'] == 'hello world', result
assert not result.get('error'), result
if mode.endswith(('usage', 'events')):
assert result['usage']['prompt_tokens_details']['cached_tokens'] == 16
assert result['usage']['completion_tokens_details']['reasoning_tokens'] == 2
wire = records[-1]['body']
assert wire['temperature'] == 0.1
assert 'request_policy' not in wire
@pytest.mark.parametrize(
'case,extra,attempts,error',
[
('retry-once', {}, 2, False),
('retry-disabled', {'num_retries': 0}, 1, True),
('retry-max-disabled', {'max_retries': 0}, 1, True),
('auth', {}, 1, True),
('first-timeout', {'request_policy': {'first_event_timeout': 0.2, 'retries': 0}}, 1, True),
('total-timeout', {'request_policy': {'total_timeout': 0.2, 'retries': 0}}, 1, True),
('idle-timeout', {'request_policy': {'stream_idle_timeout': 0.2}}, 1, True),
('disconnect', {}, 1, True),
],
)
def test_failures_retry_boundaries_and_recovery(chain, case, extra, attempts, error):
invoke, records, _ = chain
result = invoke('runner-events', case, extra)
assert bool(result.get('error')) == error, result
assert sum(row['case'] == case for row in records) == attempts
assert invoke('runner-events', 'recover-' + case)['text'] == 'hello world'
@pytest.mark.parametrize('surface', ['runner-events', 'runner-stream', 'plugin-events', 'plugin-stream'])
def test_cancellation_releases_upstream_and_next_request_works(chain, surface):
invoke, records, tmpdir = chain
case = 'cancel-' + surface
result = invoke(case, case)
assert result['text'] == 'hello ', result
time.sleep(1.5)
cancelled = [row for row in records if row['case'] == case]
with sqlite3.connect(tmpdir / 'data/langbot.db') as db:
diagnostics = db.execute('SELECT session_id, details FROM monitoring_llm_calls').fetchall()
assert len(cancelled) == 1
assert invoke('runner-events', 'after-' + case)['text'] == 'hello world'
if (
not cancelled[0]['closed']
and Version(importlib.metadata.version('langbot-plugin')) <= Version('0.7.11')
and not os.environ.get('LLM_E2E_REQUIRE_STREAM_CLOSE')
):
pytest.xfail('Published SDK <=0.7.11 omits aclosing in public proxy wrappers; requires companion SDK fix')
assert cancelled[0]['closed'], diagnostics[-2:]
@pytest.mark.parametrize('remove_think', [True, False])
def test_reasoning_visibility_through_runner(chain, remove_think):
result = chain[0]('runner-events', f'reasoning-{remove_think}', options={'remove_think': remove_think})
assert result['text'].endswith('hello world'), result
assert ('private thought' in result['text']) == (not remove_think)
def test_invalid_budget_fails_before_network_and_recovers(chain):
invoke, records, _ = chain
result = invoke('runner-events', 'invalid-budget', {'reasoning_config': {'budget_tokens': 1}})
assert result.get('error'), result
assert not any(row['case'] == 'invalid-budget' for row in records)
assert invoke('runner-events', 'after-invalid-budget')['text'] == 'hello world'
def test_monitoring_persists_complete_usage_and_structured_errors(chain):
invoke, _, tmpdir = chain
invoke('runner-events', 'monitoring')
with sqlite3.connect(tmpdir / 'data/langbot.db') as db:
rows = db.execute('SELECT details FROM monitoring_llm_calls').fetchall()
details = [json.loads(row[0]) for row in rows if row[0]]
assert any(x.get('usage', {}).get('prompt_tokens_details', {}).get('cached_tokens') == 16 for x in details)
assert any(x.get('timing', {}).get('transport_setup_ms') is not None for x in details)
assert any(x.get('error', {}).get('code') for x in details if x.get('error'))
@pytest.mark.skipif(not os.environ.get('LLM_E2E_LEGACY_API_DIR'), reason='Historical SDK sources not supplied')
@pytest.mark.parametrize('mode', ['legacy-message', 'legacy-usage', 'legacy-stream', 'legacy-events'])
def test_historical_sdk_proxy_and_message_entities(chain, mode):
result = chain[0](mode, mode)
assert result['text'] == 'hello world', result
assert not result.get('error'), result
if mode.endswith(('usage', 'events')):
assert result['usage']['prompt_tokens_details']['cached_tokens'] == 16
tests/e2e/utils/llm_probe_runner.py
+93
View File
@@ -0,0 +1,93 @@
"""Packaged into the isolated E2E plugin; calls only public SDK LLM APIs."""
from contextlib import aclosing
import json
from langbot_plugin.api.definition.components.runner.runner import Runner
from langbot_plugin.api.entities.builtin.provider.message import Message
from langbot_plugin.api.entities.builtin.runner.result import RunnerResult
from langbot_plugin.api.proxies.langbot_api import LangBotAPIProxy
class DefaultRunner(Runner):
async def run(self, ctx):
options = json.loads(ctx.input.to_text())
mode = options['mode']
cancel = mode.startswith('cancel-')
if cancel:
mode = mode.removeprefix('cancel-')
api = LangBotAPIProxy(self._plugin_runtime_handler) if mode.startswith('plugin-') else ctx.api
message_type = Message
if mode.startswith('legacy-'):
# Historical files are supplied verbatim from a released SDK tag.
# Its proxy and serializers use the real current runtime transport.
import legacy_api
import legacy_message
legacy_api.provider_message = legacy_message
api = legacy_api.LangBotAPIProxy(self._plugin_runtime_handler)
message_type = legacy_message.Message
extra = {
**options['extra'],
'extra_headers': {
**options['extra'].get('extra_headers', {}),
'x-e2e-case': options['case'],
},
}
kwargs = {
'llm_model_uuid': options['model'],
'messages': [
message_type.model_validate(m)
for m in options.get(
'messages',
[
{'role': 'user', 'content': 'Hello'},
],
)
],
'extra_args': extra,
}
if options.get('tool'):
from langbot_plugin.api.entities.builtin.resource.tool import LLMTool
async def lookup(**kwargs):
return 'result'
kwargs['funcs'] = [
LLMTool(
name='e2e_lookup',
human_desc='Lookup',
description='Lookup',
parameters={'type': 'object', 'properties': {}},
func=lookup,
)
]
result = {'text': '', 'usage': None}
if 'remove_think' in options:
kwargs['remove_think'] = options['remove_think']
try:
if mode.endswith('message'):
message = await api.invoke_llm(**kwargs)
assert isinstance(message, message_type)
result['text'] = message.content
elif mode.endswith('usage'):
response = await api.invoke_llm_with_usage(**kwargs)
result.update(
text=response.message.content, usage=response.usage.model_dump() if response.usage else None
)
else:
events = mode.endswith('events')
method = api.invoke_llm_stream_events if events else api.invoke_llm_stream
async with aclosing(method(**kwargs)) as stream:
async for item in stream:
chunk = item.chunk if events else item
if chunk:
result['text'] += chunk.content or ''
if events and item.usage:
result['usage'] = item.usage.model_dump()
if cancel and result['text']:
break
except Exception as exc:
result['error'] = type(exc).__name__ + ': ' + str(exc)
yield RunnerResult.message_completed(ctx.run_id, Message(role='assistant', content=json.dumps(result)))
yield RunnerResult.run_completed(ctx.run_id, finish_reason='stop')
tests/unit_tests/plugin/test_handler_helpers.py
+21
View File
@@ -10,6 +10,27 @@ from __future__ import annotations
import pytest
from unittest.mock import Mock, AsyncMock
from importlib import import_module
from types import SimpleNamespace
@pytest.mark.parametrize(
'call_args',
[
{'num_retries': 0},
{'max_retries': 0},
{'timeout': 3},
{'request_policy': {'stream_idle_timeout': 2}},
],
)
def test_llm_bridge_keeps_model_and_call_policy_layers_separate(call_args):
from langbot.pkg.provider.modelmgr.request_policy import RequestPolicy
model_args = {'temperature': 0.2, 'request_policy': {'retries': 2, 'read_timeout': 90}}
model = SimpleNamespace(model_entity=SimpleNamespace(extra_args=model_args))
merged = get_handler_module()._merge_model_extra_args(model, call_args, llm_policy=True)
assert RequestPolicy.resolve(model_args, merged) == RequestPolicy.resolve(model_args, call_args)
assert merged['temperature'] == 0.2
assert model_args['request_policy'] == {'retries': 2, 'read_timeout': 90}
def get_handler_module():
tests/unit_tests/provider/test_codex.py
+75
View File
@@ -1,6 +1,7 @@
"""Replay synthetic HTTP/SSE traffic through the real Codex requester."""
import json
import asyncio
from collections import OrderedDict
from types import SimpleNamespace
from unittest.mock import AsyncMock
@@ -16,6 +17,32 @@ TOKENS = {'access_token': 'access-secret', 'account_id': 'account', 'connection_
MODEL = SimpleNamespace(model_entity=SimpleNamespace(name='codex-test', extra_args={}, reasoning_config=None))
@pytest.mark.asyncio
async def test_response_headers_obey_host_first_event_deadline(monkeypatch):
from langbot.pkg.provider.modelmgr.errors import RequesterError
from langbot.pkg.provider.modelmgr.request_policy import CallState, RequestPolicy, current_call
cancelled = asyncio.Event()
async def handler(request):
try:
await asyncio.sleep(10)
finally:
cancelled.set()
obj = requester(monkeypatch, handler)
state = CallState(RequestPolicy(first_event_timeout=0.02))
token = current_call.set(state)
try:
with pytest.raises(RequesterError) as caught:
await obj.invoke_llm(None, MODEL, [])
finally:
current_call.reset(token)
assert caught.value.code == 'timeout'
assert caught.value.phase == 'first_event'
assert cancelled.is_set()
def requester(monkeypatch, handler):
real_client = httpx.AsyncClient
monkeypatch.setattr(
@@ -198,3 +225,51 @@ def test_advanced_parameters_cannot_override_transport(monkeypatch, key):
obj = requester(monkeypatch, lambda request: httpx.Response(200))
with pytest.raises(ValueError, match='Unsupported Codex advanced'):
obj._body(None, MODEL, [], None, {key: 'secret'}, TOKENS)
@pytest.mark.asyncio
async def test_early_close_releases_nested_http_stream(monkeypatch):
closed = asyncio.Event()
class Body(httpx.AsyncByteStream):
async def __aiter__(self):
yield b'data: {"type":"response.output_text.delta","delta":"first"}\n\n'
await asyncio.Event().wait()
async def aclose(self):
closed.set()
obj = requester(monkeypatch, lambda request: httpx.Response(200, stream=Body()))
stream = obj.invoke_llm_stream(None, MODEL, [])
assert (await anext(stream)).content == 'first'
await stream.aclose()
assert closed.is_set()
@pytest.mark.asyncio
async def test_host_deadline_does_not_cancel_suspended_consumer(monkeypatch, runtime_provider, runtime_llm_model):
from langbot.pkg.provider.modelmgr.errors import RequesterError
from tests.unit_tests.provider.conftest import TEST_EXECUTION_CONTEXT
runtime_provider.requester = requester(
monkeypatch,
lambda request: stream(
[
{'type': 'response.output_text.delta', 'delta': 'first'},
{'type': 'response.completed', 'response': {'status': 'completed', 'output': []}},
]
),
)
runtime_llm_model.model_entity.extra_args = {}
result = runtime_provider.invoke_llm_stream(
None,
runtime_llm_model,
[],
extra_args={'request_policy': {'total_timeout': 0.05}},
execution_context=TEST_EXECUTION_CONTEXT,
)
assert (await anext(result)).content == 'first'
await asyncio.sleep(0.07)
with pytest.raises(RequesterError) as caught:
await anext(result)
assert caught.value.phase == 'total'
tests/unit_tests/provider/test_http_pool_streams.py
+285
View File
@@ -0,0 +1,285 @@
import asyncio
import json
import threading
from types import SimpleNamespace
from unittest.mock import Mock
import httpx
import pytest
from aiohttp import web
from langbot.pkg.provider.modelmgr.http_pool import _CompletedSSEStream, RequesterHTTPPool
from langbot.pkg.provider.modelmgr.errors import RequesterError
from langbot.pkg.provider.modelmgr.requesters.litellmchat import LiteLLMRequester
from langbot.pkg.provider.modelmgr.request_policy import CallState, RequestPolicy, current_call
from tests.unit_tests.provider.test_reasoning_control import _runtime_model
from scripts.llm_latency_probe import SSEMeter, cache_tokens
def test_request_timing_uses_high_resolution_monotonic_clock(monkeypatch):
from langbot.pkg.provider.modelmgr import request_policy
ticks = iter([100.0, 100.0004])
monkeypatch.setattr(request_policy.time, 'perf_counter', lambda: next(ticks))
state = CallState(RequestPolicy())
state.mark('request_prepared_ms')
assert state.metrics['request_prepared_ms'] == 0.4
def test_meter_accepts_split_utf8_and_usage_without_content():
meter = SSEMeter()
event = {'choices': [{'delta': {'content': '你好'}}]}
data = ('data: ' + json.dumps(event, ensure_ascii=False) + '\r\n\r\n').encode()
for byte in data:
meter.feed(bytes([byte]), 1)
meter.feed(b'data: {"choices":[],"usage":{"prompt_tokens_details":{"cached_tokens":12}}}\n\n', 2)
meter.feed(b'data: [DONE]\n\n', 3)
assert meter.first_content == 1 and meter.content_chars == 2
assert cache_tokens(meter.usage) == 12 and meter.done
assert cache_tokens(None) is None
@pytest.mark.asyncio
@pytest.mark.parametrize('provider,name', [('openai', 'gpt-4o'), ('openrouter', 'openai/gpt-4o')])
@pytest.mark.parametrize('drain', [False, True])
async def test_real_stream_reuses_socket_and_preserves_late_usage(monkeypatch, provider, name, drain):
monkeypatch.setenv('NO_PROXY', '127.0.0.1,localhost')
sockets = []
async def handle(request):
await request.read()
sockets.append(request.transport)
response = web.StreamResponse(headers={'Content-Type': 'text/event-stream'})
await response.prepare(request)
events = [
{'choices': [{'index': 0, 'delta': {'role': 'assistant', 'content': 'hello'}, 'finish_reason': None}]},
{'choices': [{'index': 0, 'delta': {}, 'finish_reason': 'stop'}]},
{
'choices': [],
'usage': {
'prompt_tokens': 14,
'completion_tokens': 1,
'total_tokens': 15,
'prompt_tokens_details': {'cached_tokens': 12},
},
},
]
for event in events:
await response.write(
(
'data: '
+ json.dumps(
{
'id': 'r',
'object': 'chat.completion.chunk',
'created': 1,
'model': name,
**event,
}
)
+ '\n\n'
).encode()
)
await asyncio.sleep(0.02)
await response.write(b'data: [DO')
await response.write(b'NE]\n\n')
await response.write_eof()
return response
app = web.Application()
app.router.add_post('/chat/completions', handle)
runner = web.AppRunner(app)
await runner.setup()
site = web.TCPSite(runner, '127.0.0.1', 0)
await site.start()
adapter = LiteLLMRequester(
SimpleNamespace(logger=Mock()),
{
'custom_llm_provider': provider,
'base_url': f'http://127.0.0.1:{site._server.sockets[0].getsockname()[1]}',
},
)
await adapter.initialize()
adapter._http_pool.drain_sse = drain
model = _runtime_model(adapter, name=name)
model.provider.token_mgr.get_token = lambda: 'test'
try:
for _ in range(2):
state = CallState(RequestPolicy())
token = current_call.set(state)
try:
chunks = [chunk async for chunk in adapter.invoke_llm_stream(None, model, [])]
assert ''.join(chunk.content or '' for chunk in chunks) == 'hello'
assert state.usage['prompt_tokens_details']['cached_tokens'] == 12
finally:
current_call.reset(token)
# OpenRouter's generic HTTP adapter already consumes EOF; the OpenAI
# SDK path is the regression that needs terminal draining.
assert (sockets[0] is sockets[1]) == (drain or provider == 'openrouter')
finally:
await adapter.aclose()
await runner.cleanup()
@pytest.mark.asyncio
async def test_terminal_drain_is_bounded_but_early_close_does_not_drain():
class Stream(httpx.AsyncByteStream):
def __init__(self, terminal):
self.terminal = terminal
self.tail_read = False
self.closed = False
async def __aiter__(self):
yield b'data: [DONE]\n\n' if self.terminal else b'data: {"choices":[]}\n\n'
self.tail_read = True
await asyncio.Event().wait()
async def aclose(self):
self.closed = True
for terminal in (False, True):
raw = Stream(terminal)
wrapped = _CompletedSSEStream(raw)
iterator = wrapped.__aiter__()
async with asyncio.timeout(1):
await anext(iterator)
await wrapped.aclose()
await iterator.aclose()
assert raw.closed and raw.tail_read == terminal
@pytest.mark.asyncio
async def test_cancellation_during_terminal_drain_is_not_swallowed():
reading = asyncio.Event()
class Stream(httpx.AsyncByteStream):
async def __aiter__(self):
yield b'data: [DONE]\n\n'
reading.set()
await asyncio.Event().wait()
async def aclose(self):
pass
wrapped = _CompletedSSEStream(Stream())
iterator = wrapped.__aiter__()
pending = asyncio.create_task(anext(iterator))
await reading.wait()
pending.cancel()
with pytest.raises(asyncio.CancelledError):
await pending
await iterator.aclose()
await wrapped.aclose()
@pytest.mark.asyncio
@pytest.mark.parametrize('retire', [False, True])
async def test_concurrent_pool_initialization_loads_tls_once_off_loop(monkeypatch, retire):
context = httpx.create_ssl_context()
entered, release = threading.Event(), threading.Event()
main_thread = threading.get_ident()
calls = []
def load():
assert threading.get_ident() != main_thread
calls.append(True)
entered.set()
assert release.wait(2)
return context
monkeypatch.setattr(httpx, 'create_ssl_context', load)
pool = RequesterHTTPPool()
async def lease():
async with pool.lease() as client:
return client
first, second = asyncio.create_task(lease()), asyncio.create_task(lease())
await asyncio.to_thread(entered.wait, 2)
if retire:
await pool.aclose()
release.set()
results = await asyncio.gather(first, second, return_exceptions=True)
assert len(calls) == 1
if retire:
assert all(isinstance(result, RequesterError) for result in results)
assert pool.client is None
else:
assert results[0] is results[1] is pool.client
assert not pool.client.is_closed
await pool.aclose()
@pytest.mark.asyncio
async def test_cancelling_tls_initialization_leaks_no_http_client(monkeypatch):
context = httpx.create_ssl_context()
entered, release = threading.Event(), threading.Event()
def load():
entered.set()
assert release.wait(2)
return context
monkeypatch.setattr(httpx, 'create_ssl_context', load)
pool = RequesterHTTPPool()
async def lease():
async with pool.lease():
pytest.fail('Cancelled initializer must not publish a client')
pending = asyncio.create_task(lease())
await asyncio.to_thread(entered.wait, 2)
pending.cancel()
try:
with pytest.raises(asyncio.CancelledError):
await pending
await pool.aclose()
assert pool.client is None and pool.active == 0
finally:
release.set()
@pytest.mark.asyncio
async def test_real_openai_stream_close_disconnects_upstream(monkeypatch):
monkeypatch.setenv('NO_PROXY', '127.0.0.1,localhost')
disconnected = asyncio.Event()
async def handle(request):
await request.read()
response = web.StreamResponse(headers={'Content-Type': 'text/event-stream'})
await response.prepare(request)
await response.write(
b'data: {"id":"r","object":"chat.completion.chunk","created":1,"model":"gpt-4o","choices":[{"index":0,"delta":{"content":"hello"},"finish_reason":null}]}\n\n'
)
for _ in range(100):
if request.transport is None or request.transport.is_closing():
disconnected.set()
break
await asyncio.sleep(0.02)
return response
app = web.Application()
app.router.add_post('/chat/completions', handle)
runner = web.AppRunner(app)
await runner.setup()
site = web.TCPSite(runner, '127.0.0.1', 0)
await site.start()
adapter = LiteLLMRequester(
SimpleNamespace(logger=Mock()),
{
'custom_llm_provider': 'openai',
'base_url': f'http://127.0.0.1:{site._server.sockets[0].getsockname()[1]}',
},
)
await adapter.initialize()
model = _runtime_model(adapter, name='gpt-4o')
model.provider.token_mgr.get_token = lambda: 'test'
stream = adapter.invoke_llm_stream(None, model, [])
try:
assert (await anext(stream)).content == 'hello'
await stream.aclose()
await asyncio.wait_for(disconnected.wait(), 3)
finally:
await stream.aclose()
await adapter.aclose()
await runner.cleanup()
tests/unit_tests/provider/test_llm_preparation.py
+516
View File
@@ -0,0 +1,516 @@
"""Connection ownership, preparation reuse and provider-specific request policy."""
import asyncio
import copy
import threading
import json
from pathlib import Path
from collections import OrderedDict
from types import SimpleNamespace
from unittest.mock import AsyncMock, Mock, patch
import httpx
import pytest
import yaml
from aiohttp import web
from langbot.pkg.provider.modelmgr import reasoning
from langbot.pkg.provider.modelmgr.errors import RequesterError
from langbot.pkg.provider.modelmgr.http_pool import RequesterHTTPPool
from langbot.pkg.provider.modelmgr.prompt_caching import apply_prompt_caching
from langbot.pkg.provider.modelmgr.request_policy import CallState, RequestPolicy, current_call
from langbot.pkg.provider.modelmgr.requesters.litellmchat import LiteLLMRequester
from langbot.pkg.provider.modelmgr.requesters.litellmchat import _LeasedHTTPHandler
from langbot.pkg.provider.modelmgr.requesters import codex
from langbot_plugin.api.entities.builtin.provider.message import Message
from tests.unit_tests.provider.test_reasoning_control import _runtime_model, _requester
@pytest.mark.asyncio
async def test_pool_retirement_waits_for_stream_lease():
pool = RequesterHTTPPool()
async with pool.lease() as first:
async with pool.lease() as second:
assert first is second
await pool.aclose()
assert not first.is_closed
with pytest.raises(RequesterError, match='retired'):
async with pool.lease():
pass
assert not first.is_closed
assert first.is_closed
await pool.aclose()
@pytest.mark.asyncio
@pytest.mark.parametrize('provider,name', [('openai', 'gpt-4o'), ('openrouter', 'openai/gpt-4o')])
async def test_real_http_reuses_socket_with_rotating_key_and_timeout(monkeypatch, provider, name):
monkeypatch.setenv('NO_PROXY', '127.0.0.1,localhost')
sockets, headers, bodies = [], [], []
async def handle(request):
sockets.append(request.transport)
headers.append(request.headers['Authorization'])
bodies.append(await request.json())
return web.json_response(
{
'id': 'r',
'object': 'chat.completion',
'created': 1,
'model': 'gpt-4o',
'choices': [{'index': 0, 'message': {'role': 'assistant', 'content': 'ok'}, 'finish_reason': 'stop'}],
'usage': {'prompt_tokens': 1, 'completion_tokens': 1, 'total_tokens': 2},
}
)
app = web.Application()
app.router.add_post('/v1/chat/completions', handle)
runner = web.AppRunner(app)
await runner.setup()
site = web.TCPSite(runner, '127.0.0.1', 0)
await site.start()
port = site._server.sockets[0].getsockname()[1]
adapter = LiteLLMRequester(Mock(), {'custom_llm_provider': provider, 'base_url': f'http://127.0.0.1:{port}/v1'})
await adapter.initialize()
model = _runtime_model(adapter, name=name)
model.provider.token_mgr.get_token = Mock(side_effect=['first-key', 'second-key'])
try:
for seconds in (10, 20):
state = CallState(RequestPolicy(read_timeout=seconds), 'lb-scoped')
token = current_call.set(state)
try:
message, _ = await adapter.invoke_llm(
None,
model,
[Message(role='user', content='Hi')],
extra_args={'prompt_caching': {'mode': 'openai', 'retention': 'in_memory'}},
)
assert message.content == 'ok'
finally:
current_call.reset(token)
assert sockets[0] is sockets[1]
assert headers == ['Bearer first-key', 'Bearer second-key']
assert bodies[0]['prompt_cache_key'] == bodies[1]['prompt_cache_key']
assert bodies[0]['prompt_cache_retention'] == 'in_memory'
assert all('prompt_caching' not in body and 'reasoning_config' not in body for body in bodies)
finally:
await adapter.aclose()
await runner.cleanup()
@pytest.mark.asyncio
@pytest.mark.parametrize('cancel', [False, True])
async def test_stream_holds_pool_lease_until_consumer_close(cancel):
adapter = _requester('openai')
await adapter.initialize()
model = _runtime_model(adapter, name='gpt-4o')
model.provider.token_mgr.get_token = Mock(return_value='test')
closed = asyncio.Event()
async def raw():
try:
yield SimpleNamespace(
choices=[
SimpleNamespace(delta=SimpleNamespace(model_dump=lambda: {'content': 'hi'}), finish_reason=None)
]
)
await asyncio.Event().wait()
finally:
closed.set()
with patch('langbot.pkg.provider.modelmgr.requesters.litellmchat.acompletion', new=AsyncMock(return_value=raw())):
stream = adapter.invoke_llm_stream(None, model, [])
await anext(stream)
client = adapter._http_pool.client
await adapter.aclose()
assert not client.is_closed
if cancel:
pending = asyncio.create_task(anext(stream))
await asyncio.sleep(0)
pending.cancel()
with pytest.raises(asyncio.CancelledError):
await pending
else:
await stream.aclose()
assert closed.is_set() and client.is_closed
@pytest.mark.asyncio
async def test_preparation_reused_only_within_call_and_is_not_mutated():
adapter = _requester('openai')
adapter.ap.tool_mgr = SimpleNamespace(
generate_tools_for_openai=AsyncMock(
return_value=[{'type': 'function', 'function': {'name': 'tool', 'parameters': {'type': 'object'}}}]
)
)
model = _runtime_model(adapter, name='gpt-4o')
model.provider.token_mgr.get_token = Mock(side_effect=['key-1', 'key-2', 'key-3'])
messages, funcs = [Message(role='user', content='original')], [object()]
state = CallState(RequestPolicy(retries=1))
with patch.object(adapter, '_convert_messages', wraps=adapter._convert_messages) as convert:
token = current_call.set(state)
try:
first = await adapter._build_completion_args(model, messages, funcs)
first['messages'][0]['content'] = 'mutated'
first['tools'][0]['function']['name'] = 'mutated'
second = await adapter._build_completion_args(model, messages, funcs)
finally:
current_call.reset(token)
assert convert.call_count == 1
assert adapter.ap.tool_mgr.generate_tools_for_openai.await_count == 1
assert second['messages'][0]['content'] == 'original'
assert second['tools'][0]['function']['name'] == 'tool'
assert second['api_key'] == 'key-2'
await adapter._build_completion_args(model, messages, funcs)
assert convert.call_count == 2
assert messages[0].content == 'original'
@pytest.mark.asyncio
async def test_token_counter_runs_off_event_loop():
adapter = _requester('openai')
model = _runtime_model(adapter, name='gpt-4o')
model.provider.token_mgr.get_token = Mock(return_value='test')
main_thread = threading.get_ident()
def count(**kwargs):
assert threading.get_ident() != main_thread
return 12
with patch('litellm.token_counter', new=count):
assert await adapter.count_tokens(model, []) == 12
def test_metadata_cache_is_bounded_and_returns_copies():
adapter = _requester('openai')
compute = Mock(return_value={'levels': ['low']})
first = adapter._cached_metadata('model', compute)
first['levels'].append('bad')
assert adapter._cached_metadata('model', compute) == {'levels': ['low']}
assert compute.call_count == 1
for index in range(150):
adapter._cached_metadata(index, lambda: index)
assert len(adapter._metadata_cache) == 128
@pytest.mark.parametrize(
'config',
[
{'budget_tokens': True},
{'budget_tokens': 0},
{'budget_tokens': -1},
{'budget_tokens': 1.5},
{'level': 'high', 'budget_tokens': 1024},
{'level': 'disabled', 'budget_tokens': 1024},
],
)
def test_invalid_reasoning_budget(config):
with pytest.raises(ValueError):
reasoning.normalize_reasoning_config(config)
@pytest.mark.asyncio
@pytest.mark.parametrize(
'provider,name,budget,expected',
[
('anthropic', 'claude-sonnet-4-5', 2048, {'thinking': {'type': 'enabled', 'budget_tokens': 2048}}),
('openai', 'qwen3-32b', 2048, {'extra_body': {'enable_thinking': True, 'thinking_budget': 2048}}),
('gemini', 'gemini-2.5-flash', 2048, {'thinking': {'type': 'enabled', 'budget_tokens': 2048}}),
('openrouter', 'anthropic/claude-sonnet-4.5', 2048, {'extra_body': {'reasoning': {'max_tokens': 2048}}}),
],
)
async def test_budget_mapping_through_existing_extra_args(provider, name, budget, expected):
adapter = _requester(provider)
model = _runtime_model(adapter, name=name)
model.provider.token_mgr.get_token = Mock(return_value='test')
args = await adapter._build_completion_args(
model, [], extra_args={'reasoning_config': {'budget_tokens': budget}, 'max_tokens': 4096}
)
for key, value in expected.items():
assert args[key] == value
assert 'reasoning_config' not in args
assert model.reasoning_config_override is None
@pytest.mark.asyncio
@pytest.mark.parametrize(
'provider,name,budget',
[
('openai', 'gpt-5', 2048),
('anthropic', 'claude-opus-4-7', 2048),
('anthropic', 'claude-sonnet-4-5', 512),
('gemini', 'gemini-2.5-flash', 30000),
('gemini', 'gemini-2.5-flash-lite', 128),
],
)
async def test_budget_not_silently_dropped_on_unsupported_model(provider, name, budget):
adapter = _requester(provider)
model = _runtime_model(adapter, name=name)
model.provider.token_mgr.get_token = Mock(return_value='test')
with pytest.raises(RequesterError, match='budget'):
await adapter._build_completion_args(model, [], extra_args={'reasoning_config': {'budget_tokens': budget}})
def test_runner_level_replaces_call_and_model_budget():
adapter = _requester('anthropic')
model = _runtime_model(adapter, name='claude-sonnet-4-5')
model.model_entity.reasoning_config = {'budget_tokens': 2048}
call = {'reasoning_config': {'budget_tokens': 4096}}
assert reasoning.resolve_reasoning_config(model, call)['budget_tokens'] == 4096
model.reasoning_config_override = {'level': 'low'}
assert reasoning.resolve_reasoning_config(model, call) == {'level': 'low'}
def test_anthropic_cache_boundary_stable_and_input_unchanged():
messages = [{'role': 'system', 'content': 'stable instructions'}, {'role': 'user', 'content': 'first'}]
original = copy.deepcopy(messages)
args = {'messages': messages, 'prompt_caching': {'mode': 'anthropic', 'ttl': '1h'}}
apply_prompt_caching(args, None, provider='anthropic', family='anthropic')
prefix = args['messages'][0]
assert prefix['content'][-1]['cache_control'] == {'type': 'ephemeral', 'ttl': '1h'}
assert messages == original
assert 'cache_control' not in args['messages'][1]
args['messages'].append({'role': 'user', 'content': 'second'})
args['prompt_caching'] = {'mode': 'anthropic'}
apply_prompt_caching(args, None, provider='anthropic', family='anthropic')
assert args['messages'][0] == prefix
def test_explicit_cache_settings_win_and_default_does_nothing():
args = {
'model': 'openai/gpt-4o',
'extra_body': {'prompt_cache_key': 'explicit', 'prompt_cache_retention': '24h'},
'prompt_caching': {'mode': 'openai', 'retention': 'in_memory'},
}
apply_prompt_caching(args, CallState(RequestPolicy(), 'lb-scoped'), provider='openai', family='openai')
assert 'prompt_cache_key' not in args and 'prompt_cache_retention' not in args
args = {'messages': [{'role': 'system', 'content': 'same'}], 'prompt_caching': {'mode': 'provider_default'}}
apply_prompt_caching(args, None, provider='anthropic', family='anthropic')
assert args == {'messages': [{'role': 'system', 'content': 'same'}]}
def test_native_cache_options_bypass_model_filter_without_legacy_hints():
args = {
'model': 'openai/gpt-test',
'prompt_cache_options': {'mode': 'implicit', 'ttl': '30m'},
'prompt_caching': {'mode': 'openai', 'retention': '24h'},
}
apply_prompt_caching(args, CallState(RequestPolicy(), 'lb-scoped'), provider='openai', family='openai')
assert args == {
'model': 'openai/gpt-test',
'extra_body': {'prompt_cache_options': {'mode': 'implicit', 'ttl': '30m'}},
}
@pytest.mark.parametrize(
'config',
[
{'mode': 'typo'},
{'mode': 'anthropic', 'ttl': '24h'},
{'mode': 'openai', 'retention': 'forever'},
{'mode': 'provider_default', 'ttl': '1h'},
],
)
def test_invalid_cache_config(config):
with pytest.raises(RequesterError):
apply_prompt_caching(
{'model': 'test', 'messages': [], 'prompt_caching': config}, None, provider='openai', family='openai'
)
@pytest.mark.asyncio
async def test_gateway_transport_never_retries_or_closes_borrowed_client():
calls = 0
async def fail(request):
nonlocal calls
calls += 1
raise httpx.ConnectError('unavailable', request=request)
async with httpx.AsyncClient(transport=httpx.MockTransport(fail)) as client:
handler = _LeasedHTTPHandler(client, 10)
with pytest.raises(httpx.ConnectError):
await handler.post('https://example.com', json={})
await handler.close()
assert calls == 1 and not client.is_closed
@pytest.mark.asyncio
@pytest.mark.parametrize('provider', ['anthropic', 'openrouter', 'openrouter-metadata'])
async def test_anthropic_cache_and_budget_reach_upstream(provider):
bodies = []
async def handle(request):
bodies.append(json.loads(request.content))
if provider == 'anthropic':
payload = {
'id': 'msg_test',
'type': 'message',
'role': 'assistant',
'model': 'claude-sonnet-4-5',
'content': [{'type': 'text', 'text': 'ok'}],
'stop_reason': 'end_turn',
'usage': {'input_tokens': 1, 'output_tokens': 1, 'cache_read_input_tokens': 50},
}
else:
payload = {
'id': 'r',
'object': 'chat.completion',
'created': 1,
'model': 'anthropic/claude-sonnet-4.5',
'choices': [{'index': 0, 'message': {'role': 'assistant', 'content': 'ok'}, 'finish_reason': 'stop'}],
'usage': {
'prompt_tokens': 51,
'completion_tokens': 1,
'total_tokens': 52,
'prompt_tokens_details': {'cached_tokens': 50},
},
}
return httpx.Response(200, json=payload)
adapter = _requester('openai' if provider == 'openrouter-metadata' else provider)
if provider == 'openrouter-metadata':
metadata = yaml.safe_load(
(
Path(__file__).parents[3] / 'src/langbot/pkg/provider/modelmgr/requesters/openrouterchatcmpl.yaml'
).read_text(encoding='utf-8')
)
adapter.requester_cfg.update(
requester_name=metadata['metadata']['name'], custom_llm_provider=metadata['spec']['litellm_provider']
)
name = 'claude-sonnet-4-5' if provider == 'anthropic' else 'anthropic/claude-sonnet-4.5'
model = _runtime_model(adapter, name=name)
model.provider.token_mgr.get_token = Mock(return_value='test')
async with httpx.AsyncClient(transport=httpx.MockTransport(handle)) as client:
result, usage = await adapter.invoke_llm(
None,
model,
[
Message(role='system', content='stable'),
Message(role='user', content='hello'),
],
extra_args={
'client': _LeasedHTTPHandler(client, 10),
'reasoning_config': {'budget_tokens': 2048},
'max_tokens': 4096,
'prompt_caching': {'mode': 'anthropic', 'ttl': '1h'},
},
)
assert result.content == 'ok'
assert usage
body = bodies[0]
if provider == 'anthropic':
assert body['system'][-1]['cache_control'] == {'type': 'ephemeral', 'ttl': '1h'}
assert body['thinking']['budget_tokens'] == 2048
else:
assert body['messages'][0]['content'][-1]['cache_control'] == {'type': 'ephemeral', 'ttl': '1h'}
assert body['reasoning']['max_tokens'] == 2048
assert 'reasoning_config' not in body and 'prompt_caching' not in body
@pytest.mark.asyncio
async def test_codex_reuses_real_connection(monkeypatch):
monkeypatch.setenv('NO_PROXY', '127.0.0.1,localhost')
sockets = []
async def handle(request):
sockets.append(request.transport)
await request.read()
events = [
{'type': 'response.output_text.delta', 'delta': 'ok'},
{'type': 'response.completed', 'response': {'id': 'r', 'status': 'completed', 'output': []}},
]
return web.Response(
text=''.join('data: ' + json.dumps(event) + '\n\n' for event in events), content_type='text/event-stream'
)
app = web.Application()
app.router.add_post('/responses', handle)
runner = web.AppRunner(app)
await runner.setup()
site = web.TCPSite(runner, '127.0.0.1', 0)
await site.start()
monkeypatch.setattr(codex, 'BASE_URL', f'http://127.0.0.1:{site._server.sockets[0].getsockname()[1]}')
adapter = object.__new__(codex.CodexRequester)
adapter.workspace, adapter.provider = 'w', 'p'
adapter._replay = OrderedDict()
adapter._http_pool = RequesterHTTPPool()
adapter.auth = SimpleNamespace(
access=AsyncMock(
return_value={
'access_token': 'test',
'account_id': 'a',
'connection_id': 'c',
}
)
)
model = SimpleNamespace(model_entity=SimpleNamespace(name='test', extra_args={}, reasoning_config=None))
try:
for _ in range(2):
message, _ = await adapter.invoke_llm(None, model, [])
assert message.content == 'ok'
assert sockets[0] is sockets[1]
finally:
await adapter.aclose()
await runner.cleanup()
@pytest.mark.asyncio
async def test_explicit_client_and_redirected_provider_are_preserved():
adapter = _requester('openai')
await adapter.initialize()
explicit = object()
for args in ({'model': 'openai/gpt-4o', 'client': explicit}, {'model': 'anthropic/claude-sonnet-4-5'}):
async with adapter._client_args(args) as result:
assert result is args
assert adapter._http_pool.client is None
await adapter.aclose()
@pytest.mark.asyncio
async def test_codex_terminal_tail_is_bounded_and_releases_stream():
closed = asyncio.Event()
class HangingTail(httpx.AsyncByteStream):
async def __aiter__(self):
yield b'data: {"type":"response.completed"}\n\n'
await asyncio.Event().wait()
async def aclose(self):
closed.set()
async with httpx.AsyncClient(
transport=httpx.MockTransport(lambda request: httpx.Response(200, stream=HangingTail()))
) as client:
async with client.stream('GET', 'https://example.com') as response:
events = [event async for event in codex.sse_events(response)]
assert events == [{'type': 'response.completed'}]
assert closed.is_set()
@pytest.mark.asyncio
async def test_custom_tls_keeps_litellm_transport(monkeypatch):
import litellm
adapter = _requester('openai')
await adapter.initialize()
monkeypatch.setattr(litellm, 'ssl_verify', '/custom/ca.pem')
args = {'model': 'openai/gpt-4o'}
async with adapter._client_args(args) as result:
assert result is args
assert adapter._http_pool.client is None
await adapter.aclose()
@pytest.mark.asyncio
async def test_pooled_sdk_preserves_environment_fallbacks(monkeypatch):
monkeypatch.setenv('OPENAI_BASE_URL', 'https://gateway.example/v1')
monkeypatch.setenv('OPENAI_API_KEY', 'env-test-key')
adapter = _requester('openai')
await adapter.initialize()
try:
async with adapter._client_args({'model': 'openai/gpt-4o', 'api_key': ''}) as args:
assert str(args['client'].base_url) == 'https://gateway.example/v1/'
assert args['client'].api_key == 'env-test-key'
finally:
await adapter.aclose()
tests/unit_tests/provider/test_request_policy.py
+549
View File
@@ -0,0 +1,549 @@
import asyncio
from types import SimpleNamespace
from unittest.mock import AsyncMock, Mock, patch
import httpx
import litellm
import pytest
from langbot.pkg.provider.modelmgr.errors import RequesterError
from langbot.pkg.provider.modelmgr.request_policy import (
CallState,
RequestPolicy,
current_call,
error_metadata,
opaque_conversation_id,
)
from langbot.pkg.provider.modelmgr.requesters.litellmchat import LiteLLMRequester
from langbot_plugin.api.entities.builtin.provider.message import Message, MessageChunk
from tests.unit_tests.provider.conftest import TEST_EXECUTION_CONTEXT
from tests.unit_tests.provider.test_litellmchat import MockRuntimeModel
def test_policy_precedence_zero_and_no_mutation():
model = {'num_retries': 3, 'request_policy': {'read_timeout': 42}}
policy = RequestPolicy.resolve({'request_policy': {'retries': 2}}, model, {'num_retries': 0})
assert policy.retries == 0
assert policy.read_timeout == 42
assert model['num_retries'] == 3
@pytest.mark.parametrize(
'policy',
[
{'retries': True},
{'retries': -1},
{'retries': 6},
{'read_timeout': 0},
{'total_timeout': float('nan')},
{'connect_timeout': float('inf')},
{'attribution': 'false'},
{'conversation_header': 'authorization'},
{'typo': 1},
],
)
def test_invalid_policy(policy):
with pytest.raises(RequesterError, match='Model request failed|模型请求失败'):
RequestPolicy.resolve({'request_policy': policy})
def test_retry_metadata_does_not_include_upstream_secrets():
response = httpx.Response(429, headers={'retry-after': '2', 'x-request-id': 'req-123'})
meta = error_metadata(SimpleNamespace(status_code=429, response=response))
assert meta == {'status_code': 429, 'retry_after': 2, 'request_id': 'req-123'}
adapter = LiteLLMRequester(Mock(), {})
exc = litellm.RateLimitError('secret-key-and-prompt', llm_provider='openai', model='test', response=response)
with pytest.raises(RequesterError) as caught:
adapter._handle_litellm_error(exc)
assert caught.value.retryable
assert caught.value.request_id == 'req-123'
assert 'secret-key' not in str(caught.value)
def test_quota_is_not_retryable():
adapter = LiteLLMRequester(Mock(), {})
exc = litellm.RateLimitError('quota', llm_provider='openai', model='test')
exc.body = {'error': {'code': 'insufficient_quota'}}
with pytest.raises(RequesterError) as caught:
adapter._handle_litellm_error(exc)
assert caught.value.code == 'quota_exceeded'
assert not caught.value.retryable
@pytest.mark.asyncio
async def test_headers_and_sdk_retry_ownership():
adapter = LiteLLMRequester(Mock(), {'base_url': 'https://openrouter.ai/api/v1'})
state = CallState(RequestPolicy(retries=2), 'lb-opaque')
token = current_call.set(state)
try:
args = await adapter._build_completion_args(
MockRuntimeModel(),
[],
extra_args={
'request_policy': {'retries': 2},
'num_retries': 2,
'extra_headers': {'x-openrouter-title': 'Custom app', 'x-other': 'keep'},
},
)
finally:
current_call.reset(token)
assert 'request_policy' not in args
assert args['num_retries'] == args['max_retries'] == 0
assert isinstance(args['timeout'], httpx.Timeout)
assert args['extra_headers'] == {
'x-openrouter-title': 'Custom app',
'x-other': 'keep',
'HTTP-Referer': 'https://langbot.app',
'x-session-id': 'lb-opaque',
}
@pytest.mark.asyncio
@pytest.mark.parametrize('url', ['https://gateway.example/v1', 'https://openrouter.ai.evil.example/v1'])
async def test_no_vendor_headers_on_unrelated_hosts(url):
adapter = LiteLLMRequester(Mock(), {'base_url': url})
args = await adapter._build_completion_args(MockRuntimeModel(), [])
assert not args.get('extra_headers')
assert args['num_retries'] == 0
def test_conversation_is_stable_and_scoped():
query = SimpleNamespace(session=None, launcher_type='person', launcher_id='123', bot_uuid='bot', pipeline_uuid='p')
first = opaque_conversation_id(query, TEST_EXECUTION_CONTEXT)
assert first == opaque_conversation_id(query, TEST_EXECUTION_CONTEXT)
query.bot_uuid = 'another-bot'
assert first != opaque_conversation_id(query, TEST_EXECUTION_CONTEXT)
assert opaque_conversation_id(None, TEST_EXECUTION_CONTEXT) is None
@pytest.mark.asyncio
async def test_runtime_retries_only_before_output(runtime_provider, runtime_llm_model):
provider = runtime_provider
attempts = 0
closed = []
async def stream(**kwargs):
nonlocal attempts
attempts += 1
try:
if attempts == 1:
raise RequesterError('transient', retryable=True)
yield MessageChunk(role='assistant', content='first')
raise RequesterError('after output', retryable=True)
finally:
closed.append(attempts)
provider.requester.invoke_llm_stream = stream
result = []
with pytest.raises(RequesterError, match='after output'):
async for chunk in provider.invoke_llm_stream(
None,
runtime_llm_model,
[],
extra_args={'request_policy': {'retries': 3, 'backoff_base': 0.001}},
execution_context=TEST_EXECUTION_CONTEXT,
):
result.append(chunk.content)
assert result == ['first']
assert attempts == 2
assert closed == [1, 2]
assert current_call.get() is None
@pytest.mark.asyncio
async def test_runtime_deadline_and_close(runtime_provider, runtime_llm_model):
closed = asyncio.Event()
async def stream(**kwargs):
try:
await asyncio.Event().wait()
yield MessageChunk(role='assistant', content='never')
finally:
closed.set()
runtime_provider.requester.invoke_llm_stream = stream
with pytest.raises(RequesterError) as caught:
async for _ in runtime_provider.invoke_llm_stream(
None,
runtime_llm_model,
[],
extra_args={'request_policy': {'first_event_timeout': 0.02}},
execution_context=TEST_EXECUTION_CONTEXT,
):
pass
assert caught.value.phase == 'first_event'
assert closed.is_set()
@pytest.mark.asyncio
async def test_retry_after_cannot_extend_total_budget():
state = CallState(RequestPolicy(retries=5, total_timeout=0.1))
state.metrics['attempts'] = 1
assert not await state.retry(RequesterError('busy', retryable=True, retry_after=30))
@pytest.mark.asyncio
async def test_cancellation_is_not_retried(runtime_provider, runtime_llm_model):
invoke = AsyncMock(side_effect=asyncio.CancelledError)
runtime_provider.requester.invoke_llm = invoke
with pytest.raises(asyncio.CancelledError):
await runtime_provider.invoke_llm(
None,
runtime_llm_model,
[],
extra_args={'request_policy': {'retries': 3}},
execution_context=TEST_EXECUTION_CONTEXT,
)
assert invoke.await_count == 1
@pytest.mark.asyncio
async def test_complete_usage_and_timings_persist(runtime_provider, runtime_llm_model):
from langbot_plugin.api.entities.builtin.pipeline.query import Query
query = Query.model_construct(
query_id=1, launcher_type='person', launcher_id=1, variables={}, bot_uuid='bot', pipeline_uuid='p', session=None
)
object.__setattr__(query, '_execution_context', TEST_EXECUTION_CONTEXT)
usage = {
'prompt_tokens': 20,
'completion_tokens': 10,
'total_tokens': 30,
'prompt_tokens_details': {'cached_tokens': 15},
'completion_tokens_details': {'reasoning_tokens': 8},
'cost': 0.01,
}
runtime_provider.requester.invoke_llm = AsyncMock(return_value=(Message(role='assistant', content='ok'), usage))
with patch(
'langbot.pkg.pipeline.monitoring_helper.MonitoringHelper.record_llm_call', new_callable=AsyncMock
) as record:
result = await runtime_provider.invoke_llm(query, runtime_llm_model, [])
assert result.content == 'ok'
details = record.call_args.kwargs['details']
assert details['usage'] == usage
assert details['timing']['attempts'] == 1
assert details['timing']['total_ms'] >= details['timing']['response_completed_ms']
assert query.variables['_llm_usage'] == usage
@pytest.mark.asyncio
async def test_hidden_reasoning_marks_raw_event_and_usage():
state = CallState(RequestPolicy())
adapter = LiteLLMRequester(Mock(), {})
async def raw():
yield SimpleNamespace(
choices=[
SimpleNamespace(
delta=SimpleNamespace(model_dump=lambda: {'reasoning_content': 'hidden'}), finish_reason=None
)
]
)
yield SimpleNamespace(choices=[], usage={'prompt_tokens': 3, 'completion_tokens': 2, 'total_tokens': 5})
token = current_call.set(state)
try:
with patch(
'langbot.pkg.provider.modelmgr.requesters.litellmchat.acompletion', new=AsyncMock(return_value=raw())
):
chunks = [c async for c in adapter.invoke_llm_stream(None, MockRuntimeModel(), [], remove_think=True)]
finally:
current_call.reset(token)
assert state.raw_started
assert 'first_event_ms' in state.metrics and 'first_reasoning_ms' in state.metrics
assert 'first_content_ms' not in state.metrics
assert state.usage['total_tokens'] == 5
assert not any(c.content for c in chunks)
@pytest.mark.asyncio
async def test_real_http_retry_count_and_request_body(runtime_provider, runtime_llm_model):
"""Exercise the pinned LiteLLM/OpenAI SDK against a local HTTP server."""
from aiohttp import web
received = []
async def handler(request):
received.append(await request.json())
if len(received) == 1:
return web.json_response({'error': {'message': 'busy', 'type': 'server_error'}}, status=503)
return web.json_response(
{
'id': 'chatcmpl-local',
'object': 'chat.completion',
'created': 1,
'model': 'gpt-4o',
'choices': [{'index': 0, 'message': {'role': 'assistant', 'content': 'ok'}, 'finish_reason': 'stop'}],
'usage': {'prompt_tokens': 4, 'completion_tokens': 1, 'total_tokens': 5},
}
)
app = web.Application()
app.router.add_post('/v1/chat/completions', handler)
runner = web.AppRunner(app)
await runner.setup()
site = web.TCPSite(runner, '127.0.0.1', 0)
await site.start()
port = site._server.sockets[0].getsockname()[1]
runtime_provider.requester = LiteLLMRequester(
Mock(),
{
'base_url': f'http://127.0.0.1:{port}/v1',
'custom_llm_provider': 'openai',
},
)
runtime_llm_model.model_entity.name = 'gpt-4o'
runtime_llm_model.model_entity.extra_args = {}
try:
with pytest.raises(RequesterError):
await runtime_provider.invoke_llm(
None,
runtime_llm_model,
[Message(role='user', content='hello')],
extra_args={'num_retries': 0},
execution_context=TEST_EXECUTION_CONTEXT,
)
assert len(received) == 1
received.clear()
result = await runtime_provider.invoke_llm(
None,
runtime_llm_model,
[Message(role='user', content='hello')],
extra_args={'request_policy': {'retries': 1, 'backoff_base': 0.001}},
execution_context=TEST_EXECUTION_CONTEXT,
)
assert result.content == 'ok'
assert len(received) == 2
assert all('request_policy' not in body and 'num_retries' not in body for body in received)
finally:
await runner.cleanup()
@pytest.mark.asyncio
async def test_stream_idle_timeout_after_output(runtime_provider, runtime_llm_model):
async def stream(**kwargs):
yield MessageChunk(role='assistant', content='first')
await asyncio.Event().wait()
runtime_provider.requester.invoke_llm_stream = stream
with pytest.raises(RequesterError) as caught:
async for _ in runtime_provider.invoke_llm_stream(
None,
runtime_llm_model,
[],
extra_args={'request_policy': {'stream_idle_timeout': 0.02}},
execution_context=TEST_EXECUTION_CONTEXT,
):
pass
assert caught.value.phase == 'stream_idle'
def test_details_migration_preserves_existing_records_and_is_idempotent():
import importlib
import sqlalchemy as sa
from alembic.migration import MigrationContext
from alembic.operations import Operations
migration = importlib.import_module('langbot.pkg.persistence.alembic.versions.0038_llm_call_details')
engine = sa.create_engine('sqlite://')
try:
with engine.begin() as connection:
connection.exec_driver_sql('CREATE TABLE monitoring_llm_calls (id TEXT PRIMARY KEY)')
connection.exec_driver_sql("INSERT INTO monitoring_llm_calls (id) VALUES ('old')")
with Operations.context(MigrationContext.configure(connection)):
migration.upgrade()
migration.upgrade()
assert connection.exec_driver_sql('SELECT id, details FROM monitoring_llm_calls').one() == ('old', None)
table = sa.Table('monitoring_llm_calls', sa.MetaData(), autoload_with=connection)
payload = {'usage': {'completion_tokens_details': {'reasoning_tokens': 8}}, 'timing': {'first_event_ms': 1}}
connection.execute(table.insert().values(id='new', details=payload))
assert connection.execute(sa.select(table.c.details).where(table.c.id == 'new')).scalar() == payload
finally:
engine.dispose()
@pytest.mark.asyncio
async def test_hidden_raw_events_reset_idle_not_first_visible_deadline(runtime_provider, runtime_llm_model):
adapter = LiteLLMRequester(Mock(), {})
runtime_provider.requester = adapter
runtime_llm_model.model_entity.extra_args = {}
async def raw():
for _ in range(5):
await asyncio.sleep(0.3)
yield SimpleNamespace(
choices=[
SimpleNamespace(delta=SimpleNamespace(model_dump=lambda: {'role': 'assistant'}), finish_reason=None)
]
)
yield SimpleNamespace(
choices=[
SimpleNamespace(delta=SimpleNamespace(model_dump=lambda: {'content': 'answer'}), finish_reason='stop')
]
)
with patch('langbot.pkg.provider.modelmgr.requesters.litellmchat.acompletion', new=AsyncMock(return_value=raw())):
chunks = [
chunk
async for chunk in runtime_provider.invoke_llm_stream(
None,
runtime_llm_model,
[],
extra_args={'request_policy': {'first_event_timeout': 1, 'stream_idle_timeout': 2}},
execution_context=TEST_EXECUTION_CONTEXT,
)
]
assert ''.join(chunk.content or '' for chunk in chunks) == 'answer'
@pytest.mark.asyncio
async def test_monitoring_service_persists_full_details():
from langbot.pkg.api.http.service.monitoring import MonitoringService
persistence = SimpleNamespace(execute_async=AsyncMock())
service = MonitoringService(SimpleNamespace(persistence_mgr=persistence))
details = {
'usage': {
'prompt_tokens': 4,
'completion_tokens': 3,
'total_tokens': 9,
'cache_read_input_tokens': 2,
'provider_extra': {'cost': 0.01},
},
'timing': {'attempts': 2},
'error': None,
}
await service.record_llm_call(
TEST_EXECUTION_CONTEXT,
bot_id='b',
bot_name='b',
pipeline_id='p',
pipeline_name='p',
session_id='s',
model_name='m',
input_tokens=4,
output_tokens=3,
duration=10,
details=details,
)
values = persistence.execute_async.call_args.args[0].compile().params
assert values['details'] == details
assert values['total_tokens'] == 9
def test_codex_usage_preserves_provider_fields():
from langbot.pkg.provider.modelmgr.requesters.codex import _usage
raw = {
'input_tokens': 4,
'output_tokens': 2,
'cost': 0.02,
'output_tokens_details': {'reasoning_tokens': 1},
'future_counter': 7,
}
normalized = _usage({'usage': raw})
assert all(normalized[key] == value for key, value in raw.items())
assert normalized['total_tokens'] == 6
@pytest.mark.asyncio
async def test_nonstream_uses_response_budget_not_first_event():
adapter = LiteLLMRequester(Mock(), {})
state = CallState(RequestPolicy(first_event_timeout=0.001, total_timeout=1))
async def completion(**kwargs):
await asyncio.sleep(0.02)
return 'complete'
token = current_call.set(state)
try:
with patch('langbot.pkg.provider.modelmgr.requesters.litellmchat.acompletion', new=completion):
assert await adapter._completion({'stream': False}) == 'complete'
finally:
current_call.reset(token)
@pytest.mark.asyncio
@pytest.mark.parametrize('cancel_read', [False, True])
async def test_litellm_stream_close_and_cancel_release_upstream(runtime_provider, runtime_llm_model, cancel_read):
adapter = LiteLLMRequester(Mock(), {})
runtime_provider.requester = adapter
runtime_llm_model.model_entity.extra_args = {}
closed = asyncio.Event()
reading = asyncio.Event()
async def raw():
try:
yield SimpleNamespace(
choices=[
SimpleNamespace(delta=SimpleNamespace(model_dump=lambda: {'content': 'first'}), finish_reason=None)
]
)
reading.set()
await asyncio.Event().wait()
finally:
closed.set()
completion = AsyncMock(return_value=raw())
with patch('langbot.pkg.provider.modelmgr.requesters.litellmchat.acompletion', new=completion):
stream = runtime_provider.invoke_llm_stream(
None,
runtime_llm_model,
[],
extra_args={'request_policy': {'retries': 3}},
execution_context=TEST_EXECUTION_CONTEXT,
)
assert (await anext(stream)).content == 'first'
if cancel_read:
task = asyncio.create_task(anext(stream))
await reading.wait()
task.cancel()
with pytest.raises(asyncio.CancelledError):
await task
else:
await stream.aclose()
assert closed.is_set()
assert completion.await_count == 1
assert current_call.get() is None
@pytest.mark.asyncio
async def test_cancel_during_backoff_does_not_start_another_attempt(runtime_provider, runtime_llm_model):
runtime_provider.requester.invoke_llm = AsyncMock(side_effect=RequesterError('busy', retryable=True))
sleeping = asyncio.Event()
async def backoff(delay):
sleeping.set()
await asyncio.Event().wait()
with patch('langbot.pkg.provider.modelmgr.request_policy.asyncio.sleep', new=backoff):
task = asyncio.create_task(
runtime_provider.invoke_llm(
None,
runtime_llm_model,
[],
extra_args={'request_policy': {'retries': 3}},
execution_context=TEST_EXECUTION_CONTEXT,
)
)
await sleeping.wait()
task.cancel()
with pytest.raises(asyncio.CancelledError):
await task
assert runtime_provider.requester.invoke_llm.await_count == 1
@pytest.mark.asyncio
async def test_legacy_none_chunk_is_ignored(runtime_provider, runtime_llm_model):
async def stream(**kwargs):
yield None
yield MessageChunk(role='assistant', content='ok')
runtime_provider.requester.invoke_llm_stream = stream
chunks = [
chunk
async for chunk in runtime_provider.invoke_llm_stream(
None, runtime_llm_model, [], execution_context=TEST_EXECUTION_CONTEXT
)
]
assert [chunk.content for chunk in chunks] == ['ok']
uv.lock
Generated
+4 -4
View File
@@ -2185,7 +2185,7 @@ requires-dist = [
{ name = "gewechat-client", specifier = ">=0.1.5" },
{ name = "html2text", specifier = ">=2024.2.26" },
{ name = "httpx", extras = ["socks"], specifier = ">=0.28.1" },
{ name = "langbot-plugin", specifier = "==0.7.11" },
{ name = "langbot-plugin", specifier = "==0.7.12" },
{ name = "langchain", specifier = ">=1.3.9" },
{ name = "langchain-core", specifier = ">=1.3.3" },
{ name = "langchain-text-splitters", specifier = ">=1.1.2" },
@@ -2255,7 +2255,7 @@ dev = [
[[package]]
name = "langbot-plugin"
version = "0.7.11"
version = "0.7.12"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "aiofiles" },
@@ -2276,9 +2276,9 @@ dependencies = [
{ name = "watchdog" },
{ name = "websockets" },
]
sdist = { url = "https://files.pythonhosted.org/packages/6a/c4/ed01dd73ab5047c75a0563dd38b88146d5d509b889197d5ec47fc27541e5/langbot_plugin-0.7.11.tar.gz", hash = "sha256:19f57a020d435fd414b02cbc477a70f880a462cba73e4ec767fb7dbbffb1bc56", size = 697187, upload-time = "2026-10-14T07:41:24.85Z" }
sdist = { url = "https://files.pythonhosted.org/packages/67/4c/0286247a6caa24a21328fe5dae32b295635143cb863d9eb879384bf613ce/langbot_plugin-0.7.12.tar.gz", hash = "sha256:36a661bfe25233a5c4d74e86568df6831e9055e1b6c88bfad8dee7c2e207766e", size = 697786, upload-time = "2026-10-15T12:01:17.213Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/53/5b/959ca0b62eac62a95f7740ba3f7c36dd34a1d7ae8499eff17e417b5d2706/langbot_plugin-0.7.11-py3-none-any.whl", hash = "sha256:234d31a8ed4df7e3d8c648735062e8ceba90ac81dfa51a6c678be5a872b505fe", size = 447857, upload-time = "2026-10-14T07:41:23.14Z" },
{ url = "https://files.pythonhosted.org/packages/5c/34/2e40056497e60411e851b30320087e1385865b6357185a05e40a75ef81ca/langbot_plugin-0.7.12-py3-none-any.whl", hash = "sha256:b0bbd6bbfedf15ed48377c3904c73af62c4ee88a676b06649b470fc1bad9416a", size = 448016, upload-time = "2026-10-15T12:01:15.542Z" },
]
[[package]]