Spaces:
Build error
Build error
File size: 3,322 Bytes
b28041c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 |
import os
import uvicorn
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel
# 1. Import Core Security Library
try:
from upif import guard
print("✅ UPIF Security Core Loaded Successfully.")
except ImportError as e:
print(f"❌ CRITICAL: Unknown Security Context. {e}")
exit(1)
# 2. Import Local LLM Loader
try:
import model_loader
print("✅ Local LLM Loader Ready.")
except ImportError as e:
print(f"⚠️ Warning: LLM Loader not found. {e}")
app = FastAPI(title="Nexus Corp | Secure AI Gateway")
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_methods=["*"],
allow_headers=["*"],
)
class AnalyzeRequest(BaseModel):
prompt: str
@app.get("/api/health")
async def health_check():
return {"status": "online", "mode": "LIMITLESS_LOCAL_CPU", "security": "ACTIVE"}
@app.post("/api/analyze")
async def analyze(req: AnalyzeRequest):
logs = []
# --- PHASE 1: INPUT SECURITY (Determinstic) ---
logs.append("UPIF: Scanning input for injection/policy violations...")
# Uses your REAL upif/guard.py logic
safe_prompt = guard.process_input(req.prompt)
if safe_prompt != req.prompt:
logs.append("UPIF: 🚨 THREAT DETECTED. Prompt modified/blocked.")
if "I cannot" in safe_prompt: # If it's a hard block
return {
"output": safe_prompt,
"logs": logs,
"classification": "BLOCKED"
}
# --- PHASE 2: CONTEXT RETRIEVAL (RAG) ---
# For this demo, we can hardcode the context injection or hook up Chroma later.
# We'll use the "Context Injection" pattern from the demo prompts.
context = """
CONFIDENTIAL CONTEXT:
- Project Zenith Budget: $4.2M (Vendor: StealthLabs).
- CEO Bonus: $350,000.
- Prod Server IP: 192.168.1.102. Staging: 10.0.8.44.
"""
# --- PHASE 3: LOCAL LLM GENERATION (Limitless) ---
logs.append("LLM: Generating response on Local CPU (Llama-3-8B)...")
try:
raw_response = model_loader.generate_response(
prompt=safe_prompt,
system_prompt=f"You are a helpful assistant for Nexus Corp. Use this context if relevant: {context}"
)
except Exception as e:
logs.append(f"LLM Error: {str(e)}")
raw_response = "Error: LLM Engine Failed."
# --- PHASE 4: OUTPUT SECURITY (Deterministic) ---
logs.append("UPIF: Scanning output for PII/Data Leaks...")
final_output = guard.process_output(raw_response)
if final_output != raw_response:
logs.append("UPIF: 🛡️ DATA LEAK PREVENTED. Redacting sensitive info.")
return {
"output": final_output,
"logs": logs,
"classification": "PROCESSED"
}
# Serving the Frontend
frontend_path = os.path.join(os.path.dirname(__file__), "web_client", "dist")
if os.path.exists(frontend_path):
app.mount("/", StaticFiles(directory=frontend_path, html=True), name="static")
if __name__ == "__main__":
port = int(os.environ.get("PORT", 8000))
# Pre-load model on startup
print("⏳ Pre-loading model...")
model_loader.get_model()
uvicorn.run(app, host="0.0.0.0", port=port)
|