Spaces:
Sleeping
Sleeping
| """ | |
| 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() | |