inference.py

#!/usr/bin/env python3
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.

"""
FitScript Baseline Inference Script.

Runs an LLM agent against all 3 FitScript tasks and reports scores.
Emits structured [START] / [STEP] / [END] logs for automated evaluation.

Environment variables required:
    API_BASE_URL   — LLM API base URL (OpenAI-compatible endpoint)
    MODEL_NAME     — Model identifier (e.g. gpt-4o, meta-llama/...)
    HF_TOKEN       — Hugging Face / API key

Usage:
    export API_BASE_URL="https://api.openai.com/v1"
    export MODEL_NAME="gpt-4o"
    export HF_TOKEN="sk-..."
    python inference.py
"""

import asyncio
import json
import os
import sys
import time
from typing import List

from openai import OpenAI

# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------

API_BASE_URL: str = os.environ.get("API_BASE_URL", "https://api.openai.com/v1")
MODEL_NAME: str = os.environ.get("MODEL_NAME", "gpt-4o")
HF_TOKEN: str = os.environ.get("HF_TOKEN", "")
ENV_BASE_URL: str = os.environ.get("ENV_BASE_URL", "http://localhost:8000")

BENCHMARK = "FitScript-v1"
MAX_STEPS = 3
SUCCESS_SCORE_THRESHOLD = 0.60

# Task names for logging
TASK_NAMES = {
    1: "basic_safe_prescription",
    2: "injury_constraints",
    3: "multi_client_allocation",
}

# Maximum possible reward per step × max_steps (all tasks use 3 steps)
MAX_TOTAL_REWARD = MAX_STEPS  # each step reward is in [0,1]; we average at the end
# Per-task server ports (run 3 separate server instances)
TASK_PORTS = {1: 8001, 2: 8002, 3: 8003}

# ---------------------------------------------------------------------------
# Structured logging helpers  (STRICT format — do not modify field names)
# ---------------------------------------------------------------------------

def log_start(task: str, env: str, model: str) -> None:
    print(
        json.dumps({
            "type": "START",
            "task": task,
            "env": env,
            "model": model,
            "timestamp": time.time(),
        }),
        flush=True,
    )


def log_step(step: int, action: str, reward: float, done: bool, error) -> None:
    print(
        json.dumps({
            "type": "STEP",
            "step": step,
            "action": action[:200],   # truncate for log readability
            "reward": reward,
            "done": done,
            "error": str(error) if error else None,
            "timestamp": time.time(),
        }),
        flush=True,
    )


def log_end(success: bool, steps: int, score: float, rewards: List[float]) -> None:
    print(
        json.dumps({
            "type": "END",
            "success": success,
            "steps": steps,
            "score": score,
            "rewards": rewards,
            "timestamp": time.time(),
        }),
        flush=True,
    )


# ---------------------------------------------------------------------------
# System prompt for the agent
# ---------------------------------------------------------------------------

SYSTEM_PROMPT = """You are an expert certified fitness coach and registered dietitian.
Your task is to create complete, safe, effective, and highly personalized fitness
prescriptions for clients. Always:
- Address every specific constraint mentioned (injuries, medical conditions, equipment limits)
- Include exact workout splits, exercises, sets, reps, and rest periods
- Include caloric and macronutrient targets with rationale
- Provide progressive overload strategies
- Flag any safety concerns and recommend physician clearance when appropriate
- Tailor the plan specifically to the client — avoid generic templates
Be thorough, specific, and safety-first."""


def build_user_prompt(scenario: str, feedback: str, step: int, last_reward: float) -> str:
    """Build the user prompt for each step."""
    if step == 1:
        return (
            f"Please provide a complete fitness prescription for the following client:\n\n"
            f"{scenario}\n\n"
            f"Be thorough and address every aspect of the task."
        )
    else:
        return (
            f"CLIENT SCENARIO (same as before):\n{scenario}\n\n"
            f"PREVIOUS ATTEMPT FEEDBACK (reward={last_reward:.3f}):\n{feedback}\n\n"
            f"Please revise your prescription to address the failed checks above. "
            f"Be specific and comprehensive. Step {step}/{MAX_STEPS}."
        )


# ---------------------------------------------------------------------------
# LLM call
# ---------------------------------------------------------------------------

def get_model_message(
    client: OpenAI,
    scenario: str,
    feedback: str,
    step: int,
    last_reward: float,
    history: List[str],
) -> str:
    """Call the LLM and return its prescription text."""
    user_prompt = build_user_prompt(scenario, feedback, step, last_reward)

    messages = [{"role": "system", "content": SYSTEM_PROMPT}]

    # Include history as assistant context
    for i, h in enumerate(history):
        messages.append({"role": "assistant", "content": h})

    messages.append({"role": "user", "content": user_prompt})

    try:
        response = client.chat.completions.create(
            model=MODEL_NAME,
            messages=messages,
            max_tokens=2000,
            temperature=0.7,
        )
        return response.choices[0].message.content.strip()
    except Exception as e:
        print(f"[DEBUG] LLM call error: {e}", flush=True)
        return f"Error generating prescription: {e}"


# ---------------------------------------------------------------------------
# Run one task episode
# ---------------------------------------------------------------------------

async def run_task(client: OpenAI, task_id: int, env_client) -> float:
    """
    Run a single task episode against the FitScript environment.
    Returns the normalized score for this task (0.0–1.0).
    """
    from FitScript import FitscriptAction

    task_name = TASK_NAMES[task_id]
    rewards: List[float] = []
    steps_taken = 0
    score = 0.0
    success = False
    history: List[str] = []

    log_start(task=task_name, env=BENCHMARK, model=MODEL_NAME)

    try:
        result = await env_client.reset()
        scenario = result.observation.client_scenario
        feedback = ""
        last_reward = 0.0

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

            # Get LLM prescription
            prescription = get_model_message(
                client, scenario, feedback, step, last_reward, history
            )

            # Step the environment
            result = await env_client.step(FitscriptAction(message=prescription))
            obs = result.observation

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

            rewards.append(reward)
            steps_taken = step
            feedback = obs.feedback
            last_reward = reward

            log_step(step=step, action=prescription, reward=reward, done=done, error=error)

            history.append(prescription)

            if done:
                break

        # Score = mean reward across steps (normalized to [0, 1])
        score = sum(rewards) / len(rewards) if rewards else 0.0
        score = min(max(score, 0.0), 1.0)
        success = score >= SUCCESS_SCORE_THRESHOLD

    except Exception as e:
        print(f"[DEBUG] Task {task_id} error: {e}", flush=True)
        error = str(e)
        log_step(step=steps_taken + 1, action="", reward=0.0, done=True, error=error)
        # Still compute score from rewards collected before error
        if rewards:
            score = sum(rewards) / len(rewards)
            score = min(max(score, 0.0), 1.0)
            success = score >= SUCCESS_SCORE_THRESHOLD

    finally:
        log_end(success=success, steps=steps_taken, score=score, rewards=rewards)

    return score


# ---------------------------------------------------------------------------
# Main entrypoint
# ---------------------------------------------------------------------------

async def main() -> None:
    if not HF_TOKEN:
        print("[ERROR] HF_TOKEN environment variable not set.", file=sys.stderr)
        sys.exit(1)

    # Initialize OpenAI-compatible client
    client = OpenAI(
        api_key=HF_TOKEN,
        base_url=API_BASE_URL,
    )

    print(f"[INFO] FitScript Baseline Inference", flush=True)
    print(f"[INFO] Model:   {MODEL_NAME}", flush=True)
    print(f"[INFO] API URL: {API_BASE_URL}", flush=True)
    print(f"[INFO] Env URL: {ENV_BASE_URL}", flush=True)
    print(flush=True)

    task_scores = {}

    for task_id in [1, 2, 3]:
        print(f"[INFO] ===== Running Task {task_id}: {TASK_NAMES[task_id]} =====", flush=True)

        try:
            from FitScript import FitscriptEnv

            task_url = f"http://localhost:{TASK_PORTS[task_id]}"
            async with FitscriptEnv(base_url=task_url) as env_client:
                score = await run_task(client, task_id, env_client)

        except Exception as e:
            print(f"[DEBUG] Failed to connect to env for task {task_id}: {e}", flush=True)
            # Fallback: run against environment directly (no server needed)
            score = await run_task_local(client, task_id)

        task_scores[task_id] = score
        print(f"[INFO] Task {task_id} score: {score:.4f}", flush=True)
        print(flush=True)

    # Final summary
    overall = sum(task_scores.values()) / len(task_scores)
    print("=" * 60, flush=True)
    print(f"FINAL RESULTS — {BENCHMARK}", flush=True)
    print(f"  Task 1 (Easy):   {task_scores.get(1, 0.0):.4f}", flush=True)
    print(f"  Task 2 (Medium): {task_scores.get(2, 0.0):.4f}", flush=True)
    print(f"  Task 3 (Hard):   {task_scores.get(3, 0.0):.4f}", flush=True)
    print(f"  Overall:         {overall:.4f}", flush=True)
    print("=" * 60, flush=True)


async def run_task_local(client: OpenAI, task_id: int) -> float:
    """
    Fallback: run a task locally without a server (direct environment import).
    Used when the HTTP server is not reachable.
    """
    import sys
    import os
    sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

    from server.FitScript_environment import FitscriptEnvironment
    from models import FitscriptAction

    task_name = TASK_NAMES[task_id]
    env = FitscriptEnvironment(task_id=task_id)

    rewards: List[float] = []
    steps_taken = 0
    score = 0.0
    success = False
    history: List[str] = []

    log_start(task=task_name, env=BENCHMARK, model=MODEL_NAME)

    try:
        obs = env.reset()
        scenario = obs.client_scenario
        feedback = ""
        last_reward = 0.0
        done = obs.done

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

            prescription = get_model_message(
                client, scenario, feedback, step, last_reward, history
            )

            obs = env.step(FitscriptAction(message=prescription))

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

            rewards.append(reward)
            steps_taken = step
            feedback = obs.feedback
            last_reward = reward

            log_step(step=step, action=prescription, reward=reward, done=done, error=error)

            history.append(prescription)

            if done:
                break

        score = sum(rewards) / len(rewards) if rewards else 0.0
        score = min(max(score, 0.0), 1.0)
        success = score >= SUCCESS_SCORE_THRESHOLD

    except Exception as e:
        print(f"[DEBUG] Local task {task_id} error: {e}", flush=True)

    finally:
        log_end(success=success, steps=steps_taken, score=score, rewards=rewards)

    return score


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