diff --git a/src/kagi.rs b/src/kagi.rs deleted file mode 100644 index 72878e89..00000000 --- a/src/kagi.rs +++ /dev/null @@ -1,218 +0,0 @@ -//! Minimal Kagi Search API client -//! -//! This module provides a lightweight client for the Kagi Search API. -//! Only includes what we actually use - no bloat from auto-generated code. - -use serde::{Deserialize, Serialize}; -use std::collections::HashMap; -use std::sync::Arc; -use std::time::Duration; - -const KAGI_API_BASE: &str = "https://kagi.com/api/v1"; -const REQUEST_TIMEOUT: Duration = Duration::from_secs(10); -const CONNECT_TIMEOUT: Duration = Duration::from_secs(5); - -#[derive(Debug, thiserror::Error)] -pub enum KagiError { - #[error("HTTP request failed: {0}")] - Request(#[from] reqwest::Error), - #[error("API error: {status} - {message}")] - Api { status: u16, message: String }, -} - -/// Kagi API client with reusable HTTP client and stored API key -#[derive(Clone)] -pub struct KagiClient { - client: reqwest::Client, - api_key: Arc, -} - -impl KagiClient { - /// Create a new Kagi client with the given API key - /// The API key will automatically be prefixed with "Bot " for authorization - pub fn new(api_key: String) -> Result { - let client = reqwest::Client::builder() - .timeout(REQUEST_TIMEOUT) - .connect_timeout(CONNECT_TIMEOUT) - .pool_max_idle_per_host(100) - .user_agent("OpenAPI-Generator/0.1.0/rust") - .build() - .map_err(KagiError::Request)?; - - // Automatically prefix with "Bot " if not already present - let formatted_key = if api_key.starts_with("Bot ") { - api_key - } else { - format!("Bot {}", api_key) - }; - - Ok(Self { - client, - api_key: Arc::new(formatted_key), - }) - } - - /// Execute a search query - pub async fn search(&self, request: SearchRequest) -> Result { - let url = format!("{}/search", KAGI_API_BASE); - - let response = self - .client - .post(&url) - .header("Authorization", self.api_key.as_str()) - .json(&request) - .send() - .await?; - - let status = response.status(); - - if !status.is_success() { - let error_text = response.text().await.unwrap_or_default(); - return Err(KagiError::Api { - status: status.as_u16(), - message: error_text, - }); - } - - let search_response = response.json::().await?; - Ok(search_response) - } -} - -impl std::fmt::Debug for KagiClient { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("KagiClient") - .field("api_key", &"[REDACTED]") - .finish() - } -} - -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct SearchRequest { - pub query: String, - #[serde(skip_serializing_if = "Option::is_none")] - pub workflow: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub lens_id: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub lens: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub timeout: Option, -} - -impl SearchRequest { - pub fn new(query: String) -> Self { - Self { - query, - workflow: None, - lens_id: None, - lens: None, - timeout: None, - } - } -} - -#[derive(Debug, Clone, Serialize, Deserialize)] -#[serde(rename_all = "lowercase")] -pub enum Workflow { - Search, - Images, - Videos, - News, - Podcasts, -} - -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct LensConfig { - #[serde(skip_serializing_if = "Option::is_none")] - pub include: Option>, - #[serde(skip_serializing_if = "Option::is_none")] - pub exclude: Option>, -} - -#[derive(Debug, Clone, Deserialize)] -#[allow(dead_code)] -pub struct SearchResponse { - #[serde(skip_serializing_if = "Option::is_none")] - pub meta: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub data: Option, -} - -#[derive(Debug, Clone, Deserialize)] -#[allow(dead_code)] -pub struct Meta { - #[serde(skip_serializing_if = "Option::is_none")] - pub id: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub node: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub ms: Option, -} - -#[derive(Debug, Clone, Deserialize)] -#[allow(dead_code)] -pub struct SearchData { - #[serde(skip_serializing_if = "Option::is_none")] - pub search: Option>, - #[serde(skip_serializing_if = "Option::is_none")] - pub image: Option>, - #[serde(skip_serializing_if = "Option::is_none")] - pub video: Option>, - #[serde(skip_serializing_if = "Option::is_none")] - pub podcast: Option>, - #[serde(skip_serializing_if = "Option::is_none")] - pub podcast_creator: Option>, - #[serde(skip_serializing_if = "Option::is_none")] - pub news: Option>, - #[serde(skip_serializing_if = "Option::is_none")] - pub adjacent_question: Option>, - #[serde(skip_serializing_if = "Option::is_none")] - pub direct_answer: Option>, - #[serde(skip_serializing_if = "Option::is_none")] - pub interesting_news: Option>, - #[serde(skip_serializing_if = "Option::is_none")] - pub interesting_finds: Option>, - #[serde(skip_serializing_if = "Option::is_none")] - pub infobox: Option>, - #[serde(skip_serializing_if = "Option::is_none")] - pub code: Option>, - #[serde(skip_serializing_if = "Option::is_none")] - pub package_tracking: Option>, - #[serde(skip_serializing_if = "Option::is_none")] - pub public_records: Option>, - #[serde(skip_serializing_if = "Option::is_none")] - pub weather: Option>, - #[serde(skip_serializing_if = "Option::is_none")] - pub related_search: Option>, - #[serde(skip_serializing_if = "Option::is_none")] - pub listicle: Option>, - #[serde(skip_serializing_if = "Option::is_none")] - pub web_archive: Option>, -} - -#[derive(Debug, Clone, Deserialize)] -#[allow(dead_code)] -pub struct SearchResult { - pub url: String, - pub title: String, - #[serde(skip_serializing_if = "Option::is_none")] - pub snippet: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub time: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub image: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub props: Option>, -} - -#[derive(Debug, Clone, Deserialize)] -#[allow(dead_code)] -pub struct SearchResultImage { - #[serde(skip_serializing_if = "Option::is_none")] - pub url: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub height: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub width: Option, -} diff --git a/src/main.rs b/src/main.rs index 8f86e63c..d85c901a 100644 --- a/src/main.rs +++ b/src/main.rs @@ -77,7 +77,6 @@ mod db; mod email; mod encrypt; mod jwt; -mod kagi; mod kv; mod message_signing; mod migrations; @@ -87,6 +86,7 @@ mod os_flags; mod private_key; mod proxy_config; mod sqs; +mod tinfoil_websearch; mod tokens; mod web; @@ -107,7 +107,6 @@ const RESEND_API_KEY_NAME: &str = "resend_api_key"; const BILLING_API_KEY_NAME: &str = "billing_api_key"; const BILLING_SERVER_URL_NAME: &str = "billing_server_url"; -const KAGI_API_KEY_NAME: &str = "kagi_api_key"; const BRAVE_API_KEY_NAME: &str = "brave_api_key"; const OS_FLAGS_API_KEY_NAME: &str = "os_flags_api_key"; const OS_FLAGS_BASE_URL_NAME: &str = "os_flags_base_url"; @@ -418,8 +417,8 @@ pub struct AppState { os_flags_client: Option, apple_jwt_verifier: Arc, cancellation_broadcast: tokio::sync::broadcast::Sender, - kagi_client: Option>, brave_client: Option>, + tinfoil_web_search_client: Option>, } #[derive(Default)] @@ -445,8 +444,10 @@ pub struct AppStateBuilder { billing_server_url: Option, os_flags_base_url: Option, os_flags_api_key: Option, - kagi_api_key: Option, brave_api_key: Option, + tinfoil_web_search_api_base: Option, + tinfoil_web_search_api_key: Option, + tinfoil_web_search_allow_insecure_tls: bool, } impl AppStateBuilder { @@ -561,13 +562,32 @@ impl AppStateBuilder { self } - pub fn kagi_api_key(mut self, kagi_api_key: Option) -> Self { - self.kagi_api_key = kagi_api_key; + pub fn brave_api_key(mut self, brave_api_key: Option) -> Self { + self.brave_api_key = brave_api_key; self } - pub fn brave_api_key(mut self, brave_api_key: Option) -> Self { - self.brave_api_key = brave_api_key; + pub fn tinfoil_web_search_api_base( + mut self, + tinfoil_web_search_api_base: Option, + ) -> Self { + self.tinfoil_web_search_api_base = tinfoil_web_search_api_base; + self + } + + pub fn tinfoil_web_search_api_key( + mut self, + tinfoil_web_search_api_key: Option, + ) -> Self { + self.tinfoil_web_search_api_key = tinfoil_web_search_api_key; + self + } + + pub fn tinfoil_web_search_allow_insecure_tls( + mut self, + tinfoil_web_search_allow_insecure_tls: bool, + ) -> Self { + self.tinfoil_web_search_allow_insecure_tls = tinfoil_web_search_allow_insecure_tls; self } @@ -679,27 +699,6 @@ impl AppStateBuilder { let (cancellation_tx, _) = tokio::sync::broadcast::channel(1024); - // Initialize Kagi client if API key is provided - let kagi_client = if let Some(ref api_key) = self.kagi_api_key { - tracing::info!("Initializing Kagi client with connection pooling (max 100 idle connections, 10s timeout)"); - match crate::kagi::KagiClient::new(api_key.clone()) { - Ok(client) => { - tracing::debug!("Kagi client initialized successfully"); - Some(Arc::new(client)) - } - Err(e) => { - tracing::error!( - "Failed to initialize Kagi client: {:?}. Web search will be unavailable.", - e - ); - None - } - } - } else { - tracing::debug!("Kagi API key not configured, web search tool will be unavailable"); - None - }; - let brave_client = if let Some(ref api_key) = self.brave_api_key { tracing::info!("Initializing Brave client with connection pooling (max 100 idle connections, 10s timeout)"); match crate::brave::BraveClient::new(api_key.clone()) { @@ -720,6 +719,42 @@ impl AppStateBuilder { None }; + let tinfoil_web_search_client = match ( + self.tinfoil_web_search_api_base.as_ref(), + self.tinfoil_web_search_api_key.as_ref(), + ) { + (Some(base_url), Some(api_key)) => { + tracing::info!("Initializing Tinfoil web search client"); + match crate::tinfoil_websearch::TinfoilWebSearchClient::new( + base_url.clone(), + api_key.clone(), + self.tinfoil_web_search_allow_insecure_tls, + ) { + Ok(client) => { + tracing::debug!("Tinfoil web search client initialized successfully"); + Some(Arc::new(client)) + } + Err(e) => { + tracing::error!( + "Failed to initialize Tinfoil web search client: {:?}. Tinfoil-backed web search will be unavailable.", + e + ); + None + } + } + } + (Some(_), None) | (None, Some(_)) => { + tracing::warn!( + "Incomplete Tinfoil web search configuration; both base URL and API key are required" + ); + None + } + (None, None) => { + tracing::debug!("Tinfoil web search not configured"); + None + } + }; + Ok(AppState { app_mode, db, @@ -736,8 +771,8 @@ impl AppStateBuilder { os_flags_client, apple_jwt_verifier, cancellation_broadcast: cancellation_tx, - kagi_client, brave_client, + tinfoil_web_search_client, }) } } @@ -747,6 +782,10 @@ impl AppState { self.os_flags_client.as_ref() } + pub fn has_web_search_client(&self) -> bool { + self.tinfoil_web_search_client.is_some() || self.brave_client.is_some() + } + pub async fn is_feature_enabled(&self, user_uuid: Uuid, flag_key: &str) -> bool { let Some(client) = &self.os_flags_client else { return false; @@ -2335,48 +2374,6 @@ async fn retrieve_os_flags_base_url( } } -async fn retrieve_kagi_api_key( - aws_credential_manager: Arc>>, - db: Arc, -) -> Result, Error> { - let creds = aws_credential_manager - .read() - .await - .clone() - .expect("non-local mode should have creds") - .get_credentials() - .await - .expect("non-local mode should have creds"); - - // check if the key already exists in the db - let existing_key = db.get_enclave_secret_by_key(KAGI_API_KEY_NAME)?; - - if let Some(ref encrypted_key) = existing_key { - // Convert the stored bytes back to base64 - let base64_encrypted_key = general_purpose::STANDARD.encode(&encrypted_key.value); - - debug!("trying to decrypt base64 encrypted Kagi API key"); - - // Decrypt the existing key - let decrypted_bytes = decrypt_with_kms( - &creds.region, - &creds.access_key_id, - &creds.secret_access_key, - &creds.token, - &base64_encrypted_key, - ) - .map_err(|e| Error::EncryptionError(e.to_string()))?; - - // Convert the decrypted bytes to a UTF-8 string - String::from_utf8(decrypted_bytes) - .map_err(|e| Error::EncryptionError(format!("Failed to decode UTF-8: {}", e))) - .map(Some) - } else { - tracing::info!("Kagi API key not found in the database"); - Ok(None) - } -} - async fn retrieve_brave_api_key( aws_credential_manager: Arc>>, db: Arc, @@ -2592,6 +2589,13 @@ async fn main() -> Result<(), Error> { // Tinfoil API base is always from environment (like OpenAI API base) let tinfoil_api_base = env::var("TINFOIL_API_BASE").ok(); + let tinfoil_web_search_api_base = env::var("TINFOIL_WEB_SEARCH_API_BASE").ok(); + let tinfoil_web_search_api_key = env::var("TINFOIL_WEB_SEARCH_API_KEY") + .ok() + .or_else(|| env::var("TINFOIL_API_KEY").ok()); + let tinfoil_web_search_allow_insecure_tls = env::var("TINFOIL_WEB_SEARCH_ALLOW_INSECURE_TLS") + .map(|value| matches!(value.to_ascii_lowercase().as_str(), "1" | "true" | "yes")) + .unwrap_or(false); let jwt_secret = get_or_create_jwt_secret( &app_mode, @@ -2682,13 +2686,6 @@ async fn main() -> Result<(), Error> { std::env::var("BILLING_SERVER_URL").ok() }; - let kagi_api_key = if app_mode != AppMode::Local { - // Get from database if in enclave mode - retrieve_kagi_api_key(aws_credential_manager.clone(), db.clone()).await? - } else { - std::env::var("KAGI_API_KEY").ok() - }; - let brave_api_key = if app_mode != AppMode::Local { // Get from database if in enclave mode retrieve_brave_api_key(aws_credential_manager.clone(), db.clone()).await? @@ -2737,8 +2734,10 @@ async fn main() -> Result<(), Error> { .billing_server_url(billing_server_url) .os_flags_base_url(os_flags_base_url) .os_flags_api_key(os_flags_api_key) - .kagi_api_key(kagi_api_key) .brave_api_key(brave_api_key) + .tinfoil_web_search_api_base(tinfoil_web_search_api_base) + .tinfoil_web_search_api_key(tinfoil_web_search_api_key) + .tinfoil_web_search_allow_insecure_tls(tinfoil_web_search_allow_insecure_tls) .build() .await?; tracing::info!("App state created, app_mode: {:?}", app_mode); diff --git a/src/models/responses.rs b/src/models/responses.rs index 3bf728eb..ba20fab0 100644 --- a/src/models/responses.rs +++ b/src/models/responses.rs @@ -756,6 +756,7 @@ impl NewConversation { completion_tokens: 0, status: "in_progress".to_string(), finish_reason: None, + created_at: Utc::now(), }; placeholder_assistant.insert(tx)?; } @@ -894,6 +895,7 @@ pub struct NewToolCall { pub arguments_enc: Option>, pub argument_tokens: i32, pub status: String, + pub created_at: DateTime, } impl ToolCall { @@ -956,6 +958,7 @@ pub struct NewToolOutput { pub output_tokens: i32, pub status: String, pub error: Option, + pub created_at: DateTime, } impl NewToolOutput { @@ -1000,6 +1003,7 @@ pub struct NewAssistantMessage { pub completion_tokens: i32, pub status: String, pub finish_reason: Option, + pub created_at: DateTime, } impl AssistantMessage { diff --git a/src/tinfoil_websearch.rs b/src/tinfoil_websearch.rs new file mode 100644 index 00000000..7c8cc549 --- /dev/null +++ b/src/tinfoil_websearch.rs @@ -0,0 +1,231 @@ +use serde_json::{json, Value}; +use std::sync::Arc; +use std::time::Duration; + +const REQUEST_TIMEOUT: Duration = Duration::from_secs(30); +const CONNECT_TIMEOUT: Duration = Duration::from_secs(5); + +#[derive(Debug, thiserror::Error)] +pub enum TinfoilWebSearchError { + #[error("HTTP request failed: {0}")] + Request(#[from] reqwest::Error), + #[error("API error: {status} - {message}")] + Api { status: u16, message: String }, + #[error("JSON-RPC error: {message}")] + Rpc { message: String }, + #[error("Invalid response format: {message}")] + InvalidResponse { message: String }, +} + +#[derive(Clone)] +pub struct TinfoilWebSearchClient { + client: reqwest::Client, + base_url: Arc, + api_key: Arc, +} + +impl TinfoilWebSearchClient { + pub fn new( + base_url: String, + api_key: String, + allow_insecure_tls: bool, + ) -> Result { + let client = reqwest::Client::builder() + .timeout(REQUEST_TIMEOUT) + .connect_timeout(CONNECT_TIMEOUT) + .pool_max_idle_per_host(100) + .danger_accept_invalid_certs(allow_insecure_tls) + .user_agent("OpenSecret/0.1.0") + .build()?; + + Ok(Self { + client, + base_url: Arc::new(base_url.trim_end_matches('/').to_string()), + api_key: Arc::new(api_key), + }) + } + + pub async fn search( + &self, + query: &str, + max_results: Option, + ) -> Result { + self.call_tool( + "search", + json!({ + "query": query, + "max_results": max_results.unwrap_or(5), + }), + ) + .await + } + + async fn call_tool( + &self, + name: &str, + arguments: Value, + ) -> Result { + let response = self + .client + .post(format!("{}/mcp", self.base_url.as_str())) + .bearer_auth(self.api_key.as_str()) + .json(&json!({ + "jsonrpc": "2.0", + "id": 1, + "method": "tools/call", + "params": { + "name": name, + "arguments": arguments, + } + })) + .send() + .await?; + + let status = response.status(); + if !status.is_success() { + let message = response.text().await.unwrap_or_default(); + return Err(TinfoilWebSearchError::Api { + status: status.as_u16(), + message, + }); + } + + let body = response.text().await?; + let payload = parse_response_body(&body)?; + if let Some(error) = payload.get("error") { + return Err(TinfoilWebSearchError::Rpc { + message: error + .get("message") + .and_then(Value::as_str) + .map(ToOwned::to_owned) + .unwrap_or_else(|| error.to_string()), + }); + } + + Ok(normalize_tool_result( + payload.get("result").cloned().unwrap_or(payload), + )) + } +} + +fn parse_response_body(body: &str) -> Result { + serde_json::from_str(body).or_else(|json_error| { + let sse_data = + extract_first_sse_data(body).ok_or_else(|| TinfoilWebSearchError::InvalidResponse { + message: format!( + "could not parse JSON or SSE payload: {json_error}; body prefix: {}", + truncate_for_error(body) + ), + })?; + + serde_json::from_str(&sse_data).map_err(|sse_error| { + TinfoilWebSearchError::InvalidResponse { + message: format!( + "failed to parse SSE data as JSON: {sse_error}; data prefix: {}", + truncate_for_error(&sse_data) + ), + } + }) + }) +} + +fn extract_first_sse_data(body: &str) -> Option { + let mut current_block = Vec::new(); + + for line in body.lines() { + if let Some(data) = line.strip_prefix("data:") { + let trimmed = data.trim_start(); + if trimmed != "[DONE]" { + current_block.push(trimmed.to_string()); + } + continue; + } + + if line.trim().is_empty() && !current_block.is_empty() { + return Some(current_block.join("\n")); + } + } + + if current_block.is_empty() { + None + } else { + Some(current_block.join("\n")) + } +} + +fn normalize_tool_result(result: Value) -> Value { + if let Some(structured) = result + .get("structuredContent") + .cloned() + .or_else(|| result.get("structured_content").cloned()) + { + return structured; + } + + if let Some(text) = result + .get("content") + .and_then(Value::as_array) + .and_then(|items| { + items + .iter() + .find_map(|item| item.get("text").and_then(Value::as_str)) + }) + { + if let Ok(parsed) = serde_json::from_str::(text) { + return parsed; + } + } + + result +} + +fn truncate_for_error(value: &str) -> String { + const MAX_LEN: usize = 200; + let mut truncated = value.chars().take(MAX_LEN).collect::(); + if value.chars().count() > MAX_LEN { + truncated.push_str("..."); + } + truncated +} + +impl std::fmt::Debug for TinfoilWebSearchClient { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("TinfoilWebSearchClient") + .field("base_url", &self.base_url) + .field("api_key", &"[REDACTED]") + .finish() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_parse_response_body_accepts_plain_json() { + let payload = parse_response_body(r#"{"jsonrpc":"2.0","id":1,"result":{"results":[]}}"#) + .expect("plain json"); + assert_eq!(payload["result"]["results"], json!([])); + } + + #[test] + fn test_parse_response_body_accepts_sse_wrapped_json() { + let payload = parse_response_body( + "event: message\ndata: {\"jsonrpc\":\"2.0\",\"id\":1,\"result\":{\"results\":[]}}\n\n", + ) + .expect("sse json"); + assert_eq!(payload["result"]["results"], json!([])); + } + + #[test] + fn test_normalize_tool_result_parses_json_text_content() { + let normalized = normalize_tool_result(json!({ + "content": [{ + "type": "text", + "text": "{\"results\":[{\"title\":\"Example\"}]}" + }] + })); + + assert_eq!(normalized["results"][0]["title"], "Example"); + } +} diff --git a/src/web/responses/builders.rs b/src/web/responses/builders.rs index daa51dcc..63dd59a5 100644 --- a/src/web/responses/builders.rs +++ b/src/web/responses/builders.rs @@ -149,6 +149,10 @@ impl OutputItemBuilder { status: STATUS_IN_PROGRESS.to_string(), role: Some(ROLE_ASSISTANT.to_string()), content: Some(vec![]), + call_id: None, + name: None, + arguments: None, + output: None, }, } } diff --git a/src/web/responses/context_builder.rs b/src/web/responses/context_builder.rs index 3d6fbe3b..960fd521 100644 --- a/src/web/responses/context_builder.rs +++ b/src/web/responses/context_builder.rs @@ -41,10 +41,11 @@ pub fn build_prompt( user_key: &secp256k1::SecretKey, model: &str, override_instructions: Option<&str>, + internal_instructions: Option<&str>, ) -> Result<(Vec, usize), crate::ApiError> { // 1. Get default user instructions if they exist (unless override provided) let mut system_tokens = 0usize; - let system_msg_opt = if let Some(instruction_text) = override_instructions { + let mut system_msg_opt = if let Some(instruction_text) = override_instructions { // Use override instructions if provided let tok = count_tokens(instruction_text); system_tokens = tok; @@ -89,6 +90,16 @@ pub fn build_prompt( } }; + if let Some(internal_instructions) = internal_instructions { + let content = if let Some((_, existing_content, _)) = system_msg_opt.take() { + format!("{internal_instructions}\n\n{existing_content}") + } else { + internal_instructions.to_string() + }; + system_tokens = count_tokens(&content); + system_msg_opt = Some((ROLE_SYSTEM, content, system_tokens)); + } + // 2. PASS 1: Fetch only metadata (lightweight, no content/decryption) let metadata = db .get_conversation_context_metadata(conversation_id) @@ -995,4 +1006,275 @@ mod tests { // First message should be the first user message assert_eq!(messages[0]["content"], "First message"); } + + /// Build a tool_call ChatMsg exactly as `build_prompt` would store it once + /// a tool call is persisted and later re-read from the database. + fn create_tool_call_chat_msg( + tool_call_id: uuid::Uuid, + tool_name: &str, + arguments_json: &str, + tokens: usize, + ) -> ChatMsg { + let tool_call_msg = serde_json::json!({ + "role": "assistant", + "tool_calls": [{ + "id": tool_call_id.to_string(), + "type": "function", + "function": { + "name": tool_name, + "arguments": arguments_json + } + }] + }); + + ChatMsg { + role: ROLE_ASSISTANT, + content: serde_json::to_string(&tool_call_msg).unwrap(), + tool_call_id: None, + tok: tokens, + } + } + + fn create_tool_output_chat_msg( + tool_call_id: uuid::Uuid, + output: &str, + tokens: usize, + ) -> ChatMsg { + ChatMsg { + role: "tool", + content: output.to_string(), + tool_call_id: Some(tool_call_id), + tok: tokens, + } + } + + #[test] + fn test_build_prompt_with_single_tool_turn() { + // Single agentic turn: user → tool_call → tool_output → assistant + let tool_call_id = uuid::Uuid::new_v4(); + let msgs = vec![ + create_chat_msg("user", "What's the weather in SF?", Some(8)), + create_tool_call_chat_msg( + tool_call_id, + "web_search", + "{\"query\":\"weather San Francisco\"}", + 9, + ), + create_tool_output_chat_msg(tool_call_id, "It is 68F and sunny in SF.", 7), + create_chat_msg( + "assistant", + "The weather in San Francisco is 68F and sunny.", + Some(11), + ), + ]; + + let (messages, total_tokens) = + build_prompt_from_chat_messages(msgs, "test-model").expect("build prompt"); + + assert_eq!(messages.len(), 4, "all four items must be in the prompt"); + + // user + assert_eq!(messages[0]["role"], "user"); + assert_eq!(messages[0]["content"], "What's the weather in SF?"); + + // tool_call materializes as assistant + tool_calls array (no content field) + assert_eq!(messages[1]["role"], "assistant"); + assert!(messages[1].get("tool_calls").is_some()); + let tool_calls = messages[1]["tool_calls"].as_array().unwrap(); + assert_eq!(tool_calls.len(), 1); + assert_eq!(tool_calls[0]["id"], tool_call_id.to_string()); + assert_eq!(tool_calls[0]["function"]["name"], "web_search"); + assert_eq!( + tool_calls[0]["function"]["arguments"], + "{\"query\":\"weather San Francisco\"}" + ); + + // tool_output materializes as role=tool with tool_call_id + assert_eq!(messages[2]["role"], "tool"); + assert_eq!(messages[2]["tool_call_id"], tool_call_id.to_string()); + assert_eq!(messages[2]["content"], "It is 68F and sunny in SF."); + + // final assistant + assert_eq!(messages[3]["role"], "assistant"); + assert_eq!(messages[3]["content"], "The weather in San Francisco is 68F and sunny."); + + // Token total must be the sum of all supplied ChatMsg tokens (no double counting + // of the serialized tool_call JSON, no skipping of tool_output tokens). + assert_eq!(total_tokens, 8 + 9 + 7 + 11); + } + + #[test] + fn test_build_prompt_with_multiple_tool_turns() { + // Two back-to-back tool turns on a single user question: + // user → tool_call_a → tool_output_a → tool_call_b → tool_output_b → assistant + let tool_call_a = uuid::Uuid::new_v4(); + let tool_call_b = uuid::Uuid::new_v4(); + let msgs = vec![ + create_chat_msg("user", "Compare today's weather in SF and NYC", Some(10)), + create_tool_call_chat_msg( + tool_call_a, + "web_search", + "{\"query\":\"weather SF today\"}", + 9, + ), + create_tool_output_chat_msg(tool_call_a, "SF: 68F sunny", 5), + create_tool_call_chat_msg( + tool_call_b, + "web_search", + "{\"query\":\"weather NYC today\"}", + 9, + ), + create_tool_output_chat_msg(tool_call_b, "NYC: 52F cloudy", 5), + create_chat_msg( + "assistant", + "SF is 68F sunny, NYC is 52F cloudy.", + Some(12), + ), + ]; + + let (messages, total_tokens) = + build_prompt_from_chat_messages(msgs, "test-model").expect("build prompt"); + + assert_eq!(messages.len(), 6); + + // Two distinct tool_call ids, preserved in their assistant tool_calls entries + assert_eq!(messages[1]["role"], "assistant"); + assert_eq!( + messages[1]["tool_calls"][0]["id"], + tool_call_a.to_string() + ); + assert_eq!(messages[3]["role"], "assistant"); + assert_eq!( + messages[3]["tool_calls"][0]["id"], + tool_call_b.to_string() + ); + + // Each tool_output must link back to the correct tool_call_id (FK correctness) + assert_eq!(messages[2]["role"], "tool"); + assert_eq!(messages[2]["tool_call_id"], tool_call_a.to_string()); + assert_eq!(messages[4]["role"], "tool"); + assert_eq!(messages[4]["tool_call_id"], tool_call_b.to_string()); + + // Final assistant message + assert_eq!(messages[5]["role"], "assistant"); + assert!( + messages[5].get("tool_calls").is_none(), + "final assistant message must not have tool_calls" + ); + + // Total prompt tokens must include every item exactly once + assert_eq!(total_tokens, 10 + 9 + 5 + 9 + 5 + 12); + } + + #[test] + fn test_build_prompt_with_system_then_tool_turn() { + // System prompt + agentic turn: verify system tokens are counted and the + // system message is emitted first. + let tool_call_id = uuid::Uuid::new_v4(); + let msgs = vec![ + create_chat_msg("system", "You are Maple.", Some(4)), + create_chat_msg("user", "Who won the UFC fight last night?", Some(10)), + create_tool_call_chat_msg( + tool_call_id, + "web_search", + "{\"query\":\"UFC fight results last night\"}", + 10, + ), + create_tool_output_chat_msg(tool_call_id, "Fighter A won by KO", 6), + create_chat_msg("assistant", "Fighter A won by KO last night.", Some(8)), + ]; + + let (messages, total_tokens) = + build_prompt_from_chat_messages(msgs, "test-model").expect("build prompt"); + + assert_eq!(messages.len(), 5); + assert_eq!(messages[0]["role"], "system"); + assert_eq!(messages[0]["content"], "You are Maple."); + assert_eq!(messages[1]["role"], "user"); + assert_eq!(messages[2]["role"], "assistant"); + assert!(messages[2].get("tool_calls").is_some()); + assert_eq!(messages[3]["role"], "tool"); + assert_eq!(messages[4]["role"], "assistant"); + + assert_eq!(total_tokens, 4 + 10 + 10 + 6 + 8); + } + + #[test] + fn test_tool_turn_survives_truncation_with_last_user_intact() { + // A long-running agentic conversation that overflows the budget should still: + // - preserve the most recent user question + // - preserve the trailing tool_call/tool_output/assistant triple so the + // model has enough signal to answer + // - drop the oldest tool turns first (middle truncation) + let mut msgs = vec![create_chat_msg("user", "First question", Some(5))]; + + // Bloat the middle with many old tool turns + for i in 0..20 { + let id = uuid::Uuid::new_v4(); + msgs.push(create_tool_call_chat_msg( + id, + "web_search", + "{\"query\":\"old\"}", + 1000, + )); + msgs.push(create_tool_output_chat_msg(id, "old result", 1000)); + msgs.push(create_chat_msg( + "assistant", + &format!("old answer {}", i), + Some(1000), + )); + } + + // Recent turn that MUST survive + let recent_tool_id = uuid::Uuid::new_v4(); + msgs.push(create_chat_msg("user", "Recent question", Some(5))); + msgs.push(create_tool_call_chat_msg( + recent_tool_id, + "web_search", + "{\"query\":\"recent\"}", + 10, + )); + msgs.push(create_tool_output_chat_msg( + recent_tool_id, + "recent result", + 10, + )); + msgs.push(create_chat_msg( + "assistant", + "Recent answer", + Some(10), + )); + msgs.push(create_chat_msg("user", "Final follow-up", Some(5))); + + let (messages, total_tokens) = + build_prompt_from_chat_messages(msgs, "deepseek-r1-70b").expect("build prompt"); + + // Must end with the final follow-up user message + assert_eq!(messages.last().unwrap()["role"], "user"); + assert_eq!(messages.last().unwrap()["content"], "Final follow-up"); + + // Must contain the recent tool_call + tool_output linked by tool_call_id + let recent_tool_output = messages + .iter() + .find(|m| { + m["role"] == "tool" + && m["tool_call_id"].as_str() == Some(&recent_tool_id.to_string()) + }) + .expect("recent tool_output must survive truncation"); + assert_eq!(recent_tool_output["content"], "recent result"); + + // Must contain the truncation marker + assert!(messages + .iter() + .any(|m| m["content"].as_str().unwrap_or("").contains("truncated"))); + + // Must stay inside the budget + let budget = 64_000 - 4096 - 500; + assert!( + total_tokens <= budget, + "prompt tokens {} exceed budget {}", + total_tokens, + budget + ); + } } diff --git a/src/web/responses/handlers.rs b/src/web/responses/handlers.rs index d45f6028..6d7bf03d 100644 --- a/src/web/responses/handlers.rs +++ b/src/web/responses/handlers.rs @@ -5,14 +5,14 @@ use crate::{ billing::BillingError, db::DBError, encrypt::{decrypt_content, decrypt_string, encrypt_with_key}, - models::responses::{NewAssistantMessage, NewUserMessage, ResponseStatus, ResponsesError}, + models::responses::{NewUserMessage, ResponseStatus, ResponsesError}, models::users::User, tokens::{count_tokens, model_max_ctx}, web::{ encryption_middleware::{decrypt_request, encrypt_response, EncryptedResponse}, openai::get_chat_completion_response, responses::{ - build_prompt, build_usage, constants::*, error_mapping, prompts, storage_task, tools, + build_prompt, build_usage, constants::*, error_mapping, storage_task, tools, ContentPartBuilder, DeletedObjectResponse, MessageContent, MessageContentConverter, MessageContentPart, OutputItemBuilder, ResponseBuilder, ResponseEvent, SseEventEmitter, }, @@ -33,7 +33,7 @@ use futures::Stream; use secp256k1::SecretKey; use serde::{Deserialize, Serialize}; use serde_json::{json, Value}; -use std::sync::Arc; +use std::{collections::HashMap, sync::Arc}; use tokio::sync::mpsc; use tracing::{debug, error, info, trace, warn}; use uuid::Uuid; @@ -71,10 +71,172 @@ fn apply_responses_model_defaults(chat_request: &mut Value, model: &str) { } } +const MAPLE_SYSTEM_PROMPT: &str = "You are Maple, a friendly, concise, and helpful assistant. Give direct answers, be honest about uncertainty, and never invent tool use, search results, or sources."; +const MAPLE_WEB_SEARCH_PROMPT: &str = "If the web_search tool is available and the user explicitly asks you to search, look something up, verify, confirm, or check the web, call web_search before answering. Also use web_search when the answer depends on current or time-sensitive information. You may use web_search repeatedly across a single response when needed, but only one tool call at a time and never more than 15 tool calls for one user request. After each tool output, decide whether you have enough information to answer or whether another search is still needed. Prefer to stop searching and answer as soon as you have enough information. After receiving tool results, you must either call another tool or provide a final user-visible answer in assistant content. Do not end the turn with reasoning only. Do not place the final answer in reasoning. Never output raw tool call syntax."; + +#[derive(Debug, Clone)] +struct ModelToolCall { + name: String, + arguments: Value, +} + +#[derive(Debug, Clone, Default)] +struct StreamedToolCall { + name: Option, + arguments: String, +} + +#[derive(Debug, Clone)] +enum AssistantTurnOutcome { + ToolCall(ModelToolCall), + Final, +} + +fn should_enable_web_search_tool(state: &AppState, body: &ResponsesCreateRequest) -> bool { + is_tool_choice_allowed(&body.tool_choice) + && is_web_search_enabled(&body.tools) + && state.has_web_search_client() +} + +fn build_internal_system_prompt_for_now( + now: chrono::DateTime, + web_search_enabled: bool, +) -> String { + let current_utc_date = now.format("%A, %Y-%m-%d").to_string(); + let current_date_prompt = format!( + "Current UTC date: {current_utc_date}. Use this as today's date for any date-sensitive reasoning." + ); + + if web_search_enabled { + format!("{MAPLE_SYSTEM_PROMPT}\n\n{current_date_prompt}\n\n{MAPLE_WEB_SEARCH_PROMPT}") + } else { + format!("{MAPLE_SYSTEM_PROMPT}\n\n{current_date_prompt}") + } +} + +fn build_internal_system_prompt(web_search_enabled: bool) -> String { + build_internal_system_prompt_for_now(Utc::now(), web_search_enabled) +} + +fn build_provider_tools(request_tools: &Option) -> Vec { + let registry = tools::ToolRegistry::new(); + + request_tools + .as_ref() + .and_then(|tools| tools.as_array()) + .map(|tools| { + tools + .iter() + .filter_map(|tool| { + let tool_name = tool.get("type").and_then(|t| t.as_str())?; + registry.get_tool_schema(tool_name).map(|schema| { + json!({ + "type": "function", + "function": schema, + }) + }) + }) + .collect() + }) + .unwrap_or_default() +} + +fn build_tool_choice_value(tool_choice: &Option) -> Value { + match tool_choice.as_deref() { + Some(choice) if !choice.is_empty() => json!(choice), + _ => json!("auto"), + } +} + +fn build_model_turn_request( + body: &ResponsesCreateRequest, + prompt_messages: &[Value], + tools_enabled: bool, +) -> Value { + let mut chat_request = json!({ + "model": body.model, + "messages": prompt_messages, + "temperature": body.temperature.unwrap_or(DEFAULT_TEMPERATURE), + "top_p": body.top_p.unwrap_or(DEFAULT_TOP_P), + "max_tokens": body.max_output_tokens, + "stream": true, + "stream_options": { "include_usage": true } + }); + + if tools_enabled { + let provider_tools = build_provider_tools(&body.tools); + if !provider_tools.is_empty() { + chat_request["tools"] = Value::Array(provider_tools); + chat_request["tool_choice"] = build_tool_choice_value(&body.tool_choice); + chat_request["parallel_tool_calls"] = json!(false); + } + } + + apply_responses_model_defaults(&mut chat_request, &body.model); + chat_request +} + +fn append_streamed_tool_calls(tool_calls: &mut Vec, tool_call_delta: &Value) { + let Some(tool_call_entries) = tool_call_delta.as_array() else { + return; + }; + + if tool_call_entries.len() > 1 { + warn!( + "Model streamed {} tool calls in one chunk; only the first call will be executed in v1", + tool_call_entries.len() + ); + } + + for tool_call in tool_call_entries { + let index = tool_call + .get("index") + .and_then(|index| index.as_u64()) + .unwrap_or(0) as usize; + + while tool_calls.len() <= index { + tool_calls.push(StreamedToolCall::default()); + } + + if let Some(function) = tool_call.get("function") { + if let Some(name) = function.get("name").and_then(|name| name.as_str()) { + tool_calls[index].name = Some(name.to_string()); + } + + if let Some(arguments) = function + .get("arguments") + .and_then(|arguments| arguments.as_str()) + { + tool_calls[index].arguments.push_str(arguments); + } + } + } +} + +fn finalize_first_model_tool_call(tool_calls: &[StreamedToolCall]) -> Option { + let tool_call = tool_calls.first()?; + let name = tool_call.name.clone()?; + let arguments = serde_json::from_str(&tool_call.arguments).unwrap_or_else(|e| { + warn!( + "Failed to parse tool arguments for {} as JSON: {:?}. Using empty object.", + name, e + ); + json!({}) + }); + + Some(ModelToolCall { name, arguments }) +} + #[cfg(test)] mod tests { - use super::apply_responses_model_defaults; + use super::{ + append_streamed_tool_calls, apply_responses_model_defaults, + build_internal_system_prompt_for_now, build_provider_tools, finalize_first_model_tool_call, + ClientResponseState, StreamedToolCall, MAPLE_WEB_SEARCH_PROMPT, + }; + use chrono::{TimeZone, Utc}; use serde_json::json; + use uuid::Uuid; #[test] fn test_apply_responses_model_defaults_enables_gemma_thinking() { @@ -106,6 +268,93 @@ mod tests { assert!(chat_request.get("include_reasoning").is_none()); assert!(chat_request.get("chat_template_kwargs").is_none()); } + + #[test] + fn test_append_streamed_tool_calls_reassembles_arguments() { + let mut tool_calls = Vec::::new(); + + append_streamed_tool_calls( + &mut tool_calls, + &json!([{ + "index": 0, + "function": { + "name": "web_search", + "arguments": "{\"query\":\"Don" + } + }]), + ); + append_streamed_tool_calls( + &mut tool_calls, + &json!([{ + "index": 0, + "function": { + "arguments": "ald Trump birthday\"}" + } + }]), + ); + + let tool_call = finalize_first_model_tool_call(&tool_calls).expect("tool call"); + assert_eq!(tool_call.name, "web_search"); + assert_eq!(tool_call.arguments["query"], "Donald Trump birthday"); + } + + #[test] + fn test_build_provider_tools_filters_unknown_tools() { + let tools = build_provider_tools(&Some(json!([ + { "type": "web_search" }, + { "type": "unknown_tool" } + ]))); + + assert_eq!(tools.len(), 1); + assert_eq!(tools[0]["type"], "function"); + assert_eq!(tools[0]["function"]["name"], "web_search"); + } + + #[test] + fn test_build_internal_system_prompt_includes_current_utc_date() { + let now = Utc + .with_ymd_and_hms(2026, 4, 15, 12, 0, 0) + .single() + .expect("valid UTC timestamp"); + + let prompt = build_internal_system_prompt_for_now(now, true); + + assert!(prompt.contains("Current UTC date: Wednesday, 2026-04-15.")); + assert!(prompt.contains(MAPLE_WEB_SEARCH_PROMPT)); + } + + #[test] + fn test_client_response_state_build_output_items_uses_maple_tool_types() { + let mut state = ClientResponseState::default(); + let tool_call_id = Uuid::new_v4(); + let tool_output_id = Uuid::new_v4(); + + state.push_tool_call( + tool_call_id, + "web_search".to_string(), + json!({ "query": "ufc" }), + ); + state.push_tool_output( + tool_output_id, + tool_call_id, + "Search Results:\n\n1. Example".to_string(), + ); + + let output_items = state.build_output_items(); + let tool_call_id_str = tool_call_id.to_string(); + + assert_eq!(output_items.len(), 2); + assert_eq!(output_items[0].output_type, "tool_call"); + assert_eq!( + output_items[0].call_id.as_deref(), + Some(tool_call_id_str.as_str()) + ); + assert_eq!(output_items[1].output_type, "tool_output"); + assert_eq!( + output_items[1].call_id.as_deref(), + Some(tool_call_id_str.as_str()) + ); + } } /// Conversation parameter - can be a string UUID or an object with id field @@ -181,15 +430,15 @@ pub struct ResponsesCreateRequest { /// Maximum tokens for the response pub max_output_tokens: Option, - /// Tool choice strategy (ignored in Phase 4/5) + /// Tool choice strategy #[serde(default)] pub tool_choice: Option, - /// Tools available for the model (ignored in Phase 4/5) + /// Tools available for the model #[serde(default)] pub tools: Option, - /// Enable parallel tool calls (ignored in Phase 4/5) + /// Enable parallel tool calls #[serde(default)] pub parallel_tool_calls: bool, @@ -272,7 +521,7 @@ pub struct ResponsesCreateResponse { /// Tool choice setting pub tool_choice: String, - /// Available tools (empty array for Phase 4/5) + /// Available tools pub tools: Vec, /// Top logprobs @@ -339,6 +588,22 @@ pub struct OutputItem { /// Content array (for message type) #[serde(skip_serializing_if = "Option::is_none")] pub content: Option>, + + /// Tool call ID (for tool_call / tool_output types) + #[serde(skip_serializing_if = "Option::is_none")] + pub call_id: Option, + + /// Tool/function name (for tool_call type) + #[serde(skip_serializing_if = "Option::is_none")] + pub name: Option, + + /// Tool arguments JSON (for tool_call type) + #[serde(skip_serializing_if = "Option::is_none")] + pub arguments: Option, + + /// Tool output payload (for tool_output type) + #[serde(skip_serializing_if = "Option::is_none")] + pub output: Option, } /// Response error structure @@ -701,6 +966,9 @@ pub struct ToolCallCreatedEvent { /// Sequence number for ordering pub sequence_number: i32, + /// Index of the corresponding output item + pub output_index: i32, + /// Tool call ID pub tool_call_id: Uuid, @@ -721,6 +989,9 @@ pub struct ToolOutputCreatedEvent { /// Sequence number for ordering pub sequence_number: i32, + /// Index of the corresponding output item + pub output_index: i32, + /// Tool output ID pub tool_output_id: Uuid, @@ -734,19 +1005,34 @@ pub struct ToolOutputCreatedEvent { /// Message types for the storage task #[derive(Debug, Clone)] pub enum StorageMessage { - ContentDelta(String), + MessageStarted { + item_id: Uuid, + }, + ContentDelta { + item_id: Uuid, + delta: String, + }, + MessageDone { + item_id: Uuid, + finish_reason: String, + }, + ReasoningStarted { + item_id: Uuid, + }, /// Reasoning delta with item_id to ensure SSE and DB use the same UUID ReasoningDelta { item_id: Uuid, delta: String, }, + ReasoningDone { + item_id: Uuid, + }, Usage { prompt_tokens: i32, completion_tokens: i32, }, - Done { + ResponseDone { finish_reason: String, - message_id: Uuid, }, Error(String), Cancelled, @@ -761,8 +1047,188 @@ pub enum StorageMessage { tool_call_id: Uuid, output: String, }, - /// Signal that assistant message is about to start streaming - AssistantMessageStarting, +} + +#[derive(Debug, Clone)] +enum StreamOutputItemRecord { + Message { + id: Uuid, + text: String, + }, + Reasoning { + id: Uuid, + text: String, + }, + ToolCall { + id: Uuid, + call_id: Uuid, + name: String, + arguments: String, + }, + ToolOutput { + id: Uuid, + call_id: Uuid, + output: String, + }, +} + +#[derive(Default)] +struct ClientResponseState { + items: Vec, + indices: HashMap, +} + +impl ClientResponseState { + fn push_message(&mut self, item_id: Uuid) -> i32 { + let output_index = self.items.len(); + self.items.push(StreamOutputItemRecord::Message { + id: item_id, + text: String::new(), + }); + self.indices.insert(item_id, output_index); + output_index as i32 + } + + fn push_reasoning(&mut self, item_id: Uuid) -> i32 { + let output_index = self.items.len(); + self.items.push(StreamOutputItemRecord::Reasoning { + id: item_id, + text: String::new(), + }); + self.indices.insert(item_id, output_index); + output_index as i32 + } + + fn push_tool_call(&mut self, item_id: Uuid, name: String, arguments: Value) -> i32 { + let output_index = self.items.len(); + let arguments = serde_json::to_string(&arguments).unwrap_or_else(|_| "{}".to_string()); + self.items.push(StreamOutputItemRecord::ToolCall { + id: item_id, + call_id: item_id, + name, + arguments, + }); + self.indices.insert(item_id, output_index); + output_index as i32 + } + + fn push_tool_output(&mut self, item_id: Uuid, call_id: Uuid, output: String) -> i32 { + let output_index = self.items.len(); + self.items.push(StreamOutputItemRecord::ToolOutput { + id: item_id, + call_id, + output, + }); + self.indices.insert(item_id, output_index); + output_index as i32 + } + + fn message_output_index(&self, item_id: Uuid) -> Option { + self.indices.get(&item_id).map(|index| *index as i32) + } + + fn reasoning_output_index(&self, item_id: Uuid) -> Option { + self.indices.get(&item_id).map(|index| *index as i32) + } + + fn append_message_delta(&mut self, item_id: Uuid, delta: &str) -> Option { + let index = *self.indices.get(&item_id)?; + if let Some(StreamOutputItemRecord::Message { text, .. }) = self.items.get_mut(index) { + text.push_str(delta); + Some(index as i32) + } else { + None + } + } + + fn append_reasoning_delta(&mut self, item_id: Uuid, delta: &str) -> Option { + let index = *self.indices.get(&item_id)?; + if let Some(StreamOutputItemRecord::Reasoning { text, .. }) = self.items.get_mut(index) { + text.push_str(delta); + Some(index as i32) + } else { + None + } + } + + fn message_text(&self, item_id: Uuid) -> Option<&str> { + let index = *self.indices.get(&item_id)?; + match self.items.get(index)? { + StreamOutputItemRecord::Message { text, .. } => Some(text.as_str()), + _ => None, + } + } + + fn reasoning_text(&self, item_id: Uuid) -> Option<&str> { + let index = *self.indices.get(&item_id)?; + match self.items.get(index)? { + StreamOutputItemRecord::Reasoning { text, .. } => Some(text.as_str()), + _ => None, + } + } + + fn build_output_items(&self) -> Vec { + self.items + .iter() + .map(|item| match item { + StreamOutputItemRecord::Message { id, text } => OutputItem { + id: id.to_string(), + output_type: OUTPUT_TYPE_MESSAGE.to_string(), + status: STATUS_COMPLETED.to_string(), + role: Some(ROLE_ASSISTANT.to_string()), + content: Some(vec![ + ContentPartBuilder::new_output_text(text.clone()).build() + ]), + call_id: None, + name: None, + arguments: None, + output: None, + }, + StreamOutputItemRecord::Reasoning { id, .. } => OutputItem { + id: id.to_string(), + output_type: "reasoning".to_string(), + status: STATUS_COMPLETED.to_string(), + role: None, + content: Some(vec![]), + call_id: None, + name: None, + arguments: None, + output: None, + }, + StreamOutputItemRecord::ToolCall { + id, + call_id, + name, + arguments, + } => OutputItem { + id: id.to_string(), + output_type: "tool_call".to_string(), + status: STATUS_COMPLETED.to_string(), + role: None, + content: None, + call_id: Some(call_id.to_string()), + name: Some(name.clone()), + arguments: Some(arguments.clone()), + output: None, + }, + StreamOutputItemRecord::ToolOutput { + id, + call_id, + output, + } => OutputItem { + id: id.to_string(), + output_type: "tool_output".to_string(), + status: STATUS_COMPLETED.to_string(), + role: None, + content: None, + call_id: Some(call_id.to_string()), + name: None, + arguments: None, + output: Some(output.clone()), + }, + }) + .collect() + } } /// Validated and prepared request data @@ -1096,6 +1562,9 @@ async fn build_context_and_check_billing( user_key: &SecretKey, prepared: &PreparedRequest, ) -> Result { + let internal_system_prompt = + build_internal_system_prompt(should_enable_web_search_tool(state.as_ref(), body)); + // Extract conversation ID from the required conversation parameter let conv_uuid = match &body.conversation { ConversationParam::String(id) | ConversationParam::Object { id } => *id, @@ -1117,6 +1586,7 @@ async fn build_context_and_check_billing( user_key, &body.model, body.instructions.as_deref(), + Some(&internal_system_prompt), )?; // Add the NEW user message to the context (not yet persisted) @@ -1184,12 +1654,8 @@ async fn build_context_and_check_billing( /// - Creates Response record (status=in_progress) /// - Creates user message /// -/// Note: Assistant message is NOT created here - it's created later in Phase 6 (after tools). -/// Originally, the assistant placeholder was created here, but this caused timestamp -/// ordering issues: the assistant message would get created_at=T1 (early), then tools -/// would execute at T2/T3, making the assistant appear BEFORE its tools in queries -/// ordered by created_at. By creating the assistant message in Phase 6 (after tools), -/// we ensure the correct semantic order: user → tool_call → tool_output → assistant. +/// Note: Assistant and tool items are NOT created here - they're created later by the +/// storage task as stream events arrive so persisted ordering matches the emitted item order. async fn persist_request_data( state: &Arc, user: &User, @@ -1266,8 +1732,8 @@ async fn persist_request_data( .create_user_message(new_msg) .map_err(error_mapping::map_generic_db_error)?; - // Capture user message timestamp for reasoning item ordering - // Reasoning should appear after user message but before assistant message + // Capture the user message timestamp so persisted response items can be assigned + // a single monotonic created_at sequence immediately after the user message. let user_message_created_at = user_message.created_at; info!( @@ -1313,391 +1779,198 @@ fn is_web_search_enabled(tools: &Option) -> bool { false } -/// Phase 5: Classify intent and execute tools (optional) -/// -/// Classifies user intent and executes tools if needed. Runs after dual streams -/// are created so tool events can be sent to both client and storage. -/// -/// Flow: -/// 1. Classify intent: chat vs web_search -/// 2. If web_search: extract query and execute tool -/// 3. Send ToolCall event to streams -/// 4. Send ToolOutput event to streams (always, even on error) -/// 5. Send persistence command via dedicated channel and wait for acknowledgment -/// -/// Tool execution is best-effort: intent classification uses gpt-oss-120b and -/// query extraction uses llama3-3-70b. -struct ToolChannels<'a> { - client: &'a mpsc::Sender, - storage: &'a mpsc::Sender, -} - -async fn classify_and_execute_tools( +/// Phase 5: Let the model request tool use (optional) +/// Persist and emit a single requested tool call, then wait for storage to +/// confirm the tool output is durable before the next model turn is started. +async fn execute_tool_call_and_wait( state: &Arc, - user: &User, - prepared: &PreparedRequest, persisted: &PersistedData, - prompt_messages: &[Value], - channels: ToolChannels<'_>, - rx_tool_ack: tokio::sync::oneshot::Receiver>, -) -> Result, ApiError> { - let ToolChannels { - client: tx_client, - storage: tx_storage, - } = channels; - // Extract text from user message for classification - let user_text = - MessageContentConverter::extract_text_for_token_counting(&prepared.message_content); - - trace!( - "Classifying user intent for message: {}", - user_text.chars().take(100).collect::() - ); - debug!( - "Starting intent classification with {} conversation messages", - prompt_messages.len() - ); - - // Step 1: Classify intent using LLM with conversation history - let classification_request = - prompts::build_intent_classification_request(prompt_messages, &user_text); - - trace!( - "Intent classification request: {}", - serde_json::to_string_pretty(&classification_request) - .unwrap_or_else(|_| "failed to serialize".to_string()) - ); + tool_call: ModelToolCall, + tx_client: &mpsc::Sender, + tx_storage: &mpsc::Sender, + rx_tool_ack: &mut mpsc::Receiver>, +) -> Result<(), ApiError> { + let tool_call_id = Uuid::new_v4(); + let tool_output_id = Uuid::new_v4(); + + let tool_call_msg = StorageMessage::ToolCall { + tool_call_id, + name: tool_call.name.clone(), + arguments: tool_call.arguments.clone(), + }; - let headers = HeaderMap::new(); - let billing_context = crate::web::openai::BillingContext::new( - crate::web::openai_auth::AuthMethod::Jwt, - prompts::INTENT_CLASSIFIER_MODEL.to_string(), - ); + if let Err(e) = tx_storage.send(tool_call_msg.clone()).await { + error!( + "Critical: Storage channel closed during tool_call for response {} - {:?}", + persisted.response.uuid, e + ); + let _ = tx_client + .send(StorageMessage::Error( + "Internal storage failure - request aborted".to_string(), + )) + .await; + return Err(ApiError::InternalServerError); + } + if tx_client.try_send(tool_call_msg).is_err() { + warn!("Client channel full or closed, skipping tool_call event to client"); + } - let intent = match get_chat_completion_response( - state, - user, - classification_request, - &headers, - billing_context, + let tool_output = match tools::execute_tool( + &tool_call.name, + &tool_call.arguments, + state.tinfoil_web_search_client.as_ref(), + state.brave_client.as_ref(), ) .await { - Ok(mut completion) => { - match completion.stream.recv().await { - Some(crate::web::openai::CompletionChunk::FullResponse(response_json)) => { - trace!( - "Intent classification response: {}", - serde_json::to_string_pretty(&response_json) - .unwrap_or_else(|_| "failed to serialize".to_string()) - ); - - // Extract intent from response - if let Some(intent_str) = response_json - .get("choices") - .and_then(|c| c.get(0)) - .and_then(|c| c.get("message")) - .and_then(|m| m.get("content")) - .and_then(|c| c.as_str()) - { - let intent = intent_str.trim().to_lowercase(); - debug!("Classified intent: {}", intent); - intent - } else { - warn!( - "Failed to extract intent from classifier response, defaulting to chat" - ); - trace!("Response structure: {:?}", response_json); - "chat".to_string() - } - } - _ => { - warn!("Unexpected classifier response format, defaulting to chat"); - "chat".to_string() - } - } - } + Ok(output) => output, Err(e) => { - // Best effort - if classification fails, default to chat - warn!("Classification failed (defaulting to chat): {:?}", e); - "chat".to_string() + warn!("Tool execution failed, including error in output: {:?}", e); + format!("Error: {}", e) } }; - // Step 2: If intent is web_search, execute tool - if intent == "web_search" { - debug!("User message classified as web_search, executing tool"); + let tool_output_msg = StorageMessage::ToolOutput { + tool_output_id, + tool_call_id, + output: tool_output, + }; - // Extract search query with conversation history for context - let query_request = prompts::build_query_extraction_request(prompt_messages, &user_text); - let billing_context = crate::web::openai::BillingContext::new( - crate::web::openai_auth::AuthMethod::Jwt, - prompts::QUERY_EXTRACTOR_MODEL.to_string(), + if let Err(e) = tx_storage.send(tool_output_msg.clone()).await { + error!( + "Critical: Storage channel closed during tool_output for response {} - {:?}", + persisted.response.uuid, e ); + let _ = tx_client + .send(StorageMessage::Error( + "Internal storage failure - request aborted".to_string(), + )) + .await; + return Err(ApiError::InternalServerError); + } + if tx_client.try_send(tool_output_msg).is_err() { + warn!("Client channel full or closed, skipping tool_output event to client"); + } - let search_query = match get_chat_completion_response( - state, - user, - query_request, - &headers, - billing_context, - ) - .await - { - Ok(mut completion) => match completion.stream.recv().await { - Some(crate::web::openai::CompletionChunk::FullResponse(response_json)) => { - if let Some(query) = response_json - .get("choices") - .and_then(|c| c.get(0)) - .and_then(|c| c.get("message")) - .and_then(|m| m.get("content")) - .and_then(|c| c.as_str()) - { - let query = query.trim().to_string(); - trace!("Extracted search query: {}", query); - debug!("Search query extracted successfully"); - query - } else { - warn!("Failed to extract query, using original message"); - user_text.clone() - } - } - _ => { - warn!("Unexpected query extraction response, using original message"); - user_text.clone() - } - }, - Err(e) => { - warn!("Query extraction failed, using original message: {:?}", e); - user_text.clone() - } - }; - - // Generate UUIDs for tool_call and tool_output - let tool_call_id = Uuid::new_v4(); - let tool_output_id = Uuid::new_v4(); - - // Prepare tool arguments - let tool_arguments = json!({"query": search_query}); - - // Send tool_call event through both streams FIRST (before execution) - let tool_call_msg = StorageMessage::ToolCall { - tool_call_id, - name: "web_search".to_string(), - arguments: tool_arguments.clone(), - }; + info!( + "Successfully sent tool_call {} and tool_output {} to streams for conversation {}", + tool_call_id, tool_output_id, persisted.response.conversation_id + ); - // Send to storage (critical - must succeed) - // - // IMPORTANT: Storage channel failure means the storage task has died or the - // channel buffer (1024) is full. This is a catastrophic systemic failure, not - // a normal error. If this happens: - // 1. Nothing will be persisted to database - // 2. Continuing would waste LLM API calls for unsaved data - // 3. Client would see tool_call.created but never get completion - // - // We abort the entire request and notify the client with response.error event. - // If you see this error in production, investigate immediately - it indicates - // serious issues with the storage task or database connection. - if let Err(e) = tx_storage.send(tool_call_msg.clone()).await { - error!( - "Critical: Storage channel closed during tool_call for response {} - {:?}", - persisted.response.uuid, e - ); - // Notify client and abort - storage failure is catastrophic - let _ = tx_client - .send(StorageMessage::Error( - "Internal storage failure - request aborted".to_string(), - )) - .await; - return Err(ApiError::InternalServerError); + match tokio::time::timeout(std::time::Duration::from_secs(5), rx_tool_ack.recv()).await { + Ok(Some(Ok(()))) => { + debug!("Tools persisted successfully to database"); + Ok(()) + } + Ok(Some(Err(e))) => { + error!("Failed to persist tools to database: {}", e); + Err(ApiError::InternalServerError) } - // Send to client (best-effort) - if tx_client.try_send(tool_call_msg).is_err() { - warn!("Client channel full or closed, skipping tool_call event to client"); + Ok(None) => { + error!("Storage task dropped before sending tool acknowledgment"); + Err(ApiError::InternalServerError) } + Err(_) => { + error!("Timeout waiting for tool persistence (5s)"); + Err(ApiError::InternalServerError) + } + } +} - debug!("Sent tool_call {} to streams", tool_call_id); +async fn send_storage_message( + tx_storage: &mpsc::Sender, + tx_client: &mpsc::Sender, + msg: StorageMessage, +) -> Result<(), ApiError> { + if tx_storage.send(msg.clone()).await.is_err() { + error!("Storage channel closed unexpectedly"); + return Err(ApiError::InternalServerError); + } + if tx_client.try_send(msg).is_err() { + warn!("Client channel full or closed"); + } + Ok(()) +} - // Execute web search tool (or capture error as content) - let tool_output = match tools::execute_tool( - "web_search", - &tool_arguments, - state.brave_client.as_ref(), - state.kagi_client.as_ref(), - ) - .await - { - Ok(output) => { - debug!( - "Tool execution successful, output length: {} chars", - output.len() - ); - output - } - Err(e) => { - warn!("Tool execution failed, including error in output: {:?}", e); - // Failure becomes content, not a skip! - format!("Error: {}", e) - } - }; +fn next_assistant_message_id(next_message_id: &mut Option) -> Uuid { + next_message_id.take().unwrap_or_else(Uuid::new_v4) +} - // Send tool_output event through both streams (ALWAYS sent, even on failure) - let tool_output_msg = StorageMessage::ToolOutput { - tool_output_id, - tool_call_id, - output: tool_output.clone(), - }; +/// Lifecycle of the reasoning item within a single assistant turn. +/// +/// Models that emit `reasoning`/`reasoning_content` deltas must always close the +/// reasoning block before any final assistant content or tool-call deltas are +/// accepted, so we track the state explicitly rather than juggling several bools. +#[derive(Debug)] +enum ReasoningState { + NotStarted, + Active(Uuid), + Done, +} - // Send to storage (critical - must succeed) - // - // IMPORTANT: Storage channel failure is catastrophic (see tool_call comment above). - // At this point, client has already seen tool_call.created event. If storage fails - // here, we have an inconsistency: - // - Database has tool_call record but no tool_output - // - Client saw tool_call.created but won't see tool_output.created - // - // We abort and send response.error so the client knows the request failed rather - // than hanging indefinitely waiting for completion. The database inconsistency - // (orphaned tool_call) is acceptable given this is a catastrophic failure scenario. - if let Err(e) = tx_storage.send(tool_output_msg.clone()).await { - error!( - "Critical: Storage channel closed during tool_output for response {} - {:?}", - persisted.response.uuid, e - ); - // Notify client and abort - storage failure is catastrophic - let _ = tx_client - .send(StorageMessage::Error( - "Internal storage failure - request aborted".to_string(), - )) - .await; - return Err(ApiError::InternalServerError); - } - // Send to client (best-effort) - if tx_client.try_send(tool_output_msg).is_err() { - warn!("Client channel full or closed, skipping tool_output event to client"); +impl ReasoningState { + fn active_id(&self) -> Option { + if let ReasoningState::Active(id) = self { + Some(*id) + } else { + None } + } +} - info!( - "Successfully sent tool_call {} and tool_output {} to streams for conversation {}", - tool_call_id, tool_output_id, persisted.response.conversation_id - ); - - // Wait for storage task to confirm persistence (with timeout) - match tokio::time::timeout(std::time::Duration::from_secs(5), rx_tool_ack).await { - Ok(Ok(Ok(()))) => { - debug!("Tools persisted successfully to database"); - return Ok(Some(())); - } - Ok(Ok(Err(e))) => { - error!("Failed to persist tools to database: {}", e); - // Continue anyway - best effort - return Ok(Some(())); - } - Ok(Err(_)) => { - error!("Storage task dropped before sending acknowledgment"); - return Ok(Some(())); - } - Err(_) => { - error!("Timeout waiting for tool persistence (5s)"); - return Ok(Some(())); - } - } - } else { - debug!("User message classified as chat, skipping tool execution"); +async fn close_reasoning_if_active( + state: &mut ReasoningState, + tx_storage: &mpsc::Sender, + tx_client: &mpsc::Sender, +) -> Result<(), ApiError> { + if let Some(reasoning_id) = state.active_id() { + send_storage_message( + tx_storage, + tx_client, + StorageMessage::ReasoningDone { + item_id: reasoning_id, + }, + ) + .await?; + *state = ReasoningState::Done; } + Ok(()) +} - Ok(None) +async fn ensure_message_started( + current_message_id: &mut Option, + next_message_id: &mut Option, + tx_storage: &mpsc::Sender, + tx_client: &mpsc::Sender, +) -> Result { + if let Some(id) = *current_message_id { + return Ok(id); + } + let id = next_assistant_message_id(next_message_id); + send_storage_message( + tx_storage, + tx_client, + StorageMessage::MessageStarted { item_id: id }, + ) + .await?; + *current_message_id = Some(id); + Ok(id) } -/// Phase 6: Setup completion processor -/// -/// Gets completion stream from chat API and spawns processor task. -/// -/// Operations: -/// - Creates placeholder assistant message (AFTER tools, so timestamp is ordered correctly) -/// - Rebuilds prompt from DB if tools were executed (automatically includes tools) -/// - Calls chat API with streaming enabled -/// - Spawns processor task that converts CompletionChunks to StorageMessages -/// - Processor feeds into dual streams (storage=critical, client=best-effort) -/// - Listens for cancellation signals #[allow(clippy::too_many_arguments)] -async fn setup_completion_processor( +async fn stream_one_assistant_turn( state: &Arc, user: &User, body: &ResponsesCreateRequest, - context: &BuiltContext, - prepared: &PreparedRequest, - persisted: &PersistedData, headers: &HeaderMap, - tx_client: mpsc::Sender, - tx_storage: mpsc::Sender, - tools_executed: bool, -) -> Result { - // Create placeholder assistant message with status='in_progress' and NULL content - // - // TIMING: This happens here in Phase 6 (not earlier in Phase 3) for two reasons: - // 1. Must happen AFTER tool execution (Phase 5) to get correct timestamp ordering - // 2. Must happen BEFORE calling completion API (below) so storage task can UPDATE it - // - // Previously, this was created in Phase 3 under the assumption it needed to exist early. - // However, it only needs to exist before the storage task tries to UPDATE it (when - // streaming completes). By creating it here, we ensure proper message ordering: - // user → tool_call → tool_output → assistant (this creation) → assistant content (update) - let placeholder_assistant = NewAssistantMessage { - uuid: prepared.assistant_message_id, - conversation_id: context.conversation.id, - response_id: Some(persisted.response.id), - user_id: user.uuid, - content_enc: None, - completion_tokens: 0, - status: STATUS_IN_PROGRESS.to_string(), - finish_reason: None, - }; - state - .db - .create_assistant_message(placeholder_assistant) - .map_err(|e| { - error!("Error creating placeholder assistant message: {:?}", e); - ApiError::InternalServerError - })?; - - debug!( - "Created placeholder assistant message {} after tool execution", - prepared.assistant_message_id - ); - - // If tools were executed, rebuild prompt from DB (will now include persisted tools) - // Otherwise use the context we built earlier - let prompt_messages = if tools_executed { - debug!("Tools were executed - rebuilding prompt from DB to include tool messages"); - let (rebuilt_messages, _tokens) = build_prompt( - state.db.as_ref(), - context.conversation.id, - user.uuid, - &prepared.user_key, - &body.model, - body.instructions.as_deref(), - )?; - rebuilt_messages - } else { - // Clone out of Arc only when actually needed for the completion request - Arc::as_ref(&context.prompt_messages).clone() - }; - - // Build chat completion request - let mut chat_request = json!({ - "model": body.model, - "messages": prompt_messages, - "temperature": body.temperature.unwrap_or(DEFAULT_TEMPERATURE), - "top_p": body.top_p.unwrap_or(DEFAULT_TOP_P), - "max_tokens": body.max_output_tokens, - "stream": true, - "stream_options": { "include_usage": true } - }); - apply_responses_model_defaults(&mut chat_request, &body.model); + prompt_messages: &[Value], + tools_enabled: bool, + tx_client: &mpsc::Sender, + tx_storage: &mpsc::Sender, + next_message_id: &mut Option, +) -> Result { + let mut chat_request = build_model_turn_request(body, prompt_messages, tools_enabled); - // Log the exact request we're sending to the completions API trace!( "Chat completion request to model {}: {}", body.model, @@ -1705,195 +1978,306 @@ async fn setup_completion_processor( .unwrap_or_else(|_| "failed to serialize".to_string()) ); - // Create billing context - Responses API always uses JWT auth (not API key) let billing_context = crate::web::openai::BillingContext::new( crate::web::openai_auth::AuthMethod::Jwt, body.model.clone(), ); - // Call the chat API - billing happens automatically inside! - let completion = - get_chat_completion_response(state, user, chat_request, headers, billing_context).await?; + let mut completion = + get_chat_completion_response(state, user, chat_request.take(), headers, billing_context) + .await?; debug!( "Received completion from provider: {} (model: {})", completion.metadata.provider_name, completion.metadata.model_name ); - // Signal that assistant message is about to start streaming - // CRITICAL: Must send BEFORE spawning processor to guarantee ordering - // (processor will immediately start sending ContentDelta messages) - if let Err(e) = tx_client - .send(StorageMessage::AssistantMessageStarting) - .await - { - error!("Failed to send AssistantMessageStarting signal: {:?}", e); - // Client channel closed - not critical, continue anyway - } + let mut streamed_tool_calls = Vec::new(); + let mut finish_reason: Option = None; + let mut current_message_id: Option = None; + let mut reasoning = ReasoningState::NotStarted; + let mut saw_tool_calls = false; + + while let Some(chunk) = completion.stream.recv().await { + match chunk { + crate::web::openai::CompletionChunk::StreamChunk(json) => { + let choice = json.get("choices").and_then(|choices| choices.get(0)); + if let Some(reason) = choice + .and_then(|choice| choice.get("finish_reason")) + .and_then(|reason| reason.as_str()) + { + if !reason.is_empty() { + finish_reason = Some(reason.to_string()); + } + } - // Spawn stream processor task that converts CompletionChunks to StorageMessages - // and feeds them into the master stream channels (created in Phase 3.5) - let _processor_handle = { - let mut rx_completion = completion.stream; - let message_id = prepared.assistant_message_id; - let response_uuid = persisted.response.uuid; - let mut cancel_rx = state.cancellation_broadcast.subscribe(); + let delta = choice.and_then(|choice| choice.get("delta")); + if let Some(d) = delta { + trace!("Stream delta: {}", d); + } - tokio::spawn(async move { - trace!("Starting completion stream processor task"); - let mut client_alive = true; - // Single source of truth for reasoning item UUID - generated on first reasoning delta - let mut reasoning_item_id: Option = None; - - loop { - tokio::select! { - // Check for cancellation - Ok(cancelled_id) = cancel_rx.recv() => { - if cancelled_id == response_uuid { - debug!("Received cancellation signal for response {}", response_uuid); - let _ = tx_storage.send(StorageMessage::Cancelled).await; - if client_alive && tx_client.try_send(StorageMessage::Cancelled).is_err() { - warn!("Client channel full or closed during cancellation, terminating client stream"); - #[allow(unused_assignments)] - { client_alive = false; } + // Reasoning delta (supports both `reasoning` and legacy `reasoning_content`). + if let Some(reasoning_delta) = delta.and_then(|d| { + d.get("reasoning") + .and_then(|c| c.as_str()) + .or_else(|| d.get("reasoning_content").and_then(|c| c.as_str())) + }) { + if !reasoning_delta.is_empty() { + match &reasoning { + ReasoningState::Done => { + warn!( + "Ignoring reasoning delta after reasoning item was already closed" + ); + } + ReasoningState::NotStarted => { + let reasoning_id = Uuid::new_v4(); + send_storage_message( + tx_storage, + tx_client, + StorageMessage::ReasoningStarted { + item_id: reasoning_id, + }, + ) + .await?; + send_storage_message( + tx_storage, + tx_client, + StorageMessage::ReasoningDelta { + item_id: reasoning_id, + delta: reasoning_delta.to_string(), + }, + ) + .await?; + reasoning = ReasoningState::Active(reasoning_id); + } + ReasoningState::Active(reasoning_id) => { + send_storage_message( + tx_storage, + tx_client, + StorageMessage::ReasoningDelta { + item_id: *reasoning_id, + delta: reasoning_delta.to_string(), + }, + ) + .await?; } - break; } } + } - // Process CompletionChunks from centralized billing API - chunk_opt = rx_completion.recv() => { - let Some(chunk) = chunk_opt else { - error!("Completion stream closed unexpectedly without Done signal"); - // Explicitly notify storage and client of the failure - let msg = StorageMessage::Error("Stream closed unexpectedly".to_string()); - let _ = tx_storage.send(msg.clone()).await; - if client_alive && tx_client.try_send(msg).is_err() { - warn!("Client channel full or closed during stream error, terminating client stream"); - #[allow(unused_assignments)] - { client_alive = false; } - } - break; - }; - - match chunk { - crate::web::openai::CompletionChunk::StreamChunk(json) => { - // Extract delta for inspection - let delta = json - .get("choices") - .and_then(|c| c.get(0)) - .and_then(|c| c.get("delta")); - - // Log full delta to see reasoning fields if present - if let Some(d) = delta { - trace!("Stream delta: {}", d); - } + // Tool call deltas: close any in-flight message and reasoning, then accumulate. + if let Some(tool_call_delta) = delta.and_then(|d| d.get("tool_calls")) { + if let Some(message_id) = current_message_id.take() { + send_storage_message( + tx_storage, + tx_client, + StorageMessage::MessageDone { + item_id: message_id, + finish_reason: "tool_calls".to_string(), + }, + ) + .await?; + } + close_reasoning_if_active(&mut reasoning, tx_storage, tx_client).await?; - // Extract reasoning for thinking models while remaining - // compatible with older reasoning_content payloads. - if let Some(reasoning) = delta - .and_then(|d| { - d.get("reasoning") - .and_then(|c| c.as_str()) - .or_else(|| { - d.get("reasoning_content") - .and_then(|c| c.as_str()) - }) - }) - { - if !reasoning.is_empty() { - // Generate UUID on first reasoning delta - this ensures SSE and DB use the same ID - let item_id = *reasoning_item_id.get_or_insert_with(Uuid::new_v4); - - let msg = StorageMessage::ReasoningDelta { - item_id, - delta: reasoning.to_string(), - }; - // Send to storage for persistence - if tx_storage.send(msg.clone()).await.is_err() { - error!("Storage channel closed unexpectedly during reasoning"); - break; - } - // Send to client for streaming - if client_alive && tx_client.try_send(msg).is_err() { - warn!("Client channel full or closed during reasoning delta, terminating client stream"); - client_alive = false; - } - } - } + saw_tool_calls = true; + append_streamed_tool_calls(&mut streamed_tool_calls, tool_call_delta); + } - // Extract content from the full JSON chunk (safe chaining to avoid panics) - if let Some(content) = delta - .and_then(|d| d.get("content")) - .and_then(|c| c.as_str()) - { - // Skip empty content deltas to avoid sending unnecessary events to client - if !content.is_empty() { - let msg = StorageMessage::ContentDelta(content.to_string()); - // Must send to storage (critical, can block) - if tx_storage.send(msg.clone()).await.is_err() { - error!("Storage channel closed unexpectedly"); - break; - } - // Best-effort send to client (non-blocking, never blocks storage) - if client_alive && tx_client.try_send(msg).is_err() { - warn!("Client channel full or closed, terminating client stream"); - client_alive = false; - } - } - } - } - crate::web::openai::CompletionChunk::Usage(usage) => { - // Billing already happened in openai.rs! - // Just forward to storage for token counting - let msg = StorageMessage::Usage { - prompt_tokens: usage.prompt_tokens, - completion_tokens: usage.completion_tokens, - }; - let _ = tx_storage.send(msg.clone()).await; - if client_alive && tx_client.try_send(msg).is_err() { - warn!("Client channel full or closed during usage message, terminating client stream"); - client_alive = false; - } - } - crate::web::openai::CompletionChunk::Done => { - debug!("Received Done chunk from completion stream"); - let msg = StorageMessage::Done { - finish_reason: "stop".to_string(), - message_id, - }; - let _ = tx_storage.send(msg.clone()).await; - if client_alive && tx_client.try_send(msg).is_err() { - warn!("Client channel full or closed during done message, terminating client stream"); - #[allow(unused_assignments)] - { client_alive = false; } - } - break; - } - crate::web::openai::CompletionChunk::Error(error_msg) => { - error!("Received error from completion stream: {}", error_msg); - let msg = StorageMessage::Error(error_msg); - let _ = tx_storage.send(msg.clone()).await; - if client_alive && tx_client.try_send(msg).is_err() { - warn!("Client channel full or closed during error message, terminating client stream"); - #[allow(unused_assignments)] - { client_alive = false; } - } - break; - } - crate::web::openai::CompletionChunk::FullResponse(_) => { - // Shouldn't happen for streaming - error!("Received FullResponse in streaming mode"); - break; - } + // Assistant text content. Some providers occasionally emit a trailing sliver of + // content after tool-call deltas; we log and skip instead of aborting the turn. + if let Some(content) = delta + .and_then(|d| d.get("content")) + .and_then(|c| c.as_str()) + { + if !content.is_empty() { + if saw_tool_calls { + warn!( + "Ignoring assistant text ({} chars) after tool call deltas had already started", + content.len() + ); + } else { + close_reasoning_if_active(&mut reasoning, tx_storage, tx_client) + .await?; + + let message_id = ensure_message_started( + &mut current_message_id, + next_message_id, + tx_storage, + tx_client, + ) + .await?; + + send_storage_message( + tx_storage, + tx_client, + StorageMessage::ContentDelta { + item_id: message_id, + delta: content.to_string(), + }, + ) + .await?; } } } } + crate::web::openai::CompletionChunk::Usage(usage) => { + send_storage_message( + tx_storage, + tx_client, + StorageMessage::Usage { + prompt_tokens: usage.prompt_tokens, + completion_tokens: usage.completion_tokens, + }, + ) + .await?; + } + crate::web::openai::CompletionChunk::Done => { + close_reasoning_if_active(&mut reasoning, tx_storage, tx_client).await?; - debug!("Completion stream processor task completed"); - }) - }; + if saw_tool_calls || finish_reason.as_deref() == Some("tool_calls") { + let tool_call = finalize_first_model_tool_call(&streamed_tool_calls) + .ok_or(ApiError::InternalServerError)?; + return Ok(AssistantTurnOutcome::ToolCall(tool_call)); + } + + // Ensure a final message item exists even if the model emitted no content. + // This is rare but happens when vLLM parses a tool response into the thinking + // block or similar edge cases; we surface it explicitly rather than silently + // producing an empty row. + if current_message_id.is_none() { + warn!( + "Model produced no assistant content before Done; emitting empty message placeholder" + ); + } + let message_id = ensure_message_started( + &mut current_message_id, + next_message_id, + tx_storage, + tx_client, + ) + .await?; + + let final_finish_reason = finish_reason.unwrap_or_else(|| "stop".to_string()); + send_storage_message( + tx_storage, + tx_client, + StorageMessage::MessageDone { + item_id: message_id, + finish_reason: final_finish_reason.clone(), + }, + ) + .await?; + + send_storage_message( + tx_storage, + tx_client, + StorageMessage::ResponseDone { + finish_reason: final_finish_reason, + }, + ) + .await?; + return Ok(AssistantTurnOutcome::Final); + } + crate::web::openai::CompletionChunk::Error(error_msg) => { + error!("Received error from completion stream: {}", error_msg); + return Err(ApiError::InternalServerError); + } + crate::web::openai::CompletionChunk::FullResponse(_) => { + error!("Received FullResponse in streaming mode"); + return Err(ApiError::InternalServerError); + } + } + } + + error!("Completion stream closed unexpectedly without a terminal signal"); + Err(ApiError::InternalServerError) +} + +/// Phase 6: Run the normal assistant/tool loop. +/// +/// The selected model receives the full prompt with tool schemas, may call one +/// tool at a time, sees the tool output on the next turn, and eventually emits +/// final assistant text that is streamed to the client. +#[allow(clippy::too_many_arguments)] +async fn setup_completion_processor( + state: &Arc, + user: &User, + body: &ResponsesCreateRequest, + context: &BuiltContext, + prepared: &PreparedRequest, + persisted: &PersistedData, + headers: &HeaderMap, + tx_client: mpsc::Sender, + tx_storage: mpsc::Sender, + mut rx_tool_ack: mpsc::Receiver>, +) -> Result { + const MAX_TOOL_TURNS: usize = 15; + + let tools_enabled = should_enable_web_search_tool(state.as_ref(), body); + let internal_system_prompt = build_internal_system_prompt(tools_enabled); + let mut prompt_messages = Arc::as_ref(&context.prompt_messages).clone(); + + let loop_result: Result<(), ApiError> = async { + let mut next_message_id = Some(prepared.assistant_message_id); + let mut tool_turn_count = 0usize; + loop { + match stream_one_assistant_turn( + state, + user, + body, + headers, + &prompt_messages, + tools_enabled, + &tx_client, + &tx_storage, + &mut next_message_id, + ) + .await? + { + AssistantTurnOutcome::ToolCall(tool_call) => { + tool_turn_count += 1; + if tool_turn_count > MAX_TOOL_TURNS { + error!( + "Exceeded max tool turns ({}) for response {}", + MAX_TOOL_TURNS, persisted.response.uuid + ); + return Err(ApiError::InternalServerError); + } + + execute_tool_call_and_wait( + state, + persisted, + tool_call, + &tx_client, + &tx_storage, + &mut rx_tool_ack, + ) + .await?; + + let (rebuilt_messages, _tokens) = build_prompt( + state.db.as_ref(), + context.conversation.id, + user.uuid, + &prepared.user_key, + &body.model, + body.instructions.as_deref(), + Some(&internal_system_prompt), + )?; + prompt_messages = rebuilt_messages; + } + AssistantTurnOutcome::Final => return Ok(()), + } + } + } + .await; + + if let Err(e) = loop_result { + let msg = StorageMessage::Error(format!("Streaming failed: {:?}", e)); + let _ = tx_storage.send(msg.clone()).await; + let _ = tx_client.try_send(msg); + return Err(e); + } Ok(persisted.response.clone()) } @@ -1956,10 +2340,9 @@ async fn create_response_stream( let total_prompt_tokens = context.total_prompt_tokens; let response_id = persisted.response.id; let response_uuid = persisted.response.uuid; - // Use user_message timestamp + 1 microsecond for reasoning items to ensure proper ordering: - // user_message < reasoning < assistant_message - // Using 1 microsecond because it's the smallest safe offset - DB roundtrips always take more than this - let reasoning_base_timestamp = + // Persist all generated response items on a single monotonic timestamp sequence that + // begins immediately after the user message so retrieval order matches stream order. + let first_response_item_created_at = persisted.user_message_created_at + chrono::Duration::microseconds(1); let conversation_id = context.conversation.id; let user_id = user.uuid; @@ -2006,8 +2389,8 @@ async fn create_response_stream( let (tx_storage, rx_storage) = mpsc::channel::(STORAGE_CHANNEL_BUFFER); let (tx_client, mut rx_client) = mpsc::channel::(CLIENT_CHANNEL_BUFFER); - // Create oneshot channel for tool persistence acknowledgment - let (tx_tool_ack, rx_tool_ack) = tokio::sync::oneshot::channel(); + // Create channel for tool persistence acknowledgments (supports multiple tool loops) + let (tx_tool_ack, rx_tool_ack) = mpsc::channel::>(8); let _storage_handle = { let db = state.db.clone(); @@ -2018,11 +2401,10 @@ async fn create_response_stream( Some(tx_tool_ack), db, response_id, - reasoning_base_timestamp, + first_response_item_created_at, conversation_id, user_id, user_key, - assistant_message_id, ) .await; }) @@ -2051,62 +2433,8 @@ async fn create_response_stream( // Run phases 5-6 with cancellation support tokio::select! { _ = async { - // Phase 5: Classify intent and execute tools (if tool_choice allows it AND web_search is enabled AND Kagi client available) - let tools_executed = if is_tool_choice_allowed(&orchestrator_body.tool_choice) - && is_web_search_enabled(&orchestrator_body.tools) - && orchestrator_state.kagi_client.is_some() { - debug!("Orchestrator: tool_choice allows tools, web search enabled, and Kagi client available, proceeding with classification"); - - let prepared_for_tools = PreparedRequest { - user_key, - message_content: message_content.clone(), - user_message_tokens: 0, - content_enc: content_enc.clone(), - assistant_message_id, - }; - - let persisted_for_tools = PersistedData { - response: orchestrator_response.clone(), - decrypted_metadata: orchestrator_metadata.clone(), - user_message_created_at: orchestrator_user_message_created_at, - }; - - match classify_and_execute_tools( - &orchestrator_state, - &orchestrator_user, - &prepared_for_tools, - &persisted_for_tools, - &orchestrator_prompt_messages, - ToolChannels { - client: &orchestrator_tx_client, - storage: &orchestrator_tx_storage, - }, - rx_tool_ack, - ) - .await - { - Ok(result) => result.is_some(), - Err(e) => { - error!("Orchestrator: Critical error during tool execution, treating as cancellation: {:?}", e); - - // Treat critical errors (storage failure) same as cancellation - // Send cancellation to both channels - storage task will handle cleanup - let _ = orchestrator_tx_storage.send(StorageMessage::Cancelled).await; - let _ = orchestrator_tx_client.send(StorageMessage::Cancelled).await; - - // Abort orchestrator - don't waste resources on LLM call - // Storage task will update response status and assistant message - return; - } - } - } else { - debug!("Orchestrator: Web search tool not enabled or Kagi client not available, skipping classification"); - drop(rx_tool_ack); - false - }; - - // Phase 6: Setup completion processor - trace!("Orchestrator: Setting up completion processor"); + // Phase 5-6: Run the normal assistant/tool loop + trace!("Orchestrator: Setting up assistant/tool loop"); let context_for_completion = BuiltContext { conversation: orchestrator_conversation, @@ -2138,42 +2466,15 @@ async fn create_response_stream( &orchestrator_headers, orchestrator_tx_client.clone(), orchestrator_tx_storage.clone(), - tools_executed, + rx_tool_ack, ) .await { Ok(_) => { - trace!("Orchestrator: Completion processor setup complete"); - // AssistantMessageStarting is now sent from inside setup_completion_processor - // to guarantee it arrives before any completion deltas + trace!("Orchestrator: Assistant/tool loop completed"); } Err(e) => { - error!("Orchestrator: Failed to setup completion processor: {:?}", e); - - // Update response status to failed - if let Err(db_err) = orchestrator_state.db.update_response_status( - response_id, - ResponseStatus::Failed, - Some(Utc::now()), - ) { - error!("Orchestrator: Failed to update response status: {:?}", db_err); - } - - // Update assistant message to incomplete - if let Err(db_err) = orchestrator_state.db.update_assistant_message( - assistant_message_id, - None, - 0, - STATUS_INCOMPLETE.to_string(), - None, - ) { - error!("Orchestrator: Failed to update assistant message: {:?}", db_err); - } - - // Send error to client via channel (best-effort) - let _ = orchestrator_tx_client.try_send(StorageMessage::Error( - format!("Failed to setup streaming: {:?}", e) - )); + error!("Orchestrator: Assistant/tool loop failed: {:?}", e); } } } => { @@ -2196,164 +2497,181 @@ async fn create_response_stream( // NOW immediately start the event loop - it will receive events from orchestrator as they happen trace!("Starting event loop to receive messages from background tasks"); - let mut assistant_content = String::new(); - let mut reasoning_content = String::new(); - let mut reasoning_done_emitted = false; - let mut reasoning_item_started = false; - let mut message_item_started = false; - // Reasoning item ID comes from the message - same UUID used for SSE and DB - let mut reasoning_item_id: Option = None; + let mut client_state = ClientResponseState::default(); + let mut total_prompt_tokens_used = 0i32; let mut total_completion_tokens = 0i32; while let Some(msg) = rx_client.recv().await { trace!("Client stream received message from upstream processor"); match msg { - StorageMessage::Done { finish_reason, message_id: msg_id } => { - trace!("Client stream received Done signal with finish_reason={}, message_id={}", finish_reason, msg_id); - // Note: msg_id should match our pre-generated assistant_message_id - - // Calculate message output_index based on whether reasoning exists - let message_output_index = if reasoning_item_started { 1 } else { 0 }; - - // Fallback: emit reasoning done events if not already done (edge case) - if !reasoning_done_emitted && !reasoning_content.is_empty() { - if let Some(rid) = reasoning_item_id { - // No need to set reasoning_done_emitted = true since we break after this block - - let reasoning_done_event = ResponseReasoningTextDoneEvent { - event_type: EVENT_RESPONSE_REASONING_TEXT_DONE, - sequence_number: emitter.sequence_number(), - item_id: rid.to_string(), - output_index: 0, - content_index: 0, - text: reasoning_content.clone(), - }; - yield Ok(ResponseEvent::ReasoningTextDone(reasoning_done_event).to_sse_event(&mut emitter).await); - - let reasoning_item_done = ResponseOutputItemDoneEvent { - event_type: EVENT_RESPONSE_OUTPUT_ITEM_DONE, - sequence_number: emitter.sequence_number(), - output_index: 0, - item: OutputItem { - id: rid.to_string(), - output_type: "reasoning".to_string(), - status: STATUS_COMPLETED.to_string(), - role: None, - content: Some(vec![]), - }, - }; - yield Ok(ResponseEvent::OutputItemDone(reasoning_item_done).to_sse_event(&mut emitter).await); - } - } + StorageMessage::MessageStarted { item_id } => { + let output_index = client_state.push_message(item_id); + let output_item_added_event = ResponseOutputItemAddedEvent { + event_type: EVENT_RESPONSE_OUTPUT_ITEM_ADDED, + sequence_number: emitter.sequence_number(), + output_index, + item: OutputItemBuilder::new_message(item_id).build(), + }; + yield Ok(ResponseEvent::OutputItemAdded(output_item_added_event).to_sse_event(&mut emitter).await); - // Emit message start events if not yet started (edge case: no content deltas) - if !message_item_started { - // No need to set message_item_started = true since we break after this block - - let output_item_added_event = ResponseOutputItemAddedEvent { - event_type: EVENT_RESPONSE_OUTPUT_ITEM_ADDED, - sequence_number: emitter.sequence_number(), - output_index: message_output_index, - item: OutputItem { - id: assistant_message_id.to_string(), - output_type: OUTPUT_TYPE_MESSAGE.to_string(), - status: STATUS_IN_PROGRESS.to_string(), - role: Some(ROLE_ASSISTANT.to_string()), - content: Some(vec![]), - }, - }; - yield Ok(ResponseEvent::OutputItemAdded(output_item_added_event).to_sse_event(&mut emitter).await); - - let content_part_added_event = ResponseContentPartAddedEvent { - event_type: EVENT_RESPONSE_CONTENT_PART_ADDED, - sequence_number: emitter.sequence_number(), - item_id: assistant_message_id.to_string(), - output_index: message_output_index, - content_index: 0, - part: ContentPart { - part_type: CONTENT_PART_TYPE_OUTPUT_TEXT.to_string(), - annotations: vec![], - logprobs: vec![], - text: String::new(), - }, - }; - yield Ok(ResponseEvent::ContentPartAdded(content_part_added_event).to_sse_event(&mut emitter).await); - } + let content_part_added_event = ResponseContentPartAddedEvent { + event_type: EVENT_RESPONSE_CONTENT_PART_ADDED, + sequence_number: emitter.sequence_number(), + item_id: item_id.to_string(), + output_index, + content_index: 0, + part: ContentPart { + part_type: CONTENT_PART_TYPE_OUTPUT_TEXT.to_string(), + annotations: vec![], + logprobs: vec![], + text: String::new(), + }, + }; + yield Ok(ResponseEvent::ContentPartAdded(content_part_added_event).to_sse_event(&mut emitter).await); + } + StorageMessage::ContentDelta { item_id, delta } => { + trace!("Client stream received content delta: {}", delta); + let Some(output_index) = client_state.append_message_delta(item_id, &delta) else { + warn!("Received content delta for unknown message item {}", item_id); + continue; + }; + + let delta_event = ResponseOutputTextDeltaEvent { + event_type: EVENT_RESPONSE_OUTPUT_TEXT_DELTA, + delta, + item_id: item_id.to_string(), + output_index, + content_index: 0, + sequence_number: emitter.sequence_number(), + logprobs: vec![], + }; + + yield Ok(ResponseEvent::OutputTextDelta(delta_event).to_sse_event(&mut emitter).await); + } + StorageMessage::MessageDone { item_id, .. } => { + let Some(output_index) = client_state.message_output_index(item_id) else { + warn!("Received message done for unknown item {}", item_id); + continue; + }; + let text = client_state.message_text(item_id).unwrap_or("").to_string(); - // response.output_text.done let output_text_done_event = ResponseOutputTextDoneEvent { event_type: EVENT_RESPONSE_OUTPUT_TEXT_DONE, sequence_number: emitter.sequence_number(), - item_id: assistant_message_id.to_string(), - output_index: message_output_index, + item_id: item_id.to_string(), + output_index, content_index: 0, - text: assistant_content.clone(), + text: text.clone(), logprobs: vec![], }; yield Ok(ResponseEvent::OutputTextDone(output_text_done_event).to_sse_event(&mut emitter).await); - // response.content_part.done let content_part_done_event = ResponseContentPartDoneEvent { event_type: EVENT_RESPONSE_CONTENT_PART_DONE, sequence_number: emitter.sequence_number(), - item_id: assistant_message_id.to_string(), - output_index: message_output_index, + item_id: item_id.to_string(), + output_index, content_index: 0, - part: ContentPart { - part_type: CONTENT_PART_TYPE_OUTPUT_TEXT.to_string(), - annotations: vec![], - logprobs: vec![], - text: assistant_content.clone(), - }, + part: ContentPartBuilder::new_output_text(text.clone()).build(), }; yield Ok(ResponseEvent::ContentPartDone(content_part_done_event).to_sse_event(&mut emitter).await); - // response.output_item.done for message - let message_content_part = ContentPart { - part_type: CONTENT_PART_TYPE_OUTPUT_TEXT.to_string(), - annotations: vec![], - logprobs: vec![], - text: assistant_content.clone(), - }; - let output_item_done_event = ResponseOutputItemDoneEvent { event_type: EVENT_RESPONSE_OUTPUT_ITEM_DONE, sequence_number: emitter.sequence_number(), - output_index: message_output_index, + output_index, + item: OutputItemBuilder::new_message(item_id) + .status(STATUS_COMPLETED) + .content(vec![ContentPartBuilder::new_output_text(text).build()]) + .build(), + }; + yield Ok(ResponseEvent::OutputItemDone(output_item_done_event).to_sse_event(&mut emitter).await); + } + StorageMessage::ReasoningStarted { item_id } => { + let output_index = client_state.push_reasoning(item_id); + let reasoning_item_added = ResponseOutputItemAddedEvent { + event_type: EVENT_RESPONSE_OUTPUT_ITEM_ADDED, + sequence_number: emitter.sequence_number(), + output_index, item: OutputItem { - id: assistant_message_id.to_string(), - output_type: OUTPUT_TYPE_MESSAGE.to_string(), - status: STATUS_COMPLETED.to_string(), - role: Some(ROLE_ASSISTANT.to_string()), - content: Some(vec![message_content_part]), + id: item_id.to_string(), + output_type: "reasoning".to_string(), + status: STATUS_IN_PROGRESS.to_string(), + role: None, + content: Some(vec![]), + call_id: None, + name: None, + arguments: None, + output: None, }, }; - yield Ok(ResponseEvent::OutputItemDone(output_item_done_event).to_sse_event(&mut emitter).await); + yield Ok(ResponseEvent::OutputItemAdded(reasoning_item_added).to_sse_event(&mut emitter).await); + } + StorageMessage::ReasoningDelta { item_id, delta } => { + trace!("Client stream received reasoning delta: {}", delta); + let Some(output_index) = client_state.append_reasoning_delta(item_id, &delta) else { + warn!("Received reasoning delta for unknown item {}", item_id); + continue; + }; - // response.completed - build output array with reasoning + message if reasoning exists - let content_part = ContentPartBuilder::new_output_text(assistant_content.clone()).build(); - let message_item = OutputItemBuilder::new_message(assistant_message_id) - .status(STATUS_COMPLETED) - .content(vec![content_part]) - .build(); + let delta_event = ResponseReasoningTextDeltaEvent { + event_type: EVENT_RESPONSE_REASONING_TEXT_DELTA, + delta, + item_id: item_id.to_string(), + output_index, + content_index: 0, + sequence_number: emitter.sequence_number(), + }; - let output_items = if let Some(rid) = reasoning_item_id { - // Include reasoning item in output - let reasoning_item = OutputItem { - id: rid.to_string(), + yield Ok(ResponseEvent::ReasoningTextDelta(delta_event).to_sse_event(&mut emitter).await); + } + StorageMessage::ReasoningDone { item_id } => { + let Some(output_index) = client_state.reasoning_output_index(item_id) else { + warn!("Received reasoning done for unknown item {}", item_id); + continue; + }; + let text = client_state.reasoning_text(item_id).unwrap_or("").to_string(); + + let reasoning_done_event = ResponseReasoningTextDoneEvent { + event_type: EVENT_RESPONSE_REASONING_TEXT_DONE, + sequence_number: emitter.sequence_number(), + item_id: item_id.to_string(), + output_index, + content_index: 0, + text, + }; + yield Ok(ResponseEvent::ReasoningTextDone(reasoning_done_event).to_sse_event(&mut emitter).await); + + let reasoning_item_done = ResponseOutputItemDoneEvent { + event_type: EVENT_RESPONSE_OUTPUT_ITEM_DONE, + sequence_number: emitter.sequence_number(), + output_index, + item: OutputItem { + id: item_id.to_string(), output_type: "reasoning".to_string(), status: STATUS_COMPLETED.to_string(), role: None, content: Some(vec![]), - }; - vec![reasoning_item, message_item] - } else { - vec![message_item] + call_id: None, + name: None, + arguments: None, + output: None, + }, }; + yield Ok(ResponseEvent::OutputItemDone(reasoning_item_done).to_sse_event(&mut emitter).await); + } + StorageMessage::ResponseDone { finish_reason: _finish_reason } => { + let usage = build_usage( + if total_prompt_tokens_used > 0 { + total_prompt_tokens_used + } else { + total_prompt_tokens as i32 + }, + total_completion_tokens, + ); - let usage = build_usage(total_prompt_tokens as i32, total_completion_tokens); let done_response = ResponseBuilder::from_response(&response_for_stream) .status(STATUS_COMPLETED) - .output(output_items) + .output(client_state.build_output_items()) .usage(usage) .metadata(decrypted_metadata.clone()) .build(); @@ -2367,143 +2685,10 @@ async fn create_response_stream( yield Ok(ResponseEvent::Completed(completed_event).to_sse_event(&mut emitter).await); break; } - StorageMessage::ContentDelta(content) => { - trace!("Client stream received content delta: {}", content); - - // Calculate message output_index based on whether reasoning exists - let message_output_index = if reasoning_item_started { 1 } else { 0 }; - - // Complete reasoning phase when first content arrives - if !reasoning_done_emitted && !reasoning_content.is_empty() { - if let Some(rid) = reasoning_item_id { - reasoning_done_emitted = true; - - // response.reasoning_text.done - let reasoning_done_event = ResponseReasoningTextDoneEvent { - event_type: EVENT_RESPONSE_REASONING_TEXT_DONE, - sequence_number: emitter.sequence_number(), - item_id: rid.to_string(), - output_index: 0, - content_index: 0, - text: reasoning_content.clone(), - }; - yield Ok(ResponseEvent::ReasoningTextDone(reasoning_done_event).to_sse_event(&mut emitter).await); - - // response.output_item.done for reasoning - let reasoning_item_done = ResponseOutputItemDoneEvent { - event_type: EVENT_RESPONSE_OUTPUT_ITEM_DONE, - sequence_number: emitter.sequence_number(), - output_index: 0, - item: OutputItem { - id: rid.to_string(), - output_type: "reasoning".to_string(), - status: STATUS_COMPLETED.to_string(), - role: None, - content: Some(vec![]), // Reasoning content is in the text field, not content array - }, - }; - yield Ok(ResponseEvent::OutputItemDone(reasoning_item_done).to_sse_event(&mut emitter).await); - } - } - - // Emit message start events if not yet started - if !message_item_started { - message_item_started = true; - - // response.output_item.added for message - let output_item_added_event = ResponseOutputItemAddedEvent { - event_type: EVENT_RESPONSE_OUTPUT_ITEM_ADDED, - sequence_number: emitter.sequence_number(), - output_index: message_output_index, - item: OutputItem { - id: assistant_message_id.to_string(), - output_type: OUTPUT_TYPE_MESSAGE.to_string(), - status: STATUS_IN_PROGRESS.to_string(), - role: Some(ROLE_ASSISTANT.to_string()), - content: Some(vec![]), - }, - }; - yield Ok(ResponseEvent::OutputItemAdded(output_item_added_event).to_sse_event(&mut emitter).await); - - // response.content_part.added - let content_part_added_event = ResponseContentPartAddedEvent { - event_type: EVENT_RESPONSE_CONTENT_PART_ADDED, - sequence_number: emitter.sequence_number(), - item_id: assistant_message_id.to_string(), - output_index: message_output_index, - content_index: 0, - part: ContentPart { - part_type: CONTENT_PART_TYPE_OUTPUT_TEXT.to_string(), - annotations: vec![], - logprobs: vec![], - text: String::new(), - }, - }; - yield Ok(ResponseEvent::ContentPartAdded(content_part_added_event).to_sse_event(&mut emitter).await); - } - - assistant_content.push_str(&content); - - // Send content delta to client - let delta_event = ResponseOutputTextDeltaEvent { - event_type: EVENT_RESPONSE_OUTPUT_TEXT_DELTA, - delta: content.clone(), - item_id: assistant_message_id.to_string(), - output_index: message_output_index, - content_index: 0, - sequence_number: emitter.sequence_number(), - logprobs: vec![], - }; - - yield Ok(ResponseEvent::OutputTextDelta(delta_event).to_sse_event(&mut emitter).await); - } - StorageMessage::ReasoningDelta { item_id, delta: reasoning } => { - trace!("Client stream received reasoning delta: {}", reasoning); - - // Store the item_id from the message (same UUID used for SSE and DB) - if reasoning_item_id.is_none() { - reasoning_item_id = Some(item_id); - } - let rid = reasoning_item_id.expect("reasoning_item_id should be set"); - - // Emit reasoning item start events on first reasoning delta - if !reasoning_item_started { - reasoning_item_started = true; - - // response.output_item.added for reasoning (output_index: 0) - let reasoning_item_added = ResponseOutputItemAddedEvent { - event_type: EVENT_RESPONSE_OUTPUT_ITEM_ADDED, - sequence_number: emitter.sequence_number(), - output_index: 0, - item: OutputItem { - id: rid.to_string(), - output_type: "reasoning".to_string(), - status: STATUS_IN_PROGRESS.to_string(), - role: None, - content: Some(vec![]), - }, - }; - yield Ok(ResponseEvent::OutputItemAdded(reasoning_item_added).to_sse_event(&mut emitter).await); - } - - reasoning_content.push_str(&reasoning); - - // Send reasoning delta to client - let delta_event = ResponseReasoningTextDeltaEvent { - event_type: EVENT_RESPONSE_REASONING_TEXT_DELTA, - delta: reasoning.clone(), - item_id: rid.to_string(), - output_index: 0, - content_index: 0, - sequence_number: emitter.sequence_number(), - }; - - yield Ok(ResponseEvent::ReasoningTextDelta(delta_event).to_sse_event(&mut emitter).await); - } - - StorageMessage::Usage { prompt_tokens: _, completion_tokens } => { + StorageMessage::Usage { prompt_tokens, completion_tokens } => { trace!("Client stream received usage data"); - total_completion_tokens = completion_tokens; + total_prompt_tokens_used += prompt_tokens; + total_completion_tokens += completion_tokens; } StorageMessage::Cancelled => { debug!("Client stream received cancellation signal"); @@ -2536,35 +2721,118 @@ async fn create_response_stream( } StorageMessage::ToolCall { tool_call_id, name, arguments } => { debug!("Client stream received tool_call event: {} ({})", name, tool_call_id); + let tool_name = name.clone(); + let arguments_json = + serde_json::to_string(&arguments).unwrap_or_else(|_| "{}".to_string()); + let output_index = + client_state.push_tool_call(tool_call_id, tool_name.clone(), arguments.clone()); + + let output_item_added_event = ResponseOutputItemAddedEvent { + event_type: EVENT_RESPONSE_OUTPUT_ITEM_ADDED, + sequence_number: emitter.sequence_number(), + output_index, + item: OutputItem { + id: tool_call_id.to_string(), + output_type: "tool_call".to_string(), + status: STATUS_IN_PROGRESS.to_string(), + role: None, + content: None, + call_id: Some(tool_call_id.to_string()), + name: Some(tool_name.clone()), + arguments: Some(arguments_json.clone()), + output: None, + }, + }; + + yield Ok(ResponseEvent::OutputItemAdded(output_item_added_event).to_sse_event(&mut emitter).await); + // Send tool_call.created event let tool_call_event = ToolCallCreatedEvent { event_type: EVENT_TOOL_CALL_CREATED, sequence_number: emitter.sequence_number(), + output_index, tool_call_id, name, arguments, }; yield Ok(ResponseEvent::ToolCallCreated(tool_call_event).to_sse_event(&mut emitter).await); + + let output_item_done_event = ResponseOutputItemDoneEvent { + event_type: EVENT_RESPONSE_OUTPUT_ITEM_DONE, + sequence_number: emitter.sequence_number(), + output_index, + item: OutputItem { + id: tool_call_id.to_string(), + output_type: "tool_call".to_string(), + status: STATUS_COMPLETED.to_string(), + role: None, + content: None, + call_id: Some(tool_call_id.to_string()), + name: Some(tool_name), + arguments: Some(arguments_json), + output: None, + }, + }; + + yield Ok(ResponseEvent::OutputItemDone(output_item_done_event).to_sse_event(&mut emitter).await); } StorageMessage::ToolOutput { tool_output_id, tool_call_id, output } => { debug!("Client stream received tool_output event: {}", tool_output_id); + let output_index = client_state.push_tool_output( + tool_output_id, + tool_call_id, + output.clone(), + ); + let output_item_added_event = ResponseOutputItemAddedEvent { + event_type: EVENT_RESPONSE_OUTPUT_ITEM_ADDED, + sequence_number: emitter.sequence_number(), + output_index, + item: OutputItem { + id: tool_output_id.to_string(), + output_type: "tool_output".to_string(), + status: STATUS_IN_PROGRESS.to_string(), + role: None, + content: None, + call_id: Some(tool_call_id.to_string()), + name: None, + arguments: None, + output: Some(output.clone()), + }, + }; + + yield Ok(ResponseEvent::OutputItemAdded(output_item_added_event).to_sse_event(&mut emitter).await); + // Send tool_output.created event let tool_output_event = ToolOutputCreatedEvent { event_type: EVENT_TOOL_OUTPUT_CREATED, sequence_number: emitter.sequence_number(), + output_index, tool_output_id, tool_call_id, - output, + output: output.clone(), }; yield Ok(ResponseEvent::ToolOutputCreated(tool_output_event).to_sse_event(&mut emitter).await); - } - StorageMessage::AssistantMessageStarting => { - // Don't emit output_item.added here - we defer it until we know if there's reasoning - // This allows us to emit reasoning item at output_index 0 if present, - // then message item at output_index 1 (or 0 if no reasoning) - debug!("Client stream received assistant message starting signal (events deferred)"); + + let output_item_done_event = ResponseOutputItemDoneEvent { + event_type: EVENT_RESPONSE_OUTPUT_ITEM_DONE, + sequence_number: emitter.sequence_number(), + output_index, + item: OutputItem { + id: tool_output_id.to_string(), + output_type: "tool_output".to_string(), + status: STATUS_COMPLETED.to_string(), + role: None, + content: None, + call_id: Some(tool_call_id.to_string()), + name: None, + arguments: None, + output: Some(output), + }, + }; + + yield Ok(ResponseEvent::OutputItemDone(output_item_done_event).to_sse_event(&mut emitter).await); } } } @@ -2625,39 +2893,88 @@ async fn get_response( .await .map_err(|_| error_mapping::map_key_retrieval_error())?; - // Build output from assistant messages only (user messages are input, not output) - // TODO: Add tool_call and tool_output items to output array once we implement tool support - // Need to determine correct OpenAI format for tool items in output array let mut output_items = Vec::new(); for msg in &messages { - // Only include assistant messages in output - // TODO: When adding tool support, also handle msg.message_type == "tool_call" and "tool_output" - if msg.message_type != "assistant" { - continue; - } - - // Only include messages that have content - if let Some(content_enc) = &msg.content_enc { - if let Some(text) = decrypt_string(&user_key, Some(content_enc)).map_err(|e| { - error!("Failed to decrypt assistant message content: {:?}", e); - error_mapping::map_decryption_error("assistant message content") - })? { - // Build content part using builder - let content_part = ContentPartBuilder::new_output_text(text).build(); - - // Build output item using builder - let output_item = OutputItemBuilder::new_message(msg.uuid) - .status( - &msg.status - .clone() - .unwrap_or_else(|| STATUS_COMPLETED.to_string()), - ) - .content(vec![content_part]) - .build(); - + let status = msg + .status + .clone() + .unwrap_or_else(|| STATUS_COMPLETED.to_string()); + + match msg.message_type.as_str() { + "assistant" => { + let text = decrypt_string(&user_key, msg.content_enc.as_ref()).map_err(|e| { + error!("Failed to decrypt assistant message content: {:?}", e); + error_mapping::map_decryption_error("assistant message content") + })?; + + let output_item = if let Some(text) = text { + OutputItemBuilder::new_message(msg.uuid) + .status(&status) + .content(vec![ContentPartBuilder::new_output_text(text).build()]) + .build() + } else { + OutputItemBuilder::new_message(msg.uuid) + .status(&status) + .build() + }; output_items.push(output_item); } + "tool_call" => { + let arguments = + decrypt_string(&user_key, msg.content_enc.as_ref()).map_err(|e| { + error!("Failed to decrypt tool call arguments: {:?}", e); + error_mapping::map_decryption_error("tool call arguments") + })?; + + output_items.push(OutputItem { + id: msg.uuid.to_string(), + output_type: "tool_call".to_string(), + status, + role: None, + content: None, + call_id: Some(msg.tool_call_id.unwrap_or(msg.uuid).to_string()), + name: Some( + msg.tool_name + .clone() + .unwrap_or_else(|| "function".to_string()), + ), + arguments, + output: None, + }); + } + "tool_output" => { + let output = decrypt_string(&user_key, msg.content_enc.as_ref()).map_err(|e| { + error!("Failed to decrypt tool output: {:?}", e); + error_mapping::map_decryption_error("tool output") + })?; + + output_items.push(OutputItem { + id: msg.uuid.to_string(), + output_type: "tool_output".to_string(), + status, + role: None, + content: None, + call_id: msg.tool_call_id.map(|id| id.to_string()), + name: None, + arguments: None, + output, + }); + } + "reasoning" => { + output_items.push(OutputItem { + id: msg.uuid.to_string(), + output_type: "reasoning".to_string(), + status, + role: None, + content: Some(vec![]), + call_id: None, + name: None, + arguments: None, + output: None, + }); + } + _ => {} } } @@ -2666,13 +2983,17 @@ async fn get_response( // Sum up tokens from all messages let mut input_tokens = 0i32; let mut output_tokens = 0i32; + let mut reasoning_tokens = 0i32; for msg in &messages { if let Some(token_count) = msg.token_count { match msg.message_type.as_str() { "user" => input_tokens += token_count, "assistant" => output_tokens += token_count, - // tool_call and tool_output tokens are also considered in context + "reasoning" => { + output_tokens += token_count; + reasoning_tokens += token_count; + } "tool_call" => input_tokens += token_count, "tool_output" => input_tokens += token_count, _ => {} @@ -2684,9 +3005,7 @@ async fn get_response( input_tokens, input_tokens_details: InputTokenDetails { cached_tokens: 0 }, output_tokens, - output_tokens_details: OutputTokenDetails { - reasoning_tokens: 0, - }, + output_tokens_details: OutputTokenDetails { reasoning_tokens }, total_tokens: input_tokens + output_tokens, }) } else { diff --git a/src/web/responses/mod.rs b/src/web/responses/mod.rs index 4fea4424..2775d723 100644 --- a/src/web/responses/mod.rs +++ b/src/web/responses/mod.rs @@ -14,7 +14,6 @@ pub mod events; pub mod handlers; pub mod instructions; pub mod pagination; -pub mod prompts; pub mod storage; pub mod tools; pub mod types; diff --git a/src/web/responses/prompts.rs b/src/web/responses/prompts.rs deleted file mode 100644 index 6d4923ff..00000000 --- a/src/web/responses/prompts.rs +++ /dev/null @@ -1,320 +0,0 @@ -//! Prompt templates for the Responses API -//! -//! This module contains all prompt templates used for intent classification, -//! query extraction, and other AI-driven features of the Responses API. - -use serde_json::{json, Value}; - -pub const INTENT_CLASSIFIER_MODEL: &str = "gpt-oss-120b"; -pub const INTENT_CLASSIFIER_MAX_TOKENS: i32 = 250; -pub const QUERY_EXTRACTOR_MODEL: &str = "llama3-3-70b"; -pub const QUERY_EXTRACTOR_MAX_TOKENS: i32 = 50; - -/// System prompt for intent classification -/// -/// This prompt instructs the LLM to classify whether a user's message requires -/// web search or can be handled as a regular chat conversation. -pub const INTENT_CLASSIFIER_PROMPT: &str = "\ -Classify whether to search the web before responding: \"web_search\" or \"chat\". Return \"web_search\" if web search is needed, \"chat\" if you can answer directly. \ -Pay special attention to the context of the conversation history and the nature of the inquiry. If the user asks for current events, factual information, or specific data that may have changed recently, classify as \"web_search\". \ -For casual conversations or subjective inquiries, classify as \"chat\". Identify inquiries that may seem casual but actually require factual information, such as popular culture references or specific names that are currently relevant. \ -Recognize when a casual inquiry is part of a broader context that may require factual information, and adjust the classification accordingly."; - -/// System prompt for search query extraction -/// -/// This prompt instructs the LLM to extract a clean search query from the user's -/// natural language question, using conversation history for context. -pub const SEARCH_QUERY_EXTRACTOR_PROMPT: &str = "\ -Extract the main search query from the user's question. Use the conversation history to understand context and references. -Return only the search terms, nothing else. Be concise and specific. -If the user's question refers to something mentioned earlier in the conversation, include that context in your search query. - -Examples: -- \"What's the weather in San Francisco today?\" → weather San Francisco today -- \"Who is the current president of the United States?\" → current president United States -- \"Tell me about the latest SpaceX launch\" → latest SpaceX launch -- After discussing \"iPhone 15\", user asks \"when was it released?\" → iPhone 15 release date"; - -/// Build a chat completion request for intent classification -/// -/// Uses a fast, cheap model (gpt-oss-120b) with temperature=0 for deterministic results. -/// -/// # Arguments -/// * `conversation_history` - Recent conversation messages (will use last 6, truncated to 200 chars each) -/// * `user_message` - The current user's message to classify -/// -/// # Returns -/// A JSON request ready to be sent to `get_chat_completion_response` -pub fn build_intent_classification_request( - conversation_history: &[Value], - user_message: &str, -) -> Value { - // Format conversation history as text for context - let history_text = if !conversation_history.is_empty() { - let formatted_messages: Vec = conversation_history - .iter() - .rev() - .take(6) - .rev() - .filter_map(|msg| { - let role = msg.get("role")?.as_str()?; - let content = extract_text_from_content(msg.get("content")?); - let truncated_content: String = content.chars().take(200).collect(); - Some(format!("{}: {}", role, truncated_content)) - }) - .collect(); - - if formatted_messages.is_empty() { - String::new() - } else { - format!( - "Conversation history:\n{}\n\n", - formatted_messages.join("\n") - ) - } - } else { - String::new() - }; - - // Build single user message with history + current query - let user_prompt = format!("{}Current user query: {}", history_text, user_message); - - json!({ - "model": INTENT_CLASSIFIER_MODEL, - "messages": [ - { - "role": "system", - "content": INTENT_CLASSIFIER_PROMPT - }, - { - "role": "user", - "content": user_prompt - } - ], - "temperature": 0.0, - "max_tokens": INTENT_CLASSIFIER_MAX_TOKENS, - "stream": false - }) -} - -/// Helper function to extract text from content (handles both string and array formats) -fn extract_text_from_content(content: &Value) -> String { - match content { - Value::String(s) => s.clone(), - Value::Array(arr) => arr - .iter() - .filter_map(|part| { - part.get("text") - .and_then(|t| t.as_str()) - .map(|s| s.to_string()) - }) - .collect::>() - .join(" "), - _ => String::new(), - } -} - -/// Build a chat completion request for search query extraction -/// -/// Uses llama3-3-70b to extract a concise search query with a tighter output limit. -/// -/// # Arguments -/// * `conversation_history` - Recent conversation messages for context -/// * `user_message` - The user's message to extract a query from -/// -/// # Returns -/// A JSON request ready to be sent to `get_chat_completion_response` -pub fn build_query_extraction_request(conversation_history: &[Value], user_message: &str) -> Value { - // Format conversation history as text for context - let history_text = if !conversation_history.is_empty() { - let formatted_messages: Vec = conversation_history - .iter() - .rev() - .take(6) - .rev() - .filter_map(|msg| { - let role = msg.get("role")?.as_str()?; - let content = extract_text_from_content(msg.get("content")?); - let truncated_content: String = content.chars().take(200).collect(); - Some(format!("{}: {}", role, truncated_content)) - }) - .collect(); - - if formatted_messages.is_empty() { - String::new() - } else { - format!( - "Conversation history:\n{}\n\n", - formatted_messages.join("\n") - ) - } - } else { - String::new() - }; - - // Build single user message with history + current query - let user_prompt = format!("{}Current user question: {}", history_text, user_message); - - json!({ - "model": QUERY_EXTRACTOR_MODEL, - "messages": [ - { - "role": "system", - "content": SEARCH_QUERY_EXTRACTOR_PROMPT - }, - { - "role": "user", - "content": user_prompt - } - ], - "temperature": 0.0, - "max_tokens": QUERY_EXTRACTOR_MAX_TOKENS, - "stream": false - }) -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_build_intent_classification_request() { - let request = build_intent_classification_request(&[], "What's the weather?"); - - assert_eq!(request["model"], INTENT_CLASSIFIER_MODEL); - assert_eq!(request["temperature"], 0.0); - assert_eq!(request["stream"], false); - - let messages = request["messages"].as_array().unwrap(); - assert_eq!(messages.len(), 2); - assert_eq!(messages[0]["role"], "system"); - assert_eq!(messages[1]["role"], "user"); - // Should have formatted user query - let user_content = messages[1]["content"].as_str().unwrap(); - assert_eq!(user_content, "Current user query: What's the weather?"); - } - - #[test] - fn test_build_intent_classification_with_history() { - let history = vec![ - json!({"role": "user", "content": "Hello"}), - json!({"role": "assistant", "content": "Hi there!"}), - json!({"role": "user", "content": "How are you?"}), - ]; - - let request = build_intent_classification_request(&history, "What's the weather?"); - - let messages = request["messages"].as_array().unwrap(); - // Should have only 2 messages: system + single user message with history - assert_eq!(messages.len(), 2); - assert_eq!(messages[0]["role"], "system"); - assert_eq!(messages[1]["role"], "user"); - - let user_content = messages[1]["content"].as_str().unwrap(); - // Should contain formatted history - assert!(user_content.contains("Conversation history:")); - assert!(user_content.contains("user: Hello")); - assert!(user_content.contains("assistant: Hi there!")); - assert!(user_content.contains("user: How are you?")); - assert!(user_content.contains("Current user query: What's the weather?")); - } - - #[test] - fn test_build_intent_classification_truncates_long_messages() { - let long_text = "a".repeat(500); - let history = vec![json!({"role": "user", "content": long_text})]; - - let request = build_intent_classification_request(&history, "test"); - - let messages = request["messages"].as_array().unwrap(); - let user_content = messages[1]["content"].as_str().unwrap(); - - // History should be truncated to 200 chars per message - // Format is: "Conversation history:\nuser: <200 chars>\n\nCurrent user query: test" - assert!(user_content.contains("Conversation history:")); - assert!(user_content.contains("user: ")); - assert!(user_content.contains("Current user query: test")); - - // Extract just the history part to verify truncation - let history_part = user_content.split("Current user query:").next().unwrap(); - // The 'aaa...' part should be truncated (much less than original 500 chars) - let a_count = history_part.chars().filter(|&c| c == 'a').count(); - // Should be around 200 (within reasonable bounds, accounting for any formatting) - assert!( - (190..=210).contains(&a_count), - "Expected around 200 'a' characters, got {}", - a_count - ); - // Definitely should be much less than the original 500 - assert!( - a_count < 300, - "Truncation failed: got {} 'a' characters (original was 500)", - a_count - ); - } - - #[test] - fn test_build_intent_classification_limits_to_6_messages() { - let history: Vec = (0..10) - .map(|i| json!({"role": "user", "content": format!("Message {}", i)})) - .collect(); - - let request = build_intent_classification_request(&history, "test"); - - let messages = request["messages"].as_array().unwrap(); - // Should have only 2 messages: system + single user message - assert_eq!(messages.len(), 2); - - let user_content = messages[1]["content"].as_str().unwrap(); - // Should have last 6 messages (4-9) - assert!(user_content.contains("user: Message 4")); - assert!(user_content.contains("user: Message 9")); - // Should not have earlier messages - assert!(!user_content.contains("user: Message 0")); - assert!(!user_content.contains("user: Message 3")); - } - - #[test] - fn test_build_query_extraction_request() { - let request = build_query_extraction_request(&[], "What's the weather in New York?"); - - assert_eq!(request["model"], QUERY_EXTRACTOR_MODEL); - assert_eq!(request["max_tokens"], QUERY_EXTRACTOR_MAX_TOKENS); - - let messages = request["messages"].as_array().unwrap(); - assert_eq!(messages.len(), 2); - - let user_content = messages[1]["content"].as_str().unwrap(); - assert_eq!( - user_content, - "Current user question: What's the weather in New York?" - ); - } - - #[test] - fn test_build_query_extraction_with_context() { - let history = vec![ - json!({"role": "user", "content": "Tell me about iPhone 15"}), - json!({"role": "assistant", "content": "The iPhone 15 was announced in September 2023..."}), - ]; - - let request = build_query_extraction_request(&history, "when was it released?"); - - let messages = request["messages"].as_array().unwrap(); - assert_eq!(messages.len(), 2); - - let user_content = messages[1]["content"].as_str().unwrap(); - assert!(user_content.contains("Conversation history:")); - assert!(user_content.contains("iPhone 15")); - assert!(user_content.contains("Current user question: when was it released?")); - } - - #[test] - fn test_prompts_contain_examples() { - assert!(INTENT_CLASSIFIER_PROMPT.contains("web_search")); - assert!(INTENT_CLASSIFIER_PROMPT.contains("chat")); - assert!(INTENT_CLASSIFIER_PROMPT.contains("conversation history")); - - assert!(SEARCH_QUERY_EXTRACTOR_PROMPT.contains("Examples:")); - } -} diff --git a/src/web/responses/storage.rs b/src/web/responses/storage.rs index 1f1a919e..e026359a 100644 --- a/src/web/responses/storage.rs +++ b/src/web/responses/storage.rs @@ -1,682 +1,540 @@ -//! Storage task components for accumulating and persisting streaming responses +//! Storage task components for persisting streaming response items. use crate::{ encrypt::encrypt_with_key, - models::responses::ResponseStatus, + models::responses::{ + NewAssistantMessage, NewReasoningItem, NewToolCall, NewToolOutput, ResponseStatus, + }, tokens::count_tokens, - web::responses::constants::{FINISH_REASON_CANCELLED, STATUS_COMPLETED, STATUS_INCOMPLETE}, + web::responses::constants::{ + FINISH_REASON_CANCELLED, STATUS_COMPLETED, STATUS_INCOMPLETE, STATUS_IN_PROGRESS, + }, DBConnection, }; use chrono::Utc; use secp256k1::SecretKey; -use std::sync::Arc; +use std::{collections::HashMap, sync::Arc}; use tokio::sync::mpsc; use tracing::{debug, error, trace, warn}; use uuid::Uuid; use super::handlers::StorageMessage; -/// Accumulates streaming content and metadata -pub(crate) struct ContentAccumulator { +#[derive(Default)] +struct PendingAssistantMessage { content: String, - reasoning_content: String, - reasoning_item_id: Option, - completion_tokens: i32, + created_at: Option>, } -impl ContentAccumulator { - pub fn new() -> Self { - Self { - content: String::with_capacity(4096), - reasoning_content: String::with_capacity(4096), - reasoning_item_id: None, - completion_tokens: 0, - } - } - - /// Get the reasoning item ID if one was created - pub fn reasoning_item_id(&self) -> Option { - self.reasoning_item_id - } +#[derive(Default)] +struct PendingReasoningItem { + content: String, +} - /// Get accumulated reasoning content - pub fn reasoning_content(&self) -> &str { - &self.reasoning_content +fn clamp_token_count(text: &str, label: &str) -> i32 { + let token_count = count_tokens(text); + if token_count > i32::MAX as usize { + warn!( + "{} token count {} exceeds i32::MAX, clamping", + label, token_count + ); + i32::MAX + } else { + token_count as i32 } +} - /// Handle a storage message and return the accumulator state - pub fn handle_message(&mut self, msg: StorageMessage) -> AccumulatorState { - match msg { - StorageMessage::ContentDelta(delta) => { - trace!("Storage: received content delta: {} chars", delta.len()); - self.content.push_str(&delta); - AccumulatorState::Continue - } - StorageMessage::ReasoningDelta { item_id, delta } => { - trace!("Storage: received reasoning delta: {} chars", delta.len()); - self.reasoning_content.push_str(&delta); - - // Use the item_id from the message - same UUID as SSE events - if self.reasoning_item_id.is_none() { - self.reasoning_item_id = Some(item_id); - return AccumulatorState::CreateReasoningItem { item_id }; - } +fn allocate_created_at( + next_created_at: &mut chrono::DateTime, +) -> chrono::DateTime { + let created_at = next_created_at.to_owned(); + *next_created_at = created_at + chrono::Duration::microseconds(1); + created_at +} - AccumulatorState::Continue - } - StorageMessage::Usage { - prompt_tokens, - completion_tokens, - } => { - debug!( - "Storage: received usage - prompt_tokens={}, completion_tokens={}", - prompt_tokens, completion_tokens - ); - // Note: prompt_tokens already tracked when user message was created in DB - // Only store completion_tokens for the assistant message - self.completion_tokens = completion_tokens; - AccumulatorState::Continue - } - StorageMessage::Done { - finish_reason, - message_id, - } => { - debug!( - "Storage: received Done signal with finish_reason={}, message_id={}", - finish_reason, message_id - ); - AccumulatorState::Complete(CompleteData { - content: self.content.clone(), - completion_tokens: self.completion_tokens, - finish_reason, - message_id, - reasoning_content: self.reasoning_content.clone(), - reasoning_item_id: self.reasoning_item_id, - }) - } - StorageMessage::Cancelled => { - debug!("Storage: received cancellation signal"); - AccumulatorState::Cancelled(PartialData { - content: self.content.clone(), - completion_tokens: self.completion_tokens, - reasoning_content: self.reasoning_content.clone(), - reasoning_item_id: self.reasoning_item_id, - }) - } - StorageMessage::Error(e) => { - error!("Storage: received error: {}", e); - AccumulatorState::Failed(FailureData { - error: e, - partial_content: self.content.clone(), - completion_tokens: self.completion_tokens, - reasoning_content: self.reasoning_content.clone(), - reasoning_item_id: self.reasoning_item_id, - }) - } - StorageMessage::ToolCall { - tool_call_id, - name, - arguments, - } => { - trace!( - "Storage: received tool_call - id={}, name={}", - tool_call_id, - name - ); - // Signal immediate persistence - AccumulatorState::PersistToolCall { - tool_call_id, - name, - arguments, - } - } - StorageMessage::ToolOutput { - tool_output_id, - tool_call_id, - output, - } => { - trace!( - "Storage: received tool_output - id={}, tool_call_id={}, output_len={}", - tool_output_id, - tool_call_id, - output.len() - ); - // Signal immediate persistence - AccumulatorState::PersistToolOutput { - tool_output_id, - tool_call_id, - output, - } - } - StorageMessage::AssistantMessageStarting => { - trace!("Storage: received assistant message starting signal (no-op for storage)"); - // This is a signal for the client stream only, storage doesn't need to act on it - AccumulatorState::Continue - } - } +fn create_assistant_message_if_missing( + db: &Arc, + conversation_id: i64, + response_id: i64, + user_id: Uuid, + item_id: Uuid, + created_at: chrono::DateTime, +) -> Result<(), String> { + match db.get_assistant_message_by_uuid(item_id) { + Ok(Some(_)) => Ok(()), + Ok(None) => db + .create_assistant_message(NewAssistantMessage { + uuid: item_id, + conversation_id, + response_id: Some(response_id), + user_id, + content_enc: None, + completion_tokens: 0, + status: STATUS_IN_PROGRESS.to_string(), + finish_reason: None, + created_at, + }) + .map(|_| ()) + .map_err(|e| format!("Failed to create assistant message: {:?}", e)), + Err(e) => Err(format!("Failed to look up assistant message: {:?}", e)), } } -/// State transitions for the accumulator -pub enum AccumulatorState { - Continue, - Complete(CompleteData), - Cancelled(PartialData), - Failed(FailureData), - CreateReasoningItem { - item_id: Uuid, - }, - PersistToolCall { - tool_call_id: Uuid, - name: String, - arguments: serde_json::Value, - }, - PersistToolOutput { - tool_output_id: Uuid, - tool_call_id: Uuid, - output: String, - }, +async fn finalize_assistant_message( + db: &Arc, + user_key: &SecretKey, + item_id: Uuid, + content: String, + status: &str, + finish_reason: Option, +) -> Result<(), String> { + let content_enc = if content.is_empty() { + None + } else { + Some(encrypt_with_key(user_key, content.as_bytes()).await) + }; + + db.update_assistant_message( + item_id, + content_enc, + clamp_token_count(&content, "assistant message"), + status.to_string(), + finish_reason, + ) + .map(|_| ()) + .map_err(|e| format!("Failed to update assistant message: {:?}", e)) } -/// Data for a completed response -pub struct CompleteData { - pub content: String, - pub completion_tokens: i32, - pub finish_reason: String, - pub message_id: Uuid, - pub reasoning_content: String, - pub reasoning_item_id: Option, +fn create_reasoning_item( + db: &Arc, + conversation_id: i64, + response_id: i64, + user_id: Uuid, + item_id: Uuid, + created_at: chrono::DateTime, +) -> Result<(), String> { + db.create_reasoning_item(NewReasoningItem { + uuid: item_id, + conversation_id, + response_id: Some(response_id), + assistant_message_id: None, + user_id, + content_enc: None, + summary_enc: None, + reasoning_tokens: 0, + status: STATUS_IN_PROGRESS.to_string(), + created_at, + }) + .map(|_| ()) + .map_err(|e| format!("Failed to create reasoning item: {:?}", e)) } -/// Data for a partial/cancelled response -pub struct PartialData { - pub content: String, - pub completion_tokens: i32, - pub reasoning_content: String, - pub reasoning_item_id: Option, +async fn finalize_reasoning_item( + db: &Arc, + user_key: &SecretKey, + item_id: Uuid, + content: String, + status: &str, +) -> Result<(), String> { + let content_enc = if content.is_empty() { + None + } else { + Some(encrypt_with_key(user_key, content.as_bytes()).await) + }; + + db.update_reasoning_item( + item_id, + content_enc, + clamp_token_count(&content, "reasoning"), + status.to_string(), + ) + .map(|_| ()) + .map_err(|e| format!("Failed to update reasoning item: {:?}", e)) } -/// Data for a failed response -pub struct FailureData { - pub error: String, - pub partial_content: String, - pub completion_tokens: i32, - pub reasoning_content: String, - pub reasoning_item_id: Option, +#[allow(clippy::too_many_arguments)] +async fn persist_tool_call( + db: &Arc, + user_key: &SecretKey, + conversation_id: i64, + response_id: i64, + user_id: Uuid, + tool_call_id: Uuid, + name: String, + arguments: serde_json::Value, + created_at: chrono::DateTime, +) -> Result<(), String> { + let arguments_json = serde_json::to_string(&arguments) + .map_err(|e| format!("Failed to serialize tool arguments: {:?}", e))?; + let arguments_enc = encrypt_with_key(user_key, arguments_json.as_bytes()).await; + + db.create_tool_call(NewToolCall { + uuid: tool_call_id, + conversation_id, + response_id: Some(response_id), + user_id, + name, + arguments_enc: Some(arguments_enc), + argument_tokens: clamp_token_count(&arguments_json, "tool arguments"), + status: STATUS_COMPLETED.to_string(), + created_at, + }) + .map(|_| ()) + .map_err(|e| format!("Failed to persist tool_call: {:?}", e)) } -/// Handles persistence of responses in various states -pub(crate) struct ResponsePersister { - db: Arc, +#[allow(clippy::too_many_arguments)] +async fn persist_tool_output( + db: &Arc, + user_key: &SecretKey, + conversation_id: i64, response_id: i64, - message_id: Uuid, - user_key: SecretKey, + user_id: Uuid, + tool_output_id: Uuid, + tool_call_id: Uuid, + output: String, + created_at: chrono::DateTime, +) -> Result<(), String> { + let tool_call_fk = db + .get_tool_call_by_uuid(tool_call_id, user_id) + .map_err(|e| format!("Tool call not found in database: {:?}", e))? + .id; + let output_enc = encrypt_with_key(user_key, output.as_bytes()).await; + + db.create_tool_output(NewToolOutput { + uuid: tool_output_id, + conversation_id, + response_id: Some(response_id), + user_id, + tool_call_fk, + output_enc, + output_tokens: clamp_token_count(&output, "tool output"), + status: STATUS_COMPLETED.to_string(), + error: None, + created_at, + }) + .map(|_| ()) + .map_err(|e| format!("Failed to persist tool_output: {:?}", e)) } -impl ResponsePersister { - pub fn new( - db: Arc, - response_id: i64, - message_id: Uuid, - user_key: SecretKey, - ) -> Self { - Self { +async fn mark_pending_items_incomplete( + db: &Arc, + user_key: &SecretKey, + pending_messages: &mut HashMap, + pending_reasoning: &mut HashMap, + message_finish_reason: Option, +) { + for (item_id, pending) in pending_messages.drain() { + if let Err(e) = finalize_assistant_message( db, - response_id, - message_id, user_key, - } - } - - /// Persist a completed response - pub async fn persist_completed(&self, data: CompleteData) -> Result<(), String> { - // Fallback token counting if not provided - let completion_tokens = if data.completion_tokens == 0 && !data.content.is_empty() { - let token_count = count_tokens(&data.content); - if token_count > i32::MAX as usize { - warn!( - "Completion token count {} exceeds i32::MAX, clamping to i32::MAX", - token_count - ); - i32::MAX - } else { - token_count as i32 - } - } else { - data.completion_tokens - }; - - // Encrypt and store assistant message - let content_enc = encrypt_with_key(&self.user_key, data.content.as_bytes()).await; - - if let Err(e) = self.db.update_assistant_message( - data.message_id, - Some(content_enc), - completion_tokens, - STATUS_COMPLETED.to_string(), - Some(data.finish_reason), - ) { - error!("Failed to update assistant message: {:?}", e); - return Err(format!("Failed to update assistant message: {:?}", e)); - } - - // Update response status - if let Err(e) = self.db.update_response_status( - self.response_id, - ResponseStatus::Completed, - Some(Utc::now()), - ) { - error!("Failed to update response status to completed: {:?}", e); - return Err(format!("Failed to update response status: {:?}", e)); - } - - // Update reasoning item if present - if let Some(reasoning_item_id) = data.reasoning_item_id { - if !data.reasoning_content.is_empty() { - let reasoning_enc = - encrypt_with_key(&self.user_key, data.reasoning_content.as_bytes()).await; - let reasoning_tokens = count_tokens(&data.reasoning_content); - let reasoning_tokens_i32 = if reasoning_tokens > i32::MAX as usize { - warn!( - "Reasoning token count {} exceeds i32::MAX, clamping", - reasoning_tokens - ); - i32::MAX - } else { - reasoning_tokens as i32 - }; - - if let Err(e) = self.db.update_reasoning_item( - reasoning_item_id, - Some(reasoning_enc), - reasoning_tokens_i32, - STATUS_COMPLETED.to_string(), - ) { - error!("Failed to update reasoning item: {:?}", e); - // Non-fatal: assistant message was saved successfully - } - } - } - - debug!("Successfully persisted completed response"); - Ok(()) - } - - /// Persist a cancelled response - pub async fn persist_cancelled(&self, data: PartialData) -> Result<(), String> { - // Update response status - if let Err(e) = self.db.update_response_status( - self.response_id, - ResponseStatus::Cancelled, - Some(Utc::now()), - ) { - error!("Failed to update response status to cancelled: {:?}", e); - return Err(format!("Failed to update response status: {:?}", e)); - } - - // Update assistant message to incomplete status with partial content - let content_enc = if !data.content.is_empty() { - Some(encrypt_with_key(&self.user_key, data.content.as_bytes()).await) - } else { - None - }; - - if let Err(e) = self.db.update_assistant_message( - self.message_id, - content_enc, - data.completion_tokens, - STATUS_INCOMPLETE.to_string(), - Some(FINISH_REASON_CANCELLED.to_string()), - ) { + item_id, + pending.content, + STATUS_INCOMPLETE, + message_finish_reason.clone(), + ) + .await + { error!( - "Failed to update assistant message after cancellation: {:?}", - e + "Failed to finalize pending assistant message {}: {}", + item_id, e ); - return Err(format!("Failed to update assistant message: {:?}", e)); - } - - // Update reasoning item to incomplete if present - if let Some(reasoning_item_id) = data.reasoning_item_id { - if !data.reasoning_content.is_empty() { - let reasoning_enc = - encrypt_with_key(&self.user_key, data.reasoning_content.as_bytes()).await; - let reasoning_tokens = count_tokens(&data.reasoning_content); - let reasoning_tokens_i32 = if reasoning_tokens > i32::MAX as usize { - warn!( - "Reasoning token count {} exceeds i32::MAX, clamping", - reasoning_tokens - ); - i32::MAX - } else { - reasoning_tokens as i32 - }; - - if let Err(e) = self.db.update_reasoning_item( - reasoning_item_id, - Some(reasoning_enc), - reasoning_tokens_i32, - STATUS_INCOMPLETE.to_string(), - ) { - error!( - "Failed to update reasoning item after cancellation: {:?}", - e - ); - // Non-fatal - } - } } - - debug!( - "Persisted cancelled response {} with {} tokens", - self.response_id, data.completion_tokens - ); - Ok(()) } - /// Persist a failed response - pub async fn persist_failed(&self, data: FailureData) -> Result<(), String> { - // Update response status - if let Err(e) = self.db.update_response_status( - self.response_id, - ResponseStatus::Failed, - Some(Utc::now()), - ) { - error!("Failed to update response status to failed: {:?}", e); - return Err(format!("Failed to update response status: {:?}", e)); - } - - // Update assistant message to incomplete status with partial content - let content_enc = if !data.partial_content.is_empty() { - Some(encrypt_with_key(&self.user_key, data.partial_content.as_bytes()).await) - } else { - None - }; - - if let Err(e) = self.db.update_assistant_message( - self.message_id, - content_enc, - data.completion_tokens, - STATUS_INCOMPLETE.to_string(), - None, - ) { - error!("Failed to update assistant message to incomplete: {:?}", e); - return Err(format!("Failed to update assistant message: {:?}", e)); - } - - // Update reasoning item to incomplete if present - if let Some(reasoning_item_id) = data.reasoning_item_id { - if !data.reasoning_content.is_empty() { - let reasoning_enc = - encrypt_with_key(&self.user_key, data.reasoning_content.as_bytes()).await; - let reasoning_tokens = count_tokens(&data.reasoning_content); - let reasoning_tokens_i32 = if reasoning_tokens > i32::MAX as usize { - warn!( - "Reasoning token count {} exceeds i32::MAX, clamping", - reasoning_tokens - ); - i32::MAX - } else { - reasoning_tokens as i32 - }; - - if let Err(e) = self.db.update_reasoning_item( - reasoning_item_id, - Some(reasoning_enc), - reasoning_tokens_i32, - STATUS_INCOMPLETE.to_string(), - ) { - error!("Failed to update reasoning item after failure: {:?}", e); - // Non-fatal - } - } + for (item_id, pending) in pending_reasoning.drain() { + if let Err(e) = + finalize_reasoning_item(db, user_key, item_id, pending.content, STATUS_INCOMPLETE).await + { + error!( + "Failed to finalize pending reasoning item {}: {}", + item_id, e + ); } - - debug!("Persisted failed response with error: {}", data.error); - Ok(()) } } -/// Main storage task that orchestrates accumulation and persistence +/// Main storage task that orchestrates per-item persistence. #[allow(clippy::too_many_arguments)] pub async fn storage_task( mut rx: mpsc::Receiver, - tool_persist_ack: Option>>, + tool_persist_ack: Option>>, db: Arc, response_id: i64, - reasoning_base_timestamp: chrono::DateTime, + first_item_created_at: chrono::DateTime, conversation_id: i64, user_id: Uuid, user_key: SecretKey, - message_id: Uuid, ) { - let mut accumulator = ContentAccumulator::new(); - let persister = ResponsePersister::new(db.clone(), response_id, message_id, user_key); - - // Track tool acknowledgment channel - let mut tool_ack = tool_persist_ack; + let tool_ack = tool_persist_ack; + let mut pending_messages: HashMap = HashMap::new(); + let mut pending_reasoning: HashMap = HashMap::new(); + let mut next_item_created_at = first_item_created_at; - // Accumulate messages until completion or error while let Some(msg) = rx.recv().await { - match accumulator.handle_message(msg) { - AccumulatorState::Continue => continue, - - AccumulatorState::CreateReasoningItem { item_id } => { - // Create reasoning item record with in_progress status - use crate::models::responses::NewReasoningItem; - - // Look up assistant_message database ID from UUID - // The assistant_message should exist by now (created before streaming starts) - let assistant_message_db_id = match db.get_assistant_message_by_uuid(message_id) { - Ok(Some(am)) => Some(am.id), - Ok(None) => { - warn!( - "Assistant message {} not found when creating reasoning item", - message_id - ); - None - } - Err(e) => { - warn!( - "Failed to look up assistant message {}: {:?}", - message_id, e - ); - None - } - }; - - let new_reasoning_item = NewReasoningItem { - uuid: item_id, + match msg { + StorageMessage::MessageStarted { item_id } => { + trace!("Storage: message started {}", item_id); + let pending = pending_messages.entry(item_id).or_default(); + let created_at = pending + .created_at + .get_or_insert_with(|| allocate_created_at(&mut next_item_created_at)) + .to_owned(); + if let Err(e) = create_assistant_message_if_missing( + &db, conversation_id, - response_id: Some(response_id), - assistant_message_id: assistant_message_db_id, + response_id, user_id, - content_enc: None, - summary_enc: None, - reasoning_tokens: 0, - status: "in_progress".to_string(), - created_at: reasoning_base_timestamp, - }; - - match db.create_reasoning_item(new_reasoning_item) { - Ok(reasoning_item) => { - debug!( - "Created reasoning item {} (db id: {}, assistant_message_id: {:?})", - item_id, reasoning_item.id, assistant_message_db_id - ); - } - Err(e) => { - error!("Failed to create reasoning item {}: {:?}", item_id, e); - // Non-fatal: reasoning will be lost but response can continue - } + item_id, + created_at, + ) { + error!("{}", e); } } - - AccumulatorState::PersistToolCall { + StorageMessage::ContentDelta { item_id, delta } => { + trace!( + "Storage: content delta for {} ({} chars)", + item_id, + delta.len() + ); + pending_messages + .entry(item_id) + .or_default() + .content + .push_str(&delta); + } + StorageMessage::MessageDone { + item_id, + finish_reason, + } => { + debug!( + "Storage: message done {} with finish_reason={}", + item_id, finish_reason + ); + let pending = pending_messages.remove(&item_id).unwrap_or_default(); + let created_at = pending + .created_at + .unwrap_or_else(|| allocate_created_at(&mut next_item_created_at)); + if let Err(e) = create_assistant_message_if_missing( + &db, + conversation_id, + response_id, + user_id, + item_id, + created_at, + ) { + error!("{}", e); + } + if let Err(e) = finalize_assistant_message( + &db, + &user_key, + item_id, + pending.content, + STATUS_COMPLETED, + Some(finish_reason), + ) + .await + { + error!("{}", e); + } + } + StorageMessage::ReasoningStarted { item_id } => { + trace!("Storage: reasoning started {}", item_id); + pending_reasoning.entry(item_id).or_default(); + let created_at = allocate_created_at(&mut next_item_created_at); + if let Err(e) = create_reasoning_item( + &db, + conversation_id, + response_id, + user_id, + item_id, + created_at, + ) { + error!("{}", e); + } + } + StorageMessage::ReasoningDelta { item_id, delta } => { + trace!( + "Storage: reasoning delta for {} ({} chars)", + item_id, + delta.len() + ); + pending_reasoning + .entry(item_id) + .or_default() + .content + .push_str(&delta); + } + StorageMessage::ReasoningDone { item_id } => { + debug!("Storage: reasoning done {}", item_id); + let pending = pending_reasoning.remove(&item_id).unwrap_or_default(); + if let Err(e) = finalize_reasoning_item( + &db, + &user_key, + item_id, + pending.content, + STATUS_COMPLETED, + ) + .await + { + error!("{}", e); + } + } + StorageMessage::Usage { .. } => { + trace!("Storage: usage message ignored for item persistence"); + } + StorageMessage::ResponseDone { finish_reason } => { + debug!( + "Storage: response done {} with finish_reason={}", + response_id, finish_reason + ); + // Sender invariant: stream_one_assistant_turn always emits MessageDone / + // ReasoningDone before ResponseDone. If that invariant ever breaks we'd + // leave items as "in_progress" forever, so defensively finalize any + // stragglers as incomplete and log loudly so the bug is visible. + if !pending_messages.is_empty() || !pending_reasoning.is_empty() { + error!( + "ResponseDone received with {} pending message(s) and {} pending reasoning item(s); sender invariant violated -- finalizing as incomplete", + pending_messages.len(), + pending_reasoning.len() + ); + mark_pending_items_incomplete( + &db, + &user_key, + &mut pending_messages, + &mut pending_reasoning, + None, + ) + .await; + } + if let Err(e) = db.update_response_status( + response_id, + ResponseStatus::Completed, + Some(Utc::now()), + ) { + error!("Failed to update response status to completed: {:?}", e); + } + return; + } + StorageMessage::Cancelled => { + debug!( + "Storage: cancellation received for response {}", + response_id + ); + if let Err(e) = db.update_response_status( + response_id, + ResponseStatus::Cancelled, + Some(Utc::now()), + ) { + error!("Failed to update response status to cancelled: {:?}", e); + } + mark_pending_items_incomplete( + &db, + &user_key, + &mut pending_messages, + &mut pending_reasoning, + Some(FINISH_REASON_CANCELLED.to_string()), + ) + .await; + return; + } + StorageMessage::Error(error_msg) => { + error!("Storage: received error: {}", error_msg); + if let Err(e) = + db.update_response_status(response_id, ResponseStatus::Failed, Some(Utc::now())) + { + error!("Failed to update response status to failed: {:?}", e); + } + mark_pending_items_incomplete( + &db, + &user_key, + &mut pending_messages, + &mut pending_reasoning, + None, + ) + .await; + return; + } + StorageMessage::ToolCall { tool_call_id, name, arguments, } => { - // Persist tool call immediately to database - use crate::models::responses::NewToolCall; - - let arguments_json = match serde_json::to_string(&arguments) { - Ok(json) => json, - Err(e) => { - error!("Failed to serialize tool arguments: {:?}", e); - if let Some(ack) = tool_ack.take() { - let _ = ack - .send(Err(format!("Failed to serialize tool arguments: {:?}", e))); - } - continue; - } - }; - let arguments_enc = encrypt_with_key(&user_key, arguments_json.as_bytes()).await; - let token_count = count_tokens(&arguments_json); - let argument_tokens = if token_count > i32::MAX as usize { - warn!( - "Tool argument token count {} exceeds i32::MAX, clamping", - token_count - ); - i32::MAX - } else { - token_count as i32 - }; - - let new_tool_call = NewToolCall { - uuid: tool_call_id, + trace!("Storage: persisting tool_call {}", tool_call_id); + let created_at = allocate_created_at(&mut next_item_created_at); + match persist_tool_call( + &db, + &user_key, conversation_id, - response_id: Some(response_id), + response_id, user_id, + tool_call_id, name, - arguments_enc: Some(arguments_enc), - argument_tokens, - status: "completed".to_string(), - }; - - match db.create_tool_call(new_tool_call) { - Ok(tool_call) => { - debug!( - "Persisted tool_call {} (db id: {})", - tool_call_id, tool_call.id - ); - // No need to track the ID in memory - we'll look it up when needed - } + arguments, + created_at, + ) + .await + { + Ok(()) => debug!("Persisted tool_call {}", tool_call_id), Err(e) => { - error!("Failed to persist tool_call {}: {:?}", tool_call_id, e); - if let Some(ack) = tool_ack.take() { - let _ = ack.send(Err(format!("Failed to persist tool_call: {:?}", e))); + error!("{}", e); + if let Some(ack) = &tool_ack { + let _ = ack.send(Err(e)).await; } } } } - - AccumulatorState::PersistToolOutput { + StorageMessage::ToolOutput { tool_output_id, tool_call_id, output, } => { - // Persist tool output immediately to database - use crate::models::responses::NewToolOutput; - - // Look up the tool_call by UUID to get its database ID (primary key) - // This is more reliable than tracking in memory across async operations - // Also validates that the tool_call belongs to this user (security check) - let tool_call_fk = match db.get_tool_call_by_uuid(tool_call_id, user_id) { - Ok(tool_call) => tool_call.id, - Err(e) => { - error!( - "Failed to find tool_call {} for tool_output: {:?}", - tool_call_id, e - ); - if let Some(ack) = tool_ack.take() { - let _ = - ack.send(Err(format!("Tool call not found in database: {:?}", e))); - } - continue; - } - }; - - let output_enc = encrypt_with_key(&user_key, output.as_bytes()).await; - let token_count = count_tokens(&output); - let output_tokens = if token_count > i32::MAX as usize { - warn!( - "Tool output token count {} exceeds i32::MAX, clamping", - token_count - ); - i32::MAX - } else { - token_count as i32 - }; - - let new_tool_output = NewToolOutput { - uuid: tool_output_id, + trace!("Storage: persisting tool_output {}", tool_output_id); + let created_at = allocate_created_at(&mut next_item_created_at); + match persist_tool_output( + &db, + &user_key, conversation_id, - response_id: Some(response_id), + response_id, user_id, - tool_call_fk, - output_enc, - output_tokens, - status: "completed".to_string(), - error: None, - }; - - match db.create_tool_output(new_tool_output) { - Ok(_) => { - debug!( - "Persisted tool_output {} for tool_call {}", - tool_output_id, tool_call_id - ); - - // Send acknowledgment after tool output is persisted - if let Some(ack) = tool_ack.take() { - let _ = ack.send(Ok(())); + tool_output_id, + tool_call_id, + output, + created_at, + ) + .await + { + Ok(()) => { + debug!("Persisted tool_output {}", tool_output_id); + if let Some(ack) = &tool_ack { + let _ = ack.send(Ok(())).await; } } Err(e) => { - error!("Failed to persist tool_output {}: {:?}", tool_output_id, e); - if let Some(ack) = tool_ack.take() { - let _ = - ack.send(Err(format!("Failed to persist tool_output: {:?}", e))); + error!("{}", e); + if let Some(ack) = &tool_ack { + let _ = ack.send(Err(e)).await; } } } } - - AccumulatorState::Complete(data) => { - if let Err(e) = persister.persist_completed(data).await { - error!("Failed to persist completed response: {}", e); - } - return; - } - - AccumulatorState::Cancelled(data) => { - if let Err(e) = persister.persist_cancelled(data).await { - error!("Failed to persist cancelled response: {}", e); - } - return; - } - - AccumulatorState::Failed(data) => { - if let Err(e) = persister.persist_failed(data).await { - error!("Failed to persist failed response: {}", e); - } - return; - } } } - // Channel closed without Done or Error - treat as failure - warn!("Storage channel closed before receiving Done signal"); - if let Err(e) = persister - .persist_failed(FailureData { - error: "Channel closed prematurely".to_string(), - partial_content: String::new(), - completion_tokens: 0, - reasoning_content: accumulator.reasoning_content().to_string(), - reasoning_item_id: accumulator.reasoning_item_id(), - }) - .await + warn!("Storage channel closed before receiving ResponseDone signal"); + if let Err(e) = db.update_response_status(response_id, ResponseStatus::Failed, Some(Utc::now())) { - error!("Failed to persist incomplete response: {}", e); + error!( + "Failed to update response status after premature channel close: {:?}", + e + ); } + mark_pending_items_incomplete( + &db, + &user_key, + &mut pending_messages, + &mut pending_reasoning, + None, + ) + .await; } diff --git a/src/web/responses/tools.rs b/src/web/responses/tools.rs index 3baff002..6dc66d12 100644 --- a/src/web/responses/tools.rs +++ b/src/web/responses/tools.rs @@ -4,40 +4,69 @@ //! architecture that can be extended for additional tools in the future. use crate::brave::{BraveClient, SearchRequest as BraveSearchRequest}; -use crate::kagi::{KagiClient, SearchRequest as KagiSearchRequest}; +use crate::tinfoil_websearch::TinfoilWebSearchClient; use serde_json::{json, Value}; use std::sync::Arc; use tracing::{debug, error, info, trace, warn}; -/// Execute web search using Brave or Kagi Search API -/// -/// Prefers Brave if available, falls back to Kagi if Brave is not configured. -/// Requires at least one client to be provided (initialized at startup with connection pooling). +/// Execute web search using the configured backend. pub async fn execute_web_search( query: &str, + max_results: Option, + tinfoil_web_search_client: Option<&Arc>, brave_client: Option<&Arc>, - kagi_client: Option<&Arc>, ) -> Result { trace!("Executing web search for query: {}", query); info!("Executing web search"); - // Try Brave first, then fall back to Kagi + if let Some(client) = tinfoil_web_search_client { + return execute_tinfoil_search(query, max_results, client).await; + } + if let Some(client) = brave_client { - execute_brave_search(query, client).await - } else if let Some(client) = kagi_client { - execute_kagi_search(query, client).await + execute_brave_search(query, max_results, client).await } else { error!("No search client configured"); Err("No search client configured".to_string()) } } +async fn execute_tinfoil_search( + query: &str, + max_results: Option, + client: &Arc, +) -> Result { + trace!("Executing Tinfoil web search for query: {}", query); + + let result = client + .search(query, max_results.map(|value| value.clamp(1, 30))) + .await + .map_err(|e| { + error!("Tinfoil web search API error: {:?}", e); + format!("Search API error: {:?}", e) + })?; + + let formatted = format_tinfoil_search_result(&result); + if formatted.trim().is_empty() { + warn!("No Tinfoil web search results found"); + return Ok(format!("No results found for query: '{}'", query)); + } + + Ok(formatted) +} + /// Execute web search using Brave Search API -async fn execute_brave_search(query: &str, client: &Arc) -> Result { +async fn execute_brave_search( + query: &str, + max_results: Option, + client: &Arc, +) -> Result { trace!("Executing Brave search for query: {}", query); + let result_limit = max_results.unwrap_or(5).clamp(1, 30) as usize; // Create search request with summary enabled let mut search_request = BraveSearchRequest::new(query.to_string()); + search_request.count = Some(result_limit as u32); search_request.summary = Some(true); // Execute search @@ -53,7 +82,7 @@ async fn execute_brave_search(query: &str, client: &Arc) -> Result< if let Some(web) = response.web { if let Some(results) = web.results { result_text.push_str("Search Results:\n\n"); - for (i, result) in results.iter().take(5).enumerate() { + for (i, result) in results.iter().take(result_limit).enumerate() { result_text.push_str(&format!( "{}. {}\n URL: {}\n {}\n\n", i + 1, @@ -80,7 +109,7 @@ async fn execute_brave_search(query: &str, client: &Arc) -> Result< if let Some(news_results) = news.results { if !news_results.is_empty() { result_text.push_str("\nNews:\n\n"); - for (i, result) in news_results.iter().take(3).enumerate() { + for (i, result) in news_results.iter().take(result_limit.min(3)).enumerate() { result_text.push_str(&format!( "{}. {}\n URL: {}\n {}\n\n", i + 1, @@ -131,106 +160,89 @@ async fn execute_brave_search(query: &str, client: &Arc) -> Result< Ok(result_text) } -/// Execute web search using Kagi Search API -async fn execute_kagi_search(query: &str, client: &Arc) -> Result { - trace!("Executing Kagi search for query: {}", query); +fn format_tinfoil_search_result(result: &Value) -> String { + let content_text = result + .get("content") + .and_then(Value::as_array) + .map(|content| { + content + .iter() + .filter_map(|item| item.get("text").and_then(Value::as_str)) + .map(str::trim) + .filter(|text| !text.is_empty()) + .collect::>() + .join("\n\n") + }) + .filter(|text| !text.is_empty()); + + if let Some(text) = content_text { + return text; + } - // Create search request - let search_request = KagiSearchRequest::new(query.to_string()); + let structured_content = result + .get("structuredContent") + .or_else(|| result.get("structured_content")) + .unwrap_or(result); - // Execute search - let response = client.search(search_request).await.map_err(|e| { - error!("Kagi search API error: {:?}", e); - format!("Search API error: {:?}", e) - })?; + let formatted_results = structured_content + .get("results") + .and_then(Value::as_array) + .map(|results| format_search_results(results)) + .filter(|text| !text.is_empty()); - // Format results - let mut result_text = String::new(); + if let Some(text) = formatted_results { + return text; + } - if let Some(data) = response.data { - // Prioritize direct answers - if let Some(direct_answers) = data.direct_answer { - for answer in direct_answers { - result_text.push_str(&format!( - "Direct Answer: {}\n\n", - answer.snippet.unwrap_or_default() - )); - } - } + if let Some(text) = structured_content.as_str() { + return text.to_string(); + } - // Add weather information if available - if let Some(weather_results) = data.weather { - if !weather_results.is_empty() { - result_text.push_str("Weather:\n\n"); - for result in weather_results.iter().take(1) { - result_text.push_str(&format!( - "{}\n {}\n\n", - result.title, - result.snippet.as_ref().unwrap_or(&String::new()) - )); - } - } - } + serde_json::to_string_pretty(structured_content) + .unwrap_or_else(|_| structured_content.to_string()) +} - // Add infobox if available (detailed entity information) - if let Some(infobox_results) = data.infobox { - if !infobox_results.is_empty() { - result_text.push_str("Information:\n\n"); - for result in infobox_results.iter().take(1) { - result_text.push_str(&format!( - "{}\n {}\n", - result.title, - result.snippet.as_ref().unwrap_or(&String::new()) - )); +fn format_search_results(results: &[Value]) -> String { + let mut formatted = String::new(); - // Add URL if available for more details - if !result.url.is_empty() { - result_text.push_str(&format!(" More info: {}\n", result.url)); - } - result_text.push('\n'); - } - } - } + if !results.is_empty() { + formatted.push_str("Search Results:\n\n"); + } - // Add search results - if let Some(search_results) = data.search { - result_text.push_str("Search Results:\n\n"); - for (i, result) in search_results.iter().take(5).enumerate() { - result_text.push_str(&format!( - "{}. {}\n URL: {}\n {}\n\n", - i + 1, - result.title, - result.url, - result.snippet.as_ref().unwrap_or(&String::new()) - )); - } + for (index, result) in results.iter().enumerate() { + let title = result + .get("title") + .and_then(Value::as_str) + .unwrap_or("Untitled result"); + let url = result + .get("url") + .or_else(|| result.get("id")) + .and_then(Value::as_str) + .unwrap_or(""); + let snippet = [ + "text", + "snippet", + "summary", + "description", + "content", + "highlights", + ] + .iter() + .find_map(|field| result.get(field).and_then(Value::as_str)) + .unwrap_or(""); + + formatted.push_str(&format!("{}. {}\n", index + 1, title)); + if !url.is_empty() { + formatted.push_str(&format!(" URL: {}\n", url)); } - - // Add news results if available - if let Some(news_results) = data.news { - if !news_results.is_empty() { - result_text.push_str("\nNews:\n\n"); - for (i, result) in news_results.iter().take(3).enumerate() { - result_text.push_str(&format!( - "{}. {}\n URL: {}\n {}\n\n", - i + 1, - result.title, - result.url, - result.snippet.as_ref().unwrap_or(&String::new()) - )); - } - } + if !snippet.is_empty() { + formatted.push_str(&format!(" {}\n", snippet)); } + formatted.push('\n'); } - if result_text.is_empty() { - warn!("No search results found"); - return Ok(format!("No results found for query: '{}'", query)); - } - - Ok(result_text) + formatted.trim().to_string() } - /// Execute a tool by name with the given arguments /// /// This is the main entry point for tool execution. It routes to the appropriate @@ -239,8 +251,8 @@ async fn execute_kagi_search(query: &str, client: &Arc) -> Result) -> Result>, brave_client: Option<&Arc>, - kagi_client: Option<&Arc>, ) -> Result { trace!( "Executing tool: {} with arguments: {}", @@ -265,8 +277,12 @@ pub async fn execute_tool( .get("query") .and_then(|q| q.as_str()) .ok_or_else(|| "Missing 'query' argument for web_search".to_string())?; + let max_results = arguments + .get("max_results") + .and_then(|value| value.as_u64()) + .map(|value| value as u32); - execute_web_search(query, brave_client, kagi_client).await + execute_web_search(query, max_results, tinfoil_web_search_client, brave_client).await } _ => { error!("Unknown tool requested: {}", tool_name); @@ -304,6 +320,10 @@ impl ToolRegistry { "query": { "type": "string", "description": "The search query to execute" + }, + "max_results": { + "type": "integer", + "description": "Optional maximum number of results to return" } }, "required": ["query"] @@ -332,7 +352,7 @@ mod tests { #[tokio::test] async fn test_execute_web_search_no_client() { - let result = execute_web_search("test query", None, None).await; + let result = execute_web_search("test query", None, None, None).await; assert!(result.is_err()); assert!(result.unwrap_err().contains("No search client configured")); } @@ -354,6 +374,44 @@ mod tests { assert!(result.unwrap_err().contains("Unknown tool")); } + #[test] + fn test_format_tinfoil_search_result_prefers_text_content() { + let result = json!({ + "content": [ + { + "type": "text", + "text": "Search Results:\n\n1. Example\n URL: https://example.com" + } + ] + }); + + assert_eq!( + format_tinfoil_search_result(&result), + "Search Results:\n\n1. Example\n URL: https://example.com" + ); + } + + #[test] + fn test_format_tinfoil_search_result_formats_structured_results() { + let result = json!({ + "structuredContent": { + "results": [ + { + "title": "Example", + "url": "https://example.com", + "text": "Example summary" + } + ] + } + }); + + let formatted = format_tinfoil_search_result(&result); + assert!(formatted.contains("Search Results:")); + assert!(formatted.contains("Example")); + assert!(formatted.contains("https://example.com")); + assert!(formatted.contains("Example summary")); + } + #[test] fn test_tool_registry() { let registry = ToolRegistry::new();