from __future__ import annotations import json from pathlib import Path from typing import Any, Dict from uuid import uuid4 import pandas as pd from .experts import DataAnalystExpert, FinanceExpert, HRExpert, StrategyExpert from .graders import grade_episode, load_metric_ground_truth from .models import Brief, CoSAction, CoSObservation, CoSState, ExpertReport, RewardBreakdown TASK_ROOT = Path(__file__).resolve().parent / 'tasks' # Per-task order for “single” baselines, oracle, and any policy that should cover required experts. REQUIRED_EXPERTS_BY_TASK: dict[str, list[str]] = { 'easy_brief': ['analyst', 'finance', 'hr'], 'medium_brief': ['analyst', 'finance', 'strategy', 'hr'], 'hard_brief': ['analyst', 'finance', 'strategy', 'hr'], 'expert_brief': ['analyst', 'finance', 'strategy', 'hr'], 'risk_brief': ['analyst', 'finance', 'strategy', 'hr'], 'crisis_brief': ['analyst', 'finance', 'strategy', 'hr'], } def required_experts_for_task(task_name: str) -> list[str]: return list(REQUIRED_EXPERTS_BY_TASK.get(task_name, ['analyst', 'finance', 'hr'])) class CEOBriefEnvironment: def __init__(self, shaping: str = "default", auto_fill_required: bool = True) -> None: """Multi-agent CEO-brief env. ``shaping`` controls the dense per-step reward. The terminal grader is unchanged either way; that is what hackathon scoring uses. - ``"default"``: legacy per-step rewards. Stable; matches existing trained checkpoints and submitted runs. - ``"strict"``: anti-degenerate shaping for RL training. Adds a repetition penalty, an over-consult penalty, an early-finish bonus when all required experts are covered, and a stronger penalty for summarizing before required experts have reported. Use for new GRPO/REINFORCE runs to discourage "summarize-spam -> submit" lazy policies. ``auto_fill_required`` keeps the production/demo environment robust by filling any missing required experts before composing or grading. Turn it off only for policy-evidence runs where we want to observe what the LLM actually routed by itself. """ self.analyst = DataAnalystExpert() self.finance = FinanceExpert() self.hr = HRExpert() self.strategy = StrategyExpert() self.use_rag = False self.shaping = shaping if shaping in {"default", "strict"} else "default" self.auto_fill_required = auto_fill_required self.reset() def reset(self, task: str = 'easy_brief', episode_id: str | None = None, use_rag: bool = False) -> CoSObservation: self.use_rag = use_rag self.episode_id = episode_id or str(uuid4()) self.task_name = task if (TASK_ROOT / task).exists() else 'easy_brief' task_dir = TASK_ROOT / self.task_name self.raw_df = pd.read_csv(task_dir / 'raw.csv') self.gt_metrics = load_metric_ground_truth(str(task_dir / 'ground_truth.csv')) if (task_dir / 'ground_truth.csv').exists() else {} with open(task_dir / 'metadata.json', encoding='utf-8') as f: self.meta = json.load(f) self.step_count = 0 self.done = False self.cumulative_reward = 0.0 self.expert_reports: Dict[str, ExpertReport] = {} self.current_brief: Brief | None = None self.history: list[str] = [] self.last_reward = 0.0 self.last_terminal = None self.last_data_quality = 0.0 self.last_issues = ['No experts consulted yet.'] self._consult_counts: Dict[str, int] = {} self._last_action_key: str | None = None return self._observe(initial=True) def state(self) -> CoSState: return CoSState( episode_id=self.episode_id, task_name=self.task_name, step_count=self.step_count, done=self.done, rag_enabled=self.use_rag, consulted_experts=list(self.expert_reports.keys()), expert_reports=self.expert_reports, current_brief=self.current_brief, cumulative_reward=self.cumulative_reward, ) def _observe(self, initial: bool = False) -> CoSObservation: return CoSObservation( done=self.done, reward=0.0 if initial else self.last_reward, instruction=self.meta['instruction'], history=list(self.history), issues=list(self.last_issues), data_quality_score=self.last_data_quality, task_name=self.task_name, task_difficulty=self.meta['difficulty'], max_steps=int(self.meta.get('max_steps', 12)), step_count=self.step_count, rag_enabled=self.use_rag, consulted_experts=list(self.expert_reports.keys()), expert_reports=self.expert_reports, current_brief=self.current_brief, reward_breakdown=RewardBreakdown( immediate=self.last_reward, cumulative=self.cumulative_reward, terminal_grader=self.last_terminal, ), terminal_grader_score=self.last_terminal, ) def _compose_brief(self) -> Brief: metrics: Dict[str, Any] = {} recommendations: list[str] = [] summary_parts: list[str] = [] for expert_id in ('analyst', 'finance'): report = self.expert_reports.get(expert_id) if report: metrics.update(report.metrics) summary_parts.append(report.summary) if 'strategy' in self.expert_reports: recommendations = list(self.expert_reports['strategy'].bullet_points) summary_parts.append(self.expert_reports['strategy'].summary) hr_memo = self.expert_reports['hr'].memo if 'hr' in self.expert_reports and self.expert_reports['hr'].memo else '' summary = ' '.join(summary_parts) if summary_parts else 'No brief drafted yet.' self.current_brief = Brief( summary=summary, metrics=metrics, recommendations=recommendations, hr_memo=hr_memo, consulted_experts=list(self.expert_reports.keys()), ) return self.current_brief def _run_expert(self, expert_id: str, focused: bool = False) -> ExpertReport: question = self.meta['instruction'] if expert_id == 'analyst': report = self.analyst.run( self.task_name, question, self.raw_df, focused=focused, use_rag=self.use_rag ) self.last_data_quality = float(report.metrics.get('data_quality_score', 0.0)) self.last_issues = report.issues or ['analyst:no material issues'] return report if expert_id == 'finance': analyst = self.expert_reports.get('analyst') or self._run_expert('analyst') return self.finance.run( self.task_name, question, self.raw_df, analyst.metrics, self.meta, focused=focused, use_rag=self.use_rag, ) if expert_id == 'strategy': analyst = self.expert_reports.get('analyst') or self._run_expert('analyst') finance = self.expert_reports.get('finance') or self._run_expert('finance') return self.strategy.run( self.task_name, self.meta, analyst, finance, focused=focused, use_rag=self.use_rag ) if expert_id == 'hr': analyst = self.expert_reports.get('analyst') or self._run_expert('analyst') finance = self.expert_reports.get('finance') or self._run_expert('finance') strategy = self.expert_reports.get('strategy') return self.hr.run( self.task_name, self.meta, analyst, finance, strategy, focused=focused, use_rag=self.use_rag ) raise ValueError(f'Unknown expert {expert_id!r}') def _ensure_required_experts(self) -> list[str]: """Run any task-required experts that the policy never consulted. This guarantees the strategist (and any other required role) always contributes to the brief, so the UI / grader always has their report. Returns the list of expert ids that were auto-filled. """ if not self.auto_fill_required: return [] auto: list[str] = [] for expert_id in required_experts_for_task(self.task_name): if expert_id in self.expert_reports: continue try: self.expert_reports[expert_id] = self._run_expert(expert_id) auto.append(expert_id) except Exception: continue return auto def step(self, action: CoSAction) -> CoSObservation: if self.done: return self._observe() self.step_count += 1 immediate = -0.02 details = action.model_dump(exclude_none=True) action_key = json.dumps(details, sort_keys=True) self.history.append(action_key) strict = self.shaping == 'strict' if strict and self._last_action_key is not None and action_key == self._last_action_key: immediate -= 0.05 self._last_action_key = action_key required = list(self.meta.get('required_experts', [])) if action.action_type in {'consult', 'ask'}: if not action.expert_id: immediate -= 0.03 self.last_issues = ['action_missing_expert'] else: prior = action.expert_id in self.expert_reports report = self._run_expert(action.expert_id, focused=action.action_type == 'ask') self.expert_reports[action.expert_id] = report immediate += 0.10 if not prior and action.expert_id in required else 0.02 if prior: immediate -= 0.05 if strict: self._consult_counts[action.expert_id] = self._consult_counts.get(action.expert_id, 0) + 1 if self._consult_counts[action.expert_id] > 2: immediate -= 0.10 self.last_issues = report.issues or [f'{action.expert_id}:ok'] elif action.action_type == 'summarize': brief_already_exists = self.current_brief is not None missing_required = [e for e in required if e not in self.expert_reports] self._ensure_required_experts() self._compose_brief() immediate += 0.04 if len(self.expert_reports) >= 2 else -0.02 if strict and missing_required: immediate -= 0.05 * len(missing_required) if strict and brief_already_exists: immediate -= 0.08 self.last_issues = ['brief_composed'] elif action.action_type == 'submit': auto_filled = self._ensure_required_experts() if self.current_brief is None or auto_filled: self._compose_brief() self.done = True self.last_terminal = grade_episode( self.gt_metrics, self.meta, self.current_brief, self.expert_reports, use_rag=self.use_rag ) immediate += self.last_terminal if strict and not auto_filled: max_steps = int(self.meta.get('max_steps', 12)) steps_saved = max(0, max_steps - self.step_count) if steps_saved > 0 and all(e in self.expert_reports for e in required): immediate += min(0.10, 0.01 * steps_saved) self.last_issues = ['submitted'] + ( [f'auto_consulted:{",".join(auto_filled)}'] if auto_filled else [] ) else: self.last_issues = ['noop'] immediate -= 0.01 if not self.done and self.step_count >= int(self.meta.get('max_steps', 12)): auto_filled = self._ensure_required_experts() if self.current_brief is None or auto_filled: self._compose_brief() self.done = True self.last_terminal = grade_episode( self.gt_metrics, self.meta, self.current_brief, self.expert_reports, use_rag=self.use_rag ) immediate += self.last_terminal self.last_issues = ['forced_termination:max_steps'] + ( [f'auto_consulted:{",".join(auto_filled)}'] if auto_filled else [] ) self.last_reward = round(immediate, 4) self.cumulative_reward = round(self.cumulative_reward + self.last_reward, 4) return self._observe() def oracle_action_for_observation(obs: CoSObservation) -> CoSAction: for expert in required_experts_for_task(obs.task_name): if expert not in obs.consulted_experts: return CoSAction(action_type='consult', expert_id=expert) if obs.current_brief is None: return CoSAction(action_type='summarize') return CoSAction(action_type='submit')