""" Baseline inference script for the support triage OpenEnv. Requirements: - Uses the OpenAI API client. - Reads credentials from `OPENAI_API_KEY`. - Produces reproducible baseline scores across Task 1..3. Notes: - Sprint 1: Task 2/3 grading are stubs; the environment still supports a classify->finalize flow so the baseline can run end-to-end. """ from __future__ import annotations import json import os import re from typing import Any from openai import OpenAI try: # Normal import when used as `support_triage_env.baseline_inference` from .grader import grade_episode from .models import SupportAction, SupportObservation from .server.support_triage_env_environment import SupportTriageEnvironment except ImportError: # pragma: no cover # Fallback import when loaded as `baseline_inference` (e.g. from `server/app.py`) from grader import grade_episode from models import SupportAction, SupportObservation from server.support_triage_env_environment import SupportTriageEnvironment def _extract_first_json_object(text: str) -> dict[str, Any]: """ Best-effort JSON extraction from an LLM response. The baseline must not crash on minor formatting deviations. """ if not text: return {} # Try to locate the first {...} block. start = text.find("{") end = text.rfind("}") if start == -1 or end == -1 or end <= start: return {} blob = text[start : end + 1] try: return json.loads(blob) except json.JSONDecodeError: # Last-resort cleanup: remove trailing commas. cleaned = re.sub(r",\s*([}\]])", r"\1", blob) return json.loads(cleaned) def _heuristic_fallback_action(obs: SupportObservation, task_id: str) -> SupportAction: """ Deterministic fallback if JSON parsing fails. This keeps the baseline script robust even if the LLM response is malformed. """ if obs.task_stage == "classify": if task_id == "task_1": intent = "billing_refund" template = "tpl_refund_ack_next" elif task_id == "task_2": intent = "password_reset" template = "tpl_reset_ack_next" else: intent = "data_breach_concern" template = "tpl_safety_escalate" return SupportAction(task_id=task_id, task_stage="classify", intent_category=intent, reply_template_id=template) if obs.task_stage == "request_info": # Task 2: request missing order id + account email. if task_id == "task_2": return SupportAction( task_id=task_id, task_stage="request_info", missing_info_questions=["order id (order number)", "account email address"], finalize=False, ) # For any other task, fall back to an empty request (environment will penalize if used). return SupportAction(task_id=task_id, task_stage="request_info", missing_info_questions=[], finalize=False) # finalize if task_id == "task_1": template = "tpl_refund_ack_next" draft = "Thanks for reaching out. We received your request. Next step: please confirm the invoice details so we can process the refund." elif task_id == "task_2": template = "tpl_reset_ack_next" calc_seed = obs.episode_seed or 0 order_suffix = (calc_seed * 37 + 101) % 9000 order_id = f"ORD-{1000 + order_suffix}" email_local = 10 + ((calc_seed * 53 + 17) % 90) account_email = f"user{email_local}@example.com" draft = ( "Thanks for contacting support. We received your request. Next step: please confirm. " f"Order id is {order_id}. Account email is {account_email}." ) else: # Task 3: deterministic high-risk vs self-serve split from episode_seed. needs_escalation = (obs.episode_seed or 0) % 2 == 0 if needs_escalation: template = "tpl_safety_escalate" draft = ( "Thanks for flagging this. We received your security concern. " "Next step: we will escalate this to a human support/security team for review. " "Please do not share sensitive details; we will confirm the next steps." ) else: template = "tpl_safety_selfserve" draft = ( "Thanks for reaching out. We received your security concern. " "Next step: please change your password and enable MFA (two-factor authentication) " "to secure your account. If you still see suspicious activity, follow up with support." ) return SupportAction( task_id=task_id, task_stage="finalize", reply_template_id=template, draft_reply=draft, finalize=True, escalate_to_human=(task_id == "task_3" and ((obs.episode_seed or 0) % 2 == 0)), ) def _build_system_prompt() -> str: return ( "You are an AI customer-support triage agent inside an OpenEnv environment. " "You MUST respond with a single JSON object that matches the provided SupportAction schema. " "Do not include markdown, code fences, or extra commentary." ) def _build_user_prompt(obs: SupportObservation, action_schema: dict[str, Any]) -> str: # Keep user prompt small but sufficient. return ( f"Environment observation (JSON): {obs.model_dump()}\n\n" f"SupportAction JSON schema:\n{json.dumps(action_schema, indent=2)}\n\n" "Return ONLY a JSON object matching the schema." ) def _call_openai_for_action( client: OpenAI, *, model: str, obs: SupportObservation, task_id: str, seed: int, ) -> SupportAction: action_schema = SupportAction.model_json_schema() messages = [ {"role": "system", "content": _build_system_prompt()}, {"role": "user", "content": _build_user_prompt(obs, action_schema)}, ] kwargs: dict[str, Any] = { "model": model, "messages": messages, "temperature": 0, } # Seed support is model/backend dependent; try it but don't fail baseline if unsupported. try: kwargs["seed"] = seed except Exception: pass resp = client.chat.completions.create(**kwargs) content = resp.choices[0].message.content or "" data = _extract_first_json_object(content) # Validate/construct via Pydantic. try: return SupportAction(**data) except Exception: return _heuristic_fallback_action(obs, task_id) def run_baseline(tasks: list[str] | None = None, *, seed: int = 123, model: str = "gpt-4o-mini") -> dict[str, Any]: """ Run baseline on the environment for the given tasks. """ api_key = os.getenv("OPENAI_API_KEY") client: OpenAI | None = None if api_key: client = OpenAI(api_key=api_key) env = SupportTriageEnvironment() task_ids = tasks or ["task_1", "task_2", "task_3"] results: dict[str, Any] = {"seed": seed, "overall": 0.0, "tasks": {}} scores: list[float] = [] for i, task_id in enumerate(task_ids): obs = env.reset(seed=seed + i * 17, task_id=task_id) transcript: list[dict[str, Any]] = [] step_budget = 6 # keep baseline runtime bounded for _ in range(step_budget): if client is None: action = _heuristic_fallback_action(obs, task_id) else: action = _call_openai_for_action( client, model=model, obs=obs, task_id=task_id, seed=seed + i * 17, ) next_obs = env.step(action) transcript.append({"action": action.model_dump(), "observation": next_obs.model_dump()}) obs = next_obs if obs.done: break graded = grade_episode(task_id, seed + i * 17, transcript) score = float(graded["score"]) results["tasks"][task_id] = {"score": score, "grader": graded} scores.append(score) overall = sum(scores) / max(1, len(scores)) results["overall"] = float(overall) return results def main() -> None: tasks = ["task_1", "task_2", "task_3"] seed = int(os.getenv("BASELINE_SEED", "123")) model = os.getenv("BASELINE_MODEL", "gpt-4o-mini") out = run_baseline(tasks, seed=seed, model=model) print(json.dumps(out, indent=2)) if __name__ == "__main__": main()