""" inference.py — LLM-based agent using Scaler-injected LiteLLM proxy. Usage: python inference.py --task easy python inference.py --task all """ import os import sys import argparse import json from openai import OpenAI from models import StepName from environment import CustomerSupportEnv, STEP_ORDER from graders.base_grader import BaseGrader, HardTaskGrader from tasks import TASK_REGISTRY # ── LLM Client (uses Scaler-injected env vars) ──────────────────────────────── API_BASE_URL = os.environ.get("API_BASE_URL", "https://api.openai.com/v1") API_KEY = os.environ.get("API_KEY", "no-key") MODEL = os.environ.get("MODEL_NAME", "gpt-4o-mini") client = OpenAI(base_url=API_BASE_URL, api_key=API_KEY) # ── Step-specific system prompts ────────────────────────────────────────────── STEP_PROMPTS = { StepName.EMPATHY: ( "You are a professional AI customer support agent. " "Show genuine empathy. Apologize sincerely. Validate the customer's frustration. " "Do NOT ask for information. Do NOT give solutions yet. Max 3 sentences." ), StepName.COLLECT_INFO: ( "You are a professional AI customer support agent. " "Ask for the customer's order number or account email to look into this. " "Use: 'please provide your order number', 'may I have your email'. Max 2 sentences." ), StepName.INVESTIGATE: ( "You are a professional AI customer support agent. " "Tell the customer you are reviewing their case and share what you found. " "Use: 'I am checking', 'I can see in our records', 'I found that'. Max 3 sentences." ), StepName.RESOLUTION: ( "You are a professional AI customer support agent. " "Provide a concrete resolution: refund, replacement, or credit with a timeline. " "Personally guarantee resolution. Max 4 sentences." ), } # Fallback responses if LLM call fails FALLBACK_RESPONSES = { StepName.EMPATHY: ( "I am truly sorry to hear about your issue. I completely understand how " "frustrating this must be for you. I take full responsibility and will " "personally help resolve this immediately." ), StepName.COLLECT_INFO: ( "To assist you as quickly as possible, could you please provide me with " "your order number and the email address associated with your account so " "I can look into this right away?" ), StepName.INVESTIGATE: ( "Thank you for that information. I am checking our system right now. " "I can see your case in our records and I found the relevant details. " "Our records show the current status of your issue." ), StepName.RESOLUTION: ( "I sincerely apologize for this issue. I will personally process a full " "refund immediately, and you will receive confirmation within 24 hours. " "I will also escalate this to ensure it does not happen again. " "Thank you for your patience." ), } def call_llm(task, current_step: StepName) -> str: """Call LLM through the Scaler-injected LiteLLM proxy. Falls back gracefully on error.""" try: system_prompt = STEP_PROMPTS[current_step] user_msg = ( f"Customer message: {task.customer_message}\n" f"Context: {task.scenario_context}\n" f"Customer emotion: {task.customer_emotion}\n" f"Your task: {current_step.value.upper()}" ) response = client.chat.completions.create( model=MODEL, messages=[ {"role": "system", "content": system_prompt}, {"role": "user", "content": user_msg}, ], temperature=0.3, max_tokens=250, timeout=60, ) return response.choices[0].message.content.strip() except Exception as exc: print(f" [LLM Warning] {type(exc).__name__}: {exc} — using fallback", flush=True) return FALLBACK_RESPONSES[current_step] # ── Runner ──────────────────────────────────────────────────────────────────── def run_task(task_name: str) -> dict: try: task = TASK_REGISTRY[task_name] grader = HardTaskGrader() if task_name == "hard" else BaseGrader() env = CustomerSupportEnv(task=task, grader=grader) print(f"\n{'='*60}") print(f" TASK: {task_name.upper()} | {task.task_id}") print(f" Customer emotion: {task.customer_emotion}") print(f"{'='*60}") print(f" Customer: {task.customer_message[:120]}...") print(f"{'='*60}\n") # Required structured block print(f"[START] task={task_name}", flush=True) steps_taken = 0 for i, step in enumerate(STEP_ORDER): agent_response = call_llm(task, step) result, done = env.step(agent_response) steps_taken = i + 1 status = "CORRECT" if result.correct else "WRONG" print(f"[Step {i+1}/4] {step.value.upper()} — {status}") print(f" Agent : {agent_response[:100]}...") print(f" Detected : {result.detected_action}") print(f" Reward : {result.reward:.3f}") if result.penalty_reasons: for pr in result.penalty_reasons: print(f" Warning : {pr}") print() # Required structured block print(f"[STEP] step={i+1} reward={result.reward:.3f}", flush=True) if done: break summary = env.summary() print(f"\n{'='*60}") print(f" STATUS : {summary['status'].upper()}") print(f" REWARD : {summary['total_reward']:.3f} / 4.8 max") print(f"{'='*60}\n") # Required structured block — score must be strictly in (0, 1) MAX_SCORE = 4.8 # 4 steps × 1.2 max reward each raw_score = summary['total_reward'] normalized = raw_score / MAX_SCORE # Clamp strictly between 0 and 1 (not 0.0, not 1.0) final_score = max(0.001, min(0.999, normalized)) print( f"[END] task={task_name} score={final_score:.4f} steps={steps_taken}", flush=True, ) return summary except Exception as exc: print(f"[ERROR] run_task({task_name}) failed: {exc}", flush=True) # Emit END block with minimum valid score (strictly > 0) print(f"[END] task={task_name} score=0.001 steps=0", flush=True) return {"task_id": task_name, "status": "error", "total_reward": 0.0, "wrong_steps": 0, "fail_reason": str(exc), "steps": []} def main(): parser = argparse.ArgumentParser() parser.add_argument( "--task", choices=["easy", "medium", "hard", "all"], default="all" ) args = parser.parse_args() tasks = ["easy", "medium", "hard"] if args.task == "all" else [args.task] results = {} for t in tasks: results[t] = run_task(t) print("\n📊 FINAL SUMMARY", flush=True) print(json.dumps(results, indent=2), flush=True) sys.stdout.flush() if __name__ == "__main__": main()