9use jev_facts::{Judged, Prepared, Source, Unlearned};
10use jev_protocol::{Choice, Json, ProtocolError, Question, Response};
11use jevhooks_events::Verdict;
12use rete::{Domain, Known, Network, Rule, Test, Then};
13use serde_json::json;

15use crate::{Never, SHOWN_CHARS, head, percent, tail};

19pub const BLOCK_AT_LEAST: f64 = 0.85;

22const STOPPED_EARLY: &str = "stopped-early";

25const ENDINGS: [(&str, &str); 5] = [
26    ("finished", "Every part of the latest request was carried out, or the question it asked was answered."),
27    (
28        "waiting",
29        "The assistant needs something only the developer can give before it can go on: an answer to a \
30         question it asked, a choice between options, permission, a credential, or an action on another machine.",
31    ),
32    (
33        "blocked",
34        "The assistant tried, hit an obstacle it names plainly (a failing command, a missing tool, an error it \
35         could not get past), and reports that instead of the result.",
36    ),
37    (
38        "running",
39        "The assistant started work that is still going (a build, a background job, another agent) and says \
40         it will report when that finishes.",
41    ),
42    (
43        "stopped-early",
44        "Work the latest request asked for is left undone and the final message gives no obstacle: it \
45         describes what it will do or could do next instead of doing it, or it did part and stopped.",
46    ),
47];

50#[derive(Clone, Copy, Debug, PartialEq, Eq)]
51pub struct Stop;

54#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)]
55pub enum Fact {

57    Refused,

59    Paused,

61    Judgeable,

63    StoppedEarly,
64}

66#[derive(Clone, Copy, Debug, PartialEq, Eq)]
67pub enum Value {
68    Yes,
69    No,
70}
71
72fn yes(holds: bool) -> Value {
73    if holds { Value::Yes } else { Value::No }
74}

77#[derive(Clone, Copy, Debug, PartialEq, Eq)]
78pub enum End {

80    AlreadyRefused,

82    StillRunning,

84    NothingToJudge,

86    Refuse,

88    LetEnd,
89}

91impl End {
92    pub fn verdict(self) -> Verdict {
93        match self {
94            End::Refuse => Verdict::Ask,
95            End::AlreadyRefused | End::StillRunning | End::NothingToJudge | End::LetEnd => Verdict::Pass,
96        }
97    }

101    pub fn unasked_line(self) -> Option<&'static str> {
102        match self {
103            End::AlreadyRefused => Some("already refused once this turn"),
104            End::StillRunning => Some("background work is still running"),
105            End::NothingToJudge => Some("no prompt or final message to judge"),
106            End::Refuse | End::LetEnd => None,
107        }
108    }
109}

111impl Domain for Stop {
112    type Fact = Fact;
113    type Value = Value;
114    type Effect = Never;
115    type End = End;
116    type Note = Never;
117
118    fn facts() -> &'static [Fact] {
119        &[Fact::Refused, Fact::Paused, Fact::Judgeable, Fact::StoppedEarly]
120    }
121
122    fn fact_name(fact: Fact) -> &'static str {
123        match fact {
124            Fact::Refused => "already refused",
125            Fact::Paused => "still running",
126            Fact::Judgeable => "can be judged",
127            Fact::StoppedEarly => "stopped early",
128        }
129    }
130
131    fn values(_: Fact) -> &'static [Value] {
132        &[Value::Yes, Value::No]
133    }
134
135    fn value_name(value: Value) -> &'static str {
136        match value {
137            Value::Yes => "yes",
138            Value::No => "no",
139        }
140    }
141
142    fn asked_for(fact: Fact) -> bool {
143        fact == Fact::StoppedEarly
144    }
145
146    fn teaches(effect: Never) -> Fact {
147        match effect {}
148    }
149
150    fn effect_name(effect: Never) -> String {
151        match effect {}
152    }
153
154    fn end_name(end: End) -> String {
155        match end {
156            End::Refuse => "send it back",
157            End::AlreadyRefused | End::StillRunning | End::NothingToJudge | End::LetEnd => "let it end",
158        }
159        .to_owned()
160    }
161
162    fn note_name(note: Never) -> String {
163        match note {}
164    }
165}

170pub const RULES: [Rule<'static, Stop>; 5] = [
171    Rule { name: "refused once already", when: &[Test::Is(Fact::Refused, Value::Yes)], then: Then::End(End::AlreadyRefused) },
172    Rule { name: "paused, not ended", when: &[Test::Is(Fact::Paused, Value::Yes)], then: Then::End(End::StillRunning) },
173    Rule { name: "nothing to judge", when: &[Test::Is(Fact::Judgeable, Value::No)], then: Then::End(End::NothingToJudge) },
174    Rule { name: "surely stopped early", when: &[Test::Is(Fact::StoppedEarly, Value::Yes)], then: Then::End(End::Refuse) },
175    Rule { name: "any other ending", when: &[Test::Is(Fact::StoppedEarly, Value::No)], then: Then::End(End::LetEnd) },
176];

179pub fn network() -> Network<Stop> {
180    Network::compile(&RULES)
181}

185pub fn before(refused: bool, paused: bool, judgeable: bool) -> Known<Stop> {
186    let mut known = Known::default();
187    known.learn(Fact::Refused, yes(refused));
188    known.learn(Fact::Paused, yes(paused));
189    known.learn(Fact::Judgeable, yes(judgeable));
190    known
191}

195pub fn state(requests: &[String], final_message: &str) -> Result<Json, ProtocolError> {
196    let requests: Vec<&str> = requests.iter().map(|request| head(request, SHOWN_CHARS / 2)).collect();
197    let state = json!({
198        "developer_requests_oldest_first": requests,
199        "assistant_final_message": tail(final_message, SHOWN_CHARS),
200    });
201    Json::verbatim(&state.to_string())
202}

204impl Source for Stop {
205    type Fact = Fact;
206    type Value = Value;
207
208    fn question_id(&self, fact: Fact) -> String {
209        match fact {
210            Fact::StoppedEarly => "ending",
211            Fact::Refused => "refused",
212            Fact::Paused => "paused",
213            Fact::Judgeable => "judgeable",
214        }
215        .to_owned()
216    }
217
218    fn question(&self, fact: Fact) -> Result<Question, ProtocolError> {
219        match fact {
220            Fact::StoppedEarly => Ok(Question::Choice(Choice::new(
221                Json::text(
222                    "A developer gave a coding assistant these requests, and the assistant has now ended its turn with \
223                     this final message. Which one of these describes how the turn ended, judged against the latest request?",
224                ),
225                ENDINGS.map(|(label, definition)| (label.to_owned(), Some(Json::text(definition)))),
226            )?)),
227            Fact::Refused | Fact::Paused | Fact::Judgeable => {
228                Err(ProtocolError::Invalid(format!("{} is not a fact Jev is asked for", Stop::fact_name(fact))))
229            }
230        }
231    }
232
233    fn learned(&self, fact: Fact, judged: &Judged) -> Option<Value> {
234        match (fact, judged) {
235            (Fact::StoppedEarly, Judged::Choice(answer)) => Some(yes(stopped_early(answer) >= BLOCK_AT_LEAST)),
236            _ => None,
237        }
238    }
239}
240
241fn stopped_early(answer: &jev_protocol::ChoiceAnswer) -> f64 {
242    answer.probabilities.iter().find(|(label, _)| label == STOPPED_EARLY).map_or(0.0, |(_, p)| *p)
243}

246#[derive(Clone, Debug)]
247pub struct Asked {
248    pub known: Known<Stop>,

250    pub stopped_early: f64,

252    pub likeliest: String,
253    json: serde_json::Value,
254}

257pub fn taught(mut known: Known<Stop>, prepared: &Prepared, response: &Response, facts: &[Fact]) -> Result<Asked, Unlearned> {
258    let judged = prepared.judged(response);
259    for (fact, value) in jev_facts::learn(&Stop, prepared, &judged, facts)? {
260        known.learn(fact, value);
261    }
262    let Some(Judged::Choice(ending)) = prepared.parts.iter().position(|part| part.id == "ending").and_then(|index| judged.get(index)) else {
263        return Err(Unlearned::NotAsked { id: "ending".to_owned() });
264    };
265    Ok(Asked {
266        known,
267        stopped_early: stopped_early(ending),
268        likeliest: ending.choice.clone(),
269        json: json!({ "ending": {
270            "choice": ending.choice,
271            "confidence": ending.confidence,
272            "probabilities": ending.probabilities.iter().map(|(label, p)| (label.clone(), json!(p))).collect::<serde_json::Map<String, serde_json::Value>>(),
273        } }),
274    })
275}

277impl Asked {

279    pub fn line(&self, end: End) -> String {
280        match end {
281            End::Refuse => {
282                format!("Jev: {}% likely stopped early: work asked for is left undone and no obstacle is named", percent(self.stopped_early))
283            }
284            _ => format!("Jev: the turn ended as {}; {}% likely stopped early", self.likeliest, percent(self.stopped_early)),
285        }
286    }

289    pub fn json(&self) -> serde_json::Value {
290        self.json.clone()
291    }

295    #[cfg(feature = "made-up")]
296    pub fn made_up(stopped_early: f64) -> Self {
297        let known = before(false, false, true);
298        let facts = [Fact::StoppedEarly];
299        let model = crate::jev_model().expect("the pinned model");
300        let prepared = jev_facts::wanted(&Stop, &model, state(&["a request".to_owned()], "a message").expect("a state"), &facts).expect("a request");
301        let probabilities: Vec<(&str, f64)> = ENDINGS
302            .iter()
303            .map(|(label, _)| {
304                (*label, match *label {
305                    STOPPED_EARLY => stopped_early,
306                    "finished" => 1.0 - stopped_early,
307                    _ => 0.0,
308                })
309            })
310            .collect();
311        let body = crate::made_up_body(&[("ending", crate::made_up_choice(&probabilities))]);
312        let response = Response::parse(&model, prepared.asking().1, body.as_bytes()).expect("a response Jev could send");
313        taught(known, &prepared, &response, &facts).expect("the fact asked for")
314    }
315}

317#[cfg(test)]
318mod tests {
319    use rete::Next;
320
321    use super::*;

325    #[test]
326    fn stop_blocks_only_when_confident() {
327        let network = network();
328        assert_eq!(network.next(&Asked::made_up(0.84).known), Next::End(End::LetEnd));
329        assert_eq!(network.next(&Asked::made_up(0.85).known), Next::End(End::Refuse));
330        assert_eq!(End::LetEnd.verdict(), Verdict::Pass);
331        assert_eq!(End::Refuse.verdict(), Verdict::Ask);
332    }

336    #[test]
337    fn what_is_known_already_is_never_asked_about() {
338        let network = network();
339        assert_eq!(network.next(&before(true, false, true)), Next::End(End::AlreadyRefused));
340        assert_eq!(network.next(&before(false, true, true)), Next::End(End::StillRunning));
341        assert_eq!(network.next(&before(false, false, false)), Next::End(End::NothingToJudge));
342        assert_eq!(network.next(&before(false, false, true)), Next::Ask(vec![Fact::StoppedEarly]));
343    }

346    #[test]
347    fn a_refused_end_is_let_through_even_if_it_stopped_early() {
348        let mut known = Asked::made_up(0.99).known;
349        known.learn(Fact::Refused, Value::Yes);
350        assert_eq!(network().next(&known), Next::End(End::AlreadyRefused));
351    }

354    #[test]
355    fn the_line_carries_jevs_numbers() {
356        assert_eq!(Asked::made_up(0.9).line(End::Refuse), "Jev: 90% likely stopped early: work asked for is left undone and no obstacle is named");
357        assert_eq!(Asked::made_up(0.1).line(End::LetEnd), "Jev: the turn ended as finished; 10% likely stopped early");
358    }

362    #[test]
363    fn the_refused_ending_is_one_of_the_options() {
364        assert_eq!(ENDINGS.iter().filter(|(label, _)| *label == STOPPED_EARLY).count(), 1);
365    }
366}