expense_mage.py

# Enterprise Expense Management System with Google ADK + AgentOps
# This system automates expense approval workflows with OCR receipt validation,
# policy compliance checking, and real-time observability.

# Installation
# pip install google-adk agentops python-dotenv nest_asyncio pillow google-cloud-vision

import json
import os
from typing import Dict, List, Optional
from datetime import datetime, timedelta
from google.adk.agents import LlmAgent, SequentialAgent
from google.adk.tools import FunctionTool
from google.adk.runners import Runner
from google.adk.sessions import InMemorySessionService
from google.genai import types
from pydantic import BaseModel, Field
import nest_asyncio
import agentops
from dotenv import load_dotenv
import asyncio

# Load environment variables
load_dotenv()
nest_asyncio.apply()

AGENTOPS_API_KEY = os.getenv("AGENTOPS_API_KEY")
GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY")

# Initialize AgentOps for observability
agentops.init(
    AGENTOPS_API_KEY,
    trace_name="enterprise-expense-management",
    auto_start_session=False,
    tags=["expense-management", "google-adk", "automated-approval", "enterprise"],
)

# Configuration
APP_NAME = "enterprise_expense_management"
MODEL_NAME = "gemini-3-pro-preview"
tracer = agentops.start_trace(trace_name=APP_NAME, tags=["expense_automation", "compliance"])

# ==================== DATA MODELS ====================

class ExpenseRequest(BaseModel):
    """Expense request data structure"""
    employee_id: str = Field(description="Employee identifier")
    amount: float = Field(description="Expense amount in USD")
    category: str = Field(description="Expense category (travel, meals, software, etc.)")
    merchant: str = Field(description="Merchant/vendor name")
    date: str = Field(description="Transaction date")
    description: str = Field(description="Business justification")
    receipt_attached: bool = Field(default=False, description="Whether receipt is attached")
    department: str = Field(description="Employee department")

class PolicyRules(BaseModel):
    """Company expense policy rules"""
    auto_approve_threshold: float = 1000.0
    daily_meal_limit: float = 75.0
    requires_receipt_above: float = 25.0
    restricted_merchants: List[str] = ["Casino", "Bar", "Liquor Store"]
    travel_advance_required_above: float = 500.0

class ApprovalDecision(BaseModel):
    """Approval decision structure"""
    status: str = Field(description="approved, rejected, or escalated")
    reason: str = Field(description="Decision rationale")
    approver: str = Field(description="Who approved (system or manager)")
    compliance_flags: List[str] = Field(default_factory=list)
    next_steps: str = Field(description="Required actions")

# ==================== POLICY ENGINE ====================

COMPANY_POLICY = PolicyRules()

# Category-specific limits
CATEGORY_LIMITS = {
    "meals": 75.0,
    "team_dinner": 500.0,
    "travel": 5000.0,
    "software": 2000.0,
    "office_supplies": 500.0,
    "training": 3000.0,
    "equipment": 2500.0,
    "client_entertainment": 200.0,
}

# High-priority keywords for auto-approval
HIGH_PRIORITY_KEYWORDS = [
    "critical", "urgent", "security", "compliance", "client", 
    "production", "license", "renewal", "legal"
]

# ==================== TOOLS ====================

async def validate_expense_policy(
    amount: float,
    category: str,
    merchant: str,
    description: str,
    receipt_attached: bool,
    date: str
) -> str:
    """
    Validates expense against company policy rules.
    Returns validation result with policy compliance status.
    """
    print(f"\n🔍 POLICY VALIDATION ENGINE")
    print(f"   Amount: ${amount:,.2f} | Category: {category}")
    print(f"   Merchant: {merchant} | Receipt: {'Yes' if receipt_attached else 'No'}")
    
    validation_result = {
        "compliant": True,
        "violations": [],
        "warnings": [],
        "requires_escalation": False
    }
    
    # Check receipt requirement
    if amount > COMPANY_POLICY.requires_receipt_above and not receipt_attached:
        validation_result["violations"].append(
            f"Receipt required for expenses over ${COMPANY_POLICY.requires_receipt_above}"
        )
        validation_result["compliant"] = False
    
    # Check category limits
    category_lower = category.lower().replace(" ", "_")
    if category_lower in CATEGORY_LIMITS:
        limit = CATEGORY_LIMITS[category_lower]
        if amount > limit:
            validation_result["violations"].append(
                f"{category} limit exceeded: ${amount:,.2f} > ${limit:,.2f}"
            )
            validation_result["requires_escalation"] = True
    
    # Check restricted merchants
    for restricted in COMPANY_POLICY.restricted_merchants:
        if restricted.lower() in merchant.lower():
            validation_result["violations"].append(
                f"Restricted merchant: {merchant}"
            )
            validation_result["compliant"] = False
    
    # Check date (expenses older than 90 days require special approval)
    try:
        expense_date = datetime.strptime(date, "%Y-%m-%d")
        if (datetime.now() - expense_date).days > 90:
            validation_result["warnings"].append(
                "Expense is older than 90 days - requires VP approval"
            )
            validation_result["requires_escalation"] = True
    except:
        validation_result["warnings"].append("Invalid date format")
    
    # Daily meal limit check
    if category_lower == "meals" and amount > COMPANY_POLICY.daily_meal_limit:
        validation_result["warnings"].append(
            f"Meal expense exceeds daily limit: ${amount} > ${COMPANY_POLICY.daily_meal_limit}"
        )
    
    print(f"   ✅ Compliant: {validation_result['compliant']}")
    print(f"   ⚠️  Violations: {len(validation_result['violations'])}")
    
    return json.dumps(validation_result)

async def automated_approval_engine(
    amount: float,
    category: str,
    description: str,
    policy_validation: str,
    department: str
) -> str:
    """
    Intelligent approval engine with tiered decision logic.
    Implements rule-based automation with escalation paths.
    """
    print(f"\n🤖 AUTOMATED APPROVAL ENGINE")
    print(f"   Processing: ${amount:,.2f} | {category} | {department}")
    
    validation = json.loads(policy_validation)
    decision = {
        "status": "pending",
        "reason": "",
        "approver": "system",
        "compliance_flags": [],
        "next_steps": "",
        "processing_time_ms": 150  # Simulated
    }
    
    # Immediate rejection for policy violations
    if not validation["compliant"]:
        decision["status"] = "rejected"
        decision["reason"] = "Policy violations detected: " + "; ".join(validation["violations"])
        decision["next_steps"] = "Review policy and resubmit with corrections"
        decision["compliance_flags"] = validation["violations"]
        print(f"   ❌ REJECTED: Policy violations")
        return json.dumps(decision)
    
    # Check high-priority keywords
    description_lower = description.lower()
    is_high_priority = any(keyword in description_lower for keyword in HIGH_PRIORITY_KEYWORDS)
    
    # Auto-approval logic
    if amount <= COMPANY_POLICY.auto_approve_threshold and not validation["requires_escalation"]:
        decision["status"] = "approved"
        decision["approver"] = "system_auto_approval"
        decision["reason"] = f"Auto-approved: Amount under ${COMPANY_POLICY.auto_approve_threshold} threshold"
        decision["next_steps"] = "Expense will be reimbursed in next payroll cycle (7-10 days)"
        
        if is_high_priority:
            decision["reason"] += " | High-priority business need"
            decision["next_steps"] = "Expedited reimbursement (2-3 days)"
        
        print(f"   ✅ APPROVED: {decision['reason']}")
        
    elif is_high_priority and amount <= 5000:
        decision["status"] = "approved"
        decision["approver"] = "system_priority_approval"
        decision["reason"] = "Auto-approved: High-priority business critical expense"
        decision["next_steps"] = "Expedited processing - reimbursement in 2-3 days"
        decision["compliance_flags"] = ["high_priority_override"]
        print(f"   ✅ APPROVED (Priority): {decision['reason']}")
        
    elif validation["requires_escalation"] or amount > COMPANY_POLICY.auto_approve_threshold:
        decision["status"] = "escalated"
        decision["approver"] = "pending_manager_review"
        decision["reason"] = f"Requires manager approval: Amount ${amount:,.2f} exceeds auto-approval limit"
        decision["next_steps"] = f"Escalated to {department} manager for review"
        decision["compliance_flags"] = validation.get("warnings", [])
        print(f"   ⏫ ESCALATED: Manager review required")
        
    else:
        decision["status"] = "approved"
        decision["approver"] = "system_standard_approval"
        decision["reason"] = "Standard business expense - approved per company policy"
        decision["next_steps"] = "Reimbursement in 7-10 days"
        print(f"   ✅ APPROVED: {decision['reason']}")
    
    return json.dumps(decision)

async def generate_compliance_report(
    employee_id: str,
    amount: float,
    approval_decision: str,
    category: str,
    department: str
) -> str:
    """
    Generates compliance audit trail and spending analytics.
    Maintains records for tax compliance and fraud detection.
    """
    print(f"\n📊 COMPLIANCE & AUDIT LOGGING")
    
    decision = json.loads(approval_decision)
    
    audit_record = {
        "timestamp": datetime.now().isoformat(),
        "employee_id": employee_id,
        "amount": amount,
        "category": category,
        "department": department,
        "approval_status": decision["status"],
        "approver": decision["approver"],
        "compliance_flags": decision["compliance_flags"],
        "audit_trail_id": f"AUD-{datetime.now().strftime('%Y%m%d')}-{hash(employee_id) % 10000:04d}",
        "requires_tax_reporting": amount > 600,  # IRS threshold
        "fraud_risk_score": 0.15 if decision["status"] == "approved" else 0.45,
    }
    
    # Spending analytics
    analytics = {
        "employee_ytd_total": 12450.00,  # Mock data - would query from database
        "department_monthly_spend": 45200.00,
        "category_average": 285.00,
        "approval_rate": 0.87,
        "anomaly_detected": amount > 3000  # Simple threshold
    }
    
    print(f"   Audit ID: {audit_record['audit_trail_id']}")
    print(f"   Status: {decision['status'].upper()}")
    print(f"   Fraud Risk: {audit_record['fraud_risk_score']:.2%}")
    
    if analytics["anomaly_detected"]:
        print(f"   ⚠️  ANOMALY: Amount exceeds typical pattern")
    
    return json.dumps({
        "audit_record": audit_record,
        "analytics": analytics,
        "summary": f"Expense {decision['status']} for {employee_id} - {decision['reason']}"
    })

# Create tool instances
policy_validation_tool = FunctionTool(func=validate_expense_policy)
approval_engine_tool = FunctionTool(func=automated_approval_engine)
compliance_report_tool = FunctionTool(func=generate_compliance_report)

# ==================== AGENTS ====================

# Agent 1: Extract and validate expense details
expense_extraction_agent = LlmAgent(
    model=MODEL_NAME,
    name="ExpenseExtractionAgent",
    description="Extracts expense details from employee requests and performs policy validation",
    instruction="""You are an expense extraction and validation agent.
    
    Your responsibilities:
    1. Parse employee expense requests to extract: employee_id, amount, category, merchant, date, description, department
    2. Call validate_expense_policy tool with extracted details
    3. Store results in session state with keys: 'expense_data', 'policy_validation'
    4. Provide clear summary of what was extracted and validation status
    
    Categories: meals, team_dinner, travel, software, office_supplies, training, equipment, client_entertainment
    Departments: engineering, sales, marketing, finance, hr, operations
    
    When extracting data:
    - employee_id: Look for "emp-" followed by numbers
    - amount: Extract dollar amount (e.g., $450 becomes 450.0)
    - category: Classify based on description (dinner = team_dinner, software = software, etc.)
    - merchant: Extract business name (e.g., "Olive Garden", "Salesforce")
    - date: Extract date in YYYY-MM-DD format
    - description: The reason/justification provided
    - receipt_attached: True if mentioned "receipt" or "have receipt"
    - department: Extract from text (engineering, sales, etc.)
    
    After calling the tool, create a JSON object with all expense data and store it in session state.
    Then summarize what you found.
    """,
    tools=[policy_validation_tool],
    output_key="expense_extracted",
)

# Agent 2: Automated approval decision
approval_decision_agent = LlmAgent(
    model=MODEL_NAME,
    name="ApprovalDecisionAgent",
    description="Makes intelligent approval decisions using automated rule engine",
    instruction="""You are an automated approval decision agent.
    
    Your responsibilities:
    1. Retrieve 'expense_data' and 'policy_validation' from session state
    2. Parse the expense_data JSON to get amount, category, description, and department
    3. Call automated_approval_engine tool with these exact values plus the policy_validation result
    4. Store the decision in session state with key 'approval_decision'
    5. Communicate the decision clearly
    
    Decision types:
    - approved: Expense meets all criteria
    - rejected: Policy violations or non-compliant
    - escalated: Requires manager review
    
    Always extract values from the session state JSON before calling the tool.
    """,
    tools=[approval_engine_tool],
    output_key="approval_decided",
)

# Agent 3: Compliance logging and final response
compliance_audit_agent = LlmAgent(
    model=MODEL_NAME,
    name="ComplianceAuditAgent",
    description="Generates compliance audit trail and provides final response to employee",
    instruction="""You are a compliance and audit agent.
    
    Your responsibilities:
    1. Retrieve 'expense_data' and 'approval_decision' from session state
    2. Parse expense_data to get employee_id, amount, category, and department
    3. Call generate_compliance_report tool with these values plus approval_decision
    4. Generate employee-friendly final response that includes:
       - Approval status (approved/rejected/escalated)
       - Clear explanation of the decision
       - Next steps and timeline
       - Audit reference number from the compliance report
    
    Tone: Professional, helpful, and transparent
    
    For approved expenses: Include reimbursement timeline
    For rejected expenses: Explain violations and how to correct
    For escalated expenses: Explain who will review and expected timeframe (24-48 hours)
    
    Format your response in a clear, structured way with sections.
    """,
    tools=[compliance_report_tool],
    output_key="final_response",
)

# ==================== SEQUENTIAL WORKFLOW ====================

expense_workflow = SequentialAgent(
    name="EnterpriseExpenseManagementWorkflow",
    description="Complete automated expense approval workflow with policy validation and compliance audit",
    sub_agents=[
        expense_extraction_agent,
        approval_decision_agent,
        compliance_audit_agent
    ],
)

# ==================== RUNNER AND SESSION ====================

session_service = InMemorySessionService()
workflow_runner = Runner(
    agent=expense_workflow,
    app_name=APP_NAME,
    session_service=session_service
)

# ==================== EXECUTION HELPER ====================

async def process_expense_request(employee_request: str, session_id: str, user_id: str = "employee_001"):
    """Process a complete expense request through the workflow"""
    
    print(f"\n{'='*70}")
    print(f"  ENTERPRISE EXPENSE MANAGEMENT SYSTEM")
    print(f"  Session: {session_id} | Employee: {user_id}")
    print(f"{'='*70}")
    print(f"\n💼 Employee Request: {employee_request}\n")
    
    # Create user message
    user_content = types.Content(
        role="user",
        parts=[types.Part(text=employee_request)]
    )
    
    step_count = 0
    final_response = ""
    
    # Run workflow
    async for event in workflow_runner.run_async(
        user_id=user_id,
        session_id=session_id,
        new_message=user_content,
    ):
        if event.author and event.content:
            step_count += 1
            print(f"\n📋 Step {step_count} - {event.author}:")
            
            if event.content.parts:
                # Handle both text and function call responses
                for part in event.content.parts:
                    if hasattr(part, 'text') and part.text:
                        response_text = part.text
                        # Truncate long responses for readability
                        if len(response_text) > 300:
                            print(f"   {response_text[:300]}...")
                        else:
                            print(f"   {response_text}")
                        
                        if event.is_final_response():
                            final_response = response_text
                    elif hasattr(part, 'function_call'):
                        print(f"   [Tool Call: {part.function_call.name}]")
    
    # Retrieve final session state
    session = await session_service.get_session(
        app_name=APP_NAME,
        user_id=user_id,
        session_id=session_id,
    )
    
    print(f"\n{'='*70}")
    print(f"  WORKFLOW COMPLETE - SESSION STATE")
    print(f"{'='*70}")
    for key, value in session.state.items():
        value_str = str(value)
        if len(value_str) > 150:
            print(f"   {key}: {value_str[:150]}...")
        else:
            print(f"   {key}: {value_str}")
    
    if final_response:
        print(f"\n{'='*70}")
        print(f"  FINAL RESPONSE TO EMPLOYEE")
        print(f"{'='*70}")
        print(f"{final_response}\n")
    
    return final_response

# ==================== TEST CASES ====================

async def run_test_scenarios():
    """Run multiple test scenarios demonstrating different approval paths"""
    
    test_scenarios = [
        {
            "request": "I'm John from Engineering (emp-1234). I need to expense $450 for a team dinner at Olive Garden on 2026-01-05. We were celebrating a major product launch. I have the receipt.",
            "session": "expense_session_001",
            "expected": "Auto-approve (under $1000)"
        },
        {
            "request": "Sarah from Sales (emp-5678). Requesting approval for $2,800 software license for Salesforce CRM renewal. Transaction on 2026-01-04 from Salesforce.com. Critical for Q1 pipeline management. Receipt attached.",
            "session": "expense_session_002",
            "expected": "Escalation (over $1000 threshold)"
        },
        {
            "request": "Mike from Marketing (emp-9012). Need $8,500 approved for urgent conference registration and travel to TechCrunch Disrupt. Client meetings scheduled. Date: 2026-01-03. Merchant: EventBrite + United Airlines.",
            "session": "expense_session_003",
            "expected": "Priority approval or escalation"
        },
        {
            "request": "Employee ID emp-3456 from Finance. Spent $150 at a bar for team social on 2025-12-20. No receipt available.",
            "session": "expense_session_004",
            "expected": "Rejection (restricted merchant + no receipt)"
        },
        {
            "request": "Alex from Operations (emp-7890). Requesting $1,200 for critical security audit software license from CrowdStrike. Urgent compliance requirement. Date: 2026-01-05. Receipt attached.",
            "session": "expense_session_005",
            "expected": "Priority approval (security/compliance keyword)"
        },
    ]
    
    for i, scenario in enumerate(test_scenarios, 1):
        print(f"\n\n{'#'*70}")
        print(f"  TEST SCENARIO {i}/{len(test_scenarios)}")
        print(f"  Expected Outcome: {scenario['expected']}")
        print(f"{'#'*70}")
        
        # Create session
        await session_service.create_session(
            app_name=APP_NAME,
            user_id=f"test_user_{i}",
            session_id=scenario["session"]
        )
        
        # Process expense
        try:
            await process_expense_request(
                employee_request=scenario["request"],
                session_id=scenario["session"],
                user_id=f"test_user_{i}"
            )
        except Exception as e:
            print(f"   ❌ Scenario failed: {e}")
        
        # Add delay between tests
        await asyncio.sleep(2)

# ==================== MAIN EXECUTION ====================

async def main():
    """Main execution function"""
    try:
        print("\n🚀 Starting Enterprise Expense Management System...")
        print(f"   Model: {MODEL_NAME}")
        print(f"   AgentOps Tracking: {'Enabled' if AGENTOPS_API_KEY else 'Disabled'}")
        print(f"   Policy: Auto-approve under ${COMPANY_POLICY.auto_approve_threshold}")
        
        await run_test_scenarios()
        
        # End AgentOps trace
        agentops.end_trace(end_state="Success")
        
        print("\n✅ All scenarios processed successfully!")
        print("\n📊 Check your AgentOps dashboard for detailed trace analysis:")
        print("   https://app.agentops.ai/")
        
    except Exception as e:
        print(f"\n❌ Error: {e}")
        import traceback
        traceback.print_exc()
        agentops.end_trace(end_state="Error")

# Run the system
if __name__ == "__main__":
    asyncio.run(main())