Skip to main content

kaizen/core_loop/
query.rs

1// SPDX-License-Identifier: AGPL-3.0-or-later
2use crate::core::event::{Event, SessionRecord};
3use crate::core_loop::TraceHit;
4use crate::core_loop::query_meta::Meta;
5use crate::core_loop::query_syntax::{Field, Op, QueryExpr, Term};
6use crate::search::{kind_label, tokens_total};
7use crate::store::Store;
8use anyhow::Result;
9
10pub fn is_structured(raw: &str) -> bool {
11    crate::core_loop::query_syntax::is_structured(raw)
12}
13
14pub fn parse(raw: &str) -> Result<QueryExpr> {
15    crate::core_loop::query_syntax::parse(raw)
16}
17
18pub fn run(
19    store: &Store,
20    workspace: &str,
21    raw: &str,
22    start_ms: u64,
23    limit: usize,
24) -> Result<Vec<TraceHit>> {
25    let expr = parse(raw)?;
26    let meta = Meta::load(store, start_ms)?;
27    let rows = store.workspace_events(workspace)?;
28    Ok(rows
29        .into_iter()
30        .filter_map(|r| hit_if(&expr, &meta, r))
31        .take(limit)
32        .collect())
33}
34
35fn hit_if(expr: &QueryExpr, meta: &Meta, row: (SessionRecord, Event)) -> Option<TraceHit> {
36    let (session, event) = row;
37    expr.terms
38        .iter()
39        .all(|t| matches_term(t, meta, &session, &event))
40        .then(|| hit(session, event))
41}
42
43fn hit(session: SessionRecord, event: Event) -> TraceHit {
44    TraceHit {
45        session_id: session.id,
46        seq: Some(event.seq),
47        ts_ms: event.ts_ms,
48        agent: session.agent,
49        kind: kind_label(&event.kind).unwrap_or("unknown").into(),
50        summary: tool(&event)
51            .or_else(|| text(&event.payload))
52            .unwrap_or_default(),
53    }
54}
55
56fn matches_term(t: &Term, meta: &Meta, s: &SessionRecord, e: &Event) -> bool {
57    match t.field {
58        Field::Agent => eq(&s.agent, &t.value),
59        Field::Model => s.model.as_deref().is_some_and(|v| eq(v, &t.value)),
60        Field::Kind => kind_label(&e.kind).is_some_and(|v| eq(v, &t.value)),
61        Field::Tool => tool(e).as_deref().is_some_and(|v| eq(v, &t.value)),
62        Field::Path => paths(e).iter().any(|p| p.contains(&t.value)),
63        Field::Skill => skills(e).iter().any(|p| p.contains(&t.value)),
64        Field::TokensTotal => cmp(tokens_total(e) as f64, t),
65        Field::CostUsd => cmp(e.cost_usd_e6.unwrap_or(0) as f64 / 1_000_000.0, t),
66        Field::EvalScore => meta.eval(&s.id).is_some_and(|v| cmp(v, t)),
67        Field::FeedbackLabel => meta.feedback(&s.id).is_some_and(|v| eq(v, &t.value)),
68        Field::Prompt => s
69            .prompt_fingerprint
70            .as_deref()
71            .is_some_and(|v| v.starts_with(&t.value)),
72        Field::Status => eq(status(s), &t.value),
73        Field::SpanKind => meta.span_kind(&s.id, &t.value),
74    }
75}
76
77fn cmp(n: f64, t: &Term) -> bool {
78    let Ok(v) = t.value.parse::<f64>() else {
79        return false;
80    };
81    match t.op {
82        Op::Eq => n == v,
83        Op::Gt => n > v,
84        Op::Gte => n >= v,
85        Op::Lt => n < v,
86        Op::Lte => n <= v,
87    }
88}
89
90fn eq(a: &str, b: &str) -> bool {
91    a.eq_ignore_ascii_case(b)
92}
93
94fn status(s: &SessionRecord) -> &'static str {
95    match s.status {
96        crate::core::event::SessionStatus::Running => "running",
97        crate::core::event::SessionStatus::Waiting => "waiting",
98        crate::core::event::SessionStatus::Idle => "idle",
99        crate::core::event::SessionStatus::Done => "done",
100    }
101}
102
103fn text(v: &serde_json::Value) -> Option<String> {
104    v.get("text").and_then(|v| v.as_str()).map(str::to_string)
105}
106
107fn tool(e: &Event) -> Option<String> {
108    e.tool
109        .clone()
110        .or_else(|| crate::store::tool_span_index::hook_tool(&e.payload))
111}
112
113fn paths(e: &Event) -> Vec<String> {
114    crate::store::event_index::paths_from_event_payload(&e.payload)
115}
116
117fn skills(e: &Event) -> Vec<String> {
118    crate::store::event_index::skills_from_event_json(&e.payload)
119}