Spaces:
Running
Running
File size: 3,983 Bytes
516fa71 73781a4 516fa71 | 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 | import sqlite3
import uuid
from datetime import datetime, timezone
from typing import Optional, List
from .models import CorrectionRecord
DB_PATH = "./chroma_db/correction_store.db"
_CREATE_SQL = """
CREATE TABLE IF NOT EXISTS corrections (
correction_id TEXT PRIMARY KEY,
session_id TEXT NOT NULL,
timestamp TEXT NOT NULL,
section TEXT NOT NULL,
correction_text TEXT NOT NULL,
corrected_snippet TEXT,
root_cause TEXT NOT NULL,
action_taken TEXT NOT NULL,
affected_chunk_ids TEXT NOT NULL,
original_report TEXT NOT NULL,
features_json TEXT NOT NULL,
retrieved_context_json TEXT NOT NULL,
verified INTEGER DEFAULT 0,
verification_result TEXT
);
CREATE TABLE IF NOT EXISTS chunk_edits (
edit_id TEXT PRIMARY KEY,
correction_id TEXT NOT NULL REFERENCES corrections(correction_id),
chroma_id TEXT NOT NULL,
edit_type TEXT NOT NULL,
original_content TEXT,
new_content TEXT,
timestamp TEXT NOT NULL
);
"""
class CorrectionStore:
def __init__(self, db_path: str = DB_PATH):
self.db_path = db_path
self._init_db()
def _init_db(self) -> None:
with sqlite3.connect(self.db_path) as conn:
conn.executescript(_CREATE_SQL)
def save_correction(self, record: CorrectionRecord) -> None:
with sqlite3.connect(self.db_path) as conn:
conn.execute(
"""INSERT INTO corrections VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?,?)""",
(
record.correction_id,
record.session_id,
record.timestamp,
record.section,
record.correction_text,
record.corrected_snippet,
record.root_cause,
record.action_taken,
record.affected_chunk_ids,
record.original_report,
record.features_json,
record.retrieved_context_json,
1 if record.verified else 0,
record.verification_result,
),
)
def save_chunk_edit(
self,
correction_id: str,
chroma_id: str,
edit_type: str,
original: Optional[str],
new_content: Optional[str],
) -> None:
with sqlite3.connect(self.db_path) as conn:
conn.execute(
"INSERT INTO chunk_edits VALUES (?,?,?,?,?,?,?)",
(
str(uuid.uuid4()),
correction_id,
chroma_id,
edit_type,
original,
new_content,
datetime.now(timezone.utc).isoformat(),
),
)
def get_correction(self, correction_id: str) -> Optional[dict]:
with sqlite3.connect(self.db_path) as conn:
conn.row_factory = sqlite3.Row
row = conn.execute(
"SELECT * FROM corrections WHERE correction_id=?",
(correction_id,),
).fetchone()
return dict(row) if row else None
def list_corrections(self, limit: int = 50) -> List[dict]:
with sqlite3.connect(self.db_path) as conn:
conn.row_factory = sqlite3.Row
rows = conn.execute(
"SELECT * FROM corrections ORDER BY timestamp DESC LIMIT ?",
(limit,),
).fetchall()
return [dict(r) for r in rows]
def mark_verified(self, correction_id: str, result: str) -> None:
with sqlite3.connect(self.db_path) as conn:
conn.execute(
"UPDATE corrections SET verified=1, verification_result=? WHERE correction_id=?",
(result, correction_id),
)
|