diff --git a/docs/LLM_LATENCY_BENCHMARK.md b/docs/LLM_LATENCY_BENCHMARK.md new file mode 100644 index 000000000..c88f4ddfd --- /dev/null +++ b/docs/LLM_LATENCY_BENCHMARK.md @@ -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. diff --git a/docs/LLM_REQUEST_E2E.md b/docs/LLM_REQUEST_E2E.md new file mode 100644 index 000000000..41f18177e --- /dev/null +++ b/docs/LLM_REQUEST_E2E.md @@ -0,0 +1,107 @@ +# 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 the SDK companion fix.** Core's pinned +SDK has not been silently replaced or upgraded. The E2E suite executes the close +checks against published SDKs and reports their known failure as XFAIL. 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. diff --git a/docs/LLM_REQUEST_POLICY.md b/docs/LLM_REQUEST_POLICY.md new file mode 100644 index 000000000..1b367f6b6 --- /dev/null +++ b/docs/LLM_REQUEST_POLICY.md @@ -0,0 +1,254 @@ +# 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 the companion +[SDK fix](https://github.com/langbot-app/langbot-plugin-sdk/pull/177). +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- +``` + +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. diff --git a/docs/llm-latency-2026-10-15.json b/docs/llm-latency-2026-10-15.json new file mode 100644 index 000000000..628e02868 --- /dev/null +++ b/docs/llm-latency-2026-10-15.json @@ -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 + } + ] +} diff --git a/scripts/llm_latency_probe.py b/scripts/llm_latency_probe.py new file mode 100644 index 000000000..8696f72da --- /dev/null +++ b/scripts/llm_latency_probe.py @@ -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)) diff --git a/skills/skills/langbot-mcp-ops/SKILL.md b/skills/skills/langbot-mcp-ops/SKILL.md index edea7b762..ebe12cb47 100644 --- a/skills/skills/langbot-mcp-ops/SKILL.md +++ b/skills/skills/langbot-mcp-ops/SKILL.md @@ -226,3 +226,26 @@ are not rewritten. Embedding and rerank models are not eligible. ### Reset a bot session context Use `reset_session_context(bot_id, session_id)` only when a user asks to start a session afresh. It requires `resource.manage`, preserves monitoring records, and refuses sessions with active tasks. It clears conversation runner state and excludes earlier transcript entries from subsequent model context. Files and long-term memory are not removed. + +## 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. diff --git a/src/langbot/pkg/api/http/service/monitoring.py b/src/langbot/pkg/api/http/service/monitoring.py index dbc4eefb0..f9f924ebc 100644 --- a/src/langbot/pkg/api/http/service/monitoring.py +++ b/src/langbot/pkg/api/http/service/monitoring.py @@ -605,6 +605,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) @@ -616,7 +617,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, @@ -627,6 +628,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( diff --git a/src/langbot/pkg/api/mcp/server.py b/src/langbot/pkg/api/mcp/server.py index 77fa0ade2..ae1686229 100644 --- a/src/langbot/pkg/api/mcp/server.py +++ b/src/langbot/pkg/api/mcp/server.py @@ -165,7 +165,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'], @@ -513,7 +513,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)) diff --git a/src/langbot/pkg/entity/persistence/monitoring.py b/src/langbot/pkg/entity/persistence/monitoring.py index 9b07d4f5a..ecf032986 100644 --- a/src/langbot/pkg/entity/persistence/monitoring.py +++ b/src/langbot/pkg/entity/persistence/monitoring.py @@ -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'), diff --git a/src/langbot/pkg/persistence/alembic/versions/0038_llm_call_details.py b/src/langbot/pkg/persistence/alembic/versions/0038_llm_call_details.py new file mode 100644 index 000000000..a91469f23 --- /dev/null +++ b/src/langbot/pkg/persistence/alembic/versions/0038_llm_call_details.py @@ -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') diff --git a/src/langbot/pkg/pipeline/monitoring_helper.py b/src/langbot/pkg/pipeline/monitoring_helper.py index 14ff187db..4efb21498 100644 --- a/src/langbot/pkg/pipeline/monitoring_helper.py +++ b/src/langbot/pkg/pipeline/monitoring_helper.py @@ -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}') diff --git a/src/langbot/pkg/plugin/handler.py b/src/langbot/pkg/plugin/handler.py index 61137e0f2..12fade67b 100644 --- a/src/langbot/pkg/plugin/handler.py +++ b/src/langbot/pkg/plugin/handler.py @@ -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 [] diff --git a/src/langbot/pkg/provider/modelmgr/codex_errors.py b/src/langbot/pkg/provider/modelmgr/codex_errors.py index a7ceb0301..da32d531e 100644 --- a/src/langbot/pkg/provider/modelmgr/codex_errors.py +++ b/src/langbot/pkg/provider/modelmgr/codex_errors.py @@ -1,10 +1,19 @@ """Explicitly safe Codex failures; never construct messages from upstream bodies.""" +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__(message) + ValueError.__init__(self, 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 diff --git a/src/langbot/pkg/provider/modelmgr/errors.py b/src/langbot/pkg/provider/modelmgr/errors.py index dc3b35b6f..fda01de86 100644 --- a/src/langbot/pkg/provider/modelmgr/errors.py +++ b/src/langbot/pkg/provider/modelmgr/errors.py @@ -1,5 +1,28 @@ 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__('模型请求失败: ' + 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') + } diff --git a/src/langbot/pkg/provider/modelmgr/http_pool.py b/src/langbot/pkg/provider/modelmgr/http_pool.py new file mode 100644 index 000000000..a29efefd9 --- /dev/null +++ b/src/langbot/pkg/provider/modelmgr/http_pool.py @@ -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 diff --git a/src/langbot/pkg/provider/modelmgr/modelmgr.py b/src/langbot/pkg/provider/modelmgr/modelmgr.py index e86f1761b..e13464a9c 100644 --- a/src/langbot/pkg/provider/modelmgr/modelmgr.py +++ b/src/langbot/pkg/provider/modelmgr/modelmgr.py @@ -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': diff --git a/src/langbot/pkg/provider/modelmgr/prompt_caching.py b/src/langbot/pkg/provider/modelmgr/prompt_caching.py new file mode 100644 index 000000000..5ec55277f --- /dev/null +++ b/src/langbot/pkg/provider/modelmgr/prompt_caching.py @@ -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 diff --git a/src/langbot/pkg/provider/modelmgr/reasoning.py b/src/langbot/pkg/provider/modelmgr/reasoning.py index ecc6c38e2..dead45620 100644 --- a/src/langbot/pkg/provider/modelmgr/reasoning.py +++ b/src/langbot/pkg/provider/modelmgr/reasoning.py @@ -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) diff --git a/src/langbot/pkg/provider/modelmgr/request_policy.py b/src/langbot/pkg/provider/modelmgr/request_policy.py new file mode 100644 index 000000000..02438d8bd --- /dev/null +++ b/src/langbot/pkg/provider/modelmgr/request_policy.py @@ -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} diff --git a/src/langbot/pkg/provider/modelmgr/requester.py b/src/langbot/pkg/provider/modelmgr/requester.py index 99f6e885d..1aa75e4d9 100644 --- a/src/langbot/pkg/provider/modelmgr/requester.py +++ b/src/langbot/pkg/provider/modelmgr/requester.py @@ -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: diff --git a/src/langbot/pkg/provider/modelmgr/requesters/codex.py b/src/langbot/pkg/provider/modelmgr/requesters/codex.py index 5b793ef75..728382415 100644 --- a/src/langbot/pkg/provider/modelmgr/requesters/codex.py +++ b/src/langbot/pkg/provider/modelmgr/requesters/codex.py @@ -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: diff --git a/src/langbot/pkg/provider/modelmgr/requesters/litellmchat.py b/src/langbot/pkg/provider/modelmgr/requesters/litellmchat.py index e927a3baf..dfaaf44bf 100644 --- a/src/langbot/pkg/provider/modelmgr/requesters/litellmchat.py +++ b/src/langbot/pkg/provider/modelmgr/requesters/litellmchat.py @@ -3,10 +3,23 @@ from __future__ import annotations 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 @@ -15,6 +28,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. @@ -131,6 +175,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( @@ -179,9 +225,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.""" @@ -197,8 +317,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 @@ -297,6 +430,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) @@ -561,6 +697,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]: @@ -568,14 +718,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'}}} @@ -1027,7 +1184,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'] @@ -1035,24 +1192,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(f'上下文长度超限: {str(e)}') - if isinstance(e, litellm.BadRequestError): - raise errors.RequesterError(f'请求参数错误: {str(e)}') - if isinstance(e, litellm.AuthenticationError): - raise errors.RequesterError(f'API key 无效: {str(e)}') - if isinstance(e, litellm.NotFoundError): - raise errors.RequesterError(f'模型或路径无效: {str(e)}') - if isinstance(e, litellm.RateLimitError): - raise errors.RequesterError(f'请求过于频繁或余额不足: {str(e)}') - if isinstance(e, litellm.Timeout): - raise errors.RequesterError(f'请求超时: {str(e)}') - if isinstance(e, litellm.APIConnectionError): - raise errors.RequesterError(f'连接错误: {str(e)}') - if isinstance(e, litellm.APIError): - raise errors.RequesterError(f'API 错误: {str(e)}') - raise errors.RequesterError(f'未知错误: {str(e)}') + if isinstance(e, errors.RequesterError): + raise e + metadata = error_metadata(e) + mappings = ( + (litellm.ContextWindowExceededError, 'context_length', '上下文长度超限', False), + (litellm.AuthenticationError, 'authentication', 'API key 无效', False), + (litellm.NotFoundError, 'not_found', '模型或路径无效', False), + (litellm.RateLimitError, 'rate_limit', '请求过于频繁或余额不足', True), + (litellm.Timeout, 'timeout', '请求超时', True), + (litellm.APIConnectionError, 'connection', '连接错误', True), + (litellm.BadRequestError, 'bad_request', '请求参数错误', 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( + 'API 错误' if isinstance(e, litellm.APIError) else '未知错误', + 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, @@ -1063,14 +1298,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 = { @@ -1083,10 +1335,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: @@ -1111,8 +1367,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') @@ -1121,6 +1392,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( @@ -1142,7 +1422,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) @@ -1162,38 +1442,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, @@ -1216,141 +1499,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 = '\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\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 = '\n' + reasoning_started = True + else: + delta_content = '' + delta_content += display_reasoning_content + if delta.get('content'): + delta_content += f'\n\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 = '\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\n{delta.get("content")}' - reasoning_closed = True + if not reasoning_started: + delta_content = '\n' + reasoning_started = True + else: + delta_content = '' + delta_content += display_reasoning_content + if delta.get('content'): + delta_content += f'\n\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\n{delta_content}' - reasoning_closed = True + elif delta_content and not remove_think and reasoning_started and not reasoning_closed: + delta_content = f'\n\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\n' - reasoning_closed = True + if finish_reason and not remove_think and reasoning_started and not reasoning_closed: + delta_content = f'{delta_content}\n\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\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\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, diff --git a/src/langbot/templates/config.yaml b/src/langbot/templates/config.yaml index c04226b07..e58d48f1b 100644 --- a/src/langbot/templates/config.yaml +++ b/src/langbot/templates/config.yaml @@ -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 diff --git a/tests/e2e/test_llm_request_chain.py b/tests/e2e/test_llm_request_chain.py new file mode 100644 index 000000000..fdfabd9be --- /dev/null +++ b/tests/e2e/test_llm_request_chain.py @@ -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 diff --git a/tests/e2e/utils/llm_probe_runner.py b/tests/e2e/utils/llm_probe_runner.py new file mode 100644 index 000000000..9b1115408 --- /dev/null +++ b/tests/e2e/utils/llm_probe_runner.py @@ -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') diff --git a/tests/unit_tests/plugin/test_handler_helpers.py b/tests/unit_tests/plugin/test_handler_helpers.py index dc86b7a5f..b126a9b18 100644 --- a/tests/unit_tests/plugin/test_handler_helpers.py +++ b/tests/unit_tests/plugin/test_handler_helpers.py @@ -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(): diff --git a/tests/unit_tests/provider/test_codex.py b/tests/unit_tests/provider/test_codex.py index 14dbd6832..cb0e9e046 100644 --- a/tests/unit_tests/provider/test_codex.py +++ b/tests/unit_tests/provider/test_codex.py @@ -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' diff --git a/tests/unit_tests/provider/test_http_pool_streams.py b/tests/unit_tests/provider/test_http_pool_streams.py new file mode 100644 index 000000000..b474e19d5 --- /dev/null +++ b/tests/unit_tests/provider/test_http_pool_streams.py @@ -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() diff --git a/tests/unit_tests/provider/test_llm_preparation.py b/tests/unit_tests/provider/test_llm_preparation.py new file mode 100644 index 000000000..626003982 --- /dev/null +++ b/tests/unit_tests/provider/test_llm_preparation.py @@ -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() diff --git a/tests/unit_tests/provider/test_request_policy.py b/tests/unit_tests/provider/test_request_policy.py new file mode 100644 index 000000000..5fb873d17 --- /dev/null +++ b/tests/unit_tests/provider/test_request_policy.py @@ -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='模型请求失败'): + 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']