Bring your own framework

The dedicated pages ( LangGraph, PydanticAI, OpenAI Agents SDK) are just worked examples of one underlying idea: Union.ai is the runtime, your framework is the loop. If your agent library is written in Python, it runs on Union.ai with no special plugin.

This page is a framework-agnostic template. Drop in any library — CrewAI, AutoGen, smolagents, Atomic Agents, a raw provider SDK, or your own homegrown loop — wherever the comments say so.

The core pattern

Put your framework’s agent invocation inside an @env.task. The task gives you a container, durable inputs/outputs, retries, and a span in the dashboard. Everything inside the task is ordinary Python, so the framework behaves exactly as it does locally.

import flyte

# 1. Declare the runtime: image (with YOUR framework's deps), resources, secrets.
env = flyte.TaskEnvironment(
    name="my-agent",
    image=flyte.Image.from_debian_base(python_version=(3, 13)).with_pip_packages(
        # --> your agent framework + its provider packages go here, e.g.:
        # "crewai", "smolagents", "autogen-agentchat", ...
    ),
    resources=flyte.Resources(cpu=1, memory="1Gi"),
    secrets=[flyte.Secret(key="ANTHROPIC_API_KEY")],  # --> your model provider key(s)
)


@env.task(report=True)
async def run_agent(prompt: str) -> str:
    # 2. Build/configure your framework's agent exactly as you would locally.
    #    --> your framework setup goes here
    # agent = MyFramework.Agent(model=..., tools=[...], instructions=...)

    # 3. Run it. Use the framework's own (sync or async) entry point.
    #    --> invoke your framework here
    # result = await agent.run(prompt)

    # 4. Return a serializable value (str, pydantic model, dataclass, ...).
    # return result.output
    ...


if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(run_agent, prompt="...")  # --> your prompt / inputs
    print(run.url)

That is the whole integration. The remaining sections are optional enhancements that make the framework more durable and observable.

Make tools durable

Most frameworks let a tool be any Python callable. To make a tool durable, retryable, and independently observable, have the framework’s tool delegate to an @env.task. The framework still “owns” the tool; the heavy lifting runs on-cluster.

# A durable task that does the real work (IO, compute, external calls).
@env.task
async def fetch_data(source: str) -> dict:
    # --> your real tool implementation (API call, DB query, scrape, ...)
    ...


# Register it with your framework using whatever tool API it exposes.
# The body just awaits the durable task.
#
#   @my_framework.tool                  # --> your framework's tool decorator/registration
#   async def get_data(source: str) -> dict:
#       """Tool description the LLM sees."""
#       return await fetch_data(source)   # runs as a Flyte task, durable + traced

Reach for an @env.task when a tool does real work you want retried, cached, or run on its own hardware (GPU, more memory). For lightweight in-process helpers, a plain @flyte.trace function (below) is enough.

Trace the framework’s internals

If your framework exposes hooks, callbacks, or lets you wrap its node/step functions, decorate those with @flyte.trace to turn each LLM call, tool call, and routing decision into a span — with inputs and outputs captured and checkpointed.

@flyte.trace
async def call_model(messages: list[dict]) -> str:
    # --> wrap the framework's model call (or pass this as the framework's LLM hook)
    ...


@flyte.trace
async def route(state) -> str:
    # --> wrap a routing / decision function so the branch is visible in the dashboard
    ...

For frameworks that don’t expose hooks, wrap the whole run in flyte.group(...) to keep its trace tidy:

@env.task(report=True)
async def run_agent(prompt: str) -> str:
    with flyte.group("my-framework-run"):  # groups everything below under one span
        # --> your framework invocation
        ...

Fan out across containers

Run many independent agents in parallel — each in its own container — with asyncio.gather(). This works for any framework because each call is just an awaited task.

import asyncio


@env.task
async def run_one(task_input: str) -> str:
    # --> one self-contained agent run for a single input
    ...


@env.task
async def run_many(inputs: list[str]) -> list[str]:
    # Each run_one call lands in its own container.
    results = await asyncio.gather(*[run_one(i) for i in inputs])
    return list(results)

Checklist

To bring any Python agent framework onto Union.ai:

  1. Wrap the run — call the framework’s entry point inside an @env.task.
  2. Declare deps — add the framework + provider packages to the task’s image.
  3. Supply secrets — mount model-provider API keys via flyte.Secret.
  4. (Optional) Durable tools — have tools delegate to @env.tasks.
  5. (Optional) Observe — decorate hooks/steps with @flyte.trace, or wrap in flyte.group(...).
  6. (Optional) Scale — fan out with asyncio.gather() for parallel, per-container runs.

Next steps