File size: 15,024 Bytes
deafbd7 cb4e5b3 deafbd7 9a2a3bc deafbd7 cb4e5b3 deafbd7 9a2a3bc deafbd7 9a2a3bc deafbd7 9a2a3bc 3e2ff2f deafbd7 9a2a3bc deafbd7 9a2a3bc 3e2ff2f 9a2a3bc deafbd7 9a2a3bc deafbd7 9a2a3bc 3e2ff2f cb4e5b3 9a2a3bc deafbd7 3e2ff2f 9a2a3bc 0cc5955 9a2a3bc 3e2ff2f 9a2a3bc 3e2ff2f 9a2a3bc 0cc5955 9a2a3bc 3e2ff2f deafbd7 9a2a3bc deafbd7 3e2ff2f 0cc5955 9a2a3bc deafbd7 3e2ff2f deafbd7 cb4e5b3 9a2a3bc deafbd7 9a2a3bc 3e2ff2f 9a2a3bc 3e2ff2f cb4e5b3 3e2ff2f deafbd7 cb4e5b3 9a2a3bc 3e2ff2f fd4101e 3e2ff2f cb4e5b3 3e2ff2f fd4101e 3e2ff2f cb4e5b3 3e2ff2f cb4e5b3 3e2ff2f 0cc5955 9a2a3bc 0cc5955 9a2a3bc 3e2ff2f 9a2a3bc cb4e5b3 9a2a3bc cb4e5b3 9a2a3bc 0cc5955 9a2a3bc cb4e5b3 9a2a3bc 3e2ff2f 9a2a3bc deafbd7 d9e8945 deafbd7 9a2a3bc deafbd7 d9e8945 9a2a3bc 0cc5955 d9e8945 d8cd6bd 839bdd7 d8cd6bd fd4101e 3fe540b 9b8d24a d9e8945 9b8d24a fd4101e 9b8d24a 839bdd7 9b8d24a 839bdd7 d9e8945 9b8d24a 839bdd7 d9e8945 d8cd6bd 3e2ff2f 9a2a3bc 0cc5955 3e2ff2f 0cc5955 d9e8945 deafbd7 9a2a3bc d9e8945 0cc5955 fd4101e 839bdd7 d9e8945 839bdd7 d9e8945 839bdd7 9a2a3bc d9e8945 9a2a3bc cb4e5b3 3e2ff2f 0cc5955 d9e8945 deafbd7 9a2a3bc 0cc5955 cb4e5b3 0cc5955 cb4e5b3 fd4101e 9a2a3bc cb4e5b3 9a2a3bc 3e2ff2f cb4e5b3 0cc5955 9a2a3bc 0cc5955 3e2ff2f cb4e5b3 0cc5955 fd4101e d9e8945 839bdd7 d9e8945 839bdd7 9a2a3bc 3e2ff2f 839bdd7 9a2a3bc 0cc5955 3e2ff2f cb4e5b3 d9e8945 0cc5955 839bdd7 0cc5955 deafbd7 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 |
#!/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.
import os
import re
from typing import Optional
import tempfile
from PIL import Image as PILImage
from smolagents.agent_types import AgentAudio, AgentImage, AgentText, handle_agent_output_types
from smolagents.agents import ActionStep, MultiStepAgent
from smolagents.memory import MemoryStep
from smolagents.utils import _is_package_available
import gradio as gr
def pull_messages_from_step_dict(step_log: MemoryStep):
"""Extract messages as dicts for Gradio type='messages' Chatbot"""
if isinstance(step_log, ActionStep):
step_number_str = f"Step {step_log.step_number}" if step_log.step_number is not None else "Processing"
yield {"role": "assistant", "content": f"**{step_number_str}**"}
if hasattr(step_log, "model_output") and step_log.model_output is not None:
model_output = step_log.model_output.strip()
model_output = re.sub(r"```\s*<end_code>[\s\S]*|[\s\S]*<end_code>\s*```", "```", model_output, flags=re.DOTALL)
model_output = re.sub(r"<end_code>", "", model_output)
model_output = model_output.strip()
yield {"role": "assistant", "content": model_output}
if hasattr(step_log, "tool_calls") and step_log.tool_calls:
tc = step_log.tool_calls[0]
tool_info_md = f"🛠️ **Tool Used: {tc.name}**\n"
args = tc.arguments
if isinstance(args, dict):
args_str = str(args.get("answer", str(args)))
else:
args_str = str(args).strip()
if tc.name == "python_interpreter":
code_content = args_str
code_content = re.sub(r"^```python\s*\n?", "", code_content)
code_content = re.sub(r"\n?```\s*$", "", code_content)
code_content = re.sub(r"^\s*<end_code>\s*", "", code_content)
code_content = re.sub(r"\s*<end_code>\s*$", "", code_content)
code_content = code_content.strip()
tool_info_md += f"Executing Code:\n```python\n{code_content}\n```\n"
else:
tool_info_md += f"Arguments: `{args_str}`\n"
if hasattr(step_log, "observations") and step_log.observations and step_log.observations.strip():
obs_content = step_log.observations.strip()
obs_content = re.sub(r"^Execution logs:\s*", "", obs_content).strip()
if obs_content:
tool_info_md += f"📝 **Tool Output/Logs:**\n```text\n{obs_content}\n```\n"
if hasattr(step_log, "error") and step_log.error:
tool_info_md += f"💥 **Error:** {str(step_log.error)}\n"
yield {"role": "assistant", "content": tool_info_md.strip()}
elif hasattr(step_log, "error") and step_log.error:
yield {"role": "assistant", "content": f"💥 **Error:** {str(step_log.error)}"}
footnote_parts = []
if step_log.step_number is not None:
footnote_parts.append(f"Step {step_log.step_number}")
if hasattr(step_log, "duration") and step_log.duration is not None:
footnote_parts.append(f"Duration: {round(float(step_log.duration), 2)}s")
if hasattr(step_log, "input_token_count") and step_log.input_token_count is not None:
footnote_parts.append(f"InTokens: {step_log.input_token_count:,}")
if hasattr(step_log, "output_token_count") and step_log.output_token_count is not None:
footnote_parts.append(f"OutTokens: {step_log.output_token_count:,}")
if footnote_parts:
footnote_text = " | ".join(footnote_parts)
yield {"role": "assistant", "content": f"""<p style="color: #999; font-size: 0.8em; margin-top:0; margin-bottom:0;">{footnote_text}</p>"""}
yield {"role": "assistant", "content": "---"}
def stream_to_gradio(
agent,
task: str,
reset_agent_memory: bool = False,
additional_args: Optional[dict] = None,
):
if not _is_package_available("gradio"):
raise ModuleNotFoundError("Install 'gradio': `pip install 'smolagents[gradio]'`")
if hasattr(agent, 'interaction_logs'):
agent.interaction_logs.clear()
print("DEBUG Gradio: Cleared agent interaction_logs for new run.")
all_step_logs = []
for step_log in agent.run(task, stream=True, reset=reset_agent_memory, additional_args=additional_args):
all_step_logs.append(step_log)
if hasattr(agent.model, "last_input_token_count") and agent.model.last_input_token_count is not None:
if isinstance(step_log, ActionStep):
step_log.input_token_count = agent.model.last_input_token_count
step_log.output_token_count = agent.model.last_output_token_count
for msg_dict in pull_messages_from_step_dict(step_log):
yield msg_dict
if not all_step_logs:
yield {"role": "assistant", "content": "Agent did not produce any output."}
return
final_answer_content = all_step_logs[-1]
actual_content_for_handling = final_answer_content
if hasattr(final_answer_content, 'final_answer') and not isinstance(final_answer_content, (str, PILImage.Image, tuple)):
actual_content_for_handling = final_answer_content.final_answer
print(f"DEBUG Gradio: Extracted actual_content_for_handling from FinalAnswerStep: {type(actual_content_for_handling)}")
if isinstance(actual_content_for_handling, PILImage.Image):
print("DEBUG Gradio (stream_to_gradio): Actual content IS a raw PIL Image.")
try:
with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as tmp_file:
actual_content_for_handling.save(tmp_file, format="PNG")
image_path_for_gradio = tmp_file.name
print(f"DEBUG Gradio: Saved PIL image to temp path: {image_path_for_gradio}")
yield {"role": "assistant", "content": (image_path_for_gradio, "Generated Image")}
return
except Exception as e:
print(f"DEBUG Gradio: Error saving extracted PIL image: {e}")
yield {"role": "assistant", "content": f"**Final Answer (Error displaying image):** {e}"}
return
final_answer_processed = handle_agent_output_types(actual_content_for_handling)
print(f"DEBUG Gradio: final_answer_processed type after handle_agent_output_types: {type(final_answer_processed)}")
if isinstance(final_answer_processed, AgentText):
yield {"role": "assistant", "content": f"**Final Answer:**\n{final_answer_processed.to_string()}"}
elif isinstance(final_answer_processed, AgentImage):
image_path = final_answer_processed.to_string()
print(f"DEBUG Gradio (stream_to_gradio): final_answer_processed is AgentImage. Path: {image_path}")
if image_path and os.path.exists(image_path):
yield {"role": "assistant", "content": (image_path, "Generated Image (from AgentImage)")}
else:
err_msg = f"Error: Image path from AgentImage ('{image_path}') not found or invalid."
print(f"DEBUG Gradio: {err_msg}")
yield {"role": "assistant", "content": f"**Final Answer ({err_msg})**"}
elif isinstance(final_answer_processed, AgentAudio):
audio_path = final_answer_processed.to_string()
print(f"DEBUG Gradio (stream_to_gradio): AgentAudio path: {audio_path}")
if audio_path and os.path.exists(audio_path):
yield {"role": "assistant", "content": (audio_path, "Generated Audio")}
else:
err_msg = f"Error: Audio path from AgentAudio ('{audio_path}') not found"
print(f"DEBUG Gradio: {err_msg}")
yield {"role": "assistant", "content": f"**Final Answer ({err_msg})**"}
else:
yield {"role": "assistant", "content": f"**Final Answer:**\n{str(final_answer_processed)}"}
class GradioUI:
def __init__(self, agent: MultiStepAgent, file_upload_folder: str | None = None):
if not _is_package_available("gradio"):
raise ModuleNotFoundError("Install 'gradio': `pip install 'smolagents[gradio]'`")
self.agent = agent
self.file_upload_folder = None
self._latest_file_path_for_download = None
def _get_created_document_path(self):
if hasattr(self.agent, 'interaction_logs') and self.agent.interaction_logs:
print(f"DEBUG Gradio UI: Checking {len(self.agent.interaction_logs)} interaction log entries for created document paths.")
for log_entry in reversed(self.agent.interaction_logs):
if isinstance(log_entry, ActionStep):
observations = getattr(log_entry, 'observations', None)
tool_calls = getattr(log_entry, 'tool_calls', [])
is_python_interpreter_step = any(tc.name == "python_interpreter" for tc in tool_calls)
if is_python_interpreter_step and observations and isinstance(observations, str):
# CRITICAL DEBUG LINE: Print the exact observations string
print(f"DEBUG Gradio UI (_get_created_document_path): Python Interpreter Observations: '''{observations}'''")
match = re.search(
r"(?:Document created \((?:docx|pdf|txt)\):|Document converted to PDF:)\s*(/tmp/[a-zA-Z0-9_]+/generated_document\.(?:docx|pdf|txt))",
observations,
re.MULTILINE
)
if match:
extracted_path = match.group(1)
print(f"DEBUG Gradio UI: Regex matched. Extracted path: '{extracted_path}'")
normalized_path = os.path.normpath(extracted_path)
if os.path.exists(normalized_path):
print(f"DEBUG Gradio UI: Validated path for download: {normalized_path}")
return normalized_path
else:
print(f"DEBUG Gradio UI: Path from create_document output ('{normalized_path}') does not exist.")
print("DEBUG Gradio UI: No valid generated document path found in agent logs.")
return None
def interact_with_agent(self, prompt_text: str, current_chat_history: list):
print(f"DEBUG Gradio: interact_with_agent called with prompt: '{prompt_text}'")
updated_chat_history = current_chat_history + [{"role": "user", "content": prompt_text}]
yield updated_chat_history, gr.update(value=None, visible=False) # For file_download_display_component
agent_responses_for_history = []
for msg_dict in stream_to_gradio(self.agent, task=prompt_text, reset_agent_memory=False):
agent_responses_for_history.append(msg_dict)
yield updated_chat_history + agent_responses_for_history, gr.update(value=None, visible=False) # For file_download_display_component
final_chat_display_content = updated_chat_history + agent_responses_for_history
document_path_to_display = self._get_created_document_path()
if document_path_to_display:
print(f"DEBUG Gradio: Document found for display: {document_path_to_display}")
# CORRECTED: Use gr.update() for the File component
yield final_chat_display_content, gr.update(value=document_path_to_display,
label=os.path.basename(document_path_to_display),
visible=True)
else:
print(f"DEBUG Gradio: No document found for display.")
# CORRECTED: Use gr.update() for the File component
yield final_chat_display_content, gr.update(value=None, visible=False)
def log_user_message(self, text_input_value: str):
full_prompt = text_input_value
print(f"DEBUG Gradio: Prepared prompt for agent: {full_prompt[:300]}...")
return full_prompt, ""
# prepare_and_show_download_file is not needed if we directly update the gr.File component
def launch(self, **kwargs):
with gr.Blocks(fill_height=True, theme=gr.themes.Soft(primary_hue=gr.themes.colors.blue)) as demo:
prepared_prompt_for_agent = gr.State("")
gr.Markdown("## Smol Talk with your Agent")
with gr.Row(equal_height=False):
with gr.Column(scale=3):
chatbot_display = gr.Chatbot(
type="messages",
avatar_images=(None, "https://huggingface.co/datasets/huggingface/brand-assets/resolve/main/hf-logo-round.png"),
height=700,
show_copy_button=True,
bubble_full_width=False,
show_label=False
)
text_message_input = gr.Textbox(
lines=1,
placeholder="Type your message and press Enter, or Shift+Enter for new line...",
show_label=False
)
with gr.Column(scale=1):
# "Generated File" section directly shows the gr.File component
gr.Markdown("### Generated Document")
file_download_display_component = gr.File(
label="Downloadable Document",
visible=False,
interactive=False
)
text_message_input.submit(
self.log_user_message,
[text_message_input],
[prepared_prompt_for_agent, text_message_input]
).then(
self.interact_with_agent,
[prepared_prompt_for_agent, chatbot_display],
[chatbot_display, file_download_display_component] # Outputs update chatbot and file component
)
demo.launch(debug=True, share=kwargs.get("share", False), **kwargs)
__all__ = ["stream_to_gradio", "GradioUI"] |