Update utils/llm.py

utils/llm.py

@@ -1,49 +1,65 @@
-from langchain_community.embeddings.fastembed import FastEmbedEmbeddings
-from langchain_groq import ChatGroq
-import os
-import json
-from typing import List, Dict
-
-class LLMProcessor:
-    def __init__(self):
-        """Initialize embedding model and Groq LLM"""
-        self.api_key = os.getenv("GROQ_API_KEY")
-
-        # Use FastEmbed instead of SentenceTransformer
-        self.embed_model = FastEmbedEmbeddings()
-
-        self.llm = ChatGroq(
-            model_name="mixtral-8x7b-32768",
-            api_key=self.api_key
-        )
-
-    def format_context(self, chunks: List[Dict]) -> str:
-        """Format retrieved chunks into a structured context for the LLM"""
-        context_parts = []
-        for chunk in chunks:
-            try:
-                headings = json.loads(chunk['headings'])
-                if headings:
-                    context_parts.append(f"Section: {' > '.join(headings)}")
-            except:
-                pass
-
-            if chunk['page']:
-                context_parts.append(f"Page {chunk['page']}:")
-            
-            context_parts.append(chunk['text'])
-            context_parts.append("-" * 40)
-
-        return "\n".join(context_parts)
-
-    def generate_answer(self, context: str, question: str) -> str:
-        """Generate answer using structured context"""
-        prompt = f"""Based on the following excerpts from a document:
-
-{context}
-
-Please answer this question: {question}
-
-Make use of the section information and page numbers in your answer when relevant.
-"""
-        return self.llm.invoke(prompt)
+from langchain_community.embeddings.fastembed import FastEmbedEmbeddings
+from langchain_groq import ChatGroq
+import os
+import json
+from typing import List, Dict
+
+class LLMProcessor:
+    def __init__(self):
+        """Initialize embedding model and Groq LLM"""
+        self.api_key = os.getenv("GROQ_API_KEY")
+
+        # Use FastEmbed instead of SentenceTransformer
+        self.embed_model = FastEmbedEmbeddings()
+
+        self.llm = ChatGroq(
+            model_name="mixtral-8x7b-32768",
+            api_key=self.api_key
+        )
+
+    def format_context(self, chunks: List[Dict]) -> str:
+        """Format retrieved chunks into a structured context for the LLM"""
+        context_parts = []
+        for chunk in chunks:
+            try:
+                headings = json.loads(chunk['headings'])
+                if headings:
+                    context_parts.append(f"Section: {' > '.join(headings)}")
+            except:
+                pass
+
+            if chunk['page']:
+                context_parts.append(f"Page {chunk['page']}:")
+            
+            context_parts.append(chunk['text'])
+            context_parts.append("-" * 40)
+
+        return "\n".join(context_parts)
+
+    def generate_answer(self, context: str, question: str) -> str:
+        """Generate answer using structured context"""
+        prompt = f"""
+You are an AI assistant tasked with answering user questions based on the given document excerpts. Your goal is to provide a clear, accurate, and helpful answer using only the provided context. 
+
+If the answer is not found in the context, explicitly state that you do not know instead of making up an answer. If the question is out of context, say that it is out of context, but still try to provide the best possible response from the available information.
+
+---
+### Context:  
+{context}
+
+### User Question:  
+{question}
+
+---
+#### Instructions:
+- Use only the given context to construct your answer.
+- Reference relevant sections and page numbers where applicable.
+- Be concise yet informative, focusing on clarity and usefulness.
+- If uncertain, respond honestly (e.g., "The answer is not found in the provided context.").
+- If out of context, state so clearly (e.g., "The question is out of context, but here’s what I found in the document...").
+
+---
+### Helpful Answer:  
+"""
+
+        return self.llm.invoke(prompt)