app.py

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

@cl.on_chat_start
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()

@cl.on_settings_update
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()

@cl.on_message
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()