add 3 tasks

inference.py

@@ -23,7 +23,7 @@ from openai import OpenAI
 
 API_BASE_URL = os.getenv("API_BASE_URL", "https://api-inference.huggingface.co/v1")
 MODEL_NAME = os.getenv("MODEL_NAME", "meta-llama/Llama-3.1-8B-Instruct")
-HF_TOKEN = os.getenv("HF_TOKEN")
+HF_TOKEN = os.getenv("HF_TOKEN") or os.getenv("OPENAI_API_KEY")
 LOCAL_IMAGE_NAME = os.getenv("LOCAL_IMAGE_NAME")
 ENV_BASE_URL = os.getenv("ENV_BASE_URL", "http://localhost:7860")
 
@@ -70,12 +70,14 @@ def log_end(success: bool, steps: int, score: float, rewards: List[float]) -> No
     )
 
 
-def call_env(endpoint: str, payload: Optional[dict] = None, method: str = "POST") -> dict:
+def call_env(
+    http: requests.Session, endpoint: str, payload: Optional[dict] = None, method: str = "POST"
+) -> dict:
     url = f"{ENV_BASE_URL}/{endpoint}"
     if method == "GET":
-        response = requests.get(url, timeout=30)
+        response = http.get(url, timeout=30)
     else:
-        response = requests.post(url, json=payload or {}, timeout=30)
+        response = http.post(url, json=payload or {}, timeout=30)
     response.raise_for_status()
     return response.json()
 
@@ -105,11 +107,13 @@ def main() -> None:
 
     try:
         if not HF_TOKEN:
-            raise RuntimeError("HF_TOKEN must be set for inference.")
+            raise RuntimeError("HF_TOKEN or OPENAI_API_KEY must be set for inference.")
 
         client = OpenAI(base_url=API_BASE_URL, api_key=HF_TOKEN)
-        observation = call_env("reset", {})
+        http = requests.Session()
+        observation = call_env(http, "reset", {})
         current_task = observation.get("task_id", "unknown")
+        session_id = observation.get("metadata", {}).get("session_id")
         log_start(task=current_task, env=BENCHMARK, model=MODEL_NAME)
         start_logged = True
 
@@ -118,22 +122,31 @@ def main() -> None:
                 break
 
             current_task = observation.get("task_id", "unknown")
-            action_text = get_model_response(client, observation["scenario"])
+            scenario = observation.get("scenario")
+            if not scenario:
+                raise RuntimeError("Environment response missing 'scenario'.")
+
+            action_text = get_model_response(client, scenario)
 
             error = None
             done = False
             reward = 0.0
 
             try:
+                metadata = {}
+                if session_id:
+                    metadata["session_id"] = session_id
                 observation = call_env(
+                    http,
                     "step",
                     {
                         "response": action_text,
                         "task_id": current_task,
                         "confidence": 0.85,
-                        "metadata": {},
+                        "metadata": metadata,
                     },
                 )
+                session_id = observation.get("metadata", {}).get("session_id", session_id)
                 reward = float(observation.get("reward", 0.0) or 0.0)
                 done = bool(observation.get("done", False))
             except Exception as exc:

it_mental_health_environment.py

@@ -13,24 +13,32 @@ from typing import Optional
 from openai import OpenAI
 
 try:
-    from models import MentalHealthAction, MentalHealthObservation, MentalHealthState
+    from models import MentalHealthAction, MentalHealthObservation, MentalHealthReward, MentalHealthState
 except ImportError:
     try:
-        from server.models import MentalHealthAction, MentalHealthObservation, MentalHealthState
+        from server.models import MentalHealthAction, MentalHealthObservation, MentalHealthReward, MentalHealthState
     except ImportError:
         import sys
         sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
-        from models import MentalHealthAction, MentalHealthObservation, MentalHealthState
+        from models import MentalHealthAction, MentalHealthObservation, MentalHealthReward, MentalHealthState
 
 TASK_ORDER = ["burnout_detection", "stress_triage", "intervention_plan"]
+TASK_DIFFICULTY = {
+    "burnout_detection": "easy",
+    "stress_triage": "medium",
+    "intervention_plan": "hard",
+}
 
 # ── LLM Judge client ──────────────────────────────────────────────────────────
 _llm_client: Optional[OpenAI] = None
 
 def _get_llm_client() -> Optional[OpenAI]:
     global _llm_client
+    use_llm_judge = os.environ.get("USE_LLM_JUDGE", "").strip().lower() in {"1", "true", "yes"}
+    if not use_llm_judge:
+        return None
     if _llm_client is None:
-        api_key = os.environ.get("HF_TOKEN", "")
+        api_key = os.environ.get("HF_TOKEN", "") or os.environ.get("OPENAI_API_KEY", "")
         base_url = os.environ.get("API_BASE_URL", "https://api-inference.huggingface.co/v1")
         if api_key:
             _llm_client = OpenAI(api_key=api_key, base_url=base_url)
@@ -371,6 +379,22 @@ class ITMentalHealthEnvironment:
             "intervention_plan": generate_intervention_plan_scenario(self._rng),
         }
 
+    def _score_metadata(self, task: str, reward: float):
+        difficulty = TASK_DIFFICULTY[task]
+        task_step_count = self._state.task_step_counts.get(task, 0) + 1
+        task_cumulative_score = round(self._state.task_scores.get(task, 0.0) + reward, 3)
+        difficulty_cumulative_score = round(
+            self._state.difficulty_scores.get(difficulty, 0.0) + reward, 3
+        )
+        return {
+            "difficulty": difficulty,
+            "step_score": round(reward, 3),
+            "task_step_number": task_step_count,
+            "task_cumulative_score": task_cumulative_score,
+            "difficulty_cumulative_score": difficulty_cumulative_score,
+            "overall_cumulative_score": round(self._state.cumulative_reward + reward, 3),
+        }
+
     def reset(self, seed=None):
         actual_seed = seed if seed is not None else random.randint(0, 2**31)
         self._rng = random.Random(actual_seed)
@@ -379,13 +403,24 @@ class ITMentalHealthEnvironment:
         self._state = MentalHealthState(
             episode_id=str(uuid.uuid4()), step_count=0,
             current_task=TASK_ORDER[0], cumulative_reward=0.0, tasks_completed=[],
+            task_scores={task_id: 0.0 for task_id in TASK_ORDER},
+            task_step_counts={task_id: 0 for task_id in TASK_ORDER},
+            difficulty_scores={"easy": 0.0, "medium": 0.0, "hard": 0.0},
         )
         task = TASK_ORDER[0]
         scenario_text, _ = self._current_scenarios[task]
         return MentalHealthObservation(
             scenario=scenario_text, feedback=f"New episode (seed={actual_seed}).",
-            reward=0.0, done=False, score_breakdown={}, task_id=task,
-            metadata={"seed": actual_seed},
+            task_id=task,
+            metadata={
+                "seed": actual_seed,
+                "difficulty": TASK_DIFFICULTY[task],
+                "step_score": 0.0,
+                "task_step_number": 0,
+                "task_cumulative_score": 0.0,
+                "difficulty_cumulative_score": 0.0,
+                "overall_cumulative_score": 0.0,
+            },
         )
 
     def step(self, action):
@@ -394,12 +429,13 @@ class ITMentalHealthEnvironment:
 
         # Guard: refuse to grade once the episode is finished.
         if self._task_index >= len(TASK_ORDER):
-            return MentalHealthObservation(
+            observation = MentalHealthObservation(
                 scenario="Episode already finished. Call /reset to start a new one.",
                 feedback="No-op: episode is done.",
-                reward=0.0, done=True, score_breakdown={},
                 task_id=self._state.current_task,
             )
+            reward = MentalHealthReward(value=0.0, score_breakdown={}, feedback="No-op: episode is done.")
+            return observation, reward, True, {}
 
         self._state.step_count += 1
         task = self._state.current_task
@@ -411,7 +447,11 @@ class ITMentalHealthEnvironment:
             task_id=task, scenario=scenario_text,
             response=action.response, ground_truth=ground_truth,
         )
+        score_metadata = self._score_metadata(task, reward)
         self._state.cumulative_reward += reward
+        self._state.task_scores[task] = score_metadata["task_cumulative_score"]
+        self._state.task_step_counts[task] = score_metadata["task_step_number"]
+        self._state.difficulty_scores[TASK_DIFFICULTY[task]] = score_metadata["difficulty_cumulative_score"]
         self._state.tasks_completed.append(task)
 
         self._task_index += 1
@@ -425,11 +465,18 @@ class ITMentalHealthEnvironment:
             self._state.current_task = next_task
             next_scenario, _ = self._current_scenarios[next_task]
 
-        return MentalHealthObservation(
-            scenario=next_scenario, feedback=feedback,
-            reward=reward, done=done, score_breakdown=breakdown, task_id=next_task,
+        observation = MentalHealthObservation(
+            scenario=next_scenario,
+            feedback=feedback,
+            task_id=next_task,
+            metadata=score_metadata,
+        )
+        reward_model = MentalHealthReward(
+            value=reward,
+            score_breakdown=breakdown,
+            feedback=feedback,
         )
+        return observation, reward_model, done, score_metadata
 
-    @property
     def state(self):
         return self._state

models.py

@@ -1,66 +1,38 @@
 """
-IT Mental Health OpenEnv - Typed Models
-Action / Observation / State for the IT Burnout & Mental Wellness environment.
+IT Mental Health OpenEnv - Typed Pydantic models.
 """
 
-from dataclasses import dataclass, field
-from typing import Optional, Dict, Any, List
+from typing import Any, Dict, List, Optional
 
+from pydantic import BaseModel, Field
 
-@dataclass
-class MentalHealthAction:
-    """
-    Action submitted by the agent for each mental health scenario step.
 
-    Fields:
-        response (str): The agent's textual response / intervention recommendation.
-        task_id (str): Which task is being attempted: 'burnout_detection',
-                       'stress_triage', or 'intervention_plan'.
-        confidence (float): Agent's self-reported confidence 0.0–1.0.
-        metadata (dict): Optional extra fields (reasoning chain, flags, etc.)
-    """
+class MentalHealthAction(BaseModel):
     response: str
     task_id: str = "burnout_detection"
     confidence: float = 1.0
-    metadata: Dict[str, Any] = field(default_factory=dict)
+    metadata: Dict[str, Any] = Field(default_factory=dict)
 
 
-@dataclass
-class MentalHealthObservation:
-    """
-    Observation returned after each step.
+class MentalHealthReward(BaseModel):
+    value: float
+    score_breakdown: Dict[str, float] = Field(default_factory=dict)
+    feedback: str = ""
 
-    Fields:
-        scenario (str): The current scenario description shown to the agent.
-        feedback (str): Evaluator feedback on the last action.
-        reward (float): Reward for the last step (0.0 – 1.0).
-        done (bool): Whether the episode is finished.
-        score_breakdown (dict): Partial scores by rubric dimension.
-        task_id (str): Current task identifier.
-    """
+
+class MentalHealthObservation(BaseModel):
     scenario: str
     feedback: str
-    reward: float
-    done: bool
-    score_breakdown: Dict[str, float] = field(default_factory=dict)
     task_id: str = "burnout_detection"
-    metadata: Dict[str, Any] = field(default_factory=dict)
-
+    metadata: Dict[str, Any] = Field(default_factory=dict)
 
-@dataclass
-class MentalHealthState:
-    """
-    Episode state / metadata.
 
-    Fields:
-        episode_id (str): UUID for the current episode.
-        step_count (int): How many steps have elapsed.
-        current_task (str): Active task identifier.
-        cumulative_reward (float): Total reward accumulated so far.
-        tasks_completed (list): List of task IDs already completed.
-    """
+class MentalHealthState(BaseModel):
     episode_id: Optional[str] = None
     step_count: int = 0
     current_task: str = "burnout_detection"
     cumulative_reward: float = 0.0
-    tasks_completed: List[str] = field(default_factory=list)
+    tasks_completed: List[str] = Field(default_factory=list)
+    task_scores: Dict[str, float] = Field(default_factory=dict)
+    task_step_counts: Dict[str, int] = Field(default_factory=dict)
+    difficulty_scores: Dict[str, float] = Field(default_factory=dict)

server/app.py

@@ -13,7 +13,7 @@ from fastapi import FastAPI, Header, HTTPException, Query, Request, Response
 from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel, Field
 
-from it_mental_health_environment import ITMentalHealthEnvironment, TASK_ORDER
+from it_mental_health_environment import ITMentalHealthEnvironment, TASK_DIFFICULTY, TASK_ORDER
 from models import MentalHealthAction
 
 app = FastAPI(
@@ -106,6 +106,9 @@ class StateResponse(BaseModel):
     current_task: str
     cumulative_reward: float
     tasks_completed: list
+    task_scores: Dict[str, float]
+    task_step_counts: Dict[str, int]
+    difficulty_scores: Dict[str, float]
     session_id: str
 
 
@@ -175,9 +178,9 @@ def reset(
     return ObservationResponse(
         scenario=obs.scenario,
         feedback=obs.feedback,
-        reward=obs.reward,
-        done=obs.done,
-        score_breakdown=obs.score_breakdown,
+        reward=0.0,
+        done=False,
+        score_breakdown={},
         task_id=obs.task_id,
         metadata=_response_metadata(obs.metadata, session_id),
     )
@@ -202,17 +205,17 @@ def step(
         metadata=req.metadata or {},
     )
     try:
-        obs = env.step(action)
+        obs, reward_model, done, info = env.step(action)
     except ValueError as exc:
         raise HTTPException(status_code=400, detail=str(exc)) from exc
     return ObservationResponse(
         scenario=obs.scenario,
         feedback=obs.feedback,
-        reward=obs.reward,
-        done=obs.done,
-        score_breakdown=obs.score_breakdown,
+        reward=reward_model.value,
+        done=done,
+        score_breakdown=reward_model.score_breakdown,
         task_id=obs.task_id,
-        metadata=_response_metadata(obs.metadata, session_id),
+        metadata=_response_metadata(info or obs.metadata, session_id),
     )
 
 
@@ -228,13 +231,16 @@ def state(
         header_session_id=x_session_id,
     ) or ANONYMOUS_SESSION_ID
     env = session_store.get_or_create(resolved_session_id)
-    s = env.state
+    s = env.state()
     return StateResponse(
         episode_id=s.episode_id,
         step_count=s.step_count,
         current_task=s.current_task,
         cumulative_reward=s.cumulative_reward,
         tasks_completed=s.tasks_completed,
+        task_scores=s.task_scores,
+        task_step_counts=s.task_step_counts,
+        difficulty_scores=s.difficulty_scores,
         session_id=resolved_session_id,
     )
 
@@ -248,6 +254,11 @@ def list_tasks():
                 "task_id": tid,
                 "difficulty": TASK_METADATA[tid]["difficulty"],
                 "description": TASK_METADATA[tid]["description"],
+                "scoring": {
+                    "step_score": "score for the current graded step in this task",
+                    "task_cumulative_score": "sum of all step scores recorded for this task",
+                    "difficulty_cumulative_score": "sum of all step scores for this difficulty bucket",
+                },
             }
             for tid in TASK_ORDER
         ]
@@ -270,6 +281,19 @@ def schema():
             "done": "bool",
             "score_breakdown": "dict - partial scores per rubric dimension",
             "task_id": "str",
+            "metadata": {
+                "difficulty": "str - easy / medium / hard for the graded step",
+                "step_score": "float - score for this step",
+                "task_step_number": "int - graded step count within this task",
+                "task_cumulative_score": "float - cumulative score for this task",
+                "difficulty_cumulative_score": "float - cumulative score for this difficulty level",
+                "overall_cumulative_score": "float - episode-wide cumulative score",
+            },
+        },
+        "state": {
+            "task_scores": "dict - cumulative score per task_id",
+            "task_step_counts": "dict - graded step count per task_id",
+            "difficulty_scores": "dict - cumulative score per difficulty bucket",
         },
     }