llm_toolkit/
lib.rs

1//! 'llm-toolkit' - A low-level Rust toolkit for the LLM last mile problem.
2//!
3//! This library provides a set of sharp, reliable, and unopinionated "tools"
4//! for building robust LLM-powered applications in Rust. It focuses on solving
5//! the common and frustrating problems that occur at the boundary between a
6//! strongly-typed Rust application and the unstructured, often unpredictable
7//! string-based responses from LLM APIs.
8
9// Allow the crate to reference itself by name, which is needed for proc macros
10// to work correctly in examples, tests, and bins
11extern crate self as llm_toolkit;
12
13/// A derive macro to implement the `ToPrompt` trait for structs.
14///
15/// This macro is available only when the `derive` feature is enabled.
16/// See the [crate-level documentation](index.html#2-structured-prompts-with-derivetoprompt) for usage examples.
17#[cfg(feature = "derive")]
18pub use llm_toolkit_macros::ToPrompt;
19
20/// A derive macro to implement the `ToPromptSet` trait for structs.
21///
22/// This macro is available only when the `derive` feature is enabled.
23#[cfg(feature = "derive")]
24pub use llm_toolkit_macros::ToPromptSet;
25
26/// A derive macro to implement the `ToPromptFor` trait for structs.
27///
28/// This macro is available only when the `derive` feature is enabled.
29#[cfg(feature = "derive")]
30pub use llm_toolkit_macros::ToPromptFor;
31
32/// A macro for creating examples sections in prompts.
33///
34/// This macro is available only when the `derive` feature is enabled.
35#[cfg(feature = "derive")]
36pub use llm_toolkit_macros::examples_section;
37
38/// A procedural attribute macro for defining intent enums with automatic prompt and extractor generation.
39///
40/// This macro is available only when the `derive` feature is enabled.
41#[cfg(feature = "derive")]
42pub use llm_toolkit_macros::define_intent;
43
44/// A derive macro to implement the `Agent` trait for structs.
45///
46/// This macro is available only when the `agent` feature is enabled.
47/// It automatically generates an Agent implementation that uses ClaudeCodeAgent
48/// internally and deserializes responses into a structured output type.
49///
50/// # Example
51///
52/// ```ignore
53/// use llm_toolkit_macros::Agent;
54/// use serde::{Deserialize, Serialize};
55///
56/// #[derive(Serialize, Deserialize)]
57/// struct MyOutput {
58///     result: String,
59/// }
60///
61/// #[derive(Agent)]
62/// #[agent(expertise = "My expertise", output = "MyOutput")]
63/// struct MyAgent;
64/// ```
65#[cfg(feature = "agent")]
66pub use llm_toolkit_macros::Agent;
67
68/// An attribute macro to define agent structs with automatic trait implementations.
69///
70/// This macro is available only when the `agent` feature is enabled.
71#[cfg(feature = "agent")]
72pub use llm_toolkit_macros::agent;
73
74/// A derive macro to implement the `TypeMarker` trait for structs.
75///
76/// This macro is available only when the `agent` feature is enabled.
77/// It automatically generates a TypeMarker implementation that provides
78/// a type identifier string for type-based orchestrator output retrieval.
79///
80/// # Example
81///
82/// ```ignore
83/// use llm_toolkit::TypeMarker;
84/// use serde::{Deserialize, Serialize};
85///
86/// #[derive(Serialize, Deserialize, TypeMarker)]
87/// struct MyResponse {
88///     #[serde(default = "default_type")]
89///     __type: String,
90///     result: String,
91/// }
92///
93/// fn default_type() -> String {
94///     "MyResponse".to_string()
95/// }
96/// ```
97#[cfg(feature = "agent")]
98pub use llm_toolkit_macros::{TypeMarker, type_marker};
99
100pub mod attachment;
101pub mod extract;
102pub mod intent;
103pub mod multimodal;
104pub mod prompt;
105
106#[cfg(feature = "agent")]
107pub mod agent;
108
109#[cfg(feature = "agent")]
110pub mod orchestrator;
111
112pub use attachment::{Attachment, AttachmentSchema, ToAttachments};
113pub use extract::{FlexibleExtractor, MarkdownCodeBlockExtractor};
114pub use intent::frame::IntentFrame;
115#[allow(deprecated)]
116pub use intent::{IntentError, IntentExtractor, PromptBasedExtractor};
117pub use multimodal::ImageData;
118pub use prompt::{PromptPart, PromptSetError, ToPrompt, ToPromptFor, ToPromptSet};
119
120#[cfg(feature = "agent")]
121pub use agent::{Agent, AgentError};
122
123#[cfg(feature = "agent")]
124pub use orchestrator::{
125    BlueprintWorkflow, Orchestrator, OrchestratorError, StrategyMap, TypeMarker,
126};
127
128use extract::ParseError;
129
130/// Extracts a JSON string from a raw LLM response string.
131///
132/// This function uses a `FlexibleExtractor` with its standard strategies
133/// to find and extract a JSON object from a string that may contain extraneous
134/// text, such as explanations or Markdown code blocks.
135///
136/// For more advanced control over extraction strategies, see the `extract::FlexibleExtractor` struct.
137///
138/// # Returns
139///
140/// A `Result` containing the extracted JSON `String` on success, or a `ParseError`
141/// if no JSON could be extracted.
142pub fn extract_json(text: &str) -> Result<String, ParseError> {
143    // Try markdown code block first (common LLM output format)
144    if let Ok(content) = extract_markdown_block_with_lang(text, "json") {
145        return Ok(content);
146    }
147
148    // Also try generic markdown block (might contain JSON without language hint)
149    if let Ok(content) = extract_markdown_block(text) {
150        // Verify it's actually JSON by trying to extract JSON from it
151        let extractor = FlexibleExtractor::new();
152        if let Ok(json) = extractor.extract(&content) {
153            return Ok(json);
154        }
155    }
156
157    // Fall back to standard extraction strategies
158    let extractor = FlexibleExtractor::new();
159    extractor.extract(text)
160}
161
162/// Extracts content from any Markdown code block in the text.
163///
164/// This function searches for the first code block (delimited by triple backticks)
165/// and returns its content. The code block can have any language specifier or none at all.
166///
167/// # Returns
168///
169/// A `Result` containing the extracted code block content on success, or a `ParseError`
170/// if no code block is found.
171pub fn extract_markdown_block(text: &str) -> Result<String, ParseError> {
172    let extractor = MarkdownCodeBlockExtractor::new();
173    extractor.extract(text)
174}
175
176/// Extracts content from a Markdown code block with a specific language.
177///
178/// This function searches for a code block with the specified language hint
179/// (e.g., ```rust, ```python) and returns its content.
180///
181/// # Arguments
182///
183/// * `text` - The text containing the markdown code block
184/// * `lang` - The language specifier to match (e.g., "rust", "python")
185///
186/// # Returns
187///
188/// A `Result` containing the extracted code block content on success, or a `ParseError`
189/// if no code block with the specified language is found.
190pub fn extract_markdown_block_with_lang(text: &str, lang: &str) -> Result<String, ParseError> {
191    let extractor = MarkdownCodeBlockExtractor::with_language(lang.to_string());
192    extractor.extract(text)
193}
194
195#[cfg(test)]
196mod tests {
197    use super::*;
198
199    #[test]
200    fn test_json_extraction() {
201        let input = "Some text before {\"key\": \"value\"} and after.";
202        assert_eq!(extract_json(input).unwrap(), "{\"key\": \"value\"}");
203    }
204
205    #[test]
206    fn test_standard_extraction_from_tagged_content() {
207        let text = "<answer>{\"type\": \"success\"}</answer>";
208        let result = extract_json(text);
209        assert!(result.is_ok());
210        assert_eq!(result.unwrap(), "{\"type\": \"success\"}");
211    }
212
213    #[test]
214    fn test_markdown_extraction() {
215        // Test simple code block with no language
216        let text1 = "Here is some code:\n```\nlet x = 42;\n```\nAnd some text after.";
217        let result1 = extract_markdown_block(text1);
218        assert!(result1.is_ok());
219        assert_eq!(result1.unwrap(), "let x = 42;");
220
221        // Test code block with specific language (rust)
222        let text2 = "Here's Rust code:\n```rust\nfn main() {
223    println!(\"Hello\");
224}
225```";
226        let result2 = extract_markdown_block_with_lang(text2, "rust");
227        assert!(result2.is_ok());
228        assert_eq!(result2.unwrap(), "fn main() {\n    println!(\"Hello\");\n}");
229
230        // Test extracting rust block when json block is also present
231        let text3 = r#"\nFirst a JSON block:
232```json
233{"key": "value"}
234```
235
236Then a Rust block:
237```rust
238let data = vec![1, 2, 3];
239```
240"#;
241        let result3 = extract_markdown_block_with_lang(text3, "rust");
242        assert!(result3.is_ok());
243        assert_eq!(result3.unwrap(), "let data = vec![1, 2, 3];");
244
245        // Test case where no code block is found
246        let text4 = "This text has no code blocks at all.";
247        let result4 = extract_markdown_block(text4);
248        assert!(result4.is_err());
249
250        // Test with messy surrounding text and newlines
251        let text5 = r#"\nLots of text before...
252
253
254   ```python
255def hello():
256    print("world")
257    return True
258   ```
259
260
261And more text after with various spacing.
262"#;
263        let result5 = extract_markdown_block_with_lang(text5, "python");
264        assert!(result5.is_ok());
265        assert_eq!(
266            result5.unwrap(),
267            "def hello():\n    print(\"world\")\n    return True"
268        );
269    }
270
271    #[test]
272    fn test_extract_json_from_json_markdown_block() {
273        // Test extraction from JSON markdown block (highest priority)
274        let text = r#"Here's the response:
275```json
276{"status": "success", "count": 42}
277```
278That's the data you requested."#;
279        let result = extract_json(text);
280        assert!(result.is_ok());
281        assert_eq!(result.unwrap(), r#"{"status": "success", "count": 42}"#);
282    }
283
284    #[test]
285    fn test_extract_json_from_generic_markdown_block() {
286        // Test extraction from generic markdown block containing JSON
287        let text = r#"The output is:
288```
289{"result": "ok", "value": 123}
290```
291End of output."#;
292        let result = extract_json(text);
293        assert!(result.is_ok());
294        assert_eq!(result.unwrap(), r#"{"result": "ok", "value": 123}"#);
295    }
296
297    #[test]
298    fn test_extract_json_priority_json_block_over_inline() {
299        // When both JSON markdown block and inline JSON exist, JSON block should be preferred
300        let text = r#"Some inline {"inline": "data"} here.
301```json
302{"block": "data"}
303```
304More text."#;
305        let result = extract_json(text);
306        assert!(result.is_ok());
307        assert_eq!(result.unwrap(), r#"{"block": "data"}"#);
308    }
309
310    #[test]
311    fn test_extract_json_priority_json_block_over_generic_block() {
312        // JSON markdown block should be preferred over generic block
313        let text = r#"First a generic block:
314```
315{"generic": "block"}
316```
317
318Then a JSON block:
319```json
320{"json": "block"}
321```"#;
322        let result = extract_json(text);
323        assert!(result.is_ok());
324        assert_eq!(result.unwrap(), r#"{"json": "block"}"#);
325    }
326
327    #[test]
328    fn test_extract_json_fallback_from_non_json_markdown_block() {
329        // When markdown block contains non-JSON, fallback to inline extraction
330        let text = r#"Here's some code:
331```
332This is not JSON at all
333```
334But this is JSON: {"fallback": "value"}"#;
335        let result = extract_json(text);
336        assert!(result.is_ok());
337        assert_eq!(result.unwrap(), r#"{"fallback": "value"}"#);
338    }
339
340    #[test]
341    fn test_extract_json_from_rust_block_fallback() {
342        // When only non-JSON markdown blocks exist, fallback to inline extraction
343        let text = r#"```rust
344let x = 42;
345```
346The result is {"data": "inline"}"#;
347        let result = extract_json(text);
348        assert!(result.is_ok());
349        assert_eq!(result.unwrap(), r#"{"data": "inline"}"#);
350    }
351
352    #[test]
353    fn test_extract_json_multiline_in_markdown_block() {
354        // Test extraction of multiline JSON from markdown block
355        let text = r#"Response:
356```json
357{
358  "name": "test",
359  "values": [1, 2, 3],
360  "nested": {
361    "key": "value"
362  }
363}
364```"#;
365        let result = extract_json(text);
366        assert!(result.is_ok());
367        let json = result.unwrap();
368        // Verify it contains the expected structure
369        assert!(json.contains("\"name\": \"test\""));
370        assert!(json.contains("\"values\": [1, 2, 3]"));
371        assert!(json.contains("\"nested\""));
372    }
373}