Skip to main content

agentdb/
prompts.rs

1use crate::error::{AgentDbError, Result};
2use crate::schema::now_ms;
3use rusqlite::params;
4use rusqlite::Connection;
5use serde_json::Value;
6use std::sync::{Arc, Mutex};
7use uuid::Uuid;
8
9#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
10pub struct PromptTemplate {
11    pub id: String,
12    pub name: String,
13    pub version: i64,
14    pub template: String,
15    pub model_hint: Option<String>,
16    pub max_tokens: Option<i64>,
17    pub metadata: Option<Value>,
18    pub created_at: i64,
19}
20
21pub struct PromptStore {
22    conn: Arc<Mutex<Connection>>,
23}
24
25impl PromptStore {
26    pub(crate) fn new(conn: Arc<Mutex<Connection>>) -> Self {
27        Self { conn }
28    }
29
30    pub fn create_template(
31        &self,
32        name: &str,
33        template: &str,
34        model_hint: Option<&str>,
35        max_tokens: Option<i64>,
36        metadata: Option<Value>,
37    ) -> Result<String> {
38        let id = Uuid::new_v4().to_string();
39        let conn = self.conn.lock().unwrap();
40        let meta_str = metadata.as_ref().map(|v| v.to_string());
41        let now = now_ms();
42        let version: i64 = conn
43            .query_row(
44                "SELECT COALESCE(MAX(version) + 1, 1) FROM _adb_prompt_templates WHERE name = ?1",
45                params![name],
46                |r| r.get(0),
47            )
48            .unwrap_or(1);
49        conn.execute(
50            "INSERT INTO _adb_prompt_templates (id, name, version, template, model_hint, max_tokens, metadata, created_at)
51             VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8)",
52            params![id, name, version, template, model_hint, max_tokens, meta_str, now],
53        )?;
54        Ok(id)
55    }
56
57    /// Get the latest version of a named template.
58    pub fn get_template(&self, name: &str) -> Result<PromptTemplate> {
59        let conn = self.conn.lock().unwrap();
60        conn.query_row(
61            "SELECT id, name, version, template, model_hint, max_tokens, metadata, created_at
62             FROM _adb_prompt_templates
63             WHERE name = ?1
64             ORDER BY version DESC LIMIT 1",
65            params![name],
66            parse_template_row,
67        )
68        .map_err(|_| AgentDbError::InvalidArgument(format!("template not found: {name}")))
69    }
70
71    /// Get a specific version of a template.
72    pub fn get_template_version(&self, name: &str, version: i64) -> Result<PromptTemplate> {
73        let conn = self.conn.lock().unwrap();
74        conn.query_row(
75            "SELECT id, name, version, template, model_hint, max_tokens, metadata, created_at
76             FROM _adb_prompt_templates
77             WHERE name = ?1 AND version = ?2",
78            params![name, version],
79            parse_template_row,
80        )
81        .map_err(|_| {
82            AgentDbError::InvalidArgument(format!("template not found: {name} v{version}"))
83        })
84    }
85
86    pub fn list_templates(&self) -> Result<Vec<PromptTemplate>> {
87        let conn = self.conn.lock().unwrap();
88        let mut stmt = conn.prepare(
89            "SELECT id, name, version, template, model_hint, max_tokens, metadata, created_at
90             FROM _adb_prompt_templates
91             ORDER BY name, version DESC",
92        )?;
93        let rows = stmt.query_map([], parse_template_row)?;
94        rows.map(|r| r.map_err(AgentDbError::Sqlite)).collect()
95    }
96
97    /// Render a template by substituting `{{key}}` placeholders with values from `vars`.
98    pub fn render(
99        &self,
100        name: &str,
101        vars: &std::collections::HashMap<String, String>,
102    ) -> Result<String> {
103        let tmpl = self.get_template(name)?;
104        let mut output = tmpl.template;
105        for (key, value) in vars {
106            output = output.replace(&format!("{{{{{key}}}}}"), value);
107        }
108        Ok(output)
109    }
110
111    pub fn delete_template(&self, name: &str) -> Result<()> {
112        let conn = self.conn.lock().unwrap();
113        conn.execute(
114            "DELETE FROM _adb_prompt_templates WHERE name = ?1",
115            params![name],
116        )?;
117        Ok(())
118    }
119}
120
121fn parse_template_row(row: &rusqlite::Row) -> rusqlite::Result<PromptTemplate> {
122    let meta_str: Option<String> = row.get(6)?;
123    Ok(PromptTemplate {
124        id: row.get(0)?,
125        name: row.get(1)?,
126        version: row.get(2)?,
127        template: row.get(3)?,
128        model_hint: row.get(4)?,
129        max_tokens: row.get(5)?,
130        metadata: meta_str.and_then(|s| serde_json::from_str(&s).ok()),
131        created_at: row.get(7)?,
132    })
133}