Update app/main.py

app/main.py

@@ -8,7 +8,7 @@ from starlette.middleware.cors import CORSMiddleware
 import re
 
 # === Setup FastAPI ===
-app = FastAPI(title="Apollo AI Backend", version="1.0.0")
+app = FastAPI(title="Apollo AI Backend", version="2.0.0")
 
 # === CORS ===
 app.add_middleware(
@@ -42,15 +42,48 @@ model.eval()
 
 print("✅ Model ready!")
 
-def extract_clean_answer(full_response: str, formatted_prompt: str, user_message: str) -> str:
+def get_system_prompt(is_force_mode: bool) -> str:
     """
-    FIXED VERSION - Much gentler cleaning that preserves complete responses.
+    Returns mode-specific system prompts for proper AI behavior.
+    """
+    if is_force_mode:
+        return """You are Apollo AI in DIRECT ANSWER mode. Provide:
+- Clear, concise, direct answers
+- Complete working code when requested
+- Brief explanations (2-3 sentences max)
+- Immediate solutions without teaching moments
+- No lengthy tutorials or step-by-step guides
+- Get straight to the point
+
+Example:
+User: "How do I print hello world in Python?"
+You: "Use `print("Hello World")`. This function outputs text to the console."
+"""
+    else:
+        return """You are Apollo AI in MENTOR mode. Your role is to guide learning:
+- Ask leading questions instead of giving direct answers
+- Provide hints and concepts, never complete solutions
+- Encourage thinking: "What do you think would happen if...?"
+- Give partial code with blanks: "Try filling in the _____ part"
+- Guide discovery: "Have you considered looking into...?"
+- Make students work for understanding
+- Never give full working code - always leave something for them to figure out
+
+Example:
+User: "How do I print hello world in Python?"
+You: "Great question! What function do you think might be used to display text on screen? Think about what action you want to perform. Try looking up Python's built-in functions for output."
+"""
+
+def extract_clean_answer(full_response: str, formatted_prompt: str, user_message: str, is_force_mode: bool) -> str:
+    """
+    SINGLE POWERFUL CLEANING FUNCTION - The only place where response cleaning happens.
+    All frontend cleaning is removed, this is the source of truth.
     """
     if not full_response or len(full_response.strip()) < 5:
         return "I apologize, but I couldn't generate a response. Please try again."
     
     print(f"🔍 Raw response length: {len(full_response)}")
-    print(f"🔍 Raw response preview: {full_response[:200]}...")
+    print(f"🔍 Mode: {'FORCE' if is_force_mode else 'MENTOR'}")
     
     # Step 1: Remove the input prompt to get only generated content
     generated_text = full_response
@@ -59,10 +92,10 @@ def extract_clean_answer(full_response: str, formatted_prompt: str, user_message
         if len(parts) > 1:
             generated_text = parts[-1]
     
-    # Step 2: Extract assistant content - SIMPLIFIED approach
+    # Step 2: Extract assistant content using multiple strategies
     assistant_content = generated_text
     
-    # Look for assistant tags and extract content
+    # Strategy A: Look for assistant tags
     if "<|im_start|>assistant" in generated_text:
         assistant_parts = generated_text.split("<|im_start|>assistant")
         if len(assistant_parts) > 1:
@@ -71,53 +104,95 @@ def extract_clean_answer(full_response: str, formatted_prompt: str, user_message
             if "<|im_end|>" in assistant_content:
                 assistant_content = assistant_content.split("<|im_end|>")[0]
     
-    # Step 3: GENTLE cleaning - only remove obvious template artifacts
+    # Strategy B: Look for role-based prefixes
+    elif "assistant:" in generated_text.lower():
+        parts = generated_text.lower().split("assistant:")
+        if len(parts) > 1:
+            # Get the content after the last "assistant:" occurrence
+            assistant_content = generated_text[generated_text.lower().rfind("assistant:") + 10:]
+    
+    # Step 3: POWERFUL CLEANING - Remove all template artifacts
     clean_text = assistant_content.strip()
     
-    # Remove template tokens only
+    # Remove all chat template tokens
     clean_text = re.sub(r'<\|im_start\|>', '', clean_text)
     clean_text = re.sub(r'<\|im_end\|>', '', clean_text)
+    clean_text = re.sub(r'<\|endoftext\|>', '', clean_text)
     
-    # Remove role prefixes only at start of lines
+    # Remove role prefixes from anywhere in text
     clean_text = re.sub(r'^(system|user|assistant):\s*', '', clean_text, flags=re.MULTILINE)
+    clean_text = re.sub(r'\n(system|user|assistant):\s*', '\n', clean_text, flags=re.MULTILINE)
     
-    # REMOVED: Aggressive line-by-line filtering that was truncating responses
+    # Remove common system prompt artifacts
+    clean_text = re.sub(r'You are Apollo AI.*?mode[^\n]*\n?', '', clean_text, flags=re.IGNORECASE)
+    clean_text = re.sub(r'Guidelines?:.*?\n', '', clean_text, flags=re.IGNORECASE)
+    clean_text = re.sub(r'Example:.*?\n', '', clean_text, flags=re.IGNORECASE)
     
-    # Step 4: Final cleanup - preserve content structure
+    # Clean up excessive whitespace but preserve formatting
+    clean_text = re.sub(r'\n{4,}', '\n\n\n', clean_text)
     clean_text = clean_text.strip()
     
-    # Only apply minimal fixes
+    # Step 4: Handle edge cases and fallbacks
     if not clean_text or len(clean_text) < 10:
-        # Fallback for very short responses
-        if user_message and any(math_term in user_message.lower() for math_term in ['2+2', '2 + 2', 'calculate', 'math']):
-            return "4\n\nThe answer is 4."
-        return "I understand your question. Could you please provide more details?"
+        # Special handling for simple math questions
+        if user_message and any(term in user_message.lower() for term in ['2+2', '2 + 2', 'calculate', 'what is']):
+            if '2+2' in user_message.lower() or '2 + 2' in user_message.lower():
+                return "4" if is_force_mode else "What do you think 2 + 2 equals? Try calculating it step by step."
+        
+        # Generic fallback based on mode
+        if is_force_mode:
+            return "I understand your question. Could you please be more specific about what you need?"
+        else:
+            return "That's an interesting question! What approach do you think we should take to solve this? What's your initial thought?"
+    
+    # Step 5: Mode-specific post-processing
+    if is_force_mode:
+        # For force mode, ensure response is concise
+        if len(clean_text) > 800:  # If too long, truncate but keep it coherent
+            sentences = clean_text.split('. ')
+            if len(sentences) > 3:
+                clean_text = '. '.join(sentences[:3]) + '.'
+    else:
+        # For mentor mode, ensure it's not giving away complete solutions
+        # Check if response contains complete code without guidance
+        code_block_pattern = r'```[\w]*\n(.*?)\n```'
+        code_blocks = re.findall(code_block_pattern, clean_text, re.DOTALL)
+        
+        for code in code_blocks:
+            # If code looks complete and there's no guidance, add mentor touch
+            if len(code.strip()) > 50 and 'try' not in clean_text.lower() and '?' not in clean_text:
+                clean_text += "\n\nTry implementing this step by step. What do you think each part does?"
     
     print(f"🧹 Final cleaned answer length: {len(clean_text)}")
-    print(f"🧹 Final answer preview: {clean_text[:150]}...")
+    print(f"🧹 Preview: {clean_text[:150]}...")
+    
     return clean_text
 
-def generate_response(messages: list, max_tokens: int = 400, temperature: float = 0.7) -> str:
+def generate_response(messages: list, is_force_mode: bool = False, max_tokens: int = 400, temperature: float = 0.7) -> str:
     """
-    FIXED: Generate response with higher token limits and better settings.
+    Generate response with mode-specific system prompts and proper settings.
     """
     try:
-        # Create clean conversation
+        # Create clean conversation with mode-specific system prompt
         clean_messages = []
         
-        # Add minimal system message
+        # Add mode-specific system message
+        system_prompt = get_system_prompt(is_force_mode)
         clean_messages.append({
             "role": "system", 
-            "content": "You are Apollo AI, a helpful coding assistant. Provide clear, complete explanations with proper code formatting."
+            "content": system_prompt
         })
         
-        # Add recent conversation context (last 2-3 messages)
+        # Add recent conversation context (last 2-3 messages, but filter appropriately)
         recent_messages = messages[-3:] if len(messages) > 3 else messages
         for msg in recent_messages:
             if msg.get("role") in ["user", "assistant"]:
+                # Skip system messages from frontend to avoid conflicts
+                if msg.get("role") == "system":
+                    continue
                 clean_messages.append(msg)
         
-        print(f"🔍 Processing {len(clean_messages)} messages")
+        print(f"🔍 Processing {len(clean_messages)} messages in {'FORCE' if is_force_mode else 'MENTOR'} mode")
         
         # Build conversation using tokenizer's chat template
         formatted_prompt = tokenizer.apply_chat_template(
@@ -129,22 +204,37 @@ def generate_response(messages: list, max_tokens: int = 400, temperature: float
         # Tokenize with proper length limits
         inputs = tokenizer(formatted_prompt, return_tensors="pt", truncation=True, max_length=1500)
         
-        # FIXED: Generate with much higher token limits
+        # Mode-specific generation settings
+        generation_params = {
+            "input_ids": inputs.input_ids,
+            "attention_mask": inputs.attention_mask,
+            "pad_token_id": tokenizer.eos_token_id,
+            "eos_token_id": tokenizer.eos_token_id,
+            "do_sample": True,
+        }
+        
+        if is_force_mode:
+            # Force mode: Direct, concise answers
+            generation_params.update({
+                "max_new_tokens": min(max_tokens, 300),  # Shorter responses
+                "temperature": 0.3,  # More focused
+                "top_p": 0.8,
+                "repetition_penalty": 1.1,
+                "length_penalty": 0.8,  # Encourage shorter responses
+            })
+        else:
+            # Mentor mode: More thoughtful, questioning responses
+            generation_params.update({
+                "max_new_tokens": min(max_tokens, 500),  # Allow longer explanations
+                "temperature": 0.7,  # More creative for questions
+                "top_p": 0.9,
+                "repetition_penalty": 1.05,
+                "length_penalty": 1.0,  # Neutral length
+            })
+        
+        # Generate response
         with torch.no_grad():
-            outputs = model.generate(
-                inputs.input_ids,
-                attention_mask=inputs.attention_mask,
-                max_new_tokens=min(max_tokens, 500),  # INCREASED from 150 to 500
-                temperature=max(0.3, min(temperature, 0.9)),
-                top_p=0.9,
-                do_sample=True,
-                pad_token_id=tokenizer.eos_token_id,
-                eos_token_id=tokenizer.eos_token_id,
-                repetition_penalty=1.05,  # Reduced to allow natural repetition
-                length_penalty=1.0,  # Neutral length penalty
-                early_stopping=False,  # Don't stop early
-                no_repeat_ngram_size=2,  # Reduced to allow more natural flow
-            )
+            outputs = model.generate(**generation_params)
         
         # Decode the full response
         full_response = tokenizer.decode(outputs[0], skip_special_tokens=False)
@@ -156,28 +246,33 @@ def generate_response(messages: list, max_tokens: int = 400, temperature: float
                 user_message = msg.get("content", "")
                 break
         
-        # Clean and extract the answer with gentle approach
-        clean_answer = extract_clean_answer(full_response, formatted_prompt, user_message)
+        # Clean and extract the answer using our SINGLE POWERFUL cleaning function
+        clean_answer = extract_clean_answer(full_response, formatted_prompt, user_message, is_force_mode)
         
         return clean_answer
         
     except Exception as e:
         print(f"❌ Generation error: {e}")
-        return f"I encountered an error while processing your request. Please try again with a simpler question."
+        mode_text = "direct answer" if is_force_mode else "guided learning approach"
+        return f"I encountered an error while generating a {mode_text}. Please try rephrasing your question."
 
 # === Routes ===
 @app.get("/")
 def root():
     return {
-        "message": "🤖 Apollo AI Backend is running!",
+        "message": "🤖 Apollo AI Backend v2.0 - Mode-Specific AI",
         "model": "Qwen2-0.5B-Instruct with LoRA",
         "status": "ready",
-        "max_tokens": "500 (increased)"
+        "features": ["mentor_mode", "force_mode", "single_powerful_cleaning"],
+        "modes": {
+            "mentor": "Guides learning with questions and hints",
+            "force": "Provides direct answers and solutions"
+        }
     }
 
 @app.get("/health")
 def health():
-    return {"status": "healthy", "model_loaded": True}
+    return {"status": "healthy", "model_loaded": True, "cleaning": "single_backend_only"}
 
 @app.post("/v1/chat/completions")
 async def chat_completions(request: Request):
@@ -200,9 +295,12 @@ async def chat_completions(request: Request):
     try:
         body = await request.json()
         messages = body.get("messages", [])
-        max_tokens = body.get("max_tokens", 400)  # INCREASED default
+        max_tokens = body.get("max_tokens", 400)
         temperature = body.get("temperature", 0.7)
         
+        # NEW: Get mode information from request
+        is_force_mode = body.get("force_mode", False)  # Default to mentor mode
+        
         if not messages or not isinstance(messages, list):
             raise ValueError("Messages field is required and must be a list")
             
@@ -221,11 +319,14 @@ async def chat_completions(request: Request):
             )
 
     try:
-        # Generate response with higher limits
-        print(f"📥 Processing {len(messages)} messages with max_tokens: {max_tokens}")
+        # Generate response with mode-specific behavior
+        print(f"📥 Processing {len(messages)} messages in {'FORCE' if is_force_mode else 'MENTOR'} mode")
+        print(f"📊 Settings: max_tokens={max_tokens}, temperature={temperature}")
+        
         response_content = generate_response(
             messages=messages,
-            max_tokens=min(max_tokens, 600),  # INCREASED cap to 600
+            is_force_mode=is_force_mode,
+            max_tokens=min(max_tokens, 600),
             temperature=max(0.1, min(temperature, 1.0))
         )
         
@@ -234,7 +335,7 @@ async def chat_completions(request: Request):
             "id": f"chatcmpl-apollo-{hash(str(messages)) % 10000}",
             "object": "chat.completion",
             "created": int(torch.tensor(0).item()),
-            "model": "qwen2-0.5b-instruct-lora",
+            "model": f"qwen2-0.5b-{'force' if is_force_mode else 'mentor'}-mode",
             "choices": [
                 {
                     "index": 0,
@@ -249,7 +350,8 @@ async def chat_completions(request: Request):
                 "prompt_tokens": len(str(messages)),
                 "completion_tokens": len(response_content),
                 "total_tokens": len(str(messages)) + len(response_content)
-            }
+            },
+            "apollo_mode": "force" if is_force_mode else "mentor"
         }
         
     except Exception as e:
@@ -259,27 +361,46 @@ async def chat_completions(request: Request):
             content={"error": f"Internal server error: {str(e)}"}
         )
 
-# === Test endpoint for debugging ===
+# === Test endpoint for debugging modes ===
 @app.post("/test")
 async def test_generation(request: Request):
-    """Test endpoint for debugging the model directly"""
+    """Test endpoint for debugging both modes"""
     try:
         body = await request.json()
-        prompt = body.get("prompt", "Hello, how are you?")
+        prompt = body.get("prompt", "How do I print hello world in Python?")
         max_tokens = body.get("max_tokens", 300)
+        test_both_modes = body.get("test_both_modes", True)
+        
+        results = {}
         
-        messages = [
-            {"role": "system", "content": "You are Apollo AI, a helpful assistant."},
+        # Test mentor mode
+        messages_mentor = [
             {"role": "user", "content": prompt}
         ]
+        mentor_response = generate_response(messages_mentor, is_force_mode=False, max_tokens=max_tokens, temperature=0.7)
+        results["mentor_mode"] = {
+            "response": mentor_response,
+            "length": len(mentor_response),
+            "mode": "mentor"
+        }
         
-        response = generate_response(messages, max_tokens=max_tokens, temperature=0.7)
+        if test_both_modes:
+            # Test force mode
+            messages_force = [
+                {"role": "user", "content": prompt}
+            ]
+            force_response = generate_response(messages_force, is_force_mode=True, max_tokens=max_tokens, temperature=0.3)
+            results["force_mode"] = {
+                "response": force_response,
+                "length": len(force_response),
+                "mode": "force"
+            }
         
         return {
             "prompt": prompt,
-            "response": response,
-            "response_length": len(response),
-            "status": "success"
+            "results": results,
+            "status": "success",
+            "cleaning": "single_backend_only"
         }
         
     except Exception as e:
@@ -290,6 +411,7 @@ async def test_generation(request: Request):
 
 if __name__ == "__main__":
     import uvicorn
-    print("🚀 Starting Apollo AI Backend with FIXED response limits...")
-    print("📊 Max tokens increased to 500+ for complete responses")
+    print("🚀 Starting Apollo AI Backend v2.0...")
+    print("📊 Features: Mode-specific prompts, Single powerful cleaning")
+    print("🎯 Modes: Mentor (guides learning) vs Force (direct answers)")
     uvicorn.run(app, host="0.0.0.0", port=7860)