<#704 codegen plugin> Pull request opened by <samh...
# flyte-github
a
#704 codegen plugin Pull request opened by samhita-alla LLM-driven single-file Python code generation plugin • Adds
flyteplugins-codegen
, a plugin that uses LLMs to generate, test, and iterate on single-file Python scripts inside isolated sandboxes • Supports two execution backends: LiteLLM mode (any model with structured output support) and Claude Agent SDK mode (autonomous Claude agent) AutoCoderAgent - the main entry point: • Takes a prompt + optional samples/schema/constraints, generates a Python script, writes pytest tests, builds a sandbox image, runs tests, diagnoses failures, and iterates until tests pass •
result.run()
/
result.as_task()
execute the generated code against real data in an isolated container • Supports data extraction from DataFrames and Files, auto-generates Pandera schemas, and infers input/output types import flyte from flyte.io import File from flyte.sandbox import sandbox_environment from flyteplugins.codegen import AutoCoderAgent agent = AutoCoderAgent(model="gpt-4.1", name="summarize-sales", resources=flyte.Resources(cpu=1, memory="1Gi")) env = flyte.TaskEnvironment( name="my-env", secrets=[flyte.Secret(key="openai_key", as_env_var="OPENAI_API_KEY")], image=flyte.Image.from_debian_base().with_pip_packages( "flyteplugins-codegen", ), depends_on=[sandbox_environment], # Required ) @env.task async def process_data(csv_file: File) -> tuple[float, int, int]: result = await agent.generate.aio( prompt="Read the CSV and compute total_revenue, total_units and row_count.", samples={"sales": csv_file}, outputs={"total_revenue": float, "total_units": int, "row_count": int}, ) return await result.run.aio() Agent mode (
backend=claude
):
• Runs an autonomous Claude agent that writes files, runs tests (intercepted into sandboxes), and fixes issues • Hooks checkpoint intermediate artifacts to named remote storage for durable execution across retries agent = AutoCoderAgent( name="my-task", model="claude-sonnet-4-5-20250929", backend="claude", # Requires ANTHROPIC_API_KEY as a Flyte secret ) result = await agent.generate.aio( prompt="...", samples={"input": my_file}, outputs={"result": str}, ) flyteorg/flyte-sdk