1pub mod assistant;
79pub mod event;
81pub mod failure;
83pub mod message;
85pub mod model;
87pub mod provider;
89pub mod reasoning;
91pub mod researcher;
93pub mod retry;
95pub 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#[derive(Debug)]
138pub enum GenerateError<E> {
139 Provider(E),
141
142 Parse {
146 source: serde_json::Error,
148 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#[derive(Debug, Clone)]
181pub struct LLMRequest {
182 messages: Vec<Message>,
183 parameters: Parameters,
184 tool_definitions: Vec<tool::ToolDefinition>,
185}
186
187impl LLMRequest {
188 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 #[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 #[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 #[must_use]
216 pub fn with_parameters(mut self, parameters: Parameters) -> Self {
217 self.parameters = parameters;
218 self
219 }
220
221 #[must_use]
223 pub fn messages(&self) -> &[Message] {
224 &self.messages
225 }
226
227 pub const fn messages_mut(&mut self) -> &mut Vec<Message> {
229 &mut self.messages
230 }
231
232 #[must_use]
234 pub const fn parameters(&self) -> &Parameters {
235 &self.parameters
236 }
237
238 #[must_use]
240 pub fn tool_definitions(&self) -> &[tool::ToolDefinition] {
241 &self.tool_definitions
242 }
243
244 #[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#[derive(Debug)]
256pub struct LLMRequestWithTools<'tools> {
257 inner: LLMRequest,
258 tools: &'tools mut Tools,
259}
260
261impl LLMRequest {
262 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 #[must_use]
278 pub const fn request(&self) -> &LLMRequest {
279 &self.inner
280 }
281
282 #[must_use]
284 pub const fn tools(&mut self) -> &mut Tools {
285 self.tools
286 }
287
288 #[must_use]
290 pub fn into_parts(self) -> (LLMRequest, &'tools mut Tools) {
291 (self.inner, self.tools)
292 }
293
294 pub async fn call_tool(
300 &mut self,
301 name: &str,
302 args_json: &str,
303 cx: ToolContext,
304 ) -> crate::Result<ToolResult> {
305 self.tools.call(name, args_json, cx).await
306 }
307}
308
309pub trait LanguageModel: Sized + Send + Sync {
317 type Error: core::error::Error + Classify + Send + Sync + 'static;
319
320 fn respond(&self, request: LLMRequest)
328 -> impl Stream<Item = Result<Event, Self::Error>> + Send;
329
330 fn respond_with_tools(
335 &self,
336 request: LLMRequestWithTools<'_>,
337 ) -> impl Stream<Item = Result<Event, Self::Error>> + Send {
338 let (inner, _tools) = request.into_parts();
339 self.respond(inner)
340 }
341
342 fn generate<T: JsonSchema + DeserializeOwned + 'static>(
349 &self,
350 request: LLMRequest,
351 ) -> impl Future<Output = Result<T, GenerateError<Self::Error>>> + Send {
352 async { structured_generate(self, request).await }
353 }
354
355 fn complete(&self, prefix: &str) -> impl Stream<Item = Result<Event, Self::Error>> + Send {
357 self.respond(oneshot("Please complete the following text:", prefix))
358 }
359
360 fn summarize(&self, text: &str) -> impl Stream<Item = Result<Event, Self::Error>> + Send {
366 summarize(self, text)
367 }
368
369 fn categorize<T: JsonSchema + DeserializeOwned + 'static>(
375 &self,
376 text: &str,
377 ) -> impl Future<Output = Result<T, GenerateError<Self::Error>>> + Send {
378 async { categorize_text(self, text).await }
379 }
380
381 fn profile(&self) -> impl Future<Output = Result<Profile, Self::Error>> + Send;
391
392 fn retry_policy(&self) -> RetryPolicy {
400 RetryPolicy::default()
401 }
402}
403
404macro_rules! impl_language_model {
405 ($($name:ident),*) => {
406 $(
407 impl<T: LanguageModel> LanguageModel for $name<T> {
408 type Error = T::Error;
409
410 fn respond(
411 &self,
412 request: LLMRequest,
413 ) -> impl Stream<Item = Result<Event, Self::Error>> + Send {
414 T::respond(self, request)
415 }
416
417 fn respond_with_tools(
418 &self,
419 request: LLMRequestWithTools<'_>,
420 ) -> impl Stream<Item = Result<Event, Self::Error>> + Send {
421 T::respond_with_tools(self, request)
422 }
423
424 fn generate<U: JsonSchema + DeserializeOwned + 'static>(
425 &self,
426 request: LLMRequest,
427 ) -> impl Future<Output = Result<U, GenerateError<Self::Error>>> + Send {
428 T::generate(self, request)
429 }
430
431 fn complete(
432 &self,
433 prefix: &str,
434 ) -> impl Stream<Item = Result<Event, Self::Error>> + Send {
435 T::complete(self, prefix)
436 }
437
438 fn summarize(
439 &self,
440 text: &str,
441 ) -> impl Stream<Item = Result<Event, Self::Error>> + Send {
442 T::summarize(self, text)
443 }
444
445 fn categorize<U: JsonSchema + DeserializeOwned + 'static>(
446 &self,
447 text: &str,
448 ) -> impl Future<Output = Result<U, GenerateError<Self::Error>>> + Send {
449 T::categorize(self, text)
450 }
451
452 fn profile(&self) -> impl Future<Output = Result<Profile, Self::Error>> + Send {
453 T::profile(self)
454 }
455
456 fn retry_policy(&self) -> RetryPolicy {
457 T::retry_policy(self)
458 }
459 }
460 )*
461 };
462}
463
464impl<T: LanguageModel> LanguageModel for &T {
465 type Error = T::Error;
466
467 fn respond(
468 &self,
469 request: LLMRequest,
470 ) -> impl Stream<Item = Result<Event, Self::Error>> + Send {
471 T::respond(self, request)
472 }
473
474 fn respond_with_tools(
475 &self,
476 request: LLMRequestWithTools<'_>,
477 ) -> impl Stream<Item = Result<Event, Self::Error>> + Send {
478 T::respond_with_tools(self, request)
479 }
480
481 fn generate<U: JsonSchema + DeserializeOwned + 'static>(
482 &self,
483 request: LLMRequest,
484 ) -> impl Future<Output = Result<U, GenerateError<Self::Error>>> + Send {
485 T::generate(self, request)
486 }
487
488 fn complete(&self, prefix: &str) -> impl Stream<Item = Result<Event, Self::Error>> + Send {
489 T::complete(self, prefix)
490 }
491
492 fn summarize(&self, text: &str) -> impl Stream<Item = Result<Event, Self::Error>> + Send {
493 T::summarize(self, text)
494 }
495
496 fn categorize<U: JsonSchema + DeserializeOwned + 'static>(
497 &self,
498 text: &str,
499 ) -> impl Future<Output = Result<U, GenerateError<Self::Error>>> + Send {
500 T::categorize(self, text)
501 }
502
503 fn profile(&self) -> impl Future<Output = Result<Profile, Self::Error>> + Send {
504 T::profile(self)
505 }
506
507 fn retry_policy(&self) -> RetryPolicy {
508 T::retry_policy(self)
509 }
510}
511
512mod prompts;
513
514impl_language_model!(Arc, Box);
515
516pub async fn collect_text<S, E>(stream: S) -> Result<String, E>
522where
523 S: Stream<Item = Result<Event, E>>,
524{
525 pin!(stream);
526 let mut result = String::new();
527 while let Some(event) = stream.next().await {
528 if let Event::Text(text) = event? {
529 result.push_str(&text);
530 }
531 }
532 Ok(result)
533}
534
535async fn structured_generate<T: JsonSchema + DeserializeOwned + 'static, M: LanguageModel>(
536 model: &M,
537 mut request: LLMRequest,
538) -> Result<T, GenerateError<M::Error>> {
539 let schema = schema_for!(T);
540
541 let json = if schema.as_value().is_string() {
543 let stream = model.respond(request);
544 let response = collect_text(stream)
545 .await
546 .map_err(GenerateError::Provider)?;
547 serde_json::to_string(&response).map_err(|source| GenerateError::Parse {
549 source,
550 response: response.clone(),
551 })?
552 } else {
553 let schema =
556 serde_json::to_string_pretty(&schema).map_err(|source| GenerateError::Parse {
557 source,
558 response: String::new(),
559 })?;
560 let prompt = prompts::generate(&schema);
561 request.messages.push(Message::system(prompt));
562 request.parameters.structured_outputs = true;
563
564 let stream = model.respond(request);
565 collect_text(stream)
566 .await
567 .map_err(GenerateError::Provider)?
568 };
569
570 parse_json_with_recovery(&json).map_err(|source| GenerateError::Parse {
571 source,
572 response: truncate_for_error(&json),
573 })
574}
575
576fn truncate_for_error(response: &str) -> String {
578 const LIMIT: usize = 500;
579 response.chars().take(LIMIT).collect()
580}
581
582pub fn oneshot(system: impl Into<String>, user: impl Into<String>) -> LLMRequest {
584 let messages = vec![Message::system(system.into()), Message::user(user.into())];
585 LLMRequest::new(messages)
586}
587
588fn summarize<M: LanguageModel>(
589 model: &M,
590 text: &str,
591) -> impl Stream<Item = Result<Event, M::Error>> + Send {
592 let messages = oneshot("Summarize text:", text);
593 model.respond(messages)
594}
595
596async fn categorize_text<T: JsonSchema + DeserializeOwned + 'static, M: LanguageModel>(
597 model: &M,
598 text: &str,
599) -> Result<T, GenerateError<M::Error>> {
600 let request = oneshot("Categorize text by provided schema", text);
601 model.generate(request).await
602}
603
604fn parse_json_with_recovery<T: DeserializeOwned + 'static>(
605 json: &str,
606) -> Result<T, serde_json::Error> {
607 use serde::de::Error as _;
608
609 let trimmed = json.trim();
610 let mut last_error: Option<serde_json::Error> = None;
611 let mut last_candidate: Option<String> = None;
612
613 for candidate in build_json_candidates(trimmed) {
614 match serde_json::from_str::<T>(&candidate) {
615 Ok(value) => return Ok(value),
616 Err(err) => {
617 last_error = Some(err);
618 last_candidate = Some(candidate);
619 }
620 }
621 }
622
623 if is_string_type::<T>()
626 && let Some(candidate) = last_candidate
627 && let Ok(value) = serde_json::from_str::<serde_json::Value>(&candidate)
628 {
629 let text = match value {
630 serde_json::Value::String(s) => s,
631 other => other.to_string(),
632 };
633 let encoded = serde_json::to_string(&text)?;
634 if let Ok(value) = serde_json::from_str::<T>(&encoded) {
635 return Ok(value);
636 }
637 }
638
639 Err(last_error.unwrap_or_else(|| {
640 serde_json::Error::custom("structured output was empty or missing a JSON block")
641 }))
642}
643
644fn strip_code_fences(raw: &str) -> Option<String> {
645 let trimmed = raw.trim();
646 let fence_start = trimmed.find("```")?;
647 let after_fence = &trimmed[fence_start + 3..];
648 let mut lines = after_fence.lines();
649 let _maybe_lang = lines.next();
650 let body = lines.collect::<Vec<_>>().join("\n");
651 let content = body.rfind("```").map_or(body.as_str(), |end| &body[..end]);
652
653 let cleaned = content.trim();
654 if cleaned.is_empty() {
655 None
656 } else {
657 Some(cleaned.to_string())
658 }
659}
660
661fn extract_json_block(raw: &str) -> Option<String> {
662 if let (Some(start), Some(end)) = (raw.find('{'), raw.rfind('}'))
663 && end >= start
664 {
665 let candidate = &raw[start..=end];
666 if !candidate.trim().is_empty() {
667 return Some(candidate.trim().to_string());
668 }
669 }
670 if let (Some(start), Some(end)) = (raw.find('['), raw.rfind(']'))
671 && end >= start
672 {
673 let candidate = &raw[start..=end];
674 if !candidate.trim().is_empty() {
675 return Some(candidate.trim().to_string());
676 }
677 }
678 None
679}
680
681fn build_json_candidates(raw: &str) -> Vec<String> {
682 let mut candidates = Vec::new();
683
684 if !raw.is_empty() {
685 candidates.push(raw.to_string());
686 }
687
688 if let Some(fenced) = strip_code_fences(raw) {
689 candidates.push(fenced);
690 }
691
692 if let Some(block) = extract_json_block(raw) {
693 candidates.push(block);
694 }
695
696 if let Some(dequoted) = dequote_json_string(raw) {
697 candidates.push(dequoted);
698 }
699
700 if let Some(stripped) = strip_leading_label(raw, "json") {
701 candidates.push(stripped);
702 }
703
704 let mut deduped = Vec::new();
705 for candidate in candidates {
706 if deduped.iter().all(|seen| seen != &candidate) {
707 deduped.push(candidate);
708 }
709 }
710 deduped
711}
712
713fn dequote_json_string(raw: &str) -> Option<String> {
714 let trimmed = raw.trim();
715 if !(trimmed.starts_with('"') && trimmed.ends_with('"')) {
716 return None;
717 }
718 let inner: String = serde_json::from_str(trimmed).ok()?;
719 if inner.trim().is_empty() {
720 None
721 } else {
722 Some(inner)
723 }
724}
725
726fn strip_leading_label(raw: &str, label: &str) -> Option<String> {
727 let trimmed = raw.trim_start();
728 if !trimmed.to_ascii_lowercase().starts_with(label) {
729 return None;
730 }
731 let stripped = trimmed[label.len()..]
732 .trim_start_matches(|c: char| c.is_whitespace() || c == ':' || c == '-')
733 .trim();
734 if stripped.is_empty() {
735 None
736 } else {
737 Some(stripped.to_string())
738 }
739}
740
741fn is_string_type<T: 'static>() -> bool {
742 TypeId::of::<T>() == TypeId::of::<String>()
743}
744
745#[cfg(test)]
746mod tests {
747 use super::parse_json_with_recovery;
748 use alloc::string::String;
749 use serde::Deserialize;
750
751 #[derive(Debug, Deserialize, PartialEq, Eq)]
752 struct Foo {
753 a: u8,
754 }
755
756 #[test]
757 fn parses_plain_json() {
758 let foo: Foo = parse_json_with_recovery(r#"{"a":1}"#).unwrap();
759 assert_eq!(foo, Foo { a: 1 });
760 }
761
762 #[test]
763 fn parses_code_fence_json() {
764 let foo: Foo = parse_json_with_recovery("```json\n{\"a\":2}\n```").unwrap();
765 assert_eq!(foo, Foo { a: 2 });
766 }
767
768 #[test]
769 fn parses_embedded_block() {
770 let foo: Foo = parse_json_with_recovery("noise {\"a\":3} trailing").unwrap();
771 assert_eq!(foo, Foo { a: 3 });
772 }
773
774 #[test]
775 fn parses_quoted_json_string() {
776 let foo: Foo = parse_json_with_recovery(r#""{\"a\":4}""#).unwrap();
777 assert_eq!(foo, Foo { a: 4 });
778 }
779
780 #[test]
781 fn parses_labeled_json() {
782 let foo: Foo = parse_json_with_recovery("json {\"a\":5}").unwrap();
783 assert_eq!(foo, Foo { a: 5 });
784 }
785
786 #[test]
787 fn coerces_object_to_string() {
788 let value: String =
789 parse_json_with_recovery(r#"{"title":"summary","type":"content"}"#).unwrap();
790 assert!(
791 value.contains("\"title\":\"summary\"") && value.contains("\"type\":\"content\""),
792 "unexpected value: {value}"
793 );
794 }
795}