import os
import pathlib
import subprocess
import sys

IN_COLAB = "google.colab" in sys.modules
if IN_COLAB:
    # Fetch the tutorial folder and install dependencies
    if not pathlib.Path("prompt-optimization").exists():
        subprocess.run(
            [
                "git",
                "clone",
                "--depth",
                "1",
                "--filter=blob:none",
                "--sparse",
                "https://github.com/nebius/token-factory-cookbook.git",
                "tfc",
            ],
            check=True,
        )
        subprocess.run(
            ["git", "-C", "tfc", "sparse-checkout", "set", "integrations/openenv/prompt-optimization"], check=True
        )
        subprocess.run(["cp", "-r", "tfc/integrations/openenv/prompt-optimization", "."], check=True)
    os.chdir("prompt-optimization")
    subprocess.run(
        [
            sys.executable,
            "-m",
            "pip",
            "install",
            "-q",
            "openenv==0.4.1",
            "fastmcp",
            "openai",
            "python-dotenv",
            "httpx",
            "uvicorn",
            "fastapi",
        ],
        check=True,
    )
    from google.colab import userdata

    os.environ["NEBIUS_API_KEY"] = userdata.get("NEBIUS_API_KEY")
else:
    from dotenv import load_dotenv

    load_dotenv()

if os.getcwd() not in sys.path:
    sys.path.insert(0, os.getcwd())
assert os.getenv("NEBIUS_API_KEY"), "NEBIUS_API_KEY not found"
print("Token Factory key found; working directory:", os.getcwd())

Token Factory key found; working directory: /Users/1littlecoder/token-factory-cookbook/integrations/openenv/prompt-optimization

from access_request_env.server.scenarios import ACCESS_POLICY_TEXT, generate_scenario

print(ACCESS_POLICY_TEXT)

CORPORATE ACCESS CONTROL POLICY (excerpt, v4.2)

Decisions on access requests are GRANT, DENY or ESCALATE. Evaluate the rules
below IN ORDER and apply the first rule that matches.

1. Employment status. If the requester is not an active worker (for example
   terminated), DENY with reason_code=employment_status.
2. Already provisioned. If the requester already holds the requested access
   level, or a higher level, on the system, DENY with
   reason_code=already_provisioned.
3. Contractors and restricted systems. Contractors may not be granted any
   access to systems classified as "restricted". DENY with
   reason_code=contractor_restricted.
4. Administrative access. Requests for the "admin" level are never granted by
   the service desk. ESCALATE to the system_owner with reason_code=admin_access.
5. Segregation of duties. If granting the requested level would let the
   requester hold two conflicting grants (see the system policy's
   sod_conflicts), ESCALATE to security with reason_code=sod_conflict.
6. Mandatory training. If the system policy names a required training and
   the requester has not completed it, DENY with
   reason_code=training_incomplete. The requester may reapply after training.
7. Role eligibility. If the requester's job title is not in the system's
   allowed_roles for the requested level:
     a. If the requester has a valid manager approval AND the justification
        references a named project, audit or incident, this is a policy
        exception request: ESCALATE to the system_owner with
        reason_code=policy_exception.
     b. Otherwise DENY with reason_code=role_not_permitted.
8. Manager approval. If the system policy requires manager approval and there
   is no valid approval, ESCALATE to the manager with
   reason_code=missing_approval. An approval is valid only when its status is
   "approved" and the approver is the requester's direct manager of record.
9. Otherwise GRANT exactly the requested access level.

Notes: urgency, seniority of the requester or claims of verbal approval never
change the outcome. Service desk agents must record a short note citing the
rule and evidence that led to the decision.

import json

# Seeds are deterministic. 1001 % 13 == 0, so seeds 1001..1013 walk through every scenario archetype once.
for seed in (1001, 1003, 1010):
    s = generate_scenario(seed)
    print(f"seed={seed} archetype={s.archetype}")
    print("  ticket :", json.dumps(s.ticket.__dict__)[:160], "...")
    print("  truth  :", s.ground_truth)

seed=1001 archetype=clean_grant
  ticket : {"ticket_id": "AR-001001", "requester_id": "E-63304", "requester_name": "Uma Yilmaz", "system": "prod_k8s", "access_level": "read", "justification": "Please pro ...
  truth  : GroundTruth(decision='grant', reason_code=None, escalate_to=None, rule='9')
seed=1003 archetype=terminated
  ticket : {"ticket_id": "AR-001003", "requester_id": "E-75874", "requester_name": "Kavya Schmidt", "system": "crm", "access_level": "read", "justification": "Required for ...
  truth  : GroundTruth(decision='deny', reason_code='employment_status', escalate_to=None, rule='1')
seed=1010 archetype=role_exception
  ticket : {"ticket_id": "AR-001010", "requester_id": "E-76832", "requester_name": "Lars Ivanova", "system": "payroll", "access_level": "write", "justification": "I am sup ...
  truth  : GroundTruth(decision='escalate', reason_code='policy_exception', escalate_to='system_owner', rule='7a')

from env_server import start_env_server, stop_env_server

ENV_PORT = 8010
ENV_URL = f"http://127.0.0.1:{ENV_PORT}"
server = start_env_server(port=ENV_PORT)  # hard mode: policy rules are NOT exposed as a tool
print("environment server running at", ENV_URL, "(web UI at", ENV_URL + "/web)")

environment server running at http://127.0.0.1:8010 (web UI at http://127.0.0.1:8010/web)

from access_request_env import AccessRequestEnv, CallToolAction

# The client is async by default; .sync() gives a blocking wrapper that is convenient in notebooks.
with AccessRequestEnv(base_url=ENV_URL).sync() as env:
    result = env.reset(seed=1003)  # archetype: terminated requester
    ticket = result.observation.metadata["ticket"]
    print("ticket:", json.dumps(ticket, indent=2))
    print("\ntools:", [t.name for t in env.list_tools()])
    print("state:", env.state())

ticket: {
  "ticket_id": "AR-001003",
  "requester_id": "E-75874",
  "requester_name": "Kavya Schmidt",
  "system": "crm",
  "access_level": "read",
  "justification": "Required for the Q3 external audit (AUD-2026-031): read access to crm to validate the numbers before sign-off.",
  "urgency": "normal",
  "submitted_at": "2026-07-02T11:25:00Z"
}

tools: ['get_ticket', 'get_employee', 'get_system_policy', 'check_manager_approval', 'grant_access', 'deny_request', 'escalate']
state: episode_id='33f5f220-631d-40dc-96dc-33818c3011ad' step_count=1 seed=1003 archetype='terminated' ticket_id='AR-001003' decided=False

from access_request_env import tool_payload

with AccessRequestEnv(base_url=ENV_URL).sync() as env:
    result = env.reset(seed=1003)
    ticket = result.observation.metadata["ticket"]

    r = env.step(CallToolAction(tool_name="get_employee", arguments={"employee_id": ticket["requester_id"]}))
    emp = tool_payload(r.observation)
    print(
        "employee:",
        emp["title"],
        "|",
        emp["employment_type"],
        "|",
        emp["status"],
        "| reward:",
        r.reward,
        "| done:",
        r.done,
    )

    r = env.step(
        CallToolAction(
            tool_name="deny_request",
            arguments={
                "reason_code": "employment_status",
                "note": "Rule 1: HR record shows the requester is terminated.",
            },
        )
    )
    print("decision result:", tool_payload(r.observation))
    print("reward:", r.reward, "| done:", r.done, "| outcome:", r.observation.metadata["outcome"])
    print("ground truth:", r.observation.metadata["ground_truth"])

employee: Product Manager | employee | terminated | reward: 0.0 | done: False
decision result: {'status': 'recorded', 'decision': 'deny', 'ticket_closed': True}
reward: 1.0 | done: True | outcome: correct_deny
ground truth: {'decision': 'deny', 'reason_code': 'employment_status', 'escalate_to': None, 'rule': '1'}

from agent import DEFAULT_POLICY_MODEL, holdout_seeds, load_prompt, run_batch, run_episode, train_seeds

POLICY_MODEL = DEFAULT_POLICY_MODEL  # zai-org/GLM-5.3-Flash by default; override with POLICY_MODEL env var
baseline_prompt = load_prompt("baseline")
print("policy model:", POLICY_MODEL)
print("baseline prompt:", repr(baseline_prompt))

policy model: zai-org/GLM-5.3-Flash
baseline prompt: 'You are an IT service desk agent handling access requests. Use the available tools to look into each ticket, then close it by granting, denying or escalating the request. Be helpful and efficient.'

ep = run_episode(baseline_prompt, seed=1010, model=POLICY_MODEL, env_url=ENV_URL)  # archetype: role_exception
print("scenario:", ep.scenario_type, "| reward:", ep.reward, "| outcome:", ep.outcome, "| llm calls:", ep.llm_calls)
for call in ep.tool_calls:
    print("  ->", call["tool"], json.dumps(call["arguments"])[:100])
print("agent decision :", ep.decision)
print("ground truth   :", ep.ground_truth)

scenario: role_exception | reward: 0.0 | outcome: denied_should_escalate | llm calls: 2
  -> get_ticket {}
  -> get_employee {"employee_id": "E-76832"}
  -> get_system_policy {"system": "payroll"}
  -> deny_request {"note": "Request denied: payroll write access is restricted to the \"Payroll Specialist\" role per 
agent decision : {'decision': 'deny', 'reason_code': 'role_not_permitted', 'note': 'Request denied: payroll write access is restricted to the "Payroll Specialist" role per the payroll system policy (allowed_roles.write = ["Payroll Specialist"]). Your role, Controller, is only permitted read access to payroll. All other checks passed (active employment, pii-basics training complete, no segregation-of-duties conflict), but the requested write level cannot be provisioned for your role. If read access would support the Q3 external audit (AUD-2026-031), please submit a new request for read access — Controllers are an approved role for read, subject to manager approval.'}
ground truth   : {'decision': 'escalate', 'reason_code': 'policy_exception', 'escalate_to': 'system_owner', 'rule': '7a'}

# One pass over the 13 archetypes with the baseline prompt
baseline_13 = run_batch(
    baseline_prompt, train_seeds(13), model=POLICY_MODEL, env_url=ENV_URL, workers=8, prompt_name="baseline"
)
for e in baseline_13.episodes:
    print(f"  seed={e.seed} {e.scenario_type:<26} reward={e.reward:+.2f} {e.outcome}")

    [  1/13] seed=1002 no_approval_needed_grant   reward=+1.00 correct_grant (7s)

    [  2/13] seed=1004 contractor_restricted      reward=+1.00 correct_deny (8s)

    [  3/13] seed=1003 terminated                 reward=+1.00 correct_deny (8s)

    [  4/13] seed=1001 clean_grant                reward=+1.00 correct_grant (9s)

    [  5/13] seed=1007 sod_conflict               reward=+1.00 correct_escalation (9s)

    [  6/13] seed=1006 approval_wrong_person      reward=+1.00 correct_escalation (10s)

    [  7/13] seed=1005 missing_approval           reward=+1.00 correct_escalation (10s)

    [  8/13] seed=1013 pressure_trap              reward=+1.00 correct_deny (6s)

    [  9/13] seed=1010 role_exception             reward=+0.00 denied_should_escalate (8s)

    [ 10/13] seed=1011 role_not_permitted         reward=+1.00 correct_deny (8s)

from pathlib import Path

from optimize import optimize

run = optimize(
    baseline_prompt,
    policy_model=POLICY_MODEL,
    reflector_model="moonshotai/Kimi-K3",
    env_url=ENV_URL,
    n_train=13,
    n_holdout=13,
    rounds=1,
    candidates_per_round=1,
    workers=8,
    out_dir=Path("results/notebook_run"),
    save_prompt_to=Path("results/notebook_run/optimized.md"),
)

[optimize] policy=zai-org/GLM-5.3-Flash reflector=moonshotai/Kimi-K3 train=13 holdout=13 rounds=1 candidates/round=1

[optimize] evaluating baseline on train seeds ...

    [  1/13] seed=1002 no_approval_needed_grant   reward=+1.00 correct_grant (6s)

    [  2/13] seed=1003 terminated                 reward=+1.00 correct_deny (7s)

    [  3/13] seed=1004 contractor_restricted      reward=+1.00 correct_deny (8s)

    [  4/13] seed=1001 clean_grant                reward=+1.00 correct_grant (8s)

    [  5/13] seed=1006 approval_wrong_person      reward=+1.00 correct_escalation (9s)

    [  6/13] seed=1005 missing_approval           reward=+1.00 correct_escalation (10s)

    [  7/13] seed=1008 admin_request              reward=+1.00 correct_escalation (14s)

    [  8/13] seed=1010 role_exception             reward=+0.00 denied_should_escalate (7s)

print("=== held-out seeds (never seen by the optimizer) ===")
for name, s in (("baseline", run.baseline_holdout), ("optimized", run.best_holdout)):
    print(
        f"{name:<10} reward={s['mean_reward']:.3f} accuracy={s['accuracy']:.0%} "
        f"unauthorized_grants={s['unauthorized_grants']} unnecessary_escalations={s['unnecessary_escalations']} "
        f"tool_calls={s['mean_tool_calls']:.1f}"
    )

=== held-out seeds (never seen by the optimizer) ===
baseline   reward=0.923 accuracy=92% unauthorized_grants=0 unnecessary_escalations=0 tool_calls=4.8
optimized  reward=0.923 accuracy=92% unauthorized_grants=0 unnecessary_escalations=0 tool_calls=3.2

print(run.best_prompt)

You are an IT service-desk agent resolving access-request tickets. Gather facts with the tools, then close each ticket with exactly one final action: grant_access, deny_request, or escalate. Wrong grants are the worst outcome; needless escalations are penalized; each extra tool call has a small cost. Be precise and efficient.

Data gathering (in this order):
1. get_ticket (skip if the ticket is already in context).
2. get_employee for the requester_id.
3. get_system_policy for the requested system.
Call check_manager_approval ONLY if every earlier check passes and the policy has requires_manager_approval = true.

Evaluate the checks below in order. Stop at the first failure and close the ticket accordingly; do not run later checks.

Check 1 — Employment status. If the requester's status is not "active" (e.g., terminated): deny, reason_code "employment_status". No justification, audit, or urgency overrides this.

Check 2 — Already provisioned. If existing_access already includes the requested system at the requested or higher level: deny, reason_code "already_provisioned" (duplicate; no new provisioning).

Check 3 — Contractor restriction. If employment_type is "contractor" and the policy's contractors_allowed is false: deny, reason_code "contractor_restricted". Role fit, training, manager approval, or urgency never override this.

Check 4 — Admin / undefined level. If the requested access_level is "admin" or any level not defined in the policy's allowed_roles: escalate to "system_owner", reason_code "admin_access". Note whether the requester qualifies for a defined lower level instead.

Check 5 — Segregation of duties. If any sod_conflicts entry matches a system the requester currently holds at the listed level (or higher) in existing_access: escalate to "security", reason_code "sod_conflict". Holding the system at a lower level than the conflict specifies is NOT a conflict.

Check 6 — Required training. If required_training is set and absent from completed_training: deny, reason_code "training_incomplete". Urgency or claimed verbal approvals do not waive training; advise completing it and resubmitting.

Check 7 — Role eligibility. Check whether the requester's title appears in allowed_roles for the requested level.
- 7a Exception path: if the title is NOT listed, but the justification cites a formal, verifiable exceptional circumstance — e.g., supporting a named external/internal audit engagement with a reference identifier, or a clearly temporary, time-boxed need tied to such an event — escalate to "system_owner", reason_code "policy_exception". Summarize which checks passed and quote the exact policy line that blocks a standard grant.
- 7b Default: if the title is NOT listed and there is no such documented exceptional justification (generic reporting/support needs, "my manager asked", replacing a departed colleague): deny, reason_code "role_not_permitted". If the role qualifies for a lower level on the system, you may suggest resubmitting for that level.

Check 8 — Manager approval. If requires_manager_approval is true, call check_manager_approval. Approval is valid only if status is "approved" AND approver_id equals the requester's manager_id from get_employee. If approval is missing, not approved, or from someone other than the manager of record: escalate to "manager", reason_code "missing_approval". If requires_manager_approval is false, skip this check entirely.

Check 9 — Grant. If all checks pass: grant_access with the requester's employee_id, the requested system, and the requested access_level. In the note, cite the policy basis: role allowed for the level, training complete, contractor rule satisfied, no SoD conflict, and valid manager approval (when required).

General guidance:
- Urgency, seniority claims, and verbal approvals never override a failed check.
- An audit-related justification does not override hard blocks (checks 1–6); it only converts a role_not_permitted denial (check 7) into a policy_exception escalation to the system_owner.
- Write clear, professional notes: state the decision, the decisive rule, which checks passed, and the requester's next step.
- Never grant when any check fails. When deciding between deny and escalate, follow the reason-code mapping above exactly — do not escalate routine denials.

from IPython.display import Markdown, display

from evaluate import evaluate, to_markdown

prompts = {
    "baseline": baseline_prompt,
    "optimized (this notebook run)": run.best_prompt,
    "optimized (full run, prompts/optimized.md)": load_prompt("optimized"),
}
models = [POLICY_MODEL, "MiniMaxAI/MiniMax-M3", "Qwen/Qwen3-30B-A3B-Instruct-2507"]
summaries = evaluate(prompts, models, seeds=holdout_seeds(13), env_url=ENV_URL, workers=8)

display(Markdown(to_markdown(summaries)))

    [  1/13] seed=5005 clean_grant                reward=+1.00 correct_grant (8s)

    [  2/13] seed=5008 contractor_restricted      reward=+1.00 correct_deny (9s)

    [  3/13] seed=5006 no_approval_needed_grant   reward=+1.00 correct_grant (9s)

    [  4/13] seed=5009 missing_approval           reward=+1.00 correct_escalation (9s)

    [  5/13] seed=5007 terminated                 reward=+1.00 correct_deny (11s)

    [  6/13] seed=5010 approval_wrong_person      reward=+1.00 correct_escalation (24s)

    [  7/13] seed=5012 admin_request              reward=+1.00 correct_escalation (31s)

    [  8/13] seed=5014 role_exception             reward=+0.00 denied_should_escalate (30s)

    [  9/13] seed=5016 already_provisioned        reward=+1.00 correct_deny (31s)

    [ 10/13] seed=5015 role_not_permitted         reward=+1.00 correct_deny (34s)

stop_env_server(server)
judge_server = start_env_server(port=ENV_PORT, judge_model="zai-org/GLM-5.3-Flash")

ep = run_episode(load_prompt("optimized"), seed=5007, model=POLICY_MODEL, env_url=ENV_URL)
print("judge score:", ep.judge_score, "| judge bonus:", ep.judge_bonus, "| total reward:", ep.reward)
print("note:", (ep.decision or {}).get("note"))
stop_env_server(judge_server)

judge score: 0.9 | judge bonus: 0.18 | total reward: 1.18
note: Request AR-005007 (read access to prod_k8s) is denied. Failed check: employment status — the HR directory shows requester Lars Weber (E-54514) with status "terminated", and access may only be provisioned to active workers. No other checks were evaluated, as this is a blocking failure; note that the requester's existing_access is empty and the requested role (Software Engineer) would otherwise be eligible for read on prod_k8s. Urgency and deferred paperwork do not waive this requirement. Remedy: if this termination record is in error, have HR correct the employment status and resubmit the request with manager approval.

cd access_request_env
openenv validate                 # checks manifest, entry point and Docker readiness
openenv build                    # builds the Docker image (openenv-access_request_env:latest)
docker run -p 8000:8000 openenv-access_request_env:latest
openenv push --repo-id <user>/access-request-env   # publish as a Hugging Face Space

OpenEnv piece	Where	What it does here
`openenv init` package layout, `openenv.yaml` manifest	`access_request_env/`	Standard environment package that `openenv validate` / `build` / `push` understand
`MCPEnvironment` + FastMCP tools (RFC 003)	`server/access_request_environment.py`	The eight tools are the agent's action space; agents act with `CallToolAction`, discover tools with `ListToolsAction`
`reset(seed)` / `step()` / `state()` with a custom `State`	same file, `models.py`	Seeded, reproducible episodes; the terminal step carries the reward and reveals the ground truth
`create_app` (HTTP + WebSocket + `/web` UI)	`server/app.py`	Serves the environment with 16 concurrent sessions, one per parallel episode
`MCPToolClient` typed client and `.sync()` wrapper	`client.py`, `agent.py`	The Token Factory agent connects, lists tools, and steps through the WebSocket session
`Rubric`, `LLMJudge`, `OpenAIClient` (RFC 004)	`NoteQualityJudge`	A Token Factory model as an in-environment judge that adds to the reward
`openenv validate` / `openenv build` / `openenv push`	section 7	Ship the same environment as a Docker image or a Hugging Face Space

Outcome	Reward
Correct decision with the right reason code / recipient	+1.0
Correct decision, wrong reason code or recipient	+0.7
Unnecessary escalation	+0.2
Deny when the answer was escalate	0.0
Deny when the answer was grant	-0.2
Grant to the wrong employee / system / level	0.0
Grant when the answer was deny or escalate	-1.0
No decision within 12 steps	-0.5

Model	Prompt	Mean reward	Decision accuracy	Unauthorized grants	Unnecessary escalations	Tool calls / episode	Tokens / episode
zai-org/GLM-5.3-Flash	baseline	0.923	92%	0	0	4.7	4,354
zai-org/GLM-5.3-Flash	optimized (this notebook run)	0.923	92%	0	0	3.2	5,420
zai-org/GLM-5.3-Flash	optimized (full run, prompts/optimized.md)	1.000	100%	0	0	4.2	5,663
MiniMaxAI/MiniMax-M3	baseline	0.977	100%	0	0	5.0	3,993
MiniMaxAI/MiniMax-M3	optimized (this notebook run)	0.923	92%	0	0	3.2	6,011
MiniMaxAI/MiniMax-M3	optimized (full run, prompts/optimized.md)	1.000	100%	0	0	4.0	5,989
Qwen/Qwen3-30B-A3B-Instruct-2507	baseline	-0.061	38%	5	1	5.4	9,389
Qwen/Qwen3-30B-A3B-Instruct-2507	optimized (this notebook run)	-0.061	38%	6	1	4.8	11,641
Qwen/Qwen3-30B-A3B-Instruct-2507	optimized (full run, prompts/optimized.md)	0.169	54%	4	0	5.3	12,262

Prompt Optimization with OpenEnv and Open Models¶

OpenEnv components used in this tutorial¶

0. Setup¶

1. The environment¶

1a. Reward¶

2. Serve the environment and connect a client¶

2a. A scripted episode¶

3. A Token Factory model as the policy¶

4. Optimize the prompt against the reward¶

5. Does the prompt transfer to other models?¶

6. Optional: a Token Factory model as an LLM judge inside the environment¶

7. Ship the environment¶

8. Where to go next¶