trackingsvg/test2
0
1import os2import json3import re4import requests5import threading6import gc7from datetime import datetime8from bs4 import BeautifulSoup9from contextlib import asynccontextmanager10from fastapi import FastAPI, Request11from fastapi.responses import JSONResponse12from fastapi.middleware.cors import CORSMiddleware13from huggingface_hub import hf_hub_download14from llama_cpp import Llama 15 16# ==========================================17# 1. CONFIGURATION & CORE SETUP18# ==========================================19# IMPORTANT: Set REAPERAI_SECRET as a Secret in your Space Settings20SECRET_KEY = os.environ.get("REAPERAI_SECRET", "jan30")21 22# Optimized for 16GB RAM: Smaller, heavily quantized model23MODEL_REPO = "bartowski/Qwen2.5-1.5B-Instruct-GGUF"24MODEL_FILE = "Qwen2.5-1.5B-Instruct-Q4_K_M.gguf" # Exact filename25 26chat_memory = {}27MAX_GLOBAL_USERS = 5028memory_lock = threading.Lock()29llm = None # Initialize as None30model_semaphore = threading.Semaphore(1) # Only 1 inference at a time31 32# ==========================================33# 2. MODEL LOADING (OPTIMIZED FOR SPACES)34# ==========================================35@asynccontextmanager36async def lifespan(app: FastAPI):37 """38 Lifespan handler for FastAPI startup/shutdown.39 Downloads and loads the model.40 """41 print(f"--- [SYSTEM] Initializing ReaperAI on Hugging Face Space ---")42 global llm43 44 try:45 # Step 1: Download model to cache (will use /tmp from Dockerfile)46 print(f"--- [SYSTEM] Downloading model: {MODEL_REPO}/{MODEL_FILE} ---")47 model_path = hf_hub_download(48 repo_id=MODEL_REPO,49 filename=MODEL_FILE,50 cache_dir=os.getenv("HF_HOME", "/tmp")51 )52 53 # Step 2: Load with optimized settings for 2 vCPU / 16GB RAM54 print(f"--- [SYSTEM] Loading model into RAM (this may take a moment) ---")55 llm = Llama(56 model_path=model_path,57 n_ctx=1024, # Reduced for memory efficiency58 n_threads=2, # Matches your 2 vCPUs59 n_gpu_layers=0, # CPU only60 verbose=False61 )62 print(f"--- [SYSTEM] Model loaded successfully. ReaperAI is ready. ---")63 64 except Exception as e:65 print(f"--- [CRITICAL ERROR] Model loading failed: {str(e)} ---")66 llm = None # Ensure it's None if loading fails67 68 yield # App runs here69 70 # Cleanup on shutdown (optional)71 if llm is not None:72 del llm73 gc.collect()74 75# ==========================================76# 3. FASTAPI APP INITIALIZATION77# ==========================================78app = FastAPI(79 title="ReaperAI Secure Core",80 description="AI Assistant with Web Search Capabilities",81 version="2.0",82 lifespan=lifespan83)84 85# CORS configuration86app.add_middleware(87 CORSMiddleware,88 allow_origins=["*"],89 allow_methods=["POST", "GET"],90 allow_headers=["*"],91)92 93# ==========================================94# 4. AUTONOMOUS TOOLS (REFINED)95# ==========================================96def ddg_search(query):97 """Perform a DuckDuckGo search and return top 3 results."""98 print(f"--- [TOOL] Searching Web: {query} ---")99 try:100 headers = {"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36"}101 res = requests.get(102 "https://html.duckduckgo.com/html/",103 params={"q": query, "kl": "us-en"},104 headers=headers,105 timeout=8106 )107 res.raise_for_status()108 109 soup = BeautifulSoup(res.text, "html.parser")110 results = []111 112 for r in soup.select(".result")[:3]: # Limit to 3 results113 title_elem = r.select_one('.result__a')114 snippet_elem = r.select_one('.result__snippet')115 116 if title_elem and snippet_elem:117 title = title_elem.get_text(strip=True)118 snippet = snippet_elem.get_text(strip=True)[:200] # Truncate119 results.append(f"• {title}: {snippet}")120 121 if results:122 return f"\n[REAL-TIME SEARCH RESULTS]:\n" + "\n".join(results) + "\n"123 return ""124 125 except Exception as e:126 print(f"--- [TOOL ERROR] Search failed: {e} ---")127 return ""128 129def jina_read(url):130 """Fetch and parse content from a URL using Jina Reader."""131 print(f"--- [TOOL] Reading Source: {url} ---")132 try:133 # Clean the URL134 url = url.strip()135 if not url.startswith(('http://', 'https://')):136 url = 'https://' + url137 138 res = requests.get(139 f"https://r.jina.ai/{url}",140 headers={"x-respond-with": "text", "User-Agent": "ReaperAI/2.0"},141 timeout=10142 )143 res.raise_for_status()144 145 # Extract first 1200 chars for context146 content = res.text[:1200].strip()147 if content:148 return f"\n[SOURCE CONTENT]:\n{content}\n"149 return ""150 151 except Exception as e:152 print(f"--- [TOOL ERROR] URL read failed: {e} ---")153 return ""154 155def fast_intent_detection(message):156 """Detect user intent from message."""157 message_lower = message.lower()158 159 # Check for URL160 url_match = re.search(r"(https?://\S+)", message)161 if url_match:162 return "URL", url_match.group(1)163 164 # Check for search keywords165 search_keywords = [166 "who is", "what is", "how to", "price of", "latest", "current",167 "news", "today", "weather", "score", "stock", "update", "2024",168 "define", "explain"169 ]170 171 if any(keyword in message_lower for keyword in search_keywords):172 return "SEARCH", message173 174 return "CHAT", None175 176# ==========================================177# 5. SECURITY MIDDLEWARE & ENDPOINTS178# ==========================================179@app.middleware("http")180async def security_guard(request: Request, call_next):181 """Security middleware for API key validation."""182 # Allow root endpoint without auth183 if request.url.path == "/":184 return await call_next(request)185 186 # Check for API key in headers187 if request.headers.get("x-reaperai-key") != SECRET_KEY:188 return JSONResponse(189 status_code=403,190 content={"error": "ACCESS_DENIED", "message": "Invalid or missing API key"}191 )192 193 return await call_next(request)194 195@app.get("/")196async def root():197 """Root endpoint for health checks."""198 status = "ready" if llm is not None else "loading"199 return {200 "status": status,201 "service": "ReaperAI Secure Core",202 "version": "2.0",203 "model_loaded": llm is not None,204 "endpoints": {"/chat": "POST", "/health": "GET"}205 }206 207@app.get("/health")208async def health_check():209 """Health check endpoint for monitoring."""210 return {211 "status": "healthy" if llm is not None else "unhealthy",212 "model": MODEL_REPO if llm is not None else None,213 "memory_users": len(chat_memory),214 "timestamp": datetime.now().isoformat()215 }216 217@app.post("/chat")218async def chat_endpoint(request: Request):219 """Main chat endpoint."""220 try:221 payload = await request.json()222 except:223 return JSONResponse(224 status_code=400,225 content={"error": "INVALID_JSON", "response": "Request must be valid JSON"}226 )227 228 user_id = payload.get("userId", "default")229 message = payload.get("message", "").strip()230 current_date = datetime.now().strftime("%A, %B %d, %Y")231 232 # Validate input233 if not message:234 return JSONResponse(235 status_code=400,236 content={"error": "EMPTY_MESSAGE", "response": "Message cannot be empty"}237 )238 239 # Check if model is loaded240 if llm is None:241 return JSONResponse(242 status_code=503,243 content={244 "error": "MODEL_NOT_LOADED",245 "response": "AI model is still initializing. Please try again in 30 seconds."246 }247 )248 249 # Manage conversation history (thread-safe)250 with memory_lock:251 # Clean up old users if needed252 if len(chat_memory) > MAX_GLOBAL_USERS:253 oldest_user = next(iter(chat_memory))254 del chat_memory[oldest_user]255 256 # Get user's history (last 5 exchanges)257 if user_id not in chat_memory:258 chat_memory[user_id] = []259 history = chat_memory[user_id][-5:]260 261 # Determine intent and gather context262 intent, data = fast_intent_detection(message)263 264 context = ""265 if intent == "URL":266 context = jina_read(data)267 elif intent == "SEARCH":268 context = ddg_search(message)269 270 # Build conversation messages271 messages = [272 {273 "role": "system",274 "content": f"""You are ReaperAI, a helpful and concise AI assistant.275 Current Date: {current_date}276 Instructions:277 1. Be direct and informative278 2. Use provided context when available279 3. Keep responses under 300 words280 4. If you don't know, say so281 """282 }283 ]284 285 # Add conversation history286 for h in history:287 messages.append({"role": "user", "content": h['u']})288 messages.append({"role": "assistant", "content": h['a']})289 290 # Add current query with context291 if context:292 final_query = f"Context:\n{context}\n\nUser Query: {message}"293 else:294 final_query = message295 296 messages.append({"role": "user", "content": final_query})297 298 # Generate response (thread-safe)299 with model_semaphore:300 try:301 response = llm.create_chat_completion(302 messages=messages,303 max_tokens=400, # Limit response length304 temperature=0.7,305 stop=["###", "User:", "Assistant:"]306 )307 ai_response = response["choices"][0]["message"]["content"].strip()308 except Exception as e:309 print(f"--- [INFERENCE ERROR] {str(e)} ---")310 ai_response = f"I encountered an error processing your request. Please try again."311 312 # Update conversation history (thread-safe)313 with memory_lock:314 chat_memory[user_id].append({"u": message, "a": ai_response})315 # Keep last 8 exchanges per user316 chat_memory[user_id] = chat_memory[user_id][-8:]317 318 # Clean up319 gc.collect()320 321 return {322 "intent": intent,323 "response": ai_response,324 "context_used": bool(context),325 "user_id": user_id,326 "timestamp": datetime.now().isoformat()327 }328 329# ==========================================330# 6. MAIN EXECUTION331# ==========================================332if __name__ == "__main__":333 import uvicorn334 print("--- [SYSTEM] Starting ReaperAI Server ---")335 uvicorn.run(336 "main:app",337 host="0.0.0.0",338 port=7860,339 reload=False, # Disable reload in production340 timeout_keep_alive=60,341 access_log=True342 )