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                let created = graph
26                    .created
27                    .duration_since(std::time::UNIX_EPOCH)
28                    .unwrap()
29                    .as_secs();
30                models.push(Model {
31                    id: name.clone(),
32                    created,
33                    object: String::from("model"),
34                    owned_by: String::from("Not specified"),
35                });
36            }
37        }
38    }
39
40    {
41        if let Some(embedding_graphs) = EMBEDDING_GRAPHS.get() {
42            let embedding_graphs = embedding_graphs.lock().map_err(|e| {
43                LlamaCoreError::Operation(format!(
44                    "Fail to acquire the lock of `EMBEDDING_GRAPHS`. {e}"
45                ))
46            })?;
47
48            if !embedding_graphs.is_empty() {
49                for (name, graph) in embedding_graphs.iter() {
50                    let created = graph
51                        .created
52                        .duration_since(std::time::UNIX_EPOCH)
53                        .unwrap()
54                        .as_secs();
55                    models.push(Model {
56                        id: name.clone(),
57                        created,
58                        object: String::from("model"),
59                        owned_by: String::from("Not specified"),
60                    });
61                }
62            }
63        }
64    }
65
66    Ok(ListModelsResponse {
67        object: String::from("list"),
68        data: models,
69    })
70}