Update app.py

app.py

@@ -1,102 +1,82 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
-import torch
+from huggingface_hub import InferenceClient
 from concurrent.futures import ThreadPoolExecutor, as_completed
-import json
 import time
 from datetime import datetime
+import os
+
+# Use Hugging Face Inference API (no model loading needed!)
+# This is FREE and much faster!
 
-# Configuration for different agent types
 AGENT_CONFIGS = {
     "researcher": {
-        "model": "HuggingFaceH4/zephyr-7b-beta",
+        "model": "mistralai/Mistral-7B-Instruct-v0.2",
         "role": "Research and gather information",
-        "prompt_template": "You are a research agent. Analyze and provide detailed information about: {task}"
+        "system_prompt": "You are a research agent specialized in gathering and analyzing information. Provide detailed, well-researched responses."
     },
     "coder": {
-        "model": "Salesforce/codegen-350M-mono",
+        "model": "bigcode/starcoder2-15b",
         "role": "Generate and explain code",
-        "prompt_template": "Generate Python code for: {task}"
+        "system_prompt": "You are an expert programmer. Generate clean, efficient, well-commented code."
     },
     "analyzer": {
-        "model": "HuggingFaceH4/zephyr-7b-beta",
+        "model": "mistralai/Mistral-7B-Instruct-v0.2",
         "role": "Analyze data and provide insights",
-        "prompt_template": "Analyze the following and provide insights: {task}"
+        "system_prompt": "You are a data analyst. Provide clear insights and actionable recommendations."
     },
     "writer": {
-        "model": "HuggingFaceH4/zephyr-7b-beta",
+        "model": "mistralai/Mistral-7B-Instruct-v0.2",
         "role": "Create content and documentation",
-        "prompt_template": "Write professional content about: {task}"
+        "system_prompt": "You are a technical writer. Create clear, professional documentation and content."
     }
 }
 
 class AgentSystem:
     def __init__(self):
-        self.models = {}
-        self.tokenizers = {}
+        # No model loading! Using HF Inference API
+        self.clients = {}
         self.executor = ThreadPoolExecutor(max_workers=4)
-        self.device = "cuda" if torch.cuda.is_available() else "cpu"
-        print(f"Using device: {self.device}")
         
-    def load_model(self, agent_name, model_name):
-        """Load model for specific agent"""
-        if agent_name not in self.models:
-            print(f"Loading {agent_name} model: {model_name}")
-            try:
-                self.tokenizers[agent_name] = AutoTokenizer.from_pretrained(model_name)
-                self.models[agent_name] = AutoModelForCausalLM.from_pretrained(
-                    model_name,
-                    torch_dtype=torch.float16 if self.device == "cuda" else torch.float32,
-                    low_cpu_mem_usage=True,
-                    device_map="auto" if self.device == "cuda" else None
-                )
-                print(f"{agent_name} model loaded successfully!")
-            except Exception as e:
-                print(f"Error loading {agent_name} model: {e}")
-                # Fallback to smaller model
-                print(f"Falling back to distilgpt2 for {agent_name}")
-                self.tokenizers[agent_name] = AutoTokenizer.from_pretrained("distilgpt2")
-                self.models[agent_name] = AutoModelForCausalLM.from_pretrained("distilgpt2")
+        # Initialize inference clients for each agent
+        for agent_name in AGENT_CONFIGS.keys():
+            model = AGENT_CONFIGS[agent_name]["model"]
+            self.clients[agent_name] = InferenceClient(model=model)
+        
+        print("✅ Agent system initialized with Inference API!")
     
-    def generate_response(self, agent_name, prompt, max_length=200):
-        """Generate response for a specific agent"""
+    def generate_response(self, agent_name, task, max_tokens=300):
+        """Generate response using HF Inference API"""
         try:
             config = AGENT_CONFIGS[agent_name]
-            model_name = config["model"]
+            client = self.clients[agent_name]
             
-            # Load model if not already loaded
-            if agent_name not in self.models:
-                self.load_model(agent_name, model_name)
+            # Create prompt
+            messages = [
+                {
+                    "role": "system",
+                    "content": config["system_prompt"]
+                },
+                {
+                    "role": "user",
+                    "content": f"Task: {task}"
+                }
+            ]
             
-            tokenizer = self.tokenizers[agent_name]
-            model = self.models[agent_name]
-            
-            # Format prompt
-            formatted_prompt = config["prompt_template"].format(task=prompt)
-            
-            # Tokenize
-            inputs = tokenizer(formatted_prompt, return_tensors="pt", truncation=True, max_length=512)
-            
-            if self.device == "cuda":
-                inputs = inputs.to("cuda")
-            
-            # Generate
-            with torch.no_grad():
-                outputs = model.generate(
-                    inputs.input_ids,
-                    max_length=max_length,
-                    temperature=0.7,
-                    top_p=0.9,
-                    do_sample=True,
-                    pad_token_id=tokenizer.eos_token_id
-                )
-            
-            response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+            # Generate response
+            response_text = ""
+            for message in client.chat_completion(
+                messages=messages,
+                max_tokens=max_tokens,
+                temperature=0.7,
+                stream=True
+            ):
+                if hasattr(message.choices[0].delta, 'content'):
+                    response_text += message.choices[0].delta.content
             
             return {
                 "agent": agent_name,
                 "role": config["role"],
-                "response": response,
+                "response": response_text.strip(),
                 "status": "success"
             }
             
@@ -108,7 +88,7 @@ class AgentSystem:
                 "status": "error"
             }
     
-    def run_agents_parallel(self, task, selected_agents, max_length=200):
+    def run_agents_parallel(self, task, selected_agents, max_tokens=300):
         """Run multiple agents in parallel"""
         start_time = time.time()
         futures = {}
@@ -120,7 +100,7 @@ class AgentSystem:
                 self.generate_response,
                 agent_name,
                 task,
-                max_length
+                max_tokens
             )
             futures[future] = agent_name
         
@@ -128,14 +108,14 @@ class AgentSystem:
         for future in as_completed(futures):
             agent_name = futures[future]
             try:
-                result = future.result()
+                result = future.result(timeout=30)  # 30 second timeout per agent
                 result["time_taken"] = round(time.time() - start_time, 2)
                 results.append(result)
             except Exception as e:
                 results.append({
                     "agent": agent_name,
                     "role": AGENT_CONFIGS[agent_name]["role"],
-                    "response": f"Failed: {str(e)}",
+                    "response": f"Timeout or error: {str(e)}",
                     "status": "error",
                     "time_taken": round(time.time() - start_time, 2)
                 })
@@ -144,13 +124,14 @@ class AgentSystem:
         return results, total_time
 
 # Initialize the agent system
-print("Initializing AI Agent System...")
+print("🚀 Initializing AI Agent System...")
 agent_system = AgentSystem()
+print("✅ System ready!")
 
-def process_task(task, researcher, coder, analyzer, writer, max_length):
+def process_task(task, researcher, coder, analyzer, writer, max_tokens, progress=gr.Progress()):
     """Process task with selected agents"""
     if not task.strip():
-        return "Please enter a task!", ""
+        return "⚠️ Please enter a task!", "", ""
     
     # Determine which agents to use
     selected_agents = []
@@ -164,10 +145,14 @@ def process_task(task, researcher, coder, analyzer, writer, max_length):
         selected_agents.append("writer")
     
     if not selected_agents:
-        return "Please select at least one agent!", ""
+        return "⚠️ Please select at least one agent!", "", ""
+    
+    progress(0, desc="Starting agents...")
     
     # Run agents in parallel
-    results, total_time = agent_system.run_agents_parallel(task, selected_agents, max_length)
+    results, total_time = agent_system.run_agents_parallel(task, selected_agents, max_tokens)
+    
+    progress(1, desc="Complete!")
     
     # Format output
     output = f"# 🤖 AI Agent System Results\n\n"
@@ -176,119 +161,192 @@ def process_task(task, researcher, coder, analyzer, writer, max_length):
     output += f"**Total Time:** {total_time}s\n\n"
     output += "---\n\n"
     
-    for result in results:
+    for idx, result in enumerate(results, 1):
         status_emoji = "✅" if result["status"] == "success" else "❌"
-        output += f"## {status_emoji} {result['agent'].upper()} Agent\n"
+        output += f"## {status_emoji} Agent {idx}: {result['agent'].upper()}\n\n"
         output += f"**Role:** {result['role']}\n\n"
-        output += f"**Response:**\n```\n{result['response']}\n```\n\n"
-        output += f"*Completed in {result['time_taken']}s*\n\n"
+        output += f"**Response:**\n\n{result['response']}\n\n"
+        output += f"*⏱️ Completed in {result['time_taken']}s*\n\n"
         output += "---\n\n"
     
-    # Create summary
-    summary = {
+    # Create summary stats
+    success_count = sum(1 for r in results if r["status"] == "success")
+    stats = f"""📊 **Execution Stats**
+- Total Agents: {len(selected_agents)}
+- Successful: {success_count}
+- Failed: {len(selected_agents) - success_count}
+- Total Time: {total_time}s
+- Average per Agent: {round(total_time / len(selected_agents), 2)}s
+"""
+    
+    # Detailed JSON for download
+    details = {
         "task": task,
-        "agents_used": len(selected_agents),
+        "agents_used": selected_agents,
         "total_time": total_time,
-        "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+        "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
+        "results": results
     }
     
-    return output, json.dumps(summary, indent=2)
+    import json
+    json_output = json.dumps(details, indent=2)
+    
+    return output, stats, json_output
 
 # Create Gradio Interface
-with gr.Blocks(theme=gr.themes.Soft(), title="AI Agent System") as demo:
-    gr.Markdown(
-        """
-        # 🤖 Full-Stack AI Agent System
-        
-        **Parallel AI Processing with Multiple Specialized Agents**
-        
-        This system runs multiple AI agents simultaneously to process your tasks faster!
-        Each agent specializes in different areas and works in parallel.
-        """
-    )
+custom_css = """
+.gradio-container {
+    font-family: 'Inter', sans-serif;
+}
+.main-header {
+    text-align: center;
+    padding: 20px;
+    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
+    color: white;
+    border-radius: 10px;
+    margin-bottom: 20px;
+}
+"""
+
+with gr.Blocks(theme=gr.themes.Soft(), css=custom_css, title="AI Agent System") as demo:
+    gr.HTML("""
+        <div class="main-header">
+            <h1>🤖 Multi-Agent AI System</h1>
+            <p>Parallel AI Processing with Specialized Agents | Powered by Hugging Face Inference API</p>
+        </div>
+    """)
     
     with gr.Row():
         with gr.Column(scale=1):
-            gr.Markdown("### 📋 Task Input")
+            gr.Markdown("### 📝 Task Configuration")
+            
             task_input = gr.Textbox(
-                label="Enter Your Task",
-                placeholder="Example: Create a Python web scraper for news articles",
-                lines=4
+                label="What do you want the agents to work on?",
+                placeholder="Example: Build a user authentication system with JWT tokens",
+                lines=5
             )
             
-            gr.Markdown("### 🎯 Select Agents")
-            researcher_check = gr.Checkbox(label="🔍 Researcher Agent", value=True, info="Research and gather information")
-            coder_check = gr.Checkbox(label="💻 Coder Agent", value=True, info="Generate and explain code")
-            analyzer_check = gr.Checkbox(label="📊 Analyzer Agent", value=True, info="Analyze and provide insights")
-            writer_check = gr.Checkbox(label="✍️ Writer Agent", value=True, info="Create documentation")
+            gr.Markdown("### 🎯 Select Your Team")
+            
+            with gr.Group():
+                researcher_check = gr.Checkbox(
+                    label="🔍 Researcher Agent", 
+                    value=True,
+                    info="Gathers information and best practices"
+                )
+                coder_check = gr.Checkbox(
+                    label="💻 Coder Agent", 
+                    value=True,
+                    info="Writes production-ready code"
+                )
+                analyzer_check = gr.Checkbox(
+                    label="📊 Analyzer Agent", 
+                    value=True,
+                    info="Provides insights and recommendations"
+                )
+                writer_check = gr.Checkbox(
+                    label="✍️ Writer Agent", 
+                    value=True,
+                    info="Creates documentation"
+                )
+            
+            gr.Markdown("### ⚙️ Settings")
             
-            max_length = gr.Slider(
+            max_tokens = gr.Slider(
                 minimum=100,
                 maximum=500,
-                value=200,
+                value=300,
                 step=50,
-                label="Max Response Length",
-                info="Tokens per agent"
+                label="Response Length",
+                info="Tokens per agent response"
             )
             
-            process_btn = gr.Button("🚀 Run Agents in Parallel", variant="primary", size="lg")
+            process_btn = gr.Button(
+                "🚀 Deploy Agents", 
+                variant="primary", 
+                size="lg"
+            )
             
-            gr.Markdown(
-                """
-                ### 💡 Tips
-                - Select multiple agents for comprehensive results
-                - Agents run simultaneously for faster processing
+            gr.Markdown("""
+                ### 💡 Pro Tips
+                - Use all 4 agents for comprehensive results
+                - Agents run simultaneously = 3-4x faster!
                 - Each agent brings unique expertise
-                """
-            )
+                - No model downloads = instant startup
+            """)
         
         with gr.Column(scale=2):
-            gr.Markdown("### 📤 Agent Outputs")
-            output_display = gr.Markdown(label="Results")
+            gr.Markdown("### 📊 Results Dashboard")
+            
+            output_display = gr.Markdown(
+                value="*Results will appear here after running agents...*",
+                label="Agent Outputs"
+            )
             
-            with gr.Accordion("📊 Execution Summary", open=False):
-                summary_json = gr.Code(label="JSON Summary", language="json")
+            with gr.Accordion("📈 Execution Statistics", open=True):
+                stats_display = gr.Markdown(value="*No data yet*")
+            
+            with gr.Accordion("💾 Download Results (JSON)", open=False):
+                json_output = gr.Code(
+                    label="Complete Results", 
+                    language="json",
+                    lines=10
+                )
+    
+    gr.Markdown("### 📚 Quick Start Examples")
     
-    gr.Markdown("### 📚 Example Tasks")
     gr.Examples(
         examples=[
-            ["Create a REST API for user authentication"],
-            ["Build a machine learning model for sentiment analysis"],
-            ["Design a database schema for an e-commerce platform"],
-            ["Write a technical blog post about microservices"],
-            ["Develop a real-time chat application"]
+            ["Create a REST API for a todo list application with authentication"],
+            ["Build a machine learning pipeline for image classification"],
+            ["Design a microservices architecture for an e-commerce platform"],
+            ["Develop a real-time chat application using WebSockets"],
+            ["Create a data visualization dashboard for sales analytics"],
         ],
         inputs=task_input
     )
     
-    gr.Markdown(
-        """
+    gr.Markdown("""
         ---
         
         ## 🏗️ System Architecture
         
-        - **Parallel Processing**: All agents run simultaneously using ThreadPoolExecutor
-        - **Free Models**: Using Hugging Face hosted models (Zephyr-7B, CodeGen)
-        - **Specialized Agents**: Each agent has a specific role and expertise
-        - **Fault Tolerant**: Continues even if one agent fails
+        **How It Works:**
         
-        ## 🔧 Technology Stack
+        1. **Task Distribution** → Your task is sent to selected agents
+        2. **Parallel Processing** → All agents work simultaneously (not sequential!)
+        3. **Smart Aggregation** → Results are collected as they complete
+        4. **Instant Results** → See output from each agent in real-time
         
-        - **Frontend**: Gradio
-        - **Backend**: Python + Transformers
-        - **Models**: Hugging Face free models
-        - **Concurrency**: ThreadPoolExecutor for parallel processing
-        """
-    )
+        **Technology:**
+        - ⚡ Hugging Face Inference API (serverless, no model loading)
+        - 🔄 ThreadPoolExecutor for true parallelism
+        - 🚀 Free tier compatible
+        - 📊 Real-time progress tracking
+        
+        **Models Used:**
+        - Mistral-7B-Instruct (Researcher, Analyzer, Writer)
+        - StarCoder2-15B (Coder)
+    """)
     
-    # Connect button to processing function
+    # Connect button
     process_btn.click(
         fn=process_task,
-        inputs=[task_input, researcher_check, coder_check, analyzer_check, writer_check, max_length],
-        outputs=[output_display, summary_json]
+        inputs=[
+            task_input, 
+            researcher_check, 
+            coder_check, 
+            analyzer_check, 
+            writer_check, 
+            max_tokens
+        ],
+        outputs=[output_display, stats_display, json_output]
     )
 
-# Launch
+# Launch with optimized settings
 if __name__ == "__main__":
-    demo.queue()  # Enable queuing for better performance
-    demo.launch()
+    demo.queue(max_size=20)  # Handle multiple users
+    demo.launch(
+        show_error=True,
+        share=False
+    )