use crate::error::{Error, Result};
use serde::{Deserialize, Serialize};
use std::io::Write;
use std::process::{Command, Stdio};
pub const MAX_LLM_DRAFTS: usize = 8;
pub const MAX_GUIDANCE_LEN: usize = 600;
pub const MAX_SUMMARY_LEN: usize = 200;
pub trait LlmBackend: Send + Sync {
fn model(&self) -> &str;
fn complete(&self, request: &str) -> Result<String>;
}
#[derive(Debug, Clone, Serialize)]
pub struct FindingBrief {
pub analyzer: String,
pub summary: String,
pub target: String,
pub severity: String,
}
#[derive(Debug, Clone, Serialize)]
pub struct EvidenceItem {
pub hash: String,
pub grain_type: String,
pub text: String,
}
#[derive(Debug, Clone, Serialize)]
pub struct LlmRequest<'a> {
#[serde(rename = "loop")]
pub loop_proto: u8,
pub op: &'a str,
pub instructions: &'a str,
#[serde(skip_serializing_if = "Vec::is_empty")]
pub findings: Vec<FindingBrief>,
#[serde(skip_serializing_if = "Vec::is_empty")]
pub evidence: Vec<EvidenceItem>,
#[serde(skip_serializing_if = "Vec::is_empty")]
pub rejected: Vec<String>,
#[serde(skip_serializing_if = "Vec::is_empty")]
pub approved: Vec<String>,
}
#[derive(Debug, Clone, Deserialize, Default)]
#[serde(default)]
pub struct LlmDraft {
pub summary: String,
pub target: String,
pub guidance: String,
pub evidence: Vec<String>,
pub confidence: f64,
}
#[derive(Debug, Clone, Deserialize, Default)]
#[serde(default)]
pub struct DiscoverResponse {
pub recommendations: Vec<LlmDraft>,
}
#[derive(Debug, Clone, Deserialize, Default)]
#[serde(default)]
pub struct EnrichResponse {
pub notes: Vec<EnrichNote>,
}
#[derive(Debug, Clone, Deserialize, Default)]
#[serde(default)]
pub struct EnrichNote {
pub target: String,
pub guidance: String,
}
#[derive(Debug, Clone, Serialize)]
pub struct GroundRequest<'a> {
#[serde(rename = "loop")]
pub loop_proto: u8,
pub op: &'a str, pub instructions: &'a str,
pub claims: Vec<GroundItem>,
}
#[derive(Debug, Clone, Serialize)]
pub struct GroundItem {
pub id: usize,
pub claim: String,
pub evidence: Vec<EvidenceItem>,
}
#[derive(Debug, Clone, Deserialize, Default)]
#[serde(default)]
pub struct GroundResponse {
pub results: Vec<GroundResult>,
}
#[derive(Debug, Clone, Deserialize, Default)]
#[serde(default)]
pub struct GroundResult {
pub id: usize,
pub supported: bool,
pub reason: String,
}
#[derive(Debug, Clone, Serialize)]
pub struct VerifyRequest<'a> {
#[serde(rename = "loop")]
pub loop_proto: u8,
pub op: &'a str, pub instructions: &'a str,
pub findings: Vec<VerifyItem>,
}
#[derive(Debug, Clone, Serialize)]
pub struct VerifyItem {
pub id: usize,
pub summary: String,
pub target: String,
pub evidence: Vec<EvidenceItem>,
}
#[derive(Debug, Clone, Deserialize, Default)]
#[serde(default)]
pub struct VerifyResponse {
pub results: Vec<VerifyResult>,
}
#[derive(Debug, Clone, Deserialize, Default)]
#[serde(default)]
pub struct VerifyResult {
pub id: usize,
pub keep: bool,
pub confidence: f64,
pub reason: String,
}
#[derive(Debug, Clone, Deserialize, Default)]
#[serde(default)]
struct ProbeResponse {
model: String,
}
pub struct CommandLlm {
argv: Vec<String>,
model: String,
}
impl CommandLlm {
pub fn new(cmd: &str, model: Option<&str>) -> Result<Self> {
let argv: Vec<String> = cmd.split_whitespace().map(str::to_string).collect();
if argv.is_empty() {
return Err(Error::LlmBackend("--llm-cmd is empty".into()));
}
let mut me = CommandLlm {
argv,
model: model.unwrap_or("").to_string(),
};
let probe = me.run(r#"{"loop":1,"op":"probe"}"#)?;
let parsed: ProbeResponse = serde_json::from_str(probe.trim()).map_err(|e| {
Error::LlmBackend(format!("--llm-cmd probe did not return JSON with a model: {e}"))
})?;
if me.model.is_empty() {
me.model = if parsed.model.is_empty() {
"unspecified".to_string()
} else {
parsed.model
};
}
Ok(me)
}
fn run(&self, request: &str) -> Result<String> {
let mut child = Command::new(&self.argv[0])
.args(&self.argv[1..])
.stdin(Stdio::piped())
.stdout(Stdio::piped())
.stderr(Stdio::inherit())
.spawn()
.map_err(|e| Error::LlmBackend(format!("spawn --llm-cmd {:?}: {e}", self.argv[0])))?;
{
let mut stdin = child.stdin.take().expect("stdin piped");
stdin
.write_all(request.as_bytes())
.map_err(|e| Error::LlmBackend(format!("write to --llm-cmd: {e}")))?;
}
let out = child
.wait_with_output()
.map_err(|e| Error::LlmBackend(format!("--llm-cmd wait: {e}")))?;
if !out.status.success() {
return Err(Error::LlmBackend(format!(
"--llm-cmd exited with {}",
out.status
)));
}
String::from_utf8(out.stdout)
.map_err(|e| Error::LlmBackend(format!("--llm-cmd stdout not UTF-8: {e}")))
}
}
impl LlmBackend for CommandLlm {
fn model(&self) -> &str {
&self.model
}
fn complete(&self, request: &str) -> Result<String> {
self.run(request)
}
}
pub fn parse_discover(raw: &str) -> DiscoverResponse {
serde_json::from_str(raw.trim()).unwrap_or_default()
}
pub fn parse_enrich(raw: &str) -> EnrichResponse {
serde_json::from_str(raw.trim()).unwrap_or_default()
}
pub fn parse_ground(raw: &str) -> GroundResponse {
serde_json::from_str(raw.trim()).unwrap_or_default()
}
pub fn parse_verify(raw: &str) -> VerifyResponse {
serde_json::from_str(raw.trim()).unwrap_or_default()
}
pub fn cap(s: &str, max: usize) -> String {
if s.chars().count() <= max {
s.to_string()
} else {
s.chars().take(max).collect()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn parse_discover_drops_garbage() {
assert!(parse_discover("not json").recommendations.is_empty());
let r = parse_discover(r#"{"recommendations":[{"summary":"s","target":"entity:x/y","evidence":["h1"],"junk":1}]}"#);
assert_eq!(r.recommendations.len(), 1);
assert_eq!(r.recommendations[0].summary, "s");
assert_eq!(r.recommendations[0].evidence, vec!["h1"]);
}
#[test]
fn parse_enrich_reads_notes() {
let r = parse_enrich(r#"{"notes":[{"target":"entity:a/b","guidance":"g"}]}"#);
assert_eq!(r.notes.len(), 1);
assert_eq!(r.notes[0].guidance, "g");
}
#[test]
fn cap_respects_char_boundaries() {
assert_eq!(cap("hello", 3), "hel");
assert_eq!(cap("héllo", 2), "hé");
assert_eq!(cap("hi", 5), "hi");
}
}