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/// A derive macro to implement the `ToPrompt` trait for structs.
10///
11/// This macro is available only when the `derive` feature is enabled.
12/// See the [crate-level documentation](index.html#2-structured-prompts-with-derivetoprompt) for usage examples.
13#[cfg(feature = "derive")]
14pub use llm_toolkit_macros::ToPrompt;
15
16pub mod extract;
17pub mod intent;
18pub mod multimodal;
19pub mod prompt;
20
21pub use extract::{FlexibleExtractor, MarkdownCodeBlockExtractor};
22pub use intent::{IntentError, IntentExtractor, PromptBasedExtractor};
23pub use multimodal::ImageData;
24pub use prompt::{PromptPart, ToPrompt};
25
26use extract::ParseError;
27
28/// Extracts a JSON string from a raw LLM response string.
29///
30/// This function uses a `FlexibleExtractor` with its standard strategies
31/// to find and extract a JSON object from a string that may contain extraneous
32/// text, such as explanations or Markdown code blocks.
33///
34/// For more advanced control over extraction strategies, see the `extract::FlexibleExtractor` struct.
35///
36/// # Returns
37///
38/// A `Result` containing the extracted JSON `String` on success, or a `ParseError`
39/// if no JSON could be extracted.
40pub fn extract_json(text: &str) -> Result<String, ParseError> {
41    let extractor = FlexibleExtractor::new();
42    // Note: The standard strategies in the copied code are TaggedContent("answer"), JsonBrackets, FirstJsonObject.
43    // We will add a markdown strategy later during refactoring.
44    extractor.extract(text)
45}
46
47/// Extracts content from any Markdown code block in the text.
48///
49/// This function searches for the first code block (delimited by triple backticks)
50/// and returns its content. The code block can have any language specifier or none at all.
51///
52/// # Returns
53///
54/// A `Result` containing the extracted code block content on success, or a `ParseError`
55/// if no code block is found.
56pub fn extract_markdown_block(text: &str) -> Result<String, ParseError> {
57    let extractor = MarkdownCodeBlockExtractor::new();
58    extractor.extract(text)
59}
60
61/// Extracts content from a Markdown code block with a specific language.
62///
63/// This function searches for a code block with the specified language hint
64/// (e.g., ```rust, ```python) and returns its content.
65///
66/// # Arguments
67///
68/// * `text` - The text containing the markdown code block
69/// * `lang` - The language specifier to match (e.g., "rust", "python")
70///
71/// # Returns
72///
73/// A `Result` containing the extracted code block content on success, or a `ParseError`
74/// if no code block with the specified language is found.
75pub fn extract_markdown_block_with_lang(text: &str, lang: &str) -> Result<String, ParseError> {
76    let extractor = MarkdownCodeBlockExtractor::with_language(lang.to_string());
77    extractor.extract(text)
78}
79
80#[cfg(test)]
81mod tests {
82    use super::*;
83
84    #[test]
85    fn test_json_extraction() {
86        let input = "Some text before {\"key\": \"value\"} and after.";
87        assert_eq!(extract_json(input).unwrap(), "{\"key\": \"value\"}");
88    }
89
90    #[test]
91    fn test_standard_extraction_from_tagged_content() {
92        let text = "<answer>{\"type\": \"success\"}</answer>";
93        let result = extract_json(text);
94        assert!(result.is_ok());
95        assert_eq!(result.unwrap(), "{\"type\": \"success\"}");
96    }
97
98    #[test]
99    fn test_markdown_extraction() {
100        // Test simple code block with no language
101        let text1 = "Here is some code:\n```\nlet x = 42;\n```\nAnd some text after.";
102        let result1 = extract_markdown_block(text1);
103        assert!(result1.is_ok());
104        assert_eq!(result1.unwrap(), "let x = 42;");
105
106        // Test code block with specific language (rust)
107        let text2 = "Here's Rust code:\n```rust\nfn main() {
108    println!(\"Hello\");
109}
110```";
111        let result2 = extract_markdown_block_with_lang(text2, "rust");
112        assert!(result2.is_ok());
113        assert_eq!(result2.unwrap(), "fn main() {\n    println!(\"Hello\");\n}");
114
115        // Test extracting rust block when json block is also present
116        let text3 = r#"\nFirst a JSON block:
117```json
118{"key": "value"}
119```
120
121Then a Rust block:
122```rust
123let data = vec![1, 2, 3];
124```
125"#;
126        let result3 = extract_markdown_block_with_lang(text3, "rust");
127        assert!(result3.is_ok());
128        assert_eq!(result3.unwrap(), "let data = vec![1, 2, 3];");
129
130        // Test case where no code block is found
131        let text4 = "This text has no code blocks at all.";
132        let result4 = extract_markdown_block(text4);
133        assert!(result4.is_err());
134
135        // Test with messy surrounding text and newlines
136        let text5 = r#"\nLots of text before...
137
138
139   ```python
140def hello():
141    print("world")
142    return True
143   ```   
144
145
146And more text after with various spacing.
147"#;
148        let result5 = extract_markdown_block_with_lang(text5, "python");
149        assert!(result5.is_ok());
150        assert_eq!(
151            result5.unwrap(),
152            "def hello():\n    print(\"world\")\n    return True"
153        );
154    }
155}