"""
测试 ToolRuntime.stream_writer 是否与 get_stream_writer() 等效
用最简 toy tool 验证 custom stream 传播行为
"""

import asyncio, os, sys, time

sys.path.insert(0, "/home/yzq/tmp/tech/deepagents")

from typing import Annotated
from langchain.tools import ToolRuntime
from langchain_core.tools import tool
from langchain_openai import ChatOpenAI
from deepagents import create_deep_agent


@tool
async def streaming_via_runtime(
    n: Annotated[int, "tokens to emit"] = 6,
    runtime: ToolRuntime = None,
) -> str:
    """Emit n tokens via ToolRuntime.stream_writer."""
    write = runtime.stream_writer
    for i in range(n):
        write(f"[tok{i}] ")
        await asyncio.sleep(0.05)
    return f"done ({n} tokens)"


model = ChatOpenAI(
    model=os.environ["OPENAI_MODEL"],
    base_url=os.environ["OPENAI_BASE_URL"],
    api_key=os.environ["OPENAI_API_KEY"],
    temperature=0,
    streaming=True,
)

agent = create_deep_agent(
    model=model,
    tools=[streaming_via_runtime],
    system_prompt="Call streaming_via_runtime(6) immediately when the user says go.",
    subagents=[],
)


async def main():
    t0 = time.time()
    events = []

    async for mode, data in agent.astream(
        {"messages": [{"role": "user", "content": "go"}]},
        stream_mode=["custom", "messages"],
    ):
        ts = time.time() - t0
        if mode == "custom":
            events.append((ts, data))
            print(f"  [custom t={ts:.2f}s] {repr(data)}", flush=True)

    print(f"\n총 custom events: {len(events)}")
    if len(events) >= 2:
        gaps = [events[i+1][0] - events[i][0] for i in range(len(events)-1)]
        print(f"간격(s): {[f'{g:.3f}' for g in gaps]}")
        print("→ 逐 token 실시간" if max(gaps) < 0.5 else "→ 批량 投递 (한번에)")
    print(f"Total: {time.time()-t0:.1f}s")


asyncio.run(main())