Support one-liner UI

agents/agents.py

@@ -81,8 +81,11 @@ class AgentGenerationError(AgentError):
 
     pass
 
+class AgentStep:
+    pass
+
 @dataclass
-class ActionStep:
+class ActionStep(AgentStep):
     tool_call: str | None = None
     start_time: float | None = None
     step_end_time: float | None = None
@@ -93,18 +96,19 @@ class ActionStep:
     llm_output: str | None = None
 
 @dataclass
-class PlanningStep:
+class PlanningStep(AgentStep):
     plan: str
     facts: str
 
 @dataclass
-class TaskStep:
+class TaskStep(AgentStep):
     task: str
 
 @dataclass
-class SystemPromptStep:
+class SystemPromptStep(AgentStep):
     system_prompt: str
 
+
 def format_prompt_with_tools(toolbox: Toolbox, prompt_template: str, tool_description_template: str) -> str:
     tool_descriptions = toolbox.show_tool_descriptions(tool_description_template)
     prompt = prompt_template.replace("{{tool_descriptions}}", tool_descriptions)

agents/gradio_ui.py

@@ -0,0 +1,92 @@
+#!/usr/bin/env python
+# coding=utf-8
+
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from .agent_types import AgentAudio, AgentImage, AgentText
+from .agents import BaseAgent, AgentStep, ActionStep
+import gradio as gr
+
+def pull_messages_from_step(step_log: AgentStep, test_mode: bool = True):
+    """Extract ChatMessage objects from agent steps"""
+    if isinstance(step_log, ActionStep):
+        yield gr.ChatMessage(role="assistant", content=step_log.rationale)
+        if step_log.tool_call is not None:
+            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 gr.ChatMessage(
+                role="assistant",
+                metadata={"title": f"🛠️ Used tool {step_log.tool_call['tool_name']}"},
+                content=str(content),
+            )
+        if step_log.observation is not None:
+            yield gr.ChatMessage(role="assistant", content=f"```\n{step_log.observation}\n```")
+        if step_log.error is not None:
+            yield gr.ChatMessage(
+                role="assistant",
+                content=str(step_log.error),
+                metadata={"title": "💥 Error"},
+            )
+
+
+def stream_to_gradio(agent, task: str, test_mode: bool = False, reset_agent_memory: bool=False, **kwargs):
+    """Runs an agent with the given task and streams the messages from the agent as gradio ChatMessages."""
+
+    for step_log in agent.run(task, stream=True, reset=reset_agent_memory, **kwargs):
+        for message in pull_messages_from_step(step_log, test_mode=test_mode):
+            yield message
+
+    final_answer = step_log  # Last log is the run's final_answer
+
+    if isinstance(final_answer, AgentText):
+        yield gr.ChatMessage(role="assistant", content=f"**Final answer:**\n```\n{final_answer.to_string()}\n```")
+    elif isinstance(final_answer, AgentImage):
+        yield gr.ChatMessage(
+            role="assistant",
+            content={"path": final_answer.to_string(), "mime_type": "image/png"},
+        )
+    elif isinstance(final_answer, AgentAudio):
+        yield gr.ChatMessage(
+            role="assistant",
+            content={"path": final_answer.to_string(), "mime_type": "audio/wav"},
+        )
+    else:
+        yield gr.ChatMessage(role="assistant", content=str(final_answer))
+
+
+class GradioUI():
+    """A one-line interface to launch your agent in Gradio"""
+    def __init__(self, agent: BaseAgent):
+        self.agent = agent 
+
+    def interact_with_agent(self, prompt, messages):
+        messages.append(gr.ChatMessage(role="user", content=prompt))
+        yield messages
+        for msg in stream_to_gradio(self.agent, task=prompt, reset_agent_memory=False):
+            messages.append(msg)
+            yield messages
+        yield messages
+
+    def run(self):
+        with gr.Blocks() as demo:
+            stored_message = gr.State([])
+            chatbot = gr.Chatbot(label="Agent",
+                                type="messages",
+                                avatar_images=(None, "https://em-content.zobj.net/source/twitter/53/robot-face_1f916.png"))
+            text_input = gr.Textbox(lines=1, label="Chat Message")
+            text_input.submit(lambda s: (s, ""), [text_input], [stored_message, text_input]).then(self.interact_with_agent, [stored_message, chatbot], [chatbot])
+
+        demo.launch()

agents/monitoring.py

@@ -14,59 +14,8 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from .agent_types import AgentAudio, AgentImage, AgentText
 from .utils import console
 
-def pull_message(step_log: dict, test_mode: bool = True):
-    from gradio import ChatMessage
-
-    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=str(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, task: str, test_mode: bool = False, reset_agent_memory: bool=False, **kwargs):
-    """Runs an agent with the given task and streams the messages from the agent as gradio ChatMessages."""
-    from gradio import ChatMessage
-
-    for step_log in agent.run(task, stream=True, reset=reset_agent_memory, **kwargs):
-        if isinstance(step_log, dict):
-            for message in pull_message(step_log, test_mode=test_mode):
-                yield message
-
-    final_answer = step_log  # Last log is the run's final_answer
-
-    if isinstance(final_answer, AgentText):
-        yield ChatMessage(role="assistant", content=f"**Final answer:**\n```\n{final_answer.to_string()}\n```")
-    elif isinstance(final_answer, AgentImage):
-        yield ChatMessage(
-            role="assistant",
-            content={"path": final_answer.to_string(), "mime_type": "image/png"},
-        )
-    elif isinstance(final_answer, AgentAudio):
-        yield ChatMessage(
-            role="assistant",
-            content={"path": final_answer.to_string(), "mime_type": "audio/wav"},
-        )
-    else:
-        yield ChatMessage(role="assistant", content=str(final_answer))
 
 
 class Monitor:

examples/quick_launch_gradio.py

@@ -0,0 +1,10 @@
+from agents.gradio_ui import GradioUI
+from agents import HfApiEngine, load_tool, CodeAgent
+
+image_generation_tool = load_tool("m-ric/text-to-image")
+
+llm_engine = HfApiEngine("Qwen/Qwen2.5-72B-Instruct")
+
+agent = CodeAgent(tools=[image_generation_tool], llm_engine=llm_engine)
+
+GradioUI(agent).run()