app/main.py

import os
import torch
from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel
from starlette.middleware.cors import CORSMiddleware

# === Setup FastAPI ===
app = FastAPI(title="Apollo AI Backend - Qwen2-0.5B", version="3.0.0")

# === CORS ===
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# === Configuration ===
API_KEY = os.getenv("API_KEY", "aigenapikey1234567890")
BASE_MODEL = "Qwen/Qwen2-0.5B-Instruct"
ADAPTER_PATH = "adapter"

# === Load Model ===
print("🔧 Loading tokenizer for Qwen2-0.5B...")
tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, trust_remote_code=True)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

print("🧠 Loading Qwen2-0.5B base model...")
base_model = AutoModelForCausalLM.from_pretrained(
    BASE_MODEL,
    trust_remote_code=True,
    torch_dtype=torch.float32,
    device_map="cpu"
)

print("🔗 Applying LoRA adapter to Qwen2-0.5B...")
model = PeftModel.from_pretrained(base_model, ADAPTER_PATH)
model.eval()

print("✅ Qwen2-0.5B model ready!")

def create_conversation_prompt(messages: list, is_force_mode: bool) -> str:
    """
    Create a simple conversation prompt with appropriate system instruction
    """
    if is_force_mode:
        system_prompt = "DIRECT ANSWER MODE: Give immediate, complete solutions. Provide working code and clear explanations. Don't ask questions - just solve the problem directly."
    else:
        system_prompt = "TEACHER MODE: You are a teacher. NEVER give direct answers or solutions. ALWAYS respond with guiding questions. Ask things like 'What do you think this does?' or 'How would you approach this?' Help them discover answers themselves through questions."
    
    # Build conversation
    conversation = f"System: {system_prompt}\n\n"
    
    # Add last 6 messages (3 pairs) for context
    recent_messages = messages[-6:] if len(messages) > 6 else messages
    
    for msg in recent_messages:
        role = msg.get("role", "")
        content = msg.get("content", "")
        if role == "user":
            conversation += f"Student: {content}\n"
        elif role == "assistant":
            conversation += f"Assistant: {content}\n"
    
    conversation += "Assistant:"
    return conversation

def generate_response(messages: list, is_force_mode: bool = False, max_tokens: int = 200, temperature: float = 0.7) -> str:
    """
    Generate response using the actual AI model
    """
    try:
        # Create conversation prompt
        prompt = create_conversation_prompt(messages, is_force_mode)
        
        print(f"🎯 Generating {'FORCE' if is_force_mode else 'MENTOR'} response")
        print(f"🔍 DEBUG: force_mode = {is_force_mode}")
        print(f"📝 System prompt: {prompt.split('Student:')[0][:100]}...")
        
        # Tokenize input
        inputs = tokenizer(prompt, return_tensors="pt", max_length=1024, truncation=True)
        
        # Generate response
        with torch.no_grad():
            outputs = model.generate(
                inputs.input_ids,
                max_new_tokens=max_tokens,
                temperature=temperature,
                do_sample=True,
                pad_token_id=tokenizer.eos_token_id,
                eos_token_id=tokenizer.eos_token_id,
                top_p=0.9,
                repetition_penalty=1.1
            )
        
        # Decode response
        full_response = tokenizer.decode(outputs[0], skip_special_tokens=True)
        
        # Extract only the new generated part
        response = full_response[len(prompt):].strip()
        
        # Clean up response
        response = response.replace("Student:", "").replace("Assistant:", "").strip()
        
        # Remove any system mentions
        if response.startswith("System:"):
            response = response.split("\n", 1)[-1].strip()
        
        print(f"✅ Generated response length: {len(response)}")
        
        if not response or len(response) < 10:
            # Stronger fallback responses
            if is_force_mode:
                return "Here's the direct answer: I need you to provide a more specific question so I can give you the exact solution you need."
            else:
                return "What do you think this code is trying to do? Can you trace through it step by step and tell me what you notice?"
        
        return response
        
    except Exception as e:
        print(f"❌ Generation error: {e}")
        if is_force_mode:
            return "I encountered an error. Please try rephrasing your question."
        else:
            return "I had trouble processing that. Can you tell me what you're trying to understand?"

# === Routes ===
@app.get("/")
def root():
    return {
        "message": "🤖 Apollo AI Backend v3.0 - Qwen2-0.5B",
        "model": "Qwen/Qwen2-0.5B-Instruct with LoRA", 
        "status": "ready",
        "modes": {
            "mentor": "Guides learning with questions",
            "force": "Provides direct answers"
        }
    }

@app.get("/health")
def health():
    return {
        "status": "healthy", 
        "model_loaded": True, 
        "model_size": "0.5B"
    }

@app.post("/v1/chat/completions")
async def chat_completions(request: Request):
    # Validate API key
    auth_header = request.headers.get("Authorization", "")
    if not auth_header.startswith("Bearer "):
        return JSONResponse(
            status_code=401, 
            content={"error": "Missing or invalid Authorization header"}
        )

    token = auth_header.replace("Bearer ", "").strip()
    if token != API_KEY:
        return JSONResponse(
            status_code=401, 
            content={"error": "Invalid API key"}
        )

    # Parse request body
    try:
        body = await request.json()
        messages = body.get("messages", [])
        max_tokens = min(body.get("max_tokens", 200), 400)
        temperature = max(0.1, min(body.get("temperature", 0.7), 1.0))
        
        is_force_mode = body.get("force_mode", False)
        
        if not messages or not isinstance(messages, list):
            raise ValueError("Messages field is required and must be a list")
            
    except Exception as e:
        return JSONResponse(
            status_code=400, 
            content={"error": f"Invalid request body: {str(e)}"}
        )

    # Validate messages
    for i, msg in enumerate(messages):
        if not isinstance(msg, dict) or "role" not in msg or "content" not in msg:
            return JSONResponse(
                status_code=400,
                content={"error": f"Invalid message format at index {i}"}
            )

    try:
        print(f"📥 Processing request in {'FORCE' if is_force_mode else 'MENTOR'} mode")
        print(f"📊 Total messages: {len(messages)}")
        
        response_content = generate_response(
            messages=messages,
            is_force_mode=is_force_mode,
            max_tokens=max_tokens,
            temperature=temperature
        )
        
        return {
            "id": f"chatcmpl-apollo-{hash(str(messages)) % 10000}",
            "object": "chat.completion",
            "created": int(torch.tensor(0).item()),
            "model": f"qwen2-0.5b-{'force' if is_force_mode else 'mentor'}",
            "choices": [
                {
                    "index": 0,
                    "message": {
                        "role": "assistant",
                        "content": response_content
                    },
                    "finish_reason": "stop"
                }
            ],
            "usage": {
                "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:
        print(f"❌ Chat completion error: {e}")
        return JSONResponse(
            status_code=500,
            content={"error": f"Internal server error: {str(e)}"}
        )

if __name__ == "__main__":
    import uvicorn
    print("🚀 Starting Apollo AI Backend v3.0 - Simple & Clean...")
    print("🧠 Model: Qwen/Qwen2-0.5B-Instruct (500M parameters)")
    print("🎯 Mentor Mode: Asks guiding questions")
    print("⚡ Force Mode: Gives direct answers")
    uvicorn.run(app, host="0.0.0.0", port=7860)