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