AgentWorkflow

from osmosis_ai.rollout import AgentWorkflow, AgentWorkflowContext


class MyWorkflow(AgentWorkflow):
    async def run(self, ctx: AgentWorkflowContext) -> None:
        # Build and run your agent here.
        pass

class AgentWorkflow[TConfig: AgentWorkflowConfig](ABC):
    def __init__(self, config: TConfig | None = None):
        self.config = config

    @abstractmethod
    async def run(
        self, ctx: AgentWorkflowContext[TConfig]
    ) -> AgentWorkflowOutput | list[dict[str, Any]] | None:
        raise NotImplementedError

from osmosis_ai.rollout import (
    AgentWorkflow,
    AgentWorkflowContext,
    AgentWorkflowOutput,
)


class ExplicitWorkflow(AgentWorkflow):
    async def run(self, ctx: AgentWorkflowContext) -> AgentWorkflowOutput:
        messages = [*ctx.prompt, {"role": "assistant", "content": "done"}]
        return AgentWorkflowOutput(
            messages=messages,
            metrics={"steps": 1.0},
        )

import asyncio

from osmosis_ai.rollout import AgentWorkflow, AgentWorkflowContext


class TracingWorkflow(AgentWorkflow):
    async def run(self, ctx: AgentWorkflowContext) -> None:
        if ctx.artifacts_dir:
            log_file = ctx.artifacts_dir / "run.log"
            await asyncio.to_thread(log_file.write_text, "started\n")
        ...

~/.osmosis/<rollout_id>/
├── trajectory.json          # rollout sample 的 ATIF 文档
└── artifacts/...            # 你在 ctx.artifacts_dir 下写入的文件

from osmosis_ai.rollout import AgentWorkflow, AgentWorkflowContext, get_rollout_context
from osmosis_ai.rollout.integrations.agents.strands import (
    OsmosisRolloutModel,
    OsmosisStrandsAgent,
)


class PrivateStrandsWorkflow(AgentWorkflow):
    async def run(self, ctx: AgentWorkflowContext) -> None:
        rollout_context = get_rollout_context()
        if rollout_context is not None:
            rollout_context.preview_enabled = False
        agent = OsmosisStrandsAgent(
            model=OsmosisRolloutModel(),
            messages=ctx.prompt,
            callback_handler=None,
        )
        await agent.invoke_async()

from copy import deepcopy

from osmosis_ai.rollout import RolloutSample, SampleSource


class MySampleSource(SampleSource):
    def __init__(self, messages: list[dict]):
        self._messages = messages

    async def get_sample(self) -> RolloutSample:
        return RolloutSample(messages=list(self._messages))

    async def get_preview(self, max_messages: int) -> RolloutSample | None:
        if max_messages <= 0:
            return None
        return RolloutSample(messages=deepcopy(self._messages[-max_messages:]))

from osmosis_ai.rollout import AgentWorkflow, AgentWorkflowContext
from osmosis_ai.rollout.integrations.agents.strands import (
    OsmosisRolloutModel,
    OsmosisStrandsAgent,
)


class SimpleStrandsWorkflow(AgentWorkflow):
    async def run(self, ctx: AgentWorkflowContext) -> None:
        agent = OsmosisStrandsAgent(
            name="simple-strands-agent",
            model=OsmosisRolloutModel(params={"temperature": 1.0}),
            messages=ctx.prompt,
            callback_handler=None,
        )
        await agent.invoke_async()

from agents import ModelSettings, Runner
from osmosis_ai.rollout import AgentWorkflow, AgentWorkflowContext
from osmosis_ai.rollout.integrations.agents.openai_agents import (
    OsmosisAgent,
    OsmosisMemorySession,
    OsmosisRolloutModel,
)


class SimpleOpenAIWorkflow(AgentWorkflow):
    async def run(self, ctx: AgentWorkflowContext) -> None:
        agent = OsmosisAgent(
            name="simple-openai-agent",
            instructions="Answer the user's request clearly.",
            model=OsmosisRolloutModel(),
            model_settings=ModelSettings(temperature=1.0, max_tokens=4096),
        )
        session = OsmosisMemorySession()
        await Runner.run(
            agent,
            ctx.prompt,
            session=session,
        )

from osmosis_ai.rollout import (
    AgentWorkflow,
    AgentWorkflowConfig,
    AgentWorkflowContext,
    ConcurrencyConfig,
)


class SearchWorkflowConfig(AgentWorkflowConfig):
    name: str = "search-workflow"
    max_iterations: int = 8
    temperature: float = 1.0
    concurrency: ConcurrencyConfig = ConcurrencyConfig(max_concurrent=4)


class SearchWorkflow(AgentWorkflow[SearchWorkflowConfig]):
    async def run(self, ctx: AgentWorkflowContext[SearchWorkflowConfig]) -> None:
        config = ctx.config or SearchWorkflowConfig()
        max_iterations = config.max_iterations
        temperature = config.temperature
        # Use these values when constructing your agent.


search_workflow_config = SearchWorkflowConfig()

from typing import Any

from strands import tool
from osmosis_ai.rollout import (
    AgentWorkflow,
    AgentWorkflowConfig,
    AgentWorkflowContext,
)
from osmosis_ai.rollout.integrations.agents.strands import (
    OsmosisRolloutModel,
    OsmosisStrandsAgent,
)


@tool(name="search")
def search_tool(query: str) -> str:
    """Search for information."""
    return f"results for {query}"


class ToolWorkflowConfig(AgentWorkflowConfig):
    name: str = "tool-workflow"
    model: Any = OsmosisRolloutModel(params={"temperature": 1.0})
    tools: Any = [search_tool]
    max_iterations: int = 8


tool_workflow_config = ToolWorkflowConfig()


class ToolWorkflow(AgentWorkflow[ToolWorkflowConfig]):
    async def run(self, ctx: AgentWorkflowContext[ToolWorkflowConfig]) -> None:
        config = ctx.config or ToolWorkflowConfig()
        agent = OsmosisStrandsAgent(
            name="search-agent",
            model=config.model,
            tools=config.tools,
            messages=ctx.prompt,
            callback_handler=None,
        )

        for _ in range(config.max_iterations):
            result = await agent.invoke_async()
            content = result.message.get("content", [])
            if not any("toolUse" in block for block in content):
                break

from osmosis_ai.rollout import LocalBackend
from osmosis_ai.rollout.server import create_rollout_server

backend = LocalBackend(
    workflow=SearchWorkflow,
    workflow_config=search_workflow_config,
)
app = create_rollout_server(backend=backend)