From a8f2bfd9aa2f69ee0db5a3aca0339a5f7592a0b2 Mon Sep 17 00:00:00 2001 From: askdevai-bot Date: Tue, 24 Jun 2025 17:10:32 -0400 Subject: [PATCH] Update files via Askdev.AI bot --- tinyagent/hooks/gradio_callback.py | 381 +++++++++++++++++++++++------ 1 file changed, 310 insertions(+), 71 deletions(-) diff --git a/tinyagent/hooks/gradio_callback.py b/tinyagent/hooks/gradio_callback.py index 6e70e44..f3fa21c 100644 --- a/tinyagent/hooks/gradio_callback.py +++ b/tinyagent/hooks/gradio_callback.py @@ -5,6 +5,7 @@ import re import shutil import time +import io from pathlib import Path from typing import Any, Dict, List, Optional, Set, Union @@ -36,6 +37,7 @@ def __init__( show_thinking: bool = True, show_tool_calls: bool = True, logger: Optional[logging.Logger] = None, + log_manager: Optional[Any] = None, ): """ Initialize the Gradio callback. @@ -46,6 +48,7 @@ def __init__( show_thinking: Whether to show the thinking process show_tool_calls: Whether to show tool calls logger: Optional logger to use + log_manager: Optional LoggingManager instance to capture logs from """ self.logger = logger or logging.getLogger(__name__) self.show_thinking = show_thinking @@ -81,6 +84,37 @@ def __init__( # References to Gradio UI components (will be set in create_app) self._chatbot_component = None self._token_usage_component = None + + # Log stream for displaying logs in the UI + self.log_stream = io.StringIO() + self._log_component = None + + # Setup logging + self.log_manager = log_manager + if log_manager: + # Create a handler that writes to our StringIO stream + self.log_handler = logging.StreamHandler(self.log_stream) + self.log_handler.setFormatter( + logging.Formatter('%(asctime)s - %(levelname)s - %(name)s - %(message)s') + ) + self.log_handler.setLevel(logging.DEBUG) + + # Add the handler to the LoggingManager + log_manager.configure_handler( + self.log_handler, + format_string='%(asctime)s - %(levelname)s - %(name)s - %(message)s', + level=logging.DEBUG + ) + self.logger.debug("Added log handler to LoggingManager") + elif logger: + # Fall back to single logger if no LoggingManager is provided + self.log_handler = logging.StreamHandler(self.log_stream) + self.log_handler.setFormatter( + logging.Formatter('%(asctime)s - %(levelname)s - %(name)s - %(message)s') + ) + self.log_handler.setLevel(logging.DEBUG) + logger.addHandler(self.log_handler) + self.logger.debug("Added log handler to logger") self.logger.debug("GradioCallback initialized") @@ -157,7 +191,7 @@ async def _handle_message_add(self, agent: Any, **kwargs: Any) -> None: # Add to detailed tool call info if not already present by ID if not any(tc['id'] == tool_id for tc in self.tool_call_details): - self.tool_call_details.append({ + tool_detail = { "id": tool_id, "name": tool_name, "arguments": formatted_args, @@ -166,7 +200,25 @@ async def _handle_message_add(self, agent: Any, **kwargs: Any) -> None: "result_token_count": 0, "timestamp": current_time, "result_timestamp": None - }) + } + + # Special handling for run_python tool - extract code_lines + if tool_name == "run_python": + try: + # Look for code in different possible field names + code_content = None + for field in ['code_lines', 'code', 'script', 'python_code']: + if field in parsed_args: + code_content = parsed_args[field] + break + + if code_content is not None: + tool_detail["code_lines"] = code_content + self.logger.debug(f"Stored code content for run_python tool {tool_id}") + except Exception as e: + self.logger.error(f"Error processing run_python arguments: {e}") + + self.tool_call_details.append(tool_detail) self.logger.debug(f"Added tool call detail: {tool_name} (ID: {tool_id}, Tokens: {token_count})") # If this is a final_answer or ask_question tool, we'll handle it specially later @@ -386,20 +438,128 @@ def _get_token_usage_text(self) -> str: f"O {self.token_usage['completion_tokens']} | " + f"Total {self.token_usage['total_tokens']}") + def _format_run_python_tool(self, tool_detail: dict) -> str: + """ + Format run_python tool call with proper markdown formatting for code and output. + + Args: + tool_detail: Tool call detail dictionary + + Returns: + Formatted markdown string + """ + tool_name = tool_detail["name"] + tool_id = tool_detail.get("id", "unknown") + code_lines = tool_detail.get("code_lines", []) + result = tool_detail.get("result") + input_tokens = tool_detail.get("token_count", 0) + output_tokens = tool_detail.get("result_token_count", 0) + total_tokens = input_tokens + output_tokens + + # Start building the formatted content + parts = [] + + # Handle different code_lines formats + combined_code = "" + if code_lines: + if isinstance(code_lines, list): + # Standard case: list of code lines + combined_code = "\n".join(code_lines) + elif isinstance(code_lines, str): + # Handle case where code_lines is a single string + combined_code = code_lines + else: + # Convert other types to string + combined_code = str(code_lines) + + # If we have code content, show it as Python code block + if combined_code.strip(): + parts.append(f"**Python Code:**\n```python\n{combined_code}\n```") + else: + # Try to extract code from arguments as fallback + try: + args_dict = json.loads(tool_detail['arguments']) + # Check for different possible code field names + code_content = None + for field in ['code_lines', 'code', 'script', 'python_code']: + if field in args_dict: + code_content = args_dict[field] + break + + if code_content: + if isinstance(code_content, list): + combined_code = "\n".join(code_content) + else: + combined_code = str(code_content) + + if combined_code.strip(): + parts.append(f"**Python Code:**\n```python\n{combined_code}\n```") + else: + # Final fallback to showing raw arguments + parts.append(f"**Input Arguments:**\n```json\n{tool_detail['arguments']}\n```") + else: + # No code found, show raw arguments + parts.append(f"**Input Arguments:**\n```json\n{tool_detail['arguments']}\n```") + except (json.JSONDecodeError, KeyError): + # If we can't parse arguments, show them as-is + parts.append(f"**Input Arguments:**\n```json\n{tool_detail['arguments']}\n```") + + # Add the output if available + if result is not None: + parts.append(f"\n**Output:** ({output_tokens} tokens)") + + try: + # Try to parse result as JSON for better formatting + result_json = json.loads(result) + parts.append(f"```json\n{json.dumps(result_json, indent=2)}\n```") + except (json.JSONDecodeError, TypeError): + # Handle plain text result + if isinstance(result, str): + # Replace escaped newlines with actual newlines for better readability + formatted_result = result.replace("\\n", "\n") + parts.append(f"```\n{formatted_result}\n```") + else: + parts.append(f"```\n{str(result)}\n```") + else: + parts.append(f"\n**Status:** ⏳ Processing...") + + # Add token information + parts.append(f"\n**Token Usage:** {total_tokens} total ({input_tokens} input + {output_tokens} output)") + + return "\n".join(parts) + async def interact_with_agent(self, user_input_processed, chatbot_history): """ Process user input, interact with the agent, and stream updates to Gradio UI. Each tool call and response will be shown as a separate message. """ self.logger.info(f"Starting interaction for: {user_input_processed[:50]}...") + + # Reset state for new interaction to prevent showing previous content + self.thinking_content = "" + self.tool_calls = [] + self.tool_call_details = [] + self.assistant_text_responses = [] + self.token_usage = {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0} + self.is_running = True + self.last_update_yield_time = 0 + self.logger.debug("Reset interaction state for new conversation turn") # 1. Add user message to chatbot history as a ChatMessage chatbot_history.append( ChatMessage(role="user", content=user_input_processed) ) - # Initial yield to show user message - yield chatbot_history, self._get_token_usage_text() + # 2. Add typing indicator immediately after user message + typing_message = ChatMessage( + role="assistant", + content="πŸ€” Thinking..." + ) + chatbot_history.append(typing_message) + typing_message_index = len(chatbot_history) - 1 + + # Initial yield to show user message and typing indicator + yield chatbot_history, self._get_token_usage_text(), self.log_stream.getvalue() if self._log_component else None # Kick off the agent in the background loop = asyncio.get_event_loop() @@ -407,7 +567,7 @@ async def interact_with_agent(self, user_input_processed, chatbot_history): displayed_tool_calls = set() displayed_text_responses = set() - thinking_message_added = False + thinking_removed = False update_interval = 0.3 min_yield_interval = 0.2 @@ -420,6 +580,14 @@ async def interact_with_agent(self, user_input_processed, chatbot_history): sorted_tool_details = sorted(self.tool_call_details, key=lambda x: x.get("timestamp", 0)) sorted_text_responses = sorted(self.assistant_text_responses, key=lambda x: x.get("timestamp", 0)) + # Remove typing indicator once we have actual content to show + if not thinking_removed and (sorted_text_responses or sorted_tool_details): + # Remove the typing indicator + if typing_message_index < len(chatbot_history): + chatbot_history.pop(typing_message_index) + thinking_removed = True + self.logger.debug("Removed typing indicator") + # β†’ New assistant text chunks for resp in sorted_text_responses: content = resp["content"] @@ -430,22 +598,6 @@ async def interact_with_agent(self, user_input_processed, chatbot_history): displayed_text_responses.add(content) self.logger.debug(f"Added new text response: {content[:50]}...") - # β†’ Thinking placeholder (optional) - if self.show_thinking and self.thinking_content \ - and not thinking_message_added \ - and not displayed_text_responses: - thinking_msg = ( - "Working on it...\n\n" - "```" - f"{self.thinking_content}" - "```" - ) - chatbot_history.append( - ChatMessage(role="assistant", content=thinking_msg) - ) - thinking_message_added = True - self.logger.debug("Added thinking message") - # β†’ Show tool calls with "working..." status when they start if self.show_tool_calls: for tool in sorted_tool_details: @@ -455,11 +607,18 @@ async def interact_with_agent(self, user_input_processed, chatbot_history): # If we haven't displayed this tool call yet if tid not in displayed_tool_calls and tid not in in_progress_tool_calls: in_tok = tool.get("token_count", 0) + # Create "working..." message for this tool call - body = ( - f"**Input Arguments:**\n```json\n{tool['arguments']}\n```\n\n" - f"**Output:** ⏳ Working...\n" - ) + if tname == "run_python": + # Special formatting for run_python + body = self._format_run_python_tool(tool) + else: + # Standard formatting for other tools + body = ( + f"**Input Arguments:**\n```json\n{tool['arguments']}\n```\n\n" + f"**Output:** ⏳ Working...\n" + ) + # Add to chatbot with "working" status msg = ChatMessage( role="assistant", @@ -483,10 +642,16 @@ async def interact_with_agent(self, user_input_processed, chatbot_history): tot_tok = in_tok + out_tok # Update the message with completed status and result - body = ( - f"**Input Arguments:**\n```json\n{tool['arguments']}\n```\n\n" - f"**Output:** ({out_tok} tokens)\n```json\n{tool['result']}\n```\n" - ) + if tname == "run_python": + # Special formatting for completed run_python + body = self._format_run_python_tool(tool) + else: + # Standard formatting for other completed tools + body = ( + f"**Input Arguments:**\n```json\n{tool['arguments']}\n```\n\n" + f"**Output:** ({out_tok} tokens)\n```json\n{tool['result']}\n```\n" + ) + # Update the existing message chatbot_history[pos] = ChatMessage( role="assistant", @@ -501,13 +666,19 @@ async def interact_with_agent(self, user_input_processed, chatbot_history): del in_progress_tool_calls[tid] self.logger.debug(f"Updated tool call to completed: {tname}") - # yield updated history + token usage + # yield updated history + token usage + logs token_text = self._get_token_usage_text() - yield chatbot_history, token_text + logs = self.log_stream.getvalue() if self._log_component else None + yield chatbot_history, token_text, logs self.last_update_yield_time = now await asyncio.sleep(update_interval) + # Remove typing indicator if still present when agent finishes + if not thinking_removed and typing_message_index < len(chatbot_history): + chatbot_history.pop(typing_message_index) + self.logger.debug("Removed typing indicator at end") + # once the agent_task is done, add its final result if any try: final_text = await agent_task @@ -521,8 +692,9 @@ async def interact_with_agent(self, user_input_processed, chatbot_history): ) self.logger.debug(f"Added final result: {final_text[:50]}...") - # final token usage - yield chatbot_history, self._get_token_usage_text() + # final token usage and logs + logs = self.log_stream.getvalue() if self._log_component else None + yield chatbot_history, self._get_token_usage_text(), logs def _format_response(self, response_text): """ @@ -673,9 +845,9 @@ def create_app(self, agent: TinyAgent, title: str = "TinyAgent Chat", descriptio # Footer gr.Markdown( - "
" - "Powered by TinyAgent" - "
" + "
" + "Build your own AI Agent Today" + "
" ) # -- Right Chat Column (unchanged) -- @@ -703,6 +875,22 @@ def create_app(self, agent: TinyAgent, title: str = "TinyAgent Chat", descriptio # Clear button clear_btn = gr.Button("Clear Conversation") + # Log accordion - similar to the example provided + with gr.Accordion("Agent Logs", open=False) as log_accordion: + self._log_component = gr.Code( + label="Live Logs", + lines=15, + interactive=False, + value=self.log_stream.getvalue() + ) + refresh_logs_btn = gr.Button("πŸ”„ Refresh Logs") + refresh_logs_btn.click( + fn=lambda: self.log_stream.getvalue(), + inputs=None, + outputs=[self._log_component], + queue=False + ) + # Store processed input temporarily between steps processed_input_state = gr.State("") @@ -723,7 +911,7 @@ def create_app(self, agent: TinyAgent, title: str = "TinyAgent Chat", descriptio # 3. Run the main interaction loop (this yields updates) fn=self.interact_with_agent, inputs=[processed_input_state, self._chatbot_component], - outputs=[self._chatbot_component, self._token_usage_component], # Update chat and tokens + outputs=[self._chatbot_component, self._token_usage_component, self._log_component], # Update chat, tokens, and logs queue=True # Explicitly enable queue for this async generator ).then( # 4. Re-enable the button after interaction finishes @@ -749,7 +937,7 @@ def create_app(self, agent: TinyAgent, title: str = "TinyAgent Chat", descriptio # 3. Run the main interaction loop (this yields updates) fn=self.interact_with_agent, inputs=[processed_input_state, self._chatbot_component], - outputs=[self._chatbot_component, self._token_usage_component], # Update chat and tokens + outputs=[self._chatbot_component, self._token_usage_component, self._log_component], # Update chat, tokens, and logs queue=True # Explicitly enable queue for this async generator ).then( # 4. Re-enable the button after interaction finishes @@ -763,8 +951,8 @@ def create_app(self, agent: TinyAgent, title: str = "TinyAgent Chat", descriptio clear_btn.click( fn=self.clear_conversation, inputs=None, # No inputs needed - # Outputs: Clear chatbot and reset token text - outputs=[self._chatbot_component, self._token_usage_component], + # Outputs: Clear chatbot, reset token text, and update logs + outputs=[self._chatbot_component, self._token_usage_component, self._log_component], queue=False # Run quickly ) @@ -772,8 +960,8 @@ def create_app(self, agent: TinyAgent, title: str = "TinyAgent Chat", descriptio return app def clear_conversation(self): - """Clear the conversation history (UI + agent), reset state, and update UI.""" - self.logger.debug("Clearing conversation (UI + agent)") + """Clear the conversation history (UI + agent), reset state completely, and update UI.""" + self.logger.debug("Clearing conversation completely (UI + agent with new session)") # Reset UI‐side state self.thinking_content = "" self.tool_calls = [] @@ -781,17 +969,63 @@ def clear_conversation(self): self.assistant_text_responses = [] self.token_usage = {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0} self.is_running = False + + # Clear log stream + if hasattr(self, 'log_stream'): + self.log_stream.seek(0) + self.log_stream.truncate(0) + self.logger.info("Log stream cleared") - # Also clear the agent's conversation history + # Completely reset the agent state with a new session try: - if self.current_agent and hasattr(self.current_agent, "clear_conversation"): - self.current_agent.clear_conversation() - self.logger.debug("Cleared TinyAgent internal conversation.") + if self.current_agent: + # Generate a new session ID for a fresh start + import uuid + new_session_id = str(uuid.uuid4()) + self.current_agent.session_id = new_session_id + self.logger.debug(f"Generated new session ID: {new_session_id}") + + # Reset all agent state + # 1. Clear conversation history (preserve system message) + if self.current_agent.messages: + system_msg = self.current_agent.messages[0] + self.current_agent.messages = [system_msg] + else: + # Rebuild default system prompt if missing + default_system_prompt = ( + "You are a helpful AI assistant with access to a variety of tools. " + "Use the tools when appropriate to accomplish tasks. " + "If a tool you need isn't available, just say so." + ) + self.current_agent.messages = [{ + "role": "system", + "content": default_system_prompt + }] + + # 2. Reset session state + self.current_agent.session_state = {} + + # 3. Reset token usage in metadata + if hasattr(self.current_agent, 'metadata') and 'usage' in self.current_agent.metadata: + self.current_agent.metadata['usage'] = { + "prompt_tokens": 0, + "completion_tokens": 0, + "total_tokens": 0 + } + + # 4. Reset any other accumulated state that might affect behavior + self.current_agent.is_running = False + + # 5. Reset session load flag to prevent any deferred loading of old session + self.current_agent._needs_session_load = False + + self.logger.info(f"Completely reset TinyAgent with new session: {new_session_id}") except Exception as e: - self.logger.error(f"Failed to clear TinyAgent conversation: {e}") + self.logger.error(f"Failed to reset TinyAgent completely: {e}") - # Return cleared UI components: empty chat + fresh token usage - return [], self._get_token_usage_text() + # Return cleared UI components: empty chat + fresh token usage + empty logs + logs = self.log_stream.getvalue() if hasattr(self, 'log_stream') else "" + return [], self._get_token_usage_text(), logs def launch(self, agent, title="TinyAgent Chat", description=None, share=False, **kwargs): """ @@ -853,21 +1087,31 @@ async def run_example(): from tinyagent import TinyAgent # Assuming TinyAgent is importable from tinyagent.hooks.logging_manager import LoggingManager # Assuming LoggingManager exists - # --- Logging Setup (Simplified) --- + # --- Logging Setup (Similar to the example provided) --- log_manager = LoggingManager(default_level=logging.INFO) log_manager.set_levels({ 'tinyagent.hooks.gradio_callback': logging.DEBUG, 'tinyagent.tiny_agent': logging.DEBUG, 'tinyagent.mcp_client': logging.DEBUG, + 'tinyagent.code_agent': logging.DEBUG, }) + + # Console handler for terminal output console_handler = logging.StreamHandler(sys.stdout) log_manager.configure_handler( console_handler, format_string='%(asctime)s - %(name)s - %(levelname)s - %(message)s', level=logging.DEBUG ) + + # The Gradio UI will automatically set up its own log handler + # through the LoggingManager when we pass it to GradioCallback + + # Get loggers for different components ui_logger = log_manager.get_logger('tinyagent.hooks.gradio_callback') agent_logger = log_manager.get_logger('tinyagent.tiny_agent') + mcp_logger = log_manager.get_logger('tinyagent.mcp_client') + ui_logger.info("--- Starting GradioCallback Example ---") # --- End Logging Setup --- @@ -889,12 +1133,13 @@ async def run_example(): agent.add_tool(get_weather) - # Create the Gradio callback + # Create the Gradio callback with LoggingManager integration gradio_ui = GradioCallback( file_upload_folder=upload_folder, show_thinking=True, show_tool_calls=True, - logger=ui_logger # Pass the specific logger + logger=ui_logger, + log_manager=log_manager # Pass the LoggingManager for comprehensive logging ) agent.add_callback(gradio_ui) @@ -909,25 +1154,9 @@ async def run_example(): ui_logger.error(f"Failed to connect to MCP servers: {e}", exc_info=True) # Continue without servers - we still have the local get_weather tool - # Create the Gradio app but don't launch it yet - #app = gradio_ui.create_app( - # agent, - # title="TinyAgent Chat Interface", - # description="Chat with TinyAgent. Try asking: 'Plan a trip to Toronto for 7 days in the next month.'", - #) - - # Configure the queue without extra parameters - #app.queue() - - # Launch the app in a way that doesn't block our event loop + # Launch the Gradio interface ui_logger.info("Launching Gradio interface...") try: - # Launch without blocking - #app.launch( - # share=False, - # prevent_thread_lock=True, # Critical to not block our event loop - # show_error=True - #) gradio_ui.launch( agent, title="TinyAgent Chat Interface", @@ -938,9 +1167,19 @@ async def run_example(): ) ui_logger.info("Gradio interface launched (non-blocking).") + # Generate some log messages to demonstrate the log panel + # These will appear in both the terminal and the Gradio UI log panel + ui_logger.info("UI component initialized successfully") + agent_logger.debug("Agent ready to process requests") + mcp_logger.info("MCP connection established") + + for i in range(3): + ui_logger.info(f"Example log message {i+1} from UI logger") + agent_logger.debug(f"Example debug message {i+1} from agent logger") + mcp_logger.warning(f"Example warning {i+1} from MCP logger") + await asyncio.sleep(1) + # Keep the main event loop running to handle both Gradio and MCP operations - # This is the key part - we need to keep our main event loop running - # but also allow it to process both Gradio and MCP client operations while True: await asyncio.sleep(1) # More efficient than an Event().wait()