diff --git a/src/web/responses/handlers.rs b/src/web/responses/handlers.rs index d45f6028..12b627ea 100644 --- a/src/web/responses/handlers.rs +++ b/src/web/responses/handlers.rs @@ -1320,13 +1320,13 @@ fn is_web_search_enabled(tools: &Option) -> bool { /// /// Flow: /// 1. Classify intent: chat vs web_search -/// 2. If web_search: extract query and execute tool +/// 2. If web_search: draft the search-subagent request 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. +/// web search execution prefers TinFoil web search with Brave/Kagi fallback. struct ToolChannels<'a> { client: &'a mpsc::Sender, storage: &'a mpsc::Sender, @@ -1428,51 +1428,16 @@ async fn classify_and_execute_tools( if intent == "web_search" { debug!("User message classified as web_search, executing tool"); - // 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(), - ); - - let search_query = match get_chat_completion_response( + let web_search_context = tools::WebSearchExecutionContext { 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() - } + conversation_history: prompt_messages, + user_message: &user_text, }; + let search_query = tools::prepare_web_search_query(&web_search_context).await; + trace!("Prepared web search query: {}", search_query); + // Generate UUIDs for tool_call and tool_output let tool_call_id = Uuid::new_v4(); let tool_output_id = Uuid::new_v4(); @@ -1523,6 +1488,7 @@ async fn classify_and_execute_tools( let tool_output = match tools::execute_tool( "web_search", &tool_arguments, + Some(&web_search_context), state.brave_client.as_ref(), state.kagi_client.as_ref(), ) @@ -2051,11 +2017,11 @@ 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) + // Phase 5: Classify intent and execute tools (if tool_choice allows tools, web search is enabled, and at least one search backend is 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"); + && tools::has_web_search_backend(&orchestrator_state) { + debug!("Orchestrator: tool_choice allows tools, web search is enabled, and a search backend is available, proceeding with classification"); let prepared_for_tools = PreparedRequest { user_key, @@ -2100,7 +2066,7 @@ async fn create_response_stream( } } } else { - debug!("Orchestrator: Web search tool not enabled or Kagi client not available, skipping classification"); + debug!("Orchestrator: tool execution disabled or no search backend available, skipping classification"); drop(rx_tool_ack); false }; diff --git a/src/web/responses/prompts.rs b/src/web/responses/prompts.rs index 6d4923ff..03140e44 100644 --- a/src/web/responses/prompts.rs +++ b/src/web/responses/prompts.rs @@ -9,6 +9,10 @@ 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; +pub const WEB_SEARCH_PROMPT_DRAFTER_MODEL: &str = "llama3-3-70b"; +pub const WEB_SEARCH_PROMPT_DRAFTER_MAX_TOKENS: i32 = 120; +pub const WEB_SEARCH_SUBAGENT_MODEL: &str = "gemma4-31b"; +pub const WEB_SEARCH_SUBAGENT_MAX_TOKENS: i32 = 1200; /// System prompt for intent classification /// @@ -35,6 +39,55 @@ Examples: - \"Tell me about the latest SpaceX launch\" → latest SpaceX launch - After discussing \"iPhone 15\", user asks \"when was it released?\" → iPhone 15 release date"; +/// System prompt for drafting the request sent to the web search subagent. +pub const WEB_SEARCH_PROMPT_DRAFTER_PROMPT: &str = "\ +Rewrite the user's request into a focused brief for an internal web-search subagent. +Use the conversation history to resolve references and ambiguity. +Return only the brief to send to the search subagent. +The brief should state what factual information to gather and preserve any important constraints, entities, and time sensitivity. +Do not answer the question yourself. +Do not include citations, URLs, bullet lists, or conversational filler."; + +/// System prompt for the internal web search subagent +/// +/// This prompt instructs the LLM to use web search and return a concise, +/// grounded search report for the downstream conversational model. +pub const WEB_SEARCH_SUBAGENT_PROMPT: &str = "\ +You are an internal web search subagent for another language model. +Use web search to gather current, relevant, factual information that helps answer the user's latest question. +Return a concise search report for a downstream assistant, not a direct conversational reply to the end user. +Prioritize reliable and recent sources when the question is time-sensitive. +Do not include a 'Sources:' section or raw URLs in the body; source links are attached separately. +If the web results are incomplete, conflicting, or uncertain, say so briefly."; + +fn format_recent_conversation_history(conversation_history: &[Value]) -> String { + if conversation_history.is_empty() { + return String::new(); + } + + 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") + ) + } +} + /// Build a chat completion request for intent classification /// /// Uses a fast, cheap model (gpt-oss-120b) with temperature=0 for deterministic results. @@ -49,32 +102,7 @@ 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() - }; + let history_text = format_recent_conversation_history(conversation_history); // Build single user message with history + current query let user_prompt = format!("{}Current user query: {}", history_text, user_message); @@ -125,32 +153,7 @@ fn extract_text_from_content(content: &Value) -> String { /// # 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() - }; + let history_text = format_recent_conversation_history(conversation_history); // Build single user message with history + current query let user_prompt = format!("{}Current user question: {}", history_text, user_message); @@ -173,6 +176,56 @@ pub fn build_query_extraction_request(conversation_history: &[Value], user_messa }) } +/// Build a chat completion request for drafting the web search subagent prompt. +pub fn build_web_search_prompt_drafting_request( + conversation_history: &[Value], + user_message: &str, +) -> Value { + let history_text = format_recent_conversation_history(conversation_history); + let user_prompt = format!("{}Current user question: {}", history_text, user_message); + + json!({ + "model": WEB_SEARCH_PROMPT_DRAFTER_MODEL, + "messages": [ + { + "role": "system", + "content": WEB_SEARCH_PROMPT_DRAFTER_PROMPT + }, + { + "role": "user", + "content": user_prompt + } + ], + "temperature": 0.0, + "max_tokens": WEB_SEARCH_PROMPT_DRAFTER_MAX_TOKENS, + "stream": false + }) +} + +/// Build a chat completion request for the internal web search subagent +/// +/// Uses gemma4-31b with `web_search_options` enabled to synthesize search +/// results that can be injected back into the main conversational model. +pub fn build_web_search_subagent_request(subagent_prompt: &str) -> Value { + json!({ + "model": WEB_SEARCH_SUBAGENT_MODEL, + "messages": [ + { + "role": "system", + "content": WEB_SEARCH_SUBAGENT_PROMPT + }, + { + "role": "user", + "content": subagent_prompt + } + ], + "temperature": 0.0, + "max_tokens": WEB_SEARCH_SUBAGENT_MAX_TOKENS, + "web_search_options": {}, + "stream": false + }) +} + #[cfg(test)] mod tests { use super::*; @@ -309,6 +362,52 @@ mod tests { assert!(user_content.contains("Current user question: when was it released?")); } + #[test] + fn test_build_web_search_prompt_drafting_request() { + let history = vec![ + json!({"role": "user", "content": "Tell me about SpaceX"}), + json!({"role": "assistant", "content": "What would you like to know?"}), + ]; + + let request = + build_web_search_prompt_drafting_request(&history, "what is its latest revenue?"); + + assert_eq!(request["model"], WEB_SEARCH_PROMPT_DRAFTER_MODEL); + assert_eq!(request["max_tokens"], WEB_SEARCH_PROMPT_DRAFTER_MAX_TOKENS); + 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"); + + let user_content = messages[1]["content"].as_str().unwrap(); + assert!(user_content.contains("Conversation history:")); + assert!(user_content.contains("Tell me about SpaceX")); + assert!(user_content.contains("Current user question: what is its latest revenue?")); + } + + #[test] + fn test_build_web_search_subagent_request() { + let request = build_web_search_subagent_request( + "Find SpaceX's latest annual revenue estimate and note the source of the figure.", + ); + + assert_eq!(request["model"], WEB_SEARCH_SUBAGENT_MODEL); + assert_eq!(request["max_tokens"], WEB_SEARCH_SUBAGENT_MAX_TOKENS); + assert_eq!(request["stream"], false); + assert_eq!(request["web_search_options"], json!({})); + + let messages = request["messages"].as_array().unwrap(); + assert_eq!(messages.len(), 2); + assert_eq!(messages[0]["role"], "system"); + assert_eq!(messages[1]["role"], "user"); + assert_eq!( + messages[1]["content"], + "Find SpaceX's latest annual revenue estimate and note the source of the figure." + ); + } + #[test] fn test_prompts_contain_examples() { assert!(INTENT_CLASSIFIER_PROMPT.contains("web_search")); @@ -316,5 +415,7 @@ mod tests { assert!(INTENT_CLASSIFIER_PROMPT.contains("conversation history")); assert!(SEARCH_QUERY_EXTRACTOR_PROMPT.contains("Examples:")); + assert!(WEB_SEARCH_PROMPT_DRAFTER_PROMPT.contains("search subagent")); + assert!(WEB_SEARCH_SUBAGENT_PROMPT.contains("Do not include a 'Sources:' section")); } } diff --git a/src/web/responses/tools.rs b/src/web/responses/tools.rs index 3baff002..8edf096a 100644 --- a/src/web/responses/tools.rs +++ b/src/web/responses/tools.rs @@ -3,16 +3,108 @@ //! This module handles tool execution including web search, with a clean //! architecture that can be extended for additional tools in the future. +use super::prompts; use crate::brave::{BraveClient, SearchRequest as BraveSearchRequest}; use crate::kagi::{KagiClient, SearchRequest as KagiSearchRequest}; +use crate::models::users::User; +use crate::web::openai::{get_chat_completion_response, BillingContext, CompletionChunk}; +use crate::web::openai_auth::AuthMethod; +use crate::AppState; +use axum::http::HeaderMap; use serde_json::{json, Value}; +use std::collections::HashSet; use std::sync::Arc; use tracing::{debug, error, info, trace, warn}; -/// Execute web search using Brave or Kagi Search API +pub(crate) struct WebSearchExecutionContext<'a> { + pub state: &'a Arc, + pub user: &'a User, + pub conversation_history: &'a [Value], + pub user_message: &'a str, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +struct UrlCitation { + title: String, + url: String, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +struct BlockedSearch { + query: String, + reason: String, +} + +pub(crate) fn has_web_search_backend(state: &Arc) -> bool { + state.proxy_router.get_tinfoil_proxy().is_some() + || state.brave_client.is_some() + || state.kagi_client.is_some() +} + +pub(crate) async fn prepare_web_search_query(context: &WebSearchExecutionContext<'_>) -> String { + trace!( + "Preparing web search subagent query from user message: {}", + context.user_message + ); + + let request = prompts::build_web_search_prompt_drafting_request( + context.conversation_history, + context.user_message, + ); + let headers = HeaderMap::new(); + let billing_context = BillingContext::new( + AuthMethod::Jwt, + prompts::WEB_SEARCH_PROMPT_DRAFTER_MODEL.to_string(), + ); + + match get_chat_completion_response( + context.state, + context.user, + request, + &headers, + billing_context, + ) + .await + { + Ok(mut completion) => match completion.stream.recv().await { + Some(CompletionChunk::FullResponse(response_json)) => extract_completion_text( + response_json + .get("choices") + .and_then(|choices| choices.get(0)) + .and_then(|choice| choice.get("message")), + ) + .map(|query| query.trim().to_string()) + .filter(|query| !query.is_empty()) + .unwrap_or_else(|| { + warn!("Failed to prepare web search query, using original message"); + context.user_message.to_string() + }), + Some(CompletionChunk::Error(err)) => { + warn!( + "Web search prompt drafting returned an error, using original message: {}", + err + ); + context.user_message.to_string() + } + _ => { + warn!("Unexpected web search prompt drafting response, using original message"); + context.user_message.to_string() + } + }, + Err(e) => { + warn!( + "Web search prompt drafting failed, using original message: {:?}", + e + ); + context.user_message.to_string() + } + } +} + +/// 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). +/// This remains the fallback path when TinFoil web search is unavailable or +/// fails. Prefers Brave if available, then falls back to Kagi. pub async fn execute_web_search( query: &str, brave_client: Option<&Arc>, @@ -32,6 +124,133 @@ pub async fn execute_web_search( } } +async fn execute_web_search_with_context( + context: &WebSearchExecutionContext<'_>, + query: &str, + brave_client: Option<&Arc>, + kagi_client: Option<&Arc>, +) -> Result { + trace!("Executing context-aware web search for query: {}", query); + + let mut tinfoil_error = None; + + if context.state.proxy_router.get_tinfoil_proxy().is_some() { + match execute_tinfoil_web_search(context, query).await { + Ok(output) => return Ok(output), + Err(err) => { + warn!( + "TinFoil web search failed, falling back to legacy providers if available: {}", + err + ); + tinfoil_error = Some(err); + } + } + } + + if brave_client.is_some() || kagi_client.is_some() { + let fallback_query = extract_fallback_search_query(context).await; + return execute_web_search(&fallback_query, brave_client, kagi_client).await; + } + + if let Some(err) = tinfoil_error { + Err(err) + } else { + error!("No search client configured"); + Err("No search client configured".to_string()) + } +} + +async fn execute_tinfoil_web_search( + context: &WebSearchExecutionContext<'_>, + subagent_prompt: &str, +) -> Result { + trace!( + "Executing TinFoil web search with prepared subagent prompt: {}", + subagent_prompt + ); + + let request = prompts::build_web_search_subagent_request(subagent_prompt); + let headers = HeaderMap::new(); + let billing_context = BillingContext::new( + AuthMethod::Jwt, + prompts::WEB_SEARCH_SUBAGENT_MODEL.to_string(), + ); + + let mut completion = get_chat_completion_response( + context.state, + context.user, + request, + &headers, + billing_context, + ) + .await + .map_err(|e| format!("TinFoil web search request failed: {:?}", e))?; + + match completion.stream.recv().await { + Some(CompletionChunk::FullResponse(response_json)) => { + trace!( + "TinFoil web search response: {}", + serde_json::to_string_pretty(&response_json) + .unwrap_or_else(|_| "failed to serialize".to_string()) + ); + format_tinfoil_web_search_response(&response_json) + } + Some(CompletionChunk::Error(err)) => Err(format!("TinFoil web search error: {}", err)), + _ => Err("TinFoil web search returned unexpected response".to_string()), + } +} + +async fn extract_fallback_search_query(context: &WebSearchExecutionContext<'_>) -> String { + let query_request = + prompts::build_query_extraction_request(context.conversation_history, context.user_message); + let headers = HeaderMap::new(); + let billing_context = + BillingContext::new(AuthMethod::Jwt, prompts::QUERY_EXTRACTOR_MODEL.to_string()); + + match get_chat_completion_response( + context.state, + context.user, + query_request, + &headers, + billing_context, + ) + .await + { + Ok(mut completion) => match completion.stream.recv().await { + Some(CompletionChunk::FullResponse(response_json)) => extract_completion_text( + response_json + .get("choices") + .and_then(|choices| choices.get(0)) + .and_then(|choice| choice.get("message")), + ) + .map(|query| query.trim().to_string()) + .filter(|query| !query.is_empty()) + .unwrap_or_else(|| { + warn!("Failed to extract fallback search query, using original message"); + context.user_message.to_string() + }), + Some(CompletionChunk::Error(err)) => { + warn!( + "Fallback query extraction returned an error, using original message: {}", + err + ); + context.user_message.to_string() + } + _ => { + warn!("Unexpected fallback query extraction response, using original message"); + context.user_message.to_string() + } + }, + Err(e) => { + warn!( + "Fallback query extraction failed, using original message: {:?}", + e + ); + context.user_message.to_string() + } + } +} + /// Execute web search using Brave Search API async fn execute_brave_search(query: &str, client: &Arc) -> Result { trace!("Executing Brave search for query: {}", query); @@ -231,6 +450,167 @@ async fn execute_kagi_search(query: &str, client: &Arc) -> Result Result { + let message = response_json + .get("choices") + .and_then(|choices| choices.get(0)) + .and_then(|choice| choice.get("message")) + .ok_or_else(|| "TinFoil web search returned no assistant message".to_string())?; + + let content = + sanitize_tinfoil_summary(&extract_completion_text(Some(message)).unwrap_or_default()); + let blocked_searches = extract_blocked_searches(message); + let citations = extract_url_citations(message); + + let mut sections = Vec::new(); + + if !content.trim().is_empty() { + sections.push(format!("Search Summary:\n\n{}", content.trim())); + } + + if !blocked_searches.is_empty() { + let blocked = blocked_searches + .iter() + .map(|blocked| format!("- {} ({})", blocked.query, blocked.reason)) + .collect::>() + .join("\n"); + sections.push(format!("Blocked Searches:\n{}", blocked)); + } + + if !citations.is_empty() { + let sources = citations + .iter() + .enumerate() + .map(|(index, citation)| { + format!("{}. {} — {}", index + 1, citation.title, citation.url) + }) + .collect::>() + .join("\n"); + sections.push(format!("Sources:\n{}", sources)); + } + + if sections.is_empty() { + Err("TinFoil web search returned no content".to_string()) + } else { + Ok(sections.join("\n\n")) + } +} + +fn sanitize_tinfoil_summary(content: &str) -> String { + let mut cleaned_lines = Vec::new(); + + for line in content.lines() { + if is_sources_heading(line) { + break; + } + + cleaned_lines.push(line); + } + + cleaned_lines.join("\n").trim().to_string() +} + +fn is_sources_heading(line: &str) -> bool { + let trimmed = line.trim(); + trimmed.eq_ignore_ascii_case("sources:") + || trimmed + .get(..8) + .is_some_and(|prefix| prefix.eq_ignore_ascii_case("sources:")) +} + +fn extract_completion_text(message: Option<&Value>) -> Option { + let content = message?.get("content")?; + + match content { + Value::String(text) => Some(text.clone()), + Value::Array(parts) => { + let joined = parts + .iter() + .filter_map(|part| { + part.get("text") + .and_then(|text| text.as_str()) + .or_else(|| part.get("content").and_then(|text| text.as_str())) + .map(|text| text.to_string()) + }) + .collect::>() + .join("\n"); + + if joined.trim().is_empty() { + None + } else { + Some(joined) + } + } + _ => None, + } +} + +fn extract_url_citations(message: &Value) -> Vec { + let mut seen_urls = HashSet::new(); + let mut citations = Vec::new(); + + let Some(annotations) = message + .get("annotations") + .and_then(|annotations| annotations.as_array()) + else { + return citations; + }; + + for annotation in annotations { + let Some(citation) = annotation.get("url_citation") else { + continue; + }; + + let Some(url) = citation.get("url").and_then(|url| url.as_str()) else { + continue; + }; + + let url = url.trim(); + if url.is_empty() || !seen_urls.insert(url.to_string()) { + continue; + } + + let title = citation + .get("title") + .and_then(|title| title.as_str()) + .map(str::trim) + .filter(|title| !title.is_empty()) + .unwrap_or(url); + + citations.push(UrlCitation { + title: title.to_string(), + url: url.to_string(), + }); + } + + citations +} + +fn extract_blocked_searches(message: &Value) -> Vec { + message + .get("blocked_searches") + .and_then(|blocked_searches| blocked_searches.as_array()) + .map(|blocked_searches| { + blocked_searches + .iter() + .filter_map(|blocked| { + let query = blocked.get("query")?.as_str()?.trim(); + let reason = blocked.get("reason")?.as_str()?.trim(); + + if query.is_empty() || reason.is_empty() { + return None; + } + + Some(BlockedSearch { + query: query.to_string(), + reason: reason.to_string(), + }) + }) + .collect() + }) + .unwrap_or_default() +} + /// Execute a tool by name with the given arguments /// /// This is the main entry point for tool execution. It routes to the appropriate @@ -239,6 +619,7 @@ async fn execute_kagi_search(query: &str, client: &Arc) -> Result) -> Result>, brave_client: Option<&Arc>, kagi_client: Option<&Arc>, ) -> Result { @@ -266,7 +648,11 @@ pub async fn execute_tool( .and_then(|q| q.as_str()) .ok_or_else(|| "Missing 'query' argument for web_search".to_string())?; - execute_web_search(query, brave_client, kagi_client).await + if let Some(context) = web_search_context { + execute_web_search_with_context(context, query, brave_client, kagi_client).await + } else { + execute_web_search(query, brave_client, kagi_client).await + } } _ => { error!("Unknown tool requested: {}", tool_name); @@ -341,7 +727,7 @@ mod tests { async fn test_execute_tool_missing_args() { // Test with None client - should fail on missing args before client check let args = json!({}); - let result = execute_tool("web_search", &args, None, None).await; + let result = execute_tool("web_search", &args, None, None, None).await; assert!(result.is_err()); assert!(result.unwrap_err().contains("Missing 'query'")); } @@ -349,11 +735,99 @@ mod tests { #[tokio::test] async fn test_execute_tool_unknown() { let args = json!({"query": "test"}); - let result = execute_tool("unknown_tool", &args, None, None).await; + let result = execute_tool("unknown_tool", &args, None, None, None).await; assert!(result.is_err()); assert!(result.unwrap_err().contains("Unknown tool")); } + #[test] + fn test_format_tinfoil_web_search_response_includes_sources() { + let response = json!({ + "choices": [{ + "message": { + "content": "SpaceX generated about $15.6 billion in revenue in 2025.【1】", + "annotations": [{ + "type": "url_citation", + "url_citation": { + "title": "SpaceX Revenue Report", + "url": "https://example.com/spacex-revenue" + } + }] + } + }] + }); + + let result = format_tinfoil_web_search_response(&response).unwrap(); + assert!(result.contains("Search Summary:")); + assert!(result.contains("SpaceX generated about $15.6 billion")); + assert!(result.contains("Sources:")); + assert!(result.contains("SpaceX Revenue Report")); + assert!(result.contains("https://example.com/spacex-revenue")); + } + + #[test] + fn test_format_tinfoil_web_search_response_includes_blocked_searches() { + let response = json!({ + "choices": [{ + "message": { + "content": "I could not use web search for this request.", + "blocked_searches": [{ + "query": "search for account number 1234567890", + "reason": "Bank account number detected" + }] + } + }] + }); + + let result = format_tinfoil_web_search_response(&response).unwrap(); + assert!(result.contains("Blocked Searches:")); + assert!(result.contains("Bank account number detected")); + } + + #[test] + fn test_extract_completion_text_from_content_array() { + let message = json!({ + "content": [ + {"type": "text", "text": "First line"}, + {"type": "text", "text": "Second line"} + ] + }); + + let text = extract_completion_text(Some(&message)).unwrap(); + assert_eq!(text, "First line\nSecond line"); + } + + #[test] + fn test_sanitize_tinfoil_summary_strips_embedded_sources_section() { + let summary = sanitize_tinfoil_summary( + "It will be cloudy today.\n\nSources:\nExample Weather: https://example.com/weather", + ); + + assert_eq!(summary, "It will be cloudy today."); + } + + #[test] + fn test_format_tinfoil_web_search_response_avoids_duplicate_sources_sections() { + let response = json!({ + "choices": [{ + "message": { + "content": "It will be cloudy today.【1】\n\nSources:\nExample Weather: https://example.com/weather", + "annotations": [{ + "type": "url_citation", + "url_citation": { + "title": "Example Weather", + "url": "https://example.com/weather" + } + }] + } + }] + }); + + let result = format_tinfoil_web_search_response(&response).unwrap(); + assert_eq!(result.matches("Sources:").count(), 1); + assert!(result.contains("It will be cloudy today.")); + } + #[test] fn test_tool_registry() { let registry = ToolRegistry::new(); diff --git a/tinfoil-proxy/dist/tinfoil-proxy b/tinfoil-proxy/dist/tinfoil-proxy index 5661f9a9..9137b0aa 100755 Binary files a/tinfoil-proxy/dist/tinfoil-proxy and b/tinfoil-proxy/dist/tinfoil-proxy differ diff --git a/tinfoil-proxy/main.go b/tinfoil-proxy/main.go index 1e1e04c2..323cb2bb 100644 --- a/tinfoil-proxy/main.go +++ b/tinfoil-proxy/main.go @@ -1,7 +1,9 @@ package main import ( + "bytes" "context" + "encoding/json" "errors" "io" "log" @@ -19,11 +21,24 @@ import ( var errMissingAPIKey = errors.New("TINFOIL_API_KEY environment variable is required") const upstreamResponseStartTimeout = 120 * time.Second +const defaultWebSearchAPIBase = "https://websearch-debug.debug.tinfoil.containers.tinfoil.dev" type proxyServer struct { - httpClient *http.Client - apiKey string - enclaveHost string + httpClient *http.Client + webSearchClient *http.Client + apiKey string + enclaveHost string + webSearchBase *url.URL +} + +type upstreamTarget struct { + client *http.Client + baseURL url.URL + name string +} + +type chatCompletionEnvelope struct { + WebSearchOptions *json.RawMessage `json:"web_search_options"` } type flushWriter struct { @@ -53,23 +68,86 @@ func newProxyServer() (*proxyServer, error) { httpClient := *client.HTTPClient() httpClient.Timeout = 0 + webSearchBase := os.Getenv("TINFOIL_WEBSEARCH_API_BASE") + if webSearchBase == "" { + webSearchBase = defaultWebSearchAPIBase + } + + parsedWebSearchBase, err := url.Parse(webSearchBase) + if err != nil { + return nil, err + } + if parsedWebSearchBase.Scheme == "" || parsedWebSearchBase.Host == "" { + return nil, errors.New("TINFOIL_WEBSEARCH_API_BASE must be an absolute URL") + } + return &proxyServer{ - httpClient: &httpClient, - apiKey: apiKey, - enclaveHost: client.Enclave(), + httpClient: &httpClient, + webSearchClient: &http.Client{}, + apiKey: apiKey, + enclaveHost: client.Enclave(), + webSearchBase: parsedWebSearchBase, }, nil } func (s *proxyServer) upstreamURL(path, rawQuery string) string { - upstream := url.URL{ - Scheme: "https", - Host: s.enclaveHost, - Path: path, - RawQuery: rawQuery, - } + return buildUpstreamURL(url.URL{ + Scheme: "https", + Host: s.enclaveHost, + }, path, rawQuery) +} + +func buildUpstreamURL(base url.URL, path, rawQuery string) string { + upstream := base + upstream.Path = joinURLPath(base.Path, path) + upstream.RawQuery = rawQuery return upstream.String() } +func joinURLPath(basePath, path string) string { + trimmedPath := strings.TrimLeft(path, "/") + if basePath == "" || basePath == "/" { + return "/" + trimmedPath + } + + return strings.TrimRight(basePath, "/") + "/" + trimmedPath +} + +func shouldUseWebSearchUpstream(body []byte) bool { + if len(bytes.TrimSpace(body)) == 0 { + return false + } + + var envelope chatCompletionEnvelope + if err := json.Unmarshal(body, &envelope); err != nil { + return false + } + if envelope.WebSearchOptions == nil { + return false + } + + return !bytes.Equal(bytes.TrimSpace(*envelope.WebSearchOptions), []byte("null")) +} + +func (s *proxyServer) resolveUpstreamTarget(path string, body []byte) upstreamTarget { + if path == "/v1/chat/completions" && shouldUseWebSearchUpstream(body) { + return upstreamTarget{ + client: s.webSearchClient, + baseURL: *s.webSearchBase, + name: "websearch", + } + } + + return upstreamTarget{ + client: s.httpClient, + baseURL: url.URL{ + Scheme: "https", + Host: s.enclaveHost, + }, + name: "enclave", + } +} + func shouldSkipHeader(name string) bool { switch http.CanonicalHeaderKey(name) { case "Authorization", "Connection", "Content-Length", "Host", "Keep-Alive", @@ -171,11 +249,25 @@ func doWithResponseStartTimeout( } func (s *proxyServer) proxy(c *gin.Context, path string) { + body, err := io.ReadAll(c.Request.Body) + if err != nil { + log.Printf("failed to read request body for %s: %v", path, err) + writeProxyError(c, http.StatusBadRequest, "failed to read request body") + return + } + + target := s.resolveUpstreamTarget(path, body) + upstreamURL := buildUpstreamURL(target.baseURL, path, c.Request.URL.RawQuery) + + if target.name == "websearch" { + log.Printf("routing %s to experimental web search upstream %s", path, target.baseURL.Host) + } + req, err := http.NewRequestWithContext( c.Request.Context(), c.Request.Method, - s.upstreamURL(path, c.Request.URL.RawQuery), - c.Request.Body, + upstreamURL, + bytes.NewReader(body), ) if err != nil { log.Printf("failed to create upstream request for %s: %v", path, err) @@ -185,9 +277,9 @@ func (s *proxyServer) proxy(c *gin.Context, path string) { copyHeaders(req.Header, c.Request.Header) req.Header.Set("Authorization", "Bearer "+s.apiKey) - req.Host = s.enclaveHost + req.Host = target.baseURL.Host - resp, err := doWithResponseStartTimeout(s.httpClient, req, upstreamResponseStartTimeout) + resp, err := doWithResponseStartTimeout(target.client, req, upstreamResponseStartTimeout) if err != nil { if errors.Is(err, context.DeadlineExceeded) { log.Printf("upstream response start timed out for %s after %s", path, upstreamResponseStartTimeout) diff --git a/tinfoil-proxy/main_test.go b/tinfoil-proxy/main_test.go index 0fc998ce..422bdfa7 100644 --- a/tinfoil-proxy/main_test.go +++ b/tinfoil-proxy/main_test.go @@ -4,6 +4,7 @@ import ( "context" "io" "net/http" + "net/url" "testing" "time" ) @@ -75,6 +76,87 @@ func TestUpstreamURLPreservesQuery(t *testing.T) { } } +func TestShouldUseWebSearchUpstream(t *testing.T) { + tests := []struct { + name string + body string + want bool + }{ + { + name: "enabled", + body: `{"model":"gemma4-31b","web_search_options":{}}`, + want: true, + }, + { + name: "missing", + body: `{"model":"gemma4-31b"}`, + want: false, + }, + { + name: "null", + body: `{"model":"gemma4-31b","web_search_options":null}`, + want: false, + }, + { + name: "invalid json", + body: `{`, + want: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := shouldUseWebSearchUpstream([]byte(tt.body)); got != tt.want { + t.Fatalf("expected %v, got %v", tt.want, got) + } + }) + } +} + +func TestBuildUpstreamURLAppendsBasePath(t *testing.T) { + base, err := url.Parse("https://search.example.com/internal") + if err != nil { + t.Fatalf("failed to parse URL: %v", err) + } + + got := buildUpstreamURL(*base, "/v1/chat/completions", "stream=false") + want := "https://search.example.com/internal/v1/chat/completions?stream=false" + + if got != want { + t.Fatalf("expected %q, got %q", want, got) + } +} + +func TestResolveUpstreamTargetRoutesWebSearchRequests(t *testing.T) { + webSearchBase, err := url.Parse("https://search.example.com") + if err != nil { + t.Fatalf("failed to parse URL: %v", err) + } + + server := &proxyServer{ + httpClient: &http.Client{}, + webSearchClient: &http.Client{}, + enclaveHost: "enclave.example.com", + webSearchBase: webSearchBase, + } + + target := server.resolveUpstreamTarget("/v1/chat/completions", []byte(`{"web_search_options":{}}`)) + if target.name != "websearch" { + t.Fatalf("expected websearch target, got %q", target.name) + } + if target.baseURL.Host != "search.example.com" { + t.Fatalf("expected search host, got %q", target.baseURL.Host) + } + + target = server.resolveUpstreamTarget("/v1/chat/completions", []byte(`{"messages":[]}`)) + if target.name != "enclave" { + t.Fatalf("expected enclave target, got %q", target.name) + } + if target.baseURL.Host != "enclave.example.com" { + t.Fatalf("expected enclave host, got %q", target.baseURL.Host) + } +} + func TestDoWithResponseStartTimeoutTimesOut(t *testing.T) { client := &http.Client{ Transport: roundTripperFunc(func(req *http.Request) (*http.Response, error) {