agents/web_search_subagent.py

# =============================================================================
# WEB SEARCH SUBAGENT - Isolated context for web/search operations
#
# TABLE OF CONTENTS
# [1] Subagent State Definition
# [2] Web Search Specialist Node
# [3] Graph Builder
# [4] Tool Wrapper
# =============================================================================

from typing import TypedDict, Annotated

from langchain_core.messages import HumanMessage, SystemMessage, AnyMessage
from langchain_core.tools import tool
from langchain_google_genai import ChatGoogleGenerativeAI
from langgraph.graph import START, StateGraph
from langgraph.graph.message import add_messages
from langgraph.prebuilt import ToolNode, tools_condition

from agents.web_search_tools import get_web_search_tools


# [1] Subagent State Definition
# ---


class WebSearchState(TypedDict):
    """State for web search subagent with isolated context."""

    messages: Annotated[list[AnyMessage], add_messages]


# [2] Web Search Specialist Node
# ---


def _web_search_specialist(state: WebSearchState, llm):
    """Node that routes web/search queries to appropriate tools."""
    sys_msg = SystemMessage(
        content="""
        You are a web search specialist agent. Answer queries using:
        - wikipedia_tool: For general knowledge, people, places, historical facts
        - arxiv_tool: For research papers, scientific articles
        - duckduckgo_search: For current events, news, general web search
        - visit_webpage: When a specific URL is provided or found
        
        Use tools as needed and provide a clear, concise final answer.
        """
    )
    web_tools = get_web_search_tools()
    llm_with_tools = llm.bind_tools(web_tools)
    return {"messages": [llm_with_tools.invoke([sys_msg] + state["messages"])]}


# [3] Graph Builder
# ---


def _build_web_search_subagent(llm):
    """Build a subagent that handles web/search tasks with isolated context."""
    web_tools = get_web_search_tools()

    def web_search_specialist(state: WebSearchState):
        return _web_search_specialist(state, llm)

    builder = StateGraph(WebSearchState)
    builder.add_node("web_search_specialist", web_search_specialist)
    builder.add_node("tools", ToolNode(web_tools))
    builder.add_edge(START, "web_search_specialist")
    builder.add_conditional_edges("web_search_specialist", tools_condition)
    builder.add_edge("tools", "web_search_specialist")

    return builder.compile()


# [4] Tool Wrapper
# ---


@tool
def web_search_agent(query: str) -> str:
    """
    Intelligent web search agent with isolated context.

    Routes and executes web/search tasks. Use this for any web search,
    Wikipedia lookups, arXiv papers, or webpage visits.
    Returns only the final answer, keeping main agent context clean.

    Example queries:
    - "Who is Mercedes Sosa?" (uses Wikipedia)
    - "Latest research on transformers" (uses arXiv)
    - "Current news about AI" (uses DuckDuckGo)
    - "Visit https://example.com and summarize" (uses visit_webpage)
    """
    llm = ChatGoogleGenerativeAI(model="gemini-2.5-flash", temperature=0.1)
    subagent = _build_web_search_subagent(llm)

    result = subagent.invoke({"messages": [HumanMessage(content=query)]})

    return result["messages"][-1].content