cli_runner.py

"""Command-line runner for Agent Lightning evaluation."""

import asyncio
import os
from typing import cast
from dotenv import load_dotenv

import pandas as pd
from rich.console import Console

from agentlightning import LLM, AgentOpsTracer, InMemoryLightningStore, LitAgentRunner
from agent import loan_risk_agent, RiskTask

load_dotenv()

console = Console()


async def main():
    """Main entry point for CLI runner."""
    console.print("[bold cyan]🚀 Loan Default Early Warning System - CLI Mode[/bold cyan]")
    console.print("[dim]Using Agent Lightning Framework[/dim]\n")
    
    # Check API key
    api_key = os.getenv("OPENAI_API_KEY")
    if not api_key:
        console.print("[bold red]❌ Error:[/bold red] OPENAI_API_KEY not found in .env file!")
        return
    
    console.print(f"[green]✓ API Key loaded[/green]\n")
    
    # Initialize LLM
    llm = LLM(
        endpoint="https://api.openai.com/v1",
        model="gpt-4o-mini",
    )
    
    # Load test data
    try:
        data = pd.read_csv("data/loan_risk_samples.csv")
    except FileNotFoundError:
        console.print("[bold red]❌ Error:[/bold red] data/loan_risk_samples.csv not found!")
        return
    
    # Setup Agent Lightning components
    tracer = AgentOpsTracer()
    runner = LitAgentRunner[RiskTask](tracer=tracer)
    store = InMemoryLightningStore()
    
    # Run evaluation
    console.print(f"[bold blue]Running {len(data)} test scenarios...[/bold blue]\n")
    
    with runner.run_context(agent=loan_risk_agent, store=store):
        for index in range(min(8, len(data))):
            sample = cast(RiskTask, data.iloc[index].to_dict())
            await runner.step(sample, resources={"main_llm": llm})
            console.print()  # Blank line between tests
    
    console.print("\n[bold green]✓ All evaluations complete![/bold green]")


if __name__ == "__main__":
    asyncio.run(main())