Environment Servers

An environment server runs during each attempt. It serves the tools that the model can call, keeps the state of the attempt, and holds the grader. It is a NeMo Gym resources server in app.py, which monte env create writes from Gym’s template.

How an Attempt Runs

The agent sends the task row to the server’s seed_session endpoint, which opens a session for the attempt.
The agent sends the row’s input and tools to the model.
Each time the model calls a tool, the agent sends the arguments to the server endpoint with the same name and returns the result to the model.
When the model stops, the agent sends the whole answer and the row’s grader fields to verify.

The agent is Gym’s simple_agent, which allows a set number of tool calls in each attempt. Monte records the answer, the tool calls, and the reward as a trace.

Add a Tool

A tool is a FastAPI endpoint that setup_webserver adds, plus a declaration with the same name in the task rows. The support_escalation server has two tools:

from fastapi import FastAPI
from pydantic import BaseModel

from nemo_gym.base_resources_server import BaseResourcesServerConfig, SimpleResourcesServer


class SupportConfig(BaseResourcesServerConfig):
    top_k: int = 3  # Set in the server config.


class SearchKBRequest(BaseModel):
    query: str


class SearchKBResponse(BaseModel):
    results: list[str]


class EscalateRequest(BaseModel):
    reason: str


class EscalateResponse(BaseModel):
    status: str
    message: str


class SupportServer(SimpleResourcesServer):
    config: SupportConfig

    def setup_webserver(self) -> FastAPI:
        app = super().setup_webserver()
        app.post("/search_kb")(self.search_kb)
        app.post("/escalate_to_human")(self.escalate_to_human)
        return app

    async def search_kb(self, body: SearchKBRequest) -> SearchKBResponse:
        # rank() orders KNOWLEDGE_BASE by similarity to the query.
        return SearchKBResponse(results=rank(body.query)[: self.config.top_k])

    async def escalate_to_human(self, body: EscalateRequest) -> EscalateResponse:
        return EscalateResponse(status="escalated", message="A specialist will reply within 2 hours.")

A tool can call an outside API. The support_escalation server embeds each search_kb query with an OpenAI model, so it reads OPENAI_API_KEY from a team secret.

Keep State for Each Attempt

Many attempts run at the same time on one server, so the server keys its state by the session ID. The same session ID reaches seed_session, each tool, and verify:

from fastapi import Request
from pydantic import Field

from nemo_gym.server_utils import SESSION_ID_KEY


class SupportServer(SimpleResourcesServer):
    config: SupportConfig
    tickets: dict[str, list[str]] = Field(default_factory=dict)  # Session ID -> escalation reasons

    async def escalate_to_human(self, request: Request, body: EscalateRequest) -> EscalateResponse:
        self.tickets.setdefault(request.session[SESSION_ID_KEY], []).append(body.reason)
        return EscalateResponse(status="escalated", message="A specialist will reply within 2 hours.")

verify(self, request, body) reads self.tickets[request.session[SESSION_ID_KEY]]. To start an attempt from a state that the row sets, override seed_session, which receives the whole row. A grader that reads this state cannot regrade, because the state ends with the attempt.

Server Config

configs/support_escalation.yaml names the server, sets its config fields, connects it to the agent and the model, and declares the training rows:

support_escalation_resources_server:
  resources_servers:
    support_escalation:
      entrypoint: app.py
      domain: other              # Gym's domain field: math, coding, agent, ..., other
      top_k: 3                   # a field of SupportConfig

support_escalation_simple_agent:
  responses_api_agents:
    simple_agent:
      entrypoint: app.py
      resources_server:
        type: resources_servers
        name: support_escalation_resources_server
      model_server:
        type: responses_api_models
        name: policy_model       # the model under test or in training
      datasets:
      - name: train
        type: train
        jsonl_fpath: resources_servers/support_escalation/data/train.jsonl
        num_repeats: 1
        license: Apache 2.0

The fingerprint covers this file, so a changed field, such as a new top_k, needs a new monte env add.

Get Started

Training

Inference

Datasets

Environments

Evaluation

Observability

Security

How an Attempt Runs

Add a Tool

Keep State for Each Attempt

Server Config

​How an Attempt Runs

​Add a Tool

​Keep State for Each Attempt

​Server Config

How an Attempt Runs

Add a Tool

Keep State for Each Attempt

Server Config