backend/core/agent_execution_service.py

# -*- coding: utf-8 -*-
"""
Agent Execution Service

Provides centralized agent chat execution with:
- Full governance integration
- WebSocket streaming support
- AgentExecution audit trail
- Episode creation for memory
"""

import logging
import os
import uuid
from datetime import datetime
from typing import Any, Dict, List, Optional

from sqlalchemy.orm import Session

from core.agent_context_resolver import AgentContextResolver
from core.agent_governance_service import AgentGovernanceService
from core.chat_context_manager import get_chat_context_manager
from core.chat_session_manager import get_chat_session_manager
from core.database import get_db_session, SessionLocal
from core.episode_integration import trigger_episode_creation
from core.lancedb_handler import get_chat_history_manager
from core.llm_service import LLMService
from core.models import AgentExecution, AgentInstallation
from core.marketplace_usage_tracker import MarketplaceUsageTracker
from core.personal_budget_service import personal_budget_service
from core.websockets import manager as ws_manager

logger = logging.getLogger(__name__)


class ChatMessage:
    """Simple chat message model"""
    def __init__(self, role: str, content: str):
        self.role = role
        self.content = content


async def execute_agent_chat(
    agent_id: str,
    message: str,
    user_id: str,
    session_id: Optional[str] = None,
    workspace_id: str = "default",
    conversation_history: List[Dict[str, str]] = None,
    stream: bool = False
) -> Dict[str, Any]:
    """
    Execute agent chat with full governance and streaming support.

    This is the centralized service for executing agent chat requests,
    used by menubar, mobile, and web platforms.

    Args:
        agent_id: The ID of the agent to execute
        message: User's message to the agent
        user_id: User ID making the request
        session_id: Optional session ID for conversation continuity
        workspace_id: Workspace ID (default for single-tenant)
        conversation_history: Optional conversation history for context
        stream: Whether to stream response via WebSocket

    Returns:
        Dictionary containing:
        - success: bool
        - execution_id: str
        - response: str (full response if not streaming)
        - agent_id: str
        - agent_name: str
        - message_id: str (for WebSocket tracking)
        - error: str (if failed)

    Example:
        result = await execute_agent_chat(
            agent_id="agent_123",
            message="Hello, how can you help me?",
            user_id="user_456"
        )
        print(result["response"])
    """
    # Feature flags
    governance_enabled = os.getenv("STREAMING_GOVERNANCE_ENABLED", "true").lower() == "true"
    emergency_bypass = os.getenv("EMERGENCY_GOVERNANCE_BYPASS", "false").lower() == "true"

    agent = None
    agent_execution = None
    resolution_context = None
    governance_check = None
    db_session = None

    try:
        # ============================================
        # GOVERNANCE: Agent Resolution & Validation
        # ============================================
        if governance_enabled and not emergency_bypass:
            db_session = SessionLocal()
            resolver = AgentContextResolver(db_session)
            governance = AgentGovernanceService(db_session)

            # Resolve agent for this request
            agent, resolution_context = await resolver.resolve_agent_for_request(
                user_id=user_id,
                session_id=session_id,
                requested_agent_id=agent_id,
                action_type="chat"
            )

            if not agent:
                logger.warning(f"Agent resolution failed for agent_id={agent_id}, using system default")
                # Fall through to system default behavior

            # Perform governance check
            if agent:
                governance_check = governance.can_perform_action(
                    agent_id=agent.id,
                    action_type="chat",
                    require_approval=False
                )

                if not governance_check.get("allowed", False):
                    reason = governance_check.get("reason", "Governance policy denied this action")
                    logger.warning(f"Governance blocked agent chat: {reason}")
                    return {
                        "success": False,
                        "error": f"Action blocked by governance: {reason}",
                        "agent_id": agent_id,
                        "execution_id": None
                    }

        # ============================================
        # BUDGET: Check Budget (Warning Only, No Blocking)
        # ============================================
        # Check budget before execution (warning only, does NOT block)
        # Personal use = user's responsibility, so we only log warnings
        try:
            if personal_budget_service.is_budget_exceeded():
                logger.warning(
                    f"Budget exceeded for agent execution (agent_id={agent_id}). "
                    f"Continuing anyway (personal use = user responsibility)."
                )
                # Send alert at 100% threshold
                personal_budget_service.send_budget_alert(100.0)
            else:
                # Send alerts at 80% and 90% thresholds
                personal_budget_service.send_budget_alert(80.0)
                personal_budget_service.send_budget_alert(90.0)
        except Exception as budget_error:
            logger.error(f"Budget check failed (continuing anyway): {budget_error}")
            # Don't block execution on budget check failures

        # ============================================
        # EXECUTION: Create AgentExecution Record
        # ============================================
        execution_id = str(uuid.uuid4())

        if agent and governance_enabled:
            try:
                agent_execution = AgentExecution(
                    id=execution_id,
                    agent_id=agent.id,
                    agent_name=agent.name,
                    agent_category=agent.category,
                    user_id=user_id,
                    workspace_id=workspace_id,
                    session_id=session_id,
                    action_type="chat",
                    action_complexity=1,
                    status="running",
                    input_data={"message": message},
                    metadata={
                        "source": "menubar",
                        "governance_check": governance_check,
                        "resolution_context": resolution_context
                    }
                )

                if db_session:
                    db_session.add(agent_execution)
                    db_session.commit()
                    db_session.refresh(agent_execution)

            except Exception as exec_error:
                logger.error(f"Failed to create AgentExecution record: {exec_error}")
                # Continue anyway - don't block execution on audit failure

        # ============================================
        # LLM: Initialize LLM Service
        # ============================================
        llm_service = LLMService(tenant_id=workspace_id, db=db_session)

        # Prepare messages for LLM
        messages = []

        # Add system message
        agent_name = agent.name if agent else "ATOM"
        agent_desc = agent.description if agent else "AI Assistant"

        messages.append({
            "role": "system",
            "content": f"""You are {agent_name}, an intelligent AI assistant.

{agent_desc}

Provide helpful, concise responses. Be direct and practical."""
        })

        # Add conversation history
        if conversation_history:
            for hist_msg in conversation_history:
                messages.append({
                    "role": hist_msg.get("role", "user"),
                    "content": hist_msg.get("content", "")
                })

        # Add current message
        messages.append({
            "role": "user",
            "content": message
        })

        # Get optimal provider for this request
        complexity = llm_service.analyze_query_complexity(message, task_type="chat")
        provider_id, model = llm_service.get_optimal_provider(
            complexity,
            task_type="chat",
            prefer_cost=True,
            tenant_plan="free",
            is_managed_service=False,
            requires_tools=False
        )

        logger.info(f"Executing agent chat with {provider_id}/{model}" +
                   (f" (agent: {agent.name})" if agent else ""))

        # Create unique message ID for WebSocket tracking
        message_id = str(uuid.uuid4())

        # If streaming is requested, send initial WebSocket message
        if stream:
            user_channel = f"user:{user_id}"
            await ws_manager.broadcast(user_channel, {
                "type": "streaming:start",
                "id": message_id,
                "model": "auto",
                "agent_id": agent.id if agent else None,
                "agent_name": agent.name if agent else None,
                "execution_id": execution_id
            })

        # Execute chat (streaming or non-streaming)
        accumulated_content = ""
        tokens_count = 0
        start_time = datetime.now()

        stream_kwargs = {
            "messages": messages,
            "model": "auto",
            "temperature": 0.7,
            "max_tokens": 2000,
            "agent_id": agent.id if agent else None
        }

        # Stream response
        # Stream response via LLMService
        async for token in llm_service.stream_completion(**stream_kwargs):
            accumulated_content += token
            tokens_count += 1

            # Broadcast token via WebSocket if streaming enabled
            if stream:
                user_channel = f"user:{user_id}"
                await ws_manager.broadcast(user_channel, {
                    "type": ws_manager.STREAMING_UPDATE,
                    "id": message_id,
                    "delta": token,
                    "complete": False,
                    "metadata": {
                        "tokens_so_far": len(accumulated_content),
                        "execution_id": execution_id
                    }
                })

        # Send completion message if streaming
        if stream:
            user_channel = f"user:{user_id}"
            await ws_manager.broadcast(user_channel, {
                "type": ws_manager.STREAMING_COMPLETE,
                "id": message_id,
                "content": accumulated_content,
                "complete": True,
                "metadata": {
                    "execution_id": execution_id,
                    "tokens_total": tokens_count
                }
            })

        # ============================================
        # PERSISTENCE: Save to Chat History
        # ============================================
        try:
            chat_history = get_chat_history_manager(workspace_id)
            session_manager = get_chat_session_manager(workspace_id)

            # Create or use session
            if not session_id:
                session_id = session_manager.create_session(user_id)

            # Save messages
            chat_history.add_message(session_id, "user", message)
            chat_history.add_message(session_id, "assistant", accumulated_content)

        except Exception as persist_error:
            logger.error(f"Failed to save chat history: {persist_error}")
            # Don't fail the request on persistence errors

        # ============================================
        # GOVERNANCE: Update Execution Record
        # ============================================
        if agent_execution and governance_enabled:
            try:
                end_time = datetime.now()
                duration_ms = (end_time - start_time).total_seconds() * 1000

                agent_execution.status = "completed"
                agent_execution.output_data = {
                    "response": accumulated_content,
                    "tokens": tokens_count,
                    "model": "auto"
                }
                agent_execution.duration_ms = duration_ms
                agent_execution.end_time = end_time

                if db_session:
                    db_session.commit()

                # Marketplace Tracking
                if agent and agent.type == "marketplace":
                    try:
                        installation = db_session.query(AgentInstallation).filter(
                            AgentInstallation.instantiated_agent_id == agent.id
                        ).first()
                        if installation:
                            MarketplaceUsageTracker.track_usage(
                                item_type="agent",
                                item_id=installation.template_id,
                                success=True,
                                duration_ms=duration_ms
                            )
                    except Exception as mt_error:
                        logger.error(f"Marketplace tracking failed: {mt_error}")

            except Exception as update_error:
                logger.error(f"Failed to update AgentExecution record: {update_error}")

        # Trigger episode creation for memory
        try:
            await trigger_episode_creation(
                user_id=user_id,
                agent_id=agent.id if agent else None,
                session_id=session_id,
                workspace_id=workspace_id
            )
        except Exception as episode_error:
            logger.warning(f"Failed to trigger episode creation: {episode_error}")

        # ============================================
        # BUDGET: Track Spend After Execution
        # ============================================
        # Record spend for budget forecasting and tracking
        try:
            # Estimate cost based on tokens (rough estimation)
            # ACU cost: ~$0.0001 per token, API cost varies by provider
            estimated_cost = (tokens_count * 0.0001) + 0.001  # Base API call cost
            personal_budget_service.record_spend(estimated_cost, execution_id)
        except Exception as budget_error:
            logger.error(f"Failed to record spend (non-critical): {budget_error}")
            # Don't fail execution on budget tracking errors

        # Return success
        return {
            "success": True,
            "execution_id": execution_id,
            "response": accumulated_content,
            "agent_id": agent.id if agent else agent_id,
            "agent_name": agent.name if agent else "System",
            "message_id": message_id,
            "session_id": session_id,
            "tokens": tokens_count,
            "model": "auto"
        }

    except Exception as e:
        logger.error(f"Agent chat execution failed: {e}", exc_info=True)

        # Update execution record as failed
        if agent_execution and governance_enabled and db_session:
            try:
                agent_execution.status = "failed"
                agent_execution.error_message = str(e)
                agent_execution.end_time = datetime.now()
                db_session.commit()

                # Marketplace Tracking (Failure)
                if agent and agent.type == "marketplace":
                    try:
                        installation = db_session.query(AgentInstallation).filter(
                            AgentInstallation.instantiated_agent_id == agent.id
                        ).first()
                        if installation:
                            duration_ms = (datetime.now() - start_time).total_seconds() * 1000
                            MarketplaceUsageTracker.track_usage(
                                item_type="agent",
                                item_id=installation.template_id,
                                success=False,
                                duration_ms=duration_ms
                            )
                    except Exception as mt_error:
                        logger.error(f"Marketplace failure tracking failed: {mt_error}")

            except Exception as update_error:
                logger.error(f"Failed to update failed execution record: {update_error}")

        return {
            "success": False,
            "error": str(e),
            "agent_id": agent_id,
            "execution_id": execution_id if agent_execution else None
        }

    finally:
        # Clean up database session
        if db_session:
            try:
                db_session.close()
            except Exception:
                pass


def execute_agent_chat_sync(
    agent_id: str,
    message: str,
    user_id: str,
    session_id: Optional[str] = None,
    workspace_id: str = "default",
    conversation_history: List[Dict[str, str]] = None
) -> Dict[str, Any]:
    """
    Synchronous wrapper for execute_agent_chat.

    Use this in non-async contexts. This runs the async function in an event loop.
    Note: WebSocket streaming is disabled in sync mode.

    Args:
        Same as execute_agent_chat

    Returns:
        Same as execute_agent_chat (but without streaming support)
    """
    import asyncio

    try:
        loop = asyncio.get_event_loop()
    except RuntimeError:
        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)

    return loop.run_until_complete(
        execute_agent_chat(
            agent_id=agent_id,
            message=message,
            user_id=user_id,
            session_id=session_id,
            workspace_id=workspace_id,
            conversation_history=conversation_history,
            stream=False  # Disable streaming in sync mode
        )
    )