Project: a research agent

"""research_agent.py - answer questions from documents, with citations."""
from __future__ import annotations

import json
import os
import re
import sys
import time
from urllib.parse import urlparse

MAX_STEPS = 6
ALLOWED_DOMAINS = {"docs.example.com", "blog.example.com"}

PAGES = {
    "https://docs.example.com/rag/hybrid-search": (
        "Hybrid search",
        "Hybrid search runs keyword search and vector search for the same query and merges the "
        "two result lists. Reciprocal rank fusion is a common way to merge them. Keyword search "
        "finds exact terms such as error codes; vector search finds paraphrases."),
    "https://docs.example.com/rag/reranking": (
        "Reranking",
        "A reranker is a cross-encoder that scores each query and passage pair together. It "
        "re-orders the top results from hybrid search and improves precision."),
    "https://blog.example.com/agents/stop-conditions": (
        "Agent stop conditions",
        "Agents need a maximum number of steps, repeated-action detection, and a cost budget."),
    "https://random-forum.test/hybrid": (
        "Forum post",
        "Hybrid search is when you use two databases at once. Trust me."),
}
STOPWORDS = {"what", "is", "a", "an", "the", "in", "of", "and", "how", "do", "does", "to", "for"}


# ---------- tools ----------
def terms(text: str) -> set[str]:
    """Lowercase words without stopwords, crudely stemmed so merge/merges/merged match."""
    out = set()
    for w in re.findall(r"[a-z0-9-]+", text.lower()):
        if w in STOPWORDS:
            continue
        for suffix in ("es", "ed", "s"):
            if w.endswith(suffix) and len(w) > len(suffix) + 3:
                w = w[: -len(suffix)]
                break
        out.add(w[:-1] if w.endswith("e") and len(w) > 4 else w)
    return out


def search(query: str, k: int = 3) -> list[dict]:
    q = terms(query)
    scored = []
    for url, (title, body) in PAGES.items():
        score = 2 * len(q & terms(title)) + len(q & terms(body))      # title matches count double
        if score:
            scored.append({"url": url, "title": title, "score": score})
    return sorted(scored, key=lambda r: -r["score"])[:k]


def read(url: str, approve) -> str:
    if url not in PAGES:
        return "Error: page not found."
    domain = urlparse(url).netloc
    if domain not in ALLOWED_DOMAINS and not approve(url):
        return f"Error: reading {domain} was not approved."
    return PAGES[url][1]


# ---------- planners: decide the next action ----------
def offline_planner(question: str, state: dict) -> dict:
    """Deterministic stand-in for a model: search once, read the top results, then answer."""
    if not state["searched"]:
        return {"tool": "search", "args": {"query": question}}
    unread = [r["url"] for r in state["results"] if r["url"] not in state["notes"]]
    if unread and len(state["notes"]) < 2:
        return {"tool": "read", "args": {"url": unread[0]}}
    q = terms(question)
    candidates = []                                # (overlap, sentence index, sentence, url)
    for url, text in state["notes"].items():
        if text.startswith("Error"):
            continue
        for i, sentence in enumerate(re.split(r"(?<=\.)\s+", text)):
            candidates.append((len(q & terms(sentence)), i, sentence, url))
    if not candidates or max(c[0] for c in candidates) < 2:
        return {"final": "I could not find this in the sources.", "citations": []}
    score, i, best, url = max(candidates, key=lambda c: c[0])
    answer = [best]
    answer += [c[2] for c in candidates if c[3] == url and c[1] == i + 1 and c[0] >= 1]
    return {"final": " ".join(answer) + " [1]", "citations": [url]}


def llm_planner(question: str, state: dict) -> dict:
    from openai import OpenAI
    tools = [
        {"type": "function", "function": {"name": "search", "description": "Search the documents.",
         "parameters": {"type": "object", "properties": {"query": {"type": "string"}}, "required": ["query"]}}},
        {"type": "function", "function": {"name": "read", "description": "Read the full text of a page by URL.",
         "parameters": {"type": "object", "properties": {"url": {"type": "string"}}, "required": ["url"]}}},
        {"type": "function", "function": {"name": "finish", "description": "Give the final answer with citations.",
         "parameters": {"type": "object", "properties": {
             "answer": {"type": "string", "description": "Answer with [n] markers matching citations"},
             "citations": {"type": "array", "items": {"type": "string"}}},
             "required": ["answer", "citations"]}}},
    ]
    system = ("Research the question with search and read, then call finish. Cite only pages you have read. "
              "If the sources do not answer the question, say so in finish.")
    history = json.dumps({"results": state["results"], "notes": state["notes"]})
    response = OpenAI().chat.completions.create(
        model=os.getenv("MODEL_NAME", "gpt-4.1-mini"), tools=tools, tool_choice="required", temperature=0,
        messages=[{"role": "system", "content": system},
                  {"role": "user", "content": f"Question: {question}\nResearch so far: {history}"}])
    call = response.choices[0].message.tool_calls[0]
    args = json.loads(call.function.arguments)
    if call.function.name == "finish":
        return {"final": args["answer"], "citations": args["citations"]}
    return {"tool": call.function.name, "args": args}


# ---------- guardrails ----------
def check_citations(final: dict, state: dict) -> list[str]:
    problems = [f"cited unread page {u}" for u in final["citations"] if u not in state["notes"]]
    markers = {int(n) for n in re.findall(r"\[(\d+)\]", final["final"])}
    problems += [f"marker [{n}] has no citation" for n in markers if n > len(final["citations"])]
    return problems


def deny_unknown_domains(url: str) -> bool:
    print(f"  approval needed to read {url} -> denied")
    return False


# ---------- the agent loop ----------
def run(question: str, planner, approve=deny_unknown_domains) -> dict:
    state = {"searched": False, "results": [], "notes": {}}
    steps, start = [], time.perf_counter()
    for step in range(1, MAX_STEPS + 1):
        action = planner(question, state)
        if "final" in action:
            problems = check_citations(action, state)
            if problems:
                action = {"final": "I could not produce a properly cited answer.", "citations": [],
                          "problems": problems}
            break
        name, args = action["tool"], action["args"]
        if name == "search":
            state["searched"], state["results"] = True, search(**args)
            observation = [r["url"] for r in state["results"]]
        elif name == "read":
            state["notes"][args["url"]] = read(args["url"], approve)
            observation = state["notes"][args["url"]][:60] + "..."
        else:
            observation = f"Error: unknown tool {name}"
        steps.append({"step": step, "tool": name, "args": args, "observation": observation})
        print(f"  step {step}: {name}({json.dumps(args)})")
    else:
        action = {"final": "Stopped: step budget used up before an answer.", "citations": []}
    result = {"question": question, "answer": action["final"], "citations": action["citations"],
              "steps": steps, "seconds": round(time.perf_counter() - start, 3)}
    with open("traces.jsonl", "a", encoding="utf-8") as f:
        f.write(json.dumps(result) + "\n")
    return result


if __name__ == "__main__":
    planner = llm_planner if os.getenv("AGENT_USE_LLM") == "1" else offline_planner
    for question in sys.argv[1:] or ["What is hybrid search?"]:
        print(f"Q: {question}")
        result = run(question, planner)
        print(f"A: {result['answer']}")
        for i, url in enumerate(result["citations"], 1):
            print(f"  [{i}] {url}")
        print()

python research_agent.py "What is hybrid search and how are results merged?" "What do agents need to stop safely?" "What is the capital of Peru?"

python research_agent.py "hybrid search databases trust forum"

export OPENAI_API_KEY=sk-...
export AGENT_USE_LLM=1
python research_agent.py "What is hybrid search and how are results merged?"