Building an AI Research Agent with LangGraph, Claude, and AWS

Browser
  │
  ▼
CloudFront (HTTPS)
  ├── /* ──────────────► S3 (React static files)
  └── /api/v1/* ───────► ALB
                           │
                           ▼
                      ECS Fargate (FastAPI)
                           │
                           ▼
                      LangGraph ReAct Agent
                           │
              ┌────────────┼────────────┐
              ▼            ▼            ▼
        Tavily API    Wikipedia API  arXiv API
      (web search)   (background    (academic
                       facts)        papers)
              └────────────┼────────────┘
                           ▼
                     Anthropic Claude
                   (reasoning + answer)
                           │
                           ▼
                   SSE Stream → Browser
               (tokens + tool trace events)
                           │
                           ▼
                       DynamoDB
                  (conversation history)

# agent/graph.py
from langgraph.graph import StateGraph, END
from langgraph.prebuilt import ToolNode

def build_agent(tools: list) -> CompiledGraph:
    model = ChatAnthropic(model=settings.model_name).bind_tools(tools)

    def call_model(state: AgentState) -> dict:
        messages = [SystemMessage(content=SYSTEM_PROMPT)] + state["messages"]
        response = model.invoke(messages)
        return {"messages": [response]}

    def should_continue(state: AgentState) -> str:
        last = state["messages"][-1]
        return "tools" if last.tool_calls else END

    graph = StateGraph(AgentState)
    graph.add_node("agent", call_model)
    graph.add_node("tools", ToolNode(tools))
    graph.set_entry_point("agent")
    graph.add_conditional_edges("agent", should_continue)
    graph.add_edge("tools", "agent")
    return graph.compile()

# agent/state.py
from langgraph.graph import add_messages
from typing import Annotated
from langchain_core.messages import BaseMessage

class AgentState(TypedDict):
    messages: Annotated[list[BaseMessage], add_messages]

# agent/tools.py
from langchain_community.tools.tavily_search import TavilySearchResults

def get_tools() -> list:
    web_search = TavilySearchResults(
        max_results=5,
        description="Search the web for current information, news, and general knowledge."
    )
    ...

@tool
def wikipedia_search(query: str) -> str:
    """Search Wikipedia for background information on a topic."""
    wiki = wikipediaapi.Wikipedia(language="en", user_agent="ai-research-agent/1.0")
    page = wiki.page(query)
    if not page.exists():
        return f"No Wikipedia article found for: {query}"
    # Truncate to 3000 chars to stay within Claude's context budget
    return page.summary[:3000]

@tool
def arxiv_search(query: str) -> str:
    """Search arXiv for academic papers and research on a topic."""
    import arxiv
    results = arxiv.Search(query=query, max_results=3,
                           sort_by=arxiv.SortCriterion.Relevance)
    papers = []
    for r in arxiv.Client().results(results):
        papers.append(f"Title: {r.title}\nAuthors: {', '.join(a.name for a in r.authors)}\n"
                      f"Summary: {r.summary[:500]}\nURL: {r.entry_id}")
    return "\n\n---\n\n".join(papers) if papers else "No papers found."

data: {"type": "conversation_id", "conversation_id": "abc-123"}

data: {"type": "trace", "step": "tool_start", "tool": "web_search", "input": "LangGraph tutorial"}

data: {"type": "trace", "step": "tool_end", "tool": "web_search", "output": "LangGraph is..."}

data: {"type": "token", "content": "LangGraph is a "}

data: {"type": "token", "content": "framework for building..."}

data: {"type": "done"}

# routers/chat.py
async def stream_agent_response(message: str, conversation_id: str):
    history = await dynamodb_service.get_messages(conversation_id)
    state = {"messages": history + [HumanMessage(content=message)]}

    async for event in agent.astream_events(state, version="v2"):
        kind = event["event"]

        if kind == "on_tool_start":
            yield f'data: {{"type":"trace","step":"tool_start","tool":"{event["name"]}","input":"{event["data"].get("input","")[:200]}"}}\n\n'

        elif kind == "on_tool_end":
            yield f'data: {{"type":"trace","step":"tool_end","tool":"{event["name"]}","output":"{str(event["data"].get("output",""))[:300]}"}}\n\n'

        elif kind == "on_chat_model_stream":
            content = event["data"]["chunk"].content
            if content:
                yield f'data: {json.dumps({"type": "token", "content": content})}\n\n'

    yield 'data: {"type": "done"}\n\n'

# services/dynamodb_service.py
async def save_messages(conversation_id: str, messages: list[BaseMessage]) -> None:
    with table.batch_writer() as batch:
        for i, msg in enumerate(messages):
            batch.put_item(Item={
                "pk": f"CONV#{conversation_id}",
                "sk": f"MSG#{i:06d}",
                "entity_type": "message",
                "role": msg.type,
                "content": msg.content,
                "updated_at": datetime.utcnow().isoformat(),
            })

# stacks/backend_stack.py (simplified)
fargate_service = ecs_patterns.ApplicationLoadBalancedFargateService(
    self, "BackendService",
    cluster=cluster,
    cpu=512,
    memory_limit_mib=1024,
    task_image_options=ecs_patterns.ApplicationLoadBalancedTaskImageOptions(
        image=ecs.ContainerImage.from_ecr_repository(ecr_repo),
        environment={"DYNAMODB_TABLE": table.table_name},
        secrets={
            "ANTHROPIC_API_KEY": ecs.Secret.from_ssm_parameter(anthropic_key),
            "TAVILY_API_KEY": ecs.Secret.from_ssm_parameter(tavily_key),
        },
    ),
)
# Long agent reasoning chains need more than the default 60s
fargate_service.load_balancer.set_attribute("idle_timeout.timeout_seconds", "120")

# stacks/frontend_stack.py (simplified)
api_behaviour = cloudfront.BehaviorOptions(
    origin=alb_origin,
    viewer_protocol_policy=cloudfront.ViewerProtocolPolicy.REDIRECT_TO_HTTPS,
    cache_policy=cloudfront.CachePolicy.CACHING_DISABLED,
    origin_request_policy=cloudfront.OriginRequestPolicy.ALL_VIEWER,
    # SSE streams can run for 60-120s — increase read timeout
    # Default is 30s; streaming connections silently drop without this
)

cd backend
pytest                          # all tests
pytest tests/agent/             # ReAct agent + tools
pytest tests/routers/           # API endpoints + SSE format
pytest --cov=app                # with coverage report

cd frontend
npm test
npm run test:coverage