llama_core/
models.rs

1//! Define APIs for querying models.
2
3use crate::{error::LlamaCoreError, CHAT_GRAPHS, EMBEDDING_GRAPHS};
4use endpoints::models::{ListModelsResponse, Model};
5
6/// Lists models available
7pub async fn models() -> Result<ListModelsResponse, LlamaCoreError> {
8    #[cfg(feature = "logging")]
9    info!(target: "stdout", "List models");
10
11    let mut models = vec![];
12
13    {
14        if let Some(chat_graphs) = CHAT_GRAPHS.get() {
15            let chat_graphs = chat_graphs.lock().map_err(|e| {
16                let err_msg = format!("Fail to acquire the lock of `CHAT_GRAPHS`. {}", e);
17
18                #[cfg(feature = "logging")]
19                error!(target: "stdout", "{}", &err_msg);
20
21                LlamaCoreError::Operation(err_msg)
22            })?;
23
24            for (name, graph) in chat_graphs.iter() {
25                models.push(Model {
26                    id: name.clone(),
27                    created: graph.created.as_secs(),
28                    object: String::from("model"),
29                    owned_by: String::from("Not specified"),
30                });
31            }
32        }
33    }
34
35    {
36        if let Some(embedding_graphs) = EMBEDDING_GRAPHS.get() {
37            let embedding_graphs = embedding_graphs.lock().map_err(|e| {
38                LlamaCoreError::Operation(format!(
39                    "Fail to acquire the lock of `EMBEDDING_GRAPHS`. {}",
40                    e
41                ))
42            })?;
43
44            if !embedding_graphs.is_empty() {
45                for (name, graph) in embedding_graphs.iter() {
46                    models.push(Model {
47                        id: name.clone(),
48                        created: graph.created.as_secs(),
49                        object: String::from("model"),
50                        owned_by: String::from("Not specified"),
51                    });
52                }
53            }
54        }
55    }
56
57    Ok(ListModelsResponse {
58        object: String::from("list"),
59        data: models,
60    })
61}