Skip to main content

docling_rag/llm/
mod.rs

1//! Pluggable chat/LLM provider, used for Multi-Query and HyDE retrieval and for
2//! final answer synthesis. The only shipped backend is [`openrouter`].
3
4pub mod openrouter;
5
6use crate::{RagConfig, Result};
7use async_trait::async_trait;
8use serde::{Deserialize, Serialize};
9use std::sync::Arc;
10
11/// A single chat message.
12#[derive(Debug, Clone, Serialize, Deserialize)]
13pub struct Message {
14    /// `"system"`, `"user"`, or `"assistant"`.
15    pub role: String,
16    /// The message text.
17    pub content: String,
18}
19
20impl Message {
21    /// A system message.
22    pub fn system(content: impl Into<String>) -> Self {
23        Message {
24            role: "system".into(),
25            content: content.into(),
26        }
27    }
28    /// A user message.
29    pub fn user(content: impl Into<String>) -> Self {
30        Message {
31            role: "user".into(),
32            content: content.into(),
33        }
34    }
35}
36
37/// A chat completion model.
38#[async_trait]
39pub trait ChatModel: Send + Sync {
40    /// Complete a chat conversation, returning the assistant's reply text.
41    async fn complete(&self, messages: &[Message]) -> Result<String>;
42
43    /// Convenience: single system + user turn.
44    async fn ask(&self, system: &str, user: &str) -> Result<String> {
45        self.complete(&[Message::system(system), Message::user(user)])
46            .await
47    }
48}
49
50/// Build the configured chat model. Errors if the provider needs credentials that
51/// are not set (`OPENROUTER_API_KEY`).
52pub fn from_config(cfg: &RagConfig) -> Result<Arc<dyn ChatModel>> {
53    Ok(Arc::new(openrouter::OpenRouterClient::from_config(cfg)?))
54}