Build a Research AI Agent in Python

[ User question ]
       │
       ▼
[ Plan: identify subqueries (LLM) ]
       │
       ▼
[ Search loop ]
   ├── web_search (Tavily / Perplexity / Bing)
   ├── fetch_url (full content for promising hits)
   ├── search_internal_docs (your vector store)
   └── done?
       │
       ▼
[ Synthesize with citations ]
       │
       ▼
[ Validator: every claim has [source-N] ]

import httpx

def web_search(query: str, max_results: int = 5) -> str:
    """Search the public web. Returns titles, URLs, and snippets for top results."""
    response = httpx.post(
        "https://api.tavily.com/search",
        json={"api_key": TAVILY_KEY, "query": query, "max_results": max_results},
        timeout=10.0,
    ).json()
    return "\n\n".join(
        f"[source-{i}] {r['title']}\n{r['url']}\n{r['content']}"
        for i, r in enumerate(response.get("results", []), start=1)
    ) or "No results."

def fetch_url(url: str) -> str:
    """Fetch the full text of a URL. Use after web_search to read promising sources."""
    try:
        text = httpx.get(url, timeout=15.0, follow_redirects=True).text
        return strip_html(text)[:20_000]  # cap to keep context manageable
    except Exception as e:
        return f"Error fetching {url}: {e}"

def search_internal_docs(query: str) -> str:
    """Search internal documentation (your private knowledge base)."""
    hits = vector_client.search(query, top_k=3)
    return "\n\n".join(f"[doc-{h.id}] {h.title}\n{h.text[:500]}" for h in hits)

from tenxgraph.core.graph import Agent, StateGraph, ToolNode

tool_node = ToolNode([web_search, fetch_url, search_internal_docs])

agent = Agent(
    model="anthropic/claude-3-5-sonnet",
    system_prompt=[{"role": "system", "content": (
        "You are a research assistant. "
        "Break the question into subqueries. Use web_search and fetch_url to gather sources. "
        "Use search_internal_docs for company-specific information. "
        "When you have enough information, write a clear answer. "
        "EVERY factual claim must include a citation marker like [source-1] or [doc-2] from the tool outputs. "
        "If you cannot find sources, say so explicitly. Do not invent facts."
    )}],
    tool_node="TOOL",
)

import re

CITATION_RE = re.compile(r"\[(?:source|doc)-\d+\]")

def validate_citations(answer: str, max_uncited_sentences: int = 1) -> tuple[bool, str]:
    sentences = re.split(r"(?<=[.!?])\s+", answer)
    uncited = [s for s in sentences if not CITATION_RE.search(s)]
    if len(uncited) > max_uncited_sentences:
        return False, f"{len(uncited)} sentences without citations. Reject."
    return True, "OK"