[feat] url based rollout

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()

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 ""

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

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)

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,
+        )

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()

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 ""

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

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)

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

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",

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()