llm_daemon/
daemon_trait.rs

1use futures::Future;
2use tracing::{debug, info};
3use url::Url;
4
5pub trait LlmConfig {
6    fn endpoint(&self) -> Url;
7    fn health_url(&self) -> Url;
8}
9
10/// Represents a generic daemon capable of performing background tasks, including spawning itself,
11/// maintaining a heartbeat, and generating responses based on prompts.
12pub trait LlmDaemon {
13    type Config: LlmConfig;
14
15    fn config(&self) -> &Self::Config;
16
17    /// Spawns the daemon, initializing any necessary resources or processes.
18    /// This method is expected to be called before creation of tokio runtime, mostly
19    /// due to the use of the `fork`. User is free to use async runtime after
20    /// calling this.
21    fn fork_daemon(&self) -> anyhow::Result<()>;
22
23    /// Creates a task which maintains a periodic heartbeat to the daemon.
24    /// Daemon is expected to terminate if there's no heartbeat for a certain period of time.
25    /// Keeping this task within async runtime will ensure that the daemon is kept running
26    /// during the application.
27    /// FIXME: Extract heartbeat / ready to Async util?
28    fn heartbeat<'a, 'b>(
29        &'b self,
30    ) -> impl Future<Output = anyhow::Result<()>> + Send + 'a
31    where
32        'a: 'b;
33
34    /// Send a single heartbeat to let daemon know there is a client.
35    /// Can be useful if you don't want to use heartbeat.
36    fn ping(&self) -> anyhow::Result<()>;
37
38    /// FIXME: Extract heartbeat / ready to Async util?
39    fn ready<'a>(&self) -> impl Future<Output = ()> + Send + 'a {
40        let client = reqwest::Client::new();
41        let endpoint = self.config().health_url().clone();
42        async move {
43            loop {
44                debug!("Checking healthcheck endpoint: {}", endpoint.as_str());
45                let res = client.get(endpoint.as_str()).send().await;
46                info!("{:?}", &res);
47                match res {
48                    Ok(x) if x.status().is_success() => {
49                        break;
50                    },
51                    _ => {
52                        tokio::time::sleep(tokio::time::Duration::from_secs(1))
53                            .await;
54                    },
55                }
56            }
57        }
58    }
59}