agent-platform/agent_platform/db.py

"""SQLite-Layer: Chats, Messages, Audit, Sessions, Agents.

Agents können sein:
- source='code': aus agents/*.py, DB ist Override
- source='db':   nur in DB, via API/UI angelegt
"""
import aiosqlite
import json
import uuid
import secrets
from datetime import datetime, timedelta, timezone
from pathlib import Path
from config import settings


SCHEMA = """
CREATE TABLE IF NOT EXISTS agents (
    id TEXT PRIMARY KEY,
    name TEXT NOT NULL,
    description TEXT DEFAULT '',
    system_prompt TEXT NOT NULL,
    allowed_tools TEXT DEFAULT '[]',
    model TEXT,
    temperature REAL DEFAULT 0.7,
    max_tokens INTEGER DEFAULT 2000,
    enabled INTEGER DEFAULT 1,
    source TEXT DEFAULT 'db',
    created_at TEXT DEFAULT CURRENT_TIMESTAMP,
    updated_at TEXT DEFAULT CURRENT_TIMESTAMP
);

CREATE TABLE IF NOT EXISTS conversations (
    id TEXT PRIMARY KEY,
    agent_id TEXT NOT NULL,
    user_id TEXT,
    title TEXT,
    created_at TEXT DEFAULT CURRENT_TIMESTAMP,
    updated_at TEXT DEFAULT CURRENT_TIMESTAMP
);

CREATE TABLE IF NOT EXISTS messages (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    conversation_id TEXT NOT NULL,
    role TEXT NOT NULL,
    content TEXT,
    tool_calls TEXT,
    tool_results TEXT,
    token_count INTEGER,
    created_at TEXT DEFAULT CURRENT_TIMESTAMP
);

CREATE TABLE IF NOT EXISTS audit (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    timestamp TEXT DEFAULT CURRENT_TIMESTAMP,
    agent_id TEXT,
    user_id TEXT,
    action TEXT,
    target TEXT,
    args TEXT,
    result TEXT,
    duration_ms INTEGER
);

CREATE TABLE IF NOT EXISTS sessions (
    token TEXT PRIMARY KEY,
    user_id TEXT,
    created_at TEXT DEFAULT CURRENT_TIMESTAMP,
    expires_at TEXT
);

CREATE INDEX IF NOT EXISTS idx_messages_conv ON messages(conversation_id);
CREATE INDEX IF NOT EXISTS idx_audit_agent ON audit(agent_id);
CREATE INDEX IF NOT EXISTS idx_audit_timestamp ON audit(timestamp);
CREATE INDEX IF NOT EXISTS idx_conversations_user ON conversations(user_id);
"""


async def init_db():
    """Erstellt DB und Schema."""
    Path(settings.DB_PATH).parent.mkdir(parents=True, exist_ok=True)
    async with aiosqlite.connect(settings.DB_PATH) as db:
        await db.executescript(SCHEMA)
        await db.commit()


async def get_db_connection():
    """Gibt eine aiosqlite-Connection zurück."""
    return await aiosqlite.connect(settings.DB_PATH)


async def get_db():
    """FastAPI-Dependency: aiosqlite-Connection pro Request, mit automatischem Cleanup."""
    db = await aiosqlite.connect(settings.DB_PATH)
    try:
        yield db
    finally:
        await db.close()


# === Agents ===

def _row_to_agent(row):
    return {
        "id": row[0],
        "name": row[1],
        "description": row[2] or "",
        "system_prompt": row[3],
        "allowed_tools": json.loads(row[4]) if row[4] else [],
        "model": row[5],
        "temperature": row[6] if row[6] is not None else 0.7,
        "max_tokens": row[7] if row[7] is not None else 2000,
        "enabled": bool(row[8]),
        "source": row[9] or "db",
        "created_at": row[10],
        "updated_at": row[11],
    }


async def upsert_agent(db, agent: dict, source: str = "db"):
    """Insert oder Update. Code-Agents werden beim ersten Start hier gesynct."""
    await db.execute(
        """INSERT INTO agents (id, name, description, system_prompt, allowed_tools, model, temperature, max_tokens, enabled, source, updated_at)
           VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, CURRENT_TIMESTAMP)
           ON CONFLICT(id) DO UPDATE SET
             name=excluded.name,
             description=excluded.description,
             system_prompt=excluded.system_prompt,
             allowed_tools=excluded.allowed_tools,
             model=excluded.model,
             temperature=excluded.temperature,
             max_tokens=excluded.max_tokens,
             enabled=excluded.enabled,
             updated_at=CURRENT_TIMESTAMP""",
        (
            agent["id"],
            agent["name"],
            agent.get("description", ""),
            agent["system_prompt"],
            json.dumps(agent.get("allowed_tools", [])),
            agent.get("model"),
            agent.get("temperature", 0.7),
            agent.get("max_tokens", 2000),
            int(agent.get("enabled", True)),
            source,
        )
    )
    await db.commit()


async def get_agent(db, agent_id: str):
    async with db.execute(
        "SELECT id, name, description, system_prompt, allowed_tools, model, temperature, max_tokens, enabled, source, created_at, updated_at FROM agents WHERE id = ?",
        (agent_id,)
    ) as cur:
        row = await cur.fetchone()
        return _row_to_agent(row) if row else None


async def list_agents(db, enabled_only=False):
    query = "SELECT id, name, description, system_prompt, allowed_tools, model, temperature, max_tokens, enabled, source, created_at, updated_at FROM agents"
    params = []
    if enabled_only:
        query += " WHERE enabled = 1"
    query += " ORDER BY id ASC"
    async with db.execute(query, params) as cur:
        return [_row_to_agent(row) for row in await cur.fetchall()]


async def delete_agent(db, agent_id: str) -> bool:
    """Löscht einen Agent. Code-Agents können gelöscht werden, sind dann bis zum nächsten Sync weg."""
    async with db.execute("DELETE FROM agents WHERE id = ?", (agent_id,)) as cur:
        await db.commit()
        return cur.rowcount > 0


async def set_agent_enabled(db, agent_id: str, enabled: bool):
    await db.execute(
        "UPDATE agents SET enabled = ?, updated_at = CURRENT_TIMESTAMP WHERE id = ?",
        (int(enabled), agent_id)
    )
    await db.commit()


# === Conversations ===

async def create_conversation(db, agent_id, user_id="default", title="Neuer Chat"):
    conv_id = str(uuid.uuid4())
    await db.execute(
        "INSERT INTO conversations (id, agent_id, user_id, title) VALUES (?, ?, ?, ?)",
        (conv_id, agent_id, user_id, title)
    )
    await db.commit()
    return conv_id


async def list_conversations(db, agent_id=None, user_id="default", limit=20):
    query = "SELECT id, agent_id, title, created_at FROM conversations WHERE user_id = ?"
    params = [user_id]
    if agent_id:
        query += " AND agent_id = ?"
        params.append(agent_id)
    query += " ORDER BY updated_at DESC LIMIT ?"
    params.append(limit)
    async with db.execute(query, params) as cur:
        return await cur.fetchall()


async def touch_conversation(db, conv_id):
    await db.execute(
        "UPDATE conversations SET updated_at = CURRENT_TIMESTAMP WHERE id = ?",
        (conv_id,)
    )
    await db.commit()


# === Messages ===

async def add_message(db, conversation_id, role, content, tool_calls=None, tool_results=None, token_count=0):
    await db.execute(
        """INSERT INTO messages (conversation_id, role, content, tool_calls, tool_results, token_count)
           VALUES (?, ?, ?, ?, ?, ?)""",
        (conversation_id, role, content,
         json.dumps(tool_calls) if tool_calls else None,
         json.dumps(tool_results) if tool_results else None,
         token_count)
    )
    await db.commit()


async def get_messages(db, conversation_id, limit=None):
    query = "SELECT role, content, tool_calls, tool_results, token_count FROM messages WHERE conversation_id = ? ORDER BY id ASC"
    params = [conversation_id]
    if limit:
        query += " LIMIT ?"
        params.append(limit)
    async with db.execute(query, params) as cur:
        return await cur.fetchall()


# === Audit ===

async def log_audit(db, agent_id, user_id, action, target=None, args=None, result=None, duration_ms=None):
    await db.execute(
        """INSERT INTO audit (agent_id, user_id, action, target, args, result, duration_ms)
           VALUES (?, ?, ?, ?, ?, ?, ?)""",
        (agent_id, user_id, action, target,
         json.dumps(args) if args else None,
         json.dumps(result) if result else None,
         duration_ms)
    )
    await db.commit()


async def list_audit(db, agent_id=None, limit=100):
    query = "SELECT timestamp, agent_id, action, target, duration_ms FROM audit"
    params = []
    if agent_id:
        query += " WHERE agent_id = ?"
        params.append(agent_id)
    query += " ORDER BY id DESC LIMIT ?"
    params.append(limit)
    async with db.execute(query, params) as cur:
        return await cur.fetchall()


# === Sessions ===

async def create_session(db, user_id="default", ttl_hours=24):
    token = secrets.token_urlsafe(32)
    expires = (datetime.now(timezone.utc).replace(tzinfo=None) + timedelta(hours=ttl_hours)).isoformat()
    await db.execute(
        "INSERT INTO sessions (token, user_id, expires_at) VALUES (?, ?, ?)",
        (token, user_id, expires)
    )
    await db.commit()
    return token


async def validate_session(db, token):
    async with db.execute(
        "SELECT user_id, expires_at FROM sessions WHERE token = ?",
        (token,)
    ) as cur:
        row = await cur.fetchone()
        if not row:
            return None
        if row[1] and datetime.fromisoformat(row[1]) < datetime.now(timezone.utc).replace(tzinfo=None):
            return None
        return row[0]