import argparse
import asyncio
import tempfile
import httpx
import tinker
from tinker_agent.proxy import serve_proxy
from tinker_agent.trajectory import flush_capture, read_samples, rows_to_trajectory
from tinker_cookbook import renderers
from tinker_cookbook.capture.store import capture_to_store
from tinker_cookbook.rl.data_processing import trajectory_to_data
RUN_ID = "demo"
async def run_agent(proxy_url):
# Stand-in for your agent: any OpenAI-compatible client with base_url=proxy_url + "/v1" works.
messages = [{"role": "user", "content": "Reply with the number 1."}]
async with httpx.AsyncClient(base_url=proxy_url, timeout=300) as http:
for _ in range(2):
response = await http.post("/v1/chat/completions", json={"messages": messages, "max_tokens": 32})
response.raise_for_status()
messages += [response.json()["choices"][0]["message"], {"role": "user", "content": "Now reply with 2."}]
return 1.0 # the episode's reward, from your verifier
async def main(base_url, base_model, renderer_name, store_url):
service = tinker.ServiceClient(base_url=base_url, api_key="tml-demo")
training = await service.create_lora_training_client_async(base_model=base_model, rank=8)
saved = await (await training.save_weights_for_sampler_async(name="step-0"))
sampler = await service.create_sampling_client_async(model_path=saved.path)
renderer = renderers.get_renderer(renderer_name, sampler.get_tokenizer())
async with serve_proxy(renderer, sampler, model_label=base_model, default_max_tokens=32) as port:
reward = await run_agent(f"http://127.0.0.1:{port}/r/run/{RUN_ID}/group/0/traj/0")
await flush_capture()
rows = await read_samples(store_url, RUN_ID, group_idx=0, traj_idx=0)
trajectory = rows_to_trajectory(rows, "completed")
# With several trajectories per task, pass reward minus the group mean instead.
datums = trajectory_to_data(trajectory, traj_advantage=reward)
# "mask" is a client-side field; the cookbook trainer also drops it before forward_backward.
datums = [
tinker.Datum(
model_input=datum.model_input,
loss_fn_inputs={key: value for key, value in datum.loss_fn_inputs.items() if key != "mask"},
)
for datum in datums
]
await (await training.forward_backward_async(datums, loss_fn="ppo"))
await (await training.optim_step_async(tinker.types.AdamParams(learning_rate=1e-5)))
print(f"trained on {len(rows)} captured calls")
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--base-url", required=True)
parser.add_argument("--base-model", required=True)
parser.add_argument("--renderer", default="qwen3")
args = parser.parse_args()
with capture_to_store(RUN_ID, data_dir=tempfile.mkdtemp(prefix="tinker-capture-")) as captured:
asyncio.run(main(args.base_url, args.base_model, args.renderer, captured.base_url))
