chatty-rs 0.0.1-alpha1

A terminal-based chat client for OpenAI's GPT models.
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

#[cfg(test)]
#[path = "manager_test.rs"]
mod tests;

use crate::backend::{ArcBackend, Backend};
use crate::models::{ArcEventTx, BackendPrompt, Model};
use async_trait::async_trait;
use eyre::{Context, Result, bail};
use std::collections::HashMap;

#[derive(Default)]
pub struct Manager {
    connections: HashMap<String, ArcBackend>, /* Alias - Backend */
    models: HashMap<String, String>,          /* Model ID - Alias  */
}

impl Manager {
    pub fn len(&self) -> usize {
        self.connections.len()
    }

    pub fn is_empty(&self) -> bool {
        self.connections.is_empty()
    }

    pub async fn add_connection(&mut self, connection: ArcBackend) -> eyre::Result<()> {
        let alias = connection.name().to_string();

        if self.connections.contains_key(&alias) {
            bail!(format!("connection {} already exists", alias))
        }

        connection
            .list_models()
            .await
            .wrap_err(format!("listing models backend {}", alias))?
            .into_iter()
            .for_each(|m| {
                self.models.insert(m.id().to_string(), alias.clone());
            });

        self.connections.insert(alias, connection);
        Ok(())
    }

    pub fn get_connection(&self, model: &str) -> Option<&ArcBackend> {
        self.connections.get(self.models.get(model)?)
    }
}

#[async_trait]
impl Backend for Manager {
    fn name(&self) -> &str {
        "Manager"
    }

    async fn list_models(&self) -> Result<Vec<Model>> {
        Ok(self
            .models
            .iter()
            .map(|(id, alias)| Model::new(id).with_provider(alias))
            .collect())
    }

    async fn get_completion(&self, prompt: BackendPrompt, event_tx: ArcEventTx) -> Result<()> {
        let connection = match self.get_connection(prompt.model()) {
            Some(connection) => connection,
            None => {
                return Err(eyre::eyre!("model is not available"));
            }
        };
        connection
            .get_completion(prompt, event_tx)
            .await
            .wrap_err(format!("get completion from backend {}", connection.name()))?;
        Ok(())
    }
}