""" Direct Messages API - 1-on-1 messaging with @grimlock support """ from fastapi import APIRouter, Depends, HTTPException, status, BackgroundTasks from sqlalchemy.orm import Session from sqlalchemy import or_, and_ from pydantic import BaseModel from typing import List, Optional from datetime import datetime from core.database import get_db from core.models import DirectMessage, User from core.ai_client import AIClient from core.context_manager import ContextManager from api.auth import get_current_user from core.websocket import send_dm_notification import main router = APIRouter() class DMCreate(BaseModel): recipient_id: int content: str class DMResponse(BaseModel): id: int sender_id: int recipient_id: int content: str is_ai_message: bool read_at: Optional[str] created_at: str class Config: from_attributes = True class ConversationResponse(BaseModel): user_id: int user_name: str user_email: str last_message: str last_message_at: str unread_count: int async def handle_grimlock_dm( dm: DirectMessage, sender: User, db: Session, context_manager: ContextManager, ai_client: AIClient ): """Handle @grimlock mention in DM""" if not dm.content.lower().startswith('@grimlock'): return try: # Get conversation history history = db.query(DirectMessage)\ .filter( or_( and_(DirectMessage.sender_id == sender.id, DirectMessage.recipient_id == 1), and_(DirectMessage.sender_id == 1, DirectMessage.recipient_id == sender.id) ) )\ .filter(DirectMessage.id < dm.id)\ .order_by(DirectMessage.id.desc())\ .limit(10)\ .all() # Build conversation conversation = [] for msg in reversed(history): if msg.sender_id == sender.id: conversation.append({"role": "user", "content": msg.content}) else: conversation.append({"role": "assistant", "content": msg.content}) # Add current message query = dm.content.replace('@grimlock', '').strip() conversation.append({"role": "user", "content": query}) # Get context context = context_manager.get_context_for_query(query, role=sender.role.value) system_prompt = context_manager.get_system_prompt(role=sender.role.value) system_prompt += f"\n\nYou are in a direct message conversation with {sender.name}." if context: system_prompt += f"\n\n# Company Context\n{context}" # Get AI response response = await ai_client.chat(messages=conversation, system_prompt=system_prompt) # Create AI DM ai_dm = DirectMessage( sender_id=None, # AI has no user_id recipient_id=sender.id, content=response, is_ai_message=True ) db.add(ai_dm) db.commit() db.refresh(ai_dm) # Notify via WebSocket await send_dm_notification(sender.id, { "id": ai_dm.id, "sender_id": None, "content": response, "is_ai_message": True, "created_at": ai_dm.created_at.isoformat() }) except Exception as e: logger.error(f"Error handling Grimlock DM: {e}") @router.post("/", response_model=DMResponse) async def send_dm( dm_data: DMCreate, background_tasks: BackgroundTasks, current_user: User = Depends(get_current_user), db: Session = Depends(get_db), context_manager: ContextManager = Depends(main.get_context_manager), ai_client: AIClient = Depends(main.get_ai_client) ): """Send a direct message""" # Check recipient exists recipient = db.query(User).filter(User.id == dm_data.recipient_id).first() if not recipient: raise HTTPException(status_code=404, detail="Recipient not found") # Can't DM yourself if recipient.id == current_user.id: raise HTTPException(status_code=400, detail="Cannot send DM to yourself") # Create DM dm = DirectMessage( sender_id=current_user.id, recipient_id=dm_data.recipient_id, content=dm_data.content ) db.add(dm) db.commit() db.refresh(dm) # Notify recipient via WebSocket await send_dm_notification(recipient.id, { "id": dm.id, "sender_id": current_user.id, "sender_name": current_user.name, "content": dm.content, "created_at": dm.created_at.isoformat() }) # Check for @grimlock if '@grimlock' in dm.content.lower(): background_tasks.add_task(handle_grimlock_dm, dm, current_user, db, context_manager, ai_client) return { "id": dm.id, "sender_id": dm.sender_id, "recipient_id": dm.recipient_id, "content": dm.content, "is_ai_message": dm.is_ai_message, "read_at": dm.read_at.isoformat() if dm.read_at else None, "created_at": dm.created_at.isoformat() } @router.get("/conversations", response_model=List[ConversationResponse]) async def list_conversations( current_user: User = Depends(get_current_user), db: Session = Depends(get_db) ): """List all DM conversations""" # Get all users current user has DMed with sent = db.query(DirectMessage.recipient_id).filter(DirectMessage.sender_id == current_user.id).distinct() received = db.query(DirectMessage.sender_id).filter(DirectMessage.recipient_id == current_user.id).distinct() user_ids = set([r[0] for r in sent] + [r[0] for r in received]) conversations = [] for user_id in user_ids: if user_id is None: # Skip AI (no user_id) continue user = db.query(User).filter(User.id == user_id).first() if not user: continue # Get last message last_msg = db.query(DirectMessage)\ .filter( or_( and_(DirectMessage.sender_id == current_user.id, DirectMessage.recipient_id == user_id), and_(DirectMessage.sender_id == user_id, DirectMessage.recipient_id == current_user.id) ) )\ .order_by(DirectMessage.id.desc())\ .first() # Count unread unread = db.query(DirectMessage)\ .filter(DirectMessage.sender_id == user_id)\ .filter(DirectMessage.recipient_id == current_user.id)\ .filter(DirectMessage.read_at == None)\ .count() conversations.append({ "user_id": user.id, "user_name": user.name, "user_email": user.email, "last_message": last_msg.content[:100] if last_msg else "", "last_message_at": last_msg.created_at.isoformat() if last_msg else "", "unread_count": unread }) return conversations @router.get("/{user_id}/messages", response_model=List[DMResponse]) async def get_dm_history( user_id: int, limit: int = 50, before: Optional[int] = None, current_user: User = Depends(get_current_user), db: Session = Depends(get_db) ): """Get DM history with a user""" query = db.query(DirectMessage).filter( or_( and_(DirectMessage.sender_id == current_user.id, DirectMessage.recipient_id == user_id), and_(DirectMessage.sender_id == user_id, DirectMessage.recipient_id == current_user.id) ) ) if before: query = query.filter(DirectMessage.id < before) messages = query.order_by(DirectMessage.id.desc()).limit(limit).all() messages.reverse() # Mark as read db.query(DirectMessage)\ .filter(DirectMessage.sender_id == user_id)\ .filter(DirectMessage.recipient_id == current_user.id)\ .filter(DirectMessage.read_at == None)\ .update({"read_at": datetime.utcnow()}) db.commit() return [ { "id": msg.id, "sender_id": msg.sender_id, "recipient_id": msg.recipient_id, "content": msg.content, "is_ai_message": msg.is_ai_message, "read_at": msg.read_at.isoformat() if msg.read_at else None, "created_at": msg.created_at.isoformat() } for msg in messages ]