pv-agent/agent/audit.py

"""Persistente, strukturierte Interaktions- und Bewertungsprotokolle."""
from __future__ import annotations

import json
import sqlite3
import threading
import time
from pathlib import Path

from .config import Config

SCHEMA = """
CREATE TABLE IF NOT EXISTS interactions (
    request_id TEXT PRIMARY KEY,
    created_at INTEGER NOT NULL,
    question TEXT,
    answer TEXT,
    status TEXT NOT NULL,
    verified INTEGER NOT NULL,
    refused INTEGER NOT NULL,
    citations_json TEXT NOT NULL,
    sources_json TEXT NOT NULL,
    conflicts_json TEXT NOT NULL,
    planned_queries_json TEXT NOT NULL,
    model TEXT NOT NULL,
    latency_ms INTEGER NOT NULL,
    n_context INTEGER NOT NULL,
    regenerations INTEGER NOT NULL,
    context_json TEXT
);
CREATE TABLE IF NOT EXISTS ratings (
    request_id TEXT PRIMARY KEY,
    created_at INTEGER NOT NULL,
    updated_at INTEGER NOT NULL,
    rating TEXT NOT NULL CHECK (rating IN ('up', 'down')),
    feedback TEXT,
    FOREIGN KEY (request_id) REFERENCES interactions(request_id) ON DELETE CASCADE
);
CREATE TABLE IF NOT EXISTS comments (
    comment_id INTEGER PRIMARY KEY AUTOINCREMENT,
    request_id TEXT NOT NULL,
    created_at INTEGER NOT NULL,
    comment TEXT,
    FOREIGN KEY (request_id) REFERENCES interactions(request_id) ON DELETE CASCADE
);
CREATE INDEX IF NOT EXISTS interactions_created_at_idx
    ON interactions(created_at DESC);
CREATE INDEX IF NOT EXISTS ratings_updated_at_idx
    ON ratings(updated_at DESC);
CREATE INDEX IF NOT EXISTS comments_request_created_idx
    ON comments(request_id, created_at DESC);
"""


class AuditStore:
    def __init__(self, cfg: Config):
        self.cfg = cfg
        path = Path(cfg.audit_db_path)
        path.parent.mkdir(parents=True, exist_ok=True)
        self._lock = threading.RLock()
        self._con = sqlite3.connect(path, check_same_thread=False)
        self._con.row_factory = sqlite3.Row
        self._con.execute("PRAGMA foreign_keys = ON")
        self._con.execute("PRAGMA journal_mode = WAL")
        self._con.executescript(SCHEMA)
        self._migrate()
        self._delete_expired()

    def close(self) -> None:
        with self._lock:
            self._con.close()

    def _migrate(self) -> None:
        """Idempotente Spalten-Migration für bestehende audit.db-Dateien."""
        cols = {
            row[1]
            for row in self._con.execute("PRAGMA table_info(interactions)").fetchall()
        }
        if "context_json" not in cols:
            self._con.execute(
                "ALTER TABLE interactions ADD COLUMN context_json TEXT"
            )
            self._con.commit()

    def _delete_expired(self) -> None:
        if self.cfg.audit_retention_days <= 0:
            return
        cutoff = int(time.time()) - self.cfg.audit_retention_days * 86_400
        with self._lock:
            self._con.execute("DELETE FROM interactions WHERE created_at < ?", (cutoff,))
            self._con.commit()

    def record_interaction(self, payload: dict, context: dict | None = None) -> None:
        now = int(time.time())
        include_content = self.cfg.audit_log_content
        question = payload.get("question") if include_content else None
        answer = payload.get("answer") if include_content else None
        sources = payload.get("sources", []) if include_content else []
        conflicts = payload.get("conflicts", []) if include_content else []
        planned_queries = payload.get("planned_queries", []) if include_content else []
        context_json = (
            json.dumps(context, ensure_ascii=False)
            if include_content and context
            else None
        )
        values = (
            payload["request_id"],
            now,
            question,
            answer,
            payload["status"],
            int(bool(payload["verified"])),
            int(bool(payload["refused"])),
            json.dumps(payload.get("citations", []), ensure_ascii=False),
            json.dumps(sources, ensure_ascii=False),
            json.dumps(conflicts, ensure_ascii=False),
            json.dumps(planned_queries, ensure_ascii=False),
            payload.get("model", ""),
            int(payload.get("latency_ms", 0)),
            int(payload.get("n_context", 0)),
            int(payload.get("regenerations", 0)),
            context_json,
        )
        with self._lock:
            self._con.execute(
                "INSERT INTO interactions("
                "request_id, created_at, question, answer, status, verified, refused, "
                "citations_json, sources_json, conflicts_json, planned_queries_json, "
                "model, latency_ms, n_context, regenerations, context_json"
                ") VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) "
                "ON CONFLICT(request_id) DO UPDATE SET "
                "created_at=excluded.created_at, question=excluded.question, "
                "answer=excluded.answer, status=excluded.status, "
                "verified=excluded.verified, refused=excluded.refused, "
                "citations_json=excluded.citations_json, sources_json=excluded.sources_json, "
                "conflicts_json=excluded.conflicts_json, "
                "planned_queries_json=excluded.planned_queries_json, "
                "model=excluded.model, latency_ms=excluded.latency_ms, "
                "n_context=excluded.n_context, regenerations=excluded.regenerations",
                values,
            )
            self._con.commit()
        if self.cfg.audit_stdout:
            if include_content:
                event = {"event": "agent_interaction", "created_at": now, **payload}
                if context:
                    event["context"] = context
            else:
                event = {
                    "event": "agent_interaction",
                    "created_at": now,
                    "request_id": payload["request_id"],
                    "status": payload["status"],
                    "verified": payload["verified"],
                    "refused": payload["refused"],
                    "citations": payload.get("citations", []),
                    "model": payload.get("model", ""),
                    "latency_ms": payload.get("latency_ms", 0),
                    "n_context": payload.get("n_context", 0),
                    "regenerations": payload.get("regenerations", 0),
                }
            print("AUDIT " + json.dumps(event, ensure_ascii=False), flush=True)

    def record_rating(self, request_id: str, rating: str, feedback: str | None) -> None:
        now = int(time.time())
        stored_feedback = feedback if self.cfg.audit_log_content else None
        with self._lock:
            exists = self._con.execute(
                "SELECT 1 FROM interactions WHERE request_id = ?", (request_id,)
            ).fetchone()
            if exists is None:
                raise KeyError(request_id)
            self._con.execute(
                "INSERT INTO ratings(request_id, created_at, updated_at, rating, feedback) "
                "VALUES (?, ?, ?, ?, ?) "
                "ON CONFLICT(request_id) DO UPDATE SET "
                "updated_at=excluded.updated_at, rating=excluded.rating, "
                "feedback=excluded.feedback",
                (request_id, now, now, rating, stored_feedback),
            )
            self._con.commit()
        if self.cfg.audit_stdout:
            print(
                "AUDIT "
                + json.dumps(
                    {
                        "event": "agent_rating",
                        "created_at": now,
                        "request_id": request_id,
                        "rating": rating,
                        "feedback": stored_feedback,
                    },
                    ensure_ascii=False,
                ),
                flush=True,
            )

    def record_comment(self, request_id: str, comment: str) -> int:
        now = int(time.time())
        stored_comment = comment if self.cfg.audit_log_content else None
        with self._lock:
            exists = self._con.execute(
                "SELECT 1 FROM interactions WHERE request_id = ?", (request_id,)
            ).fetchone()
            if exists is None:
                raise KeyError(request_id)
            cursor = self._con.execute(
                "INSERT INTO comments(request_id, created_at, comment) VALUES (?, ?, ?)",
                (request_id, now, stored_comment),
            )
            self._con.commit()
            if cursor.lastrowid is None:
                raise RuntimeError("Kommentar wurde ohne ID gespeichert")
            comment_id = int(cursor.lastrowid)
        if self.cfg.audit_stdout:
            print(
                "AUDIT "
                + json.dumps(
                    {
                        "event": "agent_comment",
                        "created_at": now,
                        "request_id": request_id,
                        "comment_id": comment_id,
                        "comment": stored_comment,
                    },
                    ensure_ascii=False,
                ),
                flush=True,
            )
        return comment_id

    def recent(self, limit: int = 20) -> list[dict]:
        with self._lock:
            rows = self._con.execute(
                "SELECT i.*, r.rating, r.feedback, r.updated_at AS rating_updated_at "
                "FROM interactions i LEFT JOIN ratings r USING(request_id) "
                "ORDER BY i.created_at DESC LIMIT ?",
                (limit,),
            ).fetchall()
            comments_by_request: dict[str, list[dict]] = {}
            for row in rows:
                comment_rows = self._con.execute(
                    "SELECT comment_id, created_at, comment FROM comments "
                    "WHERE request_id = ? ORDER BY created_at, comment_id",
                    (row["request_id"],),
                ).fetchall()
                comments_by_request[row["request_id"]] = [
                    dict(comment_row) for comment_row in comment_rows
                ]
        out = []
        for row in rows:
            item = dict(row)
            for field in (
                "citations_json",
                "sources_json",
                "conflicts_json",
                "planned_queries_json",
            ):
                item[field.removesuffix("_json")] = json.loads(item.pop(field))
            item["verified"] = bool(item["verified"])
            item["refused"] = bool(item["refused"])
            raw_context = item.pop("context_json", None)
            item["context"] = json.loads(raw_context) if raw_context else None
            item["comments"] = comments_by_request[item["request_id"]]
            out.append(item)
        return out