1use objectiveai_sdk::agent::completions::message::{File, ImageUrl, InputAudio, VideoUrl};
43use serde::Serialize;
44
45use crate::db::{Error, Pool};
46
47use super::row::{MessageTable, RowValue};
48use super::shadow::WriteOp;
49
50pub async fn update_agent_token_usage(
53 pool: &Pool,
54 agent_instance_hierarchy: &str,
55 total_tokens: i64,
56) -> Result<(), Error> {
57 sqlx::query(
58 "INSERT INTO objectiveai.agent_token_usage \
59 (agent_instance_hierarchy, total_tokens) \
60 VALUES ($1, $2) \
61 ON CONFLICT (agent_instance_hierarchy) \
62 DO UPDATE SET total_tokens = EXCLUDED.total_tokens",
63 )
64 .bind(agent_instance_hierarchy)
65 .bind(total_tokens)
66 .execute(&**pool)
67 .await?;
68 Ok(())
69}
70
71pub async fn write_value<'a>(
73 pool: &Pool,
74 op: WriteOp,
75 value: &RowValue<'a>,
76 timestamp: i64,
77) -> Result<(), Error> {
78 match op {
79 WriteOp::Skip => Ok(()),
80 WriteOp::Insert => insert_value(pool, value, timestamp).await,
81 WriteOp::Update => update_value(pool, value).await,
82 }
83}
84
85async fn insert_value<'a>(
86 pool: &Pool,
87 value: &RowValue<'a>,
88 timestamp: i64,
89) -> Result<(), Error> {
90 if let RowValue::MessageQueueContent {
96 response_id,
97 agent_instance_hierarchy,
98 message_queue_content_id,
99 } = *value
100 {
101 return insert_message_queue_content_with_msg(
102 pool,
103 response_id,
104 agent_instance_hierarchy,
105 message_queue_content_id,
106 timestamp,
107 )
108 .await;
109 }
110
111 if let RowValue::ToolResponse { response_id, index, tool_call_id, .. } = *value {
117 sqlx::query(
118 "INSERT INTO objectiveai.tool_response (response_id, \"index\", tool_call_id) \
119 VALUES ($1, $2, $3)",
120 )
121 .bind(response_id)
122 .bind(index as i64)
123 .bind(tool_call_id)
124 .execute(&**pool)
125 .await?;
126 return Ok(());
127 }
128
129 let mt = value.message_table();
130 let hier = value.agent_instance_hierarchy();
131 let row_index = value.row_index();
132 let row_sub_index = value.row_sub_index();
133 let response_id = value.response_id();
134
135 match *value {
136 RowValue::MessageQueueContent { .. } => unreachable!(
137 "MessageQueueContent handled by early-return branch above"
138 ),
139 RowValue::ToolResponse { .. } => unreachable!(
140 "ToolResponse handled by early-return branch above"
141 ),
142 RowValue::AssistantResponseRefusal { text, .. } => {
143 sqlx::query(
144 "WITH data_ins AS (\
145 INSERT INTO objectiveai.assistant_response_refusal (response_id, \"index\", text) \
146 VALUES ($1, $2, $3) RETURNING response_id\
147 )\
148 INSERT INTO objectiveai.messages \
149 (response_id, \"table\", row_index, row_sub_index, \
150 agent_instance_hierarchy, \"timestamp\") \
151 SELECT $1, $4, $5, $6, $7, $8 FROM data_ins",
152 )
153 .bind(response_id)
154 .bind(row_index)
155 .bind(text)
156 .bind(mt)
157 .bind(row_index)
158 .bind(row_sub_index)
159 .bind(hier)
160 .bind(timestamp)
161 .execute(&**pool)
162 .await?;
163 }
164 RowValue::AssistantResponseReasoning { text, .. } => {
165 sqlx::query(
166 "WITH data_ins AS (\
167 INSERT INTO objectiveai.assistant_response_reasoning (response_id, \"index\", text) \
168 VALUES ($1, $2, $3) RETURNING response_id\
169 )\
170 INSERT INTO objectiveai.messages \
171 (response_id, \"table\", row_index, row_sub_index, \
172 agent_instance_hierarchy, \"timestamp\") \
173 SELECT $1, $4, $5, $6, $7, $8 FROM data_ins",
174 )
175 .bind(response_id)
176 .bind(row_index)
177 .bind(text)
178 .bind(mt)
179 .bind(row_index)
180 .bind(row_sub_index)
181 .bind(hier)
182 .bind(timestamp)
183 .execute(&**pool)
184 .await?;
185 }
186 RowValue::AssistantResponseToolCalls {
187 tool_call_index, tool_call_id, function_name, arguments, ..
188 } => {
189 sqlx::query(
190 "WITH data_ins AS (\
191 INSERT INTO objectiveai.assistant_response_tool_calls \
192 (response_id, \"index\", tool_call_index, tool_call_id, function_name, arguments) \
193 VALUES ($1, $2, $3, $4, $5, $6) RETURNING response_id\
194 )\
195 INSERT INTO objectiveai.messages \
196 (response_id, \"table\", row_index, row_sub_index, \
197 agent_instance_hierarchy, \"timestamp\") \
198 SELECT $1, $7, $8, $9, $10, $11 FROM data_ins",
199 )
200 .bind(response_id)
201 .bind(row_index)
202 .bind(tool_call_index as i64)
203 .bind(tool_call_id)
204 .bind(function_name)
205 .bind(arguments)
206 .bind(mt)
207 .bind(row_index)
208 .bind(row_sub_index)
209 .bind(hier)
210 .bind(timestamp)
211 .execute(&**pool)
212 .await?;
213 }
214 RowValue::AssistantResponseContentText { text, .. } => {
215 insert_text_part_with_msg(pool, "objectiveai.assistant_response_content_text", value, text, timestamp).await?;
216 }
217 RowValue::ToolResponseContentText { text, .. } => {
218 insert_text_part_with_msg(pool, "objectiveai.tool_response_content_text", value, text, timestamp).await?;
219 }
220 RowValue::AssistantResponseContentImage { image_url, .. } => {
221 insert_image_part_with_msg(pool, "objectiveai.assistant_response_content_image", value, image_url, timestamp).await?;
222 }
223 RowValue::ToolResponseContentImage { image_url, .. } => {
224 insert_image_part_with_msg(pool, "objectiveai.tool_response_content_image", value, image_url, timestamp).await?;
225 }
226 RowValue::AssistantResponseContentAudio { input_audio, .. } => {
227 insert_audio_part_with_msg(pool, "objectiveai.assistant_response_content_audio", value, input_audio, timestamp).await?;
228 }
229 RowValue::ToolResponseContentAudio { input_audio, .. } => {
230 insert_audio_part_with_msg(pool, "objectiveai.tool_response_content_audio", value, input_audio, timestamp).await?;
231 }
232 RowValue::AssistantResponseContentVideo { video_url, .. } => {
233 insert_video_part_with_msg(pool, "objectiveai.assistant_response_content_video", value, video_url, timestamp).await?;
234 }
235 RowValue::ToolResponseContentVideo { video_url, .. } => {
236 insert_video_part_with_msg(pool, "objectiveai.tool_response_content_video", value, video_url, timestamp).await?;
237 }
238 RowValue::AssistantResponseContentFile { file, .. } => {
239 insert_file_part_with_msg(pool, "objectiveai.assistant_response_content_file", value, file, timestamp).await?;
240 }
241 RowValue::ToolResponseContentFile { file, .. } => {
242 insert_file_part_with_msg(pool, "objectiveai.tool_response_content_file", value, file, timestamp).await?;
243 }
244 }
245 Ok(())
246}
247
248async fn update_value<'a>(pool: &Pool, value: &RowValue<'a>) -> Result<(), Error> {
249 if matches!(value, RowValue::MessageQueueContent { .. }) {
253 return Ok(());
254 }
255
256 if let RowValue::ToolResponse { response_id, index, tool_call_id, .. } = *value {
264 sqlx::query(
265 "UPDATE objectiveai.tool_response SET tool_call_id = $1 \
266 WHERE response_id = $2 AND \"index\" = $3",
267 )
268 .bind(tool_call_id)
269 .bind(response_id)
270 .bind(index as i64)
271 .execute(&**pool)
272 .await?;
273 return Ok(());
274 }
275
276 let mt = value.message_table();
277 let hier = value.agent_instance_hierarchy();
278 let row_index = value.row_index();
279 let row_sub_index = value.row_sub_index();
280 let response_id = value.response_id();
281
282 match *value {
283 RowValue::MessageQueueContent { .. } => unreachable!(
284 "MessageQueueContent handled by short-circuit above"
285 ),
286 RowValue::ToolResponse { .. } => unreachable!(
287 "ToolResponse handled by short-circuit above"
288 ),
289 RowValue::AssistantResponseRefusal { text, .. } => {
290 run_update_with_downgrade(
291 pool,
292 "UPDATE objectiveai.assistant_response_refusal SET text = $A \
293 WHERE response_id = $RESP AND \"index\" = $RI",
294 response_id, row_index, row_sub_index, mt, hier,
295 &[("A", BindVal::Str(text))],
296 &[BindIdx::Resp, BindIdx::Ri],
297 ).await?;
298 }
299 RowValue::AssistantResponseReasoning { text, .. } => {
300 run_update_with_downgrade(
301 pool,
302 "UPDATE objectiveai.assistant_response_reasoning SET text = $A \
303 WHERE response_id = $RESP AND \"index\" = $RI",
304 response_id, row_index, row_sub_index, mt, hier,
305 &[("A", BindVal::Str(text))],
306 &[BindIdx::Resp, BindIdx::Ri],
307 ).await?;
308 }
309 RowValue::AssistantResponseToolCalls { tool_call_index, tool_call_id, function_name, arguments, .. } => {
310 run_update_with_downgrade(
311 pool,
312 "UPDATE objectiveai.assistant_response_tool_calls SET tool_call_id = $A, function_name = $B, arguments = $C \
313 WHERE response_id = $RESP AND \"index\" = $RI AND tool_call_index = $RSI",
314 response_id, row_index, row_sub_index, mt, hier,
315 &[("A", BindVal::Str(tool_call_id)), ("B", BindVal::Str(function_name)), ("C", BindVal::Str(arguments))],
316 &[BindIdx::Resp, BindIdx::Ri, BindIdx::Rsi],
317 ).await?;
318 let _ = tool_call_index;
319 }
320 RowValue::AssistantResponseContentText { text, .. }
321 | RowValue::ToolResponseContentText { text, .. } => {
322 let table = match *value {
323 RowValue::AssistantResponseContentText { .. } => "objectiveai.assistant_response_content_text",
324 _ => "objectiveai.tool_response_content_text",
325 };
326 let sql = format!(
327 "UPDATE {table} SET text = $A \
328 WHERE response_id = $RESP AND \"index\" = $RI AND part_index = $RSI"
329 );
330 run_update_with_downgrade(
331 pool, &sql,
332 response_id, row_index, row_sub_index, mt, hier,
333 &[("A", BindVal::Str(text))],
334 &[BindIdx::Resp, BindIdx::Ri, BindIdx::Rsi],
335 ).await?;
336 }
337 RowValue::AssistantResponseContentImage { image_url, .. }
338 | RowValue::ToolResponseContentImage { image_url, .. } => {
339 let table = match *value {
340 RowValue::AssistantResponseContentImage { .. } => "objectiveai.assistant_response_content_image",
341 _ => "objectiveai.tool_response_content_image",
342 };
343 let detail = image_url.detail.as_ref().and_then(|d| serde_json::to_string(d).ok());
344 let sql = format!(
345 "UPDATE {table} SET url = $A, detail = $B \
346 WHERE response_id = $RESP AND \"index\" = $RI AND part_index = $RSI"
347 );
348 run_update_with_downgrade(
349 pool, &sql,
350 response_id, row_index, row_sub_index, mt, hier,
351 &[("A", BindVal::Str(image_url.url.as_str())), ("B", BindVal::OptString(detail))],
352 &[BindIdx::Resp, BindIdx::Ri, BindIdx::Rsi],
353 ).await?;
354 }
355 RowValue::AssistantResponseContentAudio { input_audio, .. }
356 | RowValue::ToolResponseContentAudio { input_audio, .. } => {
357 let table = match *value {
358 RowValue::AssistantResponseContentAudio { .. } => "objectiveai.assistant_response_content_audio",
359 _ => "objectiveai.tool_response_content_audio",
360 };
361 let sql = format!(
362 "UPDATE {table} SET data = $A, format = $B \
363 WHERE response_id = $RESP AND \"index\" = $RI AND part_index = $RSI"
364 );
365 run_update_with_downgrade(
366 pool, &sql,
367 response_id, row_index, row_sub_index, mt, hier,
368 &[
369 ("A", BindVal::Str(input_audio.data.as_str())),
370 ("B", BindVal::Str(input_audio.format.as_str())),
371 ],
372 &[BindIdx::Resp, BindIdx::Ri, BindIdx::Rsi],
373 ).await?;
374 }
375 RowValue::AssistantResponseContentVideo { video_url, .. }
376 | RowValue::ToolResponseContentVideo { video_url, .. } => {
377 let table = match *value {
378 RowValue::AssistantResponseContentVideo { .. } => "objectiveai.assistant_response_content_video",
379 _ => "objectiveai.tool_response_content_video",
380 };
381 let sql = format!(
382 "UPDATE {table} SET url = $A \
383 WHERE response_id = $RESP AND \"index\" = $RI AND part_index = $RSI"
384 );
385 run_update_with_downgrade(
386 pool, &sql,
387 response_id, row_index, row_sub_index, mt, hier,
388 &[("A", BindVal::Str(video_url.url.as_str()))],
389 &[BindIdx::Resp, BindIdx::Ri, BindIdx::Rsi],
390 ).await?;
391 }
392 RowValue::AssistantResponseContentFile { file, .. }
393 | RowValue::ToolResponseContentFile { file, .. } => {
394 let table = match *value {
395 RowValue::AssistantResponseContentFile { .. } => "objectiveai.assistant_response_content_file",
396 _ => "objectiveai.tool_response_content_file",
397 };
398 let sql = format!(
399 "UPDATE {table} SET file_data = $A, file_id = $B, filename = $C, file_url = $D \
400 WHERE response_id = $RESP AND \"index\" = $RI AND part_index = $RSI"
401 );
402 run_update_with_downgrade(
403 pool, &sql,
404 response_id, row_index, row_sub_index, mt, hier,
405 &[
406 ("A", BindVal::OptStr(file.file_data.as_deref())),
407 ("B", BindVal::OptStr(file.file_id.as_deref())),
408 ("C", BindVal::OptStr(file.filename.as_deref())),
409 ("D", BindVal::OptStr(file.file_url.as_deref())),
410 ],
411 &[BindIdx::Resp, BindIdx::Ri, BindIdx::Rsi],
412 ).await?;
413 }
414 }
415 Ok(())
416}
417
418async fn insert_text_part_with_msg<'a>(
421 pool: &Pool,
422 table: &str,
423 value: &RowValue<'a>,
424 text: &str,
425 timestamp: i64,
426) -> Result<(), Error> {
427 let sql = format!(
428 "WITH data_ins AS (\
429 INSERT INTO {table} (response_id, \"index\", part_index, text) \
430 VALUES ($1, $2, $3, $4) RETURNING response_id\
431 )\
432 INSERT INTO objectiveai.messages \
433 (response_id, \"table\", row_index, row_sub_index, \
434 agent_instance_hierarchy, \"timestamp\") \
435 SELECT $1, $5, $6, $7, $8, $9 FROM data_ins"
436 );
437 sqlx::query(&sql)
438 .bind(value.response_id())
439 .bind(value.row_index())
440 .bind(value.row_sub_index())
441 .bind(text)
442 .bind(value.message_table())
443 .bind(value.row_index())
444 .bind(value.row_sub_index())
445 .bind(value.agent_instance_hierarchy())
446 .bind(timestamp)
447 .execute(&**pool)
448 .await?;
449 Ok(())
450}
451
452async fn insert_image_part_with_msg<'a>(
453 pool: &Pool,
454 table: &str,
455 value: &RowValue<'a>,
456 image: &ImageUrl,
457 timestamp: i64,
458) -> Result<(), Error> {
459 let detail = image.detail.as_ref().and_then(|d| serde_json::to_string(d).ok());
460 let sql = format!(
461 "WITH data_ins AS (\
462 INSERT INTO {table} (response_id, \"index\", part_index, url, detail) \
463 VALUES ($1, $2, $3, $4, $5) RETURNING response_id\
464 )\
465 INSERT INTO objectiveai.messages \
466 (response_id, \"table\", row_index, row_sub_index, \
467 agent_instance_hierarchy, \"timestamp\") \
468 SELECT $1, $6, $7, $8, $9, $10 FROM data_ins"
469 );
470 sqlx::query(&sql)
471 .bind(value.response_id())
472 .bind(value.row_index())
473 .bind(value.row_sub_index())
474 .bind(image.url.as_str())
475 .bind(detail)
476 .bind(value.message_table())
477 .bind(value.row_index())
478 .bind(value.row_sub_index())
479 .bind(value.agent_instance_hierarchy())
480 .bind(timestamp)
481 .execute(&**pool)
482 .await?;
483 Ok(())
484}
485
486async fn insert_audio_part_with_msg<'a>(
487 pool: &Pool,
488 table: &str,
489 value: &RowValue<'a>,
490 audio: &InputAudio,
491 timestamp: i64,
492) -> Result<(), Error> {
493 let sql = format!(
494 "WITH data_ins AS (\
495 INSERT INTO {table} (response_id, \"index\", part_index, data, format) \
496 VALUES ($1, $2, $3, $4, $5) RETURNING response_id\
497 )\
498 INSERT INTO objectiveai.messages \
499 (response_id, \"table\", row_index, row_sub_index, \
500 agent_instance_hierarchy, \"timestamp\") \
501 SELECT $1, $6, $7, $8, $9, $10 FROM data_ins"
502 );
503 sqlx::query(&sql)
504 .bind(value.response_id())
505 .bind(value.row_index())
506 .bind(value.row_sub_index())
507 .bind(audio.data.as_str())
508 .bind(audio.format.as_str())
509 .bind(value.message_table())
510 .bind(value.row_index())
511 .bind(value.row_sub_index())
512 .bind(value.agent_instance_hierarchy())
513 .bind(timestamp)
514 .execute(&**pool)
515 .await?;
516 Ok(())
517}
518
519async fn insert_message_queue_content_with_msg(
533 pool: &Pool,
534 response_id: &str,
535 agent_instance_hierarchy: &str,
536 message_queue_content_id: i64,
537 timestamp: i64,
538) -> Result<(), Error> {
539 sqlx::query(
540 "WITH content AS (\
541 SELECT id, kind, message_queue_id \
542 FROM objectiveai.message_queue_contents \
543 WHERE id = $1 \
544 ), \
545 flip AS (\
546 UPDATE objectiveai.message_queue \
547 SET active = FALSE \
548 WHERE id = (SELECT message_queue_id FROM content) \
549 AND active = TRUE \
550 RETURNING id \
551 ) \
552 INSERT INTO objectiveai.messages \
553 (response_id, \"table\", row_index, row_sub_index, \
554 agent_instance_hierarchy, \"timestamp\") \
555 SELECT $2, \
556 CASE (SELECT kind FROM content) \
557 WHEN 'text' THEN 'message_queue_text'::objectiveai.message_table \
558 WHEN 'image' THEN 'message_queue_image'::objectiveai.message_table \
559 WHEN 'audio' THEN 'message_queue_audio'::objectiveai.message_table \
560 WHEN 'video' THEN 'message_queue_video'::objectiveai.message_table \
561 WHEN 'file' THEN 'message_queue_file'::objectiveai.message_table \
562 END, \
563 $1, NULL, $3, $4 \
564 FROM content",
565 )
566 .bind(message_queue_content_id)
567 .bind(response_id)
568 .bind(agent_instance_hierarchy)
569 .bind(timestamp)
570 .execute(&**pool)
571 .await?;
572 Ok(())
573}
574
575async fn insert_video_part_with_msg<'a>(
576 pool: &Pool,
577 table: &str,
578 value: &RowValue<'a>,
579 video: &VideoUrl,
580 timestamp: i64,
581) -> Result<(), Error> {
582 let sql = format!(
583 "WITH data_ins AS (\
584 INSERT INTO {table} (response_id, \"index\", part_index, url) \
585 VALUES ($1, $2, $3, $4) RETURNING response_id\
586 )\
587 INSERT INTO objectiveai.messages \
588 (response_id, \"table\", row_index, row_sub_index, \
589 agent_instance_hierarchy, \"timestamp\") \
590 SELECT $1, $5, $6, $7, $8, $9 FROM data_ins"
591 );
592 sqlx::query(&sql)
593 .bind(value.response_id())
594 .bind(value.row_index())
595 .bind(value.row_sub_index())
596 .bind(video.url.as_str())
597 .bind(value.message_table())
598 .bind(value.row_index())
599 .bind(value.row_sub_index())
600 .bind(value.agent_instance_hierarchy())
601 .bind(timestamp)
602 .execute(&**pool)
603 .await?;
604 Ok(())
605}
606
607async fn insert_file_part_with_msg<'a>(
608 pool: &Pool,
609 table: &str,
610 value: &RowValue<'a>,
611 file: &File,
612 timestamp: i64,
613) -> Result<(), Error> {
614 let sql = format!(
615 "WITH data_ins AS (\
616 INSERT INTO {table} (response_id, \"index\", part_index, file_data, file_id, filename, file_url) \
617 VALUES ($1, $2, $3, $4, $5, $6, $7) RETURNING response_id\
618 )\
619 INSERT INTO objectiveai.messages \
620 (response_id, \"table\", row_index, row_sub_index, \
621 agent_instance_hierarchy, \"timestamp\") \
622 SELECT $1, $8, $9, $10, $11, $12 FROM data_ins"
623 );
624 sqlx::query(&sql)
625 .bind(value.response_id())
626 .bind(value.row_index())
627 .bind(value.row_sub_index())
628 .bind(file.file_data.as_deref())
629 .bind(file.file_id.as_deref())
630 .bind(file.filename.as_deref())
631 .bind(file.file_url.as_deref())
632 .bind(value.message_table())
633 .bind(value.row_index())
634 .bind(value.row_sub_index())
635 .bind(value.agent_instance_hierarchy())
636 .bind(timestamp)
637 .execute(&**pool)
638 .await?;
639 Ok(())
640}
641
642#[derive(Clone, Copy)]
651enum BindIdx {
652 Resp,
653 Ri,
654 Rsi,
655}
656
657enum BindVal<'a> {
658 Str(&'a str),
659 OptStr(Option<&'a str>),
660 OptString(Option<String>),
661 Bool(bool),
662}
663
664#[allow(clippy::too_many_arguments)]
665async fn run_update_with_downgrade<'a>(
666 pool: &Pool,
667 update_sql_template: &str,
668 response_id: &str,
669 row_index: i64,
670 row_sub_index: Option<i64>,
671 message_table: MessageTable,
672 agent_instance_hierarchy: &str,
673 extra_binds: &[(&str, BindVal<'a>)],
674 update_where_binds: &[BindIdx],
675) -> Result<(), Error> {
676 let mut sql = update_sql_template.to_string();
685 let mut pos = 1usize;
686
687 let mut resp_pos: Option<usize> = None;
692 for slot in update_where_binds {
693 let idx = pos;
694 pos += 1;
695 match slot {
696 BindIdx::Resp => {
697 sql = sql.replace("$RESP", &format!("${idx}"));
698 resp_pos = Some(idx);
699 }
700 BindIdx::Ri => {
701 sql = sql.replace("$RI", &format!("${idx}"));
702 }
703 BindIdx::Rsi => {
704 sql = sql.replace("$RSI", &format!("${idx}"));
705 }
706 }
707 }
708 let resp_pos = resp_pos.expect("Resp bind required");
709
710 for (token, _val) in extra_binds {
712 let idx = pos;
713 pos += 1;
714 sql = sql.replace(&format!("${token}"), &format!("${idx}"));
715 }
716
717 let mt_pos = pos; pos += 1;
718 let ri_for_msg_pos = pos; pos += 1;
719 let rsi_for_msg_pos = pos; pos += 1;
720 let hier_pos = pos;
721
722 let final_sql = format!(
723 "WITH \
724 data_upd AS ({sql} RETURNING response_id),\
725 msg AS (\
726 SELECT \"index\" AS msg_index FROM objectiveai.messages \
727 WHERE response_id = ${resp_pos} \
728 AND \"table\" = ${mt_pos} \
729 AND row_index IS NOT DISTINCT FROM ${ri_for_msg_pos} \
730 AND row_sub_index IS NOT DISTINCT FROM ${rsi_for_msg_pos}\
731 )\
732 UPDATE objectiveai.messages_queue \
733 SET read_index = msg.msg_index - 1 \
734 FROM msg, data_upd \
735 WHERE spawned_agent_instance_hierarchy = ${hier_pos} \
736 AND read_index >= msg.msg_index",
737 );
738
739 let mut q = sqlx::query(&final_sql);
740 for slot in update_where_binds {
742 q = match slot {
743 BindIdx::Resp => q.bind(response_id),
744 BindIdx::Ri => q.bind(row_index),
745 BindIdx::Rsi => q.bind(row_sub_index),
746 };
747 }
748 for (_, val) in extra_binds {
750 q = match val {
751 BindVal::Str(s) => q.bind(*s),
752 BindVal::OptStr(s) => q.bind(*s),
753 BindVal::OptString(s) => q.bind(s.clone()),
754 BindVal::Bool(b) => q.bind(*b),
755 };
756 }
757 q = q.bind(message_table);
759 q = q.bind(row_index);
760 q = q.bind(row_sub_index);
761 q = q.bind(agent_instance_hierarchy);
762
763 q.execute(&**pool).await?;
764 Ok(())
765}
766
767#[derive(Debug, Clone, Copy, PartialEq, Eq)]
772pub enum Tier {
773 Agent,
774 Vector,
775 Function,
776}
777
778impl Tier {
779 pub fn request_table(self) -> &'static str {
780 match self {
781 Tier::Agent => "objectiveai.agent_completion_requests",
782 Tier::Vector => "objectiveai.vector_completion_requests",
783 Tier::Function => "objectiveai.function_execution_requests",
784 }
785 }
786 pub fn response_table(self) -> &'static str {
787 match self {
788 Tier::Agent => "objectiveai.agent_completion_responses",
789 Tier::Vector => "objectiveai.vector_completion_responses",
790 Tier::Function => "objectiveai.function_execution_responses",
791 }
792 }
793 pub fn request_message_table(self) -> MessageTable {
796 match self {
797 Tier::Agent => MessageTable::AgentCompletionRequest,
798 Tier::Vector => MessageTable::VectorCompletionRequest,
799 Tier::Function => MessageTable::FunctionExecutionRequest,
800 }
801 }
802}
803
804pub async fn insert_request_blob<P: Serialize>(
812 pool: &Pool,
813 tier: Tier,
814 response_id: &str,
815 params: &P,
816 sender_agent_instance_hierarchy: &str,
817 timestamp: i64,
818) -> Result<(), Error> {
819 let body = serde_json::to_value(params)?;
820 let sql = format!(
821 "INSERT INTO {table} \
822 (response_id, body, created_at, sender_agent_instance_hierarchy) \
823 VALUES ($1, $2, $3, $4)",
824 table = tier.request_table()
825 );
826 sqlx::query(&sql)
827 .bind(response_id)
828 .bind(sqlx::types::Json(body))
829 .bind(timestamp)
830 .bind(sender_agent_instance_hierarchy)
831 .execute(&**pool)
832 .await?;
833 Ok(())
834}
835
836pub async fn insert_request_messages_row(
845 pool: &Pool,
846 tier: Tier,
847 response_id: &str,
848 agent_instance_hierarchy: &str,
849 timestamp: i64,
850) -> Result<(), Error> {
851 sqlx::query(
852 "INSERT INTO objectiveai.messages \
853 (response_id, \"table\", row_index, row_sub_index, \
854 agent_instance_hierarchy, \"timestamp\") \
855 VALUES ($1, $2, NULL, NULL, $3, $4)",
856 )
857 .bind(response_id)
858 .bind(tier.request_message_table())
859 .bind(agent_instance_hierarchy)
860 .bind(timestamp)
861 .execute(&**pool)
862 .await?;
863 Ok(())
864}
865
866pub async fn insert_response_blob<C: Serialize>(
871 pool: &Pool,
872 tier: Tier,
873 response_id: &str,
874 chunk: &C,
875 created_at: i64,
876) -> Result<(), Error> {
877 let body = serde_json::to_value(chunk)?;
878 let sql = format!(
879 "INSERT INTO {table} (response_id, body, created_at) VALUES ($1, $2, $3)",
880 table = tier.response_table()
881 );
882 sqlx::query(&sql)
883 .bind(response_id)
884 .bind(sqlx::types::Json(body))
885 .bind(created_at)
886 .execute(&**pool)
887 .await?;
888 Ok(())
889}
890
891pub async fn update_response_blob<C: Serialize>(
893 pool: &Pool,
894 tier: Tier,
895 response_id: &str,
896 chunk: &C,
897 created_at: i64,
898) -> Result<(), Error> {
899 let body = serde_json::to_value(chunk)?;
900 let sql = format!(
901 "UPDATE {table} SET body = $2, created_at = $3 WHERE response_id = $1",
902 table = tier.response_table()
903 );
904 sqlx::query(&sql)
905 .bind(response_id)
906 .bind(sqlx::types::Json(body))
907 .bind(created_at)
908 .execute(&**pool)
909 .await?;
910 Ok(())
911}