src/api/agents/chat/router.py

"""
Chat Agent Router - API Endpoints for Chat functionality
"""
from fastapi import APIRouter, HTTPException
from pydantic import BaseModel, Field
from typing import List, Optional, Dict, Any
from enum import Enum
from src.services.groq_service import generate_response
from src.api.core.logger import logger

router = APIRouter()


class AgentType(str, Enum):
    """Available AI Agent Types"""
    GENERAL = "general"           # General assistant
    CODER = "coder"              # Code generation & review
    DATA_ANALYST = "data_analyst" # Data analysis
    RESEARCHER = "researcher"    # Research & writing
    TUTOR = "tutor"              # Teaching & explanations


class ChatMessage(BaseModel):
    """Chat message model"""
    role: str = Field(..., description="Role: 'user' or 'assistant'")
    content: str = Field(..., description="Message content")


class ChatRequest(BaseModel):
    """Request model for chat"""
    message: str = Field(..., description="User message")
    agent_type: Optional[AgentType] = Field(default=AgentType.GENERAL, description="Agent type to use")
    context: Optional[Dict[str, Any]] = Field(default=None, description="Additional context")
    conversation_id: Optional[str] = Field(default=None, description="Conversation ID for history")


class ChatResponse(BaseModel):
    """Response model for chat"""
    response: str
    conversation_id: str
    agent_type: str
    model: str = "llama-3.1-8b-instant"


# Agent System Prompts
AGENT_PROMPTS = {
    AgentType.GENERAL: """You are AmkyawDev AI Assistant, a highly capable and professional AI assistant.

## Your Role
You are a versatile AI assistant developed by AmkyawDev. You help users with various tasks including:
- Answering questions
- Providing explanations
- Problem-solving
- Creative tasks

## Guidelines
1. Always be helpful, polite, and professional
2. Provide accurate and well-structured responses
3. Use markdown formatting for clarity
4. Ask clarifying questions when needed
5. Admit when you don't know something
6. Be concise but thorough

## Response Format
- Use bullet points for lists
- Use code blocks for code
- Use headers for structure
- Be direct and clear""",

    AgentType.CODER: """You are AmkyawDev Coding Assistant, an expert software developer.

## Your Role
You are a specialized coding assistant that helps with:
- Writing clean, efficient code
- Code review and debugging
- Algorithm design
- Best practices
- Technical documentation

## Guidelines
1. Write clean, readable, and well-documented code
2. Follow language-specific best practices
3. Include code comments explaining complex logic
4. Handle errors gracefully
5. Consider performance and scalability
6. Provide multiple solutions when appropriate

## Response Format
- Always provide working code examples
- Use proper code formatting with language hints
- Explain the code after showing it
- Include time/space complexity when relevant
- Suggest improvements and alternatives""",

    AgentType.DATA_ANALYST: """You are AmkyawDev Data Analyst, an expert in data analysis and visualization.

## Your Role
You specialize in:
- Data analysis and interpretation
- Statistical analysis
- Data visualization suggestions
- Report generation
- Insights and recommendations

## Guidelines
1. Approach problems with data-driven mindset
2. Consider statistical significance
3. Suggest appropriate visualizations
4. Provide actionable insights
5. Explain your reasoning

## Response Format
- Use tables for structured data
- Use code blocks for data manipulation
- Include statistical measures
- Suggest visualization types""",

    AgentType.RESEARCHER: """You are AmkyawDev Research Assistant, an expert researcher and writer.

## Your Role
You help with:
- Research and information gathering
- Writing and content creation
- Summary and synthesis
- Citation and references
- Critical analysis

## Guidelines
1. Be thorough and well-researched
2. Cite sources when possible
3. Present multiple perspectives
4. Be objective and balanced
5. Structure content clearly

## Response Format
- Use headings for organization
- Include references section
- Provide summaries and conclusions
- Use bullet points for key findings""",

    AgentType.TUTOR: """You are AmkyawDev Tutor, a patient and knowledgeable teacher.

## Your Role
You help users learn and understand:
- Complex concepts made simple
- Step-by-step explanations
- Practical examples
- Knowledge building
- Skill development

## Guidelines
1. Be patient and encouraging
2. Break down complex topics
3. Use analogies and examples
4. Check understanding
5. Build on prior knowledge
6. Be interactive

## Response Format
- Start with fundamentals
- Build up progressively
- Use simple language
- Include practice exercises
- Ask checking questions"""
}


# Skills/Commands that users can trigger
SKILL_COMMANDS = {
    "code_review": "Review the provided code for bugs, security issues, and best practices",
    "bug_fix": "Find and fix bugs in the provided code",
    "refactor": "Improve and refactor the provided code",
    "document": "Generate documentation for the provided code",
    "explain": "Explain the provided code or concept in detail",
    "optimize": "Optimize the provided code for performance",
    "test": "Write unit tests for the provided code"
}


# Simple in-memory conversation storage
conversations: Dict[str, List[Dict[str, str]]] = {}
conversation_agents: Dict[str, AgentType] = {}  # Track agent type per conversation


def build_system_message(agent_type: AgentType) -> Dict[str, str]:
    """Build system message based on agent type"""
    base_prompt = AGENT_PROMPTS.get(agent_type, AGENT_PROMPTS[AgentType.GENERAL])
    
    # Add skill commands info
    skills_info = "\n\n## Available Skills/Commands\nYou can help with these specialized tasks:\n"
    for cmd, desc in SKILL_COMMANDS.items():
        skills_info += f"- `{cmd}`: {desc}\n"
    skills_info += "\nWhen user requests a skill, perform that task specifically."
    
    return {"role": "system", "content": base_prompt + skills_info}


@router.post("/chat", response_model=ChatResponse)
async def chat(request: ChatRequest) -> ChatResponse:
    """
    Chat endpoint - AI conversation with agent control
    
    - **message**: User's message
    - **agent_type**: Type of agent (general, coder, data_analyst, researcher, tutor)
    - **context**: Optional context data
    - **conversation_id**: Optional conversation ID for history
    """
    try:
        logger.info(f"Chat request: {request.message[:50]}... (agent: {request.agent_type})")
        
        # Get or create conversation
        conv_id = request.conversation_id or "default"
        if conv_id not in conversations:
            conversations[conv_id] = []
            conversation_agents[conv_id] = request.agent_type
        
        # Use specified agent type
        agent_type = request.agent_type or conversation_agents.get(conv_id, AgentType.GENERAL)
        conversation_agents[conv_id] = agent_type
        
        # Build messages with system prompt
        messages = [build_system_message(agent_type)]
        
        # Add context if provided
        if request.context:
            context_msg = f"Context: {request.context}"
            messages.append({"role": "system", "content": context_msg})
        
        # Add conversation history
        messages.extend(conversations[conv_id][-10:])  # Keep last 10 messages
        
        # Add current message
        messages.append({"role": "user", "content": request.message})
        
        # Generate response
        response_text = await generate_response(
            messages=messages,
            model="llama-3.1-8b-instant",
            temperature=0.7,
            max_tokens=2048
        )
        
        # Save to conversation history
        conversations[conv_id].append({"role": "user", "content": request.message})
        conversations[conv_id].append({"role": "assistant", "content": response_text})
        
        # Keep history limited
        if len(conversations[conv_id]) > 20:
            conversations[conv_id] = conversations[conv_id][-20:]
        
        return ChatResponse(
            response=response_text,
            conversation_id=conv_id,
            agent_type=agent_type.value,
            model="llama-3.1-8b-instant"
        )
        
    except Exception as e:
        logger.error(f"Chat error: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))


@router.get("/agents")
async def get_agents() -> Dict[str, Any]:
    """Get all available agents"""
    return {
        "agents": [
            {"id": agent.value, "name": agent.value.replace("_", " ").title(), "prompt": AGENT_PROMPTS[agent][:100] + "..."}
            for agent in AgentType
        ],
        "count": len(AgentType)
    }


@router.get("/agents/skills")
async def get_skills() -> Dict[str, Any]:
    """Get all available skills"""
    return {
        "skills": [
            {"id": cmd, "description": desc}
            for cmd, desc in SKILL_COMMANDS.items()
        ],
        "count": len(SKILL_COMMANDS)
    }


@router.get("/conversations")
async def get_conversations() -> Dict[str, Any]:
    """Get all conversation IDs"""
    return {
        "conversations": [
            {
                "id": conv_id,
                "agent_type": conversation_agents.get(conv_id, "general"),
                "message_count": len(conversations[conv_id])
            }
            for conv_id in conversations.keys()
        ],
        "count": len(conversations)
    }


@router.delete("/conversations/{conversation_id}")
async def delete_conversation(conversation_id: str) -> Dict[str, str]:
    """Delete a conversation"""
    if conversation_id in conversations:
        del conversations[conversation_id]
        conversation_agents.pop(conversation_id, None)
        return {"status": "deleted", "conversation_id": conversation_id}
    raise HTTPException(status_code=404, detail="Conversation not found")


@router.get("/conversations/{conversation_id}")
async def get_conversation(conversation_id: str) -> Dict[str, Any]:
    """Get conversation history"""
    if conversation_id in conversations:
        return {
            "conversation_id": conversation_id,
            "agent_type": conversation_agents.get(conversation_id, "general"),
            "messages": conversations[conversation_id]
        }
    raise HTTPException(status_code=404, detail="Conversation not found")