246 lines
7.5 KiB
Python
246 lines
7.5 KiB
Python
# -*- 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 TABLE IF NOT EXISTS documents (
|
||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||
project_id INTEGER NOT NULL,
|
||
name TEXT NOT NULL,
|
||
content TEXT DEFAULT '',
|
||
source TEXT DEFAULT 'manual', -- manual/file/url
|
||
chunk_size INTEGER DEFAULT 0,
|
||
created_at INTEGER,
|
||
updated_at INTEGER
|
||
);
|
||
|
||
CREATE TABLE IF NOT EXISTS doc_chunks (
|
||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||
document_id INTEGER NOT NULL,
|
||
idx INTEGER DEFAULT 0,
|
||
content TEXT DEFAULT '',
|
||
tokens INTEGER DEFAULT 0
|
||
);
|
||
|
||
CREATE TABLE IF NOT EXISTS alerts (
|
||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||
type TEXT DEFAULT 'system', -- budget/task_failed/worker/limit/notify/plan
|
||
level TEXT DEFAULT 'info', -- info/warn/critical
|
||
title TEXT DEFAULT '',
|
||
detail TEXT DEFAULT '',
|
||
read INTEGER DEFAULT 0,
|
||
created_at INTEGER
|
||
);
|
||
|
||
CREATE TABLE IF NOT EXISTS api_tokens (
|
||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||
name TEXT NOT NULL,
|
||
token TEXT NOT NULL UNIQUE,
|
||
created_at INTEGER,
|
||
last_used_at INTEGER
|
||
);
|
||
|
||
CREATE TABLE IF NOT EXISTS notify_channels (
|
||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||
name TEXT NOT NULL,
|
||
type TEXT NOT NULL, -- feishu/wecom/email
|
||
webhook TEXT DEFAULT '',
|
||
email TEXT DEFAULT '',
|
||
events TEXT DEFAULT '[]', -- JSON: task_review/task_done/task_failed/budget_alert/worker_alert
|
||
enabled INTEGER DEFAULT 1,
|
||
created_at INTEGER
|
||
);
|
||
|
||
CREATE TABLE IF NOT EXISTS settings (
|
||
key TEXT PRIMARY KEY,
|
||
value TEXT DEFAULT ''
|
||
);
|
||
|
||
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);
|
||
CREATE INDEX IF NOT EXISTS idx_docs_project ON documents(project_id);
|
||
CREATE INDEX IF NOT EXISTS idx_chunks_doc ON doc_chunks(document_id);
|
||
CREATE INDEX IF NOT EXISTS idx_alerts_read ON alerts(read);
|
||
"""
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 迁移：给旧表补列（V1）
|
||
# ---------------------------------------------------------------------------
|
||
def _migrate():
|
||
conn = get_conn()
|
||
cols = {r['name'] for r in conn.execute('PRAGMA table_info(tasks)')}
|
||
if 'depends_on' not in cols:
|
||
conn.execute("ALTER TABLE tasks ADD COLUMN depends_on TEXT DEFAULT '[]'")
|
||
conn.execute('CREATE INDEX IF NOT EXISTS idx_tasks_depends ON tasks(depends_on)')
|
||
conn.commit()
|
||
conn.close()
|
||
|
||
|
||
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()
|
||
_migrate()
|
||
|
||
|
||
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'])
|
||
try:
|
||
t['depends_on'] = json.loads(t.get('depends_on') or '[]')
|
||
except Exception:
|
||
t['depends_on'] = []
|
||
return t
|
||
|
||
|
||
def get_setting(key, default=''):
|
||
r = q('SELECT value FROM settings WHERE key=?', (key,), one=True)
|
||
return r['value'] if r else default
|
||
|
||
|
||
def set_setting(key, value):
|
||
conn = get_conn()
|
||
try:
|
||
conn.execute('INSERT INTO settings (key, value) VALUES (?,?) '
|
||
'ON CONFLICT(key) DO UPDATE SET value=excluded.value', (key, str(value)))
|
||
conn.commit()
|
||
finally:
|
||
conn.close()
|