Spaces:
Sleeping
Sleeping
File size: 8,924 Bytes
98ee05e e6fb135 98ee05e e6fb135 5bb23c3 e6fb135 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 | """SQLite store for runs, outputs, the golden dataset, and flagged outputs."""
import json
import os
import sqlite3
import threading
from datetime import datetime, timezone
from typing import Any, Optional
DEFAULT_SQLITE_PATH = os.getenv("SQLITE_PATH", "./agent.db")
def _utcnow() -> str:
return datetime.now(timezone.utc).isoformat()
class Store:
"""Local SQLite wrapper. One connection guarded by a lock so the API
handlers and CLI can share a Store."""
def __init__(self, path: str = DEFAULT_SQLITE_PATH):
self.path = path
self._lock = threading.Lock()
self._conn = sqlite3.connect(path, check_same_thread=False)
self._conn.row_factory = sqlite3.Row
self._conn.execute("PRAGMA journal_mode=WAL;")
self._init_schema()
def _init_schema(self) -> None:
with self._lock, self._conn:
self._conn.executescript(
"""
CREATE TABLE IF NOT EXISTS runs (
id INTEGER PRIMARY KEY AUTOINCREMENT,
brief TEXT NOT NULL, -- JSON brand brief
prompt_version TEXT NOT NULL,
timestamp TEXT NOT NULL
);
CREATE TABLE IF NOT EXISTS outputs (
id INTEGER PRIMARY KEY AUTOINCREMENT,
run_id INTEGER NOT NULL,
variant_type TEXT NOT NULL, -- headline | body | cta
content TEXT NOT NULL,
hook_strength REAL,
brand_alignment REAL,
clarity REAL,
conversion_intent REAL,
weighted_average REAL,
prompt_version TEXT NOT NULL,
timestamp TEXT NOT NULL,
FOREIGN KEY (run_id) REFERENCES runs(id)
);
CREATE TABLE IF NOT EXISTS golden (
id INTEGER PRIMARY KEY AUTOINCREMENT,
brief TEXT NOT NULL,
variant_type TEXT NOT NULL,
output TEXT NOT NULL,
hook_strength REAL,
brand_alignment REAL,
clarity REAL,
conversion_intent REAL,
weighted_average REAL NOT NULL,
prompt_version TEXT NOT NULL,
timestamp TEXT NOT NULL
);
CREATE TABLE IF NOT EXISTS flagged_outputs (
id INTEGER PRIMARY KEY AUTOINCREMENT,
run_id INTEGER,
brief TEXT NOT NULL,
variant_type TEXT NOT NULL,
output TEXT NOT NULL,
weighted_average REAL,
reason TEXT NOT NULL,
timestamp TEXT NOT NULL
);
"""
)
def create_run(self, brief: dict[str, Any], prompt_version: str) -> int:
with self._lock, self._conn:
cur = self._conn.execute(
"INSERT INTO runs (brief, prompt_version, timestamp) VALUES (?, ?, ?)",
(json.dumps(brief), prompt_version, _utcnow()),
)
return int(cur.lastrowid)
def get_run(self, run_id: int) -> Optional[dict[str, Any]]:
with self._lock:
row = self._conn.execute(
"SELECT * FROM runs WHERE id = ?", (run_id,)
).fetchone()
return dict(row) if row else None
def add_output(
self,
run_id: int,
variant_type: str,
content: str,
scores: dict[str, float],
prompt_version: str,
) -> int:
"""Persist one generated variant together with its 4 dimension scores."""
with self._lock, self._conn:
cur = self._conn.execute(
"""
INSERT INTO outputs (
run_id, variant_type, content,
hook_strength, brand_alignment, clarity, conversion_intent,
weighted_average, prompt_version, timestamp
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
""",
(
run_id,
variant_type,
content,
scores.get("hook_strength"),
scores.get("brand_alignment"),
scores.get("clarity"),
scores.get("conversion_intent"),
scores.get("weighted_average"),
prompt_version,
_utcnow(),
),
)
return int(cur.lastrowid)
def get_outputs_for_run(self, run_id: int) -> list[dict[str, Any]]:
with self._lock:
rows = self._conn.execute(
"SELECT * FROM outputs WHERE run_id = ? ORDER BY id", (run_id,)
).fetchall()
return [dict(r) for r in rows]
def add_golden(
self,
brief: dict[str, Any],
variant_type: str,
output: str,
scores: dict[str, float],
prompt_version: str,
) -> int:
with self._lock, self._conn:
cur = self._conn.execute(
"""
INSERT INTO golden (
brief, variant_type, output,
hook_strength, brand_alignment, clarity, conversion_intent,
weighted_average, prompt_version, timestamp
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
""",
(
json.dumps(brief),
variant_type,
output,
scores.get("hook_strength"),
scores.get("brand_alignment"),
scores.get("clarity"),
scores.get("conversion_intent"),
scores.get("weighted_average"),
prompt_version,
_utcnow(),
),
)
return int(cur.lastrowid)
def get_golden(self) -> list[dict[str, Any]]:
with self._lock:
rows = self._conn.execute(
"SELECT * FROM golden ORDER BY id"
).fetchall()
result = []
for r in rows:
d = dict(r)
d["brief"] = json.loads(d["brief"])
result.append(d)
return result
def golden_exists(self, brief: dict[str, Any], variant_type: str, output: str) -> bool:
"""Avoid inserting an identical golden entry twice."""
with self._lock:
row = self._conn.execute(
"SELECT 1 FROM golden WHERE brief = ? AND variant_type = ? AND output = ? LIMIT 1",
(json.dumps(brief), variant_type, output),
).fetchone()
return row is not None
def add_flagged(
self,
brief: dict[str, Any],
variant_type: str,
output: str,
weighted_average: Optional[float],
reason: str,
run_id: Optional[int] = None,
) -> int:
with self._lock, self._conn:
cur = self._conn.execute(
"""
INSERT INTO flagged_outputs (
run_id, brief, variant_type, output, weighted_average, reason, timestamp
) VALUES (?, ?, ?, ?, ?, ?, ?)
""",
(
run_id,
json.dumps(brief),
variant_type,
output,
weighted_average,
reason,
_utcnow(),
),
)
return int(cur.lastrowid)
def get_flagged(self) -> list[dict[str, Any]]:
with self._lock:
rows = self._conn.execute(
"SELECT * FROM flagged_outputs ORDER BY id"
).fetchall()
result = []
for r in rows:
d = dict(r)
d["brief"] = json.loads(d["brief"])
result.append(d)
return result
def count_runs(self) -> int:
with self._lock:
row = self._conn.execute("SELECT COUNT(*) AS n FROM runs").fetchone()
return int(row["n"])
def count_golden(self) -> int:
with self._lock:
row = self._conn.execute("SELECT COUNT(*) AS n FROM golden").fetchone()
return int(row["n"])
def count_flagged(self) -> int:
with self._lock:
row = self._conn.execute("SELECT COUNT(*) AS n FROM flagged_outputs").fetchone()
return int(row["n"])
def close(self) -> None:
with self._lock:
self._conn.close()
|