Skip to content

Prototype TinFoil-backed web search flow #169

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
AnthonyRonning wants to merge 1 commit into
base: master
Choose a base branch
from
Draft

Prototype TinFoil-backed web search flow #169

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
60 changes: 13 additions & 47 deletions src/web/responses/handlers.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1320,13 +1320,13 @@ fn is_web_search_enabled(tools: &Option<Value>) -> 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<StorageMessage>,
storage: &'a mpsc::Sender<StorageMessage>,
Expand Down Expand Up @@ -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();
Expand Down Expand Up @@ -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(),
)
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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
};
Expand Down
205 changes: 153 additions & 52 deletions src/web/responses/prompts.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
///
Expand All @@ -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<String> = 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.
Expand All @@ -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<String> = 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);
Expand Down Expand Up @@ -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<String> = 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);
Expand All @@ -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::*;
Expand Down Expand Up @@ -309,12 +362,60 @@ 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"));
assert!(INTENT_CLASSIFIER_PROMPT.contains("chat"));
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"));
}
}
Loading
Loading