zeph_memory/semantic/
trajectory.rs

1// SPDX-FileCopyrightText: 2026 Andrei G <bug-ops>
2// SPDX-License-Identifier: MIT OR Apache-2.0
3
4//! Trajectory-informed memory extraction (#2498).
5//!
6//! After each agent turn containing tool calls, a fast LLM provider analyzes the turn
7//! and produces procedural (reusable how-to patterns) and episodic (one-off event) entries.
8//! Entries are stored per-conversation so concurrent sessions do not interfere (critic S1).
9//!
10//! Extraction is always fire-and-forget (caller uses `tokio::spawn`) — no latency added to
11//! the response path (critic M3).
12
13use std::time::Duration;
14
15use serde::{Deserialize, Serialize};
16use tokio::time::timeout;
17use zeph_common::llm_response::extract_json_array_slice;
18use zeph_llm::any::AnyProvider;
19use zeph_llm::provider::{LlmProvider as _, Message, Role};
20
21use crate::error::MemoryError;
22
23const EXTRACTION_SYSTEM_PROMPT: &str = "\
24You are a trajectory memory extractor. Given messages from an agent turn that included tool \
25calls, extract reusable patterns and notable events.
26
27Rules:
281. Classify each entry as 'procedural' (a reusable how-to pattern — general technique) or \
29   'episodic' (a one-off event — specific occurrence).
302. Focus on the intent (goal), outcome (result), and tools used.
313. Confidence: 0.8-1.0 for clear outcomes, 0.4-0.7 for ambiguous ones.
324. Keep intent and outcome concise (one sentence each).
335. Return empty array if no meaningful entries can be extracted.
34
35Output JSON array:
36[
37  {
38    \"kind\": \"procedural|episodic\",
39    \"intent\": \"what the agent was trying to accomplish\",
40    \"outcome\": \"what actually happened\",
41    \"tools_used\": [\"tool_name\", ...],
42    \"confidence\": 0.0-1.0
43  }
44]";
45
46/// A single extracted trajectory entry (in-memory, before storage).
47#[derive(Debug, Clone)]
48pub struct TrajectoryEntry {
49    pub kind: String,
50    pub intent: String,
51    pub outcome: String,
52    pub tools_used: Vec<String>,
53    pub confidence: f64,
54}
55
56/// Configuration for trajectory extraction.
57pub struct TrajectoryExtractionConfig {
58    pub enabled: bool,
59    /// Maximum messages fed to the LLM per extraction pass.
60    pub max_messages: usize,
61    /// LLM timeout in seconds.
62    pub extraction_timeout_secs: u64,
63}
64
65#[derive(Debug, Deserialize, Serialize)]
66struct RawEntry {
67    kind: String,
68    intent: String,
69    outcome: String,
70    #[serde(default)]
71    tools_used: Vec<String>,
72    confidence: f64,
73}
74
75/// Extract trajectory entries from a turn's messages.
76///
77/// Returns the extracted entries. Parse failures are logged and treated as zero entries.
78///
79/// # Errors
80///
81/// Returns an error only for transport-level LLM failures.
82pub async fn extract_trajectory_entries(
83    provider: &AnyProvider,
84    messages: &[Message],
85    config: &TrajectoryExtractionConfig,
86) -> Result<Vec<TrajectoryEntry>, MemoryError> {
87    if !config.enabled || messages.is_empty() {
88        return Ok(Vec::new());
89    }
90
91    let messages_to_send: Vec<&Message> = messages
92        .iter()
93        .rev()
94        .take(config.max_messages)
95        .collect::<Vec<_>>()
96        .into_iter()
97        .rev()
98        .collect();
99
100    let user_prompt = build_extraction_prompt(&messages_to_send);
101
102    let llm_messages = [
103        Message::from_legacy(Role::System, EXTRACTION_SYSTEM_PROMPT),
104        Message::from_legacy(Role::User, user_prompt),
105    ];
106
107    let extraction_timeout = Duration::from_secs(config.extraction_timeout_secs);
108    let response = match timeout(extraction_timeout, provider.chat(&llm_messages)).await {
109        Ok(Ok(text)) => text,
110        Ok(Err(e)) => return Err(MemoryError::Llm(e)),
111        Err(_) => {
112            tracing::warn!(
113                "trajectory extraction timed out after {}s",
114                config.extraction_timeout_secs
115            );
116            return Ok(Vec::new());
117        }
118    };
119
120    let entries = parse_extraction_response(&response);
121    Ok(entries)
122}
123
124fn build_extraction_prompt(messages: &[&Message]) -> String {
125    let mut prompt = String::from("Agent turn messages:\n");
126    for (i, msg) in messages.iter().enumerate() {
127        use std::fmt::Write as _;
128        let role = format!("{:?}", msg.role);
129        let _ = writeln!(prompt, "[{}] {}: {}", i + 1, role, msg.content);
130    }
131    prompt.push_str("\nExtract trajectory entries as JSON array.");
132    prompt
133}
134
135fn parse_extraction_response(response: &str) -> Vec<TrajectoryEntry> {
136    let raw: Vec<RawEntry> = if let Ok(v) = serde_json::from_str(response) {
137        v
138    } else if let Some(slice) = extract_json_array_slice(response) {
139        serde_json::from_str(slice).unwrap_or_default()
140    } else {
141        tracing::warn!(
142            "trajectory extraction: failed to parse response (len={}): {:.200}",
143            response.len(),
144            response
145        );
146        return Vec::new();
147    };
148
149    raw.into_iter()
150        .filter(|e| !e.intent.is_empty() && !e.outcome.is_empty())
151        .filter(|e| matches!(e.kind.as_str(), "procedural" | "episodic"))
152        .map(|e| TrajectoryEntry {
153            kind: e.kind,
154            intent: e.intent,
155            outcome: e.outcome,
156            tools_used: e.tools_used,
157            confidence: e.confidence.clamp(0.0, 1.0),
158        })
159        .collect()
160}
161
162#[cfg(test)]
163mod tests {
164    use super::*;
165
166    #[test]
167    fn parse_direct_json_array() {
168        let json = r#"[{"kind":"procedural","intent":"read a file","outcome":"file read ok","tools_used":["read_file"],"confidence":0.9}]"#;
169        let entries = parse_extraction_response(json);
170        assert_eq!(entries.len(), 1);
171        assert_eq!(entries[0].kind, "procedural");
172        assert_eq!(entries[0].intent, "read a file");
173        assert_eq!(entries[0].tools_used, vec!["read_file"]);
174        assert!((entries[0].confidence - 0.9).abs() < 1e-9);
175    }
176
177    #[test]
178    fn parse_json_embedded_in_prose() {
179        let response = "Here you go:\n[{\"kind\":\"episodic\",\"intent\":\"fixed a bug\",\"outcome\":\"patch applied\",\"tools_used\":[],\"confidence\":0.8}]\nDone.";
180        let entries = parse_extraction_response(response);
181        assert_eq!(entries.len(), 1);
182        assert_eq!(entries[0].kind, "episodic");
183    }
184
185    #[test]
186    fn parse_empty_array() {
187        let entries = parse_extraction_response("[]");
188        assert!(entries.is_empty());
189    }
190
191    #[test]
192    fn parse_invalid_json_returns_empty() {
193        let entries = parse_extraction_response("not json");
194        assert!(entries.is_empty());
195    }
196
197    #[test]
198    fn parse_filters_unknown_kind() {
199        let json =
200            r#"[{"kind":"unknown","intent":"x","outcome":"y","tools_used":[],"confidence":0.5}]"#;
201        let entries = parse_extraction_response(json);
202        assert!(entries.is_empty(), "unknown kind must be filtered out");
203    }
204
205    #[test]
206    fn parse_clamps_confidence() {
207        let json = r#"[{"kind":"procedural","intent":"x","outcome":"y","tools_used":[],"confidence":1.5}]"#;
208        let entries = parse_extraction_response(json);
209        assert_eq!(entries.len(), 1);
210        assert!(
211            (entries[0].confidence - 1.0).abs() < 1e-9,
212            "confidence must be clamped to 1.0"
213        );
214    }
215
216    #[test]
217    fn parse_filters_empty_intent_or_outcome() {
218        let json =
219            r#"[{"kind":"procedural","intent":"","outcome":"y","tools_used":[],"confidence":0.8}]"#;
220        let entries = parse_extraction_response(json);
221        assert!(entries.is_empty(), "empty intent must be filtered");
222    }
223}
zeph_memory/semantic/trajectory.rs

zeph_memory/semantic/
trajectory.rs
