//! Chat endpoint for LLM-powered file editing. use axum::{ extract::{Path, State}, http::StatusCode, response::IntoResponse, Json, }; use serde::{Deserialize, Serialize}; use utoipa::ToSchema; use uuid::Uuid; use crate::db::{models::BodyElement, repository}; use crate::llm::{ claude::{self, ClaudeClient, ClaudeError, ClaudeModel}, execute_tool_call, groq::{GroqClient, GroqError, Message, ToolCallResponse}, LlmModel, ToolCall, ToolResult, AVAILABLE_TOOLS, }; use crate::server::state::{FileUpdateNotification, SharedState}; /// Maximum number of tool-calling rounds to prevent infinite loops const MAX_TOOL_ROUNDS: usize = 10; #[derive(Debug, Deserialize, ToSchema)] #[serde(rename_all = "camelCase")] pub struct ChatRequest { /// The user's message/instruction pub message: String, /// Optional model selection: "claude-sonnet" (default), "claude-opus", or "groq" #[serde(default)] pub model: Option, } #[derive(Debug, Serialize, ToSchema)] #[serde(rename_all = "camelCase")] pub struct ChatResponse { /// The LLM's response message pub response: String, /// Tool calls that were executed pub tool_calls: Vec, /// Updated file body after tool execution pub updated_body: Vec, /// Updated summary (if changed) pub updated_summary: Option, } #[derive(Debug, Serialize, ToSchema)] #[serde(rename_all = "camelCase")] pub struct ToolCallInfo { pub name: String, pub result: ToolResult, } /// Enum to hold LLM clients enum LlmClient { Groq(GroqClient), Claude(ClaudeClient), } /// Unified result from LLM call struct LlmResult { content: Option, tool_calls: Vec, raw_tool_calls: Vec, finish_reason: String, } /// Chat with a file using LLM tool calling #[utoipa::path( post, path = "/api/v1/files/{id}/chat", request_body = ChatRequest, responses( (status = 200, description = "Chat completed successfully", body = ChatResponse), (status = 404, description = "File not found"), (status = 500, description = "Internal server error") ), params( ("id" = Uuid, Path, description = "File ID") ), tag = "chat" )] pub async fn chat_handler( State(state): State, Path(id): Path, Json(request): Json, ) -> impl IntoResponse { // Check if database is configured let Some(ref pool) = state.db_pool else { return ( StatusCode::SERVICE_UNAVAILABLE, Json(serde_json::json!({ "error": "Database not configured" })), ) .into_response(); }; // Get the file let file = match repository::get_file(pool, id).await { Ok(Some(file)) => file, Ok(None) => { return ( StatusCode::NOT_FOUND, Json(serde_json::json!({ "error": "File not found" })), ) .into_response(); } Err(e) => { tracing::error!("Database error: {}", e); return ( StatusCode::INTERNAL_SERVER_ERROR, Json(serde_json::json!({ "error": format!("Database error: {}", e) })), ) .into_response(); } }; // Parse model selection (default to Claude Sonnet) let model = request .model .as_ref() .and_then(|m| LlmModel::from_str(m)) .unwrap_or_default(); tracing::info!("Using LLM model: {:?}", model); // Initialize the appropriate LLM client let llm_client = match model { LlmModel::ClaudeSonnet => { match ClaudeClient::from_env(ClaudeModel::Sonnet) { Ok(client) => LlmClient::Claude(client), Err(ClaudeError::MissingApiKey) => { return ( StatusCode::SERVICE_UNAVAILABLE, Json(serde_json::json!({ "error": "ANTHROPIC_API_KEY not configured" })), ) .into_response(); } Err(e) => { return ( StatusCode::INTERNAL_SERVER_ERROR, Json(serde_json::json!({ "error": format!("Claude client error: {}", e) })), ) .into_response(); } } } LlmModel::ClaudeOpus => { match ClaudeClient::from_env(ClaudeModel::Opus) { Ok(client) => LlmClient::Claude(client), Err(ClaudeError::MissingApiKey) => { return ( StatusCode::SERVICE_UNAVAILABLE, Json(serde_json::json!({ "error": "ANTHROPIC_API_KEY not configured" })), ) .into_response(); } Err(e) => { return ( StatusCode::INTERNAL_SERVER_ERROR, Json(serde_json::json!({ "error": format!("Claude client error: {}", e) })), ) .into_response(); } } } LlmModel::GroqKimi => { match GroqClient::from_env() { Ok(client) => LlmClient::Groq(client), Err(GroqError::MissingApiKey) => { return ( StatusCode::SERVICE_UNAVAILABLE, Json(serde_json::json!({ "error": "GROQ_API_KEY not configured" })), ) .into_response(); } Err(e) => { return ( StatusCode::INTERNAL_SERVER_ERROR, Json(serde_json::json!({ "error": format!("Groq client error: {}", e) })), ) .into_response(); } } } }; // Build context about the file let file_context = build_file_context(&file); // Build initial messages (Groq/OpenAI format - will be converted for Claude) let mut messages = vec![ Message { role: "system".to_string(), content: Some(format!( "You are a helpful assistant that helps users edit and analyze document files. \ You have access to tools to add headings, paragraphs, charts, and set summaries. \ When the user asks you to modify the file, use the appropriate tools.\n\n\ IMPORTANT: You can call multiple tools in sequence. For example, if the user provides CSV data \ and asks for a chart, first call parse_csv to convert the data to JSON, then use that JSON \ to call add_chart.\n\n\ Current file context:\n{}", file_context )), tool_calls: None, tool_call_id: None, }, Message { role: "user".to_string(), content: Some(request.message.clone()), tool_calls: None, tool_call_id: None, }, ]; // State for tracking changes let mut current_body = file.body.clone(); let mut current_summary = file.summary.clone(); let mut all_tool_call_infos: Vec = Vec::new(); let mut final_response: Option = None; // Multi-turn tool calling loop for round in 0..MAX_TOOL_ROUNDS { tracing::debug!(round = round, "LLM tool calling round"); // Call the appropriate LLM API let result = match &llm_client { LlmClient::Groq(groq) => { match groq.chat_with_tools(messages.clone(), &AVAILABLE_TOOLS).await { Ok(r) => LlmResult { content: r.content, tool_calls: r.tool_calls, raw_tool_calls: r.raw_tool_calls, finish_reason: r.finish_reason, }, Err(e) => { tracing::error!("Groq API error: {}", e); return ( StatusCode::INTERNAL_SERVER_ERROR, Json(serde_json::json!({ "error": format!("LLM API error: {}", e) })), ) .into_response(); } } } LlmClient::Claude(claude_client) => { // Convert messages to Claude format let claude_messages = claude::groq_messages_to_claude(&messages); match claude_client.chat_with_tools(claude_messages, &AVAILABLE_TOOLS).await { Ok(r) => { // Convert Claude tool uses to Groq-style ToolCallResponse for consistency let raw_tool_calls: Vec = r .tool_calls .iter() .map(|tc| ToolCallResponse { id: tc.id.clone(), call_type: "function".to_string(), function: crate::llm::groq::FunctionCall { name: tc.name.clone(), arguments: tc.arguments.to_string(), }, }) .collect(); LlmResult { content: r.content, tool_calls: r.tool_calls, raw_tool_calls, finish_reason: r.stop_reason, } } Err(e) => { tracing::error!("Claude API error: {}", e); return ( StatusCode::INTERNAL_SERVER_ERROR, Json(serde_json::json!({ "error": format!("LLM API error: {}", e) })), ) .into_response(); } } } }; // Check if there are tool calls to execute if result.tool_calls.is_empty() { // No more tool calls - capture the final response and exit loop final_response = result.content; break; } // Add assistant message with tool calls to conversation messages.push(Message { role: "assistant".to_string(), content: result.content.clone(), tool_calls: Some(result.raw_tool_calls.clone()), tool_call_id: None, }); // Execute each tool call and add results to conversation for (i, tool_call) in result.tool_calls.iter().enumerate() { let execution_result = execute_tool_call(tool_call, ¤t_body, current_summary.as_deref()); // Apply state changes if let Some(new_body) = execution_result.new_body { current_body = new_body; } if let Some(new_summary) = execution_result.new_summary { current_summary = Some(new_summary); } // Build tool result message content let result_content = if let Some(parsed_data) = &execution_result.parsed_data { // Include parsed data in the result for the LLM to use serde_json::json!({ "success": execution_result.result.success, "message": execution_result.result.message, "data": parsed_data }) .to_string() } else { serde_json::json!({ "success": execution_result.result.success, "message": execution_result.result.message }) .to_string() }; // Add tool result message // Use the appropriate ID format for each provider let tool_call_id = match &llm_client { LlmClient::Groq(_) => result.raw_tool_calls[i].id.clone(), LlmClient::Claude(_) => tool_call.id.clone(), }; messages.push(Message { role: "tool".to_string(), content: Some(result_content), tool_calls: None, tool_call_id: Some(tool_call_id), }); // Track for response all_tool_call_infos.push(ToolCallInfo { name: tool_call.name.clone(), result: execution_result.result, }); } // If finish reason indicates completion, exit loop let finish_lower = result.finish_reason.to_lowercase(); if finish_lower == "stop" || finish_lower == "end_turn" { final_response = result.content; break; } } // Save changes to database if any tools were executed if !all_tool_call_infos.is_empty() { let update_req = crate::db::models::UpdateFileRequest { name: None, description: None, transcript: None, summary: current_summary.clone(), body: Some(current_body.clone()), version: None, // Internal update, skip version check }; match repository::update_file(pool, id, update_req).await { Ok(Some(updated_file)) => { // Broadcast update notification for LLM changes let mut updated_fields = vec!["body".to_string()]; if current_summary.is_some() { updated_fields.push("summary".to_string()); } state.broadcast_file_update(FileUpdateNotification { file_id: id, version: updated_file.version, updated_fields, updated_by: "llm".to_string(), }); } Ok(None) => { // File was deleted during processing return ( StatusCode::NOT_FOUND, Json(serde_json::json!({ "error": "File not found" })), ) .into_response(); } Err(e) => { tracing::error!("Failed to save file changes: {}", e); return ( StatusCode::INTERNAL_SERVER_ERROR, Json(serde_json::json!({ "error": format!("Failed to save changes: {}", e) })), ) .into_response(); } } } // Build response let response_text = final_response.unwrap_or_else(|| { if all_tool_call_infos.is_empty() { "I couldn't understand your request. Please try rephrasing.".to_string() } else { format!( "Done! Executed {} tool{}.", all_tool_call_infos.len(), if all_tool_call_infos.len() == 1 { "" } else { "s" } ) } }); ( StatusCode::OK, Json(ChatResponse { response: response_text, tool_calls: all_tool_call_infos, updated_body: current_body, updated_summary: current_summary, }), ) .into_response() } fn build_file_context(file: &crate::db::models::File) -> String { let mut context = format!("File: {}\n", file.name); if let Some(ref desc) = file.description { context.push_str(&format!("Description: {}\n", desc)); } if let Some(ref summary) = file.summary { context.push_str(&format!("Summary: {}\n", summary)); } context.push_str(&format!("Transcript entries: {}\n", file.transcript.len())); context.push_str(&format!("Body elements: {}\n", file.body.len())); // Add body overview if !file.body.is_empty() { context.push_str("\nCurrent body elements:\n"); for (i, element) in file.body.iter().enumerate() { let desc = match element { BodyElement::Heading { level, text } => format!("H{}: {}", level, text), BodyElement::Paragraph { text } => { let preview = if text.len() > 50 { format!("{}...", &text[..50]) } else { text.clone() }; format!("Paragraph: {}", preview) } BodyElement::Chart { chart_type, title, .. } => { format!( "Chart ({:?}){}", chart_type, title.as_ref().map(|t| format!(": {}", t)).unwrap_or_default() ) } BodyElement::Image { alt, .. } => { format!("Image{}", alt.as_ref().map(|a| format!(": {}", a)).unwrap_or_default()) } }; context.push_str(&format!(" [{}] {}\n", i, desc)); } } // Add transcript preview if available if !file.transcript.is_empty() { context.push_str("\nTranscript preview (first 5 entries):\n"); for entry in file.transcript.iter().take(5) { context.push_str(&format!(" - {}: {}\n", entry.speaker, entry.text)); } if file.transcript.len() > 5 { context.push_str(&format!(" ... and {} more entries\n", file.transcript.len() - 5)); } } context }