marchwarden/cli/main.py

"""Marchwarden CLI shim.

Talks to the web researcher MCP server over stdio and pretty-prints
ResearchResult contracts to the terminal.
"""

import asyncio
import json
import sys
from typing import Optional

import click
from mcp import ClientSession, StdioServerParameters
from mcp.client.stdio import stdio_client
from rich.console import Console
from rich.panel import Panel
from rich.table import Table
from rich.text import Text

from researchers.web.models import ResearchResult


# ---------------------------------------------------------------------------
# MCP client
# ---------------------------------------------------------------------------


async def call_research_tool(
    question: str,
    depth: str,
    max_iterations: int,
    token_budget: int,
) -> ResearchResult:
    """Spawn the web researcher MCP server and call its `research` tool."""
    params = StdioServerParameters(
        command=sys.executable,
        args=["-m", "researchers.web.server"],
    )
    async with stdio_client(params) as (read, write):
        async with ClientSession(read, write) as session:
            await session.initialize()
            result = await session.call_tool(
                "research",
                arguments={
                    "question": question,
                    "depth": depth,
                    "max_iterations": max_iterations,
                    "token_budget": token_budget,
                },
            )
            # FastMCP returns the tool's string return as a TextContent block.
            payload = result.content[0].text
            return ResearchResult.model_validate_json(payload)


# ---------------------------------------------------------------------------
# Pretty printing
# ---------------------------------------------------------------------------


def render_result(result: ResearchResult, console: Console) -> None:
    """Render a ResearchResult to the console using rich."""
    # Answer
    console.print(
        Panel(
            result.answer,
            title="[bold cyan]Answer[/bold cyan]",
            border_style="cyan",
        )
    )

    # Citations
    if result.citations:
        table = Table(title="Citations", show_lines=True, expand=True)
        table.add_column("#", style="dim", width=3)
        table.add_column("Title / Locator", overflow="fold")
        table.add_column("Excerpt", overflow="fold")
        table.add_column("Conf", justify="right", width=5)
        for i, c in enumerate(result.citations, 1):
            header = f"[bold]{c.title or c.locator}[/bold]\n[dim]{c.locator}[/dim]"
            table.add_row(str(i), header, c.raw_excerpt, f"{c.confidence:.2f}")
        console.print(table)
    else:
        console.print("[dim]No citations.[/dim]")

    # Gaps grouped by category
    if result.gaps:
        gap_table = Table(title="Gaps", show_lines=True, expand=True)
        gap_table.add_column("Category", style="yellow")
        gap_table.add_column("Topic")
        gap_table.add_column("Detail", overflow="fold")
        for g in result.gaps:
            gap_table.add_row(g.category.value, g.topic, g.detail)
        console.print(gap_table)

    # Discovery events
    if result.discovery_events:
        de_table = Table(title="Discovery Events", show_lines=True, expand=True)
        de_table.add_column("Type", style="magenta")
        de_table.add_column("Suggested Researcher")
        de_table.add_column("Query", overflow="fold")
        de_table.add_column("Reason", overflow="fold")
        for d in result.discovery_events:
            de_table.add_row(
                d.type, d.suggested_researcher or "-", d.query, d.reason
            )
        console.print(de_table)

    # Open questions
    if result.open_questions:
        oq_table = Table(title="Open Questions", show_lines=True, expand=True)
        oq_table.add_column("Priority", style="green")
        oq_table.add_column("Question", overflow="fold")
        oq_table.add_column("Context", overflow="fold")
        for q in result.open_questions:
            oq_table.add_row(q.priority, q.question, q.context)
        console.print(oq_table)

    # Confidence + factors
    cf = result.confidence_factors
    conf_text = Text()
    conf_text.append(f"Overall: {result.confidence:.2f}\n", style="bold")
    conf_text.append(f"Corroborating sources: {cf.num_corroborating_sources}\n")
    conf_text.append(f"Source authority: {cf.source_authority}\n")
    conf_text.append(f"Contradiction detected: {cf.contradiction_detected}\n")
    conf_text.append(f"Query specificity match: {cf.query_specificity_match:.2f}\n")
    conf_text.append(f"Budget exhausted: {cf.budget_exhausted}\n")
    conf_text.append(f"Recency: {cf.recency or 'unknown'}")
    console.print(Panel(conf_text, title="Confidence", border_style="green"))

    # Cost
    cm = result.cost_metadata
    cost_text = Text()
    cost_text.append(f"Tokens: {cm.tokens_used}\n")
    cost_text.append(f"Iterations: {cm.iterations_run}\n")
    cost_text.append(f"Wall time: {cm.wall_time_sec:.2f}s\n")
    cost_text.append(f"Model: {cm.model_id}")
    console.print(Panel(cost_text, title="Cost", border_style="blue"))

    # Trace footer
    console.print(f"\n[dim]trace_id: {result.trace_id}[/dim]")


# ---------------------------------------------------------------------------
# Click app
# ---------------------------------------------------------------------------


@click.group()
def cli() -> None:
    """Marchwarden — agentic research CLI."""


@cli.command()
@click.argument("question")
@click.option(
    "--depth",
    type=click.Choice(["shallow", "balanced", "deep"]),
    default="balanced",
    show_default=True,
)
@click.option(
    "--budget",
    "token_budget",
    type=int,
    default=20_000,
    show_default=True,
    help="Token budget for the research loop.",
)
@click.option(
    "--max-iterations",
    type=int,
    default=5,
    show_default=True,
)
def ask(
    question: str,
    depth: str,
    token_budget: int,
    max_iterations: int,
) -> None:
    """Ask the web researcher a QUESTION."""
    console = Console()
    console.print(f"[dim]Researching:[/dim] {question}\n")
    try:
        result = asyncio.run(
            call_research_tool(
                question=question,
                depth=depth,
                max_iterations=max_iterations,
                token_budget=token_budget,
            )
        )
    except Exception as e:
        console.print(f"[bold red]Error:[/bold red] {e}")
        sys.exit(1)
    render_result(result, console)


if __name__ == "__main__":
    cli()