agent_sdk_store_postgres/
provider_arguments.rs1use agent_sdk_core::{
2 AgentError, ProviderArgumentStore, domain::ContentRef as ProviderArgumentContentRef,
3 tool_records::CanonicalToolName,
4};
5use serde_json::Value;
6use sha2::{Digest, Sha256};
7
8use crate::PostgresStoreClient;
9
10#[derive(Clone)]
11pub struct PostgresProviderArgumentStore {
12 client: PostgresStoreClient,
13}
14
15impl PostgresProviderArgumentStore {
16 pub fn new(client: PostgresStoreClient) -> Self {
17 Self { client }
18 }
19}
20
21impl ProviderArgumentStore for PostgresProviderArgumentStore {
22 fn store_provider_arguments(
23 &self,
24 provider_ref: &str,
25 call_id: &str,
26 canonical_tool_name: &CanonicalToolName,
27 raw_arguments: &str,
28 ) -> Result<Option<ProviderArgumentContentRef>, AgentError> {
29 let digest = format!("{:x}", Sha256::digest(raw_arguments.as_bytes()));
30 let content_ref = ProviderArgumentContentRef::new(format!(
31 "content.provider_arguments.{}",
32 &digest[..24]
33 ));
34 self.client.execute(
35 format!("insert into {} (store_scope, provider_ref, call_id, canonical_tool_name, raw_arguments, content_ref) values ($1, $2, $3, $4, $5, $6) on conflict (store_scope, content_ref) do update set provider_ref = excluded.provider_ref, call_id = excluded.call_id, canonical_tool_name = excluded.canonical_tool_name, raw_arguments = excluded.raw_arguments", self.client.table("agent_sdk_provider_arguments")),
36 vec![
37 self.client.scope(),
38 Value::String(provider_ref.to_string()),
39 Value::String(call_id.to_string()),
40 Value::String(canonical_tool_name.as_str().to_string()),
41 Value::String(raw_arguments.to_string()),
42 Value::String(content_ref.as_str().to_string()),
43 ],
44 )?;
45 Ok(Some(content_ref))
46 }
47
48 fn load_provider_arguments_json(
49 &self,
50 content_ref: &ProviderArgumentContentRef,
51 ) -> Result<Value, AgentError> {
52 let response = self.client.execute(
53 format!(
54 "select raw_arguments from {} where store_scope = $1 and content_ref = $2",
55 self.client.table("agent_sdk_provider_arguments")
56 ),
57 vec![
58 self.client.scope(),
59 Value::String(content_ref.as_str().to_string()),
60 ],
61 )?;
62 let raw = response
63 .rows
64 .first()
65 .and_then(|row| row.get("raw_arguments"))
66 .and_then(Value::as_str)
67 .ok_or_else(|| {
68 AgentError::contract_violation("provider argument content ref is missing")
69 })?;
70 serde_json::from_str(raw).map_err(|error| {
71 AgentError::contract_violation(format!(
72 "stored provider arguments are not valid JSON: {error}"
73 ))
74 })
75 }
76}