app.py

import os
import json
import time
import gradio as gr
from datetime import datetime
from typing import List, Dict, Any, Optional, Union
import threading
import re
import aiohttp
import asyncio

# Import Groq
from groq import Groq

class CreativeAgenticAI:
    """
    Creative Agentic AI Chat Tool using multiple providers (Groq and Chutes)
    """
    
    def __init__(self, groq_api_key: str = None, chutes_api_key: str = None, provider: str = "groq", model: str = None):
        """
        Initialize the Creative Agentic AI system.
        
        Args:
            groq_api_key: Groq API key
            chutes_api_key: Chutes API key
            provider: Which provider to use ('groq' or 'chutes')
            model: Which model to use
        """
        self.groq_api_key = groq_api_key
        self.chutes_api_key = chutes_api_key
        self.provider = provider
        self.conversation_history = []
        
        # Initialize clients based on provider
        if provider == "groq" and groq_api_key:
            if not groq_api_key:
                raise ValueError("No Groq API key provided")
            self.groq_client = Groq(api_key=groq_api_key)
            self.model = model or "compound-beta"
        elif provider == "chutes" and chutes_api_key:
            if not chutes_api_key:
                raise ValueError("No Chutes API key provided")
            self.model = model or "openai/gpt-oss-20b"
        else:
            raise ValueError(f"Invalid provider or missing API key for {provider}")
        
    async def _chutes_chat_async(self, messages: List[Dict], temperature: float = 0.7, max_tokens: int = 1024, stream: bool = False) -> Dict:
        """
        Async method for Chutes API chat with thinking support
        """
        headers = {
            "Authorization": f"Bearer {self.chutes_api_key}",
            "Content-Type": "application/json"
        }
        
        body = {
            "model": self.model,
            "messages": messages,
            "stream": stream,
            "max_tokens": max_tokens,
            "temperature": temperature
        }
        
        async with aiohttp.ClientSession() as session:
            async with session.post(
                "https://llm.chutes.ai/v1/chat/completions", 
                headers=headers,
                json=body
            ) as response:
                if response.status == 200:
                    if stream:
                        thinking_content = ""
                        final_content = ""
                        in_thinking = False
                        
                        async for line in response.content:
                            line = line.decode("utf-8").strip()
                            if line.startswith("data: "):
                                data = line[6:]
                                if data == "[DONE]":
                                    break
                                try:
                                    chunk_data = json.loads(data)
                                    if 'choices' in chunk_data and len(chunk_data['choices']) > 0:
                                        delta = chunk_data['choices'][0].get('delta', {})
                                        content = delta.get('content', '')
                                        
                                        if content:
                                            # Check for thinking tags
                                            if '<thinking>' in content:
                                                in_thinking = True
                                                thinking_content += content.replace('<thinking>', '')
                                            elif '</thinking>' in content:
                                                thinking_content += content.replace('</thinking>', '')
                                                in_thinking = False
                                            elif in_thinking:
                                                thinking_content += content
                                            else:
                                                final_content += content
                                                
                                except json.JSONDecodeError:
                                    continue
                        
                        return {
                            "thinking": thinking_content.strip(),
                            "content": final_content.strip()
                        }
                    else:
                        result = await response.json()
                        full_content = result['choices'][0]['message']['content']
                        
                        # Extract thinking and final content from non-streaming response
                        thinking_content = ""
                        final_content = full_content
                        
                        if '<thinking>' in full_content and '</thinking>' in full_content:
                            start_idx = full_content.find('<thinking>') + len('<thinking>')
                            end_idx = full_content.find('</thinking>')
                            thinking_content = full_content[start_idx:end_idx].strip()
                            final_content = full_content[end_idx + len('</thinking>'):].strip()
                        
                        return {
                            "thinking": thinking_content,
                            "content": final_content
                        }
                else:
                    error_text = await response.text()
                    raise Exception(f"Chutes API error: {response.status} - {error_text}")

    def _chutes_chat_sync(self, messages: List[Dict], temperature: float = 0.7, max_tokens: int = 1024, stream: bool = True) -> Dict:
        """
        Synchronous wrapper for Chutes API chat
        """
        try:
            loop = asyncio.get_event_loop()
        except RuntimeError:
            loop = asyncio.new_event_loop()
            asyncio.set_event_loop(loop)
        
        return loop.run_until_complete(
            self._chutes_chat_async(messages, temperature, max_tokens, stream)
        )
        
    def chat(self, message: str, 
             include_domains: List[str] = None,
             exclude_domains: List[str] = None,
             system_prompt: str = None,
             temperature: float = 0.7,
             max_tokens: int = 1024) -> Dict:
        """
        Send a message to the AI and get a response
        
        Args:
            message: User's message
            include_domains: List of domains to include for web search (Groq only)
            exclude_domains: List of domains to exclude from web search (Groq only)
            system_prompt: Custom system prompt
            temperature: Model temperature (0.0-2.0)
            max_tokens: Maximum tokens in response
            
        Returns:
            AI response with metadata
        """
        # Enhanced system prompt for better citation behavior
        if not system_prompt:
            if self.provider == "groq":
                citation_instruction = """
IMPORTANT: When you search the web and find information, you MUST:
1. Always cite your sources with clickable links in this format: [Source Title](URL)
2. Include multiple diverse sources when possible
3. Show which specific websites you used for each claim
4. At the end of your response, provide a "Sources Used" section with all the links
5. Be transparent about which information comes from which source
"""
                
                domain_context = ""
                if include_domains:
                    domain_context = f"\nYou are restricted to searching ONLY these domains: {', '.join(include_domains)}. Make sure to find and cite sources specifically from these domains."
                elif exclude_domains:
                    domain_context = f"\nAvoid searching these domains: {', '.join(exclude_domains)}. Search everywhere else on the web."
                
                system_prompt = f"""You are a creative and intelligent AI assistant with agentic capabilities. 
                You can search the web, analyze information, and provide comprehensive responses. 
                Be helpful, creative, and engaging while maintaining accuracy.
                
                {citation_instruction}
                {domain_context}
                
                Your responses should be well-structured, informative, and properly cited with working links."""
            else:
                # System prompt for Chutes thinking models
                system_prompt = """You are a creative and intelligent AI assistant with advanced reasoning capabilities. 
                Think through problems step-by-step, showing your reasoning process clearly.
                Be helpful, creative, and engaging while maintaining accuracy.
                Your responses should be well-structured, informative, and comprehensive.
                
                When solving complex problems, break them down into steps and explain your thinking process."""
        
        # Build messages
        messages = [{"role": "system", "content": system_prompt}]
        
        # Add conversation history (last 10 exchanges)
        messages.extend(self.conversation_history[-20:])  # Last 10 user-assistant pairs
        
        # Add current message with domain filtering context (Groq only)
        enhanced_message = message
        if self.provider == "groq" and (include_domains or exclude_domains):
            filter_context = []
            if include_domains:
                filter_context.append(f"ONLY search these domains: {', '.join(include_domains)}")
            if exclude_domains:
                filter_context.append(f"EXCLUDE these domains: {', '.join(exclude_domains)}")
            enhanced_message += f"\n\n[Domain Filtering: {' | '.join(filter_context)}]"
        
        messages.append({"role": "user", "content": enhanced_message})
        
        try:
            if self.provider == "groq":
                return self._handle_groq_chat(messages, include_domains, exclude_domains, temperature, max_tokens, message)
            elif self.provider == "chutes":
                return self._handle_chutes_chat(messages, temperature, max_tokens, message)
            else:
                raise ValueError(f"Unknown provider: {self.provider}")
                
        except Exception as e:
            error_msg = f"Error: {str(e)}"
            self.conversation_history.append({"role": "user", "content": message})
            self.conversation_history.append({"role": "assistant", "content": error_msg})
            
            return {
                "content": error_msg,
                "thinking": "",
                "timestamp": datetime.now().isoformat(),
                "model": self.model,
                "provider": self.provider,
                "tool_usage": None,
                "error": str(e)
            }

    def _handle_groq_chat(self, messages: List[Dict], include_domains: List[str], exclude_domains: List[str], 
                         temperature: float, max_tokens: int, original_message: str) -> Dict:
        """Handle Groq API chat"""
        # Set up API parameters
        params = {
            "messages": messages,
            "model": self.model,
            "temperature": temperature,
            "max_tokens": max_tokens
        }
        
        # Add domain filtering if specified
        if include_domains and include_domains[0].strip():
            params["include_domains"] = [domain.strip() for domain in include_domains if domain.strip()]
        if exclude_domains and exclude_domains[0].strip():
            params["exclude_domains"] = [domain.strip() for domain in exclude_domains if domain.strip()]
        
        # Make the API call
        response = self.groq_client.chat.completions.create(**params)
        content = response.choices[0].message.content
        
        # Extract tool usage information and enhance it
        tool_info = self._extract_tool_info(response)
        
        # Process content to enhance citations
        processed_content = self._enhance_citations(content, tool_info)
        
        # Add to conversation history
        self.conversation_history.append({"role": "user", "content": original_message})
        self.conversation_history.append({"role": "assistant", "content": processed_content})
        
        # Create response object
        return {
            "content": processed_content,
            "thinking": "",  # Groq doesn't have thinking process
            "timestamp": datetime.now().isoformat(),
            "model": self.model,
            "provider": "groq",
            "tool_usage": tool_info,
            "parameters": {
                "temperature": temperature,
                "max_tokens": max_tokens,
                "include_domains": include_domains,
                "exclude_domains": exclude_domains
            }
        }

    def _handle_chutes_chat(self, messages: List[Dict], temperature: float, max_tokens: int, original_message: str) -> Dict:
        """Handle Chutes API chat with thinking support"""
        result = self._chutes_chat_sync(messages, temperature, max_tokens, stream=True)
        
        thinking_content = result.get("thinking", "")
        final_content = result.get("content", "")
        
        # Add to conversation history (only store final content)
        self.conversation_history.append({"role": "user", "content": original_message})
        self.conversation_history.append({"role": "assistant", "content": final_content})
        
        # Create response object
        return {
            "content": final_content,
            "thinking": thinking_content,
            "timestamp": datetime.now().isoformat(),
            "model": self.model,
            "provider": "chutes",
            "tool_usage": None,  # Chutes doesn't have tool usage info
            "parameters": {
                "temperature": temperature,
                "max_tokens": max_tokens
            }
        }
    
    def _extract_tool_info(self, response) -> Dict:
        """Extract tool usage information in a JSON serializable format (Groq only)"""
        tool_info = {
            "tools_used": [],
            "search_queries": [],
            "sources_found": []
        }
        
        if hasattr(response.choices[0].message, 'executed_tools'):
            tools = response.choices[0].message.executed_tools
            if tools:
                for tool in tools:
                    tool_dict = {
                        "tool_type": getattr(tool, "type", "unknown"),
                        "tool_name": getattr(tool, "name", "unknown"),
                    }
                    
                    # Extract search queries and results
                    if hasattr(tool, "input"):
                        tool_input = str(tool.input)
                        tool_dict["input"] = tool_input
                        # Try to extract search query
                        if "search" in tool_dict["tool_name"].lower():
                            tool_info["search_queries"].append(tool_input)
                    
                    if hasattr(tool, "output"):
                        tool_output = str(tool.output)
                        tool_dict["output"] = tool_output
                        # Try to extract URLs from output
                        urls = self._extract_urls(tool_output)
                        tool_info["sources_found"].extend(urls)
                    
                    tool_info["tools_used"].append(tool_dict)
        
        return tool_info
    
    def _extract_urls(self, text: str) -> List[str]:
        """Extract URLs from text"""
        url_pattern = r'https?://[^\s<>"]{2,}'
        urls = re.findall(url_pattern, text)
        return list(set(urls))  # Remove duplicates
    
    def _enhance_citations(self, content: str, tool_info: Dict) -> str:
        """Enhance content with better citation formatting (Groq only)"""
        if not tool_info or not tool_info.get("sources_found"):
            return content
        
        # Add sources section if not already present
        if "Sources Used:" not in content and "sources:" not in content.lower():
            sources_section = "\n\n---\n\n### 📚 Sources Used:\n"
            for i, url in enumerate(tool_info["sources_found"][:10], 1):  # Limit to 10 sources
                # Try to extract domain name for better formatting
                domain = self._extract_domain(url)
                sources_section += f"{i}. [{domain}]({url})\n"
            
            content += sources_section
        
        return content
    
    def _extract_domain(self, url: str) -> str:
        """Extract domain name from URL for display"""
        try:
            if url.startswith(('http://', 'https://')):
                domain = url.split('/')[2]
                # Remove www. prefix if present
                if domain.startswith('www.'):
                    domain = domain[4:]
                return domain
            return url
        except:
            return url
    
    def clear_history(self):
        """Clear conversation history"""
        self.conversation_history = []
    
    def get_history_summary(self) -> str:
        """Get a summary of conversation history"""
        if not self.conversation_history:
            return "No conversation history"
        
        user_messages = [msg for msg in self.conversation_history if msg["role"] == "user"]
        assistant_messages = [msg for msg in self.conversation_history if msg["role"] == "assistant"]
        
        return f"Conversation: {len(user_messages)} user messages, {len(assistant_messages)} assistant responses"

# Global variables
ai_instance = None
current_provider = "groq"
api_key_status = {"groq": "Not Set", "chutes": "Not Set"}

def validate_api_keys(groq_api_key: str, chutes_api_key: str, provider: str, model: str) -> str:
    """Validate API keys and initialize AI instance"""
    global ai_instance, current_provider, api_key_status
    
    current_provider = provider
    
    if provider == "groq":
        if not groq_api_key or len(groq_api_key.strip()) < 10:
            api_key_status["groq"] = "Invalid ❌"
            return "❌ Please enter a valid Groq API key (should be longer than 10 characters)"
        
        try:
            # Test the Groq API key
            client = Groq(api_key=groq_api_key)
            test_response = client.chat.completions.create(
                messages=[{"role": "user", "content": "Hello"}],
                model=model,
                max_tokens=10
            )
            
            # Create AI instance
            ai_instance = CreativeAgenticAI(groq_api_key=groq_api_key, provider="groq", model=model)
            api_key_status["groq"] = "Valid ✅"
            
            return f"✅ Groq API Key Valid! Creative Agentic AI is ready.\n\n**Provider:** Groq\n**Model:** {model}\n**Status:** Connected with web search capabilities!"
            
        except Exception as e:
            api_key_status["groq"] = "Invalid ❌"
            ai_instance = None
            return f"❌ Error validating Groq API key: {str(e)}\n\nPlease check your API key and try again."
    
    elif provider == "chutes":
        if not chutes_api_key or len(chutes_api_key.strip()) < 10:
            api_key_status["chutes"] = "Invalid ❌"
            return "❌ Please enter a valid Chutes API key (should be longer than 10 characters)"
        
        try:
            # Test the Chutes API key with a simple request
            test_ai = CreativeAgenticAI(chutes_api_key=chutes_api_key, provider="chutes", model=model)
            test_response = test_ai._chutes_chat_sync(
                [{"role": "user", "content": "Hello"}], 
                temperature=0.7, 
                max_tokens=10,
                stream=False  # Use non-streaming for validation
            )
            
            # Create AI instance
            ai_instance = CreativeAgenticAI(chutes_api_key=chutes_api_key, provider="chutes", model=model)
            api_key_status["chutes"] = "Valid ✅"
            
            return f"✅ Chutes API Key Valid! Creative AI with Thinking is ready.\n\n**Provider:** Chutes\n**Model:** {model}\n**Status:** Connected with thinking model capabilities!"
            
        except Exception as e:
            api_key_status["chutes"] = "Invalid ❌"
            ai_instance = None
            return f"❌ Error validating Chutes API key: {str(e)}\n\nPlease check your API key and try again."

def get_available_models(provider: str) -> List[str]:
    """Get available models for the selected provider"""
    if provider == "groq":
        return ["compound-beta", "compound-beta-mini"]
    elif provider == "chutes":
        return ["openai/gpt-oss-20b", "zai-org/GLM-4.5-Air", "Qwen/Qwen3-8B"]
    return []

def update_model_choices(provider: str):
    """Update model choices based on provider selection"""
    models = get_available_models(provider)
    return gr.Radio(choices=models, value=models[0] if models else None, label=f"🧠 {provider.title()} Models")

def chat_with_ai(message: str, 
                include_domains: str, 
                exclude_domains: str,
                system_prompt: str,
                temperature: float,
                max_tokens: int,
                history: List,
                show_thinking: bool = True) -> tuple:
    """Main chat function with thinking support"""
    global ai_instance, current_provider
    
    if not ai_instance:
        error_msg = f"⚠️ Please set your {current_provider.title()} API key first!"
        history.append([message, error_msg])
        return history, ""
    
    if not message.strip():
        return history, ""
    
    # Process domain lists (only for Groq)
    include_list = None
    exclude_list = None
    if current_provider == "groq":
        include_list = [d.strip() for d in include_domains.split(",")] if include_domains.strip() else []
        exclude_list = [d.strip() for d in exclude_domains.split(",")] if exclude_domains.strip() else []
    
    try:
        # Get AI response
        response = ai_instance.chat(
            message=message,
            include_domains=include_list if include_list else None,
            exclude_domains=exclude_list if exclude_list else None,
            system_prompt=system_prompt if system_prompt.strip() else None,
            temperature=temperature,
            max_tokens=int(max_tokens)
        )
        
        # Format response with thinking (if available and enabled)
        ai_response = ""
        
        # Add thinking section for Chutes models
        if current_provider == "chutes" and response.get("thinking") and show_thinking:
            thinking_content = response["thinking"].strip()
            if thinking_content:
                ai_response += f"### 🤔 **Model's Thinking Process:**\n\n"
                ai_response += f"*The model is reasoning through your question...*\n\n"
                ai_response += f"```thinking\n{thinking_content}\n```\n\n"
                ai_response += "---\n\n### 💡 **Final Response:**\n\n"
        
        # Add main content
        ai_response += response["content"]
        
        # Add enhanced tool usage info (Groq only)
        if response.get("tool_usage") and current_provider == "groq":
            tool_info = response["tool_usage"]
            tool_summary = []
            
            if tool_info.get("search_queries"):
                tool_summary.append(f"🔍 Search queries: {len(tool_info['search_queries'])}")
            
            if tool_info.get("sources_found"):
                tool_summary.append(f"📄 Sources found: {len(tool_info['sources_found'])}")
            
            if tool_info.get("tools_used"):
                tool_summary.append(f"🔧 Tools used: {len(tool_info['tools_used'])}")
            
            if tool_summary:
                ai_response += f"\n\n*{' | '.join(tool_summary)}*"
        
        # Add domain filtering info (Groq only)
        if current_provider == "groq" and (include_list or exclude_list):
            filter_info = []
            if include_list:
                filter_info.append(f"✅ Included domains: {', '.join(include_list)}")
            if exclude_list:
                filter_info.append(f"❌ Excluded domains: {', '.join(exclude_list)}")
            
            ai_response += f"\n\n*🌐 Domain filtering applied: {' | '.join(filter_info)}*"
        
        # Add provider and thinking info
        provider_info = f"🤖 Powered by: {current_provider.title()} ({response.get('model', 'unknown')})"
        if current_provider == "chutes" and response.get("thinking"):
            if show_thinking:
                provider_info += " | 🤔 Thinking process shown"
            else:
                provider_info += " | 🤔 Thinking process hidden"
        
        ai_response += f"\n\n*{provider_info}*"
        
        # Add to history
        history.append([message, ai_response])
        
        return history, ""
        
    except Exception as e:
        error_msg = f"❌ Error: {str(e)}"
        history.append([message, error_msg])
        return history, ""

def clear_chat_history():
    """Clear the chat history"""
    global ai_instance
    if ai_instance:
        ai_instance.clear_history()
    return []

def create_gradio_app():
    """Create the main Gradio application"""
    
    # Custom CSS for better styling
    css = """
    .container {
        max-width: 1200px;
        margin: 0 auto;
    }
    .header {
        text-align: center;
        background: linear-gradient(to right, #00ff94, #00b4db);
        color: white;
        padding: 20px;
        border-radius: 10px;
        margin-bottom: 20px;
    }
    .status-box {
        background-color: #f8f9fa;
        border: 1px solid #dee2e6;
        border-radius: 8px;
        padding: 15px;
        margin: 10px 0;
    }
    .example-box {
        background-color: #e8f4fd;
        border-left: 4px solid #007bff;
        padding: 15px;
        margin: 10px 0;
        border-radius: 0 8px 8px 0;
    }
    .domain-info {
        background-color: #fff3cd;
        border: 1px solid #ffeaa7;
        border-radius: 8px;
        padding: 15px;
        margin: 10px 0;
    }
    .citation-info {
        background-color: #d1ecf1;
        border: 1px solid #bee5eb;
        border-radius: 8px;
        padding: 15px;
        margin: 10px 0;
    }
    .provider-info {
        background-color: #f8d7da;
        border: 1px solid #f5c6cb;
        border-radius: 8px;
        padding: 15px;
        margin: 10px 0;
    }
    .thinking-info {
        background-color: #e2e3e5;
        border: 1px solid #d6d8db;
        border-radius: 8px;
        padding: 15px;
        margin: 10px 0;
    }
    #neuroscope-accordion {
        background: linear-gradient(to right, #00ff94, #00b4db); 
        border-radius: 8px;
    }
    #neuroscope-accordion2 {
        background: linear-gradient(to right, #00ff94, #00b4db); 
        border-radius: 8px;
        margin-top: 10px;
    }
    """
    
    with gr.Blocks(css=css, title="🤖 Multi-Provider Creative Agentic AI Chat with Thinking", theme=gr.themes.Ocean()) as app:
        
        # Header
        gr.HTML("""
        <div class="header">
            <h1>🤖 NeuroScope-AI Enhanced with Thinking Models</h1>
            <p>Multi-Provider AI Chat Tool - Groq's Compound Models & Chutes Thinking Models</p>
        </div>
        """)

        # Provider Selection
        with gr.Group():
            with gr.Accordion("🤖 Multi-Provider NeuroScope AI with Thinking", open=False, elem_id="neuroscope-accordion"):
                gr.Markdown("""
                    **Enhanced with Multiple AI Providers & Thinking Models:**
                    - 🧠 Intelligence (Neuro) - Now supports Groq & Chutes Thinking Models
                    - 🔍 Advanced capabilities (Scope) - Web search with Groq, reasoning traces with Chutes
                    - 🤖 AI capabilities (AI) - Multiple model options including thinking models
                    - ⚡ Precision & Speed (Scope) - Choose the best provider for your needs
                    - 🤔 **NEW**: Thinking process visualization for Chutes models
                """)

        # Thinking Models Info
        with gr.Group():
            with gr.Accordion("🤔 About Chutes Thinking Models", open=False, elem_id="neuroscope-accordion"):
                gr.Markdown("""
                <div class="thinking-info">
                <h3>🧠 What are Thinking Models?</h3>
                <p><strong>Chutes Thinking Models</strong> are advanced AI systems that show their reasoning process before providing the final answer.</p>
                
                <h4>🔍 How They Work:</h4>
                <ul>
                    <li><strong>Step-by-Step Reasoning:</strong> Models think through problems systematically</li>
                    <li><strong>Transparent Process:</strong> You can see exactly how the AI reaches its conclusions</li>
                    <li><strong>Better Accuracy:</strong> The thinking process often leads to more accurate and well-reasoned responses</li>
                    <li><strong>Educational Value:</strong> Learn from the AI's problem-solving approach</li>
                </ul>
                
                <h4>🎯 Available Thinking Models:</h4>
                <ul>
                    <li><strong>openai/gpt-oss-20b:</strong> Large-scale reasoning and analysis</li>
                    <li><strong>meta-llama/llama-3.1-8b-instruct:</strong> Efficient thinking and instruction following</li>
                    <li><strong>anthropic/claude-3-sonnet:</strong> Advanced reasoning and creative thinking</li>
                </ul>
                
                <h4>💡 Best For:</h4>
                <ul>
                    <li>Complex problem-solving tasks</li>
                    <li>Mathematical and logical reasoning</li>
                    <li>Step-by-step analysis</li>
                    <li>Educational explanations</li>
                    <li>Creative writing with detailed planning</li>
                </ul>
                </div>
                """)

        # Provider and API Key Section
        with gr.Row():
                        
            with gr.Column():
                provider_selection = gr.Radio(
                    choices=["groq", "chutes"],
                    label="🏢 AI Provider",
                    value="groq",
                    info="Choose your AI provider"
                )
                
                # API Key inputs
                groq_api_key = gr.Textbox(
                    label="🔑 Groq API Key",
                    placeholder="Enter your Groq API key here...",
                    type="password",
                    info="Get your API key from: https://console.groq.com/",
                    visible=True
                )
                
                chutes_api_key = gr.Textbox(
                    label="🔑 Chutes API Key",
                    placeholder="Enter your Chutes API key here...",
                    type="password",
                    info="Get your API key from: https://chutes.ai/",
                    visible=False
                )
                
                model_selection = gr.Radio(
                    choices=get_available_models("groq"),
                    label="🧠 Groq Models",
                    value="compound-beta",
                    info="compound-beta: More powerful | compound-beta-mini: Faster"
                )
                
                # Thinking toggle for Chutes
                show_thinking = gr.Checkbox(
                    label="🤔 Show Thinking Process",
                    value=True,
                    info="Display the model's reasoning process (Chutes only)",
                    visible=False
                )
                
                connect_btn = gr.Button("🔗 Connect", variant="primary", size="lg")
        
        # Status display
        status_display = gr.Markdown("### 📊 Status: Not connected", elem_classes=["status-box"])

        # Provider Information
        with gr.Group():
            with gr.Accordion("🏢 Provider Comparison", open=False, elem_id="neuroscope-accordion"):
                gr.Markdown("""
                <div class="provider-info">
                <h3>🆚 Groq vs Chutes Comparison</h3>
                
                **🚀 Groq (Compound Models)**
                - ✅ **Web Search Capabilities** - Can search the internet and cite sources
                - ✅ **Agentic Tools** - Advanced tool usage and autonomous web browsing
                - ✅ **Domain Filtering** - Control which websites to search
                - ✅ **Citation System** - Automatic source linking and references
                - ⚡ **Ultra-fast inference** - Groq's hardware acceleration
                - 🧠 **Models**: compound-beta, compound-beta-mini
                - ❌ **No thinking process** - Direct responses without visible reasoning
                
                **🎯 Chutes (Thinking Models)**
                - ✅ **Multiple Model Access** - Various open-source and commercial models
                - ✅ **Thinking Process** - See the model's step-by-step reasoning
                - ✅ **High-quality reasoning** - Better accuracy through visible thinking
                - ✅ **Educational Value** - Learn from AI's problem-solving approach
                - ⚡ **Good performance** - Reliable and fast responses
                - 🧠 **Models**: GPT-OSS-20B, Llama 3.1, Claude 3 Sonnet (all with thinking)
                - ❌ **No web search** - Relies on training data only
                
                **💡 Use Groq when you need:**
                - Real-time information and web search
                - Research with source citations
                - Domain-specific searches
                - Ultra-fast responses
                
                **💡 Use Chutes when you need:**
                - Complex problem-solving with visible reasoning
                - Educational explanations
                - Mathematical and logical analysis
                - Creative planning with detailed thinking
                - Understanding AI's reasoning process
                </div>
                """)
        
        # Update UI based on provider selection
        def update_provider_ui(provider):
            groq_visible = provider == "groq"
            chutes_visible = provider == "chutes"
            models = get_available_models(provider)
            
            return (
                gr.update(visible=groq_visible),  # groq_api_key
                gr.update(visible=chutes_visible),  # chutes_api_key
                gr.update(choices=models, value=models[0] if models else None, 
                         label=f"🧠 {provider.title()} Models"),  # model_selection
                gr.update(visible=chutes_visible),  # show_thinking
                gr.update(visible=groq_visible), 
            )
        
        provider_selection.change(
            fn=update_provider_ui,
            inputs=[provider_selection],
            outputs=[groq_api_key, chutes_api_key, model_selection, show_thinking]  
        )
        
        # Connect button functionality
        connect_btn.click(
            fn=validate_api_keys,
            inputs=[groq_api_key, chutes_api_key, provider_selection, model_selection],
            outputs=[status_display]
        )
        
        # Main Chat Interface
        with gr.Tab("💬 Chat"):
            chatbot = gr.Chatbot(
                label="Multi-Provider Creative AI Assistant with Thinking",
                height=600,
                show_label=True,
                bubble_full_width=False,
                show_copy_button=True
            )
            
            with gr.Row():
                msg = gr.Textbox(
                    label="Your Message",
                    placeholder="Type your message here...",
                    lines=3
                )
                with gr.Column():
                    send_btn = gr.Button("📤 Send", variant="primary")
                    clear_btn = gr.Button("🗑️ Clear", variant="secondary")
            
        # Advanced Settings
        with gr.Accordion("⚙️ Advanced Settings", open=False, elem_id="neuroscope-accordion"):
            with gr.Row():
                temperature = gr.Slider(
                    minimum=0.0,
                    maximum=2.0,
                    value=0.7,
                    step=0.1,
                    label="🌡️ Temperature",
                    info="Higher = more creative, Lower = more focused"
                )
                max_tokens = gr.Slider(
                    minimum=100,
                    maximum=4000,
                    value=1024,
                    step=100,
                    label="📝 Max Tokens",
                    info="Maximum length of response"
                )
                
            system_prompt = gr.Textbox(
                label="🎭 Custom System Prompt",
                placeholder="Override the default system prompt...",
                lines=2,
                info="Leave empty to use provider-optimized default prompt"
            )    
        
        # Domain Filtering Section (Groq only)
        with gr.Group():
            with gr.Accordion("🌐 Domain Filtering (Groq Web Search Only)", open=False, elem_id="neuroscope-accordion"):
                gr.Markdown("""
                <div class="domain-info">
                <h4>🔍 Domain Filtering Guide (Groq Only)</h4>
                <p><strong>Note:</strong> Domain filtering only works with Groq's compound models that have web search capabilities.</p>
                <p>Control which websites the AI can search when answering your questions:</p>
                <ul>
                    <li><strong>Include Domains:</strong> Only search these domains (comma-separated)</li>
                    <li><strong>Exclude Domains:</strong> Never search these domains (comma-separated)</li>
                    <li><strong>Examples:</strong> arxiv.org, *.edu, github.com, stackoverflow.com</li>
                    <li><strong>Wildcards:</strong> Use *.edu for all educational domains</li>
                </ul>
                <p><strong>New:</strong> Domain filtering status will be shown in responses!</p>
                </div>
                """)
                
                with gr.Row():
                    include_domains = gr.Textbox(
                        label="✅ Include Domains (comma-separated)",
                        placeholder="arxiv.org, *.edu, github.com, stackoverflow.com",
                        info="Only search these domains"
                    )
                    exclude_domains = gr.Textbox(
                        label="❌ Exclude Domains (comma-separated)", 
                        placeholder="wikipedia.org, reddit.com, twitter.com",
                        info="Never search these domains"
                    )
                    
            with gr.Accordion("🔗 Common Domain Examples", open=False, elem_id="neuroscope-accordion2"):
                gr.Markdown("""
                **Academic & Research:**
                - `arxiv.org`, `*.edu`, `scholar.google.com`, `researchgate.net`
                
                **Technology & Programming:**
                - `github.com`, `stackoverflow.com`, `docs.python.org`, `developer.mozilla.org`
                
                **News & Media:**
                - `reuters.com`, `bbc.com`, `npr.org`, `apnews.com`
                
                **Business & Finance:**
                - `bloomberg.com`, `wsj.com`, `nasdaq.com`, `sec.gov`
                
                **Science & Medicine:**
                - `nature.com`, `science.org`, `pubmed.ncbi.nlm.nih.gov`, `who.int`
                """)

        # Update provider UI function with domain filtering and thinking toggle
        def update_provider_ui_complete(provider):
            groq_visible = provider == "groq"
            chutes_visible = provider == "chutes"
            models = get_available_models(provider)
            
            return (
                gr.update(visible=groq_visible),  # groq_api_key
                gr.update(visible=chutes_visible),  # chutes_api_key
                gr.update(choices=models, value=models[0] if models else None, 
                         label=f"🧠 {provider.title()} Models"),  # model_selection
                gr.update(visible=chutes_visible),  # show_thinking
                gr.update(visible=groq_visible),  
            )
        
        provider_selection.change(
            fn=update_provider_ui_complete,
            inputs=[provider_selection],
            outputs=[groq_api_key, chutes_api_key, model_selection, show_thinking]
        )
        
        # IMPORTANT Section with Citation Info and Thinking
        with gr.Group():
            with gr.Accordion("📚 IMPORTANT - Citations & Thinking Models!", open=False, elem_id="neuroscope-accordion"):
                gr.Markdown("""
                <div class="citation-info">
                <h3>🆕 Multi-Provider Enhancement with Thinking Models</h3>
                <p>This enhanced version now supports both Groq and Chutes AI providers:</p>
                <ul>
                    <li><strong>🚀 Groq Integration:</strong> Agentic AI with web search, citations, and tool usage</li>
                    <li><strong>🎯 Chutes Integration:</strong> Multiple thinking models with visible reasoning processes</li>
                    <li><strong>🤔 Thinking Process:</strong> See step-by-step reasoning from Chutes models</li>
                    <li><strong>🔄 Easy Switching:</strong> Switch between providers based on your needs</li>
                    <li><strong>📊 Provider Comparison:</strong> Clear information about each provider's strengths</li>
                </ul>
                
                <h3>🔗 Groq Enhanced Citation System</h3>
                <p>When using Groq, you get:</p>
                <ul>
                    <li><strong>Automatic Source Citations:</strong> All responses include clickable links to sources</li>
                    <li><strong>Sources Used Section:</strong> Dedicated section showing all websites referenced</li>
                    <li><strong>Domain Filtering Verification:</strong> Clear indication when domain filtering is applied</li>
                    <li><strong>Search Query Tracking:</strong> Shows what queries were made to find information</li>
                </ul>
                
                <h3>🤔 Chutes Thinking Model Features</h3>
                <p>When using Chutes, you get access to:</p>
                <ul>
                    <li><strong>Thinking Process Visualization:</strong> See exactly how the AI reasons through problems</li>
                    <li><strong>Step-by-Step Analysis:</strong> Watch the model break down complex questions</li>
                    <li><strong>Toggle Thinking Display:</strong> Choose to show or hide the reasoning process</li>
                    <li><strong>Educational Value:</strong> Learn from AI's problem-solving approaches</li>
                    <li><strong>Multiple Thinking Models:</strong> Different models with unique reasoning styles</li>
                </ul>
                </div>
                
                ### 🔍 **Web Search Behavior (Groq Only)**
                
                **No Domains Specified:**
                  - AI operates with **unrestricted web search capabilities**.
                  - Compound models autonomously search the **entire internet** for the most relevant and up-to-date information.
                  - AI has complete freedom to use its **agentic tools** and browse **any website** it finds useful.
                
                **Include Domains Specified (e.g., `arxiv.org`, `*.edu`):**
                  - AI is restricted to search **only the specified domains**.
                  - Acts as a **strict whitelist**, making the AI **laser-focused** on your chosen sources.
                  - Ensures information is sourced from **preferred or authoritative domains** (e.g., academic or research-focused).
                
                **Exclude Domains Specified (e.g., `wikipedia.org`, `reddit.com`):**
                  - AI searches the entire web **except the listed domains**.
                  - Useful for **filtering out unreliable or unwanted sources**.
                  - Allows broad search with **targeted exclusions**.
                
                ### 🤔 **Thinking Models Behavior (Chutes Only)**
                
                **How Thinking Models Work:**
                  - Models first **reason through the problem** in a thinking section
                  - You can see the **step-by-step thought process**
                  - The model then provides its **final polished answer**
                  - **Toggle thinking display** on/off as needed
                
                **Best Use Cases for Thinking Models:**
                  - Complex mathematical problems
                  - Logical reasoning tasks
                  - Creative writing with planning
                  - Educational explanations
                  - Problem-solving scenarios
                
                ---
                
                ### 🎭 **Custom System Prompt Feature**
                
                Allows complete override of the AI's **default personality and behavior** for both providers.
                You can redefine the AI to act as:
                  - A **professional business consultant**
                  - A **coding mentor**
                  - A **creative writer**
                  - A **step-by-step tutor** (especially effective with thinking models)
                  - A **specific character or persona**
                - Provides full control to **reshape the AI's tone, expertise, and conversational style** with a single prompt.
                """)
        
        # How to Use Section
        with gr.Accordion("📖 How to Use This Enhanced Multi-Provider App with Thinking", open=False, elem_id="neuroscope-accordion"):
            gr.Markdown("""
            ### 🚀 Getting Started
            1. **Choose your AI Provider** - Select between Groq (web search + agentic) or Chutes (thinking models)
            2. **Enter your API Key** - 
               - Groq: Get one from [console.groq.com](https://console.groq.com/)
               - Chutes: Get one from [chutes.ai](https://chutes.ai/)
            3. **Select a model** - Choose from provider-specific model options
            4. **Configure thinking display** - For Chutes, decide if you want to see the reasoning process
            5. **Click Connect** - Validate your key and connect to the AI
            6. **Start chatting!** - Type your message and get intelligent responses
            
            ### 🎯 Key Features
            **🚀 Groq Features:**
            - **Agentic AI**: The AI can use tools and search the web autonomously
            - **Smart Citations**: Automatic source linking and citation formatting
            - **Domain Filtering**: Control which websites the AI searches
            - **Ultra-fast**: Groq's hardware-accelerated inference
            
            **🎯 Chutes Features:**
            - **Thinking Models**: See the model's step-by-step reasoning process
            - **Multiple Models**: Access to GPT-OSS-20B, Llama 3.1, Claude 3 Sonnet
            - **Educational**: Learn from AI's problem-solving approaches
            - **Toggle Thinking**: Show/hide the reasoning process as needed
            - **High Quality**: Excellent reasoning and analysis capabilities
            
            **🔄 Universal Features:**
            - **Memory**: Maintains conversation context throughout the session
            - **Customizable**: Adjust temperature, tokens, and system prompts
            - **Provider Switching**: Easy switching between AI providers
            
            ### 💡 Tips for Best Results
            **For Groq:**
            - Be specific in your questions for better web search results
            - Use domain filtering for specialized research
            - Check the "Sources Used" section for all references
            - Try different domain combinations to see varied results
            
            **For Chutes (Thinking Models):**
            - Ask complex, multi-step questions to see rich thinking processes
            - Use for educational purposes - the thinking is very instructive
            - Try mathematical problems, logical puzzles, or creative planning tasks
            - Toggle thinking display based on whether you want to see the process
            - Different models have different thinking styles - experiment!
            
            **General:**
            - Adjust temperature: higher for creativity, lower for precision
            - Try different system prompts for different conversation styles
            - Use the provider that best fits your current task
            - For learning: use Chutes with thinking display enabled
            - For research: use Groq with appropriate domain filtering
            """)
        
        # Sample Examples Section
        with gr.Accordion("🎯 Sample Examples to Test Both Providers & Thinking", open=False, elem_id="neuroscope-accordion"):
            gr.Markdown("""
            <div class="example-box">
            <h4>🆚 Provider Comparison Examples</h4>
            <p>Try the same prompts with both providers to see the difference:</p>
            
            <h4>🤔 Perfect for Thinking Models (Chutes)</h4>
            <ul>
                <li><strong>Math & Logic:</strong> "Solve this step by step: If a train travels 120 miles in 2 hours, and then 180 miles in 3 hours, what's the average speed for the entire journey?"</li>
                <li><strong>Problem Solving:</strong> "I have a budget of $1000 for a home office setup. Help me plan the best allocation across desk, chair, computer, and lighting."</li>
                <li><strong>Creative Planning:</strong> "Plan a short story about a time traveler who accidentally changes history. Walk through the plot structure."</li>
                <li><strong>Analysis:</strong> "Compare and contrast the pros and cons of remote work vs office work, considering productivity, collaboration, and work-life balance."</li>
            </ul>
            
            <h4>🔬 Research & Real-time Info (Groq)</h4>
            <ul>
                <li><strong>Current Events:</strong> "What are the latest developments in AI research in 2024?"</li>
                <li><strong>Tech Updates:</strong> "What are the newest features in React 19?"</li>
                <li><strong>Market Analysis:</strong> "Current trends in cryptocurrency markets with sources"</li>
                <li><strong>Scientific Updates:</strong> "Recent breakthroughs in quantum computing research"</li>
            </ul>
            
            <h4>💻 Programming & Tech (Compare Both)</h4>
            <ul>
                <li><strong>Groq:</strong> "What are the current best practices for React 18 in 2024?" (with web search)</li>
                <li><strong>Chutes:</strong> "Explain how to build a React component with useState, and walk through your reasoning for the design choices"</li>
            </ul>
            
            <h4>🎨 Creative Tasks (Great for Thinking Models)</h4>
            <ul>
                <li>"Write a marketing strategy for a new eco-friendly product, showing your planning process"</li>
                <li>"Create a study plan for learning Python in 3 months, explaining your reasoning for each phase"</li>
                <li>"Design a mobile app concept for meditation, walking through your design thinking"</li>
            </ul>
            
            <h4>🧠 Model-Specific Testing (Chutes Thinking)</h4>
            <ul>
                <li><strong>GPT-OSS-20B:</strong> "Explain quantum entanglement in simple terms, showing your thought process for making it accessible"</li>
                <li><strong>Llama 3.1:</strong> "Debug this Python code and explain your debugging approach: [code with intentional errors]"</li>
                <li><strong>Claude 3 Sonnet:</strong> "Analyze this business scenario and provide strategic recommendations, showing your analytical framework"</li>
            </ul>
            
            <h4>📊 Side-by-Side Comparisons</h4>
            <ul>
                <li><strong>Same Question, Different Providers:</strong> Ask "How do neural networks work?" to both providers</li>
                <li><strong>Groq Result:</strong> Fast response with potential web sources and current information</li>
                <li><strong>Chutes Result:</strong> Detailed thinking process showing how the model breaks down the explanation</li>
            </ul>
            </div>
            """)
        
        # Event handlers - Updated to include thinking toggle
        send_btn.click(
            fn=chat_with_ai,
            inputs=[msg, include_domains, exclude_domains, system_prompt, temperature, max_tokens, chatbot, show_thinking],
            outputs=[chatbot, msg]
        )
        
        msg.submit(
            fn=chat_with_ai,
            inputs=[msg, include_domains, exclude_domains, system_prompt, temperature, max_tokens, chatbot, show_thinking],
            outputs=[chatbot, msg]
        )
        
        clear_btn.click(
            fn=clear_chat_history,
            outputs=[chatbot]
        )
        
        # Footer
        with gr.Accordion("🚀 About This Enhanced Multi-Provider Tool with Thinking", open=True, elem_id="neuroscope-accordion"):
            gr.Markdown("""
            **Enhanced Multi-Provider Creative Agentic AI Chat Tool** with thinking model support:
            
            **🆕 New Thinking Model Features:**
            - 🤔 **Thinking Process Visualization**: See step-by-step reasoning from Chutes models
            - 🧠 **Multiple Thinking Models**: GPT-OSS-20B, Llama 3.1, Claude 3 Sonnet with reasoning
            - 🔄 **Toggle Thinking Display**: Choose to show or hide the reasoning process
            - 📚 **Educational Value**: Learn from AI's problem-solving approaches
            - 🎯 **Better Accuracy**: Thinking process often leads to more accurate responses
            
            **🚀 Groq Features:**
            - 🔗 **Automatic Source Citations**: Every response includes clickable links to sources
            - 📚 **Sources Used Section**: Dedicated section showing all websites referenced  
            - 🌐 **Domain Filtering Verification**: Clear indication when filtering is applied
            - 🔍 **Search Query Tracking**: Shows what queries were made
            - ⚡ **Enhanced Tool Usage Display**: Better visibility into AI's research process
            - 🔍 Web search with domain filtering
            - 🧠 Advanced AI reasoning with tool usage
            
            **🎯 Chutes Features:**
            - 🤖 **Multiple Thinking Models**: Various AI models with visible reasoning
            - 💰 **Cost-Effective**: Competitive pricing for AI access
            - 🎨 **Creative Excellence**: Optimized for reasoning and analysis tasks
            - ⚡ **Reliable Performance**: Consistent and thoughtful responses
            - 📖 **Learning Tool**: Perfect for understanding AI reasoning
            
            **🔄 Universal Features:**
            - 💬 Conversational memory and context
            - ⚙️ Customizable parameters and prompts
            - 🎨 Creative and analytical capabilities
            - 🌟 Enhanced user interface with provider-specific optimizations
            
            **💡 Choose Your Provider:**
            - **Use Groq** when you need real-time information, web search, and citations
            - **Use Chutes** when you want to see the thinking process, need complex reasoning, or want educational explanations
            - **Toggle thinking display** in Chutes to customize your experience
            
            **🎓 Perfect for:**
            - Students learning problem-solving approaches
            - Developers wanting to understand AI reasoning
            - Researchers needing both current information (Groq) and deep analysis (Chutes)
            - Anyone curious about how AI thinks through problems
            """)
    
    return app

# Main execution
if __name__ == "__main__":
    app = create_gradio_app()
    app.launch(
        share=True
    )