From 1c508cdabde80006292d53a6603e772ae100ba3f Mon Sep 17 00:00:00 2001 From: Cameron Cordes Date: Tue, 29 Sep 2026 14:35:45 -0400 Subject: [PATCH] Tracing: LLM clients + streaming spans on tracing, drop cx threading (#140) LlmClient methods no longer take an OpenTelemetry Context. llm.request is now a tracing span (otel.kind=client, llm.provider/operation/model) that parents on the caller's ambient span; traceparent injection reads the span's OTel context via OpenTelemetrySpanExt. The src/ai feature spans (handlers, insight generation, chat, persona chat, daily summaries, TTS) move to #[instrument] / info_span!, and every LLM-bearing tokio::spawn carries its parent span with .instrument(): the streaming chat turn (chat_turn_stream_inner), async chat / persona turns, background insight jobs and TTS jobs. Sub-spans in the agentic loops (ai.chat.loop, ai.agentic.loop) instrument each LLM / tool await. DAO calls keep their explicit Context (derived from the tracing span) until the DAO migration (#141). tests/otel_llm_spans.rs pins nesting + traceparent for the blocking and streaming paths across a spawn on a multi-thread runtime. Co-Authored-By: Claude Opus 5.5 --- src/ai/daily_summary_job.rs | 101 +++----- src/ai/handlers.rs | 417 +++++++++++++++++----------------- src/ai/insight_chat.rs | 182 ++++++++------- src/ai/insight_generator.rs | 319 ++++++++++---------------- src/ai/llamacpp.rs | 75 +++--- src/ai/llm_client.rs | 17 +- src/ai/local_llm.rs | 29 +-- src/ai/mod.rs | 14 +- src/ai/nl_query.rs | 5 +- src/ai/ollama.rs | 177 ++++++--------- src/ai/openrouter.rs | 74 +++--- src/ai/persona_chat.rs | 124 +++++----- src/ai/sms_client.rs | 2 - src/ai/tts.rs | 212 ++++++++--------- src/bin/test_daily_summary.rs | 8 +- src/knowledge.rs | 3 +- src/otel.rs | 58 ++--- src/reels/script.rs | 3 +- src/unified_search.rs | 2 +- tests/otel_llm_spans.rs | 181 +++++++++++++++ 20 files changed, 973 insertions(+), 1030 deletions(-) create mode 100644 tests/otel_llm_spans.rs diff --git a/src/ai/daily_summary_job.rs b/src/ai/daily_summary_job.rs index acf7451..cd7cf0e 100644 --- a/src/ai/daily_summary_job.rs +++ b/src/ai/daily_summary_job.rs @@ -1,14 +1,13 @@ use anyhow::Result; use chrono::{NaiveDate, Utc}; -use opentelemetry::KeyValue; -use opentelemetry::trace::{Span, Status, TraceContextExt, Tracer}; +use opentelemetry::trace::Status; use std::collections::HashMap; use std::sync::{Arc, Mutex}; use tokio::time::sleep; +use tracing_opentelemetry::OpenTelemetrySpanExt; use crate::ai::{EMBEDDING_MODEL, OllamaClient, SmsApiClient, SmsMessage, user_display_name}; use crate::database::{DailySummaryDao, InsertDailySummary}; -use crate::otel::global_tracer; /// Strip boilerplate prefixes and common phrases from summaries before embedding. /// This improves embedding diversity by removing structural similarity. @@ -155,6 +154,7 @@ pub fn strip_summary_boilerplate(summary: &str) -> String { /// Generate and embed daily conversation summaries for a date range /// Default: August 2024 ±30 days (July 1 - September 30, 2024) +#[tracing::instrument(name = "ai.daily_summary.generate_batch", skip_all)] pub async fn generate_daily_summaries( contact: &str, start_date: Option, @@ -163,30 +163,16 @@ pub async fn generate_daily_summaries( sms_client: &SmsApiClient, summary_dao: Arc>>, ) -> Result<()> { - let tracer = global_tracer(); - - // Get current context (empty in background task) and start span with it - let current_cx = opentelemetry::Context::current(); - let mut span = tracer.start_with_context("ai.daily_summary.generate_batch", ¤t_cx); - span.set_attribute(KeyValue::new("contact", contact.to_string())); - - // Create context with this span for child operations - let parent_cx = current_cx.with_span(span); + let span = tracing::Span::current(); + span.set_attribute("contact", contact.to_string()); // Default to August 2024 ±30 days let start = start_date.unwrap_or_else(|| NaiveDate::from_ymd_opt(2024, 7, 1).unwrap()); let end = end_date.unwrap_or_else(|| NaiveDate::from_ymd_opt(2024, 9, 30).unwrap()); - parent_cx - .span() - .set_attribute(KeyValue::new("start_date", start.to_string())); - parent_cx - .span() - .set_attribute(KeyValue::new("end_date", end.to_string())); - parent_cx.span().set_attribute(KeyValue::new( - "date_range_days", - (end - start).num_days() + 1, - )); + span.set_attribute("start_date", start.to_string()); + span.set_attribute("end_date", end.to_string()); + span.set_attribute("date_range_days", (end - start).num_days() + 1); log::info!("========================================"); log::info!("Starting daily summary generation for {}", contact); @@ -267,15 +253,8 @@ pub async fn generate_daily_summaries( } // Generate summary for this day - match generate_and_store_daily_summary( - &parent_cx, - date, - contact, - messages, - ollama, - summary_dao.clone(), - ) - .await + match generate_and_store_daily_summary(date, contact, messages, ollama, summary_dao.clone()) + .await { Ok(_) => { processed += 1; @@ -322,55 +301,38 @@ pub async fn generate_daily_summaries( log::info!("========================================"); // Record final metrics in span - parent_cx - .span() - .set_attribute(KeyValue::new("days_processed", processed as i64)); - parent_cx - .span() - .set_attribute(KeyValue::new("days_skipped", skipped as i64)); - parent_cx - .span() - .set_attribute(KeyValue::new("days_failed", failed as i64)); - parent_cx - .span() - .set_attribute(KeyValue::new("total_days", total_days as i64)); + span.set_attribute("days_processed", processed as i64); + span.set_attribute("days_skipped", skipped as i64); + span.set_attribute("days_failed", failed as i64); + span.set_attribute("total_days", total_days as i64); if failed > 0 { - parent_cx - .span() - .set_status(Status::error(format!("{} days failed to process", failed))); + span.set_status(Status::error(format!("{} days failed to process", failed))); } else { - parent_cx.span().set_status(Status::Ok); + span.set_status(Status::Ok); } Ok(()) } /// Generate and store a single day's summary +#[tracing::instrument(name = "ai.daily_summary.generate_single", skip_all)] async fn generate_and_store_daily_summary( - parent_cx: &opentelemetry::Context, date: &NaiveDate, contact: &str, messages: &[SmsMessage], ollama: &OllamaClient, summary_dao: Arc>>, ) -> Result<()> { - let tracer = global_tracer(); - let mut span = tracer.start_with_context("ai.daily_summary.generate_single", parent_cx); - span.set_attribute(KeyValue::new("date", date.to_string())); - span.set_attribute(KeyValue::new("contact", contact.to_string())); - span.set_attribute(KeyValue::new("message_count", messages.len() as i64)); - // Parent context for the child llm.request spans; the span itself stays - // borrowed here for the attribute updates below. - let llm_cx = - opentelemetry::Context::new().with_remote_span_context(span.span_context().clone()); + let span = tracing::Span::current(); + span.set_attribute("date", date.to_string()); + span.set_attribute("contact", contact.to_string()); + span.set_attribute("message_count", messages.len() as i64); let (prompt, system_prompt) = build_daily_summary_prompt(contact, date, messages); // Generate summary with LLM - let summary = ollama - .generate(&llm_cx, &prompt, Some(system_prompt)) - .await?; + let summary = ollama.generate(&prompt, Some(system_prompt)).await?; log::debug!( "Generated summary for {}: {}", @@ -378,7 +340,7 @@ async fn generate_and_store_daily_summary( summary.chars().take(100).collect::() ); - span.set_attribute(KeyValue::new("summary_length", summary.len() as i64)); + span.set_attribute("summary_length", summary.len() as i64); // Strip boilerplate before embedding to improve vector diversity let stripped_summary = strip_summary_boilerplate(&summary); @@ -388,14 +350,9 @@ async fn generate_and_store_daily_summary( ); // Embed the stripped summary (store original summary in DB) - let embedding = ollama - .generate_embedding(&llm_cx, &stripped_summary) - .await?; + let embedding = ollama.generate_embedding(&stripped_summary).await?; - span.set_attribute(KeyValue::new( - "embedding_dimensions", - embedding.len() as i64, - )); + span.set_attribute("embedding_dimensions", embedding.len() as i64); // Store in database let insert = InsertDailySummary { @@ -408,8 +365,8 @@ async fn generate_and_store_daily_summary( model_version: EMBEDDING_MODEL.to_string(), }; - // Create context from current span for DB operation - let child_cx = opentelemetry::Context::current_with_span(span); + // DAO calls still take an explicit OTel context (#141). + let child_cx = span.context(); let mut dao = summary_dao.lock().expect("Unable to lock DailySummaryDao"); let result = dao @@ -417,8 +374,8 @@ async fn generate_and_store_daily_summary( .map_err(|e| anyhow::anyhow!("Failed to store summary: {:?}", e)); match &result { - Ok(_) => child_cx.span().set_status(Status::Ok), - Err(e) => child_cx.span().set_status(Status::error(e.to_string())), + Ok(_) => span.set_status(Status::Ok), + Err(e) => span.set_status(Status::error(e.to_string())), } result?; diff --git a/src/ai/handlers.rs b/src/ai/handlers.rs index 82a5406..7d3df91 100644 --- a/src/ai/handlers.rs +++ b/src/ai/handlers.rs @@ -1,8 +1,9 @@ -use actix_web::{HttpRequest, HttpResponse, Responder, delete, get, post, web}; +use actix_web::{HttpResponse, Responder, delete, get, post, web}; use futures::StreamExt; -use opentelemetry::KeyValue; -use opentelemetry::trace::{Span, Status, Tracer}; +use opentelemetry::trace::Status; use serde::{Deserialize, Serialize}; +use tracing::Instrument; +use tracing_opentelemetry::OpenTelemetrySpanExt; use crate::ai::insight_chat::{ChatStreamEvent, ChatTurnRequest}; use crate::ai::ollama::ChatMessage; @@ -11,7 +12,6 @@ use crate::data::Claims; use crate::database::models::{InsightGenerationType, InsightJobStatus, PhotoInsight}; use crate::database::{ExifDao, InsightDao}; use crate::libraries; -use crate::otel::{extract_context_from_request, global_tracer}; use crate::state::AppState; use crate::utils::normalize_path; @@ -416,15 +416,13 @@ fn apply_owner_defaults( /// POST /insights/generate - Generate insight for a specific photo (async) #[post("/insights/generate")] +#[tracing::instrument(name = "http.insights.generate", skip_all)] pub async fn generate_insight_handler( - http_request: HttpRequest, _claims: Claims, request: web::Json, app_state: web::Data, ) -> impl Responder { - let parent_context = extract_context_from_request(&http_request); - let tracer = global_tracer(); - let mut span = tracer.start_with_context("http.insights.generate", &parent_context); + let span = tracing::Span::current(); let mut request = request.0; // Stored owner defaults fill what the client omitted (request → pref → @@ -440,9 +438,9 @@ pub async fn generate_insight_handler( let library = app_state.primary_library(); let gen_type = InsightGenerationType::Standard; - span.set_attribute(KeyValue::new("file_path", normalized_path.clone())); + span.set_attribute("file_path", normalized_path.clone()); if let Some(ref model) = request.model { - span.set_attribute(KeyValue::new("model", model.clone())); + span.set_attribute("model", model.clone()); } log::info!( @@ -506,96 +504,102 @@ pub async fn generate_insight_handler( let job_handles = app_state.insight_job_handles.clone(); let path = normalized_path.clone(); - let handle = tokio::spawn(async move { - let timeout_secs: u64 = std::env::var("INSIGHT_GENERATION_TIMEOUT_SECS") - .ok() - .and_then(|v| v.parse().ok()) - .unwrap_or(120); + let handle = tokio::spawn( + async move { + let timeout_secs: u64 = std::env::var("INSIGHT_GENERATION_TIMEOUT_SECS") + .ok() + .and_then(|v| v.parse().ok()) + .unwrap_or(120); - let path_for_task = path.clone(); - let generator_for_task = generator.clone(); - let result = tokio::task::spawn(async move { - // Cross-model barrier: if a TTS synthesis holds the GPU, wait it - // out BEFORE the generation wall-clock starts. The per-request - // lease keeps reqwest budgets honest, but this job-level timeout - // would otherwise burn while the first chat call queues behind a - // multi-minute synthesis. Dropped immediately — holding it across - // the generation would deadlock the chat calls' own leases. - drop(crate::ai::gpu::llm_lease().await); - tokio::time::timeout( - std::time::Duration::from_secs(timeout_secs), - generator_for_task.generate_insight_for_photo_with_config( - &path_for_task, - request.model.clone(), - request.system_prompt.clone(), - request.num_ctx, - request.temperature, - request.top_p, - request.top_k, - request.min_p, - ), - ) - .await - }) - .await; + let path_for_task = path.clone(); + let generator_for_task = generator.clone(); + let result = tokio::task::spawn(async move { + // Cross-model barrier: if a TTS synthesis holds the GPU, wait it + // out BEFORE the generation wall-clock starts. The per-request + // lease keeps reqwest budgets honest, but this job-level timeout + // would otherwise burn while the first chat call queues behind a + // multi-minute synthesis. Dropped immediately — holding it across + // the generation would deadlock the chat calls' own leases. + drop(crate::ai::gpu::llm_lease().await); + tokio::time::timeout( + std::time::Duration::from_secs(timeout_secs), + generator_for_task.generate_insight_for_photo_with_config( + &path_for_task, + request.model.clone(), + request.system_prompt.clone(), + request.num_ctx, + request.temperature, + request.top_p, + request.top_k, + request.min_p, + ), + ) + .await + }) + .await; - let ctx = opentelemetry::Context::new(); - let mut dao = job_dao.lock().expect("Unable to lock InsightJobDao"); + let ctx = opentelemetry::Context::new(); + let mut dao = job_dao.lock().expect("Unable to lock InsightJobDao"); - match result { - Ok(Ok(Ok(()))) => { - let mut insight_dao = generator - .insight_dao() - .lock() - .expect("Unable to lock InsightDao"); - let insight_id = insight_dao - .get_insight(&ctx, &path) - .ok() - .flatten() - .map(|i| i.id); - if let Some(id) = insight_id { - if let Err(e) = dao.complete_job(&ctx, job_id, id) { - log::error!("Failed to mark job {} as completed: {:?}", job_id, e); + match result { + Ok(Ok(Ok(()))) => { + let mut insight_dao = generator + .insight_dao() + .lock() + .expect("Unable to lock InsightDao"); + let insight_id = insight_dao + .get_insight(&ctx, &path) + .ok() + .flatten() + .map(|i| i.id); + if let Some(id) = insight_id { + if let Err(e) = dao.complete_job(&ctx, job_id, id) { + log::error!("Failed to mark job {} as completed: {:?}", job_id, e); + } + } else if let Err(e) = + dao.fail_job(&ctx, job_id, "generation returned no insight") + { + log::error!("Failed to mark job {} as failed: {:?}", job_id, e); + } + } + Ok(Ok(Err(e))) => { + log::error!("Insight generation failed for {}: {:?}", path, e); + // `{:#}` = one-line context chain; the job's error_message is + // returned to the client verbatim, so no Debug/backtrace here. + if let Err(err) = dao.fail_job(&ctx, job_id, &format!("{:#}", e)) { + log::error!("Failed to mark job {} as failed: {:?}", job_id, err); + } + } + Ok(Err(_)) => { + log::error!( + "Insight generation timed out for {} after {}s", + path, + timeout_secs + ); + if let Err(err) = + dao.fail_job(&ctx, job_id, &format!("timeout after {}s", timeout_secs)) + { + log::error!("Failed to mark job {} as failed: {:?}", job_id, err); + } + } + Err(_) => { + log::error!("Insight generation task panicked for {}", path); + if let Err(err) = dao.fail_job(&ctx, job_id, "generation task panicked") { + log::error!("Failed to mark job {} as failed: {:?}", job_id, err); } - } else if let Err(e) = dao.fail_job(&ctx, job_id, "generation returned no insight") - { - log::error!("Failed to mark job {} as failed: {:?}", job_id, e); } } - Ok(Ok(Err(e))) => { - log::error!("Insight generation failed for {}: {:?}", path, e); - // `{:#}` = one-line context chain; the job's error_message is - // returned to the client verbatim, so no Debug/backtrace here. - if let Err(err) = dao.fail_job(&ctx, job_id, &format!("{:#}", e)) { - log::error!("Failed to mark job {} as failed: {:?}", job_id, err); - } - } - Ok(Err(_)) => { - log::error!( - "Insight generation timed out for {} after {}s", - path, - timeout_secs - ); - if let Err(err) = - dao.fail_job(&ctx, job_id, &format!("timeout after {}s", timeout_secs)) - { - log::error!("Failed to mark job {} as failed: {:?}", job_id, err); - } - } - Err(_) => { - log::error!("Insight generation task panicked for {}", path); - if let Err(err) = dao.fail_job(&ctx, job_id, "generation task panicked") { - log::error!("Failed to mark job {} as failed: {:?}", job_id, err); - } - } - } - // Remove handle from map on completion - let mut handles = job_handles - .lock() - .expect("Unable to lock InsightJobHandles"); - handles.remove(&job_id); - }); + // Remove handle from map on completion + let mut handles = job_handles + .lock() + .expect("Unable to lock InsightJobHandles"); + handles.remove(&job_id); + } + // Carry the handler span into the job so the generation (and its + // LLM / DB spans) stays in the request's trace. + .instrument(tracing::Span::current()), + ); // Store abort handle { @@ -606,7 +610,7 @@ pub async fn generate_insight_handler( handles.insert(job_id, handle.abort_handle()); } - span.set_attribute(KeyValue::new("job_id", job_id as i64)); + span.set_attribute("job_id", job_id as i64); span.set_status(Status::Ok); HttpResponse::Accepted().json(JobIdResponse { job_id }) } @@ -744,15 +748,13 @@ pub async fn get_insight_history_handler( /// POST /insights/generate/agentic - Generate insight using agentic tool-calling loop (async) #[post("/insights/generate/agentic")] +#[tracing::instrument(name = "http.insights.generate_agentic", skip_all)] pub async fn generate_agentic_insight_handler( - http_request: HttpRequest, claims: Claims, request: web::Json, app_state: web::Data, ) -> impl Responder { - let parent_context = extract_context_from_request(&http_request); - let tracer = global_tracer(); - let mut span = tracer.start_with_context("http.insights.generate_agentic", &parent_context); + let span = tracing::Span::current(); let mut request = request.0; // Stored owner defaults fill what the client omitted (request → pref → @@ -768,12 +770,12 @@ pub async fn generate_agentic_insight_handler( let library = app_state.primary_library(); let gen_type = InsightGenerationType::Agentic; - span.set_attribute(KeyValue::new("file_path", normalized_path.clone())); + span.set_attribute("file_path", normalized_path.clone()); if let Some(ref model) = request.model { - span.set_attribute(KeyValue::new("model", model.clone())); + span.set_attribute("model", model.clone()); } if let Some(ref backend) = request.backend { - span.set_attribute(KeyValue::new("backend", backend.clone())); + span.set_attribute("backend", backend.clone()); } log::info!( @@ -901,91 +903,97 @@ pub async fn generate_agentic_insight_handler( let job_handles = app_state.insight_job_handles.clone(); let path = normalized_path.clone(); - let handle = tokio::spawn(async move { - let timeout_secs: u64 = std::env::var("INSIGHT_GENERATION_TIMEOUT_SECS") - .ok() - .and_then(|v| v.parse().ok()) - .unwrap_or(180); + let handle = tokio::spawn( + async move { + let timeout_secs: u64 = std::env::var("INSIGHT_GENERATION_TIMEOUT_SECS") + .ok() + .and_then(|v| v.parse().ok()) + .unwrap_or(180); - let path_for_task = path.clone(); - let generator_for_task = generator.clone(); - let result = tokio::task::spawn(async move { - // Cross-model barrier — see generate_insight_handler: wait out any - // running TTS synthesis before the generation wall-clock starts. - drop(crate::ai::gpu::llm_lease().await); - tokio::time::timeout( - std::time::Duration::from_secs(timeout_secs), - generator_for_task.generate_agentic_insight_for_photo( - &path_for_task, - request.model.clone(), - system_prompt, - request.num_ctx, - request.temperature, - request.top_p, - request.top_k, - request.min_p, - request.enable_thinking, - max_iterations, - request.backend.clone(), - fewshot_examples, - fewshot_ids, - user_id, - persona_id, - ), - ) - .await - }) - .await; + let path_for_task = path.clone(); + let generator_for_task = generator.clone(); + let result = tokio::task::spawn(async move { + // Cross-model barrier — see generate_insight_handler: wait out any + // running TTS synthesis before the generation wall-clock starts. + drop(crate::ai::gpu::llm_lease().await); + tokio::time::timeout( + std::time::Duration::from_secs(timeout_secs), + generator_for_task.generate_agentic_insight_for_photo( + &path_for_task, + request.model.clone(), + system_prompt, + request.num_ctx, + request.temperature, + request.top_p, + request.top_k, + request.min_p, + request.enable_thinking, + max_iterations, + request.backend.clone(), + fewshot_examples, + fewshot_ids, + user_id, + persona_id, + ), + ) + .await + }) + .await; - let ctx = opentelemetry::Context::new(); - let mut dao = job_dao.lock().expect("Unable to lock InsightJobDao"); + let ctx = opentelemetry::Context::new(); + let mut dao = job_dao.lock().expect("Unable to lock InsightJobDao"); - match result { - Ok(Ok(Ok((Some(insight_id), _, _)))) => { - if let Err(e) = dao.complete_job(&ctx, job_id, insight_id) { - log::error!("Failed to mark job {} as completed: {:?}", job_id, e); + match result { + Ok(Ok(Ok((Some(insight_id), _, _)))) => { + if let Err(e) = dao.complete_job(&ctx, job_id, insight_id) { + log::error!("Failed to mark job {} as completed: {:?}", job_id, e); + } } - } - Ok(Ok(Ok((None, _, _)))) => { - if let Err(e) = dao.fail_job(&ctx, job_id, "agentic generation returned no insight") - { - log::error!("Failed to mark job {} as failed: {:?}", job_id, e); + Ok(Ok(Ok((None, _, _)))) => { + if let Err(e) = + dao.fail_job(&ctx, job_id, "agentic generation returned no insight") + { + log::error!("Failed to mark job {} as failed: {:?}", job_id, e); + } } - } - Ok(Ok(Err(e))) => { - log::error!("Agentic insight generation failed for {}: {:?}", path, e); - // `{:#}` = one-line context chain; the job's error_message is - // returned to the client verbatim, so no Debug/backtrace here. - if let Err(err) = dao.fail_job(&ctx, job_id, &format!("{:#}", e)) { - log::error!("Failed to mark job {} as failed: {:?}", job_id, err); + Ok(Ok(Err(e))) => { + log::error!("Agentic insight generation failed for {}: {:?}", path, e); + // `{:#}` = one-line context chain; the job's error_message is + // returned to the client verbatim, so no Debug/backtrace here. + if let Err(err) = dao.fail_job(&ctx, job_id, &format!("{:#}", e)) { + log::error!("Failed to mark job {} as failed: {:?}", job_id, err); + } } - } - Ok(Err(_)) => { - log::error!( - "Agentic insight generation timed out for {} after {}s", - path, - timeout_secs - ); - if let Err(err) = - dao.fail_job(&ctx, job_id, &format!("timeout after {}s", timeout_secs)) - { - log::error!("Failed to mark job {} as failed: {:?}", job_id, err); + Ok(Err(_)) => { + log::error!( + "Agentic insight generation timed out for {} after {}s", + path, + timeout_secs + ); + if let Err(err) = + dao.fail_job(&ctx, job_id, &format!("timeout after {}s", timeout_secs)) + { + log::error!("Failed to mark job {} as failed: {:?}", job_id, err); + } } - } - Err(_) => { - log::error!("Agentic insight generation task panicked for {}", path); - if let Err(err) = dao.fail_job(&ctx, job_id, "generation task panicked") { - log::error!("Failed to mark job {} as failed: {:?}", job_id, err); + Err(_) => { + log::error!("Agentic insight generation task panicked for {}", path); + if let Err(err) = dao.fail_job(&ctx, job_id, "generation task panicked") { + log::error!("Failed to mark job {} as failed: {:?}", job_id, err); + } } } + + // Remove handle from map on completion + let mut handles = job_handles + .lock() + .expect("Unable to lock InsightJobHandles"); + handles.remove(&job_id); } - - // Remove handle from map on completion - let mut handles = job_handles - .lock() - .expect("Unable to lock InsightJobHandles"); - handles.remove(&job_id); - }); + // Carry the handler span into the job so the generation (and its + // LLM / DB spans) stays in the request's trace. + .instrument(tracing::Span::current()), + ); // Store abort handle { @@ -996,7 +1004,7 @@ pub async fn generate_agentic_insight_handler( handles.insert(job_id, handle.abort_handle()); } - span.set_attribute(KeyValue::new("job_id", job_id as i64)); + span.set_attribute("job_id", job_id as i64); span.set_status(Status::Ok); HttpResponse::Accepted().json(JobIdResponse { job_id }) } @@ -1270,16 +1278,14 @@ pub struct ChatTurnHttpResponse { /// POST /insights/chat — submit a follow-up turn against an existing insight. #[post("/insights/chat")] +#[tracing::instrument(name = "http.insights.chat", skip_all)] pub async fn chat_turn_handler( - http_request: HttpRequest, claims: Claims, request: web::Json, app_state: web::Data, ) -> impl Responder { - let parent_context = extract_context_from_request(&http_request); - let tracer = global_tracer(); - let mut span = tracer.start_with_context("http.insights.chat", &parent_context); - span.set_attribute(KeyValue::new("file_path", request.file_path.clone())); + let span = tracing::Span::current(); + span.set_attribute("file_path", request.file_path.clone()); let mut request = request.0; // Stored owner defaults fill what the client omitted (request → pref → @@ -1667,12 +1673,10 @@ pub struct ChatSwitchBranchRequest { /// Returns `text/event-stream` with one event per chat stream event. #[post("/insights/chat/stream")] pub async fn chat_stream_handler( - http_request: HttpRequest, claims: Claims, request: web::Json, app_state: web::Data, ) -> HttpResponse { - let parent_context = extract_context_from_request(&http_request); let mut request = request.0; // Stored owner defaults fill what the client omitted (request → pref → // hard default; the hard default still applies downstream). @@ -1718,7 +1722,7 @@ pub async fn chat_stream_handler( }; let service = app_state.insight_chat.clone(); - let events = service.chat_turn_stream(chat_req, parent_context); + let events = service.chat_turn_stream(chat_req); // Map ChatStreamEvent → SSE frame bytes. let sse_stream = futures::stream::StreamExt::map(events, |ev| { @@ -1819,16 +1823,14 @@ fn sse_event_payload(ev: &ChatStreamEvent) -> (&'static str, serde_json::Value) /// POST /insights/chat/turn — async turn dispatch. Returns turn_id immediately, /// client then polls GET /insights/chat/turn/{turn_id} for SSE replay. #[post("/insights/chat/turn")] +#[tracing::instrument(name = "http.insights.chat_turn_async", skip_all)] pub async fn turn_async_handler( - http_request: HttpRequest, claims: Claims, request: web::Json, app_state: web::Data, ) -> impl Responder { - let parent_context = extract_context_from_request(&http_request); - let tracer = global_tracer(); - let mut span = tracer.start_with_context("http.insights.chat_turn_async", &parent_context); - span.set_attribute(KeyValue::new("file_path", request.file_path.clone())); + let span = tracing::Span::current(); + span.set_attribute("file_path", request.file_path.clone()); let mut request = request.0; // Stored owner defaults fill what the client omitted (request → pref → @@ -1877,7 +1879,7 @@ pub async fn turn_async_handler( let registry = app_state.turn_registry.clone(); let turn_id = service.chat_turn_async(registry, chat_req).await; - span.set_attribute(KeyValue::new("turn_id", turn_id.clone())); + span.set_attribute("turn_id", turn_id.clone()); span.set_status(Status::Ok); HttpResponse::Accepted().json(serde_json::json!({ "turn_id": turn_id, @@ -1896,13 +1898,13 @@ pub struct ReplayQuery { /// GET /insights/chat/turn/{turn_id} — SSE replay stream. #[get("/insights/chat/turn/{turn_id}")] +#[tracing::instrument(name = "ai.chat.turn.replay", skip_all)] pub async fn turn_replay_handler( - http_request: HttpRequest, path: web::Path, query: web::Query, app_state: web::Data, ) -> HttpResponse { - turn_replay_impl(http_request, path, query, app_state).await + turn_replay_impl(path, query, app_state).await } /// Core of the SSE replay, kept attribute-free so the persona-chat routes @@ -1910,7 +1912,6 @@ pub async fn turn_replay_handler( /// on `turn_id` only, and `render_turn_info_frame` already scopes the /// identity (persona turns carry `persona_id` and no `file_path`). pub(crate) async fn turn_replay_impl( - http_request: HttpRequest, path: web::Path, query: web::Query, app_state: web::Data, @@ -1920,11 +1921,9 @@ pub(crate) async fn turn_replay_impl( let turn_id = path.into_inner(); let skip_before = query.skip_before.unwrap_or(0); - let parent_context = extract_context_from_request(&http_request); - let tracer = global_tracer(); - let mut span = tracer.start_with_context("ai.chat.turn.replay", &parent_context); - span.set_attribute(KeyValue::new("turn_id", turn_id.clone())); - span.set_attribute(KeyValue::new("skip_before", skip_before as i64)); + let span = tracing::Span::current(); + span.set_attribute("turn_id", turn_id.clone()); + span.set_attribute("skip_before", skip_before as i64); let registry = app_state.turn_registry.clone(); let entry = match registry.get(&turn_id).await { @@ -1938,11 +1937,8 @@ pub(crate) async fn turn_replay_impl( }; let info = entry.info().await; - span.set_attribute(KeyValue::new("status", info.status.as_str())); - span.set_attribute(KeyValue::new( - "event_count", - info.total_events_pushed as i64, - )); + span.set_attribute("status", info.status.as_str()); + span.set_attribute("event_count", info.total_events_pushed as i64); let turn_info_frame = render_turn_info_frame(&info); // Initial buffered batch: events produced before this connection attached. @@ -2074,27 +2070,24 @@ fn render_turn_info_frame(info: &crate::ai::turn_registry::TurnInfo) -> String { /// DELETE /insights/chat/turn/{turn_id} — cancel a running turn. #[delete("/insights/chat/turn/{turn_id}")] +#[tracing::instrument(name = "ai.chat.turn.cancel", skip_all)] pub async fn cancel_turn_handler( - http_request: HttpRequest, path: web::Path, app_state: web::Data, ) -> impl Responder { - cancel_turn_impl(http_request, path, app_state).await + cancel_turn_impl(path, app_state).await } /// Core of the turn-cancel, attribute-free so the persona-chat DELETE route /// can reuse it under `/persona_chat/turn/{turn_id}`. pub(crate) async fn cancel_turn_impl( - http_request: HttpRequest, path: web::Path, app_state: web::Data, ) -> impl Responder { let turn_id = path.into_inner(); - let parent_context = extract_context_from_request(&http_request); - let tracer = global_tracer(); - let mut span = tracer.start_with_context("ai.chat.turn.cancel", &parent_context); - span.set_attribute(KeyValue::new("turn_id", turn_id.clone())); + let span = tracing::Span::current(); + span.set_attribute("turn_id", turn_id.clone()); let registry = app_state.turn_registry.clone(); let entry = match registry.get(&turn_id).await { @@ -2111,7 +2104,7 @@ pub(crate) async fn cancel_turn_impl( // also checks `is_running()` at each iteration boundary as a graceful // backstop in case the abort lands between await points. let aborted = entry.abort(); - span.set_attribute(KeyValue::new("aborted", aborted)); + span.set_attribute("aborted", aborted); // Push the terminal event BEFORE flipping status: a replay reader treats a // terminal status with no buffered tail as "closed", so the Done must be diff --git a/src/ai/insight_chat.rs b/src/ai/insight_chat.rs index b11a1e8..3f495f2 100644 --- a/src/ai/insight_chat.rs +++ b/src/ai/insight_chat.rs @@ -1,7 +1,6 @@ use anyhow::{Result, anyhow, bail}; use chrono::{TimeZone, Utc}; -use opentelemetry::KeyValue; -use opentelemetry::trace::{Span, Status, TraceContextExt, Tracer}; +use opentelemetry::trace::Status; use std::collections::HashMap; use std::sync::{Arc, Mutex}; use tokio::sync::Mutex as TokioMutex; @@ -15,9 +14,10 @@ use crate::ai::turn_registry::TurnEntry; use crate::ai::turn_registry::TurnRegistry; use crate::database::InsightDao; use crate::database::models::InsertPhotoInsight; -use crate::otel::global_tracer; use crate::utils::{normalize_path, retry_with_backoff}; use futures::stream::{BoxStream, StreamExt}; +use tracing::Instrument; +use tracing_opentelemetry::OpenTelemetrySpanExt; use uuid::Uuid; pub const DEFAULT_MAX_ITERATIONS: usize = 6; @@ -200,13 +200,12 @@ impl InsightChatService { }) } + #[tracing::instrument(name = "ai.insight.chat_turn", skip_all)] pub async fn chat_turn(&self, req: ChatTurnRequest) -> Result { - let tracer = global_tracer(); - let parent_cx = opentelemetry::Context::new(); - let mut span = tracer.start_with_context("ai.insight.chat_turn", &parent_cx); - span.set_attribute(KeyValue::new("file_path", req.file_path.clone())); - span.set_attribute(KeyValue::new("library_id", req.library_id as i64)); - span.set_attribute(KeyValue::new("amend", req.amend)); + let span = tracing::Span::current(); + span.set_attribute("file_path", req.file_path.clone()); + span.set_attribute("library_id", req.library_id as i64); + span.set_attribute("amend", req.amend); if req.user_message.trim().is_empty() { bail!("user_message must not be empty"); @@ -220,7 +219,7 @@ impl InsightChatService { .clone() .filter(|s| !s.trim().is_empty()) .unwrap_or_else(|| "default".to_string()); - span.set_attribute(KeyValue::new("persona_id", active_persona.clone())); + span.set_attribute("persona_id", active_persona.clone()); let normalized = normalize_path(&req.file_path); @@ -278,13 +277,13 @@ impl InsightChatService { .unwrap_or_else(|| stored_backend.clone()); validate_cross_replay(&stored_backend, &effective_backend)?; let kind = BackendKind::parse(&effective_backend)?; - span.set_attribute(KeyValue::new("backend", kind.as_str())); + span.set_attribute("backend", kind.as_str()); let max_iterations = req .max_iterations .unwrap_or(DEFAULT_MAX_ITERATIONS) .clamp(1, env_max_iterations()); - span.set_attribute(KeyValue::new("max_iterations", max_iterations as i64)); + span.set_attribute("max_iterations", max_iterations as i64); let stored_model = insight.model_version.clone(); let overrides = SamplingOverrides { @@ -302,7 +301,7 @@ impl InsightChatService { }; let backend = self.generator.resolve_backend(kind, &overrides).await?; let model_used = backend.model().to_string(); - span.set_attribute(KeyValue::new("model", model_used.clone())); + span.set_attribute("model", model_used.clone()); // 5. Decide vision + tool set. In hybrid (describe-then-inline) mode // we omit `describe_photo`. Otherwise trust the stored history: @@ -333,7 +332,7 @@ impl InsightChatService { let budget_bytes = budget_tokens.saturating_mul(BYTES_PER_TOKEN); let truncated = apply_context_budget(&mut messages, budget_bytes); if truncated { - span.set_attribute(KeyValue::new("history_truncated", true)); + span.set_attribute("history_truncated", true); } // Everything in `messages` at this point is replayed history that is // ALREADY in the tree; only what the turn appends past this mark @@ -359,12 +358,12 @@ impl InsightChatService { // accumulate across turns. let original_system_content = annotate_system_with_budget(&mut messages, max_iterations); - let insight_cx = parent_cx.with_span(span); - // 8. Agentic loop — same shape as insight_generator's, but capped // tighter and dispatching tools through the shared executor. - let loop_span = tracer.start_with_context("ai.chat.loop", &insight_cx); - let loop_cx = insight_cx.with_span(loop_span); + // Each LLM / tool await runs inside `loop_span` so their spans + // nest under it. + let loop_span = tracing::info_span!("ai.chat.loop"); + let loop_cx = loop_span.context(); let mut tool_calls_made = 0usize; let mut iterations_used = 0usize; let mut last_prompt_eval_count: Option = None; @@ -377,7 +376,8 @@ impl InsightChatService { let (response, prompt_tokens, eval_tokens) = backend .chat() - .chat_with_tools(&loop_cx, messages.clone(), tools.clone()) + .chat_with_tools(messages.clone(), tools.clone()) + .instrument(loop_span.clone()) .await?; last_prompt_eval_count = prompt_tokens; last_eval_count = eval_tokens; @@ -416,6 +416,7 @@ impl InsightChatService { &active_persona, &loop_cx, ) + .instrument(loop_span.clone()) .await; messages.push(ChatMessage::tool_result(result)); } @@ -434,7 +435,8 @@ impl InsightChatService { let synthetic_idx = push_synthetic_final_prompt(&mut messages); let (final_response, prompt_tokens, eval_tokens) = backend .chat() - .chat_with_tools(&loop_cx, messages.clone(), vec![]) + .chat_with_tools(messages.clone(), vec![]) + .instrument(loop_span.clone()) .await?; last_prompt_eval_count = prompt_tokens; last_eval_count = eval_tokens; @@ -449,7 +451,8 @@ impl InsightChatService { // return / persist as the reply (the raw transcript keeps it). let final_content = crate::ai::llm_client::strip_think_blocks(&final_content); - loop_cx.span().set_status(Status::Ok); + loop_span.set_status(Status::Ok); + drop(loop_span); // Drop the per-turn iteration-budget note from the system message // before we persist so it doesn't snowball on each subsequent turn. @@ -492,7 +495,6 @@ impl InsightChatService { let title_raw = backend .chat() .generate( - &loop_cx, &title_prompt, Some( "You are my long term memory assistant. Use only the information provided. Do not invent details.", @@ -779,14 +781,22 @@ impl InsightChatService { /// /// The stream takes ownership of the service via `Arc` (passed by /// the caller) so it can live past the handler's await boundary. + /// + /// Must be called from inside the request's span: the turn span is + /// created here, eagerly, so it parents on the caller. The returned + /// stream is polled by the HTTP body writer, outside any span. pub fn chat_turn_stream( self: Arc, req: ChatTurnRequest, - cx: opentelemetry::Context, ) -> BoxStream<'static, ChatStreamEvent> { + let span = tracing::info_span!( + "ai.chat.turn.stream", + file_path = %req.file_path, + library_id = req.library_id, + ); let svc = self; let s = async_stream::stream! { - match svc.chat_turn_stream_inner(req, cx, Ok).await { + match svc.chat_turn_stream_inner(req, span, Ok).await { Ok(mut rx) => { while let Some(ev) = rx.recv().await { yield ev; @@ -808,7 +818,7 @@ impl InsightChatService { async fn chat_turn_stream_inner( self: Arc, req: ChatTurnRequest, - cx: opentelemetry::Context, + span: tracing::Span, _ev_mapper: F, ) -> Result> where @@ -816,12 +826,22 @@ impl InsightChatService { { let (tx, rx) = tokio::sync::mpsc::channel::(64); let svc = self.clone(); - tokio::spawn(async move { - let result = svc.run_streaming_turn(req, cx, tx.clone()).await; - if let Err(e) = result { - let _ = tx.send(ChatStreamEvent::Error(format!("{}", e))).await; + // Carry the turn span into the task explicitly — tracing context is + // thread-local, so without `.instrument` every LLM / DB span inside + // the turn would orphan onto a fresh trace. + tokio::spawn( + async move { + let result = svc.run_streaming_turn(req, tx.clone()).await; + let span = tracing::Span::current(); + if let Err(e) = result { + span.set_status(Status::error(format!("{e}"))); + let _ = tx.send(ChatStreamEvent::Error(format!("{}", e))).await; + } else { + span.set_status(Status::Ok); + } } - }); + .instrument(span), + ); Ok(rx) } @@ -843,35 +863,38 @@ impl InsightChatService { let svc = self.clone(); let entry_clone = entry.clone(); - let turn_id_for_span = turn_id.clone(); - let library_id = req.library_id; - let handle = tokio::spawn(async move { - // Span covering the whole spawned turn execution. Created here (not - // in the HTTP handler) because the dispatch span ends at the 202 - // response, long before this work runs. - let tracer = global_tracer(); - let mut span = tracer.start("ai.chat.turn.execute"); - span.set_attribute(KeyValue::new("turn_id", turn_id_for_span)); - span.set_attribute(KeyValue::new("library_id", library_id as i64)); - - let result = svc - .run_streaming_turn_with_entry(req, entry_clone.clone(), &mut span) - .await; - if let Err(ref e) = result { - span.set_attribute(KeyValue::new("status", "error")); - span.set_status(Status::error(format!("{e}"))); - // Push the terminal event BEFORE flipping status: a replay - // reader treats a terminal status with no buffered tail as - // "closed", so the Error must be in the buffer first. - let _ = entry_clone - .push_event(ChatStreamEvent::Error(format!("{}", e))) + // Span covering the whole spawned turn execution. Created before the + // spawn so it parents on the dispatching request (same trace as the + // 202), and carried into the task with `.instrument`. It outlives the + // dispatch span, which ends at the 202 response. + let turn_span = tracing::info_span!( + "ai.chat.turn.execute", + turn_id = %turn_id, + library_id = req.library_id, + ); + let handle = tokio::spawn( + async move { + let span = tracing::Span::current(); + let result = svc + .run_streaming_turn_with_entry(req, entry_clone.clone(), &span) .await; - entry_clone.set_terminal_status(crate::ai::turn_registry::TurnStatus::Error); - } else { - span.set_attribute(KeyValue::new("status", "done")); - span.set_status(Status::Ok); + if let Err(ref e) = result { + span.set_attribute("status", "error"); + span.set_status(Status::error(format!("{e}"))); + // Push the terminal event BEFORE flipping status: a replay + // reader treats a terminal status with no buffered tail as + // "closed", so the Error must be in the buffer first. + let _ = entry_clone + .push_event(ChatStreamEvent::Error(format!("{}", e))) + .await; + entry_clone.set_terminal_status(crate::ai::turn_registry::TurnStatus::Error); + } else { + span.set_attribute("status", "done"); + span.set_status(Status::Ok); + } } - }); + .instrument(turn_span), + ); // Install the abort handle so DELETE can actually stop the task. entry.set_abort_handle(handle.abort_handle()); @@ -885,7 +908,7 @@ impl InsightChatService { self: Arc, req: ChatTurnRequest, entry: Arc, - turn_span: &mut opentelemetry::global::BoxedSpan, + turn_span: &tracing::Span, ) -> Result<()> { if req.user_message.trim().is_empty() { bail!("user_message must not be empty"); @@ -930,7 +953,7 @@ impl InsightChatService { normalized: String, insight: crate::database::models::PhotoInsight, entry: Arc, - turn_span: &mut opentelemetry::global::BoxedSpan, + turn_span: &tracing::Span, ) -> Result<()> { let active_persona = req .persona_id @@ -988,10 +1011,8 @@ impl InsightChatService { let model_used = backend.model().to_string(); // llm.* attributes on the turn-level span (issue #130): clients can't // know local-vs-hybrid, but the turn knows its resolved backend. - turn_span.set_attribute(KeyValue::new("llm.backend", kind.as_str())); - turn_span.set_attribute(KeyValue::new("llm.model", model_used.clone())); - let turn_cx = opentelemetry::Context::new() - .with_remote_span_context(turn_span.span_context().clone()); + turn_span.set_attribute("llm.backend", kind.as_str()); + turn_span.set_attribute("llm.model", model_used.clone()); let local_first_user_has_image = messages .iter() @@ -1041,7 +1062,6 @@ impl InsightChatService { &active_persona, max_iterations, &entry, - &turn_cx, ) .await?; let AgenticLoopOutcome { @@ -1158,7 +1178,7 @@ impl InsightChatService { req: ChatTurnRequest, normalized: String, entry: Arc, - turn_span: &mut opentelemetry::global::BoxedSpan, + turn_span: &tracing::Span, ) -> Result<()> { let active_persona = req .persona_id @@ -1185,10 +1205,8 @@ impl InsightChatService { let backend = self.generator.resolve_backend(kind, &overrides).await?; let model_used = backend.model().to_string(); // llm.* attributes on the turn-level span (issue #130). - turn_span.set_attribute(KeyValue::new("llm.backend", kind.as_str())); - turn_span.set_attribute(KeyValue::new("llm.model", model_used.clone())); - let turn_cx = opentelemetry::Context::new() - .with_remote_span_context(turn_span.span_context().clone()); + turn_span.set_attribute("llm.backend", kind.as_str()); + turn_span.set_attribute("llm.model", model_used.clone()); let image_base64: Option = self.generator.load_image_as_base64(&normalized).ok(); @@ -1208,7 +1226,7 @@ impl InsightChatService { let visual_block = if !backend.images_inline { match image_base64.as_deref() { - Some(b64) => match backend.local().describe_image(&turn_cx, b64).await { + Some(b64) => match backend.local().describe_image(b64).await { Ok(desc) => { format!("Visual description (from local vision model):\n{}\n", desc) } @@ -1263,7 +1281,6 @@ impl InsightChatService { &active_persona, max_iterations, &entry, - &turn_cx, ) .await?; let AgenticLoopOutcome { @@ -1352,7 +1369,6 @@ impl InsightChatService { active_persona: &str, max_iterations: usize, entry: &Arc, - cx: &opentelemetry::Context, ) -> Result { crate::ai::insight_chat::run_streaming_agentic_loop_with_entry( &self.generator, @@ -1365,7 +1381,6 @@ impl InsightChatService { active_persona, max_iterations, entry, - cx, ) .await } @@ -1373,7 +1388,6 @@ impl InsightChatService { async fn run_streaming_turn( self: Arc, req: ChatTurnRequest, - cx: opentelemetry::Context, tx: tokio::sync::mpsc::Sender, ) -> Result<()> { if req.user_message.trim().is_empty() { @@ -1410,10 +1424,10 @@ impl InsightChatService { }; if req.regenerate || existing_insight.is_none() { - return self.run_bootstrap_streaming(req, normalized, cx, tx).await; + return self.run_bootstrap_streaming(req, normalized, tx).await; } let insight = existing_insight.expect("just checked Some above"); - self.run_continuation_streaming(req, normalized, cx, insight, tx) + self.run_continuation_streaming(req, normalized, insight, tx) .await } @@ -1425,7 +1439,6 @@ impl InsightChatService { &self, req: ChatTurnRequest, normalized: String, - cx: opentelemetry::Context, insight: crate::database::models::PhotoInsight, tx: tokio::sync::mpsc::Sender, ) -> Result<()> { @@ -1523,7 +1536,6 @@ impl InsightChatService { &active_persona, max_iterations, &tx, - &cx, ) .await?; let AgenticLoopOutcome { @@ -1635,7 +1647,6 @@ impl InsightChatService { &self, req: ChatTurnRequest, normalized: String, - cx: opentelemetry::Context, tx: tokio::sync::mpsc::Sender, ) -> Result<()> { let active_persona = req @@ -1692,7 +1703,7 @@ impl InsightChatService { // images_inline backends send images directly to the chat model. let visual_block = if !backend.images_inline { match image_base64.as_deref() { - Some(b64) => match backend.local().describe_image(&cx, b64).await { + Some(b64) => match backend.local().describe_image(b64).await { Ok(desc) => { format!("Visual description (from local vision model):\n{}\n", desc) } @@ -1759,7 +1770,6 @@ impl InsightChatService { &active_persona, max_iterations, &tx, - &cx, ) .await?; let AgenticLoopOutcome { @@ -1842,7 +1852,6 @@ impl InsightChatService { active_persona: &str, max_iterations: usize, tx: &tokio::sync::mpsc::Sender, - cx: &opentelemetry::Context, ) -> Result { let mut tool_calls_made = 0usize; let mut iterations_used = 0usize; @@ -1861,7 +1870,7 @@ impl InsightChatService { let mut stream = backend .chat() - .chat_with_tools_stream(cx, messages.clone(), tools.clone()) + .chat_with_tools_stream(messages.clone(), tools.clone()) .await?; let mut final_message: Option = None; @@ -1955,7 +1964,7 @@ impl InsightChatService { let synthetic_idx = push_synthetic_final_prompt(messages); let mut stream = backend .chat() - .chat_with_tools_stream(cx, messages.clone(), vec![]) + .chat_with_tools_stream(messages.clone(), vec![]) .await?; let mut final_message: Option = None; while let Some(ev) = stream.next().await { @@ -2335,7 +2344,6 @@ pub async fn run_streaming_agentic_loop_with_entry( active_persona: &str, max_iterations: usize, entry: &Arc, - cx: &opentelemetry::Context, ) -> Result { let mut tool_calls_made = 0usize; let mut iterations_used = 0usize; @@ -2365,7 +2373,7 @@ pub async fn run_streaming_agentic_loop_with_entry( let mut stream = backend .chat() - .chat_with_tools_stream(cx, messages.clone(), tools.clone()) + .chat_with_tools_stream(messages.clone(), tools.clone()) .await?; let mut final_message: Option = None; @@ -2448,7 +2456,7 @@ pub async fn run_streaming_agentic_loop_with_entry( let synthetic_idx = push_synthetic_final_prompt(messages); let mut stream = backend .chat() - .chat_with_tools_stream(cx, messages.clone(), vec![]) + .chat_with_tools_stream(messages.clone(), vec![]) .await?; let mut final_message: Option = None; while let Some(ev) = stream.next().await { diff --git a/src/ai/insight_generator.rs b/src/ai/insight_generator.rs index 1546eb3..95dc47d 100644 --- a/src/ai/insight_generator.rs +++ b/src/ai/insight_generator.rs @@ -2,13 +2,14 @@ use anyhow::Result; use base64::Engine as _; use chrono::{DateTime, Local, NaiveDate, TimeZone, Utc}; use image::ImageFormat; -use opentelemetry::KeyValue; -use opentelemetry::trace::{Span, Status, TraceContextExt, Tracer}; +use opentelemetry::trace::Status; use serde::Deserialize; use std::fs::File; use std::io::Cursor; use std::str::FromStr; use std::sync::{Arc, Mutex}; +use tracing::Instrument; +use tracing_opentelemetry::OpenTelemetrySpanExt; use crate::ai::apollo_client::{ApolloClient, ApolloPlace}; use crate::ai::backend::{BackendKind, ResolvedBackend, SamplingOverrides}; @@ -25,7 +26,6 @@ use crate::database::{ }; use crate::libraries::Library; use crate::memories::extract_date_from_filename; -use crate::otel::global_tracer; use crate::tags::TagDao; use crate::utils::{earliest_fs_time, normalize_path, retry_with_backoff, retry_with_backoff_if}; @@ -502,9 +502,9 @@ impl InsightGenerator { /// Find relevant messages using RAG, excluding recent messages (>30 days ago) /// This prevents RAG from returning messages already in the immediate time window + #[tracing::instrument(name = "ai.rag.filter_historical", skip_all)] async fn find_relevant_messages_rag_historical( &self, - parent_cx: &opentelemetry::Context, date: chrono::NaiveDate, location: Option<&str>, contact: Option<&str>, @@ -512,33 +512,20 @@ impl InsightGenerator { limit: usize, extra_context: Option<&str>, ) -> Result> { - let tracer = global_tracer(); - let span = tracer.start_with_context("ai.rag.filter_historical", parent_cx); - let filter_cx = parent_cx.with_span(span); + let span = tracing::Span::current(); - filter_cx - .span() - .set_attribute(KeyValue::new("date", date.to_string())); - filter_cx - .span() - .set_attribute(KeyValue::new("limit", limit as i64)); - filter_cx - .span() - .set_attribute(KeyValue::new("exclusion_window_days", 30)); + span.set_attribute("date", date.to_string()); + span.set_attribute("limit", limit as i64); + span.set_attribute("exclusion_window_days", 30); if let Some(t) = topics { - filter_cx - .span() - .set_attribute(KeyValue::new("topics", t.join(", "))); + span.set_attribute("topics", t.join(", ")); } let query_results = self .find_relevant_messages_rag(date, location, contact, topics, limit * 2, extra_context) .await?; - filter_cx.span().set_attribute(KeyValue::new( - "rag_results_count", - query_results.len() as i64, - )); + span.set_attribute("rag_results_count", query_results.len() as i64); // Filter out messages from within 30 days of the photo date let photo_timestamp = date @@ -576,17 +563,15 @@ impl InsightGenerator { historical_only.len() ); - filter_cx.span().set_attribute(KeyValue::new( - "historical_results_count", - historical_only.len() as i64, - )); - filter_cx.span().set_status(Status::Ok); + span.set_attribute("historical_results_count", historical_only.len() as i64); + span.set_status(Status::Ok); Ok(historical_only) } /// Find relevant daily summaries using RAG (semantic search) /// Returns formatted daily summary strings for LLM context + #[tracing::instrument(name = "ai.rag.search_daily_summaries", skip_all)] async fn find_relevant_messages_rag( &self, date: chrono::NaiveDate, @@ -596,16 +581,14 @@ impl InsightGenerator { limit: usize, extra_context: Option<&str>, ) -> Result> { - let tracer = global_tracer(); - let current_cx = opentelemetry::Context::current(); - let mut span = tracer.start_with_context("ai.rag.search_daily_summaries", ¤t_cx); - span.set_attribute(KeyValue::new("date", date.to_string())); - span.set_attribute(KeyValue::new("limit", limit as i64)); + let span = tracing::Span::current(); + span.set_attribute("date", date.to_string()); + span.set_attribute("limit", limit as i64); if let Some(loc) = location { - span.set_attribute(KeyValue::new("location", loc.to_string())); + span.set_attribute("location", loc.to_string()); } if let Some(c) = contact { - span.set_attribute(KeyValue::new("contact", c.to_string())); + span.set_attribute("contact", c.to_string()); } // Build query string - prioritize topics if available (semantically meaningful) @@ -633,10 +616,10 @@ impl InsightGenerator { base_query }; - span.set_attribute(KeyValue::new("query", query.clone())); + span.set_attribute("query", query.clone()); // Create context with this span for child operations - let search_cx = current_cx.with_span(span); + let search_cx = span.context(); log::info!("========================================"); log::info!("RAG QUERY: {}", query); @@ -646,8 +629,7 @@ impl InsightGenerator { // (`LLM_BACKEND` switch). Must match the backend that populated the // daily-summary embeddings or similarity search will be garbage. let query_embedding = - crate::ai::embed_query(&search_cx, &self.ollama, self.llamacpp.as_deref(), &query) - .await?; + crate::ai::embed_query(&self.ollama, self.llamacpp.as_deref(), &query).await?; // Search for similar daily summaries with time-based weighting // This prioritizes summaries temporally close to the query date @@ -666,10 +648,7 @@ impl InsightGenerator { similar_summaries.len() ); - search_cx.span().set_attribute(KeyValue::new( - "results_count", - similar_summaries.len() as i64, - )); + span.set_attribute("results_count", similar_summaries.len() as i64); // Format daily summaries for LLM context let formatted = similar_summaries @@ -682,7 +661,7 @@ impl InsightGenerator { }) .collect(); - search_cx.span().set_status(Status::Ok); + span.set_status(Status::Ok); Ok(formatted) } @@ -691,30 +670,28 @@ impl InsightGenerator { /// tool. Embeds the caller's query as-is (no metadata boilerplate) and /// only applies time weighting when an anchor date is provided — /// without one, results rank purely by similarity across all time. + #[tracing::instrument(name = "ai.rag.search_daily_summaries", skip_all)] async fn search_summaries_semantic( &self, query: &str, date: Option, limit: usize, ) -> Result> { - let tracer = global_tracer(); - let current_cx = opentelemetry::Context::current(); - let mut span = tracer.start_with_context("ai.rag.search_daily_summaries", ¤t_cx); - span.set_attribute(KeyValue::new("query", query.to_string())); - span.set_attribute(KeyValue::new("limit", limit as i64)); - span.set_attribute(KeyValue::new("time_weighted", date.is_some())); + let span = tracing::Span::current(); + span.set_attribute("query", query.to_string()); + span.set_attribute("limit", limit as i64); + span.set_attribute("time_weighted", date.is_some()); if let Some(d) = date { - span.set_attribute(KeyValue::new("date", d.to_string())); + span.set_attribute("date", d.to_string()); } - let search_cx = current_cx.with_span(span); + let search_cx = span.context(); log::info!("RAG QUERY: {} (anchor date: {:?})", query, date); // Must use the same backend that populated the daily-summary // embeddings or similarity search is garbage (see embed_one docs). let query_embedding = - crate::ai::embed_query(&search_cx, &self.ollama, self.llamacpp.as_deref(), query) - .await?; + crate::ai::embed_query(&self.ollama, self.llamacpp.as_deref(), query).await?; let mut summary_dao = self .daily_summary_dao @@ -732,11 +709,8 @@ impl InsightGenerator { } .map_err(|e| anyhow::anyhow!("Failed to find similar summaries: {:?}", e))?; - search_cx.span().set_attribute(KeyValue::new( - "results_count", - similar_summaries.len() as i64, - )); - search_cx.span().set_status(Status::Ok); + span.set_attribute("results_count", similar_summaries.len() as i64); + span.set_status(Status::Ok); Ok(similar_summaries .into_iter() @@ -789,20 +763,17 @@ impl InsightGenerator { /// Gather calendar context for photo timestamp /// Uses hybrid time + semantic search (±7 days, ranked by relevance) + #[tracing::instrument(name = "ai.context.calendar", skip_all)] async fn gather_calendar_context( &self, - parent_cx: &opentelemetry::Context, timestamp: i64, location: Option<&str>, ) -> Result> { - let tracer = global_tracer(); - let span = tracer.start_with_context("ai.context.calendar", parent_cx); - let calendar_cx = parent_cx.with_span(span); + let span = tracing::Span::current(); + let calendar_cx = span.context(); let query_embedding = if let Some(loc) = location { - match crate::ai::embed_query(&calendar_cx, &self.ollama, self.llamacpp.as_deref(), loc) - .await - { + match crate::ai::embed_query(&self.ollama, self.llamacpp.as_deref(), loc).await { Ok(emb) => Some(emb), Err(e) => { log::warn!("Failed to generate embedding for location '{}': {}", loc, e); @@ -828,7 +799,7 @@ impl InsightGenerator { .ok() }; - calendar_cx.span().set_status(Status::Ok); + span.set_status(Status::Ok); if let Some(events) = events { if events.is_empty() { @@ -866,15 +837,14 @@ impl InsightGenerator { /// Gather location context for photo timestamp /// Finds nearest location record (±30 minutes) + #[tracing::instrument(name = "ai.context.location", skip_all)] async fn gather_location_context( &self, - parent_cx: &opentelemetry::Context, timestamp: i64, exif_gps: Option<(f64, f64)>, ) -> Result> { - let tracer = global_tracer(); - let span = tracer.start_with_context("ai.context.location", parent_cx); - let location_cx = parent_cx.with_span(span); + let span = tracing::Span::current(); + let location_cx = span.context(); let nearest = { let mut dao = self @@ -890,7 +860,7 @@ impl InsightGenerator { .flatten() }; - location_cx.span().set_status(Status::Ok); + span.set_status(Status::Ok); if let Some(loc) = nearest { // Check if this adds NEW information compared to EXIF @@ -943,17 +913,16 @@ impl InsightGenerator { /// Gather search context for photo date /// Uses semantic search on queries (±30 days, top 5 relevant) + #[tracing::instrument(name = "ai.context.search", skip_all)] async fn gather_search_context( &self, - parent_cx: &opentelemetry::Context, timestamp: i64, location: Option<&str>, contact: Option<&str>, enrichment: Option<&str>, ) -> Result> { - let tracer = global_tracer(); - let span = tracer.start_with_context("ai.context.search", parent_cx); - let search_cx = parent_cx.with_span(span); + let span = tracing::Span::current(); + let search_cx = span.context(); // Use enrichment (topics + photo description + tags) if available; // fall back to generic temporal query. @@ -973,23 +942,18 @@ impl InsightGenerator { ) }; - let query_embedding = match crate::ai::embed_query( - &search_cx, - &self.ollama, - self.llamacpp.as_deref(), - &query_text, - ) - .await - { - Ok(emb) => emb, - Err(e) => { - log::warn!("Failed to generate search embedding: {}", e); - search_cx.span().set_status(Status::Error { - description: e.to_string().into(), - }); - return Ok(None); - } - }; + let query_embedding = + match crate::ai::embed_query(&self.ollama, self.llamacpp.as_deref(), &query_text).await + { + Ok(emb) => emb, + Err(e) => { + log::warn!("Failed to generate search embedding: {}", e); + span.set_status(Status::Error { + description: e.to_string().into(), + }); + return Ok(None); + } + }; let searches = { let mut dao = self @@ -1006,7 +970,7 @@ impl InsightGenerator { .ok() }; - search_cx.span().set_status(Status::Ok); + span.set_status(Status::Ok); if let Some(searches) = searches { if searches.is_empty() { @@ -1068,6 +1032,7 @@ impl InsightGenerator { } /// Generate AI insight for a single photo with custom configuration + #[tracing::instrument(name = "ai.insight.generate", skip_all)] pub async fn generate_insight_for_photo_with_config( &self, file_path: &str, @@ -1079,20 +1044,18 @@ impl InsightGenerator { top_k: Option, min_p: Option, ) -> Result<()> { - let tracer = global_tracer(); - let current_cx = opentelemetry::Context::current(); - let mut span = tracer.start_with_context("ai.insight.generate", ¤t_cx); + let span = tracing::Span::current(); // Normalize path to ensure consistent forward slashes in database let file_path = normalize_path(file_path); log::info!("Generating insight for photo: {}", file_path); - span.set_attribute(KeyValue::new("file_path", file_path.clone())); + span.set_attribute("file_path", file_path.clone()); // Create custom Ollama client if model is specified let mut ollama_client = if let Some(model) = custom_model { log::info!("Using custom model: {}", model); - span.set_attribute(KeyValue::new("custom_model", model.clone())); + span.set_attribute("custom_model", model.clone()); OllamaClient::new( self.ollama.primary_url.clone(), self.ollama.fallback_url.clone(), @@ -1100,14 +1063,14 @@ impl InsightGenerator { Some(model), // Use the same custom model for fallback server ) } else { - span.set_attribute(KeyValue::new("model", self.ollama.primary_model.clone())); + span.set_attribute("model", self.ollama.primary_model.clone()); self.ollama.clone() }; // Set context size if specified if let Some(ctx) = num_ctx { log::info!("Using custom context size: {}", ctx); - span.set_attribute(KeyValue::new("num_ctx", ctx as i64)); + span.set_attribute("num_ctx", ctx as i64); ollama_client.set_num_ctx(Some(ctx)); } @@ -1121,22 +1084,22 @@ impl InsightGenerator { min_p ); if let Some(t) = temperature { - span.set_attribute(KeyValue::new("temperature", t as f64)); + span.set_attribute("temperature", t as f64); } if let Some(p) = top_p { - span.set_attribute(KeyValue::new("top_p", p as f64)); + span.set_attribute("top_p", p as f64); } if let Some(k) = top_k { - span.set_attribute(KeyValue::new("top_k", k as i64)); + span.set_attribute("top_k", k as i64); } if let Some(m) = min_p { - span.set_attribute(KeyValue::new("min_p", m as f64)); + span.set_attribute("min_p", m as f64); } ollama_client.set_sampling_params(temperature, top_p, top_k, min_p); } // Create context with this span for child operations - let insight_cx = current_cx.with_span(span); + let insight_cx = span.context(); // 1. Get EXIF data for the photo let exif = { @@ -1153,13 +1116,9 @@ impl InsightGenerator { let contact = Self::extract_contact_from_path(&file_path); log::info!("Extracted contact from path: {:?}", contact); - insight_cx - .span() - .set_attribute(KeyValue::new("date_taken", date_taken.to_string())); + span.set_attribute("date_taken", date_taken.to_string()); if let Some(ref c) = contact { - insight_cx - .span() - .set_attribute(KeyValue::new("contact", c.clone())); + span.set_attribute("contact", c.clone()); } // Fetch file tags (used to enrich RAG and final context) @@ -1197,9 +1156,7 @@ impl InsightGenerator { .into_iter() .next(); let combined = compose_location_string(apollo_primary, nominatim, lat, lon); - insight_cx - .span() - .set_attribute(KeyValue::new("location", combined.clone())); + span.set_attribute("location", combined.clone()); Some(combined) } else { None @@ -1234,9 +1191,7 @@ impl InsightGenerator { } }; - insight_cx - .span() - .set_attribute(KeyValue::new("model_has_vision", has_vision)); + span.set_attribute("model_has_vision", has_vision); // Load image and encode as base64 only if model supports vision let image_base64 = if has_vision { @@ -1263,10 +1218,7 @@ impl InsightGenerator { // Generate brief photo description for RAG enrichment (vision models only) let photo_description: Option = if let Some(ref img_b64) = image_base64 { - match ollama_client - .generate_photo_description(&insight_cx, img_b64) - .await - { + match ollama_client.generate_photo_description(img_b64).await { Ok(desc) => { log::info!("Photo description for RAG enrichment: {}", desc); Some(desc) @@ -1335,7 +1287,7 @@ impl InsightGenerator { if !immediate_messages.is_empty() { // Step 2: Extract topics from immediate messages to enrich RAG query let topics = self - .extract_topics_from_messages(&immediate_messages, &ollama_client, &insight_cx) + .extract_topics_from_messages(&immediate_messages, &ollama_client) .await; log::info!("Extracted topics for query enrichment: {:?}", topics); @@ -1367,7 +1319,6 @@ impl InsightGenerator { }; match self .find_relevant_messages_rag_historical( - &insight_cx, date_taken, None, contact.as_deref(), @@ -1391,7 +1342,6 @@ impl InsightGenerator { &immediate_messages, &ollama_client, custom_system_prompt.as_deref(), - &insight_cx, ) .await .unwrap_or_else(|| String::from("No immediate context")); @@ -1401,7 +1351,6 @@ impl InsightGenerator { &historical_messages, &ollama_client, custom_system_prompt.as_deref(), - &insight_cx, ) .await .unwrap_or_else(|| String::from("No historical context")); @@ -1420,7 +1369,6 @@ impl InsightGenerator { &immediate_messages, &ollama_client, custom_system_prompt.as_deref(), - &insight_cx, ) .await; } @@ -1431,7 +1379,6 @@ impl InsightGenerator { &immediate_messages, &ollama_client, custom_system_prompt.as_deref(), - &insight_cx, ) .await; } @@ -1457,7 +1404,6 @@ impl InsightGenerator { &rag_messages, &ollama_client, custom_system_prompt.as_deref(), - &insight_cx, ) .await; } @@ -1490,7 +1436,7 @@ impl InsightGenerator { if !sms_messages.is_empty() { match self .sms_client - .summarize_context(&sms_messages, &ollama_client, &insight_cx) + .summarize_context(&sms_messages, &ollama_client) .await { Ok(summary) => { @@ -1504,12 +1450,8 @@ impl InsightGenerator { } let retrieval_method = if used_rag { "RAG" } else { "time-based" }; - insight_cx - .span() - .set_attribute(KeyValue::new("retrieval_method", retrieval_method)); - insight_cx - .span() - .set_attribute(KeyValue::new("has_sms_context", sms_summary.is_some())); + span.set_attribute("retrieval_method", retrieval_method); + span.set_attribute("has_sms_context", sms_summary.is_some()); log::info!( "Photo context: date={}, location={:?}, retrieval_method={}", @@ -1520,7 +1462,7 @@ impl InsightGenerator { // 6. Gather Google Takeout context from all sources let calendar_context = self - .gather_calendar_context(&insight_cx, timestamp, location.as_deref()) + .gather_calendar_context(timestamp, location.as_deref()) .await .ok() .flatten(); @@ -1534,14 +1476,13 @@ impl InsightGenerator { }); let location_context = self - .gather_location_context(&insight_cx, timestamp, exif_gps) + .gather_location_context(timestamp, exif_gps) .await .ok() .flatten(); let search_context = self .gather_search_context( - &insight_cx, timestamp, location.as_deref(), contact.as_deref(), @@ -1574,7 +1515,6 @@ impl InsightGenerator { // 10. Generate summary first, then derive title from the summary let summary = ollama_client .generate_photo_summary( - &insight_cx, date_taken, location.as_deref(), contact.as_deref(), @@ -1585,18 +1525,14 @@ impl InsightGenerator { .await?; let title = ollama_client - .generate_photo_title(&insight_cx, &summary, custom_system_prompt.as_deref()) + .generate_photo_title(&summary, custom_system_prompt.as_deref()) .await?; log::info!("Generated title: {}", title); log::info!("Generated summary: {}", summary); - insight_cx - .span() - .set_attribute(KeyValue::new("title_length", title.len() as i64)); - insight_cx - .span() - .set_attribute(KeyValue::new("summary_length", summary.len() as i64)); + span.set_attribute("title_length", title.len() as i64); + span.set_attribute("summary_length", summary.len() as i64); // 11. Store in database. content_hash is None here — store_insight // looks it up from image_exif before persisting; reconciliation @@ -1635,11 +1571,11 @@ impl InsightGenerator { match &result { Ok(_) => { log::info!("Successfully stored insight for {}", file_path); - insight_cx.span().set_status(Status::Ok); + span.set_status(Status::Ok); } Err(e) => { log::error!("Failed to store insight: {:?}", e); - insight_cx.span().set_status(Status::error(e.to_string())); + span.set_status(Status::error(e.to_string())); } } @@ -1652,7 +1588,6 @@ impl InsightGenerator { &self, messages: &[crate::ai::SmsMessage], ollama: &OllamaClient, - cx: &opentelemetry::Context, ) -> Vec { if messages.is_empty() { return Vec::new(); @@ -1688,7 +1623,7 @@ Return ONLY the comma-separated list, nothing else."#, ); match ollama - .generate(cx, &prompt, Some("You are an entity extraction assistant. Extract proper nouns, people, places, and domain-specific terms that provide context.")) + .generate(&prompt, Some("You are an entity extraction assistant. Extract proper nouns, people, places, and domain-specific terms that provide context.")) .await { Ok(response) => { @@ -1723,7 +1658,6 @@ Return ONLY the comma-separated list, nothing else."#, messages: &[String], ollama: &OllamaClient, custom_system: Option<&str>, - cx: &opentelemetry::Context, ) -> Option { if messages.is_empty() { return None; @@ -1744,7 +1678,7 @@ Return ONLY the summary, nothing else."#, let system = custom_system .unwrap_or("You are a context summarization assistant. Be concise and factual."); - match ollama.generate(cx, &prompt, Some(system)).await { + match ollama.generate(&prompt, Some(system)).await { Ok(summary) => Some(summary), Err(e) => { log::warn!("Failed to summarize messages: {}", e); @@ -1760,7 +1694,6 @@ Return ONLY the summary, nothing else."#, messages: &[crate::ai::SmsMessage], ollama: &OllamaClient, custom_system: Option<&str>, - cx: &opentelemetry::Context, ) -> Option { if messages.is_empty() { return None; @@ -1806,7 +1739,7 @@ Return ONLY the summary, nothing else."#, "You are a context summarization assistant. Be detailed and factual, preserving important context.", ); - match ollama.generate(cx, &prompt, Some(system)).await { + match ollama.generate(&prompt, Some(system)).await { Ok(summary) => Some(summary), Err(e) => { log::warn!("Failed to summarize immediate context: {}", e); @@ -1842,7 +1775,7 @@ Return ONLY the summary, nothing else."#, let model = backend.model(); let backend_label = backend.kind.as_str(); let result = match tool_name { - "search_rag" => self.tool_search_rag(arguments, backend.local(), cx).await, + "search_rag" => self.tool_search_rag(arguments, backend.local()).await, "search_messages" => self.tool_search_messages(arguments, cx).await, "get_sms_messages" => self.tool_get_sms_messages(arguments, cx).await, "get_calendar_events" => self.tool_get_calendar_events(arguments, cx).await, @@ -1850,7 +1783,7 @@ Return ONLY the summary, nothing else."#, "get_file_tags" => self.tool_get_file_tags(arguments, cx).await, "get_faces_in_photo" => self.tool_get_faces_in_photo(arguments, cx).await, "describe_photo" => { - self.tool_describe_photo(backend.local(), image_base64, cx) + self.tool_describe_photo(backend.local(), image_base64) .await } "reverse_geocode" => self.tool_reverse_geocode(arguments).await, @@ -1897,12 +1830,7 @@ Return ONLY the summary, nothing else."#, } /// Tool: search_rag — semantic search over daily summaries - async fn tool_search_rag( - &self, - args: &serde_json::Value, - local: &dyn LlmClient, - cx: &opentelemetry::Context, - ) -> String { + async fn tool_search_rag(&self, args: &serde_json::Value, local: &dyn LlmClient) -> String { let query = match args.get("query").and_then(|v| v.as_str()) { Some(q) => q.to_string(), None => return "Error: missing required parameter 'query'".to_string(), @@ -1965,10 +1893,7 @@ Return ONLY the summary, nothing else."#, }; let final_results = if rerank_enabled && results.len() > limit { - match self - .rerank_with_llm(cx, &query, &results, limit, local) - .await - { + match self.rerank_with_llm(&query, &results, limit, local).await { Ok(reordered) => reordered, Err(e) => { log::warn!("rerank failed, using vector order: {}", e); @@ -1991,7 +1916,6 @@ Return ONLY the summary, nothing else."#, /// tool_result. On parse failure we fall back to the input order. async fn rerank_with_llm( &self, - cx: &opentelemetry::Context, query: &str, candidates: &[String], limit: usize, @@ -2035,7 +1959,7 @@ Return ONLY the summary, nothing else."#, let started = std::time::Instant::now(); let system = Some("You are a terse relevance ranker. You output only numbers separated by commas."); - let response = local.generate(cx, &prompt, system, None).await?; + let response = local.generate(&prompt, system, None).await?; log::info!( "rerank: finished in {} ms (prompt={} chars)", started.elapsed().as_millis(), @@ -2645,11 +2569,10 @@ Return ONLY the summary, nothing else."#, &self, local: &dyn LlmClient, image_base64: &Option, - cx: &opentelemetry::Context, ) -> String { log::info!("tool_describe_photo: generating visual description"); match image_base64 { - Some(img) => match local.describe_image(cx, img).await { + Some(img) => match local.describe_image(img).await { Ok(desc) => desc, Err(e) => format!("Error describing photo: {}", e), }, @@ -3057,7 +2980,6 @@ Return ONLY the summary, nothing else."#, // configured local backend. let embed_text = format!("{} {}", name, description); let embedding: Option> = match crate::ai::embed_document( - cx, &self.ollama, self.llamacpp.as_deref(), &embed_text, @@ -4175,6 +4097,7 @@ Return ONLY the summary, nothing else."#, Ok(ResolvedBackend::new(chat, local, kind, images_inline)) } + #[tracing::instrument(name = "ai.insight.generate_agentic", skip_all)] pub async fn generate_agentic_insight_for_photo( &self, file_path: &str, @@ -4193,19 +4116,17 @@ Return ONLY the summary, nothing else."#, user_id: i32, persona_id: String, ) -> Result<(Option, 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); + let span = tracing::Span::current(); let file_path = normalize_path(file_path); log::info!("Generating agentic insight for photo: {}", file_path); - span.set_attribute(KeyValue::new("file_path", file_path.clone())); - span.set_attribute(KeyValue::new("max_iterations", max_iterations as i64)); + span.set_attribute("file_path", file_path.clone()); + span.set_attribute("max_iterations", max_iterations as i64); // 1. Resolve backend + build clients. let kind = BackendKind::parse(backend.as_deref().unwrap_or("local"))?; - span.set_attribute(KeyValue::new("backend", kind.as_str())); + span.set_attribute("backend", kind.as_str()); let overrides = SamplingOverrides { model: custom_model, num_ctx, @@ -4216,10 +4137,10 @@ Return ONLY the summary, nothing else."#, enable_thinking, }; let backend = self.resolve_backend(kind, &overrides).await?; - span.set_attribute(KeyValue::new("model", backend.model().to_string())); - span.set_attribute(KeyValue::new("images_inline", backend.images_inline)); + span.set_attribute("model", backend.model().to_string()); + span.set_attribute("images_inline", backend.images_inline); - let insight_cx = current_cx.with_span(span); + let insight_cx = span.context(); // 3. Fetch EXIF let exif = { @@ -4304,7 +4225,7 @@ Return ONLY the summary, nothing else."#, // the local backend so non-text work stays off OpenRouter. let inlined_visual_description: Option = if !backend.images_inline { match image_base64.as_deref() { - Some(b64) => match backend.local().describe_image(&insight_cx, b64).await { + Some(b64) => match backend.local().describe_image(b64).await { Ok(desc) => { log::info!("{}: vision describe succeeded ({} chars)", kind, desc.len()); Some(desc) @@ -4410,8 +4331,10 @@ Return ONLY the summary, nothing else."#, let mut messages = vec![system_msg, user_msg]; - let loop_span = tracer.start_with_context("ai.agentic.loop", &insight_cx); - let loop_cx = insight_cx.with_span(loop_span); + // LLM / tool awaits below run inside `loop_span` so their spans + // nest under it. + let loop_span = tracing::info_span!("ai.agentic.loop"); + let loop_cx = loop_span.context(); let mut final_content = String::new(); let mut iterations_used = 0usize; @@ -4424,7 +4347,8 @@ Return ONLY the summary, nothing else."#, let (response, prompt_tokens, eval_tokens) = backend .chat() - .chat_with_tools(&loop_cx, messages.clone(), tools.clone()) + .chat_with_tools(messages.clone(), tools.clone()) + .instrument(loop_span.clone()) .await?; last_prompt_eval_count = prompt_tokens; @@ -4470,6 +4394,7 @@ Return ONLY the summary, nothing else."#, &persona_id, &loop_cx, ) + .instrument(loop_span.clone()) .await; messages.push(ChatMessage::tool_result(result)); } @@ -4493,7 +4418,8 @@ Return ONLY the summary, nothing else."#, ))); let (final_response, prompt_tokens, eval_tokens) = backend .chat() - .chat_with_tools(&loop_cx, messages.clone(), tools.clone()) + .chat_with_tools(messages.clone(), tools.clone()) + .instrument(loop_span.clone()) .await?; last_prompt_eval_count = prompt_tokens; last_eval_count = eval_tokens; @@ -4521,12 +4447,14 @@ Return ONLY the summary, nothing else."#, &persona_id, &loop_cx, ) + .instrument(loop_span.clone()) .await; messages.push(ChatMessage::tool_result(result)); } let (resp2, pt2, et2) = backend .chat() - .chat_with_tools(&loop_cx, messages.clone(), vec![]) + .chat_with_tools(messages.clone(), vec![]) + .instrument(loop_span.clone()) .await?; last_prompt_eval_count = pt2; last_eval_count = et2; @@ -4538,10 +4466,9 @@ Return ONLY the summary, nothing else."#, } } - loop_cx - .span() - .set_attribute(KeyValue::new("iterations_used", iterations_used as i64)); - loop_cx.span().set_status(Status::Ok); + loop_span.set_attribute("iterations_used", iterations_used as i64); + loop_span.set_status(Status::Ok); + drop(loop_span); // 13. Strip any leaked … reasoning block (thinking // models emit it ahead of the answer; the raw transcript in @@ -4608,11 +4535,11 @@ Return ONLY the summary, nothing else."#, match &stored { Ok(_) => { log::info!("Successfully stored agentic insight for {}", file_path); - insight_cx.span().set_status(Status::Ok); + span.set_status(Status::Ok); } Err(e) => { log::error!("Failed to store agentic insight: {:?}", e); - insight_cx.span().set_status(Status::error(e.to_string())); + span.set_status(Status::error(e.to_string())); } } @@ -4672,7 +4599,7 @@ Return ONLY the summary, nothing else."#, let (response, _prompt_tokens, _eval_tokens) = backend .chat() - .chat_with_tools(&loop_cx, messages.clone(), tools.clone()) + .chat_with_tools(messages.clone(), tools.clone()) .await?; // Sanitize tool call arguments before pushing back into history. @@ -4738,7 +4665,7 @@ Return ONLY the summary, nothing else."#, )); let (final_response, _, _) = backend .chat() - .chat_with_tools(&loop_cx, messages.clone(), tools.clone()) + .chat_with_tools(messages.clone(), tools.clone()) .await?; // Same bounded tool round as the insight loop: execute any last // tool call, then force a tool-less turn for the actual answer. @@ -4763,7 +4690,7 @@ Return ONLY the summary, nothing else."#, } let (resp2, _, _) = backend .chat() - .chat_with_tools(&loop_cx, messages.clone(), vec![]) + .chat_with_tools(messages.clone(), vec![]) .await?; final_content = resp2.content; } else { diff --git a/src/ai/llamacpp.rs b/src/ai/llamacpp.rs index 9d5370e..b69cf0f 100644 --- a/src/ai/llamacpp.rs +++ b/src/ai/llamacpp.rs @@ -23,12 +23,12 @@ use anyhow::Context as _; use anyhow::{Result, anyhow, bail}; use async_trait::async_trait; -use opentelemetry::trace::{Span, Status}; -use opentelemetry::{Context, KeyValue}; +use opentelemetry::trace::Status; use reqwest::Client; use serde::Deserialize; use serde_json::{Value, json}; use std::time::Duration; +use tracing_opentelemetry::OpenTelemetrySpanExt as _; use crate::ai::llm_client::{ ChatMessage, LlmClient, LlmStreamEvent, ModelCapabilities, Tool, ToolCall, ToolCallFunction, @@ -489,7 +489,7 @@ impl LlamaCppClient { async fn chat_completion_with_model( &self, model: &str, - mut span: Option<&mut opentelemetry::global::BoxedSpan>, + span: Option<&tracing::Span>, messages: Vec, tools: Vec, ) -> Result<(ChatMessage, Option, Option)> { @@ -522,13 +522,13 @@ impl LlamaCppClient { // timeout starts (see ai::gpu). let _gpu = crate::ai::gpu::llm_lease().await; let mut request = self.client.post(&url).json(&Value::Object(body)); - if let Some(s) = span.as_mut() { - request = crate::otel::with_traceparent(request, s.span_context()); + if let Some(s) = span { + request = crate::otel::with_traceparent(request, s); } let resp = match request.send().await { Ok(r) => r, Err(e) => { - if let Some(s) = span.as_mut() { + if let Some(s) = span { s.set_status(Status::error("llama-swap request failed")); } return Err(e).with_context(|| format!("POST {} failed", url)); @@ -538,14 +538,14 @@ impl LlamaCppClient { if !resp.status().is_success() { let status = resp.status(); crate::otel::record_http_failure( - span.as_deref_mut(), + span, status.as_u16(), &format!("llama-swap chat request failed: {status}"), ); let body = resp.text().await.unwrap_or_default(); bail!("llama-swap chat request failed: {} — {}", status, body); } - crate::otel::set_http_status(span.as_deref_mut(), resp.status().as_u16()); + crate::otel::set_http_status(span, resp.status().as_u16()); let parsed: Value = resp.json().await.context("parsing chat response")?; let choice = parsed @@ -578,12 +578,12 @@ impl LlamaCppClient { log_timings(&parsed, prompt_tokens, completion_tokens); - if let Some(s) = span.as_mut() { + if let Some(s) = span { if let Some(pt) = prompt_tokens { - s.set_attribute(KeyValue::new("llm.tokens.input", pt as i64)); + s.set_attribute("llm.tokens.input", pt as i64); } if let Some(ct) = completion_tokens { - s.set_attribute(KeyValue::new("llm.tokens.output", ct as i64)); + s.set_attribute("llm.tokens.output", ct as i64); } s.set_status(Status::Ok); } @@ -592,13 +592,8 @@ impl LlamaCppClient { } /// Start the `llm.request` span for an outgoing llama-swap call. - fn start_request_span( - &self, - cx: &Context, - operation: &str, - model: &str, - ) -> opentelemetry::global::BoxedSpan { - crate::otel::start_llm_request_span(cx, "llamacpp", operation, model) + fn start_request_span(&self, operation: &str, model: &str) -> tracing::Span { + crate::otel::start_llm_request_span("llamacpp", operation, model) } } @@ -606,7 +601,6 @@ impl LlamaCppClient { impl LlmClient for LlamaCppClient { async fn generate( &self, - cx: &Context, prompt: &str, system: Option<&str>, images: Option>, @@ -621,13 +615,12 @@ impl LlmClient for LlamaCppClient { // Delegates to chat_with_tools, which starts the llm.request span // (operation=chat) — one span per HTTP request. - let (reply, _, _) = self.chat_with_tools(cx, messages, Vec::new()).await?; + let (reply, _, _) = self.chat_with_tools(messages, Vec::new()).await?; Ok(reply.content) } async fn chat_with_tools( &self, - cx: &Context, messages: Vec, tools: Vec, ) -> Result<(ChatMessage, Option, Option)> { @@ -637,25 +630,19 @@ impl LlmClient for LlamaCppClient { messages.len(), tools.len() ); - let mut span = self.start_request_span(cx, "chat", &self.primary_model); - self.chat_completion_with_model( - &self.primary_model.clone(), - Some(&mut span), - messages, - tools, - ) - .await + let span = self.start_request_span("chat", &self.primary_model); + self.chat_completion_with_model(&self.primary_model.clone(), Some(&span), messages, tools) + .await } async fn chat_with_tools_stream( &self, - cx: &Context, messages: Vec, tools: Vec, ) -> Result>> { // The span outlives this method — it moves into the stream // generator and is finalized when the response drains. - let mut span = self.start_request_span(cx, "chat_stream", &self.primary_model); + let span = self.start_request_span("chat_stream", &self.primary_model); let url = format!("{}/chat/completions", self.base_url); let roles: Vec<&str> = messages.iter().map(|m| m.role.as_str()).collect(); log::debug!( @@ -690,7 +677,7 @@ impl LlmClient for LlamaCppClient { // below so the lease spans the whole generation, not just the send. let gpu = crate::ai::gpu::llm_lease().await; let mut request = self.client.post(&url).json(&Value::Object(body)); - request = crate::otel::with_traceparent(request, span.span_context()); + request = crate::otel::with_traceparent(request, &span); let resp = request .send() .await @@ -699,14 +686,14 @@ impl LlmClient for LlamaCppClient { if !resp.status().is_success() { let status = resp.status(); crate::otel::record_http_failure( - Some(&mut span), + Some(&span), status.as_u16(), &format!("llama-swap stream request failed: {status}"), ); let body = resp.text().await.unwrap_or_default(); bail!("llama-swap stream request failed: {} — {}", status, body); } - crate::otel::set_http_status(Some(&mut span), resp.status().as_u16()); + crate::otel::set_http_status(Some(&span), resp.status().as_u16()); let byte_stream = resp.bytes_stream(); let stream = async_stream::stream! { @@ -816,10 +803,10 @@ impl LlmClient for LlamaCppClient { images: None, }; if let Some(pt) = prompt_tokens { - span.set_attribute(KeyValue::new("llm.tokens.input", pt as i64)); + span.set_attribute("llm.tokens.input", pt as i64); } if let Some(ct) = completion_tokens { - span.set_attribute(KeyValue::new("llm.tokens.output", ct as i64)); + span.set_attribute("llm.tokens.output", ct as i64); } span.set_status(Status::Ok); yield Ok(LlmStreamEvent::Done { @@ -832,8 +819,8 @@ impl LlmClient for LlamaCppClient { Ok(Box::pin(stream)) } - async fn generate_embeddings(&self, cx: &Context, texts: &[&str]) -> Result>> { - let mut span = self.start_request_span(cx, "embeddings", &self.embedding_model); + async fn generate_embeddings(&self, texts: &[&str]) -> Result>> { + let span = self.start_request_span("embeddings", &self.embedding_model); // Deliberately NO GPU lease: the embed slot sits in llama-swap's // always-resident group and never participates in a model swap, so // leasing here would only stall searches behind a queued synthesis. @@ -844,7 +831,7 @@ impl LlmClient for LlamaCppClient { }); let mut request = self.client.post(&url).json(&body); - request = crate::otel::with_traceparent(request, span.span_context()); + request = crate::otel::with_traceparent(request, &span); let resp = match request.send().await { Ok(r) => r, Err(e) => { @@ -856,14 +843,14 @@ impl LlmClient for LlamaCppClient { if !resp.status().is_success() { let status = resp.status(); crate::otel::record_http_failure( - Some(&mut span), + Some(&span), status.as_u16(), &format!("llama-swap embedding request failed: {status}"), ); let body = resp.text().await.unwrap_or_default(); bail!("llama-swap embedding request failed: {} — {}", status, body); } - crate::otel::set_http_status(Some(&mut span), resp.status().as_u16()); + crate::otel::set_http_status(Some(&span), resp.status().as_u16()); #[derive(Deserialize)] struct EmbedResponse { @@ -885,9 +872,9 @@ impl LlmClient for LlamaCppClient { Ok(parsed.data.into_iter().map(|i| i.embedding).collect()) } - async fn describe_image(&self, cx: &Context, image_base64: &str) -> Result { + async fn describe_image(&self, image_base64: &str) -> Result { // Own span with operation=describe_image — routes through the vision slot. - let mut span = self.start_request_span(cx, "describe_image", &self.vision_model); + let span = self.start_request_span("describe_image", &self.vision_model); let prompt = "Briefly describe what you see in this image in 1-2 sentences. \ Focus on the people, location, and activity."; let system = "You are a scene description assistant. Be concise and factual."; @@ -905,7 +892,7 @@ impl LlmClient for LlamaCppClient { let (reply, _, _) = self .chat_completion_with_model( &self.vision_model.clone(), - Some(&mut span), + Some(&span), messages, Vec::new(), ) diff --git a/src/ai/llm_client.rs b/src/ai/llm_client.rs index e598e0e..c8b0966 100644 --- a/src/ai/llm_client.rs +++ b/src/ai/llm_client.rs @@ -1,7 +1,6 @@ use anyhow::Result; use async_trait::async_trait; use futures::stream::BoxStream; -use opentelemetry::Context; use serde::{Deserialize, Serialize}; /// Provider-agnostic surface for LLM backends (Ollama, OpenRouter, …). @@ -12,11 +11,10 @@ use serde::{Deserialize, Serialize}; /// stays as base64 strings here and is rewritten into content-parts where /// needed. /// -/// Every method that performs an HTTP call takes the caller's tracing -/// `Context` as its first parameter so the impl can start an `llm.request` -/// span parented under the caller's feature span — the codebase threads -/// contexts explicitly rather than relying on ambient `Context::current()` -/// at deep call sites. +/// Every method that performs an HTTP call starts an `llm.request` span +/// parented on the caller's ambient `tracing` span, so callers must run +/// inside their feature span (`#[instrument]` / `.instrument(span)`) — and +/// carry it across any `tokio::spawn` — for the LLM span to nest under it. // First consumer lands in a later PR (OpenRouter impl + hybrid mode routing). #[allow(dead_code)] #[async_trait] @@ -25,7 +23,6 @@ pub trait LlmClient: Send + Sync { /// base64 images (ignored by providers without vision support). async fn generate( &self, - cx: &Context, prompt: &str, system: Option<&str>, images: Option>, @@ -35,7 +32,6 @@ pub trait LlmClient: Send + Sync { /// (which may contain tool_calls) plus optional prompt/eval token counts. async fn chat_with_tools( &self, - cx: &Context, messages: Vec, tools: Vec, ) -> Result<(ChatMessage, Option, Option)>; @@ -48,17 +44,16 @@ pub trait LlmClient: Send + Sync { /// as one `Done` event. async fn chat_with_tools_stream( &self, - cx: &Context, messages: Vec, tools: Vec, ) -> Result>>; /// Batch embedding generation. Dimensionality is provider/model specific. - async fn generate_embeddings(&self, cx: &Context, texts: &[&str]) -> Result>>; + async fn generate_embeddings(&self, texts: &[&str]) -> Result>>; /// One-shot vision description of an image. Used to convert images into /// plain text for the hybrid-mode conversation flow. - async fn describe_image(&self, cx: &Context, image_base64: &str) -> Result; + async fn describe_image(&self, image_base64: &str) -> Result; /// Enumerate available models with their capabilities. async fn list_models(&self) -> Result>; diff --git a/src/ai/local_llm.rs b/src/ai/local_llm.rs index dd5c350..bf3510c 100644 --- a/src/ai/local_llm.rs +++ b/src/ai/local_llm.rs @@ -47,24 +47,12 @@ impl LocalLlm { /// pick query vs document — retrieval models treat the two sides /// differently and an unmarked embed invites prefix-mismatch bugs. pub async fn embed_query(&self, text: &str) -> Result> { - super::embed_query( - &opentelemetry::Context::new(), - &self.ollama, - self.llamacpp.as_deref(), - text, - ) - .await + super::embed_query(&self.ollama, self.llamacpp.as_deref(), text).await } /// Embed corpus text (applies `EMBED_DOCUMENT_PREFIX`). pub async fn embed_document(&self, text: &str) -> Result> { - super::embed_document( - &opentelemetry::Context::new(), - &self.ollama, - self.llamacpp.as_deref(), - text, - ) - .await + super::embed_document(&self.ollama, self.llamacpp.as_deref(), text).await } /// Single-shot local text generation via the `LLM_BACKEND`-selected @@ -72,23 +60,14 @@ impl LocalLlm { pub async fn generate(&self, prompt: &str, system: Option<&str>) -> Result { if super::local_backend_is_llamacpp() { if let Some(lc) = self.llamacpp.as_deref() { - return ::generate( - lc, - &opentelemetry::Context::new(), - prompt, - system, - None, - ) - .await; + return ::generate(lc, prompt, system, None).await; } anyhow::bail!( "LLM_BACKEND=llamacpp but LlamaCppClient is unconfigured — \ set LLAMA_SWAP_URL or switch to LLM_BACKEND=ollama" ); } - self.ollama - .generate(&opentelemetry::Context::new(), prompt, system) - .await + self.ollama.generate(prompt, system).await } /// Label identifying which backend + model produces embeddings right diff --git a/src/ai/mod.rs b/src/ai/mod.rs index 68b0153..d39645a 100644 --- a/src/ai/mod.rs +++ b/src/ai/mod.rs @@ -115,26 +115,24 @@ fn embed_prefix(key: &str) -> String { /// `Instruct: \nQuery: `. Must pair with the document prefix the /// corpus was embedded with or similarity degrades. pub async fn embed_query( - cx: &opentelemetry::Context, ollama: &OllamaClient, llamacpp: Option<&LlamaCppClient>, text: &str, ) -> anyhow::Result> { let prefixed = format!("{}{}", embed_prefix("EMBED_QUERY_PREFIX"), text); - embed_one(cx, ollama, llamacpp, &prefixed).await + embed_one(ollama, llamacpp, &prefixed).await } /// Embed corpus text (the stored side of retrieval). Applies /// `EMBED_DOCUMENT_PREFIX` (default empty; nomic v1.5 wants /// `search_document: `, Qwen3-Embedding wants none). pub async fn embed_document( - cx: &opentelemetry::Context, ollama: &OllamaClient, llamacpp: Option<&LlamaCppClient>, text: &str, ) -> anyhow::Result> { let prefixed = format!("{}{}", embed_prefix("EMBED_DOCUMENT_PREFIX"), text); - embed_one(cx, ollama, llamacpp, &prefixed).await + embed_one(ollama, llamacpp, &prefixed).await } /// Embed a batch of strings via the configured local backend. Routes @@ -142,32 +140,30 @@ pub async fn embed_document( /// configured), else Ollama. See [`local_backend_is_llamacpp`] for the /// rationale on consistency. pub async fn embed_many( - cx: &opentelemetry::Context, ollama: &OllamaClient, llamacpp: Option<&LlamaCppClient>, texts: &[&str], ) -> anyhow::Result>> { if local_backend_is_llamacpp() { if let Some(lc) = llamacpp { - return ::generate_embeddings(lc, cx, texts).await; + return ::generate_embeddings(lc, texts).await; } anyhow::bail!( "LLM_BACKEND=llamacpp but LlamaCppClient is unconfigured — \ set LLAMA_SWAP_URL or switch to LLM_BACKEND=ollama" ); } - ollama.generate_embeddings(cx, texts).await + ollama.generate_embeddings(texts).await } /// Embed one string via the configured local backend. Single-text /// convenience over [`embed_many`]. pub async fn embed_one( - cx: &opentelemetry::Context, ollama: &OllamaClient, llamacpp: Option<&LlamaCppClient>, text: &str, ) -> anyhow::Result> { - let mut vecs = embed_many(cx, ollama, llamacpp, &[text]).await?; + let mut vecs = embed_many(ollama, llamacpp, &[text]).await?; vecs.pop() .ok_or_else(|| anyhow::anyhow!("embedding backend returned no embeddings")) } diff --git a/src/ai/nl_query.rs b/src/ai/nl_query.rs index b4dabb1..d709322 100644 --- a/src/ai/nl_query.rs +++ b/src/ai/nl_query.rs @@ -259,13 +259,10 @@ pub async fn translate_nl_query( nl: &str, tag_vocab: &[(i32, String)], today: chrono::NaiveDate, - cx: &opentelemetry::Context, ) -> Result { let system = build_system_prompt(tag_vocab, today); let messages = vec![ChatMessage::system(system), ChatMessage::user(nl)]; - let (msg, _, _) = client - .chat_with_tools(cx, messages, Vec::::new()) - .await?; + let (msg, _, _) = client.chat_with_tools(messages, Vec::::new()).await?; parse_response(&msg.content, tag_vocab) } diff --git a/src/ai/ollama.rs b/src/ai/ollama.rs index 144fb01..335b047 100644 --- a/src/ai/ollama.rs +++ b/src/ai/ollama.rs @@ -2,14 +2,14 @@ use anyhow::Context as _; use anyhow::Result; use async_trait::async_trait; use chrono::NaiveDate; -use opentelemetry::trace::{Span, Status}; -use opentelemetry::{Context, KeyValue}; +use opentelemetry::trace::Status; use reqwest::Client; use serde::{Deserialize, Serialize}; use std::collections::HashMap; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::{Arc, Mutex}; use std::time::{Duration, Instant}; +use tracing_opentelemetry::OpenTelemetrySpanExt as _; use crate::ai::llm_client::{LlmClient, LlmStreamEvent}; use futures::stream::{BoxStream, StreamExt}; @@ -368,7 +368,7 @@ impl OllamaClient { async fn try_generate( &self, - mut span: Option<&mut opentelemetry::global::BoxedSpan>, + span: Option<&tracing::Span>, url: &str, model: &str, prompt: &str, @@ -390,13 +390,13 @@ impl OllamaClient { .client .post(format!("{}/api/generate", url)) .json(&request); - if let Some(s) = span.as_mut() { - request_builder = crate::otel::with_traceparent(request_builder, s.span_context()); + if let Some(s) = span { + request_builder = crate::otel::with_traceparent(request_builder, s); } let response = match request_builder.send().await { Ok(r) => r, Err(e) => { - if let Some(s) = span.as_mut() { + if let Some(s) = span { s.set_status(Status::error("Ollama request failed")); } return Err(e.into()); @@ -406,7 +406,7 @@ impl OllamaClient { if !response.status().is_success() { let status = response.status(); crate::otel::record_http_failure( - span.as_deref_mut(), + span, status.as_u16(), &format!("Ollama request failed: {status}"), ); @@ -417,23 +417,23 @@ impl OllamaClient { error_body )); } - crate::otel::set_http_status(span.as_deref_mut(), response.status().as_u16()); + crate::otel::set_http_status(span, response.status().as_u16()); let result: OllamaResponse = match response.json().await { Ok(r) => r, Err(e) => { - if let Some(s) = span.as_mut() { + if let Some(s) = span { s.set_status(Status::error("parsing Ollama response failed")); } return Err(e.into()); } }; - if let Some(s) = span.as_mut() { + if let Some(s) = span { if let Some(pt) = result.prompt_eval_count { - s.set_attribute(KeyValue::new("llm.tokens.input", pt as i64)); + s.set_attribute("llm.tokens.input", pt as i64); } if let Some(ct) = result.eval_count { - s.set_attribute(KeyValue::new("llm.tokens.output", ct as i64)); + s.set_attribute("llm.tokens.output", ct as i64); } s.set_status(Status::Ok); } @@ -447,49 +447,32 @@ impl OllamaClient { } /// Start the `llm.request` span for an outgoing Ollama call. - fn start_request_span( - &self, - cx: &Context, - operation: &str, - model: &str, - ) -> opentelemetry::global::BoxedSpan { - crate::otel::start_llm_request_span(cx, "ollama", operation, model) + fn start_request_span(&self, operation: &str, model: &str) -> tracing::Span { + crate::otel::start_llm_request_span("ollama", operation, model) } - pub async fn generate( - &self, - cx: &Context, - prompt: &str, - system: Option<&str>, - ) -> Result { - self.generate_with_images(cx, prompt, system, None).await + pub async fn generate(&self, prompt: &str, system: Option<&str>) -> Result { + self.generate_with_images(prompt, system, None).await } #[allow(dead_code)] - pub async fn generate_no_think( - &self, - cx: &Context, - prompt: &str, - system: Option<&str>, - ) -> Result { - self.generate_with_options(cx, prompt, system, None, Some(false), "generate") + pub async fn generate_no_think(&self, prompt: &str, system: Option<&str>) -> Result { + self.generate_with_options(prompt, system, None, Some(false), "generate") .await } pub async fn generate_with_images( &self, - cx: &Context, prompt: &str, system: Option<&str>, images: Option>, ) -> Result { - self.generate_with_options(cx, prompt, system, images, None, "generate") + self.generate_with_options(prompt, system, images, None, "generate") .await } async fn generate_with_options( &self, - cx: &Context, prompt: &str, system: Option<&str>, images: Option>, @@ -509,7 +492,7 @@ impl OllamaClient { // Span for the whole generate attempt; the model attribute // reflects whichever server actually answered. - let mut span = self.start_request_span(cx, operation, &self.primary_model); + let span = self.start_request_span(operation, &self.primary_model); // Try primary server first with primary model log::info!( @@ -519,7 +502,7 @@ impl OllamaClient { ); let primary_result = self .try_generate( - Some(&mut span), + Some(&span), &self.primary_url, &self.primary_model, prompt, @@ -548,10 +531,10 @@ impl OllamaClient { fallback_url, fallback_model ); - span.set_attribute(KeyValue::new("llm.model", fallback_model.to_string())); + span.set_attribute("llm.model", fallback_model.to_string()); match self .try_generate( - Some(&mut span), + Some(&span), fallback_url, fallback_model, prompt, @@ -598,7 +581,6 @@ impl OllamaClient { /// Generate a title for a single photo based on its generated summary pub async fn generate_photo_title( &self, - cx: &Context, summary: &str, custom_system: Option<&str>, ) -> Result { @@ -614,7 +596,7 @@ Capture the key moment or theme. Return ONLY the title, nothing else."#, let system = custom_system.unwrap_or("You are my long term memory assistant. Use only the information provided. Do not invent details."); let title = self - .generate_with_images(cx, &prompt, Some(system), None) + .generate_with_images(&prompt, Some(system), None) .await?; // Models decorate despite "Return ONLY the title": quotes, bold // markers, sometimes a "Title:" label. @@ -631,7 +613,6 @@ Capture the key moment or theme. Return ONLY the title, nothing else."#, /// Generate a summary for a single photo based on its context pub async fn generate_photo_summary( &self, - cx: &Context, date: NaiveDate, location: Option<&str>, contact: Option<&str>, @@ -714,32 +695,21 @@ Analyze the image and use specific details from both the visual content and the let system = custom_system.unwrap_or("You are a memory refreshing assistant who is able to provide insights through analyzing past conversations. Use only the information provided. Do not invent details."); let images = image_base64.map(|img| vec![img]); - self.generate_with_images(cx, &prompt, Some(system), images) + self.generate_with_images(&prompt, Some(system), images) .await } /// Generate a brief visual description of a photo for use in RAG query enrichment. /// Returns 1-2 sentences describing people, location, and activity visible in the image. /// Only called when the model has vision capabilities. - pub async fn generate_photo_description( - &self, - cx: &Context, - image_base64: &str, - ) -> Result { + pub async fn generate_photo_description(&self, image_base64: &str) -> Result { let prompt = "Briefly describe what you see in this image in 1-2 sentences. \ Focus on the people, location, and activity."; let system = "You are a scene description assistant. Be concise and factual."; let images = vec![image_base64.to_string()]; let description = self - .generate_with_options( - cx, - prompt, - Some(system), - Some(images), - None, - "describe_image", - ) + .generate_with_options(prompt, Some(system), Some(images), None, "describe_image") .await?; Ok(description.trim().to_string()) @@ -752,11 +722,10 @@ Analyze the image and use specific details from both the visual content and the /// iteration of the agent loop. pub async fn chat_with_tools( &self, - cx: &Context, messages: Vec, tools: Vec, ) -> Result<(ChatMessage, Option, Option)> { - let mut span = self.start_request_span(cx, "chat", &self.primary_model); + let span = self.start_request_span("chat", &self.primary_model); let order = self.attempt_order(); let mut errors: Vec = Vec::new(); @@ -768,10 +737,10 @@ Analyze the image and use specific details from both the visual content and the model ); if *label == "fallback" { - span.set_attribute(KeyValue::new("llm.model", model.to_string())); + span.set_attribute("llm.model", model.to_string()); } match self - .try_chat_with_tools(Some(&mut span), url, messages.clone(), tools.clone()) + .try_chat_with_tools(Some(&span), url, messages.clone(), tools.clone()) .await { Ok(result) => { @@ -813,7 +782,6 @@ Analyze the image and use specific details from both the visual content and the /// any, live on the final message). pub async fn chat_with_tools_stream( &self, - cx: &Context, messages: Vec, tools: Vec, ) -> Result>> { @@ -828,7 +796,7 @@ Analyze the image and use specific details from both the visual content and the for (label, url, _model) in &order { match self - .try_chat_with_tools_stream(cx, url, messages.clone(), tools.clone()) + .try_chat_with_tools_stream(url, messages.clone(), tools.clone()) .await { Ok(s) => { @@ -848,7 +816,6 @@ Analyze the image and use specific details from both the visual content and the async fn try_chat_with_tools_stream( &self, - cx: &Context, base_url: &str, messages: Vec, tools: Vec, @@ -866,7 +833,7 @@ Analyze the image and use specific details from both the visual content and the // The span outlives this method — the stream yields until the // response fully drains, so the owned span moves into the // generator and is finalized there. - let mut span = self.start_request_span(cx, "chat_stream", model); + let span = self.start_request_span("chat_stream", model); let request_body = OllamaChatRequest { model, @@ -878,7 +845,7 @@ Analyze the image and use specific details from both the visual content and the let response = { let mut rb = self.client.post(&url).json(&request_body); - rb = crate::otel::with_traceparent(rb, span.span_context()); + rb = crate::otel::with_traceparent(rb, &span); match rb.send().await { Ok(r) => r, Err(e) => { @@ -892,7 +859,7 @@ Analyze the image and use specific details from both the visual content and the if !response.status().is_success() { let status = response.status(); crate::otel::record_http_failure( - Some(&mut span), + Some(&span), status.as_u16(), &format!("Ollama stream request failed: {status}"), ); @@ -903,7 +870,7 @@ Analyze the image and use specific details from both the visual content and the body ); } - crate::otel::set_http_status(Some(&mut span), response.status().as_u16()); + crate::otel::set_http_status(Some(&span), response.status().as_u16()); // Ollama streams NDJSON: each line is a full `OllamaStreamChunk`. // We buffer partial lines across chunks from the byte stream. @@ -986,10 +953,10 @@ Analyze the image and use specific details from both the visual content and the eval_duration, ); if let Some(pt) = prompt_eval_count { - span.set_attribute(KeyValue::new("llm.tokens.input", pt as i64)); + span.set_attribute("llm.tokens.input", pt as i64); } if let Some(ct) = eval_count { - span.set_attribute(KeyValue::new("llm.tokens.output", ct as i64)); + span.set_attribute("llm.tokens.output", ct as i64); } span.set_status(Status::Ok); let message = ChatMessage { @@ -1010,7 +977,7 @@ Analyze the image and use specific details from both the visual content and the async fn try_chat_with_tools( &self, - mut span: Option<&mut opentelemetry::global::BoxedSpan>, + span: Option<&tracing::Span>, base_url: &str, messages: Vec, tools: Vec, @@ -1040,13 +1007,13 @@ Analyze the image and use specific details from both the visual content and the let response = { let mut rb = self.client.post(&url).json(&request_body); - if let Some(s) = span.as_mut() { - rb = crate::otel::with_traceparent(rb, s.span_context()); + if let Some(s) = span { + rb = crate::otel::with_traceparent(rb, s); } match rb.send().await { Ok(r) => r, Err(e) => { - if let Some(s) = span.as_mut() { + if let Some(s) = span { s.set_status(Status::error("Ollama chat connect failed")); } return Err(e) @@ -1058,7 +1025,7 @@ Analyze the image and use specific details from both the visual content and the if !response.status().is_success() { let status = response.status(); crate::otel::record_http_failure( - span.as_deref_mut(), + span, status.as_u16(), &format!("Ollama chat request failed: {status}"), ); @@ -1078,23 +1045,23 @@ Analyze the image and use specific details from both the visual content and the body ); } - crate::otel::set_http_status(span.as_deref_mut(), response.status().as_u16()); + crate::otel::set_http_status(span, response.status().as_u16()); let chat_response: OllamaChatResponse = match response.json().await { Ok(r) => r, Err(e) => { - if let Some(s) = span.as_mut() { + if let Some(s) = span { s.set_status(Status::error("parsing Ollama chat response failed")); } return Err(e).with_context(|| "Failed to parse Ollama chat response"); } }; - if let Some(s) = span.as_mut() { + if let Some(s) = span { if let Some(pt) = chat_response.prompt_eval_count { - s.set_attribute(KeyValue::new("llm.tokens.input", pt as i64)); + s.set_attribute("llm.tokens.input", pt as i64); } if let Some(ct) = chat_response.eval_count { - s.set_attribute(KeyValue::new("llm.tokens.output", ct as i64)); + s.set_attribute("llm.tokens.output", ct as i64); } s.set_status(Status::Ok); } @@ -1119,8 +1086,8 @@ Analyze the image and use specific details from both the visual content and the /// Generate an embedding vector for text using nomic-embed-text:v1.5 /// Returns a 768-dimensional vector as Vec - pub async fn generate_embedding(&self, cx: &Context, text: &str) -> Result> { - let embeddings = self.generate_embeddings(cx, &[text]).await?; + pub async fn generate_embedding(&self, text: &str) -> Result> { + let embeddings = self.generate_embeddings(&[text]).await?; embeddings .into_iter() .next() @@ -1130,10 +1097,10 @@ Analyze the image and use specific details from both the visual content and the /// Generate embeddings for multiple texts in a single API call (batch mode) /// Returns a vector of 768-dimensional vectors /// This is much more efficient than calling generate_embedding multiple times - pub async fn generate_embeddings(&self, cx: &Context, texts: &[&str]) -> Result>> { + pub async fn generate_embeddings(&self, texts: &[&str]) -> Result>> { let embedding_model = EMBEDDING_MODEL; - let mut span = self.start_request_span(cx, "embeddings", embedding_model); + let span = self.start_request_span("embeddings", embedding_model); log::debug!("=== Ollama Batch Embedding Request ==="); log::debug!("Model: {}", embedding_model); @@ -1148,7 +1115,7 @@ Analyze the image and use specific details from both the visual content and the embedding_model ); let primary_result = self - .try_generate_embeddings(Some(&mut span), &self.primary_url, embedding_model, texts) + .try_generate_embeddings(Some(&span), &self.primary_url, embedding_model, texts) .await; let embeddings = match primary_result { @@ -1171,12 +1138,7 @@ Analyze the image and use specific details from both the visual content and the embedding_model ); match self - .try_generate_embeddings( - Some(&mut span), - fallback_url, - embedding_model, - texts, - ) + .try_generate_embeddings(Some(&span), fallback_url, embedding_model, texts) .await { Ok(embeddings) => { @@ -1223,7 +1185,7 @@ Analyze the image and use specific details from both the visual content and the /// Internal helper to try generating embeddings for multiple texts from a specific server async fn try_generate_embeddings( &self, - mut span: Option<&mut opentelemetry::global::BoxedSpan>, + span: Option<&tracing::Span>, url: &str, model: &str, texts: &[&str], @@ -1238,13 +1200,13 @@ Analyze the image and use specific details from both the visual content and the .client .post(format!("{}/api/embed", url)) .json(&request); - if let Some(s) = span.as_mut() { - rb = crate::otel::with_traceparent(rb, s.span_context()); + if let Some(s) = span { + rb = crate::otel::with_traceparent(rb, s); } match rb.send().await { Ok(r) => r, Err(e) => { - if let Some(s) = span.as_mut() { + if let Some(s) = span { s.set_status(Status::error("Ollama embedding request failed")); } return Err(e.into()); @@ -1255,7 +1217,7 @@ Analyze the image and use specific details from both the visual content and the if !response.status().is_success() { let status = response.status(); crate::otel::record_http_failure( - span.as_deref_mut(), + span, status.as_u16(), &format!("Ollama batch embedding request failed: {status}"), ); @@ -1266,18 +1228,18 @@ Analyze the image and use specific details from both the visual content and the error_body )); } - crate::otel::set_http_status(span.as_deref_mut(), response.status().as_u16()); + crate::otel::set_http_status(span, response.status().as_u16()); let result: OllamaEmbedResponse = match response.json().await { Ok(r) => r, Err(e) => { - if let Some(s) = span.as_mut() { + if let Some(s) = span { s.set_status(Status::error("parsing Ollama embed response failed")); } return Err(e.into()); } }; - if let Some(s) = span.as_mut() { + if let Some(s) = span { s.set_status(Status::Ok); } Ok(result.embeddings) @@ -1288,38 +1250,35 @@ Analyze the image and use specific details from both the visual content and the impl LlmClient for OllamaClient { async fn generate( &self, - cx: &Context, prompt: &str, system: Option<&str>, images: Option>, ) -> Result { - self.generate_with_images(cx, prompt, system, images).await + self.generate_with_images(prompt, system, images).await } async fn chat_with_tools( &self, - cx: &Context, messages: Vec, tools: Vec, ) -> Result<(ChatMessage, Option, Option)> { - OllamaClient::chat_with_tools(self, cx, messages, tools).await + OllamaClient::chat_with_tools(self, messages, tools).await } async fn chat_with_tools_stream( &self, - cx: &Context, messages: Vec, tools: Vec, ) -> Result>> { - OllamaClient::chat_with_tools_stream(self, cx, messages, tools).await + OllamaClient::chat_with_tools_stream(self, messages, tools).await } - async fn generate_embeddings(&self, cx: &Context, texts: &[&str]) -> Result>> { - OllamaClient::generate_embeddings(self, cx, texts).await + async fn generate_embeddings(&self, texts: &[&str]) -> Result>> { + OllamaClient::generate_embeddings(self, texts).await } - async fn describe_image(&self, cx: &Context, image_base64: &str) -> Result { - self.generate_photo_description(cx, image_base64).await + async fn describe_image(&self, image_base64: &str) -> Result { + self.generate_photo_description(image_base64).await } async fn list_models(&self) -> Result> { diff --git a/src/ai/openrouter.rs b/src/ai/openrouter.rs index 2b67755..01e226b 100644 --- a/src/ai/openrouter.rs +++ b/src/ai/openrouter.rs @@ -5,14 +5,14 @@ use anyhow::Context as _; use anyhow::{Result, anyhow, bail}; use async_trait::async_trait; -use opentelemetry::trace::{Span, Status}; -use opentelemetry::{Context, KeyValue}; +use opentelemetry::trace::Status; use reqwest::Client; use serde::Deserialize; use serde_json::{Value, json}; use std::collections::HashMap; use std::sync::{Arc, Mutex}; use std::time::{Duration, Instant}; +use tracing_opentelemetry::OpenTelemetrySpanExt as _; use crate::ai::llm_client::{ ChatMessage, LlmClient, LlmStreamEvent, ModelCapabilities, Tool, ToolCall, ToolCallFunction, @@ -152,13 +152,8 @@ impl OpenRouterClient { /// `chat`, `embeddings`, or `describe` — the call shape, not the entry /// point (a `describe_image` funneled through `chat_with_tools` is /// still a `describe`). - fn start_request_span( - &self, - cx: &Context, - operation: &str, - model: &str, - ) -> opentelemetry::global::BoxedSpan { - crate::otel::start_llm_request_span(cx, "openrouter", operation, model) + fn start_request_span(&self, operation: &str, model: &str) -> tracing::Span { + crate::otel::start_llm_request_span("openrouter", operation, model) } /// Translate canonical messages to the OpenAI-compatible wire shape. @@ -329,7 +324,6 @@ impl OpenRouterClient { impl LlmClient for OpenRouterClient { async fn generate( &self, - cx: &Context, prompt: &str, system: Option<&str>, images: Option>, @@ -344,31 +338,29 @@ impl LlmClient for OpenRouterClient { // Delegates to chat_with_tools, which starts the llm.request span // (operation=chat) — one span per HTTP request. - let (reply, _, _) = self.chat_with_tools(cx, messages, Vec::new()).await?; + let (reply, _, _) = self.chat_with_tools(messages, Vec::new()).await?; Ok(reply.content) } async fn chat_with_tools( &self, - cx: &Context, messages: Vec, tools: Vec, ) -> Result<(ChatMessage, Option, Option)> { - let mut span = self.start_request_span(cx, "chat", &self.primary_model); - self.chat_with_tools_impl(Some(&mut span), messages, tools) + let span = self.start_request_span("chat", &self.primary_model); + self.chat_with_tools_impl(Some(&span), messages, tools) .await } async fn chat_with_tools_stream( &self, - cx: &Context, messages: Vec, tools: Vec, ) -> Result>> { // The span outlives this method — the stream yields until the // response fully drains, so the owned span moves into the // generator and is finalized there. - let mut span = self.start_request_span(cx, "chat_stream", &self.primary_model); + let span = self.start_request_span("chat_stream", &self.primary_model); let url = format!("{}/chat/completions", self.base_url); let mut body = serde_json::Map::new(); @@ -397,7 +389,7 @@ impl LlmClient for OpenRouterClient { let mut request = self .authed(self.client.post(&url)) .json(&Value::Object(body)); - request = crate::otel::with_traceparent(request, span.span_context()); + request = crate::otel::with_traceparent(request, &span); let resp = request .send() .await @@ -406,14 +398,14 @@ impl LlmClient for OpenRouterClient { if !resp.status().is_success() { let status = resp.status(); crate::otel::record_http_failure( - Some(&mut span), + Some(&span), status.as_u16(), &format!("OpenRouter stream request failed: {status}"), ); let body = resp.text().await.unwrap_or_default(); bail!("OpenRouter stream request failed: {} — {}", status, body); } - crate::otel::set_http_status(Some(&mut span), resp.status().as_u16()); + crate::otel::set_http_status(Some(&span), resp.status().as_u16()); // OpenAI-compat SSE stream. Each event is `data: \n\n`, with // `data: [DONE]` signalling completion. Tool calls arrive as @@ -489,7 +481,7 @@ impl LlmClient for OpenRouterClient { .and_then(|n| n.as_i64()) .map(|n| n as i32); if let Some(cost) = usage.get("cost").and_then(|c| c.as_f64()) { - span.set_attribute(KeyValue::new("llm.cost_usd", cost)); + span.set_attribute("llm.cost_usd", cost); } } @@ -581,10 +573,10 @@ impl LlmClient for OpenRouterClient { images: None, }; if let Some(pt) = prompt_tokens { - span.set_attribute(KeyValue::new("llm.tokens.input", pt as i64)); + span.set_attribute("llm.tokens.input", pt as i64); } if let Some(ct) = completion_tokens { - span.set_attribute(KeyValue::new("llm.tokens.output", ct as i64)); + span.set_attribute("llm.tokens.output", ct as i64); } span.set_status(Status::Ok); yield Ok(LlmStreamEvent::Done { @@ -597,8 +589,8 @@ impl LlmClient for OpenRouterClient { Ok(Box::pin(stream)) } - async fn generate_embeddings(&self, cx: &Context, texts: &[&str]) -> Result>> { - let mut span = self.start_request_span(cx, "embeddings", &self.embedding_model); + async fn generate_embeddings(&self, texts: &[&str]) -> Result>> { + let span = self.start_request_span("embeddings", &self.embedding_model); let url = format!("{}/embeddings", self.base_url); let body = json!({ "model": self.embedding_model, @@ -606,7 +598,7 @@ impl LlmClient for OpenRouterClient { }); let mut request = self.authed(self.client.post(&url)).json(&body); - request = crate::otel::with_traceparent(request, span.span_context()); + request = crate::otel::with_traceparent(request, &span); let resp = request .send() .await @@ -618,14 +610,14 @@ impl LlmClient for OpenRouterClient { if !resp.status().is_success() { let status = resp.status(); crate::otel::record_http_failure( - Some(&mut span), + Some(&span), status.as_u16(), &format!("OpenRouter embedding request failed: {status}"), ); let body = resp.text().await.unwrap_or_default(); bail!("OpenRouter embedding request failed: {} — {}", status, body); } - crate::otel::set_http_status(Some(&mut span), resp.status().as_u16()); + crate::otel::set_http_status(Some(&span), resp.status().as_u16()); #[derive(Deserialize)] struct EmbedResponse { @@ -641,10 +633,10 @@ impl LlmClient for OpenRouterClient { Ok(parsed.data.into_iter().map(|i| i.embedding).collect()) } - async fn describe_image(&self, cx: &Context, image_base64: &str) -> Result { + async fn describe_image(&self, image_base64: &str) -> Result { // Own span with operation=describe_image — the funnel into // chat_with_tools_impl keeps the request on a single span. - let mut span = self.start_request_span(cx, "describe_image", &self.primary_model); + let span = self.start_request_span("describe_image", &self.primary_model); let prompt = "Briefly describe what you see in this image in 1-2 sentences. \ Focus on the people, location, and activity."; let mut messages: Vec = Vec::new(); @@ -656,7 +648,7 @@ impl LlmClient for OpenRouterClient { messages.push(user); let (reply, _, _) = self - .chat_with_tools_impl(Some(&mut span), messages, Vec::new()) + .chat_with_tools_impl(Some(&span), messages, Vec::new()) .await?; Ok(reply.content) } @@ -719,7 +711,7 @@ impl OpenRouterClient { /// operation label. async fn chat_with_tools_impl( &self, - mut span: Option<&mut opentelemetry::global::BoxedSpan>, + span: Option<&tracing::Span>, messages: Vec, tools: Vec, ) -> Result<(ChatMessage, Option, Option)> { @@ -751,13 +743,13 @@ impl OpenRouterClient { let mut request = self .authed(self.client.post(&url)) .json(&Value::Object(body)); - if let Some(s) = span.as_mut() { - request = crate::otel::with_traceparent(request, s.span_context()); + if let Some(s) = span { + request = crate::otel::with_traceparent(request, s); } let resp = match request.send().await { Ok(r) => r, Err(e) => { - if let Some(s) = span.as_mut() { + if let Some(s) = span { s.set_status(Status::error("OpenRouter request failed")); } return Err(e).with_context(|| format!("POST {} failed", url)); @@ -767,19 +759,19 @@ impl OpenRouterClient { if !resp.status().is_success() { let status = resp.status(); crate::otel::record_http_failure( - span.as_deref_mut(), + span, status.as_u16(), &format!("OpenRouter chat request failed: {status}"), ); let body = resp.text().await.unwrap_or_default(); bail!("OpenRouter chat request failed: {} — {}", status, body); } - crate::otel::set_http_status(span.as_deref_mut(), resp.status().as_u16()); + crate::otel::set_http_status(span, resp.status().as_u16()); let parsed: Value = match resp.json().await { Ok(v) => v, Err(e) => { - if let Some(s) = span.as_mut() { + if let Some(s) = span { s.set_status(Status::error("parsing chat response failed")); } return Err(e).context("parsing chat response"); @@ -813,15 +805,15 @@ impl OpenRouterClient { .and_then(|v| v.as_i64()) .map(|n| n as i32); // OpenRouter reports per-request USD cost in `usage.cost`. - if let Some(s) = span.as_mut() { + if let Some(s) = span { if let Some(pt) = prompt_tokens { - s.set_attribute(KeyValue::new("llm.tokens.input", pt as i64)); + s.set_attribute("llm.tokens.input", pt as i64); } if let Some(ct) = completion_tokens { - s.set_attribute(KeyValue::new("llm.tokens.output", ct as i64)); + s.set_attribute("llm.tokens.output", ct as i64); } if let Some(cost) = usage.and_then(|u| u.get("cost")).and_then(|c| c.as_f64()) { - s.set_attribute(KeyValue::new("llm.cost_usd", cost)); + s.set_attribute("llm.cost_usd", cost); } s.set_status(Status::Ok); } diff --git a/src/ai/persona_chat.rs b/src/ai/persona_chat.rs index f9ff000..cbf2ae5 100644 --- a/src/ai/persona_chat.rs +++ b/src/ai/persona_chat.rs @@ -22,15 +22,16 @@ //! registry. //! - The HTTP handlers under ` /persona_chat/*`. -use actix_web::{HttpRequest, HttpResponse, Responder, delete, get, post, web}; +use actix_web::{HttpResponse, Responder, delete, get, post, web}; use anyhow::{Context, Result, anyhow, bail}; use chrono::Utc; -use opentelemetry::KeyValue; -use opentelemetry::trace::{Span, Status, TraceContextExt, Tracer}; +use opentelemetry::trace::Status; use serde::{Deserialize, Serialize}; use std::collections::HashMap; use std::sync::{Arc, Mutex as StdMutex}; use tokio::sync::Mutex as TokioMutex; +use tracing::Instrument; +use tracing_opentelemetry::OpenTelemetrySpanExt; use uuid::Uuid; use crate::ai::backend::{BackendKind, SamplingOverrides}; @@ -45,7 +46,6 @@ use crate::ai::turn_registry::{TurnEntry, TurnRegistry}; use crate::data::Claims; use crate::database::PersonaDao; use crate::libraries; -use crate::otel::{extract_context_from_request, global_tracer}; use crate::state::AppState; // ───────────────────────────────────────────────────────────────────── @@ -660,7 +660,6 @@ impl PersonaChatSession { let svc = self.clone(); let persona_for_task = persona.clone(); let entry_clone = entry.clone(); - let turn_id_for_task = turn_id.clone(); let req_for_task = PersonaChatTurnRequest { user_message: trimmed.clone(), ..req @@ -669,44 +668,50 @@ impl PersonaChatSession { in_flight: self.in_flight.clone(), conversation_id: conversation_id.clone(), }; - let conversation_for_span = conversation_id.clone(); - let handle = tokio::spawn(async move { - let tracer = global_tracer(); - let mut span = tracer.start("ai.persona_chat.turn.execute"); - span.set_attribute(KeyValue::new("turn_id", turn_id_for_task.clone())); - span.set_attribute(KeyValue::new("persona_id", persona_id.clone())); - span.set_attribute(KeyValue::new("conversation_id", conversation_for_span)); - span.set_attribute(KeyValue::new("library_id", library_id as i64)); - - let result = svc - .run_streaming_turn( - persona_for_task, - req_for_task, - entry_clone.clone(), - user_id, - library_id, - &mut span, - ) - .await; - match result { - Ok(()) => { - span.set_attribute(KeyValue::new("status", "done")); - span.set_status(Status::Ok); - } - Err(e) => { - span.set_attribute(KeyValue::new("status", "error")); - span.set_status(Status::error(format!("{e}"))); - let _ = entry_clone - .push_event(ChatStreamEvent::Error(format!("{e}"))) - .await; + // Created before the spawn so it parents on the dispatching request, + // then carried into the task with `.instrument` (tracing context is + // thread-local and does not cross `tokio::spawn` on its own). + let turn_span = tracing::info_span!( + "ai.persona_chat.turn.execute", + turn_id = %turn_id, + persona_id = %persona_id, + conversation_id = %conversation_id, + library_id = library_id, + ); + let handle = tokio::spawn( + async move { + let span = tracing::Span::current(); + let result = svc + .run_streaming_turn( + persona_for_task, + req_for_task, + entry_clone.clone(), + user_id, + library_id, + &span, + ) + .await; + match result { + Ok(()) => { + span.set_attribute("status", "done"); + span.set_status(Status::Ok); + } + Err(e) => { + span.set_attribute("status", "error"); + span.set_status(Status::error(format!("{e}"))); + let _ = entry_clone + .push_event(ChatStreamEvent::Error(format!("{e}"))) + .await; + } } + entry_clone.set_terminal_status(crate::ai::turn_registry::TurnStatus::Done); + // Frees the in-flight slot. Placed explicitly so the guard is + // captured by the spawned task: it must drop when the task ends + // (completion, error, or abort) — not when dispatch returns. + drop(in_flight_guard); } - entry_clone.set_terminal_status(crate::ai::turn_registry::TurnStatus::Done); - // Frees the in-flight slot. Placed explicitly so the guard is - // captured by the spawned task: it must drop when the task ends - // (completion, error, or abort) — not when dispatch returns. - drop(in_flight_guard); - }); + .instrument(turn_span), + ); entry.set_abort_handle(handle.abort_handle()); Ok(turn_id) @@ -722,7 +727,7 @@ impl PersonaChatSession { entry: Arc, user_id: i32, _library_id: i32, - turn_span: &mut opentelemetry::global::BoxedSpan, + turn_span: &tracing::Span, ) -> Result<()> { // Per-conversation mutex — serializes this turn against a // concurrent rewind or branch switch on the same transcript. @@ -810,12 +815,10 @@ impl PersonaChatSession { .resolve_backend(backend_kind, &overrides) .await?; let model_used = backend.model().to_string(); - // llm.* attributes on the turn-level span (issue #130), plus the - // parent context for the child llm.request spans. - turn_span.set_attribute(KeyValue::new("llm.backend", backend.kind.as_str())); - turn_span.set_attribute(KeyValue::new("llm.model", model_used.clone())); - let turn_cx = opentelemetry::Context::new() - .with_remote_span_context(turn_span.span_context().clone()); + // llm.* attributes on the turn-level span (issue #130). The child + // llm.request spans parent on it ambiently (the task is instrumented). + turn_span.set_attribute("llm.backend", backend.kind.as_str()); + turn_span.set_attribute("llm.model", model_used.clone()); // Build a no-photo tool catalog. `current_gate_opts_for_persona` // already probes each tool's backing table for presence and @@ -882,7 +885,6 @@ impl PersonaChatSession { &persona.persona_id, max_iterations, &entry, - &turn_cx, ) .await .map_err(|e| anyhow!("persona chat loop failed: {e:?}"))?; @@ -951,7 +953,7 @@ impl PersonaChatSession { // list falls back to the persona name while the title is empty. if needs_title { let title = self - .generate_title(&backend, &messages, &turn_cx) + .generate_title(&backend, &messages) .await .unwrap_or_else(|| fallback_title(&messages)); if !title.is_empty() { @@ -1056,7 +1058,6 @@ impl PersonaChatSession { &self, backend: &crate::ai::backend::ResolvedBackend, messages: &[ChatMessage], - cx: &opentelemetry::Context, ) -> Option { let user_message = messages.iter().find(|m| m.role == "user")?; // The last assistant message with actual prose — earlier ones may be @@ -1067,7 +1068,7 @@ impl PersonaChatSession { .find(|m| m.role == "assistant" && !m.content.trim().is_empty())?; let prompt = title_prompt(&user_message.content, &reply.content); - match backend.chat().generate(cx, &prompt, None, None).await { + match backend.chat().generate(&prompt, None, None).await { Ok(raw) => { let title = sanitize_title(&crate::ai::llm_client::strip_think_blocks(&raw)); if title.is_empty() { None } else { Some(title) } @@ -1276,19 +1277,14 @@ pub async fn persona_chat_history_handler( /// the `turn_id` immediately; the streaming agent loop runs in a /// background task and the client opens the SSE replay stream with it. #[post("/persona_chat/turn")] +#[tracing::instrument(name = "http.persona_chat.turn", skip_all)] pub async fn persona_chat_turn_handler( - http_request: HttpRequest, claims: Claims, request: web::Json, app_state: web::Data, ) -> impl Responder { - let parent_context = extract_context_from_request(&http_request); - let tracer = global_tracer(); - let mut span = tracer.start_with_context("http.persona_chat.turn", &parent_context); - span.set_attribute(KeyValue::new( - "conversation_id", - request.conversation_id.clone(), - )); + let span = tracing::Span::current(); + span.set_attribute("conversation_id", request.conversation_id.clone()); // The transcript is not library-scoped, but the turn's tool catalog // (memories, SMS, places) resolves against a library — mirror the @@ -1313,7 +1309,7 @@ pub async fn persona_chat_turn_handler( .await { Ok(turn_id) => { - span.set_attribute(KeyValue::new("turn_id", turn_id.clone())); + span.set_attribute("turn_id", turn_id.clone()); span.set_status(Status::Ok); HttpResponse::Accepted().json(serde_json::json!({ "turn_id": turn_id, @@ -1540,22 +1536,20 @@ pub async fn persona_chat_branches_handler( /// shared renderer. #[get("/persona_chat/turn/{turn_id}")] pub async fn persona_turn_replay_handler( - http_request: HttpRequest, path: web::Path, query: web::Query, app_state: web::Data, ) -> HttpResponse { - crate::ai::handlers::turn_replay_impl(http_request, path, query, app_state).await + crate::ai::handlers::turn_replay_impl(path, query, app_state).await } /// DELETE /persona_chat/turn/{turn_id} — cancel a running persona turn. #[delete("/persona_chat/turn/{turn_id}")] pub async fn persona_turn_cancel_handler( - http_request: HttpRequest, path: web::Path, app_state: web::Data, ) -> impl Responder { - crate::ai::handlers::cancel_turn_impl(http_request, path, app_state).await + crate::ai::handlers::cancel_turn_impl(path, app_state).await } // ───────────────────────────────────────────────────────────────────── diff --git a/src/ai/sms_client.rs b/src/ai/sms_client.rs index 02c39f1..d5e175f 100644 --- a/src/ai/sms_client.rs +++ b/src/ai/sms_client.rs @@ -324,7 +324,6 @@ impl SmsApiClient { &self, messages: &[SmsMessage], ollama: &OllamaClient, - cx: &opentelemetry::Context, ) -> Result { if messages.is_empty() { return Ok(String::from("No messages on this day")); @@ -357,7 +356,6 @@ Summary:"#, ollama .generate( - cx, &prompt, // Some("You are a summarizer for the purposes of jogging my memory and highlighting events and situations."), Some("You are the keeper of memories, ingest the context and give me a casual summary of the moment."), diff --git a/src/ai/tts.rs b/src/ai/tts.rs index d6ef89d..cd102d7 100644 --- a/src/ai/tts.rs +++ b/src/ai/tts.rs @@ -6,13 +6,12 @@ // (audio read directly; video has its audio track extracted via ffmpeg). use actix_multipart::Multipart; -use actix_web::{HttpRequest, HttpResponse, Responder, delete, get, post, web}; +use actix_web::{HttpResponse, Responder, delete, get, post, web}; use anyhow::Context; use base64::Engine; use bytes::{BufMut, BytesMut}; use futures::StreamExt; -use opentelemetry::KeyValue; -use opentelemetry::trace::{Span, Status, Tracer}; +use opentelemetry::trace::Status; use regex::Regex; use serde::{Deserialize, Serialize}; use serde_json::{Value, json}; @@ -21,6 +20,8 @@ use std::path::Path; use std::sync::{LazyLock, Mutex as StdMutex}; use std::time::{Duration, Instant}; use tokio::sync::Semaphore; +use tracing::Instrument; +use tracing_opentelemetry::OpenTelemetrySpanExt; use uuid::Uuid; use crate::ai::llamacpp::LlamaCppClient; @@ -28,7 +29,6 @@ use crate::data::Claims; use crate::file_types::{is_audio_file, is_video_file}; use crate::files::is_valid_full_path; use crate::libraries; -use crate::otel::{extract_context_from_request, global_tracer}; use crate::state::AppState; /// Hard cap on an uploaded voice-reference clip. Chatterbox itself caps the @@ -382,14 +382,13 @@ pub struct TtsSpeechResponse { /// POST /tts/speech — synthesize `text` (optionally in a named `voice`) and /// return base64-encoded audio for `data:` URI playback on the client. #[post("/tts/speech")] +#[tracing::instrument(name = "http.tts.speech", skip_all)] pub async fn tts_speech_handler( - http_request: HttpRequest, _claims: Claims, req: web::Json, app_state: web::Data, ) -> impl Responder { - let parent_context = extract_context_from_request(&http_request); - let mut span = global_tracer().start_with_context("http.tts.speech", &parent_context); + let span = tracing::Span::current(); let text = prepare_for_tts(&req.text); if text.is_empty() { @@ -414,10 +413,10 @@ pub async fn tts_speech_handler( .filter(|s| !s.is_empty()) .or(dv.as_deref()); - span.set_attribute(KeyValue::new("tts.model", client.tts_model.clone())); - span.set_attribute(KeyValue::new("tts.format", format.to_string())); - span.set_attribute(KeyValue::new("tts.has_voice", voice.is_some())); - span.set_attribute(KeyValue::new("tts.text_len", text.len() as i64)); + span.set_attribute("tts.model", client.tts_model.clone()); + span.set_attribute("tts.format", format.to_string()); + span.set_attribute("tts.has_voice", voice.is_some()); + span.set_attribute("tts.text_len", text.len() as i64); // Clamp generation knobs to Chatterbox's documented ranges before forwarding. let exaggeration = req.exaggeration.map(|x| x.clamp(0.25, 2.0)); @@ -441,7 +440,7 @@ pub async fn tts_speech_handler( .await { Ok(bytes) => { - span.set_attribute(KeyValue::new("tts.audio_bytes", bytes.len() as i64)); + span.set_attribute("tts.audio_bytes", bytes.len() as i64); span.set_status(Status::Ok); let audio_base64 = base64::engine::general_purpose::STANDARD.encode(&bytes); HttpResponse::Ok().json(TtsSpeechResponse { @@ -513,15 +512,13 @@ pub async fn synthesize_serialized( /// permit (instead of fast-failing 429) and the client polls the job until /// the audio is ready. #[post("/tts/speech/jobs")] +#[tracing::instrument(name = "http.tts.speech_job.create", skip_all)] pub async fn create_speech_job_handler( - http_request: HttpRequest, _claims: Claims, req: web::Json, app_state: web::Data, ) -> impl Responder { - let parent_context = extract_context_from_request(&http_request); - let mut span = - global_tracer().start_with_context("http.tts.speech_job.create", &parent_context); + let span = tracing::Span::current(); let text = prepare_for_tts(&req.text); if text.is_empty() { @@ -550,9 +547,9 @@ pub async fn create_speech_job_handler( let cfg_weight = req.cfg_weight.map(|x| x.clamp(0.0, 1.0)); let temperature = req.temperature.map(|x| x.clamp(0.05, 5.0)); - span.set_attribute(KeyValue::new("tts.format", format.clone())); - span.set_attribute(KeyValue::new("tts.has_voice", voice.is_some())); - span.set_attribute(KeyValue::new("tts.text_len", text.len() as i64)); + span.set_attribute("tts.format", format.clone()); + span.set_attribute("tts.has_voice", voice.is_some()); + span.set_attribute("tts.text_len", text.len() as i64); let job_id = Uuid::new_v4(); { @@ -573,75 +570,79 @@ pub async fn create_speech_job_handler( } let state = app_state.clone(); - let handle = tokio::spawn(async move { - // Queue rather than fast-fail: jobs wait their turn for the GPU. - let _permit = match TTS_PERMIT.acquire().await { - Ok(p) => p, - Err(_) => { - finish_job( - job_id, - TtsJobStatus::Error, - None, - Some("TTS queue closed".to_string()), - ); + let handle = tokio::spawn( + async move { + // Queue rather than fast-fail: jobs wait their turn for the GPU. + let _permit = match TTS_PERMIT.acquire().await { + Ok(p) => p, + Err(_) => { + finish_job( + job_id, + TtsJobStatus::Error, + None, + Some("TTS queue closed".to_string()), + ); + return; + } + }; + // Wait for the LLM side to release the GPU too (see ai::gpu) — only + // then does the job count as running. The synthesis timeout starts at + // the HTTP send below, so neither wait burns it, and the client can + // anchor its own deadline to the queued→running transition. + let _gpu = crate::ai::gpu::tts_lease().await; + + // Cancelled while queued — release the permits without synthesizing. + let cancelled = with_job(job_id, |job| { + if job.status == TtsJobStatus::Queued { + job.status = TtsJobStatus::Running; + false + } else { + true + } + }) + .unwrap_or(true); + if cancelled { return; } - }; - // Wait for the LLM side to release the GPU too (see ai::gpu) — only - // then does the job count as running. The synthesis timeout starts at - // the HTTP send below, so neither wait burns it, and the client can - // anchor its own deadline to the queued→running transition. - let _gpu = crate::ai::gpu::tts_lease().await; - // Cancelled while queued — release the permits without synthesizing. - let cancelled = with_job(job_id, |job| { - if job.status == TtsJobStatus::Queued { - job.status = TtsJobStatus::Running; - false - } else { - true - } - }) - .unwrap_or(true); - if cancelled { - return; - } - - let Some(client) = state.llamacpp.as_ref() else { - finish_job( - job_id, - TtsJobStatus::Error, - None, - Some("TTS backend not configured".to_string()), - ); - return; - }; - match client - .text_to_speech( - &text, - voice.as_deref(), - &format, - exaggeration, - cfg_weight, - temperature, - ) - .await - { - Ok(bytes) => { - let audio = base64::engine::general_purpose::STANDARD.encode(&bytes); - finish_job(job_id, TtsJobStatus::Done, Some(audio), None); - } - Err(e) => { - log::error!("TTS job {job_id} failed: {:?}", e); + let Some(client) = state.llamacpp.as_ref() else { finish_job( job_id, TtsJobStatus::Error, None, - Some(format!("TTS failed: {e}")), + Some("TTS backend not configured".to_string()), ); + return; + }; + match client + .text_to_speech( + &text, + voice.as_deref(), + &format, + exaggeration, + cfg_weight, + temperature, + ) + .await + { + Ok(bytes) => { + let audio = base64::engine::general_purpose::STANDARD.encode(&bytes); + finish_job(job_id, TtsJobStatus::Done, Some(audio), None); + } + Err(e) => { + log::error!("TTS job {job_id} failed: {:?}", e); + finish_job( + job_id, + TtsJobStatus::Error, + None, + Some(format!("TTS failed: {e}")), + ); + } } } - }); + // Carry the handler span into the job (tracing context is thread-local). + .instrument(tracing::Span::current()), + ); // Aborting an already-finished task is a no-op, so this late install is // safe even if the job raced to completion. with_job(job_id, |job| { @@ -660,14 +661,9 @@ pub async fn create_speech_job_handler( /// GET /tts/speech/jobs/{id} — poll a speech job; returns the audio once done. /// 404s after the job expires (results are kept ~10 min past completion). #[get("/tts/speech/jobs/{id}")] -pub async fn speech_job_status_handler( - http_request: HttpRequest, - _claims: Claims, - path: web::Path, -) -> impl Responder { - let parent_context = extract_context_from_request(&http_request); - let mut span = - global_tracer().start_with_context("http.tts.speech_job.status", &parent_context); +#[tracing::instrument(name = "http.tts.speech_job.status", skip_all)] +pub async fn speech_job_status_handler(_claims: Claims, path: web::Path) -> impl Responder { + let span = tracing::Span::current(); let Ok(id) = Uuid::parse_str(&path.into_inner()) else { span.set_status(Status::error("invalid job id")); @@ -701,14 +697,9 @@ pub async fn speech_job_status_handler( /// limitation as the sync path); cancelling stops the wait and discards the /// result. Cancelling an already-finished job leaves it terminal. #[delete("/tts/speech/jobs/{id}")] -pub async fn cancel_speech_job_handler( - http_request: HttpRequest, - _claims: Claims, - path: web::Path, -) -> impl Responder { - let parent_context = extract_context_from_request(&http_request); - let mut span = - global_tracer().start_with_context("http.tts.speech_job.cancel", &parent_context); +#[tracing::instrument(name = "http.tts.speech_job.cancel", skip_all)] +pub async fn cancel_speech_job_handler(_claims: Claims, path: web::Path) -> impl Responder { + let span = tracing::Span::current(); let Ok(id) = Uuid::parse_str(&path.into_inner()) else { span.set_status(Status::error("invalid job id")); @@ -749,21 +740,20 @@ pub struct ListVoicesQuery { /// in-memory cache when possible so browsing settings doesn't make llama-swap /// load the TTS model (and evict the resident LLM); see VOICES_CACHE. #[get("/tts/voices")] +#[tracing::instrument(name = "http.tts.voices.list", skip_all)] pub async fn list_voices_handler( - http_request: HttpRequest, _claims: Claims, query: web::Query, app_state: web::Data, ) -> impl Responder { - let parent_context = extract_context_from_request(&http_request); - let mut span = global_tracer().start_with_context("http.tts.voices.list", &parent_context); + let span = tracing::Span::current(); let force = query .refresh .as_deref() .is_some_and(|v| matches!(v, "1" | "true" | "yes")); if !force && let Some(v) = cached_voices() { - span.set_attribute(KeyValue::new("tts.voices_cache_hit", true)); + span.set_attribute("tts.voices_cache_hit", true); span.set_status(Status::Ok); return HttpResponse::Ok().json(v); } @@ -776,7 +766,7 @@ pub async fn list_voices_handler( match client.list_voices().await { Ok(v) => { store_voices_cache(&v); - span.set_attribute(KeyValue::new("tts.voices_cache_hit", false)); + span.set_attribute("tts.voices_cache_hit", false); span.set_status(Status::Ok); HttpResponse::Ok().json(v) } @@ -790,14 +780,13 @@ pub async fn list_voices_handler( /// DELETE /tts/voices/{name} — remove a cloned voice from the library. #[delete("/tts/voices/{name}")] +#[tracing::instrument(name = "http.tts.voices.delete", skip_all)] pub async fn delete_voice_handler( - http_request: HttpRequest, _claims: Claims, path: web::Path, app_state: web::Data, ) -> impl Responder { - let parent_context = extract_context_from_request(&http_request); - let mut span = global_tracer().start_with_context("http.tts.voices.delete", &parent_context); + let span = tracing::Span::current(); let Some(client) = app_state.llamacpp.as_ref() else { span.set_status(Status::error("tts backend not configured")); @@ -814,7 +803,7 @@ pub async fn delete_voice_handler( return HttpResponse::BadRequest().json(json!({ "error": "invalid voice name" })); } }; - span.set_attribute(KeyValue::new("tts.voice_name", name.clone())); + span.set_attribute("tts.voice_name", name.clone()); match client.delete_voice(&name).await { Ok(v) => { @@ -835,14 +824,13 @@ pub async fn delete_voice_handler( /// plus optional `start_seconds` / `duration_seconds` (text) selecting which /// window of a longer recording becomes the reference clip. #[post("/tts/voices/upload")] +#[tracing::instrument(name = "http.tts.voices.upload", skip_all)] pub async fn create_voice_upload_handler( - http_request: HttpRequest, _claims: Claims, mut payload: Multipart, app_state: web::Data, ) -> impl Responder { - let parent_context = extract_context_from_request(&http_request); - let mut span = global_tracer().start_with_context("http.tts.voices.upload", &parent_context); + let span = tracing::Span::current(); let Some(client) = app_state.llamacpp.as_ref() else { span.set_status(Status::error("tts backend not configured")); @@ -932,8 +920,8 @@ pub async fn create_voice_upload_handler( span.set_status(Status::error("voice_file is required")); return HttpResponse::BadRequest().json(json!({ "error": "voice_file is required" })); } - span.set_attribute(KeyValue::new("tts.voice_name", name.clone())); - span.set_attribute(KeyValue::new("tts.upload_bytes", file_bytes.len() as i64)); + span.set_attribute("tts.voice_name", name.clone()); + span.set_attribute("tts.upload_bytes", file_bytes.len() as i64); // Normalize to WAV so any device format (e.g. .aac / .opus, which Chatterbox // rejects by extension) is accepted. @@ -987,15 +975,13 @@ pub struct CreateVoiceFromLibraryRequest { /// WAV reference clip (window selected by start/duration_seconds, length /// capped by LLAMA_SWAP_TTS_REF_SECONDS). #[post("/tts/voices/from-library")] +#[tracing::instrument(name = "http.tts.voices.from_library", skip_all)] pub async fn create_voice_from_library_handler( - http_request: HttpRequest, _claims: Claims, req: web::Json, app_state: web::Data, ) -> impl Responder { - let parent_context = extract_context_from_request(&http_request); - let mut span = - global_tracer().start_with_context("http.tts.voices.from_library", &parent_context); + let span = tracing::Span::current(); let Some(client) = app_state.llamacpp.as_ref() else { span.set_status(Status::error("tts backend not configured")); @@ -1045,7 +1031,7 @@ pub async fn create_voice_from_library_handler( return HttpResponse::BadRequest() .json(json!({ "error": "file is not an audio or video file" })); } - span.set_attribute(KeyValue::new("tts.voice_name", voice_name.clone())); + span.set_attribute("tts.voice_name", voice_name.clone()); let wav = match prepare_reference_audio(&abs, ref_start, ref_duration).await { Ok(b) => b, diff --git a/src/bin/test_daily_summary.rs b/src/bin/test_daily_summary.rs index f497db5..d04aa32 100644 --- a/src/bin/test_daily_summary.rs +++ b/src/bin/test_daily_summary.rs @@ -220,9 +220,7 @@ async fn main() -> Result<()> { println!("Generating summary..."); - let summary = ollama - .generate(&opentelemetry::Context::new(), &prompt, Some(system_prompt)) - .await?; + let summary = ollama.generate(&prompt, Some(system_prompt)).await?; println!("\n📝 GENERATED SUMMARY:"); println!("─────────────────────────────────────────"); @@ -238,9 +236,7 @@ async fn main() -> Result<()> { ); println!("\nGenerating embedding..."); - let embedding = ollama - .generate_embedding(&opentelemetry::Context::new(), &stripped) - .await?; + let embedding = ollama.generate_embedding(&stripped).await?; println!("✓ Embedding generated ({} dimensions)", embedding.len()); println!("Saving to database..."); diff --git a/src/knowledge.rs b/src/knowledge.rs index 0783c36..2c6cf6f 100644 --- a/src/knowledge.rs +++ b/src/knowledge.rs @@ -859,8 +859,7 @@ async fn synthesize_merge( let ollama = app_state.ollama.clone(); let model_used = ollama.primary_model.clone(); - let cx = opentelemetry::Context::current(); - let proposed = match ollama.generate(&cx, &prompt, Some(system)).await { + let proposed = match ollama.generate(&prompt, Some(system)).await { Ok(out) => { // Strip the framing models reach for even with explicit // "no preamble" guidance: leading "Merged description:" diff --git a/src/otel.rs b/src/otel.rs index 3384c3b..9f7479a 100644 --- a/src/otel.rs +++ b/src/otel.rs @@ -2,17 +2,15 @@ use actix_web::HttpRequest; use actix_web::http::header::HeaderMap; use opentelemetry::global::{BoxedSpan, BoxedTracer}; use opentelemetry::propagation::TextMapPropagator; -use opentelemetry::trace::{ - Span, SpanBuilder, SpanContext, SpanKind, Status, TraceContextExt, Tracer, TracerProvider, -}; +use opentelemetry::trace::{Span, Status, TraceContextExt, Tracer, TracerProvider}; use opentelemetry::{Context, KeyValue, global}; -use tracing_subscriber::layer::SubscriberExt; -use tracing_subscriber::util::SubscriberInitExt; use opentelemetry_appender_log::OpenTelemetryLogBridge; use opentelemetry_otlp::WithExportConfig; use opentelemetry_sdk::Resource; use opentelemetry_sdk::logs::{BatchLogProcessor, SdkLoggerProvider}; use opentelemetry_sdk::propagation::TraceContextPropagator; +use tracing_subscriber::layer::SubscriberExt; +use tracing_subscriber::util::SubscriberInitExt; pub fn global_tracer() -> BoxedTracer { global::tracer("image-server") @@ -120,18 +118,19 @@ impl<'a> opentelemetry::propagation::Extractor for HeaderExtractor<'a> { /// provider / operation / model attributes) is identical across /// openrouter / llamacpp / ollama. `operation` is one of `chat`, /// `chat_stream`, `embeddings`, `describe_image`, `generate`. -pub fn start_llm_request_span( - cx: &Context, - provider: &str, - operation: &str, - model: &str, -) -> BoxedSpan { - let builder = SpanBuilder::from_name("llm.request").with_kind(SpanKind::Client); - let mut span = global_tracer().build_with_context(builder, cx); - span.set_attribute(KeyValue::new("llm.provider", provider.to_string())); - span.set_attribute(KeyValue::new("llm.operation", operation.to_string())); - span.set_attribute(KeyValue::new("llm.model", model.to_string())); - span +/// +/// The span parents on the caller's ambient `tracing` span — callers no +/// longer thread an OpenTelemetry `Context`. Code that hands work to a +/// `tokio::spawn` must carry its span across (`.instrument(span)`), or the +/// LLM span orphans onto a fresh trace. +pub fn start_llm_request_span(provider: &str, operation: &str, model: &str) -> tracing::Span { + tracing::info_span!( + "llm.request", + otel.kind = "client", + llm.provider = provider, + llm.operation = operation, + llm.model = model, + ) } /// The OpenTelemetry [`Context`] to parent this request's operation spans @@ -168,34 +167,35 @@ pub fn traceparent_headers(context: &Context) -> Vec<(String, String)> { headers.into_iter().filter(|(_, v)| !v.is_empty()).collect() } -/// Attach the W3C traceparent header to an outgoing request so a downstream -/// service parents its own spans onto this trace. No-op when the span context -/// is not valid — e.g. debug builds, where the global tracer is a no-op and an -/// all-zero traceparent would produce garbage parent spans. +/// Attach the W3C traceparent header for `span` to an outgoing request so a +/// downstream service parents its own spans onto this trace. No-op when the +/// span carries no valid OTel context — e.g. a span disabled by the +/// `EnvFilter`, where an all-zero traceparent would produce garbage parents. pub fn with_traceparent( builder: reqwest::RequestBuilder, - span_context: &SpanContext, + span: &tracing::Span, ) -> reqwest::RequestBuilder { + use tracing_opentelemetry::OpenTelemetrySpanExt as _; let mut b = builder; - for (k, v) in - traceparent_headers(&Context::new().with_remote_span_context(span_context.clone())) - { + for (k, v) in traceparent_headers(&span.context()) { b = b.header(k, v); } b } /// Record the HTTP status code on the span (success or failure path). -pub fn set_http_status(span: Option<&mut BoxedSpan>, status: u16) { +pub fn set_http_status(span: Option<&tracing::Span>, status: u16) { + use tracing_opentelemetry::OpenTelemetrySpanExt as _; if let Some(s) = span { - s.set_attribute(KeyValue::new("http.status_code", status as i64)); + s.set_attribute("http.status_code", status as i64); } } /// Record an HTTP failure on the span: status code plus an error status. -pub fn record_http_failure(span: Option<&mut BoxedSpan>, status: u16, message: &str) { +pub fn record_http_failure(span: Option<&tracing::Span>, status: u16, message: &str) { + use tracing_opentelemetry::OpenTelemetrySpanExt as _; if let Some(s) = span { - s.set_attribute(KeyValue::new("http.status_code", status as i64)); + s.set_attribute("http.status_code", status as i64); s.set_status(Status::error(message.to_string())); } } diff --git a/src/reels/script.rs b/src/reels/script.rs index 33b10b7..38ef9cc 100644 --- a/src/reels/script.rs +++ b/src/reels/script.rs @@ -281,9 +281,8 @@ pub async fn generate_script( beats: &[PlannedBeat], ) -> Result { let (system, user) = build_script_messages(meta, beats); - let cx = opentelemetry::Context::current(); let raw = client - .generate(&cx, &user, Some(&system), None) + .generate(&user, Some(&system), None) .await .context("LLM script generation failed")?; Ok(parse_script_response(&raw, beats.len())) diff --git a/src/unified_search.rs b/src/unified_search.rs index cba7987..0940a92 100644 --- a/src/unified_search.rs +++ b/src/unified_search.rs @@ -211,7 +211,7 @@ pub async fn unified_search( log::info!("unified_search: translating with model={}", backend.model()); let today = chrono::Utc::now().date_naive(); - let sq = match translate_nl_query(backend.chat(), &nl, &tag_vocab, today, &ctx).await { + let sq = match translate_nl_query(backend.chat(), &nl, &tag_vocab, today).await { Ok(sq) => sq, Err(e) => { log::warn!("unified_search: translate_nl_query failed: {e:?}"); diff --git a/tests/otel_llm_spans.rs b/tests/otel_llm_spans.rs new file mode 100644 index 0000000..b76b9a2 --- /dev/null +++ b/tests/otel_llm_spans.rs @@ -0,0 +1,181 @@ +//! Verifies the LLM-client tracing migration (#140 / map #135). +//! +//! The `LlmClient` methods no longer take an OpenTelemetry `Context`; the +//! `llm.request` span parents on the caller's ambient `tracing` span. Two +//! properties are pinned here, in the production shape of the streaming chat +//! path (`insight_chat::chat_turn_stream_inner`: a `tokio::spawn` carrying +//! the turn span via `.instrument`, on a multi-thread runtime): +//! +//! 1. **Nesting.** Both the blocking and the streaming `llm.request` spans +//! nest under the caller's span (same trace_id, parent link) and keep the +//! `llm.provider` / `llm.operation` / `llm.model` / `http.status_code` +//! attributes. +//! 2. **Propagation.** The outgoing `traceparent` header names the +//! `llm.request` span itself, so the downstream server parents onto it. +//! +//! One process-global subscriber per test binary, so this file holds a +//! single test. + +use futures::StreamExt; +use image_api::ai::llamacpp::LlamaCppClient; +use image_api::ai::llm_client::{ChatMessage, LlmClient, LlmStreamEvent}; +use opentelemetry::Value; +use opentelemetry::trace::TracerProvider; +use opentelemetry_sdk::trace::{InMemorySpanExporter, SdkTracerProvider, SpanData}; +use std::sync::{Arc, Mutex}; +use tokio::io::{AsyncReadExt, AsyncWriteExt}; +use tracing::Instrument; +use tracing_subscriber::layer::SubscriberExt; +use tracing_subscriber::registry::Registry; +use tracing_subscriber::util::SubscriberInitExt; + +const CHAT_BODY: &str = r#"{"choices":[{"message":{"role":"assistant","content":"hi"}}],"usage":{"prompt_tokens":3,"completion_tokens":1}}"#; +const STREAM_BODY: &str = "data: {\"choices\":[{\"delta\":{\"role\":\"assistant\",\"content\":\"hi\"}}]}\n\n\ +data: {\"choices\":[],\"usage\":{\"prompt_tokens\":3,\"completion_tokens\":1}}\n\n\ +data: [DONE]\n\n"; + +/// Minimal OpenAI-compatible server: records each request's `traceparent` +/// and answers with an SSE stream when the body asks for one, JSON otherwise. +async fn serve(listener: tokio::net::TcpListener, seen: Arc>>) { + loop { + let (mut sock, _) = listener.accept().await.expect("accept"); + let seen = seen.clone(); + tokio::spawn(async move { + let mut buf = Vec::new(); + let mut chunk = [0u8; 4096]; + let header_end = loop { + let n = sock.read(&mut chunk).await.expect("read"); + buf.extend_from_slice(&chunk[..n]); + if let Some(i) = buf.windows(4).position(|w| w == b"\r\n\r\n") { + break i + 4; + } + }; + let head = String::from_utf8_lossy(&buf[..header_end]).to_lowercase(); + let content_length: usize = head + .lines() + .find_map(|l| l.strip_prefix("content-length:")) + .map(|v| v.trim().parse().unwrap()) + .unwrap_or(0); + while buf.len() < header_end + content_length { + let n = sock.read(&mut chunk).await.expect("read body"); + buf.extend_from_slice(&chunk[..n]); + } + if let Some(tp) = head.lines().find_map(|l| l.strip_prefix("traceparent:")) { + seen.lock().unwrap().push(tp.trim().to_string()); + } + let body = String::from_utf8_lossy(&buf[header_end..]); + let (ctype, payload) = if body.contains("\"stream\":true") { + ("text/event-stream", STREAM_BODY) + } else { + ("application/json", CHAT_BODY) + }; + let resp = format!( + "HTTP/1.1 200 OK\r\ncontent-type: {ctype}\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{payload}", + payload.len() + ); + sock.write_all(resp.as_bytes()).await.expect("write"); + }); + } +} + +fn attr(span: &SpanData, key: &str) -> Option { + span.attributes + .iter() + .find(|kv| kv.key.as_str() == key) + .map(|kv| kv.value.clone()) +} + +#[test] +fn llm_spans_nest_under_caller_across_spawn() { + let exporter = InMemorySpanExporter::default(); + let provider = SdkTracerProvider::builder() + .with_simple_exporter(exporter.clone()) + .build(); + Registry::default() + .with(tracing_opentelemetry::OpenTelemetryLayer::new( + provider.tracer("llm-span-test"), + )) + .try_init() + .expect("global subscriber not yet set"); + + let rt = tokio::runtime::Builder::new_multi_thread() + .worker_threads(2) + .enable_all() + .build() + .unwrap(); + + let seen = Arc::new(Mutex::new(Vec::new())); + rt.block_on(async { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let base = format!("http://{}", listener.local_addr().unwrap()); + tokio::spawn(serve(listener, seen.clone())); + + let client = Arc::new(LlamaCppClient::new(Some(base), Some("chat".into()))); + let turn = tracing::info_span!("test.turn"); + tokio::spawn( + async move { + let (reply, _, _) = client + .chat_with_tools(vec![ChatMessage::user("hello")], vec![]) + .await + .expect("chat"); + assert_eq!(reply.content, "hi"); + + let mut stream = client + .chat_with_tools_stream(vec![ChatMessage::user("hello")], vec![]) + .await + .expect("stream"); + let mut done = false; + while let Some(ev) = stream.next().await { + if let LlmStreamEvent::Done { .. } = ev.expect("stream event") { + done = true; + } + } + assert!(done, "stream ended without Done"); + } + .instrument(turn), + ) + .await + .expect("spawned turn"); + }); + let _ = provider.force_flush(); + + let spans = exporter.get_finished_spans().unwrap(); + let turn = spans + .iter() + .find(|s| s.name == "test.turn") + .expect("turn span exported"); + let llm: Vec<&SpanData> = spans.iter().filter(|s| s.name == "llm.request").collect(); + assert_eq!(llm.len(), 2, "one llm.request per HTTP call"); + + let seen = seen.lock().unwrap().clone(); + assert_eq!(seen.len(), 2, "every LLM request carries a traceparent"); + + for (span, op) in llm.iter().zip(["chat", "chat_stream"]) { + assert_eq!( + span.span_context.trace_id(), + turn.span_context.trace_id(), + "{op}: llm.request must share the caller's trace" + ); + assert_eq!( + span.parent_span_id, + turn.span_context.span_id(), + "{op}: llm.request must be a child of the caller span" + ); + assert_eq!(span.span_kind, opentelemetry::trace::SpanKind::Client); + assert_eq!(attr(span, "llm.provider"), Some("llamacpp".into())); + assert_eq!(attr(span, "llm.operation"), Some(op.into())); + assert_eq!(attr(span, "llm.model"), Some("chat".into())); + assert_eq!(attr(span, "http.status_code"), Some(200i64.into())); + + // Downstream parents onto the llm.request span itself. + let expected = format!( + "00-{}-{}-01", + span.span_context.trace_id(), + span.span_context.span_id() + ); + assert!( + seen.contains(&expected), + "{op}: traceparent {expected} not sent (saw {seen:?})" + ); + } +}