Upload files to "/"

llm_router_v3.py

@@ -32,6 +32,7 @@ from io import BytesIO
 # ---------------------------------------------------------------------------
 OLLAMA_URL = "http://ollama:11434"
 CLASSIFIER_MODEL = "qwen2.5:7b"
+UNCENSORED_MODEL = "dolphin-mistral:7b"
 BRAVE_API_KEY = os.environ.get("BRAVE_API_KEY", "")
 BRAVE_SEARCH_URL = "https://api.search.brave.com/res/v1/web/search"
 
@@ -616,19 +617,22 @@ def _raw_sd_prompt(user_message: str) -> str:
 def _refine_sd_prompt(user_message: str, ollama_url: str, messages: List[dict] = None, uncensored: bool = False) -> str:
     """Use the LLM to convert a user request into an optimized SD prompt.
     Includes conversation history so the model understands context like 'generate an image of that'.
-    For uncensored mode, skips LLM entirely to avoid refusal.
+    For uncensored mode, uses dolphin-mistral (no refusal). Falls back to raw prompt on failure.
     """
-    if uncensored:
-        return _raw_sd_prompt(user_message)
-
     try:
+        # Pick model and system prompt based on mode
+        if uncensored:
+            model = UNCENSORED_MODEL
+            sys_key = "image_generation_uncensored"
+        else:
+            model = MODELS["image_generation"]
+            sys_key = "image_generation"
+
         # Build context from recent conversation history
-        sys_key = "image_generation_uncensored" if uncensored else "image_generation"
         context_messages = [{"role": "system", "content": SYSTEM_PROMPTS[sys_key]}]
         if messages:
-            # Include last few exchanges for context (trim to avoid blowing up the context)
             recent = [m for m in messages if m.get("role") in ("user", "assistant") and m.get("content")]
-            for msg in recent[-6:]:  # Last 3 exchanges
+            for msg in recent[-6:]:
                 content = msg["content"]
                 if isinstance(content, list):
                     content = " ".join(p.get("text", "") for p in content if isinstance(p, dict))
@@ -637,7 +641,7 @@ def _refine_sd_prompt(user_message: str, ollama_url: str, messages: List[dict] =
             context_messages.append({"role": "user", "content": user_message[:500]})
 
         payload = {
-            "model": MODELS["image_generation"],
+            "model": model,
             "messages": context_messages,
             "stream": False,
             "options": {"temperature": 0.7, "num_ctx": 4096},
@@ -652,8 +656,8 @@ def _refine_sd_prompt(user_message: str, ollama_url: str, messages: List[dict] =
         return refined
     except Exception as e:
         print(f"[Router] SD prompt refinement failed: {e}")
-        # Fallback: use the user message directly
-        return user_message
+        # Fallback: raw prompt with quality tags
+        return _raw_sd_prompt(user_message)
 
 
 def _negative_prompt() -> str:
@@ -924,6 +928,10 @@ class Pipeline:
         target_model = MODELS.get(category, MODELS["general"])
         system_prompt = SYSTEM_PROMPTS.get(category, SYSTEM_PROMPTS["general"])
 
+        # Override display model for uncensored mode
+        if uncensored:
+            target_model = f"{UNCENSORED_MODEL} → {SD_MODEL_UNCENSORED}"
+
         # Inject language instruction — always respond in the user's language
         if detect_finnish(user_message) and category not in ("reasoning_fi", "image_generation"):
             system_prompt = (