Tracing: migrate spans from raw OTel SDK to tracing + tracing-opentelemetry (#135) #148

Merged
cameron merged 8 commits from feature/tracing-migration into master 2026-09-29 21:52:02 +00:00
20 changed files with 716 additions and 773 deletions
Showing only changes of commit 1c508cdabd - Show all commits

No files matched your search

+28 -71
View File
@@ -1,14 +1,13 @@
use anyhow::Result; use anyhow::Result;
use chrono::{NaiveDate, Utc}; use chrono::{NaiveDate, Utc};
use opentelemetry::KeyValue; use opentelemetry::trace::Status;
use opentelemetry::trace::{Span, Status, TraceContextExt, Tracer};
use std::collections::HashMap; use std::collections::HashMap;
use std::sync::{Arc, Mutex}; use std::sync::{Arc, Mutex};
use tokio::time::sleep; use tokio::time::sleep;
use tracing_opentelemetry::OpenTelemetrySpanExt;
use crate::ai::{EMBEDDING_MODEL, OllamaClient, SmsApiClient, SmsMessage, user_display_name}; use crate::ai::{EMBEDDING_MODEL, OllamaClient, SmsApiClient, SmsMessage, user_display_name};
use crate::database::{DailySummaryDao, InsertDailySummary}; use crate::database::{DailySummaryDao, InsertDailySummary};
use crate::otel::global_tracer;
/// Strip boilerplate prefixes and common phrases from summaries before embedding. /// Strip boilerplate prefixes and common phrases from summaries before embedding.
/// This improves embedding diversity by removing structural similarity. /// 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 /// Generate and embed daily conversation summaries for a date range
/// Default: August 2024 ±30 days (July 1 - September 30, 2024) /// 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( pub async fn generate_daily_summaries(
contact: &str, contact: &str,
start_date: Option<NaiveDate>, start_date: Option<NaiveDate>,
@@ -163,30 +163,16 @@ pub async fn generate_daily_summaries(
sms_client: &SmsApiClient, sms_client: &SmsApiClient,
summary_dao: Arc<Mutex<Box<dyn DailySummaryDao>>>, summary_dao: Arc<Mutex<Box<dyn DailySummaryDao>>>,
) -> Result<()> { ) -> Result<()> {
let tracer = global_tracer(); let span = tracing::Span::current();
span.set_attribute("contact", contact.to_string());
// 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", &current_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);
// Default to August 2024 ±30 days // Default to August 2024 ±30 days
let start = start_date.unwrap_or_else(|| NaiveDate::from_ymd_opt(2024, 7, 1).unwrap()); 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()); let end = end_date.unwrap_or_else(|| NaiveDate::from_ymd_opt(2024, 9, 30).unwrap());
parent_cx span.set_attribute("start_date", start.to_string());
.span() span.set_attribute("end_date", end.to_string());
.set_attribute(KeyValue::new("start_date", start.to_string())); span.set_attribute("date_range_days", (end - start).num_days() + 1);
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,
));
log::info!("========================================"); log::info!("========================================");
log::info!("Starting daily summary generation for {}", contact); log::info!("Starting daily summary generation for {}", contact);
@@ -267,14 +253,7 @@ pub async fn generate_daily_summaries(
} }
// Generate summary for this day // Generate summary for this day
match generate_and_store_daily_summary( match generate_and_store_daily_summary(date, contact, messages, ollama, summary_dao.clone())
&parent_cx,
date,
contact,
messages,
ollama,
summary_dao.clone(),
)
.await .await
{ {
Ok(_) => { Ok(_) => {
@@ -322,55 +301,38 @@ pub async fn generate_daily_summaries(
log::info!("========================================"); log::info!("========================================");
// Record final metrics in span // Record final metrics in span
parent_cx span.set_attribute("days_processed", processed as i64);
.span() span.set_attribute("days_skipped", skipped as i64);
.set_attribute(KeyValue::new("days_processed", processed as i64)); span.set_attribute("days_failed", failed as i64);
parent_cx span.set_attribute("total_days", total_days as i64);
.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));
if failed > 0 { 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 { } else {
parent_cx.span().set_status(Status::Ok); span.set_status(Status::Ok);
} }
Ok(()) Ok(())
} }
/// Generate and store a single day's summary /// 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( async fn generate_and_store_daily_summary(
parent_cx: &opentelemetry::Context,
date: &NaiveDate, date: &NaiveDate,
contact: &str, contact: &str,
messages: &[SmsMessage], messages: &[SmsMessage],
ollama: &OllamaClient, ollama: &OllamaClient,
summary_dao: Arc<Mutex<Box<dyn DailySummaryDao>>>, summary_dao: Arc<Mutex<Box<dyn DailySummaryDao>>>,
) -> Result<()> { ) -> Result<()> {
let tracer = global_tracer(); let span = tracing::Span::current();
let mut span = tracer.start_with_context("ai.daily_summary.generate_single", parent_cx); span.set_attribute("date", date.to_string());
span.set_attribute(KeyValue::new("date", date.to_string())); span.set_attribute("contact", contact.to_string());
span.set_attribute(KeyValue::new("contact", contact.to_string())); span.set_attribute("message_count", messages.len() as i64);
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 (prompt, system_prompt) = build_daily_summary_prompt(contact, date, messages); let (prompt, system_prompt) = build_daily_summary_prompt(contact, date, messages);
// Generate summary with LLM // Generate summary with LLM
let summary = ollama let summary = ollama.generate(&prompt, Some(system_prompt)).await?;
.generate(&llm_cx, &prompt, Some(system_prompt))
.await?;
log::debug!( log::debug!(
"Generated summary for {}: {}", "Generated summary for {}: {}",
@@ -378,7 +340,7 @@ async fn generate_and_store_daily_summary(
summary.chars().take(100).collect::<String>() summary.chars().take(100).collect::<String>()
); );
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 // Strip boilerplate before embedding to improve vector diversity
let stripped_summary = strip_summary_boilerplate(&summary); 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) // Embed the stripped summary (store original summary in DB)
let embedding = ollama let embedding = ollama.generate_embedding(&stripped_summary).await?;
.generate_embedding(&llm_cx, &stripped_summary)
.await?;
span.set_attribute(KeyValue::new( span.set_attribute("embedding_dimensions", embedding.len() as i64);
"embedding_dimensions",
embedding.len() as i64,
));
// Store in database // Store in database
let insert = InsertDailySummary { let insert = InsertDailySummary {
@@ -408,8 +365,8 @@ async fn generate_and_store_daily_summary(
model_version: EMBEDDING_MODEL.to_string(), model_version: EMBEDDING_MODEL.to_string(),
}; };
// Create context from current span for DB operation // DAO calls still take an explicit OTel context (#141).
let child_cx = opentelemetry::Context::current_with_span(span); let child_cx = span.context();
let mut dao = summary_dao.lock().expect("Unable to lock DailySummaryDao"); let mut dao = summary_dao.lock().expect("Unable to lock DailySummaryDao");
let result = dao let result = dao
@@ -417,8 +374,8 @@ async fn generate_and_store_daily_summary(
.map_err(|e| anyhow::anyhow!("Failed to store summary: {:?}", e)); .map_err(|e| anyhow::anyhow!("Failed to store summary: {:?}", e));
match &result { match &result {
Ok(_) => child_cx.span().set_status(Status::Ok), Ok(_) => span.set_status(Status::Ok),
Err(e) => child_cx.span().set_status(Status::error(e.to_string())), Err(e) => span.set_status(Status::error(e.to_string())),
} }
result?; result?;
+53 -60
View File
@@ -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 futures::StreamExt;
use opentelemetry::KeyValue; use opentelemetry::trace::Status;
use opentelemetry::trace::{Span, Status, Tracer};
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};
use tracing::Instrument;
use tracing_opentelemetry::OpenTelemetrySpanExt;
use crate::ai::insight_chat::{ChatStreamEvent, ChatTurnRequest}; use crate::ai::insight_chat::{ChatStreamEvent, ChatTurnRequest};
use crate::ai::ollama::ChatMessage; use crate::ai::ollama::ChatMessage;
@@ -11,7 +12,6 @@ use crate::data::Claims;
use crate::database::models::{InsightGenerationType, InsightJobStatus, PhotoInsight}; use crate::database::models::{InsightGenerationType, InsightJobStatus, PhotoInsight};
use crate::database::{ExifDao, InsightDao}; use crate::database::{ExifDao, InsightDao};
use crate::libraries; use crate::libraries;
use crate::otel::{extract_context_from_request, global_tracer};
use crate::state::AppState; use crate::state::AppState;
use crate::utils::normalize_path; 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 - Generate insight for a specific photo (async)
#[post("/insights/generate")] #[post("/insights/generate")]
#[tracing::instrument(name = "http.insights.generate", skip_all)]
pub async fn generate_insight_handler( pub async fn generate_insight_handler(
http_request: HttpRequest,
_claims: Claims, _claims: Claims,
request: web::Json<GeneratePhotoInsightRequest>, request: web::Json<GeneratePhotoInsightRequest>,
app_state: web::Data<AppState>, app_state: web::Data<AppState>,
) -> impl Responder { ) -> impl Responder {
let parent_context = extract_context_from_request(&http_request); let span = tracing::Span::current();
let tracer = global_tracer();
let mut span = tracer.start_with_context("http.insights.generate", &parent_context);
let mut request = request.0; let mut request = request.0;
// Stored owner defaults fill what the client omitted (request → pref → // 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 library = app_state.primary_library();
let gen_type = InsightGenerationType::Standard; 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 { if let Some(ref model) = request.model {
span.set_attribute(KeyValue::new("model", model.clone())); span.set_attribute("model", model.clone());
} }
log::info!( log::info!(
@@ -506,7 +504,8 @@ pub async fn generate_insight_handler(
let job_handles = app_state.insight_job_handles.clone(); let job_handles = app_state.insight_job_handles.clone();
let path = normalized_path.clone(); let path = normalized_path.clone();
let handle = tokio::spawn(async move { let handle = tokio::spawn(
async move {
let timeout_secs: u64 = std::env::var("INSIGHT_GENERATION_TIMEOUT_SECS") let timeout_secs: u64 = std::env::var("INSIGHT_GENERATION_TIMEOUT_SECS")
.ok() .ok()
.and_then(|v| v.parse().ok()) .and_then(|v| v.parse().ok())
@@ -557,7 +556,8 @@ pub async fn generate_insight_handler(
if let Err(e) = dao.complete_job(&ctx, job_id, id) { if let Err(e) = dao.complete_job(&ctx, job_id, id) {
log::error!("Failed to mark job {} as completed: {:?}", job_id, e); 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") } 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); log::error!("Failed to mark job {} as failed: {:?}", job_id, e);
} }
@@ -595,7 +595,11 @@ pub async fn generate_insight_handler(
.lock() .lock()
.expect("Unable to lock InsightJobHandles"); .expect("Unable to lock InsightJobHandles");
handles.remove(&job_id); 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 // Store abort handle
{ {
@@ -606,7 +610,7 @@ pub async fn generate_insight_handler(
handles.insert(job_id, handle.abort_handle()); 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); span.set_status(Status::Ok);
HttpResponse::Accepted().json(JobIdResponse { job_id }) 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 - Generate insight using agentic tool-calling loop (async)
#[post("/insights/generate/agentic")] #[post("/insights/generate/agentic")]
#[tracing::instrument(name = "http.insights.generate_agentic", skip_all)]
pub async fn generate_agentic_insight_handler( pub async fn generate_agentic_insight_handler(
http_request: HttpRequest,
claims: Claims, claims: Claims,
request: web::Json<GeneratePhotoInsightRequest>, request: web::Json<GeneratePhotoInsightRequest>,
app_state: web::Data<AppState>, app_state: web::Data<AppState>,
) -> impl Responder { ) -> impl Responder {
let parent_context = extract_context_from_request(&http_request); let span = tracing::Span::current();
let tracer = global_tracer();
let mut span = tracer.start_with_context("http.insights.generate_agentic", &parent_context);
let mut request = request.0; let mut request = request.0;
// Stored owner defaults fill what the client omitted (request → pref → // 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 library = app_state.primary_library();
let gen_type = InsightGenerationType::Agentic; 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 { 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 { if let Some(ref backend) = request.backend {
span.set_attribute(KeyValue::new("backend", backend.clone())); span.set_attribute("backend", backend.clone());
} }
log::info!( log::info!(
@@ -901,7 +903,8 @@ pub async fn generate_agentic_insight_handler(
let job_handles = app_state.insight_job_handles.clone(); let job_handles = app_state.insight_job_handles.clone();
let path = normalized_path.clone(); let path = normalized_path.clone();
let handle = tokio::spawn(async move { let handle = tokio::spawn(
async move {
let timeout_secs: u64 = std::env::var("INSIGHT_GENERATION_TIMEOUT_SECS") let timeout_secs: u64 = std::env::var("INSIGHT_GENERATION_TIMEOUT_SECS")
.ok() .ok()
.and_then(|v| v.parse().ok()) .and_then(|v| v.parse().ok())
@@ -947,7 +950,8 @@ pub async fn generate_agentic_insight_handler(
} }
} }
Ok(Ok(Ok((None, _, _)))) => { Ok(Ok(Ok((None, _, _)))) => {
if let Err(e) = dao.fail_job(&ctx, job_id, "agentic generation returned no insight") 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); log::error!("Failed to mark job {} as failed: {:?}", job_id, e);
} }
@@ -985,7 +989,11 @@ pub async fn generate_agentic_insight_handler(
.lock() .lock()
.expect("Unable to lock InsightJobHandles"); .expect("Unable to lock InsightJobHandles");
handles.remove(&job_id); 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 // Store abort handle
{ {
@@ -996,7 +1004,7 @@ pub async fn generate_agentic_insight_handler(
handles.insert(job_id, handle.abort_handle()); 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); span.set_status(Status::Ok);
HttpResponse::Accepted().json(JobIdResponse { job_id }) 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 — submit a follow-up turn against an existing insight.
#[post("/insights/chat")] #[post("/insights/chat")]
#[tracing::instrument(name = "http.insights.chat", skip_all)]
pub async fn chat_turn_handler( pub async fn chat_turn_handler(
http_request: HttpRequest,
claims: Claims, claims: Claims,
request: web::Json<ChatTurnHttpRequest>, request: web::Json<ChatTurnHttpRequest>,
app_state: web::Data<AppState>, app_state: web::Data<AppState>,
) -> impl Responder { ) -> impl Responder {
let parent_context = extract_context_from_request(&http_request); let span = tracing::Span::current();
let tracer = global_tracer(); span.set_attribute("file_path", request.file_path.clone());
let mut span = tracer.start_with_context("http.insights.chat", &parent_context);
span.set_attribute(KeyValue::new("file_path", request.file_path.clone()));
let mut request = request.0; let mut request = request.0;
// Stored owner defaults fill what the client omitted (request → pref → // 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. /// Returns `text/event-stream` with one event per chat stream event.
#[post("/insights/chat/stream")] #[post("/insights/chat/stream")]
pub async fn chat_stream_handler( pub async fn chat_stream_handler(
http_request: HttpRequest,
claims: Claims, claims: Claims,
request: web::Json<ChatTurnHttpRequest>, request: web::Json<ChatTurnHttpRequest>,
app_state: web::Data<AppState>, app_state: web::Data<AppState>,
) -> HttpResponse { ) -> HttpResponse {
let parent_context = extract_context_from_request(&http_request);
let mut request = request.0; let mut request = request.0;
// Stored owner defaults fill what the client omitted (request → pref → // Stored owner defaults fill what the client omitted (request → pref →
// hard default; the hard default still applies downstream). // 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 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. // Map ChatStreamEvent → SSE frame bytes.
let sse_stream = futures::stream::StreamExt::map(events, |ev| { 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, /// POST /insights/chat/turn — async turn dispatch. Returns turn_id immediately,
/// client then polls GET /insights/chat/turn/{turn_id} for SSE replay. /// client then polls GET /insights/chat/turn/{turn_id} for SSE replay.
#[post("/insights/chat/turn")] #[post("/insights/chat/turn")]
#[tracing::instrument(name = "http.insights.chat_turn_async", skip_all)]
pub async fn turn_async_handler( pub async fn turn_async_handler(
http_request: HttpRequest,
claims: Claims, claims: Claims,
request: web::Json<ChatTurnHttpRequest>, request: web::Json<ChatTurnHttpRequest>,
app_state: web::Data<AppState>, app_state: web::Data<AppState>,
) -> impl Responder { ) -> impl Responder {
let parent_context = extract_context_from_request(&http_request); let span = tracing::Span::current();
let tracer = global_tracer(); span.set_attribute("file_path", request.file_path.clone());
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 mut request = request.0; let mut request = request.0;
// Stored owner defaults fill what the client omitted (request → pref → // 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 registry = app_state.turn_registry.clone();
let turn_id = service.chat_turn_async(registry, chat_req).await; 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); span.set_status(Status::Ok);
HttpResponse::Accepted().json(serde_json::json!({ HttpResponse::Accepted().json(serde_json::json!({
"turn_id": turn_id, "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} — SSE replay stream.
#[get("/insights/chat/turn/{turn_id}")] #[get("/insights/chat/turn/{turn_id}")]
#[tracing::instrument(name = "ai.chat.turn.replay", skip_all)]
pub async fn turn_replay_handler( pub async fn turn_replay_handler(
http_request: HttpRequest,
path: web::Path<String>, path: web::Path<String>,
query: web::Query<ReplayQuery>, query: web::Query<ReplayQuery>,
app_state: web::Data<AppState>, app_state: web::Data<AppState>,
) -> HttpResponse { ) -> 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 /// 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 /// on `turn_id` only, and `render_turn_info_frame` already scopes the
/// identity (persona turns carry `persona_id` and no `file_path`). /// identity (persona turns carry `persona_id` and no `file_path`).
pub(crate) async fn turn_replay_impl( pub(crate) async fn turn_replay_impl(
http_request: HttpRequest,
path: web::Path<String>, path: web::Path<String>,
query: web::Query<ReplayQuery>, query: web::Query<ReplayQuery>,
app_state: web::Data<AppState>, app_state: web::Data<AppState>,
@@ -1920,11 +1921,9 @@ pub(crate) async fn turn_replay_impl(
let turn_id = path.into_inner(); let turn_id = path.into_inner();
let skip_before = query.skip_before.unwrap_or(0); let skip_before = query.skip_before.unwrap_or(0);
let parent_context = extract_context_from_request(&http_request); let span = tracing::Span::current();
let tracer = global_tracer(); span.set_attribute("turn_id", turn_id.clone());
let mut span = tracer.start_with_context("ai.chat.turn.replay", &parent_context); span.set_attribute("skip_before", skip_before as i64);
span.set_attribute(KeyValue::new("turn_id", turn_id.clone()));
span.set_attribute(KeyValue::new("skip_before", skip_before as i64));
let registry = app_state.turn_registry.clone(); let registry = app_state.turn_registry.clone();
let entry = match registry.get(&turn_id).await { let entry = match registry.get(&turn_id).await {
@@ -1938,11 +1937,8 @@ pub(crate) async fn turn_replay_impl(
}; };
let info = entry.info().await; let info = entry.info().await;
span.set_attribute(KeyValue::new("status", info.status.as_str())); span.set_attribute("status", info.status.as_str());
span.set_attribute(KeyValue::new( span.set_attribute("event_count", info.total_events_pushed as i64);
"event_count",
info.total_events_pushed as i64,
));
let turn_info_frame = render_turn_info_frame(&info); let turn_info_frame = render_turn_info_frame(&info);
// Initial buffered batch: events produced before this connection attached. // 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} — cancel a running turn.
#[delete("/insights/chat/turn/{turn_id}")] #[delete("/insights/chat/turn/{turn_id}")]
#[tracing::instrument(name = "ai.chat.turn.cancel", skip_all)]
pub async fn cancel_turn_handler( pub async fn cancel_turn_handler(
http_request: HttpRequest,
path: web::Path<String>, path: web::Path<String>,
app_state: web::Data<AppState>, app_state: web::Data<AppState>,
) -> impl Responder { ) -> 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 /// Core of the turn-cancel, attribute-free so the persona-chat DELETE route
/// can reuse it under `/persona_chat/turn/{turn_id}`. /// can reuse it under `/persona_chat/turn/{turn_id}`.
pub(crate) async fn cancel_turn_impl( pub(crate) async fn cancel_turn_impl(
http_request: HttpRequest,
path: web::Path<String>, path: web::Path<String>,
app_state: web::Data<AppState>, app_state: web::Data<AppState>,
) -> impl Responder { ) -> impl Responder {
let turn_id = path.into_inner(); let turn_id = path.into_inner();
let parent_context = extract_context_from_request(&http_request); let span = tracing::Span::current();
let tracer = global_tracer(); span.set_attribute("turn_id", turn_id.clone());
let mut span = tracer.start_with_context("ai.chat.turn.cancel", &parent_context);
span.set_attribute(KeyValue::new("turn_id", turn_id.clone()));
let registry = app_state.turn_registry.clone(); let registry = app_state.turn_registry.clone();
let entry = match registry.get(&turn_id).await { 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 // also checks `is_running()` at each iteration boundary as a graceful
// backstop in case the abort lands between await points. // backstop in case the abort lands between await points.
let aborted = entry.abort(); 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 // 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 // terminal status with no buffered tail as "closed", so the Done must be
+81 -73
View File
@@ -1,7 +1,6 @@
use anyhow::{Result, anyhow, bail}; use anyhow::{Result, anyhow, bail};
use chrono::{TimeZone, Utc}; use chrono::{TimeZone, Utc};
use opentelemetry::KeyValue; use opentelemetry::trace::Status;
use opentelemetry::trace::{Span, Status, TraceContextExt, Tracer};
use std::collections::HashMap; use std::collections::HashMap;
use std::sync::{Arc, Mutex}; use std::sync::{Arc, Mutex};
use tokio::sync::Mutex as TokioMutex; use tokio::sync::Mutex as TokioMutex;
@@ -15,9 +14,10 @@ use crate::ai::turn_registry::TurnEntry;
use crate::ai::turn_registry::TurnRegistry; use crate::ai::turn_registry::TurnRegistry;
use crate::database::InsightDao; use crate::database::InsightDao;
use crate::database::models::InsertPhotoInsight; use crate::database::models::InsertPhotoInsight;
use crate::otel::global_tracer;
use crate::utils::{normalize_path, retry_with_backoff}; use crate::utils::{normalize_path, retry_with_backoff};
use futures::stream::{BoxStream, StreamExt}; use futures::stream::{BoxStream, StreamExt};
use tracing::Instrument;
use tracing_opentelemetry::OpenTelemetrySpanExt;
use uuid::Uuid; use uuid::Uuid;
pub const DEFAULT_MAX_ITERATIONS: usize = 6; 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<ChatTurnResult> { pub async fn chat_turn(&self, req: ChatTurnRequest) -> Result<ChatTurnResult> {
let tracer = global_tracer(); let span = tracing::Span::current();
let parent_cx = opentelemetry::Context::new(); span.set_attribute("file_path", req.file_path.clone());
let mut span = tracer.start_with_context("ai.insight.chat_turn", &parent_cx); span.set_attribute("library_id", req.library_id as i64);
span.set_attribute(KeyValue::new("file_path", req.file_path.clone())); span.set_attribute("amend", req.amend);
span.set_attribute(KeyValue::new("library_id", req.library_id as i64));
span.set_attribute(KeyValue::new("amend", req.amend));
if req.user_message.trim().is_empty() { if req.user_message.trim().is_empty() {
bail!("user_message must not be empty"); bail!("user_message must not be empty");
@@ -220,7 +219,7 @@ impl InsightChatService {
.clone() .clone()
.filter(|s| !s.trim().is_empty()) .filter(|s| !s.trim().is_empty())
.unwrap_or_else(|| "default".to_string()); .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); let normalized = normalize_path(&req.file_path);
@@ -278,13 +277,13 @@ impl InsightChatService {
.unwrap_or_else(|| stored_backend.clone()); .unwrap_or_else(|| stored_backend.clone());
validate_cross_replay(&stored_backend, &effective_backend)?; validate_cross_replay(&stored_backend, &effective_backend)?;
let kind = BackendKind::parse(&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 let max_iterations = req
.max_iterations .max_iterations
.unwrap_or(DEFAULT_MAX_ITERATIONS) .unwrap_or(DEFAULT_MAX_ITERATIONS)
.clamp(1, env_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 stored_model = insight.model_version.clone();
let overrides = SamplingOverrides { let overrides = SamplingOverrides {
@@ -302,7 +301,7 @@ impl InsightChatService {
}; };
let backend = self.generator.resolve_backend(kind, &overrides).await?; let backend = self.generator.resolve_backend(kind, &overrides).await?;
let model_used = backend.model().to_string(); 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 // 5. Decide vision + tool set. In hybrid (describe-then-inline) mode
// we omit `describe_photo`. Otherwise trust the stored history: // 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 budget_bytes = budget_tokens.saturating_mul(BYTES_PER_TOKEN);
let truncated = apply_context_budget(&mut messages, budget_bytes); let truncated = apply_context_budget(&mut messages, budget_bytes);
if truncated { 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 // Everything in `messages` at this point is replayed history that is
// ALREADY in the tree; only what the turn appends past this mark // ALREADY in the tree; only what the turn appends past this mark
@@ -359,12 +358,12 @@ impl InsightChatService {
// accumulate across turns. // accumulate across turns.
let original_system_content = annotate_system_with_budget(&mut messages, max_iterations); 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 // 8. Agentic loop — same shape as insight_generator's, but capped
// tighter and dispatching tools through the shared executor. // tighter and dispatching tools through the shared executor.
let loop_span = tracer.start_with_context("ai.chat.loop", &insight_cx); // Each LLM / tool await runs inside `loop_span` so their spans
let loop_cx = insight_cx.with_span(loop_span); // 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 tool_calls_made = 0usize;
let mut iterations_used = 0usize; let mut iterations_used = 0usize;
let mut last_prompt_eval_count: Option<i32> = None; let mut last_prompt_eval_count: Option<i32> = None;
@@ -377,7 +376,8 @@ impl InsightChatService {
let (response, prompt_tokens, eval_tokens) = backend let (response, prompt_tokens, eval_tokens) = backend
.chat() .chat()
.chat_with_tools(&loop_cx, messages.clone(), tools.clone()) .chat_with_tools(messages.clone(), tools.clone())
.instrument(loop_span.clone())
.await?; .await?;
last_prompt_eval_count = prompt_tokens; last_prompt_eval_count = prompt_tokens;
last_eval_count = eval_tokens; last_eval_count = eval_tokens;
@@ -416,6 +416,7 @@ impl InsightChatService {
&active_persona, &active_persona,
&loop_cx, &loop_cx,
) )
.instrument(loop_span.clone())
.await; .await;
messages.push(ChatMessage::tool_result(result)); messages.push(ChatMessage::tool_result(result));
} }
@@ -434,7 +435,8 @@ impl InsightChatService {
let synthetic_idx = push_synthetic_final_prompt(&mut messages); let synthetic_idx = push_synthetic_final_prompt(&mut messages);
let (final_response, prompt_tokens, eval_tokens) = backend let (final_response, prompt_tokens, eval_tokens) = backend
.chat() .chat()
.chat_with_tools(&loop_cx, messages.clone(), vec![]) .chat_with_tools(messages.clone(), vec![])
.instrument(loop_span.clone())
.await?; .await?;
last_prompt_eval_count = prompt_tokens; last_prompt_eval_count = prompt_tokens;
last_eval_count = eval_tokens; last_eval_count = eval_tokens;
@@ -449,7 +451,8 @@ impl InsightChatService {
// return / persist as the reply (the raw transcript keeps it). // return / persist as the reply (the raw transcript keeps it).
let final_content = crate::ai::llm_client::strip_think_blocks(&final_content); 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 // Drop the per-turn iteration-budget note from the system message
// before we persist so it doesn't snowball on each subsequent turn. // before we persist so it doesn't snowball on each subsequent turn.
@@ -492,7 +495,6 @@ impl InsightChatService {
let title_raw = backend let title_raw = backend
.chat() .chat()
.generate( .generate(
&loop_cx,
&title_prompt, &title_prompt,
Some( Some(
"You are my long term memory assistant. Use only the information provided. Do not invent details.", "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<Self>` (passed by /// The stream takes ownership of the service via `Arc<Self>` (passed by
/// the caller) so it can live past the handler's await boundary. /// 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( pub fn chat_turn_stream(
self: Arc<Self>, self: Arc<Self>,
req: ChatTurnRequest, req: ChatTurnRequest,
cx: opentelemetry::Context,
) -> BoxStream<'static, ChatStreamEvent> { ) -> 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 svc = self;
let s = async_stream::stream! { 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) => { Ok(mut rx) => {
while let Some(ev) = rx.recv().await { while let Some(ev) = rx.recv().await {
yield ev; yield ev;
@@ -808,7 +818,7 @@ impl InsightChatService {
async fn chat_turn_stream_inner<F>( async fn chat_turn_stream_inner<F>(
self: Arc<Self>, self: Arc<Self>,
req: ChatTurnRequest, req: ChatTurnRequest,
cx: opentelemetry::Context, span: tracing::Span,
_ev_mapper: F, _ev_mapper: F,
) -> Result<tokio::sync::mpsc::Receiver<ChatStreamEvent>> ) -> Result<tokio::sync::mpsc::Receiver<ChatStreamEvent>>
where where
@@ -816,12 +826,22 @@ impl InsightChatService {
{ {
let (tx, rx) = tokio::sync::mpsc::channel::<ChatStreamEvent>(64); let (tx, rx) = tokio::sync::mpsc::channel::<ChatStreamEvent>(64);
let svc = self.clone(); let svc = self.clone();
tokio::spawn(async move { // Carry the turn span into the task explicitly — tracing context is
let result = svc.run_streaming_turn(req, cx, tx.clone()).await; // 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 { if let Err(e) = result {
span.set_status(Status::error(format!("{e}")));
let _ = tx.send(ChatStreamEvent::Error(format!("{}", e))).await; let _ = tx.send(ChatStreamEvent::Error(format!("{}", e))).await;
} else {
span.set_status(Status::Ok);
} }
}); }
.instrument(span),
);
Ok(rx) Ok(rx)
} }
@@ -843,22 +863,23 @@ impl InsightChatService {
let svc = self.clone(); let svc = self.clone();
let entry_clone = entry.clone(); let entry_clone = entry.clone();
let turn_id_for_span = turn_id.clone(); // Span covering the whole spawned turn execution. Created before the
let library_id = req.library_id; // spawn so it parents on the dispatching request (same trace as the
let handle = tokio::spawn(async move { // 202), and carried into the task with `.instrument`. It outlives the
// Span covering the whole spawned turn execution. Created here (not // dispatch span, which ends at the 202 response.
// in the HTTP handler) because the dispatch span ends at the 202 let turn_span = tracing::info_span!(
// response, long before this work runs. "ai.chat.turn.execute",
let tracer = global_tracer(); turn_id = %turn_id,
let mut span = tracer.start("ai.chat.turn.execute"); library_id = req.library_id,
span.set_attribute(KeyValue::new("turn_id", turn_id_for_span)); );
span.set_attribute(KeyValue::new("library_id", library_id as i64)); let handle = tokio::spawn(
async move {
let span = tracing::Span::current();
let result = svc let result = svc
.run_streaming_turn_with_entry(req, entry_clone.clone(), &mut span) .run_streaming_turn_with_entry(req, entry_clone.clone(), &span)
.await; .await;
if let Err(ref e) = result { if let Err(ref e) = result {
span.set_attribute(KeyValue::new("status", "error")); span.set_attribute("status", "error");
span.set_status(Status::error(format!("{e}"))); span.set_status(Status::error(format!("{e}")));
// Push the terminal event BEFORE flipping status: a replay // Push the terminal event BEFORE flipping status: a replay
// reader treats a terminal status with no buffered tail as // reader treats a terminal status with no buffered tail as
@@ -868,10 +889,12 @@ impl InsightChatService {
.await; .await;
entry_clone.set_terminal_status(crate::ai::turn_registry::TurnStatus::Error); entry_clone.set_terminal_status(crate::ai::turn_registry::TurnStatus::Error);
} else { } else {
span.set_attribute(KeyValue::new("status", "done")); span.set_attribute("status", "done");
span.set_status(Status::Ok); span.set_status(Status::Ok);
} }
}); }
.instrument(turn_span),
);
// Install the abort handle so DELETE can actually stop the task. // Install the abort handle so DELETE can actually stop the task.
entry.set_abort_handle(handle.abort_handle()); entry.set_abort_handle(handle.abort_handle());
@@ -885,7 +908,7 @@ impl InsightChatService {
self: Arc<Self>, self: Arc<Self>,
req: ChatTurnRequest, req: ChatTurnRequest,
entry: Arc<TurnEntry>, entry: Arc<TurnEntry>,
turn_span: &mut opentelemetry::global::BoxedSpan, turn_span: &tracing::Span,
) -> Result<()> { ) -> Result<()> {
if req.user_message.trim().is_empty() { if req.user_message.trim().is_empty() {
bail!("user_message must not be empty"); bail!("user_message must not be empty");
@@ -930,7 +953,7 @@ impl InsightChatService {
normalized: String, normalized: String,
insight: crate::database::models::PhotoInsight, insight: crate::database::models::PhotoInsight,
entry: Arc<TurnEntry>, entry: Arc<TurnEntry>,
turn_span: &mut opentelemetry::global::BoxedSpan, turn_span: &tracing::Span,
) -> Result<()> { ) -> Result<()> {
let active_persona = req let active_persona = req
.persona_id .persona_id
@@ -988,10 +1011,8 @@ impl InsightChatService {
let model_used = backend.model().to_string(); let model_used = backend.model().to_string();
// llm.* attributes on the turn-level span (issue #130): clients can't // llm.* attributes on the turn-level span (issue #130): clients can't
// know local-vs-hybrid, but the turn knows its resolved backend. // 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("llm.backend", kind.as_str());
turn_span.set_attribute(KeyValue::new("llm.model", model_used.clone())); turn_span.set_attribute("llm.model", model_used.clone());
let turn_cx = opentelemetry::Context::new()
.with_remote_span_context(turn_span.span_context().clone());
let local_first_user_has_image = messages let local_first_user_has_image = messages
.iter() .iter()
@@ -1041,7 +1062,6 @@ impl InsightChatService {
&active_persona, &active_persona,
max_iterations, max_iterations,
&entry, &entry,
&turn_cx,
) )
.await?; .await?;
let AgenticLoopOutcome { let AgenticLoopOutcome {
@@ -1158,7 +1178,7 @@ impl InsightChatService {
req: ChatTurnRequest, req: ChatTurnRequest,
normalized: String, normalized: String,
entry: Arc<TurnEntry>, entry: Arc<TurnEntry>,
turn_span: &mut opentelemetry::global::BoxedSpan, turn_span: &tracing::Span,
) -> Result<()> { ) -> Result<()> {
let active_persona = req let active_persona = req
.persona_id .persona_id
@@ -1185,10 +1205,8 @@ impl InsightChatService {
let backend = self.generator.resolve_backend(kind, &overrides).await?; let backend = self.generator.resolve_backend(kind, &overrides).await?;
let model_used = backend.model().to_string(); let model_used = backend.model().to_string();
// llm.* attributes on the turn-level span (issue #130). // llm.* attributes on the turn-level span (issue #130).
turn_span.set_attribute(KeyValue::new("llm.backend", kind.as_str())); turn_span.set_attribute("llm.backend", kind.as_str());
turn_span.set_attribute(KeyValue::new("llm.model", model_used.clone())); turn_span.set_attribute("llm.model", model_used.clone());
let turn_cx = opentelemetry::Context::new()
.with_remote_span_context(turn_span.span_context().clone());
let image_base64: Option<String> = self.generator.load_image_as_base64(&normalized).ok(); let image_base64: Option<String> = self.generator.load_image_as_base64(&normalized).ok();
@@ -1208,7 +1226,7 @@ impl InsightChatService {
let visual_block = if !backend.images_inline { let visual_block = if !backend.images_inline {
match image_base64.as_deref() { 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) => { Ok(desc) => {
format!("Visual description (from local vision model):\n{}\n", desc) format!("Visual description (from local vision model):\n{}\n", desc)
} }
@@ -1263,7 +1281,6 @@ impl InsightChatService {
&active_persona, &active_persona,
max_iterations, max_iterations,
&entry, &entry,
&turn_cx,
) )
.await?; .await?;
let AgenticLoopOutcome { let AgenticLoopOutcome {
@@ -1352,7 +1369,6 @@ impl InsightChatService {
active_persona: &str, active_persona: &str,
max_iterations: usize, max_iterations: usize,
entry: &Arc<TurnEntry>, entry: &Arc<TurnEntry>,
cx: &opentelemetry::Context,
) -> Result<AgenticLoopOutcome> { ) -> Result<AgenticLoopOutcome> {
crate::ai::insight_chat::run_streaming_agentic_loop_with_entry( crate::ai::insight_chat::run_streaming_agentic_loop_with_entry(
&self.generator, &self.generator,
@@ -1365,7 +1381,6 @@ impl InsightChatService {
active_persona, active_persona,
max_iterations, max_iterations,
entry, entry,
cx,
) )
.await .await
} }
@@ -1373,7 +1388,6 @@ impl InsightChatService {
async fn run_streaming_turn( async fn run_streaming_turn(
self: Arc<Self>, self: Arc<Self>,
req: ChatTurnRequest, req: ChatTurnRequest,
cx: opentelemetry::Context,
tx: tokio::sync::mpsc::Sender<ChatStreamEvent>, tx: tokio::sync::mpsc::Sender<ChatStreamEvent>,
) -> Result<()> { ) -> Result<()> {
if req.user_message.trim().is_empty() { if req.user_message.trim().is_empty() {
@@ -1410,10 +1424,10 @@ impl InsightChatService {
}; };
if req.regenerate || existing_insight.is_none() { 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"); 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 .await
} }
@@ -1425,7 +1439,6 @@ impl InsightChatService {
&self, &self,
req: ChatTurnRequest, req: ChatTurnRequest,
normalized: String, normalized: String,
cx: opentelemetry::Context,
insight: crate::database::models::PhotoInsight, insight: crate::database::models::PhotoInsight,
tx: tokio::sync::mpsc::Sender<ChatStreamEvent>, tx: tokio::sync::mpsc::Sender<ChatStreamEvent>,
) -> Result<()> { ) -> Result<()> {
@@ -1523,7 +1536,6 @@ impl InsightChatService {
&active_persona, &active_persona,
max_iterations, max_iterations,
&tx, &tx,
&cx,
) )
.await?; .await?;
let AgenticLoopOutcome { let AgenticLoopOutcome {
@@ -1635,7 +1647,6 @@ impl InsightChatService {
&self, &self,
req: ChatTurnRequest, req: ChatTurnRequest,
normalized: String, normalized: String,
cx: opentelemetry::Context,
tx: tokio::sync::mpsc::Sender<ChatStreamEvent>, tx: tokio::sync::mpsc::Sender<ChatStreamEvent>,
) -> Result<()> { ) -> Result<()> {
let active_persona = req let active_persona = req
@@ -1692,7 +1703,7 @@ impl InsightChatService {
// images_inline backends send images directly to the chat model. // images_inline backends send images directly to the chat model.
let visual_block = if !backend.images_inline { let visual_block = if !backend.images_inline {
match image_base64.as_deref() { 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) => { Ok(desc) => {
format!("Visual description (from local vision model):\n{}\n", desc) format!("Visual description (from local vision model):\n{}\n", desc)
} }
@@ -1759,7 +1770,6 @@ impl InsightChatService {
&active_persona, &active_persona,
max_iterations, max_iterations,
&tx, &tx,
&cx,
) )
.await?; .await?;
let AgenticLoopOutcome { let AgenticLoopOutcome {
@@ -1842,7 +1852,6 @@ impl InsightChatService {
active_persona: &str, active_persona: &str,
max_iterations: usize, max_iterations: usize,
tx: &tokio::sync::mpsc::Sender<ChatStreamEvent>, tx: &tokio::sync::mpsc::Sender<ChatStreamEvent>,
cx: &opentelemetry::Context,
) -> Result<AgenticLoopOutcome> { ) -> Result<AgenticLoopOutcome> {
let mut tool_calls_made = 0usize; let mut tool_calls_made = 0usize;
let mut iterations_used = 0usize; let mut iterations_used = 0usize;
@@ -1861,7 +1870,7 @@ impl InsightChatService {
let mut stream = backend let mut stream = backend
.chat() .chat()
.chat_with_tools_stream(cx, messages.clone(), tools.clone()) .chat_with_tools_stream(messages.clone(), tools.clone())
.await?; .await?;
let mut final_message: Option<ChatMessage> = None; let mut final_message: Option<ChatMessage> = None;
@@ -1955,7 +1964,7 @@ impl InsightChatService {
let synthetic_idx = push_synthetic_final_prompt(messages); let synthetic_idx = push_synthetic_final_prompt(messages);
let mut stream = backend let mut stream = backend
.chat() .chat()
.chat_with_tools_stream(cx, messages.clone(), vec![]) .chat_with_tools_stream(messages.clone(), vec![])
.await?; .await?;
let mut final_message: Option<ChatMessage> = None; let mut final_message: Option<ChatMessage> = None;
while let Some(ev) = stream.next().await { while let Some(ev) = stream.next().await {
@@ -2335,7 +2344,6 @@ pub async fn run_streaming_agentic_loop_with_entry(
active_persona: &str, active_persona: &str,
max_iterations: usize, max_iterations: usize,
entry: &Arc<TurnEntry>, entry: &Arc<TurnEntry>,
cx: &opentelemetry::Context,
) -> Result<AgenticLoopOutcome> { ) -> Result<AgenticLoopOutcome> {
let mut tool_calls_made = 0usize; let mut tool_calls_made = 0usize;
let mut iterations_used = 0usize; let mut iterations_used = 0usize;
@@ -2365,7 +2373,7 @@ pub async fn run_streaming_agentic_loop_with_entry(
let mut stream = backend let mut stream = backend
.chat() .chat()
.chat_with_tools_stream(cx, messages.clone(), tools.clone()) .chat_with_tools_stream(messages.clone(), tools.clone())
.await?; .await?;
let mut final_message: Option<ChatMessage> = None; let mut final_message: Option<ChatMessage> = None;
@@ -2448,7 +2456,7 @@ pub async fn run_streaming_agentic_loop_with_entry(
let synthetic_idx = push_synthetic_final_prompt(messages); let synthetic_idx = push_synthetic_final_prompt(messages);
let mut stream = backend let mut stream = backend
.chat() .chat()
.chat_with_tools_stream(cx, messages.clone(), vec![]) .chat_with_tools_stream(messages.clone(), vec![])
.await?; .await?;
let mut final_message: Option<ChatMessage> = None; let mut final_message: Option<ChatMessage> = None;
while let Some(ev) = stream.next().await { while let Some(ev) = stream.next().await {
+114 -187
View File
@@ -2,13 +2,14 @@ use anyhow::Result;
use base64::Engine as _; use base64::Engine as _;
use chrono::{DateTime, Local, NaiveDate, TimeZone, Utc}; use chrono::{DateTime, Local, NaiveDate, TimeZone, Utc};
use image::ImageFormat; use image::ImageFormat;
use opentelemetry::KeyValue; use opentelemetry::trace::Status;
use opentelemetry::trace::{Span, Status, TraceContextExt, Tracer};
use serde::Deserialize; use serde::Deserialize;
use std::fs::File; use std::fs::File;
use std::io::Cursor; use std::io::Cursor;
use std::str::FromStr; use std::str::FromStr;
use std::sync::{Arc, Mutex}; use std::sync::{Arc, Mutex};
use tracing::Instrument;
use tracing_opentelemetry::OpenTelemetrySpanExt;
use crate::ai::apollo_client::{ApolloClient, ApolloPlace}; use crate::ai::apollo_client::{ApolloClient, ApolloPlace};
use crate::ai::backend::{BackendKind, ResolvedBackend, SamplingOverrides}; use crate::ai::backend::{BackendKind, ResolvedBackend, SamplingOverrides};
@@ -25,7 +26,6 @@ use crate::database::{
}; };
use crate::libraries::Library; use crate::libraries::Library;
use crate::memories::extract_date_from_filename; use crate::memories::extract_date_from_filename;
use crate::otel::global_tracer;
use crate::tags::TagDao; use crate::tags::TagDao;
use crate::utils::{earliest_fs_time, normalize_path, retry_with_backoff, retry_with_backoff_if}; 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) /// Find relevant messages using RAG, excluding recent messages (>30 days ago)
/// This prevents RAG from returning messages already in the immediate time window /// 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( async fn find_relevant_messages_rag_historical(
&self, &self,
parent_cx: &opentelemetry::Context,
date: chrono::NaiveDate, date: chrono::NaiveDate,
location: Option<&str>, location: Option<&str>,
contact: Option<&str>, contact: Option<&str>,
@@ -512,33 +512,20 @@ impl InsightGenerator {
limit: usize, limit: usize,
extra_context: Option<&str>, extra_context: Option<&str>,
) -> Result<Vec<String>> { ) -> Result<Vec<String>> {
let tracer = global_tracer(); let span = tracing::Span::current();
let span = tracer.start_with_context("ai.rag.filter_historical", parent_cx);
let filter_cx = parent_cx.with_span(span);
filter_cx span.set_attribute("date", date.to_string());
.span() span.set_attribute("limit", limit as i64);
.set_attribute(KeyValue::new("date", date.to_string())); span.set_attribute("exclusion_window_days", 30);
filter_cx
.span()
.set_attribute(KeyValue::new("limit", limit as i64));
filter_cx
.span()
.set_attribute(KeyValue::new("exclusion_window_days", 30));
if let Some(t) = topics { if let Some(t) = topics {
filter_cx span.set_attribute("topics", t.join(", "));
.span()
.set_attribute(KeyValue::new("topics", t.join(", ")));
} }
let query_results = self let query_results = self
.find_relevant_messages_rag(date, location, contact, topics, limit * 2, extra_context) .find_relevant_messages_rag(date, location, contact, topics, limit * 2, extra_context)
.await?; .await?;
filter_cx.span().set_attribute(KeyValue::new( span.set_attribute("rag_results_count", query_results.len() as i64);
"rag_results_count",
query_results.len() as i64,
));
// Filter out messages from within 30 days of the photo date // Filter out messages from within 30 days of the photo date
let photo_timestamp = date let photo_timestamp = date
@@ -576,17 +563,15 @@ impl InsightGenerator {
historical_only.len() historical_only.len()
); );
filter_cx.span().set_attribute(KeyValue::new( span.set_attribute("historical_results_count", historical_only.len() as i64);
"historical_results_count", span.set_status(Status::Ok);
historical_only.len() as i64,
));
filter_cx.span().set_status(Status::Ok);
Ok(historical_only) Ok(historical_only)
} }
/// Find relevant daily summaries using RAG (semantic search) /// Find relevant daily summaries using RAG (semantic search)
/// Returns formatted daily summary strings for LLM context /// Returns formatted daily summary strings for LLM context
#[tracing::instrument(name = "ai.rag.search_daily_summaries", skip_all)]
async fn find_relevant_messages_rag( async fn find_relevant_messages_rag(
&self, &self,
date: chrono::NaiveDate, date: chrono::NaiveDate,
@@ -596,16 +581,14 @@ impl InsightGenerator {
limit: usize, limit: usize,
extra_context: Option<&str>, extra_context: Option<&str>,
) -> Result<Vec<String>> { ) -> Result<Vec<String>> {
let tracer = global_tracer(); let span = tracing::Span::current();
let current_cx = opentelemetry::Context::current(); span.set_attribute("date", date.to_string());
let mut span = tracer.start_with_context("ai.rag.search_daily_summaries", &current_cx); span.set_attribute("limit", limit as i64);
span.set_attribute(KeyValue::new("date", date.to_string()));
span.set_attribute(KeyValue::new("limit", limit as i64));
if let Some(loc) = location { 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 { 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) // Build query string - prioritize topics if available (semantically meaningful)
@@ -633,10 +616,10 @@ impl InsightGenerator {
base_query base_query
}; };
span.set_attribute(KeyValue::new("query", query.clone())); span.set_attribute("query", query.clone());
// Create context with this span for child operations // 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!("========================================");
log::info!("RAG QUERY: {}", query); log::info!("RAG QUERY: {}", query);
@@ -646,8 +629,7 @@ impl InsightGenerator {
// (`LLM_BACKEND` switch). Must match the backend that populated the // (`LLM_BACKEND` switch). Must match the backend that populated the
// daily-summary embeddings or similarity search will be garbage. // daily-summary embeddings or similarity search will be garbage.
let query_embedding = let query_embedding =
crate::ai::embed_query(&search_cx, &self.ollama, self.llamacpp.as_deref(), &query) crate::ai::embed_query(&self.ollama, self.llamacpp.as_deref(), &query).await?;
.await?;
// Search for similar daily summaries with time-based weighting // Search for similar daily summaries with time-based weighting
// This prioritizes summaries temporally close to the query date // This prioritizes summaries temporally close to the query date
@@ -666,10 +648,7 @@ impl InsightGenerator {
similar_summaries.len() similar_summaries.len()
); );
search_cx.span().set_attribute(KeyValue::new( span.set_attribute("results_count", similar_summaries.len() as i64);
"results_count",
similar_summaries.len() as i64,
));
// Format daily summaries for LLM context // Format daily summaries for LLM context
let formatted = similar_summaries let formatted = similar_summaries
@@ -682,7 +661,7 @@ impl InsightGenerator {
}) })
.collect(); .collect();
search_cx.span().set_status(Status::Ok); span.set_status(Status::Ok);
Ok(formatted) Ok(formatted)
} }
@@ -691,30 +670,28 @@ impl InsightGenerator {
/// tool. Embeds the caller's query as-is (no metadata boilerplate) and /// tool. Embeds the caller's query as-is (no metadata boilerplate) and
/// only applies time weighting when an anchor date is provided — /// only applies time weighting when an anchor date is provided —
/// without one, results rank purely by similarity across all time. /// 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( async fn search_summaries_semantic(
&self, &self,
query: &str, query: &str,
date: Option<chrono::NaiveDate>, date: Option<chrono::NaiveDate>,
limit: usize, limit: usize,
) -> Result<Vec<String>> { ) -> Result<Vec<String>> {
let tracer = global_tracer(); let span = tracing::Span::current();
let current_cx = opentelemetry::Context::current(); span.set_attribute("query", query.to_string());
let mut span = tracer.start_with_context("ai.rag.search_daily_summaries", &current_cx); span.set_attribute("limit", limit as i64);
span.set_attribute(KeyValue::new("query", query.to_string())); span.set_attribute("time_weighted", date.is_some());
span.set_attribute(KeyValue::new("limit", limit as i64));
span.set_attribute(KeyValue::new("time_weighted", date.is_some()));
if let Some(d) = date { 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); log::info!("RAG QUERY: {} (anchor date: {:?})", query, date);
// Must use the same backend that populated the daily-summary // Must use the same backend that populated the daily-summary
// embeddings or similarity search is garbage (see embed_one docs). // embeddings or similarity search is garbage (see embed_one docs).
let query_embedding = let query_embedding =
crate::ai::embed_query(&search_cx, &self.ollama, self.llamacpp.as_deref(), query) crate::ai::embed_query(&self.ollama, self.llamacpp.as_deref(), query).await?;
.await?;
let mut summary_dao = self let mut summary_dao = self
.daily_summary_dao .daily_summary_dao
@@ -732,11 +709,8 @@ impl InsightGenerator {
} }
.map_err(|e| anyhow::anyhow!("Failed to find similar summaries: {:?}", e))?; .map_err(|e| anyhow::anyhow!("Failed to find similar summaries: {:?}", e))?;
search_cx.span().set_attribute(KeyValue::new( span.set_attribute("results_count", similar_summaries.len() as i64);
"results_count", span.set_status(Status::Ok);
similar_summaries.len() as i64,
));
search_cx.span().set_status(Status::Ok);
Ok(similar_summaries Ok(similar_summaries
.into_iter() .into_iter()
@@ -789,20 +763,17 @@ impl InsightGenerator {
/// Gather calendar context for photo timestamp /// Gather calendar context for photo timestamp
/// Uses hybrid time + semantic search (±7 days, ranked by relevance) /// Uses hybrid time + semantic search (±7 days, ranked by relevance)
#[tracing::instrument(name = "ai.context.calendar", skip_all)]
async fn gather_calendar_context( async fn gather_calendar_context(
&self, &self,
parent_cx: &opentelemetry::Context,
timestamp: i64, timestamp: i64,
location: Option<&str>, location: Option<&str>,
) -> Result<Option<String>> { ) -> Result<Option<String>> {
let tracer = global_tracer(); let span = tracing::Span::current();
let span = tracer.start_with_context("ai.context.calendar", parent_cx); let calendar_cx = span.context();
let calendar_cx = parent_cx.with_span(span);
let query_embedding = if let Some(loc) = location { let query_embedding = if let Some(loc) = location {
match crate::ai::embed_query(&calendar_cx, &self.ollama, self.llamacpp.as_deref(), loc) match crate::ai::embed_query(&self.ollama, self.llamacpp.as_deref(), loc).await {
.await
{
Ok(emb) => Some(emb), Ok(emb) => Some(emb),
Err(e) => { Err(e) => {
log::warn!("Failed to generate embedding for location '{}': {}", loc, e); log::warn!("Failed to generate embedding for location '{}': {}", loc, e);
@@ -828,7 +799,7 @@ impl InsightGenerator {
.ok() .ok()
}; };
calendar_cx.span().set_status(Status::Ok); span.set_status(Status::Ok);
if let Some(events) = events { if let Some(events) = events {
if events.is_empty() { if events.is_empty() {
@@ -866,15 +837,14 @@ impl InsightGenerator {
/// Gather location context for photo timestamp /// Gather location context for photo timestamp
/// Finds nearest location record (±30 minutes) /// Finds nearest location record (±30 minutes)
#[tracing::instrument(name = "ai.context.location", skip_all)]
async fn gather_location_context( async fn gather_location_context(
&self, &self,
parent_cx: &opentelemetry::Context,
timestamp: i64, timestamp: i64,
exif_gps: Option<(f64, f64)>, exif_gps: Option<(f64, f64)>,
) -> Result<Option<String>> { ) -> Result<Option<String>> {
let tracer = global_tracer(); let span = tracing::Span::current();
let span = tracer.start_with_context("ai.context.location", parent_cx); let location_cx = span.context();
let location_cx = parent_cx.with_span(span);
let nearest = { let nearest = {
let mut dao = self let mut dao = self
@@ -890,7 +860,7 @@ impl InsightGenerator {
.flatten() .flatten()
}; };
location_cx.span().set_status(Status::Ok); span.set_status(Status::Ok);
if let Some(loc) = nearest { if let Some(loc) = nearest {
// Check if this adds NEW information compared to EXIF // Check if this adds NEW information compared to EXIF
@@ -943,17 +913,16 @@ impl InsightGenerator {
/// Gather search context for photo date /// Gather search context for photo date
/// Uses semantic search on queries (±30 days, top 5 relevant) /// Uses semantic search on queries (±30 days, top 5 relevant)
#[tracing::instrument(name = "ai.context.search", skip_all)]
async fn gather_search_context( async fn gather_search_context(
&self, &self,
parent_cx: &opentelemetry::Context,
timestamp: i64, timestamp: i64,
location: Option<&str>, location: Option<&str>,
contact: Option<&str>, contact: Option<&str>,
enrichment: Option<&str>, enrichment: Option<&str>,
) -> Result<Option<String>> { ) -> Result<Option<String>> {
let tracer = global_tracer(); let span = tracing::Span::current();
let span = tracer.start_with_context("ai.context.search", parent_cx); let search_cx = span.context();
let search_cx = parent_cx.with_span(span);
// Use enrichment (topics + photo description + tags) if available; // Use enrichment (topics + photo description + tags) if available;
// fall back to generic temporal query. // fall back to generic temporal query.
@@ -973,18 +942,13 @@ impl InsightGenerator {
) )
}; };
let query_embedding = match crate::ai::embed_query( let query_embedding =
&search_cx, match crate::ai::embed_query(&self.ollama, self.llamacpp.as_deref(), &query_text).await
&self.ollama,
self.llamacpp.as_deref(),
&query_text,
)
.await
{ {
Ok(emb) => emb, Ok(emb) => emb,
Err(e) => { Err(e) => {
log::warn!("Failed to generate search embedding: {}", e); log::warn!("Failed to generate search embedding: {}", e);
search_cx.span().set_status(Status::Error { span.set_status(Status::Error {
description: e.to_string().into(), description: e.to_string().into(),
}); });
return Ok(None); return Ok(None);
@@ -1006,7 +970,7 @@ impl InsightGenerator {
.ok() .ok()
}; };
search_cx.span().set_status(Status::Ok); span.set_status(Status::Ok);
if let Some(searches) = searches { if let Some(searches) = searches {
if searches.is_empty() { if searches.is_empty() {
@@ -1068,6 +1032,7 @@ impl InsightGenerator {
} }
/// Generate AI insight for a single photo with custom configuration /// 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( pub async fn generate_insight_for_photo_with_config(
&self, &self,
file_path: &str, file_path: &str,
@@ -1079,20 +1044,18 @@ impl InsightGenerator {
top_k: Option<i32>, top_k: Option<i32>,
min_p: Option<f32>, min_p: Option<f32>,
) -> Result<()> { ) -> Result<()> {
let tracer = global_tracer(); let span = tracing::Span::current();
let current_cx = opentelemetry::Context::current();
let mut span = tracer.start_with_context("ai.insight.generate", &current_cx);
// Normalize path to ensure consistent forward slashes in database // Normalize path to ensure consistent forward slashes in database
let file_path = normalize_path(file_path); let file_path = normalize_path(file_path);
log::info!("Generating insight for photo: {}", 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 // Create custom Ollama client if model is specified
let mut ollama_client = if let Some(model) = custom_model { let mut ollama_client = if let Some(model) = custom_model {
log::info!("Using custom model: {}", 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( OllamaClient::new(
self.ollama.primary_url.clone(), self.ollama.primary_url.clone(),
self.ollama.fallback_url.clone(), self.ollama.fallback_url.clone(),
@@ -1100,14 +1063,14 @@ impl InsightGenerator {
Some(model), // Use the same custom model for fallback server Some(model), // Use the same custom model for fallback server
) )
} else { } else {
span.set_attribute(KeyValue::new("model", self.ollama.primary_model.clone())); span.set_attribute("model", self.ollama.primary_model.clone());
self.ollama.clone() self.ollama.clone()
}; };
// Set context size if specified // Set context size if specified
if let Some(ctx) = num_ctx { if let Some(ctx) = num_ctx {
log::info!("Using custom context size: {}", 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)); ollama_client.set_num_ctx(Some(ctx));
} }
@@ -1121,22 +1084,22 @@ impl InsightGenerator {
min_p min_p
); );
if let Some(t) = temperature { 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 { 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 { 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 { 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); ollama_client.set_sampling_params(temperature, top_p, top_k, min_p);
} }
// Create context with this span for child operations // 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 // 1. Get EXIF data for the photo
let exif = { let exif = {
@@ -1153,13 +1116,9 @@ impl InsightGenerator {
let contact = Self::extract_contact_from_path(&file_path); let contact = Self::extract_contact_from_path(&file_path);
log::info!("Extracted contact from path: {:?}", contact); log::info!("Extracted contact from path: {:?}", contact);
insight_cx span.set_attribute("date_taken", date_taken.to_string());
.span()
.set_attribute(KeyValue::new("date_taken", date_taken.to_string()));
if let Some(ref c) = contact { if let Some(ref c) = contact {
insight_cx span.set_attribute("contact", c.clone());
.span()
.set_attribute(KeyValue::new("contact", c.clone()));
} }
// Fetch file tags (used to enrich RAG and final context) // Fetch file tags (used to enrich RAG and final context)
@@ -1197,9 +1156,7 @@ impl InsightGenerator {
.into_iter() .into_iter()
.next(); .next();
let combined = compose_location_string(apollo_primary, nominatim, lat, lon); let combined = compose_location_string(apollo_primary, nominatim, lat, lon);
insight_cx span.set_attribute("location", combined.clone());
.span()
.set_attribute(KeyValue::new("location", combined.clone()));
Some(combined) Some(combined)
} else { } else {
None None
@@ -1234,9 +1191,7 @@ impl InsightGenerator {
} }
}; };
insight_cx span.set_attribute("model_has_vision", has_vision);
.span()
.set_attribute(KeyValue::new("model_has_vision", has_vision));
// Load image and encode as base64 only if model supports vision // Load image and encode as base64 only if model supports vision
let image_base64 = if has_vision { let image_base64 = if has_vision {
@@ -1263,10 +1218,7 @@ impl InsightGenerator {
// Generate brief photo description for RAG enrichment (vision models only) // Generate brief photo description for RAG enrichment (vision models only)
let photo_description: Option<String> = if let Some(ref img_b64) = image_base64 { let photo_description: Option<String> = if let Some(ref img_b64) = image_base64 {
match ollama_client match ollama_client.generate_photo_description(img_b64).await {
.generate_photo_description(&insight_cx, img_b64)
.await
{
Ok(desc) => { Ok(desc) => {
log::info!("Photo description for RAG enrichment: {}", desc); log::info!("Photo description for RAG enrichment: {}", desc);
Some(desc) Some(desc)
@@ -1335,7 +1287,7 @@ impl InsightGenerator {
if !immediate_messages.is_empty() { if !immediate_messages.is_empty() {
// Step 2: Extract topics from immediate messages to enrich RAG query // Step 2: Extract topics from immediate messages to enrich RAG query
let topics = self let topics = self
.extract_topics_from_messages(&immediate_messages, &ollama_client, &insight_cx) .extract_topics_from_messages(&immediate_messages, &ollama_client)
.await; .await;
log::info!("Extracted topics for query enrichment: {:?}", topics); log::info!("Extracted topics for query enrichment: {:?}", topics);
@@ -1367,7 +1319,6 @@ impl InsightGenerator {
}; };
match self match self
.find_relevant_messages_rag_historical( .find_relevant_messages_rag_historical(
&insight_cx,
date_taken, date_taken,
None, None,
contact.as_deref(), contact.as_deref(),
@@ -1391,7 +1342,6 @@ impl InsightGenerator {
&immediate_messages, &immediate_messages,
&ollama_client, &ollama_client,
custom_system_prompt.as_deref(), custom_system_prompt.as_deref(),
&insight_cx,
) )
.await .await
.unwrap_or_else(|| String::from("No immediate context")); .unwrap_or_else(|| String::from("No immediate context"));
@@ -1401,7 +1351,6 @@ impl InsightGenerator {
&historical_messages, &historical_messages,
&ollama_client, &ollama_client,
custom_system_prompt.as_deref(), custom_system_prompt.as_deref(),
&insight_cx,
) )
.await .await
.unwrap_or_else(|| String::from("No historical context")); .unwrap_or_else(|| String::from("No historical context"));
@@ -1420,7 +1369,6 @@ impl InsightGenerator {
&immediate_messages, &immediate_messages,
&ollama_client, &ollama_client,
custom_system_prompt.as_deref(), custom_system_prompt.as_deref(),
&insight_cx,
) )
.await; .await;
} }
@@ -1431,7 +1379,6 @@ impl InsightGenerator {
&immediate_messages, &immediate_messages,
&ollama_client, &ollama_client,
custom_system_prompt.as_deref(), custom_system_prompt.as_deref(),
&insight_cx,
) )
.await; .await;
} }
@@ -1457,7 +1404,6 @@ impl InsightGenerator {
&rag_messages, &rag_messages,
&ollama_client, &ollama_client,
custom_system_prompt.as_deref(), custom_system_prompt.as_deref(),
&insight_cx,
) )
.await; .await;
} }
@@ -1490,7 +1436,7 @@ impl InsightGenerator {
if !sms_messages.is_empty() { if !sms_messages.is_empty() {
match self match self
.sms_client .sms_client
.summarize_context(&sms_messages, &ollama_client, &insight_cx) .summarize_context(&sms_messages, &ollama_client)
.await .await
{ {
Ok(summary) => { Ok(summary) => {
@@ -1504,12 +1450,8 @@ impl InsightGenerator {
} }
let retrieval_method = if used_rag { "RAG" } else { "time-based" }; let retrieval_method = if used_rag { "RAG" } else { "time-based" };
insight_cx span.set_attribute("retrieval_method", retrieval_method);
.span() span.set_attribute("has_sms_context", sms_summary.is_some());
.set_attribute(KeyValue::new("retrieval_method", retrieval_method));
insight_cx
.span()
.set_attribute(KeyValue::new("has_sms_context", sms_summary.is_some()));
log::info!( log::info!(
"Photo context: date={}, location={:?}, retrieval_method={}", "Photo context: date={}, location={:?}, retrieval_method={}",
@@ -1520,7 +1462,7 @@ impl InsightGenerator {
// 6. Gather Google Takeout context from all sources // 6. Gather Google Takeout context from all sources
let calendar_context = self let calendar_context = self
.gather_calendar_context(&insight_cx, timestamp, location.as_deref()) .gather_calendar_context(timestamp, location.as_deref())
.await .await
.ok() .ok()
.flatten(); .flatten();
@@ -1534,14 +1476,13 @@ impl InsightGenerator {
}); });
let location_context = self let location_context = self
.gather_location_context(&insight_cx, timestamp, exif_gps) .gather_location_context(timestamp, exif_gps)
.await .await
.ok() .ok()
.flatten(); .flatten();
let search_context = self let search_context = self
.gather_search_context( .gather_search_context(
&insight_cx,
timestamp, timestamp,
location.as_deref(), location.as_deref(),
contact.as_deref(), contact.as_deref(),
@@ -1574,7 +1515,6 @@ impl InsightGenerator {
// 10. Generate summary first, then derive title from the summary // 10. Generate summary first, then derive title from the summary
let summary = ollama_client let summary = ollama_client
.generate_photo_summary( .generate_photo_summary(
&insight_cx,
date_taken, date_taken,
location.as_deref(), location.as_deref(),
contact.as_deref(), contact.as_deref(),
@@ -1585,18 +1525,14 @@ impl InsightGenerator {
.await?; .await?;
let title = ollama_client 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?; .await?;
log::info!("Generated title: {}", title); log::info!("Generated title: {}", title);
log::info!("Generated summary: {}", summary); log::info!("Generated summary: {}", summary);
insight_cx span.set_attribute("title_length", title.len() as i64);
.span() span.set_attribute("summary_length", summary.len() as i64);
.set_attribute(KeyValue::new("title_length", title.len() as i64));
insight_cx
.span()
.set_attribute(KeyValue::new("summary_length", summary.len() as i64));
// 11. Store in database. content_hash is None here — store_insight // 11. Store in database. content_hash is None here — store_insight
// looks it up from image_exif before persisting; reconciliation // looks it up from image_exif before persisting; reconciliation
@@ -1635,11 +1571,11 @@ impl InsightGenerator {
match &result { match &result {
Ok(_) => { Ok(_) => {
log::info!("Successfully stored insight for {}", file_path); log::info!("Successfully stored insight for {}", file_path);
insight_cx.span().set_status(Status::Ok); span.set_status(Status::Ok);
} }
Err(e) => { Err(e) => {
log::error!("Failed to store insight: {:?}", 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, &self,
messages: &[crate::ai::SmsMessage], messages: &[crate::ai::SmsMessage],
ollama: &OllamaClient, ollama: &OllamaClient,
cx: &opentelemetry::Context,
) -> Vec<String> { ) -> Vec<String> {
if messages.is_empty() { if messages.is_empty() {
return Vec::new(); return Vec::new();
@@ -1688,7 +1623,7 @@ Return ONLY the comma-separated list, nothing else."#,
); );
match ollama 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 .await
{ {
Ok(response) => { Ok(response) => {
@@ -1723,7 +1658,6 @@ Return ONLY the comma-separated list, nothing else."#,
messages: &[String], messages: &[String],
ollama: &OllamaClient, ollama: &OllamaClient,
custom_system: Option<&str>, custom_system: Option<&str>,
cx: &opentelemetry::Context,
) -> Option<String> { ) -> Option<String> {
if messages.is_empty() { if messages.is_empty() {
return None; return None;
@@ -1744,7 +1678,7 @@ Return ONLY the summary, nothing else."#,
let system = custom_system let system = custom_system
.unwrap_or("You are a context summarization assistant. Be concise and factual."); .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), Ok(summary) => Some(summary),
Err(e) => { Err(e) => {
log::warn!("Failed to summarize messages: {}", e); log::warn!("Failed to summarize messages: {}", e);
@@ -1760,7 +1694,6 @@ Return ONLY the summary, nothing else."#,
messages: &[crate::ai::SmsMessage], messages: &[crate::ai::SmsMessage],
ollama: &OllamaClient, ollama: &OllamaClient,
custom_system: Option<&str>, custom_system: Option<&str>,
cx: &opentelemetry::Context,
) -> Option<String> { ) -> Option<String> {
if messages.is_empty() { if messages.is_empty() {
return None; 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.", "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), Ok(summary) => Some(summary),
Err(e) => { Err(e) => {
log::warn!("Failed to summarize immediate context: {}", e); log::warn!("Failed to summarize immediate context: {}", e);
@@ -1842,7 +1775,7 @@ Return ONLY the summary, nothing else."#,
let model = backend.model(); let model = backend.model();
let backend_label = backend.kind.as_str(); let backend_label = backend.kind.as_str();
let result = match tool_name { 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, "search_messages" => self.tool_search_messages(arguments, cx).await,
"get_sms_messages" => self.tool_get_sms_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, "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_file_tags" => self.tool_get_file_tags(arguments, cx).await,
"get_faces_in_photo" => self.tool_get_faces_in_photo(arguments, cx).await, "get_faces_in_photo" => self.tool_get_faces_in_photo(arguments, cx).await,
"describe_photo" => { "describe_photo" => {
self.tool_describe_photo(backend.local(), image_base64, cx) self.tool_describe_photo(backend.local(), image_base64)
.await .await
} }
"reverse_geocode" => self.tool_reverse_geocode(arguments).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 /// Tool: search_rag — semantic search over daily summaries
async fn tool_search_rag( async fn tool_search_rag(&self, args: &serde_json::Value, local: &dyn LlmClient) -> String {
&self,
args: &serde_json::Value,
local: &dyn LlmClient,
cx: &opentelemetry::Context,
) -> String {
let query = match args.get("query").and_then(|v| v.as_str()) { let query = match args.get("query").and_then(|v| v.as_str()) {
Some(q) => q.to_string(), Some(q) => q.to_string(),
None => return "Error: missing required parameter 'query'".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 { let final_results = if rerank_enabled && results.len() > limit {
match self match self.rerank_with_llm(&query, &results, limit, local).await {
.rerank_with_llm(cx, &query, &results, limit, local)
.await
{
Ok(reordered) => reordered, Ok(reordered) => reordered,
Err(e) => { Err(e) => {
log::warn!("rerank failed, using vector order: {}", 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. /// tool_result. On parse failure we fall back to the input order.
async fn rerank_with_llm( async fn rerank_with_llm(
&self, &self,
cx: &opentelemetry::Context,
query: &str, query: &str,
candidates: &[String], candidates: &[String],
limit: usize, limit: usize,
@@ -2035,7 +1959,7 @@ Return ONLY the summary, nothing else."#,
let started = std::time::Instant::now(); let started = std::time::Instant::now();
let system = let system =
Some("You are a terse relevance ranker. You output only numbers separated by commas."); 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!( log::info!(
"rerank: finished in {} ms (prompt={} chars)", "rerank: finished in {} ms (prompt={} chars)",
started.elapsed().as_millis(), started.elapsed().as_millis(),
@@ -2645,11 +2569,10 @@ Return ONLY the summary, nothing else."#,
&self, &self,
local: &dyn LlmClient, local: &dyn LlmClient,
image_base64: &Option<String>, image_base64: &Option<String>,
cx: &opentelemetry::Context,
) -> String { ) -> String {
log::info!("tool_describe_photo: generating visual description"); log::info!("tool_describe_photo: generating visual description");
match image_base64 { match image_base64 {
Some(img) => match local.describe_image(cx, img).await { Some(img) => match local.describe_image(img).await {
Ok(desc) => desc, Ok(desc) => desc,
Err(e) => format!("Error describing photo: {}", e), Err(e) => format!("Error describing photo: {}", e),
}, },
@@ -3057,7 +2980,6 @@ Return ONLY the summary, nothing else."#,
// configured local backend. // configured local backend.
let embed_text = format!("{} {}", name, description); let embed_text = format!("{} {}", name, description);
let embedding: Option<Vec<u8>> = match crate::ai::embed_document( let embedding: Option<Vec<u8>> = match crate::ai::embed_document(
cx,
&self.ollama, &self.ollama,
self.llamacpp.as_deref(), self.llamacpp.as_deref(),
&embed_text, &embed_text,
@@ -4175,6 +4097,7 @@ Return ONLY the summary, nothing else."#,
Ok(ResolvedBackend::new(chat, local, kind, images_inline)) 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( pub async fn generate_agentic_insight_for_photo(
&self, &self,
file_path: &str, file_path: &str,
@@ -4193,19 +4116,17 @@ Return ONLY the summary, nothing else."#,
user_id: i32, user_id: i32,
persona_id: String, persona_id: String,
) -> Result<(Option<i32>, Option<i32>, Option<i32>)> { ) -> Result<(Option<i32>, Option<i32>, Option<i32>)> {
let tracer = global_tracer(); let span = tracing::Span::current();
let current_cx = opentelemetry::Context::current();
let mut span = tracer.start_with_context("ai.insight.generate_agentic", &current_cx);
let file_path = normalize_path(file_path); let file_path = normalize_path(file_path);
log::info!("Generating agentic insight for photo: {}", file_path); log::info!("Generating agentic insight for photo: {}", file_path);
span.set_attribute(KeyValue::new("file_path", file_path.clone())); span.set_attribute("file_path", file_path.clone());
span.set_attribute(KeyValue::new("max_iterations", max_iterations as i64)); span.set_attribute("max_iterations", max_iterations as i64);
// 1. Resolve backend + build clients. // 1. Resolve backend + build clients.
let kind = BackendKind::parse(backend.as_deref().unwrap_or("local"))?; 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 { let overrides = SamplingOverrides {
model: custom_model, model: custom_model,
num_ctx, num_ctx,
@@ -4216,10 +4137,10 @@ Return ONLY the summary, nothing else."#,
enable_thinking, enable_thinking,
}; };
let backend = self.resolve_backend(kind, &overrides).await?; let backend = self.resolve_backend(kind, &overrides).await?;
span.set_attribute(KeyValue::new("model", backend.model().to_string())); span.set_attribute("model", backend.model().to_string());
span.set_attribute(KeyValue::new("images_inline", backend.images_inline)); span.set_attribute("images_inline", backend.images_inline);
let insight_cx = current_cx.with_span(span); let insight_cx = span.context();
// 3. Fetch EXIF // 3. Fetch EXIF
let exif = { let exif = {
@@ -4304,7 +4225,7 @@ Return ONLY the summary, nothing else."#,
// the local backend so non-text work stays off OpenRouter. // the local backend so non-text work stays off OpenRouter.
let inlined_visual_description: Option<String> = if !backend.images_inline { let inlined_visual_description: Option<String> = if !backend.images_inline {
match image_base64.as_deref() { 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) => { Ok(desc) => {
log::info!("{}: vision describe succeeded ({} chars)", kind, desc.len()); log::info!("{}: vision describe succeeded ({} chars)", kind, desc.len());
Some(desc) Some(desc)
@@ -4410,8 +4331,10 @@ Return ONLY the summary, nothing else."#,
let mut messages = vec![system_msg, user_msg]; let mut messages = vec![system_msg, user_msg];
let loop_span = tracer.start_with_context("ai.agentic.loop", &insight_cx); // LLM / tool awaits below run inside `loop_span` so their spans
let loop_cx = insight_cx.with_span(loop_span); // 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 final_content = String::new();
let mut iterations_used = 0usize; let mut iterations_used = 0usize;
@@ -4424,7 +4347,8 @@ Return ONLY the summary, nothing else."#,
let (response, prompt_tokens, eval_tokens) = backend let (response, prompt_tokens, eval_tokens) = backend
.chat() .chat()
.chat_with_tools(&loop_cx, messages.clone(), tools.clone()) .chat_with_tools(messages.clone(), tools.clone())
.instrument(loop_span.clone())
.await?; .await?;
last_prompt_eval_count = prompt_tokens; last_prompt_eval_count = prompt_tokens;
@@ -4470,6 +4394,7 @@ Return ONLY the summary, nothing else."#,
&persona_id, &persona_id,
&loop_cx, &loop_cx,
) )
.instrument(loop_span.clone())
.await; .await;
messages.push(ChatMessage::tool_result(result)); messages.push(ChatMessage::tool_result(result));
} }
@@ -4493,7 +4418,8 @@ Return ONLY the summary, nothing else."#,
))); )));
let (final_response, prompt_tokens, eval_tokens) = backend let (final_response, prompt_tokens, eval_tokens) = backend
.chat() .chat()
.chat_with_tools(&loop_cx, messages.clone(), tools.clone()) .chat_with_tools(messages.clone(), tools.clone())
.instrument(loop_span.clone())
.await?; .await?;
last_prompt_eval_count = prompt_tokens; last_prompt_eval_count = prompt_tokens;
last_eval_count = eval_tokens; last_eval_count = eval_tokens;
@@ -4521,12 +4447,14 @@ Return ONLY the summary, nothing else."#,
&persona_id, &persona_id,
&loop_cx, &loop_cx,
) )
.instrument(loop_span.clone())
.await; .await;
messages.push(ChatMessage::tool_result(result)); messages.push(ChatMessage::tool_result(result));
} }
let (resp2, pt2, et2) = backend let (resp2, pt2, et2) = backend
.chat() .chat()
.chat_with_tools(&loop_cx, messages.clone(), vec![]) .chat_with_tools(messages.clone(), vec![])
.instrument(loop_span.clone())
.await?; .await?;
last_prompt_eval_count = pt2; last_prompt_eval_count = pt2;
last_eval_count = et2; last_eval_count = et2;
@@ -4538,10 +4466,9 @@ Return ONLY the summary, nothing else."#,
} }
} }
loop_cx loop_span.set_attribute("iterations_used", iterations_used as i64);
.span() loop_span.set_status(Status::Ok);
.set_attribute(KeyValue::new("iterations_used", iterations_used as i64)); drop(loop_span);
loop_cx.span().set_status(Status::Ok);
// 13. Strip any leaked <think>…</think> reasoning block (thinking // 13. Strip any leaked <think>…</think> reasoning block (thinking
// models emit it ahead of the answer; the raw transcript in // models emit it ahead of the answer; the raw transcript in
@@ -4608,11 +4535,11 @@ Return ONLY the summary, nothing else."#,
match &stored { match &stored {
Ok(_) => { Ok(_) => {
log::info!("Successfully stored agentic insight for {}", file_path); log::info!("Successfully stored agentic insight for {}", file_path);
insight_cx.span().set_status(Status::Ok); span.set_status(Status::Ok);
} }
Err(e) => { Err(e) => {
log::error!("Failed to store agentic insight: {:?}", 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 let (response, _prompt_tokens, _eval_tokens) = backend
.chat() .chat()
.chat_with_tools(&loop_cx, messages.clone(), tools.clone()) .chat_with_tools(messages.clone(), tools.clone())
.await?; .await?;
// Sanitize tool call arguments before pushing back into history. // Sanitize tool call arguments before pushing back into history.
@@ -4738,7 +4665,7 @@ Return ONLY the summary, nothing else."#,
)); ));
let (final_response, _, _) = backend let (final_response, _, _) = backend
.chat() .chat()
.chat_with_tools(&loop_cx, messages.clone(), tools.clone()) .chat_with_tools(messages.clone(), tools.clone())
.await?; .await?;
// Same bounded tool round as the insight loop: execute any last // Same bounded tool round as the insight loop: execute any last
// tool call, then force a tool-less turn for the actual answer. // 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 let (resp2, _, _) = backend
.chat() .chat()
.chat_with_tools(&loop_cx, messages.clone(), vec![]) .chat_with_tools(messages.clone(), vec![])
.await?; .await?;
final_content = resp2.content; final_content = resp2.content;
} else { } else {
+30 -43
View File
@@ -23,12 +23,12 @@
use anyhow::Context as _; use anyhow::Context as _;
use anyhow::{Result, anyhow, bail}; use anyhow::{Result, anyhow, bail};
use async_trait::async_trait; use async_trait::async_trait;
use opentelemetry::trace::{Span, Status}; use opentelemetry::trace::Status;
use opentelemetry::{Context, KeyValue};
use reqwest::Client; use reqwest::Client;
use serde::Deserialize; use serde::Deserialize;
use serde_json::{Value, json}; use serde_json::{Value, json};
use std::time::Duration; use std::time::Duration;
use tracing_opentelemetry::OpenTelemetrySpanExt as _;
use crate::ai::llm_client::{ use crate::ai::llm_client::{
ChatMessage, LlmClient, LlmStreamEvent, ModelCapabilities, Tool, ToolCall, ToolCallFunction, ChatMessage, LlmClient, LlmStreamEvent, ModelCapabilities, Tool, ToolCall, ToolCallFunction,
@@ -489,7 +489,7 @@ impl LlamaCppClient {
async fn chat_completion_with_model( async fn chat_completion_with_model(
&self, &self,
model: &str, model: &str,
mut span: Option<&mut opentelemetry::global::BoxedSpan>, span: Option<&tracing::Span>,
messages: Vec<ChatMessage>, messages: Vec<ChatMessage>,
tools: Vec<Tool>, tools: Vec<Tool>,
) -> Result<(ChatMessage, Option<i32>, Option<i32>)> { ) -> Result<(ChatMessage, Option<i32>, Option<i32>)> {
@@ -522,13 +522,13 @@ impl LlamaCppClient {
// timeout starts (see ai::gpu). // timeout starts (see ai::gpu).
let _gpu = crate::ai::gpu::llm_lease().await; let _gpu = crate::ai::gpu::llm_lease().await;
let mut request = self.client.post(&url).json(&Value::Object(body)); let mut request = self.client.post(&url).json(&Value::Object(body));
if let Some(s) = span.as_mut() { if let Some(s) = span {
request = crate::otel::with_traceparent(request, s.span_context()); request = crate::otel::with_traceparent(request, s);
} }
let resp = match request.send().await { let resp = match request.send().await {
Ok(r) => r, Ok(r) => r,
Err(e) => { Err(e) => {
if let Some(s) = span.as_mut() { if let Some(s) = span {
s.set_status(Status::error("llama-swap request failed")); s.set_status(Status::error("llama-swap request failed"));
} }
return Err(e).with_context(|| format!("POST {} failed", url)); return Err(e).with_context(|| format!("POST {} failed", url));
@@ -538,14 +538,14 @@ impl LlamaCppClient {
if !resp.status().is_success() { if !resp.status().is_success() {
let status = resp.status(); let status = resp.status();
crate::otel::record_http_failure( crate::otel::record_http_failure(
span.as_deref_mut(), span,
status.as_u16(), status.as_u16(),
&format!("llama-swap chat request failed: {status}"), &format!("llama-swap chat request failed: {status}"),
); );
let body = resp.text().await.unwrap_or_default(); let body = resp.text().await.unwrap_or_default();
bail!("llama-swap chat request failed: {} — {}", status, body); 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 parsed: Value = resp.json().await.context("parsing chat response")?;
let choice = parsed let choice = parsed
@@ -578,12 +578,12 @@ impl LlamaCppClient {
log_timings(&parsed, prompt_tokens, completion_tokens); 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 { 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 { 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); s.set_status(Status::Ok);
} }
@@ -592,13 +592,8 @@ impl LlamaCppClient {
} }
/// Start the `llm.request` span for an outgoing llama-swap call. /// Start the `llm.request` span for an outgoing llama-swap call.
fn start_request_span( fn start_request_span(&self, operation: &str, model: &str) -> tracing::Span {
&self, crate::otel::start_llm_request_span("llamacpp", operation, model)
cx: &Context,
operation: &str,
model: &str,
) -> opentelemetry::global::BoxedSpan {
crate::otel::start_llm_request_span(cx, "llamacpp", operation, model)
} }
} }
@@ -606,7 +601,6 @@ impl LlamaCppClient {
impl LlmClient for LlamaCppClient { impl LlmClient for LlamaCppClient {
async fn generate( async fn generate(
&self, &self,
cx: &Context,
prompt: &str, prompt: &str,
system: Option<&str>, system: Option<&str>,
images: Option<Vec<String>>, images: Option<Vec<String>>,
@@ -621,13 +615,12 @@ impl LlmClient for LlamaCppClient {
// Delegates to chat_with_tools, which starts the llm.request span // Delegates to chat_with_tools, which starts the llm.request span
// (operation=chat) — one span per HTTP request. // (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) Ok(reply.content)
} }
async fn chat_with_tools( async fn chat_with_tools(
&self, &self,
cx: &Context,
messages: Vec<ChatMessage>, messages: Vec<ChatMessage>,
tools: Vec<Tool>, tools: Vec<Tool>,
) -> Result<(ChatMessage, Option<i32>, Option<i32>)> { ) -> Result<(ChatMessage, Option<i32>, Option<i32>)> {
@@ -637,25 +630,19 @@ impl LlmClient for LlamaCppClient {
messages.len(), messages.len(),
tools.len() tools.len()
); );
let mut span = self.start_request_span(cx, "chat", &self.primary_model); let span = self.start_request_span("chat", &self.primary_model);
self.chat_completion_with_model( self.chat_completion_with_model(&self.primary_model.clone(), Some(&span), messages, tools)
&self.primary_model.clone(),
Some(&mut span),
messages,
tools,
)
.await .await
} }
async fn chat_with_tools_stream( async fn chat_with_tools_stream(
&self, &self,
cx: &Context,
messages: Vec<ChatMessage>, messages: Vec<ChatMessage>,
tools: Vec<Tool>, tools: Vec<Tool>,
) -> Result<BoxStream<'static, Result<LlmStreamEvent>>> { ) -> Result<BoxStream<'static, Result<LlmStreamEvent>>> {
// The span outlives this method — it moves into the stream // The span outlives this method — it moves into the stream
// generator and is finalized when the response drains. // 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 url = format!("{}/chat/completions", self.base_url);
let roles: Vec<&str> = messages.iter().map(|m| m.role.as_str()).collect(); let roles: Vec<&str> = messages.iter().map(|m| m.role.as_str()).collect();
log::debug!( log::debug!(
@@ -690,7 +677,7 @@ impl LlmClient for LlamaCppClient {
// below so the lease spans the whole generation, not just the send. // below so the lease spans the whole generation, not just the send.
let gpu = crate::ai::gpu::llm_lease().await; let gpu = crate::ai::gpu::llm_lease().await;
let mut request = self.client.post(&url).json(&Value::Object(body)); 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 let resp = request
.send() .send()
.await .await
@@ -699,14 +686,14 @@ impl LlmClient for LlamaCppClient {
if !resp.status().is_success() { if !resp.status().is_success() {
let status = resp.status(); let status = resp.status();
crate::otel::record_http_failure( crate::otel::record_http_failure(
Some(&mut span), Some(&span),
status.as_u16(), status.as_u16(),
&format!("llama-swap stream request failed: {status}"), &format!("llama-swap stream request failed: {status}"),
); );
let body = resp.text().await.unwrap_or_default(); let body = resp.text().await.unwrap_or_default();
bail!("llama-swap stream request failed: {} — {}", status, body); 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 byte_stream = resp.bytes_stream();
let stream = async_stream::stream! { let stream = async_stream::stream! {
@@ -816,10 +803,10 @@ impl LlmClient for LlamaCppClient {
images: None, images: None,
}; };
if let Some(pt) = prompt_tokens { 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 { 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); span.set_status(Status::Ok);
yield Ok(LlmStreamEvent::Done { yield Ok(LlmStreamEvent::Done {
@@ -832,8 +819,8 @@ impl LlmClient for LlamaCppClient {
Ok(Box::pin(stream)) Ok(Box::pin(stream))
} }
async fn generate_embeddings(&self, cx: &Context, texts: &[&str]) -> Result<Vec<Vec<f32>>> { async fn generate_embeddings(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>> {
let mut span = self.start_request_span(cx, "embeddings", &self.embedding_model); let span = self.start_request_span("embeddings", &self.embedding_model);
// Deliberately NO GPU lease: the embed slot sits in llama-swap's // Deliberately NO GPU lease: the embed slot sits in llama-swap's
// always-resident group and never participates in a model swap, so // always-resident group and never participates in a model swap, so
// leasing here would only stall searches behind a queued synthesis. // 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); 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 { let resp = match request.send().await {
Ok(r) => r, Ok(r) => r,
Err(e) => { Err(e) => {
@@ -856,14 +843,14 @@ impl LlmClient for LlamaCppClient {
if !resp.status().is_success() { if !resp.status().is_success() {
let status = resp.status(); let status = resp.status();
crate::otel::record_http_failure( crate::otel::record_http_failure(
Some(&mut span), Some(&span),
status.as_u16(), status.as_u16(),
&format!("llama-swap embedding request failed: {status}"), &format!("llama-swap embedding request failed: {status}"),
); );
let body = resp.text().await.unwrap_or_default(); let body = resp.text().await.unwrap_or_default();
bail!("llama-swap embedding request failed: {} — {}", status, body); 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)] #[derive(Deserialize)]
struct EmbedResponse { struct EmbedResponse {
@@ -885,9 +872,9 @@ impl LlmClient for LlamaCppClient {
Ok(parsed.data.into_iter().map(|i| i.embedding).collect()) Ok(parsed.data.into_iter().map(|i| i.embedding).collect())
} }
async fn describe_image(&self, cx: &Context, image_base64: &str) -> Result<String> { async fn describe_image(&self, image_base64: &str) -> Result<String> {
// Own span with operation=describe_image — routes through the vision slot. // 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. \ let prompt = "Briefly describe what you see in this image in 1-2 sentences. \
Focus on the people, location, and activity."; Focus on the people, location, and activity.";
let system = "You are a scene description assistant. Be concise and factual."; let system = "You are a scene description assistant. Be concise and factual.";
@@ -905,7 +892,7 @@ impl LlmClient for LlamaCppClient {
let (reply, _, _) = self let (reply, _, _) = self
.chat_completion_with_model( .chat_completion_with_model(
&self.vision_model.clone(), &self.vision_model.clone(),
Some(&mut span), Some(&span),
messages, messages,
Vec::new(), Vec::new(),
) )
+6 -11
View File
@@ -1,7 +1,6 @@
use anyhow::Result; use anyhow::Result;
use async_trait::async_trait; use async_trait::async_trait;
use futures::stream::BoxStream; use futures::stream::BoxStream;
use opentelemetry::Context;
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};
/// Provider-agnostic surface for LLM backends (Ollama, OpenRouter, …). /// 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 /// stays as base64 strings here and is rewritten into content-parts where
/// needed. /// needed.
/// ///
/// Every method that performs an HTTP call takes the caller's tracing /// Every method that performs an HTTP call starts an `llm.request` span
/// `Context` as its first parameter so the impl can start an `llm.request` /// parented on the caller's ambient `tracing` span, so callers must run
/// span parented under the caller's feature span — the codebase threads /// inside their feature span (`#[instrument]` / `.instrument(span)`) — and
/// contexts explicitly rather than relying on ambient `Context::current()` /// carry it across any `tokio::spawn` — for the LLM span to nest under it.
/// at deep call sites.
// First consumer lands in a later PR (OpenRouter impl + hybrid mode routing). // First consumer lands in a later PR (OpenRouter impl + hybrid mode routing).
#[allow(dead_code)] #[allow(dead_code)]
#[async_trait] #[async_trait]
@@ -25,7 +23,6 @@ pub trait LlmClient: Send + Sync {
/// base64 images (ignored by providers without vision support). /// base64 images (ignored by providers without vision support).
async fn generate( async fn generate(
&self, &self,
cx: &Context,
prompt: &str, prompt: &str,
system: Option<&str>, system: Option<&str>,
images: Option<Vec<String>>, images: Option<Vec<String>>,
@@ -35,7 +32,6 @@ pub trait LlmClient: Send + Sync {
/// (which may contain tool_calls) plus optional prompt/eval token counts. /// (which may contain tool_calls) plus optional prompt/eval token counts.
async fn chat_with_tools( async fn chat_with_tools(
&self, &self,
cx: &Context,
messages: Vec<ChatMessage>, messages: Vec<ChatMessage>,
tools: Vec<Tool>, tools: Vec<Tool>,
) -> Result<(ChatMessage, Option<i32>, Option<i32>)>; ) -> Result<(ChatMessage, Option<i32>, Option<i32>)>;
@@ -48,17 +44,16 @@ pub trait LlmClient: Send + Sync {
/// as one `Done` event. /// as one `Done` event.
async fn chat_with_tools_stream( async fn chat_with_tools_stream(
&self, &self,
cx: &Context,
messages: Vec<ChatMessage>, messages: Vec<ChatMessage>,
tools: Vec<Tool>, tools: Vec<Tool>,
) -> Result<BoxStream<'static, Result<LlmStreamEvent>>>; ) -> Result<BoxStream<'static, Result<LlmStreamEvent>>>;
/// Batch embedding generation. Dimensionality is provider/model specific. /// Batch embedding generation. Dimensionality is provider/model specific.
async fn generate_embeddings(&self, cx: &Context, texts: &[&str]) -> Result<Vec<Vec<f32>>>; async fn generate_embeddings(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>>;
/// One-shot vision description of an image. Used to convert images into /// One-shot vision description of an image. Used to convert images into
/// plain text for the hybrid-mode conversation flow. /// plain text for the hybrid-mode conversation flow.
async fn describe_image(&self, cx: &Context, image_base64: &str) -> Result<String>; async fn describe_image(&self, image_base64: &str) -> Result<String>;
/// Enumerate available models with their capabilities. /// Enumerate available models with their capabilities.
async fn list_models(&self) -> Result<Vec<ModelCapabilities>>; async fn list_models(&self) -> Result<Vec<ModelCapabilities>>;
+4 -25
View File
@@ -47,24 +47,12 @@ impl LocalLlm {
/// pick query vs document — retrieval models treat the two sides /// pick query vs document — retrieval models treat the two sides
/// differently and an unmarked embed invites prefix-mismatch bugs. /// differently and an unmarked embed invites prefix-mismatch bugs.
pub async fn embed_query(&self, text: &str) -> Result<Vec<f32>> { pub async fn embed_query(&self, text: &str) -> Result<Vec<f32>> {
super::embed_query( super::embed_query(&self.ollama, self.llamacpp.as_deref(), text).await
&opentelemetry::Context::new(),
&self.ollama,
self.llamacpp.as_deref(),
text,
)
.await
} }
/// Embed corpus text (applies `EMBED_DOCUMENT_PREFIX`). /// Embed corpus text (applies `EMBED_DOCUMENT_PREFIX`).
pub async fn embed_document(&self, text: &str) -> Result<Vec<f32>> { pub async fn embed_document(&self, text: &str) -> Result<Vec<f32>> {
super::embed_document( super::embed_document(&self.ollama, self.llamacpp.as_deref(), text).await
&opentelemetry::Context::new(),
&self.ollama,
self.llamacpp.as_deref(),
text,
)
.await
} }
/// Single-shot local text generation via the `LLM_BACKEND`-selected /// 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<String> { pub async fn generate(&self, prompt: &str, system: Option<&str>) -> Result<String> {
if super::local_backend_is_llamacpp() { if super::local_backend_is_llamacpp() {
if let Some(lc) = self.llamacpp.as_deref() { if let Some(lc) = self.llamacpp.as_deref() {
return <LlamaCppClient as LlmClient>::generate( return <LlamaCppClient as LlmClient>::generate(lc, prompt, system, None).await;
lc,
&opentelemetry::Context::new(),
prompt,
system,
None,
)
.await;
} }
anyhow::bail!( anyhow::bail!(
"LLM_BACKEND=llamacpp but LlamaCppClient is unconfigured — \ "LLM_BACKEND=llamacpp but LlamaCppClient is unconfigured — \
set LLAMA_SWAP_URL or switch to LLM_BACKEND=ollama" set LLAMA_SWAP_URL or switch to LLM_BACKEND=ollama"
); );
} }
self.ollama self.ollama.generate(prompt, system).await
.generate(&opentelemetry::Context::new(), prompt, system)
.await
} }
/// Label identifying which backend + model produces embeddings right /// Label identifying which backend + model produces embeddings right
+5 -9
View File
@@ -115,26 +115,24 @@ fn embed_prefix(key: &str) -> String {
/// `Instruct: <task>\nQuery: `. Must pair with the document prefix the /// `Instruct: <task>\nQuery: `. Must pair with the document prefix the
/// corpus was embedded with or similarity degrades. /// corpus was embedded with or similarity degrades.
pub async fn embed_query( pub async fn embed_query(
cx: &opentelemetry::Context,
ollama: &OllamaClient, ollama: &OllamaClient,
llamacpp: Option<&LlamaCppClient>, llamacpp: Option<&LlamaCppClient>,
text: &str, text: &str,
) -> anyhow::Result<Vec<f32>> { ) -> anyhow::Result<Vec<f32>> {
let prefixed = format!("{}{}", embed_prefix("EMBED_QUERY_PREFIX"), text); 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 corpus text (the stored side of retrieval). Applies
/// `EMBED_DOCUMENT_PREFIX` (default empty; nomic v1.5 wants /// `EMBED_DOCUMENT_PREFIX` (default empty; nomic v1.5 wants
/// `search_document: `, Qwen3-Embedding wants none). /// `search_document: `, Qwen3-Embedding wants none).
pub async fn embed_document( pub async fn embed_document(
cx: &opentelemetry::Context,
ollama: &OllamaClient, ollama: &OllamaClient,
llamacpp: Option<&LlamaCppClient>, llamacpp: Option<&LlamaCppClient>,
text: &str, text: &str,
) -> anyhow::Result<Vec<f32>> { ) -> anyhow::Result<Vec<f32>> {
let prefixed = format!("{}{}", embed_prefix("EMBED_DOCUMENT_PREFIX"), text); 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 /// 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 /// configured), else Ollama. See [`local_backend_is_llamacpp`] for the
/// rationale on consistency. /// rationale on consistency.
pub async fn embed_many( pub async fn embed_many(
cx: &opentelemetry::Context,
ollama: &OllamaClient, ollama: &OllamaClient,
llamacpp: Option<&LlamaCppClient>, llamacpp: Option<&LlamaCppClient>,
texts: &[&str], texts: &[&str],
) -> anyhow::Result<Vec<Vec<f32>>> { ) -> anyhow::Result<Vec<Vec<f32>>> {
if local_backend_is_llamacpp() { if local_backend_is_llamacpp() {
if let Some(lc) = llamacpp { if let Some(lc) = llamacpp {
return <LlamaCppClient as LlmClient>::generate_embeddings(lc, cx, texts).await; return <LlamaCppClient as LlmClient>::generate_embeddings(lc, texts).await;
} }
anyhow::bail!( anyhow::bail!(
"LLM_BACKEND=llamacpp but LlamaCppClient is unconfigured — \ "LLM_BACKEND=llamacpp but LlamaCppClient is unconfigured — \
set LLAMA_SWAP_URL or switch to LLM_BACKEND=ollama" 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 /// Embed one string via the configured local backend. Single-text
/// convenience over [`embed_many`]. /// convenience over [`embed_many`].
pub async fn embed_one( pub async fn embed_one(
cx: &opentelemetry::Context,
ollama: &OllamaClient, ollama: &OllamaClient,
llamacpp: Option<&LlamaCppClient>, llamacpp: Option<&LlamaCppClient>,
text: &str, text: &str,
) -> anyhow::Result<Vec<f32>> { ) -> anyhow::Result<Vec<f32>> {
let mut vecs = embed_many(cx, ollama, llamacpp, &[text]).await?; let mut vecs = embed_many(ollama, llamacpp, &[text]).await?;
vecs.pop() vecs.pop()
.ok_or_else(|| anyhow::anyhow!("embedding backend returned no embeddings")) .ok_or_else(|| anyhow::anyhow!("embedding backend returned no embeddings"))
} }
+1 -4
View File
@@ -259,13 +259,10 @@ pub async fn translate_nl_query(
nl: &str, nl: &str,
tag_vocab: &[(i32, String)], tag_vocab: &[(i32, String)],
today: chrono::NaiveDate, today: chrono::NaiveDate,
cx: &opentelemetry::Context,
) -> Result<StructuredQuery> { ) -> Result<StructuredQuery> {
let system = build_system_prompt(tag_vocab, today); let system = build_system_prompt(tag_vocab, today);
let messages = vec![ChatMessage::system(system), ChatMessage::user(nl)]; let messages = vec![ChatMessage::system(system), ChatMessage::user(nl)];
let (msg, _, _) = client let (msg, _, _) = client.chat_with_tools(messages, Vec::<Tool>::new()).await?;
.chat_with_tools(cx, messages, Vec::<Tool>::new())
.await?;
parse_response(&msg.content, tag_vocab) parse_response(&msg.content, tag_vocab)
} }
+68 -109
View File
@@ -2,14 +2,14 @@ use anyhow::Context as _;
use anyhow::Result; use anyhow::Result;
use async_trait::async_trait; use async_trait::async_trait;
use chrono::NaiveDate; use chrono::NaiveDate;
use opentelemetry::trace::{Span, Status}; use opentelemetry::trace::Status;
use opentelemetry::{Context, KeyValue};
use reqwest::Client; use reqwest::Client;
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};
use std::collections::HashMap; use std::collections::HashMap;
use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex}; use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant}; use std::time::{Duration, Instant};
use tracing_opentelemetry::OpenTelemetrySpanExt as _;
use crate::ai::llm_client::{LlmClient, LlmStreamEvent}; use crate::ai::llm_client::{LlmClient, LlmStreamEvent};
use futures::stream::{BoxStream, StreamExt}; use futures::stream::{BoxStream, StreamExt};
@@ -368,7 +368,7 @@ impl OllamaClient {
async fn try_generate( async fn try_generate(
&self, &self,
mut span: Option<&mut opentelemetry::global::BoxedSpan>, span: Option<&tracing::Span>,
url: &str, url: &str,
model: &str, model: &str,
prompt: &str, prompt: &str,
@@ -390,13 +390,13 @@ impl OllamaClient {
.client .client
.post(format!("{}/api/generate", url)) .post(format!("{}/api/generate", url))
.json(&request); .json(&request);
if let Some(s) = span.as_mut() { if let Some(s) = span {
request_builder = crate::otel::with_traceparent(request_builder, s.span_context()); request_builder = crate::otel::with_traceparent(request_builder, s);
} }
let response = match request_builder.send().await { let response = match request_builder.send().await {
Ok(r) => r, Ok(r) => r,
Err(e) => { Err(e) => {
if let Some(s) = span.as_mut() { if let Some(s) = span {
s.set_status(Status::error("Ollama request failed")); s.set_status(Status::error("Ollama request failed"));
} }
return Err(e.into()); return Err(e.into());
@@ -406,7 +406,7 @@ impl OllamaClient {
if !response.status().is_success() { if !response.status().is_success() {
let status = response.status(); let status = response.status();
crate::otel::record_http_failure( crate::otel::record_http_failure(
span.as_deref_mut(), span,
status.as_u16(), status.as_u16(),
&format!("Ollama request failed: {status}"), &format!("Ollama request failed: {status}"),
); );
@@ -417,23 +417,23 @@ impl OllamaClient {
error_body 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 { let result: OllamaResponse = match response.json().await {
Ok(r) => r, Ok(r) => r,
Err(e) => { Err(e) => {
if let Some(s) = span.as_mut() { if let Some(s) = span {
s.set_status(Status::error("parsing Ollama response failed")); s.set_status(Status::error("parsing Ollama response failed"));
} }
return Err(e.into()); return Err(e.into());
} }
}; };
if let Some(s) = span.as_mut() { if let Some(s) = span {
if let Some(pt) = result.prompt_eval_count { 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 { 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); s.set_status(Status::Ok);
} }
@@ -447,49 +447,32 @@ impl OllamaClient {
} }
/// Start the `llm.request` span for an outgoing Ollama call. /// Start the `llm.request` span for an outgoing Ollama call.
fn start_request_span( fn start_request_span(&self, operation: &str, model: &str) -> tracing::Span {
&self, crate::otel::start_llm_request_span("ollama", operation, model)
cx: &Context,
operation: &str,
model: &str,
) -> opentelemetry::global::BoxedSpan {
crate::otel::start_llm_request_span(cx, "ollama", operation, model)
} }
pub async fn generate( pub async fn generate(&self, prompt: &str, system: Option<&str>) -> Result<String> {
&self, self.generate_with_images(prompt, system, None).await
cx: &Context,
prompt: &str,
system: Option<&str>,
) -> Result<String> {
self.generate_with_images(cx, prompt, system, None).await
} }
#[allow(dead_code)] #[allow(dead_code)]
pub async fn generate_no_think( pub async fn generate_no_think(&self, prompt: &str, system: Option<&str>) -> Result<String> {
&self, self.generate_with_options(prompt, system, None, Some(false), "generate")
cx: &Context,
prompt: &str,
system: Option<&str>,
) -> Result<String> {
self.generate_with_options(cx, prompt, system, None, Some(false), "generate")
.await .await
} }
pub async fn generate_with_images( pub async fn generate_with_images(
&self, &self,
cx: &Context,
prompt: &str, prompt: &str,
system: Option<&str>, system: Option<&str>,
images: Option<Vec<String>>, images: Option<Vec<String>>,
) -> Result<String> { ) -> Result<String> {
self.generate_with_options(cx, prompt, system, images, None, "generate") self.generate_with_options(prompt, system, images, None, "generate")
.await .await
} }
async fn generate_with_options( async fn generate_with_options(
&self, &self,
cx: &Context,
prompt: &str, prompt: &str,
system: Option<&str>, system: Option<&str>,
images: Option<Vec<String>>, images: Option<Vec<String>>,
@@ -509,7 +492,7 @@ impl OllamaClient {
// Span for the whole generate attempt; the model attribute // Span for the whole generate attempt; the model attribute
// reflects whichever server actually answered. // 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 // Try primary server first with primary model
log::info!( log::info!(
@@ -519,7 +502,7 @@ impl OllamaClient {
); );
let primary_result = self let primary_result = self
.try_generate( .try_generate(
Some(&mut span), Some(&span),
&self.primary_url, &self.primary_url,
&self.primary_model, &self.primary_model,
prompt, prompt,
@@ -548,10 +531,10 @@ impl OllamaClient {
fallback_url, fallback_url,
fallback_model fallback_model
); );
span.set_attribute(KeyValue::new("llm.model", fallback_model.to_string())); span.set_attribute("llm.model", fallback_model.to_string());
match self match self
.try_generate( .try_generate(
Some(&mut span), Some(&span),
fallback_url, fallback_url,
fallback_model, fallback_model,
prompt, prompt,
@@ -598,7 +581,6 @@ impl OllamaClient {
/// Generate a title for a single photo based on its generated summary /// Generate a title for a single photo based on its generated summary
pub async fn generate_photo_title( pub async fn generate_photo_title(
&self, &self,
cx: &Context,
summary: &str, summary: &str,
custom_system: Option<&str>, custom_system: Option<&str>,
) -> Result<String> { ) -> Result<String> {
@@ -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 system = custom_system.unwrap_or("You are my long term memory assistant. Use only the information provided. Do not invent details.");
let title = self let title = self
.generate_with_images(cx, &prompt, Some(system), None) .generate_with_images(&prompt, Some(system), None)
.await?; .await?;
// Models decorate despite "Return ONLY the title": quotes, bold // Models decorate despite "Return ONLY the title": quotes, bold
// markers, sometimes a "Title:" label. // 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 /// Generate a summary for a single photo based on its context
pub async fn generate_photo_summary( pub async fn generate_photo_summary(
&self, &self,
cx: &Context,
date: NaiveDate, date: NaiveDate,
location: Option<&str>, location: Option<&str>,
contact: 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 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]); 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 .await
} }
/// Generate a brief visual description of a photo for use in RAG query enrichment. /// 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. /// Returns 1-2 sentences describing people, location, and activity visible in the image.
/// Only called when the model has vision capabilities. /// Only called when the model has vision capabilities.
pub async fn generate_photo_description( pub async fn generate_photo_description(&self, image_base64: &str) -> Result<String> {
&self,
cx: &Context,
image_base64: &str,
) -> Result<String> {
let prompt = "Briefly describe what you see in this image in 1-2 sentences. \ let prompt = "Briefly describe what you see in this image in 1-2 sentences. \
Focus on the people, location, and activity."; Focus on the people, location, and activity.";
let system = "You are a scene description assistant. Be concise and factual."; let system = "You are a scene description assistant. Be concise and factual.";
let images = vec![image_base64.to_string()]; let images = vec![image_base64.to_string()];
let description = self let description = self
.generate_with_options( .generate_with_options(prompt, Some(system), Some(images), None, "describe_image")
cx,
prompt,
Some(system),
Some(images),
None,
"describe_image",
)
.await?; .await?;
Ok(description.trim().to_string()) 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. /// iteration of the agent loop.
pub async fn chat_with_tools( pub async fn chat_with_tools(
&self, &self,
cx: &Context,
messages: Vec<ChatMessage>, messages: Vec<ChatMessage>,
tools: Vec<Tool>, tools: Vec<Tool>,
) -> Result<(ChatMessage, Option<i32>, Option<i32>)> { ) -> Result<(ChatMessage, Option<i32>, Option<i32>)> {
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 order = self.attempt_order();
let mut errors: Vec<String> = Vec::new(); let mut errors: Vec<String> = Vec::new();
@@ -768,10 +737,10 @@ Analyze the image and use specific details from both the visual content and the
model model
); );
if *label == "fallback" { if *label == "fallback" {
span.set_attribute(KeyValue::new("llm.model", model.to_string())); span.set_attribute("llm.model", model.to_string());
} }
match self 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 .await
{ {
Ok(result) => { 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). /// any, live on the final message).
pub async fn chat_with_tools_stream( pub async fn chat_with_tools_stream(
&self, &self,
cx: &Context,
messages: Vec<ChatMessage>, messages: Vec<ChatMessage>,
tools: Vec<Tool>, tools: Vec<Tool>,
) -> Result<BoxStream<'static, Result<LlmStreamEvent>>> { ) -> Result<BoxStream<'static, Result<LlmStreamEvent>>> {
@@ -828,7 +796,7 @@ Analyze the image and use specific details from both the visual content and the
for (label, url, _model) in &order { for (label, url, _model) in &order {
match self match self
.try_chat_with_tools_stream(cx, url, messages.clone(), tools.clone()) .try_chat_with_tools_stream(url, messages.clone(), tools.clone())
.await .await
{ {
Ok(s) => { 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( async fn try_chat_with_tools_stream(
&self, &self,
cx: &Context,
base_url: &str, base_url: &str,
messages: Vec<ChatMessage>, messages: Vec<ChatMessage>,
tools: Vec<Tool>, tools: Vec<Tool>,
@@ -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 // The span outlives this method — the stream yields until the
// response fully drains, so the owned span moves into the // response fully drains, so the owned span moves into the
// generator and is finalized there. // 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 { let request_body = OllamaChatRequest {
model, model,
@@ -878,7 +845,7 @@ Analyze the image and use specific details from both the visual content and the
let response = { let response = {
let mut rb = self.client.post(&url).json(&request_body); 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 { match rb.send().await {
Ok(r) => r, Ok(r) => r,
Err(e) => { 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() { if !response.status().is_success() {
let status = response.status(); let status = response.status();
crate::otel::record_http_failure( crate::otel::record_http_failure(
Some(&mut span), Some(&span),
status.as_u16(), status.as_u16(),
&format!("Ollama stream request failed: {status}"), &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 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`. // Ollama streams NDJSON: each line is a full `OllamaStreamChunk`.
// We buffer partial lines across chunks from the byte stream. // 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, eval_duration,
); );
if let Some(pt) = prompt_eval_count { 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 { 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); span.set_status(Status::Ok);
let message = ChatMessage { 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( async fn try_chat_with_tools(
&self, &self,
mut span: Option<&mut opentelemetry::global::BoxedSpan>, span: Option<&tracing::Span>,
base_url: &str, base_url: &str,
messages: Vec<ChatMessage>, messages: Vec<ChatMessage>,
tools: Vec<Tool>, tools: Vec<Tool>,
@@ -1040,13 +1007,13 @@ Analyze the image and use specific details from both the visual content and the
let response = { let response = {
let mut rb = self.client.post(&url).json(&request_body); let mut rb = self.client.post(&url).json(&request_body);
if let Some(s) = span.as_mut() { if let Some(s) = span {
rb = crate::otel::with_traceparent(rb, s.span_context()); rb = crate::otel::with_traceparent(rb, s);
} }
match rb.send().await { match rb.send().await {
Ok(r) => r, Ok(r) => r,
Err(e) => { Err(e) => {
if let Some(s) = span.as_mut() { if let Some(s) = span {
s.set_status(Status::error("Ollama chat connect failed")); s.set_status(Status::error("Ollama chat connect failed"));
} }
return Err(e) 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() { if !response.status().is_success() {
let status = response.status(); let status = response.status();
crate::otel::record_http_failure( crate::otel::record_http_failure(
span.as_deref_mut(), span,
status.as_u16(), status.as_u16(),
&format!("Ollama chat request failed: {status}"), &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 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 { let chat_response: OllamaChatResponse = match response.json().await {
Ok(r) => r, Ok(r) => r,
Err(e) => { Err(e) => {
if let Some(s) = span.as_mut() { if let Some(s) = span {
s.set_status(Status::error("parsing Ollama chat response failed")); s.set_status(Status::error("parsing Ollama chat response failed"));
} }
return Err(e).with_context(|| "Failed to parse Ollama chat response"); 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 { 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 { 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); 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 /// Generate an embedding vector for text using nomic-embed-text:v1.5
/// Returns a 768-dimensional vector as Vec<f32> /// Returns a 768-dimensional vector as Vec<f32>
pub async fn generate_embedding(&self, cx: &Context, text: &str) -> Result<Vec<f32>> { pub async fn generate_embedding(&self, text: &str) -> Result<Vec<f32>> {
let embeddings = self.generate_embeddings(cx, &[text]).await?; let embeddings = self.generate_embeddings(&[text]).await?;
embeddings embeddings
.into_iter() .into_iter()
.next() .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) /// Generate embeddings for multiple texts in a single API call (batch mode)
/// Returns a vector of 768-dimensional vectors /// Returns a vector of 768-dimensional vectors
/// This is much more efficient than calling generate_embedding multiple times /// This is much more efficient than calling generate_embedding multiple times
pub async fn generate_embeddings(&self, cx: &Context, texts: &[&str]) -> Result<Vec<Vec<f32>>> { pub async fn generate_embeddings(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>> {
let embedding_model = EMBEDDING_MODEL; 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!("=== Ollama Batch Embedding Request ===");
log::debug!("Model: {}", embedding_model); 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 embedding_model
); );
let primary_result = self 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; .await;
let embeddings = match primary_result { 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 embedding_model
); );
match self match self
.try_generate_embeddings( .try_generate_embeddings(Some(&span), fallback_url, embedding_model, texts)
Some(&mut span),
fallback_url,
embedding_model,
texts,
)
.await .await
{ {
Ok(embeddings) => { 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 /// Internal helper to try generating embeddings for multiple texts from a specific server
async fn try_generate_embeddings( async fn try_generate_embeddings(
&self, &self,
mut span: Option<&mut opentelemetry::global::BoxedSpan>, span: Option<&tracing::Span>,
url: &str, url: &str,
model: &str, model: &str,
texts: &[&str], texts: &[&str],
@@ -1238,13 +1200,13 @@ Analyze the image and use specific details from both the visual content and the
.client .client
.post(format!("{}/api/embed", url)) .post(format!("{}/api/embed", url))
.json(&request); .json(&request);
if let Some(s) = span.as_mut() { if let Some(s) = span {
rb = crate::otel::with_traceparent(rb, s.span_context()); rb = crate::otel::with_traceparent(rb, s);
} }
match rb.send().await { match rb.send().await {
Ok(r) => r, Ok(r) => r,
Err(e) => { Err(e) => {
if let Some(s) = span.as_mut() { if let Some(s) = span {
s.set_status(Status::error("Ollama embedding request failed")); s.set_status(Status::error("Ollama embedding request failed"));
} }
return Err(e.into()); 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() { if !response.status().is_success() {
let status = response.status(); let status = response.status();
crate::otel::record_http_failure( crate::otel::record_http_failure(
span.as_deref_mut(), span,
status.as_u16(), status.as_u16(),
&format!("Ollama batch embedding request failed: {status}"), &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 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 { let result: OllamaEmbedResponse = match response.json().await {
Ok(r) => r, Ok(r) => r,
Err(e) => { Err(e) => {
if let Some(s) = span.as_mut() { if let Some(s) = span {
s.set_status(Status::error("parsing Ollama embed response failed")); s.set_status(Status::error("parsing Ollama embed response failed"));
} }
return Err(e.into()); return Err(e.into());
} }
}; };
if let Some(s) = span.as_mut() { if let Some(s) = span {
s.set_status(Status::Ok); s.set_status(Status::Ok);
} }
Ok(result.embeddings) 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 { impl LlmClient for OllamaClient {
async fn generate( async fn generate(
&self, &self,
cx: &Context,
prompt: &str, prompt: &str,
system: Option<&str>, system: Option<&str>,
images: Option<Vec<String>>, images: Option<Vec<String>>,
) -> Result<String> { ) -> Result<String> {
self.generate_with_images(cx, prompt, system, images).await self.generate_with_images(prompt, system, images).await
} }
async fn chat_with_tools( async fn chat_with_tools(
&self, &self,
cx: &Context,
messages: Vec<ChatMessage>, messages: Vec<ChatMessage>,
tools: Vec<Tool>, tools: Vec<Tool>,
) -> Result<(ChatMessage, Option<i32>, Option<i32>)> { ) -> Result<(ChatMessage, Option<i32>, Option<i32>)> {
OllamaClient::chat_with_tools(self, cx, messages, tools).await OllamaClient::chat_with_tools(self, messages, tools).await
} }
async fn chat_with_tools_stream( async fn chat_with_tools_stream(
&self, &self,
cx: &Context,
messages: Vec<ChatMessage>, messages: Vec<ChatMessage>,
tools: Vec<Tool>, tools: Vec<Tool>,
) -> Result<BoxStream<'static, Result<LlmStreamEvent>>> { ) -> Result<BoxStream<'static, Result<LlmStreamEvent>>> {
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<Vec<Vec<f32>>> { async fn generate_embeddings(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>> {
OllamaClient::generate_embeddings(self, cx, texts).await OllamaClient::generate_embeddings(self, texts).await
} }
async fn describe_image(&self, cx: &Context, image_base64: &str) -> Result<String> { async fn describe_image(&self, image_base64: &str) -> Result<String> {
self.generate_photo_description(cx, image_base64).await self.generate_photo_description(image_base64).await
} }
async fn list_models(&self) -> Result<Vec<ModelCapabilities>> { async fn list_models(&self) -> Result<Vec<ModelCapabilities>> {
+33 -41
View File
@@ -5,14 +5,14 @@
use anyhow::Context as _; use anyhow::Context as _;
use anyhow::{Result, anyhow, bail}; use anyhow::{Result, anyhow, bail};
use async_trait::async_trait; use async_trait::async_trait;
use opentelemetry::trace::{Span, Status}; use opentelemetry::trace::Status;
use opentelemetry::{Context, KeyValue};
use reqwest::Client; use reqwest::Client;
use serde::Deserialize; use serde::Deserialize;
use serde_json::{Value, json}; use serde_json::{Value, json};
use std::collections::HashMap; use std::collections::HashMap;
use std::sync::{Arc, Mutex}; use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant}; use std::time::{Duration, Instant};
use tracing_opentelemetry::OpenTelemetrySpanExt as _;
use crate::ai::llm_client::{ use crate::ai::llm_client::{
ChatMessage, LlmClient, LlmStreamEvent, ModelCapabilities, Tool, ToolCall, ToolCallFunction, ChatMessage, LlmClient, LlmStreamEvent, ModelCapabilities, Tool, ToolCall, ToolCallFunction,
@@ -152,13 +152,8 @@ impl OpenRouterClient {
/// `chat`, `embeddings`, or `describe` — the call shape, not the entry /// `chat`, `embeddings`, or `describe` — the call shape, not the entry
/// point (a `describe_image` funneled through `chat_with_tools` is /// point (a `describe_image` funneled through `chat_with_tools` is
/// still a `describe`). /// still a `describe`).
fn start_request_span( fn start_request_span(&self, operation: &str, model: &str) -> tracing::Span {
&self, crate::otel::start_llm_request_span("openrouter", operation, model)
cx: &Context,
operation: &str,
model: &str,
) -> opentelemetry::global::BoxedSpan {
crate::otel::start_llm_request_span(cx, "openrouter", operation, model)
} }
/// Translate canonical messages to the OpenAI-compatible wire shape. /// Translate canonical messages to the OpenAI-compatible wire shape.
@@ -329,7 +324,6 @@ impl OpenRouterClient {
impl LlmClient for OpenRouterClient { impl LlmClient for OpenRouterClient {
async fn generate( async fn generate(
&self, &self,
cx: &Context,
prompt: &str, prompt: &str,
system: Option<&str>, system: Option<&str>,
images: Option<Vec<String>>, images: Option<Vec<String>>,
@@ -344,31 +338,29 @@ impl LlmClient for OpenRouterClient {
// Delegates to chat_with_tools, which starts the llm.request span // Delegates to chat_with_tools, which starts the llm.request span
// (operation=chat) — one span per HTTP request. // (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) Ok(reply.content)
} }
async fn chat_with_tools( async fn chat_with_tools(
&self, &self,
cx: &Context,
messages: Vec<ChatMessage>, messages: Vec<ChatMessage>,
tools: Vec<Tool>, tools: Vec<Tool>,
) -> Result<(ChatMessage, Option<i32>, Option<i32>)> { ) -> Result<(ChatMessage, Option<i32>, Option<i32>)> {
let mut span = self.start_request_span(cx, "chat", &self.primary_model); let span = self.start_request_span("chat", &self.primary_model);
self.chat_with_tools_impl(Some(&mut span), messages, tools) self.chat_with_tools_impl(Some(&span), messages, tools)
.await .await
} }
async fn chat_with_tools_stream( async fn chat_with_tools_stream(
&self, &self,
cx: &Context,
messages: Vec<ChatMessage>, messages: Vec<ChatMessage>,
tools: Vec<Tool>, tools: Vec<Tool>,
) -> Result<BoxStream<'static, Result<LlmStreamEvent>>> { ) -> Result<BoxStream<'static, Result<LlmStreamEvent>>> {
// The span outlives this method — the stream yields until the // The span outlives this method — the stream yields until the
// response fully drains, so the owned span moves into the // response fully drains, so the owned span moves into the
// generator and is finalized there. // 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 url = format!("{}/chat/completions", self.base_url);
let mut body = serde_json::Map::new(); let mut body = serde_json::Map::new();
@@ -397,7 +389,7 @@ impl LlmClient for OpenRouterClient {
let mut request = self let mut request = self
.authed(self.client.post(&url)) .authed(self.client.post(&url))
.json(&Value::Object(body)); .json(&Value::Object(body));
request = crate::otel::with_traceparent(request, span.span_context()); request = crate::otel::with_traceparent(request, &span);
let resp = request let resp = request
.send() .send()
.await .await
@@ -406,14 +398,14 @@ impl LlmClient for OpenRouterClient {
if !resp.status().is_success() { if !resp.status().is_success() {
let status = resp.status(); let status = resp.status();
crate::otel::record_http_failure( crate::otel::record_http_failure(
Some(&mut span), Some(&span),
status.as_u16(), status.as_u16(),
&format!("OpenRouter stream request failed: {status}"), &format!("OpenRouter stream request failed: {status}"),
); );
let body = resp.text().await.unwrap_or_default(); let body = resp.text().await.unwrap_or_default();
bail!("OpenRouter stream request failed: {} — {}", status, body); 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: <json>\n\n`, with // OpenAI-compat SSE stream. Each event is `data: <json>\n\n`, with
// `data: [DONE]` signalling completion. Tool calls arrive as // `data: [DONE]` signalling completion. Tool calls arrive as
@@ -489,7 +481,7 @@ impl LlmClient for OpenRouterClient {
.and_then(|n| n.as_i64()) .and_then(|n| n.as_i64())
.map(|n| n as i32); .map(|n| n as i32);
if let Some(cost) = usage.get("cost").and_then(|c| c.as_f64()) { 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, images: None,
}; };
if let Some(pt) = prompt_tokens { 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 { 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); span.set_status(Status::Ok);
yield Ok(LlmStreamEvent::Done { yield Ok(LlmStreamEvent::Done {
@@ -597,8 +589,8 @@ impl LlmClient for OpenRouterClient {
Ok(Box::pin(stream)) Ok(Box::pin(stream))
} }
async fn generate_embeddings(&self, cx: &Context, texts: &[&str]) -> Result<Vec<Vec<f32>>> { async fn generate_embeddings(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>> {
let mut span = self.start_request_span(cx, "embeddings", &self.embedding_model); let span = self.start_request_span("embeddings", &self.embedding_model);
let url = format!("{}/embeddings", self.base_url); let url = format!("{}/embeddings", self.base_url);
let body = json!({ let body = json!({
"model": self.embedding_model, "model": self.embedding_model,
@@ -606,7 +598,7 @@ impl LlmClient for OpenRouterClient {
}); });
let mut request = self.authed(self.client.post(&url)).json(&body); 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 let resp = request
.send() .send()
.await .await
@@ -618,14 +610,14 @@ impl LlmClient for OpenRouterClient {
if !resp.status().is_success() { if !resp.status().is_success() {
let status = resp.status(); let status = resp.status();
crate::otel::record_http_failure( crate::otel::record_http_failure(
Some(&mut span), Some(&span),
status.as_u16(), status.as_u16(),
&format!("OpenRouter embedding request failed: {status}"), &format!("OpenRouter embedding request failed: {status}"),
); );
let body = resp.text().await.unwrap_or_default(); let body = resp.text().await.unwrap_or_default();
bail!("OpenRouter embedding request failed: {} — {}", status, body); 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)] #[derive(Deserialize)]
struct EmbedResponse { struct EmbedResponse {
@@ -641,10 +633,10 @@ impl LlmClient for OpenRouterClient {
Ok(parsed.data.into_iter().map(|i| i.embedding).collect()) Ok(parsed.data.into_iter().map(|i| i.embedding).collect())
} }
async fn describe_image(&self, cx: &Context, image_base64: &str) -> Result<String> { async fn describe_image(&self, image_base64: &str) -> Result<String> {
// Own span with operation=describe_image — the funnel into // Own span with operation=describe_image — the funnel into
// chat_with_tools_impl keeps the request on a single span. // 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. \ let prompt = "Briefly describe what you see in this image in 1-2 sentences. \
Focus on the people, location, and activity."; Focus on the people, location, and activity.";
let mut messages: Vec<ChatMessage> = Vec::new(); let mut messages: Vec<ChatMessage> = Vec::new();
@@ -656,7 +648,7 @@ impl LlmClient for OpenRouterClient {
messages.push(user); messages.push(user);
let (reply, _, _) = self let (reply, _, _) = self
.chat_with_tools_impl(Some(&mut span), messages, Vec::new()) .chat_with_tools_impl(Some(&span), messages, Vec::new())
.await?; .await?;
Ok(reply.content) Ok(reply.content)
} }
@@ -719,7 +711,7 @@ impl OpenRouterClient {
/// operation label. /// operation label.
async fn chat_with_tools_impl( async fn chat_with_tools_impl(
&self, &self,
mut span: Option<&mut opentelemetry::global::BoxedSpan>, span: Option<&tracing::Span>,
messages: Vec<ChatMessage>, messages: Vec<ChatMessage>,
tools: Vec<Tool>, tools: Vec<Tool>,
) -> Result<(ChatMessage, Option<i32>, Option<i32>)> { ) -> Result<(ChatMessage, Option<i32>, Option<i32>)> {
@@ -751,13 +743,13 @@ impl OpenRouterClient {
let mut request = self let mut request = self
.authed(self.client.post(&url)) .authed(self.client.post(&url))
.json(&Value::Object(body)); .json(&Value::Object(body));
if let Some(s) = span.as_mut() { if let Some(s) = span {
request = crate::otel::with_traceparent(request, s.span_context()); request = crate::otel::with_traceparent(request, s);
} }
let resp = match request.send().await { let resp = match request.send().await {
Ok(r) => r, Ok(r) => r,
Err(e) => { Err(e) => {
if let Some(s) = span.as_mut() { if let Some(s) = span {
s.set_status(Status::error("OpenRouter request failed")); s.set_status(Status::error("OpenRouter request failed"));
} }
return Err(e).with_context(|| format!("POST {} failed", url)); return Err(e).with_context(|| format!("POST {} failed", url));
@@ -767,19 +759,19 @@ impl OpenRouterClient {
if !resp.status().is_success() { if !resp.status().is_success() {
let status = resp.status(); let status = resp.status();
crate::otel::record_http_failure( crate::otel::record_http_failure(
span.as_deref_mut(), span,
status.as_u16(), status.as_u16(),
&format!("OpenRouter chat request failed: {status}"), &format!("OpenRouter chat request failed: {status}"),
); );
let body = resp.text().await.unwrap_or_default(); let body = resp.text().await.unwrap_or_default();
bail!("OpenRouter chat request failed: {} — {}", status, body); 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 { let parsed: Value = match resp.json().await {
Ok(v) => v, Ok(v) => v,
Err(e) => { Err(e) => {
if let Some(s) = span.as_mut() { if let Some(s) = span {
s.set_status(Status::error("parsing chat response failed")); s.set_status(Status::error("parsing chat response failed"));
} }
return Err(e).context("parsing chat response"); return Err(e).context("parsing chat response");
@@ -813,15 +805,15 @@ impl OpenRouterClient {
.and_then(|v| v.as_i64()) .and_then(|v| v.as_i64())
.map(|n| n as i32); .map(|n| n as i32);
// OpenRouter reports per-request USD cost in `usage.cost`. // 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 { 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 { 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()) { 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); s.set_status(Status::Ok);
} }
+36 -42
View File
@@ -22,15 +22,16 @@
//! registry. //! registry.
//! - The HTTP handlers under ` /persona_chat/*`. //! - 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 anyhow::{Context, Result, anyhow, bail};
use chrono::Utc; use chrono::Utc;
use opentelemetry::KeyValue; use opentelemetry::trace::Status;
use opentelemetry::trace::{Span, Status, TraceContextExt, Tracer};
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};
use std::collections::HashMap; use std::collections::HashMap;
use std::sync::{Arc, Mutex as StdMutex}; use std::sync::{Arc, Mutex as StdMutex};
use tokio::sync::Mutex as TokioMutex; use tokio::sync::Mutex as TokioMutex;
use tracing::Instrument;
use tracing_opentelemetry::OpenTelemetrySpanExt;
use uuid::Uuid; use uuid::Uuid;
use crate::ai::backend::{BackendKind, SamplingOverrides}; use crate::ai::backend::{BackendKind, SamplingOverrides};
@@ -45,7 +46,6 @@ use crate::ai::turn_registry::{TurnEntry, TurnRegistry};
use crate::data::Claims; use crate::data::Claims;
use crate::database::PersonaDao; use crate::database::PersonaDao;
use crate::libraries; use crate::libraries;
use crate::otel::{extract_context_from_request, global_tracer};
use crate::state::AppState; use crate::state::AppState;
// ───────────────────────────────────────────────────────────────────── // ─────────────────────────────────────────────────────────────────────
@@ -660,7 +660,6 @@ impl PersonaChatSession {
let svc = self.clone(); let svc = self.clone();
let persona_for_task = persona.clone(); let persona_for_task = persona.clone();
let entry_clone = entry.clone(); let entry_clone = entry.clone();
let turn_id_for_task = turn_id.clone();
let req_for_task = PersonaChatTurnRequest { let req_for_task = PersonaChatTurnRequest {
user_message: trimmed.clone(), user_message: trimmed.clone(),
..req ..req
@@ -669,15 +668,19 @@ impl PersonaChatSession {
in_flight: self.in_flight.clone(), in_flight: self.in_flight.clone(),
conversation_id: conversation_id.clone(), conversation_id: conversation_id.clone(),
}; };
let conversation_for_span = conversation_id.clone(); // Created before the spawn so it parents on the dispatching request,
let handle = tokio::spawn(async move { // then carried into the task with `.instrument` (tracing context is
let tracer = global_tracer(); // thread-local and does not cross `tokio::spawn` on its own).
let mut span = tracer.start("ai.persona_chat.turn.execute"); let turn_span = tracing::info_span!(
span.set_attribute(KeyValue::new("turn_id", turn_id_for_task.clone())); "ai.persona_chat.turn.execute",
span.set_attribute(KeyValue::new("persona_id", persona_id.clone())); turn_id = %turn_id,
span.set_attribute(KeyValue::new("conversation_id", conversation_for_span)); persona_id = %persona_id,
span.set_attribute(KeyValue::new("library_id", library_id as i64)); conversation_id = %conversation_id,
library_id = library_id,
);
let handle = tokio::spawn(
async move {
let span = tracing::Span::current();
let result = svc let result = svc
.run_streaming_turn( .run_streaming_turn(
persona_for_task, persona_for_task,
@@ -685,16 +688,16 @@ impl PersonaChatSession {
entry_clone.clone(), entry_clone.clone(),
user_id, user_id,
library_id, library_id,
&mut span, &span,
) )
.await; .await;
match result { match result {
Ok(()) => { Ok(()) => {
span.set_attribute(KeyValue::new("status", "done")); span.set_attribute("status", "done");
span.set_status(Status::Ok); span.set_status(Status::Ok);
} }
Err(e) => { Err(e) => {
span.set_attribute(KeyValue::new("status", "error")); span.set_attribute("status", "error");
span.set_status(Status::error(format!("{e}"))); span.set_status(Status::error(format!("{e}")));
let _ = entry_clone let _ = entry_clone
.push_event(ChatStreamEvent::Error(format!("{e}"))) .push_event(ChatStreamEvent::Error(format!("{e}")))
@@ -706,7 +709,9 @@ impl PersonaChatSession {
// captured by the spawned task: it must drop when the task ends // captured by the spawned task: it must drop when the task ends
// (completion, error, or abort) — not when dispatch returns. // (completion, error, or abort) — not when dispatch returns.
drop(in_flight_guard); drop(in_flight_guard);
}); }
.instrument(turn_span),
);
entry.set_abort_handle(handle.abort_handle()); entry.set_abort_handle(handle.abort_handle());
Ok(turn_id) Ok(turn_id)
@@ -722,7 +727,7 @@ impl PersonaChatSession {
entry: Arc<TurnEntry>, entry: Arc<TurnEntry>,
user_id: i32, user_id: i32,
_library_id: i32, _library_id: i32,
turn_span: &mut opentelemetry::global::BoxedSpan, turn_span: &tracing::Span,
) -> Result<()> { ) -> Result<()> {
// Per-conversation mutex — serializes this turn against a // Per-conversation mutex — serializes this turn against a
// concurrent rewind or branch switch on the same transcript. // concurrent rewind or branch switch on the same transcript.
@@ -810,12 +815,10 @@ impl PersonaChatSession {
.resolve_backend(backend_kind, &overrides) .resolve_backend(backend_kind, &overrides)
.await?; .await?;
let model_used = backend.model().to_string(); let model_used = backend.model().to_string();
// llm.* attributes on the turn-level span (issue #130), plus the // llm.* attributes on the turn-level span (issue #130). The child
// parent context for the child llm.request spans. // llm.request spans parent on it ambiently (the task is instrumented).
turn_span.set_attribute(KeyValue::new("llm.backend", backend.kind.as_str())); turn_span.set_attribute("llm.backend", backend.kind.as_str());
turn_span.set_attribute(KeyValue::new("llm.model", model_used.clone())); turn_span.set_attribute("llm.model", model_used.clone());
let turn_cx = opentelemetry::Context::new()
.with_remote_span_context(turn_span.span_context().clone());
// Build a no-photo tool catalog. `current_gate_opts_for_persona` // Build a no-photo tool catalog. `current_gate_opts_for_persona`
// already probes each tool's backing table for presence and // already probes each tool's backing table for presence and
@@ -882,7 +885,6 @@ impl PersonaChatSession {
&persona.persona_id, &persona.persona_id,
max_iterations, max_iterations,
&entry, &entry,
&turn_cx,
) )
.await .await
.map_err(|e| anyhow!("persona chat loop failed: {e:?}"))?; .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. // list falls back to the persona name while the title is empty.
if needs_title { if needs_title {
let title = self let title = self
.generate_title(&backend, &messages, &turn_cx) .generate_title(&backend, &messages)
.await .await
.unwrap_or_else(|| fallback_title(&messages)); .unwrap_or_else(|| fallback_title(&messages));
if !title.is_empty() { if !title.is_empty() {
@@ -1056,7 +1058,6 @@ impl PersonaChatSession {
&self, &self,
backend: &crate::ai::backend::ResolvedBackend, backend: &crate::ai::backend::ResolvedBackend,
messages: &[ChatMessage], messages: &[ChatMessage],
cx: &opentelemetry::Context,
) -> Option<String> { ) -> Option<String> {
let user_message = messages.iter().find(|m| m.role == "user")?; let user_message = messages.iter().find(|m| m.role == "user")?;
// The last assistant message with actual prose — earlier ones may be // 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())?; .find(|m| m.role == "assistant" && !m.content.trim().is_empty())?;
let prompt = title_prompt(&user_message.content, &reply.content); 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) => { Ok(raw) => {
let title = sanitize_title(&crate::ai::llm_client::strip_think_blocks(&raw)); let title = sanitize_title(&crate::ai::llm_client::strip_think_blocks(&raw));
if title.is_empty() { None } else { Some(title) } 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 /// the `turn_id` immediately; the streaming agent loop runs in a
/// background task and the client opens the SSE replay stream with it. /// background task and the client opens the SSE replay stream with it.
#[post("/persona_chat/turn")] #[post("/persona_chat/turn")]
#[tracing::instrument(name = "http.persona_chat.turn", skip_all)]
pub async fn persona_chat_turn_handler( pub async fn persona_chat_turn_handler(
http_request: HttpRequest,
claims: Claims, claims: Claims,
request: web::Json<PersonaChatTurnRequest>, request: web::Json<PersonaChatTurnRequest>,
app_state: web::Data<AppState>, app_state: web::Data<AppState>,
) -> impl Responder { ) -> impl Responder {
let parent_context = extract_context_from_request(&http_request); let span = tracing::Span::current();
let tracer = global_tracer(); span.set_attribute("conversation_id", request.conversation_id.clone());
let mut span = tracer.start_with_context("http.persona_chat.turn", &parent_context);
span.set_attribute(KeyValue::new(
"conversation_id",
request.conversation_id.clone(),
));
// The transcript is not library-scoped, but the turn's tool catalog // The transcript is not library-scoped, but the turn's tool catalog
// (memories, SMS, places) resolves against a library — mirror the // (memories, SMS, places) resolves against a library — mirror the
@@ -1313,7 +1309,7 @@ pub async fn persona_chat_turn_handler(
.await .await
{ {
Ok(turn_id) => { 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); span.set_status(Status::Ok);
HttpResponse::Accepted().json(serde_json::json!({ HttpResponse::Accepted().json(serde_json::json!({
"turn_id": turn_id, "turn_id": turn_id,
@@ -1540,22 +1536,20 @@ pub async fn persona_chat_branches_handler(
/// shared renderer. /// shared renderer.
#[get("/persona_chat/turn/{turn_id}")] #[get("/persona_chat/turn/{turn_id}")]
pub async fn persona_turn_replay_handler( pub async fn persona_turn_replay_handler(
http_request: HttpRequest,
path: web::Path<String>, path: web::Path<String>,
query: web::Query<ReplayQuery>, query: web::Query<ReplayQuery>,
app_state: web::Data<AppState>, app_state: web::Data<AppState>,
) -> HttpResponse { ) -> 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} — cancel a running persona turn.
#[delete("/persona_chat/turn/{turn_id}")] #[delete("/persona_chat/turn/{turn_id}")]
pub async fn persona_turn_cancel_handler( pub async fn persona_turn_cancel_handler(
http_request: HttpRequest,
path: web::Path<String>, path: web::Path<String>,
app_state: web::Data<AppState>, app_state: web::Data<AppState>,
) -> impl Responder { ) -> impl Responder {
crate::ai::handlers::cancel_turn_impl(http_request, path, app_state).await crate::ai::handlers::cancel_turn_impl(path, app_state).await
} }
// ───────────────────────────────────────────────────────────────────── // ─────────────────────────────────────────────────────────────────────
-2
View File
@@ -324,7 +324,6 @@ impl SmsApiClient {
&self, &self,
messages: &[SmsMessage], messages: &[SmsMessage],
ollama: &OllamaClient, ollama: &OllamaClient,
cx: &opentelemetry::Context,
) -> Result<String> { ) -> Result<String> {
if messages.is_empty() { if messages.is_empty() {
return Ok(String::from("No messages on this day")); return Ok(String::from("No messages on this day"));
@@ -357,7 +356,6 @@ Summary:"#,
ollama ollama
.generate( .generate(
cx,
&prompt, &prompt,
// Some("You are a summarizer for the purposes of jogging my memory and highlighting events and situations."), // 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."), Some("You are the keeper of memories, ingest the context and give me a casual summary of the moment."),
+42 -56
View File
@@ -6,13 +6,12 @@
// (audio read directly; video has its audio track extracted via ffmpeg). // (audio read directly; video has its audio track extracted via ffmpeg).
use actix_multipart::Multipart; 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 anyhow::Context;
use base64::Engine; use base64::Engine;
use bytes::{BufMut, BytesMut}; use bytes::{BufMut, BytesMut};
use futures::StreamExt; use futures::StreamExt;
use opentelemetry::KeyValue; use opentelemetry::trace::Status;
use opentelemetry::trace::{Span, Status, Tracer};
use regex::Regex; use regex::Regex;
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};
use serde_json::{Value, json}; use serde_json::{Value, json};
@@ -21,6 +20,8 @@ use std::path::Path;
use std::sync::{LazyLock, Mutex as StdMutex}; use std::sync::{LazyLock, Mutex as StdMutex};
use std::time::{Duration, Instant}; use std::time::{Duration, Instant};
use tokio::sync::Semaphore; use tokio::sync::Semaphore;
use tracing::Instrument;
use tracing_opentelemetry::OpenTelemetrySpanExt;
use uuid::Uuid; use uuid::Uuid;
use crate::ai::llamacpp::LlamaCppClient; 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::file_types::{is_audio_file, is_video_file};
use crate::files::is_valid_full_path; use crate::files::is_valid_full_path;
use crate::libraries; use crate::libraries;
use crate::otel::{extract_context_from_request, global_tracer};
use crate::state::AppState; use crate::state::AppState;
/// Hard cap on an uploaded voice-reference clip. Chatterbox itself caps the /// 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 /// POST /tts/speech — synthesize `text` (optionally in a named `voice`) and
/// return base64-encoded audio for `data:` URI playback on the client. /// return base64-encoded audio for `data:` URI playback on the client.
#[post("/tts/speech")] #[post("/tts/speech")]
#[tracing::instrument(name = "http.tts.speech", skip_all)]
pub async fn tts_speech_handler( pub async fn tts_speech_handler(
http_request: HttpRequest,
_claims: Claims, _claims: Claims,
req: web::Json<TtsSpeechRequest>, req: web::Json<TtsSpeechRequest>,
app_state: web::Data<AppState>, app_state: web::Data<AppState>,
) -> impl Responder { ) -> impl Responder {
let parent_context = extract_context_from_request(&http_request); let span = tracing::Span::current();
let mut span = global_tracer().start_with_context("http.tts.speech", &parent_context);
let text = prepare_for_tts(&req.text); let text = prepare_for_tts(&req.text);
if text.is_empty() { if text.is_empty() {
@@ -414,10 +413,10 @@ pub async fn tts_speech_handler(
.filter(|s| !s.is_empty()) .filter(|s| !s.is_empty())
.or(dv.as_deref()); .or(dv.as_deref());
span.set_attribute(KeyValue::new("tts.model", client.tts_model.clone())); span.set_attribute("tts.model", client.tts_model.clone());
span.set_attribute(KeyValue::new("tts.format", format.to_string())); span.set_attribute("tts.format", format.to_string());
span.set_attribute(KeyValue::new("tts.has_voice", voice.is_some())); span.set_attribute("tts.has_voice", voice.is_some());
span.set_attribute(KeyValue::new("tts.text_len", text.len() as i64)); span.set_attribute("tts.text_len", text.len() as i64);
// Clamp generation knobs to Chatterbox's documented ranges before forwarding. // Clamp generation knobs to Chatterbox's documented ranges before forwarding.
let exaggeration = req.exaggeration.map(|x| x.clamp(0.25, 2.0)); let exaggeration = req.exaggeration.map(|x| x.clamp(0.25, 2.0));
@@ -441,7 +440,7 @@ pub async fn tts_speech_handler(
.await .await
{ {
Ok(bytes) => { 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); span.set_status(Status::Ok);
let audio_base64 = base64::engine::general_purpose::STANDARD.encode(&bytes); let audio_base64 = base64::engine::general_purpose::STANDARD.encode(&bytes);
HttpResponse::Ok().json(TtsSpeechResponse { 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 /// permit (instead of fast-failing 429) and the client polls the job until
/// the audio is ready. /// the audio is ready.
#[post("/tts/speech/jobs")] #[post("/tts/speech/jobs")]
#[tracing::instrument(name = "http.tts.speech_job.create", skip_all)]
pub async fn create_speech_job_handler( pub async fn create_speech_job_handler(
http_request: HttpRequest,
_claims: Claims, _claims: Claims,
req: web::Json<TtsSpeechRequest>, req: web::Json<TtsSpeechRequest>,
app_state: web::Data<AppState>, app_state: web::Data<AppState>,
) -> impl Responder { ) -> impl Responder {
let parent_context = extract_context_from_request(&http_request); let span = tracing::Span::current();
let mut span =
global_tracer().start_with_context("http.tts.speech_job.create", &parent_context);
let text = prepare_for_tts(&req.text); let text = prepare_for_tts(&req.text);
if text.is_empty() { 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 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)); let temperature = req.temperature.map(|x| x.clamp(0.05, 5.0));
span.set_attribute(KeyValue::new("tts.format", format.clone())); span.set_attribute("tts.format", format.clone());
span.set_attribute(KeyValue::new("tts.has_voice", voice.is_some())); span.set_attribute("tts.has_voice", voice.is_some());
span.set_attribute(KeyValue::new("tts.text_len", text.len() as i64)); span.set_attribute("tts.text_len", text.len() as i64);
let job_id = Uuid::new_v4(); let job_id = Uuid::new_v4();
{ {
@@ -573,7 +570,8 @@ pub async fn create_speech_job_handler(
} }
let state = app_state.clone(); let state = app_state.clone();
let handle = tokio::spawn(async move { let handle = tokio::spawn(
async move {
// Queue rather than fast-fail: jobs wait their turn for the GPU. // Queue rather than fast-fail: jobs wait their turn for the GPU.
let _permit = match TTS_PERMIT.acquire().await { let _permit = match TTS_PERMIT.acquire().await {
Ok(p) => p, Ok(p) => p,
@@ -641,7 +639,10 @@ pub async fn create_speech_job_handler(
); );
} }
} }
}); }
// 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 // Aborting an already-finished task is a no-op, so this late install is
// safe even if the job raced to completion. // safe even if the job raced to completion.
with_job(job_id, |job| { 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. /// 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). /// 404s after the job expires (results are kept ~10 min past completion).
#[get("/tts/speech/jobs/{id}")] #[get("/tts/speech/jobs/{id}")]
pub async fn speech_job_status_handler( #[tracing::instrument(name = "http.tts.speech_job.status", skip_all)]
http_request: HttpRequest, pub async fn speech_job_status_handler(_claims: Claims, path: web::Path<String>) -> impl Responder {
_claims: Claims, let span = tracing::Span::current();
path: web::Path<String>,
) -> 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);
let Ok(id) = Uuid::parse_str(&path.into_inner()) else { let Ok(id) = Uuid::parse_str(&path.into_inner()) else {
span.set_status(Status::error("invalid job id")); 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 /// limitation as the sync path); cancelling stops the wait and discards the
/// result. Cancelling an already-finished job leaves it terminal. /// result. Cancelling an already-finished job leaves it terminal.
#[delete("/tts/speech/jobs/{id}")] #[delete("/tts/speech/jobs/{id}")]
pub async fn cancel_speech_job_handler( #[tracing::instrument(name = "http.tts.speech_job.cancel", skip_all)]
http_request: HttpRequest, pub async fn cancel_speech_job_handler(_claims: Claims, path: web::Path<String>) -> impl Responder {
_claims: Claims, let span = tracing::Span::current();
path: web::Path<String>,
) -> 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);
let Ok(id) = Uuid::parse_str(&path.into_inner()) else { let Ok(id) = Uuid::parse_str(&path.into_inner()) else {
span.set_status(Status::error("invalid job id")); 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 /// 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. /// load the TTS model (and evict the resident LLM); see VOICES_CACHE.
#[get("/tts/voices")] #[get("/tts/voices")]
#[tracing::instrument(name = "http.tts.voices.list", skip_all)]
pub async fn list_voices_handler( pub async fn list_voices_handler(
http_request: HttpRequest,
_claims: Claims, _claims: Claims,
query: web::Query<ListVoicesQuery>, query: web::Query<ListVoicesQuery>,
app_state: web::Data<AppState>, app_state: web::Data<AppState>,
) -> impl Responder { ) -> impl Responder {
let parent_context = extract_context_from_request(&http_request); let span = tracing::Span::current();
let mut span = global_tracer().start_with_context("http.tts.voices.list", &parent_context);
let force = query let force = query
.refresh .refresh
.as_deref() .as_deref()
.is_some_and(|v| matches!(v, "1" | "true" | "yes")); .is_some_and(|v| matches!(v, "1" | "true" | "yes"));
if !force && let Some(v) = cached_voices() { 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); span.set_status(Status::Ok);
return HttpResponse::Ok().json(v); return HttpResponse::Ok().json(v);
} }
@@ -776,7 +766,7 @@ pub async fn list_voices_handler(
match client.list_voices().await { match client.list_voices().await {
Ok(v) => { Ok(v) => {
store_voices_cache(&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); span.set_status(Status::Ok);
HttpResponse::Ok().json(v) 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} — remove a cloned voice from the library.
#[delete("/tts/voices/{name}")] #[delete("/tts/voices/{name}")]
#[tracing::instrument(name = "http.tts.voices.delete", skip_all)]
pub async fn delete_voice_handler( pub async fn delete_voice_handler(
http_request: HttpRequest,
_claims: Claims, _claims: Claims,
path: web::Path<String>, path: web::Path<String>,
app_state: web::Data<AppState>, app_state: web::Data<AppState>,
) -> impl Responder { ) -> impl Responder {
let parent_context = extract_context_from_request(&http_request); let span = tracing::Span::current();
let mut span = global_tracer().start_with_context("http.tts.voices.delete", &parent_context);
let Some(client) = app_state.llamacpp.as_ref() else { let Some(client) = app_state.llamacpp.as_ref() else {
span.set_status(Status::error("tts backend not configured")); 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" })); 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 { match client.delete_voice(&name).await {
Ok(v) => { Ok(v) => {
@@ -835,14 +824,13 @@ pub async fn delete_voice_handler(
/// plus optional `start_seconds` / `duration_seconds` (text) selecting which /// plus optional `start_seconds` / `duration_seconds` (text) selecting which
/// window of a longer recording becomes the reference clip. /// window of a longer recording becomes the reference clip.
#[post("/tts/voices/upload")] #[post("/tts/voices/upload")]
#[tracing::instrument(name = "http.tts.voices.upload", skip_all)]
pub async fn create_voice_upload_handler( pub async fn create_voice_upload_handler(
http_request: HttpRequest,
_claims: Claims, _claims: Claims,
mut payload: Multipart, mut payload: Multipart,
app_state: web::Data<AppState>, app_state: web::Data<AppState>,
) -> impl Responder { ) -> impl Responder {
let parent_context = extract_context_from_request(&http_request); let span = tracing::Span::current();
let mut span = global_tracer().start_with_context("http.tts.voices.upload", &parent_context);
let Some(client) = app_state.llamacpp.as_ref() else { let Some(client) = app_state.llamacpp.as_ref() else {
span.set_status(Status::error("tts backend not configured")); 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")); span.set_status(Status::error("voice_file is required"));
return HttpResponse::BadRequest().json(json!({ "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("tts.voice_name", name.clone());
span.set_attribute(KeyValue::new("tts.upload_bytes", file_bytes.len() as i64)); span.set_attribute("tts.upload_bytes", file_bytes.len() as i64);
// Normalize to WAV so any device format (e.g. .aac / .opus, which Chatterbox // Normalize to WAV so any device format (e.g. .aac / .opus, which Chatterbox
// rejects by extension) is accepted. // rejects by extension) is accepted.
@@ -987,15 +975,13 @@ pub struct CreateVoiceFromLibraryRequest {
/// WAV reference clip (window selected by start/duration_seconds, length /// WAV reference clip (window selected by start/duration_seconds, length
/// capped by LLAMA_SWAP_TTS_REF_SECONDS). /// capped by LLAMA_SWAP_TTS_REF_SECONDS).
#[post("/tts/voices/from-library")] #[post("/tts/voices/from-library")]
#[tracing::instrument(name = "http.tts.voices.from_library", skip_all)]
pub async fn create_voice_from_library_handler( pub async fn create_voice_from_library_handler(
http_request: HttpRequest,
_claims: Claims, _claims: Claims,
req: web::Json<CreateVoiceFromLibraryRequest>, req: web::Json<CreateVoiceFromLibraryRequest>,
app_state: web::Data<AppState>, app_state: web::Data<AppState>,
) -> impl Responder { ) -> impl Responder {
let parent_context = extract_context_from_request(&http_request); let span = tracing::Span::current();
let mut span =
global_tracer().start_with_context("http.tts.voices.from_library", &parent_context);
let Some(client) = app_state.llamacpp.as_ref() else { let Some(client) = app_state.llamacpp.as_ref() else {
span.set_status(Status::error("tts backend not configured")); span.set_status(Status::error("tts backend not configured"));
@@ -1045,7 +1031,7 @@ pub async fn create_voice_from_library_handler(
return HttpResponse::BadRequest() return HttpResponse::BadRequest()
.json(json!({ "error": "file is not an audio or video file" })); .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 { let wav = match prepare_reference_audio(&abs, ref_start, ref_duration).await {
Ok(b) => b, Ok(b) => b,
+2 -6
View File
@@ -220,9 +220,7 @@ async fn main() -> Result<()> {
println!("Generating summary..."); println!("Generating summary...");
let summary = ollama let summary = ollama.generate(&prompt, Some(system_prompt)).await?;
.generate(&opentelemetry::Context::new(), &prompt, Some(system_prompt))
.await?;
println!("\n📝 GENERATED SUMMARY:"); println!("\n📝 GENERATED SUMMARY:");
println!("─────────────────────────────────────────"); println!("─────────────────────────────────────────");
@@ -238,9 +236,7 @@ async fn main() -> Result<()> {
); );
println!("\nGenerating embedding..."); println!("\nGenerating embedding...");
let embedding = ollama let embedding = ollama.generate_embedding(&stripped).await?;
.generate_embedding(&opentelemetry::Context::new(), &stripped)
.await?;
println!("✓ Embedding generated ({} dimensions)", embedding.len()); println!("✓ Embedding generated ({} dimensions)", embedding.len());
println!("Saving to database..."); println!("Saving to database...");
+1 -2
View File
@@ -859,8 +859,7 @@ async fn synthesize_merge<D: KnowledgeDao + 'static>(
let ollama = app_state.ollama.clone(); let ollama = app_state.ollama.clone();
let model_used = ollama.primary_model.clone(); let model_used = ollama.primary_model.clone();
let cx = opentelemetry::Context::current(); let proposed = match ollama.generate(&prompt, Some(system)).await {
let proposed = match ollama.generate(&cx, &prompt, Some(system)).await {
Ok(out) => { Ok(out) => {
// Strip the framing models reach for even with explicit // Strip the framing models reach for even with explicit
// "no preamble" guidance: leading "Merged description:" // "no preamble" guidance: leading "Merged description:"
+29 -29
View File
@@ -2,17 +2,15 @@ use actix_web::HttpRequest;
use actix_web::http::header::HeaderMap; use actix_web::http::header::HeaderMap;
use opentelemetry::global::{BoxedSpan, BoxedTracer}; use opentelemetry::global::{BoxedSpan, BoxedTracer};
use opentelemetry::propagation::TextMapPropagator; use opentelemetry::propagation::TextMapPropagator;
use opentelemetry::trace::{ use opentelemetry::trace::{Span, Status, TraceContextExt, Tracer, TracerProvider};
Span, SpanBuilder, SpanContext, SpanKind, Status, TraceContextExt, Tracer, TracerProvider,
};
use opentelemetry::{Context, KeyValue, global}; use opentelemetry::{Context, KeyValue, global};
use tracing_subscriber::layer::SubscriberExt;
use tracing_subscriber::util::SubscriberInitExt;
use opentelemetry_appender_log::OpenTelemetryLogBridge; use opentelemetry_appender_log::OpenTelemetryLogBridge;
use opentelemetry_otlp::WithExportConfig; use opentelemetry_otlp::WithExportConfig;
use opentelemetry_sdk::Resource; use opentelemetry_sdk::Resource;
use opentelemetry_sdk::logs::{BatchLogProcessor, SdkLoggerProvider}; use opentelemetry_sdk::logs::{BatchLogProcessor, SdkLoggerProvider};
use opentelemetry_sdk::propagation::TraceContextPropagator; use opentelemetry_sdk::propagation::TraceContextPropagator;
use tracing_subscriber::layer::SubscriberExt;
use tracing_subscriber::util::SubscriberInitExt;
pub fn global_tracer() -> BoxedTracer { pub fn global_tracer() -> BoxedTracer {
global::tracer("image-server") global::tracer("image-server")
@@ -120,18 +118,19 @@ impl<'a> opentelemetry::propagation::Extractor for HeaderExtractor<'a> {
/// provider / operation / model attributes) is identical across /// provider / operation / model attributes) is identical across
/// openrouter / llamacpp / ollama. `operation` is one of `chat`, /// openrouter / llamacpp / ollama. `operation` is one of `chat`,
/// `chat_stream`, `embeddings`, `describe_image`, `generate`. /// `chat_stream`, `embeddings`, `describe_image`, `generate`.
pub fn start_llm_request_span( ///
cx: &Context, /// The span parents on the caller's ambient `tracing` span — callers no
provider: &str, /// longer thread an OpenTelemetry `Context`. Code that hands work to a
operation: &str, /// `tokio::spawn` must carry its span across (`.instrument(span)`), or the
model: &str, /// LLM span orphans onto a fresh trace.
) -> BoxedSpan { pub fn start_llm_request_span(provider: &str, operation: &str, model: &str) -> tracing::Span {
let builder = SpanBuilder::from_name("llm.request").with_kind(SpanKind::Client); tracing::info_span!(
let mut span = global_tracer().build_with_context(builder, cx); "llm.request",
span.set_attribute(KeyValue::new("llm.provider", provider.to_string())); otel.kind = "client",
span.set_attribute(KeyValue::new("llm.operation", operation.to_string())); llm.provider = provider,
span.set_attribute(KeyValue::new("llm.model", model.to_string())); llm.operation = operation,
span llm.model = model,
)
} }
/// The OpenTelemetry [`Context`] to parent this request's operation spans /// 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() headers.into_iter().filter(|(_, v)| !v.is_empty()).collect()
} }
/// Attach the W3C traceparent header to an outgoing request so a downstream /// Attach the W3C traceparent header for `span` to an outgoing request so a
/// service parents its own spans onto this trace. No-op when the span context /// downstream service parents its own spans onto this trace. No-op when the
/// is not valid — e.g. debug builds, where the global tracer is a no-op and an /// span carries no valid OTel context — e.g. a span disabled by the
/// all-zero traceparent would produce garbage parent spans. /// `EnvFilter`, where an all-zero traceparent would produce garbage parents.
pub fn with_traceparent( pub fn with_traceparent(
builder: reqwest::RequestBuilder, builder: reqwest::RequestBuilder,
span_context: &SpanContext, span: &tracing::Span,
) -> reqwest::RequestBuilder { ) -> reqwest::RequestBuilder {
use tracing_opentelemetry::OpenTelemetrySpanExt as _;
let mut b = builder; let mut b = builder;
for (k, v) in for (k, v) in traceparent_headers(&span.context()) {
traceparent_headers(&Context::new().with_remote_span_context(span_context.clone()))
{
b = b.header(k, v); b = b.header(k, v);
} }
b b
} }
/// Record the HTTP status code on the span (success or failure path). /// 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 { 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. /// 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 { 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())); s.set_status(Status::error(message.to_string()));
} }
} }
+1 -2
View File
@@ -281,9 +281,8 @@ pub async fn generate_script(
beats: &[PlannedBeat], beats: &[PlannedBeat],
) -> Result<ReelScript> { ) -> Result<ReelScript> {
let (system, user) = build_script_messages(meta, beats); let (system, user) = build_script_messages(meta, beats);
let cx = opentelemetry::Context::current();
let raw = client let raw = client
.generate(&cx, &user, Some(&system), None) .generate(&user, Some(&system), None)
.await .await
.context("LLM script generation failed")?; .context("LLM script generation failed")?;
Ok(parse_script_response(&raw, beats.len())) Ok(parse_script_response(&raw, beats.len()))
+1 -1
View File
@@ -211,7 +211,7 @@ pub async fn unified_search<TagD: TagDao>(
log::info!("unified_search: translating with model={}", backend.model()); log::info!("unified_search: translating with model={}", backend.model());
let today = chrono::Utc::now().date_naive(); 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, Ok(sq) => sq,
Err(e) => { Err(e) => {
log::warn!("unified_search: translate_nl_query failed: {e:?}"); log::warn!("unified_search: translate_nl_query failed: {e:?}");
+181
View File
@@ -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<Mutex<Vec<String>>>) {
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<Value> {
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:?})"
);
}
}