Update app/main.py

app/main.py

@@ -44,207 +44,126 @@ print("✅ Model ready!")
 
 def extract_clean_answer(full_response: str, formatted_prompt: str, user_message: str) -> str:
     """
-    COMPLETELY REWRITTEN - Extract only the AI's actual response, removing ALL artifacts.
+    FIXED VERSION - Much gentler cleaning that preserves complete responses.
     """
     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[:400]}...")
+    print(f"🔍 Raw response preview: {full_response[:200]}...")
     
-    # Step 1: Remove the entire input prompt to get only generated content
+    # Step 1: Remove the input prompt to get only generated content
     generated_text = full_response
     if formatted_prompt in full_response:
         parts = full_response.split(formatted_prompt)
         if len(parts) > 1:
-            generated_text = parts[-1]  # Take everything after the prompt
+            generated_text = parts[-1]
     
-    print(f"🔍 After prompt removal: {generated_text[:200]}...")
+    # Step 2: Extract assistant content - SIMPLIFIED approach
+    assistant_content = generated_text
     
-    # Step 2: Extract content between assistant tags (most reliable method)
-    assistant_content = ""
-    
-    # Look for the last assistant response
+    # Look for assistant tags and extract content
     if "<|im_start|>assistant" in generated_text:
-        # Split by assistant markers and take the last one
         assistant_parts = generated_text.split("<|im_start|>assistant")
         if len(assistant_parts) > 1:
             assistant_content = assistant_parts[-1]
-            
-            # Remove everything after <|im_end|> if it exists
+            # Remove end marker if present
             if "<|im_end|>" in assistant_content:
                 assistant_content = assistant_content.split("<|im_end|>")[0]
-    else:
-        # Fallback: use the generated text as-is
-        assistant_content = generated_text
     
-    print(f"🔍 After assistant extraction: {assistant_content[:200]}...")
+    # Step 3: GENTLE cleaning - only remove obvious template artifacts
+    clean_text = assistant_content.strip()
     
-    # Step 3: Remove ALL template artifacts aggressively
-    clean_text = assistant_content
+    # Remove template tokens only
+    clean_text = re.sub(r'<\|im_start\|>', '', clean_text)
+    clean_text = re.sub(r'<\|im_end\|>', '', clean_text)
     
-    # Remove any remaining template tokens
-    template_artifacts = [
-        r'<\|im_start\|>.*?<\|im_end\|>',
-        r'<\|im_start\|>.*',
-        r'<\|im_end\|>.*',
-        r'^(system|user|assistant):\s*',
-        r'\n(system|user|assistant):\s*',
-        r'^\s*(system|user|assistant)\s*\n',
-    ]
+    # Remove role prefixes only at start of lines
+    clean_text = re.sub(r'^(system|user|assistant):\s*', '', clean_text, flags=re.MULTILINE)
     
-    for pattern in template_artifacts:
-        clean_text = re.sub(pattern, '', clean_text, flags=re.MULTILINE | re.IGNORECASE)
+    # REMOVED: Aggressive line-by-line filtering that was truncating responses
     
-    # Step 4: NUCLEAR OPTION - Remove all system prompt leaks line by line
-    lines = clean_text.split('\n')
-    final_lines = []
+    # Step 4: Final cleanup - preserve content structure
+    clean_text = clean_text.strip()
     
-    # System prompt indicators to completely remove
-    system_indicators = [
-        'you are apollo ai',
-        'you are a helpful',
-        'guidelines:',
-        'response format:',
-        '- provide clear',
-        '- use markdown',
-        '- always include',
-        '- be encouraging',
-        '- if asked about',
-        '- for project',
-        '- focus on',
-        '- keep responses',
-        '- use emojis',
-        '- use bold',
-        '- use bullet points',
-        '- never include api',
-        'vs code context:',
-        '[vs code context',
-        'what is 2+2',  # Remove question echo
-        'current request:',
-        'previous conversation:',
-    ]
+    # Only apply minimal fixes
+    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?"
     
-    skip_mode = False
-    found_real_content = False
-    
-    for line in lines:
-        line_clean = line.strip()
-        line_lower = line_clean.lower()
-        
-        # Skip empty lines at the start
-        if not line_clean and not found_real_content:
-            continue
-        
-        # Check if this line contains system prompt artifacts
-        is_system_line = any(indicator in line_lower for indicator in system_indicators)
-        
-        if is_system_line:
-            # Start skipping mode when we hit system prompts
-            skip_mode = True
-            continue
-        
-        # If we hit actual content after system prompts, stop skipping
-        if skip_mode and line_clean and not is_system_line:
-            # Check if this looks like real content (not more system stuff)
-            if (len(line_clean) > 3 and 
-                not line_lower.startswith(('what is', 'calculate', 'result =', '```')) and
-                not re.match(r'^[#*\-\d\.\s]+$', line_clean)):
-                skip_mode = False
-                found_real_content = True
-        
-        # Add line if we're not in skip mode
-        if not skip_mode:
-            final_lines.append(line)
-            found_real_content = True
-    
-    # Step 5: Reconstruct the clean response
-    final_answer = '\n'.join(final_lines).strip()
-    
-    # Step 6: Handle special cases and final cleanup
-    if not final_answer or len(final_answer) < 10:
-        # Try to extract just a simple answer if available
-        if user_message and ('2+2' in user_message or '2 + 2' in user_message):
-            return "4\n\nThe answer to 2 + 2 is 4."
-        return "I understand your question. Could you please be more specific about what you'd like to know?"
-    
-    # Remove any remaining artifacts that might have slipped through
-    final_answer = re.sub(r'```\s*$', '', final_answer)  # Remove trailing code block
-    final_answer = re.sub(r'^\s*```.*?\n', '', final_answer)  # Remove leading code block start
-    final_answer = final_answer.strip()
-    
-    print(f"🧹 Final cleaned answer: {final_answer}")
-    return final_answer
+    print(f"🧹 Final cleaned answer length: {len(clean_text)}")
+    print(f"🧹 Final answer preview: {clean_text[:150]}...")
+    return clean_text
 
-def generate_response(messages: list, max_tokens: int = 200, temperature: float = 0.7) -> str:
+def generate_response(messages: list, max_tokens: int = 400, temperature: float = 0.7) -> str:
     """
-    Generate response using the model with MINIMAL system prompts.
+    FIXED: Generate response with higher token limits and better settings.
     """
     try:
-        # Create MINIMAL conversation - just the essentials
+        # Create clean conversation
         clean_messages = []
         
         # Add minimal system message
         clean_messages.append({
             "role": "system", 
-            "content": "You are a helpful assistant."
+            "content": "You are Apollo AI, a helpful coding assistant. Provide clear, complete explanations with proper code formatting."
         })
         
-        # Add only the last few user/assistant exchanges to avoid context pollution
-        user_messages = [msg for msg in messages if msg.get("role") == "user"]
-        if user_messages:
-            # Take only the latest user message to avoid confusion
-            latest_user = user_messages[-1]
-            clean_messages.append(latest_user)
+        # Add recent conversation context (last 2-3 messages)
+        recent_messages = messages[-3:] if len(messages) > 3 else messages
+        for msg in recent_messages:
+            if msg.get("role") in ["user", "assistant"]:
+                clean_messages.append(msg)
         
-        print(f"🔍 Clean messages: {clean_messages}")
+        print(f"🔍 Processing {len(clean_messages)} messages")
         
-        # Build the conversation using tokenizer's chat template
+        # Build conversation using tokenizer's chat template
         formatted_prompt = tokenizer.apply_chat_template(
             clean_messages,
             tokenize=False,
             add_generation_prompt=True
         )
         
-        print(f"🔍 Formatted prompt: {formatted_prompt}")
-        
-        # Tokenize with truncation to prevent overlong prompts
-        inputs = tokenizer(formatted_prompt, return_tensors="pt", truncation=True, max_length=1024)
+        # Tokenize with proper length limits
+        inputs = tokenizer(formatted_prompt, return_tensors="pt", truncation=True, max_length=1500)
         
-        # Generate with conservative settings
+        # FIXED: Generate with much higher token limits
         with torch.no_grad():
             outputs = model.generate(
                 inputs.input_ids,
                 attention_mask=inputs.attention_mask,
-                max_new_tokens=min(max_tokens, 150),  # Keep responses short
-                temperature=max(0.3, min(temperature, 0.8)),  # Controlled temperature
-                top_p=0.85,
+                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.1,
-                length_penalty=0.8,
-                early_stopping=True,
-                no_repeat_ngram_size=3,
+                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
             )
         
         # Decode the full response
         full_response = tokenizer.decode(outputs[0], skip_special_tokens=False)
         
-        # Extract user message for cleaning context
+        # Extract user message for context
         user_message = ""
-        for msg in clean_messages:
+        for msg in reversed(clean_messages):
             if msg.get("role") == "user":
                 user_message = msg.get("content", "")
+                break
         
-        # Clean and extract the answer
+        # Clean and extract the answer with gentle approach
         clean_answer = extract_clean_answer(full_response, formatted_prompt, user_message)
         
         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."
+        return f"I encountered an error while processing your request. Please try again with a simpler question."
 
 # === Routes ===
 @app.get("/")
@@ -252,7 +171,8 @@ def root():
     return {
         "message": "🤖 Apollo AI Backend is running!",
         "model": "Qwen2-0.5B-Instruct with LoRA",
-        "status": "ready"
+        "status": "ready",
+        "max_tokens": "500 (increased)"
     }
 
 @app.get("/health")
@@ -280,7 +200,7 @@ async def chat_completions(request: Request):
     try:
         body = await request.json()
         messages = body.get("messages", [])
-        max_tokens = body.get("max_tokens", 200)
+        max_tokens = body.get("max_tokens", 400)  # INCREASED default
         temperature = body.get("temperature", 0.7)
         
         if not messages or not isinstance(messages, list):
@@ -301,19 +221,19 @@ async def chat_completions(request: Request):
             )
 
     try:
-        # Generate response
-        print(f"📥 Processing {len(messages)} messages")
+        # Generate response with higher limits
+        print(f"📥 Processing {len(messages)} messages with max_tokens: {max_tokens}")
         response_content = generate_response(
             messages=messages,
-            max_tokens=min(max_tokens, 300),  # Cap max tokens
-            temperature=max(0.1, min(temperature, 1.0))  # Clamp temperature
+            max_tokens=min(max_tokens, 600),  # INCREASED cap to 600
+            temperature=max(0.1, min(temperature, 1.0))
         )
         
         # Return OpenAI-compatible response
         return {
             "id": f"chatcmpl-apollo-{hash(str(messages)) % 10000}",
             "object": "chat.completion",
-            "created": int(torch.tensor(0).item()),  # Simple timestamp
+            "created": int(torch.tensor(0).item()),
             "model": "qwen2-0.5b-instruct-lora",
             "choices": [
                 {
@@ -326,8 +246,8 @@ async def chat_completions(request: Request):
                 }
             ],
             "usage": {
-                "prompt_tokens": len(str(messages)),  # Approximate
-                "completion_tokens": len(response_content),  # Approximate
+                "prompt_tokens": len(str(messages)),
+                "completion_tokens": len(response_content),
                 "total_tokens": len(str(messages)) + len(response_content)
             }
         }
@@ -346,17 +266,19 @@ async def test_generation(request: Request):
     try:
         body = await request.json()
         prompt = body.get("prompt", "Hello, how are you?")
+        max_tokens = body.get("max_tokens", 300)
         
         messages = [
-            {"role": "system", "content": "You are a helpful assistant."},
+            {"role": "system", "content": "You are Apollo AI, a helpful assistant."},
             {"role": "user", "content": prompt}
         ]
         
-        response = generate_response(messages, max_tokens=150, temperature=0.7)
+        response = generate_response(messages, max_tokens=max_tokens, temperature=0.7)
         
         return {
             "prompt": prompt,
             "response": response,
+            "response_length": len(response),
             "status": "success"
         }
         
@@ -368,5 +290,6 @@ async def test_generation(request: Request):
 
 if __name__ == "__main__":
     import uvicorn
-    print("🚀 Starting Apollo AI Backend...")
+    print("🚀 Starting Apollo AI Backend with FIXED response limits...")
+    print("📊 Max tokens increased to 500+ for complete responses")
     uvicorn.run(app, host="0.0.0.0", port=7860)