89 lines
3.3 KiB
Python
89 lines
3.3 KiB
Python
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import json
|
||
|
|
import time
|
||
|
|
from dataclasses import dataclass, field
|
||
|
|
from pathlib import Path
|
||
|
|
from typing import Any
|
||
|
|
|
||
|
|
from yuxi.utils.logging_config import logger
|
||
|
|
|
||
|
|
|
||
|
|
@dataclass
|
||
|
|
class ConversationBinding:
|
||
|
|
channel_chat_id: str
|
||
|
|
agent_id: str = "main"
|
||
|
|
label: str = ""
|
||
|
|
created_at: float = field(default_factory=time.time)
|
||
|
|
metadata: dict[str, Any] = field(default_factory=dict)
|
||
|
|
|
||
|
|
|
||
|
|
class ConversationBindingManager:
|
||
|
|
def __init__(self, storage_dir: str | None = None, account_id: str = "default"):
|
||
|
|
self._bindings: dict[str, ConversationBinding] = {}
|
||
|
|
self._storage_dir = storage_dir or str(Path.home() / ".yuxi" / "imessage")
|
||
|
|
self._account_id = account_id
|
||
|
|
|
||
|
|
def create(
|
||
|
|
self, channel_chat_id: str, agent_id: str = "main", label: str = "", metadata: dict[str, Any] | None = None
|
||
|
|
) -> ConversationBinding:
|
||
|
|
binding = ConversationBinding(
|
||
|
|
channel_chat_id=channel_chat_id,
|
||
|
|
agent_id=agent_id,
|
||
|
|
label=label,
|
||
|
|
metadata=metadata or {},
|
||
|
|
)
|
||
|
|
self._bindings[channel_chat_id] = binding
|
||
|
|
logger.info(f"[iMessage/Bindings] Created binding: {channel_chat_id} -> agent:{agent_id}")
|
||
|
|
return binding
|
||
|
|
|
||
|
|
def get(self, channel_chat_id: str) -> ConversationBinding | None:
|
||
|
|
return self._bindings.get(channel_chat_id)
|
||
|
|
|
||
|
|
def list_all(self) -> list[ConversationBinding]:
|
||
|
|
return list(self._bindings.values())
|
||
|
|
|
||
|
|
def delete(self, channel_chat_id: str) -> bool:
|
||
|
|
if channel_chat_id in self._bindings:
|
||
|
|
del self._bindings[channel_chat_id]
|
||
|
|
logger.info(f"[iMessage/Bindings] Deleted binding: {channel_chat_id}")
|
||
|
|
return True
|
||
|
|
return False
|
||
|
|
|
||
|
|
def get_agent_for_chat(self, channel_chat_id: str, default: str = "main") -> str:
|
||
|
|
binding = self._bindings.get(channel_chat_id)
|
||
|
|
return binding.agent_id if binding else default
|
||
|
|
|
||
|
|
async def save_bindings(self) -> None:
|
||
|
|
dir_path = Path(self._storage_dir) / self._account_id
|
||
|
|
dir_path.mkdir(parents=True, exist_ok=True)
|
||
|
|
data = {
|
||
|
|
cid: {
|
||
|
|
"agent_id": b.agent_id,
|
||
|
|
"label": b.label,
|
||
|
|
"created_at": b.created_at,
|
||
|
|
"metadata": b.metadata,
|
||
|
|
}
|
||
|
|
for cid, b in self._bindings.items()
|
||
|
|
}
|
||
|
|
file_path = dir_path / "conversation-bindings.json"
|
||
|
|
file_path.write_text(json.dumps(data, ensure_ascii=False, indent=2), encoding="utf-8")
|
||
|
|
|
||
|
|
async def load_bindings(self) -> None:
|
||
|
|
file_path = Path(self._storage_dir) / self._account_id / "conversation-bindings.json"
|
||
|
|
if not file_path.exists():
|
||
|
|
return
|
||
|
|
try:
|
||
|
|
data = json.loads(file_path.read_text(encoding="utf-8"))
|
||
|
|
for cid, info in data.items():
|
||
|
|
if isinstance(info, dict):
|
||
|
|
self._bindings[cid] = ConversationBinding(
|
||
|
|
channel_chat_id=cid,
|
||
|
|
agent_id=info.get("agent_id", "main"),
|
||
|
|
label=info.get("label", ""),
|
||
|
|
created_at=info.get("created_at", time.time()),
|
||
|
|
metadata=info.get("metadata", {}),
|
||
|
|
)
|
||
|
|
except (json.JSONDecodeError, OSError) as e:
|
||
|
|
logger.warning(f"[iMessage/Bindings] Failed to load bindings: {e}")
|