vent_agent_2.0/db/business_store.py

"""业务身份、任务和推送状态；复用应用 SQLite 事务连接。"""
import hashlib
import json

from db import db_conn


def ensure_state(name):
    with db_conn.transaction():
        if not db_conn.query_one("SELECT name FROM business_agent_state WHERE name = ?", (name,)):
            db_conn.execute("INSERT INTO business_agent_state (name, local_enabled, gateway_online, updated_at) VALUES (?, 1, 1, ?)",
                            (name, db_conn.fmt_dt()))


def state(name):
    ensure_state(name)
    return db_conn.query_one("SELECT * FROM business_agent_state WHERE name = ?", (name,))


def set_state(name, *, local_enabled=None, gateway_online=None):
    ensure_state(name)
    for key, value in (("local_enabled", local_enabled), ("gateway_online", gateway_online)):
        if value is not None:
            db_conn.execute(f"UPDATE business_agent_state SET {key} = ?, updated_at = ? WHERE name = ?",
                            (int(bool(value)), db_conn.fmt_dt(), name))


def session_agent(session_id):
    row = db_conn.query_one("SELECT agent_name FROM business_session_agents WHERE session_id = ?", (session_id,))
    return row["agent_name"] if row else None


def bind_session(session_id, name):
    with db_conn.transaction():
        existing = session_agent(session_id)
        # 旧会话没有绑定表记录时，已有对话仍归属主 Agent。
        if existing is None and db_conn.scalar("SELECT COUNT(*) FROM messages WHERE session_id = ?", (session_id,)):
            existing = "main"
        if existing is not None and existing != name:
            raise ValueError("会话已绑定其他智能体，请新建会话")
        if session_agent(session_id) is None:
            db_conn.execute("INSERT INTO business_session_agents (session_id, agent_name) VALUES (?, ?)", (session_id, name))


def _json(value):
    return json.dumps(value, ensure_ascii=False, sort_keys=True, allow_nan=False)


def enqueue(name, user_name, data, files, *, request_id=None, source="manual", publish=False):
    payload = _json({"input": data, "files": files, "publish": publish})
    job_id = (hashlib.sha256(f"{name}\0{user_name}\0{source}\0{request_id}".encode()).hexdigest()[:32]
              if request_id else db_conn.new_id())
    with db_conn.transaction():
        previous = get(job_id)
        if previous:
            if previous["payload"] != payload:
                raise ValueError("同一 request_id 已用于不同任务输入")
            return previous, False
        now = db_conn.fmt_dt()
        db_conn.execute(
            "INSERT INTO business_jobs (id, agent_name, user_name, source, payload, status, push_status, push_attempts, created_at, updated_at)"
            " VALUES (?, ?, ?, ?, ?, 'queued', 'none', 0, ?, ?)",
            (job_id, name, user_name, source, payload, now, now))
    return get(job_id), True


def get(job_id):
    return db_conn.query_one("SELECT * FROM business_jobs WHERE id = ?", (job_id,))


def list_jobs(name=None, user_name=None, limit=50):
    clauses, values = [], []
    for key, value in (("agent_name", name), ("user_name", user_name)):
        if value is not None:
            clauses.append(f"{key} = ?")
            values.append(value)
    where = " WHERE " + " AND ".join(clauses) if clauses else ""
    return db_conn.query("SELECT * FROM business_jobs" + where + " ORDER BY created_at DESC LIMIT ?", (*values, limit))


def claim(name):
    with db_conn.transaction():
        row = db_conn.query_one("SELECT * FROM business_jobs WHERE agent_name = ? AND status = 'queued' ORDER BY created_at, id LIMIT 1", (name,))
        if not row:
            return None
        updated = db_conn.execute("UPDATE business_jobs SET status = 'running', updated_at = ? WHERE id = ? AND status = 'queued'",
                                  (db_conn.fmt_dt(), row["id"]))
        return get(row["id"]) if updated else None


_FIELDS = {"status", "session_id", "run_id", "error", "result_json", "push_status", "push_attempts",
           "next_push_at", "receipt_json", "push_payload", "finished_at", "agent_code"}


def update(job_id, **fields):
    if not fields or set(fields) - _FIELDS:
        raise ValueError("业务任务更新字段非法")
    assignments = ", ".join(f"{key} = ?" for key in fields)
    db_conn.execute(f"UPDATE business_jobs SET {assignments}, updated_at = ? WHERE id = ?",
                    (*fields.values(), db_conn.fmt_dt(), job_id))


def pending_pushes(now):
    return db_conn.query("SELECT * FROM business_jobs WHERE push_status = 'pending' AND (next_push_at IS NULL OR next_push_at <= ?) ORDER BY updated_at LIMIT 50", (now,))


def push_target(job_id):
    return db_conn.scalar("SELECT endpoint FROM business_push_targets WHERE job_id = ?", (job_id,))


def freeze_push(job_id, payload, endpoint):
    """首次排队即冻结报文和目标；重试不会跟随配置切换到另一个平台。"""
    if not endpoint:
        raise ValueError("结果接收平台地址尚未配置")
    content = _json(payload)
    with db_conn.transaction():
        job = get(job_id)
        if not job:
            raise ValueError("结果所属任务不存在")
        if job.get("push_payload") and _json(json.loads(job["push_payload"])) != content:
            raise ValueError("已冻结的推送内容不能改写")
        previous = push_target(job_id)
        if previous and previous != endpoint:
            raise ValueError("结果接收平台地址已改变，不能改投历史结果")
        if previous is None:
            db_conn.execute("INSERT INTO business_push_targets (job_id, endpoint, created_at) VALUES (?, ?, ?)",
                            (job_id, endpoint, db_conn.fmt_dt()))
        update(job_id, push_payload=content, agent_code=payload["agentCode"])
    return payload


def record_completed_result(name, user_name, session_id, data, result, endpoint):
    """网页/工具已完成的研判直接进入结果队列，不再次执行智能体。"""
    payload = result["platform_payload"]
    job_id = payload["taskId"]
    with db_conn.transaction():
        existing = get(job_id)
        if existing:
            if (existing["agent_name"], existing["user_name"], existing["session_id"]) != (name, user_name, session_id):
                raise ValueError("结果任务身份冲突")
            freeze_push(job_id, payload, endpoint)
            return existing
        now = db_conn.fmt_dt()
        db_conn.execute(
            "INSERT INTO business_jobs (id, agent_name, user_name, source, payload, status, session_id, result_json,"
            " push_status, push_attempts, created_at, updated_at, finished_at)"
            " VALUES (?, ?, ?, 'chat', ?, ?, ?, ?, 'pending', 0, ?, ?, ?)",
            (job_id, name, user_name, _json({"input": data, "files": [], "publish": True}), result["status"],
             session_id, _json(result), now, now, payload["checkedAt"]))
        freeze_push(job_id, payload, endpoint)
    return get(job_id)


def recover():
    # 不自动重跑已经调用模型但尚未保存结果的任务，避免重复计费与副作用。
    db_conn.execute("UPDATE business_jobs SET status = 'interrupted', error = 'process_restarted', updated_at = ?, finished_at = ? WHERE status = 'running'",
                    (db_conn.fmt_dt(), db_conn.fmt_dt()))


def retry(job_id):
    return db_conn.execute("UPDATE business_jobs SET status = 'queued', error = NULL, finished_at = NULL, updated_at = ? WHERE id = ? AND status IN ('failed', 'cancelled', 'interrupted')",
                          (db_conn.fmt_dt(), job_id))


def metrics(name):
    rows = db_conn.query("SELECT status, COUNT(*) AS n FROM business_jobs WHERE agent_name = ? GROUP BY status", (name,))
    counts = {row["status"]: row["n"] for row in rows}
    executed = db_conn.query("SELECT e.status, COUNT(*) AS n FROM execution_runs e JOIN business_session_agents b ON b.session_id = e.session_id WHERE b.agent_name = ? GROUP BY e.status", (name,))
    return {"tasks": counts, "completed": counts.get("completed", 0),
            "executions": {row["status"]: row["n"] for row in executed},
            "pushed": db_conn.scalar("SELECT COUNT(*) FROM business_jobs WHERE agent_name = ? AND push_status = 'pushed'", (name,))}


def record_file(job_id, metadata):
    db_conn.execute("INSERT INTO business_files (id, job_id, file_id, version, filename, local_path, sha256, size, created_at) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)",
                    (db_conn.new_id(), job_id, metadata["file_id"], metadata["version"], metadata["filename"],
                     metadata["local_path"], metadata["sha256"], metadata["size"], db_conn.fmt_dt()))