Update main.py

main.py

@@ -17,28 +17,21 @@ import uvicorn
 # Logging configuration
 logging.basicConfig(level=logging.DEBUG)
 logger = logging.getLogger(__name__)
-logger.debug("Starting FastAPI app...")
+logger.debug("Starting application...")
 
 # Suppress warnings
-warnings.filterwarnings("ignore", message="You are using `torch.load` with `weights_only=False`")
-warnings.filterwarnings("ignore", category=FutureWarning)
-warnings.filterwarnings("ignore", category=DeprecationWarning)
+warnings.filterwarnings("ignore", message="You are using `torch.load` with `weights_only=False")
 
 # Load environment variables
 load_dotenv()
 TOGETHER_AI_API = os.getenv("TOGETHER_AI")
 HF_HOME = os.getenv("HF_HOME", "./cache")
-
-# Set cache directory for Hugging Face
 os.environ["HF_HOME"] = HF_HOME
-
-# Ensure HF_HOME exists and is writable
 if not os.path.exists(HF_HOME):
     os.makedirs(HF_HOME, exist_ok=True)
 
-# Validate environment variables
 if not TOGETHER_AI_API:
-    raise ValueError("Environment variable TOGETHER_AI_API is missing. Please set it in your .env file.")
+    raise ValueError("TOGETHER_AI_API environment variable is missing.")
 
 # Initialize embeddings
 try:
@@ -48,54 +41,43 @@ try:
     )
 except Exception as e:
     logger.error(f"Error loading embeddings: {e}")
-    raise RuntimeError("Error initializing HuggingFaceEmbeddings.")
+    raise RuntimeError("Failed to initialize embeddings.")
 
-# Ensure FAISS vectorstore is loaded or created
+# 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": 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(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": 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
+    loader = DirectoryLoader('./data')
+    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": 5, "score_threshold": 0.8})
+
+# Define prompt template
 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.
+As a legal chatbot specializing in the Indian Penal Code (IPC), provide accurate and concise answers based on the context. Respond only if the answer can be derived from the given context; otherwise, reply: "The information is not available in the provided context." Use professional language.
 
 CONTEXT: {context}
-
 CHAT HISTORY: {chat_history}
-
 QUESTION: {question}
-
 ANSWER:
 """
 prompt = PromptTemplate(template=prompt_template, input_variables=["context", "question", "chat_history"])
 
-# Initialize the Together API
+# Initialize Together API
 try:
     llm = Together(
         model="mistralai/Mistral-7B-Instruct-v0.2",
-        temperature=0.3,  # Lower temperature ensures deterministic answers
-        max_tokens=512,   # Shorter response for focus
+        temperature=0.3,
+        max_tokens=512,
         together_api_key=TOGETHER_AI_API,
     )
 except Exception as e:
     logger.error(f"Error initializing Together API: {e}")
-    raise RuntimeError("Together API could not be initialized. Check your API key and network connection.")
+    raise RuntimeError("Failed to initialize Together API.")
 
 # Initialize conversational retrieval chain
 memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)
@@ -106,42 +88,30 @@ qa = ConversationalRetrievalChain.from_llm(
     combine_docs_chain_kwargs={"prompt": prompt},
 )
 
-# Initialize FastAPI app
+# FastAPI backend
 app = FastAPI()
 
-# Define request and response models
 class ChatRequest(BaseModel):
     question: str
 
 class ChatResponse(BaseModel):
     answer: str
 
-# Health check endpoint
 @app.get("/")
 async def root():
-    return {"message": "Hello, World!"}
+    return {"message": "Legal Chatbot is running."}
 
-# Chat endpoint
 @app.post("/chat", response_model=ChatResponse)
 async def chat(request: ChatRequest):
     try:
-        logger.debug(f"User Question: {request.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}")
         raise HTTPException(status_code=500, detail="Internal server error")
 
-# Start Uvicorn server if run directly
+# Start Uvicorn if run directly
 if __name__ == "__main__":
     uvicorn.run("main:app", host="0.0.0.0", port=7860)
-