Update app.py

app.py

@@ -1,212 +1,324 @@
-"""Main application - Runs Gemma-2B with tools on CPU"""
-
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
-import warnings
-warnings.filterwarnings("ignore")
+import requests
+import re
+from datetime import datetime
+from typing import List
+import gradio as gr
+
+# ============================================
+# TOOLS
+# ============================================
+
+def web_search(query: str) -> str:
+    """Search the web"""
+    try:
+        url = f"https://api.duckduckgo.com/?q={query}&format=json&no_html=1"
+        r = requests.get(url, timeout=10)
+        data = r.json()
+        results = []
+        for topic in data.get("RelatedTopics", [])[:3]:
+            if isinstance(topic, dict) and "Text" in topic:
+                results.append(topic["Text"][:300])
+        return "\n".join(results) if results else "No results found"
+    except Exception as e:
+        return f"Search failed: {str(e)}"
+
+def wikipedia(topic: str) -> str:
+    """Get Wikipedia summary"""
+    try:
+        url = f"https://en.wikipedia.org/api/rest_v1/page/summary/{topic.replace(' ', '_')}"
+        r = requests.get(url, timeout=10)
+        data = r.json()
+        return data.get("extract", "Not found")[:1000]
+    except Exception as e:
+        return f"Wikipedia lookup failed: {str(e)}"
+
+def get_weather(city: str) -> str:
+    """Get weather for a city"""
+    try:
+        r = requests.get(f"https://wttr.in/{city}?format=%C:+%t", timeout=10)
+        return f"Weather in {city}: {r.text.strip()}"
+    except Exception as e:
+        return f"Weather lookup failed: {str(e)}"
+
+def calculate(expression: str) -> str:
+    """Calculate math expression"""
+    # Clean the expression
+    expression = re.sub(r'[^0-9+\-*/().%]', '', expression)
+    if not expression:
+        return "Please provide a math expression like '2+2' or '15 * 3'"
+    try:
+        result = eval(expression)
+        return f"{expression} = {result}"
+    except Exception as e:
+        return f"Calculation error: {str(e)}"
+
+def get_time() -> str:
+    """Get current time"""
+    return datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+
+def summarize_text(text: str) -> str:
+    """Simple summarization"""
+    if len(text) < 50:
+        return "Text too short to summarize. Please provide at least 50 characters."
+    sentences = re.split(r'[.!?]+', text)
+    key_sentences = [s.strip() for s in sentences if len(s.strip()) > 30][:3]
+    if not key_sentences:
+        return text[:300] + "..."
+    return ". ".join(key_sentences) + "."
 
-from tools import execute_tool, TOOL_REGISTRY, TOOL_DESCRIPTIONS
-from router_client import ToolRouter
+# Tool mapping
+TOOLS = {
+    "web_search": web_search,
+    "wikipedia": wikipedia,
+    "weather": get_weather,
+    "calculate": calculate,
+    "current_time": get_time,
+    "summarize": summarize_text,
+}
 
-class ToolAugmentedAgent:
-    """AI agent with tool-calling capabilities"""
+# ============================================
+# ROUTER
+# ============================================
+
+def route_task(task: str) -> List[str]:
+    """Simple keyword-based router"""
+    task_lower = task.lower()
+    routing = {
+        "web_search": ["search", "google", "find", "look up", "internet", "news", "latest"],
+        "wikipedia": ["wiki", "wikipedia", "who is", "what is", "tell me about"],
+        "weather": ["weather", "temperature", "rain", "sunny", "forecast"],
+        "calculate": ["calculate", "math", "plus", "minus", "times", "divide", "=", "+", "-", "*", "/"],
+        "current_time": ["time", "date", "today", "now"],
+        "summarize": ["summarize", "summary", "tl;dr", "shorten"],
+    }
     
-    def __init__(self):
-        print("=" * 50)
-        print("Loading Tool-Augmented Agent...")
-        print("=" * 50)
-        
-        # Load the model on CPU
-        print("\n📥 Loading Gemma-2B-Uncensored model (this will take a moment)...")
-        self.model_name = "sirev/Gemma-2b-Uncensored-v1"
-        
-        self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
-        self.model = AutoModelForCausalLM.from_pretrained(
-            self.model_name,
-            torch_dtype=torch.float32,  # Use float32 for CPU
-            low_cpu_mem_usage=True
-        )
-        
-        print(f"✅ Model loaded on CPU")
-        print(f"📊 Model parameters: {sum(p.numel() for p in self.model.parameters()) / 1e9:.1f}B")
-        
-        # Initialize the tool router
-        print("🔄 Initializing tool router...")
-        self.router = ToolRouter(use_hosted_api=True)
-        
-        # Available tools
-        self.available_tools = list(TOOL_REGISTRY.keys())
-        print(f"🔧 Available tools: {len(self.available_tools)}")
-        print(f"   {', '.join(self.available_tools[:10])}...")
-        print("=" * 50)
-    
-    def generate_response(self, prompt: str, max_length: int = 512) -> str:
-        """Generate response from the model"""
-        inputs = self.tokenizer(prompt, return_tensors="pt", truncation=True, max_length=1024)
-        
-        with torch.no_grad():
-            outputs = self.model.generate(
-                **inputs,
-                max_new_tokens=max_length,
-                temperature=0.7,
-                do_sample=True,
-                top_p=0.9,
-                pad_token_id=self.tokenizer.eos_token_id
-            )
-        
-        response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
-        # Remove the input prompt from the response
-        response = response[len(prompt):].strip()
-        return response if response else "I couldn't generate a response."
+    selected = []
+    for tool, keywords in routing.items():
+        if any(k in task_lower for k in keywords):
+            selected.append(tool)
     
-    def process_task(self, user_query: str, verbose: bool = True) -> str:
-        """Process a user query by selecting and executing appropriate tools"""
-        
-        print(f"\n📝 User: {user_query}")
-        print("-" * 40)
-        
-        # Step 1: Route to find relevant tools
-        print("🔍 Selecting tools for this task...")
-        selected_tools = self.router.route(user_query, k=3)
-        
-        if not selected_tools:
-            # If no specific tools found, use web search as default
-            selected_tools = ["web_search"]
-        
-        print(f"   Selected tools: {', '.join(selected_tools)}")
-        
-        # Step 2: Execute tools and collect results
-        print("\n🔧 Executing tools...")
-        tool_results = []
-        
-        for tool_name in selected_tools:
-            if tool_name in TOOL_REGISTRY:
-                print(f"   → Running {tool_name}...")
-                
-                # Different tools need different parameters
-                if tool_name == "web_search":
-                    result = execute_tool(tool_name, query=user_query)
-                elif tool_name == "fetch_url":
-                    # Extract URL from query if present
-                    import re
-                    urls = re.findall(r'https?://[^\s]+', user_query)
-                    if urls:
-                        result = execute_tool(tool_name, url=urls[0])
-                    else:
-                        result = {"success": False, "error": "No URL found in query"}
-                elif tool_name == "scholar_search":
-                    result = execute_tool(tool_name, query=user_query)
-                elif tool_name == "get_news":
-                    # Extract topic from query
-                    topic = user_query.replace("news", "").replace("about", "").strip()
-                    result = execute_tool(tool_name, topic=topic or "general")
-                elif tool_name == "wikipedia":
-                    result = execute_tool(tool_name, topic=user_query[:50])
-                elif tool_name == "weather":
-                    # Extract city from query
-                    import re
-                    cities = re.findall(r'in (\w+)', user_query)
-                    city = cities[0] if cities else "London"
-                    result = execute_tool(tool_name, city=city)
-                elif tool_name == "calculate":
-                    result = execute_tool(tool_name, expression=user_query)
-                elif tool_name in ["extract_dates", "extract_emails", "extract_urls", "summarize"]:
-                    # These need text input, use the query as text
-                    result = execute_tool(tool_name, text=user_query)
-                else:
-                    result = execute_tool(tool_name)
-                
-                tool_results.append({
-                    "tool": tool_name,
-                    "result": result
-                })
-                
-                # Show result preview
-                if result.get("success"):
-                    preview = str(result)[:100] + "..."
-                    print(f"      ✓ {preview}")
-                else:
-                    print(f"      ✗ Error: {result.get('error', 'Unknown error')}")
+    # Always include web_search if nothing else matches
+    return selected[:3] if selected else ["web_search"]
+
+# ============================================
+# LOAD MODEL (with better error handling)
+# ============================================
+
+print("📥 Loading Gemma-2B model... (this takes 2-3 minutes on first run)")
+
+try:
+    tokenizer = AutoTokenizer.from_pretrained("sirev/Gemma-2b-Uncensored-v1")
+    model = AutoModelForCausalLM.from_pretrained(
+        "sirev/Gemma-2b-Uncensored-v1",
+        torch_dtype=torch.float32,
+        low_cpu_mem_usage=True
+    )
+    print("✅ Model loaded successfully!")
+except Exception as e:
+    print(f"❌ Failed to load model: {e}")
+    # Fallback to a smaller model if Gemma fails
+    print("🔄 Trying fallback model...")
+    tokenizer = AutoTokenizer.from_pretrained("microsoft/phi-2")
+    model = AutoModelForCausalLM.from_pretrained(
+        "microsoft/phi-2",
+        torch_dtype=torch.float32,
+        low_cpu_mem_usage=True
+    )
+    print("✅ Fallback model loaded!")
+
+def generate_response(prompt: str) -> str:
+    """Generate response from the model"""
+    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=1024)
+    outputs = model.generate(
+        **inputs,
+        max_new_tokens=300,
+        temperature=0.7,
+        do_sample=True,
+        pad_token_id=tokenizer.eos_token_id
+    )
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    # Remove the input prompt from the response
+    response = response[len(prompt):].strip()
+    return response if response else "I couldn't generate a response."
+
+# ============================================
+# AGENT FUNCTION
+# ============================================
+
+def run_agent(user_query: str, history: list = None) -> str:
+    """Main agent function - called by Gradio"""
+    
+    if not user_query or user_query.strip() == "":
+        return "Please enter a question."
+    
+    print(f"📝 Processing: {user_query}")
+    
+    # Step 1: Route to find relevant tools
+    tools_to_use = route_task(user_query)
+    print(f"🔧 Selected tools: {tools_to_use}")
+    
+    # Step 2: Execute tools
+    tool_results = []
+    for tool in tools_to_use:
+        try:
+            if tool == "web_search":
+                result = web_search(user_query)
+            elif tool == "wikipedia":
+                # Extract topic (simple: take first few words, excluding question words)
+                words = user_query.lower().split()
+                stop_words = {"what", "who", "where", "when", "why", "how", "is", "are", "tell", "me", "about"}
+                topic_words = [w for w in user_query.split() if w.lower() not in stop_words][:4]
+                topic = ' '.join(topic_words) if topic_words else user_query[:50]
+                result = wikipedia(topic)
+            elif tool == "weather":
+                # Extract city
+                words = user_query.split()
+                city = "London"
+                if "in" in words:
+                    city = words[words.index("in") + 1]
+                elif "for" in words:
+                    city = words[words.index("for") + 1]
+                result = get_weather(city)
+            elif tool == "calculate":
+                # Extract math expression
+                math_match = re.search(r'[\d\s\+\-\*/\(\)\.\%]+', user_query)
+                expr = math_match.group(0) if math_match else user_query
+                result = calculate(expr)
+            elif tool == "current_time":
+                result = get_time()
+            elif tool == "summarize":
+                result = summarize_text(user_query)
             else:
-                print(f"   → Tool {tool_name} not found in registry")
-        
-        # Step 3: Prepare context for the LLM
-        print("\n🤖 Generating final response...")
-        
-        context = f"""You are a helpful AI assistant with access to various tools. Based on the tool results below, answer the user's question naturally.
+                result = f"Tool '{tool}' not implemented"
+            
+            tool_results.append(f"🔹 {tool}: {result[:400]}")
+            print(f"   ✓ {tool} executed")
+        except Exception as e:
+            tool_results.append(f"🔸 {tool}: Error - {str(e)[:100]}")
+            print(f"   ✗ {tool} failed: {e}")
+    
+    # Step 3: Generate response with LLM
+    tool_context = "\n".join(tool_results) if tool_results else "No tools were executed."
+    
+    context = f"""You are a helpful AI assistant. Answer the user's question based on the tool results below.
 
-User question: {user_query}
+User: {user_query}
 
 Tool results:
-"""
-        for tr in tool_results:
-            if tr["result"].get("success"):
-                context += f"\n[{tr['tool']}]: {tr['result']}\n"
-            else:
-                context += f"\n[{tr['tool']}] Error: {tr['result'].get('error', 'Failed')}\n"
+{tool_context}
 
-        context += """
+Please provide a clear, helpful answer based on the information above. If the tools didn't provide useful information, say so and offer to help differently."""
 
-Please provide a clear, helpful answer based on the information above. If the tools didn't provide useful results, say so and offer alternatives."""
+    print("🤖 Generating response...")
+    response = generate_response(context)
+    
+    # Add tool info prefix for transparency
+    if len(tools_to_use) <= 2:
+        final_response = f"🔧 {', '.join(tools_to_use)}\n\n{response}"
+    else:
+        final_response = f"🔧 {len(tools_to_use)} tools used\n\n{response}"
+    
+    print("✅ Done")
+    return final_response
 
-        # Step 4: Generate final response
-        response = self.generate_response(context, max_length=400)
-        
-        print("\n" + "=" * 40)
-        print(f"🤖 Assistant: {response}")
-        print("=" * 40)
-        
-        return response
-    
-    def interactive_mode(self):
-        """Run in interactive mode"""
-        print("\n" + "=" * 50)
-        print("🟢 Tool-Augmented Agent Ready!")
-        print("=" * 50)
-        print("Available tools:")
-        for i, tool in enumerate(self.available_tools[:15], 1):
-            print(f"  {i}. {tool}")
-        print("  ...")
-        print("\nType 'exit' to quit, 'help' for tool descriptions\n")
+# ============================================
+# GRADIO INTERFACE
+# ============================================
+
+# Example queries for users to try
+examples = [
+    "What is artificial intelligence?",
+    "Search for latest news about space exploration",
+    "What's the weather in Tokyo?",
+    "Calculate 25 * 4 + 10",
+    "What time is it?",
+    "Summarize: The quick brown fox jumps over the lazy dog. This is a test sentence. The end.",
+    "Tell me about quantum computing",
+]
+
+# CSS for better appearance
+css = """
+.gradio-container {
+    max-width: 900px !important;
+    margin: auto !important;
+}
+h1 {
+    text-align: center;
+    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
+    -webkit-background-clip: text;
+    -webkit-text-fill-color: transparent;
+}
+"""
+
+with gr.Blocks(title="Tool-Augmented AI Assistant", css=css, theme=gr.themes.Soft()) as demo:
+    gr.Markdown("""
+    # 🛠️ Tool-Augmented AI Assistant
+    
+    ### Powered by Gemma-2B-Uncensored + 6 Tools
+    """)
+    
+    with gr.Row():
+        with gr.Column(scale=4):
+            chatbot = gr.Chatbot(height=500, label="Conversation")
+        with gr.Column(scale=1):
+            gr.Markdown("""
+            ### 🚀 Available Tools
+            - 🔍 **Web Search** - Search the internet
+            - 📚 **Wikipedia** - Encyclopedia lookup  
+            - ⛅ **Weather** - Current conditions
+            - 🧮 **Calculator** - Math operations
+            - ⏰ **Current Time** - Date and time
+            - 📝 **Summarizer** - Shorten long text
+            """)
+    
+    with gr.Row():
+        msg = gr.Textbox(
+            label="Your Question",
+            placeholder="Type anything... (e.g., 'Search for AI news' or 'Weather in Paris')",
+            scale=4
+        )
+        submit_btn = gr.Button("Send", variant="primary", scale=1)
+    
+    clear_btn = gr.Button("Clear Conversation", variant="secondary")
+    
+    gr.Markdown("""
+    ### 💡 Example Questions
+    """)
+    
+    gr.Examples(
+        examples=examples,
+        inputs=msg,
+        label="Click any example to try"
+    )
+    
+    gr.Markdown("""
+    ---
+    *Note: Runs on CPU, so responses take 5-15 seconds. The model is uncensored - use responsibly.*
+    """)
+    
+    # State for conversation history
+    state = gr.State([])
+    
+    def respond(message, chat_history):
+        if not message or not message.strip():
+            return "", chat_history
         
-        while True:
-            try:
-                query = input("\n💬 You: ").strip()
-                
-                if query.lower() == 'exit':
-                    print("Goodbye! 👋")
-                    break
-                elif query.lower() == 'help':
-                    print("\n📚 Tool Descriptions:")
-                    for tool, desc in list(TOOL_DESCRIPTIONS.items())[:20]:
-                        print(f"  • {tool}: {desc}")
-                    continue
-                
-                if query:
-                    self.process_task(query)
-                    
-            except KeyboardInterrupt:
-                print("\n\nInterrupted. Goodbye!")
-                break
-            except Exception as e:
-                print(f"\n❌ Error: {e}")
-                print("Please try again.")
-
-def main():
-    """Main entry point"""
-    agent = ToolAugmentedAgent()
-    
-    # Example queries to demonstrate functionality
-    print("\n📋 Try these example queries:")
-    print("   • 'Search for recent developments in AI'")
-    print("   • 'What's the weather in Tokyo?'")
-    print("   • 'Calculate 15 * 87 + 23'")
-    print("   • 'Get today's news about space exploration'")
-    print("   • 'Find academic papers about quantum computing'")
-    print("   • 'Summarize: [paste long text]'")
-    
-    # Start interactive mode
-    agent.interactive_mode()
+        response = run_agent(message)
+        chat_history.append((message, response))
+        return "", chat_history
+    
+    def clear_chat():
+        return [], []
+    
+    msg.submit(respond, [msg, state], [msg, state])
+    submit_btn.click(respond, [msg, state], [msg, state])
+    clear_btn.click(clear_chat, None, [chatbot, state])
+
+# ============================================
+# LAUNCH
+# ============================================
 
 if __name__ == "__main__":
-    main()
+    demo.launch()