performance_memory.py

import json
import sqlite3
import time
from dataclasses import dataclass, asdict
from pathlib import Path
from typing import Optional

DB_PATH = Path(__file__).parent / "tools" / "performance.db"


@dataclass
class InvocationRecord:
    tool_name: str
    timestamp: float
    latency_ms: float
    success: bool
    error_message: Optional[str]
    input_hash: str
    output_hash: Optional[str]
    version: int


class PerformanceMemory:
    def __init__(self, db_path: Path = DB_PATH):
        self.db_path = db_path
        self.db_path.parent.mkdir(parents=True, exist_ok=True)
        self.conn = sqlite3.connect(str(self.db_path))
        self._create_tables()

    def _create_tables(self):
        self.conn.execute("""
            CREATE TABLE IF NOT EXISTS invocations (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                tool_name TEXT NOT NULL,
                timestamp REAL NOT NULL,
                latency_ms REAL NOT NULL,
                success INTEGER NOT NULL,
                error_message TEXT,
                input_hash TEXT NOT NULL,
                output_hash TEXT,
                version INTEGER NOT NULL DEFAULT 1
            )
        """)
        self.conn.execute("""
            CREATE TABLE IF NOT EXISTS tool_versions (
                tool_name TEXT NOT NULL,
                version INTEGER NOT NULL,
                source_code TEXT NOT NULL,
                created_at REAL NOT NULL,
                PRIMARY KEY (tool_name, version)
            )
        """)
        self.conn.commit()

    def record(self, rec: InvocationRecord):
        self.conn.execute(
            """INSERT INTO invocations
               (tool_name, timestamp, latency_ms, success,
                error_message, input_hash, output_hash, version)
               VALUES (?, ?, ?, ?, ?, ?, ?, ?)""",
            (
                rec.tool_name, rec.timestamp, rec.latency_ms,
                int(rec.success), rec.error_message,
                rec.input_hash, rec.output_hash, rec.version,
            ),
        )
        self.conn.commit()

    def get_recent(self, tool_name: str, limit: int = 20) -> list[dict]:
        cursor = self.conn.execute(
            """SELECT tool_name, timestamp, latency_ms, success,
                      error_message, input_hash, output_hash, version
               FROM invocations
               WHERE tool_name = ?
               ORDER BY timestamp DESC
               LIMIT ?""",
            (tool_name, limit),
        )
        columns = [
            "tool_name", "timestamp", "latency_ms", "success",
            "error_message", "input_hash", "output_hash", "version",
        ]
        return [dict(zip(columns, row)) for row in cursor.fetchall()]

    def get_current_version(self, tool_name: str) -> int:
        cursor = self.conn.execute(
            "SELECT MAX(version) FROM tool_versions WHERE tool_name = ?",
            (tool_name,),
        )
        row = cursor.fetchone()
        return row[0] if row[0] is not None else 1

    def store_version(self, tool_name: str, version: int, source_code: str):
        self.conn.execute(
            """INSERT OR REPLACE INTO tool_versions
               (tool_name, version, source_code, created_at)
               VALUES (?, ?, ?, ?)""",
            (tool_name, version, source_code, time.time()),
        )
        self.conn.commit()

    def get_version_source(self, tool_name: str, version: int) -> Optional[str]:
        cursor = self.conn.execute(
            "SELECT source_code FROM tool_versions WHERE tool_name = ? AND version = ?",
            (tool_name, version),
        )
        row = cursor.fetchone()
        return row[0] if row else None

    def summary(self, tool_name: str, last_n: int = 20) -> dict:
        records = self.get_recent(tool_name, last_n)
        if not records:
            return {"tool_name": tool_name, "invocations": 0}
        total = len(records)
        successes = sum(1 for r in records if r["success"])
        latencies = [r["latency_ms"] for r in records]
        errors = [r["error_message"] for r in records if r["error_message"]]
        return {
            "tool_name": tool_name,
            "invocations": total,
            "success_rate": successes / total,
            "avg_latency_ms": sum(latencies) / total,
            "max_latency_ms": max(latencies),
            "min_latency_ms": min(latencies),
            "recent_errors": errors[:5],
            "current_version": self.get_current_version(tool_name),
        }


memory = PerformanceMemory()