Skip to main content

agent_sdk_store_postgres/
provider_arguments.rs

1use 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}