Update rag_langgraph.py

rag_langgraph.py

@@ -1,163 +1,142 @@
-import functools, operator
-from IPython.display import Image, display
-from langchain_core.messages import (
-    AIMessage,
-    BaseMessage,
-    ToolMessage,
-    HumanMessage,
-)
-from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
-from langchain_core.tools import tool
+from typing import Annotated, List, Tuple, Union
+
 from langchain_community.tools.tavily_search import TavilySearchResults
-from langchain_experimental.utilities import PythonREPL
+from langchain_core.tools import tool
+from langchain_experimental.tools import PythonREPLTool
+
+from langchain.agents import AgentExecutor, create_openai_tools_agent
+from langchain_core.messages import BaseMessage, HumanMessage
 from langchain_openai import ChatOpenAI
-from langgraph.graph import END, StateGraph
-from langgraph.prebuilt import ToolNode
-from typing import Annotated, Literal, Sequence, TypedDict
-from typing_extensions import TypedDict
 
-def create_agent(llm, tools, system_message: str):
-    """Create an agent."""
+from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
+from langchain_core.output_parsers.openai_functions import JsonOutputFunctionsParser
+
+import operator
+from typing import Annotated, Any, Dict, List, Optional, Sequence, TypedDict
+import functools
+
+from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
+from langgraph.graph import StateGraph, END
+
+def create_agent(llm: ChatOpenAI, tools: list, system_prompt: str):
+    # Each worker node will be given a name and some tools.
     prompt = ChatPromptTemplate.from_messages(
         [
             (
                 "system",
-                "You are a helpful AI assistant, collaborating with other assistants."
-                " Use the provided tools to progress towards answering the question."
-                " If you are unable to fully answer, that's OK, another assistant with different tools "
-                " will help where you left off. Execute what you can to make progress."
-                " If you or any of the other assistants have the final answer or deliverable,"
-                " prefix your response with FINAL ANSWER so the team knows to stop."
-                " You have access to the following tools: {tool_names}.\n{system_message}",
+                system_prompt,
             ),
             MessagesPlaceholder(variable_name="messages"),
+            MessagesPlaceholder(variable_name="agent_scratchpad"),
         ]
     )
-    prompt = prompt.partial(system_message=system_message)
-    prompt = prompt.partial(tool_names=", ".join([tool.name for tool in tools]))
-    return prompt | llm.bind_tools(tools)
+    agent = create_openai_tools_agent(llm, tools, prompt)
+    executor = AgentExecutor(agent=agent, tools=tools)
+    return executor
 
-@tool
-def python_repl(code: Annotated[str, "The python code to execute to generate your chart."]):
-    """Use this to execute python code. If you want to see the output of a value,
-    you should print it out with `print(...)`. This is visible to the user."""
-    try:
-        result = repl.run(code)
-    except BaseException as e:
-        return f"Failed to execute. Error: {repr(e)}"
-    result_str = f"Successfully executed:\n```python\n{code}\n```\nStdout: {result}"
-    return (
-        result_str + "\n\nIf you have completed all tasks, respond with FINAL ANSWER."
-    )
-
-# This defines the object that is passed between each node
-# in the graph. We will create different nodes for each agent and tool
-class AgentState(TypedDict):
-    messages: Annotated[Sequence[BaseMessage], operator.add]
-    sender: str
-
-# Helper function to create a node for a given agent
 def agent_node(state, agent, name):
     result = agent.invoke(state)
-    # We convert the agent output into a format that is suitable to append to the global state
-    if isinstance(result, ToolMessage):
-        pass
-    else:
-        result = AIMessage(**result.dict(exclude={"type", "name"}), name=name)
-    return {
-        "messages": [result],
-        # Since we have a strict workflow, we can
-        # track the sender so we know who to pass to next.
-        "sender": name,
-    }
-
-def router(state) -> Literal["call_tool", "__end__", "continue"]:
-    # This is the router
-    messages = state["messages"]
-    last_message = messages[-1]
-    if last_message.tool_calls:
-        # The previous agent is invoking a tool
-        return "call_tool"
-    if "FINAL ANSWER" in last_message.content:
-        # Any agent decided the work is done
-        return "__end__"
-    return "continue"
+    return {"messages": [HumanMessage(content=result["output"], name=name)]}
 
+# The agent state is the input to each node in the graph
+class AgentState(TypedDict):
+    # The annotation tells the graph that new messages will always
+    # be added to the current states
+    messages: Annotated[Sequence[BaseMessage], operator.add]
+    # The 'next' field indicates where to route to next
+    next: str
+    
 def run_multi_agent(prompt):
     tavily_tool = TavilySearchResults(max_results=5)
     repl = PythonREPL()
 
-    llm = ChatOpenAI(model="gpt-4o")
-
-    # Research agent and node
-    research_agent = create_agent(
-        llm,
-        [tavily_tool],
-        system_message="You should provide accurate data for the chart_generator to use.",
+    members = ["Researcher", "Coder"]
+    system_prompt = (
+        "You are a supervisor 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,"
+        " respond with FINISH."
     )
+    # Our team supervisor is an LLM node. It just picks the next agent to process
+    # and decides when the work is completed
+    options = ["FINISH"] + members
+    # Using openai function calling can make output parsing easier for us
+    function_def = {
+        "name": "route",
+        "description": "Select the next role.",
+        "parameters": {
+            "title": "routeSchema",
+            "type": "object",
+            "properties": {
+                "next": {
+                    "title": "Next",
+                    "anyOf": [
+                        {"enum": options},
+                    ],
+                }
+            },
+            "required": ["next"],
+        },
+    }
+    prompt = ChatPromptTemplate.from_messages(
+        [
+            ("system", system_prompt),
+            MessagesPlaceholder(variable_name="messages"),
+            (
+                "system",
+                "Given the conversation above, who should act next?"
+                " Or should we FINISH? Select one of: {options}",
+            ),
+        ]
+    ).partial(options=str(options), members=", ".join(members))
+    
+    llm = ChatOpenAI(model="gpt-4-1106-preview")
+    
+    supervisor_chain = (
+        prompt
+        | llm.bind_functions(functions=[function_def], function_call="route")
+        | JsonOutputFunctionsParser()
+    )
+
+    research_agent = create_agent(llm, [tavily_tool], "You are a web researcher.")
     research_node = functools.partial(agent_node, agent=research_agent, name="Researcher")
     
-    # chart_generator
-    chart_agent = create_agent(
+    # NOTE: THIS PERFORMS ARBITRARY CODE EXECUTION. PROCEED WITH CAUTION
+    code_agent = create_agent(
         llm,
-        [python_repl],
-        system_message="Any charts you display will be visible by the user.",
+        [python_repl_tool],
+        "You may generate safe python code to analyze data and generate charts using matplotlib.",
     )
-    chart_node = functools.partial(agent_node, agent=chart_agent, name="chart_generator")
-
-    tools = [tavily_tool, python_repl]
-    tool_node = ToolNode(tools)
-
-    workflow = StateGraph(AgentState)
+    code_node = functools.partial(agent_node, agent=code_agent, name="Coder")
     
+    workflow = StateGraph(AgentState)
     workflow.add_node("Researcher", research_node)
-    workflow.add_node("chart_generator", chart_node)
-    workflow.add_node("call_tool", tool_node)
-    
-    workflow.add_conditional_edges(
-        "Researcher",
-        router,
-        {"continue": "chart_generator", "call_tool": "call_tool", "__end__": END},
-    )
-    workflow.add_conditional_edges(
-        "chart_generator",
-        router,
-        {"continue": "Researcher", "call_tool": "call_tool", "__end__": END},
-    )
+    workflow.add_node("Coder", code_node)
+    workflow.add_node("supervisor", supervisor_chain)
+
+    for member in members:
+        # We want our workers to ALWAYS "report back" to the supervisor when done
+        workflow.add_edge(member, "supervisor")
+    # The supervisor populates the "next" field in the graph state
+    # which routes to a node or finishes
+    conditional_map = {k: k for k in members}
+    conditional_map["FINISH"] = END
+    workflow.add_conditional_edges("supervisor", lambda x: x["next"], conditional_map)
+    # Finally, add entrypoint
+    workflow.set_entry_point("supervisor")
     
-    workflow.add_conditional_edges(
-        "call_tool",
-        # Each agent node updates the 'sender' field
-        # the tool calling node does not, meaning
-        # this edge will route back to the original agent
-        # who invoked the tool
-        lambda x: x["sender"],
-        {
-            "Researcher": "Researcher",
-            "chart_generator": "chart_generator",
-        },
-    )
-    workflow.set_entry_point("Researcher")
     graph = workflow.compile()
 
-    try:
-        display(Image(graph.get_graph(xray=True).draw_mermaid_png()))
-    except:
-        # This requires some extra dependencies and is optional
-        pass
-    
-    events = graph.stream(
+    for s in graph.stream(
         {
             "messages": [
-                HumanMessage(
-                    content=prompt
-                )
-            ],
-        },
-        # Maximum number of steps to take in the graph
-        {"recursion_limit": 150},
-    )
-    for s in events:
-        print(s)
-        print("----")
+                HumanMessage(content="Code hello world and print it to the terminal")
+            ]
+        }
+    ):
+        if "__end__" not in s:
+            print(s)
+            print("----")
+        
     return "DONE"