feat(graph): stream llm-node completions so node timeout governs long generations
Graph llm nodes (and the structured-output extractor they share with agent nodes) now obtain the model reply through a quiet streaming transport, `call_chat_completions_streaming_quiet`, instead of the non-streaming call. Because chunks arrive as they are generated, reqwest's client-level `read_timeout` (default 300s) bounds only a genuine stall between chunks and the node's own `timeout` becomes the effective bound on total generation. Previously zero bytes arrived until the whole reply finished, so a 300s read_timeout silently capped every llm node regardless of its `timeout` (seen live: code-reviewer's `synthesize` node dying on a 61-file review). The quiet transport renders nothing (silent SseHandler, no spinner), keeps tool-call handling identical to the non-streaming path (call-loop detection off), falls back to the non-streaming request for catalog `no_stream` models (raced against the abort), and never returns partial output as success. Supporting changes: - `sse_stream`/`json_stream` keep the typed `reqwest::Error` in the anyhow chain so `is_transient_error` still classifies read_timeout stalls as transient and `max_attempts` retries them (previously stringified). - `structured::extract` takes the graph/session AbortSignal and refuses to issue an extraction request for an already-aborted graph. - `run_chat_loop` gains an injectable CompletionRunner seam (owning client construction + the before-hook) so the timeout/retry/abort behaviour is unit-testable over a paused clock; a fake `impl Client` drives the real trait-default streaming for 600s of virtual time; real-socket tests pin that read_timeout fires on a stall but not on a slow continuous stream.
This commit is contained in:
6 files changed
+877
-53
No files matched your search
+26
-9
@@ -1,6 +1,6 @@
|
||||
use crate::client::call_chat_completions;
|
||||
use crate::client::call_chat_completions_streaming_quiet;
|
||||
use crate::config::{Input, RequestContext, Role, RoleLike};
|
||||
use crate::utils::create_abort_signal;
|
||||
use crate::utils::AbortSignal;
|
||||
use anyhow::{Context, Result, bail};
|
||||
use serde_json::Value;
|
||||
use std::sync::Arc;
|
||||
@@ -19,18 +19,24 @@ Rules:
|
||||
- If a field cannot be determined from the input, use `null` (when allowed) or your best inferred value.
|
||||
- Do NOT invent fields not present in the schema.";
|
||||
|
||||
pub async fn extract(raw: &str, schema: &Value, parent_ctx: &mut RequestContext) -> Result<Value> {
|
||||
pub async fn extract(
|
||||
raw: &str,
|
||||
schema: &Value,
|
||||
parent_ctx: &mut RequestContext,
|
||||
abort: &AbortSignal,
|
||||
) -> Result<Value> {
|
||||
if let Some(parsed) = try_parse_json(raw) {
|
||||
return Ok(parsed);
|
||||
}
|
||||
|
||||
extract_via_extractor(raw, schema, parent_ctx, false).await
|
||||
extract_via_extractor(raw, schema, parent_ctx, abort, false).await
|
||||
}
|
||||
|
||||
async fn extract_via_extractor(
|
||||
raw: &str,
|
||||
schema: &Value,
|
||||
parent_ctx: &mut RequestContext,
|
||||
abort: &AbortSignal,
|
||||
is_repair: bool,
|
||||
) -> Result<Value> {
|
||||
let role = build_extractor_role(parent_ctx);
|
||||
@@ -38,7 +44,7 @@ async fn extract_via_extractor(
|
||||
|
||||
let saved_role = parent_ctx.role.clone();
|
||||
parent_ctx.role = Some(role);
|
||||
let result = run_one_shot(&prompt, parent_ctx).await;
|
||||
let result = run_one_shot(&prompt, parent_ctx, abort).await;
|
||||
parent_ctx.role = saved_role;
|
||||
|
||||
let output = result.context("Structured-output extractor LLM call failed")?;
|
||||
@@ -49,7 +55,12 @@ async fn extract_via_extractor(
|
||||
"Structured-output extractor failed to produce valid JSON after repair retry. \
|
||||
Last response:\n{output}"
|
||||
),
|
||||
None => Box::pin(extract_via_extractor(&output, schema, parent_ctx, true)).await,
|
||||
None => {
|
||||
Box::pin(extract_via_extractor(
|
||||
&output, schema, parent_ctx, abort, true,
|
||||
))
|
||||
.await
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -73,15 +84,21 @@ fn build_extractor_prompt(raw: &str, schema: &Value, is_repair: bool) -> String
|
||||
}
|
||||
}
|
||||
|
||||
async fn run_one_shot(prompt: &str, ctx: &mut RequestContext) -> Result<String> {
|
||||
let abort = create_abort_signal();
|
||||
async fn run_one_shot(
|
||||
prompt: &str,
|
||||
ctx: &mut RequestContext,
|
||||
abort: &AbortSignal,
|
||||
) -> Result<String> {
|
||||
if abort.aborted() {
|
||||
bail!("Aborted.");
|
||||
}
|
||||
let app_cfg = Arc::clone(&ctx.app.config);
|
||||
let role_for_input = ctx.role.clone();
|
||||
let input = Input::from_str(ctx, prompt, role_for_input)?;
|
||||
let client = input.create_client()?;
|
||||
ctx.before_chat_completion(&input)?;
|
||||
let (output, tool_results) =
|
||||
call_chat_completions(&input, false, false, client.as_ref(), ctx, abort).await?;
|
||||
call_chat_completions_streaming_quiet(&input, client.as_ref(), ctx, abort.clone()).await?;
|
||||
ctx.after_chat_completion(app_cfg.as_ref(), &input, &output, &tool_results)?;
|
||||
|
||||
Ok(output)
|
||||
|
||||
Reference in new issue
Block a user