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

kproc_llm/
lib.rs
