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)