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}