llm 1.3.8

A Rust library unifying multiple LLM backends.
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

use std::collections::HashMap;
use std::sync::Arc;

use tokio::sync::mpsc;
use tokio_util::sync::CancellationToken;

use llm::chat::ChatMessage;
use llm::LLMProvider;

use crate::conversation::{ConversationId, MessageId};
use crate::provider::ProviderCapabilities;
use crate::runtime::AppEvent;

use super::runner::run_stream;

#[derive(Clone)]
pub struct StreamRequest {
    pub conversation_id: ConversationId,
    pub message_id: MessageId,
    pub provider: Arc<dyn LLMProvider>,
    pub messages: Vec<ChatMessage>,
    pub capabilities: ProviderCapabilities,
}

pub struct StreamManager {
    sender: mpsc::Sender<AppEvent>,
    active: HashMap<ConversationId, CancellationToken>,
}

impl StreamManager {
    pub fn new(sender: mpsc::Sender<AppEvent>) -> Self {
        Self {
            sender,
            active: HashMap::new(),
        }
    }

    pub fn cancel(&mut self, id: ConversationId) {
        if let Some(token) = self.active.remove(&id) {
            token.cancel();
        }
    }

    pub fn start(&mut self, request: StreamRequest) -> CancellationToken {
        let token = CancellationToken::new();
        self.active.insert(request.conversation_id, token.clone());
        let sender = self.sender.clone();
        tokio::spawn(run_stream(request, sender, token.clone()));
        token
    }
}