Otel: llm.request client spans + traceparent on LLM calls (#130)

Add OpenTelemetry tracing to every LLM client call so an insight/chat
turn is a single trace from HTTP entry to each downstream inference.

- LlmClient trait: thread a &opentelemetry::Context through all 5
  methods (generate, chat_with_tools, chat_with_tools_stream,
  generate_embeddings, describe_image), matching the DAO convention of
  explicit context threading rather than ambient Context::current().
- otel::start_llm_request_span: single shared helper that every
  provider (openrouter / llamacpp / ollama) funnels through, so the
  span shape is identical: name llm.request, SpanKind::Client,
  llm.provider / llm.operation / llm.model.
- Per-method spans in all three clients with the spec's operation enum
  (chat | chat_stream | embeddings | describe_image | generate); stream
  spans are owned and moved into the async_stream generator so they end
  when the response drains, not when headers arrive.
- llm.tokens.input / llm.tokens.output recorded on non-stream and
  stream completion; llm.cost_usd from usage.cost (OpenRouter);
  http.status_code + error status on failure.
- W3C traceparent injected on outgoing LLM HTTP requests (llama-swap,
  Ollama, OpenRouter) gated on span_context.is_valid() so debug
  no-op-tracer builds don't send an all-zero parent.
- Turn-level spans (ai.chat.turn.execute, ai.persona_chat.turn.execute)
  now carry llm.backend + llm.model; with_entry streaming paths thread
  the turn context (turn_cx) to child describe_image / agentic-loop
  calls instead of the ambient context.

clippy clean; 690 lib + 775 integration + 4 doc tests green.
This commit is contained in:
Cameron Cordes committed 2026-09-28 01:21:46 -04:00
1 parent e84ac82f86
commit 696d9fa8f3
17 files changed
+897 -221

No files matched your search

+10 -2
View File
@@ -360,11 +360,17 @@ async fn generate_and_store_daily_summary(
span.set_attribute(KeyValue::new("date", date.to_string()));
span.set_attribute(KeyValue::new("contact", contact.to_string()));
span.set_attribute(KeyValue::new("message_count", messages.len() as i64));
// Parent context for the child llm.request spans; the span itself stays
// borrowed here for the attribute updates below.
let llm_cx =
opentelemetry::Context::new().with_remote_span_context(span.span_context().clone());
let (prompt, system_prompt) = build_daily_summary_prompt(contact, date, messages);
// Generate summary with LLM
let summary = ollama.generate(&prompt, Some(system_prompt)).await?;
let summary = ollama
.generate(&llm_cx, &prompt, Some(system_prompt))
.await?;
log::debug!(
"Generated summary for {}: {}",
@@ -382,7 +388,9 @@ async fn generate_and_store_daily_summary(
);
// Embed the stripped summary (store original summary in DB)
let embedding = ollama.generate_embedding(&stripped_summary).await?;
let embedding = ollama
.generate_embedding(&llm_cx, &stripped_summary)
.await?;
span.set_attribute(KeyValue::new(
"embedding_dimensions",
+38 -11
View File
@@ -377,7 +377,7 @@ impl InsightChatService {
let (response, prompt_tokens, eval_tokens) = backend
.chat()
.chat_with_tools(messages.clone(), tools.clone())
.chat_with_tools(&loop_cx, messages.clone(), tools.clone())
.await?;
last_prompt_eval_count = prompt_tokens;
last_eval_count = eval_tokens;
@@ -434,7 +434,7 @@ impl InsightChatService {
let synthetic_idx = push_synthetic_final_prompt(&mut messages);
let (final_response, prompt_tokens, eval_tokens) = backend
.chat()
.chat_with_tools(messages.clone(), vec![])
.chat_with_tools(&loop_cx, messages.clone(), vec![])
.await?;
last_prompt_eval_count = prompt_tokens;
last_eval_count = eval_tokens;
@@ -492,6 +492,7 @@ impl InsightChatService {
let title_raw = backend
.chat()
.generate(
&loop_cx,
&title_prompt,
Some(
"You are my long term memory assistant. Use only the information provided. Do not invent details.",
@@ -852,7 +853,7 @@ impl InsightChatService {
span.set_attribute(KeyValue::new("library_id", library_id as i64));
let result = svc
.run_streaming_turn_with_entry(req, entry_clone.clone())
.run_streaming_turn_with_entry(req, entry_clone.clone(), &mut span)
.await;
if let Err(ref e) = result {
span.set_attribute(KeyValue::new("status", "error"));
@@ -882,6 +883,7 @@ impl InsightChatService {
self: Arc<Self>,
req: ChatTurnRequest,
entry: Arc<TurnEntry>,
turn_span: &mut opentelemetry::global::BoxedSpan,
) -> Result<()> {
if req.user_message.trim().is_empty() {
bail!("user_message must not be empty");
@@ -911,11 +913,11 @@ impl InsightChatService {
if req.regenerate || existing_insight.is_none() {
return self
.run_bootstrap_streaming_with_entry(req, normalized, entry)
.run_bootstrap_streaming_with_entry(req, normalized, entry, turn_span)
.await;
}
let insight = existing_insight.expect("just checked Some above");
self.run_continuation_streaming_with_entry(req, normalized, insight, entry)
self.run_continuation_streaming_with_entry(req, normalized, insight, entry, turn_span)
.await
}
@@ -926,6 +928,7 @@ impl InsightChatService {
normalized: String,
insight: crate::database::models::PhotoInsight,
entry: Arc<TurnEntry>,
turn_span: &mut opentelemetry::global::BoxedSpan,
) -> Result<()> {
let active_persona = req
.persona_id
@@ -981,6 +984,12 @@ impl InsightChatService {
};
let backend = self.generator.resolve_backend(kind, &overrides).await?;
let model_used = backend.model().to_string();
// llm.* attributes on the turn-level span (issue #130): clients can't
// know local-vs-hybrid, but the turn knows its resolved backend.
turn_span.set_attribute(KeyValue::new("llm.backend", kind.as_str()));
turn_span.set_attribute(KeyValue::new("llm.model", model_used.clone()));
let turn_cx = opentelemetry::Context::new()
.with_remote_span_context(turn_span.span_context().clone());
let local_first_user_has_image = messages
.iter()
@@ -1030,6 +1039,7 @@ impl InsightChatService {
&active_persona,
max_iterations,
&entry,
&turn_cx,
)
.await?;
let AgenticLoopOutcome {
@@ -1146,6 +1156,7 @@ impl InsightChatService {
req: ChatTurnRequest,
normalized: String,
entry: Arc<TurnEntry>,
turn_span: &mut opentelemetry::global::BoxedSpan,
) -> Result<()> {
let active_persona = req
.persona_id
@@ -1171,6 +1182,11 @@ impl InsightChatService {
};
let backend = self.generator.resolve_backend(kind, &overrides).await?;
let model_used = backend.model().to_string();
// llm.* attributes on the turn-level span (issue #130).
turn_span.set_attribute(KeyValue::new("llm.backend", kind.as_str()));
turn_span.set_attribute(KeyValue::new("llm.model", model_used.clone()));
let turn_cx = opentelemetry::Context::new()
.with_remote_span_context(turn_span.span_context().clone());
let image_base64: Option<String> = self.generator.load_image_as_base64(&normalized).ok();
@@ -1190,7 +1206,7 @@ impl InsightChatService {
let visual_block = if !backend.images_inline {
match image_base64.as_deref() {
Some(b64) => match backend.local().describe_image(b64).await {
Some(b64) => match backend.local().describe_image(&turn_cx, b64).await {
Ok(desc) => {
format!("Visual description (from local vision model):\n{}\n", desc)
}
@@ -1245,6 +1261,7 @@ impl InsightChatService {
&active_persona,
max_iterations,
&entry,
&turn_cx,
)
.await?;
let AgenticLoopOutcome {
@@ -1333,6 +1350,7 @@ impl InsightChatService {
active_persona: &str,
max_iterations: usize,
entry: &Arc<TurnEntry>,
cx: &opentelemetry::Context,
) -> Result<AgenticLoopOutcome> {
crate::ai::insight_chat::run_streaming_agentic_loop_with_entry(
&self.generator,
@@ -1345,6 +1363,7 @@ impl InsightChatService {
active_persona,
max_iterations,
entry,
cx,
)
.await
}
@@ -1500,6 +1519,7 @@ impl InsightChatService {
&active_persona,
max_iterations,
&tx,
&opentelemetry::Context::current(),
)
.await?;
let AgenticLoopOutcome {
@@ -1667,7 +1687,11 @@ impl InsightChatService {
// images_inline backends send images directly to the chat model.
let visual_block = if !backend.images_inline {
match image_base64.as_deref() {
Some(b64) => match backend.local().describe_image(b64).await {
Some(b64) => match backend
.local()
.describe_image(&opentelemetry::Context::current(), b64)
.await
{
Ok(desc) => {
format!("Visual description (from local vision model):\n{}\n", desc)
}
@@ -1734,6 +1758,7 @@ impl InsightChatService {
&active_persona,
max_iterations,
&tx,
&opentelemetry::Context::current(),
)
.await?;
let AgenticLoopOutcome {
@@ -1816,6 +1841,7 @@ impl InsightChatService {
active_persona: &str,
max_iterations: usize,
tx: &tokio::sync::mpsc::Sender<ChatStreamEvent>,
cx: &opentelemetry::Context,
) -> Result<AgenticLoopOutcome> {
let mut tool_calls_made = 0usize;
let mut iterations_used = 0usize;
@@ -1834,7 +1860,7 @@ impl InsightChatService {
let mut stream = backend
.chat()
.chat_with_tools_stream(messages.clone(), tools.clone())
.chat_with_tools_stream(cx, messages.clone(), tools.clone())
.await?;
let mut final_message: Option<ChatMessage> = None;
@@ -1928,7 +1954,7 @@ impl InsightChatService {
let synthetic_idx = push_synthetic_final_prompt(messages);
let mut stream = backend
.chat()
.chat_with_tools_stream(messages.clone(), vec![])
.chat_with_tools_stream(cx, messages.clone(), vec![])
.await?;
let mut final_message: Option<ChatMessage> = None;
while let Some(ev) = stream.next().await {
@@ -2308,6 +2334,7 @@ pub async fn run_streaming_agentic_loop_with_entry(
active_persona: &str,
max_iterations: usize,
entry: &Arc<TurnEntry>,
cx: &opentelemetry::Context,
) -> Result<AgenticLoopOutcome> {
let mut tool_calls_made = 0usize;
let mut iterations_used = 0usize;
@@ -2337,7 +2364,7 @@ pub async fn run_streaming_agentic_loop_with_entry(
let mut stream = backend
.chat()
.chat_with_tools_stream(messages.clone(), tools.clone())
.chat_with_tools_stream(cx, messages.clone(), tools.clone())
.await?;
let mut final_message: Option<ChatMessage> = None;
@@ -2420,7 +2447,7 @@ pub async fn run_streaming_agentic_loop_with_entry(
let synthetic_idx = push_synthetic_final_prompt(messages);
let mut stream = backend
.chat()
.chat_with_tools_stream(messages.clone(), vec![])
.chat_with_tools_stream(cx, messages.clone(), vec![])
.await?;
let mut final_message: Option<ChatMessage> = None;
while let Some(ev) = stream.next().await {
+54 -26
View File
@@ -646,7 +646,8 @@ impl InsightGenerator {
// (`LLM_BACKEND` switch). Must match the backend that populated the
// daily-summary embeddings or similarity search will be garbage.
let query_embedding =
crate::ai::embed_query(&self.ollama, self.llamacpp.as_deref(), &query).await?;
crate::ai::embed_query(&search_cx, &self.ollama, self.llamacpp.as_deref(), &query)
.await?;
// Search for similar daily summaries with time-based weighting
// This prioritizes summaries temporally close to the query date
@@ -712,7 +713,8 @@ impl InsightGenerator {
// Must use the same backend that populated the daily-summary
// embeddings or similarity search is garbage (see embed_one docs).
let query_embedding =
crate::ai::embed_query(&self.ollama, self.llamacpp.as_deref(), query).await?;
crate::ai::embed_query(&search_cx, &self.ollama, self.llamacpp.as_deref(), query)
.await?;
let mut summary_dao = self
.daily_summary_dao
@@ -798,7 +800,9 @@ impl InsightGenerator {
let calendar_cx = parent_cx.with_span(span);
let query_embedding = if let Some(loc) = location {
match crate::ai::embed_query(&self.ollama, self.llamacpp.as_deref(), loc).await {
match crate::ai::embed_query(&calendar_cx, &self.ollama, self.llamacpp.as_deref(), loc)
.await
{
Ok(emb) => Some(emb),
Err(e) => {
log::warn!("Failed to generate embedding for location '{}': {}", loc, e);
@@ -969,8 +973,13 @@ impl InsightGenerator {
)
};
let query_embedding =
match crate::ai::embed_query(&self.ollama, self.llamacpp.as_deref(), &query_text).await
let query_embedding = match crate::ai::embed_query(
&search_cx,
&self.ollama,
self.llamacpp.as_deref(),
&query_text,
)
.await
{
Ok(emb) => emb,
Err(e) => {
@@ -1254,7 +1263,10 @@ impl InsightGenerator {
// Generate brief photo description for RAG enrichment (vision models only)
let photo_description: Option<String> = if let Some(ref img_b64) = image_base64 {
match ollama_client.generate_photo_description(img_b64).await {
match ollama_client
.generate_photo_description(&insight_cx, img_b64)
.await
{
Ok(desc) => {
log::info!("Photo description for RAG enrichment: {}", desc);
Some(desc)
@@ -1323,7 +1335,7 @@ impl InsightGenerator {
if !immediate_messages.is_empty() {
// Step 2: Extract topics from immediate messages to enrich RAG query
let topics = self
.extract_topics_from_messages(&immediate_messages, &ollama_client)
.extract_topics_from_messages(&immediate_messages, &ollama_client, &insight_cx)
.await;
log::info!("Extracted topics for query enrichment: {:?}", topics);
@@ -1379,6 +1391,7 @@ impl InsightGenerator {
&immediate_messages,
&ollama_client,
custom_system_prompt.as_deref(),
&insight_cx,
)
.await
.unwrap_or_else(|| String::from("No immediate context"));
@@ -1388,6 +1401,7 @@ impl InsightGenerator {
&historical_messages,
&ollama_client,
custom_system_prompt.as_deref(),
&insight_cx,
)
.await
.unwrap_or_else(|| String::from("No historical context"));
@@ -1406,6 +1420,7 @@ impl InsightGenerator {
&immediate_messages,
&ollama_client,
custom_system_prompt.as_deref(),
&insight_cx,
)
.await;
}
@@ -1416,6 +1431,7 @@ impl InsightGenerator {
&immediate_messages,
&ollama_client,
custom_system_prompt.as_deref(),
&insight_cx,
)
.await;
}
@@ -1441,6 +1457,7 @@ impl InsightGenerator {
&rag_messages,
&ollama_client,
custom_system_prompt.as_deref(),
&insight_cx,
)
.await;
}
@@ -1473,7 +1490,7 @@ impl InsightGenerator {
if !sms_messages.is_empty() {
match self
.sms_client
.summarize_context(&sms_messages, &ollama_client)
.summarize_context(&sms_messages, &ollama_client, &insight_cx)
.await
{
Ok(summary) => {
@@ -1557,6 +1574,7 @@ impl InsightGenerator {
// 10. Generate summary first, then derive title from the summary
let summary = ollama_client
.generate_photo_summary(
&insight_cx,
date_taken,
location.as_deref(),
contact.as_deref(),
@@ -1567,7 +1585,7 @@ impl InsightGenerator {
.await?;
let title = ollama_client
.generate_photo_title(&summary, custom_system_prompt.as_deref())
.generate_photo_title(&insight_cx, &summary, custom_system_prompt.as_deref())
.await?;
log::info!("Generated title: {}", title);
@@ -1634,6 +1652,7 @@ impl InsightGenerator {
&self,
messages: &[crate::ai::SmsMessage],
ollama: &OllamaClient,
cx: &opentelemetry::Context,
) -> Vec<String> {
if messages.is_empty() {
return Vec::new();
@@ -1669,7 +1688,7 @@ Return ONLY the comma-separated list, nothing else."#,
);
match ollama
.generate(&prompt, Some("You are an entity extraction assistant. Extract proper nouns, people, places, and domain-specific terms that provide context."))
.generate(cx, &prompt, Some("You are an entity extraction assistant. Extract proper nouns, people, places, and domain-specific terms that provide context."))
.await
{
Ok(response) => {
@@ -1704,6 +1723,7 @@ Return ONLY the comma-separated list, nothing else."#,
messages: &[String],
ollama: &OllamaClient,
custom_system: Option<&str>,
cx: &opentelemetry::Context,
) -> Option<String> {
if messages.is_empty() {
return None;
@@ -1724,7 +1744,7 @@ Return ONLY the summary, nothing else."#,
let system = custom_system
.unwrap_or("You are a context summarization assistant. Be concise and factual.");
match ollama.generate(&prompt, Some(system)).await {
match ollama.generate(cx, &prompt, Some(system)).await {
Ok(summary) => Some(summary),
Err(e) => {
log::warn!("Failed to summarize messages: {}", e);
@@ -1740,6 +1760,7 @@ Return ONLY the summary, nothing else."#,
messages: &[crate::ai::SmsMessage],
ollama: &OllamaClient,
custom_system: Option<&str>,
cx: &opentelemetry::Context,
) -> Option<String> {
if messages.is_empty() {
return None;
@@ -1785,7 +1806,7 @@ Return ONLY the summary, nothing else."#,
"You are a context summarization assistant. Be detailed and factual, preserving important context.",
);
match ollama.generate(&prompt, Some(system)).await {
match ollama.generate(cx, &prompt, Some(system)).await {
Ok(summary) => Some(summary),
Err(e) => {
log::warn!("Failed to summarize immediate context: {}", e);
@@ -1829,7 +1850,7 @@ Return ONLY the summary, nothing else."#,
"get_file_tags" => self.tool_get_file_tags(arguments, cx).await,
"get_faces_in_photo" => self.tool_get_faces_in_photo(arguments, cx).await,
"describe_photo" => {
self.tool_describe_photo(backend.local(), image_base64)
self.tool_describe_photo(backend.local(), image_base64, cx)
.await
}
"reverse_geocode" => self.tool_reverse_geocode(arguments).await,
@@ -1880,7 +1901,7 @@ Return ONLY the summary, nothing else."#,
&self,
args: &serde_json::Value,
local: &dyn LlmClient,
_cx: &opentelemetry::Context,
cx: &opentelemetry::Context,
) -> String {
let query = match args.get("query").and_then(|v| v.as_str()) {
Some(q) => q.to_string(),
@@ -1944,7 +1965,10 @@ Return ONLY the summary, nothing else."#,
};
let final_results = if rerank_enabled && results.len() > limit {
match self.rerank_with_llm(&query, &results, limit, local).await {
match self
.rerank_with_llm(cx, &query, &results, limit, local)
.await
{
Ok(reordered) => reordered,
Err(e) => {
log::warn!("rerank failed, using vector order: {}", e);
@@ -1967,6 +1991,7 @@ Return ONLY the summary, nothing else."#,
/// tool_result. On parse failure we fall back to the input order.
async fn rerank_with_llm(
&self,
cx: &opentelemetry::Context,
query: &str,
candidates: &[String],
limit: usize,
@@ -2010,7 +2035,7 @@ Return ONLY the summary, nothing else."#,
let started = std::time::Instant::now();
let system =
Some("You are a terse relevance ranker. You output only numbers separated by commas.");
let response = local.generate(&prompt, system, None).await?;
let response = local.generate(cx, &prompt, system, None).await?;
log::info!(
"rerank: finished in {} ms (prompt={} chars)",
started.elapsed().as_millis(),
@@ -2620,10 +2645,11 @@ Return ONLY the summary, nothing else."#,
&self,
local: &dyn LlmClient,
image_base64: &Option<String>,
cx: &opentelemetry::Context,
) -> String {
log::info!("tool_describe_photo: generating visual description");
match image_base64 {
Some(img) => match local.describe_image(img).await {
Some(img) => match local.describe_image(cx, img).await {
Ok(desc) => desc,
Err(e) => format!("Error describing photo: {}", e),
},
@@ -3031,6 +3057,7 @@ Return ONLY the summary, nothing else."#,
// configured local backend.
let embed_text = format!("{} {}", name, description);
let embedding: Option<Vec<u8>> = match crate::ai::embed_document(
cx,
&self.ollama,
self.llamacpp.as_deref(),
&embed_text,
@@ -4277,7 +4304,7 @@ Return ONLY the summary, nothing else."#,
// the local backend so non-text work stays off OpenRouter.
let inlined_visual_description: Option<String> = if !backend.images_inline {
match image_base64.as_deref() {
Some(b64) => match backend.local().describe_image(b64).await {
Some(b64) => match backend.local().describe_image(&insight_cx, b64).await {
Ok(desc) => {
log::info!("{}: vision describe succeeded ({} chars)", kind, desc.len());
Some(desc)
@@ -4397,7 +4424,7 @@ Return ONLY the summary, nothing else."#,
let (response, prompt_tokens, eval_tokens) = backend
.chat()
.chat_with_tools(messages.clone(), tools.clone())
.chat_with_tools(&loop_cx, messages.clone(), tools.clone())
.await?;
last_prompt_eval_count = prompt_tokens;
@@ -4466,7 +4493,7 @@ Return ONLY the summary, nothing else."#,
)));
let (final_response, prompt_tokens, eval_tokens) = backend
.chat()
.chat_with_tools(messages.clone(), tools.clone())
.chat_with_tools(&loop_cx, messages.clone(), tools.clone())
.await?;
last_prompt_eval_count = prompt_tokens;
last_eval_count = eval_tokens;
@@ -4499,7 +4526,7 @@ Return ONLY the summary, nothing else."#,
}
let (resp2, pt2, et2) = backend
.chat()
.chat_with_tools(messages.clone(), vec![])
.chat_with_tools(&loop_cx, messages.clone(), vec![])
.await?;
last_prompt_eval_count = pt2;
last_eval_count = et2;
@@ -4637,6 +4664,7 @@ Return ONLY the summary, nothing else."#,
tools: Vec<Tool>,
max_iter: usize,
) -> Result<String> {
let loop_cx = opentelemetry::Context::current();
let mut final_content = String::new();
for iteration in 0..max_iter {
@@ -4644,7 +4672,7 @@ Return ONLY the summary, nothing else."#,
let (response, _prompt_tokens, _eval_tokens) = backend
.chat()
.chat_with_tools(messages.clone(), tools.clone())
.chat_with_tools(&loop_cx, messages.clone(), tools.clone())
.await?;
// Sanitize tool call arguments before pushing back into history.
@@ -4686,7 +4714,7 @@ Return ONLY the summary, nothing else."#,
"",
0,
"",
&opentelemetry::Context::new(),
&loop_cx,
)
.await;
messages.push(ChatMessage::tool_result(result));
@@ -4710,7 +4738,7 @@ Return ONLY the summary, nothing else."#,
));
let (final_response, _, _) = backend
.chat()
.chat_with_tools(messages.clone(), tools.clone())
.chat_with_tools(&loop_cx, messages.clone(), tools.clone())
.await?;
// Same bounded tool round as the insight loop: execute any last
// tool call, then force a tool-less turn for the actual answer.
@@ -4728,14 +4756,14 @@ Return ONLY the summary, nothing else."#,
"",
0,
"",
&opentelemetry::Context::new(),
&loop_cx,
)
.await;
messages.push(ChatMessage::tool_result(result));
}
let (resp2, _, _) = backend
.chat()
.chat_with_tools(messages.clone(), vec![])
.chat_with_tools(&loop_cx, messages.clone(), vec![])
.await?;
final_content = resp2.content;
} else {
+132 -25
View File
@@ -20,8 +20,11 @@
// insight_generator / insight_chat.
#![allow(dead_code)]
use anyhow::{Context, Result, anyhow, bail};
use anyhow::Context as _;
use anyhow::{Result, anyhow, bail};
use async_trait::async_trait;
use opentelemetry::trace::{Span, Status, TraceContextExt};
use opentelemetry::{Context, KeyValue};
use reqwest::Client;
use serde::Deserialize;
use serde_json::{Value, json};
@@ -486,6 +489,7 @@ impl LlamaCppClient {
async fn chat_completion_with_model(
&self,
model: &str,
mut span: Option<&mut opentelemetry::global::BoxedSpan>,
messages: Vec<ChatMessage>,
tools: Vec<Tool>,
) -> Result<(ChatMessage, Option<i32>, Option<i32>)> {
@@ -517,19 +521,41 @@ impl LlamaCppClient {
// Wait for any TTS synthesis to release the GPU before the request
// timeout starts (see ai::gpu).
let _gpu = crate::ai::gpu::llm_lease().await;
let resp = self
.client
.post(&url)
.json(&Value::Object(body))
.send()
.await
.with_context(|| format!("POST {} failed", url))?;
let mut request = self.client.post(&url).json(&Value::Object(body));
if let Some(s) = span.as_mut() {
for (k, v) in crate::otel::traceparent_headers(
&Context::new().with_remote_span_context(s.span_context().clone()),
) {
request = request.header(k, v);
}
}
let resp = match request.send().await {
Ok(r) => r,
Err(e) => {
if let Some(s) = span.as_mut() {
s.set_status(Status::error("llama-swap request failed"));
}
return Err(e).with_context(|| format!("POST {} failed", url));
}
};
if !resp.status().is_success() {
let status = resp.status();
if let Some(s) = span.as_mut() {
s.set_attribute(KeyValue::new("http.status_code", status.as_u16() as i64));
s.set_status(Status::error(format!(
"llama-swap chat request failed: {status}"
)));
}
let body = resp.text().await.unwrap_or_default();
bail!("llama-swap chat request failed: {} — {}", status, body);
}
if let Some(s) = span.as_mut() {
s.set_attribute(KeyValue::new(
"http.status_code",
resp.status().as_u16() as i64,
));
}
let parsed: Value = resp.json().await.context("parsing chat response")?;
let choice = parsed
@@ -562,14 +588,35 @@ impl LlamaCppClient {
log_timings(&parsed, prompt_tokens, completion_tokens);
if let Some(s) = span.as_mut() {
if let Some(pt) = prompt_tokens {
s.set_attribute(KeyValue::new("llm.tokens.input", pt as i64));
}
if let Some(ct) = completion_tokens {
s.set_attribute(KeyValue::new("llm.tokens.output", ct as i64));
}
s.set_status(Status::Ok);
}
Ok((chat_msg, prompt_tokens, completion_tokens))
}
/// Start the `llm.request` span for an outgoing llama-swap call.
fn start_request_span(
&self,
cx: &Context,
operation: &str,
model: &str,
) -> opentelemetry::global::BoxedSpan {
crate::otel::start_llm_request_span(cx, "llamacpp", operation, model)
}
}
#[async_trait]
impl LlmClient for LlamaCppClient {
async fn generate(
&self,
cx: &Context,
prompt: &str,
system: Option<&str>,
images: Option<Vec<String>>,
@@ -582,12 +629,15 @@ impl LlmClient for LlamaCppClient {
user.images = images;
messages.push(user);
let (reply, _, _) = self.chat_with_tools(messages, Vec::new()).await?;
// Delegates to chat_with_tools, which starts the llm.request span
// (operation=chat) — one span per HTTP request.
let (reply, _, _) = self.chat_with_tools(cx, messages, Vec::new()).await?;
Ok(reply.content)
}
async fn chat_with_tools(
&self,
cx: &Context,
messages: Vec<ChatMessage>,
tools: Vec<Tool>,
) -> Result<(ChatMessage, Option<i32>, Option<i32>)> {
@@ -597,15 +647,25 @@ impl LlmClient for LlamaCppClient {
messages.len(),
tools.len()
);
self.chat_completion_with_model(&self.primary_model.clone(), messages, tools)
let mut span = self.start_request_span(cx, "chat", &self.primary_model);
self.chat_completion_with_model(
&self.primary_model.clone(),
Some(&mut span),
messages,
tools,
)
.await
}
async fn chat_with_tools_stream(
&self,
cx: &Context,
messages: Vec<ChatMessage>,
tools: Vec<Tool>,
) -> Result<BoxStream<'static, Result<LlmStreamEvent>>> {
// The span outlives this method — it moves into the stream
// generator and is finalized when the response drains.
let mut span = self.start_request_span(cx, "chat_stream", &self.primary_model);
let url = format!("{}/chat/completions", self.base_url);
let roles: Vec<&str> = messages.iter().map(|m| m.role.as_str()).collect();
log::debug!(
@@ -639,19 +699,30 @@ impl LlmClient for LlamaCppClient {
// timeout starts (see ai::gpu). The guard is moved into the stream
// below so the lease spans the whole generation, not just the send.
let gpu = crate::ai::gpu::llm_lease().await;
let resp = self
.client
.post(&url)
.json(&Value::Object(body))
let mut request = self.client.post(&url).json(&Value::Object(body));
for (k, v) in crate::otel::traceparent_headers(
&Context::new().with_remote_span_context(span.span_context().clone()),
) {
request = request.header(k, v);
}
let resp = request
.send()
.await
.with_context(|| format!("POST {} failed", url))?;
if !resp.status().is_success() {
let status = resp.status();
span.set_attribute(KeyValue::new("http.status_code", status.as_u16() as i64));
span.set_status(Status::error(format!(
"llama-swap stream request failed: {status}"
)));
let body = resp.text().await.unwrap_or_default();
bail!("llama-swap stream request failed: {} — {}", status, body);
}
span.set_attribute(KeyValue::new(
"http.status_code",
resp.status().as_u16() as i64,
));
let byte_stream = resp.bytes_stream();
let stream = async_stream::stream! {
@@ -760,6 +831,13 @@ impl LlmClient for LlamaCppClient {
tool_calls,
images: None,
};
if let Some(pt) = prompt_tokens {
span.set_attribute(KeyValue::new("llm.tokens.input", pt as i64));
}
if let Some(ct) = completion_tokens {
span.set_attribute(KeyValue::new("llm.tokens.output", ct as i64));
}
span.set_status(Status::Ok);
yield Ok(LlmStreamEvent::Done {
message,
prompt_eval_count: prompt_tokens,
@@ -770,7 +848,8 @@ impl LlmClient for LlamaCppClient {
Ok(Box::pin(stream))
}
async fn generate_embeddings(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>> {
async fn generate_embeddings(&self, cx: &Context, texts: &[&str]) -> Result<Vec<Vec<f32>>> {
let mut span = self.start_request_span(cx, "embeddings", &self.embedding_model);
// Deliberately NO GPU lease: the embed slot sits in llama-swap's
// always-resident group and never participates in a model swap, so
// leasing here would only stall searches behind a queued synthesis.
@@ -780,19 +859,33 @@ impl LlmClient for LlamaCppClient {
"input": texts,
});
let resp = self
.client
.post(&url)
.json(&body)
.send()
.await
.with_context(|| format!("POST {} failed", url))?;
let mut request = self.client.post(&url).json(&body);
for (k, v) in crate::otel::traceparent_headers(
&Context::new().with_remote_span_context(span.span_context().clone()),
) {
request = request.header(k, v);
}
let resp = match request.send().await {
Ok(r) => r,
Err(e) => {
span.set_status(Status::error("llama-swap embedding request failed"));
return Err(e).with_context(|| format!("POST {} failed", url));
}
};
if !resp.status().is_success() {
let status = resp.status();
span.set_attribute(KeyValue::new("http.status_code", status.as_u16() as i64));
span.set_status(Status::error(format!(
"llama-swap embedding request failed: {status}"
)));
let body = resp.text().await.unwrap_or_default();
bail!("llama-swap embedding request failed: {} — {}", status, body);
}
span.set_attribute(KeyValue::new(
"http.status_code",
resp.status().as_u16() as i64,
));
#[derive(Deserialize)]
struct EmbedResponse {
@@ -803,11 +896,20 @@ impl LlmClient for LlamaCppClient {
embedding: Vec<f32>,
}
let parsed: EmbedResponse = resp.json().await.context("parsing embed response")?;
let parsed: EmbedResponse = match resp.json().await {
Ok(p) => p,
Err(e) => {
span.set_status(Status::error("parsing embed response failed"));
return Err(e).context("parsing embed response");
}
};
span.set_status(Status::Ok);
Ok(parsed.data.into_iter().map(|i| i.embedding).collect())
}
async fn describe_image(&self, image_base64: &str) -> Result<String> {
async fn describe_image(&self, cx: &Context, image_base64: &str) -> Result<String> {
// 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 prompt = "Briefly describe what you see in this image in 1-2 sentences. \
Focus on the people, location, and activity.";
let system = "You are a scene description assistant. Be concise and factual.";
@@ -823,7 +925,12 @@ impl LlmClient for LlamaCppClient {
];
let (reply, _, _) = self
.chat_completion_with_model(&self.vision_model.clone(), messages, Vec::new())
.chat_completion_with_model(
&self.vision_model.clone(),
Some(&mut span),
messages,
Vec::new(),
)
.await?;
Ok(reply.content)
}
+12 -2
View File
@@ -1,6 +1,7 @@
use anyhow::Result;
use async_trait::async_trait;
use futures::stream::BoxStream;
use opentelemetry::Context;
use serde::{Deserialize, Serialize};
/// Provider-agnostic surface for LLM backends (Ollama, OpenRouter, …).
@@ -10,6 +11,12 @@ use serde::{Deserialize, Serialize};
/// when a provider requires it (OpenAI-compatible APIs do), and `images`
/// stays as base64 strings here and is rewritten into content-parts where
/// needed.
///
/// Every method that performs an HTTP call takes the caller's tracing
/// `Context` as its first parameter so the impl can start an `llm.request`
/// span parented under the caller's feature span — the codebase threads
/// contexts explicitly rather than relying on ambient `Context::current()`
/// at deep call sites.
// First consumer lands in a later PR (OpenRouter impl + hybrid mode routing).
#[allow(dead_code)]
#[async_trait]
@@ -18,6 +25,7 @@ pub trait LlmClient: Send + Sync {
/// base64 images (ignored by providers without vision support).
async fn generate(
&self,
cx: &Context,
prompt: &str,
system: Option<&str>,
images: Option<Vec<String>>,
@@ -27,6 +35,7 @@ pub trait LlmClient: Send + Sync {
/// (which may contain tool_calls) plus optional prompt/eval token counts.
async fn chat_with_tools(
&self,
cx: &Context,
messages: Vec<ChatMessage>,
tools: Vec<Tool>,
) -> Result<(ChatMessage, Option<i32>, Option<i32>)>;
@@ -39,16 +48,17 @@ pub trait LlmClient: Send + Sync {
/// as one `Done` event.
async fn chat_with_tools_stream(
&self,
cx: &Context,
messages: Vec<ChatMessage>,
tools: Vec<Tool>,
) -> Result<BoxStream<'static, Result<LlmStreamEvent>>>;
/// Batch embedding generation. Dimensionality is provider/model specific.
async fn generate_embeddings(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>>;
async fn generate_embeddings(&self, cx: &Context, texts: &[&str]) -> Result<Vec<Vec<f32>>>;
/// One-shot vision description of an image. Used to convert images into
/// plain text for the hybrid-mode conversation flow.
async fn describe_image(&self, image_base64: &str) -> Result<String>;
async fn describe_image(&self, cx: &Context, image_base64: &str) -> Result<String>;
/// Enumerate available models with their capabilities.
async fn list_models(&self) -> Result<Vec<ModelCapabilities>>;
+25 -4
View File
@@ -47,12 +47,24 @@ impl LocalLlm {
/// pick query vs document — retrieval models treat the two sides
/// differently and an unmarked embed invites prefix-mismatch bugs.
pub async fn embed_query(&self, text: &str) -> Result<Vec<f32>> {
super::embed_query(&self.ollama, self.llamacpp.as_deref(), text).await
super::embed_query(
&opentelemetry::Context::new(),
&self.ollama,
self.llamacpp.as_deref(),
text,
)
.await
}
/// Embed corpus text (applies `EMBED_DOCUMENT_PREFIX`).
pub async fn embed_document(&self, text: &str) -> Result<Vec<f32>> {
super::embed_document(&self.ollama, self.llamacpp.as_deref(), text).await
super::embed_document(
&opentelemetry::Context::new(),
&self.ollama,
self.llamacpp.as_deref(),
text,
)
.await
}
/// Single-shot local text generation via the `LLM_BACKEND`-selected
@@ -60,14 +72,23 @@ impl LocalLlm {
pub async fn generate(&self, prompt: &str, system: Option<&str>) -> Result<String> {
if super::local_backend_is_llamacpp() {
if let Some(lc) = self.llamacpp.as_deref() {
return <LlamaCppClient as LlmClient>::generate(lc, prompt, system, None).await;
return <LlamaCppClient as LlmClient>::generate(
lc,
&opentelemetry::Context::new(),
prompt,
system,
None,
)
.await;
}
anyhow::bail!(
"LLM_BACKEND=llamacpp but LlamaCppClient is unconfigured — \
set LLAMA_SWAP_URL or switch to LLM_BACKEND=ollama"
);
}
self.ollama.generate(prompt, system).await
self.ollama
.generate(&opentelemetry::Context::new(), prompt, system)
.await
}
/// Label identifying which backend + model produces embeddings right
+9 -5
View File
@@ -115,24 +115,26 @@ fn embed_prefix(key: &str) -> String {
/// `Instruct: <task>\nQuery: `. Must pair with the document prefix the
/// corpus was embedded with or similarity degrades.
pub async fn embed_query(
cx: &opentelemetry::Context,
ollama: &OllamaClient,
llamacpp: Option<&LlamaCppClient>,
text: &str,
) -> anyhow::Result<Vec<f32>> {
let prefixed = format!("{}{}", embed_prefix("EMBED_QUERY_PREFIX"), text);
embed_one(ollama, llamacpp, &prefixed).await
embed_one(cx, ollama, llamacpp, &prefixed).await
}
/// Embed corpus text (the stored side of retrieval). Applies
/// `EMBED_DOCUMENT_PREFIX` (default empty; nomic v1.5 wants
/// `search_document: `, Qwen3-Embedding wants none).
pub async fn embed_document(
cx: &opentelemetry::Context,
ollama: &OllamaClient,
llamacpp: Option<&LlamaCppClient>,
text: &str,
) -> anyhow::Result<Vec<f32>> {
let prefixed = format!("{}{}", embed_prefix("EMBED_DOCUMENT_PREFIX"), text);
embed_one(ollama, llamacpp, &prefixed).await
embed_one(cx, ollama, llamacpp, &prefixed).await
}
/// Embed a batch of strings via the configured local backend. Routes
@@ -140,30 +142,32 @@ pub async fn embed_document(
/// configured), else Ollama. See [`local_backend_is_llamacpp`] for the
/// rationale on consistency.
pub async fn embed_many(
cx: &opentelemetry::Context,
ollama: &OllamaClient,
llamacpp: Option<&LlamaCppClient>,
texts: &[&str],
) -> anyhow::Result<Vec<Vec<f32>>> {
if local_backend_is_llamacpp() {
if let Some(lc) = llamacpp {
return <LlamaCppClient as LlmClient>::generate_embeddings(lc, texts).await;
return <LlamaCppClient as LlmClient>::generate_embeddings(lc, cx, texts).await;
}
anyhow::bail!(
"LLM_BACKEND=llamacpp but LlamaCppClient is unconfigured — \
set LLAMA_SWAP_URL or switch to LLM_BACKEND=ollama"
);
}
ollama.generate_embeddings(texts).await
ollama.generate_embeddings(cx, texts).await
}
/// Embed one string via the configured local backend. Single-text
/// convenience over [`embed_many`].
pub async fn embed_one(
cx: &opentelemetry::Context,
ollama: &OllamaClient,
llamacpp: Option<&LlamaCppClient>,
text: &str,
) -> anyhow::Result<Vec<f32>> {
let mut vecs = embed_many(ollama, llamacpp, &[text]).await?;
let mut vecs = embed_many(cx, ollama, llamacpp, &[text]).await?;
vecs.pop()
.ok_or_else(|| anyhow::anyhow!("embedding backend returned no embeddings"))
}
+4 -1
View File
@@ -259,10 +259,13 @@ pub async fn translate_nl_query(
nl: &str,
tag_vocab: &[(i32, String)],
today: chrono::NaiveDate,
cx: &opentelemetry::Context,
) -> Result<StructuredQuery> {
let system = build_system_prompt(tag_vocab, today);
let messages = vec![ChatMessage::system(system), ChatMessage::user(nl)];
let (msg, _, _) = client.chat_with_tools(messages, Vec::<Tool>::new()).await?;
let (msg, _, _) = client
.chat_with_tools(cx, messages, Vec::<Tool>::new())
.await?;
parse_response(&msg.content, tag_vocab)
}
+271 -53
View File
@@ -1,6 +1,9 @@
use anyhow::{Context, Result};
use anyhow::Context as _;
use anyhow::Result;
use async_trait::async_trait;
use chrono::NaiveDate;
use opentelemetry::trace::{Span, Status, TraceContextExt};
use opentelemetry::{Context, KeyValue};
use reqwest::Client;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
@@ -9,6 +12,7 @@ use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};
use crate::ai::llm_client::{LlmClient, LlmStreamEvent};
use crate::otel::traceparent_headers;
use futures::stream::{BoxStream, StreamExt};
// Re-export shared types so existing `crate::ai::ollama::{...}` imports
@@ -365,6 +369,7 @@ impl OllamaClient {
async fn try_generate(
&self,
mut span: Option<&mut opentelemetry::global::BoxedSpan>,
url: &str,
model: &str,
prompt: &str,
@@ -382,15 +387,29 @@ impl OllamaClient {
think,
};
let response = self
let mut request_builder = self
.client
.post(format!("{}/api/generate", url))
.json(&request)
.send()
.await?;
.json(&request);
if let Some(s) = span.as_mut() {
request_builder = self.inject_traceparent(request_builder, s);
}
let response = match request_builder.send().await {
Ok(r) => r,
Err(e) => {
if let Some(s) = span.as_mut() {
s.set_status(Status::error("Ollama request failed"));
}
return Err(e.into());
}
};
if !response.status().is_success() {
let status = response.status();
if let Some(s) = span.as_mut() {
s.set_attribute(KeyValue::new("http.status_code", status.as_u16() as i64));
s.set_status(Status::error(format!("Ollama request failed: {status}")));
}
let error_body = response.text().await.unwrap_or_default();
return Err(anyhow::anyhow!(
"Ollama request failed: {} - {}",
@@ -398,8 +417,31 @@ impl OllamaClient {
error_body
));
}
if let Some(s) = span.as_mut() {
s.set_attribute(KeyValue::new(
"http.status_code",
response.status().as_u16() as i64,
));
}
let result: OllamaResponse = response.json().await?;
let result: OllamaResponse = match response.json().await {
Ok(r) => r,
Err(e) => {
if let Some(s) = span.as_mut() {
s.set_status(Status::error("parsing Ollama response failed"));
}
return Err(e.into());
}
};
if let Some(s) = span.as_mut() {
if let Some(pt) = result.prompt_eval_count {
s.set_attribute(KeyValue::new("llm.tokens.input", pt as i64));
}
if let Some(ct) = result.eval_count {
s.set_attribute(KeyValue::new("llm.tokens.output", ct as i64));
}
s.set_status(Status::Ok);
}
log_chat_metrics(
result.prompt_eval_count,
result.prompt_eval_duration,
@@ -409,32 +451,71 @@ impl OllamaClient {
Ok(result.response)
}
pub async fn generate(&self, prompt: &str, system: Option<&str>) -> Result<String> {
self.generate_with_images(prompt, system, None).await
/// Attach the W3C traceparent header to an outgoing Ollama request so
/// a server that understands it can parent its own spans onto ours.
fn inject_traceparent(
&self,
builder: reqwest::RequestBuilder,
span: &opentelemetry::global::BoxedSpan,
) -> reqwest::RequestBuilder {
let mut builder = builder;
for (k, v) in traceparent_headers(
&Context::new().with_remote_span_context(span.span_context().clone()),
) {
builder = builder.header(k, v);
}
builder
}
/// Start the `llm.request` span for an outgoing Ollama call.
fn start_request_span(
&self,
cx: &Context,
operation: &str,
model: &str,
) -> opentelemetry::global::BoxedSpan {
crate::otel::start_llm_request_span(cx, "ollama", operation, model)
}
pub async fn generate(
&self,
cx: &Context,
prompt: &str,
system: Option<&str>,
) -> Result<String> {
self.generate_with_images(cx, prompt, system, None).await
}
#[allow(dead_code)]
pub async fn generate_no_think(&self, prompt: &str, system: Option<&str>) -> Result<String> {
self.generate_with_options(prompt, system, None, Some(false))
pub async fn generate_no_think(
&self,
cx: &Context,
prompt: &str,
system: Option<&str>,
) -> Result<String> {
self.generate_with_options(cx, prompt, system, None, Some(false), "generate")
.await
}
pub async fn generate_with_images(
&self,
cx: &Context,
prompt: &str,
system: Option<&str>,
images: Option<Vec<String>>,
) -> Result<String> {
self.generate_with_options(prompt, system, images, None)
self.generate_with_options(cx, prompt, system, images, None, "generate")
.await
}
async fn generate_with_options(
&self,
cx: &Context,
prompt: &str,
system: Option<&str>,
images: Option<Vec<String>>,
think: Option<bool>,
operation: &str,
) -> Result<String> {
log::debug!("=== Ollama Request ===");
log::debug!("Primary model: {}", self.primary_model);
@@ -447,6 +528,10 @@ impl OllamaClient {
}
log::debug!("=====================");
// Span for the whole generate attempt; the model attribute
// reflects whichever server actually answered.
let mut span = self.start_request_span(cx, operation, &self.primary_model);
// Try primary server first with primary model
log::info!(
"Attempting to generate with primary server: {} (model: {})",
@@ -455,6 +540,7 @@ impl OllamaClient {
);
let primary_result = self
.try_generate(
Some(&mut span),
&self.primary_url,
&self.primary_model,
prompt,
@@ -483,8 +569,10 @@ impl OllamaClient {
fallback_url,
fallback_model
);
span.set_attribute(KeyValue::new("llm.model", fallback_model.to_string()));
match self
.try_generate(
Some(&mut span),
fallback_url,
fallback_model,
prompt,
@@ -531,6 +619,7 @@ impl OllamaClient {
/// Generate a title for a single photo based on its generated summary
pub async fn generate_photo_title(
&self,
cx: &Context,
summary: &str,
custom_system: Option<&str>,
) -> Result<String> {
@@ -546,7 +635,7 @@ Capture the key moment or theme. Return ONLY the title, nothing else."#,
let system = custom_system.unwrap_or("You are my long term memory assistant. Use only the information provided. Do not invent details.");
let title = self
.generate_with_images(&prompt, Some(system), None)
.generate_with_images(cx, &prompt, Some(system), None)
.await?;
// Models decorate despite "Return ONLY the title": quotes, bold
// markers, sometimes a "Title:" label.
@@ -563,6 +652,7 @@ Capture the key moment or theme. Return ONLY the title, nothing else."#,
/// Generate a summary for a single photo based on its context
pub async fn generate_photo_summary(
&self,
cx: &Context,
date: NaiveDate,
location: Option<&str>,
contact: Option<&str>,
@@ -645,21 +735,32 @@ Analyze the image and use specific details from both the visual content and the
let system = custom_system.unwrap_or("You are a memory refreshing assistant who is able to provide insights through analyzing past conversations. Use only the information provided. Do not invent details.");
let images = image_base64.map(|img| vec![img]);
self.generate_with_images(&prompt, Some(system), images)
self.generate_with_images(cx, &prompt, Some(system), images)
.await
}
/// Generate a brief visual description of a photo for use in RAG query enrichment.
/// Returns 1-2 sentences describing people, location, and activity visible in the image.
/// Only called when the model has vision capabilities.
pub async fn generate_photo_description(&self, image_base64: &str) -> Result<String> {
pub async fn generate_photo_description(
&self,
cx: &Context,
image_base64: &str,
) -> Result<String> {
let prompt = "Briefly describe what you see in this image in 1-2 sentences. \
Focus on the people, location, and activity.";
let system = "You are a scene description assistant. Be concise and factual.";
let images = vec![image_base64.to_string()];
let description = self
.generate_with_images(prompt, Some(system), Some(images))
.generate_with_options(
cx,
prompt,
Some(system),
Some(images),
None,
"describe_image",
)
.await?;
Ok(description.trim().to_string())
@@ -672,9 +773,11 @@ Analyze the image and use specific details from both the visual content and the
/// iteration of the agent loop.
pub async fn chat_with_tools(
&self,
cx: &Context,
messages: Vec<ChatMessage>,
tools: Vec<Tool>,
) -> Result<(ChatMessage, Option<i32>, Option<i32>)> {
let mut span = self.start_request_span(cx, "chat", &self.primary_model);
let order = self.attempt_order();
let mut errors: Vec<String> = Vec::new();
@@ -685,8 +788,11 @@ Analyze the image and use specific details from both the visual content and the
url,
model
);
if *label == "fallback" {
span.set_attribute(KeyValue::new("llm.model", model.to_string()));
}
match self
.try_chat_with_tools(url, messages.clone(), tools.clone())
.try_chat_with_tools(Some(&mut span), url, messages.clone(), tools.clone())
.await
{
Ok(result) => {
@@ -728,18 +834,22 @@ Analyze the image and use specific details from both the visual content and the
/// any, live on the final message).
pub async fn chat_with_tools_stream(
&self,
cx: &Context,
messages: Vec<ChatMessage>,
tools: Vec<Tool>,
) -> Result<BoxStream<'static, Result<LlmStreamEvent>>> {
// Same preference logic as `chat_with_tools`. Only the initial
// connection is retried across servers — once the stream begins,
// mid-stream errors propagate to the caller.
// mid-stream errors propagate to the caller. Each attempt gets its
// own `llm.request` span, created inside
// `try_chat_with_tools_stream` so the winning attempt's span moves
// into the stream and is finalized when the response drains.
let order = self.attempt_order();
let mut last_err: Option<anyhow::Error> = None;
for (label, url, _model) in &order {
match self
.try_chat_with_tools_stream(url, messages.clone(), tools.clone())
.try_chat_with_tools_stream(cx, url, messages.clone(), tools.clone())
.await
{
Ok(s) => {
@@ -759,6 +869,7 @@ Analyze the image and use specific details from both the visual content and the
async fn try_chat_with_tools_stream(
&self,
cx: &Context,
base_url: &str,
messages: Vec<ChatMessage>,
tools: Vec<Tool>,
@@ -773,6 +884,11 @@ Analyze the image and use specific details from both the visual content and the
};
let options = self.build_options();
// The span outlives this method — the stream yields until the
// response fully drains, so the owned span moves into the
// generator and is finalized there.
let mut span = self.start_request_span(cx, "chat_stream", model);
let request_body = OllamaChatRequest {
model,
messages: &messages,
@@ -781,16 +897,25 @@ Analyze the image and use specific details from both the visual content and the
options,
};
let response = self
.client
.post(&url)
.json(&request_body)
.send()
.await
.with_context(|| format!("Failed to connect to Ollama at {}", url))?;
let response = {
let mut rb = self.client.post(&url).json(&request_body);
rb = self.inject_traceparent(rb, &span);
match rb.send().await {
Ok(r) => r,
Err(e) => {
span.set_status(Status::error("Ollama stream connect failed"));
return Err(e)
.with_context(|| format!("Failed to connect to Ollama at {}", url));
}
}
};
if !response.status().is_success() {
let status = response.status();
span.set_attribute(KeyValue::new("http.status_code", status.as_u16() as i64));
span.set_status(Status::error(format!(
"Ollama stream request failed: {status}"
)));
let body = response.text().await.unwrap_or_default();
anyhow::bail!(
"Ollama stream request failed with status {}: {}",
@@ -798,6 +923,10 @@ Analyze the image and use specific details from both the visual content and the
body
);
}
span.set_attribute(KeyValue::new(
"http.status_code",
response.status().as_u16() as i64,
));
// Ollama streams NDJSON: each line is a full `OllamaStreamChunk`.
// We buffer partial lines across chunks from the byte stream.
@@ -879,6 +1008,13 @@ Analyze the image and use specific details from both the visual content and the
eval_count,
eval_duration,
);
if let Some(pt) = prompt_eval_count {
span.set_attribute(KeyValue::new("llm.tokens.input", pt as i64));
}
if let Some(ct) = eval_count {
span.set_attribute(KeyValue::new("llm.tokens.output", ct as i64));
}
span.set_status(Status::Ok);
let message = ChatMessage {
role,
content: accumulated,
@@ -897,6 +1033,7 @@ Analyze the image and use specific details from both the visual content and the
async fn try_chat_with_tools(
&self,
mut span: Option<&mut opentelemetry::global::BoxedSpan>,
base_url: &str,
messages: Vec<ChatMessage>,
tools: Vec<Tool>,
@@ -924,16 +1061,31 @@ Analyze the image and use specific details from both the visual content and the
.unwrap_or_else(|e| format!("<serialization error: {}>", e));
log::debug!("chat_with_tools request body: {}", request_json);
let response = self
.client
.post(&url)
.json(&request_body)
.send()
.await
.with_context(|| format!("Failed to connect to Ollama at {}", url))?;
let response = {
let mut rb = self.client.post(&url).json(&request_body);
if let Some(s) = span.as_mut() {
rb = self.inject_traceparent(rb, s);
}
match rb.send().await {
Ok(r) => r,
Err(e) => {
if let Some(s) = span.as_mut() {
s.set_status(Status::error("Ollama chat connect failed"));
}
return Err(e)
.with_context(|| format!("Failed to connect to Ollama at {}", url));
}
}
};
if !response.status().is_success() {
let status = response.status();
if let Some(s) = span.as_mut() {
s.set_attribute(KeyValue::new("http.status_code", status.as_u16() as i64));
s.set_status(Status::error(format!(
"Ollama chat request failed: {status}"
)));
}
let body = response.text().await.unwrap_or_default();
// warn, not error — the outer `chat_with_tools` may recover via
// the fallback server. When both fail, the outer layer emits the
@@ -950,11 +1102,31 @@ Analyze the image and use specific details from both the visual content and the
body
);
}
if let Some(s) = span.as_mut() {
s.set_attribute(KeyValue::new(
"http.status_code",
response.status().as_u16() as i64,
));
}
let chat_response: OllamaChatResponse = response
.json()
.await
.with_context(|| "Failed to parse Ollama chat response")?;
let chat_response: OllamaChatResponse = match response.json().await {
Ok(r) => r,
Err(e) => {
if let Some(s) = span.as_mut() {
s.set_status(Status::error("parsing Ollama chat response failed"));
}
return Err(e).with_context(|| "Failed to parse Ollama chat response");
}
};
if let Some(s) = span.as_mut() {
if let Some(pt) = chat_response.prompt_eval_count {
s.set_attribute(KeyValue::new("llm.tokens.input", pt as i64));
}
if let Some(ct) = chat_response.eval_count {
s.set_attribute(KeyValue::new("llm.tokens.output", ct as i64));
}
s.set_status(Status::Ok);
}
// Log performance counters returned by Ollama. Durations are
// reported in nanoseconds; we render ms + tokens/sec for skim-ability
@@ -976,8 +1148,8 @@ Analyze the image and use specific details from both the visual content and the
/// Generate an embedding vector for text using nomic-embed-text:v1.5
/// Returns a 768-dimensional vector as Vec<f32>
pub async fn generate_embedding(&self, text: &str) -> Result<Vec<f32>> {
let embeddings = self.generate_embeddings(&[text]).await?;
pub async fn generate_embedding(&self, cx: &Context, text: &str) -> Result<Vec<f32>> {
let embeddings = self.generate_embeddings(cx, &[text]).await?;
embeddings
.into_iter()
.next()
@@ -987,9 +1159,11 @@ Analyze the image and use specific details from both the visual content and the
/// Generate embeddings for multiple texts in a single API call (batch mode)
/// Returns a vector of 768-dimensional vectors
/// This is much more efficient than calling generate_embedding multiple times
pub async fn generate_embeddings(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>> {
pub async fn generate_embeddings(&self, cx: &Context, texts: &[&str]) -> Result<Vec<Vec<f32>>> {
let embedding_model = EMBEDDING_MODEL;
let mut span = self.start_request_span(cx, "embeddings", embedding_model);
log::debug!("=== Ollama Batch Embedding Request ===");
log::debug!("Model: {}", embedding_model);
log::debug!("Batch size: {} texts", texts.len());
@@ -1003,7 +1177,7 @@ Analyze the image and use specific details from both the visual content and the
embedding_model
);
let primary_result = self
.try_generate_embeddings(&self.primary_url, embedding_model, texts)
.try_generate_embeddings(Some(&mut span), &self.primary_url, embedding_model, texts)
.await;
let embeddings = match primary_result {
@@ -1026,7 +1200,12 @@ Analyze the image and use specific details from both the visual content and the
embedding_model
);
match self
.try_generate_embeddings(fallback_url, embedding_model, texts)
.try_generate_embeddings(
Some(&mut span),
fallback_url,
embedding_model,
texts,
)
.await
{
Ok(embeddings) => {
@@ -1073,6 +1252,7 @@ Analyze the image and use specific details from both the visual content and the
/// Internal helper to try generating embeddings for multiple texts from a specific server
async fn try_generate_embeddings(
&self,
mut span: Option<&mut opentelemetry::global::BoxedSpan>,
url: &str,
model: &str,
texts: &[&str],
@@ -1082,15 +1262,33 @@ Analyze the image and use specific details from both the visual content and the
input: texts.iter().map(|s| s.to_string()).collect(),
};
let response = self
let response = {
let mut rb = self
.client
.post(format!("{}/api/embed", url))
.json(&request)
.send()
.await?;
.json(&request);
if let Some(s) = span.as_mut() {
rb = self.inject_traceparent(rb, s);
}
match rb.send().await {
Ok(r) => r,
Err(e) => {
if let Some(s) = span.as_mut() {
s.set_status(Status::error("Ollama embedding request failed"));
}
return Err(e.into());
}
}
};
if !response.status().is_success() {
let status = response.status();
if let Some(s) = span.as_mut() {
s.set_attribute(KeyValue::new("http.status_code", status.as_u16() as i64));
s.set_status(Status::error(format!(
"Ollama batch embedding request failed: {status}"
)));
}
let error_body = response.text().await.unwrap_or_default();
return Err(anyhow::anyhow!(
"Ollama batch embedding request failed: {} - {}",
@@ -1098,8 +1296,25 @@ Analyze the image and use specific details from both the visual content and the
error_body
));
}
if let Some(s) = span.as_mut() {
s.set_attribute(KeyValue::new(
"http.status_code",
response.status().as_u16() as i64,
));
}
let result: OllamaEmbedResponse = response.json().await?;
let result: OllamaEmbedResponse = match response.json().await {
Ok(r) => r,
Err(e) => {
if let Some(s) = span.as_mut() {
s.set_status(Status::error("parsing Ollama embed response failed"));
}
return Err(e.into());
}
};
if let Some(s) = span.as_mut() {
s.set_status(Status::Ok);
}
Ok(result.embeddings)
}
}
@@ -1108,35 +1323,38 @@ Analyze the image and use specific details from both the visual content and the
impl LlmClient for OllamaClient {
async fn generate(
&self,
cx: &Context,
prompt: &str,
system: Option<&str>,
images: Option<Vec<String>>,
) -> Result<String> {
self.generate_with_images(prompt, system, images).await
self.generate_with_images(cx, prompt, system, images).await
}
async fn chat_with_tools(
&self,
cx: &Context,
messages: Vec<ChatMessage>,
tools: Vec<Tool>,
) -> Result<(ChatMessage, Option<i32>, Option<i32>)> {
OllamaClient::chat_with_tools(self, messages, tools).await
OllamaClient::chat_with_tools(self, cx, messages, tools).await
}
async fn chat_with_tools_stream(
&self,
cx: &Context,
messages: Vec<ChatMessage>,
tools: Vec<Tool>,
) -> Result<BoxStream<'static, Result<LlmStreamEvent>>> {
OllamaClient::chat_with_tools_stream(self, messages, tools).await
OllamaClient::chat_with_tools_stream(self, cx, messages, tools).await
}
async fn generate_embeddings(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>> {
OllamaClient::generate_embeddings(self, texts).await
async fn generate_embeddings(&self, cx: &Context, texts: &[&str]) -> Result<Vec<Vec<f32>>> {
OllamaClient::generate_embeddings(self, cx, texts).await
}
async fn describe_image(&self, image_base64: &str) -> Result<String> {
self.generate_photo_description(image_base64).await
async fn describe_image(&self, cx: &Context, image_base64: &str) -> Result<String> {
self.generate_photo_description(cx, image_base64).await
}
async fn list_models(&self) -> Result<Vec<ModelCapabilities>> {
+224 -83
View File
@@ -2,8 +2,11 @@
// the translation helpers directly.
#![allow(dead_code)]
use anyhow::{Context, Result, anyhow, bail};
use anyhow::Context as _;
use anyhow::{Result, anyhow, bail};
use async_trait::async_trait;
use opentelemetry::trace::{Span, Status, TraceContextExt};
use opentelemetry::{Context, KeyValue};
use reqwest::Client;
use serde::Deserialize;
use serde_json::{Value, json};
@@ -14,6 +17,7 @@ use std::time::{Duration, Instant};
use crate::ai::llm_client::{
ChatMessage, LlmClient, LlmStreamEvent, ModelCapabilities, Tool, ToolCall, ToolCallFunction,
};
use futures::stream::{BoxStream, StreamExt};
const DEFAULT_BASE_URL: &str = "https://openrouter.ai/api/v1";
@@ -144,6 +148,19 @@ impl OpenRouterClient {
b
}
/// Start the `llm.request` span for an outgoing call. `operation` is
/// `chat`, `embeddings`, or `describe` — the call shape, not the entry
/// point (a `describe_image` funneled through `chat_with_tools` is
/// still a `describe`).
fn start_request_span(
&self,
cx: &Context,
operation: &str,
model: &str,
) -> opentelemetry::global::BoxedSpan {
crate::otel::start_llm_request_span(cx, "openrouter", operation, model)
}
/// Translate canonical messages to the OpenAI-compatible wire shape.
///
/// Walks in order so it can attach `tool_call_id` to `role=tool` messages
@@ -312,6 +329,7 @@ impl OpenRouterClient {
impl LlmClient for OpenRouterClient {
async fn generate(
&self,
cx: &Context,
prompt: &str,
system: Option<&str>,
images: Option<Vec<String>>,
@@ -324,90 +342,34 @@ impl LlmClient for OpenRouterClient {
user.images = images;
messages.push(user);
let (reply, _, _) = self.chat_with_tools(messages, Vec::new()).await?;
// Delegates to chat_with_tools, which starts the llm.request span
// (operation=chat) — one span per HTTP request.
let (reply, _, _) = self.chat_with_tools(cx, messages, Vec::new()).await?;
Ok(reply.content)
}
async fn chat_with_tools(
&self,
cx: &Context,
messages: Vec<ChatMessage>,
tools: Vec<Tool>,
) -> Result<(ChatMessage, Option<i32>, Option<i32>)> {
let url = format!("{}/chat/completions", self.base_url);
let mut body = serde_json::Map::new();
body.insert("model".into(), Value::String(self.primary_model.clone()));
body.insert(
"messages".into(),
Value::Array(Self::messages_to_openai(&messages)),
);
body.insert("stream".into(), Value::Bool(false));
if !tools.is_empty() {
body.insert(
"tools".into(),
serde_json::to_value(&tools).context("serializing tools")?,
);
}
for (k, v) in self.build_options() {
body.insert(k.into(), v);
}
log::info!(
"OpenRouter chat_with_tools: model={} messages={} tools={}",
self.primary_model,
messages.len(),
tools.len()
);
let resp = self
.authed(self.client.post(&url))
.json(&Value::Object(body))
.send()
let mut span = self.start_request_span(cx, "chat", &self.primary_model);
self.chat_with_tools_impl(Some(&mut span), messages, tools)
.await
.with_context(|| format!("POST {} failed", url))?;
if !resp.status().is_success() {
let status = resp.status();
let body = resp.text().await.unwrap_or_default();
bail!("OpenRouter chat request failed: {} — {}", status, body);
}
let parsed: Value = resp.json().await.context("parsing chat response")?;
let choice = parsed
.get("choices")
.and_then(|v| v.as_array())
.and_then(|a| a.first())
.ok_or_else(|| {
anyhow!(
"response missing choices[0]: {}",
extract_openrouter_error_detail(&parsed)
)
})?;
let msg = choice.get("message").ok_or_else(|| {
anyhow!(
"choices[0] missing message: {}",
extract_openrouter_error_detail(&parsed)
)
})?;
let chat_msg = Self::openai_message_to_chat(msg)?;
let usage = parsed.get("usage");
let prompt_tokens = usage
.and_then(|u| u.get("prompt_tokens"))
.and_then(|v| v.as_i64())
.map(|n| n as i32);
let completion_tokens = usage
.and_then(|u| u.get("completion_tokens"))
.and_then(|v| v.as_i64())
.map(|n| n as i32);
Ok((chat_msg, prompt_tokens, completion_tokens))
}
async fn chat_with_tools_stream(
&self,
cx: &Context,
messages: Vec<ChatMessage>,
tools: Vec<Tool>,
) -> Result<BoxStream<'static, Result<LlmStreamEvent>>> {
// The span outlives this method — the stream yields until the
// response fully drains, so the owned span moves into the
// generator and is finalized there.
let mut span = self.start_request_span(cx, "chat_stream", &self.primary_model);
let url = format!("{}/chat/completions", self.base_url);
let mut body = serde_json::Map::new();
body.insert("model".into(), Value::String(self.primary_model.clone()));
@@ -432,18 +394,32 @@ impl LlmClient for OpenRouterClient {
body.insert(k.into(), v);
}
let resp = self
let mut request = self
.authed(self.client.post(&url))
.json(&Value::Object(body))
.json(&Value::Object(body));
for (k, v) in crate::otel::traceparent_headers(
&Context::new().with_remote_span_context(span.span_context().clone()),
) {
request = request.header(k, v);
}
let resp = request
.send()
.await
.with_context(|| format!("POST {} failed", url))?;
if !resp.status().is_success() {
let status = resp.status();
span.set_attribute(KeyValue::new("http.status_code", status.as_u16() as i64));
span.set_status(Status::error(format!(
"OpenRouter stream request failed: {status}"
)));
let body = resp.text().await.unwrap_or_default();
bail!("OpenRouter stream request failed: {} — {}", status, body);
}
span.set_attribute(KeyValue::new(
"http.status_code",
resp.status().as_u16() as i64,
));
// OpenAI-compat SSE stream. Each event is `data: <json>\n\n`, with
// `data: [DONE]` signalling completion. Tool calls arrive as
@@ -467,6 +443,7 @@ impl LlmClient for OpenRouterClient {
let chunk = match chunk {
Ok(b) => b,
Err(e) => {
span.set_status(Status::error(format!("stream read failed: {e}")));
yield Err(anyhow!("stream read failed: {}", e));
return;
}
@@ -517,6 +494,9 @@ impl LlmClient for OpenRouterClient {
.get("completion_tokens")
.and_then(|n| n.as_i64())
.map(|n| n as i32);
if let Some(cost) = usage.get("cost").and_then(|c| c.as_f64()) {
span.set_attribute(KeyValue::new("llm.cost_usd", cost));
}
}
let Some(choices) = v.get("choices").and_then(|c| c.as_array())
@@ -606,6 +586,13 @@ impl LlmClient for OpenRouterClient {
tool_calls,
images: None,
};
if let Some(pt) = prompt_tokens {
span.set_attribute(KeyValue::new("llm.tokens.input", pt as i64));
}
if let Some(ct) = completion_tokens {
span.set_attribute(KeyValue::new("llm.tokens.output", ct as i64));
}
span.set_status(Status::Ok);
yield Ok(LlmStreamEvent::Done {
message,
prompt_eval_count: prompt_tokens,
@@ -616,25 +603,41 @@ impl LlmClient for OpenRouterClient {
Ok(Box::pin(stream))
}
async fn generate_embeddings(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>> {
async fn generate_embeddings(&self, cx: &Context, texts: &[&str]) -> Result<Vec<Vec<f32>>> {
let mut span = self.start_request_span(cx, "embeddings", &self.embedding_model);
let url = format!("{}/embeddings", self.base_url);
let body = json!({
"model": self.embedding_model,
"input": texts,
});
let resp = self
.authed(self.client.post(&url))
.json(&body)
let mut request = self.authed(self.client.post(&url)).json(&body);
for (k, v) in crate::otel::traceparent_headers(
&Context::new().with_remote_span_context(span.span_context().clone()),
) {
request = request.header(k, v);
}
let resp = request
.send()
.await
.with_context(|| format!("POST {} failed", url))?;
.with_context(|| format!("POST {} failed", url))
.inspect_err(|_| {
span.set_status(Status::error("OpenRouter embedding request failed"));
})?;
if !resp.status().is_success() {
let status = resp.status();
span.set_attribute(KeyValue::new("http.status_code", status.as_u16() as i64));
span.set_status(Status::error(format!(
"OpenRouter embedding request failed: {status}"
)));
let body = resp.text().await.unwrap_or_default();
bail!("OpenRouter embedding request failed: {} — {}", status, body);
}
span.set_attribute(KeyValue::new(
"http.status_code",
resp.status().as_u16() as i64,
));
#[derive(Deserialize)]
struct EmbedResponse {
@@ -646,18 +649,28 @@ impl LlmClient for OpenRouterClient {
}
let parsed: EmbedResponse = resp.json().await.context("parsing embed response")?;
span.set_status(Status::Ok);
Ok(parsed.data.into_iter().map(|i| i.embedding).collect())
}
async fn describe_image(&self, image_base64: &str) -> Result<String> {
async fn describe_image(&self, cx: &Context, image_base64: &str) -> Result<String> {
// Own span with operation=describe_image — the funnel into
// chat_with_tools_impl keeps the request on a single span.
let mut span = self.start_request_span(cx, "describe_image", &self.primary_model);
let prompt = "Briefly describe what you see in this image in 1-2 sentences. \
Focus on the people, location, and activity.";
self.generate(
prompt,
Some("You are a scene description assistant. Be concise and factual."),
Some(vec![image_base64.to_string()]),
)
.await
let mut messages: Vec<ChatMessage> = Vec::new();
messages.push(ChatMessage::system(
"You are a scene description assistant. Be concise and factual.",
));
let mut user = ChatMessage::user(prompt);
user.images = Some(vec![image_base64.to_string()]);
messages.push(user);
let (reply, _, _) = self
.chat_with_tools_impl(Some(&mut span), messages, Vec::new())
.await?;
Ok(reply.content)
}
async fn list_models(&self) -> Result<Vec<ModelCapabilities>> {
@@ -711,6 +724,134 @@ impl LlmClient for OpenRouterClient {
}
}
impl OpenRouterClient {
/// HTTP core of `chat_with_tools`. `span` is the `llm.request` span
/// started by the caller (or a funnelling entry point like
/// `describe_image`) so attributes from the response land on the right
/// operation label.
async fn chat_with_tools_impl(
&self,
mut span: Option<&mut opentelemetry::global::BoxedSpan>,
messages: Vec<ChatMessage>,
tools: Vec<Tool>,
) -> Result<(ChatMessage, Option<i32>, Option<i32>)> {
let url = format!("{}/chat/completions", self.base_url);
let mut body = serde_json::Map::new();
body.insert("model".into(), Value::String(self.primary_model.clone()));
body.insert(
"messages".into(),
Value::Array(Self::messages_to_openai(&messages)),
);
body.insert("stream".into(), Value::Bool(false));
if !tools.is_empty() {
body.insert(
"tools".into(),
serde_json::to_value(&tools).context("serializing tools")?,
);
}
for (k, v) in self.build_options() {
body.insert(k.into(), v);
}
log::info!(
"OpenRouter chat_with_tools: model={} messages={} tools={}",
self.primary_model,
messages.len(),
tools.len()
);
let mut request = self
.authed(self.client.post(&url))
.json(&Value::Object(body));
if let Some(s) = span.as_mut() {
for (k, v) in crate::otel::traceparent_headers(
&Context::new().with_remote_span_context(s.span_context().clone()),
) {
request = request.header(k, v);
}
}
let resp = match request.send().await {
Ok(r) => r,
Err(e) => {
if let Some(s) = span.as_mut() {
s.set_status(Status::error("OpenRouter request failed"));
}
return Err(e).with_context(|| format!("POST {} failed", url));
}
};
if !resp.status().is_success() {
let status = resp.status();
if let Some(s) = span.as_mut() {
s.set_attribute(KeyValue::new("http.status_code", status.as_u16() as i64));
s.set_status(Status::error(format!(
"OpenRouter chat request failed: {status}"
)));
}
let body = resp.text().await.unwrap_or_default();
bail!("OpenRouter chat request failed: {} — {}", status, body);
}
if let Some(s) = span.as_mut() {
s.set_attribute(KeyValue::new(
"http.status_code",
resp.status().as_u16() as i64,
));
}
let parsed: Value = match resp.json().await {
Ok(v) => v,
Err(e) => {
if let Some(s) = span.as_mut() {
s.set_status(Status::error("parsing chat response failed"));
}
return Err(e).context("parsing chat response");
}
};
let choice = parsed
.get("choices")
.and_then(|v| v.as_array())
.and_then(|a| a.first())
.ok_or_else(|| {
anyhow!(
"response missing choices[0]: {}",
extract_openrouter_error_detail(&parsed)
)
})?;
let msg = choice.get("message").ok_or_else(|| {
anyhow!(
"choices[0] missing message: {}",
extract_openrouter_error_detail(&parsed)
)
})?;
let chat_msg = Self::openai_message_to_chat(msg)?;
let usage = parsed.get("usage");
let prompt_tokens = usage
.and_then(|u| u.get("prompt_tokens"))
.and_then(|v| v.as_i64())
.map(|n| n as i32);
let completion_tokens = usage
.and_then(|u| u.get("completion_tokens"))
.and_then(|v| v.as_i64())
.map(|n| n as i32);
// OpenRouter reports per-request USD cost in `usage.cost`.
if let Some(s) = span.as_mut() {
if let Some(pt) = prompt_tokens {
s.set_attribute(KeyValue::new("llm.tokens.input", pt as i64));
}
if let Some(ct) = completion_tokens {
s.set_attribute(KeyValue::new("llm.tokens.output", ct as i64));
}
if let Some(cost) = usage.and_then(|u| u.get("cost")).and_then(|c| c.as_f64()) {
s.set_attribute(KeyValue::new("llm.cost_usd", cost));
}
s.set_status(Status::Ok);
}
Ok((chat_msg, prompt_tokens, completion_tokens))
}
}
/// Extract a diagnostic fragment from an OpenRouter response body that
/// doesn't match the expected `{choices: [...]}` shape. OpenRouter will
/// sometimes return 200 OK with `{"error": {"message": "...", "code": ...}}`
+13 -3
View File
@@ -26,7 +26,7 @@ use actix_web::{HttpRequest, HttpResponse, Responder, delete, get, post, web};
use anyhow::{Context, Result, anyhow, bail};
use chrono::Utc;
use opentelemetry::KeyValue;
use opentelemetry::trace::{Span, Status, Tracer};
use opentelemetry::trace::{Span, Status, TraceContextExt, Tracer};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::{Arc, Mutex as StdMutex};
@@ -685,6 +685,7 @@ impl PersonaChatSession {
entry_clone.clone(),
user_id,
library_id,
&mut span,
)
.await;
match result {
@@ -721,6 +722,7 @@ impl PersonaChatSession {
entry: Arc<TurnEntry>,
user_id: i32,
_library_id: i32,
turn_span: &mut opentelemetry::global::BoxedSpan,
) -> Result<()> {
// Per-conversation mutex — serializes this turn against a
// concurrent rewind or branch switch on the same transcript.
@@ -808,6 +810,12 @@ impl PersonaChatSession {
.resolve_backend(backend_kind, &overrides)
.await?;
let model_used = backend.model().to_string();
// llm.* attributes on the turn-level span (issue #130), plus the
// parent context for the child llm.request spans.
turn_span.set_attribute(KeyValue::new("llm.backend", backend.kind.as_str()));
turn_span.set_attribute(KeyValue::new("llm.model", model_used.clone()));
let turn_cx = opentelemetry::Context::new()
.with_remote_span_context(turn_span.span_context().clone());
// Build a no-photo tool catalog. `current_gate_opts_for_persona`
// already probes each tool's backing table for presence and
@@ -874,6 +882,7 @@ impl PersonaChatSession {
&persona.persona_id,
max_iterations,
&entry,
&turn_cx,
)
.await
.map_err(|e| anyhow!("persona chat loop failed: {e:?}"))?;
@@ -942,7 +951,7 @@ impl PersonaChatSession {
// list falls back to the persona name while the title is empty.
if needs_title {
let title = self
.generate_title(&backend, &messages)
.generate_title(&backend, &messages, &turn_cx)
.await
.unwrap_or_else(|| fallback_title(&messages));
if !title.is_empty() {
@@ -1047,6 +1056,7 @@ impl PersonaChatSession {
&self,
backend: &crate::ai::backend::ResolvedBackend,
messages: &[ChatMessage],
cx: &opentelemetry::Context,
) -> Option<String> {
let user_message = messages.iter().find(|m| m.role == "user")?;
// The last assistant message with actual prose — earlier ones may be
@@ -1057,7 +1067,7 @@ impl PersonaChatSession {
.find(|m| m.role == "assistant" && !m.content.trim().is_empty())?;
let prompt = title_prompt(&user_message.content, &reply.content);
match backend.chat().generate(&prompt, None, None).await {
match backend.chat().generate(cx, &prompt, None, None).await {
Ok(raw) => {
let title = sanitize_title(&crate::ai::llm_client::strip_think_blocks(&raw));
if title.is_empty() { None } else { Some(title) }
+2
View File
@@ -324,6 +324,7 @@ impl SmsApiClient {
&self,
messages: &[SmsMessage],
ollama: &OllamaClient,
cx: &opentelemetry::Context,
) -> Result<String> {
if messages.is_empty() {
return Ok(String::from("No messages on this day"));
@@ -356,6 +357,7 @@ Summary:"#,
ollama
.generate(
cx,
&prompt,
// Some("You are a summarizer for the purposes of jogging my memory and highlighting events and situations."),
Some("You are the keeper of memories, ingest the context and give me a casual summary of the moment."),
+6 -2
View File
@@ -220,7 +220,9 @@ async fn main() -> Result<()> {
println!("Generating summary...");
let summary = ollama.generate(&prompt, Some(system_prompt)).await?;
let summary = ollama
.generate(&opentelemetry::Context::new(), &prompt, Some(system_prompt))
.await?;
println!("\n📝 GENERATED SUMMARY:");
println!("─────────────────────────────────────────");
@@ -236,7 +238,9 @@ async fn main() -> Result<()> {
);
println!("\nGenerating embedding...");
let embedding = ollama.generate_embedding(&stripped).await?;
let embedding = ollama
.generate_embedding(&opentelemetry::Context::new(), &stripped)
.await?;
println!("✓ Embedding generated ({} dimensions)", embedding.len());
println!("Saving to database...");
+2 -1
View File
@@ -859,7 +859,8 @@ async fn synthesize_merge<D: KnowledgeDao + 'static>(
let ollama = app_state.ollama.clone();
let model_used = ollama.primary_model.clone();
let proposed = match ollama.generate(&prompt, Some(system)).await {
let cx = opentelemetry::Context::current();
let proposed = match ollama.generate(&cx, &prompt, Some(system)).await {
Ok(out) => {
// Strip the framing models reach for even with explicit
// "no preamble" guidance: leading "Merged description:"
+92 -1
View File
@@ -2,7 +2,7 @@ use actix_web::HttpRequest;
use actix_web::http::header::HeaderMap;
use opentelemetry::global::{BoxedSpan, BoxedTracer};
use opentelemetry::propagation::TextMapPropagator;
use opentelemetry::trace::{Span, Status, Tracer};
use opentelemetry::trace::{Span, SpanBuilder, SpanKind, Status, TraceContextExt, Tracer};
use opentelemetry::{Context, KeyValue, global};
use opentelemetry_appender_log::OpenTelemetryLogBridge;
use opentelemetry_otlp::WithExportConfig;
@@ -77,11 +77,46 @@ impl<'a> opentelemetry::propagation::Extractor for HeaderExtractor<'a> {
self.0.keys().map(|k| k.as_str()).collect()
}
}
/// Start a client span for a single LLM call. Every provider funnels its
/// LlmClient methods through this so the span shape (name, `SpanKind::Client`,
/// provider / operation / model attributes) is identical across
/// openrouter / llamacpp / ollama. `operation` is one of `chat`,
/// `chat_stream`, `embeddings`, `describe_image`, `generate`.
pub fn start_llm_request_span(
cx: &Context,
provider: &str,
operation: &str,
model: &str,
) -> BoxedSpan {
let builder = SpanBuilder::from_name("llm.request").with_kind(SpanKind::Client);
let mut span = global_tracer().build_with_context(builder, cx);
span.set_attribute(KeyValue::new("llm.provider", provider.to_string()));
span.set_attribute(KeyValue::new("llm.operation", operation.to_string()));
span.set_attribute(KeyValue::new("llm.model", model.to_string()));
span
}
pub fn extract_context_from_request(req: &HttpRequest) -> Context {
let propagator = TraceContextPropagator::new();
propagator.extract(&HeaderExtractor(req.headers()))
}
/// Build the W3C trace-context headers to attach to an outgoing HTTP
/// request so a downstream service (llama-swap, Ollama) can parent its own
/// spans onto this trace. Returns an empty vec when the context carries no
/// valid span — e.g. debug builds, where the global tracer is a no-op and
/// injecting an all-zero traceparent would produce garbage parent spans.
pub fn traceparent_headers(context: &Context) -> Vec<(String, String)> {
if !context.span().span_context().is_valid() {
return Vec::new();
}
let mut headers: std::collections::HashMap<String, String> = std::collections::HashMap::new();
TraceContextPropagator::new().inject_context(context, &mut headers);
// The propagator writes an empty tracestate header when there is no
// state — drop empty-valued headers rather than sending noise.
headers.into_iter().filter(|(_, v)| !v.is_empty()).collect()
}
pub fn trace_db_call<F, O>(
context: &Context,
query_type: &str,
@@ -110,3 +145,59 @@ where
result
}
#[cfg(test)]
mod traceparent_tests {
use super::*;
use opentelemetry::trace::{
SpanContext, SpanId, TraceContextExt, TraceFlags, TraceId, TraceState,
};
/// A context carrying a span context with known ids — no global
/// provider needed, so tests can't pollute process state.
fn context_with_flags(flags: TraceFlags) -> Context {
let trace_id = TraceId::from_bytes([
0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e,
0x0f, 0x10,
]);
let span_id = SpanId::from_bytes([0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08]);
let span_context = SpanContext::new(trace_id, span_id, flags, false, TraceState::default());
Context::new().with_remote_span_context(span_context)
}
#[test]
fn traceparent_headers_sampled_context() {
let cx = context_with_flags(TraceFlags::SAMPLED);
let headers = traceparent_headers(&cx);
assert_eq!(
headers,
vec![(
"traceparent".to_string(),
"00-0102030405060708090a0b0c0d0e0f10-0102030405060708-01".to_string(),
)]
);
}
#[test]
fn traceparent_headers_not_sampled_context() {
// A valid-but-unsampled context still propagates; the flags byte
// carries the state.
let cx = context_with_flags(TraceFlags::NOT_SAMPLED);
let headers = traceparent_headers(&cx);
assert_eq!(
headers,
vec![(
"traceparent".to_string(),
"00-0102030405060708090a0b0c0d0e0f10-0102030405060708-00".to_string(),
)]
);
}
#[test]
fn traceparent_headers_invalid_context_is_empty() {
// No span (debug-build no-op tracer) → no header, never an
// all-zero traceparent.
let cx = Context::new();
assert!(traceparent_headers(&cx).is_empty());
}
}
+2 -1
View File
@@ -281,8 +281,9 @@ pub async fn generate_script(
beats: &[PlannedBeat],
) -> Result<ReelScript> {
let (system, user) = build_script_messages(meta, beats);
let cx = opentelemetry::Context::current();
let raw = client
.generate(&user, Some(&system), None)
.generate(&cx, &user, Some(&system), None)
.await
.context("LLM script generation failed")?;
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());
let today = chrono::Utc::now().date_naive();
let sq = match translate_nl_query(backend.chat(), &nl, &tag_vocab, today).await {
let sq = match translate_nl_query(backend.chat(), &nl, &tag_vocab, today, &ctx).await {
Ok(sq) => sq,
Err(e) => {
log::warn!("unified_search: translate_nl_query failed: {e:?}");