From b2cf99c85712fc2a348687a020e9d9e5c8ef0d32 Mon Sep 17 00:00:00 2001 From: Cameron Date: Fri, 3 Apr 2026 17:25:35 -0400 Subject: [PATCH] feat: surface Ollama context token usage in agentic insight response Captures prompt_eval_count and eval_count from Ollama /api/chat responses during the agentic loop and returns them in POST /insights/generate/agentic so the frontend can display context window usage to the user. Co-Authored-By: Claude Sonnet 4.6 --- src/ai/handlers.rs | 12 +- src/ai/insight_generator.rs | 532 ++++++++++++++++++++++++++++++++++-- src/ai/ollama.rs | 22 +- 3 files changed, 530 insertions(+), 36 deletions(-) diff --git a/src/ai/handlers.rs b/src/ai/handlers.rs index f91268b..210aece 100644 --- a/src/ai/handlers.rs +++ b/src/ai/handlers.rs @@ -33,6 +33,10 @@ pub struct PhotoInsightResponse { pub summary: String, pub generated_at: i64, pub model_version: String, + #[serde(skip_serializing_if = "Option::is_none")] + pub prompt_eval_count: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub eval_count: Option, } #[derive(Debug, Serialize)] @@ -133,6 +137,8 @@ pub async fn get_insight_handler( summary: insight.summary, generated_at: insight.generated_at, model_version: insight.model_version, + prompt_eval_count: None, + eval_count: None, }; HttpResponse::Ok().json(response) } @@ -197,6 +203,8 @@ pub async fn get_all_insights_handler( summary: insight.summary, generated_at: insight.generated_at, model_version: insight.model_version, + prompt_eval_count: None, + eval_count: None, }) .collect(); @@ -263,7 +271,7 @@ pub async fn generate_agentic_insight_handler( .await; match result { - Ok(()) => { + Ok((prompt_eval_count, eval_count)) => { span.set_status(Status::Ok); // Fetch the stored insight to return it let otel_context = opentelemetry::Context::new(); @@ -277,6 +285,8 @@ pub async fn generate_agentic_insight_handler( summary: insight.summary, generated_at: insight.generated_at, model_version: insight.model_version, + prompt_eval_count, + eval_count, }; HttpResponse::Ok().json(response) } diff --git a/src/ai/insight_generator.rs b/src/ai/insight_generator.rs index cac03a4..e15e5ed 100644 --- a/src/ai/insight_generator.rs +++ b/src/ai/insight_generator.rs @@ -13,7 +13,8 @@ use crate::ai::ollama::{ChatMessage, OllamaClient, Tool}; use crate::ai::sms_client::SmsApiClient; use crate::database::models::InsertPhotoInsight; use crate::database::{ - CalendarEventDao, DailySummaryDao, ExifDao, InsightDao, LocationHistoryDao, SearchHistoryDao, + CalendarEventDao, DailySummaryDao, ExifDao, InsightDao, KnowledgeDao, LocationHistoryDao, + SearchHistoryDao, }; use crate::memories::extract_date_from_filename; use crate::otel::global_tracer; @@ -48,6 +49,9 @@ pub struct InsightGenerator { search_dao: Arc>>, tag_dao: Arc>>, + // Knowledge memory + knowledge_dao: Arc>>, + base_path: String, } @@ -62,6 +66,7 @@ impl InsightGenerator { location_dao: Arc>>, search_dao: Arc>>, tag_dao: Arc>>, + knowledge_dao: Arc>>, base_path: String, ) -> Self { Self { @@ -74,6 +79,7 @@ impl InsightGenerator { location_dao, search_dao, tag_dao, + knowledge_dao, base_path, } } @@ -1158,6 +1164,7 @@ impl InsightGenerator { summary, generated_at: Utc::now().timestamp(), model_version: ollama_client.primary_model.clone(), + is_current: true, }; let mut dao = self.insight_dao.lock().expect("Unable to lock InsightDao"); @@ -1345,6 +1352,7 @@ Return ONLY the summary, nothing else."#, arguments: &serde_json::Value, ollama: &OllamaClient, image_base64: &Option, + file_path: &str, cx: &opentelemetry::Context, ) -> String { let result = match tool_name { @@ -1355,16 +1363,16 @@ Return ONLY the summary, nothing else."#, "get_file_tags" => self.tool_get_file_tags(arguments, cx).await, "describe_photo" => self.tool_describe_photo(ollama, image_base64).await, "reverse_geocode" => self.tool_reverse_geocode(arguments).await, + "recall_entities" => self.tool_recall_entities(arguments, cx).await, + "recall_facts_for_photo" => self.tool_recall_facts_for_photo(arguments, cx).await, + "store_entity" => self.tool_store_entity(arguments, ollama, cx).await, + "store_fact" => self.tool_store_fact(arguments, file_path, cx).await, unknown => format!("Unknown tool: {}", unknown), }; if result.starts_with("Error") || result.starts_with("No ") { log::warn!("Tool '{}' result: {}", tool_name, result); } else { - log::info!( - "Tool '{}' result: {} chars", - tool_name, - result.len() - ); + log::info!("Tool '{}' result: {} chars", tool_name, result.len()); } result } @@ -1679,6 +1687,295 @@ Return ONLY the summary, nothing else."#, } } + /// Tool: recall_entities — search the knowledge memory for known entities + async fn tool_recall_entities( + &self, + args: &serde_json::Value, + cx: &opentelemetry::Context, + ) -> String { + use crate::database::EntityFilter; + + let name_search = args + .get("name") + .and_then(|v| v.as_str()) + .map(|s| s.to_string()); + let entity_type = args + .get("entity_type") + .and_then(|v| v.as_str()) + .map(|s| s.to_string()); + let limit = args.get("limit").and_then(|v| v.as_i64()).unwrap_or(10); + + log::info!( + "tool_recall_entities: name={:?}, type={:?}, limit={}", + name_search, + entity_type, + limit + ); + + let filter = EntityFilter { + entity_type, + status: Some("active".to_string()), + search: name_search, + limit, + offset: 0, + }; + + let mut kdao = self + .knowledge_dao + .lock() + .expect("Unable to lock KnowledgeDao"); + match kdao.list_entities(cx, filter) { + Ok((entities, _total)) if entities.is_empty() => { + "No known entities found matching the query.".to_string() + } + Ok((entities, _total)) => { + let lines: Vec = entities + .iter() + .map(|e| { + format!( + "ID:{} | {} | {} | {} | confidence:{:.2}", + e.id, e.entity_type, e.name, e.description, e.confidence + ) + }) + .collect(); + format!("Known entities:\n{}", lines.join("\n")) + } + Err(e) => format!("Error recalling entities: {:?}", e), + } + } + + /// Tool: recall_facts_for_photo — retrieve facts linked to a specific photo + async fn tool_recall_facts_for_photo( + &self, + args: &serde_json::Value, + cx: &opentelemetry::Context, + ) -> String { + let file_path = match args.get("file_path").and_then(|v| v.as_str()) { + Some(p) => p.to_string(), + None => return "Error: missing required parameter 'file_path'".to_string(), + }; + + log::info!("tool_recall_facts_for_photo: file_path={}", file_path); + + let mut kdao = self + .knowledge_dao + .lock() + .expect("Unable to lock KnowledgeDao"); + + // Fetch photo links to find which entities appear in this photo + let links = match kdao.get_links_for_photo(cx, &file_path) { + Ok(l) => l, + Err(e) => return format!("Error fetching photo links: {:?}", e), + }; + + if links.is_empty() { + return "No knowledge facts found for this photo.".to_string(); + } + + let mut output_lines = Vec::new(); + let entity_ids: Vec = links.iter().map(|l| l.entity_id).collect(); + + // For each linked entity, fetch its facts + for entity_id in entity_ids { + if let Ok(entity) = kdao.get_entity_by_id(cx, entity_id) { + if let Some(e) = entity { + let role = links + .iter() + .find(|l| l.entity_id == entity_id) + .map(|l| l.role.as_str()) + .unwrap_or("subject"); + output_lines.push(format!( + "Entity: {} ({}, role: {})", + e.name, e.entity_type, role + )); + if let Ok(facts) = kdao.get_facts_for_entity(cx, entity_id) { + for f in facts.iter().filter(|f| f.status == "active") { + let obj = if let Some(ref v) = f.object_value { + v.clone() + } else if let Some(oid) = f.object_entity_id { + kdao.get_entity_by_id(cx, oid) + .ok() + .flatten() + .map(|e| format!("{} (entity ID: {})", e.name, e.id)) + .unwrap_or_else(|| format!("entity:{}", oid)) + } else { + "(unknown)".to_string() + }; + output_lines.push(format!(" - {} {}", f.predicate, obj)); + } + } + } + } + } + + if output_lines.is_empty() { + "No active knowledge facts found for this photo.".to_string() + } else { + format!("Knowledge for this photo:\n{}", output_lines.join("\n")) + } + } + + /// Tool: store_entity — upsert an entity into the knowledge memory + async fn tool_store_entity( + &self, + args: &serde_json::Value, + ollama: &OllamaClient, + cx: &opentelemetry::Context, + ) -> String { + use crate::database::models::InsertEntity; + + let name = match args.get("name").and_then(|v| v.as_str()) { + Some(n) => n.to_string(), + None => return "Error: missing required parameter 'name'".to_string(), + }; + let entity_type = match args.get("entity_type").and_then(|v| v.as_str()) { + Some(t) => t.to_string(), + None => return "Error: missing required parameter 'entity_type'".to_string(), + }; + let description = args + .get("description") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + + log::info!( + "tool_store_entity: name='{}', type='{}', description='{}'", + name, + entity_type, + description + ); + + // Generate embedding for name + description (best-effort) + let embed_text = format!("{} {}", name, description); + let embedding: Option> = match ollama.generate_embedding(&embed_text).await { + Ok(vec) => { + let bytes: Vec = vec.iter().flat_map(|f| f.to_le_bytes()).collect(); + Some(bytes) + } + Err(e) => { + log::warn!("Embedding generation failed for entity '{}': {}", name, e); + None + } + }; + + let now = chrono::Utc::now().timestamp(); + let insert = InsertEntity { + name, + entity_type, + description, + embedding, + confidence: 0.6, + status: "active".to_string(), + created_at: now, + updated_at: now, + }; + + let mut kdao = self + .knowledge_dao + .lock() + .expect("Unable to lock KnowledgeDao"); + match kdao.upsert_entity(cx, insert) { + Ok(entity) => format!( + "Entity stored: ID:{} | {} | {} | confidence:{:.2}", + entity.id, entity.entity_type, entity.name, entity.confidence + ), + Err(e) => format!("Error storing entity: {:?}", e), + } + } + + /// Tool: store_fact — record a fact about an entity, linked to the current photo + async fn tool_store_fact( + &self, + args: &serde_json::Value, + file_path: &str, + cx: &opentelemetry::Context, + ) -> String { + use crate::database::models::{InsertEntityFact, InsertEntityPhotoLink}; + + let subject_entity_id = match args.get("subject_entity_id").and_then(|v| v.as_i64()) { + Some(id) => id as i32, + None => return "Error: missing required parameter 'subject_entity_id'".to_string(), + }; + let predicate = match args.get("predicate").and_then(|v| v.as_str()) { + Some(p) => p.to_string(), + None => return "Error: missing required parameter 'predicate'".to_string(), + }; + let object_entity_id = args + .get("object_entity_id") + .and_then(|v| v.as_i64()) + .map(|id| id as i32); + let object_value = args + .get("object_value") + .and_then(|v| v.as_str()) + .map(|s| s.to_string()); + + if object_entity_id.is_none() && object_value.is_none() { + return "Error: provide either object_entity_id or object_value".to_string(); + } + + let photo_role = args + .get("photo_role") + .and_then(|v| v.as_str()) + .unwrap_or("subject") + .to_string(); + + log::info!( + "tool_store_fact: entity_id={}, predicate='{}', object_entity_id={:?}, object_value={:?}, photo='{}'", + subject_entity_id, + predicate, + object_entity_id, + object_value, + file_path + ); + + let fact = InsertEntityFact { + subject_entity_id, + predicate, + object_entity_id, + object_value, + source_photo: Some(file_path.to_string()), + source_insight_id: None, // will be backfilled after store_insight + confidence: 0.6, + status: "active".to_string(), + created_at: chrono::Utc::now().timestamp(), + }; + + let mut kdao = self + .knowledge_dao + .lock() + .expect("Unable to lock KnowledgeDao"); + + // Upsert the fact (corroboration bumps confidence if duplicate) + let (stored_fact, is_new) = match kdao.upsert_fact(cx, fact) { + Ok(r) => r, + Err(e) => return format!("Error storing fact: {:?}", e), + }; + + // Upsert a photo link so this entity is associated with this photo + let link = InsertEntityPhotoLink { + entity_id: subject_entity_id, + file_path: file_path.to_string(), + role: photo_role, + }; + if let Err(e) = kdao.upsert_photo_link(cx, link) { + log::warn!( + "Failed to upsert photo link for entity {}: {:?}", + subject_entity_id, + e + ); + } + + let action = if is_new { + "Stored new fact" + } else { + "Corroborated existing fact" + }; + format!( + "{}: ID:{} | confidence:{:.2}", + action, stored_fact.id, stored_fact.confidence + ) + } + // ── Agentic insight generation ────────────────────────────────────── /// Build the list of tool definitions for the agentic loop @@ -1799,6 +2096,100 @@ Return ONLY the summary, nothing else."#, }), )); + // Knowledge memory tools + tools.push(Tool::function( + "recall_entities", + "Search the knowledge memory for people, places, events, or things previously learned from other photos. Use this to retrieve context about subjects appearing in this photo.", + serde_json::json!({ + "type": "object", + "properties": { + "name": { + "type": "string", + "description": "Name or partial name to search for (case-insensitive substring match)" + }, + "entity_type": { + "type": "string", + "enum": ["person", "place", "event", "thing"], + "description": "Filter by entity type (optional)" + }, + "limit": { + "type": "integer", + "description": "Maximum number of results to return (default: 10)" + } + } + }), + )); + + tools.push(Tool::function( + "recall_facts_for_photo", + "Retrieve all known facts linked to a specific photo from the knowledge memory. Use this at the start of insight generation to load any previously stored knowledge about subjects in this photo.", + serde_json::json!({ + "type": "object", + "required": ["file_path"], + "properties": { + "file_path": { + "type": "string", + "description": "The file path of the photo to retrieve facts for" + } + } + }), + )); + + tools.push(Tool::function( + "store_entity", + "Store or update a person, place, event, or thing in the knowledge memory. Call this when you identify a subject in this photo that should be remembered for future insights.", + serde_json::json!({ + "type": "object", + "required": ["name", "entity_type"], + "properties": { + "name": { + "type": "string", + "description": "The canonical name of the entity (e.g. 'John Smith', 'Banff National Park')" + }, + "entity_type": { + "type": "string", + "enum": ["person", "place", "event", "thing"], + "description": "The type of entity" + }, + "description": { + "type": "string", + "description": "A brief description of the entity" + } + } + }), + )); + + tools.push(Tool::function( + "store_fact", + "Record a fact about an entity in the knowledge memory. Provide EITHER object_entity_id (when the object is a known entity whose ID you have) OR object_value (for free-text attributes). The fact will be linked to the current photo automatically.", + serde_json::json!({ + "type": "object", + "required": ["subject_entity_id", "predicate"], + "properties": { + "subject_entity_id": { + "type": "integer", + "description": "The ID of the entity this fact is about (returned by store_entity or recall_entities)" + }, + "predicate": { + "type": "string", + "description": "The relationship or attribute (e.g. 'is_friend_of', 'located_in', 'attended_event', 'is_sibling_of')" + }, + "object_entity_id": { + "type": "integer", + "description": "Use when the object is a known entity (e.g. Cameron's entity ID for 'is_friend_of Cameron'). Takes precedence over object_value." + }, + "object_value": { + "type": "string", + "description": "Use for free-text attributes where the object is not a stored entity (e.g. 'Portland, Oregon', 'software engineer')" + }, + "photo_role": { + "type": "string", + "description": "How this entity appears in the photo (e.g. 'subject', 'background', 'location'). Defaults to 'subject'." + } + } + }), + )); + if has_vision { tools.push(Tool::function( "describe_photo", @@ -1822,7 +2213,7 @@ Return ONLY the summary, nothing else."#, custom_system_prompt: Option, num_ctx: Option, max_iterations: usize, - ) -> Result<()> { + ) -> Result<(Option, Option)> { let tracer = global_tracer(); let current_cx = opentelemetry::Context::current(); let mut span = tracer.start_with_context("ai.insight.generate_agentic", ¤t_cx); @@ -1976,7 +2367,46 @@ Return ONLY the summary, nothing else."#, .collect() }; - // 6. Load image if vision capable + // 6. Clear existing entity-photo links for this file so the run starts fresh, + // and ensure the owner entity (Cameron) exists so the agent can reference it. + let cameron_entity_id: Option = { + let mut kdao = self + .knowledge_dao + .lock() + .expect("Unable to lock KnowledgeDao"); + + if let Err(e) = kdao.delete_photo_links_for_file(&insight_cx, &file_path) { + log::warn!( + "Failed to clear entity_photo_links for {}: {:?}", + file_path, + e + ); + } + + // Upsert the owner entity so the agent always has a stable entity ID to reference. + let owner = crate::database::models::InsertEntity { + name: "Cameron".to_string(), + entity_type: "person".to_string(), + description: "The owner of this photo collection. All memories are written from Cameron's perspective.".to_string(), + embedding: None, + confidence: 1.0, + status: "active".to_string(), + created_at: Utc::now().timestamp(), + updated_at: Utc::now().timestamp(), + }; + match kdao.upsert_entity(&insight_cx, owner) { + Ok(e) => { + log::info!("Cameron entity ID: {}", e.id); + Some(e.id) + } + Err(e) => { + log::warn!("Failed to upsert Cameron entity: {:?}", e); + None + } + } + }; + + // 7. Load image if vision capable let image_base64 = if has_vision { match self.load_image_as_base64(&file_path) { Ok(b64) => { @@ -1992,21 +2422,39 @@ Return ONLY the summary, nothing else."#, None }; - // 7. Build system message - let base_system = "You are a personal photo memory assistant helping to reconstruct a memory from a photo.\n\n\ + // 8. Build system message + let cameron_id_note = match cameron_entity_id { + Some(id) => format!( + "\n\nYour identity in the knowledge store: Cameron (entity ID: {}). \ + When storing facts where you (Cameron) are the object — for example, someone is your friend, \ + sibling, or colleague — use subject_entity_id for the other person and set object_value to \ + \"Cameron\" (or use store_fact with the other person as subject). When storing facts about \ + Cameron directly, use {} as the subject_entity_id.", + id, id + ), + None => String::new(), + }; + let base_system = format!( + "You are a personal photo memory assistant helping to reconstruct a memory from a photo. \ + You are writing from the perspective of Cameron, the owner of this photo collection.{cameron_id_note}\n\n\ IMPORTANT INSTRUCTIONS:\n\ 1. You MUST call multiple tools to gather context BEFORE writing any final insight. Do not produce a final answer after only one or two tool calls.\n\ 2. Always call ALL of the following tools that are relevant: search_rag (search conversation summaries), get_sms_messages (fetch nearby messages), get_calendar_events (check what was happening that day), get_location_history (find where this was taken), get_file_tags (retrieve tags).\n\ - 3. Only produce your final insight AFTER you have gathered context from at least 3-4 tools.\n\ - 4. If a tool returns no results, that is useful information — continue calling the remaining tools anyway.\n\ - 5. Your final insight must be written in first person as Cameron, in a journal/memoir style."; + 3. Use recall_facts_for_photo to load any previously stored knowledge about subjects in this photo.\n\ + 4. Use recall_entities to look up known people, places, or things that appear in this photo.\n\ + 5. When you identify people, places, events, or notable things in this photo: use store_entity to record them and store_fact to record key facts (relationships, roles, attributes). This builds a persistent memory for future insights.\n\ + 6. Only produce your final insight AFTER you have gathered context from at least 3-4 tools.\n\ + 7. If a tool returns no results, that is useful information — continue calling the remaining tools anyway.\n\ + 8. Your final insight must be written in first person as Cameron, in a journal/memoir style.", + cameron_id_note = cameron_id_note + ); let system_content = if let Some(ref custom) = custom_system_prompt { format!("{}\n\n{}", custom, base_system) } else { base_system.to_string() }; - // 8. Build user message + // 9. Build user message let gps_info = exif .as_ref() .and_then(|e| { @@ -2045,10 +2493,10 @@ Return ONLY the summary, nothing else."#, tags_info, ); - // 9. Define tools + // 10. Define tools let tools = Self::build_tool_definitions(has_vision); - // 10. Build initial messages + // 11. Build initial messages let system_msg = ChatMessage::system(system_content); let mut user_msg = ChatMessage::user(user_content); if let Some(ref img) = image_base64 { @@ -2057,21 +2505,26 @@ Return ONLY the summary, nothing else."#, let mut messages = vec![system_msg, user_msg]; - // 11. Agentic loop + // 12. Agentic loop let loop_span = tracer.start_with_context("ai.agentic.loop", &insight_cx); let loop_cx = insight_cx.with_span(loop_span); let mut final_content = String::new(); let mut iterations_used = 0usize; + let mut last_prompt_eval_count: Option = None; + let mut last_eval_count: Option = None; for iteration in 0..max_iterations { iterations_used = iteration + 1; log::info!("Agentic iteration {}/{}", iteration + 1, max_iterations); - let response = ollama_client + let (response, prompt_tokens, eval_tokens) = ollama_client .chat_with_tools(messages.clone(), tools.clone()) .await?; + last_prompt_eval_count = prompt_tokens; + last_eval_count = eval_tokens; + // Sanitize tool call arguments before pushing back into history. // Some models occasionally return non-object arguments (bool, string, null) // which Ollama rejects when they are re-sent in a subsequent request. @@ -2107,6 +2560,7 @@ Return ONLY the summary, nothing else."#, &tool_call.function.arguments, &ollama_client, &image_base64, + &file_path, &loop_cx, ) .await; @@ -2129,7 +2583,10 @@ Return ONLY the summary, nothing else."#, messages.push(ChatMessage::user( "Based on the context gathered, please write the final photo insight: a title and a detailed personal summary. Write in first person as Cameron.", )); - let final_response = ollama_client.chat_with_tools(messages, vec![]).await?; + let (final_response, prompt_tokens, eval_tokens) = + ollama_client.chat_with_tools(messages, vec![]).await?; + last_prompt_eval_count = prompt_tokens; + last_eval_count = eval_tokens; final_content = final_response.content; } @@ -2138,7 +2595,7 @@ Return ONLY the summary, nothing else."#, .set_attribute(KeyValue::new("iterations_used", iterations_used as i64)); loop_cx.span().set_status(Status::Ok); - // 12. Generate title + // 13. Generate title let title = ollama_client .generate_photo_title(&final_content, custom_system_prompt.as_deref()) .await?; @@ -2150,21 +2607,23 @@ Return ONLY the summary, nothing else."#, &final_content[..final_content.len().min(200)] ); - // 13. Store + // 14. Store insight (returns the persisted row including its new id) let insight = InsertPhotoInsight { file_path: file_path.to_string(), title, summary: final_content, generated_at: Utc::now().timestamp(), model_version: ollama_client.primary_model.clone(), + is_current: true, }; - let mut dao = self.insight_dao.lock().expect("Unable to lock InsightDao"); - let result = dao - .store_insight(&insight_cx, insight) - .map_err(|e| anyhow::anyhow!("Failed to store agentic insight: {:?}", e)); + let stored = { + let mut dao = self.insight_dao.lock().expect("Unable to lock InsightDao"); + dao.store_insight(&insight_cx, insight) + .map_err(|e| anyhow::anyhow!("Failed to store agentic insight: {:?}", e)) + }; - match &result { + match &stored { Ok(_) => { log::info!("Successfully stored agentic insight for {}", file_path); insight_cx.span().set_status(Status::Ok); @@ -2175,8 +2634,25 @@ Return ONLY the summary, nothing else."#, } } - result?; - Ok(()) + let stored_insight = stored?; + + // 15. Backfill source_insight_id on all facts recorded for this photo during the loop + { + let mut kdao = self + .knowledge_dao + .lock() + .expect("Unable to lock KnowledgeDao"); + if let Err(e) = kdao.update_facts_insight_id(&insight_cx, &file_path, stored_insight.id) + { + log::warn!( + "Failed to backfill source_insight_id for {}: {:?}", + file_path, + e + ); + } + } + + Ok((last_prompt_eval_count, last_eval_count)) } /// Reverse geocode GPS coordinates to human-readable place names diff --git a/src/ai/ollama.rs b/src/ai/ollama.rs index 857427f..80dddcb 100644 --- a/src/ai/ollama.rs +++ b/src/ai/ollama.rs @@ -507,7 +507,7 @@ Analyze the image and use specific details from both the visual content and the &self, messages: Vec, tools: Vec, - ) -> Result { + ) -> Result<(ChatMessage, Option, Option)> { // Try primary server first log::info!( "Attempting chat_with_tools with primary server: {} (model: {})", @@ -519,9 +519,9 @@ Analyze the image and use specific details from both the visual content and the .await; match primary_result { - Ok(response) => { + Ok(result) => { log::info!("Successfully got chat_with_tools response from primary server"); - Ok(response) + Ok(result) } Err(e) => { log::warn!("Primary server chat_with_tools failed: {}", e); @@ -540,11 +540,11 @@ Analyze the image and use specific details from both the visual content and the .try_chat_with_tools(fallback_url, messages, tools) .await { - Ok(response) => { + Ok(result) => { log::info!( "Successfully got chat_with_tools response from fallback server" ); - Ok(response) + Ok(result) } Err(fallback_e) => { log::error!( @@ -571,7 +571,7 @@ Analyze the image and use specific details from both the visual content and the base_url: &str, messages: Vec, tools: Vec, - ) -> Result { + ) -> Result<(ChatMessage, Option, Option)> { let url = format!("{}/api/chat", base_url); let model = if base_url == self.primary_url { &self.primary_model @@ -623,7 +623,11 @@ Analyze the image and use specific details from both the visual content and the .await .with_context(|| "Failed to parse Ollama chat response")?; - Ok(chat_response.message) + Ok(( + chat_response.message, + chat_response.prompt_eval_count, + chat_response.eval_count, + )) } /// Generate an embedding vector for text using nomic-embed-text:v1.5 @@ -876,6 +880,10 @@ struct OllamaChatResponse { #[serde(default)] #[allow(dead_code)] done_reason: String, + #[serde(default)] + prompt_eval_count: Option, + #[serde(default)] + eval_count: Option, } #[derive(Deserialize)]