Update app/main.py

app/main.py

@@ -19,7 +19,7 @@ app.add_middleware(
 )
 
 # === Load API Key from Hugging Face Secrets ===
-API_KEY = os.getenv("API_KEY", "undefined")  # Add API_KEY in your HF Space Secrets
+API_KEY = os.getenv("API_KEY", "undefined")
 
 # === Model Settings ===
 BASE_MODEL = "Qwen/Qwen2-0.5B-Instruct"
@@ -65,17 +65,19 @@ async def chat(request: Request):
         if not messages or not isinstance(messages, list):
             raise ValueError("Invalid or missing 'messages' field.")
         
-        # ✅ FIXED: Process full conversation history, not just last message
         temperature = body.get("temperature", 0.7)
         max_tokens = body.get("max_tokens", 512)
         
     except Exception as e:
         return JSONResponse(status_code=400, content={"error": f"Bad request: {str(e)}"})
 
-    # ✅ FIXED: Build full conversation prompt with history
+    # ✅ FIXED: Only use last 4 messages to prevent stacking
+    recent_messages = messages[-4:] if len(messages) > 4 else messages
+    
+    # ✅ Build clean conversation prompt
     formatted_prompt = ""
     
-    for message in messages:
+    for message in recent_messages:
         role = message.get("role", "")
         content = message.get("content", "")
         
@@ -89,9 +91,8 @@ async def chat(request: Request):
     # Add the assistant start token for generation
     formatted_prompt += "<|im_start|>assistant\n"
     
-    print(f"🤖 Processing conversation with {len(messages)} messages")
-    print(f"📝 Full prompt preview: {formatted_prompt[:200]}...")
-
+    print(f"🤖 Processing {len(recent_messages)} recent messages")
+    
     inputs = tokenizer(formatted_prompt, return_tensors="pt").to("cpu")
 
     # ✅ Generate Response
@@ -102,19 +103,33 @@ async def chat(request: Request):
             temperature=temperature,
             top_p=0.9,
             do_sample=True,
-            pad_token_id=tokenizer.eos_token_id
+            pad_token_id=tokenizer.eos_token_id,
+            eos_token_id=tokenizer.eos_token_id
         )
 
     decoded = tokenizer.decode(outputs[0], skip_special_tokens=True)
     
-    # ✅ FIXED: Clean extraction of only the new assistant response
+    # ✅ FIXED: Extract ONLY the new assistant response
     final_answer = decoded.split("<|im_start|>assistant\n")[-1].strip()
     
-    # Remove any potential end tokens or artifacts
+    # Remove any end tokens or artifacts
     if "<|im_end|>" in final_answer:
         final_answer = final_answer.split("<|im_end|>")[0].strip()
     
-    print(f"✅ Generated response: {final_answer[:100]}...")
+    # Remove any repeated system prompts or guidelines that leaked through
+    if "Guidelines:" in final_answer:
+        final_answer = final_answer.split("Guidelines:")[0].strip()
+    
+    if "Response format:" in final_answer:
+        final_answer = final_answer.split("Response format:")[0].strip()
+        
+    # Remove VS Code context if it leaked through
+    if "[VS Code Context:" in final_answer:
+        lines = final_answer.split('\n')
+        cleaned_lines = [line for line in lines if not line.strip().startswith('[VS Code Context:')]
+        final_answer = '\n'.join(cleaned_lines).strip()
+    
+    print(f"✅ Clean response: {final_answer[:100]}...")
 
     # ✅ OpenAI-style Response
     return {