"""
DeepAgents 流式写文章测试 — 方案 B
主 Agent → jina-writer subagent → jina_deepsearch_stream tool
使用 get_stream_writer() + astream(stream_mode=["custom","messages"])
"""

import asyncio, os, sys, time

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

from langchain_openai import ChatOpenAI
from deepagents import create_deep_agent
from jina_tools import jina_writer_stream_subagent

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,
)

MAIN_PROMPT = """You are a helpful writing assistant.
When the user asks to write an article, delegate to the "jina-writer" subagent via task().
Pass the full topic and requirements. Return the result directly without modification.
"""

agent = create_deep_agent(
    model=model,
    tools=[],
    system_prompt=MAIN_PROMPT,
    subagents=[jina_writer_stream_subagent],
)

QUERY = "Write a concise 500-word article about how Jina AI DeepSearch works and its advantages over traditional RAG. Include technical details and cite sources."


async def main():
    print(f"Query: {QUERY}\n{'='*60}\n")

    t0 = time.time()
    got_custom = False   # tracks whether streaming tool output arrived

    async for mode, data in agent.astream(
        {"messages": [{"role": "user", "content": QUERY}]},
        stream_mode=["custom", "messages"],
    ):
        if mode == "custom":
            # real-time tokens from jina_deepsearch_stream via get_stream_writer()
            if not got_custom:
                print("[jina_deepsearch_stream → real-time output]\n")
                got_custom = True
            print(data, end="", flush=True)

        elif mode == "messages":
            chunk, metadata = data
            node = metadata.get("langgraph_node", "")
            content = getattr(chunk, "content", "") or ""
            # Only echo main agent's final synthesis if no custom stream occurred
            if content and node == "model" and not got_custom:
                print(content, end="", flush=True)

    print(f"\n\n{'='*60}")
    print(f"Streaming mode: {'custom (real-time)' if got_custom else 'messages (fallback)'}")
    print(f"Total time: {time.time()-t0:.1f}s")


asyncio.run(main())