orch 0.0.16

Language model orchestration library
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

#![allow(dead_code)]

use std::pin::Pin;

use async_trait::async_trait;
use dyn_clone::DynClone;
use tokio_stream::Stream;

use super::{error::LanguageModelError, LanguageModelProvider};

/// A trait for language model providers which implements text completion, embeddings, etc.
///
/// > `DynClone` is used so that there can be dynamic dispatch of the `Llm` trait,
/// > especially needed for [magic-cli](https://github.com/guywaldman/magic-cli).
#[async_trait]
pub trait LanguageModel: DynClone + Send + Sync {
    /// Generates a response from the LLM.
    ///
    /// # Arguments
    /// * `prompt` - The prompt to generate a response for.
    /// * `system_prompt` - The system prompt to use for the generation.
    /// * `options` - The options for the generation.
    ///
    /// # Returns
    /// A [Result] containing the response from the LLM or an error if there was a problem.
    ///
    async fn text_complete(
        &self,
        prompt: &str,
        system_prompt: &str,
        options: TextCompleteOptions,
    ) -> Result<TextCompleteResponse, LanguageModelError>;

    /// Generates a streaming response from the LLM.
    ///
    /// # Arguments
    /// * `prompt` - The prompt to generate a response for.
    /// * `system_prompt` - The system prompt to use for the generation.
    /// * `options` - The options for the generation.
    ///
    /// # Returns
    /// A [Result] containing the response from the LLM or an error if there was a problem.
    ///
    async fn text_complete_stream(
        &self,
        prompt: &str,
        system_prompt: &str,
        options: TextCompleteStreamOptions,
    ) -> Result<TextCompleteStreamResponse, LanguageModelError>;

    /// Generates an embedding from the LLM.
    ///
    /// # Arguments
    /// * `prompt` - The item to generate an embedding for.
    ///
    /// # Returns
    ///
    /// A [Result] containing the embedding or an error if there was a problem.
    async fn generate_embedding(&self, prompt: &str) -> Result<Vec<f32>, LanguageModelError>;

    /// Returns the provider of the LLM.
    fn provider(&self) -> LanguageModelProvider;

    /// Returns the name of the model used for text completions.
    fn text_completion_model_name(&self) -> String;

    /// Returns the name of the model used for embeddings.
    fn embedding_model_name(&self) -> String;
}

#[derive(Debug, Clone, Default)]
pub struct TextCompleteOptions {
    /// An encoding of the conversation used in this response, this can be sent in the next request to keep a conversational memory.
    /// This should be as returned from the previous response.
    pub context: Option<Vec<i64>>,
}

#[derive(Debug, Clone, Default)]
pub struct TextCompleteStreamOptions {
    pub context: Option<Vec<i64>>,
}

#[derive(Debug, Clone)]
pub struct TextCompleteResponse {
    pub text: String,
    // TODO: This is specific to Ollama, context looks differently for other LLM providers.
    pub context: Option<Vec<i64>>,
}

pub struct TextCompleteStreamResponse {
    pub stream: Pin<Box<dyn Stream<Item = Result<String, LanguageModelError>> + Send>>,
    // TODO: Handle context with streaming response.
    // pub context: Vec<i64>,
}