aither_core/llm/
mod.rs

1//! # Language Models and Conversation Management
2//!
3//! This module provides everything you need to work with language models in a provider-agnostic way.
4//! Build chat applications, generate structured output, and integrate tools without being tied to any specific AI service.
5//!
6//! ## Core Components
7//!
8//! - **[`LanguageModel`]** - The main trait for text generation and conversation
9//! - **[`LLMRequest`]** - Encapsulates messages, tools, and parameters for model calls
10//! - **[`Event`]** - Stream events from the model (text, reasoning, tool calls)
11//! - **[`Message`]** - Represents individual messages in a conversation
12//! - **[`Tool`]** - Function calling interface for extending model capabilities
13//!
14//! ## Design Philosophy
15//!
16//! The core crate provides a **low-level API** that emits events without executing tools.
17//! Tool execution is the responsibility of higher-level abstractions like `aither-agent`.
18//!
19//! This design allows:
20//! - Full control over tool execution flow
21//! - Hooks for intercepting and modifying tool calls
22//! - Proper context management between turns
23//! - Clean separation between LLM communication and agent logic
24//!
25//! ## Quick Start
26//!
27//! ### Basic Conversation
28//!
29//! ```rust,ignore
30//! use aither::llm::{LanguageModel, Event, oneshot};
31//! use futures_lite::StreamExt;
32//!
33//! async fn chat_with_model(model: impl LanguageModel) -> Result<String, Box<dyn std::error::Error>> {
34//!     let request = oneshot("You are a helpful assistant", "What's the capital of Japan?");
35//!     let mut stream = model.respond(request);
36//!     let mut full_text = String::new();
37//!
38//!     while let Some(event) = stream.next().await {
39//!         match event? {
40//!             Event::Text(chunk) => full_text.push_str(&chunk),
41//!             Event::Reasoning(thought) => eprintln!("[thinking] {}", thought),
42//!             Event::ToolCall(call) => {
43//!                 // Handle tool call (typically done by agent crate)
44//!                 println!("Tool requested: {}", call.name);
45//!             }
46//!             _ => {}
47//!         }
48//!     }
49//!
50//!     Ok(full_text)
51//! }
52//! ```
53//!
54//! ### With Tools (Agent-Controlled)
55//!
56//! ```rust,ignore
57//! use aither::llm::{LanguageModel, Event, LLMRequest, Message};
58//!
59//! // The core crate does NOT execute tools - it emits ToolCall events.
60//! // Tool execution should be handled by the agent crate.
61//! let request = LLMRequest::new([Message::user("What's the weather?")])
62//!     .with_tool_definitions(vec![weather_tool_definition()]);
63//!
64//! let mut stream = model.respond(request);
65//! while let Some(event) = stream.next().await {
66//!     match event? {
67//!         Event::ToolCall(call) => {
68//!             // Execute tool and continue conversation
69//!             let result = my_tool_executor.execute(&call).await;
70//!             // Add result to messages and send another request...
71//!         }
72//!         _ => {}
73//!     }
74//! }
75//! ```
76
77/// Assistant module for managing assistant-related functionality.
78pub mod assistant;
79/// Event types for streaming responses.
80pub mod event;
81/// Classification of provider failures.
82pub mod failure;
83/// Message types and conversation handling.
84pub mod message;
85/// Model profiles and capabilities.
86pub mod model;
87/// Provider module for managing language model providers and their configurations.
88pub mod provider;
89/// Provider-opaque reasoning state carried across turns.
90pub mod reasoning;
91/// Deep research workflows and agent capabilities.
92pub mod researcher;
93/// How a failed model request is retried.
94pub mod retry;
95/// Tool system for function calling.
96pub mod tool;
97
98use crate::llm::{model::Parameters, tool::Tools};
99use alloc::{
100    boxed::Box,
101    string::{String, ToString},
102    sync::Arc,
103    vec,
104    vec::Vec,
105};
106use core::{any::TypeId, future::Future};
107pub use event::{Event, ToolCall, Usage};
108pub use failure::{Classify, Failure, FailureKind};
109use futures_core::Stream;
110use futures_lite::{StreamExt, pin};
111pub use message::{Attachment, Message, Role};
112pub use provider::LanguageModelProvider;
113pub use reasoning::ReasoningState;
114pub use researcher::{
115    ResearchCitation, ResearchEvent, ResearchFinding, ResearchOptions, ResearchReport,
116    ResearchRequest, ResearchSource, ResearchStage, Researcher, ResearcherProfile,
117};
118pub use retry::RetryPolicy;
119use schemars::{JsonSchema, schema_for};
120use serde::de::DeserializeOwned;
121pub use tool::{
122    AgentIdentity, CallIdentity, CommandOutcome, CommandResult, CommandStream, Interruption,
123    InterruptionCause, IntoToolResult, StoredFile, SubagentIdentity, Tool, ToolCallId, ToolContext,
124    ToolResult, render_interruptions,
125};
126
127use crate::llm::model::Profile;
128
129/// Why a structured-output call failed.
130///
131/// [`LanguageModel::respond`] reports provider failures through the model's own
132/// [`LanguageModel::Error`]. Structured output adds a second, separate way to
133/// fail — the model answered, but not with the requested shape — so this keeps
134/// the two distinguishable instead of flattening both into one opaque error.
135/// Callers need that distinction: a rate limit is worth retrying, a schema
136/// mismatch is worth re-prompting, and an auth failure is worth neither.
137#[derive(Debug)]
138pub enum GenerateError<E> {
139    /// The provider itself failed: transport, rate limit, auth, and so on.
140    Provider(E),
141
142    /// The provider answered, but the response did not match the schema.
143    ///
144    /// Carries the text that could not be parsed so callers can log or retry.
145    Parse {
146        /// The underlying deserialization failure.
147        source: serde_json::Error,
148        /// The response text that failed to parse, truncated for logging.
149        response: String,
150    },
151}
152
153impl<E: core::fmt::Display> core::fmt::Display for GenerateError<E> {
154    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
155        match self {
156            Self::Provider(err) => write!(f, "language model request failed: {err}"),
157            Self::Parse { source, response } => {
158                write!(
159                    f,
160                    "structured output did not match the requested schema: {source}; response: {response}"
161                )
162            }
163        }
164    }
165}
166
167impl<E: core::error::Error + 'static> core::error::Error for GenerateError<E> {
168    fn source(&self) -> Option<&(dyn core::error::Error + 'static)> {
169        match self {
170            Self::Provider(err) => Some(err),
171            Self::Parse { source, .. } => Some(source),
172        }
173    }
174}
175
176/// Builder-style request passed into [`LanguageModel::respond`].
177///
178/// Wraps the full conversation, model parameters, and tool definitions a provider
179/// needs in order to execute a call.
180#[derive(Debug, Clone)]
181pub struct LLMRequest {
182    messages: Vec<Message>,
183    parameters: Parameters,
184    tool_definitions: Vec<tool::ToolDefinition>,
185}
186
187impl LLMRequest {
188    /// Creates a request from the provided messages using default parameters.
189    pub fn new(messages: impl Into<Vec<Message>>) -> Self {
190        Self {
191            messages: messages.into(),
192            parameters: Parameters::default(),
193            tool_definitions: Vec::new(),
194        }
195    }
196
197    /// Adds tool definitions to the request.
198    ///
199    /// These define what tools the model can request. The model will emit
200    /// [`Event::ToolCall`] events when it wants to use a tool.
201    #[must_use]
202    pub fn with_tool_definitions(mut self, definitions: Vec<tool::ToolDefinition>) -> Self {
203        self.tool_definitions = definitions;
204        self
205    }
206
207    /// Adds a single tool definition.
208    #[must_use]
209    pub fn with_tool<T: Tool>(mut self, tool: &T) -> Self {
210        self.tool_definitions.push(tool::ToolDefinition::new(tool));
211        self
212    }
213
214    /// Overrides the sampling parameters used for this call.
215    #[must_use]
216    pub fn with_parameters(mut self, parameters: Parameters) -> Self {
217        self.parameters = parameters;
218        self
219    }
220
221    /// Returns the current conversation messages.
222    #[must_use]
223    pub fn messages(&self) -> &[Message] {
224        &self.messages
225    }
226
227    /// Returns a mutable reference to messages for modification.
228    pub const fn messages_mut(&mut self) -> &mut Vec<Message> {
229        &mut self.messages
230    }
231
232    /// Returns the current parameter snapshot.
233    #[must_use]
234    pub const fn parameters(&self) -> &Parameters {
235        &self.parameters
236    }
237
238    /// Returns the tool definitions.
239    #[must_use]
240    pub fn tool_definitions(&self) -> &[tool::ToolDefinition] {
241        &self.tool_definitions
242    }
243
244    /// Breaks the request into owned components.
245    #[must_use]
246    pub fn into_parts(self) -> (Vec<Message>, Parameters, Vec<tool::ToolDefinition>) {
247        (self.messages, self.parameters, self.tool_definitions)
248    }
249}
250
251/// Legacy request builder that supports mutable tool registry.
252///
253/// This is provided for backwards compatibility with providers that
254/// handle tool execution internally (built-in tools).
255#[derive(Debug)]
256pub struct LLMRequestWithTools<'tools> {
257    inner: LLMRequest,
258    tools: &'tools mut Tools,
259}
260
261impl LLMRequest {
262    /// Attaches a mutable tool registry to the request.
263    ///
264    /// This is for providers that execute tools internally (e.g., built-in tools).
265    /// For agent-controlled tool execution, use `with_tool_definitions` instead.
266    pub fn with_tools(self, tools: &mut Tools) -> LLMRequestWithTools<'_> {
267        let definitions = tools.definitions();
268        LLMRequestWithTools {
269            inner: self.with_tool_definitions(definitions),
270            tools,
271        }
272    }
273}
274
275impl<'tools> LLMRequestWithTools<'tools> {
276    /// Returns the inner request.
277    #[must_use]
278    pub const fn request(&self) -> &LLMRequest {
279        &self.inner
280    }
281
282    /// Returns the inner request for modification, keeping its tool registry.
283    pub const fn request_mut(&mut self) -> &mut LLMRequest {
284        &mut self.inner
285    }
286
287    /// Returns the tool registry.
288    #[must_use]
289    pub const fn tools(&mut self) -> &mut Tools {
290        self.tools
291    }
292
293    /// Breaks into components.
294    #[must_use]
295    pub fn into_parts(self) -> (LLMRequest, &'tools mut Tools) {
296        (self.inner, self.tools)
297    }
298
299    /// Invokes a registered tool by name, running it with `cx` as its
300    /// [`ToolContext`].
301    ///
302    /// # Errors
303    /// Returns an error if tool is not found or the tool call fails.
304    pub async fn call_tool(
305        &mut self,
306        name: &str,
307        args_json: &str,
308        cx: ToolContext,
309    ) -> crate::Result<ToolResult> {
310        self.tools.call(name, args_json, cx).await
311    }
312}
313
314/// How a model can count the tokens a request costs, when it can.
315///
316/// [`LanguageModel::token_counting`] reports which kind of counting a model
317/// offers: `Local` counts in-process against a published vocabulary and is
318/// cheap enough to run on every assembly; `Endpoint` counts exactly over a
319/// network call and is worth its round trip only near a threshold.
320#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
321pub enum TokenCounting {
322    /// No counting: the caller falls back to its estimate.
323    #[default]
324    None,
325    /// In-process counting against an embedded vocabulary.
326    Local,
327    /// Counting through a provider endpoint.
328    Endpoint,
329}
330
331/// Language models for text generation and conversation.
332///
333/// The `respond` method returns a stream of [`Event`]s. Tool calls are emitted
334/// as events and are NOT automatically executed - this allows higher-level
335/// abstractions (like agents) to control tool execution.
336///
337/// See the [module documentation](crate::llm) for examples and usage patterns.
338pub trait LanguageModel: Sized + Send + Sync {
339    /// The error type returned by this language model.
340    type Error: core::error::Error + Classify + Send + Sync + 'static;
341
342    /// How this model counts prompt tokens, if at all. `None` by default.
343    fn token_counting(&self) -> TokenCounting {
344        TokenCounting::None
345    }
346
347    /// The prompt tokens `request` costs — messages, tool definitions, and
348    /// attachments where the counter supports them — or `None` when this
349    /// model cannot count. The default implementation counts nothing.
350    fn count_tokens(
351        &self,
352        request: &LLMRequest,
353    ) -> impl Future<Output = Result<Option<u32>, GenerateError<Self::Error>>> + Send {
354        let _ = request;
355        async { Ok(None) }
356    }
357
358    /// The tokens `text` takes under a local vocabulary, used while
359    /// streaming to measure decode speed. `None` unless
360    /// [`token_counting`](Self::token_counting) is `Local`.
361    fn count_text(&self, text: &str) -> Option<u32> {
362        let _ = text;
363        None
364    }
365
366    /// Generates a streaming response to a conversation.
367    ///
368    /// Returns a stream of [`Event`]s including:
369    /// - `Event::Text` - Visible text chunks
370    /// - `Event::Reasoning` - Internal reasoning (for reasoning models)
371    /// - `Event::ToolCall` - Requests to execute tools (NOT auto-executed)
372    /// - `Event::BuiltInToolResult` - Results from provider's built-in tools
373    fn respond(&self, request: LLMRequest)
374    -> impl Stream<Item = Result<Event, Self::Error>> + Send;
375
376    /// Generates a streaming response with a mutable tool registry.
377    ///
378    /// This is for providers that support built-in tool execution.
379    /// The default implementation ignores the tools and delegates to `respond`.
380    fn respond_with_tools(
381        &self,
382        request: LLMRequestWithTools<'_>,
383    ) -> impl Stream<Item = Result<Event, Self::Error>> + Send {
384        let (inner, _tools) = request.into_parts();
385        self.respond(inner)
386    }
387
388    /// Generates structured output conforming to JSON schema.
389    ///
390    /// # Note for Implementors
391    /// By default, we use a system prompt to instruct the model to generate structured output based on the provided JSON schema.
392    /// However, that is not efficient enough, if supported, provider should override this method to provide native structured generation support.
393    /// Native structured generation can apply decode rules in token-level, ensuring the output is always valid JSON, and reducing parsing errors.
394    fn generate<T: JsonSchema + DeserializeOwned + 'static>(
395        &self,
396        request: LLMRequest,
397    ) -> impl Future<Output = Result<T, GenerateError<Self::Error>>> + Send {
398        async { structured_generate(self, request).await }
399    }
400
401    /// Completes given text prefix.
402    fn complete(&self, prefix: &str) -> impl Stream<Item = Result<Event, Self::Error>> + Send {
403        self.respond(oneshot("Please complete the following text:", prefix))
404    }
405
406    /// Summarizes text.
407    ///
408    /// # Note for Implementors
409    /// By default, we provide a generic summarization prompt. However, some model provider, like Apple Intelligence, may have native summarization support.
410    /// It would load a summarization-specific lora at runtime, providing better quality.
411    fn summarize(&self, text: &str) -> impl Stream<Item = Result<Event, Self::Error>> + Send {
412        summarize(self, text)
413    }
414
415    /// Categorizes text.
416    ///
417    /// # Note for Implementors
418    /// By default, we use a system prompt to instruct the model to categorize text based on the provided JSON schema.
419    /// However, that is not efficient enough, if supported, provider should override this method to provide native categorization support.
420    fn categorize<T: JsonSchema + DeserializeOwned + 'static>(
421        &self,
422        text: &str,
423    ) -> impl Future<Output = Result<T, GenerateError<Self::Error>>> + Send {
424        async { categorize_text(self, text).await }
425    }
426
427    /// Returns model profile and capabilities.
428    ///
429    /// See [`Profile`] for details on model metadata.
430    ///
431    /// # Errors
432    ///
433    /// Returns an error when the profile cannot be determined — for example
434    /// when neither the provider nor the model registry knows the model's
435    /// context window. Implementations must not substitute a guessed value.
436    fn profile(&self) -> impl Future<Output = Result<Profile, Self::Error>> + Send;
437
438    /// How a request to this model that failed with a retryable failure is
439    /// sent again.
440    ///
441    /// A provider that retries establishing its response stream returns the
442    /// policy it retries under; callers reuse it for a stream that breaks
443    /// after it opened, which the provider cannot retry because the events
444    /// before the break already reached them.
445    fn retry_policy(&self) -> RetryPolicy {
446        RetryPolicy::default()
447    }
448}
449
450macro_rules! impl_language_model {
451    ($($name:ident),*) => {
452        $(
453            impl<T: LanguageModel> LanguageModel for $name<T> {
454                type Error = T::Error;
455
456                fn respond(
457                    &self,
458                    request: LLMRequest,
459                ) -> impl Stream<Item = Result<Event, Self::Error>> + Send {
460                    T::respond(self, request)
461                }
462
463                fn respond_with_tools(
464                    &self,
465                    request: LLMRequestWithTools<'_>,
466                ) -> impl Stream<Item = Result<Event, Self::Error>> + Send {
467                    T::respond_with_tools(self, request)
468                }
469
470                fn generate<U: JsonSchema + DeserializeOwned + 'static>(
471                    &self,
472                    request: LLMRequest,
473                ) -> impl Future<Output = Result<U, GenerateError<Self::Error>>> + Send {
474                    T::generate(self, request)
475                }
476
477                fn complete(
478                    &self,
479                    prefix: &str,
480                ) -> impl Stream<Item = Result<Event, Self::Error>> + Send {
481                    T::complete(self, prefix)
482                }
483
484                fn summarize(
485                    &self,
486                    text: &str,
487                ) -> impl Stream<Item = Result<Event, Self::Error>> + Send {
488                    T::summarize(self, text)
489                }
490
491                fn categorize<U: JsonSchema + DeserializeOwned + 'static>(
492                    &self,
493                    text: &str,
494                ) -> impl Future<Output = Result<U, GenerateError<Self::Error>>> + Send {
495                    T::categorize(self, text)
496                }
497
498                fn profile(&self) -> impl Future<Output = Result<Profile, Self::Error>> + Send {
499                    T::profile(self)
500                }
501
502                fn retry_policy(&self) -> RetryPolicy {
503                    T::retry_policy(self)
504                }
505
506                fn token_counting(&self) -> TokenCounting {
507                    T::token_counting(self)
508                }
509
510                fn count_tokens(
511                    &self,
512                    request: &LLMRequest,
513                ) -> impl Future<Output = Result<Option<u32>, GenerateError<Self::Error>>>
514                + Send {
515                    T::count_tokens(self, request)
516                }
517
518                fn count_text(&self, text: &str) -> Option<u32> {
519                    T::count_text(self, text)
520                }
521            }
522        )*
523    };
524}
525
526impl<T: LanguageModel> LanguageModel for &T {
527    type Error = T::Error;
528
529    fn respond(
530        &self,
531        request: LLMRequest,
532    ) -> impl Stream<Item = Result<Event, Self::Error>> + Send {
533        T::respond(self, request)
534    }
535
536    fn respond_with_tools(
537        &self,
538        request: LLMRequestWithTools<'_>,
539    ) -> impl Stream<Item = Result<Event, Self::Error>> + Send {
540        T::respond_with_tools(self, request)
541    }
542
543    fn generate<U: JsonSchema + DeserializeOwned + 'static>(
544        &self,
545        request: LLMRequest,
546    ) -> impl Future<Output = Result<U, GenerateError<Self::Error>>> + Send {
547        T::generate(self, request)
548    }
549
550    fn complete(&self, prefix: &str) -> impl Stream<Item = Result<Event, Self::Error>> + Send {
551        T::complete(self, prefix)
552    }
553
554    fn summarize(&self, text: &str) -> impl Stream<Item = Result<Event, Self::Error>> + Send {
555        T::summarize(self, text)
556    }
557
558    fn categorize<U: JsonSchema + DeserializeOwned + 'static>(
559        &self,
560        text: &str,
561    ) -> impl Future<Output = Result<U, GenerateError<Self::Error>>> + Send {
562        T::categorize(self, text)
563    }
564
565    fn profile(&self) -> impl Future<Output = Result<Profile, Self::Error>> + Send {
566        T::profile(self)
567    }
568
569    fn retry_policy(&self) -> RetryPolicy {
570        T::retry_policy(self)
571    }
572
573    fn token_counting(&self) -> TokenCounting {
574        T::token_counting(self)
575    }
576
577    fn count_tokens(
578        &self,
579        request: &LLMRequest,
580    ) -> impl Future<Output = Result<Option<u32>, GenerateError<Self::Error>>> + Send {
581        T::count_tokens(self, request)
582    }
583
584    fn count_text(&self, text: &str) -> Option<u32> {
585        T::count_text(self, text)
586    }
587}
588
589mod prompts;
590
591impl_language_model!(Arc, Box);
592
593/// Collects text from an event stream.
594///
595/// # Errors
596///
597/// Returns the first stream error encountered while collecting text chunks.
598pub async fn collect_text<S, E>(stream: S) -> Result<String, E>
599where
600    S: Stream<Item = Result<Event, E>>,
601{
602    pin!(stream);
603    let mut result = String::new();
604    while let Some(event) = stream.next().await {
605        if let Event::Text(text) = event? {
606            result.push_str(&text);
607        }
608    }
609    Ok(result)
610}
611
612async fn structured_generate<T: JsonSchema + DeserializeOwned + 'static, M: LanguageModel>(
613    model: &M,
614    mut request: LLMRequest,
615) -> Result<T, GenerateError<M::Error>> {
616    let schema = schema_for!(T);
617
618    // If it is a string, we are not required to set up structured generation and JSON schema.
619    let json = if schema.as_value().is_string() {
620        let stream = model.respond(request);
621        let response = collect_text(stream)
622            .await
623            .map_err(GenerateError::Provider)?;
624        // Let's encode it as JSON string. Serializing a String cannot fail.
625        serde_json::to_string(&response).map_err(|source| GenerateError::Parse {
626            source,
627            response: response.clone(),
628        })?
629    } else {
630        // Serializing a `Schema` is infallible in practice, but keep the error
631        // typed rather than unwrapping.
632        let schema =
633            serde_json::to_string_pretty(&schema).map_err(|source| GenerateError::Parse {
634                source,
635                response: String::new(),
636            })?;
637        let prompt = prompts::generate(&schema);
638        request.messages.push(Message::system(prompt));
639        request.parameters.structured_outputs = true;
640
641        let stream = model.respond(request);
642        collect_text(stream)
643            .await
644            .map_err(GenerateError::Provider)?
645    };
646
647    parse_json_with_recovery(&json).map_err(|source| GenerateError::Parse {
648        source,
649        response: truncate_for_error(&json),
650    })
651}
652
653/// Keeps a failed response short enough to log without dumping a whole context.
654fn truncate_for_error(response: &str) -> String {
655    const LIMIT: usize = 500;
656    response.chars().take(LIMIT).collect()
657}
658
659/// Convenience helper that creates a single system + user [`LLMRequest`].
660pub fn oneshot(system: impl Into<String>, user: impl Into<String>) -> LLMRequest {
661    let messages = vec![Message::system(system.into()), Message::user(user.into())];
662    LLMRequest::new(messages)
663}
664
665fn summarize<M: LanguageModel>(
666    model: &M,
667    text: &str,
668) -> impl Stream<Item = Result<Event, M::Error>> + Send {
669    let messages = oneshot("Summarize text:", text);
670    model.respond(messages)
671}
672
673async fn categorize_text<T: JsonSchema + DeserializeOwned + 'static, M: LanguageModel>(
674    model: &M,
675    text: &str,
676) -> Result<T, GenerateError<M::Error>> {
677    let request = oneshot("Categorize text by provided schema", text);
678    model.generate(request).await
679}
680
681fn parse_json_with_recovery<T: DeserializeOwned + 'static>(
682    json: &str,
683) -> Result<T, serde_json::Error> {
684    use serde::de::Error as _;
685
686    let trimmed = json.trim();
687    let mut last_error: Option<serde_json::Error> = None;
688    let mut last_candidate: Option<String> = None;
689
690    for candidate in build_json_candidates(trimmed) {
691        match serde_json::from_str::<T>(&candidate) {
692            Ok(value) => return Ok(value),
693            Err(err) => {
694                last_error = Some(err);
695                last_candidate = Some(candidate);
696            }
697        }
698    }
699
700    // A model asked for a plain string will often answer with an object or a
701    // bare word instead. Re-encode whatever it did produce as a JSON string.
702    if is_string_type::<T>()
703        && let Some(candidate) = last_candidate
704        && let Ok(value) = serde_json::from_str::<serde_json::Value>(&candidate)
705    {
706        let text = match value {
707            serde_json::Value::String(s) => s,
708            other => other.to_string(),
709        };
710        let encoded = serde_json::to_string(&text)?;
711        if let Ok(value) = serde_json::from_str::<T>(&encoded) {
712            return Ok(value);
713        }
714    }
715
716    Err(last_error.unwrap_or_else(|| {
717        serde_json::Error::custom("structured output was empty or missing a JSON block")
718    }))
719}
720
721fn strip_code_fences(raw: &str) -> Option<String> {
722    let trimmed = raw.trim();
723    let fence_start = trimmed.find("```")?;
724    let after_fence = &trimmed[fence_start + 3..];
725    let mut lines = after_fence.lines();
726    let _maybe_lang = lines.next();
727    let body = lines.collect::<Vec<_>>().join("\n");
728    let content = body.rfind("```").map_or(body.as_str(), |end| &body[..end]);
729
730    let cleaned = content.trim();
731    if cleaned.is_empty() {
732        None
733    } else {
734        Some(cleaned.to_string())
735    }
736}
737
738fn extract_json_block(raw: &str) -> Option<String> {
739    if let (Some(start), Some(end)) = (raw.find('{'), raw.rfind('}'))
740        && end >= start
741    {
742        let candidate = &raw[start..=end];
743        if !candidate.trim().is_empty() {
744            return Some(candidate.trim().to_string());
745        }
746    }
747    if let (Some(start), Some(end)) = (raw.find('['), raw.rfind(']'))
748        && end >= start
749    {
750        let candidate = &raw[start..=end];
751        if !candidate.trim().is_empty() {
752            return Some(candidate.trim().to_string());
753        }
754    }
755    None
756}
757
758fn build_json_candidates(raw: &str) -> Vec<String> {
759    let mut candidates = Vec::new();
760
761    if !raw.is_empty() {
762        candidates.push(raw.to_string());
763    }
764
765    if let Some(fenced) = strip_code_fences(raw) {
766        candidates.push(fenced);
767    }
768
769    if let Some(block) = extract_json_block(raw) {
770        candidates.push(block);
771    }
772
773    if let Some(dequoted) = dequote_json_string(raw) {
774        candidates.push(dequoted);
775    }
776
777    if let Some(stripped) = strip_leading_label(raw, "json") {
778        candidates.push(stripped);
779    }
780
781    let mut deduped = Vec::new();
782    for candidate in candidates {
783        if deduped.iter().all(|seen| seen != &candidate) {
784            deduped.push(candidate);
785        }
786    }
787    deduped
788}
789
790fn dequote_json_string(raw: &str) -> Option<String> {
791    let trimmed = raw.trim();
792    if !(trimmed.starts_with('"') && trimmed.ends_with('"')) {
793        return None;
794    }
795    let inner: String = serde_json::from_str(trimmed).ok()?;
796    if inner.trim().is_empty() {
797        None
798    } else {
799        Some(inner)
800    }
801}
802
803fn strip_leading_label(raw: &str, label: &str) -> Option<String> {
804    let trimmed = raw.trim_start();
805    if !trimmed.to_ascii_lowercase().starts_with(label) {
806        return None;
807    }
808    let stripped = trimmed[label.len()..]
809        .trim_start_matches(|c: char| c.is_whitespace() || c == ':' || c == '-')
810        .trim();
811    if stripped.is_empty() {
812        None
813    } else {
814        Some(stripped.to_string())
815    }
816}
817
818fn is_string_type<T: 'static>() -> bool {
819    TypeId::of::<T>() == TypeId::of::<String>()
820}
821
822#[cfg(test)]
823mod tests {
824    use super::parse_json_with_recovery;
825    use alloc::string::String;
826    use serde::Deserialize;
827
828    #[derive(Debug, Deserialize, PartialEq, Eq)]
829    struct Foo {
830        a: u8,
831    }
832
833    #[test]
834    fn parses_plain_json() {
835        let foo: Foo = parse_json_with_recovery(r#"{"a":1}"#).unwrap();
836        assert_eq!(foo, Foo { a: 1 });
837    }
838
839    #[test]
840    fn parses_code_fence_json() {
841        let foo: Foo = parse_json_with_recovery("```json\n{\"a\":2}\n```").unwrap();
842        assert_eq!(foo, Foo { a: 2 });
843    }
844
845    #[test]
846    fn parses_embedded_block() {
847        let foo: Foo = parse_json_with_recovery("noise {\"a\":3} trailing").unwrap();
848        assert_eq!(foo, Foo { a: 3 });
849    }
850
851    #[test]
852    fn parses_quoted_json_string() {
853        let foo: Foo = parse_json_with_recovery(r#""{\"a\":4}""#).unwrap();
854        assert_eq!(foo, Foo { a: 4 });
855    }
856
857    #[test]
858    fn parses_labeled_json() {
859        let foo: Foo = parse_json_with_recovery("json {\"a\":5}").unwrap();
860        assert_eq!(foo, Foo { a: 5 });
861    }
862
863    #[test]
864    fn coerces_object_to_string() {
865        let value: String =
866            parse_json_with_recovery(r#"{"title":"summary","type":"content"}"#).unwrap();
867        assert!(
868            value.contains("\"title\":\"summary\"") && value.contains("\"type\":\"content\""),
869            "unexpected value: {value}"
870        );
871    }
872}
aither_core/llm/mod.rs

aither_core/llm/
mod.rs
