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"]