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, parse_field_definitions, validate_struct_fields,
6};
7use crate::tool::{BoxFut, Tier, Tool, ToolArgs, ToolCtx, ToolResult};
8use crate::value::Value;
9
10pub struct LlmExtractTool;
11
12impl Tool for LlmExtractTool {
13    fn name(&self) -> &str {
14        "llm.extract"
15    }
16
17    fn tier(&self) -> Tier {
18        Tier::Zero
19    }
20
21    fn call<'a>(&'a self, args: ToolArgs, ctx: &'a ToolCtx) -> BoxFut<'a, ToolResult> {
22        Box::pin(async move {
23            let Some(registry) = ctx.registry.as_deref() else {
24                return Err(RuntimeError::ToolFailed(
25                    "llm.extract: no tool registry available".into(),
26                ));
27            };
28
29            let prompt = args
30                .named("prompt")
31                .and_then(|v| match v {
32                    Value::Str(s) => Some(s.clone()),
33                    _ => None,
34                })
35                .ok_or_else(|| RuntimeError::MissingArg("llm.extract: prompt".into()))?;
36
37            if prompt.trim().is_empty() {
38                return Err(RuntimeError::ToolFailed(
39                    "llm.extract: prompt is required".into(),
40                ));
41            }
42
43            // Parse fields
44            let fields_raw = match args.named("fields") {
45                Some(Value::Struct(pairs)) => pairs.clone(),
46                _ => {
47                    return Err(RuntimeError::MissingArg(
48                        "llm.extract: fields (struct of field definitions)".into(),
49                    ));
50                }
51            };
52
53            if fields_raw.is_empty() {
54                return Err(RuntimeError::ToolFailed(
55                    "llm.extract: fields cannot be empty".into(),
56                ));
57            }
58
59            let field_defs = parse_field_definitions(&fields_raw)
60                .map_err(|e| RuntimeError::ToolFailed(format!("llm.extract: {e}")))?;
61
62            // Construct prompt
63            let field_descs: Vec<String> = field_defs
64                .iter()
65                .map(|f| {
66                    if f.ty.is_empty() {
67                        format!("  - \"{}\": {}", f.name, f.description)
68                    } else {
69                        format!("  - \"{}\": {} — {}", f.name, f.ty, f.description)
70                    }
71                })
72                .collect();
73
74            let tool_prompt = format!(
75                "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{}",
76                field_descs.join("\n"),
77                prompt
78            );
79
80            let parse_retry_count = extract_retry(&args);
81
82            for attempt in 0..=parse_retry_count {
83                let mut llm_args = parse_llm_args_from_toolargs(&args, registry)?;
84                llm_args.retry_count = 0;
85                llm_args.messages_override = None;
86                llm_args.fallback_value = None;
87                llm_args.prompt = Some(if attempt == 0 {
88                    tool_prompt.clone()
89                } else {
90                    format!(
91                        "{tool_prompt}\n\nYour previous response could not be parsed. Please respond with ONLY a JSON object."
92                    )
93                });
94
95                let result = dispatch_llm(llm_args, ctx).await;
96                if let Value::Err(e) = result {
97                    return Err(e);
98                }
99
100                match parse_extract_result(&result, &field_defs) {
101                    Ok(v) => return Ok(v),
102                    Err(_) if attempt < parse_retry_count => continue,
103                    Err(e) => {
104                        return Err(RuntimeError::ToolFailed(format!("llm.extract: {e}")));
105                    }
106                }
107            }
108
109            Err(RuntimeError::ToolFailed("llm.extract: unreachable".into()))
110        })
111    }
112}
113
114fn extract_retry(args: &ToolArgs) -> u32 {
115    args.named("retry")
116        .and_then(|v| match v {
117            Value::Int(n) if *n >= 0 => Some(*n as u32),
118            _ => None,
119        })
120        .unwrap_or(0)
121}
122
123fn parse_extract_result(result: &Value, field_defs: &[FieldDef]) -> Result<Value, String> {
124    // If already Struct (assistant_message_to_value pre-parsed JSON)
125    let pairs = match result {
126        Value::Struct(pairs) => pairs.clone(),
127        Value::Str(s) => {
128            let json = extract_json_from_text(s).ok_or_else(|| {
129                format!(
130                    "could not parse response as JSON: {}",
131                    &s[..s.len().min(100)]
132                )
133            })?;
134            match Value::from_json(json) {
135                Value::Struct(pairs) => pairs,
136                other => return Err(format!("expected JSON object, got {}", other.kind_name())),
137            }
138        }
139        Value::Message(m) => {
140            let text = m.text_concat();
141            if text.is_empty() {
142                return Err("empty response from LLM".into());
143            }
144            let json = extract_json_from_text(&text).ok_or_else(|| {
145                format!(
146                    "could not parse response as JSON: {}",
147                    &text[..text.len().min(100)]
148                )
149            })?;
150            match Value::from_json(json) {
151                Value::Struct(pairs) => pairs,
152                other => return Err(format!("expected JSON object, got {}", other.kind_name())),
153            }
154        }
155        _ => return Err("unexpected value type from LLM".into()),
156    };
157
158    // Validate required fields
159    let required: Vec<&str> = field_defs.iter().map(|f| f.name.as_str()).collect();
160    validate_struct_fields(&pairs, &required)?;
161
162    // Coerce field types
163    let mut coerced: Vec<(String, Value)> = Vec::with_capacity(pairs.len());
164    for (name, val) in pairs {
165        let field_def = field_defs.iter().find(|f| f.name == name);
166        match field_def {
167            Some(def) if !def.ty.is_empty() => {
168                let coerced_val = coerce_field(val, &def.ty)?;
169                coerced.push((name, coerced_val));
170            }
171            _ => {
172                coerced.push((name, val));
173            }
174        }
175    }
176
177    Ok(Value::Struct(coerced))
178}
atman_runtime/tools/llm_extract.rs

atman_runtime/tools/
llm_extract.rs
