kproc-llm 0.7.0

Knowledge Processing library, using LLMs.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107

#![doc = include_str!("../README.MD")]
#![warn(missing_docs)]

use std::future::Future;

pub mod prelude;

#[cfg_attr(not(any(feature = "ollama", feature = "llama.cpp")), deny(warnings))]
mod error;
mod message;
mod prompts;

#[cfg(any(feature = "candle", feature = "candle-git"))]
pub mod candle;

#[cfg(feature = "llama.cpp")]
pub mod llama_cpp;

#[cfg(feature = "ollama")]
pub mod ollama;

#[cfg(feature = "simple-api")]
pub mod simple_api;

#[cfg(feature = "template")]
pub mod template;

pub use error::Error;
pub use message::{Message, Role};
pub use prompts::{ChatPrompt, Format, GenerationPrompt};

/// Export Result type.
pub type Result<T, E = Error> = std::result::Result<T, E>;

/// String stream
pub type StringStream = ccutils::streams::BoxedStream<Result<String>>;

fn accumulate<T>(stream_maker: T) -> Result<impl Future<Output = Result<String>> + Send>
where
  T: Future<Output = Result<StringStream>> + Send,
{
  use futures::stream::StreamExt;
  Ok(async {
    let mut result: String = Default::default();
    let mut stream = Box::pin(stream_maker.await?);
    while let Some(next_token) = stream.next().await
    {
      if result.is_empty()
      {
        result = next_token?;
      }
      else
      {
        result.push_str(&next_token?);
      }
    }
    Ok(result)
  })
}

#[allow(dead_code)]
pub(crate) fn generate_with_chat<LLM>(
  llm: &LLM,
  prompt: GenerationPrompt,
) -> Result<impl Future<Output = Result<StringStream>> + Send + use<'_, LLM>>
where
  LLM: LargeLanguageModel,
{
  let chat_prompt = ChatPrompt::default()
    .system_opt(prompt.system)
    .assistant_opt(prompt.assistant)
    .user(prompt.user)
    .options(prompt.options);
  llm.chat_stream(chat_prompt)
}

/// LLM
pub trait LargeLanguageModel
{
  /// Chat with a model, returning a stream
  fn chat_stream(
    &self,
    prompt: ChatPrompt,
  ) -> Result<impl Future<Output = Result<StringStream>> + Send>;
  /// Run chat on a model, return once the complete answer has been computed.
  /// The default implementation call `chat_stream` until completion of the stream.
  fn chat(&self, prompt: ChatPrompt) -> Result<impl Future<Output = Result<String>> + Send>
  {
    let stream = self.chat_stream(prompt)?;
    accumulate(stream)
  }
  /// Run inference on a model, returning a stream
  fn generate_stream(
    &self,
    prompt: GenerationPrompt,
  ) -> Result<impl Future<Output = Result<StringStream>> + Send>;
  /// Run inference on a model, return once the complete answer has been computed.
  /// The default implementation call `infer_stream` until completion of the stream.
  fn generate(
    &self,
    prompt: GenerationPrompt,
  ) -> Result<impl Future<Output = Result<String>> + Send>
  {
    let stream = self.generate_stream(prompt)?;
    accumulate(stream)
  }
}