Multi-Agent
Research System

from anthropic import Anthropic, AsyncAnthropic

DISPATCH_SUBAGENT = {
    "name": "dispatch_subagent",
    "description": (
        "하위 리서치 작업을 Subagent에게 위임한다. "
        "task에 'A 회사 매출 추이를 2020~2024 조사' 처럼 단일·구체 목표를 적어라. "
        "동시 dispatch 가능. 작업이 독립적일수록 좋다."
    ),
    "input_schema": {
        "type": "object",
        "properties": {
            "task": {"type": "string", "description": "단일 목표, 한 문장"},
            "tools": {"type": "array", "items": {"type": "string"}},
        },
        "required": ["task"],
        "additionalProperties": False,
    },
}

aclient = AsyncAnthropic()

SUBAGENT_SYSTEM = """너는 리서치 Subagent다. 단일 task에만 집중하라.
- 결과는 사실 기반으로 압축. 5문장 이내.
- 각 사실 옆에 (source: URL) 표기.
- task 범위를 넘는 추가 탐색 금지."""

async def run_subagent(task: str, tools: list) -> dict:
    messages = [{"role": "user", "content": task}]

    for step in range(8):  # max_steps 안전핀
        resp = await aclient.messages.create(
            model="claude-sonnet-4-6",
            max_tokens=2048,
            system=SUBAGENT_SYSTEM,
            tools=tools,
            messages=messages,
        )
        if resp.stop_reason == "end_turn":
            return {"task": task, "findings": resp.content[0].text}

        messages.append({"role": "assistant", "content": resp.content})
        results = []
        for block in resp.content:
            if block.type == "tool_use":
                output = await run_tool(block.name, block.input)  # 외부 도구
                results.append({
                    "type": "tool_result",
                    "tool_use_id": block.id,
                    "content": output[:8000],  # truncate — 결과 폭발 방지
                })
        messages.append({"role": "user", "content": results})

    return {"task": task, "findings": "<incomplete: max_steps>"}

import asyncio

LEAD_SYSTEM = """너는 Lead Research Agent다.
1) 사용자 질문을 받아 step-by-step으로 사고하라.
2) 독립적으로 병렬 탐색 가능한 sub-task로 분해.
3) 각 sub-task를 dispatch_subagent 도구로 한 번에 여러 개 호출.
4) findings가 모이면 synthesize.

규칙: sub-task 수는 3~8개. 강한 의존성이 있으면 분해하지 마라."""

async def lead_agent(user_query: str) -> str:
    messages = [{"role": "user", "content": user_query}]
    tools = [DISPATCH_SUBAGENT]

    for turn in range(5):
        resp = await aclient.messages.create(
            model="claude-opus-4-7",  # Lead는 비싸도 강한 모델
            max_tokens=4096,
            system=LEAD_SYSTEM,
            tools=tools,
            messages=messages,
        )
        if resp.stop_reason == "end_turn":
            return resp.content[0].text

        messages.append({"role": "assistant", "content": resp.content})

        # 같은 turn 안의 모든 tool_use를 모아 한 번에 병렬 실행
        dispatches = [b for b in resp.content if b.type == "tool_use"]
        sub_results = await asyncio.gather(
            *[run_subagent(b.input["task"], SUBAGENT_TOOLS) for b in dispatches],
            return_exceptions=True,
        )

        tool_results = []
        for b, r in zip(dispatches, sub_results):
            content = (f"<error>{r}</error>" if isinstance(r, Exception)
                       else r["findings"])
            tool_results.append({
                "type": "tool_result",
                "tool_use_id": b.id,
                "content": content,
            })
        messages.append({"role": "user", "content": tool_results})

    raise RuntimeError("lead exhausted turns without end_turn")

CITATION_SYSTEM = """너는 Citation Agent다. Lead가 작성한 본문의 각 주장 옆에
findings에 있는 (source: URL)을 [^N] 형태로 매핑하라.
- 본문은 그대로 두고 인용 표기만 추가.
- 근거가 없는 주장은 [^?] 로 표시 (사후 검토 대상)."""

async def add_citations(synthesis: str, findings: list[dict]) -> str:
    sources = "\n".join(f"- {f['task']}\n{f['findings']}" for f in findings)
    resp = await aclient.messages.create(
        model="claude-haiku-4-5",  # 후처리는 작은 모델로 충분
        max_tokens=4096,
        system=CITATION_SYSTEM,
        messages=[{"role": "user", "content":
            f"<synthesis>\n{synthesis}\n</synthesis>\n\n<findings>\n{sources}\n</findings>"}],
    )
    return resp.content[0].text

JUDGE_RUBRIC = """다음 기준으로 0~10점 채점하라.
- 정확성 (claims가 sources에 의해 뒷받침되는가)
- 완전성 (질문의 모든 측면을 커버했는가)
- 인용 품질 (claim ↔ source 매핑이 정확한가)

마지막 줄: score=N (정수)."""

async def judge(query: str, output: str) -> int:
    resp = await aclient.messages.create(
        model="claude-opus-4-7",  # judge는 강한 모델로
        max_tokens=1024,
        system=JUDGE_RUBRIC,
        messages=[{"role": "user", "content":
            f"<query>{query}</query>\n<output>{output}</output>"}],
    )
    text = resp.content[0].text
    last_line = text.splitlines()[-1]
    return int(last_line.split("=")[1].strip())

# 평가셋 n=20도 의미 있다 — early signal
async def eval_suite(queries: list[str]):
    outputs = await asyncio.gather(*[lead_agent(q) for q in queries])
    scores = await asyncio.gather(*[judge(q, o) for q, o in zip(queries, outputs)])
    return sum(scores) / len(scores)