Update app.py

app.py

@@ -14,172 +14,207 @@ _set_if_undefined("TAVILY_API_KEY")
 os.environ["LANGCHAIN_TRACING_V2"] = "true"
 os.environ["LANGCHAIN_PROJECT"] = "Multi-agent Collaboration"
 
-from typing import Annotated, List, Tuple, Union
-
-from langchain_community.tools.tavily_search import TavilySearchResults
-from langchain_core.tools import tool
-from langchain_experimental.tools import PythonREPLTool
-
-tavily_tool = TavilySearchResults(max_results=5)
-
-# This executes code locally, which can be unsafe
-python_repl_tool = PythonREPLTool()
+from langchain_core.messages import (
+    BaseMessage,
+    ToolMessage,
+    HumanMessage,
+)
+from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
+from langgraph.graph import END, StateGraph
 
-from langchain.agents import AgentExecutor, create_openai_tools_agent
-from langchain_core.messages import BaseMessage, HumanMessage
-from langchain_openai import ChatOpenAI
 
-def create_agent(llm: ChatOpenAI, tools: list, system_prompt: str):
-    # Each worker node will be given a name and some tools.
+def create_agent(llm, tools, system_message: str):
+    """Create an agent."""
     prompt = ChatPromptTemplate.from_messages(
         [
             (
                 "system",
-                system_prompt,
+                "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}",
             ),
             MessagesPlaceholder(variable_name="messages"),
-            MessagesPlaceholder(variable_name="agent_scratchpad"),
         ]
     )
-    agent = create_openai_tools_agent(llm, tools, prompt)
-    executor = AgentExecutor(agent=agent, tools=tools)
-    return executor
-
-def agent_node(state, agent, name):
-    result = agent.invoke(state)
-    return {"messages": [HumanMessage(content=result["output"], name=name)]}
+    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)
 
-from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
-from langchain_core.output_parsers.openai_functions import JsonOutputFunctionsParser
-
-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))
+from langchain_core.tools import tool
+from typing import Annotated
+from langchain_experimental.utilities import PythonREPL
+from langchain_community.tools.tavily_search import TavilySearchResults
 
-llm = ChatOpenAI(model="gpt-4-1106-preview")
+tavily_tool = TavilySearchResults(max_results=5)
 
-supervisor_chain = (
-    prompt
-    | llm.bind_functions(functions=[function_def], function_call="route")
-    | JsonOutputFunctionsParser()
-)
+# Warning: This executes code locally, which can be unsafe when not sandboxed
+
+repl = PythonREPL()
+
+@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."
+    )
 
 import operator
-from typing import Annotated, Any, Dict, List, Optional, Sequence, TypedDict
-import functools
+from typing import Annotated, Sequence, TypedDict
 
-from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
-from langgraph.graph import StateGraph, END
+from langchain_openai import ChatOpenAI
+from typing_extensions import TypedDict
 
 
-# The agent state is the input to each node in the graph
+# 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):
-    # 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
+    sender: str
 
+import functools
+from langchain_core.messages import AIMessage
 
-research_agent = create_agent(llm, [tavily_tool], "You are a web researcher.")
+
+# 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,
+    }
+
+
+llm = ChatOpenAI(model="gpt-4-1106-preview")
+
+# Research agent and node
+research_agent = create_agent(
+    llm,
+    [tavily_tool],
+    system_message="You should provide accurate data for the chart_generator to use.",
+)
 research_node = functools.partial(agent_node, agent=research_agent, name="Researcher")
 
-# NOTE: THIS PERFORMS ARBITRARY CODE EXECUTION. PROCEED WITH CAUTION
-code_agent = create_agent(
+# chart_generator
+chart_agent = create_agent(
     llm,
-    [python_repl_tool],
-    "You may generate safe python code to analyze data and generate charts using matplotlib.",
+    [python_repl],
+    system_message="Any charts you display will be visible by the user.",
 )
-code_node = functools.partial(agent_node, agent=code_agent, name="Coder")
+chart_node = functools.partial(agent_node, agent=chart_agent, name="chart_generator")
+
+from langgraph.prebuilt import ToolNode
+
+tools = [tavily_tool, python_repl]
+tool_node = ToolNode(tools)
+
+# Either agent can decide to end
+from typing import Literal
+
+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"
 
 workflow = StateGraph(AgentState)
+
 workflow.add_node("Researcher", research_node)
-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_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_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()
 
+from IPython.display import Image, display
+
+try:
+    display(Image(graph.get_graph(xray=True).draw_mermaid_png()))
+except:
+    # This requires some extra dependencies and is optional
+    pass
+
+events = graph.stream(
+    {
+        "messages": [
+            HumanMessage(
+                content="Fetch the UK's GDP over the past 5 years,"
+                " then draw a line graph of it."
+                " Once you code it up, finish."
+            )
+        ],
+    },
+    # Maximum number of steps to take in the graph
+    {"recursion_limit": 150},
+)
+for s in events:
+    print(s)
+    print("----")
+
 ###
 
-def invoke(openai_api_key, topic, word_count=500):
+def invoke(openai_api_key):
     if (openai_api_key == ""):
         raise gr.Error("OpenAI API Key is required.")
-    if (topic == ""):
-        raise gr.Error("Topic is required.")
         
-    #agentops.init(os.environ["AGENTOPS_API_KEY"])
-
     os.environ["OPENAI_API_KEY"] = openai_api_key
-
-    for s in graph.stream(
-        {
-            "messages": [
-                HumanMessage(content="Code hello world and print it to the terminal")
-            ]
-        }
-    ):
-        if "__end__" not in s:
-            print(s)
-            print("----")
-    
-    return result
+   
+    return "TODO"
 
 gr.close_all()
 
 demo = gr.Interface(fn = invoke, 
                     inputs = [gr.Textbox(label = "OpenAI API Key", type = "password", lines = 1),
-                              gr.Textbox(label = "Topic", value="TODO", lines = 1),
-                              gr.Number(label = "Word Count", value=1000, minimum=500, maximum=5000)],
-                    outputs = [gr.Markdown(label = "Generated Blog Post", value="TODO")],
-                    title = "Multi-Agent RAG: Blog Post Generation",
+                              gr.Textbox(label = "TODO", value="TODO", lines = 1)],
+                    outputs = [gr.Markdown(label = "TODO", value="TODO")],
+                    title = "Multi-Agent RAG: Chart Generation",
                     description = "TODO")
 
 demo.launch()