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