Cameron CordesandClaude Opus 5.5
e029c81274
- Thread tools are bound server-side to the chat's own thread. The model used to pass contact_id, so practice mode could read any conversation it named; a model-supplied contact_id is now only checked, never trusted, and the prompts no longer hand it one. - Practice chats need no persona. They're stored under PRACTICE_PERSONA_ID and run with an inert stand-in, so deleting or editing personas can't hide or break them. persona_id is optional on create for as_contact. - The per-thread list filters in SQL (list_persona_chats_for_contact, served by idx_persona_chat_anchor_contact) instead of decoding every transcript. - fetch_mms_image checks the content type and a 20 MB cap before reading the body, so a video attachment isn't downloaded in full. - get_photos_near_time unions tags over every path sharing the content, per the multi-library model. - look_at_photo only opens indexed photos, so an excluded folder can't be read by guessing a path. - One date parser: get_photos_near_time uses date_info::parse_center. - Practice prompt source rewrapped. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
4219 lines
167 KiB
Rust
4219 lines
167 KiB
Rust
use anyhow::{Result, anyhow, bail};
|
||
use chrono::{TimeZone, Utc};
|
||
use opentelemetry::trace::Status;
|
||
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::utils::{normalize_path, retry_with_backoff};
|
||
use futures::stream::{BoxStream, StreamExt};
|
||
use tracing::Instrument;
|
||
use tracing_opentelemetry::OpenTelemetrySpanExt;
|
||
use uuid::Uuid;
|
||
|
||
pub const DEFAULT_MAX_ITERATIONS: usize = 6;
|
||
/// Assumed context window when the request doesn't specify `num_ctx` and
|
||
/// the backend can't report the window it serves (llama-swap can; see
|
||
/// `effective_num_ctx`). Real conversations rarely pass ~16k tokens, so 32k
|
||
/// keeps the truncation pass from gutting history the server could hold
|
||
/// (which also destroys the server's KV-cache prefix reuse). Override
|
||
/// per-deploy with AGENTIC_CHAT_DEFAULT_NUM_CTX.
|
||
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.
|
||
pub(crate) 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.
|
||
pub(crate) 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,
|
||
/// Local vision descriptions of transcript images, keyed by a hash of
|
||
/// the image data, so a local-generated chat continued on a text-only
|
||
/// (hybrid) model describes its photo once rather than every turn.
|
||
image_descriptions: Arc<Mutex<HashMap<u64, String>>>,
|
||
}
|
||
|
||
/// Cap on cached image descriptions; the cache is cleared when it fills.
|
||
/// One entry per photo chatted about since the server started.
|
||
const IMAGE_DESCRIPTION_CACHE_MAX: usize = 256;
|
||
|
||
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,
|
||
image_descriptions: Arc::new(Mutex::new(HashMap::new())),
|
||
}
|
||
}
|
||
|
||
/// Prepare a replayed transcript for a backend that can't take images.
|
||
///
|
||
/// A chat generated on the local backend carries the photo itself in its
|
||
/// first user message. Continuing it on hybrid would send that photo to
|
||
/// OpenRouter — which the OpenRouter client refuses — so each image is
|
||
/// described by the local vision model and the text takes its place in
|
||
/// the copy sent to the model. The returned stash puts the images back
|
||
/// before the transcript is saved, so the stored chat keeps its photo and
|
||
/// a later local turn still sees it.
|
||
///
|
||
/// A no-op for backends that take images inline. If the local model
|
||
/// can't describe an image the turn fails: going ahead without the
|
||
/// photo would quietly change what the chat is about.
|
||
async fn inline_images_as_text(
|
||
&self,
|
||
backend: &ResolvedBackend,
|
||
messages: &mut [ChatMessage],
|
||
) -> Result<ImageStash> {
|
||
if backend.images_inline {
|
||
return Ok(ImageStash::default());
|
||
}
|
||
let mut descriptions: HashMap<String, String> = HashMap::new();
|
||
for image in transcript_images(messages) {
|
||
let key = image_cache_key(&image);
|
||
let cached = self
|
||
.image_descriptions
|
||
.lock()
|
||
.ok()
|
||
.and_then(|c| c.get(&key).cloned());
|
||
let description = match cached {
|
||
Some(d) => d,
|
||
None => {
|
||
let d = backend.local().describe_image(&image).await.map_err(|e| {
|
||
anyhow!(
|
||
"couldn't describe the photo with the local vision model, so this \
|
||
chat can't continue on a text-only model: {e}"
|
||
)
|
||
})?;
|
||
if let Ok(mut cache) = self.image_descriptions.lock() {
|
||
if cache.len() >= IMAGE_DESCRIPTION_CACHE_MAX {
|
||
cache.clear();
|
||
}
|
||
cache.insert(key, d.clone());
|
||
}
|
||
d
|
||
}
|
||
};
|
||
descriptions.insert(image, description);
|
||
}
|
||
Ok(apply_image_descriptions(messages, &descriptions))
|
||
}
|
||
|
||
/// 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 mut dao = self.insight_dao.lock().expect("Unable to lock InsightDao");
|
||
let insight = match dao
|
||
.get_current_insight_for_library(library_id, &normalized)
|
||
.map_err(|e| anyhow!("failed to load insight: {:?}", e))?
|
||
{
|
||
Some(i) => i,
|
||
None => dao
|
||
.get_insight(&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,
|
||
})
|
||
}
|
||
|
||
#[tracing::instrument(name = "ai.insight.chat_turn", skip_all)]
|
||
pub async fn chat_turn(&self, req: ChatTurnRequest) -> Result<ChatTurnResult> {
|
||
let span = tracing::Span::current();
|
||
span.set_attribute("file_path", req.file_path.clone());
|
||
span.set_attribute("library_id", req.library_id as i64);
|
||
span.set_attribute("amend", req.amend);
|
||
|
||
if req.user_message.trim().is_empty() {
|
||
bail!("user_message must not be empty");
|
||
}
|
||
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("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 mut dao = self.insight_dao.lock().expect("Unable to lock InsightDao");
|
||
dao.get_current_insight_for_library(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("backend", kind.as_str());
|
||
|
||
let max_iterations = req
|
||
.max_iterations
|
||
.unwrap_or(DEFAULT_MAX_ITERATIONS)
|
||
.clamp(1, env_max_iterations());
|
||
span.set_attribute("max_iterations", max_iterations as i64);
|
||
|
||
let stored_model = insight.model_version.clone();
|
||
let overrides = SamplingOverrides {
|
||
model: continuation_model(
|
||
req.model.clone(),
|
||
&stored_model,
|
||
&stored_backend,
|
||
kind.as_str(),
|
||
),
|
||
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?;
|
||
// Text-only backend (hybrid) on a transcript holding the photo:
|
||
// swap it for a local description. Restored before saving.
|
||
let image_stash = self.inline_images_as_text(&backend, &mut messages).await?;
|
||
let model_used = backend.model().to_string();
|
||
span.set_attribute("model", model_used.clone());
|
||
|
||
// 5. Decide vision + tool set. In hybrid (describe-then-inline) mode
|
||
// we omit `describe_photo`. Otherwise trust the stored history:
|
||
// 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 (see `apply_context_budget`), sized to
|
||
// the window the model is actually served with.
|
||
let mut fitter = ContextFitter::for_client(backend.chat(), req.num_ctx).await;
|
||
let truncated = fitter.trim(&mut messages);
|
||
if truncated {
|
||
span.set_attribute("history_truncated", true);
|
||
}
|
||
// Everything in `messages` at this point is replayed history that is
|
||
// ALREADY in the tree; only what the turn appends past this mark
|
||
// becomes new nodes. Captured after the budget pass, not from
|
||
// `path.len()`: truncation drains from the middle, so `path.len()`
|
||
// over-counts and the slice below would either swallow the new user
|
||
// turn or panic on an out-of-range start index.
|
||
let history_len = messages.len();
|
||
|
||
// 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);
|
||
|
||
// 8. Agentic loop — same shape as insight_generator's, but capped
|
||
// tighter and dispatching tools through the shared executor.
|
||
// Each LLM / tool await runs inside `loop_span` so their spans
|
||
// nest under it.
|
||
let loop_span = tracing::info_span!("ai.chat.loop");
|
||
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) = fitter
|
||
.chat(backend.chat(), &messages, tools.clone())
|
||
.instrument(loop_span.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,
|
||
None,
|
||
)
|
||
.instrument(loop_span.clone())
|
||
.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) = fitter
|
||
.chat(backend.chat(), &messages, vec![])
|
||
.instrument(loop_span.clone())
|
||
.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_span.set_status(Status::Ok);
|
||
drop(loop_span);
|
||
|
||
// Drop the per-turn iteration-budget note from the system message
|
||
// before we persist so it doesn't snowball on each subsequent turn.
|
||
restore_system_content(&mut messages, original_system_content);
|
||
// Put the photo back so the stored transcript keeps it.
|
||
image_stash.restore(&mut messages);
|
||
|
||
// 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 new_messages = messages[history_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 mut d = dao.lock().expect("Unable to lock InsightDao");
|
||
d.store_insight(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 mut d = dao.lock().expect("Unable to lock InsightDao");
|
||
d.update_training_messages(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 mut dao = self.insight_dao.lock().expect("Unable to lock InsightDao");
|
||
dao.get_current_insight_for_library(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 mut d = dao.lock().expect("Unable to lock InsightDao");
|
||
d.update_training_messages(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 mut dao = self.insight_dao.lock().expect("Unable to lock InsightDao");
|
||
dao.get_current_insight_for_library(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 mut d = dao.lock().expect("Unable to lock InsightDao");
|
||
d.update_training_messages(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 mut dao = self.insight_dao.lock().expect("Unable to lock InsightDao");
|
||
let insight = dao
|
||
.get_current_insight_for_library(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.
|
||
///
|
||
/// Must be called from inside the request's span: the turn span is
|
||
/// created here, eagerly, so it parents on the caller. The returned
|
||
/// stream is polled by the HTTP body writer, outside any span.
|
||
pub fn chat_turn_stream(
|
||
self: Arc<Self>,
|
||
req: ChatTurnRequest,
|
||
) -> BoxStream<'static, ChatStreamEvent> {
|
||
let span = tracing::info_span!(
|
||
"ai.chat.turn.stream",
|
||
file_path = %req.file_path,
|
||
library_id = req.library_id,
|
||
);
|
||
let svc = self;
|
||
let s = async_stream::stream! {
|
||
match svc.chat_turn_stream_inner(req, span, 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,
|
||
span: tracing::Span,
|
||
_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();
|
||
// Carry the turn span into the task explicitly — tracing context is
|
||
// thread-local, so without `.instrument` every LLM / DB span inside
|
||
// the turn would orphan onto a fresh trace.
|
||
tokio::spawn(
|
||
async move {
|
||
let result = svc.run_streaming_turn(req, tx.clone()).await;
|
||
let span = tracing::Span::current();
|
||
if let Err(e) = result {
|
||
span.set_status(Status::error(format!("{e}")));
|
||
let _ = tx.send(ChatStreamEvent::Error(format!("{}", e))).await;
|
||
} else {
|
||
span.set_status(Status::Ok);
|
||
}
|
||
}
|
||
.instrument(span),
|
||
);
|
||
Ok(rx)
|
||
}
|
||
|
||
/// 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();
|
||
// Span covering the whole spawned turn execution. Created before the
|
||
// spawn so it parents on the dispatching request (same trace as the
|
||
// 202), and carried into the task with `.instrument`. It outlives the
|
||
// dispatch span, which ends at the 202 response.
|
||
let turn_span = tracing::info_span!(
|
||
"ai.chat.turn.execute",
|
||
turn_id = %turn_id,
|
||
library_id = req.library_id,
|
||
);
|
||
let handle = tokio::spawn(
|
||
async move {
|
||
let span = tracing::Span::current();
|
||
let result = svc
|
||
.run_streaming_turn_with_entry(req, entry_clone.clone(), &span)
|
||
.await;
|
||
if let Err(ref e) = result {
|
||
span.set_attribute("status", "error");
|
||
span.set_status(Status::error(format!("{e}")));
|
||
// Push the terminal event BEFORE flipping status: a replay
|
||
// reader treats a terminal status with no buffered tail as
|
||
// "closed", so the Error must be in the buffer first.
|
||
let _ = entry_clone
|
||
.push_event(ChatStreamEvent::Error(format!("{}", e)))
|
||
.await;
|
||
entry_clone.set_terminal_status(crate::ai::turn_registry::TurnStatus::Error);
|
||
} else {
|
||
span.set_attribute("status", "done");
|
||
span.set_status(Status::Ok);
|
||
}
|
||
}
|
||
.instrument(turn_span),
|
||
);
|
||
|
||
// Install the abort handle so DELETE can actually stop the task.
|
||
entry.set_abort_handle(handle.abort_handle());
|
||
|
||
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>,
|
||
turn_span: &tracing::Span,
|
||
) -> 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 mut dao = self.insight_dao.lock().expect("Unable to lock InsightDao");
|
||
dao.get_current_insight_for_library(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, turn_span)
|
||
.await;
|
||
}
|
||
let insight = existing_insight.expect("just checked Some above");
|
||
self.run_continuation_streaming_with_entry(req, normalized, insight, entry, turn_span)
|
||
.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>,
|
||
turn_span: &tracing::Span,
|
||
) -> 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 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: continuation_model(
|
||
req.model.clone(),
|
||
&stored_model,
|
||
&stored_backend,
|
||
kind.as_str(),
|
||
),
|
||
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?;
|
||
// Text-only backend (hybrid) on a transcript holding the photo:
|
||
// swap it for a local description. Restored before saving.
|
||
let image_stash = self.inline_images_as_text(&backend, &mut messages).await?;
|
||
let model_used = backend.model().to_string();
|
||
// llm.* attributes on the turn-level span (issue #130): clients can't
|
||
// know local-vs-hybrid, but the turn knows its resolved backend.
|
||
turn_span.set_attribute("llm.backend", kind.as_str());
|
||
turn_span.set_attribute("llm.model", model_used.clone());
|
||
|
||
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 mut fitter = ContextFitter::for_client(backend.chat(), req.num_ctx).await;
|
||
let truncated = fitter.trim(&mut messages);
|
||
if truncated {
|
||
let _ = entry.push_event(ChatStreamEvent::Truncated).await;
|
||
}
|
||
// See the note in `chat_turn`: the new-node boundary must be read
|
||
// after the budget pass, because truncation drains replayed history
|
||
// out of the middle of `messages`.
|
||
let history_len = messages.len();
|
||
|
||
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,
|
||
&mut fitter,
|
||
)
|
||
.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);
|
||
// Put the photo back so the stored transcript keeps it.
|
||
image_stash.restore(&mut messages);
|
||
|
||
if !req.amend {
|
||
restore_system_prompt_override(&mut messages, override_stash);
|
||
}
|
||
|
||
// Append new messages as tree nodes.
|
||
let new_messages = messages[history_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 mut d = dao.lock().expect("Unable to lock InsightDao");
|
||
d.store_insight(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 mut d = dao.lock().expect("Unable to lock InsightDao");
|
||
d.update_training_messages(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>,
|
||
turn_span: &tracing::Span,
|
||
) -> 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();
|
||
// llm.* attributes on the turn-level span (issue #130).
|
||
turn_span.set_attribute("llm.backend", kind.as_str());
|
||
turn_span.set_attribute("llm.model", model_used.clone());
|
||
|
||
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,
|
||
self.generator.display_zone(),
|
||
crate::camera_tz::server_zone(),
|
||
);
|
||
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 mut fitter = ContextFitter::for_client(backend.chat(), req.num_ctx).await;
|
||
|
||
let outcome = self
|
||
.run_streaming_agentic_loop_with_entry(
|
||
&backend,
|
||
&mut messages,
|
||
tools,
|
||
&image_base64,
|
||
&normalized,
|
||
req.user_id,
|
||
&active_persona,
|
||
max_iterations,
|
||
&entry,
|
||
&mut fitter,
|
||
)
|
||
.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 mut d = dao.lock().expect("Unable to lock InsightDao");
|
||
d.store_insight(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>,
|
||
fitter: &mut ContextFitter,
|
||
) -> Result<AgenticLoopOutcome> {
|
||
crate::ai::insight_chat::run_streaming_agentic_loop_with_entry(
|
||
&self.generator,
|
||
backend,
|
||
messages,
|
||
tools,
|
||
image_base64,
|
||
normalized,
|
||
user_id,
|
||
active_persona,
|
||
None,
|
||
max_iterations,
|
||
entry,
|
||
fitter,
|
||
)
|
||
.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 mut dao = self.insight_dao.lock().expect("Unable to lock InsightDao");
|
||
dao.get_current_insight_for_library(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: continuation_model(
|
||
req.model.clone(),
|
||
&stored_model,
|
||
&stored_backend,
|
||
kind.as_str(),
|
||
),
|
||
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?;
|
||
// Text-only backend (hybrid) on a transcript holding the photo:
|
||
// swap it for a local description. Restored before saving.
|
||
let image_stash = self.inline_images_as_text(&backend, &mut messages).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 mut fitter = ContextFitter::for_client(backend.chat(), req.num_ctx).await;
|
||
let truncated = fitter.trim(&mut messages);
|
||
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,
|
||
&mut fitter,
|
||
)
|
||
.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);
|
||
// Put the photo back so the stored transcript keeps it.
|
||
image_stash.restore(&mut messages);
|
||
|
||
// 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 mut d = dao.lock().expect("Unable to lock InsightDao");
|
||
d.store_insight(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 mut d = dao.lock().expect("Unable to lock InsightDao");
|
||
d.update_training_messages(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,
|
||
self.generator.display_zone(),
|
||
crate::camera_tz::server_zone(),
|
||
);
|
||
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 mut fitter = ContextFitter::for_client(backend.chat(), req.num_ctx).await;
|
||
|
||
let outcome = self
|
||
.run_streaming_agentic_loop(
|
||
&backend,
|
||
&mut messages,
|
||
tools,
|
||
&image_base64,
|
||
&normalized,
|
||
req.user_id,
|
||
&active_persona,
|
||
max_iterations,
|
||
&tx,
|
||
&mut fitter,
|
||
)
|
||
.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 mut d = dao.lock().expect("Unable to lock InsightDao");
|
||
d.store_insight(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>,
|
||
fitter: &mut ContextFitter,
|
||
) -> 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 = fitter
|
||
.chat_stream(backend.chat(), messages, 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 result = self
|
||
.generator
|
||
.execute_tool(
|
||
&tool_call.function.name,
|
||
&tool_call.function.arguments,
|
||
backend,
|
||
image_base64,
|
||
normalized,
|
||
user_id,
|
||
active_persona,
|
||
None,
|
||
)
|
||
.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 = fitter.chat_stream(backend.chat(), messages, 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 (ADR-0001). Waterfall: the row's display wall (new
|
||
/// contract when converted, per-source convention while unconverted) →
|
||
/// filename pattern (wall patterns as-is, unix-timestamp patterns in the
|
||
/// display zone) → `None`. The fs-time fallback that insight generation
|
||
/// 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,
|
||
display_zone: chrono_tz::Tz,
|
||
server_zone: Option<chrono_tz::Tz>,
|
||
) -> Option<String> {
|
||
if let Some(e) = exif
|
||
&& let Some(wall) = crate::date_view::row_display_wall(e, display_zone, server_zone)
|
||
{
|
||
return Some(wall.format("%Y-%m-%d").to_string());
|
||
}
|
||
crate::memories::extract_date_from_filename(file_path).and_then(|fd| {
|
||
let wall = match &fd {
|
||
crate::memories::FilenameDate::Wall(w) => *w,
|
||
crate::memories::FilenameDate::Utc(dt) => display_zone
|
||
.timestamp_opt(dt.timestamp(), 0)
|
||
.single()?
|
||
.naive_local(),
|
||
};
|
||
Some(wall.format("%Y-%m-%d").to_string())
|
||
})
|
||
}
|
||
|
||
/// Validate the backend a chat continuation is switching to.
|
||
///
|
||
/// Every transition is allowed now. `local → hybrid` used to be refused
|
||
/// because a local transcript carries the raw photo; the turn now replaces
|
||
/// it with a local vision description before anything reaches OpenRouter
|
||
/// (`InsightChatService::inline_images_as_text`), and the OpenRouter client
|
||
/// refuses image data outright as a backstop. `hybrid → local` replays the
|
||
/// inlined description as text. `stored` is kept for the call sites'
|
||
/// symmetry and future rules.
|
||
///
|
||
/// 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
|
||
);
|
||
}
|
||
Ok(())
|
||
}
|
||
|
||
/// The model to continue a chat with when the request names none: the one
|
||
/// the insight was generated with, but only on the same backend. After a
|
||
/// switch, the stored id belongs to the other backend (a llama-swap slot is
|
||
/// not an OpenRouter model), so the new backend's default is used instead.
|
||
fn continuation_model(
|
||
requested: Option<String>,
|
||
stored_model: &str,
|
||
stored_backend: &str,
|
||
effective_backend: &str,
|
||
) -> Option<String> {
|
||
requested
|
||
.or_else(|| (stored_backend == effective_backend).then(|| stored_model.to_string()))
|
||
.filter(|m| !m.is_empty())
|
||
}
|
||
|
||
/// Every distinct image in a transcript, in order of first appearance.
|
||
fn transcript_images(messages: &[ChatMessage]) -> Vec<String> {
|
||
let mut seen = std::collections::HashSet::new();
|
||
messages
|
||
.iter()
|
||
.flat_map(|m| m.images.iter().flatten())
|
||
.filter(|img| seen.insert((*img).clone()))
|
||
.cloned()
|
||
.collect()
|
||
}
|
||
|
||
fn image_cache_key(image: &str) -> u64 {
|
||
use std::hash::{Hash, Hasher};
|
||
let mut h = std::collections::hash_map::DefaultHasher::new();
|
||
image.hash(&mut h);
|
||
h.finish()
|
||
}
|
||
|
||
/// What `apply_image_descriptions` took out of a transcript, so it can be
|
||
/// put back before the transcript is saved.
|
||
#[derive(Debug, Default)]
|
||
struct ImageStash {
|
||
entries: Vec<StashedImages>,
|
||
}
|
||
|
||
#[derive(Debug)]
|
||
struct StashedImages {
|
||
index: usize,
|
||
original_content: String,
|
||
rewritten_content: String,
|
||
images: Option<Vec<String>>,
|
||
}
|
||
|
||
impl ImageStash {
|
||
/// Put the original content and images back. Each entry is matched by
|
||
/// index *and* by still carrying the rewritten text, so a turn that
|
||
/// reshaped the transcript can't get the wrong message restored.
|
||
fn restore(self, messages: &mut [ChatMessage]) {
|
||
for e in self.entries {
|
||
if let Some(m) = messages.get_mut(e.index)
|
||
&& m.content == e.rewritten_content
|
||
{
|
||
m.content = e.original_content;
|
||
m.images = e.images;
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
/// Replace each image in `messages` with its description, appended to the
|
||
/// message text. Returns what was taken out, for `ImageStash::restore`.
|
||
fn apply_image_descriptions(
|
||
messages: &mut [ChatMessage],
|
||
descriptions: &HashMap<String, String>,
|
||
) -> ImageStash {
|
||
let mut stash = ImageStash::default();
|
||
for (idx, m) in messages.iter_mut().enumerate() {
|
||
let Some(images) = m.images.as_ref().filter(|i| !i.is_empty()) else {
|
||
continue;
|
||
};
|
||
let mut content = m.content.clone();
|
||
for img in images {
|
||
if let Some(desc) = descriptions.get(img) {
|
||
content.push_str(&format!(
|
||
"\n\nVisual description of the photo (from the local vision model):\n{desc}"
|
||
));
|
||
}
|
||
}
|
||
stash.entries.push(StashedImages {
|
||
index: idx,
|
||
original_content: std::mem::replace(&mut m.content, content.clone()),
|
||
rewritten_content: content,
|
||
images: m.images.take(),
|
||
});
|
||
}
|
||
stash
|
||
}
|
||
|
||
/// 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.
|
||
pub(crate) 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)
|
||
}
|
||
|
||
/// Context window (tokens) to budget a chat turn against: the smaller of what
|
||
/// the request asked for and what the server actually serves. llama-swap
|
||
/// models fix their window at launch (`-c` in config.yaml) and ignore the
|
||
/// per-request `num_ctx`, so budgeting against a bigger requested value just
|
||
/// keeps more history than the server can hold.
|
||
pub(crate) async fn effective_num_ctx(
|
||
client: &dyn crate::ai::llm_client::LlmClient,
|
||
requested: Option<i32>,
|
||
) -> usize {
|
||
let requested = requested.filter(|n| *n > 0).map(|n| n as usize);
|
||
match (requested, client.context_window().await) {
|
||
(Some(r), Some(served)) => r.min(served),
|
||
(None, Some(served)) => served,
|
||
(Some(r), None) => r,
|
||
(None, None) => env_default_num_ctx() as usize,
|
||
}
|
||
}
|
||
|
||
/// Keeps each LLM call of a chat turn inside the context budget. Every call
|
||
/// goes through [`ContextFitter::chat`] / [`ContextFitter::chat_stream`],
|
||
/// which send a trimmed copy of `messages` (the caller's vector — and so what
|
||
/// gets persisted — is untouched) and, if the server still rejects the prompt
|
||
/// as too long (the byte-per-token estimate is only a heuristic), shrink the
|
||
/// budget from the server's own token counts and retry once. The shrunk
|
||
/// budget sticks for the rest of the turn.
|
||
pub struct ContextFitter {
|
||
budget_bytes: usize,
|
||
}
|
||
|
||
impl ContextFitter {
|
||
pub(crate) fn from_num_ctx(num_ctx: usize) -> Self {
|
||
let budget_tokens = num_ctx.saturating_sub(RESPONSE_HEADROOM_TOKENS);
|
||
Self {
|
||
budget_bytes: budget_tokens.saturating_mul(BYTES_PER_TOKEN),
|
||
}
|
||
}
|
||
|
||
/// Budget for `client`'s model; see [`effective_num_ctx`].
|
||
pub(crate) async fn for_client(
|
||
client: &dyn crate::ai::llm_client::LlmClient,
|
||
requested: Option<i32>,
|
||
) -> Self {
|
||
Self::from_num_ctx(effective_num_ctx(client, requested).await)
|
||
}
|
||
|
||
/// Trim `messages` in place to the budget. Used once on the replayed
|
||
/// history before the turn starts; returns true when anything was dropped.
|
||
pub(crate) fn trim(&self, messages: &mut Vec<ChatMessage>) -> bool {
|
||
apply_context_budget(messages, self.budget_bytes)
|
||
}
|
||
|
||
fn fit(&self, messages: &[ChatMessage]) -> Vec<ChatMessage> {
|
||
let mut fitted = messages.to_vec();
|
||
apply_context_budget(&mut fitted, self.budget_bytes);
|
||
fitted
|
||
}
|
||
|
||
/// Shrink the budget after a context-overflow error. Returns false (and
|
||
/// leaves the budget alone) when `err` isn't one.
|
||
fn shrink_after(&mut self, err: &anyhow::Error, sent: &[ChatMessage]) -> bool {
|
||
let Some(counts) = parse_context_overflow(&format!("{err:#}")) else {
|
||
return false;
|
||
};
|
||
let sent_bytes = estimate_bytes(sent).min(self.budget_bytes);
|
||
let shrunk = match counts {
|
||
// Scale what we sent by the server's measured overshoot, with 10%
|
||
// slack since trimming lands on message boundaries.
|
||
Some((prompt_tokens, n_ctx)) if prompt_tokens > 0 => {
|
||
let target = n_ctx.saturating_sub(RESPONSE_HEADROOM_TOKENS);
|
||
sent_bytes.saturating_mul(target) / prompt_tokens * 9 / 10
|
||
}
|
||
_ => sent_bytes / 2,
|
||
};
|
||
tracing::warn!(
|
||
old_budget_bytes = self.budget_bytes,
|
||
new_budget_bytes = shrunk,
|
||
"chat prompt exceeded the model's context window; trimming and retrying"
|
||
);
|
||
self.budget_bytes = shrunk;
|
||
true
|
||
}
|
||
|
||
pub(crate) async fn chat(
|
||
&mut self,
|
||
client: &dyn crate::ai::llm_client::LlmClient,
|
||
messages: &[ChatMessage],
|
||
tools: Vec<Tool>,
|
||
) -> Result<(ChatMessage, Option<i32>, Option<i32>)> {
|
||
let sent = self.fit(messages);
|
||
match client.chat_with_tools(sent.clone(), tools.clone()).await {
|
||
Err(e) if self.shrink_after(&e, &sent) => {
|
||
client.chat_with_tools(self.fit(messages), tools).await
|
||
}
|
||
result => result,
|
||
}
|
||
}
|
||
|
||
pub(crate) async fn chat_stream(
|
||
&mut self,
|
||
client: &dyn crate::ai::llm_client::LlmClient,
|
||
messages: &[ChatMessage],
|
||
tools: Vec<Tool>,
|
||
) -> Result<futures::stream::BoxStream<'static, Result<LlmStreamEvent>>> {
|
||
let sent = self.fit(messages);
|
||
match client
|
||
.chat_with_tools_stream(sent.clone(), tools.clone())
|
||
.await
|
||
{
|
||
Err(e) if self.shrink_after(&e, &sent) => {
|
||
client
|
||
.chat_with_tools_stream(self.fit(messages), tools)
|
||
.await
|
||
}
|
||
result => result,
|
||
}
|
||
}
|
||
}
|
||
|
||
/// Recognise a "prompt longer than the context window" error from any
|
||
/// backend. `None` = some other error. `Some(Some((prompt_tokens, n_ctx)))`
|
||
/// when the body carries llama-server's counts
|
||
/// (`exceed_context_size_error` with `n_prompt_tokens` / `n_ctx`),
|
||
/// `Some(None)` when it's an overflow without usable numbers.
|
||
fn parse_context_overflow(err: &str) -> Option<Option<(usize, usize)>> {
|
||
let lower = err.to_lowercase();
|
||
let is_overflow = lower.contains("exceed_context_size")
|
||
|| lower.contains("context_length_exceeded")
|
||
|| lower.contains("maximum context length")
|
||
|| (lower.contains("exceed") && lower.contains("context"));
|
||
if !is_overflow {
|
||
return None;
|
||
}
|
||
let num = |key: &str| -> Option<usize> {
|
||
let re = regex::Regex::new(&format!(r#""{key}"\s*:\s*(\d+)"#)).ok()?;
|
||
re.captures(err)?.get(1)?.as_str().parse().ok()
|
||
};
|
||
Some(num("n_prompt_tokens").zip(num("n_ctx")))
|
||
}
|
||
|
||
/// 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,
|
||
thread_contact: Option<i64>,
|
||
max_iterations: usize,
|
||
entry: &Arc<TurnEntry>,
|
||
fitter: &mut ContextFitter,
|
||
) -> 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();
|
||
// Only tools in this turn's catalog may run. The catalog is how every
|
||
// gate (per-persona corrections, practice mode's narrowed set, absent
|
||
// backing tables) is expressed, and `execute_tool` dispatches on name
|
||
// alone — without this a model that names a tool it was never offered
|
||
// (one it saw in an earlier turn, or simply guessed) would still run it.
|
||
let offered: std::collections::HashSet<String> =
|
||
tools.iter().map(|t| t.function.name.clone()).collect();
|
||
|
||
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 = fitter
|
||
.chat_stream(backend.chat(), messages, 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 result = if offered.contains(&tool_call.function.name) {
|
||
generator
|
||
.execute_tool(
|
||
&tool_call.function.name,
|
||
&tool_call.function.arguments,
|
||
backend,
|
||
image_base64,
|
||
normalized,
|
||
user_id,
|
||
active_persona,
|
||
thread_contact,
|
||
)
|
||
.await
|
||
} else {
|
||
log::warn!(
|
||
"refusing tool '{}': not offered this turn",
|
||
tool_call.function.name
|
||
);
|
||
format!(
|
||
"Error: tool '{}' is not available here.",
|
||
tool_call.function.name
|
||
)
|
||
};
|
||
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 = fitter.chat_stream(backend.chat(), messages, 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.
|
||
pub(crate) 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. Always
|
||
/// 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) plus that first exchange's final reply, i.e. the
|
||
/// insight itself.
|
||
///
|
||
/// Drops, oldest first, in three phases until the budget is met:
|
||
/// 1. tool-call scaffolding (an assistant `tool_calls` turn + its tool
|
||
/// results) inside *completed* exchanges — cheapest to lose, the final
|
||
/// reply that used them stays;
|
||
/// 2. whole completed exchanges (user turn through its final reply) after
|
||
/// the first. Dropping a whole exchange keeps user/assistant alternation
|
||
/// intact, which strict chat templates (Gemma) enforce;
|
||
/// 3. tool-call scaffolding in the in-progress exchange (the current turn's
|
||
/// own tool results, which grow during the agentic loop).
|
||
///
|
||
/// An exchange is "completed" when it contains an assistant reply without
|
||
/// tool calls. The current turn never is while the loop runs, so the user's
|
||
/// question is never dropped — even after `push_synthetic_final_prompt`
|
||
/// appends a later user message.
|
||
///
|
||
/// 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 || estimate_bytes(messages) <= budget_bytes {
|
||
return false;
|
||
}
|
||
let mut dropped_any = false;
|
||
while estimate_bytes(messages) > budget_bytes {
|
||
let Some(range) = next_budget_drop(messages) else {
|
||
break;
|
||
};
|
||
messages.drain(range);
|
||
dropped_any = true;
|
||
}
|
||
dropped_any
|
||
}
|
||
|
||
/// Pick the next span of `messages` for [`apply_context_budget`] to drop,
|
||
/// following its phase order. `None` when nothing droppable remains.
|
||
fn next_budget_drop(messages: &[ChatMessage]) -> Option<std::ops::Range<usize>> {
|
||
let user_idxs: Vec<usize> = messages
|
||
.iter()
|
||
.enumerate()
|
||
.filter(|(_, m)| m.role == "user")
|
||
.map(|(i, _)| i)
|
||
.collect();
|
||
let first_user = *user_idxs.first()?;
|
||
|
||
// Exchanges: each user turn through the message before the next one.
|
||
let exchanges: Vec<std::ops::Range<usize>> = user_idxs
|
||
.iter()
|
||
.enumerate()
|
||
.map(|(n, &u)| u..user_idxs.get(n + 1).copied().unwrap_or(messages.len()))
|
||
.collect();
|
||
let completed = |r: &std::ops::Range<usize>| {
|
||
messages[r.clone()]
|
||
.iter()
|
||
.any(|m| m.role == "assistant" && !has_tool_calls(m))
|
||
};
|
||
// Oldest tool-call turn (+ trailing tool results) inside `r`.
|
||
let tool_pair_in = |r: &std::ops::Range<usize>| {
|
||
let start = (r.start + 1..r.end)
|
||
.find(|&i| messages[i].role == "assistant" && has_tool_calls(&messages[i]))?;
|
||
let mut end = start + 1;
|
||
while end < r.end && messages[end].role == "tool" {
|
||
end += 1;
|
||
}
|
||
Some(start..end)
|
||
};
|
||
|
||
// Phase 1: tool scaffolding in completed exchanges.
|
||
if let Some(r) = exchanges
|
||
.iter()
|
||
.filter(|r| completed(r))
|
||
.find_map(tool_pair_in)
|
||
{
|
||
return Some(r);
|
||
}
|
||
// Phase 2: whole completed exchanges after the first.
|
||
if let Some(r) = exchanges
|
||
.iter()
|
||
.find(|r| r.start != first_user && completed(r))
|
||
{
|
||
return Some(r.clone());
|
||
}
|
||
// Phase 3: tool scaffolding in the in-progress exchange(s).
|
||
exchanges.iter().find_map(tool_pair_in)
|
||
}
|
||
|
||
fn has_tool_calls(m: &ChatMessage) -> bool {
|
||
m.tool_calls.as_ref().is_some_and(|c| !c.is_empty())
|
||
}
|
||
|
||
/// 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);
|
||
}
|
||
|
||
fn roles(msgs: &[ChatMessage]) -> Vec<&str> {
|
||
msgs.iter().map(|m| m.role.as_str()).collect()
|
||
}
|
||
|
||
#[test]
|
||
fn truncation_drops_old_plain_exchanges_when_no_tool_pairs_left() {
|
||
// Long chat of plain Q&A: nothing but whole exchanges to drop. The
|
||
// old pass could only drop tool pairs, so this never shrank.
|
||
let filler = "x".repeat(2_000);
|
||
let mut msgs = vec![
|
||
ChatMessage::system("sys"),
|
||
ChatMessage::user("first"),
|
||
assistant_text("the insight"),
|
||
ChatMessage::user(format!("q1 {filler}")),
|
||
assistant_text(&format!("a1 {filler}")),
|
||
ChatMessage::user(format!("q2 {filler}")),
|
||
assistant_text(&format!("a2 {filler}")),
|
||
ChatMessage::user("current question"),
|
||
];
|
||
let budget = estimate_bytes(&msgs) - 1_000;
|
||
assert!(apply_context_budget(&mut msgs, budget));
|
||
assert!(estimate_bytes(&msgs) <= budget);
|
||
// Oldest exchange (q1/a1) went; the insight exchange, q2/a2 and the
|
||
// current question stay, and roles still alternate.
|
||
assert_eq!(
|
||
roles(&msgs),
|
||
["system", "user", "assistant", "user", "assistant", "user"]
|
||
);
|
||
assert_eq!(msgs[2].content, "the insight");
|
||
assert!(msgs[3].content.starts_with("q2"));
|
||
assert_eq!(msgs[5].content, "current question");
|
||
}
|
||
|
||
#[test]
|
||
fn truncation_drops_tool_scaffolding_before_whole_exchanges() {
|
||
let mut msgs = vec![
|
||
ChatMessage::system("sys"),
|
||
ChatMessage::user("first"),
|
||
assistant_text("the insight"),
|
||
ChatMessage::user("q1"),
|
||
assistant_with_tool_call("lookup"),
|
||
ChatMessage::tool_result("r".repeat(4_000)),
|
||
assistant_text("a1"),
|
||
ChatMessage::user("current question"),
|
||
];
|
||
let budget = estimate_bytes(&msgs) - 1_000;
|
||
assert!(apply_context_budget(&mut msgs, budget));
|
||
// Only the tool pair was needed; q1/a1 survive.
|
||
assert_eq!(
|
||
roles(&msgs),
|
||
["system", "user", "assistant", "user", "assistant", "user"]
|
||
);
|
||
assert_eq!(msgs[4].content, "a1");
|
||
}
|
||
|
||
#[test]
|
||
fn truncation_never_drops_the_in_progress_question() {
|
||
// Mid-loop with the synthetic final prompt appended: the real
|
||
// question's exchange has no final reply yet, so phase 2 must skip
|
||
// it; only its tool results (phase 3) can go.
|
||
let mut msgs = vec![
|
||
ChatMessage::system("sys"),
|
||
ChatMessage::user("first"),
|
||
assistant_text("the insight"),
|
||
ChatMessage::user("real question"),
|
||
assistant_with_tool_call("lookup"),
|
||
ChatMessage::tool_result("r".repeat(8_000)),
|
||
ChatMessage::user("synthetic: answer now"),
|
||
];
|
||
assert!(apply_context_budget(&mut msgs, 600));
|
||
assert!(msgs.iter().any(|m| m.content == "real question"));
|
||
assert!(msgs.iter().all(|m| m.role != "tool"));
|
||
}
|
||
|
||
#[test]
|
||
fn parse_context_overflow_reads_llama_server_counts() {
|
||
let err = r#"llama-swap chat request failed: 400 Bad Request — {"error":{"code":400,"message":"the request exceeds the available context size, try increasing it","type":"exceed_context_size_error","n_prompt_tokens":40312,"n_ctx":32768}}"#;
|
||
assert_eq!(parse_context_overflow(err), Some(Some((40312, 32768))));
|
||
assert_eq!(
|
||
parse_context_overflow("400 — This model's maximum context length is 128000 tokens"),
|
||
Some(None)
|
||
);
|
||
assert_eq!(parse_context_overflow("500 Internal Server Error"), None);
|
||
}
|
||
|
||
/// Scripted client for `ContextFitter`: fails the first `fail_first`
|
||
/// calls with `error`, records the size of every prompt it's sent.
|
||
struct ScriptedClient {
|
||
window: Option<usize>,
|
||
fail_first: usize,
|
||
error: String,
|
||
sent: Mutex<Vec<usize>>,
|
||
}
|
||
|
||
#[async_trait::async_trait]
|
||
impl crate::ai::llm_client::LlmClient for ScriptedClient {
|
||
async fn generate(
|
||
&self,
|
||
_: &str,
|
||
_: Option<&str>,
|
||
_: Option<Vec<String>>,
|
||
) -> Result<String> {
|
||
unimplemented!()
|
||
}
|
||
async fn chat_with_tools(
|
||
&self,
|
||
messages: Vec<ChatMessage>,
|
||
_: Vec<Tool>,
|
||
) -> Result<(ChatMessage, Option<i32>, Option<i32>)> {
|
||
let mut sent = self.sent.lock().unwrap();
|
||
sent.push(messages.len());
|
||
if sent.len() <= self.fail_first {
|
||
bail!("{}", self.error);
|
||
}
|
||
Ok((assistant_text("ok"), None, None))
|
||
}
|
||
async fn chat_with_tools_stream(
|
||
&self,
|
||
_: Vec<ChatMessage>,
|
||
_: Vec<Tool>,
|
||
) -> Result<futures::stream::BoxStream<'static, Result<LlmStreamEvent>>> {
|
||
unimplemented!()
|
||
}
|
||
async fn generate_embeddings(&self, _: &[&str]) -> Result<Vec<Vec<f32>>> {
|
||
unimplemented!()
|
||
}
|
||
async fn describe_image(&self, _: &str) -> Result<String> {
|
||
unimplemented!()
|
||
}
|
||
async fn list_models(&self) -> Result<Vec<crate::ai::llm_client::ModelCapabilities>> {
|
||
unimplemented!()
|
||
}
|
||
async fn model_capabilities(
|
||
&self,
|
||
_: &str,
|
||
) -> Result<crate::ai::llm_client::ModelCapabilities> {
|
||
unimplemented!()
|
||
}
|
||
fn primary_model(&self) -> &str {
|
||
"scripted"
|
||
}
|
||
async fn context_window(&self) -> Option<usize> {
|
||
self.window
|
||
}
|
||
}
|
||
|
||
fn scripted(window: Option<usize>, fail_first: usize, error: &str) -> ScriptedClient {
|
||
ScriptedClient {
|
||
window,
|
||
fail_first,
|
||
error: error.to_string(),
|
||
sent: Mutex::new(Vec::new()),
|
||
}
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn effective_num_ctx_caps_request_at_served_window() {
|
||
let served_32k = scripted(Some(32_768), 0, "");
|
||
assert_eq!(effective_num_ctx(&served_32k, Some(65_536)).await, 32_768);
|
||
assert_eq!(effective_num_ctx(&served_32k, Some(8_192)).await, 8_192);
|
||
assert_eq!(effective_num_ctx(&served_32k, None).await, 32_768);
|
||
let unknown = scripted(None, 0, "");
|
||
assert_eq!(effective_num_ctx(&unknown, Some(65_536)).await, 65_536);
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn fitter_retries_overflow_with_a_smaller_prompt() {
|
||
let filler = "x".repeat(2_000);
|
||
let mut msgs = vec![
|
||
ChatMessage::system("sys"),
|
||
ChatMessage::user("first"),
|
||
assistant_text("insight"),
|
||
];
|
||
for i in 0..10 {
|
||
msgs.push(ChatMessage::user(format!("q{i} {filler}")));
|
||
msgs.push(assistant_text(&format!("a{i} {filler}")));
|
||
}
|
||
msgs.push(ChatMessage::user("current"));
|
||
let client = scripted(
|
||
None,
|
||
1,
|
||
r#"400 — {"error":{"type":"exceed_context_size_error","n_prompt_tokens":20000,"n_ctx":10000}}"#,
|
||
);
|
||
let mut fitter = ContextFitter::from_num_ctx(1_000_000);
|
||
let (reply, _, _) = fitter.chat(&client, &msgs, vec![]).await.unwrap();
|
||
assert_eq!(reply.content, "ok");
|
||
let sent = client.sent.lock().unwrap().clone();
|
||
assert_eq!(sent.len(), 2, "one retry");
|
||
assert_eq!(sent[0], msgs.len(), "first attempt fit the generous budget");
|
||
assert!(sent[1] < sent[0], "retry sent a trimmed prompt");
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn fitter_does_not_retry_other_errors() {
|
||
let client = scripted(None, 5, "502 Bad Gateway");
|
||
let mut fitter = ContextFitter::from_num_ctx(32_768);
|
||
let msgs = vec![ChatMessage::system("s"), ChatMessage::user("u")];
|
||
assert!(fitter.chat(&client, &msgs, vec![]).await.is_err());
|
||
assert_eq!(client.sent.lock().unwrap().len(), 1);
|
||
}
|
||
|
||
#[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_allows_every_known_transition() {
|
||
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());
|
||
// Local → hybrid swaps the photo for a local description first.
|
||
assert!(validate_cross_replay("local", "hybrid").is_ok());
|
||
}
|
||
|
||
#[test]
|
||
fn continuation_model_drops_the_stored_model_across_a_switch() {
|
||
// Same backend: keep the model the insight was generated with.
|
||
assert_eq!(
|
||
continuation_model(None, "gemma4-31", "local", "local").as_deref(),
|
||
Some("gemma4-31")
|
||
);
|
||
// Switched: a llama-swap slot name is not an OpenRouter model id.
|
||
assert_eq!(
|
||
continuation_model(None, "gemma4-31", "local", "hybrid"),
|
||
None
|
||
);
|
||
// An explicit request always wins.
|
||
assert_eq!(
|
||
continuation_model(Some("x/y".into()), "gemma4-31", "local", "hybrid").as_deref(),
|
||
Some("x/y")
|
||
);
|
||
assert_eq!(continuation_model(None, "", "local", "local"), None);
|
||
}
|
||
|
||
fn with_image(text: &str, img: &str) -> ChatMessage {
|
||
let mut m = ChatMessage::user(text.to_string());
|
||
m.images = Some(vec![img.to_string()]);
|
||
m
|
||
}
|
||
|
||
#[test]
|
||
fn transcript_images_lists_each_image_once() {
|
||
let msgs = vec![
|
||
ChatMessage::system("sys"),
|
||
with_image("describe", "IMG_A"),
|
||
ChatMessage::user("more"),
|
||
with_image("again", "IMG_A"),
|
||
];
|
||
assert_eq!(transcript_images(&msgs), vec!["IMG_A".to_string()]);
|
||
}
|
||
|
||
#[test]
|
||
fn photo_is_swapped_for_text_then_restored_exactly() {
|
||
let original = vec![
|
||
ChatMessage::system("sys"),
|
||
with_image("What's in this photo?", "IMG_A"),
|
||
assistant_text("A beach."),
|
||
ChatMessage::user("who's there?"),
|
||
];
|
||
let mut msgs = original.clone();
|
||
let descriptions =
|
||
HashMap::from([("IMG_A".to_string(), "Two people on sand.".to_string())]);
|
||
|
||
let stash = apply_image_descriptions(&mut msgs, &descriptions);
|
||
// Nothing the text-only model receives carries image data.
|
||
assert!(msgs.iter().all(|m| m.images.is_none()));
|
||
assert!(msgs[1].content.starts_with("What's in this photo?"));
|
||
assert!(msgs[1].content.contains("Two people on sand."));
|
||
|
||
// The turn appends its exchange, then the stash is restored before
|
||
// saving: the stored transcript keeps the photo, unchanged.
|
||
msgs.push(assistant_text("Two friends."));
|
||
stash.restore(&mut msgs);
|
||
assert_eq!(
|
||
serde_json::to_string(&msgs[..original.len()]).unwrap(),
|
||
serde_json::to_string(&original).unwrap()
|
||
);
|
||
assert_eq!(msgs[1].images, Some(vec!["IMG_A".to_string()]));
|
||
}
|
||
|
||
#[test]
|
||
fn restore_leaves_a_message_the_turn_changed() {
|
||
let mut msgs = vec![with_image("q", "IMG_A")];
|
||
let descriptions = HashMap::from([("IMG_A".to_string(), "desc".to_string())]);
|
||
let stash = apply_image_descriptions(&mut msgs, &descriptions);
|
||
msgs[0].content = "something else entirely".into();
|
||
stash.restore(&mut msgs);
|
||
assert_eq!(msgs[0].content, "something else entirely");
|
||
assert!(msgs[0].images.is_none());
|
||
}
|
||
|
||
#[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,
|
||
date_utc: Some(1_629_028_800),
|
||
camera_offset_minutes: None,
|
||
date_taken_tz_source: Some("utc_direct".to_string()),
|
||
original_date_utc: None,
|
||
original_camera_offset_minutes: None,
|
||
original_date_taken_tz_source: None,
|
||
});
|
||
let out = resolve_date_taken_for_context(
|
||
&exif,
|
||
"Screenshot_2014-06-01.png",
|
||
chrono_tz::Tz::America__New_York,
|
||
Some(chrono_tz::Tz::America__New_York),
|
||
);
|
||
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",
|
||
chrono_tz::Tz::America__New_York,
|
||
Some(chrono_tz::Tz::America__New_York),
|
||
);
|
||
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",
|
||
chrono_tz::Tz::America__New_York,
|
||
Some(chrono_tz::Tz::America__New_York),
|
||
);
|
||
assert!(out.is_none());
|
||
}
|
||
|
||
/// AC-4 / AC-5: a Japan-offset (+9 h) photo anchored at local 21:00.
|
||
/// The AI is told the camera wall clock (instant + offset — a midnight
|
||
/// rollover relative to the owner's display zone is expected and
|
||
/// correct), while the SMS/location/calendar windows anchor on the
|
||
/// true instant with zero offset math.
|
||
#[test]
|
||
fn japan_offset_row_ai_date_is_camera_wall_and_anchor_unshifted() {
|
||
use chrono::NaiveDate;
|
||
|
||
// 2024-07-05 23:30 UTC = 2024-07-06 08:30 JST = 2024-07-05 19:30 EDT.
|
||
let instant = NaiveDate::from_ymd_opt(2024, 7, 5)
|
||
.unwrap()
|
||
.and_hms_opt(23, 30, 0)
|
||
.unwrap()
|
||
.and_utc();
|
||
let date_utc = instant.timestamp();
|
||
|
||
let row = crate::database::models::ImageExif {
|
||
id: 0,
|
||
library_id: 1,
|
||
file_path: "tokyo.jpg".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,
|
||
date_taken: Some(date_utc),
|
||
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: Some("exif".to_string()),
|
||
original_date_taken: None,
|
||
original_date_taken_source: None,
|
||
clip_embedding: None,
|
||
clip_model_version: None,
|
||
date_utc: Some(date_utc),
|
||
camera_offset_minutes: Some(540),
|
||
date_taken_tz_source: Some("offset_tag".to_string()),
|
||
original_date_utc: None,
|
||
original_camera_offset_minutes: None,
|
||
original_date_taken_tz_source: None,
|
||
};
|
||
|
||
let ny = chrono_tz::Tz::America__New_York;
|
||
|
||
// The prompt/UI date is the camera wall: 2024-07-06 08:30 JST —
|
||
// one day AFTER the owner's display-zone date (midnight rollover
|
||
// is correct here, the AI must see what the camera saw).
|
||
let ai_date = resolve_date_taken_for_context(&Some(row.clone()), "tokyo.jpg", ny, Some(ny));
|
||
assert_eq!(ai_date.as_deref(), Some("2024-07-06"));
|
||
|
||
// The proximity windows (SMS / location / calendar) anchor on the
|
||
// true instant — no +9 h, no -4 h, nothing.
|
||
let anchor = crate::date_view::row_true_instant(&row, ny, Some(ny));
|
||
assert_eq!(anchor, Some(date_utc));
|
||
}
|
||
|
||
/// AC-6: an unconverted container video row — the AI context date is
|
||
/// the display-zone rendering of the true instant (the raw UTC creation
|
||
/// time), matching the metadata/UI date. No 4–5 h offset, no midnight
|
||
/// rollover relative to the UTC-naive reading.
|
||
#[test]
|
||
fn unconverted_video_row_ai_date_matches_display_wall() {
|
||
use chrono::NaiveDate;
|
||
|
||
// The container's raw creation string is "2024-07-06 01:30:00"
|
||
// (a true UTC instant). The old server parsed that naive string
|
||
// as server-local (EDT) → legacy epoch = 05:30 UTC. The AI must
|
||
// be told the owner's evening — July 5 — not July 6 (the UTC
|
||
// date the old code would have rendered).
|
||
let legacy = NaiveDate::from_ymd_opt(2024, 7, 6)
|
||
.unwrap()
|
||
.and_hms_opt(5, 30, 0)
|
||
.unwrap()
|
||
.and_utc()
|
||
.timestamp();
|
||
|
||
let row = Some(crate::database::models::ImageExif {
|
||
id: 0,
|
||
library_id: 1,
|
||
file_path: "movie.mp4".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,
|
||
date_taken: Some(legacy),
|
||
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: Some("exiftool".to_string()),
|
||
original_date_taken: None,
|
||
original_date_taken_source: None,
|
||
clip_embedding: None,
|
||
clip_model_version: None,
|
||
date_utc: None,
|
||
camera_offset_minutes: None,
|
||
date_taken_tz_source: None,
|
||
original_date_utc: None,
|
||
original_camera_offset_minutes: None,
|
||
original_date_taken_tz_source: None,
|
||
});
|
||
|
||
let ny = chrono_tz::Tz::America__New_York;
|
||
let ai_date = resolve_date_taken_for_context(&row, "movie.mp4", ny, Some(ny));
|
||
assert_eq!(ai_date.as_deref(), Some("2024-07-05"));
|
||
}
|
||
}
|