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