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    FileReview,
37    RoleClassification,
38}
39
40pub struct LLMClient {
41    config: ResolvedConfig,
42    api_key: Option<String>,
43    client: reqwest::Client,
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 client = reqwest::Client::builder()
53            .http1_only()
54            .timeout(client_timeout())
55            .build()
56            .map_err(|err| format!("failed to build LLM HTTP client: {err}"))?;
57
58        Ok(LLMClient {
59            config,
60            api_key,
61            client,
62            sem: Semaphore::new(max_concurrency()),
63        })
64    }
65
66    pub fn new(config: ResolvedConfig, api_key: Option<String>) -> Self {
67        Self::try_new(config, api_key).expect("failed to build LLM HTTP client")
68    }
69
70    pub(crate) fn config(&self) -> &ResolvedConfig {
71        &self.config
72    }
73
74    pub async fn probe(&self) -> Result<(), String> {
75        let prompt = "Return exactly one JSON object with this shape: {\"role\":\"mixed\",\"reason\":\"probe\"}.";
76        match self
77            .call_once(prompt, ResponseSchema::RoleClassification)
78            .await
79        {
80            Ok((Some(_), _, _)) => Ok(()),
81            Ok((None, _, _)) => {
82                Err("LLM preflight failed: no valid JSON response after retries".to_string())
83            }
84            Err(err) => Err(format!("LLM preflight failed: {}", err)),
85        }
86    }
87
88    pub(super) async fn acquire_permit(&self) -> Result<tokio::sync::SemaphorePermit<'_>, String> {
89        self.sem.acquire().await.map_err(|e| e.to_string())
90    }
91
92    pub(super) async fn try_call_raw(
93        &self,
94        prompt: &str,
95    ) -> Result<(String, usize, usize), String> {
96        llm_response::try_call_raw(&self.client, &self.config, self.api_key.as_ref(), prompt).await
97    }
98
99    async fn call_once(
100        &self,
101        prompt: &str,
102        schema: ResponseSchema,
103    ) -> Result<(Option<Value>, usize, usize), String> {
104        llm_call::execute_call(self, prompt, schema).await
105    }
106
107    pub async fn call(
108        &self,
109        prompt: &str,
110        schema: ResponseSchema,
111    ) -> Result<(Option<serde_json::Value>, usize, usize), String> {
112        llm_consensus::call_with_consensus(self, prompt, schema).await
113    }
114}
115
116fn max_concurrency() -> usize {
117    env::var("SNIFF_LLM_MAX_CONCURRENCY")
118        .or_else(|_| env::var("LLM_MAX_CONCURRENCY"))
119        .ok()
120        .and_then(|value| value.trim().parse::<usize>().ok())
121        .filter(|value| *value > 0)
122        .unwrap_or(5)
123}
124
125fn client_timeout() -> Duration {
126    env::var("SNIFF_LLM_CLIENT_TIMEOUT_SECS")
127        .or_else(|_| env::var("LLM_CLIENT_TIMEOUT_SECS"))
128        .ok()
129        .and_then(|value| value.trim().parse::<u64>().ok())
130        .filter(|secs| *secs > 0)
131        .map(Duration::from_secs)
132        .unwrap_or_else(|| Duration::from_secs(600))
133}
134
135#[cfg(test)]
136#[path = "tests/llm_impl.rs"]
137mod tests;
sniff/llm_impl.rs

sniff/
llm_impl.rs
