Project Guide
Build an AI Research Agent in Python with LangGraph
Source: mortalapps.comThis guide details how to build a robust AI research agent in Python using the LangGraph framework. This agent autonomously conducts web-based research, synthesizes information from multiple sources, identifies and resolves contradictions, and generates a structured research report with citations. It's designed for developers, data scientists, and AI engineers looking to implement intelligent automation for information gathering and synthesis.
The business problem this project addresses is the manual, time-consuming, and often inconsistent process of conducting thorough research. By automating this, organizations can accelerate decision-making, improve report accuracy, and free up human resources for higher-value tasks. This agent can be adapted for market analysis, competitive intelligence, scientific literature review, or even internal knowledge base generation.
Upon completing this project, you will have a fully functional, deployable AI research agent exposed as a FastAPI service. This agent leverages the ReAct pattern for dynamic reasoning and tool use, ensuring it can adapt to various research queries. The architecture prioritizes modularity and observability, making it easy to extend and monitor in production environments.
We chose LangGraph for its explicit state management and conditional routing capabilities, which are crucial for complex, multi-step agentic workflows. Unlike simpler sequential chains, LangGraph allows for sophisticated decision-making and loop structures, essential for iterative research, contradiction resolution, and multi-source synthesis.
What You Will Build
You will build an autonomous AI research agent capable of taking a research query, breaking it down, searching the web, analyzing search results, synthesizing findings, resolving conflicting information, and generating a comprehensive, citable report. The agent interacts with a web search tool to gather information and a custom 'report writer' tool to format and finalize its output. It communicates through a simple RESTful API endpoint exposed via FastAPI, allowing external systems or users to submit research requests and retrieve results.
The agent runs as a Python application within a Docker container, making it portable and easy to deploy. Its output is a Markdown-formatted research report, complete with headings, bullet points, and source citations, delivered as a string from the API. The core architecture is a LangGraph state machine, orchestrating the LLM's reasoning, tool calls, and state updates, moving through phases like 'plan', 'search', 'analyze', 'synthesize', and 'report'. This graph-based approach provides clear visibility into the agent's decision-making process and state transitions.
Technology Stack
| Component | Technology | Why This Choice |
|---|---|---|
| LLM Provider | OpenAI GPT-4o |
Provides state-of-the-art reasoning, synthesis, and instruction following capabilities crucial for complex research tasks. Its function calling feature simplifies tool integration. |
| Orchestration Framework | LangGraph |
Enables explicit state management and conditional routing, essential for building robust, multi-step agentic workflows like research, synthesis, and contradiction resolution. Its graph-based approach offers clarity and debuggability. |
| Tool Execution | LangChain Tools, Tavily Search API |
LangChain provides a standardized interface for defining and integrating tools. Tavily offers efficient and relevant web search results, critical for comprehensive research. |
| API & Web Service | FastAPI |
A modern, fast (high-performance) web framework for building APIs with Python 3.7+ based on standard Python type hints. It's ideal for exposing the agent as a RESTful service. |
| Containerization | Docker |
Ensures consistent development, testing, and deployment environments by packaging the application and its dependencies into a portable image. |
| Observability | Python Logging |
Provides basic yet effective logging for monitoring agent execution, debugging issues, and tracking state transitions during development and production. |
| Secrets Management | Environment Variables |
A standard and secure way to inject sensitive information (API keys, credentials) into the application at runtime, preventing hardcoding. |
The choice of OpenAI GPT-4o as the LLM provider is driven by its exceptional performance in complex reasoning, summarization, and instruction following, which are paramount for an AI research agent. Its advanced understanding allows for nuanced interpretation of research queries and synthesis of diverse information. While alternatives like Anthropic's Claude or Google's Gemini offer competitive capabilities, GPT-4o's strong function calling and general availability make it a practical choice for this production-grade project.
LangGraph is selected for orchestration due to its explicit state management and powerful conditional routing. Unlike simpler sequential chains or even basic agent loops, research involves iterative steps, decision points (e.g., 'do I need more information?', 'is there a contradiction?'), and complex state transitions. LangGraph's graph-based approach provides the necessary control and transparency to model these intricate workflows, making the agent's behavior predictable and debuggable. Other frameworks like vanilla LangChain agents or CrewAI are powerful, but LangGraph's state machine paradigm is uniquely suited for the robust control flow needed here.
For tool execution, LangChain Tools combined with the Tavily Search API offer a robust solution. LangChain's tool abstraction simplifies the integration of external capabilities, allowing the agent to interact with the real world. Tavily stands out for its focused and relevant search results, which is crucial for minimizing hallucination and improving research quality compared to generic search APIs. While custom scraping tools could be built, Tavily provides a managed, high-quality solution out of the box.
FastAPI serves as the API and web service layer due to its high performance, ease of use, and strong typing support. It allows for rapid development of a robust RESTful interface for our agent, facilitating easy integration with other services or front-end applications. Its asynchronous capabilities are also beneficial for handling potentially long-running agent requests. Flask or Django are viable alternatives, but FastAPI's modern async-first design and Pydantic integration align well with modern AI application development.
Docker is used for containerization to ensure environmental consistency from development to production. This eliminates 'it works on my machine' issues and simplifies deployment across various environments. For observability, standard Python Logging is implemented initially to provide essential insights into the agent's runtime behavior. While more advanced tracing with OpenTelemetry could be integrated, basic logging offers a good starting point for monitoring critical events and debugging. Finally, Environment Variables are the standard choice for secrets management, promoting security and configurability without embedding sensitive data directly in the codebase.
Project Structure
ai_research_agent/
├── .env
├── Dockerfile
├── README.md
├── requirements.txt
├── app/
│ ├── __init__.py
│ ├── main.py
│ ├── agent.py
│ ├── tools.py
│ └── models.py
└── tests/
├── __init__.py
└── test_agent.pyThe ai_research_agent/ directory is the root of our project. The .env file stores environment variables, including API keys, ensuring sensitive information is kept out of source control. The Dockerfile defines how our application is containerized for consistent deployment. README.md provides essential project documentation and setup instructions. requirements.txt lists all Python dependencies needed for the project.
The app/ directory contains the core application logic. __init__.py marks it as a Python package. main.py initializes the FastAPI application and defines the API endpoints for interacting with our research agent. agent.py houses the LangGraph state machine, defining the agent's nodes, edges, and overall workflow. tools.py defines the custom tools the agent can use, such as web search and report generation. models.py defines Pydantic models for data validation, especially for API requests and agent state.
The tests/ directory contains unit and integration tests for the agent. test_agent.py will include tests to verify the agent's functionality and tool interactions.
Implementation
Phase 1: Environment Setup and FastAPI Scaffold
This phase sets up the project structure, defines necessary environment variables, installs dependencies, and creates a basic FastAPI application. It includes defining the data models for agent input and output, and a placeholder for the agent's API endpoint.
import os
from typing import Literal, List, Dict, Any
import logging
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel, Field
# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)
# Pydantic Models for API input/output and agent state
class ResearchRequest(BaseModel):
query: str = Field(..., min_length=1, description="The research query for the AI agent.")
class ResearchReport(BaseModel):
report_content: str = Field(..., description="The generated research report with citations.")
sources: List[str] = Field(default_factory=list, description="List of URLs or source identifiers.")
status: Literal["success", "failure"] = Field(..., description="Status of the research request.")
error_message: str | None = Field(default=None, description="Error message if status is failure.")
class AgentState(BaseModel):
query: str
intermediate_steps: List[Dict[str, Any]] = Field(default_factory=list)
research_findings: Dict[str, Any] = Field(default_factory=dict)
report: str | None = None
tool_output: str | None = None
error: str | None = None
max_retries: int = 3
current_retry: int = 0
app = FastAPI(
title="AI Research Agent API",
description="API for an autonomous AI agent capable of conducting web-based research and generating reports.",
version="1.0.0"
)
@app.get("/health")
async def health_check():
return {"status": "healthy"}
@app.post("/research", response_model=ResearchReport)
async def perform_research(request: ResearchRequest):
logger.info(f"Received research request: {request.query}")
try:
# Placeholder for agent execution logic
# In later phases, this will call the LangGraph agent
report_content = f"Research on '{request.query}' is being processed. This is a placeholder report."
sources = ["http://placeholder.com/source1", "http://placeholder.com/source2"]
return ResearchReport(report_content=report_content, sources=sources, status="success")
except Exception as e:
logger.error(f"Error processing research request: {e}", exc_info=True)
raise HTTPException(status_code=500, detail=f"Internal server error: {e}")
# app/main.py
# --- Project Root: ai_research_agent/
# .env (create this file manually)
# OPENAI_API_KEY="sk-..."
# TAVILY_API_KEY="your-tavily-api-key"
# requirements.txt (create this file manually)
# fastapi==0.111.0
# uvicorn==0.30.1
# pydantic==2.7.4
# langgraph==0.0.60
# langchain-openai==0.1.8
# langchain-community==0.0.38
# tavily-python==0.3.3
This phase establishes the foundational elements of our AI research agent. We begin by configuring Python's built-in
logging module to ensure that all critical operations, errors, and agent decisions are recorded, which is vital for debugging and monitoring in production. The FastAPI instance is initialized, providing a web server to expose our agent's capabilities. We define two Pydantic models: ResearchRequest for validating incoming API queries and ResearchReport for structuring the agent's output, including the report content, sources, and status. This strict typing enhances API reliability and developer experience.
The AgentState Pydantic model is crucial for LangGraph. It defines the schema for the state that will be passed between nodes in our agent's graph. It includes the query, a list for intermediate_steps (to track agent thoughts and actions), research_findings (to store structured data from searches), the report itself, tool_output for the most recent tool call, and error for graceful failure handling. max_retries and current_retry are added to enable robust error recovery, allowing the agent to attempt failed steps multiple times. A simple /health endpoint is added for service readiness checks, and a placeholder /research endpoint is created to demonstrate the API's structure. This endpoint currently returns a static response but will be integrated with the LangGraph agent in subsequent phases. Finally, instructions for .env and requirements.txt are provided to ensure a reproducible development environment.
Expected outcome
A runnable FastAPI application that exposes a `/health` endpoint returning 'healthy' and a `/research` endpoint that accepts a query and returns a placeholder research report. Environment variables are set up to be loaded correctly.
How to test
Save the code as app/main.py. Create .env and requirements.txt files as specified. Install dependencies with pip install -r requirements.txt. Run the app using uvicorn app.main:app --reload. Test the /health endpoint at http://127.0.0.1:8000/health and the /research endpoint using a tool like Postman or curl with a JSON payload like {"query": "latest advancements in quantum computing"}.
Phase 2: Define Tools and Initial Agent Graph State
This phase focuses on defining the external tools our agent will use, specifically a web search tool and a report generation tool. We will also initialize the LangGraph StateGraph and define the initial set of nodes and edges, laying the groundwork for the agent's operational flow.
import os
import logging
from typing import List, Dict, Any, Callable
from functools import wraps
import time
from langchain_community.tools.tavily_research import TavilySearchResults
from langchain_core.tools import tool
from langgraph.graph import StateGraph, START, END
from pydantic import BaseModel, Field
# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)
# Helper for retries with exponential backoff
def retry_with_exponential_backoff(max_retries: int = 3, initial_delay: float = 1.0):
def decorator(func: Callable):
@wraps(func)
def wrapper(*args, **kwargs):
delay = initial_delay
for i in range(max_retries):
try:
return func(*args, **kwargs)
except Exception as e:
logger.warning(f"Attempt {i+1}/{max_retries} failed for {func.__name__}: {e}")
if i < max_retries - 1:
time.sleep(delay)
delay *= 2 # Exponential backoff
else:
raise
return wrapper
return decorator
# Pydantic Models (re-define or import from app.models if split)
class ResearchRequest(BaseModel):
query: str = Field(..., min_length=1, description="The research query for the AI agent.")
class ResearchReport(BaseModel):
report_content: str = Field(..., description="The generated research report with citations.")
sources: List[str] = Field(default_factory=list, description="List of URLs or source identifiers.")
status: Literal["success", "failure"] = Field(..., description="Status of the research request.")
error_message: str | None = Field(default=None, description="Error message if status is failure.")
class AgentState(BaseModel):
query: str
intermediate_steps: List[Dict[str, Any]] = Field(default_factory=list)
research_findings: Dict[str, Any] = Field(default_factory=dict)
report: str | None = None
tool_output: str | None = None
error: str | None = None
max_retries: int = 3
current_retry: int = 0
# --- Tools Definition (app/tools.py)
class ResearchTools:
def __init__(self):
tavily_api_key = os.environ.get("TAVILY_API_KEY")
if not tavily_api_key:
logger.error("TAVILY_API_KEY not found in environment variables.")
raise ValueError("TAVILY_API_KEY is required for web search.")
self.tavily_search = TavilySearchResults(max_results=5, tavily_api_key=tavily_api_key)
@tool
@retry_with_exponential_backoff()
def web_search(self, query: str) -> str:
"""Searches the web for the given query and returns relevant snippets. """
logger.info(f"Performing web search for: {query}")
try:
results = self.tavily_search.run(query)
return str(results)
except Exception as e:
logger.error(f"Web search failed for query '{query}': {e}", exc_info=True)
raise
@tool
def format_report(self, findings: Dict[str, Any]) -> str:
"""Formats the gathered research findings into a structured markdown report. """
logger.info("Formatting research report.")
report_parts = ["# Research Report
"]
query = findings.get("query", "N/A")
report_parts.append(f"## Topic: {query}
")
if "summary" in findings:
report_parts.append("### Summary
")
report_parts.append(findings["summary"] + "
")
if "details" in findings:
report_parts.append("### Detailed Findings
")
for topic, content in findings["details"].items():
report_parts.append(f"#### {topic}
")
report_parts.append(content + "
")
if "sources" in findings and isinstance(findings["sources"], list):
report_parts.append("### Sources
")
for i, source in enumerate(findings["sources"]):
report_parts.append(f"{i+1}. {source}
")
formatted_report = "
".join(report_parts)
logger.info("Report formatting complete.")
return formatted_report
# --- Agent Graph Definition (app/agent.py - partial)
# This would typically be in app/agent.py, but for phased implementation, we put it here.
# Initialize Tools
research_tools = ResearchTools()
tools = [research_tools.web_search, research_tools.format_report]
# Define the LangGraph State Graph
workflow = StateGraph(AgentState)
# Placeholder nodes (will be filled in next phase)
def initial_plan_node(state: AgentState) -> AgentState:
logger.info(f"Initial plan for query: {state.query}")
return {"intermediate_steps": state.intermediate_steps + [{"action": "initial_plan", "observation": "Plan generated."}]}
def execute_tool_node(state: AgentState) -> AgentState:
logger.info(f"Executing tool for query: {state.query}")
return {"intermediate_steps": state.intermediate_steps + [{"action": "execute_tool", "observation": "Tool executed."}]}
workflow.add_node("initial_plan", initial_plan_node)
workflow.add_node("execute_tool", execute_tool_node)
workflow.add_edge(START, "initial_plan")
workflow.add_edge("initial_plan", "execute_tool")
workflow.add_edge("execute_tool", END)
# This build step would typically be called at the end of app/agent.py
# app_agent = workflow.compile() # Don't compile yet, will be done in main.py later
# app/tools.py content
# class ResearchTools:
# def __init__(self):
# tavily_api_key = os.environ.get("TAVILY_API_KEY")
# if not tavily_api_key:
# logger.error("TAVILY_API_KEY not found in environment variables.")
# raise ValueError("TAVILY_API_KEY is required for web search.")
# self.tavily_search = TavilySearchResults(max_results=5, tavily_api_key=tavily_api_key)
#
# @tool
# @retry_with_exponential_backoff()
# def web_search(self, query: str) -> str:
# # ... (implementation as above)
#
# @tool
# def format_report(self, findings: Dict[str, Any]) -> str:
# # ... (implementation as above)
# app/agent.py (partial)
# from langgraph.graph import StateGraph, START, END
# from app.models import AgentState
# from app.tools import research_tools
#
# tools = [research_tools.web_search, research_tools.format_report]
# workflow = StateGraph(AgentState)
#
# # ... (nodes and edges as above)
In this phase, we establish the tools our agent will wield and begin defining its operational graph using LangGraph. First, we create a
ResearchTools class to encapsulate our web search and report formatting functionalities. The web_search tool leverages TavilySearchResults to perform internet queries, requiring a TAVILY_API_KEY from environment variables for security and configuration. Crucially, a retry_with_exponential_backoff decorator is applied to web_search to enhance reliability against transient network issues or API rate limits, a common challenge in production systems. This decorator automatically retries failed tool calls with increasing delays, preventing immediate failures and improving resilience. Error logging is also integrated to capture any search failures.
The format_report tool is a custom function that takes structured findings and transforms them into a readable Markdown report. This modular approach allows for easy modification of the report's structure without altering the core agent logic. Both tools are decorated with @tool from langchain_core.tools, making them discoverable and usable by LangChain-compatible LLMs.
Next, we initialize a StateGraph from LangGraph, explicitly typing its state with our AgentState Pydantic model. This ensures type safety and clarity in how information flows through the agent. We add two placeholder nodes, initial_plan_node and execute_tool_node, along with basic edges connecting START to initial_plan, then to execute_tool, and finally to END. These placeholders serve as the skeleton for our agent's workflow, demonstrating how nodes and edges are defined. While these nodes don't yet contain complex LLM logic, they illustrate the graph structure and how state updates are returned. For clarity and adherence to project structure, the code is presented as it would be split into app/tools.py and app/agent.py in a real project, though combined here for sequential explanation.
Expected outcome
A Python application with properly initialized web search and report formatting tools. A basic LangGraph `StateGraph` is defined with initial placeholder nodes and edges, ready for the LLM integration.
How to test
Save the tool definitions in app/tools.py and the graph definition in app/agent.py. Ensure your .env has TAVILY_API_KEY set. You can add a temporary main block to app/agent.py to instantiate ResearchTools and print tools to verify they are loaded correctly. You can also instantiate workflow and print its structure (though it won't run yet).
Phase 3: Integrate LLM and Implement ReAct Logic
This phase integrates the Large Language Model (LLM) into the agent's workflow and implements the core ReAct (Reasoning and Acting) pattern. We will create an LLM-powered node that can decide whether to use a tool or generate a final answer, based on the current state and prompt. The graph will be updated to reflect this dynamic decision-making.
import os
import logging
from typing import List, Dict, Any, Callable, Tuple, Literal
from functools import wraps
import time
from langchain_community.tools.tavily_research import TavilySearchResults
from langchain_core.tools import tool
from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.messages import BaseMessage, HumanMessage, AIMessage, FunctionMessage
from langgraph.graph import StateGraph, START, END
from pydantic import BaseModel, Field
# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)
# Helper for retries with exponential backoff
def retry_with_exponential_backoff(max_retries: int = 3, initial_delay: float = 1.0):
def decorator(func: Callable):
@wraps(func)
def wrapper(*args, **kwargs):
delay = initial_delay
for i in range(max_retries):
try:
return func(*args, **kwargs)
except Exception as e:
logger.warning(f"Attempt {i+1}/{max_retries} failed for {func.__name__}: {e}")
if i < max_retries - 1:
time.sleep(delay)
delay *= 2 # Exponential backoff
else:
raise
return wrapper
return decorator
# Pydantic Models (re-define or import from app.models if split)
class ResearchRequest(BaseModel):
query: str = Field(..., min_length=1, description="The research query for the AI agent.")
class ResearchReport(BaseModel):
report_content: str = Field(..., description="The generated research report with citations.")
sources: List[str] = Field(default_factory=list, description="List of URLs or source identifiers.")
status: Literal["success", "failure"] = Field(..., description="Status of the research request.")
error_message: str | None = Field(default=None, description="Error message if status is failure.")
class AgentState(BaseModel):
query: str
intermediate_steps: List[Dict[str, Any]] = Field(default_factory=list)
research_findings: Dict[str, Any] = Field(default_factory=dict)
report: str | None = None
tool_output: str | None = None
error: str | None = None
max_retries: int = 3
current_retry: int = 0
# --- Tools Definition (app/tools.py)
class ResearchTools:
def __init__(self):
tavily_api_key = os.environ.get("TAVILY_API_KEY")
if not tavily_api_key:
logger.error("TAVILY_API_KEY not found in environment variables.")
raise ValueError("TAVILY_API_KEY is required for web search.")
self.tavily_search = TavilySearchResults(max_results=5, tavily_api_key=tavily_api_key)
@tool
@retry_with_exponential_backoff()
def web_search(self, query: str) -> str:
"""Searches the web for the given query and returns relevant snippets. """
logger.info(f"Performing web search for: {query}")
try:
results = self.tavily_search.run(query)
return str(results)
except Exception as e:
logger.error(f"Web search failed for query '{query}': {e}", exc_info=True)
raise
@tool
def format_report(self, findings: Dict[str, Any]) -> str:
"""Formats the gathered research findings into a structured markdown report. """
logger.info("Formatting research report.")
report_parts = ["# Research Report
"]
query = findings.get("query", "N/A")
report_parts.append(f"## Topic: {query}
")
if "summary" in findings:
report_parts.append("### Summary
")
report_parts.append(findings["summary"] + "
")
if "details" in findings:
report_parts.append("### Detailed Findings
")
for topic, content in findings["details"].items():
report_parts.append(f"#### {topic}
")
report_parts.append(content + "
")
if "sources" in findings and isinstance(findings["sources"], list):
report_parts.append("### Sources
")
for i, source in enumerate(findings["sources"]):
report_parts.append(f"{i+1}. {source}
")
formatted_report = "
".join(report_parts)
logger.info("Report formatting complete.")
return formatted_report
# --- Agent Graph Definition (app/agent.py)
# Initialize LLM
openai_api_key = os.environ.get("OPENAI_API_KEY")
if not openai_api_key:
logger.error("OPENAI_API_KEY not found in environment variables.")
raise ValueError("OPENAI_API_KEY is required for LLM.")
llm = ChatOpenAI(model="gpt-4o", temperature=0, api_key=openai_api_key)
# Initialize Tools
research_tools_instance = ResearchTools()
tools = [research_tools_instance.web_search, research_tools_instance.format_report]
# Bind tools to LLM for function calling
llm_with_tools = llm.bind_tools(tools)
# Define the LangGraph State Graph
workflow = StateGraph(AgentState)
# --- Nodes for the ReAct Agent ---
def agent_node(state: AgentState) -> AgentState:
logger.info(f"Entering agent_node for query: {state.query}")
messages: List[BaseMessage] = [
HumanMessage(content=state.query)
]
# Add previous intermediate steps (thoughts, actions, observations)
for step in state.intermediate_steps:
if "action" in step and "action_input" in step:
messages.append(AIMessage(content="", tool_calls=[{
"name": step["action"],
"args": {"query": step["action_input"]}
}]))
if "observation" in step:
messages.append(FunctionMessage(name=step["action"], content=step["observation"]))
# Add current tool output if available
if state.tool_output:
# Assuming the last action was a tool call that produced this output
if state.intermediate_steps and "action" in state.intermediate_steps[-1]:
messages.append(FunctionMessage(name=state.intermediate_steps[-1]["action"], content=state.tool_output))
else:
logger.warning("Tool output available but no prior action recorded in intermediate steps.")
system_prompt = (
"You are an expert AI research assistant. Your goal is to conduct thorough research "
"on the user's query, synthesize information, resolve contradictions, and produce a "
"comprehensive report with citations. You have access to the following tools:
"
+ "
".join([f"- {t.name}: {t.description}" for t in tools]) + "
"
"When you need to search the web, use the 'web_search' tool. "
"Once you have sufficient information to answer the query and resolve any contradictions, "
"synthesize the findings into a structured dictionary for the 'format_report' tool. "
"The dictionary should include 'query', 'summary', 'details' (a dict of sub-topics to content), and 'sources' (list of URLs). "
"You must cite your sources. If you believe you have gathered enough information and formatted the report, call 'format_report'. "
"If you encounter an error, try to rephrase your search or analyze the problem. "
"Keep track of your thoughts and observations in intermediate steps."
)
# Prepend system prompt to messages
messages.insert(0, HumanMessage(content=system_prompt))
try:
response = llm_with_tools.invoke(messages)
logger.info(f"LLM response: {response.content}, Tool calls: {response.tool_calls}")
new_state = {"intermediate_steps": state.intermediate_steps + [{"thought": response.content}]}
if response.tool_calls:
tool_call = response.tool_calls[0] # Assuming one tool call for simplicity in this phase
tool_name = tool_call["name"]
tool_args = tool_call["args"]
new_state["intermediate_steps"][-1]["action"] = tool_name
new_state["intermediate_steps"][-1]["action_input"] = tool_args
# We don't execute the tool here, just record the decision
elif "format_report" in response.content.lower(): # Heuristic for final report call if not tool_call
# This is a fallback/simplification. In a real system, LLM should use tool_calls for format_report.
new_state["intermediate_steps"][-1]["action"] = "format_report"
# Assume the LLM provides the findings as a JSON string in its content
# This requires careful prompt engineering to ensure correct format
# For this phase, we'll just mark it as ready for report.
logger.warning("LLM indicated report formatting without explicit tool_call. Expecting findings in content.")
else:
# If no tool call and not a report format, it's likely a final answer or thinking step.
# For now, we'll assume it's a final answer if no tool call.
new_state["report"] = response.content # Store as temporary report
new_state["intermediate_steps"][-1]["action"] = "final_answer"
return new_state
except Exception as e:
logger.error(f"Error in agent_node: {e}", exc_info=True)
return {"error": str(e), "current_retry": state.current_retry + 1}
def tool_node(state: AgentState) -> AgentState:
logger.info(f"Entering tool_node. Current intermediate steps: {state.intermediate_steps}")
last_step = state.intermediate_steps[-1]
tool_name = last_step.get("action")
tool_args = last_step.get("action_input")
if not tool_name or not tool_args:
logger.error(f"Tool name or arguments missing in state: {last_step}")
return {"error": "Missing tool name or arguments", "current_retry": state.current_retry + 1}
tool_output = None
try:
if tool_name == "web_search":
query = tool_args.get("query")
if not query:
raise ValueError("Web search tool requires a 'query' argument.")
tool_output = research_tools_instance.web_search.invoke(query)
# Attempt to parse search results into structured findings
# This is a simplification; a dedicated LLM call might be needed for robust parsing
if isinstance(tool_output, str):
search_results = eval(tool_output) # WARNING: eval is dangerous, use safely or parse JSON/XML
if isinstance(search_results, list):
# Extract content and sources from Tavily results
content_snippets = "
".join([r['content'] for r in search_results if 'content' in r])
sources_list = [r['url'] for r in search_results if 'url' in r]
new_findings = {"query": state.query, "details": {"search_results": content_snippets}, "sources": sources_list}
return {"tool_output": tool_output, "research_findings": {**state.research_findings, **new_findings}}
else:
logger.warning(f"Unexpected search result format: {type(search_results)}")
tool_output = f"Web search completed. Result: {tool_output}"
elif tool_name == "format_report":
findings_for_report = state.research_findings
if not findings_for_report:
raise ValueError("No research findings available to format report.")
report_content = research_tools_instance.format_report.invoke(findings_for_report)
tool_output = report_content
return {"report": report_content, "tool_output": tool_output}
else:
raise ValueError(f"Unknown tool: {tool_name}")
except Exception as e:
logger.error(f"Tool '{tool_name}' execution failed: {e}", exc_info=True)
return {"error": str(e), "current_retry": state.current_retry + 1}
return {"tool_output": tool_output, "intermediate_steps": state.intermediate_steps + [{"observation": tool_output}]}
def decide_next_step(state: AgentState) -> str:
logger.info(f"Deciding next step. Current report status: {state.report is not None}, Error: {state.error}")
if state.error and state.current_retry < state.max_retries:
logger.warning(f"Retrying due to error: {state.error}")
# For simplicity, retry the agent node if an error occurred in either agent or tool node.
# A more sophisticated approach would retry the specific failed node.
return "agent_node"
elif state.error and state.current_retry >= state.max_retries:
logger.error(f"Max retries reached. Failing research request. Last error: {state.error}")
return "END"
elif state.report:
logger.info("Report successfully generated. Ending workflow.")
return "END"
else:
# Check if the last action was a tool call that needs execution
if state.intermediate_steps and "action" in state.intermediate_steps[-1] and \
state.intermediate_steps[-1]["action"] not in ["final_answer", "format_report"]:
logger.info(f"LLM decided to use tool: {state.intermediate_steps[-1]['action']}. Executing tool.")
return "tool_node"
else:
logger.info("Continuing research, returning to agent_node for further reasoning.")
return "agent_node"
# Add nodes to the workflow
workflow.add_node("agent_node", agent_node)
workflow.add_node("tool_node", tool_node)
# Set entry point
workflow.set_entry_point("agent_node")
# Add conditional edges
workflow.add_conditional_edges(
"agent_node",
decide_next_step,
{
"tool_node": "tool_node",
"agent_node": "agent_node", # Loop back to agent if more reasoning needed
"END": END
}
)
workflow.add_conditional_edges(
"tool_node",
decide_next_step,
{
"agent_node": "agent_node", # After tool execution, go back to agent to reason on results
"END": END
}
)
# Compile the graph
app_agent = workflow.compile()
This phase is the core of our AI research agent, introducing the LLM and the ReAct pattern. We initialize
ChatOpenAI with gpt-4o and bind our previously defined tools to it. This llm_with_tools object is now capable of performing function calls based on its internal reasoning, which is the foundation of ReAct.
The agent_node function is where the LLM's reasoning takes place. It constructs a messages list, including the system prompt (which instructs the LLM on its role, available tools, and desired output format) and the agent's historical intermediate_steps. This history is crucial for the ReAct loop, allowing the LLM to reflect on past thoughts, actions, and observations. The LLM then invokes these messages and its response is processed. If the LLM decides to call a tool (indicated by response.tool_calls), the tool's name and arguments are extracted and stored in the intermediate_steps of the state. If the LLM generates a final answer or indicates a report without a tool call (a less preferred but handled fallback), it's also recorded.
The tool_node is responsible for executing the tool that the agent_node decided to use. It retrieves the tool name and arguments from the AgentState's intermediate_steps, then invokes the corresponding tool (e.g., web_search or format_report). The output of the tool is then stored in state.tool_output and potentially processed further (e.g., parsing web_search results into research_findings).
The decide_next_step function is the conditional router in our LangGraph. After agent_node or tool_node execution, this function determines the next node. It checks for errors and retry limits, allowing the agent to attempt failed steps. If a report is ready, the workflow ENDs. Otherwise, if the LLM decided to use a tool, it routes to tool_node; after a tool execution, it routes back to agent_node for the LLM to reason on the tool's output. This creates the iterative ReAct loop: Agent (Reason) -> Tool (Act) -> Agent (Observe/Reason) -> ....
Finally, we add these new nodes to the workflow, set the agent_node as the entry_point, and define conditional edges using add_conditional_edges. This dynamically routes the execution based on the decide_next_step function, completing the ReAct architecture.
Expected outcome
A fully functional LangGraph agent that can take a query, reason about it, decide to use a tool (web search or report format), execute the tool, and then reason again on the tool's output. It can also handle basic retries on errors and produce a final report.
How to test
Integrate the app_agent = workflow.compile() into app/main.py. Update the /research endpoint in app/main.py to call app_agent.invoke({'query': request.query}). Run the FastAPI app and send a research query. Observe the logs to see the agent's reasoning, tool calls, and state transitions. Expect multiple log messages indicating the agent reasoning, performing web search, and eventually formatting a report.
Phase 4: Refine Agent Logic, State Management, and Error Handling
This phase refines the agent's core logic, focusing on robust state management, improved error handling, and sophisticated information synthesis. We'll enhance the agent_node to parse and store research findings more effectively, manage conversation history for better context, and ensure the agent can identify and resolve contradictions.
import os
import logging
from typing import List, Dict, Any, Callable, Tuple, Literal, TypedDict
from functools import wraps
import time
import json
from langchain_community.tools.tavily_research import TavilySearchResults
from langchain_core.tools import tool
from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.messages import BaseMessage, HumanMessage, AIMessage, FunctionMessage
from langgraph.graph import StateGraph, START, END
from pydantic import BaseModel, Field, ValidationError
from fastapi import FastAPI, HTTPException
# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)
# Helper for retries with exponential backoff
def retry_with_exponential_backoff(max_retries: int = 3, initial_delay: float = 1.0):
def decorator(func: Callable):
@wraps(func)
def wrapper(*args, **kwargs):
delay = initial_delay
for i in range(max_retries):
try:
return func(*args, **kwargs)
except Exception as e:
logger.warning(f"Attempt {i+1}/{max_retries} failed for {func.__name__}: {e}")
if i < max_retries - 1:
time.sleep(delay)
delay *= 2 # Exponential backoff
else:
raise
return wrapper
return decorator
# Pydantic Models for API input/output and agent state
class ResearchRequest(BaseModel):
query: str = Field(..., min_length=1, description="The research query for the AI agent.")
class ResearchReport(BaseModel):
report_content: str = Field(..., description="The generated research report with citations.")
sources: List[str] = Field(default_factory=list, description="List of URLs or source identifiers.")
status: Literal["success", "failure"] = Field(..., description="Status of the research request.")
error_message: str | None = Field(default=None, description="Error message if status is failure.")
class AgentState(TypedDict):
query: str
chat_history: List[BaseMessage]
research_findings: Dict[str, Any]
report: str | None
tool_output: str | None
error: str | None
max_retries: int
current_retry: int
num_steps: int # To prevent infinite loops
# --- Tools Definition (app/tools.py)
class ResearchTools:
def __init__(self):
tavily_api_key = os.environ.get("TAVILY_API_KEY")
if not tavily_api_key:
logger.error("TAVILY_API_KEY not found in environment variables.")
raise ValueError("TAVILY_API_KEY is required for web search.")
self.tavily_search = TavilySearchResults(max_results=5, tavily_api_key=tavily_api_key)
@tool
@retry_with_exponential_backoff()
def web_search(self, query: str) -> str:
"""Searches the web for the given query and returns relevant snippets. """
logger.info(f"Performing web search for: {query}")
try:
results = self.tavily_search.run(query)
return str(results)
except Exception as e:
logger.error(f"Web search failed for query '{query}': {e}", exc_info=True)
raise
@tool
def format_report(self, findings: Dict[str, Any]) -> str:
"""Formats the gathered research findings into a structured markdown report. """
logger.info("Formatting research report.")
report_parts = ["# Research Report
"]
query = findings.get("query", "N/A")
report_parts.append(f"## Topic: {query}
")
if "summary" in findings and findings["summary"]:
report_parts.append("### Summary
")
report_parts.append(findings["summary"] + "
")
if "details" in findings and findings["details"]:
report_parts.append("### Detailed Findings
")
for topic, content in findings["details"].items():
report_parts.append(f"#### {topic}
")
report_parts.append(content + "
")
if "sources" in findings and isinstance(findings["sources"], list) and findings["sources"]:
report_parts.append("### Sources
")
for i, source in enumerate(findings["sources"]):
report_parts.append(f"{i+1}. {source}
")
formatted_report = "
".join(report_parts)
logger.info("Report formatting complete.")
return formatted_report
# --- Agent Graph Definition (app/agent.py)
# Initialize LLM
openai_api_key = os.environ.get("OPENAI_API_KEY")
if not openai_api_key:
logger.error("OPENAI_API_KEY not found in environment variables.")
raise ValueError("OPENAI_API_KEY is required for LLM.")
llm = ChatOpenAI(model="gpt-4o", temperature=0.2, api_key=openai_api_key)
# Initialize Tools
research_tools_instance = ResearchTools()
tools = [research_tools_instance.web_search, research_tools_instance.format_report]
# Bind tools to LLM for function calling
llm_with_tools = llm.bind_tools(tools)
# Define the LangGraph State Graph - using TypedDict now
workflow = StateGraph(AgentState)
# --- Nodes for the ReAct Agent ---
def agent_node(state: AgentState) -> AgentState:
logger.info(f"Entering agent_node for query: {state['query']}")
state['num_steps'] += 1
if state['num_steps'] > state['max_retries'] * 3: # Heuristic to prevent infinite loops
return {"error": "Agent exceeded maximum reasoning steps, potential infinite loop.", "current_retry": state['max_retries']}
messages: List[BaseMessage] = state['chat_history']
system_prompt = (
"You are an expert AI research assistant. Your goal is to conduct thorough research "
"on the user's query, synthesize information from various sources, identify and resolve "
"contradictions, and produce a comprehensive, well-cited report. "
"You have access to the following tools: " + "
".join([f"- {t.name}: {t.description}" for t in tools]) + "
"
"Current Research Findings:
" + json.dumps(state['research_findings'], indent=2) + "
"
"Instructions:
"
"1. **Search**: If more information is needed, use `web_search` with precise queries."
"2. **Analyze**: After searching, analyze the 'tool_output'. Extract key facts, data, and potential contradictions. Update `research_findings` in a structured JSON format (e.g., {'summary': '...', 'details': {'topic1': '...', 'topic2': '...'}, 'sources': ['url1', 'url2']})."
"3. **Synthesize & Resolve**: If multiple sources present conflicting information, use `web_search` again to find clarifying data or state the contradiction clearly in your findings. Continuously refine `research_findings`."
"4. **Report**: Once you have sufficient, consistent information to answer the original query, and you have structured the `research_findings` well, call `format_report` with the full `research_findings` dictionary. Ensure all claims are supported by sources. If you believe the report is complete, call `format_report`."
"5. **Reflect**: Before calling `format_report`, ensure your findings are comprehensive and address all aspects of the original query. If not, continue searching and analyzing."
"Your response should either be a tool call or a final thought process leading to a tool call."
)
# Replace the initial HumanMessage with the system_prompt and query
if not messages or messages[0].content != state['query']:
messages = [HumanMessage(content=system_prompt), HumanMessage(content=state['query'])]
else:
# Prepend system prompt to existing history if it's not already there or needs update
if not messages or (isinstance(messages[0], HumanMessage) and not messages[0].content.startswith("You are an expert AI research assistant")):
messages.insert(0, HumanMessage(content=system_prompt))
else:
messages[0].content = system_prompt # Update system prompt with latest findings
try:
response = llm_with_tools.invoke(messages)
logger.info(f"LLM response: {response.content}, Tool calls: {response.tool_calls}")
new_chat_history = state['chat_history'] + [response]
new_state = {"chat_history": new_chat_history, "error": None}
if response.tool_calls:
tool_call = response.tool_calls[0]
tool_name = tool_call["name"]
tool_args = tool_call["args"]
new_state["tool_output"] = None # Clear previous tool output
new_state['chat_history'].append(AIMessage(content="", tool_calls=[tool_call]))
# Store action for tool_node
new_state['intermediate_steps'] = state.get('intermediate_steps', []) + [{
"action": tool_name,
"action_input": tool_args,
"thought": response.content # Store LLM's thought as well
}]
else:
# If LLM doesn't call a tool, it's either a final answer or it's stuck.
# For research agent, it should always call a tool until report is done.
logger.warning("LLM did not call a tool. Assuming final answer or unexpected state.")
new_state["report"] = response.content # Treat as final report if no tool call
new_state['chat_history'].append(AIMessage(content=response.content))
return new_state
except Exception as e:
logger.error(f"Error in agent_node: {e}", exc_info=True)
return {"error": str(e), "current_retry": state['current_retry'] + 1}
def tool_node(state: AgentState) -> AgentState:
logger.info(f"Entering tool_node. Current intermediate steps: {state.get('intermediate_steps', [])}")
last_step = state.get('intermediate_steps', [])[-1]
tool_name = last_step.get("action")
tool_args = last_step.get("action_input")
if not tool_name or not tool_args:
logger.error(f"Tool name or arguments missing in state: {last_step}")
return {"error": "Missing tool name or arguments", "current_retry": state['current_retry'] + 1}
tool_output_str = None
try:
if tool_name == "web_search":
query = tool_args.get("query")
if not query:
raise ValueError("Web search tool requires a 'query' argument.")
tool_output_str = research_tools_instance.web_search.invoke(query)
# Parse Tavily search results more robustly
parsed_results = json.loads(tool_output_str) if isinstance(tool_output_str, str) else tool_output_str
if isinstance(parsed_results, list):
content_snippets = "
".join([r['content'] for r in parsed_results if 'content' in r])
sources_list = list(set([r['url'] for r in parsed_results if 'url' in r])) # Unique sources
# Update research findings: LLM will synthesize this later
# For now, just append raw content to a 'search_results' key
current_findings = state.get('research_findings', {})
current_details = current_findings.get('details', {})
current_details['search_results'] = current_details.get('search_results', '') + "
" + content_snippets
current_findings['details'] = current_details
current_findings['sources'] = list(set(current_findings.get('sources', []) + sources_list))
return {"tool_output": tool_output_str, "research_findings": current_findings, "error": None}
else:
logger.warning(f"Unexpected search result format: {type(parsed_results)}")
tool_output_str = f"Web search completed. Raw result: {tool_output_str}"
elif tool_name == "format_report":
findings_for_report = state.get('research_findings')
if not findings_for_report:
raise ValueError("No research findings available to format report.")
# Ensure findings have query, summary, details, sources before formatting
if not all(k in findings_for_report for k in ['query', 'summary', 'details', 'sources']):
# Attempt to extract from tool_args if LLM put it there (less robust)
if 'findings' in tool_args and isinstance(tool_args['findings'], dict):
findings_for_report = tool_args['findings']
else:
# Use a fallback LLM call to structure findings if they are not structured yet
logger.warning("Research findings not fully structured for report. Attempting to structure with LLM.")
structure_prompt = ChatPromptTemplate.from_messages([
("system", "You are an expert at structuring raw research into a JSON format with 'query', 'summary', 'details' (dict of subtopics), and 'sources' (list of urls). Given the raw findings, structure them."),
("human", f"Raw findings: {json.dumps(findings_for_report)}")
])
structured_findings_str = llm.invoke(structure_prompt).content
try:
findings_for_report = json.loads(structured_findings_str)
if not isinstance(findings_for_report, dict):
raise ValueError("LLM did not return a dictionary.")
except json.JSONDecodeError:
logger.error(f"Failed to parse LLM structured findings: {structured_findings_str}")
raise ValueError("LLM failed to structure findings into JSON.")
report_content = research_tools_instance.format_report.invoke(findings_for_report)
return {"report": report_content, "tool_output": report_content, "error": None}
else:
raise ValueError(f"Unknown tool: {tool_name}")
except Exception as e:
logger.error(f"Tool '{tool_name}' execution failed: {e}", exc_info=True)
return {"error": str(e), "current_retry": state['current_retry'] + 1}
# Update chat history with tool observation for LLM to reason on
new_chat_history = state['chat_history'] + [FunctionMessage(name=tool_name, content=tool_output_str)]
return {"tool_output": tool_output_str, "chat_history": new_chat_history, "error": None}
def decide_next_step(state: AgentState) -> str:
logger.info(f"Deciding next step. Current report status: {state['report'] is not None}, Error: {state['error']}")
if state['error'] and state['current_retry'] < state['max_retries']:
logger.warning(f"Retrying due to error: {state['error']}. Attempt {state['current_retry']}/{state['max_retries']}")
# Reset error and increment retry count, then go back to agent_node for re-planning
return "agent_node"
elif state['error'] and state['current_retry'] >= state['max_retries']:
logger.error(f"Max retries reached. Failing research request. Last error: {state['error']}")
return "END"
elif state['report']:
logger.info("Report successfully generated. Ending workflow.")
return "END"
else:
# Check if the last message from LLM was a tool call
last_message = state['chat_history'][-1] if state['chat_history'] else None
if last_message and isinstance(last_message, AIMessage) and last_message.tool_calls:
logger.info(f"LLM decided to use tool: {last_message.tool_calls[0]['name']}. Executing tool.")
return "tool_node"
else:
logger.info("Continuing research, returning to agent_node for further reasoning.")
return "agent_node"
# Configure LangGraph with TypedDict for state
workflow = StateGraph(AgentState)
workflow.add_node("agent_node", agent_node)
workflow.add_node("tool_node", tool_node)
workflow.set_entry_point("agent_node")
workflow.add_conditional_edges(
"agent_node",
decide_next_step,
{
"tool_node": "tool_node",
"agent_node": "agent_node",
"END": END
}
)
workflow.add_conditional_edges(
"tool_node",
decide_next_step,
{
"agent_node": "agent_node",
"END": END
}
)
app_agent = workflow.compile()
# --- app/main.py (updated to integrate compiled agent)
app = FastAPI(
title="AI Research Agent API",
description="API for an autonomous AI agent capable of conducting web-based research and generating reports.",
version="1.0.0"
)
@app.get("/health")
async def health_check():
return {"status": "healthy"}
@app.post("/research", response_model=ResearchReport)
async def perform_research(request: ResearchRequest):
logger.info(f"Received research request: {request.query}")
initial_state: AgentState = {
"query": request.query,
"chat_history": [HumanMessage(content=request.query)],
"research_findings": {"query": request.query},
"report": None,
"tool_output": None,
"error": None,
"max_retries": 3,
"current_retry": 0,
"num_steps": 0
}
try:
# Execute the LangGraph agent
final_state = app_agent.invoke(initial_state)
if final_state.get("report"):
sources = final_state.get('research_findings', {}).get('sources', [])
return ResearchReport(
report_content=final_state["report"],
sources=sources,
status="success"
)
else:
error_msg = final_state.get("error", "Agent failed to produce a report without a specific error message.")
logger.error(f"Agent finished without a report: {error_msg}")
raise HTTPException(status_code=500, detail=f"Research agent failed: {error_msg}")
except Exception as e:
logger.error(f"Error processing research request: {e}", exc_info=True)
raise HTTPException(status_code=500, detail=f"Internal server error: {e}")
This phase significantly refines the agent's intelligence and robustness. First, the
AgentState is refactored to use TypedDict, which offers better integration with LangGraph and ensures type consistency. A chat_history field is added to the state to maintain the full conversation (including LLM thoughts, tool calls, and observations), providing better context for the LLM across turns. A num_steps counter is introduced to prevent infinite loops, an important safety mechanism in autonomous agents.
The agent_node is enhanced with a more detailed system_prompt that explicitly guides the LLM on research strategy, including steps for searching, analyzing, synthesizing, resolving contradictions, and finally reporting. This prompt emphasizes structuring research_findings as a JSON dictionary, which is critical for the format_report tool. The chat_history is now passed to the LLM, and the LLM's responses (both thoughts and tool calls) are appended to this history, allowing for long-term reasoning. Error handling within the agent_node is also improved to capture exceptions and update the state's error field.
The tool_node now handles web_search results by parsing the JSON output from Tavily and progressively updating the research_findings in the state. Instead of directly summarizing, it appends raw content and sources, allowing the LLM in the agent_node to perform the complex synthesis. For format_report, a crucial improvement is added: if the research_findings are not fully structured (e.g., missing 'summary' or 'details'), a fallback LLM call is made to explicitly structure them into the required JSON format. This makes the report generation more resilient to imperfect intermediate states. The tool_node also appends FunctionMessages to chat_history after tool execution, giving the LLM the observation it needs for the next reasoning step.
Finally, app/main.py is updated to integrate the compiled app_agent. The /research endpoint now initializes the AgentState and invokes the LangGraph agent. It then processes the final_state to extract the generated report_content and sources for the ResearchReport response, or raises an HTTPException if the agent fails to produce a report after exhausting retries or encountering an unrecoverable error. This robust error handling and structured output ensure the API is production-ready.
Expected outcome
A more intelligent and robust AI research agent capable of iterative research, better state management, and improved error recovery. The FastAPI endpoint will now trigger the full LangGraph workflow and return a structured report or an error.
How to test
Update app/main.py, app/agent.py, and app/tools.py with the new code. Ensure your .env has OPENAI_API_KEY and TAVILY_API_KEY. Run the FastAPI app and send a complex research query, e.g., "Compare and contrast the ethical implications of large language models and autonomous driving systems.". Observe the detailed logs to see the agent's multi-step reasoning, multiple web searches, synthesis, and final report generation. Verify the report quality and source citations.
Phase 5: Production Readiness - Observability and Dockerization
This phase focuses on preparing the agent for production deployment by enhancing observability with structured logging and containerizing the application using Docker. We will create a Dockerfile and update the FastAPI application to include a uvicorn entry point, ensuring it runs reliably in a containerized environment.
import os
import logging
from typing import List, Dict, Any, Callable, Tuple, Literal, TypedDict
from functools import wraps
import time
import json
import uvicorn
from langchain_community.tools.tavily_research import TavilySearchResults
from langchain_core.tools import tool
from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.messages import BaseMessage, HumanMessage, AIMessage, FunctionMessage
from langgraph.graph import StateGraph, START, END
from pydantic import BaseModel, Field, ValidationError
from fastapi import FastAPI, HTTPException
# Configure structured logging for production
# This simple configuration logs JSON-like output, suitable for log aggregators
class JsonFormatter(logging.Formatter):
def format(self, record):
log_entry = {
"timestamp": self.formatTime(record, self.datefmt),
"level": record.levelname,
"name": record.name,
"message": record.getMessage(),
"file": record.filename,
"line": record.lineno
}
if hasattr(record, 'extra_data'):
log_entry.update(record.extra_data)
if record.exc_info:
log_entry["exception"] = self.formatException(record.exc_info)
return json.dumps(log_entry)
log_handler = logging.StreamHandler()
log_handler.setFormatter(JsonFormatter())
root_logger = logging.getLogger()
root_logger.setLevel(os.environ.get("LOG_LEVEL", "INFO").upper())
root_logger.handlers = [] # Clear existing handlers
root_logger.addHandler(log_handler)
logger = logging.getLogger(__name__)
# Helper for retries with exponential backoff
def retry_with_exponential_backoff(max_retries: int = 3, initial_delay: float = 1.0):
def decorator(func: Callable):
@wraps(func)
def wrapper(*args, **kwargs):
delay = initial_delay
for i in range(max_retries):
try:
return func(*args, **kwargs)
except Exception as e:
logger.warning(f"Attempt {i+1}/{max_retries} failed for {func.__name__}: {e}", extra_data={"function": func.__name__,"attempt": i+1})
if i < max_retries - 1:
time.sleep(delay)
delay *= 2 # Exponential backoff
else:
raise
return wrapper
return decorator
# Pydantic Models for API input/output and agent state (as defined in Phase 4)
class ResearchRequest(BaseModel):
query: str = Field(..., min_length=1, description="The research query for the AI agent.")
class ResearchReport(BaseModel):
report_content: str = Field(..., description="The generated research report with citations.")
sources: List[str] = Field(default_factory=list, description="List of URLs or source identifiers.")
status: Literal["success", "failure"] = Field(..., description="Status of the research request.")
error_message: str | None = Field(default=None, description="Error message if status is failure.")
class AgentState(TypedDict):
query: str
chat_history: List[BaseMessage]
research_findings: Dict[str, Any]
report: str | None
tool_output: str | None
error: str | None
max_retries: int
current_retry: int
num_steps: int
# --- Tools Definition (app/tools.py - identical to Phase 4)
class ResearchTools:
def __init__(self):
tavily_api_key = os.environ.get("TAVILY_API_KEY")
if not tavily_api_key:
logger.error("TAVILY_API_KEY not found in environment variables.")
raise ValueError("TAVILY_API_KEY is required for web search.")
self.tavily_search = TavilySearchResults(max_results=5, tavily_api_key=tavily_api_key)
@tool
@retry_with_exponential_backoff()
def web_search(self, query: str) -> str:
"""Searches the web for the given query and returns relevant snippets. """
logger.info(f"Performing web search for: {query}", extra_data={"query": query})
try:
results = self.tavily_search.run(query)
return str(results)
except Exception as e:
logger.error(f"Web search failed for query '{query}': {e}", exc_info=True, extra_data={"query": query})
raise
@tool
def format_report(self, findings: Dict[str, Any]) -> str:
"""Formats the gathered research findings into a structured markdown report. """
logger.info("Formatting research report.", extra_data={"findings_summary": list(findings.keys())})
report_parts = ["# Research Report
"]
query = findings.get("query", "N/A")
report_parts.append(f"## Topic: {query}
")
if "summary" in findings and findings["summary"]:
report_parts.append("### Summary
")
report_parts.append(findings["summary"] + "
")
if "details" in findings and findings["details"]:
report_parts.append("### Detailed Findings
")
for topic, content in findings["details"].items():
report_parts.append(f"#### {topic}
")
report_parts.append(content + "
")
if "sources" in findings and isinstance(findings["sources"], list) and findings["sources"]:
report_parts.append("### Sources
")
for i, source in enumerate(findings["sources"]):
report_parts.append(f"{i+1}. {source}
")
formatted_report = "
".join(report_parts)
logger.info("Report formatting complete.")
return formatted_report
# --- Agent Graph Definition (app/agent.py - identical to Phase 4, with updated logging)
# Initialize LLM
openai_api_key = os.environ.get("OPENAI_API_KEY")
if not openai_api_key:
logger.error("OPENAI_API_KEY not found in environment variables.")
raise ValueError("OPENAI_API_KEY is required for LLM.")
llm = ChatOpenAI(model="gpt-4o", temperature=0.2, api_key=openai_api_key)
# Initialize Tools
research_tools_instance = ResearchTools()
tools = [research_tools_instance.web_search, research_tools_instance.format_report]
# Bind tools to LLM for function calling
llm_with_tools = llm.bind_tools(tools)
# Define the LangGraph State Graph - using TypedDict now
workflow = StateGraph(AgentState)
# --- Nodes for the ReAct Agent ---
def agent_node(state: AgentState) -> AgentState:
logger.info(f"Entering agent_node for query: {state['query']}", extra_data={"num_steps": state['num_steps'], "query": state['query']})
state['num_steps'] += 1
if state['num_steps'] > state['max_retries'] * 3: # Heuristic to prevent infinite loops
logger.error("Agent exceeded maximum reasoning steps, potential infinite loop.", extra_data={"query": state['query'], "num_steps": state['num_steps']})
return {"error": "Agent exceeded maximum reasoning steps, potential infinite loop.", "current_retry": state['max_retries']}
messages: List[BaseMessage] = state['chat_history']
system_prompt = (
"You are an expert AI research assistant. Your goal is to conduct thorough research "
"on the user's query, synthesize information from various sources, identify and resolve "
"contradictions, and produce a comprehensive, well-cited report. "
"You have access to the following tools: " + "
".join([f"- {t.name}: {t.description}" for t in tools]) + "
"
"Current Research Findings:
" + json.dumps(state['research_findings'], indent=2) + "
"
"Instructions:
"
"1. **Search**: If more information is needed, use `web_search` with precise queries."
"2. **Analyze**: After searching, analyze the 'tool_output'. Extract key facts, data, and potential contradictions. Update `research_findings` in a structured JSON format (e.g., {'summary': '...', 'details': {'topic1': '...', 'topic2': '...'}, 'sources': ['url1', 'url2']})."
"3. **Synthesize & Resolve**: If multiple sources present conflicting information, use `web_search` again to find clarifying data or state the contradiction clearly in your findings. Continuously refine `research_findings`."
"4. **Report**: Once you have sufficient, consistent information to answer the original query, and you have structured the `research_findings` well, call `format_report` with the full `research_findings` dictionary. Ensure all claims are supported by sources. If you believe the report is complete, call `format_report`."
"5. **Reflect**: Before calling `format_report`, ensure your findings are comprehensive and address all aspects of the original query. If not, continue searching and analyzing."
"Your response should either be a tool call or a final thought process leading to a tool call."
)
if not messages or (isinstance(messages[0], HumanMessage) and not messages[0].content.startswith("You are an expert AI research assistant")):
messages.insert(0, HumanMessage(content=system_prompt))
else:
messages[0].content = system_prompt # Update system prompt with latest findings
try:
response = llm_with_tools.invoke(messages)
logger.info(f"LLM response generated.", extra_data={"tool_calls_count": len(response.tool_calls) if response.tool_calls else 0, "response_content_length": len(response.content)})
new_chat_history = state['chat_history'] + [response]
new_state = {"chat_history": new_chat_history, "error": None}
if response.tool_calls:
tool_call = response.tool_calls[0]
tool_name = tool_call["name"]
tool_args = tool_call["args"]
new_state["tool_output"] = None
new_state['chat_history'].append(AIMessage(content="", tool_calls=[tool_call]))
new_state['intermediate_steps'] = state.get('intermediate_steps', []) + [{
"action": tool_name,
"action_input": tool_args,
"thought": response.content
}]
else:
logger.warning("LLM did not call a tool, assuming final report or unexpected state.", extra_data={"llm_response": response.content})
new_state["report"] = response.content
new_state['chat_history'].append(AIMessage(content=response.content))
return new_state
except Exception as e:
logger.error(f"Error in agent_node: {e}", exc_info=True, extra_data={"query": state['query']})
return {"error": str(e), "current_retry": state['current_retry'] + 1}
def tool_node(state: AgentState) -> AgentState:
logger.info(f"Entering tool_node.", extra_data={"current_step_count": len(state.get('intermediate_steps', []))})
last_step = state.get('intermediate_steps', [])[-1]
tool_name = last_step.get("action")
tool_args = last_step.get("action_input")
if not tool_name or not tool_args:
logger.error(f"Tool name or arguments missing in state: {last_step}", extra_data={"last_step": last_step})
return {"error": "Missing tool name or arguments", "current_retry": state['current_retry'] + 1}
tool_output_str = None
try:
if tool_name == "web_search":
query = tool_args.get("query")
if not query:
raise ValueError("Web search tool requires a 'query' argument.")
tool_output_str = research_tools_instance.web_search.invoke(query)
parsed_results = json.loads(tool_output_str) if isinstance(tool_output_str, str) else tool_output_str
if isinstance(parsed_results, list):
content_snippets = "
".join([r['content'] for r in parsed_results if 'content' in r])
sources_list = list(set([r['url'] for r in parsed_results if 'url' in r]))
current_findings = state.get('research_findings', {})
current_details = current_findings.get('details', {})
current_details['search_results'] = current_details.get('search_results', '') + "
" + content_snippets
current_findings['details'] = current_details
current_findings['sources'] = list(set(current_findings.get('sources', []) + sources_list))
return {"tool_output": tool_output_str, "research_findings": current_findings, "error": None}
else:
logger.warning(f"Unexpected search result format: {type(parsed_results)}", extra_data={"raw_output": tool_output_str})
tool_output_str = f"Web search completed. Raw result: {tool_output_str}"
elif tool_name == "format_report":
findings_for_report = state.get('research_findings')
if not findings_for_report:
raise ValueError("No research findings available to format report.")
if not all(k in findings_for_report for k in ['query', 'summary', 'details', 'sources']):
logger.warning("Research findings not fully structured for report. Attempting to structure with LLM.", extra_data={"current_findings_keys": list(findings_for_report.keys())})
structure_prompt = ChatPromptTemplate.from_messages([
("system", "You are an expert at structuring raw research into a JSON format with 'query', 'summary', 'details' (dict of subtopics), and 'sources' (list of urls). Given the raw findings, structure them."),
("human", f"Raw findings: {json.dumps(findings_for_report)}")
])
structured_findings_str = llm.invoke(structure_prompt).content
try:
findings_for_report = json.loads(structured_findings_str)
if not isinstance(findings_for_report, dict):
raise ValueError("LLM did not return a dictionary.")
except json.JSONDecodeError:
logger.error(f"Failed to parse LLM structured findings: {structured_findings_str}")
raise ValueError("LLM failed to structure findings into JSON.")
report_content = research_tools_instance.format_report.invoke(findings_for_report)
return {"report": report_content, "tool_output": report_content, "error": None}
else:
raise ValueError(f"Unknown tool: {tool_name}")
except Exception as e:
logger.error(f"Tool '{tool_name}' execution failed: {e}", exc_info=True, extra_data={"tool_name": tool_name, "tool_args": tool_args})
return {"error": str(e), "current_retry": state['current_retry'] + 1}
new_chat_history = state['chat_history'] + [FunctionMessage(name=tool_name, content=tool_output_str)]
return {"tool_output": tool_output_str, "chat_history": new_chat_history, "error": None}
def decide_next_step(state: AgentState) -> str:
logger.info(f"Deciding next step.", extra_data={"report_generated": state['report'] is not None, "error_present": state['error'] is not None, "current_retry": state['current_retry']})
if state['error'] and state['current_retry'] < state['max_retries']:
logger.warning(f"Retrying due to error: {state['error']}. Attempt {state['current_retry']}/{state['max_retries']}")
return "agent_node"
elif state['error'] and state['current_retry'] >= state['max_retries']:
logger.error(f"Max retries reached. Failing research request. Last error: {state['error']}")
return "END"
elif state['report']:
logger.info("Report successfully generated. Ending workflow.")
return "END"
else:
last_message = state['chat_history'][-1] if state['chat_history'] else None
if last_message and isinstance(last_message, AIMessage) and last_message.tool_calls:
logger.info(f"LLM decided to use tool: {last_message.tool_calls[0]['name']}. Executing tool.")
return "tool_node"
else:
logger.info("Continuing research, returning to agent_node for further reasoning.")
return "agent_node"
# Configure LangGraph with TypedDict for state
workflow = StateGraph(AgentState)
workflow.add_node("agent_node", agent_node)
workflow.add_node("tool_node", tool_node)
workflow.set_entry_point("agent_node")
workflow.add_conditional_edges(
"agent_node",
decide_next_step,
{
"tool_node": "tool_node",
"agent_node": "agent_node",
"END": END
}
)
workflow.add_conditional_edges(
"tool_node",
decide_next_step,
{
"agent_node": "agent_node",
"END": END
}
)
app_agent = workflow.compile()
# --- app/main.py (fully updated and ready for Docker)
app = FastAPI(
title="AI Research Agent API",
description="API for an autonomous AI agent capable of conducting web-based research and generating reports.",
version="1.0.0"
)
@app.get("/health")
async def health_check():
return {"status": "healthy"}
@app.post("/research", response_model=ResearchReport)
async def perform_research(request: ResearchRequest):
logger.info(f"Received research request.", extra_data={"query": request.query})
initial_state: AgentState = {
"query": request.query,
"chat_history": [HumanMessage(content=request.query)],
"research_findings": {"query": request.query},
"report": None,
"tool_output": None,
"error": None,
"max_retries": int(os.environ.get("AGENT_MAX_RETRIES", "3")),
"current_retry": 0,
"num_steps": 0
}
try:
final_state = app_agent.invoke(initial_state)
if final_state.get("report"):
sources = final_state.get('research_findings', {}).get('sources', [])
logger.info("Research completed successfully.", extra_data={"query": request.query, "report_length": len(final_state['report'])})
return ResearchReport(
report_content=final_state["report"],
sources=sources,
status="success"
)
else:
error_msg = final_state.get("error", "Agent failed to produce a report without a specific error message.")
logger.error(f"Agent finished without a report: {error_msg}", extra_data={"query": request.query, "final_state_error": error_msg})
raise HTTPException(status_code=500, detail=f"Research agent failed: {error_msg}")
except Exception as e:
logger.error(f"Error processing research request: {e}", exc_info=True, extra_data={"query": request.query})
raise HTTPException(status_code=500, detail=f"Internal server error: {e}")
# Entry point for uvicorn when running directly or via Docker
if __name__ == "__main__":
uvicorn.run("app.main:app", host="0.0.0.0", port=8000, reload=True if os.environ.get("ENV") == "development" else False)
# Dockerfile content (create at project root)
# FROM python:3.11-slim-bookworm
#
# WORKDIR /app
#
# COPY requirements.txt .
# RUN pip install --no-cache-dir -r requirements.txt
#
# COPY ./app .
#
# EXPOSE 8000
#
# CMD ["python", "main.py"]
This final implementation phase focuses on making our AI research agent production-ready. A custom
JsonFormatter is introduced for the Python logging module. This formatter outputs logs in a structured JSON format, which is significantly easier for centralized logging systems (like ELK stack, Splunk, Datadog) to parse, filter, and analyze compared to plain text logs. It also allows for adding extra_data to log records, providing more context (e.g., query, tool_name) to each log entry, which is invaluable for debugging and monitoring in a distributed environment. The LOG_LEVEL is made configurable via an environment variable.
All existing log statements in agent_node, tool_node, and decide_next_step are updated to use the new structured logging approach, providing richer context for each step of the agent's execution. This ensures that when the agent runs in production, engineers can quickly understand its behavior, identify bottlenecks, or diagnose failures.
The app/main.py is updated with an if __name__ == "__main__": block, allowing it to be directly executed by uvicorn. The AGENT_MAX_RETRIES value is now dynamically pulled from environment variables, making the agent's retry behavior configurable without code changes. The initial_state for the agent is carefully constructed, including the chat_history and research_findings to ensure the agent starts with the correct context.
Finally, a Dockerfile is provided. This Dockerfile defines a lightweight Python environment, copies the requirements.txt and installs dependencies, then copies the application code. It exposes port 8000 and sets the CMD to run main.py using Python. This containerization ensures that our application runs consistently across different environments, from local development to staging and production, by bundling all its dependencies and configuration into a single, portable unit. It's a critical step for reliable deployment.
Expected outcome
A fully containerized AI research agent that logs structured JSON output and can be deployed consistently. The agent will execute complex research queries, handle errors, and produce high-quality reports within a Docker container.
How to test
Create the Dockerfile at the root of your ai_research_agent directory. Build the Docker image: docker build -t ai-research-agent .. Run the container: docker run -p 8000:8000 --env-file ./.env ai-research-agent. Send a research request via curl or Postman. Observe the Docker container logs, which should now be in structured JSON format. Test with both simple and complex queries to verify functionality and error handling.
Testing & Evaluation
Testing and evaluating an AI research agent goes beyond simple unit tests; it requires assessing the quality, accuracy, and completeness of its outputs, as well as its resilience to failures. For this project, we'll cover functional testing, LLM-as-a-Judge evaluation, failure scenario testing, edge case coverage, and performance validation.
Functional Testing
Functional tests verify that individual components and the overall system behave as expected. For our agent, this includes:
- Tool Unit Tests: Test
web_searchandformat_reporttools in isolation. Verifyweb_searchreturns relevant results for specific queries andformat_reportcorrectly structures differentfindingsdictionaries. - LangGraph Node Tests: Test
agent_nodeandtool_nodewith mocked LLM responses and tool outputs to ensure they correctly update theAgentStateand make appropriate routing decisions. - End-to-End API Tests: Send various research queries to the
/researchendpoint and assert that aResearchReportis returned withstatus="success", and thereport_contentis non-empty and contains expected keywords. Verifysourcesare present and valid URLs.
LLM-as-a-Judge Evaluation
Evaluating the quality of generated reports is subjective, making LLM-as-a-Judge a powerful technique. We can use a separate, higher-capacity LLM to score our agent's output against a golden standard or a set of criteria.
Here's an example using LangChain and a custom prompt:
import os
import json
from typing import Dict, Any, List
from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate
from pydantic import BaseModel, Field
# Assume OPENAI_API_KEY is in environment
llm_judge = ChatOpenAI(model="gpt-4o", temperature=0)
class EvaluationResult(BaseModel):
score: int = Field(..., description="Score from 1 (poor) to 5 (excellent).")
reasoning: str = Field(..., description="Detailed reasoning for the score.")
completeness: bool = Field(..., description="True if the report fully addresses the query.")
accuracy: bool = Field(..., description="True if all facts in the report are accurate and cited.")
contradiction_resolution: bool = Field(..., description="True if identified contradictions were resolved or clearly stated.")
citations_present: bool = Field(..., description="True if all claims are supported by citations.")
def evaluate_report_with_llm(query: str, generated_report: str) -> EvaluationResult:
evaluation_prompt = ChatPromptTemplate.from_messages([
("system", "You are an expert AI report evaluator. Your task is to critically assess a research report generated by another AI agent. Score the report from 1 (very poor) to 5 (excellent) based on completeness, accuracy, contradiction resolution, and citation quality. Provide detailed reasoning."),
("human", f"Original Query: {query}
Generated Report:
{generated_report}
---
Evaluate this report and provide a JSON output with the following schema:
{{"score": int, "reasoning": str, "completeness": bool, "accuracy": bool, "contradiction_resolution": bool, "citations_present": bool}}")
])
parser = EvaluationResult.with_structured_output(llm_judge)
evaluation_output = parser.invoke(evaluation_prompt.format_messages(query=query, generated_report=generated_report))
return evaluation_output
# Example usage:
# research_query = "What are the main challenges in developing truly autonomous AI agents?"
# agent_report = "# Research Report
## Topic: Autonomous AI Challenges
... (actual generated report content) ..."
#
# try:
# eval_result = evaluate_report_with_llm(research_query, agent_report)
# print(json.dumps(eval_result.dict(), indent=2))
# except Exception as e:
# print(f"Evaluation failed: {e}")Failure Scenario Testing
- API Key Missing/Invalid: Test the agent's behavior when
OPENAI_API_KEYorTAVILY_API_KEYare missing or invalid. Expect appropriate error logging and HTTP 500 responses. - Tool Failures: Simulate
web_searchreturning empty results, API errors, or rate limits. Verify the retry mechanism and graceful error handling. - LLM Hallucinations/Misinterpretation: Provide ambiguous queries or deliberately contradictory initial information to see if the agent can identify and flag/resolve contradictions.
- Infinite Loop Detection: Test queries that might lead the agent into a loop (e.g., extremely broad searches, queries with no clear answer). Verify the
num_stepssafeguard triggers.
Edge Case Coverage
- Very Short/Long Queries: Ensure the agent handles queries of varying lengths.
- No Search Results: Test queries for which
TavilySearchResultsyields no relevant information. The agent should report this gracefully. - Complex/Niche Topics: Provide highly specialized or multi-faceted queries to stress the agent's ability to decompose tasks and synthesize complex information.
- Empty
research_findings: Test if theformat_reporttool or the LLM's structuring logic handles empty or incomplete findings gracefully.
Performance Validation
- Latency: Measure the end-to-end response time for research queries under various loads. Identify bottlenecks (e.g., LLM inference time, tool call latency). Logging with timestamps helps here.
- Token Usage: Monitor token usage per request to understand cost implications. Integrate token tracking if possible.
- Concurrency: Test how the FastAPI service handles multiple simultaneous research requests. Use tools like
locustorabto simulate load.
Regression Testing
As the agent evolves, maintain a suite of diverse test cases (queries and expected outcomes) to ensure new features or bug fixes don't introduce regressions. Automate these tests within a CI/CD pipeline, ideally including LLM-as-a-Judge evaluations for critical quality metrics.
Deployment
Deploying the AI research agent involves moving it from a local development environment to a production system, ensuring it's scalable, reliable, and secure. We'll leverage Docker for containerization and outline a deployment strategy for Kubernetes, a popular choice for orchestrating containerized applications.
1. Dockerfile
Our Dockerfile (located at the project root) defines the build process for our application's container image. It ensures a consistent environment.
# Use a lightweight Python base image
FROM python:3.11-slim-bookworm
# Set the working directory inside the container
WORKDIR /app
# Copy the requirements file and install dependencies
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt
# Copy the application code
COPY ./app .
# Expose the port FastAPI runs on
EXPOSE 8000
# Set environment variables for production (can be overridden at runtime)
ENV LOG_LEVEL=INFO
ENV AGENT_MAX_RETRIES=3
# Command to run the application using uvicorn
CMD ["python", "main.py"]Build and Run Locally with Docker:
- Create
.envfile at the project root:OPENAI_API_KEY="sk-...",TAVILY_API_KEY="your-tavily-api-key" - Build the Docker image:
docker build -t ai-research-agent:latest . - Run the container, mounting the
.envfile:docker run -p 8000:8000 --env-file ./.env ai-research-agent:latest
Alternatively, pass individual secrets: docker run -p 8000:8000 -e OPENAI_API_KEY=$OPENAI_API_KEY -e TAVILY_API_KEY=$TAVILY_API_KEY ai-research-agent:latest
2. Kubernetes Deployment
For production, Kubernetes provides orchestration, scaling, and self-healing. Here's a basic k8s-deployment.yaml:
apiVersion: apps/v1
kind: Deployment
metadata:
name: ai-research-agent
labels:
app: ai-research-agent
spec:
replicas: 2 # Start with 2 replicas for high availability
selector:
matchLabels:
app: ai-research-agent
template:
metadata:
labels:
app: ai-research-agent
spec:
containers:
- name: agent
image: your_docker_registry/ai-research-agent:latest # Replace with your pushed image
ports:
- containerPort: 8000
env:
- name: LOG_LEVEL
value: "INFO"
- name: AGENT_MAX_RETRIES
value: "3"
- name: OPENAI_API_KEY # Use Kubernetes Secrets for sensitive data
valueFrom:
secretKeyRef:
name: ai-research-agent-secrets
key: openai_api_key
- name: TAVILY_API_KEY
valueFrom:
secretKeyRef:
name: ai-research-agent-secrets
key: tavily_api_key
livenessProbe:
httpGet:
path: /health
port: 8000
initialDelaySeconds: 15
periodSeconds: 20
readinessProbe:
httpGet:
path: /health
port: 8000
initialDelaySeconds: 5
periodSeconds: 10
timeoutSeconds: 5
resources:
requests:
memory: "512Mi"
cpu: "500m"
limits:
memory: "1Gi"
cpu: "1000m"
---
apiVersion: v1
kind: Service
metadata:
name: ai-research-agent-service
spec:
selector:
app: ai-research-agent
ports:
- protocol: TCP
port: 80
targetPort: 8000
type: ClusterIP # Use LoadBalancer for external access in cloud environmentsSecrets Management (Kubernetes): Create a Kubernetes Secret to store your API keys securely:
kubectl create secret generic ai-research-agent-secrets \
--from-literal=openai_api_key=$OPENAI_API_KEY \
--from-literal=tavily_api_key=$TAVILY_API_KEYApply the deployment: kubectl apply -f k8s-deployment.yaml
3. Health Check and Readiness Probe
Both livenessProbe and readinessProbe are configured to hit the /health endpoint. The livenessProbe ensures the application is running and restarts the container if it becomes unresponsive. The readinessProbe ensures the container is ready to serve traffic before sending requests to it, preventing traffic from being sent to an uninitialized or unhealthy pod.
4. Scaling Considerations
- Horizontal Pod Autoscaler (HPA): Configure HPA based on CPU utilization or custom metrics (e.g., requests per second) to automatically scale the number of agent replicas up or down.
kubectl autoscale deployment ai-research-agent --cpu-percent=70 --min=2 --max=10. - Resource Requests/Limits: Define appropriate CPU and memory requests and limits in the deployment manifest to ensure fair resource allocation and prevent resource exhaustion.
- LLM Rate Limits: Be mindful of LLM provider rate limits. If scaling up, ensure your LLM plan can handle increased concurrent requests or implement client-side rate limiting/queueing.
5. Rollback Strategy
Kubernetes deployments automatically support rollbacks. If a new deployment introduces issues, you can revert to a previous stable version:
- Check deployment history:
kubectl rollout history deployment/ai-research-agent - Rollback to the previous version:
kubectl rollout undo deployment/ai-research-agent - Rollback to a specific revision:
kubectl rollout undo deployment/ai-research-agent --to-revision=N
Always tag your Docker images with meaningful version numbers (e.g., v1.0.0, v1.0.1) and use these tags in your Kubernetes manifests to ensure deterministic rollbacks.
Production Hardening
Hardening the AI research agent for production involves implementing best practices across security, observability, reliability, and cost management.
Security
- Secrets Management: Never hardcode API keys or sensitive credentials. Always use environment variables, and for deployment, leverage secure secrets management systems like Kubernetes Secrets, AWS Secrets Manager, or Azure Key Vault. This prevents sensitive data from being committed to source control. Our project uses
os.environ.get()and Kubernetes Secrets in the deployment. - Input Validation and Sanitization: Thoroughly validate all incoming API requests (
ResearchRequestPydantic model) to prevent injection attacks or unexpected inputs that could destabilize the agent or lead to prompt injection. While Pydantic handles basic validation, consider more advanced sanitization for complex string inputs. - Least Privilege: Ensure the container and underlying infrastructure run with the minimum necessary permissions. Avoid running as root. If interacting with other services, use fine-grained IAM roles or service accounts.
- Dependency Scanning: Regularly scan
requirements.txtfor known vulnerabilities using tools like Snyk or Dependabot to mitigate supply chain risks. - Prompt Injection Defense: Design prompts to be resilient against malicious attempts to hijack the agent's goal. Techniques include clear role definition, input sanitization, and potentially using a separate LLM to guardrail inputs.
Observability
- Structured Logging: As implemented in Phase 5, use structured (JSON) logging. This makes logs easily parsable by log aggregation tools (e.g., ELK Stack, Grafana Loki, Datadog), enabling efficient searching, filtering, and alerting. Ensure logs capture key events: agent start/end, tool calls, LLM decisions, errors, and state transitions.
- Application Performance Monitoring (APM): Integrate with an APM tool (e.g., OpenTelemetry, Datadog, New Relic) to collect metrics like latency, error rates, and resource utilization. This provides real-time insights into agent performance and health.
- Tracing: Implement distributed tracing (e.g., with OpenTelemetry) to visualize the flow of a single research request through the LangGraph nodes, LLM calls, and tool executions. This is crucial for debugging complex, multi-step agentic workflows and identifying performance bottlenecks.
- Alerting: Set up alerts on critical metrics (e.g., high error rates, increased latency,
num_stepsexceeding thresholds, LLM token usage spikes) to notify operations teams proactively of issues.
Reliability
- Retry Mechanisms: Implement robust retry logic with exponential backoff for external API calls (LLMs, web search). This is crucial for handling transient network issues and API rate limits, as demonstrated in our
retry_with_exponential_backoffdecorator. - Circuit Breakers: Consider implementing circuit breakers for external services to prevent cascading failures when a dependency is consistently unhealthy. This can temporarily stop calling a failing service to allow it to recover.
- Redundancy and High Availability: Deploy multiple replicas of the agent (e.g., 2+ Kubernetes pods) behind a load balancer to ensure availability if one instance fails. Kubernetes probes (
livenessProbe,readinessProbe) are essential here. - Rate Limiting (Internal): Implement internal rate limiting for LLM calls if your budget or provider limits require it, preventing excessive token consumption or exceeding API quotas.
- Graceful Shutdown: Ensure the application can gracefully shut down, completing ongoing requests before terminating, especially important in containerized environments.
Rate Limiting
- API Gateway Rate Limiting: Deploy an API Gateway (e.g., Nginx, AWS API Gateway, Azure API Management) in front of the FastAPI service to enforce global rate limits on incoming requests, protecting the backend from overload and abuse.
- LLM Provider Limits: Be aware of and respect your LLM provider's rate limits. The exponential backoff on
web_searchhelps, but for LLM calls, you might need a dedicated token bucket or leaky bucket algorithm if facing high concurrency.
Authentication
- API Key / OAuth: For external access, secure the FastAPI endpoint with API keys or OAuth 2.0. This prevents unauthorized usage of your research agent. (Not implemented in this guide for brevity, but critical for production).
Governance
- Audit Logs: Enhance logging to include audit trails for who initiated a research request and when, important for compliance and accountability.
- Content Moderation: Implement content moderation for both input queries and generated reports to prevent misuse or generation of harmful content. This could involve using moderation APIs from LLM providers or custom solutions.
Cost Optimisation
- Token Usage Monitoring: Continuously monitor and log token usage for LLM calls. This data is critical for understanding operational costs and identifying inefficient prompts or agent behaviors.
- Model Selection: Evaluate if smaller, cheaper LLMs can achieve sufficient quality for specific sub-tasks or less critical queries. Use
temperaturesettings judiciously (lower for factual, higher for creative). - Caching: Implement caching for frequently requested research topics or repeated tool calls to reduce LLM and external API costs. For example, cache
web_searchresults for a short TTL.
Compliance
- Data Privacy: Ensure that no personally identifiable information (PII) or sensitive data is inadvertently processed, stored, or exposed by the agent. If PII must be handled, implement strict data anonymization, encryption, and access controls.
- Regional Deployment: Deploy the agent in data centers compliant with regional data residency requirements (e.g., GDPR, CCPA) if processing sensitive data.
Learning Objectives
- Implement the ReAct pattern for sophisticated reasoning and tool use within an AI agent.
- Design and build a multi-step agent workflow using LangGraph's state machine capabilities.
- Integrate external tools, such as web search, into an agent's operational loop.
- Handle and resolve contradictions or conflicting information from multiple sources within an agent's reasoning process.
- Expose an AI agent's functionality as a robust and scalable FastAPI service.
- Implement basic observability and error handling for production-ready AI agents.
- Containerize an AI agent application using Docker for consistent deployment.
Conceptual Prerequisites
What Is Agentic AI?
Understanding the core concepts of agentic AI, including autonomy, reasoning, and tool use, is fundamental to grasping why this architecture is effective.
ReAct Pattern: Reasoning and Acting in Agents
The ReAct pattern is central to this agent's intelligence. A solid understanding of how agents reason, plan, and act with tools is crucial.
State Machines for AI Agents
LangGraph is a state machine framework. Familiarity with state transitions, nodes, and edges will help you understand the agent's flow.
Prompt Engineering for AI Agents
Effective agent behavior relies heavily on well-crafted prompts. Knowing how to structure prompts for reasoning and tool use will be beneficial.
Technical Prerequisites
Python 3.9+
All code examples are in Python. Basic proficiency in Python programming, including virtual environments and package management, is required.
Docker
To containerize and deploy the application, you'll need Docker installed and a basic understanding of Dockerfiles and images.
OpenAI API Key (or equivalent)
Access to a large language model API (e.g., OpenAI, Anthropic, Google) is necessary for the agent's intelligence. Ensure you have an API key configured.
Tavily API Key (or equivalent web search API)
The agent uses a web search tool. A Tavily API key is recommended for robust search capabilities, or you can adapt to another search API.