fix: rewrite inference.py to use LiteLLM proxy via API_BASE_URL/API_KEY env vars

inference.py

@@ -1,133 +1,91 @@
 """
-inference.py — Rule-based agent + CLI test runner
+inference.py — LLM-based agent using Scaler-injected LiteLLM proxy.
 Usage:
-    python inference.py --task easy   --agent rule
-    python inference.py --task medium --agent rule
-    python inference.py --task hard   --agent rule
-    python inference.py --task all    --agent rule
+    python inference.py --task easy
+    python inference.py --task all
 """
 
+import os
 import argparse
 import json
-from typing import List
+from openai import OpenAI
 
-from models import StepName, DifficultyLevel
-from environment import CustomerSupportEnv
+from models import StepName, EpisodeStatus
+from environment import CustomerSupportEnv, STEP_ORDER
 from graders.base_grader import BaseGrader, HardTaskGrader
 from tasks import TASK_REGISTRY
 
-
-# ── Rule-based agent responses ─────────────────────────────────────────────────
-
-RULE_BASED_RESPONSES = {
-    "easy": [
-        # Step 1: Empathy
-        (
-            "I am truly sorry to hear that your laptop has not arrived yet. "
-            "I completely understand how frustrating and urgent this must be for you, "
-            "especially when you need it for work. Please allow me to help resolve this immediately."
-        ),
-        # Step 2: Collect Info
-        (
-            "To assist you as quickly as possible, could you please provide me with "
-            "your order number or order ID? Additionally, may I have your full name "
-            "and the email address associated with the order?"
-        ),
-        # Step 3: Investigate
-        (
-            "Thank you for that information. I am checking our system right now. "
-            "I found that order #ORD-8821 is currently tracked and it appears the package "
-            "is stuck at a sorting facility. Our records show the delay is due to a "
-            "logistical hold-up that we are actively investigating."
-        ),
-        # Step 4: Resolution
-        (
-            "I sincerely apologize for this unacceptable delay. I will personally "
-            "expedite your delivery and escalate this to our logistics partner immediately. "
-            "You will receive your laptop within 2 business days. As compensation, "
-            "I will also credit ₹500 to your account for the inconvenience caused."
-        ),
-    ],
-
-    "medium": [
-        # Step 1: Empathy
-        (
-            "I sincerely apologize for the double charge on your account. "
-            "I completely understand how upsetting and inconvenient this is. "
-            "This should never have happened and I take full responsibility. "
-            "I am going to resolve this for you right away."
-        ),
-        # Step 2: Collect Info
-        (
-            "To look into this immediately, may I please have your account number "
-            "or the email address linked to your subscription? "
-            "Could you also confirm the transaction dates and the card used for billing?"
-        ),
-        # Step 3: Investigate
-        (
-            "Thank you. I am reviewing your billing history right now. "
-            "I can see in our billing system that there was indeed a duplicate charge "
-            "on both the 1st and the 15th. It appears this was caused by a payment "
-            "gateway retry issue on our end. I found the duplicate transaction clearly."
-        ),
-        # Step 4: Resolution
-        (
-            "I sincerely apologize for this error. I will process a full refund of "
-            "$49.99 immediately. You should see this credited back to your card "
-            "within 3 to 5 business days. I am also escalating this to our billing "
-            "team to ensure this does not happen again. Thank you for your patience."
-        ),
-    ],
-
-    "hard": [
-        # Step 1: Empathy (de-escalation required)
-        (
-            "Mr. Mehta, I am personally and deeply sorry for what you have experienced. "
-            "This is completely unacceptable and you have every right to be angry. "
-            "I take full responsibility for the failure of our team to respond properly "
-            "over the past three weeks. This should not have happened, especially for a "
-            "valued VIP customer like yourself. I hear your frustration and I assure you "
-            "that I will personally ensure this is resolved today."
-        ),
-        # Step 2: Collect Info
-        (
-            "To immediately escalate this as our top priority, could you please confirm "
-            "your order number and VIP account details? I also need your email and "
-            "contact number so I can personally follow up with you today."
-        ),
-        # Step 3: Investigate
-        (
-            "Thank you, Mr. Mehta. I am reviewing your tickets right now. "
-            "I can see five tickets — all marked pending — with no technician dispatched. "
-            "I found that this is unacceptable internally. I have already escalated "
-            "this to our senior operations manager and our VIP support head. "
-            "Our records show the replacement unit is available and ready to ship."
-        ),
-        # Step 4: Resolution
-        (
-            "Mr. Mehta, here is what I am doing right now: "
-            "1) A replacement Premium Standing Desk Pro will be dispatched today with "
-            "priority delivery — you will receive it within 24 hours. "
-            "2) As per our VIP policy, I am authorizing a compensation of ₹9,000 "
-            "(20% of your order value) credited to your account immediately. "
-            "3) I will personally ensure a senior technician contacts you within 2 hours. "
-            "4) I am waiving your next month's subscription fee as an additional apology. "
-            "I personally guarantee this will be resolved to your satisfaction today."
-        ),
-    ],
+# ── LLM Client (uses Scaler-injected env vars) ────────────────────────────────
+client = OpenAI(
+    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")
+
+# ── Step-specific system prompts ──────────────────────────────────────────────
+STEP_PROMPTS = {
+    StepName.EMPATHY: (
+        "You are a professional AI customer support agent. "
+        "Your ONLY task: Show genuine empathy. Apologize sincerely, validate the "
+        "customer's frustration, take responsibility, assure them you will help. "
+        "Do NOT ask for information. Do NOT give solutions yet. "
+        "Use: 'I am deeply sorry', 'I completely understand', "
+        "'This should not have happened', 'I take full responsibility'. "
+        "Tone: Warm, sincere. Max 3-4 sentences."
+    ),
+    StepName.COLLECT_INFO: (
+        "You are a professional AI customer support agent. "
+        "Your ONLY task: Collect the customer's details to investigate their case. "
+        "Ask for their order number or account email. "
+        "Use phrases like: 'please provide', 'your order number', "
+        "'so I can look into this', 'I will need'. "
+        "Tone: Professional, direct. Max 2 sentences."
+    ),
+    StepName.INVESTIGATE: (
+        "You are a professional AI customer support agent. "
+        "Your ONLY task: Investigate and share findings. Say you are reviewing the "
+        "case and describe what you found. "
+        "Use: 'I am checking', 'I can see in our records', 'I found that', "
+        "'Our system shows'. "
+        "Do NOT give the final resolution yet. Max 3-4 sentences."
+    ),
+    StepName.RESOLUTION: (
+        "You are a professional AI customer support agent. "
+        "Your ONLY task: Provide a clear, concrete resolution with a specific action "
+        "(refund/replacement/credit/expedite) and timeline. "
+        "For VIP/hard tasks, mention 20%% compensation. "
+        "Personally guarantee resolution. Max 4-5 sentences."
+    ),
 }
 
 
-# ── Runner ─────────────────────────────────────────────────────────────────────
+def call_llm(task, current_step: StepName) -> str:
+    """Call LLM through the Scaler-injected LiteLLM proxy."""
+    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 current 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,
+    )
+    return response.choices[0].message.content.strip()
 
-def run_task(task_name: str) -> dict:
-    task      = TASK_REGISTRY[task_name]
-    grader    = HardTaskGrader() if task_name == "hard" else BaseGrader()
-    env       = CustomerSupportEnv(task=task, grader=grader)
-    responses = RULE_BASED_RESPONSES[task_name]
 
-    # ── Required structured output: START block ────────────────────────────
-    print(f"[START] task={task_name}", flush=True)
+# ── Runner ────────────────────────────────────────────────────────────────────
+
+def run_task(task_name: str) -> dict:
+    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}")
@@ -136,14 +94,18 @@ def run_task(task_name: str) -> dict:
     print(f"  Customer: {task.customer_message[:120]}...")
     print(f"{'='*60}\n")
 
+    # Required structured output: START block
+    print(f"[START] task={task_name}", flush=True)
+
     steps_taken = 0
-    for i, response in enumerate(responses):
-        step_name = list(StepName)[i]
-        result, done = env.step(response)
-        steps_taken = i + 1
+    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_name.value.upper()} — {status}")
+        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}  "
               f"(base={result.base_score:.2f}, bonus={result.step_bonus:.2f}, "
@@ -153,7 +115,7 @@ def run_task(task_name: str) -> dict:
                 print(f"  ⚠ {pr}")
         print()
 
-        # ── Required structured output: STEP block ─────────────────────────
+        # Required structured output: STEP block
         print(f"[STEP] step={i+1} reward={result.reward:.3f}", flush=True)
 
         if done:
@@ -168,7 +130,7 @@ def run_task(task_name: str) -> dict:
         print(f"  FAIL REASON     : {summary['fail_reason']}")
     print(f"{'='*60}\n")
 
-    # ── Required structured output: END block ──────────────────────────────
+    # Required structured output: END block
     print(
         f"[END] task={task_name} score={summary['total_reward']:.3f} steps={steps_taken}",
         flush=True,
@@ -179,19 +141,19 @@ def run_task(task_name: str) -> dict:
 
 def main():
     parser = argparse.ArgumentParser()
-    parser.add_argument("--task",  choices=["easy", "medium", "hard", "all"],
-                        default="all")
-    parser.add_argument("--agent", choices=["rule"], default="rule")
-    args = parser.parse_args()
+    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]
 
-    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")
-    print(json.dumps(results, indent=2))
+    print(json.dumps(results, indent=2), flush=True)
 
 
 if __name__ == "__main__":
-    main()
+    main()

pyproject.toml

@@ -12,6 +12,7 @@ dependencies = [
     "uvicorn>=0.23.0",
     "pydantic>=2.0.0",
     "pyyaml>=6.0",
+    "openai>=1.0.0",
     "openenv-core>=0.2.0",
 ]
 

requirements.txt

@@ -1,5 +1,5 @@
 fastapi>=0.100.0
 uvicorn>=0.23.0
-groq>=0.4.0
-pyyaml>=6.0
 pydantic>=2.0.0
+pyyaml>=6.0
+openai>=1.0.0

uv.lock

@@ -541,6 +541,7 @@ version = "1.0.0"
 source = { editable = "." }
 dependencies = [
     { name = "fastapi" },
+    { name = "openai" },
     { name = "openenv-core" },
     { name = "pydantic" },
     { name = "pyyaml" },
@@ -557,6 +558,7 @@ dev = [
 requires-dist = [
     { name = "fastapi", specifier = ">=0.100.0" },
     { name = "httpx", marker = "extra == 'dev'", specifier = ">=0.24.0" },
+    { name = "openai", specifier = ">=1.0.0" },
     { name = "openenv-core", specifier = ">=0.2.0" },
     { name = "pydantic", specifier = ">=2.0.0" },
     { name = "pytest", marker = "extra == 'dev'", specifier = ">=7.0" },