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 = "You are a helpful coding assistant. Give direct, clear answers with code examples when needed. Be concise and practical."
    else:
        system_prompt = "You are a teacher helping a student learn programming. Don't give direct answers. Instead, ask guiding questions to help them think and discover the solution themselves. Guide them step by step with questions like 'What do you think...?' or 'How would you...?'"
    
    # 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"📝 Prompt length: {len(prompt)}")
        
        # 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:
            # Fallback responses
            if is_force_mode:
                return "I need more specific information to provide a direct answer. Please clarify your question."
            else:
                return "That's an interesting question! What do you think the answer might be? Try to break it down step by step."
        
        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)