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 tool;
95
96use crate::llm::{model::Parameters, tool::Tools};
97use alloc::{
98 boxed::Box,
99 string::{String, ToString},
100 sync::Arc,
101 vec,
102 vec::Vec,
103};
104use core::{any::TypeId, future::Future};
105pub use event::{Event, ToolCall, Usage};
106pub use failure::{Classify, Failure, FailureKind};
107use futures_core::Stream;
108use futures_lite::{StreamExt, pin};
109pub use message::{Attachment, Message, Role};
110pub use provider::LanguageModelProvider;
111pub use reasoning::ReasoningState;
112pub use researcher::{
113 ResearchCitation, ResearchEvent, ResearchFinding, ResearchOptions, ResearchReport,
114 ResearchRequest, ResearchSource, ResearchStage, Researcher, ResearcherProfile,
115};
116use schemars::{JsonSchema, schema_for};
117use serde::de::DeserializeOwned;
118pub use tool::{
119 CommandOutcome, CommandResult, CommandStream, Interruption, InterruptionCause, IntoToolResult,
120 StoredFile, Tool, ToolContext, ToolResult, render_interruptions,
121};
122
123use crate::llm::model::Profile;
124
125#[derive(Debug)]
134pub enum GenerateError<E> {
135 Provider(E),
137
138 Parse {
142 source: serde_json::Error,
144 response: String,
146 },
147}
148
149impl<E: core::fmt::Display> core::fmt::Display for GenerateError<E> {
150 fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
151 match self {
152 Self::Provider(err) => write!(f, "language model request failed: {err}"),
153 Self::Parse { source, response } => {
154 write!(
155 f,
156 "structured output did not match the requested schema: {source}; response: {response}"
157 )
158 }
159 }
160 }
161}
162
163impl<E: core::error::Error + 'static> core::error::Error for GenerateError<E> {
164 fn source(&self) -> Option<&(dyn core::error::Error + 'static)> {
165 match self {
166 Self::Provider(err) => Some(err),
167 Self::Parse { source, .. } => Some(source),
168 }
169 }
170}
171
172#[derive(Debug, Clone)]
177pub struct LLMRequest {
178 messages: Vec<Message>,
179 parameters: Parameters,
180 tool_definitions: Vec<tool::ToolDefinition>,
181}
182
183impl LLMRequest {
184 pub fn new(messages: impl Into<Vec<Message>>) -> Self {
186 Self {
187 messages: messages.into(),
188 parameters: Parameters::default(),
189 tool_definitions: Vec::new(),
190 }
191 }
192
193 #[must_use]
198 pub fn with_tool_definitions(mut self, definitions: Vec<tool::ToolDefinition>) -> Self {
199 self.tool_definitions = definitions;
200 self
201 }
202
203 #[must_use]
205 pub fn with_tool<T: Tool>(mut self, tool: &T) -> Self {
206 self.tool_definitions.push(tool::ToolDefinition::new(tool));
207 self
208 }
209
210 #[must_use]
212 pub fn with_parameters(mut self, parameters: Parameters) -> Self {
213 self.parameters = parameters;
214 self
215 }
216
217 #[must_use]
219 pub fn messages(&self) -> &[Message] {
220 &self.messages
221 }
222
223 pub const fn messages_mut(&mut self) -> &mut Vec<Message> {
225 &mut self.messages
226 }
227
228 #[must_use]
230 pub const fn parameters(&self) -> &Parameters {
231 &self.parameters
232 }
233
234 #[must_use]
236 pub fn tool_definitions(&self) -> &[tool::ToolDefinition] {
237 &self.tool_definitions
238 }
239
240 #[must_use]
242 pub fn into_parts(self) -> (Vec<Message>, Parameters, Vec<tool::ToolDefinition>) {
243 (self.messages, self.parameters, self.tool_definitions)
244 }
245}
246
247#[derive(Debug)]
252pub struct LLMRequestWithTools<'tools> {
253 inner: LLMRequest,
254 tools: &'tools mut Tools,
255}
256
257impl LLMRequest {
258 pub fn with_tools(self, tools: &mut Tools) -> LLMRequestWithTools<'_> {
263 let definitions = tools.definitions();
264 LLMRequestWithTools {
265 inner: self.with_tool_definitions(definitions),
266 tools,
267 }
268 }
269}
270
271impl<'tools> LLMRequestWithTools<'tools> {
272 #[must_use]
274 pub const fn request(&self) -> &LLMRequest {
275 &self.inner
276 }
277
278 #[must_use]
280 pub const fn tools(&mut self) -> &mut Tools {
281 self.tools
282 }
283
284 #[must_use]
286 pub fn into_parts(self) -> (LLMRequest, &'tools mut Tools) {
287 (self.inner, self.tools)
288 }
289
290 pub async fn call_tool(
296 &mut self,
297 name: &str,
298 args_json: &str,
299 cx: ToolContext,
300 ) -> crate::Result<ToolResult> {
301 self.tools.call(name, args_json, cx).await
302 }
303}
304
305pub trait LanguageModel: Sized + Send + Sync {
313 type Error: core::error::Error + Classify + Send + Sync + 'static;
315
316 fn respond(&self, request: LLMRequest)
324 -> impl Stream<Item = Result<Event, Self::Error>> + Send;
325
326 fn respond_with_tools(
331 &self,
332 request: LLMRequestWithTools<'_>,
333 ) -> impl Stream<Item = Result<Event, Self::Error>> + Send {
334 let (inner, _tools) = request.into_parts();
335 self.respond(inner)
336 }
337
338 fn generate<T: JsonSchema + DeserializeOwned + 'static>(
345 &self,
346 request: LLMRequest,
347 ) -> impl Future<Output = Result<T, GenerateError<Self::Error>>> + Send {
348 async { structured_generate(self, request).await }
349 }
350
351 fn complete(&self, prefix: &str) -> impl Stream<Item = Result<Event, Self::Error>> + Send {
353 self.respond(oneshot("Please complete the following text:", prefix))
354 }
355
356 fn summarize(&self, text: &str) -> impl Stream<Item = Result<Event, Self::Error>> + Send {
362 summarize(self, text)
363 }
364
365 fn categorize<T: JsonSchema + DeserializeOwned + 'static>(
371 &self,
372 text: &str,
373 ) -> impl Future<Output = Result<T, GenerateError<Self::Error>>> + Send {
374 async { categorize_text(self, text).await }
375 }
376
377 fn profile(&self) -> impl Future<Output = Result<Profile, Self::Error>> + Send;
387}
388
389macro_rules! impl_language_model {
390 ($($name:ident),*) => {
391 $(
392 impl<T: LanguageModel> LanguageModel for $name<T> {
393 type Error = T::Error;
394
395 fn respond(
396 &self,
397 request: LLMRequest,
398 ) -> impl Stream<Item = Result<Event, Self::Error>> + Send {
399 T::respond(self, request)
400 }
401
402 fn respond_with_tools(
403 &self,
404 request: LLMRequestWithTools<'_>,
405 ) -> impl Stream<Item = Result<Event, Self::Error>> + Send {
406 T::respond_with_tools(self, request)
407 }
408
409 fn generate<U: JsonSchema + DeserializeOwned + 'static>(
410 &self,
411 request: LLMRequest,
412 ) -> impl Future<Output = Result<U, GenerateError<Self::Error>>> + Send {
413 T::generate(self, request)
414 }
415
416 fn complete(
417 &self,
418 prefix: &str,
419 ) -> impl Stream<Item = Result<Event, Self::Error>> + Send {
420 T::complete(self, prefix)
421 }
422
423 fn summarize(
424 &self,
425 text: &str,
426 ) -> impl Stream<Item = Result<Event, Self::Error>> + Send {
427 T::summarize(self, text)
428 }
429
430 fn categorize<U: JsonSchema + DeserializeOwned + 'static>(
431 &self,
432 text: &str,
433 ) -> impl Future<Output = Result<U, GenerateError<Self::Error>>> + Send {
434 T::categorize(self, text)
435 }
436
437 fn profile(&self) -> impl Future<Output = Result<Profile, Self::Error>> + Send {
438 T::profile(self)
439 }
440 }
441 )*
442 };
443}
444
445impl<T: LanguageModel> LanguageModel for &T {
446 type Error = T::Error;
447
448 fn respond(
449 &self,
450 request: LLMRequest,
451 ) -> impl Stream<Item = Result<Event, Self::Error>> + Send {
452 T::respond(self, request)
453 }
454
455 fn respond_with_tools(
456 &self,
457 request: LLMRequestWithTools<'_>,
458 ) -> impl Stream<Item = Result<Event, Self::Error>> + Send {
459 T::respond_with_tools(self, request)
460 }
461
462 fn generate<U: JsonSchema + DeserializeOwned + 'static>(
463 &self,
464 request: LLMRequest,
465 ) -> impl Future<Output = Result<U, GenerateError<Self::Error>>> + Send {
466 T::generate(self, request)
467 }
468
469 fn complete(&self, prefix: &str) -> impl Stream<Item = Result<Event, Self::Error>> + Send {
470 T::complete(self, prefix)
471 }
472
473 fn summarize(&self, text: &str) -> impl Stream<Item = Result<Event, Self::Error>> + Send {
474 T::summarize(self, text)
475 }
476
477 fn categorize<U: JsonSchema + DeserializeOwned + 'static>(
478 &self,
479 text: &str,
480 ) -> impl Future<Output = Result<U, GenerateError<Self::Error>>> + Send {
481 T::categorize(self, text)
482 }
483
484 fn profile(&self) -> impl Future<Output = Result<Profile, Self::Error>> + Send {
485 T::profile(self)
486 }
487}
488
489mod prompts;
490
491impl_language_model!(Arc, Box);
492
493pub async fn collect_text<S, E>(stream: S) -> Result<String, E>
499where
500 S: Stream<Item = Result<Event, E>>,
501{
502 pin!(stream);
503 let mut result = String::new();
504 while let Some(event) = stream.next().await {
505 if let Event::Text(text) = event? {
506 result.push_str(&text);
507 }
508 }
509 Ok(result)
510}
511
512async fn structured_generate<T: JsonSchema + DeserializeOwned + 'static, M: LanguageModel>(
513 model: &M,
514 mut request: LLMRequest,
515) -> Result<T, GenerateError<M::Error>> {
516 let schema = schema_for!(T);
517
518 let json = if schema.as_value().is_string() {
520 let stream = model.respond(request);
521 let response = collect_text(stream)
522 .await
523 .map_err(GenerateError::Provider)?;
524 serde_json::to_string(&response).map_err(|source| GenerateError::Parse {
526 source,
527 response: response.clone(),
528 })?
529 } else {
530 let schema =
533 serde_json::to_string_pretty(&schema).map_err(|source| GenerateError::Parse {
534 source,
535 response: String::new(),
536 })?;
537 let prompt = prompts::generate(&schema);
538 request.messages.push(Message::system(prompt));
539 request.parameters.structured_outputs = true;
540
541 let stream = model.respond(request);
542 collect_text(stream)
543 .await
544 .map_err(GenerateError::Provider)?
545 };
546
547 parse_json_with_recovery(&json).map_err(|source| GenerateError::Parse {
548 source,
549 response: truncate_for_error(&json),
550 })
551}
552
553fn truncate_for_error(response: &str) -> String {
555 const LIMIT: usize = 500;
556 response.chars().take(LIMIT).collect()
557}
558
559pub fn oneshot(system: impl Into<String>, user: impl Into<String>) -> LLMRequest {
561 let messages = vec![Message::system(system.into()), Message::user(user.into())];
562 LLMRequest::new(messages)
563}
564
565fn summarize<M: LanguageModel>(
566 model: &M,
567 text: &str,
568) -> impl Stream<Item = Result<Event, M::Error>> + Send {
569 let messages = oneshot("Summarize text:", text);
570 model.respond(messages)
571}
572
573async fn categorize_text<T: JsonSchema + DeserializeOwned + 'static, M: LanguageModel>(
574 model: &M,
575 text: &str,
576) -> Result<T, GenerateError<M::Error>> {
577 let request = oneshot("Categorize text by provided schema", text);
578 model.generate(request).await
579}
580
581fn parse_json_with_recovery<T: DeserializeOwned + 'static>(
582 json: &str,
583) -> Result<T, serde_json::Error> {
584 use serde::de::Error as _;
585
586 let trimmed = json.trim();
587 let mut last_error: Option<serde_json::Error> = None;
588 let mut last_candidate: Option<String> = None;
589
590 for candidate in build_json_candidates(trimmed) {
591 match serde_json::from_str::<T>(&candidate) {
592 Ok(value) => return Ok(value),
593 Err(err) => {
594 last_error = Some(err);
595 last_candidate = Some(candidate);
596 }
597 }
598 }
599
600 if is_string_type::<T>()
603 && let Some(candidate) = last_candidate
604 && let Ok(value) = serde_json::from_str::<serde_json::Value>(&candidate)
605 {
606 let text = match value {
607 serde_json::Value::String(s) => s,
608 other => other.to_string(),
609 };
610 let encoded = serde_json::to_string(&text)?;
611 if let Ok(value) = serde_json::from_str::<T>(&encoded) {
612 return Ok(value);
613 }
614 }
615
616 Err(last_error.unwrap_or_else(|| {
617 serde_json::Error::custom("structured output was empty or missing a JSON block")
618 }))
619}
620
621fn strip_code_fences(raw: &str) -> Option<String> {
622 let trimmed = raw.trim();
623 let fence_start = trimmed.find("```")?;
624 let after_fence = &trimmed[fence_start + 3..];
625 let mut lines = after_fence.lines();
626 let _maybe_lang = lines.next();
627 let body = lines.collect::<Vec<_>>().join("\n");
628 let content = body.rfind("```").map_or(body.as_str(), |end| &body[..end]);
629
630 let cleaned = content.trim();
631 if cleaned.is_empty() {
632 None
633 } else {
634 Some(cleaned.to_string())
635 }
636}
637
638fn extract_json_block(raw: &str) -> Option<String> {
639 if let (Some(start), Some(end)) = (raw.find('{'), raw.rfind('}'))
640 && end >= start
641 {
642 let candidate = &raw[start..=end];
643 if !candidate.trim().is_empty() {
644 return Some(candidate.trim().to_string());
645 }
646 }
647 if let (Some(start), Some(end)) = (raw.find('['), raw.rfind(']'))
648 && end >= start
649 {
650 let candidate = &raw[start..=end];
651 if !candidate.trim().is_empty() {
652 return Some(candidate.trim().to_string());
653 }
654 }
655 None
656}
657
658fn build_json_candidates(raw: &str) -> Vec<String> {
659 let mut candidates = Vec::new();
660
661 if !raw.is_empty() {
662 candidates.push(raw.to_string());
663 }
664
665 if let Some(fenced) = strip_code_fences(raw) {
666 candidates.push(fenced);
667 }
668
669 if let Some(block) = extract_json_block(raw) {
670 candidates.push(block);
671 }
672
673 if let Some(dequoted) = dequote_json_string(raw) {
674 candidates.push(dequoted);
675 }
676
677 if let Some(stripped) = strip_leading_label(raw, "json") {
678 candidates.push(stripped);
679 }
680
681 let mut deduped = Vec::new();
682 for candidate in candidates {
683 if deduped.iter().all(|seen| seen != &candidate) {
684 deduped.push(candidate);
685 }
686 }
687 deduped
688}
689
690fn dequote_json_string(raw: &str) -> Option<String> {
691 let trimmed = raw.trim();
692 if !(trimmed.starts_with('"') && trimmed.ends_with('"')) {
693 return None;
694 }
695 let inner: String = serde_json::from_str(trimmed).ok()?;
696 if inner.trim().is_empty() {
697 None
698 } else {
699 Some(inner)
700 }
701}
702
703fn strip_leading_label(raw: &str, label: &str) -> Option<String> {
704 let trimmed = raw.trim_start();
705 if !trimmed.to_ascii_lowercase().starts_with(label) {
706 return None;
707 }
708 let stripped = trimmed[label.len()..]
709 .trim_start_matches(|c: char| c.is_whitespace() || c == ':' || c == '-')
710 .trim();
711 if stripped.is_empty() {
712 None
713 } else {
714 Some(stripped.to_string())
715 }
716}
717
718fn is_string_type<T: 'static>() -> bool {
719 TypeId::of::<T>() == TypeId::of::<String>()
720}
721
722#[cfg(test)]
723mod tests {
724 use super::parse_json_with_recovery;
725 use alloc::string::String;
726 use serde::Deserialize;
727
728 #[derive(Debug, Deserialize, PartialEq, Eq)]
729 struct Foo {
730 a: u8,
731 }
732
733 #[test]
734 fn parses_plain_json() {
735 let foo: Foo = parse_json_with_recovery(r#"{"a":1}"#).unwrap();
736 assert_eq!(foo, Foo { a: 1 });
737 }
738
739 #[test]
740 fn parses_code_fence_json() {
741 let foo: Foo = parse_json_with_recovery("```json\n{\"a\":2}\n```").unwrap();
742 assert_eq!(foo, Foo { a: 2 });
743 }
744
745 #[test]
746 fn parses_embedded_block() {
747 let foo: Foo = parse_json_with_recovery("noise {\"a\":3} trailing").unwrap();
748 assert_eq!(foo, Foo { a: 3 });
749 }
750
751 #[test]
752 fn parses_quoted_json_string() {
753 let foo: Foo = parse_json_with_recovery(r#""{\"a\":4}""#).unwrap();
754 assert_eq!(foo, Foo { a: 4 });
755 }
756
757 #[test]
758 fn parses_labeled_json() {
759 let foo: Foo = parse_json_with_recovery("json {\"a\":5}").unwrap();
760 assert_eq!(foo, Foo { a: 5 });
761 }
762
763 #[test]
764 fn coerces_object_to_string() {
765 let value: String =
766 parse_json_with_recovery(r#"{"title":"summary","type":"content"}"#).unwrap();
767 assert!(
768 value.contains("\"title\":\"summary\"") && value.contains("\"type\":\"content\""),
769 "unexpected value: {value}"
770 );
771 }
772}