"""
方案 B 修复版：jina_deepsearch_stream 作为主 agent 直接 tool
去掉 subagent 中间层 → get_stream_writer() 可传播到父 graph
"""

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_deepsearch_stream, jina_search, jina_reader

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

# jina_deepsearch_stream 直接挂在主 agent — 不经过 subagent
MAIN_PROMPT = """You are a research writer.
When asked to write an article:
1. Call jina_deepsearch_stream with the topic to gather research
2. Use the returned research to write a well-structured Markdown article with citations
"""

agent = create_deep_agent(
    model=model,
    tools=[jina_deepsearch_stream, jina_search, jina_reader],
    system_prompt=MAIN_PROMPT,
    subagents=[],
)

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


async def main():
    print(f"Query: {QUERY}\n{'='*60}\n")
    t0 = time.time()
    phase = {"custom": 0, "messages": 0}

    async for mode, data in agent.astream(
        {"messages": [{"role": "user", "content": QUERY}]},
        stream_mode=["custom", "messages"],
    ):
        if mode == "custom":
            if phase["custom"] == 0:
                print("\n[DeepSearch 实时流]\n", flush=True)
            phase["custom"] += 1
            print(data, end="", flush=True)

        elif mode == "messages":
            chunk, metadata = data
            node = metadata.get("langgraph_node", "")
            content = getattr(chunk, "content", "") or ""
            if content and node == "model":
                if phase["messages"] == 0:
                    print("\n\n[主 Agent 写文章]\n", flush=True)
                phase["messages"] += 1
                print(content, end="", flush=True)

    print(f"\n\n{'='*60}")
    print(f"custom chunks : {phase['custom']}")
    print(f"message chunks: {phase['messages']}")
    print(f"Total time    : {time.time()-t0:.1f}s")


asyncio.run(main())