Spaces:
Running
Running
| import os | |
| import sys | |
| import chainlit as cl | |
| from src.research_agent import run_agent, create_system_message | |
| from langchain_core.messages import ( | |
| AIMessage, | |
| HumanMessage, | |
| SystemMessage, | |
| ToolMessage, | |
| ) | |
| # Set domain as global variable | |
| DOMAIN = "General Research" | |
| DEBUG_MODE = True # Set to True to show detailed tool usage | |
| # List of available domains | |
| AVAILABLE_DOMAINS = [ | |
| "General Research", | |
| "Computer Science", | |
| "Artificial Intelligence", | |
| ] | |
| async def on_chat_start(): | |
| """ | |
| Initialize the chat session | |
| """ | |
| # Set up the session state | |
| cl.user_session.set("messages", [create_system_message(DOMAIN)]) | |
| cl.user_session.set("debug_mode", DEBUG_MODE) | |
| # Send a welcome message | |
| await cl.Message( | |
| content=f"Hello! I'm your research assistant specialized in {DOMAIN}. How can I help you?\n\n" | |
| f"I have access to the following tools:\n" | |
| f"1. 🔎 Web Search - For recent information and general queries\n" | |
| f"2. 📄 Research Papers - For academic and scientific knowledge\n" | |
| f"3. 📚 Wikipedia - For background information and factual summaries\n" | |
| f"4. 📊 Data Analysis - For analyzing data you provide\n\n" | |
| f"You can change the research domain by typing `/domain` followed by one of these options:\n" | |
| f"{', '.join(AVAILABLE_DOMAINS)}\n\n" | |
| f"You can toggle debug mode with `/debug on` or `/debug off` to see detailed tool usage.", | |
| author="Research Assistant" | |
| ).send() | |
| async def on_settings_update(settings): | |
| """ | |
| Handle settings updates | |
| """ | |
| global DOMAIN | |
| if "domain" in settings: | |
| DOMAIN = settings["domain"] | |
| # Reset messages with new domain | |
| cl.user_session.set("messages", [create_system_message(DOMAIN)]) | |
| # Notify user of domain change | |
| await cl.Message( | |
| content=f"Domain changed to {DOMAIN}. My knowledge is now specialized for this domain.", | |
| author="System" | |
| ).send() | |
| async def on_message(message: cl.Message): | |
| """ | |
| Process incoming messages | |
| """ | |
| global DOMAIN, DEBUG_MODE | |
| # Check for debug mode command | |
| if message.content.startswith("/debug"): | |
| command_parts = message.content.split() | |
| if len(command_parts) > 1: | |
| if command_parts[1].lower() == "on": | |
| DEBUG_MODE = True | |
| cl.user_session.set("debug_mode", True) | |
| await cl.Message( | |
| content="Debug mode turned ON. You'll see detailed tool usage information.", | |
| author="System" | |
| ).send() | |
| elif command_parts[1].lower() == "off": | |
| DEBUG_MODE = False | |
| cl.user_session.set("debug_mode", False) | |
| await cl.Message( | |
| content="Debug mode turned OFF.", | |
| author="System" | |
| ).send() | |
| else: | |
| current_state = "ON" if DEBUG_MODE else "OFF" | |
| await cl.Message( | |
| content=f"Debug mode is currently {current_state}. Use `/debug on` or `/debug off` to change.", | |
| author="System" | |
| ).send() | |
| return | |
| # Check for domain change command | |
| if message.content.startswith("/domain"): | |
| command_parts = message.content.split() | |
| if len(command_parts) == 1: | |
| # Display available domains | |
| domains_list = "\n".join([f"- {domain}" for domain in AVAILABLE_DOMAINS]) | |
| await cl.Message( | |
| content=f"Please specify a domain. Available domains:\n{domains_list}\n\nExample: `/domain Artificial Intelligence`", | |
| author="System" | |
| ).send() | |
| return | |
| # Get the requested domain | |
| requested_domain = " ".join(command_parts[1:]) | |
| # Check if it's in the available domains (case insensitive) | |
| found_domain = None | |
| for domain in AVAILABLE_DOMAINS: | |
| if domain.lower() == requested_domain.lower(): | |
| found_domain = domain | |
| break | |
| if found_domain: | |
| DOMAIN = found_domain | |
| # Reset messages with new domain | |
| cl.user_session.set("messages", [create_system_message(DOMAIN)]) | |
| # Notify user of domain change | |
| await cl.Message( | |
| content=f"Domain changed to {DOMAIN}. My knowledge is now specialized for this domain.", | |
| author="Research Assistant" | |
| ).send() | |
| else: | |
| # Domain not found | |
| domains_list = "\n".join([f"- {domain}" for domain in AVAILABLE_DOMAINS]) | |
| await cl.Message( | |
| content=f"Domain '{requested_domain}' not found. Available domains:\n{domains_list}", | |
| author="System" | |
| ).send() | |
| return | |
| # Get current message history and debug mode setting | |
| messages = cl.user_session.get("messages") | |
| debug_mode = cl.user_session.get("debug_mode", DEBUG_MODE) | |
| # Add user message to history | |
| user_message = HumanMessage(content=message.content) | |
| # Create a temporary thinking message | |
| thinking = cl.Message(content="Researching your query...", author="Research Assistant") | |
| await thinking.send() | |
| try: | |
| # Call the agent | |
| response_messages = run_agent( | |
| user_input=message.content, | |
| domain=DOMAIN, | |
| messages=messages | |
| ) | |
| # Update the thinking message | |
| await thinking.remove() | |
| # Track tool usage for summary | |
| tool_usage = [] | |
| # Process and display messages | |
| for msg in response_messages: | |
| if not msg in messages: # Only process new messages | |
| if isinstance(msg, AIMessage): | |
| await cl.Message( | |
| content=msg.content, | |
| author="Research Assistant" | |
| ).send() | |
| elif isinstance(msg, ToolMessage): | |
| # Add to tool usage list | |
| tool_usage.append(msg.name) | |
| # Create elements for tool output | |
| elements = [] | |
| # Format the tool output for better readability | |
| formatted_content = msg.content | |
| # Try to detect JSON and format it | |
| if msg.content.strip().startswith('{') or msg.content.strip().startswith('['): | |
| try: | |
| import json | |
| content_obj = json.loads(msg.content) | |
| formatted_content = json.dumps(content_obj, indent=2) | |
| except: | |
| pass | |
| # Add tool output as an element | |
| elements.append( | |
| cl.Text( | |
| name=f"Tool Result: {msg.name}", | |
| content=formatted_content, | |
| display="inline" | |
| ) | |
| ) | |
| # Only show detailed tool messages in debug mode | |
| if debug_mode: | |
| # Send tool message with more prominent styling | |
| await cl.Message( | |
| content=f"🔍 **TOOL CALL**: `{msg.name}`\n\nI used this tool to find information about your query.", | |
| author="Research Assistant", | |
| elements=elements | |
| ).send() | |
| # Display tool usage summary if tools were used | |
| if tool_usage and debug_mode: | |
| tool_counts = {} | |
| for tool in tool_usage: | |
| if tool in tool_counts: | |
| tool_counts[tool] += 1 | |
| else: | |
| tool_counts[tool] = 1 | |
| summary = "📊 **RESEARCH SUMMARY**\n\nTo answer your question, I used:\n" | |
| for tool, count in tool_counts.items(): | |
| # Add emoji based on tool type | |
| emoji = "🔎" if tool == "web_search" else "📄" if tool == "research_paper_search" else "📚" if tool == "wikipedia_search" else "📊" | |
| summary += f"- {emoji} `{tool}`: {count} time{'s' if count > 1 else ''}\n" | |
| await cl.Message( | |
| content=summary, | |
| author="System" | |
| ).send() | |
| # Update session with new messages | |
| cl.user_session.set("messages", response_messages) | |
| except Exception as e: | |
| await thinking.remove() | |
| await cl.Message( | |
| content=f"Sorry, I encountered an error: {str(e)}", | |
| author="System" | |
| ).send() |