Skip to main content

agent_sdk_store_sqlite/
provider_arguments.rs

1use std::path::{Path, PathBuf};
2
3use agent_sdk_core::{
4    AgentError, ProviderArgumentStore, domain::ContentRef, tool_records::CanonicalToolName,
5};
6use rusqlite::{OptionalExtension, params};
7
8use crate::util::{open, sha256_hex, sqlite_error};
9
10const SCHEMA: &str = "
11CREATE TABLE IF NOT EXISTS provider_arguments (
12    content_ref TEXT PRIMARY KEY,
13    provider_ref TEXT NOT NULL,
14    call_id TEXT NOT NULL,
15    canonical_tool_name TEXT NOT NULL,
16    raw_arguments TEXT NOT NULL
17);
18";
19
20#[derive(Clone, Debug)]
21/// SQLite-backed raw provider argument store.
22pub struct SqliteProviderArgumentStore {
23    path: PathBuf,
24}
25
26impl SqliteProviderArgumentStore {
27    /// Opens or creates a SQLite provider-argument store.
28    pub fn open(path: impl AsRef<Path>) -> Result<Self, AgentError> {
29        crate::util::init(path.as_ref(), SCHEMA)?;
30        Ok(Self {
31            path: path.as_ref().to_path_buf(),
32        })
33    }
34}
35
36impl ProviderArgumentStore for SqliteProviderArgumentStore {
37    fn store_provider_arguments(
38        &self,
39        provider_ref: &str,
40        call_id: &str,
41        canonical_tool_name: &CanonicalToolName,
42        raw_arguments: &str,
43    ) -> Result<Option<ContentRef>, AgentError> {
44        let digest = sha256_hex(raw_arguments.as_bytes());
45        let content_ref = ContentRef::new(format!("content.provider_arguments.{}", &digest[..24]));
46        let connection = open(&self.path)?;
47        connection
48            .execute(
49                "INSERT OR REPLACE INTO provider_arguments
50                 (content_ref, provider_ref, call_id, canonical_tool_name, raw_arguments)
51                 VALUES (?1, ?2, ?3, ?4, ?5)",
52                params![
53                    content_ref.as_str(),
54                    provider_ref,
55                    call_id,
56                    canonical_tool_name.as_str(),
57                    raw_arguments,
58                ],
59            )
60            .map_err(sqlite_error)?;
61        Ok(Some(content_ref))
62    }
63
64    fn load_provider_arguments_json(
65        &self,
66        content_ref: &ContentRef,
67    ) -> Result<serde_json::Value, AgentError> {
68        let connection = open(&self.path)?;
69        let raw = connection
70            .query_row(
71                "SELECT raw_arguments FROM provider_arguments WHERE content_ref = ?1",
72                params![content_ref.as_str()],
73                |row| row.get::<_, String>(0),
74            )
75            .optional()
76            .map_err(sqlite_error)?
77            .ok_or_else(|| {
78                AgentError::contract_violation("provider argument content ref is missing")
79            })?;
80        serde_json::from_str(&raw).map_err(|error| {
81            AgentError::contract_violation(format!(
82                "stored provider arguments are not valid JSON: {error}"
83            ))
84        })
85    }
86}