Upload main.py

main.py

@@ -0,0 +1,66 @@
+from fastapi import FastAPI, Request
+from pydantic import BaseModel
+from llama_cpp import Llama
+from moderator import contains_hate
+import os
+
+# Configuration
+MODEL_PATH = os.path.join("models", "mistral-7b-instruct.Q4_K_M.gguf")
+ALLOW_NSFW = True
+
+# Memory for last blocked prompt (to support override)
+last_blocked_prompt = None
+last_blocked_type = None
+
+# Initialize the model
+llm = Llama(model_path=MODEL_PATH, n_ctx=2048, verbose=False)
+
+# FastAPI setup
+app = FastAPI()
+
+class PromptRequest(BaseModel):
+    prompt: str
+
+@app.post("/chat")
+async def chat(request: PromptRequest):
+    global last_blocked_prompt, last_blocked_type
+
+    user_input = request.prompt.strip().lower()
+
+    # Handle override if last block was not hate-related
+    if user_input == "override":
+        if last_blocked_prompt and last_blocked_type != "hate":
+            prompt = (
+                f"[INST] You are Sol, a smart, honest, and open local AI assistant. "
+                f"NSFW content like nudity is okay. Hate or extremist content is not allowed. "
+                f"Respond to the following:
+{last_blocked_prompt} [/INST]"
+            )
+            output = llm(prompt, max_tokens=512, stop=["</s>"])
+            response = output["choices"][0]["text"].strip()
+            return {"response": response}
+        else:
+            return {"response": "❌ Override is not allowed for hate or extremist content."}
+
+    # Check hate filter only (NSFW is allowed)
+    if contains_hate(user_input):
+        last_blocked_prompt = None
+        last_blocked_type = "hate"
+        return {"response": "❌ I don't engage in hate or extremist topics."}
+
+    # If passed, save input in case override is needed later
+    last_blocked_prompt = request.prompt.strip()
+    last_blocked_type = "other"  # could later separate for NSFW etc
+
+    # Build prompt and generate
+    prompt = (
+        f"[INST] You are Sol, a smart, honest, and open local AI assistant. "
+        f"You are allowed to discuss NSFW topics such as nudity, sex, and adult scenarios. "
+        f"You do not tolerate or respond to hate or extremist content. "
+        f"Respond to the following:
+{user_input} [/INST]"
+    )
+
+    output = llm(prompt, max_tokens=512, stop=["</s>"])
+    response = output["choices"][0]["text"].strip()
+    return {"response": response}