Skip to main content

toolpath_gemini/
query.rs

1//! Query/filter operations over a loaded `Conversation`.
2
3use crate::types::{Conversation, GeminiMessage, GeminiRole};
4use chrono::{DateTime, Utc};
5
6pub struct ConversationQuery<'a> {
7    conversation: &'a Conversation,
8}
9
10impl<'a> ConversationQuery<'a> {
11    pub fn new(conversation: &'a Conversation) -> Self {
12        Self { conversation }
13    }
14
15    /// Messages with the given role.
16    pub fn by_role(&self, role: GeminiRole) -> Vec<&'a GeminiMessage> {
17        self.conversation
18            .all_messages()
19            .filter(|m| m.role == role)
20            .collect()
21    }
22
23    /// Messages whose parsed timestamp falls in `[start, end]`.
24    pub fn by_time_range(
25        &self,
26        start: DateTime<Utc>,
27        end: DateTime<Utc>,
28    ) -> Vec<&'a GeminiMessage> {
29        self.conversation
30            .all_messages()
31            .filter(|m| {
32                if let Ok(ts) = m.timestamp.parse::<DateTime<Utc>>() {
33                    ts >= start && ts <= end
34                } else {
35                    false
36                }
37            })
38            .collect()
39    }
40
41    /// Messages that invoke a tool with the given `name`.
42    pub fn tool_uses_by_name(&self, tool_name: &str) -> Vec<&'a GeminiMessage> {
43        self.conversation
44            .all_messages()
45            .filter(|m| m.tool_calls().iter().any(|t| t.name == tool_name))
46            .collect()
47    }
48
49    /// Messages whose text or tool-call result text contains `search`
50    /// (case-insensitive).
51    pub fn contains_text(&self, search: &str) -> Vec<&'a GeminiMessage> {
52        let needle = search.to_lowercase();
53        self.conversation
54            .all_messages()
55            .filter(|m| {
56                if m.content.text().to_lowercase().contains(&needle) {
57                    return true;
58                }
59                m.tool_calls()
60                    .iter()
61                    .any(|t| t.result_text().to_lowercase().contains(&needle))
62            })
63            .collect()
64    }
65
66    /// Messages where at least one tool call reported an error.
67    pub fn errors(&self) -> Vec<&'a GeminiMessage> {
68        self.conversation
69            .all_messages()
70            .filter(|m| m.tool_calls().iter().any(|t| t.is_error()))
71            .collect()
72    }
73}
74
75#[cfg(test)]
76mod tests {
77    use super::*;
78    use crate::types::{ChatFile, Conversation, GeminiContent, GeminiMessage, ToolCall};
79
80    fn msg(role: GeminiRole, ts: &str, text: &str) -> GeminiMessage {
81        GeminiMessage {
82            id: ts.into(),
83            timestamp: ts.into(),
84            role,
85            content: GeminiContent::Text(text.into()),
86            thoughts: None,
87            tokens: None,
88            model: None,
89            tool_calls: None,
90            extra: Default::default(),
91        }
92    }
93
94    fn tool_msg(role: GeminiRole, ts: &str, tool_name: &str, error: bool) -> GeminiMessage {
95        let mut m = msg(role, ts, "");
96        m.tool_calls = Some(vec![ToolCall {
97            id: "t".into(),
98            name: tool_name.into(),
99            args: serde_json::Value::Null,
100            status: if error {
101                "error".into()
102            } else {
103                "success".into()
104            },
105            timestamp: ts.into(),
106            result: vec![],
107            result_display: None,
108            description: None,
109            display_name: None,
110            extra: Default::default(),
111        }]);
112        m
113    }
114
115    fn build_convo(messages: Vec<GeminiMessage>) -> Conversation {
116        let chat = ChatFile {
117            session_id: "s".into(),
118            project_hash: "h".into(),
119            start_time: None,
120            last_updated: None,
121            directories: None,
122            kind: None,
123            summary: None,
124            messages,
125            extra: Default::default(),
126        };
127        Conversation::new("session-uuid".into(), chat)
128    }
129
130    #[test]
131    fn test_by_role() {
132        let convo = build_convo(vec![
133            msg(GeminiRole::User, "2026-04-17T10:00:00Z", "Hi"),
134            msg(GeminiRole::Gemini, "2026-04-17T10:00:01Z", "Hey"),
135        ]);
136        let q = ConversationQuery::new(&convo);
137        assert_eq!(q.by_role(GeminiRole::User).len(), 1);
138        assert_eq!(q.by_role(GeminiRole::Gemini).len(), 1);
139    }
140
141    #[test]
142    fn test_tool_uses_by_name() {
143        let convo = build_convo(vec![
144            msg(GeminiRole::User, "2026-04-17T10:00:00Z", "read"),
145            tool_msg(
146                GeminiRole::Gemini,
147                "2026-04-17T10:00:01Z",
148                "read_file",
149                false,
150            ),
151            tool_msg(
152                GeminiRole::Gemini,
153                "2026-04-17T10:00:02Z",
154                "write_file",
155                false,
156            ),
157        ]);
158        let q = ConversationQuery::new(&convo);
159        assert_eq!(q.tool_uses_by_name("read_file").len(), 1);
160        assert_eq!(q.tool_uses_by_name("missing").len(), 0);
161    }
162
163    #[test]
164    fn test_contains_text_case_insensitive() {
165        let convo = build_convo(vec![msg(
166            GeminiRole::User,
167            "2026-04-17T10:00:00Z",
168            "Look at AUTH.rs",
169        )]);
170        let q = ConversationQuery::new(&convo);
171        assert_eq!(q.contains_text("auth").len(), 1);
172        assert_eq!(q.contains_text("db").len(), 0);
173    }
174
175    #[test]
176    fn test_errors() {
177        let convo = build_convo(vec![
178            tool_msg(
179                GeminiRole::Gemini,
180                "2026-04-17T10:00:00Z",
181                "run_shell_command",
182                false,
183            ),
184            tool_msg(
185                GeminiRole::Gemini,
186                "2026-04-17T10:00:01Z",
187                "run_shell_command",
188                true,
189            ),
190        ]);
191        let q = ConversationQuery::new(&convo);
192        let errs = q.errors();
193        assert_eq!(errs.len(), 1);
194    }
195
196    #[test]
197    fn test_by_time_range() {
198        let convo = build_convo(vec![
199            msg(GeminiRole::User, "2026-04-17T10:00:00Z", "early"),
200            msg(GeminiRole::User, "2026-04-17T12:00:00Z", "mid"),
201            msg(GeminiRole::User, "2026-04-17T14:00:00Z", "late"),
202        ]);
203        let q = ConversationQuery::new(&convo);
204        let start: DateTime<Utc> = "2026-04-17T11:00:00Z".parse().unwrap();
205        let end: DateTime<Utc> = "2026-04-17T13:00:00Z".parse().unwrap();
206        let hits = q.by_time_range(start, end);
207        assert_eq!(hits.len(), 1);
208        assert_eq!(hits[0].content.text(), "mid");
209    }
210
211    #[test]
212    fn test_contains_text_matches_tool_result() {
213        let mut m = tool_msg(GeminiRole::Gemini, "ts", "read_file", false);
214        if let Some(calls) = m.tool_calls.as_mut() {
215            calls[0].result = vec![crate::types::FunctionResponse {
216                function_response: crate::types::FunctionResponseBody {
217                    id: "t".into(),
218                    name: "read_file".into(),
219                    response: serde_json::json!({"output": "hello AUTH module"}),
220                },
221            }];
222        }
223        let convo = build_convo(vec![m]);
224        let q = ConversationQuery::new(&convo);
225        assert_eq!(q.contains_text("auth").len(), 1);
226    }
227
228    #[test]
229    fn test_spans_main_and_sub_agents() {
230        let main_chat = ChatFile {
231            session_id: "main".into(),
232            project_hash: "".into(),
233            start_time: None,
234            last_updated: None,
235            directories: None,
236            kind: None,
237            summary: None,
238            messages: vec![msg(GeminiRole::User, "ts1", "main text")],
239            extra: Default::default(),
240        };
241        let sub_chat = ChatFile {
242            session_id: "sub".into(),
243            project_hash: "".into(),
244            start_time: None,
245            last_updated: None,
246            directories: None,
247            kind: Some("subagent".into()),
248            summary: None,
249            messages: vec![msg(GeminiRole::User, "ts2", "sub text")],
250            extra: Default::default(),
251        };
252        let mut convo = Conversation::new("uuid".into(), main_chat);
253        convo.sub_agents.push(sub_chat);
254        let q = ConversationQuery::new(&convo);
255        assert_eq!(q.contains_text("sub text").len(), 1);
256        assert_eq!(q.contains_text("main text").len(), 1);
257        assert_eq!(q.by_role(GeminiRole::User).len(), 2);
258    }
259}