362 lines
13 KiB
Python
362 lines
13 KiB
Python
"""Chat Room Service - orchestrates multi-agent chat"""
|
|
import json
|
|
import re
|
|
import uuid
|
|
import asyncio
|
|
import logging
|
|
from typing import List, Dict, Any, Optional, AsyncGenerator
|
|
from dataclasses import dataclass
|
|
|
|
from luxx.core.database import SessionLocal
|
|
from luxx.models.room import ChatRoom, Agent, ChatRoomAgent, ChatRoomMessage
|
|
from luxx.agents.base import BaseAgent
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
# ==================== Dispatcher ====================
|
|
|
|
@dataclass
|
|
class DispatchResult:
|
|
"""Result of message dispatch"""
|
|
triggered_agents: List[BaseAgent]
|
|
mentions: List[str]
|
|
should_respond: bool
|
|
|
|
|
|
class MessageDispatcher:
|
|
"""Dispatcher for routing messages to agents"""
|
|
|
|
@staticmethod
|
|
def parse_mentions(content: str) -> List[str]:
|
|
"""Parse @mentions from message content"""
|
|
pattern = r'@(\w+)'
|
|
return re.findall(pattern, content)
|
|
|
|
@staticmethod
|
|
def get_agents_by_names(names: List[str], room_agents: List[BaseAgent]) -> List[BaseAgent]:
|
|
"""Get agents by their names (case-insensitive)"""
|
|
name_lower_map = {a.name.lower(): a for a in room_agents}
|
|
matched = []
|
|
for name in names:
|
|
agent = name_lower_map.get(name.lower())
|
|
if agent:
|
|
matched.append(agent)
|
|
return matched
|
|
|
|
@staticmethod
|
|
def get_agents_by_ids(agent_ids: List[str], room_agents: List[BaseAgent]) -> List[BaseAgent]:
|
|
"""Get agents by their IDs"""
|
|
id_set = set(agent_ids)
|
|
return [a for a in room_agents if a.agent_id in id_set]
|
|
|
|
def dispatch(self, content: str, room_agents: List[BaseAgent], sender_id: str, sender_type: str = "user") -> DispatchResult:
|
|
"""Dispatch a message to appropriate agents."""
|
|
available_agents = [a for a in room_agents if a.agent_id != sender_id]
|
|
mentions = self.parse_mentions(content)
|
|
|
|
if mentions:
|
|
triggered = self.get_agents_by_names(mentions, available_agents)
|
|
logger.info(f"Message with mentions: {mentions} -> triggered: {[a.name for a in triggered]}")
|
|
return DispatchResult(triggered_agents=triggered, mentions=mentions, should_respond=len(triggered) > 0)
|
|
|
|
auto_agents = [a for a in available_agents if a.auto_response]
|
|
auto_agents.sort(key=lambda a: a.priority)
|
|
logger.info(f"Auto-response agents triggered: {[a.name for a in auto_agents]}")
|
|
return DispatchResult(triggered_agents=auto_agents, mentions=[], should_respond=len(auto_agents) > 0)
|
|
|
|
|
|
# ==================== Aggregator ====================
|
|
|
|
class ResponseAggregator:
|
|
"""Aggregates responses from multiple agents"""
|
|
|
|
def __init__(self, room_id: str):
|
|
self.room_id = room_id
|
|
self._agent_responses: Dict[str, Dict[str, Any]] = {}
|
|
|
|
async def aggregate_stream(self, agent_streams: Dict[str, AsyncGenerator]) -> AsyncGenerator[Dict[str, Any], None]:
|
|
"""Aggregate streaming responses from multiple agents."""
|
|
if not agent_streams:
|
|
return
|
|
|
|
async def collect_agent_stream(agent_id: str, stream):
|
|
try:
|
|
async for event in stream:
|
|
event["agent_id"] = agent_id
|
|
yield event
|
|
except Exception as e:
|
|
logger.error(f"Agent {agent_id} stream error: {e}")
|
|
yield {"event": "error", "agent_id": agent_id, "data": {"content": str(e)}}
|
|
|
|
tasks = [collect_agent_stream(agent_id, stream) for agent_id, stream in agent_streams.items()]
|
|
async for event in asyncio.merge(*tasks):
|
|
yield event
|
|
|
|
def aggregate_final(self, responses: Dict[str, Dict[str, Any]]) -> List[Dict[str, Any]]:
|
|
"""Aggregate final responses from agents."""
|
|
results = []
|
|
for agent_id, response in responses.items():
|
|
if response.get("event") == "done":
|
|
results.append({
|
|
"agent_id": agent_id,
|
|
"agent_name": response.get("agent_name"),
|
|
"message_id": response.get("message_id"),
|
|
"content": response.get("content"),
|
|
"token_count": response.get("token_count", 0)
|
|
})
|
|
return results
|
|
|
|
|
|
# ==================== Chat Room Service ====================
|
|
|
|
class ChatRoomService:
|
|
"""Service for managing chat rooms with multi-agent support"""
|
|
|
|
def __init__(self):
|
|
self.dispatcher = MessageDispatcher()
|
|
|
|
def get_room(self, room_id: str) -> Optional[ChatRoom]:
|
|
"""Get a chat room by ID"""
|
|
db = SessionLocal()
|
|
try:
|
|
return db.query(ChatRoom).filter(ChatRoom.id == room_id).first()
|
|
finally:
|
|
db.close()
|
|
|
|
def get_room_agents(self, room_id: str) -> List[BaseAgent]:
|
|
"""Get all active agents in a chat room"""
|
|
db = SessionLocal()
|
|
try:
|
|
room_agents = db.query(ChatRoomAgent).filter(
|
|
ChatRoomAgent.chat_room_id == room_id,
|
|
ChatRoomAgent.is_active == True
|
|
).all()
|
|
|
|
agents = []
|
|
for ra in room_agents:
|
|
agent_db = db.query(Agent).filter(Agent.id == ra.agent_id, Agent.is_active == True).first()
|
|
if agent_db:
|
|
agents.append(BaseAgent.from_model(agent_db))
|
|
|
|
agents.sort(key=lambda a: a.priority)
|
|
return agents
|
|
finally:
|
|
db.close()
|
|
|
|
def get_agent(self, agent_id: str) -> Optional[BaseAgent]:
|
|
"""Get an agent by ID"""
|
|
db = SessionLocal()
|
|
try:
|
|
agent_db = db.query(Agent).filter(Agent.id == agent_id).first()
|
|
if agent_db:
|
|
return BaseAgent.from_model(agent_db)
|
|
return None
|
|
finally:
|
|
db.close()
|
|
|
|
def list_rooms(self, user_id: int = None, include_agents: bool = True) -> List[Dict]:
|
|
"""List all chat rooms"""
|
|
db = SessionLocal()
|
|
try:
|
|
query = db.query(ChatRoom)
|
|
if user_id:
|
|
query = query.filter(ChatRoom.owner_id == user_id)
|
|
rooms = query.order_by(ChatRoom.updated_at.desc()).all()
|
|
return [r.to_dict(include_agents=include_agents) for r in rooms]
|
|
finally:
|
|
db.close()
|
|
|
|
def create_room(self, name: str, owner_id: int, description: str = None, agent_ids: List[str] = None) -> Dict:
|
|
"""Create a new chat room"""
|
|
db = SessionLocal()
|
|
try:
|
|
room = ChatRoom(
|
|
id=str(uuid.uuid4()),
|
|
name=name,
|
|
description=description,
|
|
owner_id=owner_id
|
|
)
|
|
db.add(room)
|
|
|
|
if agent_ids:
|
|
for agent_id in agent_ids:
|
|
room_agent = ChatRoomAgent(
|
|
id=str(uuid.uuid4()),
|
|
chat_room_id=room.id,
|
|
agent_id=agent_id
|
|
)
|
|
db.add(room_agent)
|
|
|
|
db.commit()
|
|
return room.to_dict(include_agents=True)
|
|
finally:
|
|
db.close()
|
|
|
|
def update_room(self, room_id: str, name: str = None, description: str = None, is_active: bool = None) -> Optional[Dict]:
|
|
"""Update a chat room"""
|
|
db = SessionLocal()
|
|
try:
|
|
room = db.query(ChatRoom).filter(ChatRoom.id == room_id).first()
|
|
if not room:
|
|
return None
|
|
|
|
if name is not None:
|
|
room.name = name
|
|
if description is not None:
|
|
room.description = description
|
|
if is_active is not None:
|
|
room.is_active = is_active
|
|
|
|
db.commit()
|
|
return room.to_dict(include_agents=True)
|
|
finally:
|
|
db.close()
|
|
|
|
def delete_room(self, room_id: str) -> bool:
|
|
"""Delete a chat room"""
|
|
db = SessionLocal()
|
|
try:
|
|
room = db.query(ChatRoom).filter(ChatRoom.id == room_id).first()
|
|
if not room:
|
|
return False
|
|
db.delete(room)
|
|
db.commit()
|
|
return True
|
|
finally:
|
|
db.close()
|
|
|
|
def add_agent_to_room(self, room_id: str, agent_id: str) -> bool:
|
|
"""Add an agent to a chat room"""
|
|
db = SessionLocal()
|
|
try:
|
|
existing = db.query(ChatRoomAgent).filter(
|
|
ChatRoomAgent.chat_room_id == room_id,
|
|
ChatRoomAgent.agent_id == agent_id
|
|
).first()
|
|
|
|
if existing:
|
|
existing.is_active = True
|
|
else:
|
|
room_agent = ChatRoomAgent(
|
|
id=str(uuid.uuid4()),
|
|
chat_room_id=room_id,
|
|
agent_id=agent_id
|
|
)
|
|
db.add(room_agent)
|
|
|
|
db.commit()
|
|
return True
|
|
finally:
|
|
db.close()
|
|
|
|
def remove_agent_from_room(self, room_id: str, agent_id: str) -> bool:
|
|
"""Remove an agent from a chat room"""
|
|
db = SessionLocal()
|
|
try:
|
|
room_agent = db.query(ChatRoomAgent).filter(
|
|
ChatRoomAgent.chat_room_id == room_id,
|
|
ChatRoomAgent.agent_id == agent_id
|
|
).first()
|
|
|
|
if room_agent:
|
|
room_agent.is_active = False
|
|
db.commit()
|
|
return True
|
|
return False
|
|
finally:
|
|
db.close()
|
|
|
|
def get_messages(self, room_id: str, limit: int = 50, before_id: str = None) -> List[Dict]:
|
|
"""Get messages from a chat room"""
|
|
db = SessionLocal()
|
|
try:
|
|
query = db.query(ChatRoomMessage).filter(
|
|
ChatRoomMessage.room_id == room_id
|
|
).order_by(ChatRoomMessage.created_at.desc())
|
|
|
|
if before_id:
|
|
before_msg = db.query(ChatRoomMessage).filter(
|
|
ChatRoomMessage.id == before_id
|
|
).first()
|
|
if before_msg:
|
|
query = query.filter(ChatRoomMessage.created_at < before_msg.created_at)
|
|
|
|
messages = query.limit(limit).all()
|
|
return [m.to_dict() for m in reversed(messages)]
|
|
finally:
|
|
db.close()
|
|
|
|
def save_message(self, room_id: str, sender_type: str, sender_id: str, sender_name: str, content: str,
|
|
mentions: List[str] = None, parent_id: str = None, token_count: int = 0) -> Dict:
|
|
"""Save a message to a chat room"""
|
|
db = SessionLocal()
|
|
try:
|
|
msg = ChatRoomMessage(
|
|
id=str(uuid.uuid4()),
|
|
room_id=room_id,
|
|
sender_type=sender_type,
|
|
sender_id=sender_id,
|
|
sender_name=sender_name,
|
|
content=content,
|
|
mentions=json.dumps(mentions) if mentions else None,
|
|
parent_id=parent_id,
|
|
token_count=token_count
|
|
)
|
|
db.add(msg)
|
|
|
|
room = db.query(ChatRoom).filter(ChatRoom.id == room_id).first()
|
|
if room:
|
|
from datetime import datetime
|
|
room.updated_at = datetime.now()
|
|
|
|
db.commit()
|
|
return msg.to_dict()
|
|
finally:
|
|
db.close()
|
|
|
|
async def process_message(self, room_id: str, user_message: str, user_id: str, user_name: str, context: Dict = None) -> AsyncGenerator[Dict[str, Any], None]:
|
|
"""Process a user message and dispatch to appropriate agents."""
|
|
room = self.get_room(room_id)
|
|
if not room:
|
|
yield {"event": "error", "data": {"content": "Chat room not found"}}
|
|
return
|
|
|
|
room_agents = self.get_room_agents(room_id)
|
|
if not room_agents:
|
|
yield {"event": "error", "data": {"content": "No agents available in this room"}}
|
|
return
|
|
|
|
dispatch_result = self.dispatcher.dispatch(
|
|
content=user_message,
|
|
room_agents=room_agents,
|
|
sender_id=user_id,
|
|
sender_type="user"
|
|
)
|
|
|
|
if not dispatch_result.should_respond:
|
|
yield {"event": "no_response", "data": {"message": "No agents triggered"}}
|
|
return
|
|
|
|
messages = self.get_messages(room_id, limit=20)
|
|
|
|
agent_streams = {}
|
|
for agent in dispatch_result.triggered_agents:
|
|
stream = agent.stream_response(
|
|
user_message=user_message,
|
|
conversation_history=messages,
|
|
context=context
|
|
)
|
|
agent_streams[agent.agent_id] = stream
|
|
|
|
aggregator = ResponseAggregator(room_id)
|
|
async for event in aggregator.aggregate_stream(agent_streams):
|
|
yield event
|
|
|
|
|
|
# Global service instance
|
|
chat_room_service = ChatRoomService()
|
|
|
|
# Export for backward compatibility
|
|
dispatcher = chat_room_service.dispatcher
|