Spaces:
Sleeping
Sleeping
| from __future__ import annotations | |
| import argparse | |
| import json | |
| import os | |
| from typing import Dict, List | |
| from dotenv import load_dotenv | |
| from openai import OpenAI | |
| load_dotenv() | |
| from models import SupportAction | |
| from server.customer_support_environment import CustomerSupportEnvironment | |
| from tasks import grade_task | |
| def log_start(task: str, env: str, model: str) -> None: | |
| print(f"[START] task={task} env={env} model={model}", flush=True) | |
| def log_step(step: int, action: str, reward: float, done: bool, error: str | None = None) -> None: | |
| error_value = "null" if error is None else error.replace("\n", " ") | |
| print( | |
| f"[STEP] step={step} action={action!r} reward={reward:.4f} done={done} error={error_value}", | |
| flush=True, | |
| ) | |
| def log_end(success: bool, steps: int, score: float, rewards: List[float]) -> None: | |
| reward_text = ",".join(f"{r:.4f}" for r in rewards) | |
| print( | |
| f"[END] success={str(success).lower()} steps={steps} score={score:.4f} rewards=[{reward_text}]", | |
| flush=True, | |
| ) | |
| def get_model_plan(client: OpenAI, model_name: str, task_prompt: str) -> Dict[str, str | bool]: | |
| completion = client.chat.completions.create( | |
| model=model_name, | |
| temperature=0, | |
| messages=[ | |
| { | |
| "role": "system", | |
| "content": ( | |
| "You are a customer support policy agent. Reply with strict JSON only and no markdown. " | |
| "Schema: {\"category\": string, \"search_kb\": bool, \"escalate\": bool, \"response\": string}." | |
| ), | |
| }, | |
| {"role": "user", "content": task_prompt}, | |
| ], | |
| response_format={"type": "json_object"}, | |
| ) | |
| raw = completion.choices[0].message.content or "{}" | |
| data = json.loads(raw) | |
| return { | |
| "category": str(data.get("category", "technical_issue")), | |
| "search_kb": bool(data.get("search_kb", True)), | |
| "escalate": bool(data.get("escalate", False)), | |
| "response": str(data.get("response", "Thank you for contacting support. We will assist shortly.")), | |
| } | |
| def heuristic_plan(observation: object) -> Dict[str, str | bool]: | |
| query = str(getattr(observation, "query")).lower() | |
| requires_escalation = bool(getattr(observation, "requires_escalation")) | |
| category = "technical_issue" | |
| if any(t in query for t in ("refund", "charge", "invoice", "billing")): | |
| category = "billing" | |
| elif any(t in query for t in ("cancel", "termination", "close account", "cancellation")): | |
| category = "cancellation" | |
| elif any(t in query for t in ("warranty", "buy", "purchase", "availability", "product")): | |
| category = "product_inquiry" | |
| elif any(t in query for t in ("refund not received", "return", "reimburse")): | |
| category = "refund" | |
| return { | |
| "category": category, | |
| "search_kb": True, | |
| "escalate": requires_escalation, | |
| "response": "Thank you for contacting us. We will review your case and assist you shortly.", | |
| } | |
| def build_actions(observation: object, plan: Dict[str, str | bool]) -> List[SupportAction]: | |
| actions = [ | |
| SupportAction(action_type="classify", content=str(plan["category"])), | |
| ] | |
| if bool(plan.get("search_kb", True)): | |
| actions.append(SupportAction(action_type="search_kb", content=str(getattr(observation, "kb_id")))) | |
| if bool(plan.get("escalate", False)): | |
| actions.append(SupportAction(action_type="escalate", content="Escalating to specialist support for manual review.")) | |
| else: | |
| actions.append(SupportAction(action_type="respond", content=str(plan.get("response", "Thank you for contacting support.")))) | |
| return actions | |
| def main() -> int: | |
| parser = argparse.ArgumentParser(description="OpenEnv baseline inference for customer support benchmark") | |
| parser.add_argument("--csv", default="dataset.csv") | |
| parser.add_argument("--limit-per-task", type=int, default=5) | |
| parser.add_argument("--max-steps", type=int, default=4) | |
| parser.add_argument("--offline", action="store_true", help="Use heuristic policy without API calls") | |
| args = parser.parse_args() | |
| api_base_url = os.getenv("API_BASE_URL") | |
| model_name = os.getenv("MODEL_NAME") | |
| hf_token = os.getenv("HF_TOKEN") | |
| if not args.offline: | |
| missing = [name for name, value in (("API_BASE_URL", api_base_url), ("MODEL_NAME", model_name), ("HF_TOKEN", hf_token)) if not value] | |
| if missing: | |
| raise RuntimeError(f"Missing required env vars: {', '.join(missing)}") | |
| client = OpenAI(base_url=api_base_url, api_key=hf_token) if not args.offline else None | |
| env = CustomerSupportEnvironment(csv_path=args.csv) | |
| all_rewards: List[float] = [] | |
| task_scores: List[float] = [] | |
| task_labels: List[str] = [] | |
| total_steps = 0 | |
| log_start(task="all", env="customer_support_benchmark", model=model_name or "offline-heuristic") | |
| difficulties = ["easy", "medium", "hard"] | |
| for difficulty in difficulties: | |
| for idx in range(args.limit_per_task): | |
| obs = env.reset(difficulty=difficulty, index=idx) | |
| rewards: List[float] = [] | |
| done = False | |
| prompt = ( | |
| "Return the best support plan for this ticket. " | |
| "Prefer policy-safe responses and escalate only when needed.\n" | |
| f"TicketID={obs.ticket_id}\n" | |
| f"Difficulty={obs.difficulty}\n" | |
| f"Task={obs.task_id}\n" | |
| f"Query={obs.query}\n" | |
| f"KB={obs.kb_id}\n" | |
| f"RequiresEscalation={obs.requires_escalation}\n" | |
| ) | |
| error = None | |
| try: | |
| plan = heuristic_plan(obs) | |
| if client is not None and model_name is not None: | |
| plan = get_model_plan(client, model_name, prompt) | |
| except Exception as exc: | |
| error = str(exc) | |
| plan = heuristic_plan(obs) | |
| planned_actions = build_actions(obs, plan) | |
| for step, action in enumerate(planned_actions[: args.max_steps], start=1): | |
| if done: | |
| break | |
| obs = env.step(action) | |
| reward = float(obs.reward or 0.0) | |
| rewards.append(reward) | |
| all_rewards.append(reward) | |
| total_steps += 1 | |
| done = bool(obs.done) | |
| step_error = error if step == 1 else None | |
| log_step(step=step, action=f"{action.action_type}|{action.content}", reward=reward, done=done, error=step_error) | |
| final_score = grade_task(obs.task_id, obs.metadata, obs.history) | |
| task_scores.append(final_score) | |
| task_labels.append(difficulty) | |
| log_step( | |
| step=total_steps, | |
| action=f"episode_score|difficulty={difficulty}|index={idx}", | |
| reward=final_score, | |
| done=True, | |
| error=None, | |
| ) | |
| benchmark_score = sum(task_scores) / len(task_scores) if task_scores else 0.0 | |
| by_task: dict[str, list[float]] = {"easy": [], "medium": [], "hard": []} | |
| for label, score in zip(task_labels, task_scores): | |
| by_task[label].append(score) | |
| for label in ("easy", "medium", "hard"): | |
| values = by_task[label] | |
| avg = sum(values) / len(values) if values else 0.0 | |
| log_step( | |
| step=total_steps, | |
| action=f"task_average|difficulty={label}", | |
| reward=avg, | |
| done=False, | |
| error=None, | |
| ) | |
| print( | |
| f"[END] success={str(benchmark_score >= 0.7).lower()} steps={total_steps} score={benchmark_score:.4f} rewards_count={len(all_rewards)}", | |
| flush=True, | |
| ) | |
| return 0 | |
| if __name__ == "__main__": | |
| raise SystemExit(main()) |