kproc_llm/
lib.rs

1#![doc = include_str!("../README.MD")]
2#![warn(missing_docs)]
3#[cfg_attr(not(any(feature = "ollama", feature = "llama.cpp")), deny(warnings))]
4mod error;
5pub mod prelude;
6
7#[cfg(feature = "llama.cpp")]
8pub mod llama_cpp;
9#[cfg(feature = "ollama")]
10pub mod ollama;
11
12#[cfg(feature = "template")]
13pub mod template;
14
15use std::{future::Future, pin::Pin};
16
17/// Export Error enum.
18pub use error::Error;
19use futures::Stream;
20
21/// Export Result type.
22pub type Result<T> = std::result::Result<T, Error>;
23
24/// Format
25pub enum Format
26{
27  /// Output text
28  Text,
29  /// Output Json
30  Json,
31}
32
33/// Prompt
34pub struct Prompt
35{
36  /// User prompt
37  pub prompt: String,
38  /// Assistant message
39  pub assistant: Option<String>,
40  /// System message
41  pub system: Option<String>,
42  /// Requested output format
43  pub format: Format,
44  /// Optional prompt image (for multi modal models)
45  #[cfg(feature = "image")]
46  pub image: Option<kproc_values::Image>,
47}
48
49impl Prompt
50{
51  /// Create a new prompt, from the given string.
52  pub fn prompt(prompt: impl Into<String>) -> Self
53  {
54    let prompt = prompt.into();
55    Self {
56      prompt,
57      assistant: None,
58      system: None,
59      format: Format::Text,
60      #[cfg(feature = "image")]
61      image: None,
62    }
63  }
64  /// Set the system hint.
65  pub fn system(mut self, system: impl Into<String>) -> Self
66  {
67    self.system = Some(system.into());
68    self
69  }
70  /// Set the result format
71  pub fn format(mut self, format: impl Into<Format>) -> Self
72  {
73    self.format = format.into();
74    self
75  }
76  /// Set an input image
77  #[cfg(feature = "image")]
78  pub fn image(mut self, image: impl Into<kproc_values::Image>) -> Self
79  {
80    self.image = Some(image.into());
81    self
82  }
83}
84
85/// String stream
86pub type StringStream = Pin<Box<dyn Stream<Item = Result<String>> + Send>>;
87
88/// Convenient function for coercing a box pin stream to a StringStream
89pub(crate) fn pin_stream<T: 'static + Send + Stream<Item = Result<String>>>(t: T) -> StringStream
90{
91  Box::pin(t)
92}
93
94/// LLM
95pub trait LargeLanguageModel
96{
97  /// Run inference on a model, returning a stream
98  fn infer_stream(
99    &self,
100    prompt: Prompt,
101  ) -> Result<impl Future<Output = Result<StringStream>> + Send>;
102  /// Run inference on a model, return once the complete answer has been computed.
103  /// The default implementation call `infer_stream` until completion of the stream.
104  fn infer(&self, prompt: Prompt) -> Result<impl Future<Output = Result<String>> + Send>
105  {
106    use futures::stream::StreamExt;
107    let stream = self.infer_stream(prompt)?;
108    Ok(async {
109      let mut result: String = Default::default();
110      let mut stream = Box::pin(stream.await?);
111      while let Some(next_token) = stream.next().await
112      {
113        if result.is_empty()
114        {
115          result = next_token?;
116        }
117        else
118        {
119          result.push_str(&next_token?);
120        }
121      }
122      Ok(result)
123    })
124  }
125}
kproc_llm/lib.rs

kproc_llm/
lib.rs
