Update main.py

main.py

@@ -1,7 +1,6 @@
 import logging
-from langchain_community.document_loaders import PyPDFLoader, DirectoryLoader
+from langchain_community.document_loaders import DirectoryLoader
 from langchain_huggingface import HuggingFaceEmbeddings
-from sentence_transformers import SentenceTransformer
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain_community.vectorstores import FAISS
 from langchain.prompts import PromptTemplate
@@ -22,7 +21,6 @@ logger.debug("Starting FastAPI app...")
 
 # Suppress warnings
 warnings.filterwarnings("ignore", message="You are using `torch.load` with `weights_only=False`")
-warnings.filterwarnings("ignore", message="Tried to instantiate class '__path__._path'")
 warnings.filterwarnings("ignore", category=FutureWarning)
 warnings.filterwarnings("ignore", category=DeprecationWarning)
 
@@ -55,37 +53,44 @@ except Exception as e:
 # Ensure FAISS vectorstore is loaded or created
 try:
     db = FAISS.load_local("ipc_vector_db", embeddings, allow_dangerous_deserialization=True)
-    db_retriever = db.as_retriever(search_type="similarity", search_kwargs={"k": 2, "max_length": 512})
+    db_retriever = db.as_retriever(search_type="similarity", search_kwargs={"k": 5, "score_threshold": 0.8})
 except Exception as e:
     logger.error(f"Error loading FAISS vectorstore: {e}")
     # If not found, create a new vectorstore
     try:
         loader = DirectoryLoader('./data')
-        text_splitter = RecursiveCharacterTextSplitter()
+        text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
         documents = text_splitter.split_documents(loader.load())
         db = FAISS.from_documents(documents, embeddings)
         db.save_local("ipc_vector_db")
-        db_retriever = db.as_retriever(search_type="similarity", search_kwargs={"k": 2, "max_length": 512})
+        db_retriever = db.as_retriever(search_type="similarity", search_kwargs={"k": 5, "score_threshold": 0.8})
     except Exception as inner_e:
         logger.error(f"Error creating FAISS vectorstore: {inner_e}")
         raise RuntimeError("FAISS vectorstore could not be created or loaded.")
 
 # Define the prompt template
-prompt_template = """<s>[INST]As a legal chatbot specializing in the Indian Penal Code, provide a concise and accurate answer based on the given context. Avoid unnecessary details or unrelated content. Only respond if the answer can be derived from the provided context; otherwise, say "The information is not available in the provided context." 
-    CONTEXT: {context}
-    CHAT HISTORY: {chat_history}
-    QUESTION: {question}
-    ANSWER:
-    </s>[INST]
-    """
+prompt_template = """
+As a legal chatbot specializing in the Indian Penal Code (IPC), provide precise, fact-based answers to the user’s question based on the provided context. 
+Respond only if the answer can be derived from the given context; otherwise, say: 
+"The information is not available in the provided context."
+Use plain, professional language in your response.
+
+CONTEXT: {context}
+
+CHAT HISTORY: {chat_history}
+
+QUESTION: {question}
+
+ANSWER:
+"""
 prompt = PromptTemplate(template=prompt_template, input_variables=["context", "question", "chat_history"])
 
 # Initialize the Together API
 try:
     llm = Together(
         model="mistralai/Mistral-7B-Instruct-v0.2",
-        temperature=0.5,
-        max_tokens=1024,
+        temperature=0.3,  # Lower temperature ensures deterministic answers
+        max_tokens=512,   # Shorter response for focus
         together_api_key=TOGETHER_AI_API,
     )
 except Exception as e:
@@ -120,9 +125,17 @@ async def root():
 @app.post("/chat", response_model=ChatResponse)
 async def chat(request: ChatRequest):
     try:
-        # Pass the user question
+        logger.debug(f"User Question: {request.question}")
         result = qa.invoke(input=request.question)
+        logger.debug(f"Retrieved Context: {result.get('context', '')}")
+        logger.debug(f"Model Response: {result.get('answer', '')}")
+
         answer = result.get("answer", "The chatbot could not generate a response.")
+        confidence_score = result.get("score", 0)  # Assuming LLM provides a score
+
+        if confidence_score < 0.7:
+            answer = "The answer is uncertain. Please consult a professional."
+
         return ChatResponse(answer=answer)
     except Exception as e:
         logger.error(f"Error during chat invocation: {e}")