diff --git a/examples/multiply_openai_agents/main.py b/examples/multiply_openai_agents/main.py new file mode 100644 index 00000000..582728e6 --- /dev/null +++ b/examples/multiply_openai_agents/main.py @@ -0,0 +1,27 @@ +"""Local OpenAI Agents multiply rollout server.""" + +from __future__ import annotations + +import os + +import uvicorn + +from multiply_rollout.grader import MultiplyGrader, multiply_grader_config +from multiply_rollout.workflow import MultiplyWorkflow, multiply_workflow_config +from osmosis_ai.rollout.backend.local import LocalBackend +from osmosis_ai.rollout.server import create_rollout_server + + +def main() -> None: + backend = LocalBackend( + workflow=MultiplyWorkflow, + workflow_config=multiply_workflow_config, + grader=MultiplyGrader, + grader_config=multiply_grader_config, + ) + app = create_rollout_server(backend=backend) + uvicorn.run(app, host="0.0.0.0", port=int(os.environ.get("_OSMOSIS_ROLLOUT_PORT", "8000"))) + + +if __name__ == "__main__": + main() diff --git a/examples/multiply_openai_agents/multiply_rollout/__init__.py b/examples/multiply_openai_agents/multiply_rollout/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/examples/multiply_openai_agents/multiply_rollout/grader.py b/examples/multiply_openai_agents/multiply_rollout/grader.py new file mode 100644 index 00000000..03000e37 --- /dev/null +++ b/examples/multiply_openai_agents/multiply_rollout/grader.py @@ -0,0 +1,55 @@ +import logging +from typing import Any + +from multiply_rollout.utils import extract_solution +from osmosis_ai.rollout.context import GraderContext +from osmosis_ai.rollout.grader import Grader +from osmosis_ai.rollout.types import GraderConfig + +logger = logging.getLogger(__name__) + + +class MultiplyGraderConfig(GraderConfig): + name: str = "MultiplyOpenAIAgentsGrader" + description: str = "Grades multiplication rollouts using OpenAI Agents" + + +multiply_grader_config = MultiplyGraderConfig() + + +class MultiplyGrader(Grader): + def compute_reward(self, solution_str: str, ground_truth: str) -> float: + extracted = extract_solution(solution_str) + try: + sol_val = float(extracted) + gt_val = float(ground_truth) + except (TypeError, ValueError): + return 0.0 + + if abs(gt_val - sol_val) < 1e-2: + return 1.0 + return 0.0 + + async def grade(self, ctx: GraderContext) -> None: + if ctx.sample is None: + ctx.set_reward(0.0) + return + + last_message = ctx.sample.messages[-1] + content = last_message.get("content", "") + reward = self.compute_reward(_extract_text(content), ctx.label or "") + ctx.set_reward(reward) + logger.info("[MultiplyGrader] reward = %.1f", reward) + + +def _extract_text(content: Any) -> str: + if isinstance(content, str): + return content + if isinstance(content, list): + parts = [ + item.get("text", "") + for item in content + if isinstance(item, dict) and item.get("type") == "output_text" + ] + return "\n".join(parts) + return "" diff --git a/examples/multiply_openai_agents/multiply_rollout/tools.py b/examples/multiply_openai_agents/multiply_rollout/tools.py new file mode 100644 index 00000000..b6ce0c65 --- /dev/null +++ b/examples/multiply_openai_agents/multiply_rollout/tools.py @@ -0,0 +1,7 @@ +from agents import function_tool + + +@function_tool +async def multiply_tool(a: float, b: float) -> float: + """Multiply two numbers.""" + return a * b diff --git a/examples/multiply_openai_agents/multiply_rollout/utils.py b/examples/multiply_openai_agents/multiply_rollout/utils.py new file mode 100644 index 00000000..c008ed12 --- /dev/null +++ b/examples/multiply_openai_agents/multiply_rollout/utils.py @@ -0,0 +1,9 @@ +import re + + +def extract_solution(solution_str: str) -> str | None: + """Extract a final answer from the expected #### format.""" + solution = re.search(r"####\s*([-+]?\d*\.?\d+)", solution_str) + if not solution: + return None + return solution.group(1) diff --git a/examples/multiply_openai_agents/multiply_rollout/workflow.py b/examples/multiply_openai_agents/multiply_rollout/workflow.py new file mode 100644 index 00000000..d44ada0f --- /dev/null +++ b/examples/multiply_openai_agents/multiply_rollout/workflow.py @@ -0,0 +1,50 @@ +from typing import Any + +from agents import ModelSettings, Runner +from agents.models.interface import Model + +from multiply_rollout.tools import multiply_tool +from osmosis_ai.rollout.agent_workflow import AgentWorkflow +from osmosis_ai.rollout.context import AgentWorkflowContext +from osmosis_ai.rollout.integrations.agents.openai_agents import ( + OsmosisAgent, + OsmosisMemorySession, + OsmosisRolloutModel, +) +from osmosis_ai.rollout.types import AgentWorkflowConfig + +MAX_TURNS = 8 + + +class MultiplyAgentWorkflowConfig(AgentWorkflowConfig): + name: str = "MultiplyAgentWorkflow" + description: str = "Multiply two numbers using OpenAI Agents" + model: Model + model_settings: ModelSettings + tools: Any + + +multiply_workflow_config = MultiplyAgentWorkflowConfig( + model=OsmosisRolloutModel(), + model_settings=ModelSettings(temperature=1.0, top_p=1.0, max_tokens=4096), + tools=[multiply_tool], +) + + +class MultiplyWorkflow(AgentWorkflow): + async def run(self, ctx: AgentWorkflowContext) -> None: + config = ctx.config + agent = OsmosisAgent( + name="multiply", + model=config.model, + model_settings=config.model_settings, + tools=config.tools, + ) + + session = OsmosisMemorySession() + await Runner.run( + agent, + ctx.prompt, + session=session, + max_turns=MAX_TURNS, + ) diff --git a/examples/multiply_strands/main.py b/examples/multiply_strands/main.py new file mode 100644 index 00000000..5c0ff66d --- /dev/null +++ b/examples/multiply_strands/main.py @@ -0,0 +1,27 @@ +"""Local Strands multiply rollout server.""" + +from __future__ import annotations + +import os + +import uvicorn + +from multiply_rollout.grader import MultiplyGrader, multiply_grader_config +from multiply_rollout.workflow import MultiplyWorkflow, multiply_workflow_config +from osmosis_ai.rollout.backend.local import LocalBackend +from osmosis_ai.rollout.server import create_rollout_server + + +def main() -> None: + backend = LocalBackend( + workflow=MultiplyWorkflow, + workflow_config=multiply_workflow_config, + grader=MultiplyGrader, + grader_config=multiply_grader_config, + ) + app = create_rollout_server(backend=backend) + uvicorn.run(app, host="0.0.0.0", port=int(os.environ.get("_OSMOSIS_ROLLOUT_PORT", "8000"))) + + +if __name__ == "__main__": + main() diff --git a/examples/multiply_strands/multiply_rollout/__init__.py b/examples/multiply_strands/multiply_rollout/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/examples/multiply_strands/multiply_rollout/grader.py b/examples/multiply_strands/multiply_rollout/grader.py new file mode 100644 index 00000000..39eec48e --- /dev/null +++ b/examples/multiply_strands/multiply_rollout/grader.py @@ -0,0 +1,55 @@ +import logging +from typing import Any + +from multiply_rollout.utils import extract_solution +from osmosis_ai.rollout.context import GraderContext +from osmosis_ai.rollout.grader import Grader +from osmosis_ai.rollout.types import GraderConfig + +logger = logging.getLogger(__name__) + + +class MultiplyGraderConfig(GraderConfig): + name: str = "MultiplyStrandsGrader" + description: str = "Grades multiplication rollouts using Strands" + + +multiply_grader_config = MultiplyGraderConfig() + + +class MultiplyGrader(Grader): + def compute_reward(self, solution_str: str, ground_truth: str) -> float: + extracted = extract_solution(solution_str) + try: + sol_val = float(extracted) + gt_val = float(ground_truth) + except (TypeError, ValueError): + return 0.0 + + if abs(gt_val - sol_val) < 1e-2: + return 1.0 + return 0.0 + + async def grade(self, ctx: GraderContext) -> None: + if ctx.sample is None: + ctx.set_reward(0.0) + return + + last_message = ctx.sample.messages[-1] + content = last_message.get("content", "") + reward = self.compute_reward(_extract_text(content), ctx.label or "") + ctx.set_reward(reward) + logger.info("[MultiplyGrader] reward = %.1f", reward) + + +def _extract_text(content: Any) -> str: + if isinstance(content, str): + return content + if isinstance(content, list): + parts = [ + item.get("text", "") + for item in content + if isinstance(item, dict) and "text" in item + ] + return "\n".join(parts) + return "" diff --git a/examples/multiply_strands/multiply_rollout/tools.py b/examples/multiply_strands/multiply_rollout/tools.py new file mode 100644 index 00000000..a97079a5 --- /dev/null +++ b/examples/multiply_strands/multiply_rollout/tools.py @@ -0,0 +1,7 @@ +from strands import tool + + +@tool +def multiply_tool(a: float, b: float) -> float: + """Multiply two numbers.""" + return a * b diff --git a/examples/multiply_strands/multiply_rollout/utils.py b/examples/multiply_strands/multiply_rollout/utils.py new file mode 100644 index 00000000..c008ed12 --- /dev/null +++ b/examples/multiply_strands/multiply_rollout/utils.py @@ -0,0 +1,9 @@ +import re + + +def extract_solution(solution_str: str) -> str | None: + """Extract a final answer from the expected #### format.""" + solution = re.search(r"####\s*([-+]?\d*\.?\d+)", solution_str) + if not solution: + return None + return solution.group(1) diff --git a/examples/multiply_strands/multiply_rollout/workflow.py b/examples/multiply_strands/multiply_rollout/workflow.py new file mode 100644 index 00000000..dabc4822 --- /dev/null +++ b/examples/multiply_strands/multiply_rollout/workflow.py @@ -0,0 +1,54 @@ +from typing import Any + +from strands.agent.agent_result import AgentResult +from strands.models.model import Model + +from multiply_rollout.tools import multiply_tool +from multiply_rollout.utils import extract_solution +from osmosis_ai.rollout.agent_workflow import AgentWorkflow +from osmosis_ai.rollout.context import AgentWorkflowContext +from osmosis_ai.rollout.integrations.agents.strands import ( + OsmosisRolloutModel, + OsmosisStrandsAgent, +) +from osmosis_ai.rollout.types import AgentWorkflowConfig + +MAX_ITERATIONS = 8 + + +class MultiplyAgentWorkflowConfig(AgentWorkflowConfig): + name: str = "MultiplyAgentWorkflow" + description: str = "Multiply two numbers using Strands" + model: Model + tools: Any + + +multiply_workflow_config = MultiplyAgentWorkflowConfig( + model=OsmosisRolloutModel(params={"temperature": 1.0, "top_p": 1.0, "max_tokens": 4096}), + tools=[multiply_tool], +) + + +class MultiplyWorkflow(AgentWorkflow): + async def check_done(self, result: AgentResult) -> bool: + content = result.message.get("content", "") + if not any("toolUse" in block for block in content): + return True + text_content = next((block for block in content if block.get("text")), None) + if text_content: + return extract_solution(text_content.get("text", "")) is not None + return False + + async def run(self, ctx: AgentWorkflowContext) -> None: + config = ctx.config + agent = OsmosisStrandsAgent( + name="multiply", + model=config.model, + tools=config.tools, + messages=ctx.prompt, + callback_handler=None, + ) + for _ in range(MAX_ITERATIONS): + result = await agent.invoke_async() + if await self.check_done(result): + break diff --git a/osmosis_ai/rollout/__init__.py b/osmosis_ai/rollout/__init__.py index 4805a45a..257da4fa 100644 --- a/osmosis_ai/rollout/__init__.py +++ b/osmosis_ai/rollout/__init__.py @@ -23,7 +23,6 @@ GraderCompleteRequest, GraderConfig, GraderStatus, - MultiTurnMode, RolloutCompleteRequest, RolloutErrorCategory, RolloutInitRequest, @@ -48,7 +47,6 @@ "GraderStatus", "HarborAgentWorkflowContext", "LocalBackend", - "MultiTurnMode", "OsmosisRolloutModel", "OsmosisStrandsAgent", "RolloutCompleteRequest", diff --git a/osmosis_ai/rollout/backend/harbor/agent_runner.py b/osmosis_ai/rollout/backend/harbor/agent_runner.py index b0b5476e..9e08b0d2 100644 --- a/osmosis_ai/rollout/backend/harbor/agent_runner.py +++ b/osmosis_ai/rollout/backend/harbor/agent_runner.py @@ -67,16 +67,19 @@ async def run_workflow( with rollout_ctx: await workflow.run(ctx) - samples = await rollout_ctx.get_samples() - samples_data = {sid: s.model_dump() for sid, s in samples.items()} + sample = await rollout_ctx.get_sample() + sample_data = sample.model_dump() if sample is not None else None - (AGENT_LOGS_DIR / "samples.json").write_text( - json.dumps(samples_data, default=str) + (AGENT_LOGS_DIR / "sample.json").write_text( + json.dumps(sample_data, default=str) ) meta["status"] = "success" - meta["samples"] = samples_data - print(f"Agent runner complete: {len(samples)} samples collected") + meta["sample"] = sample_data + print( + "Agent runner complete: " + + ("1 sample collected" if sample is not None else "no sample collected") + ) except Exception as e: traceback.print_exc() diff --git a/osmosis_ai/rollout/backend/harbor/backend.py b/osmosis_ai/rollout/backend/harbor/backend.py index d9a3ab1d..3ecb9b18 100644 --- a/osmosis_ai/rollout/backend/harbor/backend.py +++ b/osmosis_ai/rollout/backend/harbor/backend.py @@ -39,7 +39,7 @@ RolloutStatus, ) from osmosis_ai.rollout.utils.imports import to_import_path -from osmosis_ai.rollout.utils.rewards import validate_samples_have_rewards +from osmosis_ai.rollout.utils.rewards import validate_sample_has_reward logger: logging.Logger = logging.getLogger(__name__) @@ -53,7 +53,7 @@ set -e PYTHONPATH=/workspace python -m osmosis_ai.rollout.backend.harbor.grader_runner \ --config /logs/agent/rollout_config.json \ - --samples /logs/agent/samples.json + --sample /logs/agent/sample.json """ @@ -424,8 +424,8 @@ async def on_verification_start(self, event: TrialHookEvent) -> None: metadata = get_agent_metadata(event) if metadata and metadata.get("status") == "success": - samples = parse_samples(metadata.get("samples", {})) - result = ExecutionResult(status=RolloutStatus.SUCCESS, samples=samples) + sample = parse_sample(metadata.get("sample")) + result = ExecutionResult(status=RolloutStatus.SUCCESS, sample=sample) elif event.result and event.result.exception_info: err = event.result.exception_info result = ExecutionResult( @@ -489,17 +489,22 @@ async def on_trial_end(self, event: TrialHookEvent) -> None: if pending.on_grader_complete: metadata = get_agent_metadata(event) - samples = parse_samples(metadata.get("samples", {})) if metadata else {} + sample = parse_sample(metadata.get("sample") if metadata else None) # Read before the trial-dir cleanup further down. artifacts = self.read_grader_artifacts(rollout_id) if event.result and event.result.verifier_result: + # Harbor surfaces verifier output as ``rewards: dict[str, float]`` + # regardless of what shape the verifier wrote — for the single-sample + # contract we take any one value and apply it to our single sample. rewards = event.result.verifier_result.rewards or {} - for sid, reward in rewards.items(): - if sid in samples: - samples[sid].reward = float(reward) + reward_values = [ + float(v) for v in rewards.values() if v is not None + ] + if sample is not None and reward_values: + sample.reward = reward_values[0] try: - validate_samples_have_rewards(samples) + validate_sample_has_reward(sample) except ValueError as e: if not rewards: logger.warning( @@ -515,24 +520,24 @@ async def on_trial_end(self, event: TrialHookEvent) -> None: ) result = ExecutionResult( status=RolloutStatus.FAILURE, - samples=samples, + sample=sample, err_message=str(e), err_category=RolloutErrorCategory.VALIDATION_ERROR, ) else: result = ExecutionResult( - status=RolloutStatus.SUCCESS, samples=samples + status=RolloutStatus.SUCCESS, sample=sample ) elif event.result and event.result.exception_info: err = event.result.exception_info result = ExecutionResult( status=RolloutStatus.FAILURE, - samples=samples, + sample=sample, err_message=err.exception_message, err_category=RolloutErrorCategory.AGENT_ERROR, ) else: - result = ExecutionResult(status=RolloutStatus.FAILURE, samples=samples) + result = ExecutionResult(status=RolloutStatus.FAILURE, sample=sample) result.artifacts = artifacts await pending.on_grader_complete(result) @@ -566,11 +571,14 @@ def get_agent_metadata(event: TrialHookEvent) -> dict[str, Any] | None: return None -def parse_samples(raw: dict[str, Any]) -> dict[str, RolloutSample]: - return { - sid: RolloutSample.model_validate(data) if isinstance(data, dict) else data - for sid, data in raw.items() - } +def parse_sample(raw: Any) -> RolloutSample | None: + if raw is None: + return None + if isinstance(raw, RolloutSample): + return raw + if isinstance(raw, dict): + return RolloutSample.model_validate(raw) + return None def rewrite_url_for_docker(url: str) -> str: diff --git a/osmosis_ai/rollout/backend/harbor/grader_runner.py b/osmosis_ai/rollout/backend/harbor/grader_runner.py index f093e293..8b3665ee 100644 --- a/osmosis_ai/rollout/backend/harbor/grader_runner.py +++ b/osmosis_ai/rollout/backend/harbor/grader_runner.py @@ -1,7 +1,7 @@ -"""Grade agent samples inside a Harbor container. +"""Grade the rollout's single sample inside a Harbor container. Usage: - osmosis-grader-runner --config /workspace/rollout_config.json --samples /logs/agent/samples.json + osmosis-grader-runner --config /workspace/rollout_config.json --sample /logs/agent/sample.json """ from __future__ import annotations @@ -21,12 +21,12 @@ def parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser(description="Grade agent samples") + parser = argparse.ArgumentParser(description="Grade the rollout sample") parser.add_argument( "--config", type=Path, required=True, help="Path to rollout_config.json" ) parser.add_argument( - "--samples", type=Path, required=True, help="Path to samples.json" + "--sample", type=Path, required=True, help="Path to sample.json" ) return parser.parse_args() @@ -35,14 +35,16 @@ def load_config(path: Path) -> dict[str, Any]: return json.loads(path.read_text()) -def load_samples(path: Path) -> dict[str, RolloutSample]: - raw: dict[str, Any] = json.loads(path.read_text()) - return {sid: RolloutSample.model_validate(data) for sid, data in raw.items()} +def load_sample(path: Path) -> RolloutSample | None: + raw: Any = json.loads(path.read_text()) + if raw is None: + return None + return RolloutSample.model_validate(raw) -def write_rewards(rewards: dict[str, float | None]) -> None: +def write_reward(reward: float | None) -> None: VERIFIER_LOGS_DIR.mkdir(parents=True, exist_ok=True) - (VERIFIER_LOGS_DIR / "reward.json").write_text(json.dumps(rewards)) + (VERIFIER_LOGS_DIR / "reward.json").write_text(json.dumps({"reward": reward})) def write_artifacts(artifacts: dict[str, Any]) -> None: @@ -72,36 +74,38 @@ def main() -> None: if not label and metadata is None: print("No label or metadata in config, skipping grading") - write_rewards({}) + write_reward(None) return - if not args.samples.exists(): - print("No samples found, skipping grading") - write_rewards({}) + if not args.sample.exists(): + print("No sample found, skipping grading") + write_reward(None) return if "grader" not in config: print("No grader in config, skipping grading") - write_rewards({}) + write_reward(None) + return + + sample = load_sample(args.sample) + if sample is None: + print("Sample file is empty, skipping grading") + write_reward(None) return - samples = load_samples(args.samples) grader_cls = resolve_object(config["grader"]) grader_config = ( resolve_object(config["grader_config"]) if "grader_config" in config else None ) - ctx = GraderContext(label=label, samples=samples, metadata=metadata) + ctx = GraderContext(label=label, sample=sample, metadata=metadata) grader = grader_cls(grader_config) asyncio.run(grader.grade(ctx)) - graded = ctx.get_samples() - for sid, sample in graded.items(): - if sample.reward is None: - raise RuntimeError(f"Sample {sid} has no reward after grading") + if ctx.sample is None or ctx.sample.reward is None: + raise RuntimeError("Sample has no reward after grading") - rewards = {sid: sample.reward for sid, sample in graded.items()} - write_rewards(rewards) + write_reward(ctx.sample.reward) if ctx.artifacts is not None: # Artifacts are best-effort and must never block reward delivery, even # if the verifier dir is unwritable after reward.json is persisted. @@ -109,7 +113,7 @@ def main() -> None: write_artifacts(ctx.artifacts) except OSError as e: print(f"Warning: failed to write grader artifacts: {e}") - print(f"Grading complete: {rewards}") + print(f"Grading complete: reward={ctx.sample.reward}") if __name__ == "__main__": diff --git a/osmosis_ai/rollout/backend/local/backend.py b/osmosis_ai/rollout/backend/local/backend.py index 405dfd1b..f53df2c2 100644 --- a/osmosis_ai/rollout/backend/local/backend.py +++ b/osmosis_ai/rollout/backend/local/backend.py @@ -22,7 +22,7 @@ ) from osmosis_ai.rollout.utils.concurrency import ConcurrencyLimiter from osmosis_ai.rollout.utils.imports import resolve_object -from osmosis_ai.rollout.utils.rewards import validate_samples_have_rewards +from osmosis_ai.rollout.utils.rewards import validate_sample_has_reward logger: logging.Logger = logging.getLogger(__name__) @@ -115,7 +115,7 @@ async def run_workflow(self, request: ExecutionRequest) -> ExecutionResult: return ExecutionResult( status=RolloutStatus.SUCCESS, - samples=await rollout_ctx.get_samples(), + sample=await rollout_ctx.get_sample(), ) async def run_grader( @@ -126,24 +126,23 @@ async def run_grader( grader_ctx = GraderContext( label=request.label, - samples=result.samples, + sample=result.sample, metadata=request.metadata, ) try: grader = self.grader_cls(self.grader_config) await grader.grade(grader_ctx) - samples = grader_ctx.get_samples() - validate_samples_have_rewards(samples) + validate_sample_has_reward(grader_ctx.sample) return ExecutionResult( status=RolloutStatus.SUCCESS, - samples=samples, + sample=grader_ctx.sample, artifacts=grader_ctx.artifacts, ) except Exception as e: logger.error(traceback.format_exc()) return ExecutionResult( status=RolloutStatus.FAILURE, - samples=result.samples, + sample=result.sample, artifacts=grader_ctx.artifacts, err_message=str(e), err_category=_categorize_exception(e), diff --git a/osmosis_ai/rollout/context.py b/osmosis_ai/rollout/context.py index 7ddc102f..57786d6e 100644 --- a/osmosis_ai/rollout/context.py +++ b/osmosis_ai/rollout/context.py @@ -11,22 +11,18 @@ class SampleSource(ABC): - """Produces a ``RolloutSample`` for grading from a framework-specific object. + """Produces the rollout's ``RolloutSample`` for grading. - Each integration ships a small source class that wraps its underlying - object (an Agent, a Session, etc.) and produces a sample on demand. - The source is decoupled from the integration's framework classes so - they do not need to inherit from this ABC. - - RolloutContext indexes registered sources by name and calls - ``get_sample`` lazily at sample collection time, passing the - registration name so the source can stamp it onto the returned - ``RolloutSample.id``. + A rollout produces exactly one sample (one agent run, one reward). The + active ``RolloutContext`` holds at most one ``SampleSource``; the + integration registers a source that wraps its underlying object + (typically an Agent or a Session) and the backend pulls the sample at + grading time. """ @abstractmethod - async def get_sample(self, name: str) -> RolloutSample: - """Return the current rollout sample for the given registration name.""" + async def get_sample(self) -> RolloutSample: + """Return the rollout's current sample.""" rollout_contextvar: ContextVar["RolloutContext | None"] = ContextVar( @@ -50,7 +46,7 @@ class RolloutContext: chat_completions_url: str = "" api_key: str | None = None rollout_id: str = "" - sample_sources: dict[str, SampleSource] = field(default_factory=dict) + sample_source: SampleSource | None = None def __post_init__(self) -> None: if not self.chat_completions_url: @@ -67,24 +63,28 @@ def __enter__(self) -> "RolloutContext": def __exit__(self, *_: Any) -> None: rollout_contextvar.set(None) - def register_sample_source(self, name: str, source: SampleSource) -> None: - """Register a sample source under the given name. The context calls - ``source.get_sample(name)`` lazily when ``get_samples()`` is invoked, - so the source can keep reflecting state that mutates until the - rollout ends. + def set_sample_source(self, source: SampleSource) -> None: + """Register the single sample source for this rollout. + + A rollout has exactly one sample; integrations call this from their + agent/session constructors so the backend can pull the sample at + grading time. Attempting to register a second source raises, which + catches the common bug of putting two ``OsmosisMemorySession``s or + two ``OsmosisStrandsAgent``s in the same workflow. """ - if name in self.sample_sources: + if self.sample_source is not None: raise ValueError( - f"Session with {name} already exists, please give your session " - "or agent a unique name in the workflow" + "RolloutContext already has a sample source registered. " + "A rollout produces one sample; construct a single agent/" + "session per rollout." ) - self.sample_sources[name] = source + self.sample_source = source - async def get_samples(self) -> dict[str, RolloutSample]: - out: dict[str, RolloutSample] = {} - for name, source in self.sample_sources.items(): - out[name] = await source.get_sample(name) - return out + async def get_sample(self) -> RolloutSample | None: + """Return the rollout's sample, or ``None`` if no source was registered.""" + if self.sample_source is None: + return None + return await self.sample_source.get_sample() def get_rollout_context() -> RolloutContext | None: @@ -93,19 +93,22 @@ def get_rollout_context() -> RolloutContext | None: @dataclass class GraderContext: + """Context passed to ``Grader.grade``. + + Carries the single sample produced by the rollout. Graders set its + reward via :meth:`set_reward`. + """ + label: str | None = None - samples: dict[str, RolloutSample] = field(default_factory=dict) + sample: RolloutSample | None = None project_path: str | None = None metadata: dict[str, Any] | None = None artifacts: dict[str, Any] | None = None - def get_samples(self) -> dict[str, RolloutSample]: - return self.samples - - def set_sample_reward(self, sample_id: str, reward: float) -> None: - if sample_id not in self.samples: - raise ValueError(f"Sample {sample_id} not found") - self.samples[sample_id].reward = reward + def set_reward(self, reward: float) -> None: + if self.sample is None: + raise ValueError("GraderContext has no sample to reward") + self.sample.reward = reward def set_artifacts(self, artifacts: dict[str, Any]) -> None: """Set the rollout-level artifacts object (replaces any prior value).""" diff --git a/osmosis_ai/rollout/driver.py b/osmosis_ai/rollout/driver.py index 84f24e82..8cfac4ea 100644 --- a/osmosis_ai/rollout/driver.py +++ b/osmosis_ai/rollout/driver.py @@ -17,22 +17,19 @@ class RolloutOutcome: """Result of a single rollout execution. - This is the eval-facing contract. Eval doesn't need to know - whether this came from in-process execution or an HTTP call. + Single-sample by design: one rollout = one agent run = one reward. + ``sample`` carries the conversation and reward (when grading succeeded), + ``rollout_id`` identifies the rollout in logs/cache rows, and the rest + are run-level metrics + diagnostics. """ status: RolloutStatus - samples: dict[str, RolloutSample] = field(default_factory=dict) + sample: RolloutSample | None = None error: str | None = None duration_ms: float = 0.0 tokens: int = 0 systemic_error: str | None = None rollout_id: str | None = None - controller_created_sample_ids: list[str] = field(default_factory=list) - completion_counts: dict[str, int] = field(default_factory=dict) - full_callback_sample_ids: list[str] = field(default_factory=list) - scored_sample_ids: list[str] = field(default_factory=list) - skipped_sample_ids: list[str] = field(default_factory=list) callback_diagnostics: dict[str, Any] = field(default_factory=dict) skipped: bool = False diff --git a/osmosis_ai/rollout/integrations/agents/openai_agents.py b/osmosis_ai/rollout/integrations/agents/openai_agents.py index 78f8b55b..5969f854 100644 --- a/osmosis_ai/rollout/integrations/agents/openai_agents.py +++ b/osmosis_ai/rollout/integrations/agents/openai_agents.py @@ -1,6 +1,5 @@ import uuid from collections.abc import AsyncIterator -from contextvars import ContextVar from typing import Any from agents import Agent @@ -17,36 +16,26 @@ from osmosis_ai.rollout.context import SampleSource, get_rollout_context from osmosis_ai.rollout.types import RolloutSample -current_sample_id: ContextVar[str | None] = ContextVar( - "osmosis_current_sample_id", default=None -) - class OsmosisMemorySession(SessionABC): - """In-memory session that doubles as the sample source for an Osmosis rollout. - - Behavior: - - Inside a ``RolloutContext``: registers itself for sample collection - (via ``SessionSampleSource``) and publishes its sample id - to a ContextVar each time the runner reads from or writes to it. - ``OsmosisRolloutModel`` reads that ContextVar to stamp per-rollout - headers. - - Outside a ``RolloutContext``: behaves as a plain in-memory - ``SessionABC`` implementation, so the same workflow code can run - locally without an Osmosis rollout. + """In-memory session that doubles as the rollout's sample source. + + Inside a ``RolloutContext`` this registers itself with the context so + the backend can pull the conversation at grading time. Outside a + ``RolloutContext`` it behaves as a plain in-memory ``SessionABC``, so + the same workflow code can run locally without an Osmosis rollout. Pass exactly one ``OsmosisMemorySession`` per ``Runner.run`` when using - ``OsmosisRolloutModel``. + ``OsmosisRolloutModel`` — a rollout produces one sample. """ - def __init__(self, name: str | None = None) -> None: - self.name: str = name or uuid.uuid4().hex - self.session_id: str = self.name + def __init__(self) -> None: + self.session_id: str = uuid.uuid4().hex self.items: list[TResponseInputItem] = [] self._registered_context_id: int | None = None ctx = get_rollout_context() if ctx is not None: - ctx.register_sample_source(self.name, SessionSampleSource(self)) + ctx.set_sample_source(SessionSampleSource(self)) self._registered_context_id = id(ctx) def _raise_if_unregistered_in_rollout_context(self) -> None: @@ -63,14 +52,12 @@ def _raise_if_unregistered_in_rollout_context(self) -> None: async def get_items(self, limit: int | None = None) -> list[TResponseInputItem]: self._raise_if_unregistered_in_rollout_context() - current_sample_id.set(self.name) if limit is None: return list(self.items) return list(self.items[-limit:]) async def add_items(self, items: list[TResponseInputItem]) -> None: self._raise_if_unregistered_in_rollout_context() - current_sample_id.set(self.name) self.items.extend(items) async def pop_item(self) -> TResponseInputItem | None: @@ -81,7 +68,7 @@ async def clear_session(self) -> None: class SessionSampleSource(SampleSource): - """Produces a ``RolloutSample`` from any OpenAI Agents SDK ``Session``. + """Produces the rollout's sample from any OpenAI Agents SDK ``Session``. Works against any ``SessionABC`` by returning the items the runner persisted via ``add_items`` (canonical Responses-API ``TResponseInputItem`` @@ -97,22 +84,23 @@ class SessionSampleSource(SampleSource): def __init__(self, session: SessionABC) -> None: self.session = session - async def get_sample(self, name: str) -> RolloutSample: + async def get_sample(self) -> RolloutSample: items = await self.session.get_items() - return RolloutSample(id=name, messages=items) + return RolloutSample(messages=items) class OsmosisRolloutModel(LitellmModel): - """Placeholder ``Model`` that carries litellm kwargs for use in workflow configs. + """Placeholder ``Model`` that carries litellm kwargs for workflow configs. - This is *not* a usable model on its own: ``OsmosisAgent`` replaces it - with a real :class:`OsmosisLitellmModel` (wired to the active - ``RolloutContext``) at agent construction time. Any direct call into - ``stream_response`` / ``get_response`` raises ``NotImplementedError`` - because the placeholder has no connection params. + Not a usable model on its own: ``OsmosisAgent`` replaces it with a real + :class:`OsmosisLitellmModel` (wired to the active ``RolloutContext``) + at agent construction time. Any direct call into ``stream_response`` / + ``get_response`` raises ``NotImplementedError`` because the placeholder + has no connection params. - Subclassing ``LitellmModel`` (without invoking its ``__init__``) is purely - a typing convenience so the placeholder satisfies ``Agent.model: Model``. + Subclassing ``LitellmModel`` (without invoking its ``__init__``) is + purely a typing convenience so the placeholder satisfies + ``Agent.model: Model``. """ def __init__(self, **litellm_kwargs: Any) -> None: @@ -135,19 +123,21 @@ async def get_response(self, *args: Any, **kwargs: Any) -> ModelResponse: class OsmosisLitellmModel(LitellmModel): - """Streaming-only LitellmModel pre-wired for the Osmosis completions server. - - Stateless across calls: the per-call session name is read from a - ContextVar that ``OsmosisMemorySession`` publishes when the runner interacts - with it. One instance can be safely shared across agents and runs (within - the rollout context that constructed it). - - Caveats this class handles for the user: - - The Osmosis completions server is streaming-only, so ``get_response`` - is implemented by aggregating ``stream_response`` events. This makes - ``Runner.run`` and ``Runner.run_sync`` work without forcing the user - to switch to ``Runner.run_streamed``. - - Per-rollout bookkeeping headers are injected on every model call. + """Streaming-only LitellmModel pre-wired to the active rollout's chat endpoint. + + The model server identifies the rollout entirely via the URL it was + handed (rollout id is baked into the ``chat_completions_url`` path), + so this class no longer stamps per-call headers. It does, however, + refuse to make a request when no ``OsmosisMemorySession`` has been + registered with the active ``RolloutContext`` — that catches the + common bug of calling ``Runner.run`` without + ``session=OsmosisMemorySession()`` (which would otherwise leave the + grader with no conversation to score). + + The Osmosis completions server is streaming-only, so ``get_response`` + is implemented by aggregating ``stream_response`` events. This makes + ``Runner.run`` and ``Runner.run_sync`` work without forcing the user + to switch to ``Runner.run_streamed``. """ def __init__(self, **litellm_kwargs: Any) -> None: @@ -158,7 +148,6 @@ def __init__(self, **litellm_kwargs: Any) -> None: "Construct it inside your workflow run, where the execution " "backend has set up the context." ) - self.rollout_id: str = ctx.rollout_id super().__init__( model="openai/osmosis-rollout", base_url=ctx.chat_completions_url, @@ -167,19 +156,18 @@ def __init__(self, **litellm_kwargs: Any) -> None: ) def _merge_headers(self, model_settings: ModelSettings) -> dict[str, str]: - sample_id = current_sample_id.get() - if sample_id is None: + # No per-call routing headers anymore — the URL carries rollout + # identity. We still gate the call on a registered sample source + # so users who forget ``session=OsmosisMemorySession()`` get a + # loud error instead of a silently-empty sample at grading time. + ctx = get_rollout_context() + if ctx is None or ctx.sample_source is None: raise RuntimeError( "OsmosisLitellmModel was called without an active OsmosisMemorySession.\n" "Pass `session=OsmosisMemorySession()` to Runner.run() so the " - "model can stamp per-rollout headers and the runner can " - "persist the conversation for grading." + "runner persists the conversation for grading." ) - return { - **super()._merge_headers(model_settings), - "x-sample-id": sample_id, - "x-rollout-id": self.rollout_id, - } + return super()._merge_headers(model_settings) async def get_response(self, *args: Any, **kwargs: Any) -> ModelResponse: output_items: list[Any] = [] diff --git a/osmosis_ai/rollout/integrations/agents/strands.py b/osmosis_ai/rollout/integrations/agents/strands.py index 56bcd8b9..2ae5c084 100644 --- a/osmosis_ai/rollout/integrations/agents/strands.py +++ b/osmosis_ai/rollout/integrations/agents/strands.py @@ -1,94 +1,57 @@ -import uuid -from collections.abc import AsyncGenerator -from typing import Any, TypeVar, cast +from typing import Any, cast from strands import Agent as StrandsAgent from strands.models.litellm import LiteLLMModel from strands.models.model import Model -from strands.types.content import Messages, SystemContentBlock -from strands.types.streaming import StreamEvent -from strands.types.tools import ToolChoice, ToolSpec +from strands.types.content import Messages from osmosis_ai.rollout.context import ( - RolloutContext, SampleSource, get_rollout_context, ) from osmosis_ai.rollout.types import RolloutSample from osmosis_ai.rollout.utils.messages import map_initial_messages_to_content_blocks -T = TypeVar("T") - class StrandsAgentSampleSource(SampleSource): - """Produces a ``RolloutSample`` from a Strands agent's ``messages`` field. + """Produces the rollout sample from a Strands agent's ``messages`` field. - Strands agents accumulate the conversation history on ``agent.messages`` - in chat-completion format, so the source just wraps that list. + Strands accumulates the conversation on ``agent.messages`` in + chat-completion format, so this just wraps that list. """ def __init__(self, agent: StrandsAgent) -> None: self.agent = agent - async def get_sample(self, name: str) -> RolloutSample: - return RolloutSample(id=name, messages=list(self.agent.messages)) + async def get_sample(self) -> RolloutSample: + return RolloutSample(messages=list(self.agent.messages)) class OsmosisRolloutModel(LiteLLMModel): - """Placeholder model that carries sampling params. - - At runtime, OsmosisStrandsAgent calls for_sample() which reads connection - info from the active RolloutContext to create a concrete LiteLLMModel. - """ - - def for_sample(self, sample_id: str, rollout_ctx: RolloutContext) -> LiteLLMModel: - headers = { - "x-sample-id": sample_id, - "x-rollout-id": rollout_ctx.rollout_id, - } - client_args = { - "api_base": rollout_ctx.chat_completions_url, - "api_key": rollout_ctx.api_key, - "extra_headers": headers, - } - return LiteLLMModel( - client_args=client_args, - model_id="openai/osmosis-rollout", - **self.config, # type: ignore - ) - - def update_config(self, **model_config: Any) -> None: - raise NotImplementedError("This should not be called for OsmosisRolloutModel") + """Placeholder ``Model`` that carries litellm kwargs for workflow configs. - def get_config(self) -> Any: - raise NotImplementedError("This should not be called for OsmosisRolloutModel") + Not a usable model on its own: ``OsmosisStrandsAgent`` replaces it with + a real ``LiteLLMModel`` (wired to the active ``RolloutContext``) at + agent construction time. It has no connection params, so any direct + call into the ``Model`` API will fail. - def structured_output( - self, - output_model: type[T], - prompt: Messages, - system_prompt: str | None = None, - **kwargs: Any, - ) -> AsyncGenerator[dict[str, T | Any], None]: - raise NotImplementedError("This should not be called for OsmosisRolloutModel") + Subclassing ``LiteLLMModel`` (without invoking its ``__init__``) is + purely a typing convenience so the placeholder satisfies + ``StrandsAgent.model: Model`` -- same pattern as + ``integrations.agents.openai_agents.OsmosisRolloutModel``. + """ - def stream( - self, - messages: Messages, - tool_specs: list[ToolSpec] | None = None, - system_prompt: str | None = None, - *, - tool_choice: ToolChoice | None = None, - system_prompt_content: list[SystemContentBlock] | None = None, - invocation_state: dict[str, Any] | None = None, - **kwargs: Any, - ) -> AsyncGenerator[StreamEvent, None]: - raise NotImplementedError("This should not be called for OsmosisRolloutModel") + def __init__(self, **litellm_kwargs: Any) -> None: + self.litellm_kwargs: dict[str, Any] = litellm_kwargs class OsmosisStrandsAgent(StrandsAgent): - """Drop-in replacement for StrandsAgent that handles rollout model swap and - sample registration transparently. The constructor signature matches StrandsAgent. + """Drop-in ``StrandsAgent`` that wires itself into the active rollout. + + If ``model`` is an ``OsmosisRolloutModel`` placeholder, this materializes + a real ``LiteLLMModel`` against the active ``RolloutContext`` and + registers the agent as the rollout's sample source. One + ``OsmosisStrandsAgent`` per rollout (matching the single-sample model). """ def __init__( @@ -106,11 +69,19 @@ def __init__( if rollout_ctx is None: raise RuntimeError( "OsmosisRolloutModel requires an active RolloutContext. " - "Ensure the execution backend sets up the context before running the workflow." + "Ensure the execution backend sets up the context before " + "running the workflow." ) - name = kwargs.get("name") or kwargs.get("agent_id") or uuid.uuid4().hex - model = model.for_sample(name, rollout_ctx) - rollout_ctx.register_sample_source(name, StrandsAgentSampleSource(self)) + litellm_model: Model = LiteLLMModel( + client_args={ + "api_base": rollout_ctx.chat_completions_url, + "api_key": rollout_ctx.api_key, + }, + model_id="openai/osmosis-rollout", + **model.litellm_kwargs, + ) + rollout_ctx.set_sample_source(StrandsAgentSampleSource(self)) + model = litellm_model super().__init__( *args, diff --git a/osmosis_ai/rollout/server/app.py b/osmosis_ai/rollout/server/app.py index 5a651ba6..f1fbd264 100644 --- a/osmosis_ai/rollout/server/app.py +++ b/osmosis_ai/rollout/server/app.py @@ -1,5 +1,6 @@ import logging import traceback +import uuid from typing import Any from fastapi import BackgroundTasks, FastAPI, HTTPException @@ -49,20 +50,24 @@ async def rollout( async def _handle_rollout( backend: ExecutionBackend, request: RolloutInitRequest ) -> None: + # Routing identity is in the URLs; ``rollout_id`` in the body is debug + # metadata. We prefer the caller's id (so logs/cache rows correlate + # across systems) and synthesize one only if the caller omits it. + rollout_id = request.rollout_id or uuid.uuid4().hex auth = ControllerAuth(api_key=request.controller_api_key) rollout_ctx = RolloutContext( chat_completions_url=request.chat_completions_url, api_key=request.controller_api_key, - rollout_id=request.rollout_id, + rollout_id=rollout_id, ) async def on_workflow_complete(result: ExecutionResult) -> None: await post_json_with_retry( url=request.completion_callback_url, payload=RolloutCompleteRequest( - rollout_id=request.rollout_id, status=result.status, + rollout_id=rollout_id, err_message=result.err_message, err_category=result.err_category, ).model_dump(), @@ -73,23 +78,23 @@ async def on_grader_complete(result: ExecutionResult) -> None: if not request.grader_callback_url: logger.info( "Skipping grader callback for %s: no grader_callback_url", - request.rollout_id, + rollout_id, ) return logger.info( - "Posting grader callback for %s to %s (status=%s, samples=%d)", - request.rollout_id, + "Posting grader callback for %s to %s (status=%s, has_sample=%s)", + rollout_id, request.grader_callback_url, result.status, - len(result.samples), + result.sample is not None, ) artifacts = sanitize_artifacts(result.artifacts) payload = GraderCompleteRequest( - rollout_id=request.rollout_id, + rollout_id=rollout_id, status=GraderStatus.SUCCESS if result.status == RolloutStatus.SUCCESS else GraderStatus.FAILURE, - samples=result.samples, + sample=result.sample, artifacts=artifacts, err_message=result.err_message, err_category=result.err_category, @@ -104,7 +109,7 @@ async def on_grader_complete(result: ExecutionResult) -> None: ) logger.info( "Grader callback for %s completed: status=%d", - request.rollout_id, + rollout_id, resp.status_code, ) @@ -112,7 +117,7 @@ async def on_grader_complete(result: ExecutionResult) -> None: with rollout_ctx: await backend.execute( ExecutionRequest( - id=request.rollout_id, + id=rollout_id, prompt=request.initial_messages, label=request.label, metadata=request.metadata, @@ -124,17 +129,15 @@ async def on_grader_complete(result: ExecutionResult) -> None: if request.grader_callback_url else None, ) - logger.info("Rollout %s completed successfully", request.rollout_id) + logger.info("Rollout %s completed successfully", rollout_id) except Exception: - logger.error( - "Rollout %s failed: %s", request.rollout_id, traceback.format_exc() - ) + logger.error("Rollout %s failed: %s", rollout_id, traceback.format_exc()) try: await post_json_with_retry( url=request.completion_callback_url, payload=RolloutCompleteRequest( - rollout_id=request.rollout_id, status=RolloutStatus.FAILURE, + rollout_id=rollout_id, err_message="Internal server error", ).model_dump(), headers=auth.as_bearer_headers(), diff --git a/osmosis_ai/rollout/types/__init__.py b/osmosis_ai/rollout/types/__init__.py index 8e24b9d9..2b117342 100644 --- a/osmosis_ai/rollout/types/__init__.py +++ b/osmosis_ai/rollout/types/__init__.py @@ -17,7 +17,6 @@ ExecutionRequest, ExecutionResult, MessageDict, - MultiTurnMode, RolloutErrorCategory, RolloutSample, RolloutStatus, @@ -35,7 +34,6 @@ "GraderInitResponse", "GraderStatus", "MessageDict", - "MultiTurnMode", "RolloutCompleteRequest", "RolloutErrorCategory", "RolloutInitRequest", diff --git a/osmosis_ai/rollout/types/protocol.py b/osmosis_ai/rollout/types/protocol.py index 3f44a433..c8f0cf7e 100644 --- a/osmosis_ai/rollout/types/protocol.py +++ b/osmosis_ai/rollout/types/protocol.py @@ -12,20 +12,27 @@ class RolloutInitRequest(BaseModel): - rollout_id: str + """Body of the POST to the SDK's /rollout endpoint. + + Routing identity lives in the URLs the caller hands us: the rollout + id is baked into ``chat_completions_url`` (session-scoped) and into + both callback URLs as a path segment. ``rollout_id`` is repeated in + the body for debug logging on the rollout-server side and for + correlation in user-side dashboards; the SDK does not rely on it for + routing, so it is optional. + """ - # These come from slime samples initial_messages: list[MessageDict] label: str | None = None metadata: dict[str, Any] | None = None - # These come from the rollout controller + rollout_id: str | None = None + chat_completions_url: str controller_api_key: str | None = None completion_callback_url: str grader_callback_url: str | None = None - # Per agent and grader timeout agent_timeout_sec: float | None = None grader_timeout_sec: float | None = None @@ -36,8 +43,15 @@ class RolloutInitResponse(BaseModel): ... class RolloutCompleteRequest(BaseModel): - rollout_id: str + """Body of the rollout-complete callback. + + ``rollout_id`` mirrors the id embedded in the callback URL purely for + debug/log correlation on the controller side. The receiver identifies + the rollout from the URL path, not this field. + """ + status: RolloutStatus + rollout_id: str | None = None extra_fields: dict[str, Any] | None = None @@ -46,8 +60,15 @@ class RolloutCompleteRequest(BaseModel): class GraderInitRequest(BaseModel): - rollout_id: str - samples: dict[str, RolloutSample] + """Body of a POST to a remote grader endpoint. + + A rollout produces a single sample (the agent's conversation), so this + carries that one sample directly rather than the legacy + ``dict[str, RolloutSample]``. + """ + + sample: RolloutSample + rollout_id: str | None = None completion_callback_url: str extra_fields: dict[str, Any] | None = None @@ -64,9 +85,17 @@ class GraderStatus(StrEnum): class GraderCompleteRequest(BaseModel): - rollout_id: str + """Body of the grader-complete callback. + + Carries the single graded sample (with its ``reward`` populated on + success) and nothing on failure. ``rollout_id`` mirrors the URL path + segment for debug/log correlation; the controller resolves the + rollout from the URL. + """ + status: GraderStatus - samples: dict[str, RolloutSample] + rollout_id: str | None = None + sample: RolloutSample | None = None artifacts: dict[str, Any] | None = None err_message: str | None = None err_category: RolloutErrorCategory | None = None diff --git a/osmosis_ai/rollout/types/sample.py b/osmosis_ai/rollout/types/sample.py index 748b622d..97c49fb2 100644 --- a/osmosis_ai/rollout/types/sample.py +++ b/osmosis_ai/rollout/types/sample.py @@ -9,7 +9,15 @@ class RolloutSample(BaseModel): - id: str + """The conversation + grading artefacts produced by one rollout. + + A rollout produces exactly one sample (one agent run, one reward). There + used to be a per-sample id so the wire protocol could carry a + ``dict[str, RolloutSample]``; with the URL-routed single-sample wire + protocol the id is gone and callers identify rollouts via the URL paths + they hand the SDK. + """ + messages: Sequence[SampleMessage] = Field(default_factory=list) label: str | None = None reward: float | None = None @@ -33,11 +41,6 @@ class RolloutErrorCategory(StrEnum): AGENT_ERROR = "agent_error" -class MultiTurnMode(StrEnum): - MULTI_SAMPLE = "multi_sample" - SINGLE_SAMPLE = "single_sample" - - class ExecutionRequest(BaseModel): id: str prompt: list[MessageDict] @@ -49,7 +52,7 @@ class ExecutionRequest(BaseModel): class ExecutionResult(BaseModel): status: RolloutStatus - samples: dict[str, RolloutSample] = Field(default_factory=dict) + sample: RolloutSample | None = None artifacts: dict[str, Any] | None = None err_message: str | None = None err_category: RolloutErrorCategory | None = None diff --git a/osmosis_ai/rollout/utils/rewards.py b/osmosis_ai/rollout/utils/rewards.py index d736d638..98dc4b19 100644 --- a/osmosis_ai/rollout/utils/rewards.py +++ b/osmosis_ai/rollout/utils/rewards.py @@ -1,10 +1,18 @@ from osmosis_ai.rollout.types import RolloutSample -def validate_samples_have_rewards(samples: dict[str, RolloutSample]) -> None: - if not samples: - raise ValueError("No samples after grading") +def validate_sample_has_reward(sample: RolloutSample | None) -> None: + """Raise if the rollout did not produce a graded sample. - for sid, sample in samples.items(): - if sample.reward is None: - raise ValueError(f"Sample {sid} has no reward after grading") + A rollout has at most one sample; an ungraded or missing sample is a + grader contract violation (or a workflow that forgot to register an + agent/session, in which case ``sample`` is ``None``). + """ + if sample is None: + raise ValueError( + "No sample to grade. The workflow must register a sample " + "source (e.g. via OsmosisStrandsAgent or OsmosisMemorySession) " + "before grading." + ) + if not sample.remove_sample and sample.reward is None: + raise ValueError("Sample has no reward after grading")