atman_runtime/tools/
llm_extract.rs

1use crate::error::RuntimeError;
2use crate::eval::llm_args::parse_llm_args_from_toolargs;
3use crate::eval::llm_dispatch::dispatch_llm;
4use crate::eval::llm_parse::{
5    FieldDef, coerce_field, extract_json_from_text, find_first_json_object,
6    parse_field_definitions, validate_struct_fields,
7};
8use crate::tool::{BoxFut, Tier, Tool, ToolArgs, ToolCtx, ToolResult};
9use crate::value::Value;
10
11pub struct LlmExtractTool;
12
13impl Tool for LlmExtractTool {
14    fn name(&self) -> &str {
15        "llm.extract"
16    }
17
18    fn tier(&self) -> Tier {
19        Tier::Zero
20    }
21
22    fn call<'a>(&'a self, args: ToolArgs, ctx: &'a ToolCtx) -> BoxFut<'a, ToolResult> {
23        Box::pin(async move {
24            let Some(registry) = ctx.registry.as_deref() else {
25                return Err(RuntimeError::ToolFailed(
26                    "llm.extract: no tool registry available".into(),
27                ));
28            };
29
30            let prompt = args
31                .named("prompt")
32                .and_then(|v| match v {
33                    Value::Str(s) => Some(s.clone()),
34                    _ => None,
35                })
36                .ok_or_else(|| RuntimeError::MissingArg("llm.extract: prompt".into()))?;
37
38            if prompt.trim().is_empty() {
39                return Err(RuntimeError::ToolFailed(
40                    "llm.extract: prompt is required".into(),
41                ));
42            }
43
44            // Parse fields
45            let fields_raw = match args.named("fields") {
46                Some(Value::Struct(pairs)) => pairs.clone(),
47                _ => {
48                    return Err(RuntimeError::MissingArg(
49                        "llm.extract: fields (struct of field definitions)".into(),
50                    ));
51                }
52            };
53
54            if fields_raw.is_empty() {
55                return Err(RuntimeError::ToolFailed(
56                    "llm.extract: fields cannot be empty".into(),
57                ));
58            }
59
60            let field_defs = parse_field_definitions(&fields_raw)
61                .map_err(|e| RuntimeError::ToolFailed(format!("llm.extract: {e}")))?;
62
63            // Construct prompt
64            let field_descs: Vec<String> = field_defs
65                .iter()
66                .map(|f| {
67                    if f.ty.is_empty() {
68                        format!("  - \"{}\": {}", f.name, f.description)
69                    } else {
70                        format!("  - \"{}\": {} — {}", f.name, f.ty, f.description)
71                    }
72                })
73                .collect();
74
75            let tool_prompt = format!(
76                "Extract structured information. Respond as a JSON object with exactly these fields:\n\n{}\n\nRespond with ONLY the JSON object, no markdown, no explanation.\n\n{}",
77                field_descs.join("\n"),
78                prompt
79            );
80
81            let parse_retry_count = extract_retry(&args);
82            let mut last_parse_error: Option<String> = None;
83
84            for attempt in 0..=parse_retry_count {
85                let mut llm_args = parse_llm_args_from_toolargs(&args, registry)?;
86                llm_args.retry_count = 0;
87                llm_args.messages_override = None;
88                llm_args.fallback_value = None;
89                llm_args.prompt = Some(if attempt == 0 {
90                    tool_prompt.clone()
91                } else {
92                    format!(
93                        "{tool_prompt}\n\nYour previous response was invalid: {}. Return one complete, closed JSON object and nothing else.",
94                        last_parse_error.as_deref().unwrap_or("parse failed")
95                    )
96                });
97
98                let result = dispatch_llm(llm_args, ctx).await;
99                if let Value::Err(e) = result {
100                    return Err(e);
101                }
102
103                match parse_extract_result(&result, &field_defs) {
104                    Ok(v) => return Ok(v),
105                    Err(error) if attempt < parse_retry_count => {
106                        last_parse_error = Some(error);
107                        continue;
108                    }
109                    Err(e) => {
110                        return Err(RuntimeError::ToolFailed(format!("llm.extract: {e}")));
111                    }
112                }
113            }
114
115            Err(RuntimeError::ToolFailed("llm.extract: unreachable".into()))
116        })
117    }
118}
119
120fn extract_retry(args: &ToolArgs) -> u32 {
121    args.named("retry")
122        .and_then(|v| match v {
123            Value::Int(n) if *n >= 0 => Some(*n as u32),
124            _ => None,
125        })
126        .unwrap_or(0)
127}
128
129fn json_parse_error(text: &str) -> String {
130    let preview: String = text.chars().take(100).collect();
131    if text.contains('{') && find_first_json_object(text).is_none() {
132        format!("incomplete JSON object in response: {preview}")
133    } else {
134        format!("could not parse response as JSON: {preview}")
135    }
136}
137
138fn parse_extract_result(result: &Value, field_defs: &[FieldDef]) -> Result<Value, String> {
139    // If already Struct (assistant_message_to_value pre-parsed JSON)
140    let pairs = match result {
141        Value::Struct(pairs) => pairs.clone(),
142        Value::Str(s) => {
143            let json = extract_json_from_text(s).ok_or_else(|| json_parse_error(s))?;
144            match Value::from_json(json) {
145                Value::Struct(pairs) => pairs,
146                other => return Err(format!("expected JSON object, got {}", other.kind_name())),
147            }
148        }
149        Value::Message(m) => {
150            let text = m.text_concat();
151            if text.is_empty() {
152                return Err("empty response from LLM".into());
153            }
154            let json = extract_json_from_text(&text).ok_or_else(|| json_parse_error(&text))?;
155            match Value::from_json(json) {
156                Value::Struct(pairs) => pairs,
157                other => return Err(format!("expected JSON object, got {}", other.kind_name())),
158            }
159        }
160        _ => return Err("unexpected value type from LLM".into()),
161    };
162
163    // Validate required fields
164    let required: Vec<&str> = field_defs.iter().map(|f| f.name.as_str()).collect();
165    validate_struct_fields(&pairs, &required)?;
166
167    // Coerce field types
168    let mut coerced: Vec<(String, Value)> = Vec::with_capacity(pairs.len());
169    for (name, val) in pairs {
170        let field_def = field_defs.iter().find(|f| f.name == name);
171        match field_def {
172            Some(def) if !def.ty.is_empty() => {
173                let coerced_val = coerce_field(val, &def.ty)?;
174                coerced.push((name, coerced_val));
175            }
176            _ => {
177                coerced.push((name, val));
178            }
179        }
180    }
181
182    Ok(Value::Struct(coerced))
183}
184
185#[cfg(test)]
186mod tests {
187    use super::*;
188
189    #[test]
190    fn json_parse_error_preview_is_utf8_safe() {
191        let text = format!("{}中", "a".repeat(99));
192        let error = json_parse_error(&text);
193        assert!(error.ends_with('中'));
194    }
195
196    #[test]
197    fn truncated_multibyte_json_reports_incomplete_object() {
198        let text = r#"{"rule_names":["<project> — 项目红线 (AGENTS.md)","Global Rules"],"confession_triggers":["extract 工具存在 utf8 截断问题""#;
199        let error = json_parse_error(text);
200        assert!(error.starts_with("incomplete JSON object in response:"));
201    }
202
203    #[test]
204    fn complete_multibyte_json_still_parses() {
205        let json = extract_json_from_text(r#"{"name":"项目红线"}"#).unwrap();
206        assert_eq!(json["name"], "项目红线");
207    }
208}
atman_runtime/tools/llm_extract.rs

atman_runtime/tools/
llm_extract.rs
