Add conversation branching to insight chat (rewind preserves forks)
Store chat history as a tree in training_messages instead of a flat
array. Rewind now sets active_leaf_id to the target node rather than
truncating, so discarded paths survive as alternate branches. Fork
indicators ("X/Y") mark divergence points and let the client load or
switch to alternate paths.
- llm_client: StoredChatNode/ChatHistoryStore with tree traversal
helpers (path_to_leaf, children_of, fork_at_node, leaves*, etc.) and
from_flat_array for backward-compatible reads of the old flat format.
- insight_chat: load_history/chat_turn/rewind_history operate on the
tree; new switch_branch and get_branches. render_tree_path walks every
node (including non-rendered tool-dispatch nodes) when computing fork
info, so a fork whose diverging child is a tool call still surfaces on
the following rendered message.
- handlers: GET /insights/chat/branches, POST /insights/chat/switch-branch,
and a branch_id param on the history endpoint.
Backward-compatible: flat arrays are converted to a tree lazily on read.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
+146
-23
@@ -1310,6 +1310,11 @@ pub struct ChatHistoryQuery {
|
|||||||
pub path: String,
|
pub path: String,
|
||||||
#[serde(default)]
|
#[serde(default)]
|
||||||
pub library: Option<String>,
|
pub library: Option<String>,
|
||||||
|
/// When set, load the branch anchored at this leaf ID instead of the
|
||||||
|
/// active branch. Allows the client to preview alternate conversation
|
||||||
|
/// paths without making them active.
|
||||||
|
#[serde(default)]
|
||||||
|
pub branch_id: Option<u64>,
|
||||||
}
|
}
|
||||||
|
|
||||||
#[derive(Debug, Serialize)]
|
#[derive(Debug, Serialize)]
|
||||||
@@ -1318,6 +1323,18 @@ pub struct ChatHistoryHttpResponse {
|
|||||||
pub turn_count: usize,
|
pub turn_count: usize,
|
||||||
pub model_version: String,
|
pub model_version: String,
|
||||||
pub backend: String,
|
pub backend: String,
|
||||||
|
pub active_leaf_id: u64,
|
||||||
|
/// The branch leaf ID being viewed. Equals `active_leaf_id` when viewing
|
||||||
|
/// the active branch, or the `branch_id` query param for alternates.
|
||||||
|
pub viewing_branch_id: u64,
|
||||||
|
#[serde(skip_serializing_if = "Vec::is_empty")]
|
||||||
|
pub fork_info: Vec<Option<ChatForkInfo>>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize)]
|
||||||
|
pub struct ChatForkInfo {
|
||||||
|
pub position: usize,
|
||||||
|
pub total: usize,
|
||||||
}
|
}
|
||||||
|
|
||||||
#[derive(Debug, Serialize)]
|
#[derive(Debug, Serialize)]
|
||||||
@@ -1414,31 +1431,49 @@ pub async fn chat_history_handler(
|
|||||||
.flatten()
|
.flatten()
|
||||||
.unwrap_or_else(|| app_state.primary_library());
|
.unwrap_or_else(|| app_state.primary_library());
|
||||||
|
|
||||||
match app_state.insight_chat.load_history(library.id, &query.path) {
|
match app_state
|
||||||
Ok(view) => HttpResponse::Ok().json(ChatHistoryHttpResponse {
|
.insight_chat
|
||||||
messages: view
|
.load_history(library.id, &query.path, query.branch_id)
|
||||||
.messages
|
{
|
||||||
|
Ok(view) => {
|
||||||
|
let fork_info: Vec<Option<ChatForkInfo>> = view
|
||||||
|
.fork_info
|
||||||
.into_iter()
|
.into_iter()
|
||||||
.map(|m| RenderedHistoryMessage {
|
.map(|f| {
|
||||||
role: m.role,
|
f.map(|fi| ChatForkInfo {
|
||||||
content: m.content,
|
position: fi.position,
|
||||||
is_initial: m.is_initial,
|
total: fi.total,
|
||||||
tools: m
|
})
|
||||||
.tools
|
|
||||||
.into_iter()
|
|
||||||
.map(|t| HistoryToolInvocation {
|
|
||||||
name: t.name,
|
|
||||||
arguments: t.arguments,
|
|
||||||
result: t.result,
|
|
||||||
result_truncated: t.result_truncated,
|
|
||||||
})
|
|
||||||
.collect(),
|
|
||||||
})
|
})
|
||||||
.collect(),
|
.collect();
|
||||||
turn_count: view.turn_count,
|
HttpResponse::Ok().json(ChatHistoryHttpResponse {
|
||||||
model_version: view.model_version,
|
messages: view
|
||||||
backend: view.backend,
|
.messages
|
||||||
}),
|
.into_iter()
|
||||||
|
.map(|m| RenderedHistoryMessage {
|
||||||
|
role: m.role,
|
||||||
|
content: m.content,
|
||||||
|
is_initial: m.is_initial,
|
||||||
|
tools: m
|
||||||
|
.tools
|
||||||
|
.into_iter()
|
||||||
|
.map(|t| HistoryToolInvocation {
|
||||||
|
name: t.name,
|
||||||
|
arguments: t.arguments,
|
||||||
|
result: t.result,
|
||||||
|
result_truncated: t.result_truncated,
|
||||||
|
})
|
||||||
|
.collect(),
|
||||||
|
})
|
||||||
|
.collect(),
|
||||||
|
turn_count: view.turn_count,
|
||||||
|
model_version: view.model_version,
|
||||||
|
backend: view.backend,
|
||||||
|
active_leaf_id: view.active_leaf_id,
|
||||||
|
viewing_branch_id: view.viewing_branch_id,
|
||||||
|
fork_info,
|
||||||
|
})
|
||||||
|
}
|
||||||
Err(e) => {
|
Err(e) => {
|
||||||
let msg = format!("{}", e);
|
let msg = format!("{}", e);
|
||||||
if msg.contains("no insight found") {
|
if msg.contains("no insight found") {
|
||||||
@@ -1452,6 +1487,94 @@ pub async fn chat_history_handler(
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// GET /insights/chat/branches — return the list of branch leaf IDs for a
|
||||||
|
/// photo's conversation tree. Each branch is a leaf node; the active branch
|
||||||
|
/// is the one matching `active_leaf_id`. Use with `branch_id` on the
|
||||||
|
/// history endpoint to load an alternate conversation path.
|
||||||
|
#[get("/insights/chat/branches")]
|
||||||
|
pub async fn chat_branches_handler(
|
||||||
|
_claims: Claims,
|
||||||
|
query: web::Query<ChatHistoryQuery>,
|
||||||
|
app_state: web::Data<AppState>,
|
||||||
|
) -> impl Responder {
|
||||||
|
let library = libraries::resolve_library_param_state(&app_state, query.library.as_deref())
|
||||||
|
.ok()
|
||||||
|
.flatten()
|
||||||
|
.unwrap_or_else(|| app_state.primary_library());
|
||||||
|
|
||||||
|
let (branches, active_leaf_id) =
|
||||||
|
match app_state.insight_chat.get_branches(library.id, &query.path) {
|
||||||
|
Ok(result) => result,
|
||||||
|
Err(e) => {
|
||||||
|
let msg = format!("{}", e);
|
||||||
|
if msg.contains("no insight found") {
|
||||||
|
return HttpResponse::NotFound().json(serde_json::json!({ "error": msg }));
|
||||||
|
} else if msg.contains("no chat history") {
|
||||||
|
return HttpResponse::Conflict().json(serde_json::json!({ "error": msg }));
|
||||||
|
} else {
|
||||||
|
return HttpResponse::InternalServerError()
|
||||||
|
.json(serde_json::json!({ "error": msg }));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
HttpResponse::Ok().json(serde_json::json!({
|
||||||
|
"branches": branches,
|
||||||
|
"active_leaf_id": active_leaf_id,
|
||||||
|
}))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// POST /insights/chat/switch-branch — switch the active branch to a
|
||||||
|
/// different conversation path. The new branch becomes the active
|
||||||
|
/// conversation; the previous active branch becomes a regular fork.
|
||||||
|
#[post("/insights/chat/switch-branch")]
|
||||||
|
pub async fn chat_switch_branch_handler(
|
||||||
|
_claims: Claims,
|
||||||
|
request: web::Json<ChatSwitchBranchRequest>,
|
||||||
|
app_state: web::Data<AppState>,
|
||||||
|
) -> impl Responder {
|
||||||
|
let library =
|
||||||
|
match libraries::resolve_library_param_state(&app_state, request.library.as_deref()) {
|
||||||
|
Ok(Some(lib)) => lib,
|
||||||
|
Ok(None) => app_state.primary_library(),
|
||||||
|
Err(e) => {
|
||||||
|
return HttpResponse::BadRequest().json(serde_json::json!({
|
||||||
|
"error": format!("invalid library: {}", e)
|
||||||
|
}));
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
match app_state
|
||||||
|
.insight_chat
|
||||||
|
.switch_branch(library.id, &request.file_path, request.branch_id)
|
||||||
|
.await
|
||||||
|
{
|
||||||
|
Ok(()) => HttpResponse::Ok().json(serde_json::json!({ "success": true })),
|
||||||
|
Err(e) => {
|
||||||
|
let msg = format!("{}", e);
|
||||||
|
log::error!("Switch branch failed: {}", msg);
|
||||||
|
if msg.contains("no insight found") {
|
||||||
|
HttpResponse::NotFound().json(serde_json::json!({ "error": msg }))
|
||||||
|
} else if msg.contains("no chat history") {
|
||||||
|
HttpResponse::Conflict().json(serde_json::json!({ "error": msg }))
|
||||||
|
} else if msg.contains("not found") || msg.contains("not a leaf") {
|
||||||
|
HttpResponse::BadRequest().json(serde_json::json!({ "error": msg }))
|
||||||
|
} else {
|
||||||
|
HttpResponse::InternalServerError().json(serde_json::json!({ "error": msg }))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Deserialize)]
|
||||||
|
pub struct ChatSwitchBranchRequest {
|
||||||
|
pub file_path: String,
|
||||||
|
#[serde(default)]
|
||||||
|
pub library: Option<String>,
|
||||||
|
/// The leaf node ID to switch to. Must be a valid leaf in the tree.
|
||||||
|
pub branch_id: u64,
|
||||||
|
}
|
||||||
|
|
||||||
/// POST /insights/chat/stream — streaming variant of /insights/chat.
|
/// POST /insights/chat/stream — streaming variant of /insights/chat.
|
||||||
/// Returns `text/event-stream` with one event per chat stream event.
|
/// Returns `text/event-stream` with one event per chat stream event.
|
||||||
#[post("/insights/chat/stream")]
|
#[post("/insights/chat/stream")]
|
||||||
|
|||||||
+394
-114
@@ -8,7 +8,9 @@ use tokio::sync::Mutex as TokioMutex;
|
|||||||
|
|
||||||
use crate::ai::backend::{BackendKind, ResolvedBackend, SamplingOverrides};
|
use crate::ai::backend::{BackendKind, ResolvedBackend, SamplingOverrides};
|
||||||
use crate::ai::insight_generator::InsightGenerator;
|
use crate::ai::insight_generator::InsightGenerator;
|
||||||
use crate::ai::llm_client::{ChatMessage, LlmStreamEvent, Tool};
|
use crate::ai::llm_client::{
|
||||||
|
BranchLeafInfo, ChatHistoryStore, ChatMessage, LlmStreamEvent, StoredChatNode, Tool,
|
||||||
|
};
|
||||||
use crate::ai::turn_registry::TurnEntry;
|
use crate::ai::turn_registry::TurnEntry;
|
||||||
use crate::ai::turn_registry::TurnRegistry;
|
use crate::ai::turn_registry::TurnRegistry;
|
||||||
use crate::database::InsightDao;
|
use crate::database::InsightDao;
|
||||||
@@ -137,19 +139,20 @@ impl InsightChatService {
|
|||||||
&self.insight_dao
|
&self.insight_dao
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Load the rendered transcript for chat-UI display. Filters internal
|
/// Load the rendered transcript for chat-UI display. Deserializes the
|
||||||
/// scaffolding (system message, tool turns, tool-dispatch-only assistant
|
/// `training_messages` tree, traverses to `active_leaf_id`, and renders
|
||||||
/// messages) and drops base64 images from user turns to keep payloads
|
/// the path from root to that leaf. Backward-compatible: if the stored
|
||||||
/// small. The first remaining user message is flagged `is_initial`.
|
/// value is a flat array (old format), converts to a tree automatically.
|
||||||
///
|
///
|
||||||
/// `library_id` scopes the lookup to one library — without it, a
|
/// `library_id` scopes the lookup to one library — without it, a
|
||||||
/// regenerate on lib1 can be shadowed on the next refresh by an
|
/// regenerate on lib1 can be shadowed on the next refresh by an
|
||||||
/// untouched `is_current=true` row in lib2 for the same rel_path.
|
/// untouched `is_current=true` row in lib2 for the same rel_path.
|
||||||
/// Falls back to the cross-library `get_insight` only when the
|
pub fn load_history(
|
||||||
/// scoped lookup misses, preserving the cross-library "show this
|
&self,
|
||||||
/// photo's primary insight" merge for the case where the active
|
library_id: i32,
|
||||||
/// library has no insight but another library does.
|
file_path: &str,
|
||||||
pub fn load_history(&self, library_id: i32, file_path: &str) -> Result<HistoryView> {
|
branch_id: Option<u64>,
|
||||||
|
) -> Result<HistoryView> {
|
||||||
let normalized = normalize_path(file_path);
|
let normalized = normalize_path(file_path);
|
||||||
let cx = opentelemetry::Context::new();
|
let cx = opentelemetry::Context::new();
|
||||||
let mut dao = self.insight_dao.lock().expect("Unable to lock InsightDao");
|
let mut dao = self.insight_dao.lock().expect("Unable to lock InsightDao");
|
||||||
@@ -168,98 +171,32 @@ impl InsightChatService {
|
|||||||
.training_messages
|
.training_messages
|
||||||
.as_ref()
|
.as_ref()
|
||||||
.ok_or_else(|| anyhow!("insight has no chat history (pre-agentic insight)"))?;
|
.ok_or_else(|| anyhow!("insight has no chat history (pre-agentic insight)"))?;
|
||||||
let messages: Vec<ChatMessage> = serde_json::from_str(raw)
|
|
||||||
.map_err(|e| anyhow!("failed to deserialize chat history: {}", e))?;
|
|
||||||
|
|
||||||
let mut rendered = Vec::new();
|
// Backward-compatible deserialization: flat array (old) vs tree (new).
|
||||||
let mut user_turns_seen = 0usize;
|
let store: ChatHistoryStore = if let Ok(arr) = serde_json::from_str::<Vec<ChatMessage>>(raw)
|
||||||
let mut assistant_turns_seen = 0usize;
|
{
|
||||||
|
ChatHistoryStore::from_flat_array(arr)
|
||||||
|
} else {
|
||||||
|
serde_json::from_str(raw)
|
||||||
|
.map_err(|e| anyhow!("failed to deserialize chat history: {}", e))?
|
||||||
|
};
|
||||||
|
|
||||||
// Accumulate tool invocations seen since the last user turn. An
|
// Use branch_id if provided, otherwise use active_leaf_id.
|
||||||
// invocation is: one assistant tool_call message (which may hold
|
let target_leaf = branch_id.unwrap_or(store.active_leaf_id);
|
||||||
// multiple calls) + the N following tool-role messages (one per call,
|
let path = store
|
||||||
// in order). They attach to the next assistant-with-content, which
|
.path_to_leaf(target_leaf)
|
||||||
// is the "final" reply for the current turn.
|
.ok_or_else(|| anyhow!("branch_id {} not found in tree", target_leaf))?;
|
||||||
//
|
|
||||||
// Wire shape from the model:
|
|
||||||
// assistant { tool_calls: [A, B], content: "" }
|
|
||||||
// tool { content: "result of A" }
|
|
||||||
// tool { content: "result of B" }
|
|
||||||
// assistant { content: "here's the answer" } ← rendered as final
|
|
||||||
let mut pending_tools: Vec<ToolInvocation> = Vec::new();
|
|
||||||
// Queue of (name, arguments) awaiting a tool_result to pair with.
|
|
||||||
let mut pending_calls: std::collections::VecDeque<(String, serde_json::Value)> =
|
|
||||||
std::collections::VecDeque::new();
|
|
||||||
|
|
||||||
for msg in &messages {
|
let (rendered, turn_count, _node_ids, fork_info) = render_tree_path(&store, &path);
|
||||||
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,
|
|
||||||
});
|
|
||||||
}
|
|
||||||
// If there's no pending call, the tool message is an
|
|
||||||
// orphan (shouldn't happen in practice) — skip silently.
|
|
||||||
}
|
|
||||||
"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() {
|
|
||||||
// Tool-dispatch turn: enqueue calls, wait for tool
|
|
||||||
// results on subsequent messages.
|
|
||||||
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;
|
|
||||||
}
|
|
||||||
// Final assistant reply for this turn — drain accumulated
|
|
||||||
// tools into it.
|
|
||||||
assistant_turns_seen += 1;
|
|
||||||
let tools = std::mem::take(&mut pending_tools);
|
|
||||||
pending_calls.clear(); // any leftover unpaired calls are dropped
|
|
||||||
rendered.push(RenderedMessage {
|
|
||||||
role: "assistant".to_string(),
|
|
||||||
content: msg.content.clone(),
|
|
||||||
is_initial: false,
|
|
||||||
tools,
|
|
||||||
});
|
|
||||||
}
|
|
||||||
"user" => {
|
|
||||||
let is_initial = user_turns_seen == 0;
|
|
||||||
user_turns_seen += 1;
|
|
||||||
// New user turn resets any in-flight tool state.
|
|
||||||
pending_tools.clear();
|
|
||||||
pending_calls.clear();
|
|
||||||
rendered.push(RenderedMessage {
|
|
||||||
role: "user".to_string(),
|
|
||||||
content: msg.content.clone(),
|
|
||||||
is_initial,
|
|
||||||
tools: Vec::new(),
|
|
||||||
});
|
|
||||||
}
|
|
||||||
_ => continue,
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
Ok(HistoryView {
|
Ok(HistoryView {
|
||||||
messages: rendered,
|
messages: rendered,
|
||||||
turn_count: assistant_turns_seen,
|
turn_count,
|
||||||
model_version: insight.model_version,
|
model_version: insight.model_version,
|
||||||
backend: insight.backend,
|
backend: insight.backend,
|
||||||
|
active_leaf_id: store.active_leaf_id,
|
||||||
|
viewing_branch_id: target_leaf,
|
||||||
|
fork_info,
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -315,8 +252,21 @@ impl InsightChatService {
|
|||||||
anyhow!("insight has no chat history; regenerate this insight in agentic mode")
|
anyhow!("insight has no chat history; regenerate this insight in agentic mode")
|
||||||
})?
|
})?
|
||||||
.clone();
|
.clone();
|
||||||
let mut messages: Vec<ChatMessage> = serde_json::from_str(&raw_history)
|
|
||||||
.map_err(|e| anyhow!("failed to deserialize chat history: {}", e))?;
|
// 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.
|
// 3. Resolve effective backend. Reject the unsupported switch.
|
||||||
let stored_backend = insight.backend.clone();
|
let stored_backend = insight.backend.clone();
|
||||||
@@ -509,8 +459,22 @@ impl InsightChatService {
|
|||||||
// 9. Persist. Append mode rewrites the JSON blob in place; amend
|
// 9. Persist. Append mode rewrites the JSON blob in place; amend
|
||||||
// mode regenerates the title and inserts a new insight row,
|
// mode regenerates the title and inserts a new insight row,
|
||||||
// relying on store_insight to flip prior rows' is_current=false.
|
// relying on store_insight to flip prior rows' is_current=false.
|
||||||
let json = serde_json::to_string(&messages)
|
// Append new messages (this turn's user + assistant exchanges) as
|
||||||
.map_err(|e| anyhow!("failed to serialize chat history: {}", e))?;
|
// tree nodes chained from the previous active_leaf_id.
|
||||||
|
let path_len = path.len();
|
||||||
|
let new_messages = messages[path_len..].to_vec();
|
||||||
|
let mut parent_id = Some(store.active_leaf_id);
|
||||||
|
for msg in &new_messages {
|
||||||
|
let new_id = store.append_node(parent_id, msg.clone());
|
||||||
|
parent_id = Some(new_id);
|
||||||
|
}
|
||||||
|
// Update active_leaf_id to the last node appended this turn.
|
||||||
|
if let Some(last_id) = parent_id {
|
||||||
|
store.active_leaf_id = last_id;
|
||||||
|
}
|
||||||
|
|
||||||
|
let json = serde_json::to_string(&store)
|
||||||
|
.map_err(|e| anyhow!("failed to serialize chat tree: {}", e))?;
|
||||||
|
|
||||||
let mut amended_insight_id: Option<i32> = None;
|
let mut amended_insight_id: Option<i32> = None;
|
||||||
if req.amend {
|
if req.amend {
|
||||||
@@ -597,11 +561,14 @@ impl InsightChatService {
|
|||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Truncate the stored conversation so the rendered message at
|
/// Rewind the conversation to the rendered message at
|
||||||
/// `discard_from_rendered_index` (and everything after it — including
|
/// `discard_from_rendered_index` by setting `active_leaf_id` to the
|
||||||
/// the tool-call scaffolding that produced a discarded assistant reply)
|
/// node just before the discarded message. Unlike the legacy flat-array
|
||||||
/// is removed. The initial user turn cannot be discarded; attempting to
|
/// approach, this preserves all fork branches — the discarded path
|
||||||
/// do so returns an error.
|
/// 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`.
|
/// Holds the per-file chat mutex so it serialises with `chat_turn`.
|
||||||
pub async fn rewind_history(
|
pub async fn rewind_history(
|
||||||
@@ -636,15 +603,30 @@ impl InsightChatService {
|
|||||||
.training_messages
|
.training_messages
|
||||||
.as_ref()
|
.as_ref()
|
||||||
.ok_or_else(|| anyhow!("insight has no chat history"))?;
|
.ok_or_else(|| anyhow!("insight has no chat history"))?;
|
||||||
let messages: Vec<ChatMessage> = serde_json::from_str(raw_history)
|
|
||||||
.map_err(|e| anyhow!("failed to deserialize chat history: {}", e))?;
|
|
||||||
|
|
||||||
let cut_at = find_raw_cut(&messages, discard_from_rendered_index)
|
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"))?;
|
.ok_or_else(|| anyhow!("discard_from_rendered_index out of range"))?;
|
||||||
|
|
||||||
let truncated = &messages[..cut_at];
|
store.active_leaf_id = new_active_leaf_id;
|
||||||
let json = serde_json::to_string(truncated)
|
let json = serde_json::to_string(&store)
|
||||||
.map_err(|e| anyhow!("failed to serialize truncated history: {}", e))?;
|
.map_err(|e| anyhow!("failed to serialize tree: {}", e))?;
|
||||||
|
|
||||||
let dao = self.insight_dao.clone();
|
let dao = self.insight_dao.clone();
|
||||||
let rows = retry_with_backoff("update_training_messages", 3, || {
|
let rows = retry_with_backoff("update_training_messages", 3, || {
|
||||||
@@ -652,7 +634,7 @@ impl InsightChatService {
|
|||||||
let mut d = dao.lock().expect("Unable to lock InsightDao");
|
let mut d = dao.lock().expect("Unable to lock InsightDao");
|
||||||
d.update_training_messages(&cx, library_id, &normalized, &json)
|
d.update_training_messages(&cx, library_id, &normalized, &json)
|
||||||
})
|
})
|
||||||
.map_err(|e| anyhow!("failed to persist truncated history: {:?}", e))?;
|
.map_err(|e| anyhow!("failed to persist rewound history: {:?}", e))?;
|
||||||
if rows == 0 {
|
if rows == 0 {
|
||||||
log::warn!(
|
log::warn!(
|
||||||
"update_training_messages (rewind) updated 0 rows for {} (lib {}), \
|
"update_training_messages (rewind) updated 0 rows for {} (lib {}), \
|
||||||
@@ -664,6 +646,108 @@ impl InsightChatService {
|
|||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Switch the active branch to the given leaf ID. The new branch becomes
|
||||||
|
/// the active conversation path; the previous active branch becomes a
|
||||||
|
/// regular fork. If the leaf ID doesn't exist, returns an error.
|
||||||
|
pub async fn switch_branch(
|
||||||
|
&self,
|
||||||
|
library_id: i32,
|
||||||
|
file_path: &str,
|
||||||
|
leaf_id: u64,
|
||||||
|
) -> Result<()> {
|
||||||
|
let normalized = normalize_path(file_path);
|
||||||
|
|
||||||
|
let lock_key = (library_id, normalized.clone());
|
||||||
|
let entry_lock = {
|
||||||
|
let mut locks = self.chat_locks.lock().await;
|
||||||
|
locks
|
||||||
|
.entry(lock_key.clone())
|
||||||
|
.or_insert_with(|| Arc::new(TokioMutex::new(())))
|
||||||
|
.clone()
|
||||||
|
};
|
||||||
|
let _guard = entry_lock.lock().await;
|
||||||
|
|
||||||
|
let insight = {
|
||||||
|
let cx = opentelemetry::Context::new();
|
||||||
|
let mut dao = self.insight_dao.lock().expect("Unable to lock InsightDao");
|
||||||
|
dao.get_current_insight_for_library(&cx, library_id, &normalized)
|
||||||
|
.map_err(|e| anyhow!("failed to load insight: {:?}", e))?
|
||||||
|
.ok_or_else(|| anyhow!("no insight found for path"))?
|
||||||
|
};
|
||||||
|
let raw_history = insight
|
||||||
|
.training_messages
|
||||||
|
.as_ref()
|
||||||
|
.ok_or_else(|| anyhow!("insight has no chat history"))?;
|
||||||
|
|
||||||
|
let mut store: ChatHistoryStore =
|
||||||
|
if let Ok(arr) = serde_json::from_str::<Vec<ChatMessage>>(raw_history) {
|
||||||
|
ChatHistoryStore::from_flat_array(arr)
|
||||||
|
} else {
|
||||||
|
serde_json::from_str(raw_history)
|
||||||
|
.map_err(|e| anyhow!("failed to deserialize chat history: {}", e))?
|
||||||
|
};
|
||||||
|
|
||||||
|
// Validate the leaf_id exists and is a leaf.
|
||||||
|
if !store.nodes.iter().any(|n| n.id == leaf_id) {
|
||||||
|
bail!("branch_id {} not found in tree", leaf_id);
|
||||||
|
}
|
||||||
|
if !store.children_of(leaf_id).is_empty() {
|
||||||
|
bail!("branch_id {} is not a leaf node", leaf_id);
|
||||||
|
}
|
||||||
|
|
||||||
|
store.active_leaf_id = leaf_id;
|
||||||
|
let json = serde_json::to_string(&store)
|
||||||
|
.map_err(|e| anyhow!("failed to serialize tree: {}", e))?;
|
||||||
|
|
||||||
|
let dao = self.insight_dao.clone();
|
||||||
|
let rows = retry_with_backoff("update_training_messages", 3, || {
|
||||||
|
let cx = opentelemetry::Context::new();
|
||||||
|
let mut d = dao.lock().expect("Unable to lock InsightDao");
|
||||||
|
d.update_training_messages(&cx, library_id, &normalized, &json)
|
||||||
|
})
|
||||||
|
.map_err(|e| anyhow!("failed to persist branch switch: {:?}", e))?;
|
||||||
|
if rows == 0 {
|
||||||
|
log::warn!(
|
||||||
|
"update_training_messages (switch_branch) updated 0 rows for {} (lib {}), \
|
||||||
|
concurrent regenerate likely flipped is_current",
|
||||||
|
normalized,
|
||||||
|
library_id
|
||||||
|
);
|
||||||
|
}
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Return branch leaf metadata (id, snippet, message_count) and the
|
||||||
|
/// active leaf ID for a photo's conversation tree.
|
||||||
|
pub fn get_branches(
|
||||||
|
&self,
|
||||||
|
library_id: i32,
|
||||||
|
file_path: &str,
|
||||||
|
) -> Result<(Vec<BranchLeafInfo>, u64)> {
|
||||||
|
let normalized = normalize_path(file_path);
|
||||||
|
let cx = opentelemetry::Context::new();
|
||||||
|
let mut dao = self.insight_dao.lock().expect("Unable to lock InsightDao");
|
||||||
|
let insight = dao
|
||||||
|
.get_current_insight_for_library(&cx, library_id, &normalized)
|
||||||
|
.map_err(|e| anyhow!("failed to load insight: {:?}", e))?
|
||||||
|
.ok_or_else(|| anyhow!("no insight found for path"))?;
|
||||||
|
|
||||||
|
let raw = insight
|
||||||
|
.training_messages
|
||||||
|
.as_ref()
|
||||||
|
.ok_or_else(|| anyhow!("insight has no chat history"))?;
|
||||||
|
|
||||||
|
let store: ChatHistoryStore = if let Ok(arr) = serde_json::from_str::<Vec<ChatMessage>>(raw)
|
||||||
|
{
|
||||||
|
ChatHistoryStore::from_flat_array(arr)
|
||||||
|
} else {
|
||||||
|
serde_json::from_str(raw)
|
||||||
|
.map_err(|e| anyhow!("failed to deserialize chat history: {}", e))?
|
||||||
|
};
|
||||||
|
|
||||||
|
Ok((store.leaves_with_info(), store.active_leaf_id))
|
||||||
|
}
|
||||||
|
|
||||||
/// Streaming variant of `chat_turn`. Emits user-facing events as the
|
/// Streaming variant of `chat_turn`. Emits user-facing events as the
|
||||||
/// conversation progresses: iteration starts, tool dispatch + result,
|
/// conversation progresses: iteration starts, tool dispatch + result,
|
||||||
/// text deltas from the final assistant reply, and a terminal `Done`
|
/// text deltas from the final assistant reply, and a terminal `Done`
|
||||||
@@ -828,8 +912,21 @@ impl InsightChatService {
|
|||||||
let raw_history = insight.training_messages.as_ref().ok_or_else(|| {
|
let raw_history = insight.training_messages.as_ref().ok_or_else(|| {
|
||||||
anyhow!("insight has no chat history; regenerate this insight in agentic mode")
|
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))?;
|
// Backward-compatible: flat array (old) vs tree (new).
|
||||||
|
let mut store: ChatHistoryStore =
|
||||||
|
if let Ok(arr) = serde_json::from_str::<Vec<ChatMessage>>(raw_history) {
|
||||||
|
ChatHistoryStore::from_flat_array(arr)
|
||||||
|
} else {
|
||||||
|
serde_json::from_str(raw_history)
|
||||||
|
.map_err(|e| anyhow!("failed to deserialize chat history: {}", e))?
|
||||||
|
};
|
||||||
|
|
||||||
|
let path = store
|
||||||
|
.path_to_leaf(store.active_leaf_id)
|
||||||
|
.ok_or_else(|| anyhow!("active_leaf_id {} not found in tree", store.active_leaf_id))?;
|
||||||
|
let path_len = path.len();
|
||||||
|
let mut messages: Vec<ChatMessage> = path.iter().map(|n| n.message.clone()).collect();
|
||||||
|
|
||||||
let stored_backend = insight.backend.clone();
|
let stored_backend = insight.backend.clone();
|
||||||
let effective_backend = req
|
let effective_backend = req
|
||||||
@@ -931,8 +1028,19 @@ impl InsightChatService {
|
|||||||
restore_system_prompt_override(&mut messages, override_stash);
|
restore_system_prompt_override(&mut messages, override_stash);
|
||||||
}
|
}
|
||||||
|
|
||||||
let json = serde_json::to_string(&messages)
|
// Append new messages as tree nodes.
|
||||||
.map_err(|e| anyhow!("failed to serialize chat history: {}", e))?;
|
let new_messages = messages[path_len..].to_vec();
|
||||||
|
let mut parent_id = Some(store.active_leaf_id);
|
||||||
|
for msg in &new_messages {
|
||||||
|
let new_id = store.append_node(parent_id, msg.clone());
|
||||||
|
parent_id = Some(new_id);
|
||||||
|
}
|
||||||
|
if let Some(last_id) = parent_id {
|
||||||
|
store.active_leaf_id = last_id;
|
||||||
|
}
|
||||||
|
|
||||||
|
let json = serde_json::to_string(&store)
|
||||||
|
.map_err(|e| anyhow!("failed to serialize chat tree: {}", e))?;
|
||||||
|
|
||||||
let mut amended_insight_id: Option<i32> = None;
|
let mut amended_insight_id: Option<i32> = None;
|
||||||
if req.amend {
|
if req.amend {
|
||||||
@@ -2164,6 +2272,7 @@ pub enum ChatStreamEvent {
|
|||||||
/// Is this raw message visible in the rendered transcript? Must match
|
/// Is this raw message visible in the rendered transcript? Must match
|
||||||
/// `load_history`'s filter exactly — `find_raw_cut` depends on it to map
|
/// `load_history`'s filter exactly — `find_raw_cut` depends on it to map
|
||||||
/// rendered indices back to raw positions.
|
/// rendered indices back to raw positions.
|
||||||
|
#[allow(dead_code)]
|
||||||
fn is_rendered(m: &ChatMessage) -> bool {
|
fn is_rendered(m: &ChatMessage) -> bool {
|
||||||
match m.role.as_str() {
|
match m.role.as_str() {
|
||||||
"user" => true,
|
"user" => true,
|
||||||
@@ -2189,6 +2298,7 @@ fn is_rendered(m: &ChatMessage) -> bool {
|
|||||||
/// regenerating after a failed turn — its optimistic user bubble lives at
|
/// regenerating after a failed turn — its optimistic user bubble lives at
|
||||||
/// the index just past the server's persisted history. Strictly past the end
|
/// the index just past the server's persisted history. Strictly past the end
|
||||||
/// (`discard > rendered_count`) returns `None`.
|
/// (`discard > rendered_count`) returns `None`.
|
||||||
|
#[allow(dead_code)]
|
||||||
pub(crate) fn find_raw_cut(
|
pub(crate) fn find_raw_cut(
|
||||||
messages: &[ChatMessage],
|
messages: &[ChatMessage],
|
||||||
discard_from_rendered_index: usize,
|
discard_from_rendered_index: usize,
|
||||||
@@ -2349,6 +2459,112 @@ pub(crate) fn restore_system_prompt_override(
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// 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,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Render a path of tree nodes into the UI-friendly `RenderedMessage` format,
|
||||||
|
/// matching the same filtering logic as the legacy flat-array `load_history`.
|
||||||
|
/// Returns `(rendered_messages, turn_count, node_ids, fork_info)` where
|
||||||
|
/// `node_ids[i]` is the tree node ID that produced `rendered_messages[i]` and
|
||||||
|
/// `fork_info[i]` carries the nearest divergence point at or upstream of it.
|
||||||
|
///
|
||||||
|
/// Fork detection walks *every* node in the path — including the non-rendered
|
||||||
|
/// tool-dispatch/tool/system nodes — so a fork whose diverging child is a
|
||||||
|
/// tool-dispatch assistant (empty content + tool_calls) is still attributed to
|
||||||
|
/// the next rendered message. Detecting forks only on rendered nodes would miss
|
||||||
|
/// regenerations where the model replied with a tool call.
|
||||||
|
fn render_tree_path(
|
||||||
|
store: &ChatHistoryStore,
|
||||||
|
path: &[&StoredChatNode],
|
||||||
|
) -> (Vec<RenderedMessage>, usize, Vec<u64>, Vec<Option<ForkInfo>>) {
|
||||||
|
let mut rendered = Vec::new();
|
||||||
|
let mut user_turns_seen = 0usize;
|
||||||
|
let mut assistant_turns_seen = 0usize;
|
||||||
|
let mut pending_tools: Vec<ToolInvocation> = Vec::new();
|
||||||
|
let mut pending_calls: std::collections::VecDeque<(String, serde_json::Value)> =
|
||||||
|
std::collections::VecDeque::new();
|
||||||
|
let mut node_ids: Vec<u64> = Vec::new();
|
||||||
|
let mut fork_info: Vec<Option<ForkInfo>> = Vec::new();
|
||||||
|
// Nearest divergence seen so far, carried forward across nodes (rendered or
|
||||||
|
// not) so every message at/after a fork shows the indicator.
|
||||||
|
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.
|
||||||
|
if let Some((position, total)) = store.fork_at_node(node.id) {
|
||||||
|
last_fork = Some(ForkInfo { position, total });
|
||||||
|
}
|
||||||
|
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.clone());
|
||||||
|
}
|
||||||
|
"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.clone());
|
||||||
|
}
|
||||||
|
_ => continue,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
(rendered, assistant_turns_seen, node_ids, fork_info)
|
||||||
|
}
|
||||||
|
|
||||||
/// View returned to clients for chat-UI rendering.
|
/// View returned to clients for chat-UI rendering.
|
||||||
#[derive(Debug)]
|
#[derive(Debug)]
|
||||||
pub struct HistoryView {
|
pub struct HistoryView {
|
||||||
@@ -2356,6 +2572,15 @@ pub struct HistoryView {
|
|||||||
pub turn_count: usize,
|
pub turn_count: usize,
|
||||||
pub model_version: String,
|
pub model_version: String,
|
||||||
pub backend: 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)]
|
#[derive(Debug)]
|
||||||
@@ -2581,6 +2806,61 @@ mod tests {
|
|||||||
assert!(!dropped);
|
assert!(!dropped);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn render_tree_path_detects_fork_at_tool_dispatch_child() {
|
||||||
|
// Regression: a regeneration where the model replies with a tool call
|
||||||
|
// forks at the user turn, but the diverging children are non-rendered
|
||||||
|
// tool-dispatch assistant nodes. The indicator must still surface on
|
||||||
|
// the rendered final reply that follows — earlier code only inspected
|
||||||
|
// rendered nodes and silently dropped the fork, orphaning the branch.
|
||||||
|
let mut store = ChatHistoryStore::empty();
|
||||||
|
let u = store.append_node(None, ChatMessage::user("q1"));
|
||||||
|
// Original branch: tool-dispatch → tool result → final reply.
|
||||||
|
let d1 = store.append_node(Some(u), assistant_with_tool_call("lookup"));
|
||||||
|
let t1 = store.append_node(Some(d1), ChatMessage::tool_result("data1"));
|
||||||
|
let _a1 = store.append_node(Some(t1), assistant_text("answer 1"));
|
||||||
|
// Regenerated branch: also a tool reply, forking at the user turn.
|
||||||
|
let d2 = store.append_node(Some(u), assistant_with_tool_call("lookup"));
|
||||||
|
let t2 = store.append_node(Some(d2), ChatMessage::tool_result("data2"));
|
||||||
|
let a2 = store.append_node(Some(t2), assistant_text("answer 2"));
|
||||||
|
store.active_leaf_id = a2;
|
||||||
|
|
||||||
|
let path = store.path_to_leaf(a2).expect("path to active leaf");
|
||||||
|
let (rendered, _turns, node_ids, fork_info) = render_tree_path(&store, &path);
|
||||||
|
|
||||||
|
assert_eq!(rendered.len(), 2, "user turn + final reply are rendered");
|
||||||
|
assert_eq!(rendered[1].content, "answer 2");
|
||||||
|
assert_eq!(node_ids, vec![u, a2]);
|
||||||
|
assert!(
|
||||||
|
fork_info[0].is_none(),
|
||||||
|
"user turn sits before the divergence"
|
||||||
|
);
|
||||||
|
let f = fork_info[1]
|
||||||
|
.as_ref()
|
||||||
|
.expect("fork indicator carried onto the rendered reply");
|
||||||
|
assert_eq!((f.position, f.total), (2, 2));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[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));
|
||||||
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn rewind_strips_assistant_and_tool_scaffolding() {
|
fn rewind_strips_assistant_and_tool_scaffolding() {
|
||||||
// Rendered: [user1, asst1, user2, asst2] → cut at rendered index 3
|
// Rendered: [user1, asst1, user2, asst2] → cut at rendered index 3
|
||||||
|
|||||||
@@ -171,6 +171,178 @@ pub struct ModelCapabilities {
|
|||||||
pub has_tool_calling: bool,
|
pub has_tool_calling: bool,
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// A single node in the conversation tree. Wraps ChatMessage with structural
|
||||||
|
/// metadata so the tree can be persisted as JSON in `training_messages`.
|
||||||
|
#[derive(Serialize, Deserialize, Clone, Debug)]
|
||||||
|
pub struct StoredChatNode {
|
||||||
|
pub id: u64,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
pub parent_id: Option<u64>,
|
||||||
|
pub message: ChatMessage,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Top-level container for tree-based chat history. Stored as JSON in the
|
||||||
|
/// `training_messages` column. Backward-compatible reader detects whether
|
||||||
|
/// the stored value is a flat array (old format) or this object (new format).
|
||||||
|
#[derive(Serialize, Deserialize, Clone, Debug)]
|
||||||
|
pub struct ChatHistoryStore {
|
||||||
|
pub nodes: Vec<StoredChatNode>,
|
||||||
|
pub active_leaf_id: u64,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl ChatHistoryStore {
|
||||||
|
/// Create an empty store. Test-only: production stores are always built
|
||||||
|
/// from `from_flat_array` or deserialized from `training_messages`.
|
||||||
|
#[cfg(test)]
|
||||||
|
pub fn empty() -> Self {
|
||||||
|
Self {
|
||||||
|
nodes: Vec::new(),
|
||||||
|
active_leaf_id: 0,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Convert a flat `Vec<ChatMessage>` (old format) into a tree.
|
||||||
|
/// Each message becomes a node chained by parent_id.
|
||||||
|
pub fn from_flat_array(messages: Vec<ChatMessage>) -> Self {
|
||||||
|
let mut nodes = Vec::with_capacity(messages.len());
|
||||||
|
let mut prev_id: Option<u64> = None;
|
||||||
|
for (i, msg) in messages.into_iter().enumerate() {
|
||||||
|
let id = (i + 1) as u64;
|
||||||
|
nodes.push(StoredChatNode {
|
||||||
|
id,
|
||||||
|
parent_id: prev_id,
|
||||||
|
message: msg,
|
||||||
|
});
|
||||||
|
prev_id = Some(id);
|
||||||
|
}
|
||||||
|
let active_leaf_id = prev_id.unwrap_or(0);
|
||||||
|
Self {
|
||||||
|
nodes,
|
||||||
|
active_leaf_id,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Allocate the next node ID (max existing + 1, or 1 if empty).
|
||||||
|
pub fn next_id(&self) -> u64 {
|
||||||
|
self.nodes.iter().map(|n| n.id).max().unwrap_or(0) + 1
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Append a new node as a child of the given parent. Returns the new node's id.
|
||||||
|
pub fn append_node(&mut self, parent_id: Option<u64>, message: ChatMessage) -> u64 {
|
||||||
|
let id = self.next_id();
|
||||||
|
self.nodes.push(StoredChatNode {
|
||||||
|
id,
|
||||||
|
parent_id,
|
||||||
|
message,
|
||||||
|
});
|
||||||
|
id
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Traverse from root to the given leaf_id, returning the path in order.
|
||||||
|
/// Returns None if the leaf_id doesn't exist.
|
||||||
|
pub fn path_to_leaf(&self, leaf_id: u64) -> Option<Vec<&StoredChatNode>> {
|
||||||
|
// Build a map from id to node for fast lookup.
|
||||||
|
let node_map: std::collections::HashMap<u64, usize> = self
|
||||||
|
.nodes
|
||||||
|
.iter()
|
||||||
|
.enumerate()
|
||||||
|
.map(|(i, n)| (n.id, i))
|
||||||
|
.collect();
|
||||||
|
let &_leaf_idx = node_map.get(&leaf_id)?;
|
||||||
|
|
||||||
|
let mut path_indices = Vec::new();
|
||||||
|
let mut current_id = leaf_id;
|
||||||
|
while let Some(&idx) = node_map.get(¤t_id) {
|
||||||
|
path_indices.push(idx);
|
||||||
|
match self.nodes[idx].parent_id {
|
||||||
|
Some(pid) => current_id = pid,
|
||||||
|
None => break,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
path_indices.reverse();
|
||||||
|
Some(path_indices.into_iter().map(|i| &self.nodes[i]).collect())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// All direct children of a given node, ordered by id (chronological).
|
||||||
|
pub fn children_of(&self, node_id: u64) -> Vec<&StoredChatNode> {
|
||||||
|
self.nodes
|
||||||
|
.iter()
|
||||||
|
.filter(|n| n.parent_id == Some(node_id))
|
||||||
|
.collect()
|
||||||
|
}
|
||||||
|
|
||||||
|
/// For a given node, return (position, total) where:
|
||||||
|
/// - total = number of children of the parent of this node (i.e., siblings + self)
|
||||||
|
/// - position = 1-based rank of this node among siblings, ordered by id
|
||||||
|
///
|
||||||
|
/// Returns None if the node has no parent (root) or the parent has 0 or 1 child.
|
||||||
|
pub fn fork_at_node(&self, node_id: u64) -> Option<(usize, usize)> {
|
||||||
|
let node = self.nodes.iter().find(|n| n.id == node_id)?;
|
||||||
|
let parent_id = node.parent_id?;
|
||||||
|
let children = self.children_of(parent_id);
|
||||||
|
let total = children.len();
|
||||||
|
if total <= 1 {
|
||||||
|
return None;
|
||||||
|
}
|
||||||
|
let position = children.iter().position(|c| c.id == node_id)? + 1;
|
||||||
|
Some((position, total))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// All leaf node IDs (nodes with no children).
|
||||||
|
pub fn leaves(&self) -> Vec<u64> {
|
||||||
|
let child_ids: std::collections::HashSet<u64> =
|
||||||
|
self.nodes.iter().filter_map(|n| n.parent_id).collect();
|
||||||
|
self.nodes
|
||||||
|
.iter()
|
||||||
|
.filter(|n| !child_ids.contains(&n.id))
|
||||||
|
.map(|n| n.id)
|
||||||
|
.collect()
|
||||||
|
}
|
||||||
|
|
||||||
|
/// The parent of the given node, if any. Test-only helper.
|
||||||
|
#[cfg(test)]
|
||||||
|
pub fn parent_of(&self, node_id: u64) -> Option<&StoredChatNode> {
|
||||||
|
let node = self.nodes.iter().find(|n| n.id == node_id)?;
|
||||||
|
let parent_id = node.parent_id?;
|
||||||
|
self.nodes.iter().find(|n| n.id == parent_id)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Return all leaf nodes with metadata: id, snippet (first 100 chars of
|
||||||
|
/// the first user-visible message after the divergence point), and
|
||||||
|
/// message_count (number of nodes in the path from root to leaf).
|
||||||
|
pub fn leaves_with_info(&self) -> Vec<BranchLeafInfo> {
|
||||||
|
let leaf_ids = self.leaves();
|
||||||
|
leaf_ids
|
||||||
|
.into_iter()
|
||||||
|
.map(|leaf_id| {
|
||||||
|
let path = self.path_to_leaf(leaf_id).unwrap_or_default();
|
||||||
|
let message_count = path.len();
|
||||||
|
// Snippet: first 100 chars of the first user/assistant message
|
||||||
|
// (skip system messages).
|
||||||
|
let snippet = path
|
||||||
|
.iter()
|
||||||
|
.find(|n| n.message.role != "system")
|
||||||
|
.map(|n| n.message.content.chars().take(100).collect::<String>())
|
||||||
|
.unwrap_or_default();
|
||||||
|
BranchLeafInfo {
|
||||||
|
id: leaf_id,
|
||||||
|
snippet,
|
||||||
|
message_count,
|
||||||
|
}
|
||||||
|
})
|
||||||
|
.collect()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Metadata about a branch leaf, returned by the branches endpoint.
|
||||||
|
#[derive(Serialize, Clone, Debug)]
|
||||||
|
pub struct BranchLeafInfo {
|
||||||
|
pub id: u64,
|
||||||
|
#[serde(skip_serializing_if = "String::is_empty")]
|
||||||
|
pub snippet: String,
|
||||||
|
pub message_count: usize,
|
||||||
|
}
|
||||||
|
|
||||||
/// Strip a leading `<think>…</think>` reasoning block from model output.
|
/// Strip a leading `<think>…</think>` reasoning block from model output.
|
||||||
///
|
///
|
||||||
/// Thinking models sometimes emit chain-of-thought inside think tags before
|
/// Thinking models sometimes emit chain-of-thought inside think tags before
|
||||||
@@ -221,4 +393,135 @@ mod tests {
|
|||||||
let raw = "<think>thinking forever";
|
let raw = "<think>thinking forever";
|
||||||
assert_eq!(strip_think_blocks(raw), raw);
|
assert_eq!(strip_think_blocks(raw), raw);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// ── ChatHistoryStore tree traversal tests ──────────────────────────
|
||||||
|
|
||||||
|
fn mk_msg(role: &str, content: &str) -> ChatMessage {
|
||||||
|
match role {
|
||||||
|
"system" => ChatMessage::system(content),
|
||||||
|
"user" => ChatMessage::user(content),
|
||||||
|
_ => ChatMessage {
|
||||||
|
role: role.to_string(),
|
||||||
|
content: content.to_string(),
|
||||||
|
tool_calls: None,
|
||||||
|
images: None,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn store_from_flat_array_chains_parent_ids() {
|
||||||
|
let msgs = vec![
|
||||||
|
mk_msg("system", "sys"),
|
||||||
|
mk_msg("user", "hello"),
|
||||||
|
mk_msg("assistant", "hi there"),
|
||||||
|
];
|
||||||
|
let store = ChatHistoryStore::from_flat_array(msgs);
|
||||||
|
assert_eq!(store.nodes.len(), 3);
|
||||||
|
assert_eq!(store.nodes[0].id, 1);
|
||||||
|
assert_eq!(store.nodes[0].parent_id, None);
|
||||||
|
assert_eq!(store.nodes[1].parent_id, Some(1));
|
||||||
|
assert_eq!(store.nodes[2].parent_id, Some(2));
|
||||||
|
assert_eq!(store.active_leaf_id, 3);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn store_path_to_leaf_returns_ordered_path() {
|
||||||
|
let mut store = ChatHistoryStore::empty();
|
||||||
|
let n1 = store.append_node(None, mk_msg("user", "u1"));
|
||||||
|
let n2 = store.append_node(Some(n1), mk_msg("assistant", "a1"));
|
||||||
|
let n3 = store.append_node(Some(n2), mk_msg("user", "u2"));
|
||||||
|
let path = store.path_to_leaf(n3).unwrap();
|
||||||
|
assert_eq!(path.len(), 3);
|
||||||
|
assert_eq!(path[0].id, n1);
|
||||||
|
assert_eq!(path[1].id, n2);
|
||||||
|
assert_eq!(path[2].id, n3);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn store_path_to_leaf_returns_none_for_missing_id() {
|
||||||
|
let store = ChatHistoryStore::empty();
|
||||||
|
assert!(store.path_to_leaf(999).is_none());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn store_children_of_returns_direct_children() {
|
||||||
|
let mut store = ChatHistoryStore::empty();
|
||||||
|
let root = store.append_node(None, mk_msg("user", "root"));
|
||||||
|
let c1 = store.append_node(Some(root), mk_msg("assistant", "c1"));
|
||||||
|
let c2 = store.append_node(Some(root), mk_msg("user", "c2"));
|
||||||
|
let _d1 = store.append_node(Some(c1), mk_msg("assistant", "d1"));
|
||||||
|
let children = store.children_of(root);
|
||||||
|
assert_eq!(children.len(), 2);
|
||||||
|
assert_eq!(children[0].id, c1);
|
||||||
|
assert_eq!(children[1].id, c2);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn store_fork_at_node_detects_fork() {
|
||||||
|
let mut store = ChatHistoryStore::empty();
|
||||||
|
let root = store.append_node(None, mk_msg("user", "root"));
|
||||||
|
let c1 = store.append_node(Some(root), mk_msg("assistant", "c1"));
|
||||||
|
let c2 = store.append_node(Some(root), mk_msg("user", "c2"));
|
||||||
|
let c3 = store.append_node(Some(root), mk_msg("user", "c3"));
|
||||||
|
// c1 is position 1 of 3 siblings
|
||||||
|
assert_eq!(store.fork_at_node(c1), Some((1, 3)));
|
||||||
|
assert_eq!(store.fork_at_node(c2), Some((2, 3)));
|
||||||
|
assert_eq!(store.fork_at_node(c3), Some((3, 3)));
|
||||||
|
// root has no parent, so no fork
|
||||||
|
assert!(store.fork_at_node(root).is_none());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn store_fork_at_node_returns_none_for_single_child() {
|
||||||
|
let mut store = ChatHistoryStore::empty();
|
||||||
|
let root = store.append_node(None, mk_msg("user", "root"));
|
||||||
|
let _only = store.append_node(Some(root), mk_msg("assistant", "only"));
|
||||||
|
assert!(store.fork_at_node(2).is_none());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn store_leaves_returns_leaf_ids() {
|
||||||
|
let mut store = ChatHistoryStore::empty();
|
||||||
|
let root = store.append_node(None, mk_msg("user", "root"));
|
||||||
|
let c1 = store.append_node(Some(root), mk_msg("assistant", "c1"));
|
||||||
|
let c2 = store.append_node(Some(root), mk_msg("user", "c2"));
|
||||||
|
let _d1 = store.append_node(Some(c1), mk_msg("assistant", "d1"));
|
||||||
|
let leaves = store.leaves();
|
||||||
|
assert_eq!(leaves.len(), 2);
|
||||||
|
assert!(leaves.contains(&c2));
|
||||||
|
assert!(leaves.contains(&_d1));
|
||||||
|
assert!(!leaves.contains(&root));
|
||||||
|
assert!(!leaves.contains(&c1));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn store_parent_of_returns_parent() {
|
||||||
|
let mut store = ChatHistoryStore::empty();
|
||||||
|
let root = store.append_node(None, mk_msg("user", "root"));
|
||||||
|
let child = store.append_node(Some(root), mk_msg("assistant", "child"));
|
||||||
|
let parent = store.parent_of(child).unwrap();
|
||||||
|
assert_eq!(parent.id, root);
|
||||||
|
assert!(store.parent_of(root).is_none());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn store_next_id_increments() {
|
||||||
|
let mut store = ChatHistoryStore::empty();
|
||||||
|
assert_eq!(store.next_id(), 1);
|
||||||
|
store.append_node(None, mk_msg("user", "a"));
|
||||||
|
assert_eq!(store.next_id(), 2);
|
||||||
|
store.append_node(Some(1), mk_msg("assistant", "b"));
|
||||||
|
assert_eq!(store.next_id(), 3);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn store_append_node_sets_parent() {
|
||||||
|
let mut store = ChatHistoryStore::empty();
|
||||||
|
let parent = store.append_node(None, mk_msg("user", "parent"));
|
||||||
|
let child = store.append_node(Some(parent), mk_msg("assistant", "child"));
|
||||||
|
assert_eq!(store.nodes.len(), 2);
|
||||||
|
assert_eq!(store.nodes[1].parent_id, Some(parent));
|
||||||
|
assert_eq!(store.nodes[1].id, child);
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
+6
-6
@@ -25,12 +25,12 @@ pub use daily_summary_job::{
|
|||||||
generate_daily_summaries, strip_summary_boilerplate,
|
generate_daily_summaries, strip_summary_boilerplate,
|
||||||
};
|
};
|
||||||
pub use handlers::{
|
pub use handlers::{
|
||||||
cancel_generation_handler, cancel_turn_handler, chat_history_handler, chat_rewind_handler,
|
cancel_generation_handler, cancel_turn_handler, chat_branches_handler, chat_history_handler,
|
||||||
chat_stream_handler, chat_turn_handler, delete_insight_handler, export_training_data_handler,
|
chat_rewind_handler, chat_stream_handler, chat_switch_branch_handler, chat_turn_handler,
|
||||||
generate_agentic_insight_handler, generate_insight_handler, generation_status_handler,
|
delete_insight_handler, export_training_data_handler, generate_agentic_insight_handler,
|
||||||
get_all_insights_handler, get_available_models_handler, get_insight_handler,
|
generate_insight_handler, generation_status_handler, get_all_insights_handler,
|
||||||
get_insight_history_handler, get_openrouter_models_handler, rate_insight_handler,
|
get_available_models_handler, get_insight_handler, get_insight_history_handler,
|
||||||
turn_async_handler, turn_replay_handler,
|
get_openrouter_models_handler, rate_insight_handler, turn_async_handler, turn_replay_handler,
|
||||||
};
|
};
|
||||||
pub use insight_generator::InsightGenerator;
|
pub use insight_generator::InsightGenerator;
|
||||||
pub use llamacpp::LlamaCppClient;
|
pub use llamacpp::LlamaCppClient;
|
||||||
|
|||||||
@@ -377,6 +377,8 @@ fn main() -> std::io::Result<()> {
|
|||||||
.service(ai::chat_stream_handler)
|
.service(ai::chat_stream_handler)
|
||||||
.service(ai::chat_history_handler)
|
.service(ai::chat_history_handler)
|
||||||
.service(ai::chat_rewind_handler)
|
.service(ai::chat_rewind_handler)
|
||||||
|
.service(ai::chat_branches_handler)
|
||||||
|
.service(ai::chat_switch_branch_handler)
|
||||||
.service(ai::turn_async_handler)
|
.service(ai::turn_async_handler)
|
||||||
.service(ai::turn_replay_handler)
|
.service(ai::turn_replay_handler)
|
||||||
.service(ai::cancel_turn_handler)
|
.service(ai::cancel_turn_handler)
|
||||||
|
|||||||
Reference in New Issue
Block a user