Upload main.py

main.py

@@ -0,0 +1,44 @@
+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  # ✅ Flip this to False if you ever want to re-enable NSFW filtering
+
+# Initialize the model
+llm = Llama(model_path=MODEL_PATH, n_ctx=2048, verbose=False)
+
+# FastAPI setup
+app = FastAPI()
+
+# Request body schema
+class PromptRequest(BaseModel):
+    prompt: str
+
+# Main route
+@app.post("/chat")
+async def chat(request: PromptRequest):
+    user_input = request.prompt.strip()
+
+    # Hate speech check
+    if contains_hate(user_input):
+        return {"response": "❌ I don't engage in hate or extremist topics."}
+
+    # Optional: Add NSFW check logic here if needed
+    if not ALLOW_NSFW:
+        if "nsfw" in user_input.lower() or "sex" in user_input.lower():
+            return {"response": "⚠️ NSFW topics are currently disabled."}
+
+    # Instruction prompt with optional NSFW allowance
+    prompt = (
+        f"[INST] You are Sol, a smart, honest, and free-talking local AI assistant. "
+        f"You answer clearly and do not censor NSFW topics unless they are hateful or extreme. "
+        f"Now respond to the following:\n{user_input} [/INST]"
+    )
+
+    output = llm(prompt, max_tokens=512, stop=["</s>"])
+    response = output["choices"][0]["text"].strip()
+    return {"response": response}