Cameron Cordes
3b9d985025
- PersonaChatSession dispatches turns through the shared agent loop, persisting transcripts keyed on (user_id, persona_id) in the new persona_chat table (migration included). - GET /persona_chat/history returns a rendered transcript (tool invocations folded, is_initial flag) matching the file-chat shape. - POST /persona_chat/turn returns 202 with a turn_id; SSE replay and cancel reuse turn_replay_impl/cancel_turn_impl, extracted from the actix-attributed handlers so both route families share the logic. - POST /persona_chat/reset clears the persona transcript. - Concurrent dispatches for the same (user, persona) are rejected with 409 via an in-flight gate (InFlightPersonaTurns) whose RAII guard drops with the spawned turn task, freeing the slot on completion, error, or abort.
3392 lines
135 KiB
Rust
3392 lines
135 KiB
Rust
use anyhow::{Result, anyhow, bail};
|
|
use chrono::Utc;
|
|
use opentelemetry::KeyValue;
|
|
use opentelemetry::trace::{Span, Status, TraceContextExt, Tracer};
|
|
use std::collections::HashMap;
|
|
use std::sync::{Arc, Mutex};
|
|
use tokio::sync::Mutex as TokioMutex;
|
|
|
|
use crate::ai::backend::{BackendKind, ResolvedBackend, SamplingOverrides};
|
|
use crate::ai::insight_generator::InsightGenerator;
|
|
use crate::ai::llm_client::{
|
|
BranchLeafInfo, ChatHistoryStore, ChatMessage, LlmStreamEvent, StoredChatNode, Tool,
|
|
};
|
|
use crate::ai::turn_registry::TurnEntry;
|
|
use crate::ai::turn_registry::TurnRegistry;
|
|
use crate::database::InsightDao;
|
|
use crate::database::models::InsertPhotoInsight;
|
|
use crate::otel::global_tracer;
|
|
use crate::utils::{normalize_path, retry_with_backoff};
|
|
use futures::stream::{BoxStream, StreamExt};
|
|
use uuid::Uuid;
|
|
|
|
pub const DEFAULT_MAX_ITERATIONS: usize = 6;
|
|
/// Assumed context window when the request doesn't specify `num_ctx`.
|
|
/// The llama-swap chat slots serve 20k-131k contexts and real conversations
|
|
/// rarely pass ~16k tokens, so 32k keeps the truncation pass from gutting
|
|
/// history that the server could comfortably hold (which also destroys the
|
|
/// server's KV-cache prefix reuse). Override per-deploy with
|
|
/// AGENTIC_CHAT_DEFAULT_NUM_CTX if the serving models change shape.
|
|
const DEFAULT_NUM_CTX: i32 = 32768;
|
|
/// Headroom reserved for the model's response, deducted from the context
|
|
/// budget when deciding whether to truncate the replayed history.
|
|
const RESPONSE_HEADROOM_TOKENS: usize = 2048;
|
|
/// Cheap byte-to-token approximation used by the truncation pass. The real
|
|
/// tokenization is model-specific; this avoids carrying tiktoken just for a
|
|
/// soft bound.
|
|
const BYTES_PER_TOKEN: usize = 4;
|
|
/// Flat token cost charged per inlined image in the truncation budget. A
|
|
/// 1024px-longest-edge JPEG (see `load_image_as_base64`) costs vision models on
|
|
/// the order of ~1.3K tokens. Crucially, the raw base64 (hundreds of KB of
|
|
/// characters) must NOT be counted as text bytes — doing so dwarfs the entire
|
|
/// text budget and forces spurious truncation on every turn.
|
|
const IMAGE_TOKENS_EACH: usize = 1300;
|
|
/// User prompt injected when the agentic loop exhausts its iteration budget
|
|
/// without producing a tool-free reply. Internal scaffolding only — it is
|
|
/// stripped from the transcript before persistence (see
|
|
/// [`push_synthetic_final_prompt`] / [`remove_synthetic_final_prompt`]).
|
|
const SYNTHETIC_FINAL_ANSWER_PROMPT: &str =
|
|
"Please write your final answer now without calling any more tools.";
|
|
|
|
pub type ChatLockMap = Arc<TokioMutex<HashMap<(i32, String), Arc<TokioMutex<()>>>>>;
|
|
|
|
#[derive(Debug)]
|
|
pub struct ChatTurnRequest {
|
|
pub library_id: i32,
|
|
/// Author's user_id, extracted from Claims at the handler. Tagged
|
|
/// onto every entity_fact row written this turn so the composite FK
|
|
/// (user_id, persona_id) → personas holds and so cross-user reads
|
|
/// stay isolated. Service token claims that don't parse as i32
|
|
/// fall through to user_id=1 (operator convention).
|
|
pub user_id: i32,
|
|
pub file_path: String,
|
|
pub user_message: String,
|
|
/// Override the model id. Local mode: an Ollama model name. Hybrid:
|
|
/// an OpenRouter id. None defers to the stored insight's `model_version`.
|
|
pub model: Option<String>,
|
|
/// Override the backend used for this turn. None defers to the stored
|
|
/// insight's `backend`. Switching `local -> hybrid` is rejected in v1.
|
|
pub backend: Option<String>,
|
|
pub num_ctx: Option<i32>,
|
|
pub temperature: Option<f32>,
|
|
pub top_p: Option<f32>,
|
|
pub top_k: Option<i32>,
|
|
pub min_p: Option<f32>,
|
|
/// Reasoning toggle for thinking-capable models. Forwarded to the
|
|
/// llama.cpp backend as `chat_template_kwargs.enable_thinking`; ignored
|
|
/// by other backends. None defers to the model/template default.
|
|
pub enable_thinking: Option<bool>,
|
|
pub max_iterations: Option<usize>,
|
|
/// Per-turn system-prompt override. In append mode (default), applied
|
|
/// ephemerally — original system message restored before persistence.
|
|
/// In amend mode, persisted into the new insight row's system message.
|
|
/// None / empty = no change.
|
|
pub system_prompt: Option<String>,
|
|
/// Active persona id for this turn. Tools that write to
|
|
/// `entity_facts` tag the new rows with it; `recall_facts_for_photo`
|
|
/// scopes its read to it. None defaults to `"default"`.
|
|
pub persona_id: Option<String>,
|
|
/// When true, write a new insight row (regenerating title) instead of
|
|
/// updating training_messages on the existing row.
|
|
pub amend: bool,
|
|
/// When true, force the bootstrap path even if an insight already exists:
|
|
/// flip prior rows to `is_current=false` and create a new insight row
|
|
/// from `system_prompt` + `user_message` + photo. Takes precedence over
|
|
/// `amend`. With no existing insight, collapses to a normal bootstrap
|
|
/// (the row-flip step is a no-op).
|
|
pub regenerate: bool,
|
|
}
|
|
|
|
#[derive(Debug)]
|
|
pub struct ChatTurnResult {
|
|
pub assistant_message: String,
|
|
pub tool_calls_made: usize,
|
|
pub iterations_used: usize,
|
|
pub truncated: bool,
|
|
pub prompt_eval_count: Option<i32>,
|
|
pub eval_count: Option<i32>,
|
|
/// Set when `amend=true` and the new insight row was inserted.
|
|
pub amended_insight_id: Option<i32>,
|
|
/// Backend used for this turn — useful when the client overrode the
|
|
/// stored value.
|
|
pub backend_used: String,
|
|
/// Model identifier the chat backend ran with.
|
|
pub model_used: String,
|
|
}
|
|
|
|
#[derive(Clone)]
|
|
pub struct InsightChatService {
|
|
generator: Arc<InsightGenerator>,
|
|
insight_dao: Arc<Mutex<Box<dyn InsightDao>>>,
|
|
chat_locks: ChatLockMap,
|
|
}
|
|
|
|
impl InsightChatService {
|
|
pub fn new(
|
|
generator: Arc<InsightGenerator>,
|
|
insight_dao: Arc<Mutex<Box<dyn InsightDao>>>,
|
|
chat_locks: ChatLockMap,
|
|
) -> Self {
|
|
Self {
|
|
generator,
|
|
insight_dao,
|
|
chat_locks,
|
|
}
|
|
}
|
|
|
|
/// Accessor for the insight DAO (used by async job completion).
|
|
pub fn insight_dao(&self) -> &Arc<Mutex<Box<dyn InsightDao>>> {
|
|
&self.insight_dao
|
|
}
|
|
|
|
/// Load the rendered transcript for chat-UI display. Deserializes the
|
|
/// `training_messages` tree, traverses to `active_leaf_id`, and renders
|
|
/// the path from root to that leaf. Backward-compatible: if the stored
|
|
/// value is a flat array (old format), converts to a tree automatically.
|
|
///
|
|
/// `library_id` scopes the lookup to one library — without it, a
|
|
/// regenerate on lib1 can be shadowed on the next refresh by an
|
|
/// untouched `is_current=true` row in lib2 for the same rel_path.
|
|
pub fn load_history(
|
|
&self,
|
|
library_id: i32,
|
|
file_path: &str,
|
|
branch_id: Option<u64>,
|
|
) -> Result<HistoryView> {
|
|
let normalized = normalize_path(file_path);
|
|
let cx = opentelemetry::Context::new();
|
|
let mut dao = self.insight_dao.lock().expect("Unable to lock InsightDao");
|
|
let insight = match dao
|
|
.get_current_insight_for_library(&cx, library_id, &normalized)
|
|
.map_err(|e| anyhow!("failed to load insight: {:?}", e))?
|
|
{
|
|
Some(i) => i,
|
|
None => dao
|
|
.get_insight(&cx, &normalized)
|
|
.map_err(|e| anyhow!("failed to load insight: {:?}", e))?
|
|
.ok_or_else(|| anyhow!("no insight found for path"))?,
|
|
};
|
|
|
|
let raw = insight
|
|
.training_messages
|
|
.as_ref()
|
|
.ok_or_else(|| anyhow!("insight has no chat history (pre-agentic insight)"))?;
|
|
|
|
// Backward-compatible deserialization: flat array (old) vs tree (new).
|
|
let store: ChatHistoryStore = if let Ok(arr) = serde_json::from_str::<Vec<ChatMessage>>(raw)
|
|
{
|
|
ChatHistoryStore::from_flat_array(arr)
|
|
} else {
|
|
serde_json::from_str(raw)
|
|
.map_err(|e| anyhow!("failed to deserialize chat history: {}", e))?
|
|
};
|
|
|
|
// Use branch_id if provided, otherwise use active_leaf_id.
|
|
let target_leaf = branch_id.unwrap_or(store.active_leaf_id);
|
|
let path = store
|
|
.path_to_leaf(target_leaf)
|
|
.ok_or_else(|| anyhow!("branch_id {} not found in tree", target_leaf))?;
|
|
|
|
let (rendered, turn_count, _node_ids, fork_info) = render_tree_path(&store, &path);
|
|
|
|
Ok(HistoryView {
|
|
messages: rendered,
|
|
turn_count,
|
|
model_version: insight.model_version,
|
|
backend: insight.backend,
|
|
active_leaf_id: store.active_leaf_id,
|
|
viewing_branch_id: target_leaf,
|
|
fork_info,
|
|
})
|
|
}
|
|
|
|
pub async fn chat_turn(&self, req: ChatTurnRequest) -> Result<ChatTurnResult> {
|
|
let tracer = global_tracer();
|
|
let parent_cx = opentelemetry::Context::new();
|
|
let mut span = tracer.start_with_context("ai.insight.chat_turn", &parent_cx);
|
|
span.set_attribute(KeyValue::new("file_path", req.file_path.clone()));
|
|
span.set_attribute(KeyValue::new("library_id", req.library_id as i64));
|
|
span.set_attribute(KeyValue::new("amend", req.amend));
|
|
|
|
if req.user_message.trim().is_empty() {
|
|
bail!("user_message must not be empty");
|
|
}
|
|
if req.user_message.len() > 8192 {
|
|
bail!("user_message exceeds 8192 chars");
|
|
}
|
|
|
|
let active_persona = req
|
|
.persona_id
|
|
.clone()
|
|
.filter(|s| !s.trim().is_empty())
|
|
.unwrap_or_else(|| "default".to_string());
|
|
span.set_attribute(KeyValue::new("persona_id", active_persona.clone()));
|
|
|
|
let normalized = normalize_path(&req.file_path);
|
|
|
|
// 1. Acquire the per-(library, file) async mutex. Two concurrent
|
|
// chat turns on the same insight would race on the JSON blob —
|
|
// the lock serialises them.
|
|
let lock_key = (req.library_id, normalized.clone());
|
|
let entry_lock = {
|
|
let mut locks = self.chat_locks.lock().await;
|
|
locks
|
|
.entry(lock_key.clone())
|
|
.or_insert_with(|| Arc::new(TokioMutex::new(())))
|
|
.clone()
|
|
};
|
|
let _guard = entry_lock.lock().await;
|
|
|
|
// 2. Load the current insight + history.
|
|
let insight = {
|
|
let cx = opentelemetry::Context::new();
|
|
let mut dao = self.insight_dao.lock().expect("Unable to lock InsightDao");
|
|
dao.get_current_insight_for_library(&cx, req.library_id, &normalized)
|
|
.map_err(|e| anyhow!("failed to load insight: {:?}", e))?
|
|
.ok_or_else(|| anyhow!("no insight found for path"))?
|
|
};
|
|
let raw_history = insight
|
|
.training_messages
|
|
.as_ref()
|
|
.ok_or_else(|| {
|
|
anyhow!("insight has no chat history; regenerate this insight in agentic mode")
|
|
})?
|
|
.clone();
|
|
|
|
// Backward-compatible: flat array (old) vs tree (new).
|
|
let mut store: ChatHistoryStore =
|
|
if let Ok(arr) = serde_json::from_str::<Vec<ChatMessage>>(&raw_history) {
|
|
ChatHistoryStore::from_flat_array(arr)
|
|
} else {
|
|
serde_json::from_str(&raw_history)
|
|
.map_err(|e| anyhow!("failed to deserialize chat history: {}", e))?
|
|
};
|
|
|
|
// Build messages from the path to the active leaf.
|
|
let path = store
|
|
.path_to_leaf(store.active_leaf_id)
|
|
.ok_or_else(|| anyhow!("active_leaf_id {} not found in tree", store.active_leaf_id))?;
|
|
let mut messages: Vec<ChatMessage> = path.iter().map(|n| n.message.clone()).collect();
|
|
|
|
// 3. Resolve effective backend. Reject the unsupported switch.
|
|
let stored_backend = insight.backend.clone();
|
|
let effective_backend = req
|
|
.backend
|
|
.as_deref()
|
|
.map(|s| s.trim().to_lowercase())
|
|
.filter(|s| !s.is_empty())
|
|
.unwrap_or_else(|| stored_backend.clone());
|
|
validate_cross_replay(&stored_backend, &effective_backend)?;
|
|
let kind = BackendKind::parse(&effective_backend)?;
|
|
span.set_attribute(KeyValue::new("backend", kind.as_str()));
|
|
|
|
let max_iterations = req
|
|
.max_iterations
|
|
.unwrap_or(DEFAULT_MAX_ITERATIONS)
|
|
.clamp(1, env_max_iterations());
|
|
span.set_attribute(KeyValue::new("max_iterations", max_iterations as i64));
|
|
|
|
let stored_model = insight.model_version.clone();
|
|
let overrides = SamplingOverrides {
|
|
model: req
|
|
.model
|
|
.clone()
|
|
.or_else(|| Some(stored_model.clone()))
|
|
.filter(|m| !m.is_empty()),
|
|
num_ctx: req.num_ctx,
|
|
temperature: req.temperature,
|
|
top_p: req.top_p,
|
|
top_k: req.top_k,
|
|
min_p: req.min_p,
|
|
enable_thinking: req.enable_thinking,
|
|
};
|
|
let backend = self.generator.resolve_backend(kind, &overrides).await?;
|
|
let model_used = backend.model().to_string();
|
|
span.set_attribute(KeyValue::new("model", model_used.clone()));
|
|
|
|
// 5. Decide vision + tool set. In hybrid (describe-then-inline) mode
|
|
// we omit `describe_photo`. Otherwise trust the stored history:
|
|
// if the first user message carries images, describe_photo stays.
|
|
let local_first_user_has_image = messages
|
|
.iter()
|
|
.find(|m| m.role == "user")
|
|
.and_then(|m| m.images.as_ref())
|
|
.map(|imgs| !imgs.is_empty())
|
|
.unwrap_or(false);
|
|
let offer_describe_tool = backend.images_inline && local_first_user_has_image;
|
|
let gate_opts = self.generator.current_gate_opts_for_persona(
|
|
offer_describe_tool,
|
|
Some((req.user_id, &active_persona)),
|
|
);
|
|
let tools = InsightGenerator::build_tool_definitions(gate_opts);
|
|
|
|
let image_base64: Option<String> = if offer_describe_tool {
|
|
self.generator.load_image_as_base64(&normalized).ok()
|
|
} else {
|
|
None
|
|
};
|
|
|
|
// 6. Apply truncation budget. Drops oldest tool_call+tool pairs
|
|
// (preserves system + first user including any images).
|
|
let budget_tokens = (req.num_ctx.unwrap_or_else(env_default_num_ctx) as usize)
|
|
.saturating_sub(RESPONSE_HEADROOM_TOKENS);
|
|
let budget_bytes = budget_tokens.saturating_mul(BYTES_PER_TOKEN);
|
|
let truncated = apply_context_budget(&mut messages, budget_bytes);
|
|
if truncated {
|
|
span.set_attribute(KeyValue::new("history_truncated", true));
|
|
}
|
|
|
|
// 7. Append the new user turn.
|
|
messages.push(ChatMessage::user(req.user_message.clone()));
|
|
|
|
// Apply per-turn system-prompt override BEFORE the budget annotation
|
|
// so the budget note attaches to the override, not the original.
|
|
// The stash is consumed below before persistence (append mode) or
|
|
// dropped (amend mode, where the override stays in place).
|
|
let override_stash =
|
|
apply_system_prompt_override(&mut messages, req.system_prompt.as_deref());
|
|
|
|
// Temporarily annotate the system message with this turn's iteration
|
|
// budget so the model knows how many tool-calling rounds it has. We
|
|
// restore the original content before persistence so the note doesn't
|
|
// accumulate across turns.
|
|
let original_system_content = annotate_system_with_budget(&mut messages, max_iterations);
|
|
|
|
let insight_cx = parent_cx.with_span(span);
|
|
|
|
// 8. Agentic loop — same shape as insight_generator's, but capped
|
|
// tighter and dispatching tools through the shared executor.
|
|
let loop_span = tracer.start_with_context("ai.chat.loop", &insight_cx);
|
|
let loop_cx = insight_cx.with_span(loop_span);
|
|
let mut tool_calls_made = 0usize;
|
|
let mut iterations_used = 0usize;
|
|
let mut last_prompt_eval_count: Option<i32> = None;
|
|
let mut last_eval_count: Option<i32> = None;
|
|
let mut final_content = String::new();
|
|
|
|
for iteration in 0..max_iterations {
|
|
iterations_used = iteration + 1;
|
|
log::info!("Chat iteration {}/{}", iterations_used, max_iterations);
|
|
|
|
let (response, prompt_tokens, eval_tokens) = backend
|
|
.chat()
|
|
.chat_with_tools(messages.clone(), tools.clone())
|
|
.await?;
|
|
last_prompt_eval_count = prompt_tokens;
|
|
last_eval_count = eval_tokens;
|
|
|
|
let mut response = response;
|
|
if let Some(ref mut tcs) = response.tool_calls {
|
|
for tc in tcs.iter_mut() {
|
|
if !tc.function.arguments.is_object() {
|
|
tc.function.arguments = serde_json::Value::Object(Default::default());
|
|
}
|
|
}
|
|
}
|
|
|
|
messages.push(response.clone());
|
|
|
|
if let Some(ref tool_calls) = response.tool_calls
|
|
&& !tool_calls.is_empty()
|
|
{
|
|
for tool_call in tool_calls {
|
|
tool_calls_made += 1;
|
|
log::info!(
|
|
"Chat tool call [{}]: {} {:?}",
|
|
iteration,
|
|
tool_call.function.name,
|
|
tool_call.function.arguments
|
|
);
|
|
let result = self
|
|
.generator
|
|
.execute_tool(
|
|
&tool_call.function.name,
|
|
&tool_call.function.arguments,
|
|
&backend,
|
|
&image_base64,
|
|
&normalized,
|
|
req.user_id,
|
|
&active_persona,
|
|
&loop_cx,
|
|
)
|
|
.await;
|
|
messages.push(ChatMessage::tool_result(result));
|
|
}
|
|
continue;
|
|
}
|
|
|
|
final_content = response.content;
|
|
break;
|
|
}
|
|
|
|
if final_content.is_empty() {
|
|
log::info!(
|
|
"Chat loop exhausted after {} iterations, requesting final answer",
|
|
iterations_used
|
|
);
|
|
let synthetic_idx = push_synthetic_final_prompt(&mut messages);
|
|
let (final_response, prompt_tokens, eval_tokens) = backend
|
|
.chat()
|
|
.chat_with_tools(messages.clone(), vec![])
|
|
.await?;
|
|
last_prompt_eval_count = prompt_tokens;
|
|
last_eval_count = eval_tokens;
|
|
final_content = final_response.content.clone();
|
|
messages.push(final_response);
|
|
// Drop the synthetic prompt before persistence — internal
|
|
// scaffolding only (mirrors both streaming variants).
|
|
remove_synthetic_final_prompt(&mut messages, synthetic_idx);
|
|
}
|
|
|
|
// Strip any leaked <think> reasoning block from the content we
|
|
// return / persist as the reply (the raw transcript keeps it).
|
|
let final_content = crate::ai::llm_client::strip_think_blocks(&final_content);
|
|
|
|
loop_cx.span().set_status(Status::Ok);
|
|
|
|
// Drop the per-turn iteration-budget note from the system message
|
|
// before we persist so it doesn't snowball on each subsequent turn.
|
|
restore_system_content(&mut messages, original_system_content);
|
|
|
|
// Append mode: undo the per-turn system-prompt override so the
|
|
// stored transcript keeps the original baked persona. Amend mode:
|
|
// keep the override in place — it becomes the new insight row's
|
|
// system message.
|
|
if !req.amend {
|
|
restore_system_prompt_override(&mut messages, override_stash);
|
|
}
|
|
|
|
// 9. Persist. Append mode rewrites the JSON blob in place; amend
|
|
// mode regenerates the title and inserts a new insight row,
|
|
// relying on store_insight to flip prior rows' is_current=false.
|
|
// Append new messages (this turn's user + assistant exchanges) as
|
|
// tree nodes chained from the previous active_leaf_id.
|
|
let path_len = path.len();
|
|
let new_messages = messages[path_len..].to_vec();
|
|
let mut parent_id = Some(store.active_leaf_id);
|
|
for msg in &new_messages {
|
|
let new_id = store.append_node(parent_id, msg.clone());
|
|
parent_id = Some(new_id);
|
|
}
|
|
// Update active_leaf_id to the last node appended this turn.
|
|
if let Some(last_id) = parent_id {
|
|
store.active_leaf_id = last_id;
|
|
}
|
|
|
|
let json = serde_json::to_string(&store)
|
|
.map_err(|e| anyhow!("failed to serialize chat tree: {}", e))?;
|
|
|
|
let mut amended_insight_id: Option<i32> = None;
|
|
if req.amend {
|
|
let title_prompt = format!(
|
|
"Create a short title (maximum 8 words) for the following journal entry:\n\n{}\n\n\
|
|
Capture the key moment or theme. Return ONLY the title, nothing else.",
|
|
final_content
|
|
);
|
|
let title_raw = backend
|
|
.chat()
|
|
.generate(
|
|
&title_prompt,
|
|
Some(
|
|
"You are my long term memory assistant. Use only the information provided. Do not invent details.",
|
|
),
|
|
None,
|
|
)
|
|
.await?;
|
|
let title = title_raw.trim().trim_matches('"').to_string();
|
|
|
|
// Amended rows intentionally do not inherit the parent's
|
|
// `fewshot_source_ids`. The parent's few-shot influence is still
|
|
// present in this row's content; if you want strict lineage
|
|
// tracking for training-set filtering, fetch the parent here and
|
|
// copy its value forward.
|
|
let new_row = InsertPhotoInsight {
|
|
library_id: req.library_id,
|
|
file_path: normalized.clone(),
|
|
title,
|
|
summary: final_content.clone(),
|
|
generated_at: Utc::now().timestamp(),
|
|
model_version: model_used.clone(),
|
|
is_current: true,
|
|
training_messages: Some(json),
|
|
backend: kind.as_str().to_string(),
|
|
fewshot_source_ids: None,
|
|
content_hash: None,
|
|
num_ctx: req.num_ctx,
|
|
temperature: req.temperature,
|
|
top_p: req.top_p,
|
|
top_k: req.top_k,
|
|
min_p: req.min_p,
|
|
system_prompt: req.system_prompt.clone(),
|
|
persona_id: req.persona_id.clone(),
|
|
prompt_eval_count: None,
|
|
eval_count: None,
|
|
};
|
|
let dao = self.insight_dao.clone();
|
|
let stored = retry_with_backoff("store_insight", 3, || {
|
|
let cx = opentelemetry::Context::new();
|
|
let mut d = dao.lock().expect("Unable to lock InsightDao");
|
|
d.store_insight(&cx, new_row.clone())
|
|
})
|
|
.map_err(|e| anyhow!("failed to store amended insight: {:?}", e))?;
|
|
amended_insight_id = Some(stored.id);
|
|
} else {
|
|
let dao = self.insight_dao.clone();
|
|
let rows = retry_with_backoff("update_training_messages", 3, || {
|
|
let cx = opentelemetry::Context::new();
|
|
let mut d = dao.lock().expect("Unable to lock InsightDao");
|
|
d.update_training_messages(&cx, req.library_id, &normalized, &json)
|
|
})
|
|
.map_err(|e| anyhow!("failed to persist chat history: {:?}", e))?;
|
|
if rows == 0 {
|
|
log::warn!(
|
|
"update_training_messages updated 0 rows for {} (lib {}), \
|
|
concurrent regenerate likely flipped is_current",
|
|
normalized,
|
|
req.library_id
|
|
);
|
|
}
|
|
}
|
|
|
|
Ok(ChatTurnResult {
|
|
assistant_message: final_content,
|
|
tool_calls_made,
|
|
iterations_used,
|
|
truncated,
|
|
prompt_eval_count: last_prompt_eval_count,
|
|
eval_count: last_eval_count,
|
|
amended_insight_id,
|
|
backend_used: kind.as_str().to_string(),
|
|
model_used,
|
|
})
|
|
}
|
|
|
|
/// Rewind the conversation to the rendered message at
|
|
/// `discard_from_rendered_index` by setting `active_leaf_id` to the
|
|
/// node just before the discarded message. Unlike the legacy flat-array
|
|
/// approach, this preserves all fork branches — the discarded path
|
|
/// remains accessible as an alternate branch.
|
|
///
|
|
/// The initial user turn cannot be discarded; attempting to do so
|
|
/// returns an error.
|
|
///
|
|
/// Holds the per-file chat mutex so it serialises with `chat_turn`.
|
|
pub async fn rewind_history(
|
|
&self,
|
|
library_id: i32,
|
|
file_path: &str,
|
|
discard_from_rendered_index: usize,
|
|
) -> Result<()> {
|
|
if discard_from_rendered_index == 0 {
|
|
bail!("cannot discard the initial user message");
|
|
}
|
|
let normalized = normalize_path(file_path);
|
|
|
|
let lock_key = (library_id, normalized.clone());
|
|
let entry_lock = {
|
|
let mut locks = self.chat_locks.lock().await;
|
|
locks
|
|
.entry(lock_key.clone())
|
|
.or_insert_with(|| Arc::new(TokioMutex::new(())))
|
|
.clone()
|
|
};
|
|
let _guard = entry_lock.lock().await;
|
|
|
|
let insight = {
|
|
let cx = opentelemetry::Context::new();
|
|
let mut dao = self.insight_dao.lock().expect("Unable to lock InsightDao");
|
|
dao.get_current_insight_for_library(&cx, library_id, &normalized)
|
|
.map_err(|e| anyhow!("failed to load insight: {:?}", e))?
|
|
.ok_or_else(|| anyhow!("no insight found for path"))?
|
|
};
|
|
let raw_history = insight
|
|
.training_messages
|
|
.as_ref()
|
|
.ok_or_else(|| anyhow!("insight has no chat history"))?;
|
|
|
|
let mut store: ChatHistoryStore =
|
|
if let Ok(arr) = serde_json::from_str::<Vec<ChatMessage>>(raw_history) {
|
|
ChatHistoryStore::from_flat_array(arr)
|
|
} else {
|
|
serde_json::from_str(raw_history)
|
|
.map_err(|e| anyhow!("failed to deserialize chat history: {}", e))?
|
|
};
|
|
|
|
let path = store
|
|
.path_to_leaf(store.active_leaf_id)
|
|
.ok_or_else(|| anyhow!("active_leaf_id not found in tree"))?;
|
|
|
|
let (_rendered, _turn_count, node_ids, _fork_info) = render_tree_path(&store, &path);
|
|
|
|
// The last kept rendered message is at index `discard_from_rendered_index - 1`.
|
|
let last_kept_idx = discard_from_rendered_index - 1;
|
|
let new_active_leaf_id = *node_ids
|
|
.get(last_kept_idx)
|
|
.ok_or_else(|| anyhow!("discard_from_rendered_index out of range"))?;
|
|
|
|
store.active_leaf_id = new_active_leaf_id;
|
|
let json = serde_json::to_string(&store)
|
|
.map_err(|e| anyhow!("failed to serialize tree: {}", e))?;
|
|
|
|
let dao = self.insight_dao.clone();
|
|
let rows = retry_with_backoff("update_training_messages", 3, || {
|
|
let cx = opentelemetry::Context::new();
|
|
let mut d = dao.lock().expect("Unable to lock InsightDao");
|
|
d.update_training_messages(&cx, library_id, &normalized, &json)
|
|
})
|
|
.map_err(|e| anyhow!("failed to persist rewound history: {:?}", e))?;
|
|
if rows == 0 {
|
|
log::warn!(
|
|
"update_training_messages (rewind) updated 0 rows for {} (lib {}), \
|
|
concurrent regenerate likely flipped is_current",
|
|
normalized,
|
|
library_id
|
|
);
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
/// Switch the active branch to the given leaf ID. The new branch becomes
|
|
/// the active conversation path; the previous active branch becomes a
|
|
/// regular fork. If the leaf ID doesn't exist, returns an error.
|
|
pub async fn switch_branch(
|
|
&self,
|
|
library_id: i32,
|
|
file_path: &str,
|
|
leaf_id: u64,
|
|
) -> Result<()> {
|
|
let normalized = normalize_path(file_path);
|
|
|
|
let lock_key = (library_id, normalized.clone());
|
|
let entry_lock = {
|
|
let mut locks = self.chat_locks.lock().await;
|
|
locks
|
|
.entry(lock_key.clone())
|
|
.or_insert_with(|| Arc::new(TokioMutex::new(())))
|
|
.clone()
|
|
};
|
|
let _guard = entry_lock.lock().await;
|
|
|
|
let insight = {
|
|
let cx = opentelemetry::Context::new();
|
|
let mut dao = self.insight_dao.lock().expect("Unable to lock InsightDao");
|
|
dao.get_current_insight_for_library(&cx, library_id, &normalized)
|
|
.map_err(|e| anyhow!("failed to load insight: {:?}", e))?
|
|
.ok_or_else(|| anyhow!("no insight found for path"))?
|
|
};
|
|
let raw_history = insight
|
|
.training_messages
|
|
.as_ref()
|
|
.ok_or_else(|| anyhow!("insight has no chat history"))?;
|
|
|
|
let mut store: ChatHistoryStore =
|
|
if let Ok(arr) = serde_json::from_str::<Vec<ChatMessage>>(raw_history) {
|
|
ChatHistoryStore::from_flat_array(arr)
|
|
} else {
|
|
serde_json::from_str(raw_history)
|
|
.map_err(|e| anyhow!("failed to deserialize chat history: {}", e))?
|
|
};
|
|
|
|
// Validate the leaf_id exists and is a leaf.
|
|
if !store.nodes.iter().any(|n| n.id == leaf_id) {
|
|
bail!("branch_id {} not found in tree", leaf_id);
|
|
}
|
|
if !store.children_of(leaf_id).is_empty() {
|
|
bail!("branch_id {} is not a leaf node", leaf_id);
|
|
}
|
|
|
|
store.active_leaf_id = leaf_id;
|
|
let json = serde_json::to_string(&store)
|
|
.map_err(|e| anyhow!("failed to serialize tree: {}", e))?;
|
|
|
|
let dao = self.insight_dao.clone();
|
|
let rows = retry_with_backoff("update_training_messages", 3, || {
|
|
let cx = opentelemetry::Context::new();
|
|
let mut d = dao.lock().expect("Unable to lock InsightDao");
|
|
d.update_training_messages(&cx, library_id, &normalized, &json)
|
|
})
|
|
.map_err(|e| anyhow!("failed to persist branch switch: {:?}", e))?;
|
|
if rows == 0 {
|
|
log::warn!(
|
|
"update_training_messages (switch_branch) updated 0 rows for {} (lib {}), \
|
|
concurrent regenerate likely flipped is_current",
|
|
normalized,
|
|
library_id
|
|
);
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
/// Return branch metadata and the active leaf ID for a photo's
|
|
/// conversation tree.
|
|
///
|
|
/// With `node_id = None` this lists every leaf in the tree. With
|
|
/// `node_id = Some(fork)` (from `ForkInfo.node_id`) it lists only the
|
|
/// sibling branches diverging at that node, position-ranked; options in
|
|
/// the same subtree as `viewing_leaf` anchor to that leaf so the client
|
|
/// can recognise "the branch I'm on".
|
|
pub fn get_branches(
|
|
&self,
|
|
library_id: i32,
|
|
file_path: &str,
|
|
node_id: Option<u64>,
|
|
viewing_leaf: Option<u64>,
|
|
) -> Result<(Vec<BranchLeafInfo>, u64)> {
|
|
let normalized = normalize_path(file_path);
|
|
let cx = opentelemetry::Context::new();
|
|
let mut dao = self.insight_dao.lock().expect("Unable to lock InsightDao");
|
|
let insight = dao
|
|
.get_current_insight_for_library(&cx, library_id, &normalized)
|
|
.map_err(|e| anyhow!("failed to load insight: {:?}", e))?
|
|
.ok_or_else(|| anyhow!("no insight found for path"))?;
|
|
|
|
let raw = insight
|
|
.training_messages
|
|
.as_ref()
|
|
.ok_or_else(|| anyhow!("insight has no chat history"))?;
|
|
|
|
let store: ChatHistoryStore = if let Ok(arr) = serde_json::from_str::<Vec<ChatMessage>>(raw)
|
|
{
|
|
ChatHistoryStore::from_flat_array(arr)
|
|
} else {
|
|
serde_json::from_str(raw)
|
|
.map_err(|e| anyhow!("failed to deserialize chat history: {}", e))?
|
|
};
|
|
|
|
let list = match node_id {
|
|
Some(fork_node) => {
|
|
if !store.nodes.iter().any(|n| n.id == fork_node) {
|
|
bail!("node_id {} not found in tree", fork_node);
|
|
}
|
|
store.branch_options_at(fork_node, viewing_leaf.unwrap_or(store.active_leaf_id))
|
|
}
|
|
None => store.leaves_with_info(),
|
|
};
|
|
Ok((list, store.active_leaf_id))
|
|
}
|
|
|
|
/// Streaming variant of `chat_turn`. Emits user-facing events as the
|
|
/// conversation progresses: iteration starts, tool dispatch + result,
|
|
/// text deltas from the final assistant reply, and a terminal `Done`
|
|
/// frame. Persistence happens inside the stream after the loop ends.
|
|
///
|
|
/// The stream takes ownership of the service via `Arc<Self>` (passed by
|
|
/// the caller) so it can live past the handler's await boundary.
|
|
pub fn chat_turn_stream(
|
|
self: Arc<Self>,
|
|
req: ChatTurnRequest,
|
|
) -> BoxStream<'static, ChatStreamEvent> {
|
|
let svc = self;
|
|
let s = async_stream::stream! {
|
|
match svc.chat_turn_stream_inner(req, Ok).await {
|
|
Ok(mut rx) => {
|
|
while let Some(ev) = rx.recv().await {
|
|
yield ev;
|
|
}
|
|
}
|
|
Err(e) => {
|
|
yield ChatStreamEvent::Error(format!("{}", e));
|
|
}
|
|
}
|
|
};
|
|
Box::pin(s)
|
|
}
|
|
|
|
/// Internal: drives the streaming loop on a background task, returning
|
|
/// a receiver the caller drains. Keeping the work on a spawned task
|
|
/// decouples the HTTP request lifetime from the chat execution, which
|
|
/// matters because the chat may run longer than any single network hop
|
|
/// and we want clean cancellation semantics via the channel close.
|
|
async fn chat_turn_stream_inner<F>(
|
|
self: Arc<Self>,
|
|
req: ChatTurnRequest,
|
|
_ev_mapper: F,
|
|
) -> Result<tokio::sync::mpsc::Receiver<ChatStreamEvent>>
|
|
where
|
|
F: Fn(ChatStreamEvent) -> Result<ChatStreamEvent> + Send + 'static,
|
|
{
|
|
let (tx, rx) = tokio::sync::mpsc::channel::<ChatStreamEvent>(64);
|
|
let svc = self.clone();
|
|
tokio::spawn(async move {
|
|
let result = svc.run_streaming_turn(req, tx.clone()).await;
|
|
if let Err(e) = result {
|
|
let _ = tx.send(ChatStreamEvent::Error(format!("{}", e))).await;
|
|
}
|
|
});
|
|
Ok(rx)
|
|
}
|
|
|
|
/// Async turn dispatch: creates a TurnEntry in the registry, spawns the
|
|
/// agentic loop on a Tokio task, and returns the turn_id immediately.
|
|
/// Events are buffered in the TurnEntry for SSE replay.
|
|
pub async fn chat_turn_async(
|
|
self: Arc<Self>,
|
|
registry: Arc<TurnRegistry>,
|
|
req: ChatTurnRequest,
|
|
) -> String {
|
|
let turn_id = Uuid::new_v4().to_string();
|
|
let entry = Arc::new(TurnEntry::new(
|
|
turn_id.clone(),
|
|
req.file_path.clone(),
|
|
req.library_id,
|
|
));
|
|
registry.insert(entry.clone()).await;
|
|
|
|
let svc = self.clone();
|
|
let entry_clone = entry.clone();
|
|
let turn_id_for_span = turn_id.clone();
|
|
let library_id = req.library_id;
|
|
let handle = tokio::spawn(async move {
|
|
// Span covering the whole spawned turn execution. Created here (not
|
|
// in the HTTP handler) because the dispatch span ends at the 202
|
|
// response, long before this work runs.
|
|
let tracer = global_tracer();
|
|
let mut span = tracer.start("ai.chat.turn.execute");
|
|
span.set_attribute(KeyValue::new("turn_id", turn_id_for_span));
|
|
span.set_attribute(KeyValue::new("library_id", library_id as i64));
|
|
|
|
let result = svc
|
|
.run_streaming_turn_with_entry(req, entry_clone.clone())
|
|
.await;
|
|
if let Err(ref e) = result {
|
|
span.set_attribute(KeyValue::new("status", "error"));
|
|
span.set_status(Status::error(format!("{e}")));
|
|
// Push the terminal event BEFORE flipping status: a replay
|
|
// reader treats a terminal status with no buffered tail as
|
|
// "closed", so the Error must be in the buffer first.
|
|
let _ = entry_clone
|
|
.push_event(ChatStreamEvent::Error(format!("{}", e)))
|
|
.await;
|
|
entry_clone.set_terminal_status(crate::ai::turn_registry::TurnStatus::Error);
|
|
} else {
|
|
span.set_attribute(KeyValue::new("status", "done"));
|
|
span.set_status(Status::Ok);
|
|
}
|
|
});
|
|
|
|
// Install the abort handle so DELETE can actually stop the task.
|
|
entry.set_abort_handle(handle.abort_handle());
|
|
|
|
turn_id
|
|
}
|
|
|
|
/// Variant of `run_streaming_turn` that pushes events to a `TurnEntry`
|
|
/// buffer instead of an `mpsc::Sender`.
|
|
async fn run_streaming_turn_with_entry(
|
|
self: Arc<Self>,
|
|
req: ChatTurnRequest,
|
|
entry: Arc<TurnEntry>,
|
|
) -> Result<()> {
|
|
if req.user_message.trim().is_empty() {
|
|
bail!("user_message must not be empty");
|
|
}
|
|
if req.user_message.len() > 8192 {
|
|
bail!("user_message exceeds 8192 chars");
|
|
}
|
|
let normalized = normalize_path(&req.file_path);
|
|
|
|
let lock_key = (req.library_id, normalized.clone());
|
|
let entry_lock = {
|
|
let mut locks = self.chat_locks.lock().await;
|
|
locks
|
|
.entry(lock_key.clone())
|
|
.or_insert_with(|| Arc::new(TokioMutex::new(())))
|
|
.clone()
|
|
};
|
|
let _guard = entry_lock.lock().await;
|
|
|
|
// Look up existing insight scoped to this turn's library_id.
|
|
let existing_insight = {
|
|
let cx = opentelemetry::Context::new();
|
|
let mut dao = self.insight_dao.lock().expect("Unable to lock InsightDao");
|
|
dao.get_current_insight_for_library(&cx, req.library_id, &normalized)
|
|
.map_err(|e| anyhow!("failed to load insight: {:?}", e))?
|
|
};
|
|
|
|
if req.regenerate || existing_insight.is_none() {
|
|
return self
|
|
.run_bootstrap_streaming_with_entry(req, normalized, entry)
|
|
.await;
|
|
}
|
|
let insight = existing_insight.expect("just checked Some above");
|
|
self.run_continuation_streaming_with_entry(req, normalized, insight, entry)
|
|
.await
|
|
}
|
|
|
|
/// Continuation path with TurnEntry buffer.
|
|
async fn run_continuation_streaming_with_entry(
|
|
&self,
|
|
req: ChatTurnRequest,
|
|
normalized: String,
|
|
insight: crate::database::models::PhotoInsight,
|
|
entry: Arc<TurnEntry>,
|
|
) -> Result<()> {
|
|
let active_persona = req
|
|
.persona_id
|
|
.clone()
|
|
.filter(|s| !s.trim().is_empty())
|
|
.unwrap_or_else(|| "default".to_string());
|
|
let raw_history = insight.training_messages.as_ref().ok_or_else(|| {
|
|
anyhow!("insight has no chat history; regenerate this insight in agentic mode")
|
|
})?;
|
|
|
|
// Backward-compatible: flat array (old) vs tree (new).
|
|
let mut store: ChatHistoryStore =
|
|
if let Ok(arr) = serde_json::from_str::<Vec<ChatMessage>>(raw_history) {
|
|
ChatHistoryStore::from_flat_array(arr)
|
|
} else {
|
|
serde_json::from_str(raw_history)
|
|
.map_err(|e| anyhow!("failed to deserialize chat history: {}", e))?
|
|
};
|
|
|
|
let path = store
|
|
.path_to_leaf(store.active_leaf_id)
|
|
.ok_or_else(|| anyhow!("active_leaf_id {} not found in tree", store.active_leaf_id))?;
|
|
let path_len = path.len();
|
|
let mut messages: Vec<ChatMessage> = path.iter().map(|n| n.message.clone()).collect();
|
|
|
|
let stored_backend = insight.backend.clone();
|
|
let effective_backend = req
|
|
.backend
|
|
.as_deref()
|
|
.map(|s| s.trim().to_lowercase())
|
|
.filter(|s| !s.is_empty())
|
|
.unwrap_or_else(|| stored_backend.clone());
|
|
let kind = BackendKind::parse(&effective_backend)?;
|
|
validate_cross_replay(&stored_backend, kind.as_str())?;
|
|
|
|
let max_iterations = req
|
|
.max_iterations
|
|
.unwrap_or(DEFAULT_MAX_ITERATIONS)
|
|
.clamp(1, env_max_iterations());
|
|
|
|
let stored_model = insight.model_version.clone();
|
|
let overrides = SamplingOverrides {
|
|
model: req
|
|
.model
|
|
.clone()
|
|
.or_else(|| Some(stored_model.clone()))
|
|
.filter(|m| !m.is_empty()),
|
|
num_ctx: req.num_ctx,
|
|
temperature: req.temperature,
|
|
top_p: req.top_p,
|
|
top_k: req.top_k,
|
|
min_p: req.min_p,
|
|
enable_thinking: req.enable_thinking,
|
|
};
|
|
let backend = self.generator.resolve_backend(kind, &overrides).await?;
|
|
let model_used = backend.model().to_string();
|
|
|
|
let local_first_user_has_image = messages
|
|
.iter()
|
|
.find(|m| m.role == "user")
|
|
.and_then(|m| m.images.as_ref())
|
|
.map(|imgs| !imgs.is_empty())
|
|
.unwrap_or(false);
|
|
let offer_describe_tool = backend.images_inline && local_first_user_has_image;
|
|
let gate_opts = self.generator.current_gate_opts_for_persona(
|
|
offer_describe_tool,
|
|
Some((req.user_id, &active_persona)),
|
|
);
|
|
let tools = InsightGenerator::build_tool_definitions(gate_opts);
|
|
|
|
let image_base64: Option<String> = if offer_describe_tool {
|
|
self.generator.load_image_as_base64(&normalized).ok()
|
|
} else {
|
|
None
|
|
};
|
|
|
|
let budget_tokens = (req.num_ctx.unwrap_or_else(env_default_num_ctx) as usize)
|
|
.saturating_sub(RESPONSE_HEADROOM_TOKENS);
|
|
let budget_bytes = budget_tokens.saturating_mul(BYTES_PER_TOKEN);
|
|
let truncated = apply_context_budget(&mut messages, budget_bytes);
|
|
if truncated {
|
|
let _ = entry.push_event(ChatStreamEvent::Truncated).await;
|
|
}
|
|
|
|
messages.push(ChatMessage::user(req.user_message.clone()));
|
|
|
|
let override_stash =
|
|
apply_system_prompt_override(&mut messages, req.system_prompt.as_deref());
|
|
let original_system_content = annotate_system_with_budget(&mut messages, max_iterations);
|
|
|
|
let outcome = self
|
|
.run_streaming_agentic_loop_with_entry(
|
|
&backend,
|
|
&mut messages,
|
|
tools,
|
|
&image_base64,
|
|
&normalized,
|
|
req.user_id,
|
|
&active_persona,
|
|
max_iterations,
|
|
&entry,
|
|
)
|
|
.await?;
|
|
let AgenticLoopOutcome {
|
|
tool_calls_made,
|
|
iterations_used,
|
|
last_prompt_eval_count,
|
|
last_eval_count,
|
|
final_content,
|
|
cancelled,
|
|
} = outcome;
|
|
|
|
// Turn was cancelled mid-flight: the DELETE handler already pushed the
|
|
// terminal event and flipped status. Don't persist a partial turn or
|
|
// push a second terminal event.
|
|
if cancelled {
|
|
return Ok(());
|
|
}
|
|
|
|
restore_system_content(&mut messages, original_system_content);
|
|
|
|
if !req.amend {
|
|
restore_system_prompt_override(&mut messages, override_stash);
|
|
}
|
|
|
|
// Append new messages as tree nodes.
|
|
let new_messages = messages[path_len..].to_vec();
|
|
let mut parent_id = Some(store.active_leaf_id);
|
|
for msg in &new_messages {
|
|
let new_id = store.append_node(parent_id, msg.clone());
|
|
parent_id = Some(new_id);
|
|
}
|
|
if let Some(last_id) = parent_id {
|
|
store.active_leaf_id = last_id;
|
|
}
|
|
|
|
let json = serde_json::to_string(&store)
|
|
.map_err(|e| anyhow!("failed to serialize chat tree: {}", e))?;
|
|
|
|
let mut amended_insight_id: Option<i32> = None;
|
|
if req.amend {
|
|
let (title, body) = crate::ai::insight_generator::parse_title_body(&final_content);
|
|
let final_content = body;
|
|
|
|
let new_row = InsertPhotoInsight {
|
|
library_id: req.library_id,
|
|
file_path: normalized.clone(),
|
|
title,
|
|
summary: final_content.clone(),
|
|
generated_at: Utc::now().timestamp(),
|
|
model_version: model_used.clone(),
|
|
is_current: true,
|
|
training_messages: Some(json),
|
|
backend: kind.as_str().to_string(),
|
|
fewshot_source_ids: None,
|
|
content_hash: None,
|
|
num_ctx: req.num_ctx,
|
|
temperature: req.temperature,
|
|
top_p: req.top_p,
|
|
top_k: req.top_k,
|
|
min_p: req.min_p,
|
|
system_prompt: req.system_prompt.clone(),
|
|
persona_id: req.persona_id.clone(),
|
|
prompt_eval_count: None,
|
|
eval_count: None,
|
|
};
|
|
let dao = self.insight_dao.clone();
|
|
let stored = retry_with_backoff("store_insight", 3, || {
|
|
let cx = opentelemetry::Context::new();
|
|
let mut d = dao.lock().expect("Unable to lock InsightDao");
|
|
d.store_insight(&cx, new_row.clone())
|
|
})
|
|
.map_err(|e| anyhow!("failed to store amended insight: {:?}", e))?;
|
|
amended_insight_id = Some(stored.id);
|
|
} else {
|
|
let dao = self.insight_dao.clone();
|
|
let rows = retry_with_backoff("update_training_messages", 3, || {
|
|
let cx = opentelemetry::Context::new();
|
|
let mut d = dao.lock().expect("Unable to lock InsightDao");
|
|
d.update_training_messages(&cx, req.library_id, &normalized, &json)
|
|
})
|
|
.map_err(|e| anyhow!("failed to persist chat history: {:?}", e))?;
|
|
if rows == 0 {
|
|
log::warn!(
|
|
"update_training_messages (stream) updated 0 rows for {} (lib {}), \
|
|
concurrent regenerate likely flipped is_current",
|
|
normalized,
|
|
req.library_id
|
|
);
|
|
}
|
|
}
|
|
|
|
let _ = entry
|
|
.push_event(ChatStreamEvent::Done {
|
|
tool_calls_made,
|
|
iterations_used,
|
|
truncated,
|
|
prompt_tokens: last_prompt_eval_count,
|
|
eval_tokens: last_eval_count,
|
|
num_ctx: req.num_ctx,
|
|
amended_insight_id,
|
|
backend_used: kind.as_str().to_string(),
|
|
model_used,
|
|
cancelled: false,
|
|
})
|
|
.await;
|
|
|
|
entry.set_terminal_status(crate::ai::turn_registry::TurnStatus::Done);
|
|
Ok(())
|
|
}
|
|
|
|
/// Bootstrap path with TurnEntry buffer.
|
|
async fn run_bootstrap_streaming_with_entry(
|
|
&self,
|
|
req: ChatTurnRequest,
|
|
normalized: String,
|
|
entry: Arc<TurnEntry>,
|
|
) -> Result<()> {
|
|
let active_persona = req
|
|
.persona_id
|
|
.clone()
|
|
.filter(|s| !s.trim().is_empty())
|
|
.unwrap_or_else(|| "default".to_string());
|
|
let effective_backend = resolve_bootstrap_backend(req.backend.as_deref())?;
|
|
let kind = BackendKind::parse(&effective_backend)?;
|
|
|
|
let max_iterations = req
|
|
.max_iterations
|
|
.unwrap_or(DEFAULT_MAX_ITERATIONS)
|
|
.clamp(1, env_max_iterations());
|
|
|
|
let overrides = SamplingOverrides {
|
|
model: req.model.clone().filter(|m| !m.is_empty()),
|
|
num_ctx: req.num_ctx,
|
|
temperature: req.temperature,
|
|
top_p: req.top_p,
|
|
top_k: req.top_k,
|
|
min_p: req.min_p,
|
|
enable_thinking: req.enable_thinking,
|
|
};
|
|
let backend = self.generator.resolve_backend(kind, &overrides).await?;
|
|
let model_used = backend.model().to_string();
|
|
|
|
let image_base64: Option<String> = self.generator.load_image_as_base64(&normalized).ok();
|
|
|
|
let exif = self.generator.fetch_exif(&normalized);
|
|
let date_taken_str = resolve_date_taken_for_context(&exif, &normalized);
|
|
let gps = exif
|
|
.as_ref()
|
|
.and_then(|e| match (e.gps_latitude, e.gps_longitude) {
|
|
(Some(lat), Some(lon)) => Some((lat as f64, lon as f64)),
|
|
_ => None,
|
|
});
|
|
|
|
let visual_block = if !backend.images_inline {
|
|
match image_base64.as_deref() {
|
|
Some(b64) => match backend.local().describe_image(b64).await {
|
|
Ok(desc) => {
|
|
format!("Visual description (from local vision model):\n{}\n", desc)
|
|
}
|
|
Err(e) => {
|
|
log::warn!("{} bootstrap: describe_image failed: {}", kind.as_str(), e);
|
|
String::new()
|
|
}
|
|
},
|
|
None => String::new(),
|
|
}
|
|
} else {
|
|
String::new()
|
|
};
|
|
|
|
let offer_describe_tool = backend.images_inline && image_base64.is_some();
|
|
let gate_opts = self.generator.current_gate_opts_for_persona(
|
|
offer_describe_tool,
|
|
Some((req.user_id, &active_persona)),
|
|
);
|
|
let tools = InsightGenerator::build_tool_definitions(gate_opts);
|
|
|
|
// Server-side persona resolution: explicit client system_prompt wins;
|
|
// else the active persona's stored prompt; else the neutral default.
|
|
let persona_prompt = self
|
|
.generator
|
|
.persona_system_prompt(req.user_id, &active_persona);
|
|
let persona = resolve_bootstrap_system_prompt(req.system_prompt.as_deref(), persona_prompt);
|
|
let system_content = build_bootstrap_system_message(
|
|
&persona,
|
|
&normalized,
|
|
date_taken_str.as_deref(),
|
|
gps,
|
|
&visual_block,
|
|
);
|
|
let system_msg = ChatMessage::system(system_content);
|
|
let mut user_msg = ChatMessage::user(req.user_message.clone());
|
|
if backend.images_inline
|
|
&& let Some(ref img) = image_base64
|
|
{
|
|
user_msg.images = Some(vec![img.clone()]);
|
|
}
|
|
let mut messages = vec![system_msg, user_msg];
|
|
|
|
let outcome = self
|
|
.run_streaming_agentic_loop_with_entry(
|
|
&backend,
|
|
&mut messages,
|
|
tools,
|
|
&image_base64,
|
|
&normalized,
|
|
req.user_id,
|
|
&active_persona,
|
|
max_iterations,
|
|
&entry,
|
|
)
|
|
.await?;
|
|
let AgenticLoopOutcome {
|
|
tool_calls_made,
|
|
iterations_used,
|
|
last_prompt_eval_count,
|
|
last_eval_count,
|
|
final_content,
|
|
cancelled,
|
|
} = outcome;
|
|
|
|
// Turn was cancelled mid-flight: the DELETE handler already pushed the
|
|
// terminal event and flipped status. Don't persist a partial turn or
|
|
// push a second terminal event.
|
|
if cancelled {
|
|
return Ok(());
|
|
}
|
|
|
|
let (title, body) = crate::ai::insight_generator::parse_title_body(&final_content);
|
|
|
|
let json = serde_json::to_string(&messages)
|
|
.map_err(|e| anyhow!("failed to serialize chat history: {}", e))?;
|
|
let new_row = InsertPhotoInsight {
|
|
library_id: req.library_id,
|
|
file_path: normalized.clone(),
|
|
title,
|
|
summary: body,
|
|
generated_at: Utc::now().timestamp(),
|
|
model_version: model_used.clone(),
|
|
is_current: true,
|
|
training_messages: Some(json),
|
|
backend: kind.as_str().to_string(),
|
|
fewshot_source_ids: None,
|
|
content_hash: None,
|
|
num_ctx: req.num_ctx,
|
|
temperature: req.temperature,
|
|
top_p: req.top_p,
|
|
top_k: req.top_k,
|
|
min_p: req.min_p,
|
|
system_prompt: req.system_prompt.clone(),
|
|
persona_id: req.persona_id.clone(),
|
|
prompt_eval_count: None,
|
|
eval_count: None,
|
|
};
|
|
let dao = self.insight_dao.clone();
|
|
let stored = retry_with_backoff("store_insight", 3, || {
|
|
let cx = opentelemetry::Context::new();
|
|
let mut d = dao.lock().expect("Unable to lock InsightDao");
|
|
d.store_insight(&cx, new_row.clone())
|
|
})
|
|
.map_err(|e| anyhow!("failed to store bootstrap insight: {:?}", e))?;
|
|
|
|
let _ = entry
|
|
.push_event(ChatStreamEvent::Done {
|
|
tool_calls_made,
|
|
iterations_used,
|
|
truncated: false,
|
|
prompt_tokens: last_prompt_eval_count,
|
|
eval_tokens: last_eval_count,
|
|
num_ctx: req.num_ctx,
|
|
amended_insight_id: Some(stored.id),
|
|
backend_used: kind.as_str().to_string(),
|
|
model_used,
|
|
cancelled: false,
|
|
})
|
|
.await;
|
|
|
|
entry.set_terminal_status(crate::ai::turn_registry::TurnStatus::Done);
|
|
Ok(())
|
|
}
|
|
|
|
/// Agentic loop variant that pushes events to a `TurnEntry` buffer.
|
|
/// Same as `run_streaming_agentic_loop` but emits events to a
|
|
/// `TurnEntry` for SSE replay. Thin wrapper around the free function
|
|
/// `run_streaming_agentic_loop_with_entry` so the persona chat
|
|
/// (and any future chat-without-file surface) can reuse the loop
|
|
/// body without instantiating an `InsightChatService`.
|
|
pub async fn run_streaming_agentic_loop_with_entry(
|
|
&self,
|
|
backend: &ResolvedBackend,
|
|
messages: &mut Vec<ChatMessage>,
|
|
tools: Vec<Tool>,
|
|
image_base64: &Option<String>,
|
|
normalized: &str,
|
|
user_id: i32,
|
|
active_persona: &str,
|
|
max_iterations: usize,
|
|
entry: &Arc<TurnEntry>,
|
|
) -> Result<AgenticLoopOutcome> {
|
|
crate::ai::insight_chat::run_streaming_agentic_loop_with_entry(
|
|
&self.generator,
|
|
backend,
|
|
messages,
|
|
tools,
|
|
image_base64,
|
|
normalized,
|
|
user_id,
|
|
active_persona,
|
|
max_iterations,
|
|
entry,
|
|
)
|
|
.await
|
|
}
|
|
|
|
async fn run_streaming_turn(
|
|
self: Arc<Self>,
|
|
req: ChatTurnRequest,
|
|
tx: tokio::sync::mpsc::Sender<ChatStreamEvent>,
|
|
) -> Result<()> {
|
|
if req.user_message.trim().is_empty() {
|
|
bail!("user_message must not be empty");
|
|
}
|
|
if req.user_message.len() > 8192 {
|
|
bail!("user_message exceeds 8192 chars");
|
|
}
|
|
let normalized = normalize_path(&req.file_path);
|
|
|
|
let lock_key = (req.library_id, normalized.clone());
|
|
let entry_lock = {
|
|
let mut locks = self.chat_locks.lock().await;
|
|
locks
|
|
.entry(lock_key.clone())
|
|
.or_insert_with(|| Arc::new(TokioMutex::new(())))
|
|
.clone()
|
|
};
|
|
let _guard = entry_lock.lock().await;
|
|
|
|
// Look up existing insight scoped to this turn's library_id.
|
|
// Path-only lookup would let an `is_current=true` row in
|
|
// another library route us into the continuation path against
|
|
// a transcript we'd then update_training_messages on — corrupting
|
|
// the other library's curated insight. Library-scoped lookup
|
|
// means a fresh chat on a photo that has no insight in this
|
|
// library bootstraps cleanly, even when another library has
|
|
// an insight for the same rel_path.
|
|
let existing_insight = {
|
|
let cx = opentelemetry::Context::new();
|
|
let mut dao = self.insight_dao.lock().expect("Unable to lock InsightDao");
|
|
dao.get_current_insight_for_library(&cx, req.library_id, &normalized)
|
|
.map_err(|e| anyhow!("failed to load insight: {:?}", e))?
|
|
};
|
|
|
|
if req.regenerate || existing_insight.is_none() {
|
|
return self.run_bootstrap_streaming(req, normalized, tx).await;
|
|
}
|
|
let insight = existing_insight.expect("just checked Some above");
|
|
self.run_continuation_streaming(req, normalized, insight, tx)
|
|
.await
|
|
}
|
|
|
|
/// Continuation path: photo has an existing agentic insight with
|
|
/// `training_messages` populated. Replay the transcript, append a new
|
|
/// turn, run the agentic loop, persist (UPDATE for append; INSERT new
|
|
/// row for amend).
|
|
async fn run_continuation_streaming(
|
|
&self,
|
|
req: ChatTurnRequest,
|
|
normalized: String,
|
|
insight: crate::database::models::PhotoInsight,
|
|
tx: tokio::sync::mpsc::Sender<ChatStreamEvent>,
|
|
) -> Result<()> {
|
|
let active_persona = req
|
|
.persona_id
|
|
.clone()
|
|
.filter(|s| !s.trim().is_empty())
|
|
.unwrap_or_else(|| "default".to_string());
|
|
let raw_history = insight.training_messages.as_ref().ok_or_else(|| {
|
|
anyhow!("insight has no chat history; regenerate this insight in agentic mode")
|
|
})?;
|
|
let mut messages: Vec<ChatMessage> = serde_json::from_str(raw_history)
|
|
.map_err(|e| anyhow!("failed to deserialize chat history: {}", e))?;
|
|
|
|
// Backend selection — defer to stored insight's backend unless the
|
|
// request supplies an override.
|
|
let stored_backend = insight.backend.clone();
|
|
let effective_backend = req
|
|
.backend
|
|
.as_deref()
|
|
.map(|s| s.trim().to_lowercase())
|
|
.filter(|s| !s.is_empty())
|
|
.unwrap_or_else(|| stored_backend.clone());
|
|
let kind = BackendKind::parse(&effective_backend)?;
|
|
validate_cross_replay(&stored_backend, kind.as_str())?;
|
|
|
|
let max_iterations = req
|
|
.max_iterations
|
|
.unwrap_or(DEFAULT_MAX_ITERATIONS)
|
|
.clamp(1, env_max_iterations());
|
|
|
|
let stored_model = insight.model_version.clone();
|
|
let overrides = SamplingOverrides {
|
|
model: req
|
|
.model
|
|
.clone()
|
|
.or_else(|| Some(stored_model.clone()))
|
|
.filter(|m| !m.is_empty()),
|
|
num_ctx: req.num_ctx,
|
|
temperature: req.temperature,
|
|
top_p: req.top_p,
|
|
top_k: req.top_k,
|
|
min_p: req.min_p,
|
|
enable_thinking: req.enable_thinking,
|
|
};
|
|
let backend = self.generator.resolve_backend(kind, &overrides).await?;
|
|
let model_used = backend.model().to_string();
|
|
|
|
// Tool set — images_inline mode + first user turn carries an image →
|
|
// offer describe_photo. Describe-then-inline mode (hybrid only):
|
|
// visual description was inlined at bootstrap, no describe tool needed.
|
|
let local_first_user_has_image = messages
|
|
.iter()
|
|
.find(|m| m.role == "user")
|
|
.and_then(|m| m.images.as_ref())
|
|
.map(|imgs| !imgs.is_empty())
|
|
.unwrap_or(false);
|
|
let offer_describe_tool = backend.images_inline && local_first_user_has_image;
|
|
let gate_opts = self.generator.current_gate_opts_for_persona(
|
|
offer_describe_tool,
|
|
Some((req.user_id, &active_persona)),
|
|
);
|
|
let tools = InsightGenerator::build_tool_definitions(gate_opts);
|
|
|
|
let image_base64: Option<String> = if offer_describe_tool {
|
|
self.generator.load_image_as_base64(&normalized).ok()
|
|
} else {
|
|
None
|
|
};
|
|
|
|
// Truncate before appending the new user turn.
|
|
let budget_tokens = (req.num_ctx.unwrap_or_else(env_default_num_ctx) as usize)
|
|
.saturating_sub(RESPONSE_HEADROOM_TOKENS);
|
|
let budget_bytes = budget_tokens.saturating_mul(BYTES_PER_TOKEN);
|
|
let truncated = apply_context_budget(&mut messages, budget_bytes);
|
|
if truncated {
|
|
let _ = tx.send(ChatStreamEvent::Truncated).await;
|
|
}
|
|
|
|
messages.push(ChatMessage::user(req.user_message.clone()));
|
|
|
|
// Mirror chat_turn: per-turn override goes on first, budget note next.
|
|
let override_stash =
|
|
apply_system_prompt_override(&mut messages, req.system_prompt.as_deref());
|
|
let original_system_content = annotate_system_with_budget(&mut messages, max_iterations);
|
|
|
|
let outcome = self
|
|
.run_streaming_agentic_loop(
|
|
&backend,
|
|
&mut messages,
|
|
tools,
|
|
&image_base64,
|
|
&normalized,
|
|
req.user_id,
|
|
&active_persona,
|
|
max_iterations,
|
|
&tx,
|
|
)
|
|
.await?;
|
|
let AgenticLoopOutcome {
|
|
tool_calls_made,
|
|
iterations_used,
|
|
last_prompt_eval_count,
|
|
last_eval_count,
|
|
final_content,
|
|
// The mpsc (legacy) path has no cancellation channel.
|
|
cancelled: _,
|
|
} = outcome;
|
|
|
|
// Drop the per-turn iteration-budget note before persisting so it
|
|
// doesn't snowball on subsequent turns.
|
|
restore_system_content(&mut messages, original_system_content);
|
|
|
|
// Append mode: undo the per-turn system-prompt override.
|
|
// Amend mode: keep it — it becomes the new row's system message.
|
|
if !req.amend {
|
|
restore_system_prompt_override(&mut messages, override_stash);
|
|
}
|
|
|
|
let json = serde_json::to_string(&messages)
|
|
.map_err(|e| anyhow!("failed to serialize chat history: {}", e))?;
|
|
|
|
let mut amended_insight_id: Option<i32> = None;
|
|
if req.amend {
|
|
let (title, body) = crate::ai::insight_generator::parse_title_body(&final_content);
|
|
let final_content = body;
|
|
|
|
// Amended rows intentionally do not inherit the parent's
|
|
// `fewshot_source_ids`. The parent's few-shot influence is still
|
|
// present in this row's content; if you want strict lineage
|
|
// tracking for training-set filtering, fetch the parent here and
|
|
// copy its value forward.
|
|
let new_row = InsertPhotoInsight {
|
|
library_id: req.library_id,
|
|
file_path: normalized.clone(),
|
|
title,
|
|
summary: final_content.clone(),
|
|
generated_at: Utc::now().timestamp(),
|
|
model_version: model_used.clone(),
|
|
is_current: true,
|
|
training_messages: Some(json),
|
|
backend: kind.as_str().to_string(),
|
|
fewshot_source_ids: None,
|
|
content_hash: None,
|
|
num_ctx: req.num_ctx,
|
|
temperature: req.temperature,
|
|
top_p: req.top_p,
|
|
top_k: req.top_k,
|
|
min_p: req.min_p,
|
|
system_prompt: req.system_prompt.clone(),
|
|
persona_id: req.persona_id.clone(),
|
|
prompt_eval_count: None,
|
|
eval_count: None,
|
|
};
|
|
let dao = self.insight_dao.clone();
|
|
let stored = retry_with_backoff("store_insight", 3, || {
|
|
let cx = opentelemetry::Context::new();
|
|
let mut d = dao.lock().expect("Unable to lock InsightDao");
|
|
d.store_insight(&cx, new_row.clone())
|
|
})
|
|
.map_err(|e| anyhow!("failed to store amended insight: {:?}", e))?;
|
|
amended_insight_id = Some(stored.id);
|
|
} else {
|
|
let dao = self.insight_dao.clone();
|
|
let rows = retry_with_backoff("update_training_messages", 3, || {
|
|
let cx = opentelemetry::Context::new();
|
|
let mut d = dao.lock().expect("Unable to lock InsightDao");
|
|
d.update_training_messages(&cx, req.library_id, &normalized, &json)
|
|
})
|
|
.map_err(|e| anyhow!("failed to persist chat history: {:?}", e))?;
|
|
if rows == 0 {
|
|
log::warn!(
|
|
"update_training_messages (stream) updated 0 rows for {} (lib {}), \
|
|
concurrent regenerate likely flipped is_current",
|
|
normalized,
|
|
req.library_id
|
|
);
|
|
}
|
|
}
|
|
|
|
let _ = tx
|
|
.send(ChatStreamEvent::Done {
|
|
tool_calls_made,
|
|
iterations_used,
|
|
truncated,
|
|
prompt_tokens: last_prompt_eval_count,
|
|
eval_tokens: last_eval_count,
|
|
num_ctx: req.num_ctx,
|
|
amended_insight_id,
|
|
backend_used: kind.as_str().to_string(),
|
|
model_used,
|
|
cancelled: false,
|
|
})
|
|
.await;
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Bootstrap path: no insight row yet (or `regenerate=true`). Build a
|
|
/// fresh transcript from `req.system_prompt` + `req.user_message` +
|
|
/// the photo, run the agentic loop, generate a title, and INSERT a
|
|
/// new insight row. `store_insight` flips any prior rows for the same
|
|
/// `(library_id, file_path)` to `is_current=false` — that's how
|
|
/// `regenerate` shadows the previous insight.
|
|
async fn run_bootstrap_streaming(
|
|
&self,
|
|
req: ChatTurnRequest,
|
|
normalized: String,
|
|
tx: tokio::sync::mpsc::Sender<ChatStreamEvent>,
|
|
) -> Result<()> {
|
|
let active_persona = req
|
|
.persona_id
|
|
.clone()
|
|
.filter(|s| !s.trim().is_empty())
|
|
.unwrap_or_else(|| "default".to_string());
|
|
let effective_backend = resolve_bootstrap_backend(req.backend.as_deref())?;
|
|
let kind = BackendKind::parse(&effective_backend)?;
|
|
|
|
let max_iterations = req
|
|
.max_iterations
|
|
.unwrap_or(DEFAULT_MAX_ITERATIONS)
|
|
.clamp(1, env_max_iterations());
|
|
|
|
let overrides = SamplingOverrides {
|
|
model: req.model.clone().filter(|m| !m.is_empty()),
|
|
num_ctx: req.num_ctx,
|
|
temperature: req.temperature,
|
|
top_p: req.top_p,
|
|
top_k: req.top_k,
|
|
min_p: req.min_p,
|
|
enable_thinking: req.enable_thinking,
|
|
};
|
|
let backend = self.generator.resolve_backend(kind, &overrides).await?;
|
|
let model_used = backend.model().to_string();
|
|
|
|
// Load image bytes once. RAW preview fallback is handled inside
|
|
// load_image_as_base64. Errors degrade silently — a chat that
|
|
// discusses metadata-only is still useful.
|
|
let image_base64: Option<String> = self.generator.load_image_as_base64(&normalized).ok();
|
|
|
|
// EXIF lookup once — date_taken and GPS go into the photo
|
|
// context block in the system message. Without these the model
|
|
// hallucinates dates / GPS-keyed tool args (`get_sms_messages`
|
|
// would otherwise default to today's date and miss every
|
|
// historical photo).
|
|
let exif = self.generator.fetch_exif(&normalized);
|
|
let date_taken_str = resolve_date_taken_for_context(&exif, &normalized);
|
|
let gps = exif
|
|
.as_ref()
|
|
.and_then(|e| match (e.gps_latitude, e.gps_longitude) {
|
|
(Some(lat), Some(lon)) => Some((lat as f64, lon as f64)),
|
|
_ => None,
|
|
});
|
|
|
|
// Describe-then-inline (hybrid only): pre-describe the image so a
|
|
// text-only chat model gets the visual description inline.
|
|
// 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 {
|
|
Ok(desc) => {
|
|
format!("Visual description (from local vision model):\n{}\n", desc)
|
|
}
|
|
Err(e) => {
|
|
log::warn!("{} bootstrap: describe_image failed: {}", kind.as_str(), e);
|
|
String::new()
|
|
}
|
|
},
|
|
None => String::new(),
|
|
}
|
|
} else {
|
|
String::new()
|
|
};
|
|
|
|
// Tool gates. images_inline + image present → expose describe_photo so
|
|
// the chat model can re-look at the photo on demand. Non-inline:
|
|
// already inlined, no tool needed.
|
|
let offer_describe_tool = backend.images_inline && image_base64.is_some();
|
|
let gate_opts = self.generator.current_gate_opts_for_persona(
|
|
offer_describe_tool,
|
|
Some((req.user_id, &active_persona)),
|
|
);
|
|
let tools = InsightGenerator::build_tool_definitions(gate_opts);
|
|
|
|
// System message = persona + photo context block. Photo context
|
|
// is in the system message — not the user turn — so the user's
|
|
// bubble in the rendered transcript shows only what they typed.
|
|
// Several agentic tools (recall_facts_for_photo, get_file_tags,
|
|
// get_faces_in_photo, etc.) take a `file_path` arg the model
|
|
// can't know without being told. `Date taken:` and `GPS:` give
|
|
// get_sms_messages / reverse_geocode / get_personal_place_at
|
|
// the args they need. In hybrid mode the visual description
|
|
// belongs here for the same reason.
|
|
// Server-side persona resolution: explicit client system_prompt wins;
|
|
// else the active persona's stored prompt; else the neutral default.
|
|
let persona_prompt = self
|
|
.generator
|
|
.persona_system_prompt(req.user_id, &active_persona);
|
|
let persona = resolve_bootstrap_system_prompt(req.system_prompt.as_deref(), persona_prompt);
|
|
let system_content = build_bootstrap_system_message(
|
|
&persona,
|
|
&normalized,
|
|
date_taken_str.as_deref(),
|
|
gps,
|
|
&visual_block,
|
|
);
|
|
let system_msg = ChatMessage::system(system_content);
|
|
let mut user_msg = ChatMessage::user(req.user_message.clone());
|
|
if backend.images_inline
|
|
&& let Some(ref img) = image_base64
|
|
{
|
|
user_msg.images = Some(vec![img.clone()]);
|
|
}
|
|
let mut messages = vec![system_msg, user_msg];
|
|
|
|
let outcome = self
|
|
.run_streaming_agentic_loop(
|
|
&backend,
|
|
&mut messages,
|
|
tools,
|
|
&image_base64,
|
|
&normalized,
|
|
req.user_id,
|
|
&active_persona,
|
|
max_iterations,
|
|
&tx,
|
|
)
|
|
.await?;
|
|
let AgenticLoopOutcome {
|
|
tool_calls_made,
|
|
iterations_used,
|
|
last_prompt_eval_count,
|
|
last_eval_count,
|
|
final_content,
|
|
// The mpsc (legacy) path has no cancellation channel.
|
|
cancelled: _,
|
|
} = outcome;
|
|
|
|
let (title, body) = crate::ai::insight_generator::parse_title_body(&final_content);
|
|
|
|
let json = serde_json::to_string(&messages)
|
|
.map_err(|e| anyhow!("failed to serialize chat history: {}", e))?;
|
|
let new_row = InsertPhotoInsight {
|
|
library_id: req.library_id,
|
|
file_path: normalized.clone(),
|
|
title,
|
|
summary: body,
|
|
generated_at: Utc::now().timestamp(),
|
|
model_version: model_used.clone(),
|
|
is_current: true,
|
|
training_messages: Some(json),
|
|
backend: kind.as_str().to_string(),
|
|
fewshot_source_ids: None,
|
|
content_hash: None,
|
|
num_ctx: req.num_ctx,
|
|
temperature: req.temperature,
|
|
top_p: req.top_p,
|
|
top_k: req.top_k,
|
|
min_p: req.min_p,
|
|
system_prompt: req.system_prompt.clone(),
|
|
persona_id: req.persona_id.clone(),
|
|
prompt_eval_count: None,
|
|
eval_count: None,
|
|
};
|
|
let dao = self.insight_dao.clone();
|
|
let stored = retry_with_backoff("store_insight", 3, || {
|
|
let cx = opentelemetry::Context::new();
|
|
let mut d = dao.lock().expect("Unable to lock InsightDao");
|
|
d.store_insight(&cx, new_row.clone())
|
|
})
|
|
.map_err(|e| anyhow!("failed to store bootstrap insight: {:?}", e))?;
|
|
|
|
// amended_insight_id semantics broaden on bootstrap/regenerate:
|
|
// populated whenever this turn produced a new insight row, so
|
|
// clients don't need a separate field to learn the new id.
|
|
let _ = tx
|
|
.send(ChatStreamEvent::Done {
|
|
tool_calls_made,
|
|
iterations_used,
|
|
truncated: false,
|
|
prompt_tokens: last_prompt_eval_count,
|
|
eval_tokens: last_eval_count,
|
|
num_ctx: req.num_ctx,
|
|
amended_insight_id: Some(stored.id),
|
|
backend_used: kind.as_str().to_string(),
|
|
model_used,
|
|
cancelled: false,
|
|
})
|
|
.await;
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Drive the agentic loop with streaming SSE events. Shared between
|
|
/// bootstrap and continuation. Mutates `messages` in place (response
|
|
/// turns + tool results are appended) and returns counters + the
|
|
/// final assistant content.
|
|
async fn run_streaming_agentic_loop(
|
|
&self,
|
|
backend: &ResolvedBackend,
|
|
messages: &mut Vec<ChatMessage>,
|
|
tools: Vec<Tool>,
|
|
image_base64: &Option<String>,
|
|
normalized: &str,
|
|
user_id: i32,
|
|
active_persona: &str,
|
|
max_iterations: usize,
|
|
tx: &tokio::sync::mpsc::Sender<ChatStreamEvent>,
|
|
) -> Result<AgenticLoopOutcome> {
|
|
let mut tool_calls_made = 0usize;
|
|
let mut iterations_used = 0usize;
|
|
let mut last_prompt_eval_count: Option<i32> = None;
|
|
let mut last_eval_count: Option<i32> = None;
|
|
let mut final_content = String::new();
|
|
|
|
for iteration in 0..max_iterations {
|
|
iterations_used = iteration + 1;
|
|
let _ = tx
|
|
.send(ChatStreamEvent::IterationStart {
|
|
n: iterations_used,
|
|
max: max_iterations,
|
|
})
|
|
.await;
|
|
|
|
let mut stream = backend
|
|
.chat()
|
|
.chat_with_tools_stream(messages.clone(), tools.clone())
|
|
.await?;
|
|
|
|
let mut final_message: Option<ChatMessage> = None;
|
|
while let Some(ev) = stream.next().await {
|
|
let ev = ev?;
|
|
match ev {
|
|
LlmStreamEvent::TextDelta(delta) => {
|
|
let _ = tx.send(ChatStreamEvent::TextDelta(delta)).await;
|
|
}
|
|
LlmStreamEvent::Done {
|
|
message,
|
|
prompt_eval_count,
|
|
eval_count,
|
|
} => {
|
|
last_prompt_eval_count = prompt_eval_count;
|
|
last_eval_count = eval_count;
|
|
final_message = Some(message);
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
let mut response =
|
|
final_message.ok_or_else(|| anyhow!("stream ended without a Done event"))?;
|
|
|
|
// Normalize non-object tool arguments (some models occasionally
|
|
// return null/string/bool which Ollama rejects on the next turn).
|
|
if let Some(ref mut tcs) = response.tool_calls {
|
|
for tc in tcs.iter_mut() {
|
|
if !tc.function.arguments.is_object() {
|
|
tc.function.arguments = serde_json::Value::Object(Default::default());
|
|
}
|
|
}
|
|
}
|
|
|
|
messages.push(response.clone());
|
|
|
|
if let Some(ref tool_calls) = response.tool_calls
|
|
&& !tool_calls.is_empty()
|
|
{
|
|
for tool_call in tool_calls {
|
|
tool_calls_made += 1;
|
|
let call_index = tool_calls_made - 1;
|
|
let _ = tx
|
|
.send(ChatStreamEvent::ToolCall {
|
|
index: call_index,
|
|
name: tool_call.function.name.clone(),
|
|
arguments: tool_call.function.arguments.clone(),
|
|
})
|
|
.await;
|
|
let cx = opentelemetry::Context::new();
|
|
let result = self
|
|
.generator
|
|
.execute_tool(
|
|
&tool_call.function.name,
|
|
&tool_call.function.arguments,
|
|
backend,
|
|
image_base64,
|
|
normalized,
|
|
user_id,
|
|
active_persona,
|
|
&cx,
|
|
)
|
|
.await;
|
|
let (result_preview, result_truncated) = truncate_tool_result(&result);
|
|
let _ = tx
|
|
.send(ChatStreamEvent::ToolResult {
|
|
index: call_index,
|
|
name: tool_call.function.name.clone(),
|
|
result: result_preview,
|
|
result_truncated,
|
|
})
|
|
.await;
|
|
messages.push(ChatMessage::tool_result(result));
|
|
}
|
|
continue;
|
|
}
|
|
|
|
final_content = response.content;
|
|
break;
|
|
}
|
|
|
|
// No-tools fallback: loop exhausted iterations without producing a
|
|
// tool-free reply. Ask once more, with no tools attached.
|
|
if final_content.is_empty() {
|
|
// Index of the synthetic "please write your final answer"
|
|
// user message — we strip it from history after the model
|
|
// responds, so it never appears in the rendered transcript
|
|
// and load_history's user-turn handler doesn't reset
|
|
// pending_tools at this position (wiping the prior tool
|
|
// calls from the final assistant render).
|
|
let synthetic_idx = push_synthetic_final_prompt(messages);
|
|
let mut stream = backend
|
|
.chat()
|
|
.chat_with_tools_stream(messages.clone(), vec![])
|
|
.await?;
|
|
let mut final_message: Option<ChatMessage> = None;
|
|
while let Some(ev) = stream.next().await {
|
|
let ev = ev?;
|
|
match ev {
|
|
LlmStreamEvent::TextDelta(delta) => {
|
|
let _ = tx.send(ChatStreamEvent::TextDelta(delta)).await;
|
|
}
|
|
LlmStreamEvent::Done {
|
|
message,
|
|
prompt_eval_count,
|
|
eval_count,
|
|
} => {
|
|
last_prompt_eval_count = prompt_eval_count;
|
|
last_eval_count = eval_count;
|
|
final_message = Some(message);
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
let final_response =
|
|
final_message.ok_or_else(|| anyhow!("final stream ended without a Done event"))?;
|
|
final_content = final_response.content.clone();
|
|
messages.push(final_response);
|
|
// Drop the synthetic prompt — internal scaffolding only. The
|
|
// model's final_response (now at the end) was generated with
|
|
// it in context and reads coherently without it on replay.
|
|
remove_synthetic_final_prompt(messages, synthetic_idx);
|
|
}
|
|
|
|
Ok(AgenticLoopOutcome {
|
|
tool_calls_made,
|
|
iterations_used,
|
|
last_prompt_eval_count,
|
|
last_eval_count,
|
|
// Strip any leaked <think> reasoning block from the content the
|
|
// caller persists as title/summary (the raw transcript keeps it).
|
|
final_content: crate::ai::llm_client::strip_think_blocks(&final_content),
|
|
cancelled: false,
|
|
})
|
|
}
|
|
}
|
|
|
|
/// Default system prompt for bootstrap when the client didn't supply one
|
|
/// (or supplied an empty string). Apollo's frontend always sends a persona
|
|
/// prompt today, so this is a fallback for clients that don't.
|
|
const BOOTSTRAP_DEFAULT_SYSTEM_PROMPT: &str = "You are a helpful AI assistant analyzing the user's photo. \
|
|
Use the available tools to gather context and answer their questions \
|
|
in a conversational tone.";
|
|
|
|
/// Pick the system prompt for bootstrap. Precedence: trimmed-non-empty
|
|
/// `supplied` (the client's explicit `system_prompt`) wins; else
|
|
/// `persona_prompt` (the active persona's stored prompt, resolved
|
|
/// server-side from the persona store); else
|
|
/// [`BOOTSTRAP_DEFAULT_SYSTEM_PROMPT`]. Returns an owned `String` because
|
|
/// the bootstrap caller persists it on the new insight row.
|
|
fn resolve_bootstrap_system_prompt(
|
|
supplied: Option<&str>,
|
|
persona_prompt: Option<String>,
|
|
) -> String {
|
|
supplied
|
|
.map(str::trim)
|
|
.filter(|s| !s.is_empty())
|
|
.map(str::to_string)
|
|
.or_else(|| persona_prompt.filter(|s| !s.trim().is_empty()))
|
|
.unwrap_or_else(|| BOOTSTRAP_DEFAULT_SYSTEM_PROMPT.to_string())
|
|
}
|
|
|
|
/// Compose the bootstrap system message: the persona on top, followed
|
|
/// by a photo-context block carrying the file path, date taken (when
|
|
/// known), GPS (when present), and — in hybrid mode — the local-vision
|
|
/// visual description. Lives in the system message — not the user
|
|
/// turn — so the rendered transcript shows only what the user typed.
|
|
fn build_bootstrap_system_message(
|
|
persona: &str,
|
|
normalized_path: &str,
|
|
date_taken: Option<&str>,
|
|
gps: Option<(f64, f64)>,
|
|
visual_block: &str,
|
|
) -> String {
|
|
let mut out = persona.trim_end().to_string();
|
|
out.push_str("\n\n--- PHOTO CONTEXT ---\n");
|
|
out.push_str(&format!("Photo file path: {}\n", normalized_path));
|
|
out.push_str(&format!(
|
|
"Date taken: {}\n",
|
|
date_taken.unwrap_or("unknown")
|
|
));
|
|
if let Some((lat, lon)) = gps {
|
|
// Four decimal places ≈ 11 m of precision — plenty for any
|
|
// place-lookup tool, and keeps the prompt short.
|
|
out.push_str(&format!("GPS: {:.4}, {:.4}\n", lat, lon));
|
|
}
|
|
if !visual_block.is_empty() {
|
|
// visual_block already ends with a newline; no extra separator
|
|
// needed.
|
|
out.push_str(visual_block);
|
|
}
|
|
out
|
|
}
|
|
|
|
/// Resolve a human-readable `YYYY-MM-DD` date string for the photo
|
|
/// context block. Waterfall: EXIF `date_taken` → filename pattern →
|
|
/// `None`. The fs-time fallback that `generate_agentic_insight_for_photo`
|
|
/// uses is intentionally NOT applied here — for chat we'd rather show
|
|
/// "unknown" than a misleading inode mtime as the photo's date.
|
|
fn resolve_date_taken_for_context(
|
|
exif: &Option<crate::database::models::ImageExif>,
|
|
file_path: &str,
|
|
) -> Option<String> {
|
|
let from_exif = exif
|
|
.as_ref()
|
|
.and_then(|e| e.date_taken)
|
|
.and_then(|ts| chrono::DateTime::from_timestamp(ts, 0))
|
|
.map(|dt| dt.format("%Y-%m-%d").to_string());
|
|
if from_exif.is_some() {
|
|
return from_exif;
|
|
}
|
|
crate::memories::extract_date_from_filename(file_path)
|
|
.map(|dt| dt.format("%Y-%m-%d").to_string())
|
|
}
|
|
|
|
/// Validate a stored→effective backend transition for a chat continuation.
|
|
/// Continuation runs against a transcript that was generated with a specific
|
|
/// backend; the only blocked transition is `local → hybrid`, because the
|
|
/// stored transcript has images embedded in the first user message and the
|
|
/// hybrid path (OpenRouter chat with describe-then-inline) can't replay
|
|
/// raw image bytes through OpenRouter consistently across providers.
|
|
/// `hybrid → local` is allowed (the inlined description replays verbatim
|
|
/// as text).
|
|
///
|
|
/// Whether "local" routes through Ollama or llama-swap is decided at
|
|
/// startup by `LLM_BACKEND`; both share the same transcript shape from
|
|
/// the chat-replay perspective.
|
|
fn validate_cross_replay(stored: &str, effective: &str) -> Result<()> {
|
|
if !matches!(effective, "local" | "hybrid") {
|
|
bail!(
|
|
"unknown backend '{}'; expected 'local' or 'hybrid'",
|
|
effective
|
|
);
|
|
}
|
|
if stored == "local" && effective == "hybrid" {
|
|
bail!(
|
|
"switching from local to hybrid mid-chat isn't supported; \
|
|
regenerate the insight in hybrid mode if you want OpenRouter chat"
|
|
);
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
/// Pick the backend label for bootstrap. Bootstrap has no stored insight
|
|
/// to defer to (that's continuation's behaviour), so the default is
|
|
/// `"local"`. Returns an error if the supplied label is non-empty but
|
|
/// not one of the recognised values — same surface as continuation's
|
|
/// validation.
|
|
fn resolve_bootstrap_backend(supplied: Option<&str>) -> Result<String> {
|
|
let lower = supplied
|
|
.map(|s| s.trim().to_lowercase())
|
|
.filter(|s| !s.is_empty())
|
|
.unwrap_or_else(|| "local".to_string());
|
|
if !matches!(lower.as_str(), "local" | "hybrid") {
|
|
bail!("unknown backend '{}'; expected 'local' or 'hybrid'", lower);
|
|
}
|
|
Ok(lower)
|
|
}
|
|
|
|
/// Outcome of one streaming agentic loop pass. Shared between bootstrap
|
|
/// and continuation. `pub` so the persona chat surface (and any future
|
|
/// chat-without-file surface) can read the result of
|
|
/// `run_streaming_agentic_loop_with_entry` without reaching back into
|
|
/// private fields.
|
|
pub struct AgenticLoopOutcome {
|
|
pub tool_calls_made: usize,
|
|
pub iterations_used: usize,
|
|
pub last_prompt_eval_count: Option<i32>,
|
|
pub last_eval_count: Option<i32>,
|
|
pub final_content: String,
|
|
/// True when the loop exited early because the turn was cancelled
|
|
/// (status flipped out of `Running`). Callers skip persistence and the
|
|
/// terminal `Done` push — the cancel handler owns the terminal event.
|
|
pub cancelled: bool,
|
|
}
|
|
|
|
/// Events emitted by `chat_turn_stream`. One stream per turn; ends after
|
|
/// `Done` or `Error`.
|
|
#[derive(Debug, Clone)]
|
|
pub enum ChatStreamEvent {
|
|
/// Starting iteration `n` of up to `max` (1-based).
|
|
IterationStart { n: usize, max: usize },
|
|
/// History was trimmed to fit the context budget before the turn ran.
|
|
/// Emitted at most once, before any tool or text events.
|
|
Truncated,
|
|
/// Incremental content from the final assistant reply. Concatenate to
|
|
/// reconstruct the reply body. Tool-dispatch turns don't produce these.
|
|
TextDelta(String),
|
|
/// The model requested this tool call. Emitted just before execution.
|
|
/// `index` is a monotonically-increasing counter across the turn so the
|
|
/// client can pair `ToolCall` with its matching `ToolResult`.
|
|
ToolCall {
|
|
index: usize,
|
|
name: String,
|
|
arguments: serde_json::Value,
|
|
},
|
|
/// The tool finished; `result` is the (possibly truncated) output.
|
|
ToolResult {
|
|
index: usize,
|
|
name: String,
|
|
result: String,
|
|
result_truncated: bool,
|
|
},
|
|
/// Terminal success event with counters + persistence result.
|
|
Done {
|
|
tool_calls_made: usize,
|
|
iterations_used: usize,
|
|
truncated: bool,
|
|
/// Renamed from `prompt_eval_count` to match the wire-name Apollo's
|
|
/// frontend (and the mobile client) consume on the `done` SSE
|
|
/// payload.
|
|
prompt_tokens: Option<i32>,
|
|
/// Renamed from `eval_count`. Same rationale.
|
|
eval_tokens: Option<i32>,
|
|
/// The configured context-window ceiling that ran this turn — echoes
|
|
/// the request's `num_ctx` (or the server default when none was
|
|
/// supplied). Lets clients render `prompt_tokens / num_ctx` without
|
|
/// remembering what they asked for.
|
|
num_ctx: Option<i32>,
|
|
/// Populated when this turn produced a NEW insight row — bootstrap
|
|
/// (no prior insight), regenerate (old row flipped to is_current=
|
|
/// false), or amend. Null on append. Clients should target this id
|
|
/// for any subsequent operation that previously needed a known
|
|
/// insight (e.g., a follow-up `/insights/chat` against the just-
|
|
/// created row).
|
|
amended_insight_id: Option<i32>,
|
|
backend_used: String,
|
|
model_used: String,
|
|
/// True only for the synthetic terminal event emitted by the cancel
|
|
/// handler, so clients can distinguish a user-cancelled turn from a
|
|
/// natural completion. Always false on the normal success path.
|
|
cancelled: bool,
|
|
},
|
|
/// Terminal failure event. No further events follow.
|
|
Error(String),
|
|
}
|
|
|
|
/// Is this raw message visible in the rendered transcript? Must match
|
|
/// `load_history`'s filter exactly — `find_raw_cut` depends on it to map
|
|
/// rendered indices back to raw positions.
|
|
#[allow(dead_code)]
|
|
fn is_rendered(m: &ChatMessage) -> bool {
|
|
match m.role.as_str() {
|
|
"user" => true,
|
|
"assistant" => {
|
|
let has_tool_calls = m
|
|
.tool_calls
|
|
.as_ref()
|
|
.map(|c| !c.is_empty())
|
|
.unwrap_or(false);
|
|
!(has_tool_calls && m.content.trim().is_empty())
|
|
}
|
|
_ => false,
|
|
}
|
|
}
|
|
|
|
/// Given a rendered index to start discarding from, find the raw index at
|
|
/// which to truncate. The cut position is the raw length after all prior
|
|
/// rendered messages — which also strips any tool-call scaffolding that
|
|
/// immediately precedes the discarded rendered message.
|
|
///
|
|
/// Discarding *at* the end (`discard == rendered_count`) is a no-op success:
|
|
/// returns `Some(messages.len())`. The mobile client hits this when
|
|
/// regenerating after a failed turn — its optimistic user bubble lives at
|
|
/// the index just past the server's persisted history. Strictly past the end
|
|
/// (`discard > rendered_count`) returns `None`.
|
|
#[allow(dead_code)]
|
|
pub(crate) fn find_raw_cut(
|
|
messages: &[ChatMessage],
|
|
discard_from_rendered_index: usize,
|
|
) -> Option<usize> {
|
|
let mut rendered_count = 0usize;
|
|
let mut last_kept_raw_end = 0usize;
|
|
for (i, m) in messages.iter().enumerate() {
|
|
if !is_rendered(m) {
|
|
continue;
|
|
}
|
|
if rendered_count == discard_from_rendered_index {
|
|
return Some(last_kept_raw_end);
|
|
}
|
|
rendered_count += 1;
|
|
last_kept_raw_end = i + 1;
|
|
}
|
|
if discard_from_rendered_index == rendered_count {
|
|
return Some(messages.len());
|
|
}
|
|
None
|
|
}
|
|
|
|
/// Read AGENTIC_CHAT_MAX_ITERATIONS once per call. Cheap; keeps the code
|
|
/// free of static globals and lets the operator change the cap by env without
|
|
/// a restart in test harnesses (the running server still caches via Default).
|
|
pub fn env_max_iterations() -> usize {
|
|
std::env::var("AGENTIC_CHAT_MAX_ITERATIONS")
|
|
.ok()
|
|
.and_then(|s| s.parse::<usize>().ok())
|
|
.unwrap_or(DEFAULT_MAX_ITERATIONS)
|
|
.max(1)
|
|
}
|
|
|
|
/// Read AGENTIC_CHAT_DEFAULT_NUM_CTX once per call — the assumed context
|
|
/// window for the truncation budget when the request omits `num_ctx`. Same
|
|
/// no-static-global rationale as `env_max_iterations` above.
|
|
fn env_default_num_ctx() -> i32 {
|
|
std::env::var("AGENTIC_CHAT_DEFAULT_NUM_CTX")
|
|
.ok()
|
|
.and_then(|s| s.parse::<i32>().ok())
|
|
.unwrap_or(DEFAULT_NUM_CTX)
|
|
.max(RESPONSE_HEADROOM_TOKENS as i32 + 1024)
|
|
}
|
|
|
|
/// Append a per-turn iteration-budget reminder to the replayed system
|
|
/// message so the model knows how many tool-calling rounds this turn gets.
|
|
/// Returns the original `content` so the caller can restore it before
|
|
/// persistence — otherwise the note would accumulate across turns.
|
|
///
|
|
/// No-op (returns `None`) when `messages` has no leading system message.
|
|
fn annotate_system_with_budget(
|
|
messages: &mut [ChatMessage],
|
|
max_iterations: usize,
|
|
) -> Option<String> {
|
|
let first = messages.first_mut()?;
|
|
if first.role != "system" {
|
|
return None;
|
|
}
|
|
let original = first.content.clone();
|
|
first.content = format!(
|
|
"{}\n\n(Budget for this chat turn: up to {} tool-calling iterations. Produce your final reply before the budget is exhausted.)",
|
|
first.content, max_iterations
|
|
);
|
|
Some(original)
|
|
}
|
|
|
|
/// Restore a system-message content previously captured by
|
|
/// [`annotate_system_with_budget`]. No-op when `original` is `None` or the
|
|
/// first message isn't a system message.
|
|
fn restore_system_content(messages: &mut [ChatMessage], original: Option<String>) {
|
|
let Some(original) = original else { return };
|
|
if let Some(first) = messages.first_mut()
|
|
&& first.role == "system"
|
|
{
|
|
first.content = original;
|
|
}
|
|
}
|
|
|
|
/// Append the synthetic "write your final answer" user prompt, returning the
|
|
/// Free-function form of `InsightChatService::run_streaming_agentic_loop_with_entry`.
|
|
/// Persona chat (and any future chat-without-file surface) doesn't need an
|
|
/// `InsightChatService` — it just needs the agent loop. This is the same
|
|
/// body lifted out of the `InsightChatService` impl, taking
|
|
/// `&InsightGenerator` so any caller in the crate can drive it.
|
|
///
|
|
/// Public so the persona chat session can call into it directly.
|
|
pub async fn run_streaming_agentic_loop_with_entry(
|
|
generator: &crate::ai::insight_generator::InsightGenerator,
|
|
backend: &ResolvedBackend,
|
|
messages: &mut Vec<ChatMessage>,
|
|
tools: Vec<Tool>,
|
|
image_base64: &Option<String>,
|
|
normalized: &str,
|
|
user_id: i32,
|
|
active_persona: &str,
|
|
max_iterations: usize,
|
|
entry: &Arc<TurnEntry>,
|
|
) -> Result<AgenticLoopOutcome> {
|
|
let mut tool_calls_made = 0usize;
|
|
let mut iterations_used = 0usize;
|
|
let mut last_prompt_eval_count: Option<i32> = None;
|
|
let mut last_eval_count: Option<i32> = None;
|
|
let mut final_content = String::new();
|
|
|
|
for iteration in 0..max_iterations {
|
|
if !entry.is_running() {
|
|
return Ok(AgenticLoopOutcome {
|
|
tool_calls_made,
|
|
iterations_used,
|
|
last_prompt_eval_count,
|
|
last_eval_count,
|
|
final_content,
|
|
cancelled: true,
|
|
});
|
|
}
|
|
|
|
iterations_used = iteration + 1;
|
|
let _ = entry
|
|
.push_event(ChatStreamEvent::IterationStart {
|
|
n: iterations_used,
|
|
max: max_iterations,
|
|
})
|
|
.await;
|
|
|
|
let mut stream = backend
|
|
.chat()
|
|
.chat_with_tools_stream(messages.clone(), tools.clone())
|
|
.await?;
|
|
|
|
let mut final_message: Option<ChatMessage> = None;
|
|
while let Some(ev) = stream.next().await {
|
|
let ev = ev?;
|
|
match ev {
|
|
LlmStreamEvent::TextDelta(delta) => {
|
|
let _ = entry.push_event(ChatStreamEvent::TextDelta(delta)).await;
|
|
}
|
|
LlmStreamEvent::Done {
|
|
message,
|
|
prompt_eval_count,
|
|
eval_count,
|
|
} => {
|
|
last_prompt_eval_count = prompt_eval_count;
|
|
last_eval_count = eval_count;
|
|
final_message = Some(message);
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
let mut response =
|
|
final_message.ok_or_else(|| anyhow!("stream ended without a Done event"))?;
|
|
|
|
if let Some(ref mut tcs) = response.tool_calls {
|
|
for tc in tcs.iter_mut() {
|
|
if !tc.function.arguments.is_object() {
|
|
tc.function.arguments = serde_json::Value::Object(Default::default());
|
|
}
|
|
}
|
|
}
|
|
|
|
messages.push(response.clone());
|
|
|
|
if let Some(ref tool_calls) = response.tool_calls
|
|
&& !tool_calls.is_empty()
|
|
{
|
|
for tool_call in tool_calls {
|
|
tool_calls_made += 1;
|
|
let call_index = tool_calls_made - 1;
|
|
let _ = entry
|
|
.push_event(ChatStreamEvent::ToolCall {
|
|
index: call_index,
|
|
name: tool_call.function.name.clone(),
|
|
arguments: tool_call.function.arguments.clone(),
|
|
})
|
|
.await;
|
|
let cx = opentelemetry::Context::new();
|
|
let result = generator
|
|
.execute_tool(
|
|
&tool_call.function.name,
|
|
&tool_call.function.arguments,
|
|
backend,
|
|
image_base64,
|
|
normalized,
|
|
user_id,
|
|
active_persona,
|
|
&cx,
|
|
)
|
|
.await;
|
|
let (result_preview, result_truncated) = truncate_tool_result(&result);
|
|
let _ = entry
|
|
.push_event(ChatStreamEvent::ToolResult {
|
|
index: call_index,
|
|
name: tool_call.function.name.clone(),
|
|
result: result_preview,
|
|
result_truncated,
|
|
})
|
|
.await;
|
|
messages.push(ChatMessage::tool_result(result));
|
|
}
|
|
continue;
|
|
}
|
|
|
|
final_content = response.content;
|
|
break;
|
|
}
|
|
|
|
if final_content.is_empty() {
|
|
let synthetic_idx = push_synthetic_final_prompt(messages);
|
|
let mut stream = backend
|
|
.chat()
|
|
.chat_with_tools_stream(messages.clone(), vec![])
|
|
.await?;
|
|
let mut final_message: Option<ChatMessage> = None;
|
|
while let Some(ev) = stream.next().await {
|
|
let ev = ev?;
|
|
match ev {
|
|
LlmStreamEvent::TextDelta(delta) => {
|
|
let _ = entry.push_event(ChatStreamEvent::TextDelta(delta)).await;
|
|
}
|
|
LlmStreamEvent::Done {
|
|
message,
|
|
prompt_eval_count,
|
|
eval_count,
|
|
} => {
|
|
last_prompt_eval_count = prompt_eval_count;
|
|
last_eval_count = eval_count;
|
|
final_message = Some(message);
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
let final_response =
|
|
final_message.ok_or_else(|| anyhow!("final stream ended without a Done event"))?;
|
|
final_content = final_response.content.clone();
|
|
messages.push(final_response);
|
|
remove_synthetic_final_prompt(messages, synthetic_idx);
|
|
}
|
|
|
|
Ok(AgenticLoopOutcome {
|
|
tool_calls_made,
|
|
iterations_used,
|
|
last_prompt_eval_count,
|
|
last_eval_count,
|
|
final_content: crate::ai::llm_client::strip_think_blocks(&final_content),
|
|
cancelled: false,
|
|
})
|
|
}
|
|
|
|
/// index the caller must later hand to [`remove_synthetic_final_prompt`].
|
|
/// Used when the agentic loop exhausts its budget: the model gets one more
|
|
/// (tool-free) request, but the nudge itself must never persist — it would
|
|
/// render as a user bubble in the transcript and reset `load_history`'s
|
|
/// pending-tools tracking at that position.
|
|
fn push_synthetic_final_prompt(messages: &mut Vec<ChatMessage>) -> usize {
|
|
let idx = messages.len();
|
|
messages.push(ChatMessage::user(SYNTHETIC_FINAL_ANSWER_PROMPT));
|
|
idx
|
|
}
|
|
|
|
/// Remove the synthetic prompt inserted by [`push_synthetic_final_prompt`].
|
|
/// Defensive no-op when the message at `idx` isn't the synthetic prompt —
|
|
/// guards against index drift if the surrounding code is reordered.
|
|
fn remove_synthetic_final_prompt(messages: &mut Vec<ChatMessage>, idx: usize) {
|
|
if messages
|
|
.get(idx)
|
|
.is_some_and(|m| m.role == "user" && m.content == SYNTHETIC_FINAL_ANSWER_PROMPT)
|
|
{
|
|
messages.remove(idx);
|
|
}
|
|
}
|
|
|
|
/// Receipt produced by [`apply_system_prompt_override`] so the caller can
|
|
/// undo the override before persistence. Two variants because we either
|
|
/// replaced an existing system message (need its original content) or
|
|
/// prepended a synthetic one (need to pop it).
|
|
#[derive(Debug)]
|
|
pub(crate) enum SystemPromptStash {
|
|
Replaced { original: String },
|
|
Prepended,
|
|
}
|
|
|
|
/// Apply a per-turn `system_prompt` override to `messages` so the model
|
|
/// sees the requested persona for this turn. Returns a stash the caller
|
|
/// must pass to [`restore_system_prompt_override`] before persisting the
|
|
/// transcript — without that step, append-mode chat would silently
|
|
/// rewrite the stored persona.
|
|
///
|
|
/// No-op (returns `None`) when `override_prompt` is `None` or empty.
|
|
pub(crate) fn apply_system_prompt_override(
|
|
messages: &mut Vec<ChatMessage>,
|
|
override_prompt: Option<&str>,
|
|
) -> Option<SystemPromptStash> {
|
|
let prompt = override_prompt
|
|
.map(str::trim)
|
|
.filter(|s| !s.is_empty())?
|
|
.to_string();
|
|
if let Some(first) = messages.first_mut()
|
|
&& first.role == "system"
|
|
{
|
|
let original = std::mem::replace(&mut first.content, prompt);
|
|
return Some(SystemPromptStash::Replaced { original });
|
|
}
|
|
messages.insert(0, ChatMessage::system(prompt));
|
|
Some(SystemPromptStash::Prepended)
|
|
}
|
|
|
|
/// Undo an override previously applied by [`apply_system_prompt_override`].
|
|
/// No-op when `stash` is `None`.
|
|
pub(crate) fn restore_system_prompt_override(
|
|
messages: &mut Vec<ChatMessage>,
|
|
stash: Option<SystemPromptStash>,
|
|
) {
|
|
let Some(stash) = stash else { return };
|
|
match stash {
|
|
SystemPromptStash::Replaced { original } => {
|
|
if let Some(first) = messages.first_mut()
|
|
&& first.role == "system"
|
|
{
|
|
first.content = original;
|
|
}
|
|
}
|
|
SystemPromptStash::Prepended => {
|
|
if matches!(messages.first(), Some(m) if m.role == "system") {
|
|
messages.remove(0);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Fork indicator for a rendered message: position and total branches at
|
|
/// the nearest divergence point upstream in the conversation tree.
|
|
#[derive(Debug, Clone, serde::Serialize)]
|
|
pub struct ForkInfo {
|
|
pub position: usize,
|
|
pub total: usize,
|
|
/// The divergence node itself (the parent whose children fork). Pass as
|
|
/// `node_id` to the branches endpoint to list the sibling branches at
|
|
/// this exact point rather than every leaf in the tree.
|
|
pub node_id: u64,
|
|
}
|
|
|
|
/// Render a path of tree nodes into the UI-friendly `RenderedMessage` format,
|
|
/// matching the same filtering logic as the legacy flat-array `load_history`.
|
|
/// Returns `(rendered_messages, turn_count, node_ids, fork_info)` where
|
|
/// `node_ids[i]` is the tree node ID that produced `rendered_messages[i]` and
|
|
/// `fork_info[i]` carries the nearest divergence point at or upstream of it.
|
|
///
|
|
/// Fork detection walks *every* node in the path — including the non-rendered
|
|
/// tool-dispatch/tool/system nodes — so a fork whose diverging child is a
|
|
/// tool-dispatch assistant (empty content + tool_calls) is still attributed to
|
|
/// the next rendered message. Detecting forks only on rendered nodes would miss
|
|
/// regenerations where the model replied with a tool call.
|
|
fn render_tree_path(
|
|
store: &ChatHistoryStore,
|
|
path: &[&StoredChatNode],
|
|
) -> (Vec<RenderedMessage>, usize, Vec<u64>, Vec<Option<ForkInfo>>) {
|
|
let mut rendered = Vec::new();
|
|
let mut user_turns_seen = 0usize;
|
|
let mut assistant_turns_seen = 0usize;
|
|
let mut pending_tools: Vec<ToolInvocation> = Vec::new();
|
|
let mut pending_calls: std::collections::VecDeque<(String, serde_json::Value)> =
|
|
std::collections::VecDeque::new();
|
|
let mut node_ids: Vec<u64> = Vec::new();
|
|
let mut fork_info: Vec<Option<ForkInfo>> = Vec::new();
|
|
// Nearest divergence seen so far, carried forward across non-rendered
|
|
// nodes (e.g. tool-dispatch) so the fork is attributed to the next
|
|
// rendered message. Cleared after the first rendered message consumes it,
|
|
// so the chip appears only at the actual divergence point.
|
|
let mut last_fork: Option<ForkInfo> = None;
|
|
|
|
for node in path {
|
|
// Update the carried fork BEFORE rendering: a non-rendered fork child
|
|
// (e.g. a tool-dispatch assistant) must colour the rendered message
|
|
// that follows it.
|
|
// (fork_at_node only fires for nodes with a parent, so the pair
|
|
// pattern below can't drop a real fork.)
|
|
if let (Some(node_id), Some((position, total))) =
|
|
(node.parent_id, store.fork_at_node(node.id))
|
|
{
|
|
last_fork = Some(ForkInfo {
|
|
position,
|
|
total,
|
|
node_id,
|
|
});
|
|
}
|
|
let msg = &node.message;
|
|
match msg.role.as_str() {
|
|
"system" => continue,
|
|
"tool" => {
|
|
if let Some((name, arguments)) = pending_calls.pop_front() {
|
|
let (result, result_truncated) = truncate_tool_result(&msg.content);
|
|
pending_tools.push(ToolInvocation {
|
|
name,
|
|
arguments,
|
|
result,
|
|
result_truncated,
|
|
});
|
|
}
|
|
}
|
|
"assistant" => {
|
|
let has_tool_calls = msg
|
|
.tool_calls
|
|
.as_ref()
|
|
.map(|c| !c.is_empty())
|
|
.unwrap_or(false);
|
|
if has_tool_calls && msg.content.trim().is_empty() {
|
|
if let Some(ref tcs) = msg.tool_calls {
|
|
for tc in tcs {
|
|
pending_calls.push_back((
|
|
tc.function.name.clone(),
|
|
tc.function.arguments.clone(),
|
|
));
|
|
}
|
|
}
|
|
continue;
|
|
}
|
|
assistant_turns_seen += 1;
|
|
let tools = std::mem::take(&mut pending_tools);
|
|
pending_calls.clear();
|
|
rendered.push(RenderedMessage {
|
|
role: "assistant".to_string(),
|
|
content: msg.content.clone(),
|
|
is_initial: false,
|
|
tools,
|
|
});
|
|
node_ids.push(node.id);
|
|
fork_info.push(last_fork.take());
|
|
}
|
|
"user" => {
|
|
let is_initial = user_turns_seen == 0;
|
|
user_turns_seen += 1;
|
|
pending_tools.clear();
|
|
pending_calls.clear();
|
|
rendered.push(RenderedMessage {
|
|
role: "user".to_string(),
|
|
content: msg.content.clone(),
|
|
is_initial,
|
|
tools: Vec::new(),
|
|
});
|
|
node_ids.push(node.id);
|
|
fork_info.push(last_fork.take());
|
|
}
|
|
_ => continue,
|
|
}
|
|
}
|
|
|
|
(rendered, assistant_turns_seen, node_ids, fork_info)
|
|
}
|
|
|
|
/// View returned to clients for chat-UI rendering.
|
|
#[derive(Debug)]
|
|
pub struct HistoryView {
|
|
pub messages: Vec<RenderedMessage>,
|
|
pub turn_count: usize,
|
|
pub model_version: String,
|
|
pub backend: String,
|
|
/// ID of the leaf node the current view is anchored to.
|
|
pub active_leaf_id: u64,
|
|
/// ID of the branch leaf being viewed. Equals `active_leaf_id` when
|
|
/// viewing the active branch, or the `branch_id` parameter when
|
|
/// viewing an alternate branch.
|
|
pub viewing_branch_id: u64,
|
|
/// Fork info for each rendered message. `None` means no divergence at
|
|
/// or before that point in the tree.
|
|
pub fork_info: Vec<Option<ForkInfo>>,
|
|
}
|
|
|
|
#[derive(Debug)]
|
|
pub struct RenderedMessage {
|
|
pub role: String,
|
|
pub content: String,
|
|
pub is_initial: bool,
|
|
/// Tools invoked during this turn (only populated for assistant replies).
|
|
/// Empty for user messages and for assistant replies that didn't involve
|
|
/// tool calls.
|
|
pub tools: Vec<ToolInvocation>,
|
|
}
|
|
|
|
#[derive(Debug, Clone)]
|
|
pub struct ToolInvocation {
|
|
pub name: String,
|
|
pub arguments: serde_json::Value,
|
|
pub result: String,
|
|
/// True when `result` was trimmed for payload size. Full value remains
|
|
/// available in the raw training_messages blob.
|
|
pub result_truncated: bool,
|
|
}
|
|
|
|
/// Soft cap for tool-result bodies returned via the history API. Keeps
|
|
/// payloads small for the mobile client — verbose SMS / geocoding responses
|
|
/// don't need to ship in full for inspection.
|
|
pub(crate) const TOOL_RESULT_PREVIEW_MAX: usize = 2000;
|
|
|
|
pub(crate) fn truncate_tool_result(s: &str) -> (String, bool) {
|
|
if s.len() <= TOOL_RESULT_PREVIEW_MAX {
|
|
(s.to_string(), false)
|
|
} else {
|
|
// Cut on a char boundary.
|
|
let mut cut = TOOL_RESULT_PREVIEW_MAX;
|
|
while !s.is_char_boundary(cut) && cut > 0 {
|
|
cut -= 1;
|
|
}
|
|
(s[..cut].to_string(), true)
|
|
}
|
|
}
|
|
|
|
/// Trim history to fit within `budget_bytes` of serialized JSON. Preserves
|
|
/// the system message and the first user message (with its base64 images
|
|
/// intact, since dropping those would invalidate the model's prior visual
|
|
/// reasoning). Drops the oldest assistant-tool_call + corresponding
|
|
/// tool-result pair on each pass until the budget is met or only the
|
|
/// preserved prefix remains.
|
|
///
|
|
/// Returns true when at least one message was dropped.
|
|
pub(crate) fn apply_context_budget(messages: &mut Vec<ChatMessage>, budget_bytes: usize) -> bool {
|
|
if budget_bytes == 0 {
|
|
return false;
|
|
}
|
|
if estimate_bytes(messages) <= budget_bytes {
|
|
return false;
|
|
}
|
|
|
|
// Find the index past the protected prefix: system messages + the first
|
|
// user message. Everything after is droppable in pairs.
|
|
let first_user_idx = messages.iter().position(|m| m.role == "user");
|
|
let preserve_through = match first_user_idx {
|
|
Some(i) => i, // keep [0..=i]
|
|
None => return false,
|
|
};
|
|
|
|
let mut dropped_any = false;
|
|
loop {
|
|
if estimate_bytes(messages) <= budget_bytes {
|
|
break;
|
|
}
|
|
// Find the oldest assistant-with-tool_calls strictly after the
|
|
// preserved prefix. Drop it together with the following tool turn(s)
|
|
// until we hit the next assistant or user turn.
|
|
let drop_start = (preserve_through + 1..messages.len()).find(|&i| {
|
|
let m = &messages[i];
|
|
m.role == "assistant"
|
|
&& m.tool_calls
|
|
.as_ref()
|
|
.map(|c| !c.is_empty())
|
|
.unwrap_or(false)
|
|
});
|
|
let Some(start) = drop_start else { break };
|
|
// Determine end: drop the assistant turn plus any contiguous tool
|
|
// result turns that follow.
|
|
let mut end = start + 1;
|
|
while end < messages.len() && messages[end].role == "tool" {
|
|
end += 1;
|
|
}
|
|
// Stop if dropping these would leave the just-appended user turn at
|
|
// the end alone with no preceding context — we still want it kept.
|
|
if end > messages.len() {
|
|
break;
|
|
}
|
|
messages.drain(start..end);
|
|
dropped_any = true;
|
|
}
|
|
|
|
dropped_any
|
|
}
|
|
|
|
/// Estimate the serialized byte size of `messages` for the truncation budget,
|
|
/// EXCLUDING inlined base64 image payloads. Images are charged a flat
|
|
/// `IMAGE_TOKENS_EACH` instead: their base64 is hundreds of KB of characters
|
|
/// that have no relation to the text token pressure we're budgeting against,
|
|
/// and counting them verbatim makes a single photo exceed the entire budget,
|
|
/// spuriously trimming all history on every turn.
|
|
fn estimate_bytes(messages: &[ChatMessage]) -> usize {
|
|
let mut image_count = 0usize;
|
|
// Clone with image payloads stripped so they don't inflate the byte count.
|
|
// We still account for the (small) non-image fields verbatim.
|
|
let stripped: Vec<ChatMessage> = messages
|
|
.iter()
|
|
.map(|m| {
|
|
if let Some(imgs) = m.images.as_ref() {
|
|
image_count += imgs.len();
|
|
}
|
|
ChatMessage {
|
|
images: None,
|
|
..m.clone()
|
|
}
|
|
})
|
|
.collect();
|
|
let text_bytes = serde_json::to_string(&stripped)
|
|
.map(|s| s.len())
|
|
.unwrap_or(0);
|
|
text_bytes + image_count * IMAGE_TOKENS_EACH * BYTES_PER_TOKEN
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use crate::ai::llm_client::{ToolCall, ToolCallFunction};
|
|
|
|
fn assistant_with_tool_call(name: &str) -> ChatMessage {
|
|
ChatMessage {
|
|
role: "assistant".to_string(),
|
|
content: String::new(),
|
|
tool_calls: Some(vec![ToolCall {
|
|
id: None,
|
|
function: ToolCallFunction {
|
|
name: name.to_string(),
|
|
arguments: serde_json::Value::Object(Default::default()),
|
|
},
|
|
}]),
|
|
images: None,
|
|
}
|
|
}
|
|
|
|
fn assistant_text(text: &str) -> ChatMessage {
|
|
ChatMessage {
|
|
role: "assistant".to_string(),
|
|
content: text.to_string(),
|
|
tool_calls: None,
|
|
images: None,
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn truncation_preserves_system_and_first_user() {
|
|
let mut msgs = vec![
|
|
ChatMessage::system("sys"),
|
|
ChatMessage::user("first user with lots of context".repeat(50)),
|
|
assistant_with_tool_call("get_x"),
|
|
ChatMessage::tool_result("x result ".repeat(200)),
|
|
assistant_with_tool_call("get_y"),
|
|
ChatMessage::tool_result("y result ".repeat(200)),
|
|
assistant_text("final answer"),
|
|
];
|
|
let original_len = msgs.len();
|
|
let dropped = apply_context_budget(&mut msgs, 500);
|
|
assert!(dropped, "should drop something at this small budget");
|
|
assert!(msgs.len() < original_len);
|
|
// First two messages preserved.
|
|
assert_eq!(msgs[0].role, "system");
|
|
assert_eq!(msgs[1].role, "user");
|
|
}
|
|
|
|
#[test]
|
|
fn truncation_no_op_when_under_budget() {
|
|
let mut msgs = vec![ChatMessage::system("s"), ChatMessage::user("u")];
|
|
let dropped = apply_context_budget(&mut msgs, 1_000_000);
|
|
assert!(!dropped);
|
|
assert_eq!(msgs.len(), 2);
|
|
}
|
|
|
|
#[test]
|
|
fn image_payload_excluded_from_budget() {
|
|
// First user message carries a ~400KB base64 image but only a little
|
|
// text. Counting the base64 verbatim (old behavior) dwarfs the budget
|
|
// and forces all tool history to be dropped on every turn. The image
|
|
// must instead be charged a flat per-image cost so a short
|
|
// conversation comfortably fits.
|
|
let mut user = ChatMessage::user("describe this");
|
|
user.images = Some(vec!["A".repeat(400_000)]);
|
|
let mut msgs = vec![
|
|
ChatMessage::system("sys"),
|
|
user,
|
|
assistant_with_tool_call("get_x"),
|
|
ChatMessage::tool_result("small x result"),
|
|
assistant_text("here is the answer"),
|
|
];
|
|
|
|
// Default budget: (32768 - 2048) * 4 bytes ≈ 120KB. The text easily
|
|
// fits; only the (excluded) image bytes could blow it.
|
|
let budget_bytes = (DEFAULT_NUM_CTX as usize - RESPONSE_HEADROOM_TOKENS) * BYTES_PER_TOKEN;
|
|
let original_len = msgs.len();
|
|
let dropped = apply_context_budget(&mut msgs, budget_bytes);
|
|
|
|
assert!(
|
|
!dropped,
|
|
"short conversation with one image must not truncate"
|
|
);
|
|
assert_eq!(msgs.len(), original_len, "no messages should be dropped");
|
|
// Sanity: the flat image charge is accounted for but stays well under budget.
|
|
assert!(estimate_bytes(&msgs) <= budget_bytes);
|
|
}
|
|
|
|
#[test]
|
|
fn truncation_returns_false_with_no_droppable_pairs() {
|
|
// Only system + user, no tool-call turns to drop.
|
|
let mut msgs = vec![ChatMessage::system("s"), ChatMessage::user("u")];
|
|
let dropped = apply_context_budget(&mut msgs, 1);
|
|
assert!(!dropped);
|
|
}
|
|
|
|
#[test]
|
|
fn render_tree_path_detects_fork_at_tool_dispatch_child() {
|
|
// Regression: a regeneration where the model replies with a tool call
|
|
// forks at the user turn, but the diverging children are non-rendered
|
|
// tool-dispatch assistant nodes. The indicator must still surface on
|
|
// the rendered final reply that follows — earlier code only inspected
|
|
// rendered nodes and silently dropped the fork, orphaning the branch.
|
|
let mut store = ChatHistoryStore::empty();
|
|
let u = store.append_node(None, ChatMessage::user("q1"));
|
|
// Original branch: tool-dispatch → tool result → final reply.
|
|
let d1 = store.append_node(Some(u), assistant_with_tool_call("lookup"));
|
|
let t1 = store.append_node(Some(d1), ChatMessage::tool_result("data1"));
|
|
let _a1 = store.append_node(Some(t1), assistant_text("answer 1"));
|
|
// Regenerated branch: also a tool reply, forking at the user turn.
|
|
let d2 = store.append_node(Some(u), assistant_with_tool_call("lookup"));
|
|
let t2 = store.append_node(Some(d2), ChatMessage::tool_result("data2"));
|
|
let a2 = store.append_node(Some(t2), assistant_text("answer 2"));
|
|
store.active_leaf_id = a2;
|
|
|
|
let path = store.path_to_leaf(a2).expect("path to active leaf");
|
|
let (rendered, _turns, node_ids, fork_info) = render_tree_path(&store, &path);
|
|
|
|
assert_eq!(rendered.len(), 2, "user turn + final reply are rendered");
|
|
assert_eq!(rendered[1].content, "answer 2");
|
|
assert_eq!(node_ids, vec![u, a2]);
|
|
assert!(
|
|
fork_info[0].is_none(),
|
|
"user turn sits before the divergence"
|
|
);
|
|
let f = fork_info[1]
|
|
.as_ref()
|
|
.expect("fork indicator carried onto the rendered reply");
|
|
assert_eq!((f.position, f.total), (2, 2));
|
|
assert_eq!(f.node_id, u, "divergence node is the forked user turn");
|
|
}
|
|
|
|
#[test]
|
|
fn render_tree_path_detects_fork_at_rendered_child() {
|
|
// Non-tool case: two direct assistant replies to the same user turn.
|
|
// The diverging child is itself rendered, so the indicator lands on it.
|
|
let mut store = ChatHistoryStore::empty();
|
|
let u = store.append_node(None, ChatMessage::user("q1"));
|
|
let _a1 = store.append_node(Some(u), assistant_text("answer 1"));
|
|
let a2 = store.append_node(Some(u), assistant_text("answer 2"));
|
|
store.active_leaf_id = a2;
|
|
|
|
let path = store.path_to_leaf(a2).expect("path to active leaf");
|
|
let (rendered, _turns, _node_ids, fork_info) = render_tree_path(&store, &path);
|
|
|
|
assert_eq!(rendered.len(), 2);
|
|
let f = fork_info[1]
|
|
.as_ref()
|
|
.expect("fork indicator on the divergent rendered reply");
|
|
assert_eq!((f.position, f.total), (2, 2));
|
|
assert_eq!(f.node_id, u, "divergence node is the forked user turn");
|
|
}
|
|
|
|
#[test]
|
|
fn render_tree_path_nested_forks_only_emits_at_divergence() {
|
|
// Regression: with .take() instead of .clone(), fork_info should only
|
|
// appear at the actual divergence point, not propagate downstream.
|
|
// Structure:
|
|
// u (user "q1")
|
|
// ├── a1 (assistant "answer 1")
|
|
// └── a2 (assistant "answer 2")
|
|
// ├── q2 (user "q2")
|
|
// └── r2 (assistant "reply 2") <- active leaf
|
|
//
|
|
// Expected fork_info: [None, Some, None, None]
|
|
// (only a2 has fork info since it's the divergence from a1)
|
|
let mut store = ChatHistoryStore::empty();
|
|
let u = store.append_node(None, ChatMessage::user("q1"));
|
|
let _a1 = store.append_node(Some(u), assistant_text("answer 1"));
|
|
let a2 = store.append_node(Some(u), assistant_text("answer 2"));
|
|
let q2 = store.append_node(Some(a2), ChatMessage::user("q2"));
|
|
let r2 = store.append_node(Some(q2), assistant_text("reply 2"));
|
|
store.active_leaf_id = r2;
|
|
|
|
let path = store.path_to_leaf(r2).expect("path to active leaf");
|
|
let (rendered, _turns, _node_ids, fork_info) = render_tree_path(&store, &path);
|
|
|
|
assert_eq!(rendered.len(), 4, "user + asst + user + asst");
|
|
assert!(fork_info[0].is_none(), "initial user has no fork");
|
|
assert!(fork_info[1].is_some(), "a2 is the fork point");
|
|
assert!(fork_info[2].is_none(), "q2 should NOT carry fork forward");
|
|
assert!(fork_info[3].is_none(), "r2 should NOT carry fork forward");
|
|
}
|
|
|
|
#[test]
|
|
fn rewind_strips_assistant_and_tool_scaffolding() {
|
|
// Rendered: [user1, asst1, user2, asst2] → cut at rendered index 3
|
|
// (the final asst2) should drop the tool-call scaffolding + asst2,
|
|
// leaving raw up through user2.
|
|
let msgs = vec![
|
|
ChatMessage::system("sys"),
|
|
ChatMessage::user("q1"),
|
|
assistant_text("a1"),
|
|
ChatMessage::user("q2"),
|
|
assistant_with_tool_call("lookup"),
|
|
ChatMessage::tool_result("data"),
|
|
assistant_text("a2 final"),
|
|
];
|
|
let cut = find_raw_cut(&msgs, 3).expect("cut found");
|
|
// raw[0..cut] should end at user("q2") — indices 0..=3.
|
|
assert_eq!(cut, 4);
|
|
assert_eq!(msgs[cut - 1].role, "user");
|
|
assert_eq!(msgs[cut - 1].content, "q2");
|
|
}
|
|
|
|
#[test]
|
|
fn rewind_at_second_rendered_cuts_after_first_user() {
|
|
// Rendered index 1 = the first assistant reply → dropping it should
|
|
// leave just the initial user message.
|
|
let msgs = vec![
|
|
ChatMessage::system("s"),
|
|
ChatMessage::user("q1"),
|
|
assistant_with_tool_call("tool"),
|
|
ChatMessage::tool_result("r"),
|
|
assistant_text("a1"),
|
|
];
|
|
let cut = find_raw_cut(&msgs, 1).expect("cut found");
|
|
assert_eq!(cut, 2); // sys + user("q1")
|
|
}
|
|
|
|
#[test]
|
|
fn rewind_beyond_range_returns_none() {
|
|
let msgs = vec![ChatMessage::user("q1"), assistant_text("a1")];
|
|
assert!(find_raw_cut(&msgs, 5).is_none());
|
|
}
|
|
|
|
#[test]
|
|
fn rewind_at_end_is_noop_success() {
|
|
// Mobile client retries after a failed turn that never persisted —
|
|
// its optimistic user bubble's index equals the server's rendered
|
|
// count. Should resolve to "no cut" rather than an out-of-range error.
|
|
let msgs = vec![
|
|
ChatMessage::system("s"),
|
|
ChatMessage::user("q1"),
|
|
assistant_text("a1"),
|
|
];
|
|
let cut = find_raw_cut(&msgs, 2).expect("boundary cut should succeed");
|
|
assert_eq!(cut, msgs.len());
|
|
}
|
|
|
|
#[test]
|
|
fn apply_override_replaces_existing_system_message() {
|
|
let mut msgs = vec![
|
|
ChatMessage::system("original persona"),
|
|
ChatMessage::user("hi"),
|
|
];
|
|
let stash = apply_system_prompt_override(&mut msgs, Some("new persona"));
|
|
assert_eq!(msgs[0].content, "new persona");
|
|
match stash {
|
|
Some(SystemPromptStash::Replaced { original }) => {
|
|
assert_eq!(original, "original persona");
|
|
}
|
|
other => panic!("expected Replaced, got {:?}", other),
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn apply_override_prepends_synthetic_when_missing() {
|
|
let mut msgs = vec![ChatMessage::user("hi")];
|
|
let stash = apply_system_prompt_override(&mut msgs, Some("new persona"));
|
|
assert_eq!(msgs.len(), 2);
|
|
assert_eq!(msgs[0].role, "system");
|
|
assert_eq!(msgs[0].content, "new persona");
|
|
assert!(matches!(stash, Some(SystemPromptStash::Prepended)));
|
|
}
|
|
|
|
#[test]
|
|
fn apply_override_no_op_when_none() {
|
|
let mut msgs = vec![ChatMessage::system("sys"), ChatMessage::user("hi")];
|
|
let stash = apply_system_prompt_override(&mut msgs, None);
|
|
assert!(stash.is_none());
|
|
assert_eq!(msgs[0].content, "sys");
|
|
}
|
|
|
|
#[test]
|
|
fn apply_override_no_op_for_empty_string() {
|
|
let mut msgs = vec![ChatMessage::system("sys")];
|
|
let stash = apply_system_prompt_override(&mut msgs, Some(""));
|
|
assert!(stash.is_none());
|
|
assert_eq!(msgs[0].content, "sys");
|
|
}
|
|
|
|
#[test]
|
|
fn restore_override_replaces_back() {
|
|
let mut msgs = vec![ChatMessage::system("new"), ChatMessage::user("hi")];
|
|
restore_system_prompt_override(
|
|
&mut msgs,
|
|
Some(SystemPromptStash::Replaced {
|
|
original: "original".to_string(),
|
|
}),
|
|
);
|
|
assert_eq!(msgs[0].content, "original");
|
|
assert_eq!(msgs.len(), 2);
|
|
}
|
|
|
|
#[test]
|
|
fn restore_override_pops_synthetic() {
|
|
let mut msgs = vec![ChatMessage::system("new"), ChatMessage::user("hi")];
|
|
restore_system_prompt_override(&mut msgs, Some(SystemPromptStash::Prepended));
|
|
assert_eq!(msgs.len(), 1);
|
|
assert_eq!(msgs[0].role, "user");
|
|
}
|
|
|
|
#[test]
|
|
fn override_round_trip_preserves_original_system_message() {
|
|
let mut msgs = vec![
|
|
ChatMessage::system("original persona"),
|
|
ChatMessage::user("first user"),
|
|
assistant_text("first reply"),
|
|
];
|
|
let stash = apply_system_prompt_override(&mut msgs, Some("ephemeral persona"));
|
|
assert_eq!(msgs[0].content, "ephemeral persona");
|
|
restore_system_prompt_override(&mut msgs, stash);
|
|
assert_eq!(msgs[0].content, "original persona");
|
|
assert_eq!(msgs.len(), 3);
|
|
assert_eq!(msgs[1].role, "user");
|
|
assert_eq!(msgs[2].role, "assistant");
|
|
}
|
|
|
|
#[test]
|
|
fn override_with_synthetic_round_trip_drops_extra_message() {
|
|
let mut msgs = vec![ChatMessage::user("first user")];
|
|
let stash = apply_system_prompt_override(&mut msgs, Some("ephemeral"));
|
|
assert_eq!(msgs.len(), 2);
|
|
assert_eq!(msgs[0].role, "system");
|
|
restore_system_prompt_override(&mut msgs, stash);
|
|
assert_eq!(msgs.len(), 1);
|
|
assert_eq!(msgs[0].role, "user");
|
|
}
|
|
|
|
// ── Bootstrap prompt / backend resolution ─────────────────────────
|
|
// Resolution helpers run on every bootstrap turn — they pick the
|
|
// persisted system prompt and the chosen backend label. Bugs here
|
|
// would silently swap the persona or miscategorise a backend.
|
|
|
|
#[test]
|
|
fn bootstrap_system_prompt_falls_back_to_default_for_none() {
|
|
let out = resolve_bootstrap_system_prompt(None, None);
|
|
assert_eq!(out, BOOTSTRAP_DEFAULT_SYSTEM_PROMPT);
|
|
}
|
|
|
|
#[test]
|
|
fn bootstrap_system_prompt_falls_back_to_default_for_empty_string() {
|
|
// Apollo currently sends `''` when no persona is selected.
|
|
let out = resolve_bootstrap_system_prompt(Some(""), None);
|
|
assert_eq!(out, BOOTSTRAP_DEFAULT_SYSTEM_PROMPT);
|
|
}
|
|
|
|
#[test]
|
|
fn bootstrap_system_prompt_falls_back_to_default_for_whitespace() {
|
|
let out = resolve_bootstrap_system_prompt(Some(" \n\t "), None);
|
|
assert_eq!(out, BOOTSTRAP_DEFAULT_SYSTEM_PROMPT);
|
|
}
|
|
|
|
#[test]
|
|
fn bootstrap_system_prompt_uses_supplied_when_non_empty() {
|
|
let out = resolve_bootstrap_system_prompt(Some("you are a journal"), None);
|
|
assert_eq!(out, "you are a journal");
|
|
}
|
|
|
|
#[test]
|
|
fn bootstrap_system_prompt_does_not_strip_inner_whitespace() {
|
|
// Trim only happens at the edges — interior newlines and spacing
|
|
// (which Apollo's persona uses for tool listings) must survive.
|
|
let prompt = "line one\nline two\n bullet";
|
|
let out = resolve_bootstrap_system_prompt(Some(prompt), None);
|
|
assert_eq!(out, prompt);
|
|
}
|
|
|
|
#[test]
|
|
fn bootstrap_system_prompt_explicit_wins_over_persona_store() {
|
|
let out = resolve_bootstrap_system_prompt(
|
|
Some("explicit prompt"),
|
|
Some("stored persona prompt".to_string()),
|
|
);
|
|
assert_eq!(out, "explicit prompt");
|
|
}
|
|
|
|
#[test]
|
|
fn bootstrap_system_prompt_uses_persona_store_when_no_explicit() {
|
|
// Request carried persona_id but no system_prompt — the persona's
|
|
// stored prompt must be used, not the neutral default.
|
|
let out = resolve_bootstrap_system_prompt(None, Some("stored persona prompt".to_string()));
|
|
assert_eq!(out, "stored persona prompt");
|
|
|
|
// Empty explicit prompt behaves like None.
|
|
let out =
|
|
resolve_bootstrap_system_prompt(Some(""), Some("stored persona prompt".to_string()));
|
|
assert_eq!(out, "stored persona prompt");
|
|
}
|
|
|
|
#[test]
|
|
fn bootstrap_system_prompt_blank_persona_prompt_falls_to_default() {
|
|
let out = resolve_bootstrap_system_prompt(None, Some(" ".to_string()));
|
|
assert_eq!(out, BOOTSTRAP_DEFAULT_SYSTEM_PROMPT);
|
|
}
|
|
|
|
// ── Synthetic final-answer prompt scaffolding ──────────────────────
|
|
|
|
#[test]
|
|
fn synthetic_final_prompt_round_trip_leaves_no_scaffolding() {
|
|
// Exhausted-loop fallback: nudge pushed, model reply appended, nudge
|
|
// removed — the persisted transcript must contain the reply but not
|
|
// the synthetic user prompt (all three loop variants rely on this).
|
|
let mut msgs = vec![
|
|
ChatMessage::system("sys"),
|
|
ChatMessage::user("q"),
|
|
assistant_with_tool_call("lookup"),
|
|
ChatMessage::tool_result("data"),
|
|
];
|
|
let idx = push_synthetic_final_prompt(&mut msgs);
|
|
assert_eq!(msgs[idx].content, SYNTHETIC_FINAL_ANSWER_PROMPT);
|
|
|
|
msgs.push(assistant_text("final answer"));
|
|
remove_synthetic_final_prompt(&mut msgs, idx);
|
|
|
|
assert_eq!(msgs.len(), 5);
|
|
assert!(
|
|
msgs.iter()
|
|
.all(|m| m.content != SYNTHETIC_FINAL_ANSWER_PROMPT),
|
|
"synthetic prompt must not persist"
|
|
);
|
|
assert_eq!(msgs.last().unwrap().content, "final answer");
|
|
}
|
|
|
|
#[test]
|
|
fn remove_synthetic_final_prompt_is_noop_on_index_mismatch() {
|
|
// Defensive guard: if the message at idx isn't the synthetic prompt
|
|
// (index drift), nothing is removed.
|
|
let mut msgs = vec![ChatMessage::user("q"), assistant_text("a")];
|
|
remove_synthetic_final_prompt(&mut msgs, 0);
|
|
remove_synthetic_final_prompt(&mut msgs, 5);
|
|
assert_eq!(msgs.len(), 2);
|
|
}
|
|
|
|
#[test]
|
|
fn bootstrap_backend_defaults_to_local_when_none() {
|
|
let out = resolve_bootstrap_backend(None).unwrap();
|
|
assert_eq!(out, "local");
|
|
}
|
|
|
|
#[test]
|
|
fn bootstrap_backend_defaults_to_local_when_empty() {
|
|
let out = resolve_bootstrap_backend(Some("")).unwrap();
|
|
assert_eq!(out, "local");
|
|
}
|
|
|
|
#[test]
|
|
fn bootstrap_backend_accepts_local_and_hybrid_case_insensitively() {
|
|
assert_eq!(resolve_bootstrap_backend(Some("LOCAL")).unwrap(), "local");
|
|
assert_eq!(resolve_bootstrap_backend(Some("Hybrid")).unwrap(), "hybrid");
|
|
assert_eq!(
|
|
resolve_bootstrap_backend(Some(" local ")).unwrap(),
|
|
"local"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn bootstrap_backend_rejects_unknown_label() {
|
|
// `llamacpp` is no longer a per-request backend value — it's chosen
|
|
// at deploy time via `LLM_BACKEND`.
|
|
for label in &["openrouter", "llamacpp", "ollama"] {
|
|
let err = resolve_bootstrap_backend(Some(label)).unwrap_err();
|
|
let msg = format!("{}", err);
|
|
assert!(msg.contains("unknown backend"), "label={}", label);
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn cross_replay_rejects_local_to_hybrid() {
|
|
let err = validate_cross_replay("local", "hybrid").unwrap_err();
|
|
assert!(format!("{}", err).contains("local to hybrid"));
|
|
}
|
|
|
|
#[test]
|
|
fn cross_replay_allows_supported_transitions() {
|
|
assert!(validate_cross_replay("local", "local").is_ok());
|
|
assert!(validate_cross_replay("hybrid", "hybrid").is_ok());
|
|
// Hybrid → local replays the inlined description as plain text.
|
|
assert!(validate_cross_replay("hybrid", "local").is_ok());
|
|
}
|
|
|
|
#[test]
|
|
fn cross_replay_rejects_unknown_effective() {
|
|
// Both "openrouter" and the former "llamacpp" value are unknown now.
|
|
for label in &["openrouter", "llamacpp"] {
|
|
let err = validate_cross_replay("local", label).unwrap_err();
|
|
assert!(
|
|
format!("{}", err).contains("unknown backend"),
|
|
"label={}",
|
|
label
|
|
);
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn bootstrap_system_message_includes_path_and_persona() {
|
|
let out = build_bootstrap_system_message("you are helpful", "pics/IMG.jpg", None, None, "");
|
|
assert!(out.starts_with("you are helpful"));
|
|
assert!(out.contains("--- PHOTO CONTEXT ---"));
|
|
assert!(out.contains("Photo file path: pics/IMG.jpg"));
|
|
// No date supplied → "unknown" so the model doesn't guess.
|
|
assert!(out.contains("Date taken: unknown"));
|
|
assert!(!out.contains("GPS:"));
|
|
assert!(!out.contains("Visual description"));
|
|
}
|
|
|
|
#[test]
|
|
fn bootstrap_system_message_includes_date_when_supplied() {
|
|
let out =
|
|
build_bootstrap_system_message("voice", "pics/IMG.jpg", Some("2014-11-08"), None, "");
|
|
assert!(out.contains("Date taken: 2014-11-08"));
|
|
assert!(!out.contains("Date taken: unknown"));
|
|
}
|
|
|
|
#[test]
|
|
fn bootstrap_system_message_includes_gps_when_present() {
|
|
let out = build_bootstrap_system_message(
|
|
"voice",
|
|
"p.jpg",
|
|
Some("2020-01-01"),
|
|
Some((42.36123, -71.05789)),
|
|
"",
|
|
);
|
|
// Four decimals — enough for place lookup, short enough to
|
|
// not bloat the system prompt.
|
|
assert!(out.contains("GPS: 42.3612, -71.0579"));
|
|
}
|
|
|
|
#[test]
|
|
fn bootstrap_system_message_omits_gps_when_none() {
|
|
let out = build_bootstrap_system_message("voice", "p.jpg", Some("2020-01-01"), None, "");
|
|
assert!(!out.contains("GPS:"));
|
|
}
|
|
|
|
#[test]
|
|
fn bootstrap_system_message_includes_visual_block_when_supplied() {
|
|
let visual = "Visual description (from local vision model):\nA dog in a park.\n";
|
|
let out =
|
|
build_bootstrap_system_message("voice", "p.jpg", Some("2020-01-01"), None, visual);
|
|
assert!(out.contains("Photo file path: p.jpg"));
|
|
assert!(out.contains("A dog in a park"));
|
|
// Path before date before visual.
|
|
let path_pos = out.find("Photo file path:").unwrap();
|
|
let date_pos = out.find("Date taken:").unwrap();
|
|
let visual_pos = out.find("A dog in a park").unwrap();
|
|
assert!(path_pos < date_pos);
|
|
assert!(date_pos < visual_pos);
|
|
}
|
|
|
|
#[test]
|
|
fn bootstrap_system_message_trims_persona_trailing_whitespace() {
|
|
let out = build_bootstrap_system_message("voice \n\n\n", "p.jpg", None, None, "");
|
|
assert!(out.contains("voice\n\n--- PHOTO CONTEXT ---"));
|
|
}
|
|
|
|
#[test]
|
|
fn date_taken_for_context_prefers_exif_over_filename() {
|
|
// EXIF wins when both are present (matches the canonical
|
|
// date_resolver waterfall — EXIF is more reliable than
|
|
// import-named filenames).
|
|
let exif = Some(crate::database::models::ImageExif {
|
|
id: 0,
|
|
library_id: 1,
|
|
file_path: "Screenshot_2014-06-01.png".to_string(),
|
|
camera_make: None,
|
|
camera_model: None,
|
|
lens_model: None,
|
|
width: None,
|
|
height: None,
|
|
orientation: None,
|
|
gps_latitude: None,
|
|
gps_longitude: None,
|
|
gps_altitude: None,
|
|
focal_length: None,
|
|
aperture: None,
|
|
shutter_speed: None,
|
|
iso: None,
|
|
// 2021-08-15 12:00:00 UTC
|
|
date_taken: Some(1_629_028_800),
|
|
created_time: 0,
|
|
last_modified: 0,
|
|
content_hash: None,
|
|
size_bytes: None,
|
|
phash_64: None,
|
|
dhash_64: None,
|
|
duplicate_of_hash: None,
|
|
duplicate_decided_at: None,
|
|
date_taken_source: None,
|
|
original_date_taken: None,
|
|
original_date_taken_source: None,
|
|
clip_embedding: None,
|
|
clip_model_version: None,
|
|
});
|
|
let out = resolve_date_taken_for_context(&exif, "Screenshot_2014-06-01.png");
|
|
assert_eq!(out.as_deref(), Some("2021-08-15"));
|
|
}
|
|
|
|
#[test]
|
|
fn date_taken_for_context_falls_back_to_filename_when_no_exif() {
|
|
// memories::extract_date_from_filename requires date+time in
|
|
// the filename — date-only patterns aren't matched. Use the
|
|
// canonical screenshot pattern for the regression case.
|
|
let out = resolve_date_taken_for_context(&None, "Screenshot_2014-06-01-20-44-50.png");
|
|
assert_eq!(out.as_deref(), Some("2014-06-01"));
|
|
}
|
|
|
|
#[test]
|
|
fn date_taken_for_context_returns_none_when_neither_source() {
|
|
let out = resolve_date_taken_for_context(&None, "DSC_5171.JPG");
|
|
assert!(out.is_none());
|
|
}
|
|
}
|