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 finenv.client import FinenvEnv
from finenv.models import FinenvAction

# ==============================
# ENV VARIABLES (MANDATORY)
# ==============================
IMAGE_NAME = os.getenv("LOCAL_IMAGE_NAME", "finenv")
API_KEY = os.getenv("HF_TOKEN") or os.getenv("OPENAI_API_KEY")

API_BASE_URL = os.getenv("API_BASE_URL")
MODEL_NAME = os.getenv("MODEL_NAME")

TASK_NAME = os.getenv("FINENV_TASK", "easy")
BENCHMARK = "finenv"

MAX_STEPS = 15
SUCCESS_SCORE_THRESHOLD = 0.2  # normalized score in [0, 1]

SYSTEM_PROMPT = textwrap.dedent(
    """
    You are interacting with a stock trading environment. 
    Your goal is to maximize profit by buying, selling, or holding shares of a stock over a series of steps.
    At each step, you can choose one of the following actions:  
    - buy: purchase 1 share of the stock at the current price 
    - sell: sell 1 share of the stock at the current price (only if you have shares to sell)
    - hold: take no action
    The environment will provide feedback in the form of rewards based on the change in your portfolio value.
    Your objective is to achieve the highest possible return by the end of the episode.
    """
).strip()

# ==============================
# LOGGING (STRICT FORMAT)
# ==============================
def log_start(task: str, env: str, model: str) -> None:
    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]) -> None:
    error_val = error if error else "null"
    done_val = str(done).lower()
    print(
        f"[STEP] step={step} action={action} reward={reward:.2f} done={done_val} error={error_val}",
        flush=True,
    )


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


# ==============================
# MODEL DECISION LOGIC
# ==============================
def get_model_action(client: OpenAI, step: int, last_reward: float, history: List[str]) -> str:
    """
    Uses LLM to decide trading action.
    """

    prompt = f"""
You are a trading agent.

Step: {step}
Last reward: {last_reward}
Recent history:
{history[-3:] if history else "None"}

Choose ONE:
buy
sell
hold

Respond with only one word.
"""

    try:
        response = client.chat.completions.create(
            model=MODEL_NAME,
            messages=[{"role": "user", "content": prompt}],
            temperature=0.2,
        )

        action = (response.choices[0].message.content or "").strip().lower()

        if action not in ["buy", "sell", "hold"]:
            return "hold"

        return action

    except Exception as exc:
        print(f"[DEBUG] Model request failed: {exc}", flush=True)
        return "hold"


# ==============================
# MAIN EXECUTION
# ==============================
async def main() -> None:
    client = OpenAI(base_url=API_BASE_URL, api_key=API_KEY)

    env = await FinenvEnv.from_docker_image(IMAGE_NAME)

    history: List[str] = []
    rewards: List[float] = []

    steps_taken = 0
    score = 0.0
    success = False

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

    try:
        # Reset environment
        result = await env.reset()
        last_reward = 0.0

        # INIT STEP (MANDATORY for your env)
        init_action = FinenvAction(
            type="init",
            stock="RELIANCE",
            market="NSE",
            initial_cash=10000,
            max_steps=MAX_STEPS
        )

        result = await env.step(init_action)

        # Log init step as step 0
        log_step(
            step=0,
            action="init",
            reward=0.00,
            done=False,
            error=None
        )

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

            action_type = get_model_action(client, step, last_reward, history)

            action = FinenvAction(
                type=action_type,
                quantity=1
            )

            error = None

            try:
                result = await env.step(action)
            except Exception as e:
                error = str(e)
                result = result  # keep last safe state

            reward = float(result.reward or 0.0)
            done = result.done

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

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

            history.append(f"{action_type}:{reward:.2f}")

            if done:
                break

        # ==============================
        # SCORE CALCULATION
        # ==============================
        if len(rewards) > 0:
            score = sum(rewards) / len(rewards)
        else:
            score = 0.0

        score = min(max(score, 0.0), 1.0)
        success = score >= SUCCESS_SCORE_THRESHOLD

    finally:
        try:
            await env.close()
        except Exception as e:
            print(f"[DEBUG] env.close() error: {e}", flush=True)

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


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