import asyncio
import os
from typing import Any

import httpx
from pydantic import BaseModel
from timbal import Agent, Tool, Workflow
from timbal.errors import bail
from timbal.state import get_run_context
from timbal.tools import WebSearch
from timbal.types import File

# os.environ.pop("ANTHROPIC_API_KEY", None)
os.environ["TIMBAL_API_HOST"] = "api.timbal.ai"
os.environ["TIMBAL_ORG_ID"] = "1"
os.environ["TIMBAL_APP_ID"] = "177"
# os.environ["TIMBAL_PROJECT_ID"] = "60"
os.environ["TIMBAL_LOG_EVENTS"] = "START,OUTPUT,DELTA"


def step_a():
    return {"value": 5}


def step_b(y):
    print(get_run_context().parent_span().input.get("x"))
    print(y)


sub_workflow = (
    Workflow(name="sub_test").step(step_a).step(step_b, y=lambda: get_run_context().parent_span().input.get("x"))
)


main_workflow = Workflow(name="main_test").step(sub_workflow, x=3)


async def main():
    await main_workflow().collect()


if __name__ == "__main__":
    asyncio.run(main())
