Multi-Agent Research Assistant with LangGraph

%%capture
!pip install langgraph langchain langchain-openai arxiv scholarly python-dotenv

# LangGraph: multi-agent orchestration framework
# Provides StateGraph for defining agent workflows with nodes and edges
from langgraph.graph import StateGraph, START, END

# LangChain: core framework for LLM applications
from langchain_core.messages import HumanMessage, AIMessage, SystemMessage # define roles
from langchain_core.prompts import ChatPromptTemplate

# OpenAI integration for LangChain (used with OpenRouter)
from langchain_openai import ChatOpenAI

# Academic paper search
import arxiv  # arXiv API client

# Standard library
import os

# # Load API key from Colab Secrets into environment variable if running in Colab
# from google.colab import userdata

# key = "OPENROUTER_API_KEY"
# value = userdata.get(key)
# assert value is not None, f"{key} not found in Colab Secrets or access is disabled"
# os.environ[key] = value

# print("API key successfully loaded")

from dotenv import load_dotenv

# Load API key from .env if running locally - see local install instructions in the repo README
load_dotenv()

# Sanity check
print("OPENROUTER_API_KEY present:", bool(os.getenv("OPENROUTER_API_KEY")))

class AgentState(TypedDict):
    """
    Shared state that flows through the graph and is updated by each node.
    
    When a node returns {"field": value}, LangGraph merges it into the state:
    - Regular fields: new value REPLACES the old value
    - Annotated fields with reducer: new value is COMBINED with old value using the reducer
    """
    
    # The current search query, replaced when Query Refinement Agent updates it
    query: str
    
    # Raw search results from the current iteration, replaced on each search
    # (we only need the latest batch to filter)
    papers: List[dict]
    
    # Papers that passed relevance filtering, ACCUMULATES across iterations
    # Using operator.add as reducer means new papers are appended, not replaced
    filtered_papers: Annotated[List[dict], operator.add]
    
    # All evaluations from the Filter Agent (for debugging/inspection)
    # Each entry contains: title, score, justification - regardless of whether it passed
    #
    # We use a custom "replace" reducer (lambda old, new: new) instead of operator.add because:
    # 1. LangGraph requires Annotated fields for proper state tracking between nodes
    # 2. But we only need the CURRENT iteration's evaluations for query refinement context
    # 3. Accumulating all evaluations across iterations would waste memory/tokens
    # The lambda simply returns the new value, effectively replacing instead of appending
    all_evaluations: Annotated[List[dict], lambda old, new: new]
    
    # Tracks how many search iterations we've done (to prevent excessively long iterations or infinite loops)
    iteration: int
    
    # Decision from the Supervisor Agent: "end" or "refine"
    decision: str

# Initialize a test state that we can pass to agent functions for testing
# This lets us test each agent independently before wiring them into the graph

test_state: AgentState = {
    "query": "multi-agent reinforcement learning",  # The research topic to search for
    "papers": [],                                   # Will be populated by search_agent
    "filtered_papers": [],                          # Will be populated by filter_agent
    "all_evaluations": [],                          # Will store all LLM evaluations for debugging
    "iteration": 0,                                 # Starting iteration
    "decision": ""                                  # Will be set by supervisor_agent
}

print("Test state initialized:")
print(f"  query: '{test_state['query']}'")
print(f"  papers: {len(test_state['papers'])} items")
print(f"  filtered_papers: {len(test_state['filtered_papers'])} items")
print(f"  all_evaluations: {len(test_state['all_evaluations'])} items")
print(f"  iteration: {test_state['iteration']}")
print(f"  decision: '{test_state['decision']}'")

# Create a single arXiv client to reuse across all searches
# This ensures proper rate limiting (the client tracks request timestamps internally)
arxiv_client = arxiv.Client()

def search_agent(state: dict) -> dict:
    """
    Search for academic papers on arXiv based on the current query.
    
    Args:
        state: Current graph state containing 'query'
        
    Returns:
        Updated state with 'papers' list containing search results
    """
    query = state["query"]
    max_results = 10
    
    # Search arXiv using the shared client (handles rate limiting internally)
    search = arxiv.Search(
        query=query,
        max_results=max_results,
        sort_by=arxiv.SortCriterion.Relevance
    )
    
    papers = []
    for result in arxiv_client.results(search):
        paper = {
            "title": result.title,
            "authors": [author.name for author in result.authors],
            "abstract": result.summary,
            "url": result.entry_id,
            "published": result.published.strftime("%Y-%m-%d"),
            "source": "arxiv"
        }
        papers.append(paper)
    
    print(f"Search Agent: Found {len(papers)} papers for query '{query}'")
    
    return {"papers": papers}

papers = search_agent(test_state)["papers"]  
len(papers)

titles = [paper["title"] for paper in papers]                                              
dates = [paper["published"] for paper in papers]                                           
                                                                                            
for title, date in zip(titles, dates):                                                     
    print(f"{date}: {title}") 

# add papers to test_state for testing
test_state["papers"] = papers
len(test_state["papers"])

import json

# Initialize the LLM for agents that need reasoning capabilities
llm = ChatOpenAI(
    model="openai/gpt-4o-mini",
    temperature=0,  # Deterministic output for consistent evaluations
    openai_api_base="https://openrouter.ai/api/v1",
    openai_api_key=os.getenv("OPENROUTER_API_KEY")
)

def filter_agent(state: dict) -> dict:
    """
    Evaluate the relevance of each paper to the research query using an LLM.
    
    Args:
        state: Current graph state containing 'query' and 'papers'
        
    Returns:
        Updated state with:
        - 'filtered_papers': papers that scored >= 0.7
        - 'all_evaluations': all papers with their scores and justifications (for debugging)
    """
    query = state["query"]
    papers = state["papers"]
    
    # SystemMessage: Defines the AI's role, behavior, and output format
    # These are persistent instructions that apply to ALL evaluations
    system_prompt = SystemMessage(content="""You are an academic paper relevance evaluator.
Your task is to assess how relevant a given paper is to a research query.
Be objective and base your assessment on the paper's title, abstract, and publication date.

When evaluating relevance, consider:
- How directly the paper addresses the research query
- The recency of the paper (more recent papers are preferred when content relevance is similar)

You must respond with ONLY a valid JSON object in this exact format:
{"relevance_score": 0.0, "justification": "brief explanation"}

The relevance_score must be between 0.0 and 1.0 where:
- 0.0-0.3: Not relevant (paper does not address the research query)
- 0.4-0.6: Somewhat relevant (paper touches on related topics)
- 0.7-1.0: Highly relevant (paper directly addresses the research query)""")
    
    filtered = []
    all_evaluations = []  # Track ALL evaluations for debugging
    
    for paper in papers:
        # HumanMessage: Contains ONLY the variable data for this specific evaluation
        # No instructions here - just the inputs that change per paper
        user_prompt = HumanMessage(content=f"""Research Query: {query}

Paper Title: {paper['title']}

Publication Date: {paper['published']}

Abstract: {paper['abstract']}""")
        
        try:
            # Pass both SystemMessage and HumanMessage to the LLM
            # SystemMessage sets the behavior, HumanMessage provides the specific data
            response = llm.invoke([system_prompt, user_prompt])
            result = json.loads(response.content)
            
            score = result.get("relevance_score", 0)
            justification = result.get("justification", "")
            
            # Record this evaluation (regardless of whether it passes)
            evaluation = {
                "title": paper["title"],
                "published": paper["published"],
                "relevance_score": score,
                "justification": justification,
                "passed": score >= 0.7
            }
            all_evaluations.append(evaluation)
            
            # Keep papers that meet the relevance threshold
            if score >= 0.7:
                paper_with_score = paper.copy()
                paper_with_score["relevance_score"] = score
                paper_with_score["justification"] = justification
                filtered.append(paper_with_score)
                
        except (json.JSONDecodeError, Exception) as e:
            # If parsing fails, record the error and skip this paper
            all_evaluations.append({
                "title": paper["title"],
                "published": paper["published"],
                "relevance_score": None,
                "justification": f"Error: {e}",
                "passed": False
            })
            print(f"Filter Agent: Error evaluating '{paper['title'][:50]}...': {e}")
            continue
    
    print(f"Filter Agent: {len(filtered)}/{len(papers)} papers passed relevance threshold (>= 0.7)")
    
    # Return both filtered papers and all evaluations
    return {
        "filtered_papers": filtered,
        "all_evaluations": all_evaluations
    }

filter_results = filter_agent(test_state)

len(filter_results["filtered_papers"])

filter_results["filtered_papers"][0]["title"]

len(filter_results["all_evaluations"])

for paper in filter_results["all_evaluations"]:                                                     
    print(paper["title"], paper["relevance_score"], paper["justification"]) 

# Configuration constants for the Supervisor's decision logic
MIN_RELEVANT_PAPERS = 3  # Minimum papers needed to consider search successful
MAX_ITERATIONS = 3       # Maximum search attempts before giving up

def supervisor_agent(state: dict) -> dict:
    """
    Decide whether to end the workflow or continue with query refinement.
    
    This agent uses simple conditional logic (no LLM) to make routing decisions
    based on the current state of the search.
    
    Args:
        state: Current graph state containing 'filtered_papers' and 'iteration'
        
    Returns:
        Updated state with 'decision' field: "end" or "refine"
    """
    filtered_papers = state["filtered_papers"]
    iteration = state["iteration"]
    
    num_relevant = len(filtered_papers)
    stop_reason = ""
    
    # Decision logic
    if num_relevant >= MIN_RELEVANT_PAPERS:
        # Success: we have enough relevant papers
        decision = "end"
        stop_reason = f"Success: Found {num_relevant} relevant papers (>= {MIN_RELEVANT_PAPERS} required)"
    elif iteration >= MAX_ITERATIONS:
        # Max attempts reached: return what we have
        decision = "end"
        stop_reason = f"Max iterations ({MAX_ITERATIONS}) reached with only {num_relevant} relevant papers"
    else:
        # Need more results: refine query and try again
        decision = "refine"
    
    print(f"Supervisor Agent: {decision.upper()}" + (f" - {stop_reason}" if stop_reason else f" - Refining query (iteration {iteration + 1})"))
    
    return {"decision": decision}

def query_refinement_agent(state: dict) -> dict:
    """
    Refine the search query based on feedback from previous evaluations.
    
    Analyzes why papers were rejected and generates an improved query
    that is more likely to find relevant results.
    
    Args:
        state: Current graph state containing 'query', 'all_evaluations', and 'iteration'
        
    Returns:
        Updated state with:
        - 'query': refined search query
        - 'iteration': incremented iteration count
    """
    current_query = state["query"]
    all_evaluations = state["all_evaluations"]
    iteration = state["iteration"]
    
    # Format evaluation feedback for the LLM
    feedback_lines = []
    for eval in all_evaluations:
        status = "PASSED" if eval["passed"] else "REJECTED"
        feedback_lines.append(
            f"- [{status}] \"{eval['title']}\" (score: {eval['relevance_score']}) - {eval['justification']}"
        )
    feedback_summary = "\n".join(feedback_lines)
    
    # SystemMessage: Defines the AI's role and output format
    system_prompt = SystemMessage(content="""You are a search query optimization expert.
    Your task is to refine academic search queries based on feedback from previous search results.

    Analyze why papers were rejected and craft a more targeted query that will find more relevant results.
    Consider:
    - Adding specific technical terms that were missing
    - Removing overly broad or ambiguous terms
    - Including synonyms or related concepts
    - Narrowing the scope if results were too general

    You must respond with ONLY the refined query string, nothing else.
    Do not include quotes around the query. Just output the query text directly.""")
    
    # HumanMessage: Contains the specific data for this refinement
    user_prompt = HumanMessage(content=f"""Current Query: {current_query}

    Iteration: {iteration + 1}

    Previous Search Results Feedback:
    {feedback_summary}

    Based on this feedback, generate an improved search query that will find more relevant papers.""")
    
    try:
        response = llm.invoke([system_prompt, user_prompt])
        refined_query = response.content.strip()
        
        # Clean up the query (remove quotes if LLM added them)
        refined_query = refined_query.strip('"\'')
        
        print(f"Query Refinement Agent: '{current_query}' → '{refined_query}'")
        
    except Exception as e:
        # If refinement fails, add "survey" to find overview papers
        refined_query = f"{current_query} survey"
        print(f"Query Refinement Agent: Error ({e}), using fallback: '{refined_query}'")
    
    return {
        "query": refined_query,
        "iteration": iteration + 1
    }

# Test the Query Refinement Agent
# The test_state already has all_evaluations from the filter agent test

print(f"Current state before refinement:")
print(f"  query: '{test_state['query']}'")
print(f"  iteration: {test_state['iteration']}")
print(f"  all_evaluations: {len(test_state['all_evaluations'])} items")
print()

# Test query refinement agent
refinement_result = query_refinement_agent(test_state)

print()
print(f"Refinement result:")
print(f"  new query: '{refinement_result['query']}'")
print(f"  new iteration: {refinement_result['iteration']}")

# Create the StateGraph with our state schema
graph = StateGraph(AgentState)

# Add nodes — each agent function becomes a node
graph.add_node("search_agent", search_agent)
graph.add_node("filter_agent", filter_agent)
graph.add_node("supervisor_agent", supervisor_agent)
graph.add_node("query_refinement_agent", query_refinement_agent)

# Add edges — define the linear flow
graph.add_edge(START, "search_agent")           # Entry point
graph.add_edge("search_agent", "filter_agent")  # Search → Filter
graph.add_edge("filter_agent", "supervisor_agent")  # Filter → Supervisor

# Add conditional edge — Supervisor decides next step based on 'decision' field
def route_supervisor(state: dict) -> str:
    """Route based on supervisor's decision."""
    if state["decision"] == "end":
        return END
    else:
        return "query_refinement_agent"

graph.add_conditional_edges(
    "supervisor_agent",
    route_supervisor,
    {END: END, "query_refinement_agent": "query_refinement_agent"}
)

# Query refinement loops back to search
graph.add_edge("query_refinement_agent", "search_agent")

# Compile the graph into a runnable workflow
workflow = graph.compile()

print("Graph compiled successfully!")

# Define the initial state with our research query
initial_state = {
    "query": "multi-llm-agent reinforcement learning",
    "papers": [],
    "filtered_papers": [],
    "all_evaluations": [],
    "iteration": 0,
    "decision": ""
}

print(f"Starting workflow with query: '{initial_state['query']}'")
print("=" * 60)

# Run the workflow
final_state = workflow.invoke(initial_state)

print("=" * 60)
print(f"\nWorkflow complete!")

# Determine stop reason from final state values
num_found = len(final_state['filtered_papers'])
if num_found >= MIN_RELEVANT_PAPERS:
    stop_reason = f"Success: Found {num_found} relevant papers"
else:
    stop_reason = f"Max iterations reached with only {num_found} relevant papers"

print(f"Result: {stop_reason}")
print(f"Final query: '{final_state['query']}'")
print(f"Total iterations: {final_state['iteration']}")
print(f"Relevant papers found: {num_found}")

# Display the relevant papers found
print("Relevant Papers Found:")
print("-" * 60)

for i, paper in enumerate(final_state["filtered_papers"], 1):
    print(f"\n{i}. {paper['title']}")
    print(f"   Published: {paper['published']}")
    print(f"   Relevance: {paper['relevance_score']}")
    print(f"   URL: {paper['url']}")