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)