Update main.py

main.py

@@ -1,30 +1,25 @@
 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 dotenv import load_dotenv
 from langchain_huggingface import HuggingFaceEmbeddings
-from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain_community.vectorstores import FAISS
 from langchain.prompts import PromptTemplate
 from langchain_together import Together
-from langchain.memory import ConversationBufferMemory
-from langchain.chains import ConversationalRetrievalChain
+import uvicorn
 
 # ==========================
-# Logging Configuration
+# Logging Setup
 # ==========================
-logging.basicConfig(level=logging.DEBUG)
-logger = logging.getLogger("LegalChatbot")
-logger.debug("Initializing Legal Chatbot application...")
+logging.basicConfig(level=logging.DEBUG, format="%(asctime)s - %(levelname)s - %(message)s")
+logger = logging.getLogger(__name__)
 
 # ==========================
 # Suppress Warnings
 # ==========================
-warnings.filterwarnings("ignore", message="You are using `torch.load` with `weights_only=False")
+warnings.filterwarnings("ignore")
 
 # ==========================
 # Load Environment Variables
@@ -33,16 +28,19 @@ load_dotenv()
 TOGETHER_AI_API = os.getenv("TOGETHER_AI")
 HF_HOME = os.getenv("HF_HOME", "./cache")
 os.environ["HF_HOME"] = HF_HOME
-
-# Ensure the HF_HOME directory exists
 os.makedirs(HF_HOME, exist_ok=True)
 
-# Validate required environment variables
 if not TOGETHER_AI_API:
-    raise ValueError("The TOGETHER_AI_API environment variable is missing. Please set it in your .env file.")
+    logger.error("TOGETHER_AI_API key is missing. Please set it in the environment variables.")
+    raise RuntimeError("API key not found. Set TOGETHER_AI_API in .env.")
+
+# ==========================
+# App Initialization
+# ==========================
+app = FastAPI()
 
 # ==========================
-# Initialize Embeddings
+# Load Existing IPC Vectorstore
 # ==========================
 try:
     embeddings = HuggingFaceEmbeddings(
@@ -50,33 +48,31 @@ try:
         model_kwargs={"trust_remote_code": True, "revision": "289f532e14dbbbd5a04753fa58739e9ba766f3c7"},
     )
     logger.info("Embeddings successfully initialized.")
-except Exception as e:
-    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:
+    # Load the pre-existing IPC vector store directly
+    logger.info("Loading existing IPC vectorstore.")
     db = FAISS.load_local("ipc_vector_db", embeddings, allow_dangerous_deserialization=True)
-    db_retriever = db.as_retriever(search_type="mmr", search_kwargs={"k": 5, "max-length": 512})
-    logger.info("Vectorstore successfully loaded.")
+
+    db_retriever = db.as_retriever(search_type="similarity", search_kwargs={"k": 5})
+    logger.info("IPC 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.")
+    logger.error(f"Error during vectorstore setup: {e}")
+    raise RuntimeError("Initialization failed. Please check your embeddings or vectorstore setup.")
 
 # ==========================
-# Define Prompt Template
+# Prompt Template (Context-Only)
 # ==========================
-prompt_template = """<s>[INST]You are a legal chatbot specializing in the Indian Penal Code. Provide concise, context-aware answers in a conversational tone. Avoid presenting the response as a question-answer format unless explicitly required. 
-If the answer cannot be derived from the given context, respond with: "I'm sorry, I couldn't find relevant information for your query."
+prompt_template = """<s>[INST]
+You are a legal assistant specializing in the Indian Penal Code (IPC). Use only the provided CONTEXT to answer questions. 
+If the information is not found in the CONTEXT, respond with: "I don't have enough information yet." 
+Do not use any outside knowledge.
 CONTEXT: {context}
-CHAT HISTORY: {chat_history}
-QUESTION: {question}
-ANSWER:
-</s>[INST]"""
+USER QUERY: {question}
+RESPONSE:
+</s>[INST]
+"""
 
-prompt = PromptTemplate(template=prompt_template, input_variables=["context", "question", "chat_history"])
+prompt = PromptTemplate(template=prompt_template, input_variables=["context", "question"])
 
 # ==========================
 # Initialize Together API
@@ -91,25 +87,48 @@ try:
     logger.info("Together API successfully initialized.")
 except Exception as e:
     logger.error(f"Error initializing Together API: {e}")
-    raise RuntimeError("Something went wrong with the Together API setup. Please verify your API key and configuration.")
+    raise RuntimeError("Something went wrong with the Together API setup. Please verify your API key.")
 
 # ==========================
-# Conversational Retrieval Chain
+# Chat Processing Function
 # ==========================
-memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)
-qa = ConversationalRetrievalChain.from_llm(
-    llm=llm,
-    memory=memory,
-    retriever=db_retriever,
-    combine_docs_chain_kwargs={"prompt": prompt},
-)
-logger.info("Conversational Retrieval Chain initialized.")
+def generate_response(user_query: str) -> str:
+    try:
+        # Retrieve relevant documents
+        retrieved_docs = db_retriever.get_relevant_documents(user_query)
+        
+        # Log retrieved documents
+        logger.info(f"User Query: {user_query}")
+        for i, doc in enumerate(retrieved_docs):
+            logger.info(f"Document {i + 1}: {doc.page_content[:500]}...")
+
+        # Prepare context for the LLM
+        context = "\n\n".join(doc.page_content for doc in retrieved_docs)
+
+        # Check if context is empty
+        if not context.strip():
+            return "I don't have enough information yet."
+
+        # Construct LLM prompt input
+        prompt_input = {"context": context, "question": user_query}
+        logger.debug(f"Payload sent to LLM: {prompt_input}")
+
+        # Generate response using the LLM
+        response = llm(prompt.format(**prompt_input))
+        
+        # Check if response is empty
+        if not response.strip():
+            return "I don't have enough information yet."
+        
+        return response
+
+    except Exception as e:
+        logger.error(f"Error generating response: {e}")
+        return "An error occurred while generating the response."
 
 # ==========================
-# FastAPI Backend
+# FastAPI Models and Endpoints
 # ==========================
-app = FastAPI()
-
 class ChatRequest(BaseModel):
     question: str
 
@@ -118,41 +137,23 @@ class ChatResponse(BaseModel):
 
 @app.get("/")
 async def root():
-    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?"}
+    return {
+        "message": "Welcome to the Legal Chatbot! Ask me questions about the Indian Penal Code (IPC)."
+    }
 
 @app.post("/chat", response_model=ChatResponse)
 async def chat(request: ChatRequest):
     try:
-        logger.debug(f"Received user question: {request.question}")
-        
-        # Retrieve documents and log them with similarity scores
-        retrieved_docs = db_retriever.invoke({"query": request.question})
-        logger.debug("Retrieved Documents and Scores:")
-        for i, doc in enumerate(retrieved_docs["documents"], start=1):
-            logger.debug(f"Document {i}: {doc.page_content[:500]}...")
-            logger.debug(f"Score: {retrieved_docs['scores'][i-1]}")
-        
-        # Invoke the QA chain with the user question
-        result = qa.invoke(input=request.question)
-        if isinstance(result, dict) and "answer" in result:
-            answer = result["answer"]
-        else:
-            answer = "I'm sorry, I couldn't find relevant information for your query."
-
-        
-        if not answer or "The information is not available in the provided context" in answer:
-            answer = "I'm sorry, I couldn't find relevant information for your query. Please try rephrasing or providing more details."
-        
-        # Log the final answer
-        logger.debug(f"Chatbot Answer: {answer}")
-        
+        logger.debug(f"User question received: {request.question}")
+        answer = generate_response(request.question)
+        logger.debug(f"Chatbot response: {answer}")
         return ChatResponse(answer=answer)
     except Exception as e:
-        logger.error(f"Error during chat invocation: {e}")
-        raise HTTPException(status_code=500, detail="Oops! Something went wrong on our end. Please try again later.")
+        logger.error(f"Error processing chat request: {e}")
+        raise HTTPException(status_code=500, detail="An internal error occurred. Please try again later.")
 
 # ==========================
 # Run Uvicorn Server
 # ==========================
 if __name__ == "__main__":
-    uvicorn.run("main:app", host="0.0.0.0", port=7860)
+    uvicorn.run("main:app", host="0.0.0.0", port=7860, reload=True)