Update message format and llama version

app.py

@@ -4,18 +4,14 @@ import gradio as gr
 from transformers import ReactCodeAgent, HfEngine, Tool
 import pandas as pd
 
-from gradio_agentchatbot import (
-    AgentChatbot,
-    stream_from_transformers_agent,
-    ChatMessage,
-    ChatFileMessage,
-)
+from gradio import Chatbot
+from streaming import stream_to_gradio
 from huggingface_hub import login
 from gradio.data_classes import FileData
 
 login(os.getenv("HUGGINGFACEHUB_API_TOKEN"))
 
-llm_engine = HfEngine("meta-llama/Meta-Llama-3-70B-Instruct")
+llm_engine = HfEngine("meta-llama/Meta-Llama-3.1-70B-Instruct")
 
 agent = ReactCodeAgent(
     tools=[],
@@ -25,7 +21,6 @@ agent = ReactCodeAgent(
 )
 
 base_prompt = """You are an expert data analyst.
-Please load the source file with pandas (you cannot use 'os' module).
 According to the features you have and the dta structure given below, determine which feature should be the target.
 Then list 3 interesting questions that could be asked on this data, for instance about specific correlations with target variable.
 Then answer these questions one by one, by finding the relevant numbers.
@@ -35,9 +30,11 @@ In your final answer: summarize these correlations and trends
 After each number derive real worlds insights, for instance: "Correlation between is_december and boredness is 1.3453, which suggest people are more bored in winter".
 Your final answer should be a long string with at least 3 numbered and detailed parts.
 
-Source file for the data = {source_file}
 Structure of the data:
 {structure_notes}
+
+The data file is passed to you as the variable data_file, it is a pandas dataframe, you can use it directly.
+DO NOT try to load data_file, it is already a dataframe pre-loaded in your python interpreter!
 """
 
 example_notes="""This data is about the Titanic wreck in 1912.
@@ -69,34 +66,36 @@ def interact_with_agent(file_input, additional_notes):
     shutil.rmtree("./figures")
     os.makedirs("./figures")
 
-    read_file = pd.read_csv(file_input)
+    data_file = pd.read_csv(file_input)
     data_structure_notes = f"""- Description (output of .describe()):
-    {read_file.describe()}
+    {data_file.describe()}
     - Columns with dtypes:
-    {read_file.dtypes}"""
+    {data_file.dtypes}"""
 
-    prompt = base_prompt.format(source_file=file_input, structure_notes=data_structure_notes)
+    prompt = base_prompt.format(structure_notes=data_structure_notes)
 
     if additional_notes and len(additional_notes) > 0:
         prompt += "\nAdditional notes on the data:\n" + additional_notes
 
-    messages = [ChatMessage(role="user", content=prompt, thought=True)]
-    yield messages
+    messages = [gr.ChatMessage(role="user", content=prompt)]
+    yield messages + [
+        gr.ChatMessage(role="assistant", content="⏳ _Task not finished yet!_")
+    ]
 
     plot_image_paths = {}
-    for msg in stream_from_transformers_agent(agent, prompt):
+    for msg in stream_to_gradio(agent, prompt, data_file=data_file):
         messages.append(msg)
         for image_path in get_images_in_directory("./figures"):
             if image_path not in plot_image_paths:
-                image_message = ChatFileMessage(
+                image_message = gr.ChatMessage(
                     role="assistant",
-                    file=FileData(path=image_path, mime_type="image/png"),
-                    content="",
-                    thought=True,
+                    content=FileData(path=image_path, mime_type="image/png"),
                 )
                 plot_image_paths[image_path] = True
                 messages.append(image_message)
-        yield messages
+        yield messages + [
+            gr.ChatMessage(role="assistant", content="⏳ _Task not finished yet!_")
+        ]
     yield messages
 
 
@@ -109,7 +108,14 @@ Drop a `.csv` file below, add notes to describe this data if needed, and **Llama
         label="Additional notes to support the analysis"
     )
     submit = gr.Button("Run analysis!")
-    chatbot = AgentChatbot(label="Agent")
+    chatbot = gr.Chatbot(
+        label="Agent",
+        type="messages",
+        avatar_images=(
+            None,
+            "https://em-content.zobj.net/source/twitter/53/robot-face_1f916.png",
+        ),
+    )
     gr.Examples(
         examples=[["./example/titanic.csv", example_notes]],
         inputs=[file_input, text_input],

requirements.txt

@@ -1,4 +1,3 @@
-gradio_agentchatbot
 git+https://github.com/huggingface/transformers.git#egg=transformers[agents]
 matplotlib
 seaborn

streaming.py

@@ -0,0 +1,64 @@
+from transformers.agents.agent_types import AgentAudio, AgentImage, AgentText, AgentType
+from transformers.agents import ReactAgent
+
+
+def pull_message(step_log: dict):
+    try:
+        from gradio import ChatMessage
+    except ImportError:
+        raise ImportError("Gradio should be installed in order to launch a gradio demo.")
+
+    if step_log.get("rationale"):
+        yield ChatMessage(role="assistant", content=step_log["rationale"])
+    if step_log.get("tool_call"):
+        used_code = step_log["tool_call"]["tool_name"] == "code interpreter"
+        content = step_log["tool_call"]["tool_arguments"]
+        if used_code:
+            content = f"```py\n{content}\n```"
+        yield ChatMessage(
+            role="assistant",
+            metadata={"title": f"🛠️ Used tool {step_log['tool_call']['tool_name']}"},
+            content=content,
+        )
+    if step_log.get("observation"):
+        yield ChatMessage(role="assistant", content=f"```\n{step_log['observation']}\n```")
+    if step_log.get("error"):
+        yield ChatMessage(
+            role="assistant",
+            content=str(step_log["error"]),
+            metadata={"title": "💥 Error"},
+        )
+
+
+def stream_to_gradio(agent: ReactAgent, task: str, **kwargs):
+    """Runs an agent with the given task and streams the messages from the agent as gradio ChatMessages."""
+
+    try:
+        from gradio import ChatMessage
+    except ImportError:
+        raise ImportError("Gradio should be installed in order to launch a gradio demo.")
+
+    class Output:
+        output: AgentType | str = None
+
+    for step_log in agent.run(task, stream=True, **kwargs):
+        if isinstance(step_log, dict):
+            for message in pull_message(step_log):
+                print("message", message)
+                yield message
+
+    Output.output = step_log
+    if isinstance(Output.output, AgentText):
+        yield ChatMessage(role="assistant", content=f"**Final answer:**\n```\n{Output.output.to_string()}\n```")
+    elif isinstance(Output.output, AgentImage):
+        yield ChatMessage(
+            role="assistant",
+            content={"path": Output.output.to_string(), "mime_type": "image/png"},
+        )
+    elif isinstance(Output.output, AgentAudio):
+        yield ChatMessage(
+            role="assistant",
+            content={"path": Output.output.to_string(), "mime_type": "audio/wav"},
+        )
+    else:
+        yield ChatMessage(role="assistant", content=Output.output)