322 lines
12 KiB
Python
322 lines
12 KiB
Python
"""绑定路由器"""
|
||
|
||
import hashlib
|
||
import json
|
||
import re
|
||
from dataclasses import dataclass
|
||
from typing import TYPE_CHECKING, Any
|
||
|
||
from sqlalchemy.ext.asyncio import AsyncSession
|
||
|
||
from yuxi.channel.plugins.protocol import (
|
||
BindingConversationRef,
|
||
BindingRoute,
|
||
ChannelPlugin,
|
||
InboundMessage,
|
||
SessionConversationRef,
|
||
)
|
||
from yuxi.channel.routing.cache import RouteCache, get_default_route_cache
|
||
from yuxi.repositories.channel_binding_repository import ChannelBindingRepository
|
||
from yuxi.repositories.conversation_repository import ConversationRepository
|
||
|
||
if TYPE_CHECKING:
|
||
from yuxi.storage.postgres.model_channel import ChannelSession
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class _BindingContext:
|
||
"""路由匹配所需的会话上下文，兼容 ChannelSession 与 SessionConversationRef。"""
|
||
|
||
session_key: str
|
||
channel_type: str
|
||
account_id: str
|
||
chat_type: str | None
|
||
peer_id: str | None
|
||
conversation_id: int | None = None
|
||
|
||
|
||
class _MatchTier:
|
||
def __init__(self, name: str, matcher: Any) -> None:
|
||
self.name = name
|
||
self._matcher = matcher
|
||
|
||
def matches(self, rule: dict, ctx: _BindingContext, inbound: InboundMessage) -> bool:
|
||
return self._matcher(rule, ctx, inbound)
|
||
|
||
|
||
def _peer_id(ctx: _BindingContext, inbound: InboundMessage) -> str | None:
|
||
return inbound.peer_id if inbound.peer_id is not None else ctx.peer_id
|
||
|
||
|
||
def _chat_type(ctx: _BindingContext, inbound: InboundMessage) -> str | None:
|
||
return inbound.chat_type if inbound.chat_type is not None else ctx.chat_type
|
||
|
||
|
||
class BindingRouter:
|
||
"""解析渠道会话应路由到的 agent_id。
|
||
|
||
支持配置绑定（八级匹配）、运行时绑定、默认路由与 LRU 缓存。
|
||
"""
|
||
|
||
def __init__(self, route_cache: RouteCache | None = None) -> None:
|
||
self._cache = route_cache or get_default_route_cache()
|
||
self._match_tiers: tuple[_MatchTier, ...] = (
|
||
_MatchTier(
|
||
"binding.session_key",
|
||
lambda rule, ctx, _inbound: rule.get("session_key") == ctx.session_key,
|
||
),
|
||
_MatchTier(
|
||
"binding.channel_account_chat_peer",
|
||
lambda rule, ctx, inbound: (
|
||
rule.get("channel_type") == ctx.channel_type
|
||
and rule.get("account_id") == ctx.account_id
|
||
and rule.get("chat_type") == _chat_type(ctx, inbound)
|
||
and rule.get("peer_id") == _peer_id(ctx, inbound)
|
||
),
|
||
),
|
||
_MatchTier(
|
||
"binding.channel_account_peer",
|
||
lambda rule, ctx, inbound: (
|
||
rule.get("channel_type") == ctx.channel_type
|
||
and rule.get("account_id") == ctx.account_id
|
||
and rule.get("peer_id") == _peer_id(ctx, inbound)
|
||
),
|
||
),
|
||
_MatchTier(
|
||
"binding.channel_account_chat",
|
||
lambda rule, ctx, inbound: (
|
||
rule.get("channel_type") == ctx.channel_type
|
||
and rule.get("account_id") == ctx.account_id
|
||
and rule.get("chat_type") == _chat_type(ctx, inbound)
|
||
),
|
||
),
|
||
_MatchTier(
|
||
"binding.channel_chat_peer",
|
||
lambda rule, ctx, inbound: (
|
||
rule.get("channel_type") == ctx.channel_type
|
||
and rule.get("chat_type") == _chat_type(ctx, inbound)
|
||
and rule.get("peer_id") == _peer_id(ctx, inbound)
|
||
),
|
||
),
|
||
_MatchTier(
|
||
"binding.session_key_regex",
|
||
lambda rule, ctx, _inbound: (
|
||
rule.get("session_key_regex") is not None
|
||
and re.fullmatch(str(rule["session_key_regex"]), ctx.session_key) is not None
|
||
),
|
||
),
|
||
_MatchTier(
|
||
"binding.channel_account",
|
||
lambda rule, ctx, _inbound: (
|
||
rule.get("channel_type") == ctx.channel_type and rule.get("account_id") == ctx.account_id
|
||
),
|
||
),
|
||
_MatchTier(
|
||
"binding.channel_type",
|
||
lambda rule, ctx, _inbound: rule.get("channel_type") == ctx.channel_type,
|
||
),
|
||
)
|
||
|
||
def _build_cache_key(self, ctx: _BindingContext) -> str:
|
||
return f"{ctx.channel_type}:{ctx.account_id}:{ctx.session_key}"
|
||
|
||
@staticmethod
|
||
def _compute_rule_hash(rule: dict) -> str:
|
||
"""对规则字典计算稳定 sha256 hash。"""
|
||
payload = json.dumps(rule, sort_keys=True, ensure_ascii=False, default=str)
|
||
return hashlib.sha256(payload.encode("utf-8")).hexdigest()
|
||
|
||
@staticmethod
|
||
def _parse_session_key(session_key: str) -> tuple[str | None, str | None]:
|
||
"""安全解析 session_key，返回 (channel_type, account_id)。"""
|
||
if not session_key:
|
||
return None, None
|
||
parts = session_key.split(":", 2)
|
||
if len(parts) < 2:
|
||
return parts[0] if parts else None, None
|
||
return parts[0], parts[1]
|
||
|
||
def _make_context(
|
||
self,
|
||
config: dict,
|
||
plugin: ChannelPlugin,
|
||
ref: SessionConversationRef,
|
||
conversation_id: int | None = None,
|
||
) -> _BindingContext:
|
||
"""从 config 或 session_key 构建路由上下文。"""
|
||
channel_type = config.get("channel_type")
|
||
account_id = config.get("account_id")
|
||
if not channel_type or not account_id:
|
||
parsed_type, parsed_account = self._parse_session_key(ref.session_key)
|
||
channel_type = channel_type or parsed_type
|
||
account_id = account_id or parsed_account
|
||
if not channel_type or not account_id:
|
||
raise ValueError(f"Unable to resolve channel_type/account_id for session {ref.session_key}")
|
||
return _BindingContext(
|
||
session_key=ref.session_key,
|
||
channel_type=channel_type,
|
||
account_id=account_id,
|
||
chat_type=ref.chat_type,
|
||
peer_id=ref.channel_sender_id,
|
||
conversation_id=conversation_id,
|
||
)
|
||
|
||
async def resolve_runtime(
|
||
self,
|
||
config: dict,
|
||
plugin: ChannelPlugin,
|
||
ref: SessionConversationRef,
|
||
db_session: AsyncSession,
|
||
) -> BindingRoute:
|
||
"""在创建 Conversation 前解析 agent_id，优先配置绑定，再运行时绑定，最后默认路由。"""
|
||
ctx = self._make_context(config, plugin, ref)
|
||
cache_key = self._build_cache_key(ctx)
|
||
cached = await self._cache.get(cache_key)
|
||
if cached is not None:
|
||
return cached
|
||
|
||
inbound = InboundMessage(
|
||
channel_type=ctx.channel_type,
|
||
account_id=ctx.account_id,
|
||
chat_type=ctx.chat_type,
|
||
peer_id=ctx.peer_id,
|
||
session_key=ctx.session_key,
|
||
)
|
||
|
||
route = await self._resolve_configured_binding(ctx, config, plugin, inbound)
|
||
if route is not None:
|
||
if route.agent_id:
|
||
await self._cache.set(cache_key, route)
|
||
return route
|
||
|
||
runtime_route = await self._resolve_runtime_binding(ctx, db_session)
|
||
if runtime_route is not None:
|
||
if runtime_route.agent_id:
|
||
await self._cache.set(cache_key, runtime_route)
|
||
return runtime_route
|
||
|
||
default_route = BindingRoute(
|
||
agent_id=config.get("default_agent_id"),
|
||
session_key=ctx.session_key,
|
||
matched_by="default",
|
||
)
|
||
if default_route.agent_id:
|
||
await self._cache.set(cache_key, default_route)
|
||
return default_route
|
||
|
||
async def resolve(
|
||
self,
|
||
session: "ChannelSession",
|
||
config: dict,
|
||
plugin: ChannelPlugin,
|
||
inbound: InboundMessage,
|
||
db_session: AsyncSession,
|
||
) -> BindingRoute:
|
||
"""完整解析：先查缓存，再配置绑定（八级匹配），再运行时绑定，最后默认路由。"""
|
||
ctx = _BindingContext(
|
||
session_key=session.session_key,
|
||
channel_type=session.channel_type,
|
||
account_id=session.account_id,
|
||
chat_type=session.chat_type,
|
||
peer_id=inbound.peer_id if inbound.peer_id is not None else session.channel_sender_id,
|
||
conversation_id=session.conversation_id,
|
||
)
|
||
cache_key = self._build_cache_key(ctx)
|
||
cached = await self._cache.get(cache_key)
|
||
if cached is not None:
|
||
return cached
|
||
|
||
config_route = await self._resolve_configured_binding(ctx, config, plugin, inbound)
|
||
if config_route is not None:
|
||
if config_route.agent_id:
|
||
await self._cache.set(cache_key, config_route)
|
||
return config_route
|
||
|
||
runtime_route = await self._resolve_runtime_binding(ctx, db_session)
|
||
if runtime_route is not None:
|
||
if runtime_route.agent_id:
|
||
await self._cache.set(cache_key, runtime_route)
|
||
return runtime_route
|
||
|
||
conversation_repo = ConversationRepository(db_session)
|
||
default_agent_id = config.get("default_agent_id")
|
||
if not default_agent_id and ctx.conversation_id:
|
||
conversation = await conversation_repo.get_conversation_by_id(ctx.conversation_id)
|
||
if conversation is not None:
|
||
default_agent_id = conversation.agent_id
|
||
|
||
default_route = BindingRoute(
|
||
agent_id=default_agent_id,
|
||
session_key=ctx.session_key,
|
||
matched_by="default",
|
||
)
|
||
if default_route.agent_id:
|
||
await self._cache.set(cache_key, default_route)
|
||
return default_route
|
||
|
||
async def _resolve_configured_binding(
|
||
self,
|
||
ctx: _BindingContext,
|
||
config: dict,
|
||
plugin: ChannelPlugin,
|
||
inbound: InboundMessage,
|
||
) -> BindingRoute | None:
|
||
"""按八级匹配解析配置绑定；插件可通过 compile_binding/match_binding 扩展匹配逻辑。"""
|
||
bindings = config.get("bindings", [])
|
||
|
||
# 1. 插件自定义绑定编译与匹配
|
||
if isinstance(plugin.compile_binding(config, inbound), BindingConversationRef):
|
||
for rule in bindings:
|
||
if plugin.match_binding(config, rule, inbound):
|
||
return BindingRoute(
|
||
agent_id=rule.get("agent_id"),
|
||
session_key=ctx.session_key,
|
||
matched_by="binding.plugin",
|
||
binding_rule_hash=self._compute_rule_hash(rule),
|
||
)
|
||
|
||
if not bindings:
|
||
return None
|
||
|
||
# 2. 通用八级匹配
|
||
for tier in self._match_tiers:
|
||
for rule in bindings:
|
||
if tier.matches(rule, ctx, inbound):
|
||
return BindingRoute(
|
||
agent_id=rule.get("agent_id"),
|
||
session_key=ctx.session_key,
|
||
matched_by=tier.name,
|
||
binding_rule_hash=self._compute_rule_hash(rule),
|
||
)
|
||
return None
|
||
|
||
async def _resolve_runtime_binding(
|
||
self,
|
||
ctx: _BindingContext,
|
||
db_session: AsyncSession,
|
||
) -> BindingRoute | None:
|
||
"""查询 channel_bindings 运行时绑定记录。"""
|
||
binding_repo = ChannelBindingRepository(db_session)
|
||
bindings = await binding_repo.find_runtime_bindings(
|
||
ctx.channel_type,
|
||
ctx.account_id,
|
||
ctx.session_key,
|
||
)
|
||
if not bindings:
|
||
return None
|
||
binding = bindings[0]
|
||
return BindingRoute(
|
||
agent_id=binding.agent_id,
|
||
session_key=ctx.session_key,
|
||
matched_by="runtime.binding",
|
||
binding_rule_hash=binding.binding_rule_hash,
|
||
)
|
||
|
||
async def invalidate_channel_cache(self, channel_type: str) -> int:
|
||
"""当配置变更时，按渠道类型失效路由缓存。"""
|
||
return await self._cache.invalidate_by_channel(channel_type)
|
||
|
||
async def invalidate_account_cache(self, channel_type: str, account_id: str) -> int:
|
||
"""当配置变更时，按渠道账户失效路由缓存。"""
|
||
return await self._cache.invalidate_by_account(channel_type, account_id)
|