1extern crate self as llm_toolkit;
12
13#[cfg(feature = "derive")]
18pub use llm_toolkit_macros::ToPrompt;
19
20#[cfg(feature = "derive")]
24pub use llm_toolkit_macros::ToPromptSet;
25
26#[cfg(feature = "derive")]
30pub use llm_toolkit_macros::ToPromptFor;
31
32#[cfg(feature = "derive")]
36pub use llm_toolkit_macros::examples_section;
37
38#[cfg(feature = "derive")]
42pub use llm_toolkit_macros::define_intent;
43
44#[cfg(feature = "agent")]
66pub use llm_toolkit_macros::Agent;
67
68#[cfg(feature = "agent")]
72pub use llm_toolkit_macros::agent;
73
74#[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
130pub fn extract_json(text: &str) -> Result<String, ParseError> {
143    if let Ok(content) = extract_markdown_block_with_lang(text, "json") {
145        return Ok(content);
146    }
147
148    if let Ok(content) = extract_markdown_block(text) {
150        let extractor = FlexibleExtractor::new();
152        if let Ok(json) = extractor.extract(&content) {
153            return Ok(json);
154        }
155    }
156
157    let extractor = FlexibleExtractor::new();
159    extractor.extract(text)
160}
161
162pub fn extract_markdown_block(text: &str) -> Result<String, ParseError> {
172    let extractor = MarkdownCodeBlockExtractor::new();
173    extractor.extract(text)
174}
175
176pub 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        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        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        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        let text4 = "This text has no code blocks at all.";
247        let result4 = extract_markdown_block(text4);
248        assert!(result4.is_err());
249
250        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        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        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        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        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        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        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        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        assert!(json.contains("\"name\": \"test\""));
370        assert!(json.contains("\"values\": [1, 2, 3]"));
371        assert!(json.contains("\"nested\""));
372    }
373}