llmvm_core_lib/
lib.rs

1//! The core components for [llmvm](https://github.com/djandries/llmvm).
2//!
3//! Assumes the following responsibilities:
4//! - Sending generation requests to backends
5//! - Managing message threads
6//! - Model presets
7//! - Prompt templates
8//! - Projects
9
10pub mod error;
11pub mod generation;
12pub mod presets;
13pub mod prompts;
14pub mod service;
15mod threads;
16
17use llmvm_protocol::http::client::{HttpClient, HttpClientConfig};
18use llmvm_protocol::service::{BackendRequest, BackendResponse};
19use llmvm_protocol::stdio::client::{StdioClient, StdioClientConfig};
20use llmvm_protocol::{BoxedService, ConfigExampleSnippet, ModelDescription};
21use serde::Deserialize;
22use threads::clean_old_threads;
23
24use std::collections::HashMap;
25use tokio::sync::Mutex;
26use tracing::debug;
27
28use crate::error::CoreError;
29
30const BACKEND_COMMAND_PREFIX: &str = "llmvm-";
31const PROJECT_DIR_NAME: &str = ".llmvm";
32const DEFAULT_THREAD_TTL_SECS: u64 = 14 * 24 * 3600;
33
34pub type Result<T> = std::result::Result<T, CoreError>;
35
36/// Configuration structure for the llmvm core.
37#[derive(Deserialize, Default)]
38#[serde(default)]
39pub struct LLMVMCoreConfig {
40    /// Configuration for all backend stdio clients.
41    pub stdio_client: Option<StdioClientConfig>,
42    /// The max time-to-live for threads in seconds.
43    /// Threads with a last modified time older than the TTL will be removed.
44    pub thread_ttl_secs: Option<u64>,
45    /// HTTP client configurations for remote backends.
46    pub http_backends: HashMap<String, HttpClientConfig>,
47}
48
49impl ConfigExampleSnippet for LLMVMCoreConfig {
50    fn config_example_snippet() -> String {
51        format!(
52            r#"# TTL in seconds for message threads (based on last modified time)
53# thread_ttl_secs = 604800
54
55# Stdio client configuration
56# [stdio_client]
57{}
58
59# HTTP backend client config
60# [http_backends.llmrs-http]
61{}"#,
62            StdioClientConfig::config_example_snippet(),
63            HttpClientConfig::config_example_snippet(),
64        )
65    }
66}
67
68/// The llmvm core which handles sending generation requests to the backend,
69/// managing message threads, model presets, and prompt templates. Frontends
70/// request text generation by invoking the core.
71pub struct LLMVMCore {
72    clients: Mutex<HashMap<String, BoxedService<BackendRequest, BackendResponse>>>,
73    config: LLMVMCoreConfig,
74}
75
76impl LLMVMCore {
77    pub async fn new(mut config: LLMVMCoreConfig) -> Result<Self> {
78        clean_old_threads(config.thread_ttl_secs.unwrap_or(DEFAULT_THREAD_TTL_SECS)).await?;
79
80        let mut clients: HashMap<String, BoxedService<BackendRequest, BackendResponse>> =
81            HashMap::new();
82
83        for (name, config) in config.http_backends.drain() {
84            debug!("loading {} http backend", name);
85            clients.insert(
86                name,
87                Box::new(HttpClient::new(config).map_err(|_| CoreError::HttpServiceCreate)?),
88            );
89        }
90
91        Ok(Self {
92            clients: Mutex::new(clients),
93            config,
94        })
95    }
96
97    async fn get_client<'a>(
98        &self,
99        clients_guard: &'a mut HashMap<String, BoxedService<BackendRequest, BackendResponse>>,
100        model_description: &ModelDescription,
101    ) -> Result<&'a mut BoxedService<BackendRequest, BackendResponse>> {
102        let command = format!("{}{}", BACKEND_COMMAND_PREFIX, model_description.backend);
103        if !clients_guard.contains_key(&model_description.backend) {
104            debug!(
105                "starting backend {command} in {:?}",
106                self.config.stdio_client.as_ref().map(|c| &c.bin_path)
107            );
108            let backend = model_description.backend.as_str();
109            clients_guard.insert(
110                backend.to_string(),
111                Box::new(
112                    StdioClient::new(
113                        &command,
114                        &["--log-to-file"],
115                        self.config.stdio_client.clone().unwrap_or_default(),
116                    )
117                    .await
118                    .map_err(|e| CoreError::StdioBackendStart(e))?,
119                ),
120            );
121        }
122        Ok(clients_guard.get_mut(&model_description.backend).unwrap())
123    }
124
125    /// Close the active stdio (local child process) or HTTP client (remote process)
126    /// for a given backend. If the backend client is over stdio, the backend child process
127    /// will be killed.
128    pub async fn close_client(&self, backend_name: &str) {
129        self.clients.lock().await.remove(backend_name);
130    }
131}