Upload 2 files

app.py

@@ -0,0 +1,516 @@
+#!/usr/bin/env python3
+"""
+Just search - A Smart Search Agent using Menlo/Lucy-128k
+Part of the Just, AKA Simple series
+Built with Gradio, DuckDuckGo Search, and Hugging Face Transformers
+"""
+
+import gradio as gr
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+from duckduckgo_search import DDGS
+import json
+import re
+import time
+from typing import List, Dict, Tuple
+import spaces
+
+# Initialize the model and tokenizer globally for efficiency
+MODEL_NAME = "Menlo/Lucy-128k"
+tokenizer = None
+model = None
+search_pipeline = None
+
+def initialize_model():
+    """Initialize the Menlo/Lucy-128k model and tokenizer"""
+    global tokenizer, model, search_pipeline
+    try:
+        tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
+        model = AutoModelForCausalLM.from_pretrained(
+            MODEL_NAME,
+            torch_dtype=torch.float16,
+            device_map="auto",
+            trust_remote_code=True
+        )
+        search_pipeline = pipeline(
+            "text-generation",
+            model=model,
+            tokenizer=tokenizer,
+            torch_dtype=torch.float16,
+            device_map="auto",
+            max_new_tokens=2048,
+            temperature=0.7,
+            do_sample=True,
+            pad_token_id=tokenizer.eos_token_id
+        )
+        return True
+    except Exception as e:
+        print(f"Error initializing model: {e}")
+        return False
+
+def extract_thinking_and_response(text: str) -> Tuple[str, str]:
+    """Extract thinking process and clean response from AI output"""
+    thinking = ""
+    response = text
+    
+    # Extract thinking content
+    thinking_match = re.search(r'<think>(.*?)</think>', text, re.DOTALL)
+    if thinking_match:
+        thinking = thinking_match.group(1).strip()
+        response = re.sub(r'<think>.*?</think>', '', text, flags=re.DOTALL)
+    
+    # Clean up the response
+    response = re.sub(r'^(Assistant:|AI:|Response:|Answer:)\s*', '', response.strip())
+    response = re.sub(r'\[INST\].*?\[\/INST\]', '', response)
+    response = re.sub(r'<\|.*?\|>', '', response)
+    
+    return thinking.strip(), response.strip()
+
+def clean_response(text: str) -> str:
+    """Clean up the AI response to extract just the relevant content"""
+    _, response = extract_thinking_and_response(text)
+    return response
+
+@spaces.GPU
+def generate_search_queries(user_query: str) -> Tuple[List[str], str]:
+    """Generate multiple search queries based on user input using AI"""
+    prompt = f"""<|begin_of_text|><|start_header_id|>system<|end_header_id|>
+You are a search query generator. Given a user's question, generate 3-5 different search queries that would help find comprehensive information to answer their question. Return only the search queries, one per line, without numbering or bullet points.
+
+Example:
+User: "What are the latest developments in AI?"
+latest AI developments 2024
+artificial intelligence breakthroughs recent
+AI technology advances news
+machine learning innovations 2024
+
+<|eot_id|><|start_header_id|>user<|end_header_id|>
+{user_query}
+<|eot_id|><|start_header_id|>assistant<|end_header_id|>"""
+
+    try:
+        response = search_pipeline(prompt, max_new_tokens=200, temperature=0.3)
+        generated_text = response[0]['generated_text']
+        
+        # Extract just the assistant's response
+        assistant_response = generated_text.split('<|start_header_id|>assistant<|end_header_id|>')[-1]
+        thinking, cleaned_response = extract_thinking_and_response(assistant_response)
+        
+        # Split into individual queries and clean them
+        queries = [q.strip() for q in cleaned_response.split('\n') if q.strip()]
+        # Filter out any non-query text
+        queries = [q for q in queries if len(q) > 5 and not q.startswith('Note:') and not q.startswith('Example:')]
+        
+        return queries[:5], thinking  # Return max 5 queries and thinking
+    except Exception as e:
+        print(f"Error generating queries: {e}")
+        # Fallback to simple query variations
+        return [user_query, f"{user_query} 2024", f"{user_query} latest"], ""
+
+def search_web(queries: List[str]) -> List[Dict]:
+    """Search the web using DuckDuckGo with multiple queries"""
+    all_results = []
+    ddgs = DDGS()
+    
+    for query in queries:
+        try:
+            results = ddgs.text(query, max_results=5, region='wt-wt', safesearch='moderate')
+            for result in results:
+                result['search_query'] = query
+                all_results.append(result)
+            time.sleep(0.5)  # Rate limiting
+        except Exception as e:
+            print(f"Error searching for '{query}': {e}")
+            continue
+    
+    # Remove duplicates based on URL
+    seen_urls = set()
+    unique_results = []
+    for result in all_results:
+        if result['href'] not in seen_urls:
+            seen_urls.add(result['href'])
+            unique_results.append(result)
+    
+    return unique_results[:15]  # Return max 15 results
+
+@spaces.GPU
+def filter_relevant_results(user_query: str, search_results: List[Dict]) -> Tuple[List[Dict], str]:
+    """Use AI to filter and rank search results by relevance"""
+    if not search_results:
+        return [], ""
+    
+    # Prepare results summary for AI
+    results_text = ""
+    for i, result in enumerate(search_results[:12]):  # Limit to avoid token overflow
+        results_text += f"{i+1}. Title: {result.get('title', 'No title')}\n"
+        results_text += f"   URL: {result.get('href', 'No URL')}\n"
+        results_text += f"   Snippet: {result.get('body', 'No description')[:200]}...\n\n"
+    
+    prompt = f"""<|begin_of_text|><|start_header_id|>system<|end_header_id|>
+You are a search result evaluator. Given a user's question and search results, identify which results are most relevant and helpful for answering the question. 
+
+Return only the numbers of the most relevant results (1-5 results maximum), separated by commas. Consider:
+- Direct relevance to the question
+- Credibility of the source
+- Recency of information
+- Comprehensiveness of content
+
+Example response: 1, 3, 7
+
+<|eot_id|><|start_header_id|>user<|end_header_id|>
+Question: {user_query}
+
+Search Results:
+{results_text}
+
+<|eot_id|><|start_header_id|>assistant<|end_header_id|>"""
+
+    try:
+        response = search_pipeline(prompt, max_new_tokens=100, temperature=0.1)
+        generated_text = response[0]['generated_text']
+        
+        # Extract assistant's response
+        assistant_response = generated_text.split('<|start_header_id|>assistant<|end_header_id|>')[-1]
+        thinking, cleaned_response = extract_thinking_and_response(assistant_response)
+        
+        # Extract numbers
+        numbers = re.findall(r'\d+', cleaned_response)
+        selected_indices = [int(n) - 1 for n in numbers if int(n) <= len(search_results)]
+        
+        return [search_results[i] for i in selected_indices if 0 <= i < len(search_results)][:5], thinking
+    except Exception as e:
+        print(f"Error filtering results: {e}")
+        return search_results[:5], ""  # Fallback to first 5 results
+
+@spaces.GPU
+def generate_final_answer(user_query: str, selected_results: List[Dict]) -> Tuple[str, str]:
+    """Generate final answer based on selected search results"""
+    if not selected_results:
+        return "I couldn't find relevant information to answer your question. Please try rephrasing your query.", ""
+    
+    # Prepare context from selected results
+    context = ""
+    for i, result in enumerate(selected_results):
+        context += f"Source {i+1}: {result.get('title', 'Unknown')}\n"
+        context += f"Content: {result.get('body', 'No content available')}\n"
+        context += f"URL: {result.get('href', 'No URL')}\n\n"
+    
+    prompt = f"""<|begin_of_text|><|start_header_id|>system<|end_header_id|>
+You are a helpful research assistant. Based on the provided search results, give a comprehensive answer to the user's question. 
+
+Guidelines:
+- Synthesize information from multiple sources
+- Be accurate and factual
+- Cite sources when possible
+- If information is conflicting, mention it
+- Keep the answer well-structured and easy to read
+- Include relevant URLs for further reading
+
+<|eot_id|><|start_header_id|>user<|end_header_id|>
+Question: {user_query}
+
+Search Results:
+{context}
+
+Please provide a comprehensive answer based on these sources.
+
+<|eot_id|><|start_header_id|>assistant<|end_header_id|>"""
+
+    try:
+        response = search_pipeline(prompt, max_new_tokens=1024, temperature=0.2)
+        generated_text = response[0]['generated_text']
+        
+        # Extract assistant's response
+        assistant_response = generated_text.split('<|start_header_id|>assistant<|end_header_id|>')[-1]
+        thinking, answer = extract_thinking_and_response(assistant_response)
+        
+        return answer, thinking
+    except Exception as e:
+        print(f"Error generating final answer: {e}")
+        return "I encountered an error while processing the search results. Please try again.", ""
+
+def search_agent_workflow(user_query: str, progress=gr.Progress()) -> Tuple[str, str, str]:
+    """Main workflow that orchestrates the search agent"""
+    if not user_query.strip():
+        return "Please enter a search query.", "", ""
+    
+    progress(0.1, desc="Initializing...")
+    all_thinking = []
+    
+    # Step 1: Generate search queries
+    progress(0.2, desc="Generating search queries...")
+    queries, thinking1 = generate_search_queries(user_query)
+    if thinking1:
+        all_thinking.append(f"**Query Generation:**\n{thinking1}")
+    queries_text = "Generated queries:\n" + "\n".join(f"• {q}" for q in queries)
+    
+    # Step 2: Search the web
+    progress(0.4, desc="Searching the web...")
+    search_results = search_web(queries)
+    
+    if not search_results:
+        return "No search results found. Please try a different query.", queries_text, "\n\n".join(all_thinking)
+    
+    # Step 3: Filter relevant results
+    progress(0.6, desc="Filtering relevant results...")
+    relevant_results, thinking2 = filter_relevant_results(user_query, search_results)
+    if thinking2:
+        all_thinking.append(f"**Result Filtering:**\n{thinking2}")
+    
+    # Step 4: Generate final answer
+    progress(0.8, desc="Generating comprehensive answer...")
+    final_answer, thinking3 = generate_final_answer(user_query, relevant_results)
+    if thinking3:
+        all_thinking.append(f"**Answer Generation:**\n{thinking3}")
+    
+    progress(1.0, desc="Complete!")
+    
+    # Prepare debug info
+    debug_info = f"{queries_text}\n\nSelected {len(relevant_results)} relevant sources:\n"
+    for i, result in enumerate(relevant_results):
+        debug_info += f"{i+1}. {result.get('title', 'No title')} - {result.get('href', 'No URL')}\n"
+    
+    thinking_display = "\n\n".join(all_thinking) if all_thinking else "No thinking process recorded."
+    
+    return final_answer, debug_info, thinking_display
+
+# Custom CSS for dark blue theme and mobile responsiveness
+custom_css = """
+/* Dark blue theme */
+:root {
+    --primary-bg: #0a1628;
+    --secondary-bg: #1e3a5f;
+    --accent-bg: #2563eb;
+    --text-primary: #f8fafc;
+    --text-secondary: #cbd5e1;
+    --border-color: #334155;
+    --input-bg: #1e293b;
+    --button-bg: #3b82f6;
+    --button-hover: #2563eb;
+}
+
+/* Global styles */
+.gradio-container {
+    background: linear-gradient(135deg, var(--primary-bg) 0%, var(--secondary-bg) 100%) !important;
+    color: var(--text-primary) !important;
+    font-family: 'Inter', 'Segoe UI', system-ui, sans-serif !important;
+}
+
+/* Mobile responsiveness */
+@media (max-width: 768px) {
+    .gradio-container {
+        padding: 10px !important;
+    }
+    
+    .gr-form {
+        gap: 15px !important;
+    }
+    
+    .gr-button {
+        font-size: 16px !important;
+        padding: 12px 20px !important;
+    }
+}
+
+/* Input styling */
+.gr-textbox textarea, .gr-textbox input {
+    background: var(--input-bg) !important;
+    border: 1px solid var(--border-color) !important;
+    color: var(--text-primary) !important;
+    border-radius: 8px !important;
+}
+
+/* Button styling */
+.gr-button {
+    background: linear-gradient(135deg, var(--button-bg) 0%, var(--accent-bg) 100%) !important;
+    color: white !important;
+    border: none !important;
+    border-radius: 8px !important;
+    font-weight: 600 !important;
+    transition: all 0.3s ease !important;
+}
+
+.gr-button:hover {
+    background: linear-gradient(135deg, var(--button-hover) 0%, var(--button-bg) 100%) !important;
+    transform: translateY(-1px) !important;
+    box-shadow: 0 4px 12px rgba(59, 130, 246, 0.3) !important;
+}
+
+/* Output styling */
+.gr-markdown, .gr-textbox {
+    background: var(--input-bg) !important;
+    border: 1px solid var(--border-color) !important;
+    border-radius: 8px !important;
+    color: var(--text-primary) !important;
+}
+
+/* Header styling */
+.gr-markdown h1 {
+    color: var(--accent-bg) !important;
+    text-align: center !important;
+    margin-bottom: 20px !important;
+    font-size: 2.5rem !important;
+    font-weight: 700 !important;
+}
+
+/* Thinking section styling */
+#thinking-output {
+    background: var(--secondary-bg) !important;
+    border: 1px solid var(--border-color) !important;
+    border-radius: 8px !important;
+    padding: 15px !important;
+    font-family: 'Fira Code', 'Monaco', monospace !important;
+    font-size: 0.9rem !important;
+    line-height: 1.4 !important;
+}
+
+/* Loading animation */
+.gr-loading {
+    background: var(--secondary-bg) !important;
+    border-radius: 8px !important;
+}
+
+/* Scrollbar styling */
+::-webkit-scrollbar {
+    width: 8px;
+}
+
+::-webkit-scrollbar-track {
+    background: var(--primary-bg);
+}
+
+::-webkit-scrollbar-thumb {
+    background: var(--accent-bg);
+    border-radius: 4px;
+}
+
+::-webkit-scrollbar-thumb:hover {
+    background: var(--button-hover);
+}
+"""
+
+def create_interface():
+    """Create the Gradio interface"""
+    with gr.Blocks(
+        theme=gr.themes.Base(
+            primary_hue="blue",
+            secondary_hue="slate",
+            neutral_hue="slate",
+            text_size="lg",
+            spacing_size="lg",
+            radius_size="md"
+        ),
+        css=custom_css,
+        title="Just search - AI Search Agent",
+        head="<meta name='viewport' content='width=device-width, initial-scale=1.0'>"
+    ) as interface:
+        
+        gr.Markdown("# 🔍 Just search", elem_id="header")
+        gr.Markdown(
+            "*Part of the Just, AKA Simple series*\n\n"
+            "**Intelligent search agent powered by Menlo/Lucy-128k**\n\n"
+            "Ask any question and get comprehensive answers from the web.",
+            elem_id="description"
+        )
+        
+        with gr.Row():
+            with gr.Column(scale=4):
+                query_input = gr.Textbox(
+                    label="Your Question",
+                    placeholder="Ask me anything... (e.g., 'What are the latest developments in AI?')",
+                    lines=2,
+                    elem_id="query-input"
+                )
+            with gr.Column(scale=1):
+                search_btn = gr.Button(
+                    "🔎 Search",
+                    variant="primary",
+                    size="lg",
+                    elem_id="search-button"
+                )
+        
+        with gr.Row():
+            answer_output = gr.Markdown(
+                label="Answer",
+                elem_id="answer-output",
+                height=400
+            )
+        
+        with gr.Accordion("🤔 AI Thinking Process", open=False):
+            thinking_output = gr.Markdown(
+                label="Model's Chain of Thought",
+                elem_id="thinking-output",
+                height=300
+            )
+        
+        with gr.Accordion("🔧 Debug Info", open=False):
+            debug_output = gr.Textbox(
+                label="Search Process Details",
+                lines=8,
+                elem_id="debug-output"
+            )
+        
+        # Event handlers
+        search_btn.click(
+            fn=search_agent_workflow,
+            inputs=[query_input],
+            outputs=[answer_output, debug_output, thinking_output],
+            show_progress=True
+        )
+        
+        query_input.submit(
+            fn=search_agent_workflow,
+            inputs=[query_input],
+            outputs=[answer_output, debug_output, thinking_output],
+            show_progress=True
+        )
+        
+        # Example queries
+        gr.Examples(
+            examples=[
+                ["What are the latest breakthroughs in quantum computing?"],
+                ["How does climate change affect ocean currents?"],
+                ["What are the best practices for sustainable agriculture?"],
+                ["Explain the recent developments in renewable energy technology"],
+                ["What are the health benefits of the Mediterranean diet?"]
+            ],
+            inputs=query_input,
+            outputs=[answer_output, debug_output, thinking_output],
+            fn=search_agent_workflow,
+            cache_examples=False
+        )
+        
+        gr.Markdown(
+            "---\n**Note:** This search agent generates multiple queries, searches the web, "
+            "filters results for relevance, and provides comprehensive answers. "
+            "Results are sourced from DuckDuckGo search."
+        )
+    
+    return interface
+
+def main():
+    """Main function to initialize and launch the app"""
+    print("🚀 Initializing Just search...")
+    
+    # Initialize the model
+    if not initialize_model():
+        print("❌ Failed to initialize model. Please check your setup.")
+        return
+    
+    print("✅ Model initialized successfully!")
+    print("🌐 Creating interface...")
+    
+    # Create and launch the interface
+    interface = create_interface()
+    
+    print("🎉 Just search is ready!")
+    interface.launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        share=True,
+        show_error=True,
+        debug=True
+    )
+
+if __name__ == "__main__":
+    main()

test_thinking.py

@@ -0,0 +1,51 @@
+#!/usr/bin/env python3
+"""
+Test script to verify thinking extraction works correctly
+"""
+
+import re
+from typing import Tuple
+
+def extract_thinking_and_response(text: str) -> Tuple[str, str]:
+    """Extract thinking process and clean response from AI output"""
+    thinking = ""
+    response = text
+    
+    # Extract thinking content
+    thinking_match = re.search(r'<think>(.*?)</think>', text, re.DOTALL)
+    if thinking_match:
+        thinking = thinking_match.group(1).strip()
+        response = re.sub(r'<think>.*?</think>', '', text, flags=re.DOTALL)
+    
+    # Clean up the response
+    response = re.sub(r'^(Assistant:|AI:|Response:|Answer:)\s*', '', response.strip())
+    response = re.sub(r'\[INST\].*?\[\/INST\]', '', response)
+    response = re.sub(r'<\|.*?\|>', '', response)
+    
+    return thinking.strip(), response.strip()
+
+# Test cases
+test_cases = [
+    # Basic thinking extraction
+    "<think>I need to analyze this query and generate search terms.</think>Here are the search queries:\n1. AI developments 2024\n2. artificial intelligence news",
+    
+    # No thinking
+    "Here are the search queries:\n1. AI developments 2024\n2. artificial intelligence news",
+    
+    # Complex thinking
+    "<think>The user is asking about quantum computing breakthroughs. I should focus on recent developments, key research areas, and practical applications. Let me think about the best search terms...</think>Based on your question, here are optimized search queries for quantum computing breakthroughs.",
+    
+    # Thinking with newlines
+    "<think>\nThis is a complex question about climate change.\nI need to consider multiple aspects:\n1. Ocean currents\n2. Temperature changes\n3. Recent research\n</think>Climate change affects ocean currents in several ways..."
+]
+
+print("Testing thinking extraction...")
+for i, test in enumerate(test_cases, 1):
+    thinking, response = extract_thinking_and_response(test)
+    print(f"\nTest {i}:")
+    print(f"Original: {test[:100]}...")
+    print(f"Thinking: {thinking[:100]}..." if thinking else "Thinking: None")
+    print(f"Response: {response[:100]}...")
+    print("-" * 50)
+
+print("✅ Thinking extraction test completed!")