inference.py

"""
Inference Script Example
===================================
MANDATORY
- Before submitting, ensure the following variables are defined in your environment configuration:
    API_BASE_URL   The API endpoint for the LLM.
    MODEL_NAME     The model identifier to use for inference.
    HF_TOKEN       Your Hugging Face / API key.
    LOCAL_IMAGE_NAME The name of the local image to use for the environment if you are using from_docker_image()
                     method

- Defaults are set only for API_BASE_URL and MODEL_NAME 
    (and should reflect your active inference setup):
    API_BASE_URL = os.getenv("API_BASE_URL", "<your-active-endpoint>")
    MODEL_NAME = os.getenv("MODEL_NAME", "<your-active-model>")
    
- The inference script must be named `inference.py` and placed in the root directory of the project
- Participants must use OpenAI Client for all LLM calls using above variables

STDOUT FORMAT
- The script must emit exactly three line types to stdout, in this order:

    [START] task=<task_name> env=<benchmark> model=<model_name>
    [STEP]  step=<n> action=<action_str> reward=<0.00> done=<true|false> error=<msg|null>
    [END]   success=<true|false> steps=<n> score=<score> rewards=<r1,r2,...,rn>

  Rules:
    - One [START] line at episode begin.
    - One [STEP] line per step, immediately after env.step() returns.
    - One [END] line after env.close(), always emitted (even on exception).
    - reward and rewards are formatted to 2 decimal places.
    - done and success are lowercase booleans: true or false.
    - error is the raw last_action_error string, or null if none.
    - All fields on a single line with no newlines within a line.
    - Each tasks should return score in [0, 1]

  Example:
    [START] task=click-test env=miniwob model=Qwen3-VL-30B
    [STEP] step=1 action=click('123') reward=0.00 done=false error=null
    [STEP] step=2 action=fill('456','text') reward=0.00 done=false error=null
    [STEP] step=3 action=click('789') reward=1.00 done=true error=null
    [END] success=true steps=3 score=1.00 rewards=0.00,0.00,1.00
"""


import asyncio
import os
import textwrap
from typing import List, Optional

from openai import OpenAI

from client import Mentalhealthpatientenv
from models import MentalhealthpatientenvAction


# =========================
# ENV VARIABLES (MANDATORY)
# =========================
API_KEY = os.getenv("HF_TOKEN") or os.getenv("API_KEY")
API_BASE_URL = os.getenv("API_BASE_URL") or "https://router.huggingface.co/v1"
MODEL_NAME = os.getenv("MODEL_NAME") or "Qwen/Qwen2.5-72B-Instruct"

TASK_NAME = "mental_health"
BENCHMARK = "mentalHealthPatientenv"

MAX_STEPS = 20
TEMPERATURE = 0.7
MAX_TOKENS = 120

SUCCESS_SCORE_THRESHOLD = 0.6


# =========================
# LOGGING (STRICT FORMAT)
# =========================
def log_start(task: str, env: str, model: str):
    print(f"[START] task={task} env={env} model={model}", flush=True)


def log_step(step: int, action: str, reward: float, done: bool, error: Optional[str]):
    error_val = error if error else "null"
    print(
        f"[STEP] step={step} action={action} reward={reward:.2f} done={str(done).lower()} error={error_val}",
        flush=True,
    )


def log_end(success: bool, steps: int, score: float, rewards: List[float]):
    rewards_str = ",".join(f"{r:.2f}" for r in rewards)
    print(
        f"[END] success={str(success).lower()} steps={steps} score={score:.2f} rewards={rewards_str}",
        flush=True,
    )


# =========================
# PROMPT
# =========================
SYSTEM_PROMPT = """
You are a mental health therapist interacting with a patient.

Available actions:
- ask_open
- ask_direct
- ask_risk
- reflect
- diagnose

Respond ONLY in format:
action_type|message
"""


def build_prompt(observation, step):
    return f"""
Step: {step}

Patient says: {observation.response}
Trust level: {observation.trust_level}
Emotional state: {observation.emotional_state}
Risk flag: {observation.risk_flag}

What should you do next?
"""


# =========================
# LLM CALL (HF ROUTER)
# =========================
def get_action(client, observation, step):
    prompt = build_prompt(observation, step)

    try:
        completion = client.chat.completions.create(
            model=MODEL_NAME,
            messages=[
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user", "content": prompt},
            ],
            temperature=TEMPERATURE,
            max_tokens=MAX_TOKENS,
        )

        text = (completion.choices[0].message.content or "").strip()

        if "|" in text:
            action_type, message = text.split("|", 1)
        else:
            action_type = "ask_open"
            message = text

        return action_type.strip(), message.strip()

    except Exception:
        return "ask_open", "How have you been feeling lately?"


# =========================
# MAIN LOOP
# =========================
async def main():
    client = OpenAI(base_url=API_BASE_URL, api_key=API_KEY)

    env = Mentalhealthpatientenv(base_url="http://localhost:8000")

    rewards = []
    steps_taken = 0
    success = False

    log_start(TASK_NAME, BENCHMARK, MODEL_NAME)

    try:
        result = await env.reset()

        for step in range(1, MAX_STEPS + 1):
            if result.done:
                break

            obs = result.observation

            action_type, message = get_action(client, obs, step)

            action = MentalhealthpatientenvAction(
                action_type=action_type,
                message=message
            )

            result = await env.step(action)

            reward = result.reward or 0.0
            done = result.done
            error = None

            rewards.append(reward)
            steps_taken = step

            log_step(step, f"{action_type}|{message}", reward, done, error)

            if done:
                break

        # Normalize score [0,1]
        score = sum(rewards) / len(rewards) if rewards else 0.0
        score = max(0.0, min(1.0, score))

        success = score >= SUCCESS_SCORE_THRESHOLD

    finally:
        try:
            await env.close()
        except:
            pass

        log_end(success, steps_taken, score, rewards)


if __name__ == "__main__":
    asyncio.run(main())