use crate::error::{Error, Result};
use serde::{Deserialize, Serialize};
use serde_json::Value;
pub const MAX_LLM_DRAFTS: usize = 8;
pub const MAX_GUIDANCE_LEN: usize = 600;
pub const MAX_SUMMARY_LEN: usize = 200;
pub const MAX_LESSON_LEN: usize = 240;
pub const MAX_RELATION_LEN: usize = 64;
pub const MAX_OBJECT_LEN: usize = 480;
pub const MAX_QUERY_BODY_LEN: usize = 2_000;
pub const MAX_PLAN_EDITS: usize = 8;
pub const MAX_CODE_LEN: usize = 20_000;
pub trait LlmBackend: Send + Sync {
fn model(&self) -> &str;
fn complete(&self, request: &str) -> Result<String>;
}
impl<T: LlmBackend + ?Sized> LlmBackend for Box<T> {
fn model(&self) -> &str {
(**self).model()
}
fn complete(&self, request: &str) -> Result<String> {
(**self).complete(request)
}
}
#[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,
pub lesson: String,
pub proposal: Option<Value>,
}
#[derive(Debug, Clone, Deserialize, Default, PartialEq)]
#[serde(default)]
pub struct PlanEdit {
pub path: String,
pub from: Value,
pub to: Value,
}
#[derive(Debug, Clone, Deserialize, PartialEq)]
#[serde(tag = "kind", rename_all = "snake_case")]
pub enum DraftProposal {
Lesson {
#[serde(default)]
lesson: String,
},
Fact {
#[serde(default)]
relation: String,
#[serde(default)]
object: String,
},
QueryRevision {
#[serde(default)]
body: String,
},
PlanRevision {
#[serde(default)]
edits: Vec<PlanEdit>,
},
CodeRevision {
#[serde(default)]
source: String,
},
}
impl LlmDraft {
pub fn parsed_proposal(&self) -> Option<DraftProposal> {
if let Some(v) = &self.proposal {
return serde_json::from_value::<DraftProposal>(v.clone()).ok();
}
if self.lesson.trim().is_empty() {
None
} else {
Some(DraftProposal::Lesson {
lesson: self.lesson.clone(),
})
}
}
}
#[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 out = crate::proc::run_argv(&self.argv, request, Some(crate::proc::DEFAULT_TIMEOUT))
.map_err(|e| Error::LlmBackend(format!("spawn --llm-cmd {:?}: {e}", self.argv[0])))?;
if let Some(why) = out.failure("--llm-cmd") {
return Err(Error::LlmBackend(why));
}
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 lesson_desugars_when_no_proposal_is_sent() {
let r = parse_discover(
r#"{"recommendations":[{"summary":"s","target":"entity:a/b","evidence":["h"],"lesson":"Do the thing"}]}"#,
);
assert_eq!(
r.recommendations[0].parsed_proposal(),
Some(DraftProposal::Lesson { lesson: "Do the thing".into() })
);
}
#[test]
fn an_explicit_proposal_wins_over_lesson() {
let r = parse_discover(
r#"{"recommendations":[{"summary":"s","target":"entity:a/b","evidence":["h"],
"lesson":"ignored","proposal":{"kind":"fact","relation":"alias_of","object":"Cobalt Cloud"}}]}"#,
);
assert_eq!(
r.recommendations[0].parsed_proposal(),
Some(DraftProposal::Fact {
relation: "alias_of".into(),
object: "Cobalt Cloud".into()
})
);
}
#[test]
fn an_unparseable_proposal_is_advisory_not_reinterpreted() {
for body in [
r#""proposal":{"kind":"teleport","x":1}"#,
r#""proposal":{"kind":"plan_revision","edits":"not-a-list"}"#,
r#""proposal":42"#,
] {
let raw = format!(
r#"{{"recommendations":[{{"summary":"s","target":"entity:a/b","evidence":["h"],"lesson":"L",{body}}}]}}"#
);
let r = parse_discover(&raw);
assert_eq!(r.recommendations.len(), 1, "{body}");
assert_eq!(r.recommendations[0].parsed_proposal(), None, "{body}");
}
}
#[test]
fn one_bad_proposal_does_not_drop_its_siblings() {
let r = parse_discover(
r#"{"recommendations":[
{"summary":"a","target":"entity:a/b","evidence":["h"],"proposal":{"kind":"nope"}},
{"summary":"b","target":"entity:a/c","evidence":["h"],"proposal":{"kind":"lesson","lesson":"Keep me"}}
]}"#,
);
assert_eq!(r.recommendations.len(), 2);
assert_eq!(r.recommendations[0].parsed_proposal(), None);
assert_eq!(
r.recommendations[1].parsed_proposal(),
Some(DraftProposal::Lesson { lesson: "Keep me".into() })
);
}
#[test]
fn plan_edits_carry_their_staleness_check() {
let r = parse_discover(
r#"{"recommendations":[{"summary":"s","target":"grain:abc","evidence":["h"],
"proposal":{"kind":"plan_revision","edits":[{"path":"retries.fetch","from":null,"to":3}]}}]}"#,
);
let Some(DraftProposal::PlanRevision { edits }) =
r.recommendations[0].parsed_proposal()
else {
panic!("expected a plan revision");
};
assert_eq!(edits.len(), 1);
assert_eq!(edits[0].path, "retries.fetch");
assert_eq!(edits[0].from, Value::Null);
assert_eq!(edits[0].to, Value::from(3));
}
#[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");
}
}