48 lines
1.7 KiB
Python
48 lines
1.7 KiB
Python
|
from __future__ import annotations
|
||
|
|
|
||
|
|
from typing import Any
|
||
|
|
|
||
|
|
|
||
|
|
class TopicRouter:
|
||
|
|
def __init__(self, config: dict[str, Any] | None = None):
|
||
|
|
cfg = config or {}
|
||
|
|
self._groups_config = cfg.get("groups", {})
|
||
|
|
self._default_agent_id = cfg.get("default_agent_id", "default")
|
||
|
|
|
||
|
|
def resolve_topic_agent(self, chat_id: str, thread_id: str) -> str:
|
||
|
|
chat_config = self._groups_config.get(chat_id, {})
|
||
|
|
topics_config = chat_config.get("topics", {})
|
||
|
|
|
||
|
|
topic_config = topics_config.get(thread_id)
|
||
|
|
if topic_config and "agent_id" in topic_config:
|
||
|
|
return topic_config["agent_id"]
|
||
|
|
|
||
|
|
group_agent = chat_config.get("agent_id")
|
||
|
|
return group_agent or self._default_agent_id
|
||
|
|
|
||
|
|
def resolve_group_agent(self, chat_id: str) -> str:
|
||
|
|
chat_config = self._groups_config.get(chat_id, {})
|
||
|
|
group_agent = chat_config.get("agent_id")
|
||
|
|
return group_agent or self._default_agent_id
|
||
|
|
|
||
|
|
def resolve_route(
|
||
|
|
self,
|
||
|
|
chat_id: str,
|
||
|
|
chat_type: str,
|
||
|
|
thread_id: str | None = None,
|
||
|
|
) -> str:
|
||
|
|
agent_id = self._resolve_agent(chat_id, chat_type, thread_id)
|
||
|
|
|
||
|
|
if chat_type == "private":
|
||
|
|
return f"agent:{agent_id}:telegram:direct:{chat_id}"
|
||
|
|
if thread_id:
|
||
|
|
return f"agent:{agent_id}:telegram:group:{chat_id}:topic:{thread_id}"
|
||
|
|
return f"agent:{agent_id}:telegram:group:{chat_id}"
|
||
|
|
|
||
|
|
def _resolve_agent(self, chat_id: str, chat_type: str, thread_id: str | None) -> str:
|
||
|
|
if chat_type == "private":
|
||
|
|
return self._default_agent_id
|
||
|
|
if thread_id:
|
||
|
|
return self.resolve_topic_agent(chat_id, thread_id)
|
||
|
|
return self.resolve_group_agent(chat_id)
|