From 2d75667b9ab1482ef3b387b86e550dd6dbcb025f Mon Sep 17 00:00:00 2001 From: runstone Date: Fri, 12 Jun 2026 11:43:09 +0800 Subject: [PATCH] feat(ai): upgrade agent loop tool calling to true SSE streaming MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Move tool calling from non-streaming (complete) API calls to real streaming SSE, so text, reasoning, and tool call arguments all arrive incrementally — users see the AI think and act in real time. - Extract ~370 lines of non-streaming tool-call functions from agent_loop.rs (call_with_tools, call_openai_with_tools, call_claude_with_tools, call_gemini_with_tools) and replace with a single ai::stream_with_tools() dispatch call. - Add streaming-with-tools infrastructure to ai.rs (~550 lines): - StreamToolEvent enum: Chunk, ToolCallStart, ToolCallDelta, ToolCallComplete — provider-agnostic, feeds into the StreamingToolCallAccumulator - stream_claude_with_tools(): parse SSE content_block_start/delta/ stop events, incremental input_json_delta for tool arguments - stream_openai_with_tools(): parse SSE delta.tool_calls with incremental function.arguments - stream_gemini_with_tools(): parse SSE functionCall (Gemini sends complete objects, not deltas — emit as one chunk) - Cancellation support via tokio::select! in every streaming loop - Extended timeout for reasoning models (600s vs 120s) - StreamingToolCallAccumulator: collect partial tool call fragments from streaming deltas, then deserialize into complete ToolCall objects in index order once the stream ends. --- .../src/components/editor/AiAssistant.vue | 22 +- apps/desktop/src/lib/ai.ts | 4 +- crates/dbx-core/src/agent_loop.rs | 340 +---------- crates/dbx-core/src/agent_tools.rs | 4 +- crates/dbx-core/src/ai.rs | 534 ++++++++++++++++++ 5 files changed, 565 insertions(+), 339 deletions(-) diff --git a/apps/desktop/src/components/editor/AiAssistant.vue b/apps/desktop/src/components/editor/AiAssistant.vue index cc303399aa..d59ba3ea7f 100644 --- a/apps/desktop/src/components/editor/AiAssistant.vue +++ b/apps/desktop/src/components/editor/AiAssistant.vue @@ -253,8 +253,10 @@ function agentStepClass(tone: AiAgentStepTone): string { } function agentStepTitle(step: AiAgentStepItem): string { - if (!step.titleKey) return t(step.labelKey); - return t(step.titleKey, step.titleParams || {}); + if (step.titleKey) return t(step.titleKey, step.titleParams || {}); + const tool = step.titleParams?.tool; + if (tool) return `${t(step.labelKey)}: ${tool}`; + return t(step.labelKey); } function toggleReasoning(index: number) { @@ -513,6 +515,22 @@ async function send() { if (event.type === "reasoning_delta" && event.delta) { appendAssistantReasoning(assistantIdx, event.delta); } + // Real-time agent step rendering + if (event.type === "tool_call_start" || event.type === "tool_call_end") { + const msg = messages.value[assistantIdx]; + if (msg) { + const steps = agentEvents + .filter((e) => e.type === "tool_call_start" || e.type === "tool_call_end") + .map((e) => ({ + key: `${e.tool_call_id || ""}-${e.type}`, + labelKey: e.type === "tool_call_start" ? "ai.agentSteps.callingTool" : e.is_error ? "ai.agentSteps.toolError" : "ai.agentSteps.toolDone", + tone: (e.type === "tool_call_start" ? "active" : e.is_error ? "danger" : "success") as AiAgentStepTone, + titleKey: undefined, + titleParams: { tool: e.tool_name || "" }, + })); + msg.agentSteps = steps; + } + } scrollToBottom(); }, sessionId, diff --git a/apps/desktop/src/lib/ai.ts b/apps/desktop/src/lib/ai.ts index 65bdf2c99a..ed4ef08024 100644 --- a/apps/desktop/src/lib/ai.ts +++ b/apps/desktop/src/lib/ai.ts @@ -178,8 +178,8 @@ function buildBasePromptLines(isZh: boolean): string[] { isZh ? "精确、保守，根据当前数据库方言生成 SQL。" : "Be precise, conservative, and adapt SQL to the active database dialect.", isZh ? "严格使用当前数据库方言；标识符引用、分页、日期函数、字符串拼接、LIMIT/TOP/OFFSET 语法必须匹配数据库类型。" : "Strictly use the active database dialect; identifier quoting, pagination, date functions, string concatenation, and LIMIT/TOP/OFFSET syntax must match the database type.", isZh - ? "对于普通数据查询，优先使用下面已加载的 Schema 上下文，不要为了重复确认已给出的结构而查询 information_schema 或系统表。" - : "For ordinary data queries, prefer the loaded schema context below. Do not query information_schema or system tables merely to rediscover structure already provided.", + ? "对于普通数据查询，优先使用下面已加载的 Schema 上下文，不要为了重复确认已给出的结构而查询 information_schema 或系统表。但当用户询问某表的字段详情、列信息时，应使用 get_columns 工具获取最权威完整的定义。" + : "For ordinary data queries, prefer the loaded schema context below. Do not query information_schema or system tables merely to rediscover structure already provided. However, when the user asks for detailed column/field information of a specific table, use the get_columns tool for the authoritative and complete definition.", isZh ? "例外：当用户明确询问当前有哪些表/Schema、某张表是否存在、或需要盘点数据库对象时，应生成符合当前方言的只读元数据查询（例如 SHOW TABLES、information_schema、sqlite_master 等）。" : "Exception: when the user explicitly asks what tables/schemas exist, whether a table exists, or asks for database object inventory, generate a read-only metadata query appropriate for the active dialect (for example SHOW TABLES, information_schema, sqlite_master).", diff --git a/crates/dbx-core/src/agent_loop.rs b/crates/dbx-core/src/agent_loop.rs index 59ab95141a..f3a1e0776b 100644 --- a/crates/dbx-core/src/agent_loop.rs +++ b/crates/dbx-core/src/agent_loop.rs @@ -149,7 +149,7 @@ pub async fn run_agent_loop( // Add tool result to conversation for the next LLM call // Uses "tool" role per OpenAI convention; provider-specific conversion - // happens in call_*_with_tools functions when building provider API requests. + // happens in ai::stream_with_tools when building provider API requests. conversation_messages.push(AiMessage { role: "tool".to_string(), content: result.content.clone(), @@ -170,7 +170,7 @@ fn build_tool_request( config: &AiConfig, system_prompt: &str, messages: &[AiMessage], - _tools: &[ToolDefinition], // Tools are injected in stream_with_tools / call_*_with_tools, not via AiCompletionRequest. + _tools: &[ToolDefinition], // Tools are injected in ai::stream_with_tools, not via AiCompletionRequest. max_tokens: Option, temperature: Option, ) -> AiCompletionRequest { @@ -185,15 +185,10 @@ fn build_tool_request( } } -/// Collect an LLM response with tool support, parsing tool_calls. +/// Stream an LLM response with tool support, parsing tool_calls from SSE deltas. /// -/// Named "stream" for consistency with the agent loop's streaming-API contract, -/// but currently uses a non-streaming (complete) LLM call internally, then emits -/// the full text as a single delta. True streaming tool_calls is deferred to -/// Phase 2 (provider-specific delta accumulation). -/// -/// Check `cancelled` before the LLM call; if the LLM request has already started -/// it will run to completion (non-streaming requests cannot be interrupted mid-flight). +/// True streaming: text, reasoning, and tool call arguments are all emitted +/// incrementally as they arrive from the provider. async fn stream_with_tools( config: &AiConfig, request: &AiCompletionRequest, @@ -208,330 +203,7 @@ async fn stream_with_tools( return Err("Agent loop cancelled".to_string()); } - // TODO(agent-loop-phase2): Implement streaming tool_calls for each provider. - // Ref: .trellis/tasks/06-11-feat-ai-agent-loop-tool-calling-upgrade PRD Phase 2 - let client = ai::build_ai_http_client(config, 120)?; - let response_text = call_with_tools(&client, config, request, &tools).await?; - - // Emit the full text as a single delta for the UI - if !response_text.text.is_empty() { - on_chunk(AiStreamChunk { - session_id: session_id.to_string(), - delta: response_text.text.clone(), - reasoning_delta: None, - done: false, - }); - } - - on_chunk(AiStreamChunk { - session_id: session_id.to_string(), - delta: String::new(), - reasoning_delta: None, - done: true, - }); - - Ok(response_text.tool_calls) -} - -/// Response from an LLM call that may include tool calls. -struct LlmResponse { - text: String, - tool_calls: Vec, -} - -/// Make a non-streaming LLM call with tools and parse the response. -async fn call_with_tools( - client: &reqwest::Client, - config: &AiConfig, - request: &AiCompletionRequest, - tools: &[ToolDefinition], -) -> Result { - match config.provider { - AiProvider::Claude => call_claude_with_tools(client, config, request, tools).await, - AiProvider::Gemini => call_gemini_with_tools(client, config, request, tools).await, - _ => call_openai_with_tools(client, config, request, tools).await, - } -} - -/// OpenAI / compatible providers: non-streaming call with tools. -async fn call_openai_with_tools( - client: &reqwest::Client, - config: &AiConfig, - request: &AiCompletionRequest, - tools: &[ToolDefinition], -) -> Result { - let headers = ai::maybe_bearer_headers(config)?; - - let mut messages = vec![json!({ "role": "system", "content": request.system_prompt })]; - messages.extend(request.messages.iter().map(|m| { - let mut msg = json!({ "role": m.role, "content": m.content }); - if m.role == "tool" { - if let Some(ref tc_id) = m.tool_call_id { - msg["tool_call_id"] = json!(tc_id); - } - } else if m.role == "assistant" && !m.tool_calls.is_empty() { - let calls: Vec = m - .tool_calls - .iter() - .map(|tc| { - json!({ - "id": tc.id, - "type": "function", - "function": { - "name": tc.name, - "arguments": tc.arguments.to_string() - } - }) - }) - .collect(); - msg["tool_calls"] = json!(calls); - } - msg - })); - - let tool_json: Vec = tools.iter().map(|t| t.to_openai_tool()).collect(); - - let mut body = json!({ - "model": config.model, - "messages": messages, - "max_tokens": request.max_tokens.unwrap_or(4096), - "tools": tool_json, - "tool_choice": "auto", - "stream": false, - }); - ai::add_temperature_if_supported(&mut body, request); - - // Responses API tool calling is deferred to a future phase. - // Use the completions endpoint regardless of api_style for now. - let endpoint = ai::resolve_endpoint(config); - - let res = client - .post(&endpoint) - .headers(headers) - .json(&body) - .send() - .await - .map_err(|e| format!("AI request failed: {e}"))?; - - if !res.status().is_success() { - let data: serde_json::Value = res.json().await.map_err(|e| e.to_string())?; - return Err(ai::extract_error(&data).unwrap_or_else(|| "API error".to_string())); - } - - let data: serde_json::Value = res.json().await.map_err(|e| e.to_string())?; - - let text = data["choices"].get(0).and_then(|c| c["message"]["content"].as_str()).unwrap_or("").to_string(); - - let mut tool_calls = Vec::new(); - if let Some(calls) = data["choices"].get(0).and_then(|c| c["message"]["tool_calls"].as_array()) { - for tc in calls { - let id = tc["id"].as_str().unwrap_or_default().to_string(); - let name = tc["function"]["name"].as_str().unwrap_or_default().to_string(); - let args_str = tc["function"]["arguments"].as_str().unwrap_or("{}"); - let args: serde_json::Value = serde_json::from_str(args_str).unwrap_or(json!({})); - tool_calls.push(ToolCall { id, name, arguments: args }); - } - } - - Ok(LlmResponse { text, tool_calls }) -} - -/// Claude: non-streaming call with tools. -async fn call_claude_with_tools( - client: &reqwest::Client, - config: &AiConfig, - request: &AiCompletionRequest, - tools: &[ToolDefinition], -) -> Result { - let headers = ai::claude_headers(config)?; - - let mut messages: Vec = Vec::new(); - for m in &request.messages { - if m.role == "tool" { - // Convert tool results to Anthropic tool_result content blocks - messages.push(json!({ - "role": "user", - "content": [{ - "type": "tool_result", - "tool_use_id": m.tool_call_id.as_deref().unwrap_or_default(), - "content": m.content - }] - })); - } else if m.role == "assistant" && !m.tool_calls.is_empty() { - // Reconstruct assistant message with tool_use content blocks - let mut content_blocks: Vec = Vec::new(); - if !m.content.is_empty() { - content_blocks.push(json!({ - "type": "text", - "text": m.content - })); - } - for tc in &m.tool_calls { - content_blocks.push(json!({ - "type": "tool_use", - "id": tc.id, - "name": tc.name, - "input": tc.arguments - })); - } - messages.push(json!({ - "role": "assistant", - "content": content_blocks - })); - } else { - messages.push(json!({ "role": m.role, "content": m.content })); - } - } - - let tool_json: Vec = tools.iter().map(|t| t.to_anthropic_tool()).collect(); - - let body = json!({ - "model": config.model, - "max_tokens": request.max_tokens.unwrap_or(4096), - "system": request.system_prompt, - "messages": messages, - "tools": tool_json, - }); - - let res = client - .post(ai::resolve_endpoint(config)) - .headers(headers) - .json(&body) - .send() - .await - .map_err(|e| format!("AI request failed: {e}"))?; - - if !res.status().is_success() { - let data: serde_json::Value = res.json().await.map_err(|e| e.to_string())?; - return Err(ai::extract_error(&data).unwrap_or_else(|| "API error".to_string())); - } - - let data: serde_json::Value = res.json().await.map_err(|e| e.to_string())?; - - let mut text = String::new(); - let mut tool_calls = Vec::new(); - - if let Some(content) = data["content"].as_array() { - for block in content { - match block["type"].as_str() { - Some("text") => { - if let Some(t) = block["text"].as_str() { - text.push_str(t); - } - } - Some("tool_use") => { - let id = block["id"].as_str().unwrap_or_default().to_string(); - let name = block["name"].as_str().unwrap_or_default().to_string(); - let args = block["input"].clone(); - tool_calls.push(ToolCall { id, name, arguments: args }); - } - _ => {} - } - } - } - - Ok(LlmResponse { text, tool_calls }) -} - -/// Gemini: non-streaming call with tools. -async fn call_gemini_with_tools( - client: &reqwest::Client, - config: &AiConfig, - request: &AiCompletionRequest, - tools: &[ToolDefinition], -) -> Result { - let mut contents: Vec = Vec::new(); - for m in &request.messages { - if m.role == "tool" { - // Extract tool name from id "gemini-tc--" - let tool_name = m - .tool_call_id - .as_deref() - .and_then(|s| s.strip_prefix("gemini-tc-")) - .and_then(|s| s.rsplitn(2, '-').nth(1)) - .unwrap_or("unknown"); - contents.push(json!({ - "role": "user", - "parts": [{ - "functionResponse": { - "name": tool_name, - "response": { "content": m.content } - } - }] - })); - } else if m.role == "assistant" && !m.tool_calls.is_empty() { - let mut parts: Vec = Vec::new(); - if !m.content.is_empty() { - parts.push(json!({ "text": m.content })); - } - for tc in &m.tool_calls { - parts.push(json!({ - "functionCall": { - "name": tc.name, - "args": tc.arguments - } - })); - } - contents.push(json!({ - "role": "model", - "parts": parts - })); - } else { - let role = if m.role == "assistant" { "model" } else { "user" }; - contents.push(json!({ - "role": role, - "parts": [{ "text": m.content }] - })); - } - } - - let tool_declarations: Vec = tools.iter().map(|t| t.to_gemini_tool()).collect(); - - let body = json!({ - "contents": contents, - "systemInstruction": { "parts": [{ "text": request.system_prompt }] }, - "tools": [{ "functionDeclarations": tool_declarations }], - "generationConfig": { - "maxOutputTokens": request.max_tokens.unwrap_or(4096), - } - }); - - let res = client - .post(ai::resolve_endpoint(config)) - .query(&[("key", config.api_key.as_str())]) - .header(reqwest::header::CONTENT_TYPE, "application/json") - .json(&body) - .send() - .await - .map_err(|e| format!("AI request failed: {e}"))?; - - if !res.status().is_success() { - let data: serde_json::Value = res.json().await.map_err(|e| e.to_string())?; - return Err(ai::extract_error(&data).unwrap_or_else(|| "API error".to_string())); - } - - let data: serde_json::Value = res.json().await.map_err(|e| e.to_string())?; - - let mut text = String::new(); - let mut tool_calls = Vec::new(); - - if let Some(candidates) = data["candidates"].as_array() { - if let Some(parts) = candidates[0]["content"]["parts"].as_array() { - for part in parts { - if let Some(t) = part["text"].as_str() { - text.push_str(t); - } - if let Some(fc) = part.get("functionCall") { - let name = fc["name"].as_str().unwrap_or_default().to_string(); - let args = fc["args"].clone(); - let id = format!("gemini-tc-{name}-{}", tool_calls.len()); - tool_calls.push(ToolCall { id, name, arguments: args }); - } - } - } - } - - Ok(LlmResponse { text, tool_calls }) + ai::stream_with_tools(config, request, session_id, &tools, cancelled, on_chunk).await } /// Text-only fallback for providers that don't support function calling. diff --git a/crates/dbx-core/src/agent_tools.rs b/crates/dbx-core/src/agent_tools.rs index c23ed8ac84..c8d8f5b4bc 100644 --- a/crates/dbx-core/src/agent_tools.rs +++ b/crates/dbx-core/src/agent_tools.rs @@ -38,7 +38,9 @@ fn get_columns_tool() -> ToolDefinition { ToolDefinition { name: "get_columns", description: - "Get column definitions for a table: names, types, primary keys, nullable, defaults, and comments.", + "Get column definitions for a table: names, types, primary keys, nullable, defaults, and comments. \ + Use this when the user asks about table structure, column details, or field information — \ + even if some schema context was provided, this tool returns the authoritative and complete column list.", parameters: json!({ "type": "object", "properties": { diff --git a/crates/dbx-core/src/ai.rs b/crates/dbx-core/src/ai.rs index ff216a7b39..247b32d67b 100644 --- a/crates/dbx-core/src/ai.rs +++ b/crates/dbx-core/src/ai.rs @@ -1120,6 +1120,540 @@ async fn stream_gemini( Ok(()) } +// --------------------------------------------------------------------------- +// Streaming with tools (agent loop) +// --------------------------------------------------------------------------- + +/// Events emitted by provider-specific streaming-with-tools functions. +/// The public `stream_with_tools` entry point uses these to accumulate +/// tool calls and forward text/reasoning chunks to the caller. +pub enum StreamToolEvent { + /// A text or reasoning delta for the frontend. + Chunk(AiStreamChunk), + /// A tool_use / function_call block has started. + ToolCallStart { index: u32, id: String, name: String }, + /// An argument fragment for an in-progress tool call. + ToolCallDelta { index: u32, fragment: String }, + /// A tool_use / function_call block has ended. + ToolCallComplete { index: u32 }, +} + +/// Partially accumulated tool call during streaming. +#[derive(Debug)] +struct PartialToolCall { + id: String, + name: String, + arguments: String, +} + +/// Accumulates streaming tool-call events into complete `ToolCall` objects. +#[derive(Debug)] +pub struct StreamingToolCallAccumulator { + calls: std::collections::HashMap, + ordered_indices: Vec, +} + +impl StreamingToolCallAccumulator { + pub fn new() -> Self { + Self { calls: std::collections::HashMap::new(), ordered_indices: Vec::new() } + } + + pub fn process(&mut self, event: StreamToolEvent, on_chunk: &impl Fn(AiStreamChunk)) { + match event { + StreamToolEvent::Chunk(chunk) => on_chunk(chunk), + StreamToolEvent::ToolCallStart { index, id, name } => { + self.calls.insert(index, PartialToolCall { id, name, arguments: String::new() }); + if !self.ordered_indices.contains(&index) { + self.ordered_indices.push(index); + } + } + StreamToolEvent::ToolCallDelta { index, fragment } => { + if let Some(tc) = self.calls.get_mut(&index) { + tc.arguments.push_str(&fragment); + } + } + StreamToolEvent::ToolCallComplete { index: _ } => { + // Nothing extra to do — the call is already accumulated. + } + } + } + + pub fn finalize(self) -> Vec { + let mut result = Vec::new(); + for idx in &self.ordered_indices { + if let Some(tc) = self.calls.get(idx) { + let args: serde_json::Value = serde_json::from_str(&tc.arguments).unwrap_or(json!({})); + result.push(crate::agent_events::ToolCall { + id: tc.id.clone(), + name: tc.name.clone(), + arguments: args, + }); + } + } + result + } +} + +/// Streaming Claude call with tool support. +/// Returns a stream of `StreamToolEvent` via the `on_event` callback. +async fn stream_claude_with_tools( + client: &reqwest::Client, + session_id: &str, + request: &AiCompletionRequest, + tools: &[crate::agent_events::ToolDefinition], + cancelled: &Notify, + on_event: &impl Fn(StreamToolEvent), +) -> Result<(), String> { + let mut messages: Vec = Vec::new(); + for m in &request.messages { + if m.role == "tool" { + messages.push(json!({ + "role": "user", + "content": [{ + "type": "tool_result", + "tool_use_id": m.tool_call_id.as_deref().unwrap_or_default(), + "content": m.content + }] + })); + } else if m.role == "assistant" && !m.tool_calls.is_empty() { + let mut content_blocks: Vec = Vec::new(); + if !m.content.is_empty() { + content_blocks.push(json!({ "type": "text", "text": m.content })); + } + for tc in &m.tool_calls { + content_blocks.push(json!({ "type": "tool_use", "id": tc.id, "name": tc.name, "input": tc.arguments })); + } + messages.push(json!({ "role": "assistant", "content": content_blocks })); + } else { + messages.push(json!({ "role": m.role, "content": m.content })); + } + } + + let tool_json: Vec = tools.iter().map(|t| t.to_anthropic_tool()).collect(); + + let mut body = json!({ + "model": request.config.model, + "max_tokens": request.max_tokens.unwrap_or(4096), + "system": request.system_prompt, + "messages": messages, + "tools": tool_json, + "stream": true, + }); + add_temperature_if_supported(&mut body, request); + + let res = client + .post(resolve_endpoint(&request.config)) + .headers(claude_headers(&request.config)?) + .json(&body) + .send() + .await + .map_err(|e| format!("Claude request failed: {e}"))?; + + if !res.status().is_success() { + let data: serde_json::Value = res.json().await.map_err(|e| e.to_string())?; + return Err(extract_error(&data).unwrap_or_else(|| "Claude API error".to_string())); + } + + let mut byte_stream = res.bytes_stream(); + let mut buf = String::new(); + // Track the current content block index and type for tool_use blocks + let mut current_block_index: Option = None; + let mut current_block_type: Option = None; + + loop { + tokio::select! { + chunk = byte_stream.next() => { + let Some(chunk) = chunk else { break }; + let chunk = chunk.map_err(|e| e.to_string())?; + buf.push_str(&String::from_utf8_lossy(&chunk)); + + let mut finished = false; + while let Some(pos) = buf.find('\n') { + let line = buf[..pos].to_string(); + buf = buf[pos + 1..].to_string(); + + let Some(data) = stream_data_payload(&line) else { continue }; + if data == "[DONE]" { + finished = true; + break; + } + + if let Ok(event) = serde_json::from_str::(data) { + let event_type = event["type"].as_str().unwrap_or(""); + + match event_type { + "content_block_start" => { + let idx = event["index"].as_u64().unwrap_or(0) as u32; + let block_type = event["content_block"]["type"].as_str().unwrap_or(""); + current_block_index = Some(idx); + current_block_type = Some(block_type.to_string()); + + if block_type == "tool_use" { + let id = event["content_block"]["id"].as_str().unwrap_or_default().to_string(); + let name = event["content_block"]["name"].as_str().unwrap_or_default().to_string(); + on_event(StreamToolEvent::ToolCallStart { index: idx, id, name }); + } + } + "content_block_delta" => { + let idx = event["index"].as_u64().unwrap_or(0) as u32; + let delta_type = event["delta"]["type"].as_str().unwrap_or(""); + + match delta_type { + "text_delta" => { + if let Some(text) = event["delta"]["text"].as_str() { + on_event(StreamToolEvent::Chunk(AiStreamChunk { + session_id: session_id.to_string(), + delta: text.to_string(), + reasoning_delta: None, + done: false, + })); + } + } + "thinking_delta" => { + if let Some(thinking) = event["delta"]["thinking"].as_str() { + on_event(StreamToolEvent::Chunk(AiStreamChunk { + session_id: session_id.to_string(), + delta: String::new(), + reasoning_delta: Some(thinking.to_string()), + done: false, + })); + } + } + "input_json_delta" => { + if let Some(fragment) = event["delta"]["partial_json"].as_str() { + on_event(StreamToolEvent::ToolCallDelta { + index: idx, + fragment: fragment.to_string(), + }); + } + } + _ => {} + } + } + "content_block_stop" => { + if let Some(idx) = current_block_index.take() { + if current_block_type.as_deref() == Some("tool_use") { + on_event(StreamToolEvent::ToolCallComplete { index: idx }); + } + } + current_block_type = None; + } + _ => {} + } + } + } + + if finished { break; } + } + _ = cancelled.notified() => { break; } + } + } + + Ok(()) +} + +/// Streaming OpenAI-compatible call with tool support. +async fn stream_openai_with_tools( + client: &reqwest::Client, + session_id: &str, + request: &AiCompletionRequest, + tools: &[crate::agent_events::ToolDefinition], + cancelled: &Notify, + on_event: &impl Fn(StreamToolEvent), +) -> Result<(), String> { + let headers = maybe_bearer_headers(&request.config)?; + + let mut messages = vec![json!({ "role": "system", "content": request.system_prompt })]; + messages.extend(request.messages.iter().map(|m| { + let mut msg = json!({ "role": m.role, "content": m.content }); + if m.role == "tool" { + if let Some(ref tc_id) = m.tool_call_id { + msg["tool_call_id"] = json!(tc_id); + } + } else if m.role == "assistant" && !m.tool_calls.is_empty() { + let calls: Vec = m + .tool_calls + .iter() + .map(|tc| { + json!({ + "id": tc.id, + "type": "function", + "function": { + "name": tc.name, + "arguments": tc.arguments.to_string() + } + }) + }) + .collect(); + msg["tool_calls"] = json!(calls); + } + msg + })); + + let tool_json: Vec = tools.iter().map(|t| t.to_openai_tool()).collect(); + + let mut body = json!({ + "model": request.config.model, + "messages": messages, + "max_tokens": request.max_tokens.unwrap_or(4096), + "tools": tool_json, + "tool_choice": "auto", + "stream": true, + }); + add_temperature_if_supported(&mut body, request); + + let res = client + .post(resolve_endpoint(&request.config)) + .headers(headers) + .json(&body) + .send() + .await + .map_err(|e| format!("AI request failed: {e}"))?; + + if !res.status().is_success() { + let data: serde_json::Value = res.json().await.map_err(|e| e.to_string())?; + return Err(extract_error(&data).unwrap_or_else(|| "API error".to_string())); + } + + let mut byte_stream = res.bytes_stream(); + let mut buf = String::new(); + + loop { + tokio::select! { + chunk = byte_stream.next() => { + let Some(chunk) = chunk else { break }; + let chunk = chunk.map_err(|e| e.to_string())?; + buf.push_str(&String::from_utf8_lossy(&chunk)); + + let mut finished = false; + while let Some(pos) = buf.find('\n') { + let line = buf[..pos].to_string(); + buf = buf[pos + 1..].to_string(); + + let Some(data) = stream_data_payload(&line) else { continue }; + if data == "[DONE]" { + finished = true; + break; + } + + if let Ok(event) = serde_json::from_str::(data) { + // Reasoning + if let Some(reasoning) = openai_stream_reasoning(&event) { + on_event(StreamToolEvent::Chunk(AiStreamChunk { + session_id: session_id.to_string(), + delta: String::new(), + reasoning_delta: Some(reasoning.to_string()), + done: false, + })); + } + // Text + if let Some(text) = openai_stream_text(&event) { + on_event(StreamToolEvent::Chunk(AiStreamChunk { + session_id: session_id.to_string(), + delta: text, + reasoning_delta: None, + done: false, + })); + } + // Tool calls + if let Some(tool_calls) = event["choices"].get(0).and_then(|c| c["delta"]["tool_calls"].as_array()) { + for tc in tool_calls { + let idx = tc["index"].as_u64().unwrap_or(0) as u32; + // First chunk for this tool call has id and name + if let Some(id) = tc["id"].as_str() { + let name = tc["function"]["name"].as_str().unwrap_or_default().to_string(); + on_event(StreamToolEvent::ToolCallStart { index: idx, id: id.to_string(), name }); + } + // Argument fragments + if let Some(fragment) = tc["function"]["arguments"].as_str() { + on_event(StreamToolEvent::ToolCallDelta { index: idx, fragment: fragment.to_string() }); + } + } + } + } + } + + if finished { break; } + } + _ = cancelled.notified() => { break; } + } + } + + Ok(()) +} + +/// Streaming Gemini call with tool support. +async fn stream_gemini_with_tools( + client: &reqwest::Client, + session_id: &str, + request: &AiCompletionRequest, + tools: &[crate::agent_events::ToolDefinition], + cancelled: &Notify, + on_event: &impl Fn(StreamToolEvent), +) -> Result<(), String> { + let mut contents: Vec = Vec::new(); + for m in &request.messages { + if m.role == "tool" { + let tool_name = m + .tool_call_id + .as_deref() + .and_then(|s| s.strip_prefix("gemini-tc-")) + .and_then(|s| s.rsplitn(2, '-').nth(1)) + .unwrap_or("unknown"); + contents.push(json!({ + "role": "user", + "parts": [{ + "functionResponse": { + "name": tool_name, + "response": { "content": m.content } + } + }] + })); + } else if m.role == "assistant" && !m.tool_calls.is_empty() { + let mut parts: Vec = Vec::new(); + if !m.content.is_empty() { + parts.push(json!({ "text": m.content })); + } + for tc in &m.tool_calls { + parts.push(json!({ "functionCall": { "name": tc.name, "args": tc.arguments } })); + } + contents.push(json!({ "role": "model", "parts": parts })); + } else { + let role = if m.role == "assistant" { "model" } else { "user" }; + contents.push(json!({ "role": role, "parts": [{ "text": m.content }] })); + } + } + + let tool_declarations: Vec = tools.iter().map(|t| t.to_gemini_tool()).collect(); + + let body = json!({ + "contents": contents, + "systemInstruction": { "parts": [{ "text": request.system_prompt }] }, + "tools": [{ "functionDeclarations": tool_declarations }], + "generationConfig": { + "maxOutputTokens": request.max_tokens.unwrap_or(4096), + } + }); + + let res = client + .post(resolve_gemini_stream_endpoint(&request.config)) + .query(&[("key", request.config.api_key.as_str()), ("alt", "sse")]) + .header(CONTENT_TYPE, "application/json") + .json(&body) + .send() + .await + .map_err(|e| format!("Gemini request failed: {e}"))?; + + if !res.status().is_success() { + let data: serde_json::Value = res.json().await.map_err(|e| e.to_string())?; + return Err(extract_error(&data).unwrap_or_else(|| "Gemini API error".to_string())); + } + + let mut byte_stream = res.bytes_stream(); + let mut buf = String::new(); + let mut tool_call_idx: u32 = 0; + + loop { + tokio::select! { + chunk = byte_stream.next() => { + let Some(chunk) = chunk else { break }; + let chunk = chunk.map_err(|e| e.to_string())?; + buf.push_str(&String::from_utf8_lossy(&chunk)); + + while let Some(pos) = buf.find('\n') { + let line = buf[..pos].to_string(); + buf = buf[pos + 1..].to_string(); + + let Some(data) = stream_data_payload(&line) else { continue }; + if let Ok(event) = serde_json::from_str::(data) { + if let Some(candidates) = event["candidates"].as_array() { + if let Some(parts) = candidates[0]["content"]["parts"].as_array() { + for part in parts { + // Text + if let Some(text) = part["text"].as_str() { + on_event(StreamToolEvent::Chunk(AiStreamChunk { + session_id: session_id.to_string(), + delta: text.to_string(), + reasoning_delta: None, + done: false, + })); + } + // Function call (Gemini sends complete objects, not deltas) + if let Some(fc) = part.get("functionCall") { + let name = fc["name"].as_str().unwrap_or_default().to_string(); + let args = fc["args"].clone(); + let id = format!("gemini-tc-{name}-{tool_call_idx}"); + let args_str = args.to_string(); + on_event(StreamToolEvent::ToolCallStart { + index: tool_call_idx, + id: id.clone(), + name: name.clone(), + }); + on_event(StreamToolEvent::ToolCallDelta { + index: tool_call_idx, + fragment: args_str, + }); + on_event(StreamToolEvent::ToolCallComplete { index: tool_call_idx }); + tool_call_idx += 1; + } + } + } + } + } + } + } + _ = cancelled.notified() => { break; } + } + } + + Ok(()) +} + +/// Public entry point: stream an LLM call with tool support, accumulating tool calls. +/// Returns completed tool calls when the stream finishes. +pub async fn stream_with_tools( + config: &AiConfig, + request: &AiCompletionRequest, + session_id: &str, + tools: &[crate::agent_events::ToolDefinition], + cancelled: &Notify, + on_chunk: impl Fn(AiStreamChunk), +) -> Result, String> { + validate_config(config)?; + + let stream_timeout = if config.enable_thinking { 600 } else { 120 }; + let client = build_ai_http_client(config, stream_timeout)?; + + let accumulator = Arc::new(std::sync::Mutex::new(StreamingToolCallAccumulator::new())); + + match config.provider { + AiProvider::Claude => { + stream_claude_with_tools(&client, session_id, request, tools, cancelled, &|event| { + accumulator.lock().unwrap().process(event, &on_chunk); + }) + .await? + } + AiProvider::Gemini => { + stream_gemini_with_tools(&client, session_id, request, tools, cancelled, &|event| { + accumulator.lock().unwrap().process(event, &on_chunk); + }) + .await? + } + _ => { + stream_openai_with_tools(&client, session_id, request, tools, cancelled, &|event| { + accumulator.lock().unwrap().process(event, &on_chunk); + }) + .await? + } + } + + Ok(Arc::try_unwrap(accumulator) + .expect("stream_with_tools: accumulator Arc should have single owner") + .into_inner() + .expect("stream_with_tools: accumulator Mutex should not be poisoned") + .finalize()) +} + // --------------------------------------------------------------------------- // Conversation persistence (path-based) // ---------------------------------------------------------------------------