diff --git a/.env.example b/.env.example index 8cb0676..e040040 100644 --- a/.env.example +++ b/.env.example @@ -1,3 +1,7 @@ +# OneBot 11 适配器(可选):NapCat / Lagrange 等实现反向 WS 连入 +# ws://<本服务>:3000/onebot/v11/ws 留空则不校验令牌 +OB_ACCESS_TOKEN= + # PostgreSQL 连接串(必需),例:postgresql://postgres:密码@localhost:5432/ai_chat # 数据库需要提前建好:CREATE DATABASE ai_chat; PG_DSN= diff --git a/app/main.py b/app/main.py index ffdb0df..d9b9dc4 100644 --- a/app/main.py +++ b/app/main.py @@ -27,6 +27,7 @@ from . import db, tools from .connector import (ASSISTANT_NAME, ASSISTANT_PERSONA, API_KEY, BATCH_CONFIGURED, complete, complete_night, embed, route, run_model_tools) +from .onebot import router as onebot_router, send_to_session load_dotenv() @@ -70,6 +71,7 @@ async def lifespan(_: FastAPI): app = FastAPI(lifespan=lifespan) +app.include_router(onebot_router) # 当前在线的连接:session -> {ws, ...},用于主动消息推送 connected: dict[str, set[WebSocket]] = {} @@ -91,20 +93,61 @@ def _daily_slot(session: str, key: str, lo: int, hi: int, today: str) -> int: return lo + h % (hi - lo) +async def chat_turn(session: str, text: str, notify) -> str: + """一次完整对话回合(TUI WS 与 OneBot 适配器共用): + 路由 → 分支 → 生成回复 → 落库 → 后台提取。返回完整回复文本。""" + sig = await route(text, datetime.now(), + [m["text"] for m in await db.get_memories(session)]) + await db.add_message(session, "user", text) + full = "" + try: + if sig.get("letter"): + days, lbody = sig["letter"] + if days == "ask": + full = ("想给未来的自己留封信呀。完整地跟我说一句:" + "「给7天后的自己写信:<想说的话>」,我替你收着。") + else: + deliver_at = (datetime.now() + + timedelta(days=days)).timestamp() + await db.add_letter(session, lbody[:2000], deliver_at) + full = f"✉️ 信我替你收好了,{days} 天后亲手交给你。" + elif sig.get("view"): + full = await compose_view(session, sig["view"]) + elif not API_KEY: + full = f"(回声模式,未配置 AI_API_KEY)你说的是:{text}" + else: + qvec = (await embed([text]))[0] + relevant = await db.search_memories(session, qvec) + messages = [{"role": "system", + "content": await system_prompt(session, relevant)}] + messages += await db.get_history(session) + full = await run_model_tools( + messages, tool_executor(session, sig.get("city_hint", ""))) + except Exception as e: # 连接器或上游报错,作为一条普通回复返回 + full = f"(出了点小差错:{e})" + if full: + # 注意:助手消息由各传输层在"确认送达"后自己落库 + # (TUI 发完就记,OneBot 在 retcode==0 后记,避免重复) + asyncio.create_task(remember_later(notify, session, text, sig)) + return full + + async def push_proactive(session: str, text: str) -> bool: - """把主动消息写进历史并推给在线的连接;不在线返回 False""" + """主动消息:优先推在线 TUI,OneBot 会话走 QQ;都不在线返回 False""" targets = {w for w in connected.get(session, set()) if w.client_state.name == "CONNECTED"} - if not targets: - return False - await db.add_message(session, "assistant", text) - payload = json.dumps({"type": "proactive", "text": text}) - for w in targets: - try: - await w.send_text(payload) - except Exception: - connected.get(session, set()).discard(w) - return True + if targets: + await db.add_message(session, "assistant", text) + payload = json.dumps({"type": "proactive", "text": text}) + for w in targets: + try: + await w.send_text(payload) + except Exception: + connected.get(session, set()).discard(w) + return True + if session.startswith("ob-"): + return await send_to_session(session, text) + return False async def lore_block(session: str) -> str: @@ -130,8 +173,7 @@ async def system_prompt(session: str, relevant: list[dict]) -> str: return "\n".join(parts) -async def remember_later(ws: WebSocket, session: str, text: str, - sig: dict) -> None: +async def remember_later(notify, session: str, text: str, sig: dict) -> None: """后台落库:路由器判出的 记忆/事件/作息/情绪/承诺/身世 + 连续天数""" now = datetime.now() facts = sig.get("facts") or [] @@ -143,7 +185,7 @@ async def remember_later(ws: WebSocket, session: str, text: str, async def send(payload: dict) -> None: try: - await ws.send_json(payload) + await notify(payload) except Exception: pass # 连接断了,数据已入库,不碍事 @@ -231,41 +273,29 @@ async def proactive_text(ev: dict) -> str: async def deliver_event(ev: dict, ws: WebSocket | None = None) -> bool: """推送一次性事件;谁都不在线则保持 pending,等下次上线补发""" - targets = ({ws} if ws - else {w for w in connected.get(ev["session"], set()) - if w.client_state.name == "CONNECTED"}) - if not targets: - return False text = await proactive_text(ev) - payload = json.dumps({"type": "proactive", "text": text, - "event_id": ev["id"]}) - for w in targets: - try: - await w.send_text(payload) - except Exception: - connected.get(ev["session"], set()).discard(w) - await db.add_message(ev["session"], "assistant", text) - await db.mark_event_done(ev["id"]) - return True + if ws is not None: # 补发场景:只发这条连接 + if ws.client_state.name != "CONNECTED": + return False + payload = json.dumps({"type": "proactive", "text": text, + "event_id": ev["id"]}) + await ws.send_text(payload) + await db.add_message(ev["session"], "assistant", text) + await db.mark_event_done(ev["id"]) + return True + if await push_proactive(ev["session"], text): + await db.mark_event_done(ev["id"]) + return True + return False async def deliver_daily_event(ev: dict) -> bool: """每日事件到点推送;不在线就保留,下次上线补发""" - targets = {w for w in connected.get(ev["session"], set()) - if w.client_state.name == "CONNECTED"} - if not targets: - return False text = await proactive_text(ev) - payload = json.dumps({"type": "proactive", "text": text, - "event_id": ev["id"]}) - for w in targets: - try: - await w.send_text(payload) - except Exception: - connected.get(ev["session"], set()).discard(w) - await db.add_message(ev["session"], "assistant", text) - await db.mark_event_fired(ev["id"], datetime.now().strftime("%Y-%m-%d")) - return True + if await push_proactive(ev["session"], text): + await db.mark_event_fired(ev["id"], datetime.now().strftime("%Y-%m-%d")) + return True + return False # ---- 作息问候 ---- @@ -787,7 +817,9 @@ async def scheduler() -> None: @app.get("/") async def index(): - return {"app": "小暖", "hint": "终端里运行 uv run python tui.py 开始聊天"} + return {"app": "小暖", + "hint": "终端聊天 uv run python tui.py;" + "OneBot 反向WS /onebot/v11/ws(可带 ?access_token=)"} @app.websocket("/ws") @@ -926,43 +958,11 @@ async def chat(ws: WebSocket): if not text: continue - # 小模型一次路由:想看哪份档案 / 是不是慢信 / 情绪 / 各类提取 - sig = await route(text, datetime.now(), - [m["text"] for m in - await db.get_memories(session)]) - await db.add_message(session, "user", text) await ws.send_json({"type": "start"}) - full = "" - try: - if sig.get("letter"): - days, lbody = sig["letter"] - if days == "ask": - full = ("想给未来的自己留封信呀。完整地跟我说一句:" - "「给7天后的自己写信:<想说的话>」,我替你收着。") - else: - deliver_at = (datetime.now() - + timedelta(days=days)).timestamp() - await db.add_letter(session, lbody[:2000], deliver_at) - full = f"✉️ 信我替你收好了,{days} 天后亲手交给你。" - elif sig.get("view"): - full = await compose_view(session, sig["view"]) - elif not API_KEY: - full = f"(回声模式,未配置 AI_API_KEY)你说的是:{text}" - else: - qvec = (await embed([text]))[0] - relevant = await db.search_memories(session, qvec) - messages = [{"role": "system", - "content": await system_prompt(session, relevant)}] - messages += await db.get_history(session) - full = await run_model_tools( - messages, tool_executor(session, - sig.get("city_hint", ""))) - await send_in_chunks(ws, full) - except Exception as e: # 连接器或上游报错,回传给前端 - await ws.send_json({"type": "error", "text": str(e)}) + full = await chat_turn(session, text, ws.send_json) + await send_in_chunks(ws, full) if full: await db.add_message(session, "assistant", full) - asyncio.create_task(remember_later(ws, session, text, sig)) await ws.send_json({"type": "done"}) except WebSocketDisconnect: pass diff --git a/app/onebot.py b/app/onebot.py new file mode 100644 index 0000000..002aebe --- /dev/null +++ b/app/onebot.py @@ -0,0 +1,185 @@ +"""OneBot 11 适配器:反向 WebSocket(NapCat / Lagrange 等实现连进来)。 + +接入方式(以 NapCat 为例):在 OneBot 实现里配置反向 WebSocket 地址 + ws://<本服务>:3000/onebot/v11/ws (配了 OB_ACCESS_TOKEN 时带 ?access_token=) + +- 私聊:所有消息都回 +- 群聊:只回应 @她 的消息,回复自动 @ 回去 +- 会话 id:ob--p / ob--g, + 与 TUI 会话隔离,各自有独立的记忆、作息和时间表 +- 主动消息(叫早 / 提醒 / 回访 / 慢信 / 连续天数…)经适配器发到对应 QQ +""" + +import asyncio +import json +import os +import re +import uuid + +from fastapi import APIRouter, WebSocket, WebSocketDisconnect + +from . import db + +router = APIRouter() + +# 可选访问令牌:OneBot 实现侧配置同样的 token 即可 +ACCESS_TOKEN = os.getenv("OB_ACCESS_TOKEN", "") + + +class OneBotConn: + def __init__(self, ws: WebSocket, self_id): + self.ws = ws + self.self_id = self_id + self.pending = {} # echo -> Future + self.targets = {} # session -> {"type", "group_id", "user_id"} + + +connections: dict = {} # self_id -> OneBotConn + + +def _extract_text(message, self_id) -> tuple[str, bool]: + """返回 (纯文本, 是否@了她)。兼容 CQ 码字符串与消息段数组两种格式""" + if isinstance(message, str): + at_me = f"[CQ:at,qq={self_id}]" in message + return re.sub(r"\[CQ:[^\]]*\]", "", message).strip(), at_me + parts, at_me = [], False + for seg in message or []: + t, d = seg.get("type"), seg.get("data") or {} + if t == "text": + parts.append(d.get("text", "")) + elif t == "at" and str(d.get("qq")) == str(self_id): + at_me = True + return "".join(parts).strip(), at_me + + +async def call_action(conn: OneBotConn, action: str, params: dict, + timeout: float = 30) -> dict: + """向 OneBot 实现发一个 action 并等它的 echo 响应""" + echo = uuid.uuid4().hex + fut = asyncio.get_running_loop().create_future() + conn.pending[echo] = fut + await conn.ws.send_text(json.dumps( + {"action": action, "params": params, "echo": echo})) + try: + return await asyncio.wait_for(fut, timeout) + finally: + conn.pending.pop(echo, None) + + +async def reply_message(conn: OneBotConn, session: str, text: str) -> bool: + """往会话对应的 QQ 对象发消息;发送成功才记进对话历史""" + tgt = conn.targets.get(session) + if not tgt: + return False + if tgt["type"] == "group": + action = "send_group_msg" + params = {"group_id": tgt["group_id"], + "message": f"[CQ:at,qq={tgt['user_id']}] {text}"} + else: + action = "send_private_msg" + params = {"user_id": tgt["user_id"], "message": text} + try: + res = await call_action(conn, action, params) + except Exception: + return False + if res.get("retcode") != 0: + return False + await db.add_message(session, "assistant", text) + return True + + +async def send_to_session(session: str, text: str) -> bool: + """主动消息入口(push_proactive 调用):找挂着这个会话的连接发出去""" + for conn in list(connections.values()): + if session in conn.targets: + return await reply_message(conn, session, text) + return False + + +def _notify_for(conn: OneBotConn, session: str): + """把 remember_later 的推送转成 QQ 消息(只挑用户真正关心的)""" + async def notify(payload: dict) -> None: + t = payload.get("type") + try: + if t == "event": + kind = "叫早" if payload.get("kind") == "alarm" else "提醒" + extra = ",每天都算" if payload.get("repeat") == "daily" else "" + await reply_message( + conn, session, + f"好呀,我记下了:{payload['fire_at'].replace('T', ' ')} " + f"{kind}({payload['text']}){extra}") + elif t == "care": + await reply_message( + conn, session, + "🫂 嗯,我记在心里了。明天这个时候我来问问你好不好些。" + if payload.get("kind") == "emotion" + else f"好,「{payload['note']}」这事儿我惦记着!" + "到时候我来问你结果~") + elif t == "streak": + await reply_message( + conn, session, + f"🔥 对了,今天是我们连续陪伴的第 {payload['days']} 天") + elif t == "lore_saved": + await reply_message( + conn, session, + "🧬 这段身世我刻进脑海里了(" + + "、".join(payload["topics"]) + "),永远记得,绝不瞎说。") + # memories / routines 等后台数据推送不发 QQ,静默 + except Exception: + pass # 连接断了就算了,数据已入库 + return notify + + +async def _process_message(conn: OneBotConn, data: dict) -> None: + mtype = data.get("message_type") + user_id = data.get("user_id") + text, at_me = _extract_text(data.get("message"), conn.self_id) + if mtype == "private": + session = f"ob-{conn.self_id}-p{user_id}" + conn.targets[session] = {"type": "private", "user_id": user_id} + else: + session = f"ob-{conn.self_id}-g{data.get('group_id')}" + conn.targets[session] = {"type": "group", + "group_id": data.get("group_id"), + "user_id": user_id} + if not text or (mtype == "group" and not at_me): + return # 空消息 / 群里没 @ 她 + from .main import chat_turn # 延迟导入避免循环 + reply = await chat_turn(session, text, _notify_for(conn, session)) + if reply: + await reply_message(conn, session, reply) + + +@router.websocket("/onebot/v11/ws") +async def onebot_ws(ws: WebSocket): + if ACCESS_TOKEN: + token = ws.query_params.get("access_token") or "" + auth = ws.headers.get("authorization", "") + if auth.startswith("Bearer "): + token = auth[7:] + if token != ACCESS_TOKEN: + await ws.close(code=4401) + return + await ws.accept() + conn = None + try: + while True: + data = json.loads(await ws.receive_text()) + if "post_type" in data: # 事件上报 + self_id = data.get("self_id") + if conn is None or conn.self_id != self_id: + conn = OneBotConn(ws, self_id) + connections[self_id] = conn + if data.get("post_type") == "message": + # 后台处理,读循环继续收(事件里可能夹着 action 的 echo 响应) + asyncio.create_task(_process_message(conn, data)) + elif conn is not None and "echo" in data: # action 的响应 + fut = conn.pending.pop(data.get("echo"), None) + if fut and not fut.done(): + fut.set_result(data) + # 其余(对端主动调 action 等)忽略 + except WebSocketDisconnect: + pass + finally: + if conn is not None: + connections.pop(conn.self_id, None) diff --git a/tests/conftest.py b/tests/conftest.py index 4600883..116025b 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -297,8 +297,11 @@ def offline(monkeypatch): @pytest.fixture def client(offline, fdb): """不带 lifespan 的测试应用(绕过 PG 连接)""" + from app import onebot + app = FastAPI() app.websocket("/ws")(main.chat) + app.include_router(onebot.router) return TestClient(app) diff --git a/tests/test_chat_flow.py b/tests/test_chat_flow.py index 2242b2e..359fab2 100644 --- a/tests/test_chat_flow.py +++ b/tests/test_chat_flow.py @@ -76,7 +76,7 @@ def test_streak_increments(fdb, offline): async def send_json(self, payload): pass - asyncio.run(main.remember_later(FakeWS(), "s1", "我来了", sig)) + asyncio.run(main.remember_later(FakeWS().send_json, "s1", "我来了", sig)) assert asyncio.run(fdb.get_meta("s1", "streak")) == "4" @@ -89,5 +89,5 @@ def test_streak_resets_after_gap(fdb, offline): async def send_json(self, payload): pass - asyncio.run(main.remember_later(FakeWS(), "s1", "好久不见", sig)) + asyncio.run(main.remember_later(FakeWS().send_json, "s1", "好久不见", sig)) assert asyncio.run(fdb.get_meta("s1", "streak")) == "1" diff --git a/tests/test_onebot.py b/tests/test_onebot.py new file mode 100644 index 0000000..9e06ef1 --- /dev/null +++ b/tests/test_onebot.py @@ -0,0 +1,140 @@ +"""OneBot 11 适配器测试(TestClient 模拟 OneBot 实现的反向 WS)""" +import json +import time + +import pytest +from starlette.websockets import WebSocketDisconnect + +from app import onebot + + +@pytest.fixture(autouse=True) +def clean_conn(): + onebot.connections.clear() + yield + onebot.connections.clear() + + +def _msg_event(user_id=123, message="你好呀", message_type="private", + group_id=None, self_id=10000): + ev = {"time": 0, "self_id": self_id, "post_type": "message", + "message_type": message_type, + "sub_type": "friend" if message_type == "private" else "normal", + "message_id": 1, "user_id": user_id, "message": message, + "raw_message": str(message), + "sender": {"user_id": user_id, "nickname": "tester"}} + if group_id: + ev["group_id"] = group_id + return ev + + +def _drain_action(ws): + """收一个 action 并替 OneBot 实现回一个成功 echo,返回该 action""" + action = json.loads(ws.receive_text()) + ws.send_text(json.dumps({"status": "ok", "retcode": 0, "data": {}, + "echo": action["echo"]})) + return action + + +def test_private_reply(client, route_sig, fdb): + route_sig({}) + with client.websocket_connect("/onebot/v11/ws") as ws: + ws.send_text(json.dumps(_msg_event(message="你好呀"))) + action = _drain_action(ws) + assert action["action"] == "send_private_msg" + assert action["params"]["user_id"] == 123 + assert "回声模式" in action["params"]["message"] + assert "你好呀" in action["params"]["message"] + # 会话映射注册成功 + assert "ob-10000-p123" in onebot.connections[10000].targets + time.sleep(0.3) + # 对话进了历史(user + assistant) + assert len(fdb.history) == 2 + + +def test_cq_code_stripped(client, route_sig, fdb): + route_sig({}) + with client.websocket_connect("/onebot/v11/ws") as ws: + ws.send_text(json.dumps( + _msg_event(message="[CQ:face,id=1]你叫什么名字"))) + action = _drain_action(ws) + assert "你叫什么名字" in action["params"]["message"] + + +def test_group_requires_at(client, route_sig, fdb): + route_sig({}) + with client.websocket_connect("/onebot/v11/ws") as ws: + # 群里没 @ 她 -> 不回应 + ws.send_text(json.dumps(_msg_event( + message_type="group", group_id=777, message="大家好啊"))) + # 紧跟一条私聊 -> 第一条收到的 action 应该是私聊回复 + ws.send_text(json.dumps(_msg_event(message="私聊一下"))) + action = _drain_action(ws) + assert action["action"] == "send_private_msg" + assert fdb.history[0]["session"] == "ob-10000-p123" + + +def test_group_at_replies_with_at(client, route_sig, fdb): + route_sig({}) + with client.websocket_connect("/onebot/v11/ws") as ws: + ws.send_text(json.dumps(_msg_event( + message_type="group", group_id=777, + message=[{"type": "at", "data": {"qq": 10000}}, + {"type": "text", "data": {"text": " 在吗"}}]))) + action = _drain_action(ws) + assert action["action"] == "send_group_msg" + assert action["params"]["group_id"] == 777 + assert action["params"]["message"].startswith("[CQ:at,qq=123]") + assert "回声模式" in action["params"]["message"] + # 群会话独立于私聊 + sessions = {h["session"] for h in fdb.history} + assert "ob-10000-g777" in sessions + + +def test_notify_event_becomes_qq_message(client, route_sig, fdb): + from datetime import datetime, timedelta + route_sig({"events": [{"kind": "alarm", + "fire_at": datetime.now() + timedelta(hours=14), + "text": "起床", "repeat": "daily"}]}) + with client.websocket_connect("/onebot/v11/ws") as ws: + ws.send_text(json.dumps(_msg_event(message="每天早上7点叫我起床"))) + a1 = _drain_action(ws) # 回复 + assert "回声模式" in a1["params"]["message"] + a2 = _drain_action(ws) # 事件确认 -> QQ 消息 + assert a2["action"] == "send_private_msg" + assert "我记下了" in a2["params"]["message"] + assert "每天都算" in a2["params"]["message"] + assert len(fdb.events) == 1 + + +def test_streak_pushed_proactively_via_onebot(client, route_sig, fdb): + route_sig({}) + from datetime import datetime, timedelta + import asyncio + yesterday = (datetime.now() - timedelta(days=1)).strftime("%Y-%m-%d") + # 用 FakeDB 直接种下"昨天来过、已连续 4 天" -> 这次该庆祝第 5 天 + asyncio.run(fdb.set_meta("ob-10000-p123", "last_active", yesterday)) + asyncio.run(fdb.set_meta("ob-10000-p123", "streak", "4")) + with client.websocket_connect("/onebot/v11/ws") as ws: + ws.send_text(json.dumps(_msg_event(message="我来了"))) + actions = [_drain_action(ws), _drain_action(ws)] + texts = [a["params"]["message"] for a in actions] + assert any("连续陪伴的第 5 天" in t for t in texts), texts + assert any("回声模式" in t for t in texts), texts + + +def test_access_token_enforced(client, route_sig, monkeypatch): + monkeypatch.setattr(onebot, "ACCESS_TOKEN", "secret") + with pytest.raises(WebSocketDisconnect): + with client.websocket_connect("/onebot/v11/ws") as ws: + ws.receive_text() + with client.websocket_connect("/onebot/v11/ws?access_token=secret") as ws: + ws.send_text(json.dumps(_msg_event(message="在吗"))) + action = _drain_action(ws) + assert "回声模式" in action["params"]["message"] + + +def test_disconnected_target_returns_false(client, route_sig): + import asyncio + route_sig({}) + assert asyncio.run(onebot.send_to_session("ob-x-p1", "hi")) is False