1use objectiveai_sdk::cli::command::agents::logs::list::{
31 AssistantResponsePart, ClientNotificationPart, ClientNotificationPartType, ResponseItem,
32 ToolResponsePart, ToolResponsePartType,
33};
34use sqlx::Row as _;
35
36use super::super::time::unix_to_rfc3339;
37use super::super::{Error, Pool};
38use super::row::MessageTable;
39
40struct MsgRow {
43 id: i64,
46 response_id: String,
47 table_kind: MessageTable,
48 agent_instance_hierarchy: String,
49 timestamp_delivered: i64,
50 sender_agent_instance_hierarchy: Option<String>,
56 message_queue_id: Option<i64>,
61 timestamp_queued: Option<i64>,
65 message_queue_key: Option<String>,
72 function_name: Option<String>,
78 tool_call_id: Option<String>,
86 row_sub_index: Option<i64>,
91}
92
93#[derive(Debug, Clone, Copy, PartialEq, Eq)]
97enum BlockClass {
98 AgentCompletionRequest,
99 VectorCompletionRequest,
100 FunctionExecutionRequest,
101 ClientNotification,
102 AssistantResponse,
103 ToolResponse,
104}
105
106fn block_class(t: MessageTable) -> BlockClass {
107 match t {
108 MessageTable::AgentCompletionRequest => BlockClass::AgentCompletionRequest,
109 MessageTable::VectorCompletionRequest => BlockClass::VectorCompletionRequest,
110 MessageTable::FunctionExecutionRequest => BlockClass::FunctionExecutionRequest,
111 MessageTable::MessageQueueText
112 | MessageTable::MessageQueueImage
113 | MessageTable::MessageQueueAudio
114 | MessageTable::MessageQueueVideo
115 | MessageTable::MessageQueueFile => BlockClass::ClientNotification,
116 MessageTable::ToolResponseContentText
117 | MessageTable::ToolResponseContentImage
118 | MessageTable::ToolResponseContentAudio
119 | MessageTable::ToolResponseContentVideo
120 | MessageTable::ToolResponseContentFile => BlockClass::ToolResponse,
121 MessageTable::AssistantResponseRefusal
122 | MessageTable::AssistantResponseReasoning
123 | MessageTable::AssistantResponseToolCalls
124 | MessageTable::AssistantResponseContentText
125 | MessageTable::AssistantResponseContentImage
126 | MessageTable::AssistantResponseContentAudio
127 | MessageTable::AssistantResponseContentVideo
128 | MessageTable::AssistantResponseContentFile => BlockClass::AssistantResponse,
129 }
130}
131
132fn client_notification_kind(t: MessageTable) -> Option<ClientNotificationPartType> {
133 match t {
134 MessageTable::MessageQueueText => Some(ClientNotificationPartType::Text),
135 MessageTable::MessageQueueImage => Some(ClientNotificationPartType::Image),
136 MessageTable::MessageQueueAudio => Some(ClientNotificationPartType::Audio),
137 MessageTable::MessageQueueVideo => Some(ClientNotificationPartType::Video),
138 MessageTable::MessageQueueFile => Some(ClientNotificationPartType::File),
139 _ => None,
140 }
141}
142
143fn assistant_response_part(row: &MsgRow) -> Option<AssistantResponsePart> {
148 let id = row.id;
149 let delivered_at = unix_to_rfc3339(row.timestamp_delivered);
150 Some(match row.table_kind {
151 MessageTable::AssistantResponseToolCalls => AssistantResponsePart::ToolCall {
152 id,
153 delivered_at,
154 function_name: row.function_name.clone().unwrap_or_default(),
155 tool_call_id: row.tool_call_id.clone().unwrap_or_default(),
156 tool_call_index: row.row_sub_index.unwrap_or_default(),
157 },
158 MessageTable::AssistantResponseRefusal => {
159 AssistantResponsePart::Refusal { id, delivered_at }
160 }
161 MessageTable::AssistantResponseReasoning => {
162 AssistantResponsePart::Reasoning { id, delivered_at }
163 }
164 MessageTable::AssistantResponseContentText => {
165 AssistantResponsePart::Text { id, delivered_at }
166 }
167 MessageTable::AssistantResponseContentImage => {
168 AssistantResponsePart::Image { id, delivered_at }
169 }
170 MessageTable::AssistantResponseContentAudio => {
171 AssistantResponsePart::Audio { id, delivered_at }
172 }
173 MessageTable::AssistantResponseContentVideo => {
174 AssistantResponsePart::Video { id, delivered_at }
175 }
176 MessageTable::AssistantResponseContentFile => {
177 AssistantResponsePart::File { id, delivered_at }
178 }
179 _ => return None,
180 })
181}
182
183fn tool_response_kind(t: MessageTable) -> Option<ToolResponsePartType> {
184 match t {
185 MessageTable::ToolResponseContentText => Some(ToolResponsePartType::Text),
186 MessageTable::ToolResponseContentImage => Some(ToolResponsePartType::Image),
187 MessageTable::ToolResponseContentAudio => Some(ToolResponsePartType::Audio),
188 MessageTable::ToolResponseContentVideo => Some(ToolResponsePartType::Video),
189 MessageTable::ToolResponseContentFile => Some(ToolResponsePartType::File),
190 _ => None,
191 }
192}
193
194const SELECT_SHAPE: &str = "SELECT \
200 m.\"index\" AS id, \
201 m.response_id, \
202 m.\"table\" AS table_kind, \
203 m.agent_instance_hierarchy, \
204 m.\"timestamp\" AS timestamp_delivered, \
205 CASE m.\"table\" \
206 WHEN 'message_queue_text' THEN mq.sender_agent_instance_hierarchy \
207 WHEN 'message_queue_image' THEN mq.sender_agent_instance_hierarchy \
208 WHEN 'message_queue_audio' THEN mq.sender_agent_instance_hierarchy \
209 WHEN 'message_queue_video' THEN mq.sender_agent_instance_hierarchy \
210 WHEN 'message_queue_file' THEN mq.sender_agent_instance_hierarchy \
211 WHEN 'agent_completion_request' THEN acr.sender_agent_instance_hierarchy \
212 WHEN 'vector_completion_request' THEN vcr.sender_agent_instance_hierarchy \
213 WHEN 'function_execution_request' THEN fer.sender_agent_instance_hierarchy \
214 ELSE NULL \
215 END AS sender_agent_instance_hierarchy, \
216 mq.id AS message_queue_id, \
217 mq.enqueued_at AS timestamp_queued, \
218 mq.key AS message_queue_key, \
219 atc.function_name AS function_name, \
220 COALESCE(atc.tool_call_id, tr.tool_call_id) AS tool_call_id, \
221 m.row_sub_index AS row_sub_index";
222
223const FROM_JOINS: &str = "FROM objectiveai.messages m \
224 LEFT JOIN objectiveai.message_queue_contents mqc \
225 ON m.row_index = mqc.id \
226 AND m.\"table\" IN ( \
227 'message_queue_text', 'message_queue_image', 'message_queue_audio', \
228 'message_queue_video', 'message_queue_file' \
229 ) \
230 LEFT JOIN objectiveai.message_queue mq ON mqc.message_queue_id = mq.id \
231 LEFT JOIN objectiveai.agent_completion_requests acr \
232 ON m.response_id = acr.response_id \
233 AND m.\"table\" = 'agent_completion_request' \
234 LEFT JOIN objectiveai.vector_completion_requests vcr \
235 ON m.response_id = vcr.response_id \
236 AND m.\"table\" = 'vector_completion_request' \
237 LEFT JOIN objectiveai.function_execution_requests fer \
238 ON m.response_id = fer.response_id \
239 AND m.\"table\" = 'function_execution_request' \
240 LEFT JOIN objectiveai.assistant_response_tool_calls atc \
241 ON m.response_id = atc.response_id \
242 AND m.row_index = atc.\"index\" \
243 AND m.row_sub_index = atc.tool_call_index \
244 AND m.\"table\" = 'assistant_response_tool_calls' \
245 LEFT JOIN objectiveai.tool_response tr \
246 ON m.response_id = tr.response_id \
247 AND m.row_index = tr.\"index\" \
248 AND m.\"table\" IN ( \
249 'tool_response_content_text', 'tool_response_content_image', \
250 'tool_response_content_audio', 'tool_response_content_video', \
251 'tool_response_content_file' \
252 )";
253
254fn row_into_msg(r: &sqlx::postgres::PgRow) -> Result<MsgRow, Error> {
255 Ok(MsgRow {
256 id: r.try_get("id")?,
257 response_id: r.try_get("response_id")?,
258 table_kind: r.try_get("table_kind")?,
259 agent_instance_hierarchy: r.try_get("agent_instance_hierarchy")?,
260 timestamp_delivered: r.try_get("timestamp_delivered")?,
261 sender_agent_instance_hierarchy: r.try_get("sender_agent_instance_hierarchy")?,
262 message_queue_id: r.try_get("message_queue_id")?,
263 timestamp_queued: r.try_get("timestamp_queued")?,
264 message_queue_key: r.try_get("message_queue_key")?,
265 function_name: r.try_get("function_name")?,
266 tool_call_id: r.try_get("tool_call_id")?,
267 row_sub_index: r.try_get("row_sub_index")?,
268 })
269}
270
271fn coalesce_into_blocks(rows: Vec<MsgRow>) -> Vec<ResponseItem> {
274 let mut out: Vec<ResponseItem> = Vec::new();
275 let mut cur_class: Option<BlockClass> = None;
276 let mut cur_aih: String = String::new();
277 let mut cur_rid: String = String::new();
278 let mut cur_sender: Option<String> = None;
281 let mut cur_mq_id: Option<i64> = None;
286 let mut cur_timestamp_queued: Option<i64> = None;
289 let mut cur_key: Option<String> = None;
293 let mut cur_tool_call_id: Option<String> = None;
297 let mut cur_notification_parts: Vec<ClientNotificationPart> = Vec::new();
298 let mut cur_assistant_parts: Vec<AssistantResponsePart> = Vec::new();
299 let mut cur_tool_parts: Vec<ToolResponsePart> = Vec::new();
300
301 let flush = |class: Option<BlockClass>,
302 aih: &mut String,
303 rid: &mut String,
304 sender: &mut Option<String>,
305 mq_id: &mut Option<i64>,
306 timestamp_queued: &mut Option<i64>,
307 key: &mut Option<String>,
308 tool_call_id: &mut Option<String>,
309 notification_parts: &mut Vec<ClientNotificationPart>,
310 assistant_parts: &mut Vec<AssistantResponsePart>,
311 tool_parts: &mut Vec<ToolResponsePart>,
312 out: &mut Vec<ResponseItem>| {
313 match class {
314 Some(BlockClass::ClientNotification) if !notification_parts.is_empty() => {
315 out.push(ResponseItem::ClientNotification {
316 agent_instance_hierarchy: std::mem::take(aih),
317 sender_agent_instance_hierarchy: sender.take().unwrap_or_default(),
318 response_id: std::mem::take(rid),
319 queued_at: unix_to_rfc3339(timestamp_queued.take().unwrap_or_default()),
320 key: key.take(),
321 parts: std::mem::take(notification_parts),
322 });
323 *mq_id = None;
324 }
325 Some(BlockClass::AssistantResponse) if !assistant_parts.is_empty() => {
326 out.push(ResponseItem::AssistantResponse {
327 agent_instance_hierarchy: std::mem::take(aih),
328 response_id: std::mem::take(rid),
329 parts: std::mem::take(assistant_parts),
330 });
331 }
332 Some(BlockClass::ToolResponse) if !tool_parts.is_empty() => {
333 out.push(ResponseItem::ToolResponse {
334 agent_instance_hierarchy: std::mem::take(aih),
335 response_id: std::mem::take(rid),
336 tool_call_id: tool_call_id.take().unwrap_or_default(),
337 parts: std::mem::take(tool_parts),
338 });
339 }
340 _ => {
341 aih.clear();
342 rid.clear();
343 *sender = None;
344 *mq_id = None;
345 *timestamp_queued = None;
346 *key = None;
347 *tool_call_id = None;
348 notification_parts.clear();
349 assistant_parts.clear();
350 tool_parts.clear();
351 }
352 }
353 };
354
355 for row in rows {
356 let class = block_class(row.table_kind);
357
358 match class {
360 BlockClass::AgentCompletionRequest => {
361 flush(
362 cur_class, &mut cur_aih, &mut cur_rid, &mut cur_sender,
363 &mut cur_mq_id, &mut cur_timestamp_queued, &mut cur_key,
364 &mut cur_tool_call_id,
365 &mut cur_notification_parts, &mut cur_assistant_parts,
366 &mut cur_tool_parts, &mut out,
367 );
368 out.push(ResponseItem::AgentCompletionRequest {
369 id: row.id,
370 agent_instance_hierarchy: row.agent_instance_hierarchy,
371 sender_agent_instance_hierarchy: row
372 .sender_agent_instance_hierarchy
373 .unwrap_or_default(),
374 delivered_at: unix_to_rfc3339(row.timestamp_delivered),
375 response_id: row.response_id,
376 });
377 cur_class = None;
378 continue;
379 }
380 BlockClass::VectorCompletionRequest => {
381 flush(
382 cur_class, &mut cur_aih, &mut cur_rid, &mut cur_sender,
383 &mut cur_mq_id, &mut cur_timestamp_queued, &mut cur_key,
384 &mut cur_tool_call_id,
385 &mut cur_notification_parts, &mut cur_assistant_parts,
386 &mut cur_tool_parts, &mut out,
387 );
388 out.push(ResponseItem::VectorCompletionRequest {
389 id: row.id,
390 agent_instance_hierarchy: row.agent_instance_hierarchy,
391 sender_agent_instance_hierarchy: row
392 .sender_agent_instance_hierarchy
393 .unwrap_or_default(),
394 delivered_at: unix_to_rfc3339(row.timestamp_delivered),
395 response_id: row.response_id,
396 });
397 cur_class = None;
398 continue;
399 }
400 BlockClass::FunctionExecutionRequest => {
401 flush(
402 cur_class, &mut cur_aih, &mut cur_rid, &mut cur_sender,
403 &mut cur_mq_id, &mut cur_timestamp_queued, &mut cur_key,
404 &mut cur_tool_call_id,
405 &mut cur_notification_parts, &mut cur_assistant_parts,
406 &mut cur_tool_parts, &mut out,
407 );
408 out.push(ResponseItem::FunctionExecutionRequest {
409 id: row.id,
410 agent_instance_hierarchy: row.agent_instance_hierarchy,
411 sender_agent_instance_hierarchy: row
412 .sender_agent_instance_hierarchy
413 .unwrap_or_default(),
414 delivered_at: unix_to_rfc3339(row.timestamp_delivered),
415 response_id: row.response_id,
416 });
417 cur_class = None;
418 continue;
419 }
420 _ => {}
421 }
422
423 let boundary = cur_class != Some(class)
429 || cur_aih != row.agent_instance_hierarchy
430 || cur_rid != row.response_id
431 || (class == BlockClass::ClientNotification
432 && (cur_sender.as_deref() != row.sender_agent_instance_hierarchy.as_deref()
433 || cur_mq_id != row.message_queue_id))
434 || (class == BlockClass::ToolResponse
435 && cur_tool_call_id.as_deref() != row.tool_call_id.as_deref());
436 if boundary {
437 flush(
438 cur_class, &mut cur_aih, &mut cur_rid, &mut cur_sender,
439 &mut cur_mq_id, &mut cur_timestamp_queued, &mut cur_key,
440 &mut cur_tool_call_id,
441 &mut cur_notification_parts, &mut cur_assistant_parts,
442 &mut cur_tool_parts, &mut out,
443 );
444 cur_class = Some(class);
445 cur_aih = row.agent_instance_hierarchy.clone();
446 cur_rid = row.response_id.clone();
447 if class == BlockClass::ClientNotification {
448 cur_sender = row.sender_agent_instance_hierarchy.clone();
449 cur_mq_id = row.message_queue_id;
450 cur_timestamp_queued = row.timestamp_queued;
451 cur_key = row.message_queue_key.clone();
452 cur_tool_call_id = None;
453 } else if class == BlockClass::ToolResponse {
454 cur_sender = None;
455 cur_mq_id = None;
456 cur_timestamp_queued = None;
457 cur_key = None;
458 cur_tool_call_id = row.tool_call_id.clone();
459 } else {
460 cur_sender = None;
461 cur_mq_id = None;
462 cur_timestamp_queued = None;
463 cur_key = None;
464 cur_tool_call_id = None;
465 }
466 }
467
468 match class {
469 BlockClass::ClientNotification => {
470 let r#type = client_notification_kind(row.table_kind)
471 .expect("class invariant: ClientNotification maps to message_queue_*");
472 cur_notification_parts.push(ClientNotificationPart {
473 id: row.id,
474 delivered_at: unix_to_rfc3339(row.timestamp_delivered),
475 r#type,
476 });
477 }
478 BlockClass::AssistantResponse => {
479 let part = assistant_response_part(&row)
480 .expect("class invariant: AssistantResponse maps to assistant_response_*");
481 cur_assistant_parts.push(part);
482 }
483 BlockClass::ToolResponse => {
484 let r#type = tool_response_kind(row.table_kind)
485 .expect("class invariant: ToolResponse maps to tool_response*");
486 cur_tool_parts.push(ToolResponsePart {
487 id: row.id,
488 delivered_at: unix_to_rfc3339(row.timestamp_delivered),
489 r#type,
490 });
491 }
492 _ => unreachable!("request classes handled above"),
493 }
494 }
495
496 flush(
497 cur_class, &mut cur_aih, &mut cur_rid, &mut cur_sender,
498 &mut cur_mq_id, &mut cur_timestamp_queued, &mut cur_key,
499 &mut cur_tool_call_id,
500 &mut cur_notification_parts, &mut cur_assistant_parts,
501 &mut cur_tool_parts, &mut out,
502 );
503
504 out
505}
506
507pub async fn read_all_for_hierarchy(
511 pool: &Pool,
512 agent_instance_hierarchy: &str,
513 after_id: Option<i64>,
514 limit: Option<i64>,
515) -> Result<Vec<ResponseItem>, Error> {
516 let sql = format!(
517 "{select} {from} \
518 WHERE m.agent_instance_hierarchy = $1 \
519 AND m.\"index\" > COALESCE($2, 0) \
520 ORDER BY m.\"index\" ASC \
521 LIMIT $3",
522 select = SELECT_SHAPE,
523 from = FROM_JOINS,
524 );
525 let rows = sqlx::query(&sql)
526 .bind(agent_instance_hierarchy)
527 .bind(after_id)
528 .bind(limit)
529 .fetch_all(&**pool)
530 .await?;
531
532 let msg_rows: Vec<MsgRow> = rows.iter().map(row_into_msg).collect::<Result<_, _>>()?;
533 Ok(coalesce_into_blocks(msg_rows))
534}
535
536pub async fn read_pending_for_parent(
543 pool: &Pool,
544 parent_agent_instance_hierarchy: &str,
545 after_id: Option<i64>,
546 limit: Option<i64>,
547) -> Result<Vec<ResponseItem>, Error> {
548 let sql = format!(
560 "WITH selected AS ( \
561 {select} \
562 {from} \
563 JOIN objectiveai.messages_queue q \
564 ON q.spawned_agent_instance_hierarchy = m.agent_instance_hierarchy \
565 WHERE q.parent_agent_instance_hierarchy = $1 \
566 AND m.\"index\" > GREATEST(q.read_index, COALESCE($2, 0)) \
567 ORDER BY m.\"index\" ASC \
568 LIMIT $3 \
569 ), \
570 maxes AS ( \
571 SELECT agent_instance_hierarchy AS spawned, MAX(id) AS max_id \
572 FROM selected \
573 GROUP BY agent_instance_hierarchy \
574 ), \
575 bump AS ( \
576 UPDATE objectiveai.messages_queue qq \
577 SET read_index = GREATEST(qq.read_index, mx.max_id) \
578 FROM maxes mx \
579 WHERE qq.parent_agent_instance_hierarchy = $1 \
580 AND qq.spawned_agent_instance_hierarchy = mx.spawned \
581 RETURNING 1 \
582 ) \
583 SELECT s.id, s.response_id, s.table_kind, \
584 s.agent_instance_hierarchy, s.timestamp_delivered, \
585 s.sender_agent_instance_hierarchy, \
586 s.message_queue_id, s.timestamp_queued, \
587 s.message_queue_key, s.function_name, \
588 s.tool_call_id, s.row_sub_index \
589 FROM selected s \
590 ORDER BY s.id ASC",
591 select = SELECT_SHAPE,
592 from = FROM_JOINS,
593 );
594 let rows = sqlx::query(&sql)
595 .bind(parent_agent_instance_hierarchy)
596 .bind(after_id)
597 .bind(limit)
598 .fetch_all(&**pool)
599 .await?;
600
601 let msg_rows: Vec<MsgRow> = rows.iter().map(row_into_msg).collect::<Result<_, _>>()?;
602 Ok(coalesce_into_blocks(msg_rows))
603}
604
605pub async fn any_pending_matching_kinds(
621 pool: &Pool,
622 parent_agent_instance_hierarchy: &str,
623 after_id: Option<i64>,
624 kinds: Option<&[MessageTable]>,
625) -> Result<bool, Error> {
626 let kinds_clause = match kinds {
627 Some(ks) if !ks.is_empty() => {
628 let list = ks
629 .iter()
630 .map(|k| format!("'{}'", k.schema_name()))
631 .collect::<Vec<_>>()
632 .join(", ");
633 format!("AND m.\"table\" IN ({list})")
634 }
635 _ => String::new(),
636 };
637 let sql = format!(
638 "SELECT EXISTS( \
639 SELECT 1 FROM objectiveai.messages m \
640 JOIN objectiveai.messages_queue q \
641 ON q.spawned_agent_instance_hierarchy = m.agent_instance_hierarchy \
642 WHERE q.parent_agent_instance_hierarchy = $1 \
643 AND m.\"index\" > GREATEST(q.read_index, COALESCE($2, 0)) \
644 {kinds_clause} \
645 )"
646 );
647 let exists: bool = sqlx::query_scalar(&sql)
648 .bind(parent_agent_instance_hierarchy)
649 .bind(after_id)
650 .fetch_one(&**pool)
651 .await?;
652 Ok(exists)
653}