Update rag_langgraph.py

rag_langgraph.py

@@ -17,6 +17,8 @@ import functools
 from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
 from langgraph.graph import StateGraph, END
 
+LLM = "gpt-4o"
+
 class AgentState(TypedDict):
     messages: Annotated[Sequence[BaseMessage], operator.add]
     next: str
@@ -24,10 +26,7 @@ class AgentState(TypedDict):
 def create_agent(llm: ChatOpenAI, tools: list, system_prompt: str):
     prompt = ChatPromptTemplate.from_messages(
         [
-            (
-                "system", 
-                system_prompt
-            ),
+            ("system", system_prompt),
             MessagesPlaceholder(variable_name="messages"),
             MessagesPlaceholder(variable_name="agent_scratchpad"),
         ]
@@ -46,7 +45,7 @@ def create_graph(topic, word_count):
     members = ["Blogger"]
     
     system_prompt = (
-        "You are a supervisor tasked with managing a conversation between the"
+        "You are a manager tasked with managing a conversation between the"
         " following workers: {members}. Given the following user request,"
         " respond with the worker to act next. Each worker will perform a"
         " task and respond with their results and status. When finished,"
@@ -85,7 +84,7 @@ def create_graph(topic, word_count):
         ]
     ).partial(options=str(options), members=", ".join(members))
     
-    llm = ChatOpenAI(model="gpt-4o")
+    llm = ChatOpenAI(model=LLM)
     
     supervisor_chain = (
         prompt
@@ -93,14 +92,14 @@ def create_graph(topic, word_count):
         | JsonOutputFunctionsParser()
     )
 
-    #research_agent = create_agent(llm, [tavily_tool], f"Research content on topic {topic}")
-    #research_node = functools.partial(agent_node, agent=research_agent, name="Researcher")
+    research_agent = create_agent(llm, [tavily_tool], f"Research content on topic {topic}, prioritizing research papers.")
+    research_node = functools.partial(agent_node, agent=research_agent, name="Researcher")
     
-    blogger_agent = create_agent(llm, [tavily_tool], f"Based on research papers, write a {word_count}-word blog post on topic {topic}. Add a references section.")
+    blogger_agent = create_agent(llm, f"Write a {word_count}-word blog post on topic {topic}. Add a references section with research papers.")
     blogger_node = functools.partial(agent_node, agent=blogger_agent, name="Blogger")
     
     workflow = StateGraph(AgentState)
-    #workflow.add_node("Researcher", research_node)
+    workflow.add_node("Researcher", research_node)
     workflow.add_node("Blogger", blogger_node)
     workflow.add_node("Manager", supervisor_chain)
 
@@ -119,16 +118,7 @@ def run_multi_agent(topic, word_count):
     graph = create_graph(topic, word_count)
     result = graph.invoke({
         "messages": [
-            HumanMessage(content="Evolution of Retrieval-Augmented Generation from Naive RAG to Agentic RAG")
+            HumanMessage(content=topic)
         ]
     })
-    print("###")
-    print(result)
-    print("###")
-    print(result['messages'])
-    print("###")
-    print(result['messages'][1])
-    print("###")
-    print(result['messages'][1].content)
-    print("###")
     return result['messages'][1].content