ai-worker-platform/db.py

# -*- coding: utf-8 -*-
"""
数据库层：SQLite + WAL，轻量直连封装
"""
import sqlite3
import json
import time
from config import DB_PATH

SCHEMA = """
CREATE TABLE IF NOT EXISTS projects (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    name TEXT NOT NULL,
    description TEXT DEFAULT '',
    objective TEXT DEFAULT '',
    acceptance_criteria TEXT DEFAULT '',
    status TEXT DEFAULT 'active',          -- planning/active/done/archived
    budget_limit REAL DEFAULT 0,           -- 项目预算上限（元），0=不限
    created_at INTEGER,
    updated_at INTEGER
);

CREATE TABLE IF NOT EXISTS workers (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    name TEXT NOT NULL,
    description TEXT DEFAULT '',
    provider TEXT NOT NULL,
    model TEXT NOT NULL,
    base_url TEXT DEFAULT '',
    api_key TEXT DEFAULT '',               -- 留空则用供应商全局 key
    system_prompt TEXT DEFAULT '',
    temperature REAL DEFAULT 0.7,
    max_tokens INTEGER DEFAULT 2000,
    task_cost_limit REAL DEFAULT 0,        -- 单任务成本上限（元），0=不限
    monthly_cost_limit REAL DEFAULT 0,     -- 月度成本上限（元），0=不限
    status TEXT DEFAULT 'enabled',         -- enabled/disabled
    created_at INTEGER,
    updated_at INTEGER
);

CREATE TABLE IF NOT EXISTS tasks (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    project_id INTEGER NOT NULL,
    worker_id INTEGER,                     -- NULL = 自动路由
    title TEXT NOT NULL,
    description TEXT DEFAULT '',
    status TEXT DEFAULT 'todo',            -- todo/running/review/done/rejected/failed/cancelled
    priority TEXT DEFAULT 'medium',        -- high/medium/low
    review_required INTEGER DEFAULT 1,     -- 完成后是否需要人工审核
    output_text TEXT DEFAULT '',
    output_version INTEGER DEFAULT 0,
    rejection_count INTEGER DEFAULT 0,
    error TEXT DEFAULT '',
    deadline TEXT DEFAULT '',
    created_at INTEGER,
    updated_at INTEGER,
    started_at INTEGER,
    finished_at INTEGER
);

CREATE TABLE IF NOT EXISTS task_logs (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    task_id INTEGER NOT NULL,
    level TEXT DEFAULT 'info',             -- info/success/warn/error
    message TEXT DEFAULT '',
    created_at INTEGER
);

CREATE TABLE IF NOT EXISTS cost_records (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    task_id INTEGER,
    project_id INTEGER,
    worker_id INTEGER,
    provider TEXT DEFAULT '',
    model TEXT DEFAULT '',
    prompt_tokens INTEGER DEFAULT 0,
    completion_tokens INTEGER DEFAULT 0,
    total_tokens INTEGER DEFAULT 0,
    cost REAL DEFAULT 0,
    created_at INTEGER
);

CREATE INDEX IF NOT EXISTS idx_tasks_project ON tasks(project_id);
CREATE INDEX IF NOT EXISTS idx_tasks_status ON tasks(status);
CREATE INDEX IF NOT EXISTS idx_logs_task ON task_logs(task_id);
CREATE INDEX IF NOT EXISTS idx_cost_task ON cost_records(task_id);
CREATE INDEX IF NOT EXISTS idx_cost_project ON cost_records(project_id);
"""


def get_conn():
    conn = sqlite3.connect(DB_PATH, timeout=30)
    conn.row_factory = sqlite3.Row
    conn.execute('PRAGMA journal_mode=WAL')
    conn.execute('PRAGMA foreign_keys=ON')
    return conn


def init_db():
    conn = get_conn()
    conn.executescript(SCHEMA)
    conn.commit()
    conn.close()


def q(sql, args=(), one=False):
    """查询"""
    conn = get_conn()
    try:
        cur = conn.execute(sql, args)
        rows = [dict(r) for r in cur.fetchall()]
        return (rows[0] if rows else None) if one else rows
    finally:
        conn.close()


def w(sql, args=()):
    """写入，返回 lastrowid"""
    conn = get_conn()
    try:
        cur = conn.execute(sql, args)
        conn.commit()
        return cur.lastrowid
    finally:
        conn.close()


def now():
    return int(time.time())


# ---------------------------------------------------------------------------
# 成本统计辅助
# ---------------------------------------------------------------------------
def monthly_worker_cost(worker_id, month_ts=None):
    """某 Worker 当月累计成本（元）"""
    if month_ts is None:
        import datetime
        month_ts = int(datetime.datetime.now().replace(day=1, hour=0, minute=0, second=0, microsecond=0).timestamp())
    rows = q(
        'SELECT COALESCE(SUM(cost),0) AS total FROM cost_records '
        'WHERE worker_id=? AND created_at>=?', (worker_id, month_ts))
    return rows[0]['total'] if rows else 0.0


def task_worker_cost(worker_id):
    """某 Worker 最近一次任务成本（用于单任务上限判断前先看历史，不作为硬限制）"""
    rows = q(
        'SELECT COALESCE(SUM(cost),0) AS total FROM cost_records WHERE worker_id=?',
        (worker_id,))
    return rows[0]['total'] if rows else 0.0


def serialize_task(t):
    t = dict(t)
    t['review_required'] = bool(t['review_required'])
    return t