1use 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 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 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 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 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 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}