llm-ai-agent / app.py
Yadav122's picture
Fix: Handle torch import errors with smart fallback mode
352df25 verified
raw
history blame
13.1 kB
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("Real LLM 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 a better conversational model
model_name = os.getenv("MODEL_NAME", "microsoft/DialoGPT-small") # Use small for better compatibility
logger.info(f"Loading real LLM model: {model_name}")
# Load tokenizer
tokenizer = AutoTokenizer.from_pretrained(model_name)
if tokenizer.pad_token is None:
tokenizer.pad_token = tokenizer.eos_token
# Load model with optimizations
model = AutoModelForCausalLM.from_pretrained(
model_name,
torch_dtype=torch.float32,
low_cpu_mem_usage=True,
pad_token_id=tokenizer.eos_token_id
)
model_loaded = True
logger.info("Real LLM model loaded successfully!")
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 LLM model: {e}")
logger.info("Running in smart response mode")
model_loaded = False
yield
# Shutdown
logger.info("AI Assistant shutting down...")
# Initialize FastAPI app with lifespan
app = FastAPI(
title="Real LLM AI Agent API",
description="AI Agent powered by actual LLM models with fallback",
version="4.1.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(200, ge=50, le=500)
temperature: Optional[float] = Field(0.8, 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)
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_smart_fallback_response(message: str) -> str:
"""Smart fallback responses when LLM 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 AI assistant. I'm currently running in smart mode while the full LLM model loads.
I can still help you with questions about:
• Machine Learning and AI concepts
• Programming and Python
• Data Science topics
• Technology explanations
• General conversations
What would you like to know about? I'll do my best to provide helpful information!"""
elif any(word in message_lower for word in ["machine learning", "ml"]):
return """Machine learning is a fascinating field! It's a subset of artificial intelligence where computers learn to make predictions or decisions by finding patterns in data, rather than being explicitly programmed for every scenario.
Key concepts:
• **Training**: The model learns from example data
• **Patterns**: It identifies relationships and trends
• **Prediction**: It applies learned patterns to new data
• **Improvement**: Performance gets better with more data
Common applications include recommendation systems (like Netflix suggestions), image recognition, natural language processing, and autonomous vehicles.
Would you like me to explain any specific aspect of machine learning in more detail?"""
elif any(word in message_lower for word in ["ai", "artificial intelligence"]):
return """Artificial Intelligence is the simulation of human intelligence in machines! It's about creating systems that can think, learn, and solve problems.
Current AI can:
• Understand and generate human language
• Recognize images and objects
• Play complex games at superhuman levels
• Drive cars autonomously
• Discover new medicines
Types of AI:
• **Narrow AI**: Specialized for specific tasks (what we have today)
• **General AI**: Human-level intelligence across all domains (future goal)
• **Super AI**: Beyond human intelligence (theoretical)
AI is transforming every industry and changing how we work, learn, and live. What aspect of AI interests you most?"""
elif any(word in message_lower for word in ["python", "programming"]):
return """Python is an excellent choice for AI and programming! It's known for its simple, readable syntax and powerful capabilities.
Why Python is great:
• **Easy to learn**: Clear, English-like syntax
• **Versatile**: Web development, AI, data science, automation
• **Rich ecosystem**: Thousands of libraries and frameworks
• **Community**: Large, helpful developer community
For AI/ML specifically:
• **NumPy**: Numerical computing
• **Pandas**: Data manipulation
• **Scikit-learn**: Machine learning algorithms
• **TensorFlow/PyTorch**: Deep learning
Python lets you focus on solving problems rather than wrestling with complex syntax. Are you interested in learning Python for a specific purpose?"""
else:
return f"""I understand you're asking about: "{message}"
I'm currently running in smart mode while the full LLM model loads. I can provide helpful information on topics like:
• **Technology**: AI, machine learning, programming
• **Science**: Data science, computer science concepts
• **Learning**: Programming languages, career advice
• **General**: Explanations, discussions, problem-solving
Could you be more specific about what you'd like to know? I'm here to help and will provide the most useful information I can!
If you're looking for creative writing, storytelling, or very specific technical details, the full LLM model will provide even better responses once it's loaded."""
def generate_llm_response(message: str, max_length: int = 200, temperature: float = 0.8, top_p: float = 0.9, do_sample: bool = True) -> tuple:
"""Generate response using actual LLM model or smart fallback"""
global model, tokenizer, model_loaded, torch_available
if not torch_available:
return get_smart_fallback_response(message), "smart_fallback_mode", len(message.split())
if not model_loaded or model is None or tokenizer is None:
return get_smart_fallback_response(message), "smart_fallback_mode", len(message.split())
try:
import torch
# Prepare input with conversation format
input_text = f"Human: {message}\nAssistant:"
# Tokenize input
inputs = tokenizer.encode(input_text, return_tensors="pt")
# Generate response
with torch.no_grad():
outputs = model.generate(
inputs,
max_length=inputs.shape[1] + 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,
num_return_sequences=1,
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 "Assistant:" in response:
response = response.split("Assistant:")[-1].strip()
# Remove the input text if it's still there
if input_text.replace("Assistant:", "").strip() in response:
response = response.replace(input_text.replace("Assistant:", "").strip(), "").strip()
# Clean up the response
response = response.strip()
if not response or len(response) < 10:
return get_smart_fallback_response(message), "smart_fallback_mode", len(message.split())
# Count tokens
tokens_used = len(tokenizer.encode(response))
return response, os.getenv("MODEL_NAME", "microsoft/DialoGPT-small"), tokens_used
except Exception as e:
logger.error(f"Error generating LLM response: {str(e)}")
return get_smart_fallback_response(message), "smart_fallback_mode", 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 real LLM model or smart fallback"""
start_time = datetime.now()
try:
# Generate response using actual LLM or smart fallback
response_text, model_used, tokens_used = generate_llm_response(
request.message,
request.max_length,
request.temperature,
request.top_p,
request.do_sample
)
# 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)}")
# Even if there's an error, provide a helpful 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", "microsoft/DialoGPT-small"),
"model_loaded": model_loaded,
"torch_available": torch_available,
"status": "active" if model_loaded else "smart_fallback_mode",
"capabilities": [
"Real LLM text generation" if model_loaded else "Smart fallback responses",
"Conversational AI responses",
"Dynamic response generation" if model_loaded else "Contextual smart responses",
"Adjustable temperature and top_p" if model_loaded else "Fixed high-quality responses",
"Natural language understanding"
],
"version": "4.1.0",
"type": "Real LLM Model" if model_loaded else "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
)