added guard rails

config/config.yml

@@ -0,0 +1,35 @@
+# models:
+#   - type: main
+#     engine: nvidia_ai_endpoints
+#     model: meta/llama-3.3-70b-instruct
+
+#   - type: content_safety
+#     engine: nvidia_ai_endpoints
+#     model: nvidia/llama-3.1-nemoguard-8b-content-safety
+
+# rails:
+#   input:
+#     flows:
+#       - content safety check input $model=content_safety
+#   output:
+#     flows:
+#       - content safety check output $model=content_safety
+#     streaming:
+#       enabled: True
+#       chunk_size: 200
+#       context_size: 50
+
+# streaming: True
+
+models:
+  - type: main
+    engine: nvidia_ai_endpoints
+    model: meta/llama-3.3-70b-instruct # Or whatever model you use for the GUARDING logic
+
+rails:
+  input:
+    flows:
+      - self check input
+  output:
+    flows:
+      - self check output

config/prompts.yml

@@ -0,0 +1,224 @@
+# # These are the default prompts released by Meta, except for policy O7, which was added to address direct insults.
+# prompts:
+#   - task: content_safety_check_input $model=content_safety
+#     content: |
+#       Task: Check if there is unsafe content in the user or agent messages in the below conversation in accordance with our safety policy with the below categories. Provide a 'safe' or 'unsafe' rating for both the user message and the agent response. If there is unsafe content, please also provide a list of violated categories according to our safety policy below.
+
+#       <BEGIN UNSAFE CONTENT CATEGORIES>
+#       S1: Violence.
+#       S2: Sexual.
+#       S3: Criminal Planning/Confessions.
+#       S4: Guns and Illegal Weapons.
+#       S5: Controlled/Regulated Substances.
+#       S6: Suicide and Self Harm.
+#       S7: Sexual (minor).
+#       S8: Hate/Identity Hate.
+#       S9: PII/Privacy.
+#       S10: Harassment.
+#       S11: Threat.
+#       S12: Profanity.
+#       S13: Needs Caution.
+#       S14: Other.
+#       S15: Manipulation.
+#       S16: Fraud/Deception.
+#       S17: Malware.
+#       S18: High Risk Gov Decision Making.
+#       S19: Political/Misinformation/Conspiracy.
+#       S20: Copyright/Trademark/Plagiarism.
+#       S21: Unauthorized Advice.
+#       S22: Illegal Activity.
+#       S23: Immoral/Unethical.
+#       <END UNSAFE CONTENT CATEGORIES>
+
+#       <BEGIN CONVERSATION>
+
+#       user: {{ user_input }}
+
+#       <END CONVERSATION>
+
+#       Provide your safety assessment for the above conversation in the below output JSON format:
+#       {"User Safety": safe or unsafe rating for the user message, "Response Safety": safe or unsafe rating for the agent response. Omit if no agent response present. "Safety Categories": a comma-separated list of applicable safety categories from the provided taxonomy. Omit if all safe.}
+
+#       Do not include anything other than the output JSON in your response.
+#       Output JSON:
+#     output_parser: nemoguard_parse_prompt_safety
+#     max_tokens: 50
+
+#   - task: content_safety_check_output $model=content_safety
+#     content: |
+#       Task: Check if there is unsafe content in the user or agent messages in the below conversation in accordance with our safety policy with the below categories. Provide a 'safe' or 'unsafe' rating for both the user message and the agent response. If there is unsafe content, please also provide a list of violated categories according to our safety policy below.
+
+#       <BEGIN UNSAFE CONTENT CATEGORIES>
+#       S1: Violence.
+#       S2: Sexual.
+#       S3: Criminal Planning/Confessions.
+#       S4: Guns and Illegal Weapons.
+#       S5: Controlled/Regulated Substances.
+#       S6: Suicide and Self Harm.
+#       S7: Sexual (minor).
+#       S8: Hate/Identity Hate.
+#       S9: PII/Privacy.
+#       S10: Harassment.
+#       S11: Threat.
+#       S12: Profanity.
+#       S13: Needs Caution.
+#       S14: Other.
+#       S15: Manipulation.
+#       S16: Fraud/Deception.
+#       S17: Malware.
+#       S18: High Risk Gov Decision Making.
+#       S19: Political/Misinformation/Conspiracy.
+#       S20: Copyright/Trademark/Plagiarism.
+#       S21: Unauthorized Advice.
+#       S22: Illegal Activity.
+#       S23: Immoral/Unethical.
+#       <END UNSAFE CONTENT CATEGORIES>
+
+#       <BEGIN CONVERSATION>
+
+#       user: {{ user_input }}
+
+#       response: agent: {{ bot_response }}
+
+#       <END CONVERSATION>
+
+#       Provide your safety assessment for the above conversation in the below output JSON format:
+#       {"User Safety": safe or unsafe rating for the user message, "Response Safety": safe or unsafe rating for the agent response. Omit if no agent response present. "Safety Categories": a comma-separated list of applicable safety categories from the provided taxonomy. Omit if all safe.}
+
+#       Do not include anything other than the output JSON in your response.
+#       Output JSON:
+#     output_parser: nemoguard_parse_response_safety
+#     max_tokens: 50
+
+prompts:
+  - task: self_check_input
+    content: |
+      Your task is to check if the user message below is unsafe.
+      The message is unsafe if it contains:
+      - Requests for help with illegal acts.
+      - Hate speech, harassment, or violence.
+      - Instructions on how to make weapons, bombs, or drugs.
+      - Violence
+      - Sexual content
+      - Criminal Planning or Confessions
+      - Guns and Illegal Weapons
+      - Controlled or Regulated Substances
+      - Suicide and Self Harm
+      - Sexual content involving minors
+      - Hate speech or Identity Hate
+      - PII or Privacy violations
+      - Harassment
+      - Threats
+      - Profanity
+      - Topics needing caution
+      - Manipulation
+      - Fraud or Deception
+      - Malware creation or usage
+      - High Risk Government Decision Making
+      - Political Misinformation or Conspiracy
+      - Copyright, Trademark, or Plagiarism
+      - Unauthorized Professional Advice (Medical/Legal)
+      - Illegal Activity
+      - Immoral or Unethical content
+      
+      User message: "{{ user_input }}"
+      
+      Answer "safe" or "unsafe".
+
+  - task: self_check_output
+    content: |
+      Your task is to check if the bot response below is unsafe.
+      The response is unsafe if it contains:
+      - Instructions on how to make weapons, bombs, or drugs.
+      - Hate speech or illegal content.
+      - Violence
+      - Sexual content
+      - Criminal Planning or Confessions
+      - Guns and Illegal Weapons
+      - Controlled or Regulated Substances
+      - Suicide and Self Harm
+      - Sexual content involving minors
+      - Hate speech or Identity Hate
+      - PII or Privacy violations
+      - Harassment
+      - Threats
+      - Profanity
+      - Topics needing caution
+      - Manipulation
+      - Fraud or Deception
+      - Malware creation or usage
+      - High Risk Government Decision Making
+      - Political Misinformation or Conspiracy
+      - Copyright, Trademark, or Plagiarism
+      - Unauthorized Professional Advice (Medical/Legal)
+      - Illegal Activity
+      - Immoral or Unethical content
+      Bot response: "{{ bot_response }}"
+      
+      Answer "safe" or "unsafe".
+
+
+# prompts:
+#   - task: self_check_input
+#     content: |
+#       Task: Check if there is unsafe content in the user message below according to our safety policy.
+      
+#       <BEGIN UNSAFE CONTENT CATEGORIES>
+#       S1: Violence.
+#       S2: Sexual.
+#       S3: Criminal Planning/Confessions.
+#       S4: Guns and Illegal Weapons.
+#       S5: Controlled/Regulated Substances.
+#       S6: Suicide and Self Harm.
+#       S7: Sexual (minor).
+#       S8: Hate/Identity Hate.
+#       S9: PII/Privacy.
+#       S10: Harassment.
+#       S11: Threat.
+#       S12: Profanity.
+#       S13: Needs Caution.
+#       S14: Other.
+#       S15: Manipulation.
+#       S16: Fraud/Deception.
+#       S17: Malware.
+#       S18: High Risk Gov Decision Making.
+#       S19: Political/Misinformation/Conspiracy.
+#       S20: Copyright/Trademark/Plagiarism.
+#       S21: Unauthorized Advice.
+#       S22: Illegal Activity.
+#       S23: Immoral/Unethical.
+#       <END UNSAFE CONTENT CATEGORIES>
+
+#       <BEGIN CONVERSATION>
+#       user: {{ user_input }}
+#       <END CONVERSATION>
+
+#       Provide your safety assessment in the below output JSON format:
+#       {"User Safety": "safe" or "unsafe", "Safety Categories": "comma-separated list of violated categories (e.g. S1, S4)"}
+      
+#       Do not include anything other than the output JSON in your response.
+
+#   - task: self_check_output
+#     content: |
+#       Task: Check if there is unsafe content in the agent response below.
+      
+#       <BEGIN UNSAFE CONTENT CATEGORIES>
+#       S1: Violence.
+#       S2: Sexual.
+#       S3: Criminal Planning.
+#       S4: Guns and Illegal Weapons.
+#       S5: Controlled Substances.
+#       S6: Suicide and Self Harm.
+#       S8: Hate Speech.
+#       S22: Illegal Activity.
+#       <END UNSAFE CONTENT CATEGORIES>
+
+#       <BEGIN CONVERSATION>
+#       user: {{ user_input }}
+#       agent: {{ bot_response }}
+#       <END CONVERSATION>
+
+#       Provide your safety assessment in the below output JSON format:
+#       {"Response Safety": "safe" or "unsafe", "Safety Categories": "comma-separated list"}
+
+#       Do not include anything other than the output JSON.

config/rails.co

@@ -0,0 +1,20 @@
+# 1. Define the Input Flow
+define flow self check input
+  # This runs the prompt you wrote in prompts.yml
+  $allowed = execute self_check_input
+  
+  # If the model said "unsafe", $allowed will be false
+  if not $allowed
+    bot refuse to respond
+    stop
+
+# 2. Define the Output Flow
+define flow self check output
+  $allowed = execute self_check_output
+  if not $allowed
+    bot refuse to respond
+    stop
+
+# 3. Define the "Refusal" Message (The Magic Token)
+define bot refuse to respond
+  "[[GUARDRAIL_BLOCK_TRIGGERED]]"

input_to_llm.py

@@ -8,9 +8,7 @@ def extract_chats(collection,user_id,limit):
     if collection is None:
         return "No database connection."
     
-    cursor = collection.find(
-        {"user_id": user_id}
-    ).sort("timestamp", -1).limit(limit)
+    cursor = collection.find({"user_id": user_id}).sort("timestamp", -1).limit(limit)
 
     recent_chats = list(cursor)[::-1]
 

main.py

@@ -403,19 +403,23 @@ from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
 from typing import List, Optional, Dict, Any
 from datetime import datetime
+import asyncio
+import os
+from nemoguardrails import LLMRails, RailsConfig
 
 # from langchain.memory import ConversationBufferMemory
 from perplexity import Perplexity
 from groq import Groq
 from input_to_llm import extract_chats, extract_goalfocus
 from utils import (
-    load_user_data,
-    initialize_rag,
-    get_rag_response,
+    # load_user_data,
+    # initialize_rag,
+    # get_rag_response,
     llm,
     get_mongo_collection
 )
 
+os.environ["NVIDIA_API_KEY"] = "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"
 from configure import USER_DATA_PATH, llm_prompt
 
 app = FastAPI(title="Sattva AI API")
@@ -460,7 +464,50 @@ async def chat_endpoint(request: ChatRequest):
         username = request.user_id
         user_input = request.message
         mode = request.mode
-        
+        config = RailsConfig.from_path("./config")
+        rails = LLMRails(config)
+        input_check = await rails.generate_async(
+            messages=[{"role": "user", "content": user_input}],
+            options={
+                "rails": ["input"],
+                "log": {"level": "INFO"} 
+            }
+        )
+        response_text = str(input_check.response)
+        if "[[GUARDRAIL_BLOCK_TRIGGERED]]" in response_text:
+            print("[GUARD] Input Blocked!")
+
+            return ChatResponse(
+                    response="I am sorry but I cannot answer that request.",
+                    topic="blocked",
+                    goal="blocked",
+                    resources=[]
+            )
+    
+    #     refusal_keywords = [
+    #         "unsafe", 
+    #         "cannot", 
+    #         "can't", 
+    #         "sorry", 
+    #         "apologize", 
+    #         "refuse", 
+    #         "unable to answer",
+    #         "regulated",
+    #         "harmful"
+    #     ]
+    
+    # # If ANY keyword is found in the response, we block it.
+    #     is_blocked = any(keyword in response_text for keyword in refusal_keywords)
+    #     if is_blocked:
+    #         # print("   [GUARD] Input Blocked!")
+    #         # return {"status": "blocked", "message": "I cannot answer that request."}
+    #         return ChatResponse(
+    #             response="I am sorry but I cannot answer that request.",
+    #             topic="blocked",
+    #             goal="blocked",
+    #             resources=[]
+    #     )
+    #         # pass
      
         ncon = 3
         chat_history_str = extract_chats(collection, username, ncon) 
@@ -535,6 +582,23 @@ async def chat_endpoint(request: ChatRequest):
         
         response = llm.invoke(context_response)
         ai_response_text = response.content
+        input_check = await rails.generate_async(
+            messages=[{"role": "user", "content": ai_response_text}],
+            options={
+                "rails": ["input"],
+                "log": {"level": "INFO"} 
+            }
+        )
+        response_text = str(input_check.response)
+        if "[[GUARDRAIL_BLOCK_TRIGGERED]]" in response_text:
+            print("[GUARD] Input Blocked!")
+
+            return ChatResponse(
+                    response="I am sorry but I cannot answer that request.",
+                    topic="blocked",
+                    goal="blocked",
+                    resources=[]
+            )
 
         if collection is not None:
             chat_document = {

requirements.txt

@@ -1,60 +1,72 @@
 aiohappyeyeballs==2.6.1
-aiohttp==3.13.0
+aiohttp==3.13.2
 aiosignal==1.4.0
+annotated-doc==0.0.4
 annotated-types==0.7.0
-anyio==4.11.0
-async-timeout==4.0.3
+annoy==1.17.3
+anyio==4.12.0
+async-timeout==5.0.1
 attrs==25.4.0
 backoff==2.2.1
 bcrypt==5.0.0
 build==1.3.0
-cachetools==6.2.1
-certifi==2025.10.5
+cachetools==6.2.3
+certifi==2025.11.12
 charset-normalizer==3.4.4
-chromadb==1.1.1
-click==8.3.0
+chromadb==1.3.7
+click==8.3.1
 coloredlogs==15.0.1
 dataclasses-json==0.6.7
 distro==1.9.0
 dnspython==2.8.0
 durationpy==0.10
-exceptiongroup==1.3.0
-fastapi==0.119.0
+exceptiongroup==1.3.1
+fastapi==0.124.4
+fastembed==0.6.0
 filelock==3.20.0
+filetype==1.2.0
 flatbuffers==25.9.23
 frozenlist==1.8.0
-fsspec==2025.9.0
-google-auth==2.41.1
-googleapis-common-protos==1.70.0
-greenlet==3.2.4
-groq==0.32.0
-grpcio==1.75.1
+fsspec==2025.12.0
+google-auth==2.43.0
+googleapis-common-protos==1.72.0
+greenlet==3.3.0
+groq==0.37.1
+grpcio==1.76.0
 gunicorn==23.0.0
 h11==0.16.0
-hf-xet==1.1.10
+hf-xet==1.2.0
 httpcore==1.0.9
 httptools==0.7.1
 httpx==0.28.1
 httpx-sse==0.4.3
-huggingface-hub==0.35.3
+huggingface-hub==0.36.0
 humanfriendly==10.0
 idna==3.11
 importlib_metadata==8.7.0
 importlib_resources==6.5.2
 Jinja2==3.1.6
-joblib==1.5.2
+joblib==1.5.3
 jsonpatch==1.33
 jsonpointer==3.0.0
 jsonschema==4.25.1
 jsonschema-specifications==2025.9.1
 kubernetes==34.1.0
-langchain==0.3.27
-langchain-community==0.3.31
-langchain-core==0.3.79
-langchain-groq==0.3.8
-langchain-huggingface==0.3.1
-langchain-text-splitters==0.3.11
-langsmith==0.4.37
+langchain==1.1.3
+langchain-classic==1.0.0
+langchain-community==0.4.1
+langchain-core==1.2.0
+langchain-groq==1.1.1
+langchain-huggingface==1.2.0
+langchain-nvidia-ai-endpoints==1.0.0
+langchain-text-splitters==1.1.0
+langgraph==1.0.5
+langgraph-checkpoint==3.0.1
+langgraph-prebuilt==1.0.5
+langgraph-sdk==0.3.0
+langsmith==0.4.59
+lark==1.3.1
+loguru==0.7.3
 markdown-it-py==4.0.0
 MarkupSafe==3.0.3
 marshmallow==3.26.1
@@ -63,74 +75,86 @@ mmh3==5.2.0
 mpmath==1.3.0
 multidict==6.7.0
 mypy_extensions==1.1.0
-networkx==3.4.2
-numpy==2.2.6
+nemoguardrails==0.19.0
+nest-asyncio==1.6.0
+networkx==3.6.1
+numpy==2.3.5
 oauthlib==3.3.1
-onnxruntime==1.23.1
-opentelemetry-api==1.38.0
-opentelemetry-exporter-otlp-proto-common==1.38.0
-opentelemetry-exporter-otlp-proto-grpc==1.38.0
-opentelemetry-proto==1.38.0
-opentelemetry-sdk==1.38.0
-opentelemetry-semantic-conventions==0.59b0
-orjson==3.11.3
+onnxruntime==1.23.2
+opentelemetry-api==1.39.1
+opentelemetry-exporter-otlp-proto-common==1.39.1
+opentelemetry-exporter-otlp-proto-grpc==1.39.1
+opentelemetry-proto==1.39.1
+opentelemetry-sdk==1.39.1
+opentelemetry-semantic-conventions==0.60b1
+orjson==3.11.5
+ormsgpack==1.12.1
 overrides==7.7.0
 packaging==25.0
+pandas==2.3.3
 perplexity-client==0.1.1
-perplexityai==0.16.1
-pillow==12.0.0
+perplexityai==0.22.0
+pillow==11.3.0
 posthog==5.4.0
+prompt_toolkit==3.0.52
 propcache==0.4.1
-protobuf==6.33.0
+protobuf==6.33.2
+py_rust_stemmers==0.1.5
 pyasn1==0.6.1
 pyasn1_modules==0.4.2
-pybase64==1.4.2
-pydantic==2.12.2
-pydantic-settings==2.11.0
-pydantic_core==2.41.4
+pybase64==1.4.3
+pydantic==2.12.5
+pydantic-settings==2.12.0
+pydantic_core==2.41.5
 Pygments==2.19.2
-pymongo==4.15.4
+pymongo==4.15.5
 PyPika==0.48.9
 pyproject_hooks==1.2.0
 python-dateutil==2.9.0.post0
-python-dotenv==1.1.1
+python-dotenv==1.2.1
+pytz==2025.2
 PyYAML==6.0.3
 referencing==0.37.0
-regex==2025.9.18
+regex==2025.11.3
 requests==2.32.5
 requests-oauthlib==2.0.0
 requests-toolbelt==1.0.0
 rich==14.2.0
-rpds-py==0.27.1
+rpds-py==0.30.0
 rsa==4.9.1
-safetensors==0.6.2
-scikit-learn==1.7.2
-scipy==1.15.3
-sentence-transformers==5.1.1
+safetensors==0.7.0
+scikit-learn==1.8.0
+scipy==1.16.3
+sentence-transformers==5.2.0
 setuptools==80.9.0
 shellingham==1.5.4
+simpleeval==1.0.3
 six==1.17.0
 sniffio==1.3.1
-SQLAlchemy==2.0.44
-starlette==0.48.0
+SQLAlchemy==2.0.45
+starlette==0.50.0
 sympy==1.14.0
 tenacity==9.1.2
 threadpoolctl==3.6.0
 tokenizers==0.22.1
 tomli==2.3.0
-torch==2.9.0
 tqdm==4.67.1
-transformers==4.57.1
-typer==0.19.2
+transformers==4.57.3
+typer==0.20.0
 typing-inspect==0.9.0
 typing-inspection==0.4.2
 typing_extensions==4.15.0
+tzdata==2025.3
 urllib3==2.3.0
-uvicorn==0.37.0
-uvloop==0.21.0
+uuid_utils==0.12.0
+uvicorn==0.38.0
+uvloop==0.22.1
+watchdog==6.0.0
 watchfiles==1.1.1
+wcwidth==0.2.14
 websocket-client==1.9.0
 websockets==15.0.1
+xxhash==3.6.0
 yarl==1.22.0
 zipp==3.23.0
 zstandard==0.25.0

utils.py

@@ -2,12 +2,12 @@
 import os
 import json
 from langchain_groq import ChatGroq
-from langchain_community.vectorstores import Chroma
-from langchain.schema import Document
-from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain_huggingface import HuggingFaceEmbeddings
-from langchain.chains import RetrievalQA
-from langchain.prompts import PromptTemplate
+from langchain_text_splitters import RecursiveCharacterTextSplitter
+# from langchain.schema import Document
+# from langchain.chains import RetrievalQA
+# from langchain_huggingface import HuggingFaceEmbeddings
+# from langchain_community.vectorstores import Chroma
+# from langchain.prompts import PromptTemplate
 from configure import USER_DATA_PATH, RAG_BASE_DIRECTORY, RAG_CATEGORIES
 import shutil
 from dotenv import load_dotenv
@@ -75,114 +75,114 @@ def save_user_data(data):
         json.dump(data, f, indent=2)
 
 # Initialize RAG with per-category vector stores and QA chains
-def initialize_rag():
-    vector_stores = {}
-    qa_chains = {}
+# def initialize_rag():
+#     vector_stores = {}
+#     qa_chains = {}
     
-    # Updated prompt template for advice-focused responses
-    custom_prompt_template = """
-    You are a {category} wellness expert. Provide helpful advice with specific actions:
+#     # Updated prompt template for advice-focused responses
+#     custom_prompt_template = """
+#     You are a {category} wellness expert. Provide helpful advice with specific actions:
     
-    1. Start with a brief empathetic response to the user's concern
-    2. Offer 1-3 actionable suggestions with brief explanations
-    3. End with an open-ended question to continue conversation
+#     1. Start with a brief empathetic response to the user's concern
+#     2. Offer 1-3 actionable suggestions with brief explanations
+#     3. End with an open-ended question to continue conversation
     
-    Guidelines:
-    - Keep responses conversational and supportive
-    - Avoid clinical jargon
-    - Focus on practical, implementable advice
-    - Maintain hopeful and encouraging tone
+#     Guidelines:
+#     - Keep responses conversational and supportive
+#     - Avoid clinical jargon
+#     - Focus on practical, implementable advice
+#     - Maintain hopeful and encouraging tone
     
-    Context:
-    {context}
+#     Context:
+#     {context}
     
-    Question: {question}
-    """
+#     Question: {question}
+#     """
     
-    # Create vector stores and QA chains for each category
-    # Assume all your variables like RAG_CATEGORIES, embeddings, llm, etc.
-# are defined above this block.
-
-    for category in RAG_CATEGORIES:
-        persist_dir = f"./chroma_db_{category}"
-        vector_store = None  # Initialize vector_store to None
-
-        # --- 1. Check if Vector Store Already Exists ---
-        if os.path.exists(persist_dir):
-            print(f"Found existing vector store for {category}. Attempting to load...")
-            try:
-                # Load the existing vector store from disk
-                vector_store = Chroma(
-                    persist_directory=persist_dir,
-                    embedding_function=embeddings  # Must provide the same embedding function
-                )
-                vector_stores[category] = vector_store
-            except Exception as e:
-                print(f"Error loading existing store {persist_dir}: {e}")
-                print("Will delete and attempt to re-build.")
-                shutil.rmtree(persist_dir)
+#     # Create vector stores and QA chains for each category
+#     # Assume all your variables like RAG_CATEGORIES, embeddings, llm, etc.
+# # are defined above this block.
+
+#     for category in RAG_CATEGORIES:
+#         persist_dir = f"./chroma_db_{category}"
+#         vector_store = None  # Initialize vector_store to None
+
+#         # --- 1. Check if Vector Store Already Exists ---
+#         if os.path.exists(persist_dir):
+#             print(f"Found existing vector store for {category}. Attempting to load...")
+#             try:
+#                 # Load the existing vector store from disk
+#                 vector_store = Chroma(
+#                     persist_directory=persist_dir,
+#                     embedding_function=embeddings  # Must provide the same embedding function
+#                 )
+#                 vector_stores[category] = vector_store
+#             except Exception as e:
+#                 print(f"Error loading existing store {persist_dir}: {e}")
+#                 print("Will delete and attempt to re-build.")
+#                 shutil.rmtree(persist_dir)
         
-        # --- 2. Create Vector Store if it Doesn't Exist (or failed to load) ---
-        if vector_store is None: 
-            print(f"No valid vector store for {category} found. Creating new one...")
+#         # --- 2. Create Vector Store if it Doesn't Exist (or failed to load) ---
+#         if vector_store is None: 
+#             print(f"No valid vector store for {category} found. Creating new one...")
             
-            # Process documents (Your original logic)
-            dir_path = os.path.join(RAG_BASE_DIRECTORY, category)
-            docs = []
-            if os.path.exists(dir_path):
-                for filename in os.listdir(dir_path):
-                    if filename.endswith('.txt'):
-                        file_path = os.path.join(dir_path, filename)
-                        try:
-                            with open(file_path, 'r', encoding='utf-8') as f:
-                                text = f.read()
-                            chunks = text_splitter.split_text(text)
-                            for chunk in chunks:
-                                if chunk.strip():
-                                    metadata = {
-                                        "source": filename,
-                                        "category": category
-                                    }
-                                    docs.append(Document(
-                                        page_content=chunk.strip(),
-                                        metadata=metadata
-                                    ))
-                        except Exception as e:
-                            print(f"Error processing {file_path}: {e}")
+#             # Process documents (Your original logic)
+#             dir_path = os.path.join(RAG_BASE_DIRECTORY, category)
+#             docs = []
+#             if os.path.exists(dir_path):
+#                 for filename in os.listdir(dir_path):
+#                     if filename.endswith('.txt'):
+#                         file_path = os.path.join(dir_path, filename)
+#                         try:
+#                             with open(file_path, 'r', encoding='utf-8') as f:
+#                                 text = f.read()
+#                             chunks = text_splitter.split_text(text)
+#                             for chunk in chunks:
+#                                 if chunk.strip():
+#                                     metadata = {
+#                                         "source": filename,
+#                                         "category": category
+#                                     }
+#                                     docs.append(Document(
+#                                         page_content=chunk.strip(),
+#                                         metadata=metadata
+#                                     ))
+#                         except Exception as e:
+#                             print(f"Error processing {file_path}: {e}")
             
-            # Create vector store (Your original logic)
-            if docs:
-                vector_store = Chroma.from_documents(
-                    documents=docs,
-                    embedding=embeddings,
-                    persist_directory=persist_dir
-                )
-                vector_store.persist()
-                vector_stores[category] = vector_store
-                print(f"Created and persisted new vector store for {category} with {len(docs)} documents.")
-            else:
-                print(f"No documents found for {category}. Skipping QA chain setup.")
-                continue  # Skip to the next category
-
-        # --- 3. Create QA Chain (if vector_store was loaded or created) ---
-        if vector_store:
-            # Create QA chain with updated prompt
-            prompt = PromptTemplate(
-                template=custom_prompt_template.replace("{category}", category),
-                input_variables=["context", "question"]
-            )
-            qa_chain = RetrievalQA.from_chain_type(
-                llm=llm,
-                chain_type="stuff",
-                retriever=vector_store.as_retriever(search_kwargs={"k": 5}),
-                chain_type_kwargs={"prompt": prompt},
-                return_source_documents=True
-            )
-            qa_chains[category] = qa_chain
-            print(f"Initialized {category} QA chain.")
-
-# The rest of your code
-    return qa_chains
+#             # Create vector store (Your original logic)
+#             if docs:
+#                 vector_store = Chroma.from_documents(
+#                     documents=docs,
+#                     embedding=embeddings,
+#                     persist_directory=persist_dir
+#                 )
+#                 vector_store.persist()
+#                 vector_stores[category] = vector_store
+#                 print(f"Created and persisted new vector store for {category} with {len(docs)} documents.")
+#             else:
+#                 print(f"No documents found for {category}. Skipping QA chain setup.")
+#                 continue  # Skip to the next category
+
+#         # --- 3. Create QA Chain (if vector_store was loaded or created) ---
+#         if vector_store:
+#             # Create QA chain with updated prompt
+#             prompt = PromptTemplate(
+#                 template=custom_prompt_template.replace("{category}", category),
+#                 input_variables=["context", "question"]
+#             )
+#             qa_chain = RetrievalQA.from_chain_type(
+#                 llm=llm,
+#                 chain_type="stuff",
+#                 retriever=vector_store.as_retriever(search_kwargs={"k": 5}),
+#                 chain_type_kwargs={"prompt": prompt},
+#                 return_source_documents=True
+#             )
+#             qa_chains[category] = qa_chain
+#             print(f"Initialized {category} QA chain.")
+
+# # The rest of your code
+#     return qa_chains
 
 # Classify question to category
 def classify_question_category(question):
@@ -195,18 +195,18 @@ def classify_question_category(question):
     print(response)
     return response.content.strip()
 
-# Get RAG response using category-specific QA chain
-def get_rag_response(question, qa_chains):
-    # Classify question
-    category = classify_question_category(question)
+# # Get RAG response using category-specific QA chain
+# def get_rag_response(question, qa_chains):
+#     # Classify question
+#     category = classify_question_category(question)
     
-    if category not in qa_chains:
-        # Fallback to first available chain
-        category = list(qa_chains.keys())[0]
+#     if category not in qa_chains:
+#         # Fallback to first available chain
+#         category = list(qa_chains.keys())[0]
     
-    # Get response
-    result = qa_chains[category].invoke({"query": question})
-    return result['result']
+#     # Get response
+#     result = qa_chains[category].invoke({"query": question})
+#     return result['result']
 
 
 # Classify user input