llama_core/
models.rs

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
//! Define APIs for querying models.

use crate::{error::LlamaCoreError, CHAT_GRAPHS, EMBEDDING_GRAPHS};
use endpoints::models::{ListModelsResponse, Model};

/// Lists models available
pub async fn models() -> Result<ListModelsResponse, LlamaCoreError> {
    #[cfg(feature = "logging")]
    info!(target: "stdout", "List models");

    let mut models = vec![];

    {
        if let Some(chat_graphs) = CHAT_GRAPHS.get() {
            let chat_graphs = chat_graphs.lock().map_err(|e| {
                let err_msg = format!("Fail to acquire the lock of `CHAT_GRAPHS`. {}", e);

                #[cfg(feature = "logging")]
                error!(target: "stdout", "{}", &err_msg);

                LlamaCoreError::Operation(err_msg)
            })?;

            for (name, graph) in chat_graphs.iter() {
                models.push(Model {
                    id: name.clone(),
                    created: graph.created.as_secs(),
                    object: String::from("model"),
                    owned_by: String::from("Not specified"),
                });
            }
        }
    }

    {
        if let Some(embedding_graphs) = EMBEDDING_GRAPHS.get() {
            let embedding_graphs = embedding_graphs.lock().map_err(|e| {
                LlamaCoreError::Operation(format!(
                    "Fail to acquire the lock of `EMBEDDING_GRAPHS`. {}",
                    e
                ))
            })?;

            if !embedding_graphs.is_empty() {
                for (name, graph) in embedding_graphs.iter() {
                    models.push(Model {
                        id: name.clone(),
                        created: graph.created.as_secs(),
                        object: String::from("model"),
                        owned_by: String::from("Not specified"),
                    });
                }
            }
        }
    }

    Ok(ListModelsResponse {
        object: String::from("list"),
        data: models,
    })
}