Update main.py

main.py

@@ -1,4 +1,10 @@
 import logging
+import os
+import warnings
+from dotenv import load_dotenv
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
+import uvicorn
 from langchain_community.document_loaders import DirectoryLoader
 from langchain_huggingface import HuggingFaceEmbeddings
 from langchain.text_splitter import RecursiveCharacterTextSplitter
@@ -7,59 +13,73 @@ from langchain.prompts import PromptTemplate
 from langchain_together import Together
 from langchain.memory import ConversationBufferMemory
 from langchain.chains import ConversationalRetrievalChain
-from fastapi import FastAPI, HTTPException
-from pydantic import BaseModel
-import os
-from dotenv import load_dotenv
-import warnings
-import uvicorn
 
-# Logging configuration
+# ==========================
+# Logging Configuration
+# ==========================
 logging.basicConfig(level=logging.DEBUG)
-logger = logging.getLogger(__name__)
-logger.debug("Starting application...")
+logger = logging.getLogger("LegalChatbot")
+logger.debug("Initializing Legal Chatbot application...")
 
-# Suppress warnings
+# ==========================
+# Suppress Warnings
+# ==========================
 warnings.filterwarnings("ignore", message="You are using `torch.load` with `weights_only=False")
 
-# Load environment variables
+# ==========================
+# Load Environment Variables
+# ==========================
 load_dotenv()
 TOGETHER_AI_API = os.getenv("TOGETHER_AI")
 HF_HOME = os.getenv("HF_HOME", "./cache")
 os.environ["HF_HOME"] = HF_HOME
-if not os.path.exists(HF_HOME):
-    os.makedirs(HF_HOME, exist_ok=True)
 
+# Ensure the HF_HOME directory exists
+os.makedirs(HF_HOME, exist_ok=True)
+
+# Validate required environment variables
 if not TOGETHER_AI_API:
-    raise ValueError("TOGETHER_AI_API environment variable is missing.")
+    raise ValueError("The TOGETHER_AI_API environment variable is missing. Please set it in your .env file.")
 
-# Initialize embeddings
+# ==========================
+# Initialize Embeddings
+# ==========================
 try:
     embeddings = HuggingFaceEmbeddings(
         model_name="nomic-ai/nomic-embed-text-v1",
-        model_kwargs={"trust_remote_code": True,"revision":"289f532e14dbbbd5a04753fa58739e9ba766f3c7"},
+        model_kwargs={"trust_remote_code": True, "revision": "289f532e14dbbbd5a04753fa58739e9ba766f3c7"},
     )
+    logger.info("Embeddings successfully initialized.")
 except Exception as e:
-    logger.error(f"Error loading embeddings: {e}")
-    raise RuntimeError("Failed to initialize embeddings.")
-
-# Load FAISS vectorstore
-
-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})
+    logger.error(f"Error initializing embeddings: {e}")
+    raise RuntimeError("Oops! Something went wrong while setting up embeddings. Please check the configuration and try again.")
 
+# ==========================
+# Load FAISS Vectorstore
+# ==========================
+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})
+    logger.info("Vectorstore successfully loaded.")
+except Exception as e:
+    logger.error(f"Error loading FAISS vectorstore: {e}")
+    raise RuntimeError("We couldn't load the vector database. Please ensure the database file is available and try again.")
 
-# Define 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." 
+# ==========================
+# Define 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]
-"""
+</s>[INST]"""
+
 prompt = PromptTemplate(template=prompt_template, input_variables=["context", "question", "chat_history"])
 
+# ==========================
 # Initialize Together API
+# ==========================
 try:
     llm = Together(
         model="mistralai/Mistral-7B-Instruct-v0.2",
@@ -67,11 +87,14 @@ try:
         max_tokens=1024,
         together_api_key=TOGETHER_AI_API,
     )
+    logger.info("Together API successfully initialized.")
 except Exception as e:
     logger.error(f"Error initializing Together API: {e}")
-    raise RuntimeError("Failed to initialize Together API.")
+    raise RuntimeError("Something went wrong with the Together API setup. Please verify your API key and configuration.")
 
-# Initialize conversational retrieval chain
+# ==========================
+# Conversational Retrieval Chain
+# ==========================
 memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)
 qa = ConversationalRetrievalChain.from_llm(
     llm=llm,
@@ -79,8 +102,11 @@ qa = ConversationalRetrievalChain.from_llm(
     retriever=db_retriever,
     combine_docs_chain_kwargs={"prompt": prompt},
 )
+logger.info("Conversational Retrieval Chain initialized.")
 
-# FastAPI backend
+# ==========================
+# FastAPI Backend
+# ==========================
 app = FastAPI()
 
 class ChatRequest(BaseModel):
@@ -91,19 +117,24 @@ class ChatResponse(BaseModel):
 
 @app.get("/")
 async def root():
-    return {"message": "Legal Chatbot is running."}
+    return {"message": "Hello! Welcome to the Legal Chatbot. I'm here to assist you with your legal queries related to the Indian Penal Code. How can I help you today?"}
 
 @app.post("/chat", response_model=ChatResponse)
 async def chat(request: ChatRequest):
     try:
-        logger.debug(f"User question: {request.question}")
+        logger.debug(f"Received user question: {request.question}")
         result = qa.invoke(input=request.question)
-        answer = result.get("answer", "The chatbot could not generate a response.")
+        answer = result.get("answer", "I'm sorry, but I couldn't generate a response to your query. Please try rephrasing or providing more details.")
         return ChatResponse(answer=answer)
     except Exception as e:
         logger.error(f"Error during chat invocation: {e}")
-        raise HTTPException(status_code=500, detail="Internal server error")
+        raise HTTPException(status_code=500, detail="Oops! Something went wrong on our end. Please try again later.")
 
-# Start Uvicorn if run directly
+# ==========================
+# Run Uvicorn Server
+# ==========================
 if __name__ == "__main__":
     uvicorn.run("main:app", host="0.0.0.0", port=7860)
+
+    
+