agent_tracker/audit_analyzer.py

import os
import json
import re
import asyncio
import logging
from datetime import datetime
from bson import ObjectId

from google import genai

from database import agents_collection, audit_history_collection

logger = logging.getLogger("audit_analyzer")
logger.setLevel(logging.INFO)
if not logger.handlers:
    handler = logging.StreamHandler()
    handler.setFormatter(logging.Formatter("%(asctime)s [AUDIT] %(message)s"))
    logger.addHandler(handler)


def is_gemini_configured() -> bool:
    return bool(os.getenv("GOOGLE_API_KEY"))


def _get_client():
    api_key = os.getenv("GOOGLE_API_KEY")
    return genai.Client(api_key=api_key)


def _get_model_name():
    return os.getenv("AUDIT_GEMINI_MODEL", "gemini-2.5-pro")


SYSTEM_INSTRUCTION = """You are an AI agent compliance analyst and classifier. Analyse AI agents based on their
system prompts/instructions and metadata. You must:
1. Classify into business risk categories
2. Assign a business discipline
3. Infer the department/team from the instructions
4. Detect whether the agent is used for client-specific work

Respond ONLY with valid JSON matching this schema:
{
  "category": "1" | "1B" | "2" | "3",
  "category_reasoning": "string",
  "discipline": "Strategy" | "Creative" | "Oversight including delivery" | "Optimization" | "Back Office including operations" | "Pencil Agents",
  "discipline_reasoning": "string - why this discipline was chosen",
  "department": "string or null - inferred team/department from instructions (e.g. Project Management, Creative Services, Finance, Media, Strategy). null if not determinable",
  "is_client_work": true | false,
  "client_work_reasoning": "string - evidence from instructions that this agent handles client-specific work, references client names, brands, deliverables, or external-facing outputs. Empty string if not client work",
  "client_name_detected": "string or null - specific client or brand name found in instructions, null if none",
  "flags": ["array", "of", "strings"],
  "summary": "2-3 sentence analysis",
  "recommendations": "which team(s) should review and why",
  "risk_level": "low" | "medium" | "high" | "critical"
}

CATEGORY DEFINITIONS:
Cat 1 - Internal Sandbox/Experimentation (Oliver AI Sandbox, behind the scenes, needs IT/Compliance)
Cat 1B - High Cost Internal (may incur large cost, not Cat 2 or 3)
Cat 2 - Client-Exposed Not Sold (Pencil platform, exposed to clients, needs Legal)
Cat 3 - Client-Sold (Pencil platform, sold to clients, needs Commercial team)

DISCIPLINE DEFINITIONS (pick the best fit):
- Strategy: Agents focused on strategic planning, research, market analysis, insights
- Creative: Agents focused on creative work, content creation, design, copywriting
- Oversight including delivery: Agents focused on project management, delivery, QA, oversight, compliance
- Optimization: Agents focused on performance optimization, data analysis, efficiency, media planning
- Back Office including operations: Agents focused on internal operations, HR, finance, IT support, admin tasks
- Pencil Agents: Agents built on or for the Pencil platform specifically

DEPARTMENT INFERENCE:
Look for clues in the instructions about which team or role uses this agent. Examples:
- "I am a project manager" -> department: "Project Management"
- "help the media team" -> department: "Media"
- "creative brief" -> department: "Creative"
- If no department/team clues found, set department to null

CLIENT WORK DETECTION:
An agent is client work (is_client_work = true) if the instructions reference:
- Specific client names or brand names
- Client deliverables, client presentations, client reports
- External-facing outputs meant for clients
- Client briefs, client feedback, client approvals
- Work explicitly described as "for the client" or "client-facing"
Do NOT flag as client work if the agent merely mentions "users" or "stakeholders" generically.

RISK LEVELS:
- low: Internal-only, limited capabilities
- medium: Internal with external tool access or moderate cost
- high: Client-facing or accesses sensitive data
- critical: Client-sold or handles financial/legal/PII data

FLAGS TO CONSIDER:
internal_only, experimental, sandbox, client_facing, pencil_platform,
revenue_generating, not_for_sale, uses_external_tools, uses_code_interpreter,
uses_file_search, accesses_sensitive_data, handles_pii, high_cost,
resource_intensive, legal_review_needed, commercial_review_needed,
compliance_review_needed, no_instructions"""


USER_PROMPT_TEMPLATE = """Analyse this AI agent:

AGENT NAME: {name}
DESCRIPTION: {description}
AUTHOR: {author}

INSTRUCTIONS (SYSTEM PROMPT):
---
{instructions}
---

TOOLS: {tools}"""


def _parse_json_response(text: str) -> dict:
    """Parse JSON from Gemini response, with regex fallback."""
    # Try direct parse first
    try:
        return json.loads(text)
    except json.JSONDecodeError:
        pass

    # Regex fallback: extract JSON block
    match = re.search(r'\{[\s\S]*\}', text)
    if match:
        try:
            return json.loads(match.group())
        except json.JSONDecodeError:
            pass

    # Last resort: return raw text as summary
    return {
        "category": "1",
        "category_reasoning": "Could not parse LLM response",
        "discipline": "Back Office including operations",
        "discipline_reasoning": "Default — LLM response unparseable",
        "department": None,
        "is_client_work": False,
        "client_work_reasoning": "",
        "client_name_detected": None,
        "flags": ["parse_error"],
        "summary": text[:500],
        "recommendations": "Manual review required — automated analysis failed to produce structured output",
        "risk_level": "medium"
    }


async def analyze_single_agent(agent_name: str, instructions: str, tools: str,
                                description: str, author: str, max_retries: int = 3) -> dict:
    """Analyse a single agent using Gemini and return structured results.
    Includes retry with exponential backoff for rate limits."""
    client = _get_client()
    model_name = _get_model_name()

    prompt = USER_PROMPT_TEMPLATE.format(
        name=agent_name,
        description=description or "No description provided",
        author=author or "Unknown",
        instructions=instructions or "No instructions available",
        tools=tools or "None specified"
    )

    for attempt in range(max_retries):
        try:
            response = await asyncio.to_thread(
                client.models.generate_content,
                model=model_name,
                contents=prompt,
                config=genai.types.GenerateContentConfig(
                    system_instruction=SYSTEM_INSTRUCTION
                )
            )
            result = _parse_json_response(response.text)
            result["agent_name"] = agent_name
            return result
        except Exception as e:
            error_str = str(e).lower()
            is_rate_limit = "429" in error_str or "rate" in error_str or "quota" in error_str or "resource_exhausted" in error_str
            if is_rate_limit and attempt < max_retries - 1:
                wait_time = (2 ** attempt) * 10  # 10s, 20s, 40s
                logger.warning(f"Rate limited on '{agent_name}', retrying in {wait_time}s (attempt {attempt + 1}/{max_retries})")
                await asyncio.sleep(wait_time)
                continue
            # Final attempt failed or non-rate-limit error
            logger.error(f"Failed to analyse '{agent_name}': {e}")
            e_final = e

    return {
        "error": str(e_final),
        "agent_name": agent_name,
        "category": "1",
        "category_reasoning": f"Analysis failed: {e_final}",
        "discipline": None,
        "discipline_reasoning": f"Analysis failed: {e_final}",
        "department": None,
        "is_client_work": False,
        "client_work_reasoning": "",
        "client_name_detected": None,
        "flags": ["analysis_error"],
        "summary": f"Automated analysis failed: {e_final}",
        "recommendations": "Manual review required",
        "risk_level": "medium"
    }


async def store_audit_result(agent_id: str, audit_data: dict):
    """Store audit results on the agent document and in audit_history."""
    now = datetime.utcnow()

    # Update agent document with audit fields
    update_fields = {
        "audit_status": "flagged",
        "audit_date": now.isoformat(),
        "audit_category": audit_data.get("category"),
        "audit_risk_level": audit_data.get("risk_level"),
        "audit_summary": audit_data.get("summary"),
        "audit_flags": audit_data.get("flags", []),
        "audit_recommendations": audit_data.get("recommendations"),
        "audit_category_reasoning": audit_data.get("category_reasoning"),
        "audit_discipline": audit_data.get("discipline"),
        "audit_discipline_reasoning": audit_data.get("discipline_reasoning"),
        "audit_department": audit_data.get("department"),
        "audit_is_client_work": audit_data.get("is_client_work", False),
        "audit_client_work_reasoning": audit_data.get("client_work_reasoning", ""),
        "audit_client_name_detected": audit_data.get("client_name_detected"),
    }

    await agents_collection.update_one(
        {"_id": ObjectId(agent_id)},
        {"$set": update_fields}
    )

    # Insert into audit_history for historical tracking
    history_entry = {
        "agent_id": agent_id,
        "audit_date": now,
        **audit_data
    }
    await audit_history_collection.insert_one(history_entry)


async def apply_classification_fields(agent_id: str, audit_data: dict):
    """Apply discipline, department, and client detection to the agent document.

    Only overwrites fields that are currently empty/null, to respect manual edits.
    """
    agent = await agents_collection.find_one({"_id": ObjectId(agent_id)})
    if not agent:
        return

    update_fields = {}

    # Auto-assign discipline if not already set
    discipline = audit_data.get("discipline")
    if discipline and not agent.get("discipline"):
        update_fields["discipline"] = discipline

    # Auto-assign department if not already set
    department = audit_data.get("department")
    if department and not agent.get("agent_department"):
        update_fields["agent_department"] = department

    # Auto-flag client work if not already manually set
    is_client_work = audit_data.get("is_client_work", False)
    if is_client_work and not agent.get("client"):
        update_fields["client"] = "yes"
        update_fields["verification_status"] = "needs_verification"
        detected_name = audit_data.get("client_name_detected")
        if detected_name and not agent.get("client_name"):
            update_fields["client_name"] = detected_name

    if update_fields:
        await agents_collection.update_one(
            {"_id": ObjectId(agent_id)},
            {"$set": update_fields}
        )


async def classify_single_agent(agent_id: str):
    """Convenience function for post-sync: load agent, analyse, store, apply fields."""
    if not is_gemini_configured():
        return

    try:
        agent = await agents_collection.find_one({"_id": ObjectId(agent_id)})
        if not agent:
            logger.warning(f"Agent {agent_id} not found")
            return

        instructions = agent.get("instructions")
        if not instructions:
            logger.info(f"Agent '{agent.get('agent_name')}' has no instructions, skipping")
            return

        # Build tools string from agent metadata
        tools = ", ".join(agent.get("agent_capabilities", []) or [])

        audit_data = await analyze_single_agent(
            agent_name=agent.get("agent_name", "Unknown"),
            instructions=instructions,
            tools=tools,
            description=agent.get("agent_description", ""),
            author=agent.get("created_by", "")
        )

        await store_audit_result(agent_id, audit_data)
        await apply_classification_fields(agent_id, audit_data)

        logger.info(f"Classified '{agent.get('agent_name')}' as Cat {audit_data.get('category')}, "
                    f"discipline={audit_data.get('discipline')}, client={audit_data.get('is_client_work')}")

    except Exception as e:
        logger.error(f"Failed to classify agent {agent_id}: {e}")


async def run_audit_batch(unclassified_only: bool = False, single_agent_id: str = None,
                          concurrency: int = None) -> dict:
    """Run audit on multiple agents with concurrency control.

    Returns summary dict with counts and results.
    """
    if not is_gemini_configured():
        return {"error": "GOOGLE_API_KEY not configured"}

    if concurrency is None:
        concurrency = int(os.getenv("AUDIT_CONCURRENCY", "2"))

    # Build query
    query = {}
    if single_agent_id:
        query["_id"] = ObjectId(single_agent_id)
    elif unclassified_only:
        query["audit_status"] = {"$exists": False}

    agents = await agents_collection.find(query).to_list(length=None)

    total = len(agents)
    audited = 0
    failed = 0
    skipped = 0
    results = []

    # Filter out agents without instructions first
    agents_with_instructions = []
    for agent in agents:
        if not agent.get("instructions"):
            skipped += 1
        else:
            agents_with_instructions.append(agent)

    logger.info(f"Audit batch: {total} total, {len(agents_with_instructions)} with instructions, {skipped} skipped")

    # Process agents sequentially in small batches to respect Gemini rate limits
    batch_size = concurrency
    for i in range(0, len(agents_with_instructions), batch_size):
        batch = agents_with_instructions[i:i + batch_size]

        for agent in batch:
            agent_id = str(agent["_id"])
            tools = ", ".join(agent.get("agent_capabilities", []) or [])

            try:
                audit_data = await analyze_single_agent(
                    agent_name=agent.get("agent_name", "Unknown"),
                    instructions=agent.get("instructions"),
                    tools=tools,
                    description=agent.get("agent_description", ""),
                    author=agent.get("created_by", "")
                )

                if audit_data.get("error"):
                    failed += 1
                    results.append({"agent_name": agent.get("agent_name"), "error": audit_data["error"]})
                else:
                    await store_audit_result(agent_id, audit_data)
                    await apply_classification_fields(agent_id, audit_data)
                    audited += 1
                    results.append({
                        "agent_name": agent.get("agent_name"),
                        "category": audit_data.get("category"),
                        "discipline": audit_data.get("discipline"),
                        "risk_level": audit_data.get("risk_level"),
                        "is_client_work": audit_data.get("is_client_work")
                    })
                    logger.info(f"[{audited + failed}/{len(agents_with_instructions)}] "
                                f"Classified '{agent.get('agent_name')}' -> Cat {audit_data.get('category')}")

            except Exception as e:
                failed += 1
                results.append({"agent_name": agent.get("agent_name"), "error": str(e)})
                logger.error(f"[{audited + failed}/{len(agents_with_instructions)}] "
                             f"Failed '{agent.get('agent_name')}': {e}")

        # Pause between batches to avoid rate limits (4 seconds per batch)
        if i + batch_size < len(agents_with_instructions):
            await asyncio.sleep(4)

    logger.info(f"Audit batch complete: {audited} audited, {failed} failed, {skipped} skipped")

    return {
        "status": "completed",
        "total": total,
        "audited_count": audited,
        "failed_count": failed,
        "skipped_count": skipped,
        "results_summary": results[:50]
    }


async def get_all_audit_results() -> list:
    """Get all agents with audit data for the audit tab."""
    agents = await agents_collection.find(
        {},
        {
            "agent_name": 1, "audit_status": 1, "audit_date": 1,
            "audit_category": 1, "audit_risk_level": 1, "audit_summary": 1,
            "audit_flags": 1, "audit_recommendations": 1,
            "audit_category_reasoning": 1,
            "audit_discipline": 1, "audit_discipline_reasoning": 1,
            "audit_department": 1,
            "audit_is_client_work": 1, "audit_client_work_reasoning": 1,
            "audit_client_name_detected": 1,
            "discipline": 1, "agent_department": 1,
            "client": 1, "client_name": 1, "verification_status": 1,
            "instructions": 1, "agent_capabilities": 1,
            "created_by": 1, "agent_description": 1,
            "audit_reviewer": 1, "audit_reviewer_notes": 1,
            "audit_reviewed_date": 1
        }
    ).to_list(length=None)

    # Convert ObjectId to string
    for agent in agents:
        agent["_id"] = str(agent["_id"])

    return agents


async def update_audit_review(agent_id: str, status: str, notes: str, reviewer_email: str):
    """Mark an agent's audit as reviewed/cleared."""
    now = datetime.utcnow()

    await agents_collection.update_one(
        {"_id": ObjectId(agent_id)},
        {"$set": {
            "audit_status": status,
            "audit_reviewer": reviewer_email,
            "audit_reviewer_notes": notes,
            "audit_reviewed_date": now.isoformat()
        }}
    )

    # Also record in history
    await audit_history_collection.insert_one({
        "agent_id": agent_id,
        "audit_date": now,
        "action": "review",
        "status": status,
        "reviewer": reviewer_email,
        "notes": notes
    })