docling_rag/queue/
mod.rs

1//! Pluggable message queue for decoupling document discovery from ingestion.
2//!
3//! The default is an in-process [`memory`] channel; RabbitMQ and Redis pub/sub are
4//! available behind the `rabbitmq` and `redis` features. Payloads are opaque byte
5//! blobs (the pipeline publishes JSON-encoded [`crate::source::SourceRef`]s).
6
7pub mod memory;
8
9#[cfg(feature = "rabbitmq")]
10pub mod rabbitmq;
11#[cfg(feature = "redis")]
12pub mod redis;
13
14use crate::config::QueueKind;
15use crate::{RagConfig, RagError, Result};
16use async_trait::async_trait;
17use std::sync::Arc;
18
19/// The queue/topic name used across backends.
20pub const TOPIC: &str = "docling_rag_ingest";
21
22/// A subscription that yields published payloads until the queue closes.
23#[async_trait]
24pub trait QueueReceiver: Send {
25    /// Await the next payload; `None` when the queue is closed / drained.
26    async fn recv(&mut self) -> Option<Vec<u8>>;
27}
28
29/// A pluggable publish/subscribe message queue.
30#[async_trait]
31pub trait MessageQueue: Send + Sync {
32    /// Publish one payload to all current subscribers.
33    async fn publish(&self, payload: &[u8]) -> Result<()>;
34
35    /// Open a new subscription.
36    async fn subscribe(&self) -> Result<Box<dyn QueueReceiver>>;
37}
38
39/// Build the queue selected by `cfg.queue`.
40pub async fn from_config(cfg: &RagConfig) -> Result<Arc<dyn MessageQueue>> {
41    match cfg.queue {
42        QueueKind::Memory => Ok(Arc::new(memory::MemoryQueue::new())),
43        QueueKind::RabbitMq => {
44            #[cfg(feature = "rabbitmq")]
45            {
46                let url = cfg.rabbitmq_url.clone().ok_or_else(|| {
47                    RagError::config("RABBITMQ_URL is required for the rabbitmq queue")
48                })?;
49                Ok(Arc::new(rabbitmq::RabbitMqQueue::connect(&url).await?))
50            }
51            #[cfg(not(feature = "rabbitmq"))]
52            {
53                Err(RagError::FeatureDisabled(
54                    "rabbitmq".into(),
55                    "rabbitmq".into(),
56                ))
57            }
58        }
59        QueueKind::Redis => {
60            #[cfg(feature = "redis")]
61            {
62                let url = cfg
63                    .redis_url
64                    .clone()
65                    .ok_or_else(|| RagError::config("REDIS_URL is required for the redis queue"))?;
66                Ok(Arc::new(redis::RedisQueue::connect(&url).await?))
67            }
68            #[cfg(not(feature = "redis"))]
69            {
70                Err(RagError::FeatureDisabled("redis".into(), "redis".into()))
71            }
72        }
73    }
74}
docling_rag/queue/mod.rs

docling_rag/queue/
mod.rs
