CoolFace
Apppublic

trackingsvg/test2

sourceHugging Faceupdated 9mo agoView on Hugging Face
0likes
main.py342 linesDownload Raw Back to root
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    )