utils.py

# utils.py
import os
import json
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, Optional
from sqlalchemy import create_engine, Engine
from sqlalchemy.orm import sessionmaker
from sqlalchemy.ext.declarative import declarative_base
from contextlib import contextmanager
from langchain_core.messages import (
    SystemMessage,
    HumanMessage, 
)
from langgraph.graph.state import CompiledStateGraph


Base = declarative_base()

# Logging configuration
LOGS_DIR = Path("logs")
LOGS_DIR.mkdir(exist_ok=True)
AGENT_LOG_FILE = LOGS_DIR / "agent_decisions.jsonl"

def reset_db(db_path: str, echo: bool = True):
    """Drops the existing udahub.db file and recreates all tables."""

    # Remove the file if it exists
    if os.path.exists(db_path):
        os.remove(db_path)
        print(f"✅ Removed existing {db_path}")

    # Create a new engine and recreate tables
    engine = create_engine(f"sqlite:///{db_path}", echo=echo)
    Base.metadata.create_all(engine)
    print(f"✅ Recreated {db_path} with fresh schema")


@contextmanager
def get_session(engine: Engine):
    Session = sessionmaker(bind=engine)
    session = Session()
    try:
        yield session
        session.commit()
    except:
        session.rollback()
        raise
    finally:
        session.close()


def model_to_dict(instance):
    """Convert a SQLAlchemy model instance to a dictionary."""
    return {
        column.name: getattr(instance, column.name)
        for column in instance.__table__.columns
    }


def log_decision(
    ticket_id: str,
    agent: str,
    action: str,
    input_data: Optional[Dict[str, Any]] = None,
    output_data: Optional[Dict[str, Any]] = None,
    confidence: Optional[float] = None,
    duration_ms: Optional[float] = None,
    error: Optional[str] = None,
    metadata: Optional[Dict[str, Any]] = None
) -> None:
    """
    Log agent decisions and actions to structured JSONL file.
    
    Args:
        ticket_id: Unique ticket identifier
        agent: Agent name (classifier, supervisor, knowledge, action, escalation)
        action: Action performed (invoked, completed, failed, routing, tool_call)
        input_data: Input parameters/context
        output_data: Output results
        confidence: Confidence score (0-1)
        duration_ms: Execution time in milliseconds
        error: Error message if action failed
        metadata: Additional metadata (tool names, reasoning, etc.)
    """
    log_entry = {
        "timestamp": datetime.utcnow().isoformat() + "Z",
        "ticket_id": ticket_id,
        "agent": agent,
        "action": action,
        "input": input_data or {},
        "output": output_data or {},
    }
    
    if confidence is not None:
        log_entry["confidence"] = confidence
    
    if duration_ms is not None:
        log_entry["duration_ms"] = duration_ms
    
    if error:
        log_entry["error"] = error
    
    if metadata:
        log_entry["metadata"] = metadata
    
    # Write to JSONL file
    with open(AGENT_LOG_FILE, "a", encoding="utf-8") as f:
        f.write(json.dumps(log_entry) + "\n")


def log_tool_usage(
    ticket_id: str,
    agent: str,
    tool_name: str,
    parameters: Dict[str, Any],
    result: Any,
    success: bool = True,
    duration_ms: Optional[float] = None
) -> None:
    """
    Log tool invocation details.
    
    Args:
        ticket_id: Unique ticket identifier
        agent: Agent using the tool
        tool_name: Name of the tool/function called
        parameters: Tool input parameters
        result: Tool execution result
        success: Whether tool execution succeeded
        duration_ms: Execution time in milliseconds
    """
    log_decision(
        ticket_id=ticket_id,
        agent=agent,
        action="tool_call",
        input_data={"tool_name": tool_name, "parameters": parameters},
        output_data={"result": result, "success": success},
        duration_ms=duration_ms,
        metadata={"tool_name": tool_name}
    )


def log_escalation(
    ticket_id: str,
    reason: str,
    summary: str,
    context: Dict[str, Any]
) -> None:
    """
    Log escalation events with full context.
    
    Args:
        ticket_id: Unique ticket identifier
        reason: Reason for escalation
        summary: Escalation summary
        context: Additional escalation context
    """
    log_decision(
        ticket_id=ticket_id,
        agent="escalation_agent",
        action="escalated",
        input_data={"reason": reason},
        output_data={"summary": summary},
        metadata=context
    )


def log_routing_decision(
    ticket_id: str,
    from_agent: str,
    to_agent: str,
    reasoning: str,
    confidence: Optional[float] = None
) -> None:
    """
    Log supervisor routing decisions.
    
    Args:
        ticket_id: Unique ticket identifier
        from_agent: Agent making the routing decision
        to_agent: Target agent to route to
        reasoning: Explanation for routing decision
        confidence: Confidence in the routing decision
    """
    log_decision(
        ticket_id=ticket_id,
        agent=from_agent,
        action="routing",
        output_data={"next_agent": to_agent, "reasoning": reasoning},
        confidence=confidence,
        metadata={"routing": f"{from_agent} -> {to_agent}"}
    )


def get_logs_for_ticket(ticket_id: str) -> list[Dict[str, Any]]:
    """
    Retrieve all logs for a specific ticket.
    
    Args:
        ticket_id: Unique ticket identifier
        
    Returns:
        List of log entries for the ticket
    """
    if not AGENT_LOG_FILE.exists():
        return []
    
    logs = []
    with open(AGENT_LOG_FILE, "r", encoding="utf-8") as f:
        for line in f:
            try:
                entry = json.loads(line.strip())
                if entry.get("ticket_id") == ticket_id:
                    logs.append(entry)
            except json.JSONDecodeError:
                continue
    
    return logs


def chat_interface(agent:CompiledStateGraph, ticket_id:str):
    is_first_iteration = False
    messages = [SystemMessage(content = f"ThreadId: {ticket_id}")]
    while True:
        user_input = input("User: ")
        print("User:", user_input)
        if user_input.lower() in ["quit", "exit", "q"]:
            print("Assistant: Goodbye!")
            break
        messages = [HumanMessage(content=user_input)]
        if is_first_iteration:
            messages.append(HumanMessage(content=user_input))
        trigger = {
            "messages": messages
        }
        config = {
            "configurable": {
                "thread_id": ticket_id,
            }
        }
        
        result = agent.invoke(input=trigger, config=config)
        print("Assistant:", result["messages"][-1].content)
        is_first_iteration = False


# Structured logging functions
import json
import os
from datetime import datetime
from typing import Optional, Any

# Log file paths
LOGS_DIR = os.path.join(os.path.dirname(__file__), "logs")
DECISIONS_LOG = os.path.join(LOGS_DIR, "agent_decisions.jsonl")
TOOLS_LOG = os.path.join(LOGS_DIR, "tool_usage.jsonl")
ROUTING_LOG = os.path.join(LOGS_DIR, "routing_decisions.jsonl")
ESCALATIONS_LOG = os.path.join(LOGS_DIR, "escalations.jsonl")

# Ensure logs directory exists
os.makedirs(LOGS_DIR, exist_ok=True)


def log_decision(
    ticket_id: str,
    agent: str,
    action: str,
    input_data: Optional[dict] = None,
    output_data: Optional[dict] = None,
    confidence: Optional[float] = None,
    duration_ms: Optional[float] = None,
    error: Optional[str] = None,
    metadata: Optional[dict] = None
):
    """
    Log agent decision to structured JSONL file.
    
    Args:
        ticket_id: Ticket identifier
        agent: Agent name (classifier, supervisor, knowledge_agent, etc.)
        action: Action taken (invoked, completed, failed, etc.)
        input_data: Input data for the decision
        output_data: Output data from the decision
        confidence: Confidence score if applicable
        duration_ms: Duration in milliseconds
        error: Error message if failed
        metadata: Additional metadata
    """
    log_entry = {
        "timestamp": datetime.now().isoformat(),
        "ticket_id": ticket_id,
        "agent": agent,
        "action": action,
        "input": input_data or {},
        "output": output_data or {},
        "confidence": confidence,
        "duration_ms": duration_ms,
        "error": error,
        "metadata": metadata or {}
    }
    
    try:
        with open(DECISIONS_LOG, 'a') as f:
            f.write(json.dumps(log_entry) + '\n')
    except Exception as e:
        print(f"Error writing to decisions log: {e}")


def log_tool_usage(
    ticket_id: str,
    agent: str,
    tool_name: str,
    parameters: dict,
    result: Any,
    success: bool,
    duration_ms: Optional[float] = None,
    error: Optional[str] = None
):
    """
    Log tool invocation to structured JSONL file.
    
    Args:
        ticket_id: Ticket identifier
        agent: Agent that invoked the tool
        tool_name: Name of the tool
        parameters: Tool parameters
        result: Tool result
        success: Whether tool call succeeded
        duration_ms: Duration in milliseconds
        error: Error message if failed
    """
    log_entry = {
        "timestamp": datetime.now().isoformat(),
        "ticket_id": ticket_id,
        "agent": agent,
        "tool_name": tool_name,
        "parameters": parameters,
        "result": result if isinstance(result, (dict, list, str, int, float, bool)) else str(result),
        "success": success,
        "duration_ms": duration_ms,
        "error": error
    }
    
    try:
        with open(TOOLS_LOG, 'a') as f:
            f.write(json.dumps(log_entry) + '\n')
    except Exception as e:
        print(f"Error writing to tools log: {e}")


def log_routing(
    ticket_id: str,
    from_agent: str,
    to_agent: str,
    reasoning: str,
    classification: Optional[dict] = None,
    confidence: Optional[float] = None,
    metadata: Optional[dict] = None
):
    """
    Log routing decision to structured JSONL file.
    
    Args:
        ticket_id: Ticket identifier
        from_agent: Agent making routing decision
        to_agent: Target agent
        reasoning: Reasoning for routing decision
        classification: Ticket classification if available
        confidence: Confidence in routing decision
        metadata: Additional metadata
    """
    log_entry = {
        "timestamp": datetime.now().isoformat(),
        "ticket_id": ticket_id,
        "from_agent": from_agent,
        "to_agent": to_agent,
        "reasoning": reasoning,
        "classification": classification or {},
        "confidence": confidence,
        "metadata": metadata or {}
    }
    
    try:
        with open(ROUTING_LOG, 'a') as f:
            f.write(json.dumps(log_entry) + '\n')
    except Exception as e:
        print(f"Error writing to routing log: {e}")


def log_escalation(
    ticket_id: str,
    reason: str,
    summary: str,
    classification: dict,
    attempted_actions: list,
    confidence: float,
    metadata: Optional[dict] = None
):
    """
    Log escalation to structured JSONL file.
    
    Args:
        ticket_id: Ticket identifier
        reason: Escalation reason
        summary: Full escalation summary
        classification: Ticket classification
        attempted_actions: List of actions attempted before escalation
        confidence: Final confidence score
        metadata: Additional metadata
    """
    log_entry = {
        "timestamp": datetime.now().isoformat(),
        "ticket_id": ticket_id,
        "reason": reason,
        "summary": summary,
        "classification": classification,
        "attempted_actions": attempted_actions,
        "confidence": confidence,
        "metadata": metadata or {}
    }
    
    try:
        with open(ESCALATIONS_LOG, 'a') as f:
            f.write(json.dumps(log_entry) + '\n')
    except Exception as e:
        print(f"Error writing to escalations log: {e}")


def search_logs(
    log_type: str = "decisions",
    ticket_id: Optional[str] = None,
    agent: Optional[str] = None,
    start_time: Optional[str] = None,
    end_time: Optional[str] = None
) -> list[dict]:
    """
    Search structured logs with filters.
    
    Args:
        log_type: Type of log (decisions, tools, routing, escalations)
        ticket_id: Filter by ticket ID
        agent: Filter by agent name
        start_time: Filter by start time (ISO format)
        end_time: Filter by end time (ISO format)
        
    Returns:
        List of matching log entries
    """
    log_file_map = {
        "decisions": DECISIONS_LOG,
        "tools": TOOLS_LOG,
        "routing": ROUTING_LOG,
        "escalations": ESCALATIONS_LOG
    }
    
    log_file = log_file_map.get(log_type, DECISIONS_LOG)
    
    if not os.path.exists(log_file):
        return []
    
    results = []
    
    try:
        with open(log_file, 'r') as f:
            for line in f:
                entry = json.loads(line.strip())
                
                # Apply filters
                if ticket_id and entry.get("ticket_id") != ticket_id:
                    continue
                if agent and entry.get("agent") != agent:
                    continue
                if start_time and entry.get("timestamp") < start_time:
                    continue
                if end_time and entry.get("timestamp") > end_time:
                    continue
                
                results.append(entry)
                
    except Exception as e:
        print(f"Error searching logs: {e}")
    
    return results


def get_ticket_log_summary(ticket_id: str, latest_only: bool = False) -> dict:
    """
    Get comprehensive log summary for a ticket.
    
    Args:
        ticket_id: Ticket identifier
        latest_only: If True, return only the most recent run
        
    Returns:
        Dictionary with all logs for the ticket
    """
    summary = {
        "ticket_id": ticket_id,
        "decisions": search_logs("decisions", ticket_id=ticket_id),
        "tools": search_logs("tools", ticket_id=ticket_id),
        "routing": search_logs("routing", ticket_id=ticket_id),
        "escalations": search_logs("escalations", ticket_id=ticket_id)
    }
    
    if latest_only and summary["decisions"]:
        # Find the latest classifier invocation timestamp
        classifier_invocations = [
            entry for entry in summary["decisions"]
            if entry.get("agent") == "classifier" and entry.get("action") == "invoked"
        ]
        
        if classifier_invocations:
            # Get the most recent classifier invocation
            latest_start = max(entry["timestamp"] for entry in classifier_invocations)
            
            # Filter all logs to only those after this timestamp
            summary["decisions"] = [e for e in summary["decisions"] if e["timestamp"] >= latest_start]
            summary["tools"] = [e for e in summary["tools"] if e["timestamp"] >= latest_start]
            summary["routing"] = [e for e in summary["routing"] if e["timestamp"] >= latest_start]
            summary["escalations"] = [e for e in summary["escalations"] if e["timestamp"] >= latest_start]
    
    return summary