Skip to main content

sniff/
llm_impl.rs

1use crate::config::ResolvedConfig;
2use crate::env_value;
3use serde_json::Value;
4use std::env;
5use std::time::Duration;
6use tokio::sync::Semaphore;
7
8#[path = "llm_call.rs"]
9mod llm_call;
10#[path = "llm_consensus.rs"]
11mod llm_consensus;
12#[path = "llm_content.rs"]
13mod llm_content;
14#[path = "llm_json.rs"]
15mod llm_json;
16#[path = "llm_payload.rs"]
17mod llm_payload;
18#[path = "llm_repair.rs"]
19mod llm_repair;
20#[path = "llm_response.rs"]
21mod llm_response;
22#[path = "llm_retry.rs"]
23mod llm_retry;
24#[path = "llm_schema.rs"]
25mod llm_schema;
26#[path = "llm_text.rs"]
27mod llm_text;
28#[path = "llm_transport.rs"]
29mod llm_transport;
30#[path = "llm_usage.rs"]
31mod llm_usage;
32
33#[derive(Debug, Clone, Copy, PartialEq, Eq)]
34pub enum ResponseSchema {
35    MethodReview,
36    SemanticMethodReview,
37    FileReview,
38    RoleClassification,
39}
40
41pub struct LLMClient {
42    config: ResolvedConfig,
43    api_key: Option<String>,
44    sem: Semaphore,
45}
46
47impl LLMClient {
48    pub fn try_new(config: ResolvedConfig, api_key: Option<String>) -> Result<Self, String> {
49        let api_key = api_key
50            .map(|value| env_value::normalize(&value))
51            .filter(|value| !value.is_empty());
52        let timeout = client_timeout();
53        build_http_client(timeout)?;
54
55        Ok(LLMClient {
56            config,
57            api_key,
58            sem: Semaphore::new(max_concurrency()),
59        })
60    }
61
62    pub fn new(config: ResolvedConfig, api_key: Option<String>) -> Self {
63        Self::try_new(config, api_key).expect("failed to build LLM HTTP client")
64    }
65
66    #[cfg(test)]
67    pub(crate) fn config(&self) -> &ResolvedConfig {
68        &self.config
69    }
70
71    pub(crate) fn review_context_key(&self) -> String {
72        format!(
73            "model={}\nendpoint={}\nsystem_context={}",
74            self.config.model, self.config.llm.endpoint, self.config.llm.system_context
75        )
76    }
77
78    pub async fn probe(&self) -> Result<(), String> {
79        let prompt = "Return exactly one JSON object with this shape: {\"role\":\"mixed\",\"reason\":\"probe\"}.";
80        match self
81            .call_once(prompt, ResponseSchema::RoleClassification)
82            .await
83        {
84            Ok((Some(_), _, _)) => Ok(()),
85            Ok((None, _, _)) => {
86                Err("LLM preflight failed: no valid JSON response after retries".to_string())
87            }
88            Err(err) => Err(format!("LLM preflight failed: {}", err)),
89        }
90    }
91
92    pub(super) async fn acquire_permit(&self) -> Result<tokio::sync::SemaphorePermit<'_>, String> {
93        self.sem.acquire().await.map_err(|e| e.to_string())
94    }
95
96    pub(super) async fn try_call_raw(
97        &self,
98        prompt: &str,
99    ) -> Result<(String, usize, usize), String> {
100        // Do not reuse a pooled connection across thousands of sequential
101        // reviews. A fresh client makes each retry independent of a stale
102        // provider socket and gives the watchdog a clean request boundary.
103        let client = build_http_client(client_timeout())?;
104        llm_response::try_call_raw(&client, &self.config, self.api_key.as_ref(), prompt).await
105    }
106
107    async fn call_once(
108        &self,
109        prompt: &str,
110        schema: ResponseSchema,
111    ) -> Result<(Option<Value>, usize, usize), String> {
112        llm_call::execute_call(self, prompt, schema).await
113    }
114
115    pub async fn call(
116        &self,
117        prompt: &str,
118        schema: ResponseSchema,
119    ) -> Result<(Option<serde_json::Value>, usize, usize), String> {
120        llm_consensus::call_with_consensus(self, prompt, schema).await
121    }
122
123    pub(crate) async fn call_single(
124        &self,
125        prompt: &str,
126        schema: ResponseSchema,
127    ) -> Result<(Option<serde_json::Value>, usize, usize), String> {
128        self.call_once(prompt, schema).await
129    }
130}
131
132fn max_concurrency() -> usize {
133    env::var("SNIFF_LLM_MAX_CONCURRENCY")
134        .or_else(|_| env::var("LLM_MAX_CONCURRENCY"))
135        .ok()
136        .and_then(|value| value.trim().parse::<usize>().ok())
137        .filter(|value| *value > 0)
138        .unwrap_or(5)
139}
140
141fn client_timeout() -> Duration {
142    env::var("SNIFF_LLM_CLIENT_TIMEOUT_SECS")
143        .or_else(|_| env::var("SNIFF_LLM_REQUEST_TIMEOUT_SECS"))
144        .or_else(|_| env::var("LLM_CLIENT_TIMEOUT_SECS"))
145        .ok()
146        .and_then(|value| value.trim().parse::<u64>().ok())
147        .filter(|secs| *secs > 0)
148        .map(Duration::from_secs)
149        .unwrap_or_else(|| Duration::from_secs(600))
150}
151
152fn build_http_client(timeout: Duration) -> Result<reqwest::Client, String> {
153    reqwest::Client::builder()
154        .http1_only()
155        .pool_max_idle_per_host(0)
156        .connect_timeout(timeout)
157        .read_timeout(timeout)
158        .timeout(timeout)
159        .build()
160        .map_err(|err| format!("failed to build LLM HTTP client: {err}"))
161}
162
163#[cfg(test)]
164#[path = "tests/llm_impl.rs"]
165mod tests;