Fix Phase 2 validation: Integrate LiteLLM proxy and implement LLM-powered baseline inference

PROJECT_SUMMARY.md

@@ -26,7 +26,7 @@ Implemented three specific scenarios with automated graders:
     - `/baseline`: Automated trigger for the inference script.
 
 ### 4. Baseline Inference (`inference.py`)
-- Created a reproducible baseline script that uses medical heuristics to evaluate the environment. It ensures the environment is "solveable" and provides a benchmark score for comparison.
+- Created a reproducible baseline script that leverages the Hackathon's LiteLLM proxy to evaluate the environment. It ensures the environment is "solveable" using a real LLM and serves as a benchmark for API-based agent interactions.
 
 ## 🚀 Technical Improvements & Fixes
 To ensure the project meets the highest submission standards, we performed the following:

client.py

@@ -6,12 +6,5 @@ from openenv.core.mcp_client import MCPToolClient
 class MedTriageEnv(MCPToolClient):
     """
     Client for the MedTriage Environment.
-    
-    Example:
-        >>> with MedTriageEnv(base_url="http://localhost:8000") as env:
-        ...     obs = env.reset(task_id="TASK_HARD")
-        ...     print(obs.symptoms_text)
-        ...     result = env.call_tool("triage_patient", level=3, reasoning="High BP and atypical symptoms in elderly patient.")
-        ...     print(f"Reward: {result.reward}")
     """
     pass

inference.py

@@ -7,13 +7,21 @@ import json
 import requests
 import asyncio
 from typing import List, Dict, Any
+import openai
 from client import MedTriageEnv
 
-# --- Mandatory Environment Configuration ---
+# --- Mandatory Environment Configuration (Injected by Hackathon Validator) ---
 API_BASE_URL = os.getenv("API_BASE_URL", "http://localhost:8002/v1")
+API_KEY = os.getenv("API_KEY", "dummy-key")
 MODEL_NAME = os.getenv("MODEL_NAME", "med-triage-baseline")
 HF_TOKEN = os.getenv("HF_TOKEN", "dummy-token")
 
+# Initialize OpenAI client with the injected LiteLLM proxy credentials
+llm_client = openai.OpenAI(
+    base_url=API_BASE_URL,
+    api_key=API_KEY
+)
+
 # --- Structured Logging Functions ---
 def log_start(task: str, env: str, model: str):
     """Emit structured [START] log."""
@@ -68,6 +76,67 @@ def wait_for_server(url: str, timeout: int = 30):
         time.sleep(1)
     return False
 
+def get_llm_triage(obs: Any) -> Dict[str, Any]:
+    """Call the LLM via proxy to get triage level and reasoning."""
+    # Heuristic fallback for local development when proxy is not active
+    is_local = "localhost" in API_BASE_URL and API_KEY == "dummy-key"
+    
+    # Extract data safely from observation
+    # Sometimes openenv-core puts custom fields in .metadata if they don't fit the schema
+    metadata = getattr(obs, "metadata", {})
+    
+    age = getattr(obs, "age", metadata.get("age", 0))
+    gender = getattr(obs, "gender", metadata.get("gender", "unknown"))
+    symptoms = getattr(obs, "symptoms_text", metadata.get("symptoms_text", "No symptoms provided"))
+    vitals = getattr(obs, "vitals", metadata.get("vitals", {}))
+    history = getattr(obs, "history", metadata.get("history", []))
+
+    if is_local:
+        # Heuristic Logic (mimes model output for local testing)
+        bp_sys = int(vitals.get("bp", "120/80").split("/")[0]) if isinstance(vitals.get("bp"), str) else 120
+        if bp_sys > 150 or age > 65:
+            level = 3 # EMERGENCY
+        elif "severe pain" in symptoms.lower():
+            level = 2 # URGENT_CARE
+        else:
+            level = 0 # SELF_CARE
+        return {"level": level, "reasoning": "Heuristic fallback (Local testing)."}
+
+    # Actual LLM call through the LiteLLM proxy
+    prompt = f"""
+    You are an expert medical triage officer. Based on the patient data below, categorize the patient into the most appropriate care level.
+
+    PATIENT DATA:
+    - Age: {age}
+    - Gender: {gender}
+    - Symptoms: {symptoms}
+    - Vitals: {vitals}
+    - Medical History: {history}
+
+    TRIAGE LEVELS (0-3):
+    0: SELF_CARE (Rest, OTC medication, home monitoring)
+    1: CLINIC (Non-urgent appointment within 48 hours)
+    2: URGENT_CARE (Immediate attention for non-life-threatening conditions)
+    3: EMERGENCY (Life-threatening symptoms, immediate ER visit)
+
+    Assign a level (0, 1, 2, or 3) and provide a concise medical reasoning.
+    Respond ONLY in valid JSON format:
+    {{"level": <int>, "reasoning": "<string>"}}
+    """
+
+    try:
+        response = llm_client.chat.completions.create(
+            model=MODEL_NAME,
+            messages=[{"role": "user", "content": prompt}],
+            response_format={"type": "json_object"}
+        )
+        result = json.loads(response.choices[0].message.content)
+        return result
+    except Exception as e:
+        print(f"LLM Error encountered: {e}")
+        # Final safety fallback: triage to Emergency if unsure/error
+        return {"level": 3, "reasoning": f"Emergency fallback due to triage system error: {str(e)}"}
+
 def run_baseline(base_url: str = "http://localhost:8002"):
     """Run baseline agent against all 3 tasks and return results."""
     # Ensure base_url uses the port from environment if available
@@ -96,32 +165,49 @@ def run_baseline(base_url: str = "http://localhost:8002"):
         
         try:
             with MedTriageEnv(base_url=base_url).sync() as env:
-                obs = env.reset(task_id=task_id)
+                res = env.reset(task_id=task_id)
+                obs = res.observation
                 steps_taken = 1
                 
-                # Heuristic Logic (mimes model output)
-                bp_sys = int(obs.vitals.get("bp", "120/80").split("/")[0])
-                if bp_sys > 150 or obs.age > 65:
-                    level = 3 # EMERGENCY
-                elif "severe pain" in obs.symptoms_text.lower():
-                    level = 2 # URGENT_CARE
+                # Import CallToolAction for proper serialization
+                try:
+                    from openenv.core.env_server.mcp_types import CallToolAction
+                except ImportError:
+                    # Fallback if needed, but it should be available
+                    CallToolAction = None
+
+                decision = get_llm_triage(obs)
+                level = int(decision.get("level", 0))
+                reasoning = decision.get("reasoning", "Baseline triage decision.")
+                
+                if CallToolAction:
+                    action = CallToolAction(
+                        tool_name="triage_patient", 
+                        arguments={"level": level, "reasoning": reasoning}
+                    )
                 else:
-                    level = 0 # SELF_CARE
+                    action = {"tool_name": "triage_patient", "arguments": {"level": level, "reasoning": reasoning}}
                 
-                action = {"tool_name": "triage_patient", "arguments": {"level": level, "reasoning": "Heuristic baseline."}}
                 result = env.step(action)
                 
                 reward = result.reward or 0.0
                 done = result.done
                 rewards.append(reward)
                 
-                log_step(step=1, action=action, reward=reward, done=done)
+                if CallToolAction:
+                    action_to_log = action.model_dump() if hasattr(action, "model_dump") else action
+                else:
+                    action_to_log = action
+                
+                log_step(step=1, action=action_to_log, reward=reward, done=done)
                 
                 final_score = reward
                 success = reward >= 1.0
                 all_scores[task_id] = reward
                 
         except Exception as e:
+            import traceback
+            traceback.print_exc()
             log_step(step=steps_taken, action=None, reward=0.0, done=True, error=str(e))
             all_scores[task_id] = 0.0
             

pyproject.toml

@@ -11,6 +11,7 @@ dependencies = [
     "pydantic",
     "fastmcp",
     "requests",
+    "openai",
     "openenv-core>=0.2.0",
 ]
 

requirements.txt

@@ -3,4 +3,5 @@ uvicorn
 pydantic
 fastmcp
 requests
+openai
 openenv-core>=0.2.0

server/app.py

@@ -9,14 +9,14 @@ from typing import List, Dict, Any
 
 try:
     from server.triage_environment import MedTriageEnvironment, TASKS
-    from models import TriageAction
+    from models import TriageAction, TriageObservation
 except ImportError:
     try:
         from .triage_environment import MedTriageEnvironment, TASKS
-        from .models import TriageAction
+        from .models import TriageAction, TriageObservation
     except ImportError:
         from triage_environment import MedTriageEnvironment, TASKS
-        from models import TriageAction
+        from models import TriageAction, TriageObservation
 
 # Initialize the environment instance to be used by the app
 env_instance = MedTriageEnvironment()
@@ -25,7 +25,7 @@ env_instance = MedTriageEnvironment()
 app = create_app(
     MedTriageEnvironment, 
     CallToolAction, 
-    CallToolObservation, 
+    TriageObservation, 
     env_name="med_triage_env"
 )
 

server/requirements.txt

@@ -3,4 +3,5 @@ uvicorn
 pydantic
 fastmcp
 requests
+openai
 openenv-core>=0.2.0

server/triage_environment.py

@@ -126,6 +126,9 @@ class MedTriageEnvironment(MCPEnvironment):
         """
         Process the agent's triage decision and return a score.
         """
+        print(f"DEBUG: Received action type: {type(action)}")
+        if hasattr(action, "tool_name"):
+            print(f"DEBUG: tool_name: {action.tool_name}")
         self._state.step_count += 1
         
         # If the action is an MCP CallToolAction
@@ -133,12 +136,18 @@ class MedTriageEnvironment(MCPEnvironment):
         
         if isinstance(action, CallToolAction) and action.tool_name == "triage_patient":
             agent_level = action.arguments.get("level")
-            reward = self._calculate_reward(TriageLevel(agent_level), self._state.ground_truth_level)
+            reward = self._calculate_reward(TriageLevel(int(agent_level)), self._state.ground_truth_level)
             self._last_reward = reward
             
             patient = self._current_task["patient"]
+            # Ensure we return the model type expected by the app
             return TriageObservation(
-                **patient,
+                patient_id=patient["patient_id"],
+                age=patient["age"],
+                gender=patient["gender"],
+                symptoms_text=patient["symptoms_text"],
+                vitals=patient["vitals"],
+                history=patient["history"],
                 done=True,
                 reward=reward,
                 message=f"Episode complete. Agent Triage: {agent_level}. Ground Truth: {self._state.ground_truth_level.value}. Score: {reward}"