124 lines
3.6 KiB
Python
124 lines
3.6 KiB
Python
|
from __future__ import annotations
|
||
|
|
|
||
|
|
from collections.abc import Callable
|
||
|
|
from dataclasses import dataclass, field
|
||
|
|
from datetime import datetime, timedelta
|
||
|
|
from enum import StrEnum
|
||
|
|
from typing import Any
|
||
|
|
|
||
|
|
from yuxi.utils.datetime_utils import utc_now_naive
|
||
|
|
|
||
|
|
|
||
|
|
class BindingType(StrEnum):
|
||
|
|
AGENT = "agent"
|
||
|
|
SUBAGENT = "subagent"
|
||
|
|
ACP = "acp"
|
||
|
|
CONVERSATION = "conversation"
|
||
|
|
|
||
|
|
|
||
|
|
@dataclass
|
||
|
|
class ThreadBinding:
|
||
|
|
thread_id: str
|
||
|
|
binding_type: BindingType
|
||
|
|
target_id: str
|
||
|
|
created_at: datetime = field(default_factory=utc_now_naive)
|
||
|
|
expires_at: datetime | None = None
|
||
|
|
metadata: dict[str, Any] = field(default_factory=dict)
|
||
|
|
|
||
|
|
@property
|
||
|
|
def is_expired(self) -> bool:
|
||
|
|
if self.expires_at is None:
|
||
|
|
return False
|
||
|
|
return utc_now_naive() > self.expires_at
|
||
|
|
|
||
|
|
|
||
|
|
class ThreadBindingManager:
|
||
|
|
def __init__(
|
||
|
|
self,
|
||
|
|
default_ttl_hours: int = 24,
|
||
|
|
):
|
||
|
|
self._bindings: dict[str, ThreadBinding] = {}
|
||
|
|
self._default_ttl = timedelta(hours=default_ttl_hours) if default_ttl_hours else None
|
||
|
|
self._listeners: list[Callable] = []
|
||
|
|
|
||
|
|
def bind(
|
||
|
|
self,
|
||
|
|
thread_id: str,
|
||
|
|
binding_type: BindingType,
|
||
|
|
target_id: str,
|
||
|
|
ttl_hours: int | None = None,
|
||
|
|
metadata: dict[str, Any] | None = None,
|
||
|
|
) -> ThreadBinding:
|
||
|
|
expires = None
|
||
|
|
if ttl_hours is not None:
|
||
|
|
expires = utc_now_naive() + timedelta(hours=ttl_hours)
|
||
|
|
elif self._default_ttl:
|
||
|
|
expires = utc_now_naive() + self._default_ttl
|
||
|
|
|
||
|
|
binding = ThreadBinding(
|
||
|
|
thread_id=thread_id,
|
||
|
|
binding_type=binding_type,
|
||
|
|
target_id=target_id,
|
||
|
|
expires_at=expires,
|
||
|
|
metadata=metadata or {},
|
||
|
|
)
|
||
|
|
|
||
|
|
self._bindings[f"{thread_id}:{binding_type.value}"] = binding
|
||
|
|
self._notify("bind", binding)
|
||
|
|
return binding
|
||
|
|
|
||
|
|
def unbind(self, thread_id: str, binding_type: BindingType) -> bool:
|
||
|
|
key = f"{thread_id}:{binding_type.value}"
|
||
|
|
if key in self._bindings:
|
||
|
|
binding = self._bindings.pop(key)
|
||
|
|
self._notify("unbind", binding)
|
||
|
|
return True
|
||
|
|
return False
|
||
|
|
|
||
|
|
def get_binding(self, thread_id: str, binding_type: BindingType) -> ThreadBinding | None:
|
||
|
|
key = f"{thread_id}:{binding_type.value}"
|
||
|
|
binding = self._bindings.get(key)
|
||
|
|
|
||
|
|
if binding and binding.is_expired:
|
||
|
|
self.unbind(thread_id, binding_type)
|
||
|
|
return None
|
||
|
|
|
||
|
|
return binding
|
||
|
|
|
||
|
|
def list_bindings(
|
||
|
|
self,
|
||
|
|
thread_id: str | None = None,
|
||
|
|
binding_type: BindingType | None = None,
|
||
|
|
) -> list[ThreadBinding]:
|
||
|
|
results: list[ThreadBinding] = []
|
||
|
|
expired_keys: list[str] = []
|
||
|
|
|
||
|
|
for key, binding in self._bindings.items():
|
||
|
|
if binding.is_expired:
|
||
|
|
expired_keys.append(key)
|
||
|
|
continue
|
||
|
|
if thread_id and binding.thread_id != thread_id:
|
||
|
|
continue
|
||
|
|
if binding_type and binding.binding_type != binding_type:
|
||
|
|
continue
|
||
|
|
results.append(binding)
|
||
|
|
|
||
|
|
for key in expired_keys:
|
||
|
|
self._bindings.pop(key, None)
|
||
|
|
|
||
|
|
return results
|
||
|
|
|
||
|
|
def add_listener(self, listener: Callable[[str, ThreadBinding], None]) -> None:
|
||
|
|
self._listeners.append(listener)
|
||
|
|
|
||
|
|
def remove_listener(self, listener: Callable[[str, ThreadBinding], None]) -> None:
|
||
|
|
if listener in self._listeners:
|
||
|
|
self._listeners.remove(listener)
|
||
|
|
|
||
|
|
def _notify(self, event: str, binding: ThreadBinding) -> None:
|
||
|
|
for listener in self._listeners:
|
||
|
|
try:
|
||
|
|
listener(event, binding)
|
||
|
|
except Exception:
|
||
|
|
pass
|