Capstone Project: AI Agent System

1# requirements.txt
2# langchain==0.2.0
3# langchain-openai==0.1.0
4# langgraph==0.2.0
5# python-dotenv==1.0.0
6# tavily-python==0.4.0  # Web search
7
8import os
9from dotenv import load_dotenv
10load_dotenv()
11
12from typing import TypedDict, Annotated
13from langgraph.graph import StateGraph, END
14from langgraph.checkpoint.memory import MemorySaver
15from langchain_openai import ChatOpenAI
16from langchain_core.prompts import ChatPromptTemplate
17import operator

1class ResearchState(TypedDict):
2    messages: Annotated[list, operator.add]
3    
4    # Input
5    topic: str
6    requirements: str
7    
8    # Research
9    search_queries: list
10    search_results: list
11    
12    # Analysis
13    key_findings: str
14    
15    # Writing
16    outline: str
17    draft: str
18    
19    # Review
20    review_feedback: str
21    revision_count: int
22    max_revisions: int
23    
24    # Control
25    next_agent: str
26    final_report: str

1from langchain_core.tools import tool
2
3@tool
4def web_search(query: str) -> str:
5    """Search the web for information.
6    
7    Args:
8        query: Search query
9    """
10    # Use Tavily or similar
11    from tavily import TavilyClient
12    client = TavilyClient()
13    results = client.search(query, max_results=3)
14    
15    output = ""
16    for r in results.get("results", []):
17        output += f"Title: {r['title']}\n"
18        output += f"Content: {r['content'][:300]}\n"
19        output += f"URL: {r['url']}\n\n"
20    
21    return output
22
23llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)
24
25def search_agent(state):
26    """Search agent: generate queries and search."""
27    topic = state["topic"]
28    
29    # Generate search queries
30    queries_result = llm.invoke(
31        f"Generate 3 specific search queries to research:\n{topic}\n\n"
32        f"Requirements: {state.get('requirements', 'comprehensive overview')}\n"
33        f"Return each query on a new line."
34    ).content
35    
36    queries = [q.strip() for q in queries_result.strip().split("\n") if q.strip()][:3]
37    
38    # Execute searches
39    all_results = []
40    for query in queries:
41        try:
42            result = web_search.invoke(query)
43            all_results.append(f"Query: {query}\n{result}")
44        except Exception as e:
45            all_results.append(f"Query: {query}\nError: {str(e)}")
46    
47    return {
48        "search_queries": queries,
49        "search_results": all_results,
50        "messages": [{"role": "search_agent", "content": f"Searched {len(queries)} queries"}]
51    }

1def analysis_agent(state):
2    """Analyze search results and extract key findings."""
3    results = "\n\n---\n\n".join(state.get("search_results", []))
4    
5    findings = llm.invoke(
6        f"Analyze these search results about: {state['topic']}\n\n"
7        f"Results:\n{results}\n\n"
8        f"Extract:\n"
9        f"1. Key facts and statistics\n"
10        f"2. Main trends or themes\n"
11        f"3. Different perspectives\n"
12        f"4. Notable sources\n"
13        f"Be specific and cite data points."
14    ).content
15    
16    return {
17        "key_findings": findings,
18        "messages": [{"role": "analysis_agent", "content": "Analysis complete"}]
19    }

1def writer_agent(state):
2    """Write the research report."""
3    feedback = state.get("review_feedback", "")
4    
5    if state.get("draft") and feedback:
6        # Revision mode
7        prompt = (
8            f"Revise this report based on feedback:\n\n"
9            f"Current draft:\n{state['draft']}\n\n"
10            f"Feedback:\n{feedback}\n\n"
11            f"Provide the complete revised report."
12        )
13    else:
14        # First draft
15        prompt = (
16            f"Write a comprehensive research report.\n\n"
17            f"Topic: {state['topic']}\n"
18            f"Requirements: {state.get('requirements', 'comprehensive overview')}\n\n"
19            f"Key Findings:\n{state.get('key_findings', '')}\n\n"
20            f"Report structure:\n"
21            f"1. Executive Summary\n"
22            f"2. Introduction\n"
23            f"3. Key Findings (3-5 sections)\n"
24            f"4. Analysis\n"
25            f"5. Conclusion\n"
26            f"6. Sources\n\n"
27            f"Write in Vietnamese. Be thorough and well-structured."
28        )
29    
30    draft = llm.invoke(prompt).content
31    
32    return {
33        "draft": draft,
34        "messages": [{"role": "writer_agent", "content": "Draft complete"}]
35    }

1def review_agent(state):
2    """Review the draft and provide feedback."""
3    review = llm.invoke(
4        f"Critically review this research report:\n\n"
5        f"Topic: {state['topic']}\n"
6        f"Report:\n{state.get('draft', '')}\n\n"
7        f"Evaluate:\n"
8        f"1. Accuracy: Are claims supported by evidence?\n"
9        f"2. Completeness: Are all aspects covered?\n"
10        f"3. Structure: Is it well-organized?\n"
11        f"4. Clarity: Is it easy to understand?\n"
12        f"5. Citations: Are sources mentioned?\n\n"
13        f"Rate: EXCELLENT, GOOD, or NEEDS_WORK\n"
14        f"Provide specific feedback for improvement."
15    ).content
16    
17    return {
18        "review_feedback": review,
19        "revision_count": state.get("revision_count", 0) + 1,
20        "messages": [{"role": "review_agent", "content": review}]
21    }

1def supervisor(state):
2    """Orchestrate the research workflow."""
3    # Determine what to do next
4    has_search = bool(state.get("search_results"))
5    has_findings = bool(state.get("key_findings"))
6    has_draft = bool(state.get("draft"))
7    has_review = bool(state.get("review_feedback"))
8    revision_count = state.get("revision_count", 0)
9    max_revisions = state.get("max_revisions", 2)
10    
11    if not has_search:
12        next_agent = "search"
13    elif not has_findings:
14        next_agent = "analyze"
15    elif not has_draft:
16        next_agent = "write"
17    elif not has_review or (has_review and revision_count < max_revisions):
18        if has_review and "EXCELLENT" in state.get("review_feedback", ""):
19            next_agent = "finalize"
20        elif has_review and "GOOD" in state.get("review_feedback", ""):
21            next_agent = "finalize"
22        elif not has_review:
23            next_agent = "review"
24        else:
25            # Needs work and can still revise
26            next_agent = "write"
27    else:
28        next_agent = "finalize"
29    
30    return {"next_agent": next_agent}
31
32def finalize(state):
33    """Finalize the report."""
34    return {
35        "final_report": state.get("draft", ""),
36        "messages": [{"role": "supervisor", "content": "Report finalized!"}]
37    }
38
39def route(state):
40    return state.get("next_agent", "search")

1# Assemble the multi-agent graph
2graph = StateGraph(ResearchState)
3
4# Add nodes
5graph.add_node("supervisor", supervisor)
6graph.add_node("search", search_agent)
7graph.add_node("analyze", analysis_agent)
8graph.add_node("write", writer_agent)
9graph.add_node("review", review_agent)
10graph.add_node("finalize", finalize)
11
12# Routing
13graph.set_entry_point("supervisor")
14graph.add_conditional_edges("supervisor", route, {
15    "search": "search",
16    "analyze": "analyze",
17    "write": "write",
18    "review": "review",
19    "finalize": "finalize"
20})
21
22# All agents return to supervisor
23graph.add_edge("search", "supervisor")
24graph.add_edge("analyze", "supervisor")
25graph.add_edge("write", "supervisor")
26graph.add_edge("review", "supervisor")
27graph.add_edge("finalize", END)
28
29# Compile with checkpointing
30checkpointer = MemorySaver()
31research_app = graph.compile(checkpointer=checkpointer)

1def research(topic, requirements=""):
2    """Run the research assistant."""
3    config = {"configurable": {"thread_id": f"research_{hash(topic)}"}}
4    
5    result = research_app.invoke({
6        "topic": topic,
7        "requirements": requirements,
8        "messages": [],
9        "revision_count": 0,
10        "max_revisions": 2,
11        "search_results": [],
12        "search_queries": [],
13    }, config=config)
14    
15    # Print workflow
16    print("=== Workflow Log ===")
17    for msg in result.get("messages", []):
18        role = msg.get("role", "unknown")
19        content = msg.get("content", "")[:100]
20        print(f"  [{role}] {content}")
21    
22    print(f"\n=== Final Report ===")
23    print(result.get("final_report", "No report generated"))
24    
25    return result
26
27# Example run
28result = research(
29    topic="AI adoption in Vietnamese enterprises 2025",
30    requirements="Focus on SMEs, include statistics, compare with SEA region"
31)