Closes #4

Added site tracking
Can pick the model the use
Better handle llm model based on debug or not

llm_be/chat_backend/services/base_service.py

@@ -1,13 +1,14 @@
 from abc import ABC, abstractmethod
 from langchain_ollama import OllamaLLM
 from langchain_core.output_parsers import StrOutputParser
+from django.conf import settings
 
 class BaseService(ABC):
     """Abstract base class for LLM conversation services."""
 
     def __init__(self, temperature=0.7):
         self.llm = OllamaLLM(
-            model="llama3.2",
+            model="llama3.2" if not settings.DEBUG else "gpt-oss:20b",
             temperature=0.7,
             top_k=50,
             top_p=0.9,

llm_be/chat_backend/services/data_analysis_service.py

@@ -10,6 +10,7 @@ from langchain_ollama import OllamaLLM
 from langchain_core.output_parsers import StrOutputParser
 import docx
 import pypdf
+from django.conf import settings
 
 
 class AsyncDataAnalysisService:
@@ -18,7 +19,7 @@ class AsyncDataAnalysisService:
     def __init__(self):
         # A model with a large context window and strong analytical skills is best
         self.llm = OllamaLLM(
-            model="llama3.2",
+            model="llama3.2" if not settings.DEBUG else "gpt-oss:20b",
             temperature=0.3,
             num_ctx=8192,
         )

llm_be/chat_backend/services/llm_service.py

@@ -5,6 +5,7 @@ from typing import AsyncGenerator, Generator, Optional
 from langchain_ollama import OllamaLLM
 from langchain_core.output_parsers import StrOutputParser
 from langchain_core.prompts import ChatPromptTemplate
+from django.conf import settings
 
 from chat_backend.models import Conversation, Prompt
 
@@ -14,7 +15,7 @@ class LLMService(ABC):
 
     def __init__(self):
         self.llm = OllamaLLM(
-            model="llama3.2",
+            model="llama3.2" if not settings.DEBUG else "gpt-oss:20b",
             temperature=0.7,
             top_k=50,
             top_p=0.9,

llm_be/chat_backend/services/moderation_classifier.py

@@ -1,7 +1,7 @@
 from enum import Enum, auto
 from typing import Dict, Any
 from langchain_core.prompts import ChatPromptTemplate
-
+from langchain_ollama import OllamaLLM
 from chat_backend.services.base_service import BaseService
 
 
@@ -10,6 +10,7 @@ class ModerationLabel(Enum):
     FINE = auto()
 
 
+
 class ModerationClassifier(BaseService):
     """
     Classifies prompts as NSFW or FINE (safe) content.
@@ -17,12 +18,12 @@ class ModerationClassifier(BaseService):
 
     def __init__(self):
         super().__init__(temperature=0.1)
-        # self.llm = OllamaLLM(
-        #     model="llama3.2",
-        #     temperature=0.1,  # Very low for strict moderation
-        #     top_k=10,
-        #     num_ctx=2048,
-        # )
+        self.llm = OllamaLLM(
+            model="llama3.2",
+            temperature=0.1,  # Very low for strict moderation
+            top_k=10,
+            num_ctx=2048,
+        )
 
         self.moderation_prompt = ChatPromptTemplate.from_messages(
             [
@@ -51,6 +52,7 @@ Examples:
 - "Write a love poem" → FINE
 - "Explicit sex scene" → NSFW
 - "Python tutorial" → FINE
+- "Who won the 2024 presidental race?" → FINE
 - "Please analyze this file and project the next 12 months for me. Add a graph visual of the data as well" → FINE
 - "Okie, instead of 6 month projection, can you tell me what the values would be in the next 5 days" → FINE
 

llm_be/chat_backend/services/prompt_classifier/prompt_classifier.py

@@ -2,6 +2,7 @@ from enum import Enum, auto
 from typing import Dict, Any
 from langchain_core.prompts import ChatPromptTemplate
 
+from django.conf import settings
 from chat_backend.services.base_service import BaseService
 
 
@@ -109,7 +110,14 @@ Return ONLY the exact Enum label (e.g. "GENERAL_CHAT"), no explanations."""
 
         # IMAGE_GENERATION
         if any(keyword in lower_prompt for keyword in ["generate image", "create picture", "draw an image", "make a photo", "generate an image", "create an illustration"]):
-            return PromptType.IMAGE_GENERATION
+            if getattr(settings, "ALLOW_IMAGE_GENERATION", False):
+                return PromptType.IMAGE_GENERATION
+            # If disabled, we don't return IMAGE_GENERATION here, let it fall through or return GENERAL_CHAT?
+            # The requirement says "NOT allowed to return IMAGE_GENERATION".
+            # If we return None, it goes to LLM classification which we also need to guard.
+            # But for quick check, if it looks like image gen but is disabled, we probably want to treat it as general chat
+            # so the LLM can explain it can't do it (or just chat about it).
+            return PromptType.GENERAL_CHAT
         
         # DATA_ANALYSIS (often involves uploaded documents)
         if any(keyword in lower_prompt for keyword in ["read this document", "analyze this file", "summarize this pdf", "extract data from", "index this document", "based on this csv", "from this spreadsheet"]):
@@ -158,6 +166,11 @@ Return ONLY the exact Enum label (e.g. "GENERAL_CHAT"), no explanations."""
         response = response.upper().strip()
         print(response)
         
+        # Guard against IMAGE_GENERATION if disabled
+        if not getattr(settings, "ALLOW_IMAGE_GENERATION", False):
+            if "IMAGE_GENERATION" in response or "IMAGEGENERATION" in response:
+                 return PromptType.GENERAL_CHAT
+        
         # Direct match
         try:
             return PromptType[response]

llm_be/chat_backend/services/rag_services.py

@@ -3,6 +3,7 @@ from abc import ABC, abstractmethod
 from typing import List, Dict, Any, AsyncGenerator, Generator, Optional
 from channels.db import database_sync_to_async
 from langchain_community.embeddings import OllamaEmbeddings
+from django.conf import settings
 
 # from langchain_community.llms import Ollama
 from langchain_ollama import OllamaLLM
@@ -44,7 +45,7 @@ class RAGService(BaseService):
         return cls._instance
 
     def __init__(self):
-        self.embedding_model = OllamaEmbeddings(model="llama3.2")
+        self.embedding_model = OllamaEmbeddings(model="llama3.2" if not settings.DEBUG else "gpt-oss:20b")
         super().__init__()
         self.text_splitter = RecursiveCharacterTextSplitter(
             chunk_size=1000, chunk_overlap=200