ctrlcode 0.1.0__py3-none-any.whl

ctrlcode/storage/history_db.py

@@ -0,0 +1,708 @@
+"""SQLite-based historical knowledge database for fuzzing sessions."""
+
+import json
+import logging
+import sqlite3
+from dataclasses import dataclass
+from datetime import datetime
+from pathlib import Path
+from typing import Optional
+
+import numpy as np
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class FuzzingSession:
+    """Record of a complete fuzzing session."""
+
+    session_id: str
+    user_request: str
+    generated_code: str
+    oracle: str
+    timestamp: datetime
+    num_tests: int
+    num_failures: int
+    oracle_reused: bool = False
+    reused_from: Optional[str] = None
+    quality_score: Optional[float] = None
+
+
+@dataclass
+class CodeRecord:
+    """Code snippet with embedding."""
+
+    code_id: str
+    session_id: str
+    code: str
+    embedding: np.ndarray
+    timestamp: datetime
+
+
+@dataclass
+class OracleRecord:
+    """Oracle with embedding and versioning."""
+
+    oracle_id: str
+    session_id: str
+    oracle: str
+    embedding: np.ndarray
+    quality_score: float
+    timestamp: datetime
+    oracle_version: int = 1
+    parent_oracle_id: Optional[str] = None
+    reuse_count: int = 0
+
+
+@dataclass
+class BugPattern:
+    """Bug pattern with embedding."""
+
+    bug_id: str
+    session_id: str
+    bug_description: str
+    code_snippet: str
+    embedding: np.ndarray
+    severity: str
+    timestamp: datetime
+
+
+@dataclass
+class StoredTest:
+    """Test case with embedding."""
+
+    test_id: str
+    session_id: str
+    test_code: str
+    embedding: np.ndarray
+    passed: bool
+    timestamp: datetime
+
+
+class HistoryDB:
+    """SQLite database for persistent fuzzing history and knowledge base.
+
+    Schema:
+    - fuzzing_sessions: High-level session metadata
+    - code_embeddings: Code snippets with embeddings
+    - oracle_embeddings: Oracles with embeddings and quality scores
+    - bug_patterns: Bug patterns for pattern matching
+    - test_cases: Test cases with pass/fail status
+    """
+
+    def __init__(self, db_path: str | Path = ":memory:"):
+        """Initialize database connection.
+
+        Args:
+            db_path: Path to SQLite database file (or :memory: for in-memory)
+        """
+        self.db_path = Path(db_path) if db_path != ":memory:" else db_path
+        self.conn: Optional[sqlite3.Connection] = None
+        self._initialize_db()
+
+    def _initialize_db(self) -> None:
+        """Create database schema if not exists."""
+        self.conn = sqlite3.connect(
+            str(self.db_path) if self.db_path != ":memory:" else ":memory:",
+            check_same_thread=False,
+        )
+        self.conn.row_factory = sqlite3.Row
+
+        cursor = self.conn.cursor()
+
+        # Fuzzing sessions table
+        cursor.execute("""
+            CREATE TABLE IF NOT EXISTS fuzzing_sessions (
+                session_id TEXT PRIMARY KEY,
+                user_request TEXT NOT NULL,
+                generated_code TEXT NOT NULL,
+                oracle TEXT NOT NULL,
+                timestamp TEXT NOT NULL,
+                num_tests INTEGER NOT NULL,
+                num_failures INTEGER NOT NULL,
+                oracle_reused INTEGER NOT NULL DEFAULT 0,
+                reused_from TEXT,
+                quality_score REAL
+            )
+        """)
+
+        # Code embeddings table
+        cursor.execute("""
+            CREATE TABLE IF NOT EXISTS code_embeddings (
+                code_id TEXT PRIMARY KEY,
+                session_id TEXT NOT NULL,
+                code TEXT NOT NULL,
+                embedding BLOB NOT NULL,
+                timestamp TEXT NOT NULL,
+                FOREIGN KEY (session_id) REFERENCES fuzzing_sessions(session_id)
+            )
+        """)
+
+        # Oracle embeddings table
+        cursor.execute("""
+            CREATE TABLE IF NOT EXISTS oracle_embeddings (
+                oracle_id TEXT PRIMARY KEY,
+                session_id TEXT NOT NULL,
+                oracle TEXT NOT NULL,
+                embedding BLOB NOT NULL,
+                quality_score REAL NOT NULL,
+                timestamp TEXT NOT NULL,
+                oracle_version INTEGER NOT NULL DEFAULT 1,
+                parent_oracle_id TEXT,
+                reuse_count INTEGER NOT NULL DEFAULT 0,
+                FOREIGN KEY (session_id) REFERENCES fuzzing_sessions(session_id),
+                FOREIGN KEY (parent_oracle_id) REFERENCES oracle_embeddings(oracle_id)
+            )
+        """)
+
+        # Migration: Add versioning columns to existing oracle_embeddings table
+        cursor.execute("PRAGMA table_info(oracle_embeddings)")
+        columns = {row[1] for row in cursor.fetchall()}
+
+        if "oracle_version" not in columns:
+            cursor.execute("""
+                ALTER TABLE oracle_embeddings
+                ADD COLUMN oracle_version INTEGER NOT NULL DEFAULT 1
+            """)
+            logger.info("Added oracle_version column to oracle_embeddings table")
+
+        if "parent_oracle_id" not in columns:
+            cursor.execute("""
+                ALTER TABLE oracle_embeddings
+                ADD COLUMN parent_oracle_id TEXT
+            """)
+            logger.info("Added parent_oracle_id column to oracle_embeddings table")
+
+        if "reuse_count" not in columns:
+            cursor.execute("""
+                ALTER TABLE oracle_embeddings
+                ADD COLUMN reuse_count INTEGER NOT NULL DEFAULT 0
+            """)
+            logger.info("Added reuse_count column to oracle_embeddings table")
+
+        # Bug patterns table
+        cursor.execute("""
+            CREATE TABLE IF NOT EXISTS bug_patterns (
+                bug_id TEXT PRIMARY KEY,
+                session_id TEXT NOT NULL,
+                bug_description TEXT NOT NULL,
+                code_snippet TEXT NOT NULL,
+                embedding BLOB NOT NULL,
+                severity TEXT NOT NULL,
+                timestamp TEXT NOT NULL,
+                FOREIGN KEY (session_id) REFERENCES fuzzing_sessions(session_id)
+            )
+        """)
+
+        # Test cases table
+        cursor.execute("""
+            CREATE TABLE IF NOT EXISTS test_cases (
+                test_id TEXT PRIMARY KEY,
+                session_id TEXT NOT NULL,
+                test_code TEXT NOT NULL,
+                embedding BLOB NOT NULL,
+                passed INTEGER NOT NULL,
+                timestamp TEXT NOT NULL,
+                FOREIGN KEY (session_id) REFERENCES fuzzing_sessions(session_id)
+            )
+        """)
+
+        # Create indexes for faster queries
+        cursor.execute("""
+            CREATE INDEX IF NOT EXISTS idx_sessions_timestamp
+            ON fuzzing_sessions(timestamp)
+        """)
+        cursor.execute("""
+            CREATE INDEX IF NOT EXISTS idx_code_session
+            ON code_embeddings(session_id)
+        """)
+        cursor.execute("""
+            CREATE INDEX IF NOT EXISTS idx_oracle_quality
+            ON oracle_embeddings(quality_score DESC)
+        """)
+        cursor.execute("""
+            CREATE INDEX IF NOT EXISTS idx_bugs_severity
+            ON bug_patterns(severity)
+        """)
+
+        self.conn.commit()
+        logger.debug(f"Initialized history database at {self.db_path}")
+
+    def store_session(self, session: FuzzingSession) -> None:
+        """Store fuzzing session record.
+
+        Args:
+            session: Fuzzing session to store
+        """
+        cursor = self.conn.cursor()
+        cursor.execute(
+            """
+            INSERT OR REPLACE INTO fuzzing_sessions
+            (session_id, user_request, generated_code, oracle, timestamp,
+             num_tests, num_failures, oracle_reused, reused_from, quality_score)
+            VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
+        """,
+            (
+                session.session_id,
+                session.user_request,
+                session.generated_code,
+                session.oracle,
+                session.timestamp.isoformat(),
+                session.num_tests,
+                session.num_failures,
+                1 if session.oracle_reused else 0,
+                session.reused_from,
+                session.quality_score,
+            ),
+        )
+        self.conn.commit()
+        logger.debug(f"Stored session {session.session_id}")
+
+    def store_code(self, record: CodeRecord) -> None:
+        """Store code with embedding.
+
+        Args:
+            record: Code record to store
+        """
+        cursor = self.conn.cursor()
+        cursor.execute(
+            """
+            INSERT OR REPLACE INTO code_embeddings
+            (code_id, session_id, code, embedding, timestamp)
+            VALUES (?, ?, ?, ?, ?)
+        """,
+            (
+                record.code_id,
+                record.session_id,
+                record.code,
+                record.embedding.tobytes(),
+                record.timestamp.isoformat(),
+            ),
+        )
+        self.conn.commit()
+
+    def store_oracle(self, record: OracleRecord) -> None:
+        """Store oracle with embedding and versioning info.
+
+        Args:
+            record: Oracle record to store
+        """
+        cursor = self.conn.cursor()
+        cursor.execute(
+            """
+            INSERT OR REPLACE INTO oracle_embeddings
+            (oracle_id, session_id, oracle, embedding, quality_score, timestamp,
+             oracle_version, parent_oracle_id, reuse_count)
+            VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
+        """,
+            (
+                record.oracle_id,
+                record.session_id,
+                record.oracle,
+                record.embedding.tobytes(),
+                record.quality_score,
+                record.timestamp.isoformat(),
+                record.oracle_version,
+                record.parent_oracle_id,
+                record.reuse_count,
+            ),
+        )
+        self.conn.commit()
+
+    def store_bug(self, record: BugPattern) -> None:
+        """Store bug pattern with embedding.
+
+        Args:
+            record: Bug pattern to store
+        """
+        cursor = self.conn.cursor()
+        cursor.execute(
+            """
+            INSERT OR REPLACE INTO bug_patterns
+            (bug_id, session_id, bug_description, code_snippet, embedding, severity, timestamp)
+            VALUES (?, ?, ?, ?, ?, ?, ?)
+        """,
+            (
+                record.bug_id,
+                record.session_id,
+                record.bug_description,
+                record.code_snippet,
+                record.embedding.tobytes(),
+                record.severity,
+                record.timestamp.isoformat(),
+            ),
+        )
+        self.conn.commit()
+
+    def store_test(self, record: StoredTest) -> None:
+        """Store test case with embedding.
+
+        Args:
+            record: Test case to store
+        """
+        cursor = self.conn.cursor()
+        cursor.execute(
+            """
+            INSERT OR REPLACE INTO test_cases
+            (test_id, session_id, test_code, embedding, passed, timestamp)
+            VALUES (?, ?, ?, ?, ?, ?)
+        """,
+            (
+                record.test_id,
+                record.session_id,
+                record.test_code,
+                record.embedding.tobytes(),
+                1 if record.passed else 0,
+                record.timestamp.isoformat(),
+            ),
+        )
+        self.conn.commit()
+
+    def get_all_code_embeddings(self, limit: Optional[int] = None) -> list[CodeRecord]:
+        """Retrieve all code embeddings for vector store initialization.
+
+        Args:
+            limit: Maximum number of records to retrieve
+
+        Returns:
+            List of code records
+        """
+        cursor = self.conn.cursor()
+        query = "SELECT * FROM code_embeddings ORDER BY timestamp DESC"
+        if limit:
+            query += f" LIMIT {limit}"
+
+        rows = cursor.execute(query).fetchall()
+        return [self._row_to_code_record(row) for row in rows]
+
+    def get_all_oracle_embeddings(self, limit: Optional[int] = None) -> list[OracleRecord]:
+        """Retrieve all oracle embeddings.
+
+        Args:
+            limit: Maximum number of records
+
+        Returns:
+            List of oracle records
+        """
+        cursor = self.conn.cursor()
+        query = "SELECT * FROM oracle_embeddings ORDER BY quality_score DESC"
+        if limit:
+            query += f" LIMIT {limit}"
+
+        rows = cursor.execute(query).fetchall()
+        return [self._row_to_oracle_record(row) for row in rows]
+
+    def get_all_bug_embeddings(self, limit: Optional[int] = None) -> list[BugPattern]:
+        """Retrieve all bug pattern embeddings.
+
+        Args:
+            limit: Maximum number of records
+
+        Returns:
+            List of bug patterns
+        """
+        cursor = self.conn.cursor()
+        query = "SELECT * FROM bug_patterns ORDER BY timestamp DESC"
+        if limit:
+            query += f" LIMIT {limit}"
+
+        rows = cursor.execute(query).fetchall()
+        return [self._row_to_bug_record(row) for row in rows]
+
+    def get_all_test_embeddings(self, limit: Optional[int] = None) -> list[StoredTest]:
+        """Retrieve all test case embeddings.
+
+        Args:
+            limit: Maximum number of records
+
+        Returns:
+            List of test cases
+        """
+        cursor = self.conn.cursor()
+        query = "SELECT * FROM test_cases ORDER BY timestamp DESC"
+        if limit:
+            query += f" LIMIT {limit}"
+
+        rows = cursor.execute(query).fetchall()
+        return [self._row_to_test_record(row) for row in rows]
+
+    def get_session(self, session_id: str) -> Optional[FuzzingSession]:
+        """Retrieve session by ID.
+
+        Args:
+            session_id: Session identifier
+
+        Returns:
+            Fuzzing session or None
+        """
+        cursor = self.conn.cursor()
+        row = cursor.execute(
+            "SELECT * FROM fuzzing_sessions WHERE session_id = ?", (session_id,)
+        ).fetchone()
+
+        if not row:
+            return None
+
+        return FuzzingSession(
+            session_id=row["session_id"],
+            user_request=row["user_request"],
+            generated_code=row["generated_code"],
+            oracle=row["oracle"],
+            timestamp=datetime.fromisoformat(row["timestamp"]),
+            num_tests=row["num_tests"],
+            num_failures=row["num_failures"],
+            oracle_reused=bool(row["oracle_reused"]),
+            reused_from=row["reused_from"],
+            quality_score=row["quality_score"],
+        )
+
+    def get_stats(self) -> dict:
+        """Get database statistics.
+
+        Returns:
+            Dictionary of statistics
+        """
+        cursor = self.conn.cursor()
+
+        stats = {}
+
+        # Total sessions
+        stats["total_sessions"] = cursor.execute(
+            "SELECT COUNT(*) FROM fuzzing_sessions"
+        ).fetchone()[0]
+
+        # Oracle reuse rate
+        reused = cursor.execute(
+            "SELECT COUNT(*) FROM fuzzing_sessions WHERE oracle_reused = 1"
+        ).fetchone()[0]
+        stats["oracle_reuse_count"] = reused
+        stats["oracle_reuse_rate"] = (
+            reused / stats["total_sessions"] if stats["total_sessions"] > 0 else 0.0
+        )
+
+        # Total bugs
+        stats["total_bugs"] = cursor.execute("SELECT COUNT(*) FROM bug_patterns").fetchone()[0]
+
+        # Total tests
+        stats["total_tests"] = cursor.execute("SELECT COUNT(*) FROM test_cases").fetchone()[0]
+
+        # Test pass rate
+        passed = cursor.execute(
+            "SELECT COUNT(*) FROM test_cases WHERE passed = 1"
+        ).fetchone()[0]
+        stats["test_pass_rate"] = (
+            passed / stats["total_tests"] if stats["total_tests"] > 0 else 0.0
+        )
+
+        # Average quality score
+        avg_quality = cursor.execute(
+            "SELECT AVG(quality_score) FROM oracle_embeddings"
+        ).fetchone()[0]
+        stats["avg_oracle_quality"] = avg_quality if avg_quality else 0.0
+
+        # Total code embeddings
+        stats["total_code_embeddings"] = cursor.execute(
+            "SELECT COUNT(*) FROM code_embeddings"
+        ).fetchone()[0]
+
+        # Total oracle embeddings
+        stats["total_oracle_embeddings"] = cursor.execute(
+            "SELECT COUNT(*) FROM oracle_embeddings"
+        ).fetchone()[0]
+
+        return stats
+
+    def clear(self) -> None:
+        """Clear all data from database."""
+        cursor = self.conn.cursor()
+        cursor.execute("DELETE FROM test_cases")
+        cursor.execute("DELETE FROM bug_patterns")
+        cursor.execute("DELETE FROM oracle_embeddings")
+        cursor.execute("DELETE FROM code_embeddings")
+        cursor.execute("DELETE FROM fuzzing_sessions")
+        self.conn.commit()
+        logger.info("Cleared all data from history database")
+
+    def close(self) -> None:
+        """Close database connection."""
+        if self.conn:
+            self.conn.close()
+            self.conn = None
+
+    def _row_to_code_record(self, row: sqlite3.Row) -> CodeRecord:
+        """Convert database row to CodeRecord."""
+        embedding = np.frombuffer(row["embedding"], dtype=np.float32)
+        return CodeRecord(
+            code_id=row["code_id"],
+            session_id=row["session_id"],
+            code=row["code"],
+            embedding=embedding,
+            timestamp=datetime.fromisoformat(row["timestamp"]),
+        )
+
+    def _row_to_oracle_record(self, row: sqlite3.Row) -> OracleRecord:
+        """Convert database row to OracleRecord."""
+        embedding = np.frombuffer(row["embedding"], dtype=np.float32)
+
+        # Handle optional columns with try/except (sqlite3.Row doesn't support .get())
+        try:
+            oracle_version = row["oracle_version"]
+        except (KeyError, IndexError):
+            oracle_version = 1
+
+        try:
+            parent_oracle_id = row["parent_oracle_id"]
+        except (KeyError, IndexError):
+            parent_oracle_id = None
+
+        try:
+            reuse_count = row["reuse_count"]
+        except (KeyError, IndexError):
+            reuse_count = 0
+
+        return OracleRecord(
+            oracle_id=row["oracle_id"],
+            session_id=row["session_id"],
+            oracle=row["oracle"],
+            embedding=embedding,
+            quality_score=row["quality_score"],
+            timestamp=datetime.fromisoformat(row["timestamp"]),
+            oracle_version=oracle_version,
+            parent_oracle_id=parent_oracle_id,
+            reuse_count=reuse_count,
+        )
+
+    def _row_to_bug_record(self, row: sqlite3.Row) -> BugPattern:
+        """Convert database row to BugPattern."""
+        embedding = np.frombuffer(row["embedding"], dtype=np.float32)
+        return BugPattern(
+            bug_id=row["bug_id"],
+            session_id=row["session_id"],
+            bug_description=row["bug_description"],
+            code_snippet=row["code_snippet"],
+            embedding=embedding,
+            severity=row["severity"],
+            timestamp=datetime.fromisoformat(row["timestamp"]),
+        )
+
+    def _row_to_test_record(self, row: sqlite3.Row) -> StoredTest:
+        """Convert database row to StoredTest."""
+        embedding = np.frombuffer(row["embedding"], dtype=np.float32)
+        return StoredTest(
+            test_id=row["test_id"],
+            session_id=row["session_id"],
+            test_code=row["test_code"],
+            embedding=embedding,
+            passed=bool(row["passed"]),
+            timestamp=datetime.fromisoformat(row["timestamp"]),
+        )
+
+    def increment_oracle_reuse(self, oracle_id: str) -> None:
+        """Increment reuse count for an oracle.
+
+        Args:
+            oracle_id: Oracle ID to increment
+        """
+        cursor = self.conn.cursor()
+        cursor.execute(
+            """
+            UPDATE oracle_embeddings
+            SET reuse_count = reuse_count + 1
+            WHERE oracle_id = ?
+        """,
+            (oracle_id,),
+        )
+        self.conn.commit()
+        logger.debug(f"Incremented reuse count for oracle {oracle_id}")
+
+    def get_golden_oracles(
+        self, min_quality: float = 0.8, min_reuse: int = 3, limit: int = 10
+    ) -> list[OracleRecord]:
+        """Get high-quality, frequently reused oracles.
+
+        These are "golden oracles" that work well and are reused often.
+
+        Args:
+            min_quality: Minimum quality score (default: 0.8)
+            min_reuse: Minimum reuse count (default: 3)
+            limit: Maximum results (default: 10)
+
+        Returns:
+            List of golden oracle records, sorted by quality * reuse_count
+        """
+        cursor = self.conn.cursor()
+        rows = cursor.execute(
+            """
+            SELECT *
+            FROM oracle_embeddings
+            WHERE quality_score >= ? AND reuse_count >= ?
+            ORDER BY (quality_score * reuse_count) DESC
+            LIMIT ?
+        """,
+            (min_quality, min_reuse, limit),
+        ).fetchall()
+
+        return [self._row_to_oracle_record(row) for row in rows]
+
+    def get_oracle_lineage(self, oracle_id: str) -> list[OracleRecord]:
+        """Get oracle lineage (parent chain).
+
+        Args:
+            oracle_id: Oracle ID to trace
+
+        Returns:
+            List of oracles from newest to oldest (child to ancestor)
+        """
+        lineage = []
+        current_id = oracle_id
+        visited = set()
+
+        while current_id and current_id not in visited:
+            visited.add(current_id)
+
+            cursor = self.conn.cursor()
+            row = cursor.execute(
+                "SELECT * FROM oracle_embeddings WHERE oracle_id = ?",
+                (current_id,),
+            ).fetchone()
+
+            if not row:
+                break
+
+            oracle = self._row_to_oracle_record(row)
+            lineage.append(oracle)
+
+            # Move to parent
+            current_id = oracle.parent_oracle_id
+
+        return lineage
+
+    def get_oracle_descendants(self, oracle_id: str) -> list[OracleRecord]:
+        """Get all oracles derived from this oracle.
+
+        Args:
+            oracle_id: Parent oracle ID
+
+        Returns:
+            List of descendant oracles
+        """
+        cursor = self.conn.cursor()
+        rows = cursor.execute(
+            """
+            SELECT * FROM oracle_embeddings
+            WHERE parent_oracle_id = ?
+            ORDER BY timestamp ASC
+        """,
+            (oracle_id,),
+        ).fetchall()
+
+        return [self._row_to_oracle_record(row) for row in rows]
+
+    def __enter__(self):
+        """Context manager entry."""
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """Context manager exit."""
+        self.close()