|
|
@@ -1,335 +1,10 @@
|
|
|
-import asyncio
|
|
|
-import json
|
|
|
-import logging
|
|
|
-import os
|
|
|
-import tempfile
|
|
|
-import uuid
|
|
|
-from datetime import datetime
|
|
|
-from pathlib import Path
|
|
|
+"""Compatibility module for the packaged :mod:`agent.im_client.client`."""
|
|
|
|
|
|
-import websockets
|
|
|
-from filelock import FileLock
|
|
|
+from _framework_import import load
|
|
|
|
|
|
-from protocol import IMMessage, IMResponse
|
|
|
-from notifier import AgentNotifier, ConsoleNotifier
|
|
|
+_module = load("agent.im_client.client")
|
|
|
|
|
|
-logging.basicConfig(level=logging.INFO, format="%(asctime)s [CLIENT:%(name)s] %(message)s")
|
|
|
-
|
|
|
-
|
|
|
-class ChatWindow:
|
|
|
- """单个聊天窗口的数据管理。"""
|
|
|
-
|
|
|
- def __init__(self, chat_id: str, data_dir: Path):
|
|
|
- self.chat_id = chat_id
|
|
|
- self.data_dir = data_dir
|
|
|
- self.data_dir.mkdir(parents=True, exist_ok=True)
|
|
|
-
|
|
|
- self.chatbox_path = data_dir / "chatbox.jsonl"
|
|
|
- self.in_pending_path = data_dir / "in_pending.json"
|
|
|
- self.out_pending_path = data_dir / "out_pending.jsonl"
|
|
|
-
|
|
|
- # 文件锁
|
|
|
- self._in_pending_lock = FileLock(str(data_dir / ".in_pending.lock"))
|
|
|
- self._out_pending_lock = FileLock(str(data_dir / ".out_pending.lock"))
|
|
|
- self._chatbox_lock = FileLock(str(data_dir / ".chatbox.lock"))
|
|
|
-
|
|
|
- # 初始化文件
|
|
|
- if not self.chatbox_path.exists():
|
|
|
- self.chatbox_path.write_text("")
|
|
|
- if not self.in_pending_path.exists():
|
|
|
- self.in_pending_path.write_text("[]")
|
|
|
- if not self.out_pending_path.exists():
|
|
|
- self.out_pending_path.write_text("")
|
|
|
-
|
|
|
- def append_to_in_pending(self, msg: dict):
|
|
|
- with self._in_pending_lock:
|
|
|
- pending = self._load_json_array(self.in_pending_path)
|
|
|
- pending.append(msg)
|
|
|
- self._atomic_write_json(self.in_pending_path, pending)
|
|
|
-
|
|
|
- def read_in_pending(self) -> list[dict]:
|
|
|
- with self._in_pending_lock:
|
|
|
- return self._load_json_array(self.in_pending_path)
|
|
|
-
|
|
|
- def clear_in_pending(self):
|
|
|
- with self._in_pending_lock:
|
|
|
- self._atomic_write_json(self.in_pending_path, [])
|
|
|
-
|
|
|
- def append_to_chatbox(self, msg: dict):
|
|
|
- with self._chatbox_lock:
|
|
|
- with open(self.chatbox_path, "a", encoding="utf-8") as f:
|
|
|
- f.write(json.dumps(msg, ensure_ascii=False) + "\n")
|
|
|
-
|
|
|
- def append_to_out_pending(self, msg: dict):
|
|
|
- with self._out_pending_lock:
|
|
|
- with open(self.out_pending_path, "a", encoding="utf-8") as f:
|
|
|
- f.write(json.dumps(msg, ensure_ascii=False) + "\n")
|
|
|
-
|
|
|
- @staticmethod
|
|
|
- def _load_json_array(path: Path) -> list:
|
|
|
- if not path.exists():
|
|
|
- return []
|
|
|
- text = path.read_text(encoding="utf-8").strip()
|
|
|
- if not text:
|
|
|
- return []
|
|
|
- try:
|
|
|
- data = json.loads(text)
|
|
|
- return data if isinstance(data, list) else []
|
|
|
- except json.JSONDecodeError:
|
|
|
- return []
|
|
|
-
|
|
|
- @staticmethod
|
|
|
- def _atomic_write_json(path: Path, data):
|
|
|
- tmp_fd, tmp_path = tempfile.mkstemp(dir=str(path.parent), suffix=".tmp")
|
|
|
- try:
|
|
|
- with os.fdopen(tmp_fd, "w", encoding="utf-8") as f:
|
|
|
- json.dump(data, f, ensure_ascii=False, indent=2)
|
|
|
- os.replace(tmp_path, str(path))
|
|
|
- except Exception:
|
|
|
- if os.path.exists(tmp_path):
|
|
|
- os.unlink(tmp_path)
|
|
|
- raise
|
|
|
-
|
|
|
-
|
|
|
-class IMClient:
|
|
|
- """IM Client - 一个实例管理多个聊天窗口。
|
|
|
-
|
|
|
- 一个 Agent (contact_id) 对应一个 IMClient 实例。
|
|
|
- 该实例可以管理多个 chat_id(窗口),每个窗口有独立的消息存储。
|
|
|
- """
|
|
|
-
|
|
|
- def __init__(
|
|
|
- self,
|
|
|
- contact_id: str,
|
|
|
- server_url: str = "ws://localhost:8000",
|
|
|
- data_dir: str | None = None,
|
|
|
- notify_interval: float = 30.0,
|
|
|
- ):
|
|
|
- self.contact_id = contact_id
|
|
|
- self.server_url = server_url
|
|
|
- self.notify_interval = notify_interval
|
|
|
-
|
|
|
- self.base_dir = Path(data_dir) if data_dir else Path("data") / contact_id
|
|
|
- self.base_dir.mkdir(parents=True, exist_ok=True)
|
|
|
-
|
|
|
- # 窗口管理
|
|
|
- self._windows: dict[str, ChatWindow] = {}
|
|
|
- self._notifiers: dict[str, AgentNotifier] = {}
|
|
|
-
|
|
|
- self.ws = None
|
|
|
- self.log = logging.getLogger(contact_id)
|
|
|
- self._send_queue = asyncio.Queue()
|
|
|
-
|
|
|
- # 消息回调:{chat_id: callback} 或 {"*": callback} 匹配所有窗口
|
|
|
- self._on_message_callbacks: dict[str, callable] = {}
|
|
|
-
|
|
|
- def open_window(self, chat_id: str | None = None, notifier: AgentNotifier | None = None) -> str:
|
|
|
- """打开一个新窗口。
|
|
|
-
|
|
|
- Args:
|
|
|
- chat_id: 窗口 ID(留空自动生成)
|
|
|
- notifier: 该窗口的通知器
|
|
|
-
|
|
|
- Returns:
|
|
|
- 窗口的 chat_id
|
|
|
- """
|
|
|
- if chat_id is None:
|
|
|
- chat_id = datetime.now().strftime("%Y%m%d_%H%M%S_") + uuid.uuid4().hex[:6]
|
|
|
-
|
|
|
- if chat_id in self._windows:
|
|
|
- return chat_id
|
|
|
-
|
|
|
- window_dir = self.base_dir / "windows" / chat_id
|
|
|
- self._windows[chat_id] = ChatWindow(chat_id, window_dir)
|
|
|
- self._notifiers[chat_id] = notifier or ConsoleNotifier()
|
|
|
-
|
|
|
- self.log.info(f"打开窗口: {chat_id}")
|
|
|
- return chat_id
|
|
|
-
|
|
|
- def close_window(self, chat_id: str):
|
|
|
- """关闭一个窗口。"""
|
|
|
- self._windows.pop(chat_id, None)
|
|
|
- self._notifiers.pop(chat_id, None)
|
|
|
-
|
|
|
- def on_message(self, callback: callable, chat_id: str = "*"):
|
|
|
- """注册消息回调。收到消息时立即触发,无需轮询。
|
|
|
-
|
|
|
- Args:
|
|
|
- callback: 回调函数,签名 async def callback(msg: dict)
|
|
|
- chat_id: 监听的窗口 ID,"*" 表示所有窗口
|
|
|
- """
|
|
|
- self._on_message_callbacks[chat_id] = callback
|
|
|
- self.log.info(f"关闭窗口: {chat_id}")
|
|
|
-
|
|
|
- def list_windows(self) -> list[str]:
|
|
|
- """列出所有打开的窗口。"""
|
|
|
- return list(self._windows.keys())
|
|
|
-
|
|
|
- async def run(self):
|
|
|
- """启动 Client 服务,自动重连。"""
|
|
|
- while True:
|
|
|
- try:
|
|
|
- # 连接时不带 chat_id,因为一个实例管理多个窗口
|
|
|
- ws_url = f"{self.server_url}/ws?contact_id={self.contact_id}&chat_id=__multi__"
|
|
|
- self.log.info(f"连接 {ws_url} ...")
|
|
|
- async with websockets.connect(ws_url) as ws:
|
|
|
- self.ws = ws
|
|
|
- self.log.info("已连接")
|
|
|
- await asyncio.gather(
|
|
|
- self._ws_listener(),
|
|
|
- self._send_worker(),
|
|
|
- self._pending_notifier(),
|
|
|
- )
|
|
|
- except (websockets.ConnectionClosed, ConnectionRefusedError, OSError) as e:
|
|
|
- self.log.warning(f"连接断开: {e}, 5 秒后重连...")
|
|
|
- self.ws = None
|
|
|
- await asyncio.sleep(5)
|
|
|
- except asyncio.CancelledError:
|
|
|
- self.log.info("服务停止")
|
|
|
- break
|
|
|
-
|
|
|
- async def _ws_listener(self):
|
|
|
- """监听 WebSocket,根据 receiver_chat_id 分发到对应窗口。"""
|
|
|
- async for raw in self.ws:
|
|
|
- try:
|
|
|
- data = json.loads(raw)
|
|
|
- except json.JSONDecodeError:
|
|
|
- self.log.warning(f"收到无效 JSON: {raw}")
|
|
|
- continue
|
|
|
-
|
|
|
- if "sender" in data and "receiver" in data:
|
|
|
- # 聊天消息
|
|
|
- receiver_chat_id = data.get("receiver_chat_id")
|
|
|
-
|
|
|
- if receiver_chat_id and receiver_chat_id in self._windows:
|
|
|
- # 定向发送到指定窗口
|
|
|
- window = self._windows[receiver_chat_id]
|
|
|
- window.append_to_in_pending(data)
|
|
|
- window.append_to_chatbox(data)
|
|
|
- self.log.info(f"收到消息 -> 窗口 {receiver_chat_id}: {data['sender']}")
|
|
|
- await self._fire_on_message(receiver_chat_id, data)
|
|
|
- elif not receiver_chat_id:
|
|
|
- # 广播到所有窗口
|
|
|
- for chat_id, window in self._windows.items():
|
|
|
- window.append_to_in_pending(data)
|
|
|
- window.append_to_chatbox(data)
|
|
|
- await self._fire_on_message(chat_id, data)
|
|
|
- self.log.info(f"收到消息 -> 广播到 {len(self._windows)} 个窗口: {data['sender']}")
|
|
|
- else:
|
|
|
- self.log.warning(f"收到消息但窗口 {receiver_chat_id} 不存在")
|
|
|
-
|
|
|
- elif "status" in data:
|
|
|
- # 发送回执
|
|
|
- resp = IMResponse(**data)
|
|
|
- if resp.status == "success":
|
|
|
- self.log.info(f"消息 {resp.msg_id} 发送成功")
|
|
|
- else:
|
|
|
- self.log.warning(f"消息 {resp.msg_id} 发送失败: {resp.error}")
|
|
|
-
|
|
|
- async def _fire_on_message(self, chat_id: str, data: dict):
|
|
|
- """触发消息回调。"""
|
|
|
- # 精确匹配
|
|
|
- cb = self._on_message_callbacks.get(chat_id)
|
|
|
- if cb is None:
|
|
|
- # 通配符匹配
|
|
|
- cb = self._on_message_callbacks.get("*")
|
|
|
- if cb:
|
|
|
- try:
|
|
|
- await cb(data)
|
|
|
- except Exception as e:
|
|
|
- self.log.error(f"on_message 回调异常: {e}")
|
|
|
-
|
|
|
- async def _send_worker(self):
|
|
|
- """从队列取消息并发送。"""
|
|
|
- while True:
|
|
|
- msg_data = await self._send_queue.get()
|
|
|
- msg = IMMessage(sender=self.contact_id, **msg_data)
|
|
|
- try:
|
|
|
- await self.ws.send(msg.model_dump_json())
|
|
|
- self.log.info(f"发送消息: -> {msg.receiver}:{msg.receiver_chat_id or '*'}")
|
|
|
- # 记录到发送方窗口的 chatbox
|
|
|
- if msg.sender_chat_id and msg.sender_chat_id in self._windows:
|
|
|
- self._windows[msg.sender_chat_id].append_to_chatbox(msg.model_dump())
|
|
|
- except Exception as e:
|
|
|
- self.log.error(f"发送失败: {e}")
|
|
|
- if msg.sender_chat_id and msg.sender_chat_id in self._windows:
|
|
|
- self._windows[msg.sender_chat_id].append_to_out_pending(msg.model_dump())
|
|
|
-
|
|
|
- async def _pending_notifier(self):
|
|
|
- """轮询各窗口的 in_pending,有新消息就调通知回调。"""
|
|
|
- while True:
|
|
|
- for chat_id, window in list(self._windows.items()):
|
|
|
- pending = window.read_in_pending()
|
|
|
- if pending:
|
|
|
- senders = list(set(m.get("sender", "unknown") for m in pending))
|
|
|
- count = len(pending)
|
|
|
- notifier = self._notifiers.get(chat_id)
|
|
|
- if notifier:
|
|
|
- try:
|
|
|
- await notifier.notify(count=count, from_contacts=senders)
|
|
|
- except Exception as e:
|
|
|
- self.log.error(f"窗口 {chat_id} 通知回调异常: {e}")
|
|
|
- await asyncio.sleep(self.notify_interval)
|
|
|
-
|
|
|
- # ── Agent 调用的工具方法 ──
|
|
|
-
|
|
|
- def read_pending(self, chat_id: str) -> list[dict]:
|
|
|
- """读取某个窗口的待处理消息,并清空。"""
|
|
|
- window = self._windows.get(chat_id)
|
|
|
- if window is None:
|
|
|
- return []
|
|
|
- pending = window.read_in_pending()
|
|
|
- if pending:
|
|
|
- window.clear_in_pending()
|
|
|
- return pending
|
|
|
-
|
|
|
- def send_message(
|
|
|
- self,
|
|
|
- chat_id: str,
|
|
|
- receiver: str,
|
|
|
- content: str,
|
|
|
- msg_type: str = "chat",
|
|
|
- receiver_chat_id: str | None = None,
|
|
|
- ):
|
|
|
- """从某个窗口发送消息。"""
|
|
|
- msg_data = {
|
|
|
- "sender_chat_id": chat_id,
|
|
|
- "receiver": receiver,
|
|
|
- "content": content,
|
|
|
- "msg_type": msg_type,
|
|
|
- "receiver_chat_id": receiver_chat_id,
|
|
|
- }
|
|
|
- self._send_queue.put_nowait(msg_data)
|
|
|
-
|
|
|
- def get_chat_history(self, chat_id: str, peer_id: str | None = None, limit: int = 20) -> list[dict]:
|
|
|
- """查询某个窗口的聊天历史。"""
|
|
|
- window = self._windows.get(chat_id)
|
|
|
- if window is None or not window.chatbox_path.exists():
|
|
|
- return []
|
|
|
-
|
|
|
- lines = window.chatbox_path.read_text(encoding="utf-8").strip().splitlines()
|
|
|
- messages = []
|
|
|
- for line in reversed(lines):
|
|
|
- if not line.strip():
|
|
|
- continue
|
|
|
- try:
|
|
|
- m = json.loads(line)
|
|
|
- except json.JSONDecodeError:
|
|
|
- continue
|
|
|
-
|
|
|
- if peer_id and m.get("sender") != peer_id and m.get("receiver") != peer_id:
|
|
|
- continue
|
|
|
-
|
|
|
- messages.append({
|
|
|
- "sender": m.get("sender", "unknown"),
|
|
|
- "receiver": m.get("receiver", "unknown"),
|
|
|
- "content": m.get("content", ""),
|
|
|
- "msg_type": m.get("msg_type", "chat"),
|
|
|
- })
|
|
|
-
|
|
|
- if len(messages) >= limit:
|
|
|
- break
|
|
|
-
|
|
|
- messages.reverse()
|
|
|
- return messages
|
|
|
+ChatWindow = _module.ChatWindow
|
|
|
+IMClient = _module.IMClient
|
|
|
|
|
|
+__all__ = ["ChatWindow", "IMClient"]
|