mdx_rust_core/
llm.rs

1//! Thin LLM client abstraction for the optimizer.
2//!
3//! Currently backed by Rig for convenience. Later we can support direct HTTP
4//! for more control over "heavy reasoning" models.
5
6use crate::optimizer::ModelProvenance;
7use rig::completion::Prompt;
8use rig::providers::openai;
9use serde::Serialize;
10
11/// Very simple diagnosis request.
12#[derive(Serialize)]
13pub struct DiagnosisRequest {
14    pub policy: String,
15    pub bundle_summary: String, // path count + key files
16    pub traces_summary: String,
17    pub scores: Vec<f32>,
18}
19
20/// Structured candidate returned by the LLM.
21#[derive(Debug, Clone, serde::Deserialize, serde::Serialize)]
22pub struct StructuredCandidate {
23    pub focus: String, // "system_prompt", "tool_descriptions", "reasoning", "logic"
24    pub description: String,
25    pub expected_improvement: String,
26}
27
28/// Result of asking the model for diagnosis + candidates.
29#[derive(Debug, Clone)]
30pub struct DiagnosisResult {
31    pub summary: String,
32    pub candidates: Vec<StructuredCandidate>,
33}
34
35/// Basic LLM client for diagnosis.
36pub struct LlmClient {
37    model: String,
38}
39
40impl LlmClient {
41    pub fn new(model: impl Into<String>) -> Self {
42        Self {
43            model: model.into(),
44        }
45    }
46
47    pub fn provenance(&self, used: bool) -> ModelProvenance {
48        ModelProvenance {
49            role: "diagnosis".to_string(),
50            provider: "openai".to_string(),
51            model: self.model.clone(),
52            used,
53        }
54    }
55
56    /// Ask a strong model for diagnosis and candidate ideas.
57    pub async fn diagnose(&self, req: DiagnosisRequest) -> anyhow::Result<DiagnosisResult> {
58        // Only attempt if we have a key (avoids panic on Client::from_env)
59        if std::env::var("OPENAI_API_KEY").is_err() {
60            return Err(anyhow::anyhow!("No OPENAI_API_KEY"));
61        }
62
63        let client = openai::Client::from_env();
64        let agent = client
65            .agent(&self.model)
66            .preamble(
67                "You are an expert at debugging and improving LLM agents written in Rust. \
68                 Always respond with a JSON object: {\"summary\": \"short diagnosis\", \"candidates\": [{\"focus\": \"system_prompt|tool_description|fallback_logic|output_schema|model_config\", \"description\": \"...\", \"expected_improvement\": \"...\"}]}. \
69                 Be concise and actionable. Focus on minimal, high-impact changes to prompts, tool usage, schemas, model config, or control flow.",
70            )
71            .build();
72
73        let prompt = format!(
74            "Policy:\n{}\n\nCode analysis (extracted preambles, tools, structure):\n{}\n\nTraces summary:\n{}\n\nCurrent scores: {:?}\n\n\
75             Return ONLY a JSON object with keys \"summary\" and \"candidates\" (array of objects with focus, description, expected_improvement). No prose outside the JSON.",
76            req.policy, req.bundle_summary, req.traces_summary, req.scores
77        );
78
79        let response = agent.prompt(prompt).await?;
80
81        // Try to parse structured JSON from the model response.
82        // We instruct the model to return a JSON object with "summary" and "candidates".
83        #[derive(serde::Deserialize)]
84        struct LlmOutput {
85            summary: Option<String>,
86            candidates: Option<Vec<StructuredCandidate>>,
87        }
88
89        let parsed: Option<LlmOutput> = serde_json::from_str(&response).ok().or_else(|| {
90            // Sometimes the model wraps it in ```json ... ```
91            if let Some(start) = response.find('{') {
92                if let Some(end) = response.rfind('}') {
93                    serde_json::from_str(&response[start..=end]).ok()
94                } else {
95                    None
96                }
97            } else {
98                None
99            }
100        });
101
102        if let Some(out) = parsed {
103            let summary = out
104                .summary
105                .unwrap_or_else(|| "LLM diagnosis completed.".to_string());
106            let candidates = out.candidates.unwrap_or_default();
107            if !candidates.is_empty() {
108                return Ok(DiagnosisResult {
109                    summary,
110                    candidates,
111                });
112            }
113        }
114
115        // Fallback: treat the whole response as a single textual candidate
116        let summary = response
117            .lines()
118            .next()
119            .unwrap_or("LLM returned free text")
120            .to_string();
121        let fallback = vec![StructuredCandidate {
122            focus: "system_prompt".to_string(),
123            description: response.chars().take(280).collect(),
124            expected_improvement: "Model suggestion".to_string(),
125        }];
126
127        Ok(DiagnosisResult {
128            summary,
129            candidates: fallback,
130        })
131    }
132}
133
134impl Default for LlmClient {
135    fn default() -> Self {
136        Self::new("gpt-4o")
137    }
138}
mdx_rust_core/llm.rs

mdx_rust_core/
llm.rs
