Update app.py

app.py

@@ -1,30 +1,16 @@
 import os
-import json
-import re
 import base64
-import streamlit as st
 from io import BytesIO
-from langchain_core.utils.function_calling import convert_to_openai_function
-from langchain_core.messages import (
-    AIMessage,
-    BaseMessage,
-    ChatMessage,
-    FunctionMessage,
-    HumanMessage,
-)
-from langchain.tools.render import format_tool_to_openai_function
-from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
-from langgraph.graph import END, StateGraph
-from langgraph.prebuilt.tool_executor import ToolExecutor, ToolInvocation
+
+import streamlit as st
+import matplotlib.pyplot as plt
 from langchain_core.tools import tool
+from langchain_core.utils.function_calling import convert_to_openai_function
+from langchain_openai import ChatOpenAI
 from langchain_community.tools.tavily_search import TavilySearchResults
 from langchain_experimental.utilities import PythonREPL
-from langchain_openai import ChatOpenAI
-from typing import Annotated, Sequence
-from typing_extensions import TypedDict
-import operator
-import functools
-import matplotlib.pyplot as plt
+from langgraph.graph import StateGraph, END
+from langgraph.prebuilt.tool_executor import ToolExecutor, ToolInvocation
 
 # Set up environment variables for API keys
 TAVILY_API_KEY = os.getenv("TAVILY_API_KEY")
@@ -35,20 +21,14 @@ if not TAVILY_API_KEY or not OPENAI_API_KEY:
     st.error("API keys are missing. Please set TAVILY_API_KEY and OPENAI_API_KEY as secrets.")
     st.stop()
 
-# Define the AgentState class
-class AgentState(TypedDict):
-    messages: Annotated[Sequence[BaseMessage], operator.add]
-    sender: str
-
 # Initialize tools
 tavily_tool = TavilySearchResults(max_results=5)
-repl = PythonREPL()
 
 @tool
-def python_repl(code: Annotated[str, "The python code to execute to generate your chart."]):
+def python_repl(code: str):
     """Executes Python code to generate a chart and returns the chart as a base64-encoded image."""
     try:
-        # Execute the code
+        # Execute the provided Python code
         exec_globals = {"plt": plt}
         exec_locals = {}
         exec(code, exec_globals, exec_locals)
@@ -66,130 +46,53 @@ def python_repl(code: Annotated[str, "The python code to execute to generate you
         encoded_image = base64.b64encode(buf.getvalue()).decode("utf-8")
         return {"status": "success", "image": encoded_image}
     except Exception as e:
-        return {"status": "failed", "error": repr(e)}
+        return {"status": "failed", "error": str(e)}
 
 tools = [tavily_tool, python_repl]
-
-# Define a tool executor
 tool_executor = ToolExecutor(tools)
 
-# Define tool node
-def tool_node(state):
-    """Executes tools in the graph."""
-    messages = state["messages"]
-    last_message = messages[-1]
-    tool_input = json.loads(last_message.additional_kwargs["function_call"]["arguments"])
-    if len(tool_input) == 1 and "__arg1" in tool_input:
-        tool_input = next(iter(tool_input.values()))
-    tool_name = last_message.additional_kwargs["function_call"]["name"]
-    action = ToolInvocation(tool=tool_name, tool_input=tool_input)
-    response = tool_executor.invoke(action)
-    if isinstance(response, dict) and response.get("status") == "success" and "image" in response:
-        return {
-            "messages": [
-                {
-                    "role": "assistant",
-                    "content": "Image generated successfully.",
-                    "image": response["image"],
-                }
-            ]
-        }
-    else:
-        function_message = FunctionMessage(
-            content=f"{tool_name} response: {str(response)}", name=action.tool
-        )
-        return {"messages": [function_message]}
-
-# Define router
-def router(state):
-    """Determines the next step in the workflow."""
-    messages = state["messages"]
-    last_message = messages[-1]
-    if "function_call" in last_message.additional_kwargs:
-        return "call_tool"
-    if "FINAL ANSWER" in last_message.content:
-        return "end"
-    return "continue"
-
-# Define agent creation function
-def create_agent(llm, tools, system_message: str):
-    """Creates an agent."""
-    functions = [convert_to_openai_function(t) for t in 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}",
-            ),
-            MessagesPlaceholder(variable_name="messages"),
-        ]
-    )
-    prompt = prompt.partial(system_message=system_message)
-    prompt = prompt.partial(tool_names=", ".join([tool.name for tool in tools]))
-    return prompt | llm.bind_functions(functions)
-
-# Define agent node
-def agent_node(state, agent, name):
-    result = agent.invoke(state)
-    if isinstance(result, FunctionMessage):
-        pass
-    else:
-        # Sanitize the name field to match OpenAI's naming conventions
-        sanitized_name = re.sub(r"[^a-zA-Z0-9_-]", "_", name)
-        result = HumanMessage(**result.dict(exclude={"type", "name"}), name=sanitized_name)
-    return {"messages": [result], "sender": name}
-
-# Initialize LLM
-llm = ChatOpenAI(api_key=OPENAI_API_KEY)
-
-# Create agents
-research_agent = create_agent(
-    llm, [tavily_tool], system_message="You should provide accurate data for the chart generator to use."
-)
-chart_agent = create_agent(
-    llm, [python_repl], system_message="Any charts you display will be visible by the user."
-)
-
-# Define workflow graph
-workflow = StateGraph(AgentState)
-workflow.add_node("Researcher", functools.partial(agent_node, agent=research_agent, name="Researcher"))
-workflow.add_node("Chart Generator", functools.partial(agent_node, agent=chart_agent, name="Chart Generator"))
-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", lambda x: x["sender"], {"Researcher": "Researcher", "Chart Generator": "Chart Generator"})
-workflow.set_entry_point("Researcher")
-graph = workflow.compile()
+# Define the multi-agent workflow
+workflow = StateGraph()
+workflow.add_node("call_tool", lambda state: tool_executor.invoke(ToolInvocation(**state)))
+workflow.set_entry_point("call_tool")
 
 # Streamlit UI
 st.title("Multi-Agent Workflow")
-user_query = st.text_area("Enter your query:", "Fetch Malaysia's GDP over the past 5 years and draw a line graph.")
-if st.button("Run Workflow"):
-    st.write("Running workflow...")
+st.markdown("### Generate a Chart from Python Code")
+code_input = st.text_area(
+    "Enter Python code for the chart:",
+    """
+import matplotlib.pyplot as plt
+years = [2019, 2020, 2021, 2022, 2023]
+gdp = [300, 310, 330, 360, 399]
+plt.figure(figsize=(10, 6))
+plt.plot(years, gdp, marker='o', color='b', linestyle='-')
+plt.title('Malaysia GDP Over 5 Years')
+plt.xlabel('Year')
+plt.ylabel('GDP (in billion USD)')
+plt.grid(True)
+"""
+)
+
+if st.button("Generate Chart"):
+    st.write("Generating chart...")
     with st.spinner("Processing..."):
-        try:
-            messages = [HumanMessage(content=user_query)]
-            for step in graph.stream({"messages": messages}, {"recursion_limit": 150}):
-                st.write("Step Details:", step)
-                if "messages" in step:
-                    for message in step["messages"]:
-                        if "image" in message:
-                            try:
-                                # Decode the base64-encoded image
-                                encoded_image = message["image"]
-                                decoded_image = BytesIO(base64.b64decode(encoded_image))
-                                # Display the image
-                                st.image(decoded_image, caption="Generated Chart", use_column_width=True)
-                            except Exception as e:
-                                st.error(f"Failed to decode and display the image: {repr(e)}")
-                        elif "content" in message:
-                            # Display any text content
-                            st.write(message["content"])
-        except Exception as e:
-            st.error(f"An error occurred: {e}")
+        # Invoke the Python REPL tool with the code
+        response = python_repl(code_input)
+
+        # Check response status
+        if response["status"] == "success" and "image" in response:
+            encoded_image = response["image"]
+            try:
+                # Decode the base64-encoded image
+                decoded_image = BytesIO(base64.b64decode(encoded_image))
+                # Display the image
+                st.image(decoded_image, caption="Generated Chart", use_column_width=True)
+            except Exception as e:
+                st.error(f"Failed to decode and display the image: {str(e)}")
+        else:
+            st.error(f"Failed to generate chart: {response.get('error', 'Unknown error')}")
+
+st.markdown("### Example Queries")
+st.write("1. Plot GDP of Malaysia over 5 years")
+st.write("2. Create a bar chart of sales data")