from cabeza import Agent
from cabeza.tools.search import SearchTool

agent = Agent(
    model="deepseek-v4-flash",
    base_url="https://api.deepseek.com/v1",
    api_key="sk-...",                      # your provider key
    family="deepseek",                     # qwen / glm / kimi / deepseek / gpt / gpt_oss
    tools=[SearchTool(api_key="serper-...")],
    max_steps=8,
    verbose=True,                          # stream step/tool/final to stdout
)

print(agent.run(
    "Which 2017 paper introduced the Transformer architecture? Use web "
    "search to confirm and reply with just the title."
))
# [cabeza] step 1 → search({'query': ['2017 paper Transformer architecture introduced']})
# [Serper]  Searching: 2017 paper Transformer architecture introduced
# [cabeza] step 1 ← search: A Google search for ... found 9 results …
# [cabeza] step 2 ✓ final: "Attention Is All You Need"
# Attention Is All You Need

agent.on("after_llm",  lambda state, parsed: ...)
agent.on("after_tool", lambda state, name, args, result: ...)

# Local Qwen via vLLM
agent = Agent(
    model="Qwen3-30B-A3B-Instruct-2507",
    base_url="http://localhost:8001/v1",
    api_key="EMPTY",
    family="qwen",
    tools=[SearchTool(api_key="serper-...")],
)

# Equivalent — pick whichever fits your shell setup:
agent = Agent(model="deepseek-v4-flash", base_url="https://api.deepseek.com/v1",
              api_key="sk-...", family="deepseek")

# With DEEPSEEK_API_KEY set in the environment, the kwarg can be dropped:
agent = Agent(model="deepseek-v4-flash", base_url="https://api.deepseek.com/v1",
              family="deepseek")

# Tools follow the same pattern:
SearchTool(api_key="serper-...")    # explicit
SearchTool()                         # reads SERPER_API_KEY
VisitTool(jina_api_key="...", llm_base_url=..., llm_model=...)
VisitTool(llm_base_url=..., llm_model=...)   # reads JINA_API_KEY

from cabeza import Agent
from cabeza.tools.search import SearchTool
from cabeza.tools.visit import VisitTool

deepseek_url = "https://api.deepseek.com/v1"
deepseek_key = "sk-..."           # or set DEEPSEEK_API_KEY and drop the kwarg

agent = Agent(
    model="deepseek-v4-flash",
    base_url=deepseek_url,
    api_key=deepseek_key,
    family="deepseek",
    tools=[
        SearchTool(api_key="serper-..."),        # or SearchTool() + $SERPER_API_KEY
        VisitTool(                                # or VisitTool() + $JINA_API_KEY
            jina_api_key="jina_...",
            llm_api_key=deepseek_key,             # extractor LLM for visited pages
            llm_base_url=deepseek_url,
            llm_model="deepseek-v4-flash",
        ),
    ],
    context_management="summary",
    context_management_tokens=64_000,
    context_window_tokens=128_000,
    # Same DeepSeek endpoint reused as the rolling-summary compressor:
    summary_model="deepseek-v4-flash",
    summary_base_url=deepseek_url,
    summary_api_key=deepseek_key,
    max_steps=20,
    timeout=120.0,
    verbose=True,
)
print(agent.run("Which 2024 paper introduced …?"))

from cabeza.agent import MemoryConfig, TeamConfig

agent = Agent(
    model="Qwen3-30B-A3B-Instruct-2507",
    base_url="http://localhost:8001/v1",
    api_key="EMPTY",
    family="qwen",
    team="fugue",
    team_config=TeamConfig(
        mode="fugue",
        size=2,
        # Optional: each peer can use a different family / endpoint.
        # families=["qwen", "deepseek"],
        # models=["Qwen3-30B-A3B-Instruct-2507", "deepseek-v4-flash"],
        # base_urls=["http://localhost:8001/v1", "https://api.deepseek.com/v1"],
        # api_keys=["EMPTY", "sk-..."],
    ),
    context_management="page_memory",
    memory_config=MemoryConfig(
        summarizer_model="gpt-4o-mini",
        summarizer_base_url="https://api.openai.com/v1",
        summarizer_api_key="sk-...",
    ),
    tools=[SearchTool(api_key="serper-...")],
)

from cabeza.datasets import load, JSONLDataset

ds = load("bc")                              # all 1266 items
ds = load("bc", limit=100)                   # first 100 only (--eval-num)
ds = load("dsqa")                            # all 900 DeepSearchQA items
ds = load("gisa")                            # all 373 GISA items
ds = load("bc", path="/path/to/my.jsonl")    # use a different file
ds = JSONLDataset("/path/to/eval.jsonl")     # anything JSONL-shaped

# Override the bundled location entirely:
#   export CABEZA_DATA_ROOT=/somewhere/else

from cabeza.runner import evaluate

summary = evaluate(
    agent,
    load("bc", limit=100),
    out="results/bc_qwen.jsonl",
    workers=4,
    log_dir="logs/bc_qwen",
)
print(summary)   # {"completed": ..., "skipped": ..., "total": ..., ...}

from cabeza.eval import score

score(
    benchmark="bc",
    input_files=["results/bc_qwen.jsonl"],
    judge_model="deepseek-v4-flash",
    judge_base_url="https://api.deepseek.com/v1",
    judge_api_key="sk-...",
    extra_argv=[
        "--scored_output", "results/bc_qwen.scored.jsonl",
        "--summary_output", "results/bc_qwen.summary.json",
        "--num_workers", "16",
    ],
)