docling_rag/llm/
openrouter.rs1use super::{ChatModel, Message};
5use crate::{RagConfig, RagError, Result};
6use async_trait::async_trait;
7use serde::{Deserialize, Serialize};
8
9#[derive(Debug, Clone)]
11pub struct OpenRouterClient {
12 client: reqwest::Client,
13 base_url: String,
14 api_key: String,
15 model: String,
16}
17
18#[derive(Serialize)]
19struct ChatReq<'a> {
20 model: &'a str,
21 messages: &'a [Message],
22 temperature: f32,
23}
24
25#[derive(Deserialize)]
26struct ChatResp {
27 #[serde(default)]
28 choices: Vec<Choice>,
29}
30
31#[derive(Deserialize)]
32struct Choice {
33 message: RespMessage,
34}
35
36#[derive(Deserialize)]
37struct RespMessage {
38 #[serde(default)]
39 content: String,
40}
41
42impl OpenRouterClient {
43 pub fn from_config(cfg: &RagConfig) -> Result<Self> {
45 let api_key = cfg.openrouter_api_key.clone().ok_or_else(|| {
46 RagError::config("OPENROUTER_API_KEY is required for LLM-backed retrieval/synthesis")
47 })?;
48 Ok(OpenRouterClient {
49 client: reqwest::Client::new(),
50 base_url: cfg.openrouter_base_url.trim_end_matches('/').to_string(),
51 api_key,
52 model: cfg.llm_model.clone(),
53 })
54 }
55}
56
57#[async_trait]
58impl ChatModel for OpenRouterClient {
59 async fn complete(&self, messages: &[Message]) -> Result<String> {
60 let url = format!("{}/chat/completions", self.base_url);
61 tracing::debug!(url = %url, model = %self.model, "llm chat request");
62 let resp = self
63 .client
64 .post(&url)
65 .bearer_auth(&self.api_key)
66 .header(
68 "HTTP-Referer",
69 "https://github.com/docling-project/docling.rs",
70 )
71 .header("X-Title", "docling-rag")
72 .json(&ChatReq {
73 model: &self.model,
74 messages,
75 temperature: 0.2,
76 })
77 .send()
78 .await?;
79 let status = resp.status();
82 if !status.is_success() {
83 let body = resp.text().await.unwrap_or_default();
84 let snippet: String = body.chars().take(400).collect();
85 let hint = if status.as_u16() == 401 {
86 " (hint: OpenRouter keys start with 'sk-or-'; for a native \
87 DeepSeek key set OPENROUTER_BASE_URL=https://api.deepseek.com \
88 and RAG_LLM_MODEL=deepseek-chat)"
89 } else {
90 ""
91 };
92 return Err(RagError::Llm(format!(
93 "{url} returned {status}: {snippet}{hint}"
94 )));
95 }
96 let body: ChatResp = resp.json().await?;
97 body.choices
98 .into_iter()
99 .next()
100 .map(|c| c.message.content)
101 .ok_or_else(|| RagError::Llm("llm returned no choices".into()))
102 }
103}