agents/research_graph.py

from typing import Any, TypedDict

from agents.services import ResearchResult
from rag.retrieve import extract_response_text, generate_sub_queries
from rag.web_ingest import fetch_urls_parallel, web_search
from rag.notes import save_note

try:
    from langgraph.graph import END, StateGraph
except ImportError:  # pragma: no cover
    END = None
    StateGraph = None


class ResearchGraphState(TypedDict, total=False):
    topic: str
    search_queries: list[str]
    search_results: list[dict]   # {url, title, snippet, content}
    synthesis: str | None
    save_to_notes: bool
    note_id: str | None
    error: str | None
    debug_mode: bool
    debug_data: dict[str, Any] | None


def plan_search_node(state: ResearchGraphState, *, llm, max_queries: int) -> ResearchGraphState:
    print(f"[Research Agent] Node: plan_search | Topic: '{state['topic']}'")
    updated_state = dict(state)
    updated_state["search_queries"] = generate_sub_queries(
        state["topic"],
        llm,
        max_queries=max_queries,
    )
    print(f"[Research Agent] Planned {len(updated_state['search_queries'])} query(s): {updated_state['search_queries']}")
    
    if state.get("debug_mode") and updated_state.get("debug_data") is not None:
        updated_state["debug_data"]["search_queries"] = updated_state["search_queries"]
    
    return updated_state


def execute_search_node(state: ResearchGraphState) -> ResearchGraphState:
    print("[Research Agent] Node: execute_search | Starting web search execution")
    updated_state = dict(state)
    all_results = []
    seen_urls = set()
    
    for query in state.get("search_queries", []):
        try:
            print(f"[Research Agent] Search Query: '{query}'")
            results = web_search(query, max_results=3)
            print(f"[Research Agent] Search returned {len(results)} raw result(s) for '{query}'")
            for r in results:
                if r["url"] not in seen_urls:
                    all_results.append(r)
                    seen_urls.add(r["url"])
        except Exception as exc:
            print(f"[Research Agent] Search FAILED for '{query}': {exc}")
            continue

    if not all_results:
        print("[Research Agent] No search results found across all queries.")
        updated_state["error"] = "no_search_results"
        return updated_state

    # Fetch full content for the top 3 results in parallel
    top_urls = [res["url"] for res in all_results[:3]]
    print(f"[Research Agent] Fetching {len(top_urls)} URL(s) in parallel")
    fetched = {p["source_url"]: p["page_text"] for p in fetch_urls_parallel(top_urls)}
    for res in all_results[:3]:
        res["content"] = fetched.get(res["url"]) or res["snippet"]
        print(f"[Research Agent] Content for {res['url']}: chars={len(res['content'])}")

    updated_state["search_results"] = all_results
    print(f"[Research Agent] Search execution complete | unique results={len(all_results)}")
    
    if state.get("debug_mode") and updated_state.get("debug_data") is not None:
        updated_state["debug_data"]["results_found"] = len(all_results)
        
    return updated_state


def save_note_node(state: ResearchGraphState) -> ResearchGraphState:
    print("[Research Agent] Node: save_note | Checking whether note should be saved")
    updated_state = dict(state)
    if not state.get("save_to_notes") or not state.get("synthesis"):
        print("[Research Agent] Save note skipped.")
        return updated_state

    # For research notes, we use the topic as the query and the synthesis as the answer
    citations = [
        {
            "number": i + 1,
            "source": r["url"],
            "content": r["snippet"]
        }
        for i, r in enumerate(state.get("search_results", []))
    ]
    
    try:
        note_id = save_note(
            query=state["topic"],
            answer=state["synthesis"],
            citations=citations
        )
        updated_state["note_id"] = note_id
        print(f"[Research Agent] Note saved successfully: {note_id}")
    except Exception as exc:
        print(f"[Research Agent] Save note FAILED: {exc}")
        
    return updated_state


def finalize_research_node(state: ResearchGraphState) -> ResearchGraphState:
    print("[Research Agent] Node: finalize")
    updated_state = dict(state)
    if state.get("error") == "no_search_results":
        updated_state["synthesis"] = f"I was unable to find any web search results for the topic: {state['topic']}"
        print("[Research Agent] Finalized with no search results.")
    elif not state.get("synthesis"):
        search_results = state.get("search_results", [])
        if search_results:
            top_titles = ", ".join(result.get("title", "Untitled") for result in search_results[:3])
            updated_state["synthesis"] = f"Collected web sources for {state['topic']}: {top_titles}"
        else:
            updated_state["synthesis"] = f"Collected web sources for {state['topic']}."
    
    return updated_state


def route_research(state: ResearchGraphState) -> str:
    return "finalize"


def build_research_graph(*, llm, max_queries: int = 4):
    if StateGraph is None:
        raise RuntimeError("langgraph is not installed.")

    graph = StateGraph(ResearchGraphState)
    
    graph.add_node("plan_search", lambda s: plan_search_node(s, llm=llm, max_queries=max_queries))
    graph.add_node("execute_search", execute_search_node)
    graph.add_node("save_note", save_note_node)
    graph.add_node("finalize", finalize_research_node)
    
    graph.set_entry_point("plan_search")
    graph.add_edge("plan_search", "execute_search")
    
    graph.add_conditional_edges(
        "execute_search",
        route_research,
        {
            "finalize": "finalize"
        }
    )
    
    graph.add_edge("finalize", END)
    
    return graph.compile()


def run_research_agent(
    topic: str,
    *,
    llm,
    save_to_notes: bool = False,
    max_queries: int = 4,
    debug_mode: bool = False,
) -> ResearchResult:
    print(f"[Research Agent] Starting session for topic: '{topic}'")
    graph = build_research_graph(llm=llm, max_queries=max_queries)
    
    initial_state = {
        "topic": topic,
        "search_queries": [],
        "search_results": [],
        "synthesis": None,
        "save_to_notes": save_to_notes,
        "note_id": None,
        "error": None,
        "debug_mode": debug_mode,
        "debug_data": {} if debug_mode else None,
    }
    
    final_state = graph.invoke(initial_state)
    print(
        f"[Research Agent] Session complete | "
        f"error={final_state.get('error')} | "
        f"results={len(final_state.get('search_results', []))} | "
        f"note_id={final_state.get('note_id')}"
    )
    
    # Format sources for the result
    sources = [
        {
            "url": r["url"],
            "title": r["title"],
            "snippet": r["snippet"],
            "content": r.get("content", r["snippet"]),
        }
        for r in final_state.get("search_results", [])
    ]
    
    return ResearchResult(
        synthesis=final_state.get("synthesis") or "Research failed.",
        sources=sources,
        note_id=final_state.get("note_id"),
        debug_data=final_state.get("debug_data")
    )