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