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 pub const fn request_mut(&mut self) -> &mut LLMRequest {
284 &mut self.inner
285 }
286
287 #[must_use]
289 pub const fn tools(&mut self) -> &mut Tools {
290 self.tools
291 }
292
293 #[must_use]
295 pub fn into_parts(self) -> (LLMRequest, &'tools mut Tools) {
296 (self.inner, self.tools)
297 }
298
299 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
314pub trait LanguageModel: Sized + Send + Sync {
322 type Error: core::error::Error + Classify + Send + Sync + 'static;
324
325 fn respond(&self, request: LLMRequest)
333 -> impl Stream<Item = Result<Event, Self::Error>> + Send;
334
335 fn respond_with_tools(
340 &self,
341 request: LLMRequestWithTools<'_>,
342 ) -> impl Stream<Item = Result<Event, Self::Error>> + Send {
343 let (inner, _tools) = request.into_parts();
344 self.respond(inner)
345 }
346
347 fn generate<T: JsonSchema + DeserializeOwned + 'static>(
354 &self,
355 request: LLMRequest,
356 ) -> impl Future<Output = Result<T, GenerateError<Self::Error>>> + Send {
357 async { structured_generate(self, request).await }
358 }
359
360 fn complete(&self, prefix: &str) -> impl Stream<Item = Result<Event, Self::Error>> + Send {
362 self.respond(oneshot("Please complete the following text:", prefix))
363 }
364
365 fn summarize(&self, text: &str) -> impl Stream<Item = Result<Event, Self::Error>> + Send {
371 summarize(self, text)
372 }
373
374 fn categorize<T: JsonSchema + DeserializeOwned + 'static>(
380 &self,
381 text: &str,
382 ) -> impl Future<Output = Result<T, GenerateError<Self::Error>>> + Send {
383 async { categorize_text(self, text).await }
384 }
385
386 fn profile(&self) -> impl Future<Output = Result<Profile, Self::Error>> + Send;
396
397 fn retry_policy(&self) -> RetryPolicy {
405 RetryPolicy::default()
406 }
407}
408
409macro_rules! impl_language_model {
410 ($($name:ident),*) => {
411 $(
412 impl<T: LanguageModel> LanguageModel for $name<T> {
413 type Error = T::Error;
414
415 fn respond(
416 &self,
417 request: LLMRequest,
418 ) -> impl Stream<Item = Result<Event, Self::Error>> + Send {
419 T::respond(self, request)
420 }
421
422 fn respond_with_tools(
423 &self,
424 request: LLMRequestWithTools<'_>,
425 ) -> impl Stream<Item = Result<Event, Self::Error>> + Send {
426 T::respond_with_tools(self, request)
427 }
428
429 fn generate<U: JsonSchema + DeserializeOwned + 'static>(
430 &self,
431 request: LLMRequest,
432 ) -> impl Future<Output = Result<U, GenerateError<Self::Error>>> + Send {
433 T::generate(self, request)
434 }
435
436 fn complete(
437 &self,
438 prefix: &str,
439 ) -> impl Stream<Item = Result<Event, Self::Error>> + Send {
440 T::complete(self, prefix)
441 }
442
443 fn summarize(
444 &self,
445 text: &str,
446 ) -> impl Stream<Item = Result<Event, Self::Error>> + Send {
447 T::summarize(self, text)
448 }
449
450 fn categorize<U: JsonSchema + DeserializeOwned + 'static>(
451 &self,
452 text: &str,
453 ) -> impl Future<Output = Result<U, GenerateError<Self::Error>>> + Send {
454 T::categorize(self, text)
455 }
456
457 fn profile(&self) -> impl Future<Output = Result<Profile, Self::Error>> + Send {
458 T::profile(self)
459 }
460
461 fn retry_policy(&self) -> RetryPolicy {
462 T::retry_policy(self)
463 }
464 }
465 )*
466 };
467}
468
469impl<T: LanguageModel> LanguageModel for &T {
470 type Error = T::Error;
471
472 fn respond(
473 &self,
474 request: LLMRequest,
475 ) -> impl Stream<Item = Result<Event, Self::Error>> + Send {
476 T::respond(self, request)
477 }
478
479 fn respond_with_tools(
480 &self,
481 request: LLMRequestWithTools<'_>,
482 ) -> impl Stream<Item = Result<Event, Self::Error>> + Send {
483 T::respond_with_tools(self, request)
484 }
485
486 fn generate<U: JsonSchema + DeserializeOwned + 'static>(
487 &self,
488 request: LLMRequest,
489 ) -> impl Future<Output = Result<U, GenerateError<Self::Error>>> + Send {
490 T::generate(self, request)
491 }
492
493 fn complete(&self, prefix: &str) -> impl Stream<Item = Result<Event, Self::Error>> + Send {
494 T::complete(self, prefix)
495 }
496
497 fn summarize(&self, text: &str) -> impl Stream<Item = Result<Event, Self::Error>> + Send {
498 T::summarize(self, text)
499 }
500
501 fn categorize<U: JsonSchema + DeserializeOwned + 'static>(
502 &self,
503 text: &str,
504 ) -> impl Future<Output = Result<U, GenerateError<Self::Error>>> + Send {
505 T::categorize(self, text)
506 }
507
508 fn profile(&self) -> impl Future<Output = Result<Profile, Self::Error>> + Send {
509 T::profile(self)
510 }
511
512 fn retry_policy(&self) -> RetryPolicy {
513 T::retry_policy(self)
514 }
515}
516
517mod prompts;
518
519impl_language_model!(Arc, Box);
520
521pub async fn collect_text<S, E>(stream: S) -> Result<String, E>
527where
528 S: Stream<Item = Result<Event, E>>,
529{
530 pin!(stream);
531 let mut result = String::new();
532 while let Some(event) = stream.next().await {
533 if let Event::Text(text) = event? {
534 result.push_str(&text);
535 }
536 }
537 Ok(result)
538}
539
540async fn structured_generate<T: JsonSchema + DeserializeOwned + 'static, M: LanguageModel>(
541 model: &M,
542 mut request: LLMRequest,
543) -> Result<T, GenerateError<M::Error>> {
544 let schema = schema_for!(T);
545
546 let json = if schema.as_value().is_string() {
548 let stream = model.respond(request);
549 let response = collect_text(stream)
550 .await
551 .map_err(GenerateError::Provider)?;
552 serde_json::to_string(&response).map_err(|source| GenerateError::Parse {
554 source,
555 response: response.clone(),
556 })?
557 } else {
558 let schema =
561 serde_json::to_string_pretty(&schema).map_err(|source| GenerateError::Parse {
562 source,
563 response: String::new(),
564 })?;
565 let prompt = prompts::generate(&schema);
566 request.messages.push(Message::system(prompt));
567 request.parameters.structured_outputs = true;
568
569 let stream = model.respond(request);
570 collect_text(stream)
571 .await
572 .map_err(GenerateError::Provider)?
573 };
574
575 parse_json_with_recovery(&json).map_err(|source| GenerateError::Parse {
576 source,
577 response: truncate_for_error(&json),
578 })
579}
580
581fn truncate_for_error(response: &str) -> String {
583 const LIMIT: usize = 500;
584 response.chars().take(LIMIT).collect()
585}
586
587pub fn oneshot(system: impl Into<String>, user: impl Into<String>) -> LLMRequest {
589 let messages = vec![Message::system(system.into()), Message::user(user.into())];
590 LLMRequest::new(messages)
591}
592
593fn summarize<M: LanguageModel>(
594 model: &M,
595 text: &str,
596) -> impl Stream<Item = Result<Event, M::Error>> + Send {
597 let messages = oneshot("Summarize text:", text);
598 model.respond(messages)
599}
600
601async fn categorize_text<T: JsonSchema + DeserializeOwned + 'static, M: LanguageModel>(
602 model: &M,
603 text: &str,
604) -> Result<T, GenerateError<M::Error>> {
605 let request = oneshot("Categorize text by provided schema", text);
606 model.generate(request).await
607}
608
609fn parse_json_with_recovery<T: DeserializeOwned + 'static>(
610 json: &str,
611) -> Result<T, serde_json::Error> {
612 use serde::de::Error as _;
613
614 let trimmed = json.trim();
615 let mut last_error: Option<serde_json::Error> = None;
616 let mut last_candidate: Option<String> = None;
617
618 for candidate in build_json_candidates(trimmed) {
619 match serde_json::from_str::<T>(&candidate) {
620 Ok(value) => return Ok(value),
621 Err(err) => {
622 last_error = Some(err);
623 last_candidate = Some(candidate);
624 }
625 }
626 }
627
628 if is_string_type::<T>()
631 && let Some(candidate) = last_candidate
632 && let Ok(value) = serde_json::from_str::<serde_json::Value>(&candidate)
633 {
634 let text = match value {
635 serde_json::Value::String(s) => s,
636 other => other.to_string(),
637 };
638 let encoded = serde_json::to_string(&text)?;
639 if let Ok(value) = serde_json::from_str::<T>(&encoded) {
640 return Ok(value);
641 }
642 }
643
644 Err(last_error.unwrap_or_else(|| {
645 serde_json::Error::custom("structured output was empty or missing a JSON block")
646 }))
647}
648
649fn strip_code_fences(raw: &str) -> Option<String> {
650 let trimmed = raw.trim();
651 let fence_start = trimmed.find("```")?;
652 let after_fence = &trimmed[fence_start + 3..];
653 let mut lines = after_fence.lines();
654 let _maybe_lang = lines.next();
655 let body = lines.collect::<Vec<_>>().join("\n");
656 let content = body.rfind("```").map_or(body.as_str(), |end| &body[..end]);
657
658 let cleaned = content.trim();
659 if cleaned.is_empty() {
660 None
661 } else {
662 Some(cleaned.to_string())
663 }
664}
665
666fn extract_json_block(raw: &str) -> Option<String> {
667 if let (Some(start), Some(end)) = (raw.find('{'), raw.rfind('}'))
668 && end >= start
669 {
670 let candidate = &raw[start..=end];
671 if !candidate.trim().is_empty() {
672 return Some(candidate.trim().to_string());
673 }
674 }
675 if let (Some(start), Some(end)) = (raw.find('['), raw.rfind(']'))
676 && end >= start
677 {
678 let candidate = &raw[start..=end];
679 if !candidate.trim().is_empty() {
680 return Some(candidate.trim().to_string());
681 }
682 }
683 None
684}
685
686fn build_json_candidates(raw: &str) -> Vec<String> {
687 let mut candidates = Vec::new();
688
689 if !raw.is_empty() {
690 candidates.push(raw.to_string());
691 }
692
693 if let Some(fenced) = strip_code_fences(raw) {
694 candidates.push(fenced);
695 }
696
697 if let Some(block) = extract_json_block(raw) {
698 candidates.push(block);
699 }
700
701 if let Some(dequoted) = dequote_json_string(raw) {
702 candidates.push(dequoted);
703 }
704
705 if let Some(stripped) = strip_leading_label(raw, "json") {
706 candidates.push(stripped);
707 }
708
709 let mut deduped = Vec::new();
710 for candidate in candidates {
711 if deduped.iter().all(|seen| seen != &candidate) {
712 deduped.push(candidate);
713 }
714 }
715 deduped
716}
717
718fn dequote_json_string(raw: &str) -> Option<String> {
719 let trimmed = raw.trim();
720 if !(trimmed.starts_with('"') && trimmed.ends_with('"')) {
721 return None;
722 }
723 let inner: String = serde_json::from_str(trimmed).ok()?;
724 if inner.trim().is_empty() {
725 None
726 } else {
727 Some(inner)
728 }
729}
730
731fn strip_leading_label(raw: &str, label: &str) -> Option<String> {
732 let trimmed = raw.trim_start();
733 if !trimmed.to_ascii_lowercase().starts_with(label) {
734 return None;
735 }
736 let stripped = trimmed[label.len()..]
737 .trim_start_matches(|c: char| c.is_whitespace() || c == ':' || c == '-')
738 .trim();
739 if stripped.is_empty() {
740 None
741 } else {
742 Some(stripped.to_string())
743 }
744}
745
746fn is_string_type<T: 'static>() -> bool {
747 TypeId::of::<T>() == TypeId::of::<String>()
748}
749
750#[cfg(test)]
751mod tests {
752 use super::parse_json_with_recovery;
753 use alloc::string::String;
754 use serde::Deserialize;
755
756 #[derive(Debug, Deserialize, PartialEq, Eq)]
757 struct Foo {
758 a: u8,
759 }
760
761 #[test]
762 fn parses_plain_json() {
763 let foo: Foo = parse_json_with_recovery(r#"{"a":1}"#).unwrap();
764 assert_eq!(foo, Foo { a: 1 });
765 }
766
767 #[test]
768 fn parses_code_fence_json() {
769 let foo: Foo = parse_json_with_recovery("```json\n{\"a\":2}\n```").unwrap();
770 assert_eq!(foo, Foo { a: 2 });
771 }
772
773 #[test]
774 fn parses_embedded_block() {
775 let foo: Foo = parse_json_with_recovery("noise {\"a\":3} trailing").unwrap();
776 assert_eq!(foo, Foo { a: 3 });
777 }
778
779 #[test]
780 fn parses_quoted_json_string() {
781 let foo: Foo = parse_json_with_recovery(r#""{\"a\":4}""#).unwrap();
782 assert_eq!(foo, Foo { a: 4 });
783 }
784
785 #[test]
786 fn parses_labeled_json() {
787 let foo: Foo = parse_json_with_recovery("json {\"a\":5}").unwrap();
788 assert_eq!(foo, Foo { a: 5 });
789 }
790
791 #[test]
792 fn coerces_object_to_string() {
793 let value: String =
794 parse_json_with_recovery(r#"{"title":"summary","type":"content"}"#).unwrap();
795 assert!(
796 value.contains("\"title\":\"summary\"") && value.contains("\"type\":\"content\""),
797 "unexpected value: {value}"
798 );
799 }
800}