openai_api_dispatch/worker/
mod.rs1use 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"))]
16pub type NatsOpenaiWorker = Worker<NatsWorker, ExecutorAsyncOpenai>;
18
19#[derive(Debug, Clone)]
20pub 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 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 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}