agent_works/focus/
core.rs

1use std::sync::Arc;
2use std::time::Duration;
3
4use agent_base::{ChatMessage, LlmClient, ResponseFormat};
5use serde::de::DeserializeOwned;
6use serde_json::Value;
7
8// ── FocusInput ───────────────────────────────────────────────────────────────
9
10/// Input for a Focus call — either a simple string or a structured context.
11pub trait FocusInput {
12    /// Format the input into the user prompt text sent to the LLM.
13    fn to_prompt(&self) -> String;
14}
15
16/// Simple case: pass a string directly.
17impl FocusInput for str {
18    fn to_prompt(&self) -> String {
19        self.to_string()
20    }
21}
22
23impl FocusInput for String {
24    fn to_prompt(&self) -> String {
25        self.clone()
26    }
27}
28
29// ── Context ──────────────────────────────────────────────────────────────────
30
31/// Structured context for multi-field input scenarios.
32///
33/// Fields are formatted as `【key】\nvalue` when sent to the LLM,
34/// where the key acts as a label to help the LLM understand the context.
35///
36/// # Usage
37///
38/// ```ignore
39/// let ctx = Context::new()
40///     .add("command", "apt install nginx")
41///     .add("screen", screen_content);
42/// ```
43pub struct Context {
44    entries: Vec<(String, String)>,
45}
46
47impl Context {
48    pub fn new() -> Self {
49        Self {
50            entries: Vec::new(),
51        }
52    }
53
54    /// Add a context field. The key is used as a label when sent to the LLM.
55    pub fn add(mut self, key: &str, value: &str) -> Self {
56        self.entries.push((key.to_string(), value.to_string()));
57        self
58    }
59}
60
61impl Default for Context {
62    fn default() -> Self {
63        Self::new()
64    }
65}
66
67impl FocusInput for Context {
68    fn to_prompt(&self) -> String {
69        self.entries
70            .iter()
71            .map(|(key, value)| format!("【{}】\n{}", key, value))
72            .collect::<Vec<_>>()
73            .join("\n\n")
74    }
75}
76
77// ── FocusOutput ──────────────────────────────────────────────────────────────
78
79/// Output wrapper for a Focus call.
80///
81/// Contains both the structured result and the raw LLM response,
82/// useful for debugging when something goes wrong.
83pub struct FocusOutput<T> {
84    /// Deserialized structured result.
85    pub result: T,
86    /// Raw LLM response text (JSON string), for logging and debugging.
87    pub raw_response: String,
88}
89
90// ── Focus ────────────────────────────────────────────────────────────────────
91
92/// A focused LLM call.
93///
94/// Each instance is bound to a system prompt and dedicated to one specific
95/// judgment question. Use `ask()` to send input and receive a structured
96/// JSON answer.
97///
98/// # Usage
99///
100/// ```ignore
101/// // Simple case: single string input
102/// let classify = Focus::new(client, "You are a task complexity classifier...");
103/// let output = classify.ask::<TaskComplexity>(&user_input, 5s).await?;
104///
105/// // Complex case: multiple context fields
106/// let status_focus = Focus::new(client, "You are a task status judge...");
107/// let ctx = Context::new()
108///     .add("command", command)
109///     .add("screen", screen);
110/// let output = status_focus.ask::<TaskStatus>(&ctx, 5s).await?;
111/// ```
112pub struct Focus {
113    client: Arc<dyn LlmClient>,
114    system_prompt: String,
115}
116
117impl std::fmt::Debug for Focus {
118    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
119        f.debug_struct("Focus").finish_non_exhaustive()
120    }
121}
122
123impl Focus {
124    /// Create a new Focus instance.
125    ///
126    /// - `client`: LLM client (shared; multiple Focus instances can reuse the same client)
127    /// - `system_prompt`: The role and judgment rules for this Focus (bound at creation, never changes)
128    pub fn new(client: Arc<dyn LlmClient>, system_prompt: impl Into<String>) -> Self {
129        Self {
130            client,
131            system_prompt: system_prompt.into(),
132        }
133    }
134
135    /// Make a focused LLM call.
136    ///
137    /// Sends the system prompt (bound at creation) + user input (this call),
138    /// forces JSON output, and deserializes into `T`.
139    ///
140    /// # Arguments
141    /// - `input`: User input — can be `&str` or `Context`
142    /// - `timeout`: Call timeout
143    ///
144    /// # Returns
145    /// `FocusOutput<T>` containing the structured result and raw response.
146    pub async fn ask<T: DeserializeOwned>(
147        &self,
148        input: &impl FocusInput,
149        timeout: Duration,
150    ) -> Result<FocusOutput<T>, FocusError> {
151        let user_prompt = input.to_prompt();
152
153        // Logging: first line of prompt + char count (privacy-friendly)
154        let prompt_first_line = user_prompt.lines().next().unwrap_or("(empty)");
155        let prompt_char_count = user_prompt.chars().count();
156        let sys_first_line = self.system_prompt.lines().next().unwrap_or("(empty)");
157        let target_type = std::any::type_name::<T>();
158
159        tracing::info!(
160            target_type = target_type,
161            system_prompt = %sys_first_line,
162            user_prompt_first_line = %prompt_first_line,
163            user_prompt_chars = prompt_char_count,
164            timeout_secs = timeout.as_secs(),
165            "[Focus] calling LLM"
166        );
167
168        let start = std::time::Instant::now();
169        let messages = vec![
170            ChatMessage::system(self.system_prompt.clone()),
171            ChatMessage::user(user_prompt),
172        ];
173
174        let response = tokio::time::timeout(
175            timeout,
176            self.client
177                .chat(&messages, &[], None, Some(&ResponseFormat::JsonObject)),
178        )
179        .await
180        .map_err(|_| FocusError::Timeout(timeout))?
181        .map_err(|e| FocusError::Llm(e.to_string()))?;
182
183        let elapsed_ms = start.elapsed().as_millis();
184        let raw_response = extract_content(&response).to_string();
185
186        let result: T = serde_json::from_str(&raw_response).map_err(|e| {
187            tracing::warn!(
188                error = %e,
189                raw_response = %raw_response,
190                elapsed_ms = elapsed_ms,
191                "[Focus] failed to parse LLM response as JSON"
192            );
193            FocusError::Parse {
194                error: e.to_string(),
195                raw: raw_response.clone(),
196            }
197        })?;
198
199        tracing::info!(
200            target_type = target_type,
201            raw_response_chars = raw_response.chars().count(),
202            elapsed_ms = elapsed_ms,
203            "[Focus] call succeeded"
204        );
205
206        Ok(FocusOutput {
207            result,
208            raw_response,
209        })
210    }
211}
212
213// ── FocusError ───────────────────────────────────────────────────────────────
214
215/// Error type for Focus calls.
216#[derive(Debug)]
217pub enum FocusError {
218    /// LLM call timed out.
219    Timeout(Duration),
220    /// LLM call failed (network error, API error, etc.).
221    Llm(String),
222    /// LLM response could not be parsed into the expected JSON type.
223    Parse { error: String, raw: String },
224}
225
226impl std::fmt::Display for FocusError {
227    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
228        match self {
229            FocusError::Timeout(d) => write!(f, "Focus timeout after {:?}", d),
230            FocusError::Llm(e) => write!(f, "Focus LLM error: {}", e),
231            FocusError::Parse { error, .. } => write!(f, "Focus parse error: {}", error),
232        }
233    }
234}
235
236impl std::error::Error for FocusError {}
237
238// ── Internal helpers ─────────────────────────────────────────────────────────
239
240/// Extract the `content` field from an LLM response.
241///
242/// Supports OpenAI-compatible format: `choices[0].message.content`.
243/// Falls back to the full response string if extraction fails.
244fn extract_content(response: &Value) -> &str {
245    response
246        .get("choices")
247        .and_then(|c| c.get(0))
248        .and_then(|c| c.get("message"))
249        .and_then(|m| m.get("content"))
250        .and_then(|c| c.as_str())
251        .unwrap_or_else(|| {
252            tracing::warn!(
253                response = %response,
254                "Focus: could not extract choices[0].message.content, using full response"
255            );
256            response.as_str().unwrap_or("{}")
257        })
258}
259
260// ── Tests ────────────────────────────────────────────────────────────────────
261
262#[cfg(test)]
263mod tests {
264    use super::*;
265    use serde::Deserialize;
266
267    // ── Context tests ──
268
269    #[test]
270    fn context_single_field() {
271        let ctx = Context::new().add("command", "df -h");
272        assert_eq!(ctx.to_prompt(), "【command】\ndf -h");
273    }
274
275    #[test]
276    fn context_multiple_fields() {
277        let ctx = Context::new()
278            .add("command", "apt install nginx")
279            .add("elapsed", "30s")
280            .add("screen", "Reading package lists...");
281        let expected = "【command】\napt install nginx\n\n【elapsed】\n30s\n\n【screen】\nReading package lists...";
282        assert_eq!(ctx.to_prompt(), expected);
283    }
284
285    #[test]
286    fn context_empty() {
287        let ctx = Context::new();
288        assert_eq!(ctx.to_prompt(), "");
289    }
290
291    // ── FocusInput tests ──
292
293    #[test]
294    fn str_input() {
295        let input: &str = "hello";
296        assert_eq!(input.to_prompt(), "hello");
297    }
298
299    #[test]
300    fn string_input() {
301        let input = String::from("hello");
302        assert_eq!(input.to_prompt(), "hello");
303    }
304
305    // ── extract_content tests ──
306
307    #[derive(Deserialize, Debug, PartialEq)]
308    struct MockResult {
309        status: String,
310        reason: String,
311    }
312
313    #[test]
314    fn extract_content_openai_format() {
315        let response = serde_json::json!({
316            "choices": [{
317                "message": {
318                    "content": "{\"status\": \"finished\"}"
319                }
320            }]
321        });
322        assert_eq!(extract_content(&response), "{\"status\": \"finished\"}");
323    }
324
325    #[test]
326    fn extract_content_missing_choices() {
327        let response = serde_json::json!({"error": "something"});
328        assert_eq!(extract_content(&response), "{}");
329    }
330
331    #[test]
332    fn extract_content_empty_choices() {
333        let response = serde_json::json!({"choices": []});
334        assert_eq!(extract_content(&response), "{}");
335    }
336
337    #[test]
338    fn focus_output_deserialize() {
339        let raw = r#"{"status":"finished","reason":"done"}"#;
340        let result: MockResult = serde_json::from_str(raw).unwrap();
341        assert_eq!(result.status, "finished");
342        assert_eq!(result.reason, "done");
343    }
344
345    // ── FocusError tests ──
346
347    #[test]
348    fn focus_error_display() {
349        let err = FocusError::Timeout(Duration::from_secs(5));
350        assert_eq!(format!("{}", err), "Focus timeout after 5s");
351
352        let err = FocusError::Llm("network error".to_string());
353        assert_eq!(format!("{}", err), "Focus LLM error: network error");
354
355        let err = FocusError::Parse {
356            error: "unexpected token".to_string(),
357            raw: "not json".to_string(),
358        };
359        assert_eq!(format!("{}", err), "Focus parse error: unexpected token");
360    }
361}
agent_works/focus/core.rs

agent_works/focus/
core.rs
