235 lines
6.9 KiB
Python
235 lines
6.9 KiB
Python
#!/usr/bin/env python3
|
|
"""SQLite 任务持久化"""
|
|
import json
|
|
import sqlite3
|
|
import time
|
|
import uuid
|
|
|
|
from config import DB_PATH
|
|
|
|
|
|
def _conn():
|
|
c = sqlite3.connect(DB_PATH, timeout=30)
|
|
c.row_factory = sqlite3.Row
|
|
return c
|
|
|
|
|
|
def init_db():
|
|
c = _conn()
|
|
c.execute('''
|
|
CREATE TABLE IF NOT EXISTS tasks (
|
|
id TEXT PRIMARY KEY,
|
|
url TEXT NOT NULL,
|
|
goal TEXT NOT NULL,
|
|
status TEXT DEFAULT 'queued',
|
|
result TEXT DEFAULT 'pending',
|
|
max_steps INTEGER DEFAULT 30,
|
|
timeout INTEGER DEFAULT 600,
|
|
created_at REAL,
|
|
started_at REAL,
|
|
finished_at REAL,
|
|
steps INTEGER DEFAULT 0,
|
|
summary TEXT,
|
|
error TEXT,
|
|
report_path TEXT
|
|
)
|
|
''')
|
|
# 大模型配置表
|
|
c.execute('''
|
|
CREATE TABLE IF NOT EXISTS llm_configs (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
name TEXT NOT NULL,
|
|
base_url TEXT NOT NULL,
|
|
api_key TEXT NOT NULL,
|
|
model TEXT NOT NULL,
|
|
vision INTEGER DEFAULT 0,
|
|
temperature REAL DEFAULT 0.2,
|
|
timeout INTEGER DEFAULT 120,
|
|
is_default INTEGER DEFAULT 0,
|
|
created_at REAL
|
|
)
|
|
''')
|
|
# tasks 表迁移：新增模型相关列（老库升级）
|
|
cols = [r[1] for r in c.execute('PRAGMA table_info(tasks)')]
|
|
if 'llm_config_id' not in cols:
|
|
c.execute('ALTER TABLE tasks ADD COLUMN llm_config_id INTEGER')
|
|
if 'llm_name' not in cols:
|
|
c.execute('ALTER TABLE tasks ADD COLUMN llm_name TEXT')
|
|
if 'vision' not in cols:
|
|
c.execute('ALTER TABLE tasks ADD COLUMN vision INTEGER DEFAULT 0')
|
|
# 首次启动：写入默认模型配置
|
|
n = c.execute('SELECT COUNT(*) FROM llm_configs').fetchone()[0]
|
|
if n == 0:
|
|
from config import (LLM_BASE_URL, LLM_API_KEY, LLM_MODEL,
|
|
LLM_TEMPERATURE, LLM_TIMEOUT, DEFAULT_LLM_NAME)
|
|
c.execute(
|
|
'INSERT INTO llm_configs (name, base_url, api_key, model, vision, '
|
|
'temperature, timeout, is_default, created_at) VALUES (?,?,?,?,?,?,?,?,?)',
|
|
(DEFAULT_LLM_NAME, LLM_BASE_URL, LLM_API_KEY, LLM_MODEL, 0,
|
|
LLM_TEMPERATURE, LLM_TIMEOUT, 1, time.time()))
|
|
c.commit()
|
|
c.close()
|
|
|
|
|
|
def create_task(url, goal, max_steps, timeout, llm_config_id=None, llm_name=None, vision=0):
|
|
tid = uuid.uuid4().hex[:12]
|
|
c = _conn()
|
|
c.execute(
|
|
'INSERT INTO tasks (id, url, goal, status, max_steps, timeout, created_at, '
|
|
'llm_config_id, llm_name, vision) VALUES (?,?,?,?,?,?,?,?,?,?)',
|
|
(tid, url, goal, 'queued', max_steps, timeout, time.time(),
|
|
llm_config_id, llm_name, 1 if vision else 0))
|
|
c.commit()
|
|
c.close()
|
|
return tid
|
|
|
|
|
|
def update_task(tid, **fields):
|
|
allowed = {'status', 'result', 'started_at', 'finished_at', 'steps',
|
|
'summary', 'error', 'report_path'}
|
|
sets = [f'{k}=?' for k in fields if k in allowed]
|
|
vals = [fields[k] for k in fields if k in allowed]
|
|
if not sets:
|
|
return
|
|
c = _conn()
|
|
c.execute(f'UPDATE tasks SET {", ".join(sets)} WHERE id=?', (*vals, tid))
|
|
c.commit()
|
|
c.close()
|
|
|
|
|
|
def get_task(tid):
|
|
c = _conn()
|
|
row = c.execute('SELECT * FROM tasks WHERE id=?', (tid,)).fetchone()
|
|
c.close()
|
|
return dict(row) if row else None
|
|
|
|
|
|
def list_tasks(limit=50):
|
|
c = _conn()
|
|
rows = c.execute(
|
|
'SELECT * FROM tasks ORDER BY created_at DESC LIMIT ?', (limit,)
|
|
).fetchall()
|
|
c.close()
|
|
return [dict(r) for r in rows]
|
|
|
|
|
|
def save_step_log(tid, step):
|
|
"""把单个步骤 JSON 追加到任务目录的 steps.jsonl"""
|
|
from config import TASKS_DIR
|
|
import os
|
|
p = os.path.join(TASKS_DIR, tid, 'steps.jsonl')
|
|
with open(p, 'a', encoding='utf-8') as f:
|
|
f.write(json.dumps(step, ensure_ascii=False) + '\n')
|
|
|
|
|
|
def load_step_logs(tid):
|
|
from config import TASKS_DIR
|
|
import os
|
|
p = os.path.join(TASKS_DIR, tid, 'steps.jsonl')
|
|
if not os.path.exists(p):
|
|
return []
|
|
out = []
|
|
with open(p, encoding='utf-8') as f:
|
|
for line in f:
|
|
line = line.strip()
|
|
if line:
|
|
try:
|
|
out.append(json.loads(line))
|
|
except json.JSONDecodeError:
|
|
pass
|
|
return out
|
|
|
|
|
|
# ========== 大模型配置 ==========
|
|
|
|
|
|
def list_llm_configs():
|
|
c = _conn()
|
|
rows = c.execute(
|
|
'SELECT * FROM llm_configs ORDER BY is_default DESC, id ASC'
|
|
).fetchall()
|
|
c.close()
|
|
return [dict(r) for r in rows]
|
|
|
|
|
|
def get_llm_config(cid):
|
|
c = _conn()
|
|
row = c.execute('SELECT * FROM llm_configs WHERE id=?', (cid,)).fetchone()
|
|
c.close()
|
|
return dict(row) if row else None
|
|
|
|
|
|
def get_default_llm_config():
|
|
c = _conn()
|
|
row = c.execute('SELECT * FROM llm_configs WHERE is_default=1').fetchone()
|
|
if not row:
|
|
row = c.execute('SELECT * FROM llm_configs ORDER BY id ASC').fetchone()
|
|
c.close()
|
|
return dict(row) if row else None
|
|
|
|
|
|
def create_llm_config(name, base_url, api_key, model, vision=0,
|
|
temperature=0.2, timeout=120, is_default=False):
|
|
c = _conn()
|
|
if is_default:
|
|
c.execute('UPDATE llm_configs SET is_default=0')
|
|
c.execute(
|
|
'INSERT INTO llm_configs (name, base_url, api_key, model, vision, '
|
|
'temperature, timeout, is_default, created_at) VALUES (?,?,?,?,?,?,?,?,?)',
|
|
(name, base_url, api_key, model, 1 if vision else 0,
|
|
temperature, timeout, 1 if is_default else 0, time.time()))
|
|
c.commit()
|
|
cid = c.execute('SELECT last_insert_rowid()').fetchone()[0]
|
|
c.close()
|
|
return cid
|
|
|
|
|
|
def update_llm_config(cid, **fields):
|
|
allowed = {'name', 'base_url', 'api_key', 'model', 'vision',
|
|
'temperature', 'timeout', 'is_default'}
|
|
c = _conn()
|
|
if fields.get('is_default'):
|
|
c.execute('UPDATE llm_configs SET is_default=0')
|
|
sets = []
|
|
vals = []
|
|
for k, v in fields.items():
|
|
if k not in allowed:
|
|
continue
|
|
if k == 'vision':
|
|
v = 1 if v else 0
|
|
if k == 'is_default':
|
|
v = 1 if v else 0
|
|
sets.append(f'{k}=?')
|
|
vals.append(v)
|
|
if sets:
|
|
c.execute(f'UPDATE llm_configs SET {", ".join(sets)} WHERE id=?', (*vals, cid))
|
|
c.commit()
|
|
c.close()
|
|
|
|
|
|
def delete_llm_config(cid):
|
|
c = _conn()
|
|
c.execute('DELETE FROM llm_configs WHERE id=?', (cid,))
|
|
c.commit()
|
|
c.close()
|
|
|
|
|
|
def set_default_llm_config(cid):
|
|
c = _conn()
|
|
c.execute('UPDATE llm_configs SET is_default=0')
|
|
c.execute('UPDATE llm_configs SET is_default=1 WHERE id=?', (cid,))
|
|
c.commit()
|
|
c.close()
|
|
|
|
|
|
def llm_config_to_dict(row):
|
|
"""把配置行转成 llm.py 可用的 dict（不含数据库元信息）"""
|
|
return {
|
|
'base_url': row['base_url'],
|
|
'api_key': row['api_key'],
|
|
'model': row['model'],
|
|
'temperature': row['temperature'],
|
|
'timeout': row['timeout'],
|
|
'vision': bool(row['vision']),
|
|
}
|