Update main.py

main.py

@@ -60,12 +60,12 @@ except Exception as e:
     raise RuntimeError("Initialization failed. Please check your embeddings or vectorstore setup.")
 
 # ==========================
-# Prompt Template (Context-Only, Strict)
+# Prompt Template (Context-Only)
 # ==========================
 prompt_template = """<s>[INST]
-You are a legal assistant specializing in the Indian Penal Code (IPC). Use only the provided CONTEXT to answer questions. 
+You are a legal assistant specializing in the Indian Penal Code (IPC).  Provide precise, context-specific responses based solely on the given CONTEXT. 
 If the information is not found in the CONTEXT, respond with: "I don't have enough information yet." 
-Do not use any outside knowledge or provide additional explanations.
+
 CONTEXT: {context}
 USER QUERY: {question}
 RESPONSE:
@@ -90,7 +90,7 @@ except Exception as e:
     raise RuntimeError("Something went wrong with the Together API setup. Please verify your API key.")
 
 # ==========================
-# Chat Processing Function (Strict Context-Only)
+# Chat Processing Function
 # ==========================
 def generate_response(user_query: str) -> str:
     try:
@@ -114,10 +114,10 @@ def generate_response(user_query: str) -> str:
         logger.debug(f"Payload sent to LLM: {prompt_input}")
 
         # Generate response using the LLM
-        response = llm(prompt.format(**prompt_input)).strip()
+        response = llm(prompt.format(**prompt_input))
         
-        # Check if response is empty or contains outside knowledge
-        if not response or response.lower() != "i don't have enough information yet.":
+        # Check if response is empty
+        if not response.strip():
             return "I don't have enough information yet."
         
         return response