Luxx/luxx/routes/conversations.py

"""Conversation routes"""
from typing import Optional, List
from fastapi import APIRouter, Depends
from pydantic import BaseModel
from sqlalchemy.orm import Session

from luxx.database import get_db
from luxx.models import Conversation, Message, User
from luxx.routes.auth import get_current_user
from luxx.utils.helpers import generate_id, success_response, error_response, paginate


router = APIRouter(prefix="/conversations", tags=["Conversations"])


class ConversationCreate(BaseModel):
    """Create conversation model"""
    project_id: Optional[str] = None
    provider_id: Optional[int] = None
    title: Optional[str] = None
    model: Optional[str] = None
    system_prompt: str = "You are a helpful assistant."
    temperature: float = 0.7
    max_tokens: int = 2000
    thinking_enabled: bool = False


class ConversationUpdate(BaseModel):
    """Update conversation model"""
    title: Optional[str] = None
    model: Optional[str] = None
    system_prompt: Optional[str] = None
    temperature: Optional[float] = None
    max_tokens: Optional[int] = None
    thinking_enabled: Optional[bool] = None


@router.get("/", response_model=dict)
def list_conversations(
    page: int = 1,
    page_size: int = 20,
    current_user: User = Depends(get_current_user),
    db: Session = Depends(get_db)
):
    """Get conversation list"""
    import json
    query = db.query(Conversation).filter(Conversation.user_id == current_user.id)
    result = paginate(query.order_by(Conversation.updated_at.desc()), page, page_size)
    
    items = []
    for c in result["items"]:
        conv_dict = c.to_dict()
        # Get first user message for fallback title
        first_msg = db.query(Message).filter(
            Message.conversation_id == c.id,
            Message.role == 'user'
        ).order_by(Message.created_at).first()
        if first_msg:
            conv_dict['first_message'] = first_msg.content[:50] + ('...' if len(first_msg.content) > 50 else '')
        
        # Calculate total tokens from all assistant messages in this conversation
        assistant_messages = db.query(Message).filter(
            Message.conversation_id == c.id,
            Message.role == 'assistant'
        ).all()
        total_tokens = 0
        for msg in assistant_messages:
            total_tokens += msg.token_count or 0
            # Also try to get usage from the usage field
            if msg.usage:
                try:
                    usage_obj = json.loads(msg.usage)
                    total_tokens = usage_obj.get("total_tokens", total_tokens)
                except:
                    pass
        conv_dict['token_count'] = total_tokens
        items.append(conv_dict)
    
    return success_response(data={
        "items": items,
        "total": result["total"],
        "page": result["page"],
        "page_size": result["page_size"]
    })


@router.post("/", response_model=dict)
def create_conversation(
    data: ConversationCreate,
    current_user: User = Depends(get_current_user),
    db: Session = Depends(get_db)
):
    """Create conversation"""
    from luxx.models import LLMProvider
    
    # Get provider info - use default provider if not specified
    provider_id = data.provider_id
    model = data.model
    
    if not provider_id:
        # Find default provider
        default_provider = db.query(LLMProvider).filter(
            LLMProvider.user_id == current_user.id,
            LLMProvider.is_default == True
        ).first()
        if default_provider:
            provider_id = default_provider.id
    
    if provider_id and not model:
        provider = db.query(LLMProvider).filter(
            LLMProvider.id == provider_id,
            LLMProvider.user_id == current_user.id
        ).first()
        if provider:
            model = provider.default_model
    
    if not model:
        model = "gpt-4"
    
    conversation = Conversation(
        id=generate_id("conv"),
        user_id=current_user.id,
        project_id=data.project_id,
        provider_id=provider_id,
        title=data.title or "New Conversation",
        model=model,
        system_prompt=data.system_prompt,
        temperature=data.temperature,
        max_tokens=data.max_tokens,
        thinking_enabled=data.thinking_enabled
    )
    
    db.add(conversation)
    db.commit()
    db.refresh(conversation)
    
    return success_response(data=conversation.to_dict(), message="Conversation created successfully")


@router.get("/{conversation_id}", response_model=dict)
def get_conversation(
    conversation_id: str,
    current_user: User = Depends(get_current_user),
    db: Session = Depends(get_db)
):
    """Get conversation details"""
    conversation = db.query(Conversation).filter(
        Conversation.id == conversation_id,
        Conversation.user_id == current_user.id
    ).first()
    
    if not conversation:
        return error_response("Conversation not found", 404)
    
    return success_response(data=conversation.to_dict())


@router.put("/{conversation_id}", response_model=dict)
def update_conversation(
    conversation_id: str,
    data: ConversationUpdate,
    current_user: User = Depends(get_current_user),
    db: Session = Depends(get_db)
):
    """Update conversation"""
    conversation = db.query(Conversation).filter(
        Conversation.id == conversation_id,
        Conversation.user_id == current_user.id
    ).first()
    
    if not conversation:
        return error_response("Conversation not found", 404)
    
    update_data = data.dict(exclude_unset=True)
    for key, value in update_data.items():
        setattr(conversation, key, value)
    
    db.commit()
    db.refresh(conversation)
    
    return success_response(data=conversation.to_dict(), message="Conversation updated successfully")


@router.delete("/{conversation_id}", response_model=dict)
def delete_conversation(
    conversation_id: str,
    current_user: User = Depends(get_current_user),
    db: Session = Depends(get_db)
):
    """Delete conversation"""
    conversation = db.query(Conversation).filter(
        Conversation.id == conversation_id,
        Conversation.user_id == current_user.id
    ).first()
    
    if not conversation:
        return error_response("Conversation not found", 404)
    
    db.delete(conversation)
    db.commit()
    
    return success_response(message="Conversation deleted successfully")