
# Allow direct execution from subdirectories while keeping imports package-relative.
import sys
from pathlib import Path

_REPO_ROOT = Path(__file__).resolve().parents[1]
if str(_REPO_ROOT) not in sys.path:
    sys.path.insert(0, str(_REPO_ROOT))

# ---
# pytest: false
# ---

# # Run OpenAI's gpt-oss model with vLLM

# ## Background

# [gpt-oss](https://openai.com/index/introducing-gpt-oss/) is a reasoning model
# that comes in two flavors: `gpt-oss-120B` and `gpt-oss-20B`. They are both Mixture
# of Experts (MoE) models with a low number of active parameters, ensuring they
# combine good world knowledge and capabilities with fast inference.

# We describe a few of its notable features below.

# ### MXFP4

# OpenAI's gpt-oss models use a fairly uncommon 4bit [`mxfp4`](https://arxiv.org/abs/2310.10537) floating point
# format for the MoE layers. This "block" quantization format combines `e2m1` floating point numbers
# with blockwise scaling factors. The attention operations are not quantized.

# ### Attention Sinks

# Attention sink models allow for longer context lengths without sacrificing output quality. The vLLM team
# added [attention sink support](https://huggingface.co/kernels-community/vllm-flash-attn3)
# for Flash Attention 3 (FA3) in preparation for this release.

# ### Response Format

# GPT-OSS is trained with the [harmony response format](https://github.com/openai/harmony) which enables models
# to output to multiple channels for chain-of-thought (CoT) and input tool-calling preambles along with regular text responses.
# We'll stick to a simpler format here, but see [this cookbook](https://cookbook.openai.com/articles/openai-harmony)
# for details on the new format.

# ## Set up the container image

# We'll start by defining a [custom container `Image`](https://modal.com/docs/guide/custom-container) that
# installs all the necessary dependencies to run vLLM and the model.

import json
import os
import time
from datetime import datetime, timezone
from typing import Any, List

import aiohttp
import modal

from openai_harmony import (
    Conversation,
    DeveloperContent,
    HarmonyEncodingName,
    Message,
    ReasoningEffort,
    Role,
    SystemContent,
    RenderConversationConfig,
    load_harmony_encoding,
)

# Load harmony encoding once at module level
HARMONY_ENC = load_harmony_encoding(HarmonyEncodingName.HARMONY_GPT_OSS)

vllm_image = (
    modal.Image.from_registry(
        "nvidia/cuda:12.8.1-devel-ubuntu22.04",
        add_python="3.12",
    )
    .entrypoint([])
    .uv_pip_install(
        "vllm==0.13.0",
        "huggingface_hub[hf_transfer]==0.36.0",
    )
    .env({"VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8": "1"})
)


# ## Download the model weights

# We'll be downloading OpenAI's model from Hugging Face. We're running
# the 20B parameter model by default but you can easily switch to [the 120B model](https://huggingface.co/openai/gpt-oss-120b),
# which also fits in a single H100 or H200 GPU.

MODEL_NAME = os.environ.get("HARNESS1_HF_MODEL", "harness-1")
MODEL_REVISION = os.environ.get("HARNESS1_HF_REVISION", "main")

# Although vLLM will download weights from Hugging Face on-demand, we want to
# cache them so we don't do it every time our server starts. We'll use [Modal Volumes](https://modal.com/docs/guide/volumes)
# for our cache. Modal Volumes are essentially a "shared disk" that all Modal
# Functions can access like it's a regular disk. For more on storing model
# weights on Modal, see [this guide](https://modal.com/docs/guide/model-weights).

hf_cache_vol = modal.Volume.from_name("huggingface-cache", create_if_missing=True)

# The first time you run a new model or configuration with vLLM on a fresh machine,
# a number of artifacts are created. We also cache these artifacts.

vllm_cache_vol = modal.Volume.from_name("vllm-cache", create_if_missing=True)
flashinfer_cache_vol = modal.Volume.from_name(
    "flashinfer-cache", create_if_missing=True
)

# ## Configuring vLLM to serve GPT-OSS

# The vLLM docs include an [excellent resource on tuning GPT-OSS](https://docs.vllm.ai/projects/recipes/en/latest/OpenAI/GPT-OSS.html).
# We mostly use the configuration values reported there, but try to explain the reasoning as we go.

VLLM_CONFIG = {  # return tokens in chunks of 20, save on host overhead
    "stream-interval": 20
}

# One of the most important choices is to use speculative decoding,
# which attempts to generate multiple tokens per forward pass
# by means of a separate "speculator" model.
# We here use RedHatAI's open source, generic EAGLE3-based speculator for this model.
# We recommend using the EAGLE3 technique to train a custom speculator on your own traffic.

SPECULATIVE_CONFIG = {
    "model": "RedHatAI/gpt-oss-20b-speculator.eagle3",
    "num_speculative_tokens": 7,
    "method": "eagle3",
}

# Speculative decoding acclerates inference without changing model behavior.
# We can also accelerate inference by further quantizing the model.
# Here, we reduce the size of KV cache entries by quantizing them to FP8.

VLLM_CONFIG |= {"kv-cache-dtype": "fp8"}

# There are a number of compilation settings for vLLM. Compilation improves inference performance
# but incurs extra latency at engine start time. When iterating on and developing a server,
# we recommend turning compilation off to speed up development cycles, which we here control
# with a global variable.

FAST_BOOT = False

# Otherwise, we use the values suggested in the recipe:

COMPILATION_CONFIG = {
    "pass_config": {"fuse_allreduce_rms": True, "eliminate_noops": True}
}

# As part of compilation, vLLM collects up sequences (really, DAGs)
# of CUDA kernel launches into CUDA graphs.
# We set the maximum batch size for the CUDA graph capture step to the
# maximum number of inputs we want to handle per replica,
# which also shows up in our autoscaling configuration below.

MAX_INPUTS = 32  # how many requests can one replica handle? tune carefully!
VLLM_CONFIG |= {"max-cudagraph-capture-size": MAX_INPUTS}

# Lastly, there are a few knobs we can tune based on the typical lengths
# of sequences we expect to observe.
# For many agentic tasks to which this model is well-suited,
# those lengths can go into the tens of thousands of tokens.
# Let's assume they're never longer than 2 ^ 15 tokens.

VLLM_CONFIG |= {
    "max-num-batched-tokens": 16384,
    "max-model-len": 32768,
}

# ## Build a vLLM engine and serve it

# The function below spawns a vLLM instance listening at port 8000, serving requests to our model.

app = modal.App("example-gpt-oss-inference")

N_GPU = 4  # increase for more tensor parallelism (2, 4, or 8 GPUs)
MINUTES = 60  # seconds
VLLM_PORT = 8000


@app.function(
    image=vllm_image,
    gpu=f"B200:{N_GPU}",
    scaledown_window=10 * MINUTES,  # how long should we stay up with no requests?
    timeout=30 * MINUTES,  # how long should we wait for container start?
    volumes={
        "/root/.cache/huggingface": hf_cache_vol,
        "/root/.cache/vllm": vllm_cache_vol,
        "/root/.cache/flashinfer": flashinfer_cache_vol,
    },
)
@modal.concurrent(max_inputs=MAX_INPUTS)
@modal.web_server(port=VLLM_PORT, startup_timeout=30 * MINUTES)
def serve():
    import subprocess

    cmd = [
        "vllm",
        "serve",
        MODEL_NAME,
        "--revision",
        MODEL_REVISION,
        "--served-model-name",
        "llm",  # serve as "llm" for simpler API calls
        "--host",
        "0.0.0.0",
        "--port",
        str(VLLM_PORT),
        "--async-scheduling",  # reduces host overhead, but might not be compatible with all features
    ]

    # enforce-eager disables both Torch compilation and CUDA graph capture
    # default is no-enforce-eager. see the --compilation-config flag for tighter control
    cmd += ["--enforce-eager" if FAST_BOOT else "--no-enforce-eager"]

    # assume multiple GPUs are for splitting up large matrix multiplications
    cmd += ["--tensor-parallel-size", str(N_GPU)]

    # add complex configuration objects
    cmd += ["--compilation-config", json.dumps(COMPILATION_CONFIG)]
    # cmd += ["--speculative-config", json.dumps(SPECULATIVE_CONFIG)]

    cmd += [  # add assorted config
        item for k, v in VLLM_CONFIG.items() for item in (f"--{k}", str(v))
    ]

    print(*cmd)

    subprocess.Popen(cmd)


# ## Deploy the server

# To deploy the API on Modal, just run

# ```bash
# modal deploy gpt_oss_inference.py
# ```

# This will create a new app on Modal, build the container image for it if it hasn't been built yet,
# and deploy the app.

# ## Test the server

# To make it easier to test the server setup, we also include a `local_entrypoint`
# that does a healthcheck and then hits the server.

# If you execute the command

# ```bash
# modal run gpt_oss_inference.py
# ```

# a fresh replica of the server will be spun up on Modal while
# the code below executes on your local machine.

# We set up the system prompt with low reasoning effort to run
# inference a bit faster. For the best ergonomics we recommend using
# the [harmony API](https://cookbook.openai.com/articles/openai-harmony#example-system-message),
# which can be installed with `pip install openai-harmony`.


@app.local_entrypoint()
async def test(test_timeout=30 * MINUTES, user_content=None, twice=True):
    url = serve.get_web_url()
    system_prompt = """"
You are ChatGPT, a large language model trained by OpenAI.
Knowledge cutoff: 2024-06
Current date: 2026-01-14

Reasoning: high

# Valid channels: analysis, commentary, final. Channel must be included for every message.
Calls to these tools must go to the commentary channel: 'functions'.<|end|><|start|>developer<|message|># Tools

## functions

namespace functions {

// Searches the corpus for relevant documents based on the input query. Returns a section of the document that is relevant to the query.
type search_corpus = (_: {
// The search query to find relevant documents in the corpus.
query: string,
}) => any;

// Performs a regex search on the corpus to find documents matching the query.
type grep_corpus = (_: {
// The regex query to search for in the corpus.
pattern: string,
}) => any;

// Reads the content of a document based on its ID.
type read_document = (_: {
// The unique identifier of the document to read.
doc_id: string,
}) => any;

// Allows the agent to use multiple tools in parallel to gather information.
type multi_tool_use = (_: {
// List of tool calls to execute in parallel.
tool_calls: {
    tool_name: string,
    parameters: {
        },
    }[],
}) => any;

// Prunes the chunks by id that are not relevant to the main question from the history of the conversation.
type prune_chunks = (_: {
chunk_ids: string[],
}) => any;

} // namespace functions"""

    user_content = "\n\n    You are a retrieval subagent in a multi-agent system. Your specific role is to identify and retrieve the most relevant documents from a large corpus to help another agent answer questions. You do NOT answer questions yourself - you only find and retrieve relevant documents.\n\n    Here is the query you need to find documents for:\n\n    <query>\n    A paper was published well into the 20th century, and by December 2023, it had many citations. One of the authors was affiliated with an institution founded in the early twentieth century and was only granted full university status between 1940 and 1960. This author contributed by improving laboratory techniques, addressing a problem that had long hindered progress in their field. The other author not only discovered a major class of compound but also participated in a major competition representing their country between 1920 and 1940. What's the name of the variety mentioned in the abstract used in the experiments?\n    </query>\n\n    **Available Tools:**\n    - SearchTool: Hybrid semantic and keyword search\n    - GrepTool: Text pattern matching\n    - ReadDocument: Read specific document snippets that look promising but incomplete\n    - PruneChunksTool: Remove irrelevant chunks to free up context space\n\n    **Your Process:**\n    - Break down the query into its key concepts and information needs (list each one explicitly)\n    - For each key concept, develop a specific search strategy that targets that concept\n    - Consider what types of documents and evidence would be most helpful for answering this query\n    - Plan several distinct, non-overlapping search strategies that approach the question from different angles\n    - Then execute your searches using multiple parallel tool calls.\n\n    **Your Thinking:**\n    After each round of searches, in your thinking:\n    - Consider the following:\n        - **What do I know?**: List the key topics, themes, or aspects of the question that your currently retrieved documents address. What specific information do you have?\n        - **What should I search for next?**: Systematically consider what search approaches, keywords, or document types you haven't yet tried that might yield valuable information.\n        - **What should I prune?**: If you were to prune chunks, what would you remove and what new searches would you prioritize? Would this likely yield significantly better or more complete information than what you currently have?\n        - **Do I have enough information?**: Given the question's complexity and requirements, do you have sufficient information to help answer it, or are there critical gaps?\n    - Decide if additional searches are needed (and if so, ensure they use genuinely different approaches and do not duplicate or redundant searches)\n    - Avoid getting stuck on a single search strategy - if one approach isn't yielding results, prune and backtrack and try different approaches\n\n    **Tactics to Consider:**\n    - When queries fail, try different approaches or keywords to improve the results\n    - Avoid duplicate or redundant searches\n    - Execute multiple tool calls in parallel when possible\n    - It's OK for this section to be quite long.\n    - If you notice your token budget is approaching the threshold, prune irrelevant chunks proactively to avoid running out of context.\n    - Focus on gathering as much relevant information as possible, it is useful to get multiple perspectives on the same topic or redundant information to confirm the information you have found is correct.\n    - Follow explicit textual evidence rather than speculation\n\n    **Output Format:**\n    Present your final results in order from most relevant to least relevant using this structure:\n\n    <Document id={document_id}>\n    <Justification>\n    Brief explanation (1-3 sentences) of why this document is relevant to the query.\n    </Justification>\n    </Document>\n\n    Example:\n    <Document id=doc_123>\n    <Justification>\n    This document contains detailed analysis of the specific topic mentioned in the query and provides quantitative data that directly supports answering the question.\n    </Justification>\n    </Document>\n\n    Your final output should consist only of the up to 30 ranked document results in the specified format and should not duplicate or rehash any of the search planning or evaluation work you did in the thinking block.\n`\n"

    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": user_content},
    ]

    # messages = [  # OpenAI chat format
    #     system_prompt,
    #     {"role": "user", "content": user_content},
    # ]

    async with aiohttp.ClientSession(base_url=url) as session:
        print(f"Running health check for server at {url}")
        async with session.get("/health", timeout=test_timeout - 1 * MINUTES) as resp:
            up = resp.status == 200
        assert up, f"Failed health check for server at {url}"
        print(f"Successful health check for server at {url}")

        print(f"Sending messages to {url}:", *messages, sep="\n\t")
        await _send_request(session, "llm", messages)

        if twice:
            messages[0]["content"] += "\nTalk like a pirate, matey."
            print(f"Re-sending messages to {url}:", *messages, sep="\n\t")
            await _send_request(session, "llm", messages)


async def _send_request(
    session: aiohttp.ClientSession, model: str, messages: list
) -> None:
    # `stream=True` tells an OpenAI-compatible backend to stream chunks
    payload: dict[str, Any] = {"messages": messages, "model": model, "stream": True}

    headers = {"Content-Type": "application/json", "Accept": "text/event-stream"}

    t = time.perf_counter()
    async with session.post(
        "/v1/chat/completions", json=payload, headers=headers, timeout=10 * MINUTES
    ) as resp:
        async for raw in resp.content:
            resp.raise_for_status()
            # extract new content and stream it
            line = raw.decode().strip()
            if not line or line == "data: [DONE]":
                continue
            if line.startswith("data: "):  # SSE prefix
                line = line[len("data: ") :]

            chunk = json.loads(line)
            assert (
                chunk["object"] == "chat.completion.chunk"
            )  # or something went horribly wrong
            delta = chunk["choices"][0]["delta"]

            if "content" in delta:
                print(delta["content"], end="")  # print the content as it comes in
            elif "reasoning_content" in delta:
                print(delta["reasoning_content"], end="")
            elif not delta:
                print()
            else:
                raise ValueError(f"Unsupported response delta: {delta}")
    print("")
    print(f"Time to Last Token: {time.perf_counter() - t:.2f} seconds")


# ## Standalone function to hit deployed endpoint

DEPLOYED_URL = "https://chroma-core--example-gpt-oss-inference-serve.modal.run"


async def query_deployed(
    user_content: str,
    url: str = DEPLOYED_URL,
    system_prompt: str | None = None,
    model: str = "llm",
    stream: bool = True,
) -> str:
    """
    Query the deployed GPT-OSS endpoint.

    Args:
        user_content: The user message to send
        url: The deployed Modal endpoint URL
        system_prompt: Optional custom system prompt
        model: Model name to use (default: "llm")
        stream: Whether to stream the response

    Returns:
        The full response content as a string
    """
    if system_prompt is None:
        system_prompt = """"
You are ChatGPT, a large language model trained by OpenAI.
Knowledge cutoff: 2024-06
Current date: 2026-01-14

Reasoning: high

# Valid channels: analysis, commentary, final. Channel must be included for every message.
Calls to these tools must go to the commentary channel: 'functions'.<|end|><|start|>developer<|message|># Tools

## functions

namespace functions {

// Searches the corpus for relevant documents based on the input query. Returns a section of the document that is relevant to the query.
type search_corpus = (_: {
// The search query to find relevant documents in the corpus.
query: string,
}) => any;

// Performs a regex search on the corpus to find documents matching the query.
type grep_corpus = (_: {
// The regex query to search for in the corpus.
pattern: string,
}) => any;

// Reads the content of a document based on its ID.
type read_document = (_: {
// The unique identifier of the document to read.
doc_id: string,
}) => any;

// Allows the agent to use multiple tools in parallel to gather information.
type multi_tool_use = (_: {
// List of tool calls to execute in parallel.
tool_calls: {
    tool_name: string,
    parameters: {
        },
    }[],
}) => any;

// Prunes the chunks by id that are not relevant to the main question from the history of the conversation.
type prune_chunks = (_: {
chunk_ids: string[],
}) => any;

} // namespace functions"""

    user_content = "\n\n    You are a retrieval subagent in a multi-agent system. Your specific role is to identify and retrieve the most relevant documents from a large corpus to help another agent answer questions. You do NOT answer questions yourself - you only find and retrieve relevant documents.\n\n    Here is the query you need to find documents for:\n\n    <query>\n    A paper was published well into the 20th century, and by December 2023, it had many citations. One of the authors was affiliated with an institution founded in the early twentieth century and was only granted full university status between 1940 and 1960. This author contributed by improving laboratory techniques, addressing a problem that had long hindered progress in their field. The other author not only discovered a major class of compound but also participated in a major competition representing their country between 1920 and 1940. What's the name of the variety mentioned in the abstract used in the experiments?\n    </query>\n\n    **Available Tools:**\n    - SearchTool: Hybrid semantic and keyword search\n    - GrepTool: Text pattern matching\n    - ReadDocument: Read specific document snippets that look promising but incomplete\n    - PruneChunksTool: Remove irrelevant chunks to free up context space\n\n    **Your Process:**\n    - Break down the query into its key concepts and information needs (list each one explicitly)\n    - For each key concept, develop a specific search strategy that targets that concept\n    - Consider what types of documents and evidence would be most helpful for answering this query\n    - Plan several distinct, non-overlapping search strategies that approach the question from different angles\n    - Then execute your searches using multiple parallel tool calls.\n\n    **Your Thinking:**\n    After each round of searches, in your thinking:\n    - Consider the following:\n        - **What do I know?**: List the key topics, themes, or aspects of the question that your currently retrieved documents address. What specific information do you have?\n        - **What should I search for next?**: Systematically consider what search approaches, keywords, or document types you haven't yet tried that might yield valuable information.\n        - **What should I prune?**: If you were to prune chunks, what would you remove and what new searches would you prioritize? Would this likely yield significantly better or more complete information than what you currently have?\n        - **Do I have enough information?**: Given the question's complexity and requirements, do you have sufficient information to help answer it, or are there critical gaps?\n    - Decide if additional searches are needed (and if so, ensure they use genuinely different approaches and do not duplicate or redundant searches)\n    - Avoid getting stuck on a single search strategy - if one approach isn't yielding results, prune and backtrack and try different approaches\n\n    **Tactics to Consider:**\n    - When queries fail, try different approaches or keywords to improve the results\n    - Avoid duplicate or redundant searches\n    - Execute multiple tool calls in parallel when possible\n    - It's OK for this section to be quite long.\n    - If you notice your token budget is approaching the threshold, prune irrelevant chunks proactively to avoid running out of context.\n    - Focus on gathering as much relevant information as possible, it is useful to get multiple perspectives on the same topic or redundant information to confirm the information you have found is correct.\n    - Follow explicit textual evidence rather than speculation\n\n    **Output Format:**\n    Present your final results in order from most relevant to least relevant using this structure:\n\n    <Document id={document_id}>\n    <Justification>\n    Brief explanation (1-3 sentences) of why this document is relevant to the query.\n    </Justification>\n    </Document>\n\n    Example:\n    <Document id=doc_123>\n    <Justification>\n    This document contains detailed analysis of the specific topic mentioned in the query and provides quantitative data that directly supports answering the question.\n    </Justification>\n    </Document>\n\n    Your final output should consist only of the up to 30 ranked document results in the specified format and should not duplicate or rehash any of the search planning or evaluation work you did in the thinking block.\n`\n"

    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": user_content},
    ]

    payload: dict[str, Any] = {
        "messages": messages,
        "model": model,
        "stream": stream,
        "temperature": 1,
        "stream_options": {"include_usage": True} if stream else {},
    }
    headers = {"Content-Type": "application/json", "Accept": "text/event-stream"}

    full_response = []
    usage_info = None
    t = time.perf_counter()

    async with aiohttp.ClientSession(base_url=url) as session:
        async with session.post(
            "/v1/chat/completions", json=payload, headers=headers, timeout=10 * MINUTES
        ) as resp:
            resp.raise_for_status()
            if stream:
                async for raw in resp.content:
                    line = raw.decode().strip()
                    if not line or line == "data: [DONE]":
                        continue
                    if line.startswith("data: "):
                        line = line[len("data: ") :]

                    chunk = json.loads(line)
                    # Capture usage from final chunk
                    if chunk.get("usage"):
                        usage_info = chunk["usage"]
                    # Check choices exists and has content
                    choices = chunk.get("choices", [])
                    if choices and chunk.get("object") == "chat.completion.chunk":
                        delta = choices[0].get("delta", {})
                        content = delta.get("content")
                        reasoning = delta.get("reasoning_content")
                        if content:
                            print(content, end="", flush=True)
                            full_response.append(content)
                        elif reasoning:
                            print(reasoning, end="", flush=True)
                            full_response.append(reasoning)
            else:
                data = await resp.json()
                message = data["choices"][0]["message"]
                content = (
                    message.get("content") or message.get("reasoning_content") or ""
                )
                if content:
                    full_response.append(content)
                    print(content)
                usage_info = data.get("usage")

    elapsed = time.perf_counter() - t
    print(f"\n\nTime: {elapsed:.2f}s")
    if usage_info:
        completion_tokens = usage_info.get("completion_tokens", 0)
        prompt_tokens = usage_info.get("prompt_tokens", 0)
        print(f"Prompt tokens: {prompt_tokens}, Completion tokens: {completion_tokens}")
        print(f"Server throughput: {completion_tokens / elapsed:.1f} tok/s")
    return "".join(full_response)


def query_deployed_sync(
    user_content: str,
    url: str = DEPLOYED_URL,
    system_prompt: str | None = None,
    model: str = "llm",
    stream: bool = True,
) -> str:
    """Synchronous wrapper for query_deployed."""
    import asyncio

    return asyncio.run(query_deployed(user_content, url, system_prompt, model, stream))


# =============================================================================
# Native Harmony Format Functions
# =============================================================================


async def query_harmony_native(
    conversation: Conversation,
    url: str = DEPLOYED_URL,
    model: str = "llm",
    stream: bool = True,
    max_tokens: int = 4096,
    temperature: float = 1.0,
) -> tuple[str, List[Message]]:
    """
    Query the vLLM endpoint using native Harmony format via /v1/completions.

    This sends pre-rendered Harmony tokens directly to the completions endpoint,
    bypassing the chat completions API wrapper.

    Args:
        conversation: A Harmony Conversation object (from openai_harmony)
        url: The deployed Modal endpoint URL
        model: Model name to use (default: "llm")
        stream: Whether to stream the response
        max_tokens: Maximum tokens to generate
        temperature: Sampling temperature

    Returns:
        Tuple of (raw_text_output, parsed_messages)
    """
    # Render conversation to tokens for completion (includes assistant prefill)
    prompt_tokens = HARMONY_ENC.render_conversation_for_completion(
        conversation,
        next_turn_role=Role.ASSISTANT,
    )

    # Get stop token IDs for proper termination
    stop_token_ids = list(HARMONY_ENC.stop_tokens_for_assistant_actions())

    # Build payload for /v1/completions - prompt accepts array of tokens directly
    payload: dict[str, Any] = {
        "model": model,
        "prompt": list(prompt_tokens),  # Array of token IDs
        "max_tokens": max_tokens,
        "temperature": temperature,
        "stream": stream,
        "stop_token_ids": stop_token_ids,  # vLLM extension for stop tokens
        "return_token_ids": True,  # Get token IDs back directly
        "stream_options": {"include_usage": True} if stream else {},
    }

    headers = {"Content-Type": "application/json", "Accept": "text/event-stream"}
    full_response: list[str] = []
    output_token_ids: list[int] = []
    usage_info = None
    t = time.perf_counter()

    async with aiohttp.ClientSession(base_url=url) as session:
        async with session.post(
            "/v1/completions",  # Raw completions, not chat!
            json=payload,
            headers=headers,
            timeout=10 * MINUTES,
        ) as resp:
            if resp.status >= 400:
                error_body = await resp.text()
                print(f"Error {resp.status}: {error_body}")
                resp.raise_for_status()
            if stream:
                async for raw in resp.content:
                    line = raw.decode().strip()
                    if not line or line == "data: [DONE]":
                        continue
                    if line.startswith("data: "):
                        line = line[len("data: ") :]

                    chunk = json.loads(line)
                    if chunk.get("usage"):
                        usage_info = chunk["usage"]

                    choices = chunk.get("choices", [])
                    if choices:
                        # Get text for display
                        text = choices[0].get("text", "")
                        if text:
                            print(text, end="", flush=True)
                            full_response.append(text)
                        # Collect token IDs (delta tokens in streaming mode)
                        token_ids = choices[0].get("token_ids", [])
                        if token_ids:
                            output_token_ids.extend(token_ids)
            else:
                data = await resp.json()
                choice = data["choices"][0]
                text = choice.get("text", "")
                if text:
                    full_response.append(text)
                    print(text)
                # Get full token IDs in non-streaming mode
                output_token_ids = choice.get("token_ids", [])
                usage_info = data.get("usage")

    elapsed = time.perf_counter() - t
    raw_output = "".join(full_response)

    print(f"\n\nTime: {elapsed:.2f}s")
    if usage_info:
        completion_tokens = usage_info.get("completion_tokens", 0)
        prompt_tokens_count = usage_info.get("prompt_tokens", 0)
        print(
            f"Prompt tokens: {prompt_tokens_count}, Completion tokens: {completion_tokens}"
        )
        print(f"Server throughput: {completion_tokens / elapsed:.1f} tok/s")

    # Parse the output tokens directly into structured Harmony messages
    print(f"Output token IDs ({len(output_token_ids)}): {output_token_ids[:20]}...")
    parsed_messages = HARMONY_ENC.parse_messages_from_completion_tokens(
        output_token_ids, role=Role.ASSISTANT
    )

    return raw_output, parsed_messages


def build_harmony_conversation(
    user_content: str,
    developer_instructions: str | None = None,
    tools: list[dict] | None = None,
    reasoning_effort: str = "high",
) -> Conversation:
    """
    Build a Harmony Conversation object from user content and optional tools.

    Args:
        user_content: The user message
        developer_instructions: Optional developer/system instructions
        tools: Optional list of tool definitions (OpenAI function format)
        reasoning_effort: Reasoning effort level ("low", "medium", "high")

    Returns:
        A Harmony Conversation object ready to be rendered
    """
    # Build system content
    system_content = SystemContent.new().with_reasoning_effort(ReasoningEffort.HIGH)

    messages = [Message.from_role_and_content(Role.SYSTEM, system_content)]

    # Add developer message if provided
    if developer_instructions:
        developer_content = DeveloperContent.new().with_instructions(
            developer_instructions
        )
        if tools:
            developer_content = developer_content.with_function_tools(tools)
        messages.append(
            Message.from_role_and_content(Role.DEVELOPER, developer_content)
        )
    elif tools:
        # Tools without developer instructions
        developer_content = DeveloperContent.new().with_function_tools(tools)
        messages.append(
            Message.from_role_and_content(Role.DEVELOPER, developer_content)
        )

    # Add user message
    messages.append(Message.from_role_and_content(Role.USER, user_content))

    return Conversation.from_messages(messages)


async def query_harmony_simple(
    user_content: str,
    developer_instructions: str | None = None,
    tools: list[dict] | None = None,
    url: str = DEPLOYED_URL,
    model: str = "llm",
    stream: bool = True,
    max_tokens: int = 4096,
    temperature: float = 1.0,
    reasoning_effort: str = "high",
) -> tuple[str, List[Message]]:
    """
    Simple interface to query with native Harmony format.

    Args:
        user_content: The user message
        developer_instructions: Optional developer/system instructions
        tools: Optional list of tool definitions
        url: The deployed endpoint URL
        model: Model name
        stream: Whether to stream
        max_tokens: Max tokens to generate
        temperature: Sampling temperature
        reasoning_effort: Reasoning effort level

    Returns:
        Tuple of (raw_text_output, parsed_messages)
    """
    conversation = build_harmony_conversation(
        user_content=user_content,
        developer_instructions=developer_instructions,
        tools=tools,
        reasoning_effort=reasoning_effort,
    )
    return await query_harmony_native(
        conversation=conversation,
        url=url,
        model=model,
        stream=stream,
        max_tokens=max_tokens,
        temperature=temperature,
    )


def query_harmony_simple_sync(
    user_content: str,
    developer_instructions: str | None = None,
    tools: list[dict] | None = None,
    url: str = DEPLOYED_URL,
    model: str = "llm",
    stream: bool = True,
    max_tokens: int = 4096,
    temperature: float = 1.0,
    reasoning_effort: str = "high",
) -> tuple[str, List[Message]]:
    """Synchronous wrapper for query_harmony_simple."""
    import asyncio

    return asyncio.run(
        query_harmony_simple(
            user_content=user_content,
            developer_instructions=developer_instructions,
            tools=tools,
            url=url,
            model=model,
            stream=stream,
            max_tokens=max_tokens,
            temperature=temperature,
            reasoning_effort=reasoning_effort,
        )
    )


if __name__ == "__main__":
    # Example usage when running directly (not via modal)
    import sys
    import asyncio

    if len(sys.argv) > 1 and sys.argv[1] == "query":
        # Streaming test: uv run inference/gpt_oss_inference.py query "prompt"
        user_msg = sys.argv[2] if len(sys.argv) > 2 else "What is 2+2?"
        for i in range(10):
            print(f"\n{'='*50}")
            print(f"Request {i+1}/10 (streaming)")
            print("=" * 50)
            query_deployed_sync(user_msg, stream=True)

    elif len(sys.argv) > 1 and sys.argv[1] == "nostream":
        # Non-streaming test: uv run inference/gpt_oss_inference.py nostream "prompt"
        user_msg = sys.argv[2] if len(sys.argv) > 2 else "What is 2+2?"
        print(f"Searching for {user_msg}")
        for i in range(10):
            print(f"\n{'='*50}")
            print(f"Request {i+1}/10 (non-streaming)")
            print("=" * 50)
            query_deployed_sync(user_msg, stream=False)

    elif len(sys.argv) > 1 and sys.argv[1] == "harmony":
        # Native Harmony test: uv run inference/gpt_oss_inference.py harmony "prompt"
        user_msg = sys.argv[2] if len(sys.argv) > 2 else "What is 2+2?"
        print(f"\n{'='*50}")
        print("Testing Native Harmony Format")
        print("=" * 50)

        raw_output, parsed_messages = query_harmony_simple_sync(
            user_content=user_msg,
            developer_instructions="You are a helpful assistant. Be concise.",
            stream=True,
            reasoning_effort="low",
        )

        print(f"\n{'='*50}")
        print("Parsed Messages:")
        print("=" * 50)
        for msg in parsed_messages:
            print(f"  {msg.to_dict()}")

    elif len(sys.argv) > 1 and sys.argv[1] == "harmony-tools":
        # Native Harmony with tools: uv run inference/gpt_oss_inference.py harmony-tools "prompt"
        user_msg = (
            sys.argv[2] if len(sys.argv) > 2 else "Search for information about Python"
        )

        # Example tool definitions
        tools = [
            {
                "type": "function",
                "function": {
                    "name": "search_corpus",
                    "description": "Searches the corpus for relevant documents.",
                    "parameters": {
                        "type": "object",
                        "properties": {
                            "query": {
                                "type": "string",
                                "description": "The search query",
                            }
                        },
                        "required": ["query"],
                    },
                },
            }
        ]

        print(f"\n{'='*50}")
        print("Testing Native Harmony Format with Tools")
        print("=" * 50)

        raw_output, parsed_messages = query_harmony_simple_sync(
            user_content=user_msg,
            developer_instructions="You are a search agent. Use the search_corpus tool to find information.",
            tools=tools,
            stream=True,
            reasoning_effort="high",
        )

        print(f"\n{'='*50}")
        print("Parsed Messages:")
        print("=" * 50)
        for msg in parsed_messages:
            print(f"  {msg.to_dict()}")
