dprajjwal/supportops-env
0
1"""2SupportOps-Env: inference.py3Baseline inference script using OpenAI client against all 3 tasks.4 5Required environment variables:6 API_BASE_URL - The API endpoint for the LLM (e.g., https://api.openai.com/v1)7 MODEL_NAME - The model identifier (e.g., gpt-4o-mini)8 HF_TOKEN - Your HuggingFace / API key (used as OPENAI_API_KEY if set)9 OPENAI_API_KEY - OpenAI API key (overrides HF_TOKEN)10 ENV_BASE_URL - URL of the SupportOps-Env server (default: http://localhost:8000)11 12Output format (strict):13 [START] task=<name> max_steps=<n>14 [STEP] step=<n> action=<json> reward=<float> done=<bool> error=<str|None>15 [END] success=<bool> steps=<n> score=<float> rewards=<list>16"""17 18from __future__ import annotations19 20import asyncio21import json22import os23import re24import sys25from typing import Any, Dict, List, Optional, Tuple26 27import httpx28try:29 from openai import AsyncOpenAI30except ModuleNotFoundError:31 print("CRITICAL ERROR: Failed to import 'openai'. This script requires the 'openai' package.", flush=True)32 print("Please install it using: pip install openai>=1.0.0", flush=True)33 sys.exit(1)34 35# ─────────────────────────────────────────────36# Configuration (from environment variables)37# ─────────────────────────────────────────────38API_BASE_URL = os.getenv("API_BASE_URL", "https://api.openai.com/v1")39MODEL_NAME = os.getenv("MODEL_NAME", "gpt-4o-mini")40HF_TOKEN = os.getenv("HF_TOKEN")41ENV_BASE_URL = os.getenv("ENV_BASE_URL", "http://localhost:8000").rstrip("/")42 43# If HF_TOKEN is missing, look for OPENAI_API_KEY as a fallback44API_KEY = os.getenv("OPENAI_API_KEY", HF_TOKEN)45 46 47# Task parameters48TASKS = ["ticket_classification", "priority_sorting", "draft_response"]49MAX_STEPS: Dict[str, int] = {50 "ticket_classification": 5,51 "priority_sorting": 10,52 "draft_response": 15,53}54MAX_TOTAL_REWARD: Dict[str, float] = {55 "ticket_classification": 1.0,56 "priority_sorting": 1.0,57 "draft_response": 1.0,58}59SUCCESS_SCORE_THRESHOLD = 0.560 61 62# ─────────────────────────────────────────────63# Logging (strict format required by validator)64# ─────────────────────────────────────────────65def log_start(task: str, max_steps: int) -> None:66 print(f"[START] task={task} max_steps={max_steps}", flush=True)67 68 69def log_step(70 step: int,71 action: Any,72 reward: float,73 done: bool,74 error: Optional[str] = None,75) -> None:76 action_str = json.dumps(action) if not isinstance(action, str) else action77 print(78 f"[STEP] step={step} action={action_str!r} "79 f"reward={reward} done={done} error={error}",80 flush=True,81 )82 83 84def log_end(85 success: bool,86 steps: int,87 score: float,88 rewards: List[float],89) -> None:90 print(91 f"[END] success={success} steps={steps} "92 f"score={score:.4f} rewards={rewards}",93 flush=True,94 )95 96 97# ─────────────────────────────────────────────98# Environment HTTP helpers99# ─────────────────────────────────────────────100async def env_reset(101 client: httpx.AsyncClient,102 task_name: str,103 seed: int = 42,104) -> Dict[str, Any]:105 try:106 resp = await client.post(107 f"{ENV_BASE_URL}/reset",108 json={"task_name": task_name, "seed": seed},109 timeout=30.0,110 )111 resp.raise_for_status()112 return resp.json()113 except Exception as e:114 print(f"[ERROR] env_reset failed: {e}", flush=True)115 return {"done": True, "observation": {}, "error": str(e)}116 117 118async def env_step(119 client: httpx.AsyncClient,120 action: Dict[str, Any],121) -> Dict[str, Any]:122 try:123 resp = await client.post(124 f"{ENV_BASE_URL}/step",125 json={"action": action},126 timeout=30.0,127 )128 resp.raise_for_status()129 return resp.json()130 except Exception as e:131 print(f"[ERROR] env_step failed: {e}", flush=True)132 return {"done": True, "observation": {}, "reward": 0.0, "error": str(e)}133 134 135# ─────────────────────────────────────────────136# LLM Prompt builders137# ─────────────────────────────────────────────138SYSTEM_PROMPT = """You are an expert customer support operations agent.139You interact with a support ticket environment by outputting structured JSON actions.140ALWAYS respond with ONLY a valid JSON object (no markdown, no code blocks, no extra text).141 142For ticket_classification tasks, respond with:143{"action_type": "classify", "payload": {"category": "Bug|Feature Request|Billing|Account|General"}}144 145For priority_sorting tasks, respond with:146{"action_type": "set_priority", "payload": {"ticket_id": "<ID>", "priority": "critical|high|medium|low|minimal"}}147 148For draft_response tasks, respond with ONE of:149{"action_type": "search_kb", "payload": {"query": "<relevant search query>"}}150{"action_type": "draft_response", "payload": {"response_text": "<your professional response>"}}151{"action_type": "mark_resolved", "payload": {"reason": "Response drafted and sent"}}152 153Rules:154- Always pick the most logical next action given the current observation155- For draft_response: search KB first, then draft a response, then mark_resolved156- For priority_sorting: assign each ticket one at a time with set_priority157- Output ONLY the JSON action, nothing else158"""159 160 161def build_user_message(obs: Dict[str, Any], history: List[str], task_name: str) -> str:162 """Build user message from current observation."""163 lines = [164 f"Task: {obs.get('task_name', task_name)}",165 f"Steps remaining: {obs.get('steps_remaining', '?')}",166 "",167 ]168 169 # Task description (first call only, via step_context)170 step_ctx = obs.get("step_context", "")171 if step_ctx:172 lines.append(f"Context: {step_ctx}")173 lines.append("")174 175 # Ticket(s)176 if obs.get("ticket"):177 t = obs["ticket"]178 lines += [179 "=== TICKET ===",180 f"ID: {t.get('ticket_id')}",181 f"Subject: {t.get('subject')}",182 f"Body: {t.get('body')}",183 f"Customer tier: {t.get('customer_tier')}",184 f"Sentiment: {t.get('sentiment_score', 0):.1f} (-1=angry, +1=happy)",185 f"SLA hours: {t.get('sla_hours')}",186 "==============",187 ]188 189 if obs.get("tickets"):190 lines.append("=== TICKETS TO PRIORITIZE ===")191 for t in obs["tickets"]:192 lines += [193 f" [{t.get('ticket_id')}] {t.get('subject')}",194 f" Tier: {t.get('customer_tier')} | Sentiment: {t.get('sentiment_score', 0):.1f} | SLA: {t.get('sla_hours')}h",195 f" Body snippet: {t.get('body', '')[:120]}...",196 "",197 ]198 lines.append("=============================")199 200 # KB results201 if obs.get("kb_results"):202 lines.append("=== KB RESULTS ===")203 for r in obs["kb_results"]:204 lines.append(r[:300])205 lines.append("==================")206 207 # Recent history208 if history:209 lines.append("\nRecent history (last 5 steps):")210 for h in history[-5:]:211 lines.append(f" {h}")212 213 # Task-specific instructions214 if task_name == "ticket_classification":215 lines.append("\nClassify this ticket now.")216 elif task_name == "priority_sorting":217 assigned = [h for h in history if "set_priority" in h]218 lines.append(f"\nAssigned so far: {len(assigned)} tickets. Assign the next one.")219 elif task_name == "draft_response":220 if not any("search_kb" in h for h in history):221 lines.append("\nStart by searching the knowledge base for relevant information.")222 elif not any("draft_response" in h for h in history):223 lines.append("\nNow draft your response to the customer.")224 else:225 lines.append("\nMark the ticket as resolved.")226 227 return "\n".join(lines)228 229 230def parse_action_from_llm(text: str, task_name: str) -> Dict[str, Any]:231 """232 Parse a JSON action from LLM output.233 Falls back to a default action if parsing fails.234 """235 # Try to extract JSON from the response236 text = text.strip()237 238 # Remove markdown code blocks if present239 text = re.sub(r"```(?:json)?\s*", "", text)240 text = re.sub(r"```", "", text)241 text = text.strip()242 243 # Try direct parse244 try:245 data = json.loads(text)246 if "action_type" in data:247 return data248 except json.JSONDecodeError:249 pass250 251 # Try to find JSON in the text252 json_match = re.search(r'\{[^{}]*"action_type"[^{}]*\}', text, re.DOTALL)253 if json_match:254 try:255 return json.loads(json_match.group())256 except json.JSONDecodeError:257 pass258 259 # Fallback defaults per task260 fallbacks = {261 "ticket_classification": {"action_type": "classify", "payload": {"category": "General"}},262 "priority_sorting": {"action_type": "set_priority", "payload": {"ticket_id": "T001", "priority": "medium"}},263 "draft_response": {"action_type": "mark_resolved", "payload": {"reason": "Unable to parse action"}},264 }265 return fallbacks.get(task_name, {"action_type": "mark_resolved", "payload": {"reason": "parse_error"}})266 267 268# ─────────────────────────────────────────────269# Main task runner270# ─────────────────────────────────────────────271async def run_task(272 openai_client: AsyncOpenAI,273 http_client: httpx.AsyncClient,274 task_name: str,275 seed: int = 42,276) -> Tuple[float, bool, List[float]]:277 """Run one task episode and return (score, success, rewards)."""278 max_steps = MAX_STEPS[task_name]279 max_total = MAX_TOTAL_REWARD[task_name]280 281 rewards: List[float] = []282 steps_taken = 0283 success = False284 score = 0.0285 history: List[str] = []286 287 log_start(task_name, max_steps)288 289 # Reset290 result = await env_reset(http_client, task_name, seed=seed)291 obs = result.get("observation", {})292 done = result.get("done", False)293 last_reward = 0.0294 295 messages = [{"role": "system", "content": SYSTEM_PROMPT}]296 297 try:298 for step in range(1, max_steps + 1):299 if done:300 break301 302 # Build user message303 user_msg = build_user_message(obs, history, task_name)304 messages.append({"role": "user", "content": user_msg})305 306 # Get LLM action307 error = None308 action = {}309 try:310 completion = await openai_client.chat.completions.create(311 model=MODEL_NAME,312 messages=messages,313 temperature=0.0,314 max_tokens=512,315 )316 llm_text = completion.choices[0].message.content or ""317 action = parse_action_from_llm(llm_text, task_name)318 messages.append({"role": "assistant", "content": llm_text})319 except Exception as e:320 error = str(e)321 action = {"action_type": "mark_resolved", "payload": {"reason": f"LLM error: {e}"}}322 323 # Step environment324 try:325 result = await env_step(http_client, action)326 obs = result.get("observation", {})327 reward = float(result.get("reward") or 0.0)328 done = result.get("done", False)329 except Exception as e:330 error = str(e)331 reward = 0.0332 done = True333 334 rewards.append(reward)335 steps_taken = step336 last_reward = reward337 338 log_step(step=step, action=action, reward=reward, done=done, error=error)339 history.append(340 f"Step {step}: {action.get('action_type')} "341 f"-> reward {reward:+.3f} done={done}"342 )343 344 if done:345 break346 347 # Score348 score = sum(rewards) / max_total if max_total > 0 else 0.0349 score = min(max(score, 0.0), 1.0)350 success = score >= SUCCESS_SCORE_THRESHOLD351 352 finally:353 log_end(success=success, steps=steps_taken, score=score, rewards=rewards)354 355 return score, success, rewards356 357 358# ─────────────────────────────────────────────359# Entry point360# ─────────────────────────────────────────────361async def main() -> None:362 if not API_KEY:363 print("ERROR: Set HF_TOKEN or OPENAI_API_KEY environment variable.", flush=True)364 sys.exit(1)365 366 openai_client = AsyncOpenAI(367 api_key=API_KEY,368 base_url=API_BASE_URL,369 )370 371 all_scores: Dict[str, float] = {}372 373 async with httpx.AsyncClient() as http_client:374 # Verify server is running375 try:376 resp = await http_client.get(f"{ENV_BASE_URL}/health", timeout=10.0)377 resp.raise_for_status()378 print(f"[INFO] Connected to SupportOps-Env at {ENV_BASE_URL}", flush=True)379 except Exception as e:380 print(f"[ERROR] Cannot connect to environment server at {ENV_BASE_URL}: {e}", flush=True)381 print("[INFO] Start the server with: uvicorn server.app:app --port 8000", flush=True)382 sys.exit(1)383 384 for task_name in TASKS:385 print(f"\n{'='*60}", flush=True)386 print(f"Running task: {task_name}", flush=True)387 print(f"{'='*60}", flush=True)388 score, success, rewards = await run_task(389 openai_client, http_client, task_name, seed=42390 )391 all_scores[task_name] = score392 393 # Final summary394 print(f"\n{'='*60}", flush=True)395 print("BASELINE SCORES SUMMARY", flush=True)396 print(f"{'='*60}", flush=True)397 for task, s in all_scores.items():398 status = "✓ PASS" if s >= SUCCESS_SCORE_THRESHOLD else "✗ FAIL"399 print(f" {task:<30} score={s:.4f} {status}", flush=True)400 avg = sum(all_scores.values()) / len(all_scores) if all_scores else 0.0401 print(f"\n {'AVERAGE':<30} score={avg:.4f}", flush=True)402 print(f"{'='*60}", flush=True)403 404 405if __name__ == "__main__":406 asyncio.run(main())407