microfish/backend/app/api/agent_group.py

"""
Agent Grouping API
Groups simulation agents into categories and filters them
"""

import json
import os
from flask import Blueprint, request, jsonify
from ..utils.llm_client import LLMClient
from ..utils.locale import t, get_language_instruction
from ..utils.logger import get_logger

logger = get_logger('crowdsight.agent_group')

agent_group_bp = Blueprint('agent_group', __name__)


@agent_group_bp.route('/categorize', methods=['POST'])
def categorize_agents():
    """
    Categorize agents into groups based on their profiles.
    
    Request JSON:
        {
            "agents": [
                {
                    "agent_id": 0,
                    "name": "...",
                    "profession": "...",
                    "bio": "...",
                    "persona": "...",
                    "interested_topics": [...]
                }
            ],
            "simulation_requirement": "..."
        }
    
    Response JSON:
        {
            "success": true,
            "groups": [
                {
                    "group_id": "target_audience",
                    "group_name": "กลุ่มเป้าหมาย",
                    "default_enabled": true,
                    "agents": [0, 2, 5]
                },
                {
                    "group_id": "advertiser",
                    "group_name": "ผู้โฆษณา/แบรนด์",
                    "default_enabled": false,
                    "agents": [1, 3]
                }
            ]
        }
    """
    try:
        data = request.get_json() or {}
        agents = data.get('agents', [])
        simulation_requirement = data.get('simulation_requirement', '')
        
        if not agents:
            return jsonify({'success': False, 'error': 'No agents provided'}), 400
        
        # Build agent summary for LLM
        agent_summaries = []
        for i, agent in enumerate(agents):
            summary = f"[{i}] {agent.get('name', 'Unknown')} - {agent.get('profession', 'N/A')} - {agent.get('bio', '')[:100]}"
            agent_summaries.append(summary)
        
        agents_text = '\n'.join(agent_summaries)
        lang_instruction = get_language_instruction()
        
        system_prompt = f"""You are an expert at analyzing social simulation agents. Your task is to categorize agents into groups and determine which groups are useful for the simulation.

{lang_instruction}

Return JSON format:
{{
    "groups": [
        {{
            "group_id": "unique_english_id",
            "group_name": "group name in user's language",
            "description": "brief description",
            "default_enabled": true/false,
            "agent_indices": [0, 1, 2]
        }}
    ]
}}

Guidelines:
- Group agents by their ROLE in the simulation context
- Groups that represent the TARGET AUDIENCE, INFLUENCERS, MEDIA, COMPETITORS should be default_enabled=true
- Groups that represent THE ADVERTISER/BRAND ITSELF, ABSTRACT CONCEPTS, MARKETING METADATA should be default_enabled=false
- If the simulation is about an ad/campaign, the company that made the ad should be in a disabled group
- Each agent should be in exactly one group
- Agent indices must match the input indices"""

        user_prompt = f"""Simulation requirement: {simulation_requirement}

Agents to categorize:
{agents_text}

Categorize these agents into groups. Mark groups as default_enabled=false if they represent the entity that created the content being simulated (e.g., the advertiser in an ad scenario)."""

        llm = LLMClient()
        result = llm.chat_json(
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_prompt}
            ],
            temperature=0.3
        )
        
        # Validate and clean up
        groups = result.get('groups', [])
        if not groups:
            # Fallback: put all agents in one enabled group
            groups = [{
                'group_id': 'all',
                'group_name': 'All Agents',
                'description': 'All agents',
                'default_enabled': True,
                'agent_indices': list(range(len(agents)))
            }]
        
        return jsonify({
            'success': True,
            'groups': groups
        })
    
    except Exception as e:
        logger.error(f"Agent categorization failed: {e}")
        return jsonify({'success': False, 'error': str(e)}), 500


@agent_group_bp.route('/filter', methods=['POST'])
def filter_agents():
    """
    Filter agents based on selected groups.
    
    Request JSON:
        {
            "agents": [...],  // full agent list
            "selected_groups": [0, 2]  // indices of enabled groups
        }
    
    Response JSON:
        {
            "success": true,
            "selected_agent_ids": [0, 2, 5]
        }
    """
    try:
        data = request.get_json() or {}
        agents = data.get('agents', [])
        groups = data.get('groups', [])
        selected_group_ids = data.get('selected_group_ids', [])
        
        # Collect agent indices from selected groups
        selected_indices = set()
        for group in groups:
            if group.get('group_id') in selected_group_ids:
                selected_indices.update(group.get('agent_indices', []))
        
        return jsonify({
            'success': True,
            'selected_agent_ids': sorted(selected_indices)
        })
    
    except Exception as e:
        logger.error(f"Agent filtering failed: {e}")
        return jsonify({'success': False, 'error': str(e)}), 500