diff --git a/crates/dbx-core/src/ai.rs b/crates/dbx-core/src/ai.rs index acf0721f68..521c7a6ade 100644 --- a/crates/dbx-core/src/ai.rs +++ b/crates/dbx-core/src/ai.rs @@ -447,6 +447,11 @@ pub fn openai_stream_reasoning(event: &serde_json::Value) -> Option<&str> { } pub fn responses_stream_text(event: &serde_json::Value) -> Option<&str> { + let event_type = event["type"].as_str().unwrap_or_default(); + if !event_type.is_empty() && event_type != "response.output_text.delta" { + return None; + } + event["delta"].as_str().filter(|s| !s.is_empty()) } @@ -454,6 +459,13 @@ fn responses_max_output_tokens(max_tokens: Option) -> u32 { max_tokens.unwrap_or(2048).max(16) } +fn responses_token_usage(event: &serde_json::Value) -> Option { + let usage = event.get("usage").or_else(|| event.get("response").and_then(|response| response.get("usage")))?; + let input = usage.get("input_tokens").and_then(|v| v.as_u64())?; + let output = usage.get("output_tokens").and_then(|v| v.as_u64())?; + Some(TokenUsage { input_tokens: input as u32, output_tokens: output as u32 }) +} + fn is_openai_api_config(config: &AiConfig) -> bool { matches!(config.provider, AiProvider::Openai) || config.endpoint.to_ascii_lowercase().contains("api.openai.com") } @@ -548,6 +560,111 @@ pub fn build_responses_input(system_prompt: &str, messages: &[AiMessage]) -> ser json!(input) } +fn build_responses_input_with_tools(system_prompt: &str, messages: &[AiMessage]) -> serde_json::Value { + let mut input = Vec::new(); + if !system_prompt.is_empty() { + input.push(json!({ + "role": "developer", + "content": system_prompt, + })); + } + + for message in messages { + if message.role == "tool" { + input.push(json!({ + "type": "function_call_output", + "call_id": message.tool_call_id.as_deref().unwrap_or_default(), + "output": message.content, + })); + continue; + } + + if message.role == "assistant" && !message.tool_calls.is_empty() { + if !message.content.is_empty() { + input.push(json!({ + "role": "assistant", + "content": message.content, + })); + } + for tool_call in &message.tool_calls { + input.push(json!({ + "type": "function_call", + "call_id": tool_call.id, + "name": tool_call.name, + "arguments": tool_call.arguments.to_string(), + })); + } + continue; + } + + input.push(json!({ + "role": message.role, + "content": message.content, + })); + } + + json!(input) +} + +fn responses_function_tool(tool: &crate::agent_events::ToolDefinition) -> serde_json::Value { + json!({ + "type": "function", + "name": tool.name, + "description": tool.description, + "parameters": tool.parameters, + }) +} + +fn responses_tool_index( + event: &serde_json::Value, + item_indices: &mut HashMap, + next_index: &mut u32, +) -> (String, u32) { + let item = &event["item"]; + let item_id = item["id"] + .as_str() + .or_else(|| event["item_id"].as_str()) + .map(ToString::to_string) + .unwrap_or_else(|| format!("responses-tool-{next_index}")); + let index = item_indices.get(&item_id).copied().unwrap_or_else(|| { + let index = event["output_index"].as_u64().map(|i| i as u32).unwrap_or(*next_index); + *next_index = (*next_index).max(index + 1); + item_indices.insert(item_id.clone(), index); + index + }); + (item_id, index) +} + +fn emit_responses_function_call_item( + event: &serde_json::Value, + item_indices: &mut HashMap, + started_indices: &mut HashSet, + argument_indices: &mut HashSet, + next_index: &mut u32, + on_event: &impl Fn(StreamToolEvent), +) -> Option { + let item = &event["item"]; + if item["type"].as_str() != Some("function_call") { + return None; + } + + let (_item_id, index) = responses_tool_index(event, item_indices, next_index); + if started_indices.insert(index) { + let id = item["call_id"].as_str().or_else(|| item["id"].as_str()).unwrap_or_default().to_string(); + let name = item["name"].as_str().unwrap_or_default().to_string(); + on_event(StreamToolEvent::ToolCallStart { index, id, name }); + } + + if !argument_indices.contains(&index) { + if let Some(arguments) = item["arguments"].as_str().filter(|s| !s.is_empty()) { + argument_indices.insert(index); + on_event(StreamToolEvent::ToolCallDelta { index, fragment: arguments.to_string() }); + } + } + + Some(index) +} + // --------------------------------------------------------------------------- // Validation helper // --------------------------------------------------------------------------- @@ -1024,15 +1141,24 @@ pub async fn test_connection_core(config: &AiConfig) -> Result { // OpenAI-compatible providers - let messages = vec![json!({ "role": "user", "content": TEST_PROMPT })]; - let mut body_obj = json!({ - "model": &model, - "messages": messages, - "max_tokens": 16, - "stream": true, - }); + let mut body_obj = if config.api_style == AiApiStyle::Responses { + json!({ + "model": &model, + "input": [{ "role": "user", "content": TEST_PROMPT }], + "max_output_tokens": 16, + "stream": true, + }) + } else { + let messages = vec![json!({ "role": "user", "content": TEST_PROMPT })]; + json!({ + "model": &model, + "messages": messages, + "max_tokens": 16, + "stream": true, + }) + }; add_temperature_if_supported_for_config(&mut body_obj, config, Some(0.0)); - if !config.enable_thinking && !is_kimi_model(&config.model) { + if config.api_style != AiApiStyle::Responses && !config.enable_thinking && !is_kimi_model(&config.model) { body_obj["extra_body"] = json!({ "chat_template_kwargs": { "enable_thinking": false } }); @@ -1922,6 +2048,133 @@ async fn stream_openai_with_tools( Ok(token_usage) } +async fn stream_responses_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 tool_json: Vec = tools.iter().map(responses_function_tool).collect(); + + let mut body = json!({ + "model": request.config.model, + "input": build_responses_input_with_tools(&request.system_prompt, &request.messages), + "max_output_tokens": responses_max_output_tokens(request.max_tokens), + "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 = Vec::new(); + let mut item_indices: HashMap = HashMap::new(); + let mut started_indices: HashSet = HashSet::new(); + let mut argument_indices: HashSet = HashSet::new(); + let mut next_index: u32 = 0; + let mut token_usage: 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.extend_from_slice(&chunk); + + let mut finished = false; + while let Some(line) = drain_next_stream_line(&mut buf)? { + let Some(data) = stream_data_payload(&line) else { continue }; + if data == "[DONE]" { + finished = true; + break; + } + + if let Ok(event) = serde_json::from_str::(data) { + if let Some(usage) = responses_token_usage(&event) { + token_usage = Some(usage); + } + + if let Some(text) = responses_stream_text(&event) { + on_event(StreamToolEvent::Chunk(AiStreamChunk { + session_id: session_id.to_string(), + delta: text.to_string(), + reasoning_delta: None, + done: false, + })); + } + + match event["type"].as_str().unwrap_or_default() { + "response.output_item.added" => { + emit_responses_function_call_item( + &event, + &mut item_indices, + &mut started_indices, + &mut argument_indices, + &mut next_index, + on_event, + ); + } + "response.output_item.done" => { + if let Some(index) = emit_responses_function_call_item( + &event, + &mut item_indices, + &mut started_indices, + &mut argument_indices, + &mut next_index, + on_event, + ) { + on_event(StreamToolEvent::ToolCallComplete { index }); + } + } + "response.function_call_arguments.delta" => { + let index = event["item_id"] + .as_str() + .and_then(|id| item_indices.get(id).copied()) + .or_else(|| event["output_index"].as_u64().map(|i| i as u32)) + .unwrap_or(0); + if let Some(fragment) = event["delta"].as_str() { + argument_indices.insert(index); + on_event(StreamToolEvent::ToolCallDelta { index, fragment: fragment.to_string() }); + } + } + "response.function_call_arguments.done" => { + let index = event["item_id"] + .as_str() + .and_then(|id| item_indices.get(id).copied()) + .or_else(|| event["output_index"].as_u64().map(|i| i as u32)) + .unwrap_or(0); + on_event(StreamToolEvent::ToolCallComplete { index }); + } + _ => {} + } + } + } + + if finished { break; } + } + _ = cancelled.notified() => { break; } + } + } + + Ok(token_usage) +} + /// Streaming Gemini call with tool support. async fn stream_gemini_with_tools( client: &reqwest::Client, @@ -2104,6 +2357,12 @@ pub async fn stream_with_tools( }) .await? } + _ if config.api_style == AiApiStyle::Responses => { + stream_responses_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); @@ -2175,13 +2434,18 @@ pub fn load_config(path: &Path) -> Result, String> { #[cfg(test)] mod tests { + use std::cell::RefCell; + use std::collections::{HashMap, HashSet}; + use super::{ - add_temperature_if_supported_for_config, build_ai_http_client, claude_headers, claude_system_prompt, - drain_next_stream_line, gemini_text, is_kimi_model, openai_response_text, openai_stream_reasoning, - openai_stream_text, parse_model_list_response, resolve_endpoint, resolve_model_list_endpoint, - responses_max_output_tokens, responses_text, stream_data_payload, supports_temperature, temperature_value, - uses_anthropic_messages_api, validate_config, AiApiStyle, AiAuthMethod, AiConfig, AiModelInfo, AiProvider, - AiReasoningLevel, AUTHORIZATION, CLAUDE_DEFAULT_SYSTEM, TEST_PROMPT, + add_temperature_if_supported_for_config, build_ai_http_client, build_responses_input_with_tools, + claude_headers, claude_system_prompt, drain_next_stream_line, emit_responses_function_call_item, gemini_text, + is_kimi_model, openai_response_text, openai_stream_reasoning, openai_stream_text, parse_model_list_response, + resolve_endpoint, resolve_model_list_endpoint, responses_function_tool, responses_max_output_tokens, + responses_stream_text, responses_text, responses_token_usage, stream_data_payload, supports_temperature, + temperature_value, uses_anthropic_messages_api, validate_config, AiApiStyle, AiAuthMethod, AiConfig, AiMessage, + AiModelInfo, AiProvider, AiReasoningLevel, StreamToolEvent, StreamingToolCallAccumulator, ToolCallRef, + AUTHORIZATION, CLAUDE_DEFAULT_SYSTEM, TEST_PROMPT, }; #[test] @@ -2540,6 +2804,208 @@ mod tests { assert_eq!(responses_max_output_tokens(None), 2048); } + #[test] + fn responses_stream_text_reads_current_delta_shapes() { + assert_eq!( + responses_stream_text(&serde_json::json!({ + "type": "response.output_text.delta", + "delta": "SELECT" + })), + Some("SELECT") + ); + assert_eq!( + responses_stream_text(&serde_json::json!({ + "type": "response.output_text.done", + "text": "SELECT 1;" + })), + None + ); + } + + #[test] + fn responses_token_usage_reads_stream_completed_response_usage() { + let completed_usage = responses_token_usage(&serde_json::json!({ + "type": "response.completed", + "response": { + "usage": { + "input_tokens": 12, + "output_tokens": 34 + } + } + })) + .unwrap(); + assert_eq!(completed_usage.input_tokens, 12); + assert_eq!(completed_usage.output_tokens, 34); + + let top_level_usage = responses_token_usage(&serde_json::json!({ + "usage": { + "input_tokens": 56, + "output_tokens": 78 + } + })) + .unwrap(); + assert_eq!(top_level_usage.input_tokens, 56); + assert_eq!(top_level_usage.output_tokens, 78); + } + + #[test] + fn responses_tools_use_responses_schema() { + let input = build_responses_input_with_tools( + "system", + &[ + AiMessage { + role: "user".to_string(), + content: "inspect db".to_string(), + tool_call_id: None, + tool_calls: Vec::new(), + }, + AiMessage { + role: "assistant".to_string(), + content: String::new(), + tool_call_id: None, + tool_calls: vec![ToolCallRef { + id: "call_1".to_string(), + name: "list_tables".to_string(), + arguments: serde_json::json!({"schema": "public"}), + }], + }, + AiMessage { + role: "tool".to_string(), + content: "users".to_string(), + tool_call_id: Some("call_1".to_string()), + tool_calls: Vec::new(), + }, + ], + ); + + assert_eq!(input[0]["role"], "developer"); + assert_eq!(input[2]["type"], "function_call"); + assert_eq!(input[2]["call_id"], "call_1"); + assert_eq!(input[2]["arguments"], "{\"schema\":\"public\"}"); + assert_eq!(input[3]["type"], "function_call_output"); + assert_eq!(input[3]["call_id"], "call_1"); + + let tool = crate::agent_events::ToolDefinition { + name: "list_tables", + description: "List tables", + parameters: serde_json::json!({"type": "object"}), + read_only: true, + parallel_ok: true, + }; + let tool_json = responses_function_tool(&tool); + assert_eq!(tool_json["type"], "function"); + assert_eq!(tool_json["name"], "list_tables"); + assert!(tool_json.get("function").is_none()); + } + + #[test] + fn responses_tool_done_item_can_supply_complete_function_call() { + let mut accumulator = StreamingToolCallAccumulator::new(); + let mut item_indices = HashMap::new(); + let mut started_indices = HashSet::new(); + let mut argument_indices = HashSet::new(); + let mut next_index = 0; + let event = serde_json::json!({ + "type": "response.output_item.done", + "output_index": 0, + "item": { + "id": "fc_1", + "type": "function_call", + "call_id": "call_1", + "name": "list_tables", + "arguments": "{\"schema\":\"public\"}" + } + }); + let events = RefCell::new(Vec::new()); + if let Some(index) = emit_responses_function_call_item( + &event, + &mut item_indices, + &mut started_indices, + &mut argument_indices, + &mut next_index, + &|event| events.borrow_mut().push(event), + ) { + events.borrow_mut().push(StreamToolEvent::ToolCallComplete { index }); + } + for event in events.into_inner() { + accumulator.process(event, &|_| {}); + } + + let calls = accumulator.finalize(); + assert_eq!(calls.len(), 1); + assert_eq!(calls[0].id, "call_1"); + assert_eq!(calls[0].name, "list_tables"); + assert_eq!(calls[0].arguments["schema"], "public"); + } + + #[test] + fn responses_tool_arguments_are_not_duplicated_when_done_follows_delta() { + let mut accumulator = StreamingToolCallAccumulator::new(); + let mut item_indices = HashMap::new(); + let mut started_indices = HashSet::new(); + let mut argument_indices = HashSet::new(); + let mut next_index = 0; + let added = serde_json::json!({ + "type": "response.output_item.added", + "output_index": 0, + "item": { + "id": "fc_1", + "type": "function_call", + "call_id": "call_1", + "name": "list_tables" + } + }); + let done = serde_json::json!({ + "type": "response.output_item.done", + "output_index": 0, + "item": { + "id": "fc_1", + "type": "function_call", + "call_id": "call_1", + "name": "list_tables", + "arguments": "{\"schema\":\"public\"}" + } + }); + + let events = RefCell::new(Vec::new()); + emit_responses_function_call_item( + &added, + &mut item_indices, + &mut started_indices, + &mut argument_indices, + &mut next_index, + &|event| events.borrow_mut().push(event), + ); + for event in events.take() { + accumulator.process(event, &|_| {}); + } + + let delta_index = item_indices.get("fc_1").copied().unwrap(); + argument_indices.insert(delta_index); + accumulator.process( + StreamToolEvent::ToolCallDelta { index: delta_index, fragment: "{\"schema\":\"public\"}".to_string() }, + &|_| {}, + ); + + if let Some(index) = emit_responses_function_call_item( + &done, + &mut item_indices, + &mut started_indices, + &mut argument_indices, + &mut next_index, + &|event| events.borrow_mut().push(event), + ) { + events.borrow_mut().push(StreamToolEvent::ToolCallComplete { index }); + } + for event in events.into_inner() { + accumulator.process(event, &|_| {}); + } + + let calls = accumulator.finalize(); + assert_eq!(calls.len(), 1); + assert_eq!(calls[0].arguments, serde_json::json!({"schema": "public"})); + } + #[test] fn temperature_value_rounds_f32_to_provider_safe_precision() { assert_eq!(temperature_value(Some(0.15)), serde_json::json!(0.15));