"""
最简 DeepAgents + Jina DeepSearch subagent 集成
流式主 Agent → jina-writer subagent → 写文章
"""

import 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, jina_search, jina_reader, jina_writer_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,
)

# ── 主 Agent ──────────────────────────────────────────────────
MAIN_PROMPT = """You are a helpful writing assistant.
When the user asks to write an article, report, or research piece:
1. Delegate the writing task to the "jina-writer" subagent using task().
2. Include the full topic and any requirements in the task prompt.
3. Return the completed article from the subagent to the user without modification.
Do not write the article yourself — always use the jina-writer subagent.
"""

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

# ── 流式调用 ──────────────────────────────────────────────────
query = "Write a concise article (about 500 words) about how Jina AI DeepSearch works and its advantages over traditional RAG. Include real technical details and cite sources."

print(f"User: {query}\n")
print("=" * 60)
print("Streaming output:\n")

t0 = time.time()
current_node = None

for chunk, metadata in agent.stream(
    {"messages": [{"role": "user", "content": query}]},
    stream_mode="messages",
):
    node = metadata.get("langgraph_node", "")

    # 节点切换提示
    if node != current_node:
        if current_node is not None:
            print()
        current_node = node
        label = {"agent": "🤖 Main Agent", "subagraph:jina-writer": "✍️  jina-writer"}.get(node, f"[{node}]")
        print(f"\n── {label} ──")

    content = getattr(chunk, "content", "") or ""
    if content:
        print(content, end="", flush=True)

print(f"\n\n{'='*60}")
print(f"Total time: {time.time()-t0:.1f}s")