Spaces:
Runtime error
Runtime error
Download app.py from Evil-paradox007/chatbase-api: direct link, hf CLI and curl.
- Browser
- Download file 16.9 kB
-
https://huggingface.co/spaces/Evil-paradox007/chatbase-api/resolve/main/app.py
- Command line
-
hf download hf://spaces/Evil-paradox007/chatbase-api/app.py
-
curl -L -o app.py https://huggingface.co/spaces/Evil-paradox007/chatbase-api/resolve/main/app.py
16.9 kB
| # ============================================================================= | |
| # 🎓 FMU AI ASSISTANT — BACKEND API | |
| # ============================================================================= | |
| # Deploy to Hugging Face Spaces (Docker SDK) | |
| # RAG-powered chat using FAISS + Gemini + OpenRouter fallback | |
| # ============================================================================= | |
| import os | |
| import json | |
| import hashlib | |
| import time | |
| import shutil | |
| import httpx | |
| from pathlib import Path | |
| from typing import Optional, List, Dict, Any | |
| import numpy as np | |
| import faiss | |
| from fastapi import FastAPI, HTTPException, Request | |
| from fastapi.middleware.cors import CORSMiddleware | |
| from fastapi.responses import StreamingResponse, JSONResponse | |
| from pydantic import BaseModel, Field | |
| from sentence_transformers import SentenceTransformer | |
| import google.generativeai as genai | |
| # ============================================================================= | |
| # Configuration | |
| # ============================================================================= | |
| DATA_DIR = os.environ.get("DATA_DIR", "/data") | |
| GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY", "") | |
| OPENROUTER_API_KEY = os.environ.get("OPENROUTER_API_KEY", "") | |
| EMBED_MODEL_NAME = "all-MiniLM-L6-v2" | |
| TOP_K = 5 # Number of chunks to retrieve | |
| # Gemini fallback chain | |
| GEMINI_MODELS = [ | |
| "gemini-2.5-flash", | |
| "gemini-2.5-flash-lite", | |
| "gemini-3.1-flash-lite-preview" | |
| ] | |
| # OpenRouter free models (used after Gemini exhausted) | |
| OPENROUTER_MODELS = [ | |
| "google/gemini-2.0-flash-exp:free", | |
| "deepseek/deepseek-chat-v3-0324:free", | |
| "meta-llama/llama-4-maverick:free", | |
| "qwen/qwen3-235b-a22b:free", | |
| ] | |
| OPENROUTER_URL = "https://openrouter.ai/api/v1/chat/completions" | |
| # ============================================================================= | |
| # Initialize | |
| # ============================================================================= | |
| app = FastAPI( | |
| title="FMU AI Assistant API", | |
| description="RAG-powered chatbot backend using FAISS + Gemini + OpenRouter", | |
| version="2.0.0", | |
| ) | |
| # CORS — allow all origins for development | |
| app.add_middleware( | |
| CORSMiddleware, | |
| allow_origins=["*"], | |
| allow_credentials=True, | |
| allow_methods=["*"], | |
| allow_headers=["*"], | |
| ) | |
| # Load embedding model at startup | |
| print("Loading embedding model...") | |
| embed_model = SentenceTransformer(EMBED_MODEL_NAME) | |
| EMBED_DIM = embed_model.get_sentence_embedding_dimension() | |
| # Configure Gemini models | |
| gemini_models = [] | |
| if GEMINI_API_KEY: | |
| genai.configure(api_key=GEMINI_API_KEY) | |
| for model_name in GEMINI_MODELS: | |
| try: | |
| gemini_models.append({ | |
| 'name': model_name, | |
| 'provider': 'gemini', | |
| 'model': genai.GenerativeModel(model_name), | |
| }) | |
| print(f" ✅ Gemini: {model_name}") | |
| except Exception as e: | |
| print(f" ⚠️ Gemini failed: {model_name}: {e}") | |
| else: | |
| print("⚠️ GEMINI_API_KEY not set.") | |
| # Configure OpenRouter models | |
| openrouter_models = [] | |
| if OPENROUTER_API_KEY: | |
| for model_name in OPENROUTER_MODELS: | |
| openrouter_models.append({ | |
| 'name': model_name, | |
| 'provider': 'openrouter', | |
| }) | |
| print(f" ✅ OpenRouter: {model_name}") | |
| else: | |
| print("ℹ️ OPENROUTER_API_KEY not set (optional fallback).") | |
| # Unified fallback chain: Gemini first, then OpenRouter | |
| all_models = gemini_models + openrouter_models | |
| print(f"\n📦 Total models in fallback chain: {len(all_models)}") | |
| if not all_models: | |
| print("⚠️ No LLM models configured! Set GEMINI_API_KEY and/or OPENROUTER_API_KEY.") | |
| # In-memory store for chatbot data | |
| chatbots: Dict[str, Dict[str, Any]] = {} | |
| # ============================================================================= | |
| # Pydantic Models | |
| # ============================================================================= | |
| class ChunkData(BaseModel): | |
| chunk_id: int | |
| source_url: str = "" | |
| page_title: str = "" | |
| chunk_index: int = 0 | |
| total_chunks: int = 0 | |
| content: str | |
| class CreateChatbotRequest(BaseModel): | |
| chatbot_id: str | |
| source_url: str = "" | |
| total_chunks: int = 0 | |
| chunk_config: dict = {} | |
| chunks: List[ChunkData] | |
| class ChatMessage(BaseModel): | |
| role: str # "user" or "assistant" | |
| content: str | |
| class ChatRequest(BaseModel): | |
| message: str | |
| history: List[ChatMessage] = [] | |
| stream: bool = False | |
| class ChatbotSettings(BaseModel): | |
| name: Optional[str] = None | |
| system_prompt: Optional[str] = None | |
| temperature: Optional[float] = None | |
| initial_message: Optional[str] = None | |
| theme_color: Optional[str] = None | |
| # ============================================================================= | |
| # Helper Functions | |
| # ============================================================================= | |
| def get_chatbot_dir(chatbot_id: str) -> Path: | |
| """Get the directory for a specific chatbot's data.""" | |
| safe_id = hashlib.md5(chatbot_id.encode()).hexdigest()[:12] | |
| return Path(DATA_DIR) / f"chatbot_{safe_id}" | |
| def build_faiss_index(chunks: List[ChunkData]) -> tuple: | |
| """Build a FAISS index from text chunks.""" | |
| texts = [c.content for c in chunks] | |
| print(f" Embedding {len(texts)} chunks...") | |
| embeddings = embed_model.encode(texts, show_progress_bar=True, batch_size=32) | |
| embeddings = np.array(embeddings, dtype='float32') | |
| # Normalize for cosine similarity | |
| faiss.normalize_L2(embeddings) | |
| # Create index | |
| index = faiss.IndexFlatIP(EMBED_DIM) # Inner product = cosine on normalized | |
| index.add(embeddings) | |
| return index, texts | |
| def search_chunks(chatbot_id: str, query: str, top_k: int = TOP_K) -> List[Dict]: | |
| """Search for relevant chunks using FAISS.""" | |
| bot = chatbots.get(chatbot_id) | |
| if not bot: | |
| return [] | |
| query_embedding = embed_model.encode([query]) | |
| query_embedding = np.array(query_embedding, dtype='float32') | |
| faiss.normalize_L2(query_embedding) | |
| scores, indices = bot['index'].search(query_embedding, top_k) | |
| results = [] | |
| for score, idx in zip(scores[0], indices[0]): | |
| if idx < len(bot['texts']): | |
| results.append({ | |
| 'content': bot['texts'][idx], | |
| 'score': float(score), | |
| }) | |
| return results | |
| def build_rag_prompt(query: str, context_chunks: List[Dict], history: List[ChatMessage], system_prompt: str = "") -> str: | |
| """Build the RAG prompt for Gemini.""" | |
| if not system_prompt: | |
| system_prompt = ( | |
| "You are a helpful AI assistant. Answer questions based on the provided context. " | |
| "If the context doesn't contain relevant information, say so honestly. " | |
| "Be concise, accurate, and helpful. Use markdown formatting when appropriate." | |
| ) | |
| context_text = "\n\n---\n\n".join([c['content'] for c in context_chunks]) | |
| prompt = f"""{system_prompt} | |
| ## Context from Knowledge Base: | |
| {context_text} | |
| ## Conversation History: | |
| """ | |
| for msg in history[-6:]: # Last 6 messages for context | |
| role = "User" if msg.role == "user" else "Assistant" | |
| prompt += f"{role}: {msg.content}\n" | |
| prompt += f"\nUser: {query}\n\nAssistant:" | |
| return prompt | |
| # ============================================================================= | |
| # API Endpoints | |
| # ============================================================================= | |
| async def root(): | |
| return {"status": "ok", "service": "FMU AI Assistant API", "version": "1.0.0"} | |
| async def health(): | |
| return { | |
| "status": "healthy", | |
| "gemini_models": [m['name'] for m in gemini_models], | |
| "openrouter_models": [m['name'] for m in openrouter_models], | |
| "total_fallback_models": len(all_models), | |
| "chatbots_loaded": len(chatbots), | |
| "embed_model": EMBED_MODEL_NAME, | |
| } | |
| # --- Chatbot CRUD --- | |
| async def create_chatbot(request: CreateChatbotRequest): | |
| """Create a new chatbot from chunked data.""" | |
| chatbot_id = request.chatbot_id | |
| if not request.chunks: | |
| raise HTTPException(status_code=400, detail="No chunks provided") | |
| print(f"📦 Creating chatbot: {chatbot_id} with {len(request.chunks)} chunks") | |
| # Build FAISS index | |
| index, texts = build_faiss_index(request.chunks) | |
| # Store in memory | |
| chatbots[chatbot_id] = { | |
| 'id': chatbot_id, | |
| 'source_url': request.source_url, | |
| 'total_chunks': len(request.chunks), | |
| 'index': index, | |
| 'texts': texts, | |
| 'created_at': time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()), | |
| 'settings': { | |
| 'name': chatbot_id, | |
| 'system_prompt': '', | |
| 'temperature': 0.7, | |
| 'initial_message': f"Hi! I'm an AI assistant trained on {request.source_url}. How can I help you?", | |
| 'theme_color': '#6C63FF', | |
| }, | |
| 'chat_count': 0, | |
| 'message_count': 0, | |
| } | |
| # Save to disk for persistence | |
| bot_dir = get_chatbot_dir(chatbot_id) | |
| bot_dir.mkdir(parents=True, exist_ok=True) | |
| faiss.write_index(index, str(bot_dir / "index.faiss")) | |
| # Save metadata | |
| meta = {k: v for k, v in chatbots[chatbot_id].items() if k not in ('index', 'texts')} | |
| meta['texts'] = texts | |
| with open(bot_dir / "meta.json", 'w') as f: | |
| json.dump(meta, f, indent=2) | |
| print(f"✅ Chatbot '{chatbot_id}' created successfully!") | |
| return { | |
| "status": "success", | |
| "chatbot_id": chatbot_id, | |
| "total_chunks": len(request.chunks), | |
| "message": f"Chatbot '{chatbot_id}' created with {len(request.chunks)} chunks", | |
| } | |
| async def list_chatbots(): | |
| """List all chatbots.""" | |
| bots = [] | |
| for bot_id, bot in chatbots.items(): | |
| bots.append({ | |
| 'id': bot['id'], | |
| 'name': bot['settings']['name'], | |
| 'source_url': bot['source_url'], | |
| 'total_chunks': bot['total_chunks'], | |
| 'created_at': bot['created_at'], | |
| 'chat_count': bot.get('chat_count', 0), | |
| 'message_count': bot.get('message_count', 0), | |
| 'settings': bot['settings'], | |
| }) | |
| return {"chatbots": bots} | |
| async def get_chatbot(chatbot_id: str): | |
| """Get chatbot details.""" | |
| if chatbot_id not in chatbots: | |
| raise HTTPException(status_code=404, detail="Chatbot not found") | |
| bot = chatbots[chatbot_id] | |
| return { | |
| 'id': bot['id'], | |
| 'name': bot['settings']['name'], | |
| 'source_url': bot['source_url'], | |
| 'total_chunks': bot['total_chunks'], | |
| 'created_at': bot['created_at'], | |
| 'chat_count': bot.get('chat_count', 0), | |
| 'message_count': bot.get('message_count', 0), | |
| 'settings': bot['settings'], | |
| } | |
| async def update_settings(chatbot_id: str, settings: ChatbotSettings): | |
| """Update chatbot settings.""" | |
| if chatbot_id not in chatbots: | |
| raise HTTPException(status_code=404, detail="Chatbot not found") | |
| bot = chatbots[chatbot_id] | |
| update_data = settings.dict(exclude_none=True) | |
| bot['settings'].update(update_data) | |
| return {"status": "success", "settings": bot['settings']} | |
| async def delete_chatbot(chatbot_id: str): | |
| """Delete a chatbot.""" | |
| if chatbot_id not in chatbots: | |
| raise HTTPException(status_code=404, detail="Chatbot not found") | |
| del chatbots[chatbot_id] | |
| # Remove from disk | |
| bot_dir = get_chatbot_dir(chatbot_id) | |
| if bot_dir.exists(): | |
| shutil.rmtree(bot_dir) | |
| return {"status": "success", "message": f"Chatbot '{chatbot_id}' deleted"} | |
| # --- Chat --- | |
| def _call_gemini(model_entry, prompt, temperature): | |
| """Call a Gemini model and return response text.""" | |
| response = model_entry['model'].generate_content( | |
| prompt, | |
| generation_config=genai.GenerationConfig( | |
| temperature=temperature, | |
| max_output_tokens=2048, | |
| ), | |
| ) | |
| return response.text | |
| def _call_openrouter(model_name, prompt, temperature, history=None): | |
| """Call an OpenRouter model (OpenAI-compatible API) and return response text.""" | |
| messages = [{"role": "user", "content": prompt}] | |
| headers = { | |
| "Authorization": f"Bearer {OPENROUTER_API_KEY}", | |
| "Content-Type": "application/json", | |
| "HTTP-Referer": "https://fmu-ai-assistant.app", | |
| "X-Title": "FMU AI Assistant", | |
| } | |
| payload = { | |
| "model": model_name, | |
| "messages": messages, | |
| "temperature": temperature, | |
| "max_tokens": 2048, | |
| } | |
| resp = httpx.post(OPENROUTER_URL, headers=headers, json=payload, timeout=60) | |
| if resp.status_code == 429: | |
| raise Exception(f"429 Rate limit exceeded for {model_name}") | |
| if resp.status_code != 200: | |
| raise Exception(f"{resp.status_code} {resp.text[:200]}") | |
| data = resp.json() | |
| return data['choices'][0]['message']['content'] | |
| async def chat(chatbot_id: str, request: ChatRequest): | |
| """Send a message with automatic Gemini → OpenRouter fallback.""" | |
| if chatbot_id not in chatbots: | |
| raise HTTPException(status_code=404, detail="Chatbot not found") | |
| if not all_models: | |
| raise HTTPException(status_code=500, detail="No LLM models configured. Set GEMINI_API_KEY and/or OPENROUTER_API_KEY.") | |
| bot = chatbots[chatbot_id] | |
| bot['chat_count'] = bot.get('chat_count', 0) + 1 | |
| bot['message_count'] = bot.get('message_count', 0) + 1 | |
| # Retrieve relevant chunks | |
| context_chunks = search_chunks(chatbot_id, request.message, TOP_K) | |
| # Build prompt | |
| prompt = build_rag_prompt( | |
| query=request.message, | |
| context_chunks=context_chunks, | |
| history=request.history, | |
| system_prompt=bot['settings'].get('system_prompt', ''), | |
| ) | |
| temperature = bot['settings'].get('temperature', 0.7) | |
| # Try each model in unified fallback chain (Gemini → OpenRouter) | |
| last_error = None | |
| for model_entry in all_models: | |
| model_name = model_entry['name'] | |
| provider = model_entry['provider'] | |
| try: | |
| if provider == 'gemini': | |
| response_text = _call_gemini(model_entry, prompt, temperature) | |
| elif provider == 'openrouter': | |
| response_text = _call_openrouter(model_name, prompt, temperature) | |
| else: | |
| continue | |
| print(f" ✅ Response from [{provider}] {model_name}") | |
| return { | |
| "response": response_text, | |
| "model_used": f"{provider}/{model_name}", | |
| "sources": [{"content": c['content'][:200], "score": c['score']} for c in context_chunks[:3]], | |
| } | |
| except Exception as e: | |
| error_str = str(e) | |
| last_error = error_str | |
| is_quota = any(kw in error_str.lower() for kw in ['429', 'quota', 'rate limit', 'exceeded']) | |
| if is_quota: | |
| print(f" ⚠️ [{provider}] {model_name} — quota/rate limit, trying next...") | |
| continue | |
| else: | |
| raise HTTPException(status_code=500, detail=f"LLM error ({provider}/{model_name}): {error_str}") | |
| # All models exhausted | |
| raise HTTPException(status_code=429, detail=f"All models exhausted (tried {len(all_models)}). Try again later. Last error: {last_error}") | |
| # --- Startup: Load existing chatbots --- | |
| async def startup_load(): | |
| """Load previously saved chatbots from disk.""" | |
| data_path = Path(DATA_DIR) | |
| if not data_path.exists(): | |
| data_path.mkdir(parents=True, exist_ok=True) | |
| return | |
| for bot_dir in data_path.iterdir(): | |
| if bot_dir.is_dir() and bot_dir.name.startswith("chatbot_"): | |
| meta_path = bot_dir / "meta.json" | |
| index_path = bot_dir / "index.faiss" | |
| if meta_path.exists() and index_path.exists(): | |
| try: | |
| with open(meta_path) as f: | |
| meta = json.load(f) | |
| index = faiss.read_index(str(index_path)) | |
| texts = meta.pop('texts', []) | |
| chatbots[meta['id']] = { | |
| **meta, | |
| 'index': index, | |
| 'texts': texts, | |
| } | |
| print(f" ✅ Loaded chatbot: {meta['id']}") | |
| except Exception as e: | |
| print(f" ⚠️ Failed to load {bot_dir.name}: {e}") | |
| print(f"📦 Loaded {len(chatbots)} chatbots from disk") | |
| # ============================================================================= | |
| # Run with: uvicorn app:app --host 0.0.0.0 --port 7860 | |
| # ============================================================================= | |
| if __name__ == "__main__": | |
| import uvicorn | |
| uvicorn.run(app, host="0.0.0.0", port=7860) | |