Tiered Workflow Factory

Build subscription-specific workflow pipelines with a WorkflowFactory.

02_tiered_workflow_factory.py

"""Tiered Workflow Factory -- pipeline depth based on subscription.

Free-tier tenants get a 2-step pipeline (draft + edit).
Enterprise tenants get a 3-step pipeline (research + draft + edit) with a better model.

Run:
    .venvs/demo/bin/python cookbook/05_agent_os/factories/workflow/02_tiered_workflow_factory.py

Test:
    # Free tier (2 steps, cheaper model)
    curl -X POST http://localhost:7777/workflows/article-pipeline/runs \
        -H "Authorization: Bearer <FREE_TOKEN>" \
        -F 'message=Write an article about remote work trends' \
        -F 'stream=false'

    # Enterprise tier (3 steps, best model)
    curl -X POST http://localhost:7777/workflows/article-pipeline/runs \
        -H "Authorization: Bearer <ENTERPRISE_TOKEN>" \
        -F 'message=Write an article about remote work trends' \
        -F 'stream=false'
"""

from datetime import UTC, datetime, timedelta

import jwt as pyjwt
from agno.agent import Agent
from agno.db.postgres import PostgresDb
from agno.factory import RequestContext
from agno.models.openai import OpenAIResponses
from agno.os import AgentOS
from agno.os.middleware import JWTMiddleware
from agno.workflow.factory import WorkflowFactory
from agno.workflow.step import Step
from agno.workflow.workflow import Workflow

# ---------------------------------------------------------------------------
# Config
# ---------------------------------------------------------------------------

JWT_SECRET = "a-string-secret-at-least-256-bits-long"

db = PostgresDb(
    id="tiered-workflow-db",
    db_url="postgresql+psycopg://ai:ai@localhost:5532/ai",
)

TIER_MODELS = {
    "free": "gpt-4.1-mini",
    "enterprise": "gpt-5.4",
}

# ---------------------------------------------------------------------------
# Factory
# ---------------------------------------------------------------------------


def build_article_pipeline(ctx: RequestContext) -> Workflow:
    """Build an article pipeline whose depth depends on subscription tier."""
    claims = ctx.trusted.claims
    tier = claims.get("tier", "free")
    model_id = TIER_MODELS.get(tier, TIER_MODELS["free"])

    steps = []

    # Enterprise gets a research step first
    if tier == "enterprise":
        researcher = Agent(
            name="Researcher",
            model=OpenAIResponses(id=model_id),
            instructions="Research the topic thoroughly. Provide key facts, statistics, and sources.",
        )
        steps.append(
            Step(name="research", description="Research the topic", agent=researcher)
        )

    drafter = Agent(
        name="Drafter",
        model=OpenAIResponses(id=model_id),
        instructions="Write a well-structured article draft based on the input. Be thorough but readable.",
    )
    steps.append(
        Step(name="draft", description="Write the article draft", agent=drafter)
    )

    editor = Agent(
        name="Editor",
        model=OpenAIResponses(id=model_id),
        instructions="Edit the article for clarity, flow, and correctness. Output the final polished version.",
    )
    steps.append(Step(name="edit", description="Edit and finalize", agent=editor))

    return Workflow(
        name="Article Pipeline",
        description=f"Article pipeline ({len(steps)} steps, {tier} tier)",
        db=db,
        steps=steps,
    )


article_pipeline_factory = WorkflowFactory(
    db=db,
    id="article-pipeline",
    name="Article Pipeline",
    description="Article pipeline -- depth and model quality scale with subscription tier",
    factory=build_article_pipeline,
)

# ---------------------------------------------------------------------------
# AgentOS
# ---------------------------------------------------------------------------

agent_os = AgentOS(
    id="tiered-workflow-demo",
    description="Demo: tiered workflow factory with JWT",
    workflows=[article_pipeline_factory],
)
app = agent_os.get_app()

app.add_middleware(
    JWTMiddleware,
    verification_keys=[JWT_SECRET],
    algorithm="HS256",
    user_id_claim="sub",
    validate=False,
)

# ---------------------------------------------------------------------------
# Run
# ---------------------------------------------------------------------------

if __name__ == "__main__":

    def make_token(tier: str, user_id: str = "user_1") -> str:
        payload = {
            "sub": user_id,
            "tier": tier,
            "exp": datetime.now(UTC) + timedelta(hours=24),
            "iat": datetime.now(UTC),
        }
        return pyjwt.encode(payload, JWT_SECRET, algorithm="HS256")

    print("Test tokens (valid for 24h):")
    print()
    print(f"  FREE:        {make_token('free')}")
    print(f"  ENTERPRISE:  {make_token('enterprise')}")
    print()

    agent_os.serve(app="02_tiered_workflow_factory:app", port=7777, reload=True)

Run the Example

1. Set up your virtual environment

   Mac

   uv venv --python 3.12
   source .venv/bin/activate
   

   uv venv --python 3.12
   .venv\Scripts\activate
   
2. Install dependencies

   uv pip install -U "agno[os]" "psycopg[binary]" openai
   
3. Export your OpenAI API key

   Mac/Linux

   export OPENAI_API_KEY="your_openai_api_key_here"
   

   $Env:OPENAI_API_KEY="your_openai_api_key_here"
   
4. Run PgVector

   macOS / Linux

   docker run -d \
     -e POSTGRES_DB=ai \
     -e POSTGRES_USER=ai \
     -e POSTGRES_PASSWORD=ai \
     -e PGDATA=/var/lib/postgresql \
     -v pgvolume:/var/lib/postgresql \
     -p 5532:5432 \
     --name pgvector \
     agnohq/pgvector:18
   

   docker run -d `
     -e POSTGRES_DB=ai `
     -e POSTGRES_USER=ai `
     -e POSTGRES_PASSWORD=ai `
     -e PGDATA=/var/lib/postgresql `
     -v pgvolume:/var/lib/postgresql `
     -p 5532:5432 `
     --name pgvector `
     agnohq/pgvector:18
   
5. Run the example

   Save the code above as 02_tiered_workflow_factory.py, then run:

   python 02_tiered_workflow_factory.py
   

Full source: cookbook/05_agent_os/factories/workflow/02_tiered_workflow_factory.py