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