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

use futures::Future;
use url::Url;

pub trait LlmConfig {
    fn endpoint(&self) -> Url;
}

/// Represents a generic daemon capable of performing background tasks, including spawning itself,
/// maintaining a heartbeat, and generating responses based on prompts.
pub trait LlmDaemon {
    type Config: LlmConfig;
    
    fn config(&self) -> &Self::Config;

    /// Spawns the daemon, initializing any necessary resources or processes.
    /// This method is expected to be called before creation of tokio runtime, mostly
    /// due to the use of the `fork`. User is free to use async runtime after
    /// calling this.
    fn fork_daemon(&self) -> anyhow::Result<()>;

    /// Creates a task which maintains a periodic heartbeat to the daemon.
    /// Daemon is expected to terminate if there's no heartbeat for a certain period of time.
    /// Keeping this task within async runtime will ensure that the daemon is kept running
    /// during the application.
    fn heartbeat(
        &self,
    ) -> impl Future<Output = anyhow::Result<()>> + Send + 'static;
    
    fn ready<'a>(&self) -> impl Future<Output = ()> + Send + 'a {
        let client = reqwest::Client::new();
        let endpoint = self.config().endpoint().clone();
        async move {
            for _ in 0..300 {
                let res = client.get(endpoint.as_str())
                    .send()
                    .await;
                match res {
                    Ok(_) => {
                        break;
                    },
                    Err(_) => {
                        tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
                    }
                }
            }
        }
    }
}