Skip to main content

openai_api_dispatch/worker/
mod.rs

1use alloc::{string::String, sync::Arc};
2
3use crate::{
4    executor::Executor,
5    queue::{QueueWorker, WrappedTask},
6    task::{Response, Task},
7};
8
9#[cfg(test)]
10mod tests;
11
12#[cfg(all(feature = "nats-queue", feature = "std"))]
13use crate::{executor::openai::ExecutorAsyncOpenai, queue::nats::NatsWorker};
14
15#[cfg(all(feature = "nats-queue", feature = "std"))]
16/// A Core NATS consumer backed by the OpenAI-compatible chat executor.
17pub type NatsOpenaiWorker = Worker<NatsWorker, ExecutorAsyncOpenai>;
18
19#[derive(Debug, Clone)]
20/// Connects a queue consumer to an executor, processing one task at a time.
21///
22/// Only `NatsOpenaiWorker::from_env_or_default` currently has a public
23/// constructor (with `nats-queue` enabled). Other queue/executor combinations
24/// require a caller-managed loop using their respective traits.
25pub struct Worker<Q: QueueWorker, E: Executor> {
26    queue: Q,
27    executor: E,
28}
29
30#[cfg(all(feature = "nats-queue", feature = "std"))]
31impl NatsOpenaiWorker {
32    /// Builds the NATS consumer and API executor from their environment settings.
33    ///
34    /// Requires `OPENAI_API_DEFAULT_MODEL` and a reachable NATS server. See
35    /// [`NatsWorker::from_env_or_default`] and
36    /// [`ExecutorAsyncOpenai::from_env_or_default`] for defaults.
37    pub async fn from_env_or_default() -> anyhow::Result<Self> {
38        let queue = NatsWorker::from_env_or_default().await?;
39        let executor = ExecutorAsyncOpenai::from_env_or_default();
40
41        Ok(Self { queue, executor })
42    }
43}
44
45impl<Q: QueueWorker, E: Executor> Worker<Q, E> {
46    /// Processes tasks sequentially until the queue returns `None` or an error.
47    ///
48    /// Queue and executor errors stop the loop immediately; execution errors
49    /// are not converted into replies, and no application-level retry occurs.
50    /// Responses with `success == false` are sent normally and do not stop it.
51    /// A polling queue returning `None` ends the loop even if more work may arrive.
52    /// There is no shutdown signal; callers must arrange cancellation themselves.
53    pub async fn run(self) -> anyhow::Result<()> {
54        tracing::info!("openai-api worker subscribed to queue");
55
56        while let Some(WrappedTask { message, task }) = self.queue.receive_task().await? {
57            let id = task.id;
58
59            tracing::debug!("received task id `{id}`");
60
61            let response = self.executor.execute(task).await?;
62
63            tracing::debug!(
64                "responding task id `{id}` with success({})",
65                response.success
66            );
67
68            if !response.success {
69                tracing::debug!(
70                    "task id `{id}` responded with error `{}`",
71                    response.contents
72                );
73            }
74
75            self.queue.send_response(message, response).await?;
76
77            tracing::debug!("task id `{id}` response sent to queue");
78        }
79
80        tracing::info!("openai-api worker terminated");
81
82        Ok(())
83    }
84}
85
86impl<Q: QueueWorker, E: Executor> QueueWorker for Worker<Q, E> {
87    type Message = Q::Message;
88
89    async fn receive_task(&self) -> anyhow::Result<Option<WrappedTask<Self::Message>>> {
90        self.queue.receive_task().await
91    }
92
93    async fn send_response(
94        &self,
95        message: Self::Message,
96        response: Response,
97    ) -> anyhow::Result<()> {
98        self.queue.send_response(message, response).await
99    }
100}
101
102impl<Q: QueueWorker, E: Executor> Executor for Worker<Q, E> {
103    fn default_model(&self) -> &Arc<Option<String>> {
104        self.executor.default_model()
105    }
106
107    async fn execute(&self, task: Task) -> anyhow::Result<Response> {
108        self.executor.execute(task).await
109    }
110}