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