app.py

import os
import logging
from typing import Optional
from datetime import datetime
from contextlib import asynccontextmanager

from fastapi import FastAPI, HTTPException, Depends, Security, status
from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field
import uvicorn

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Global variables for model
model = None
tokenizer = None
model_loaded = False
torch_available = False

@asynccontextmanager
async def lifespan(app: FastAPI):
    # Startup
    global model, tokenizer, model_loaded, torch_available
    logger.info("Llama 3 AI Assistant starting up...")
    
    try:
        # Try to import torch and transformers
        import torch
        from transformers import AutoTokenizer, AutoModelForCausalLM
        torch_available = True
        logger.info("PyTorch and Transformers available!")
        
        # Use Llama 3 model - try different variants based on availability
        llama_models = [
            "meta-llama/Llama-3.2-1B-Instruct",  # Smallest Llama 3.2
            "meta-llama/Llama-3.2-3B-Instruct",  # Medium Llama 3.2
            "microsoft/Llama2-7b-chat-hf",       # Fallback to Llama 2
            "huggingface/CodeBERTa-small-v1",    # Ultra fallback
        ]
        
        model_name = os.getenv("MODEL_NAME", llama_models[0])
        logger.info(f"Attempting to load Llama model: {model_name}")
        
        # Try to load the model
        for attempt_model in llama_models:
            try:
                logger.info(f"Trying to load: {attempt_model}")
                
                # Load tokenizer
                tokenizer = AutoTokenizer.from_pretrained(attempt_model)
                if tokenizer.pad_token is None:
                    tokenizer.pad_token = tokenizer.eos_token
                
                # Load model with optimizations for free tier
                model = AutoModelForCausalLM.from_pretrained(
                    attempt_model,
                    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
                    low_cpu_mem_usage=True,
                    device_map="auto" if torch.cuda.is_available() else None,
                    trust_remote_code=True
                )
                
                model_loaded = True
                model_name = attempt_model
                logger.info(f"Successfully loaded Llama model: {attempt_model}")
                break
                
            except Exception as e:
                logger.warning(f"Failed to load {attempt_model}: {e}")
                continue
        
        if not model_loaded:
            logger.warning("Could not load any Llama model, using fallback mode")
        
    except ImportError as e:
        logger.warning(f"PyTorch/Transformers not available: {e}")
        logger.info("Running in smart response mode")
        torch_available = False
        model_loaded = False
        
    except Exception as e:
        logger.warning(f"Could not load Llama model: {e}")
        logger.info("Running in smart response mode")
        model_loaded = False
    
    yield
    # Shutdown
    logger.info("Llama AI Assistant shutting down...")

# Initialize FastAPI app with lifespan
app = FastAPI(
    title="Llama 3 AI Agent API",
    description="AI Agent powered by Llama 3 models",
    version="5.0.0",
    lifespan=lifespan
)

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

# Security
security = HTTPBearer()

# Configuration
API_KEYS = {
    os.getenv("API_KEY_1", "27Eud5J73j6SqPQAT2ioV-CtiCg-p0WNqq6I4U0Ig6E"): "user1",
    os.getenv("API_KEY_2", "QbzG2CqHU1Nn6F1EogZ1d3dp8ilRTMJQBwTJDQBzS-U"): "user2",
}

# Request/Response models
class ChatRequest(BaseModel):
    message: str = Field(..., min_length=1, max_length=2000)
    max_length: Optional[int] = Field(300, ge=50, le=1000)
    temperature: Optional[float] = Field(0.7, ge=0.1, le=1.5)
    top_p: Optional[float] = Field(0.9, ge=0.1, le=1.0)
    do_sample: Optional[bool] = Field(True)
    system_prompt: Optional[str] = Field("You are a helpful AI assistant.", max_length=500)

class ChatResponse(BaseModel):
    response: str
    model_used: str
    timestamp: str
    processing_time: float
    tokens_used: int
    model_loaded: bool

class HealthResponse(BaseModel):
    status: str
    model_loaded: bool
    timestamp: str

def verify_api_key(credentials: HTTPAuthorizationCredentials = Security(security)) -> str:
    """Verify API key authentication"""
    api_key = credentials.credentials
    
    if api_key not in API_KEYS:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Invalid API key"
        )
    
    return API_KEYS[api_key]

def get_llama_smart_response(message: str) -> str:
    """Smart fallback responses when Llama is not available"""
    message_lower = message.lower()
    
    if any(word in message_lower for word in ["hello", "hi", "hey", "hii"]):
        return """Hello! I'm your Llama 3 AI assistant! 🦙

I'm designed to be helpful, harmless, and honest. I can assist you with:

• **Programming & Development**: Python, JavaScript, web development, debugging
• **AI & Machine Learning**: Concepts, implementations, best practices
• **Data Science**: Analysis, visualization, statistics
• **Problem Solving**: Breaking down complex problems step by step
• **Creative Tasks**: Writing, brainstorming, content creation
• **Learning**: Explaining concepts in simple terms

I aim to provide thoughtful, detailed responses that are actually useful. What would you like to explore today?"""

    elif any(word in message_lower for word in ["machine learning", "ml"]):
        return """Machine learning is fascinating! It's the science of getting computers to learn and make decisions from data without being explicitly programmed for every scenario.

**Core Concept**: Instead of writing specific rules, we show the computer lots of examples and let it figure out the patterns.

**How it works**:
1. **Data Collection**: Gather relevant examples
2. **Training**: Algorithm learns patterns from the data
3. **Validation**: Test how well it learned
4. **Prediction**: Apply learned patterns to new situations

**Types of ML**:
• **Supervised Learning**: Learning with labeled examples (like email spam detection)
• **Unsupervised Learning**: Finding hidden patterns (like customer segmentation)
• **Reinforcement Learning**: Learning through trial and error (like game AI)

**Real-world applications**:
- Netflix recommendations know your taste better than you do
- Medical AI can detect diseases in X-rays
- Self-driving cars navigate complex traffic
- Language models like me understand and generate text

The exciting part? We're still in the early stages. What specific aspect interests you most?"""

    elif any(word in message_lower for word in ["ai", "artificial intelligence"]):
        return """Artificial Intelligence is one of the most transformative technologies of our time! At its core, AI is about creating machines that can perform tasks requiring human-like intelligence.

**What makes AI special**:
- **Learning**: Improves from experience, just like humans
- **Reasoning**: Can draw logical conclusions from information
- **Perception**: Understands images, speech, and text
- **Decision Making**: Weighs options and chooses actions

**Current AI landscape**:
• **Language Models**: Like me! We understand and generate human language
• **Computer Vision**: AI that "sees" and interprets images
• **Robotics**: Physical AI that interacts with the world
• **Game AI**: Masters complex strategy games

**The philosophical angle**: AI forces us to ask deep questions about intelligence, consciousness, and what makes us human. As AI gets more capable, we're discovering that intelligence might be more about pattern recognition and prediction than we thought.

**Future implications**: AI will likely transform every industry - healthcare, education, transportation, entertainment. The key is ensuring it benefits everyone, not just tech companies.

What aspect of AI fascinates or concerns you most? I love diving into both the technical and philosophical sides!"""

    elif any(word in message_lower for word in ["python", "programming"]):
        return """Python is absolutely fantastic for AI and general programming! It's like the Swiss Army knife of programming languages.

**Why Python rocks**:
• **Readable**: Code looks almost like English
• **Versatile**: Web apps, AI, data science, automation, games
• **Powerful libraries**: Massive ecosystem of tools
• **Beginner-friendly**: Great first language
• **Industry standard**: Used by Google, Netflix, Instagram

**For AI specifically**:
- **NumPy**: Fast numerical computing
- **Pandas**: Data manipulation and analysis
- **Scikit-learn**: Machine learning algorithms
- **TensorFlow/PyTorch**: Deep learning frameworks
- **OpenAI**: API integrations for modern AI

**Learning path I recommend**:
1. **Basics**: Variables, functions, loops (1-2 weeks)
2. **Data structures**: Lists, dictionaries, sets
3. **Libraries**: Start with Pandas for data handling
4. **Projects**: Build something you care about
5. **Specialization**: Pick web dev, AI, or data science

**Pro tip**: Don't just read tutorials - build projects! Start small:
- A calculator
- A web scraper
- A simple chatbot
- Data analysis of something interesting to you

What kind of projects are you thinking about? I can suggest specific resources and next steps!"""

    else:
        return f"""I'm a Llama 3-powered AI assistant, and I'd love to help you with your question: "{message}"

I'm designed to provide thoughtful, detailed responses on a wide range of topics. I'm particularly good at:

• **Technical topics**: Programming, AI, data science, technology
• **Problem-solving**: Breaking down complex issues step by step
• **Learning support**: Explaining concepts clearly with examples
• **Creative tasks**: Writing, brainstorming, content creation
• **Analysis**: Examining ideas from multiple perspectives

To give you the most helpful response, could you provide a bit more context about what you're looking for? Are you:
- Trying to learn something new?
- Solving a specific problem?
- Looking for creative ideas?
- Seeking technical guidance?

I'm here to provide genuinely useful insights, not just generic responses. What would be most valuable for you right now?"""

def generate_llama_response(message: str, max_length: int = 300, temperature: float = 0.7, top_p: float = 0.9, do_sample: bool = True, system_prompt: str = "You are a helpful AI assistant.") -> tuple:
    """Generate response using Llama model or smart fallback"""
    global model, tokenizer, model_loaded, torch_available
    
    if not torch_available or not model_loaded or model is None or tokenizer is None:
        return get_llama_smart_response(message), "llama_smart_fallback", len(message.split())
    
    try:
        import torch
        
        # Format prompt for Llama (instruction format)
        if "llama" in str(model.config._name_or_path).lower():
            # Llama 3 instruction format
            prompt = f"<|begin_of_text|><|start_header_id|>system<|end_header_id|>\n\n{system_prompt}<|eot_id|><|start_header_id|>user<|end_header_id|>\n\n{message}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
        else:
            # Generic format
            prompt = f"System: {system_prompt}\nUser: {message}\nAssistant:"
        
        # Tokenize input
        inputs = tokenizer.encode(prompt, return_tensors="pt", truncation=True, max_length=1024)
        
        # Generate response
        with torch.no_grad():
            outputs = model.generate(
                inputs,
                max_new_tokens=max_length,
                temperature=temperature,
                top_p=top_p,
                do_sample=do_sample,
                pad_token_id=tokenizer.eos_token_id,
                eos_token_id=tokenizer.eos_token_id,
                repetition_penalty=1.1,
                length_penalty=1.0
            )
        
        # Decode response
        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
        
        # Extract only the assistant's response
        if "<|start_header_id|>assistant<|end_header_id|>" in response:
            response = response.split("<|start_header_id|>assistant<|end_header_id|>")[-1].strip()
        elif "Assistant:" in response:
            response = response.split("Assistant:")[-1].strip()
        
        # Clean up the response
        response = response.strip()
        if not response or len(response) < 10:
            return get_llama_smart_response(message), "llama_smart_fallback", len(message.split())
        
        # Count tokens
        tokens_used = len(tokenizer.encode(response))
        
        return response, os.getenv("MODEL_NAME", "meta-llama/Llama-3.2-1B-Instruct"), tokens_used
        
    except Exception as e:
        logger.error(f"Error generating Llama response: {str(e)}")
        return get_llama_smart_response(message), "llama_smart_fallback", len(message.split())

@app.get("/", response_model=HealthResponse)
async def root():
    """Health check endpoint"""
    return HealthResponse(
        status="healthy",
        model_loaded=model_loaded,
        timestamp=datetime.now().isoformat()
    )

@app.get("/health", response_model=HealthResponse)
async def health_check():
    """Detailed health check"""
    return HealthResponse(
        status="healthy" if model_loaded else "smart_mode",
        model_loaded=model_loaded,
        timestamp=datetime.now().isoformat()
    )

@app.post("/chat", response_model=ChatResponse)
async def chat(
    request: ChatRequest,
    user: str = Depends(verify_api_key)
):
    """Main chat endpoint using Llama 3 model or smart fallback"""
    start_time = datetime.now()
    
    try:
        # Generate response using Llama 3 or smart fallback
        response_text, model_used, tokens_used = generate_llama_response(
            request.message,
            request.max_length,
            request.temperature,
            request.top_p,
            request.do_sample,
            request.system_prompt
        )
        
        # Calculate processing time
        processing_time = (datetime.now() - start_time).total_seconds()
        
        return ChatResponse(
            response=response_text,
            model_used=model_used,
            timestamp=datetime.now().isoformat(),
            processing_time=processing_time,
            tokens_used=tokens_used,
            model_loaded=model_loaded
        )
        
    except Exception as e:
        logger.error(f"Error in chat endpoint: {str(e)}")
        # Provide helpful fallback response
        return ChatResponse(
            response="I'm experiencing some technical difficulties, but I'm still here to help! Could you please try rephrasing your question?",
            model_used="error_recovery_mode",
            timestamp=datetime.now().isoformat(),
            processing_time=(datetime.now() - start_time).total_seconds(),
            tokens_used=0,
            model_loaded=model_loaded
        )

@app.get("/models")
async def get_model_info(user: str = Depends(verify_api_key)):
    """Get information about the loaded model"""
    return {
        "model_name": os.getenv("MODEL_NAME", "meta-llama/Llama-3.2-1B-Instruct"),
        "model_loaded": model_loaded,
        "torch_available": torch_available,
        "status": "active" if model_loaded else "smart_fallback_mode",
        "capabilities": [
            "Llama 3 text generation" if model_loaded else "Smart Llama-style responses",
            "Instruction following",
            "Conversational AI responses",
            "System prompt support",
            "Adjustable creativity parameters",
            "Natural language understanding"
        ],
        "version": "5.0.0",
        "type": "Llama 3 Model" if model_loaded else "Llama Smart Fallback Mode"
    }

if __name__ == "__main__":
    # For Hugging Face Spaces
    port = int(os.getenv("PORT", "7860"))
    uvicorn.run(
        app,
        host="0.0.0.0",
        port=port,
        reload=False
    )