main.py

import logging
from langchain_community.document_loaders import DirectoryLoader
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
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import os
from dotenv import load_dotenv
import warnings
import uvicorn

# Logging configuration
logging.basicConfig(level=logging.DEBUG)
logger = logging.getLogger(__name__)
logger.debug("Starting application...")

# Suppress warnings
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")
os.environ["HF_HOME"] = HF_HOME
if not os.path.exists(HF_HOME):
    os.makedirs(HF_HOME, exist_ok=True)

if not TOGETHER_AI_API:
    raise ValueError("TOGETHER_AI_API environment variable is missing.")

# Initialize embeddings
try:
    embeddings = HuggingFaceEmbeddings(
        model_name="nomic-ai/nomic-embed-text-v1",
        model_kwargs={"trust_remote_code": True,"revision":"289f532e14dbbbd5a04753fa58739e9ba766f3c7"},
    )
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})


# Define prompt template
prompt_template = """
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 Together API
try:
    llm = Together(
        model="mistralai/Mistral-7B-Instruct-v0.2",
        temperature=0.5,
        max_tokens=1024,
        together_api_key=TOGETHER_AI_API,
    )
except Exception as e:
    logger.error(f"Error initializing Together API: {e}")
    raise RuntimeError("Failed to initialize Together API.")

# Initialize conversational retrieval chain
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},
)

# FastAPI backend
app = FastAPI()

class ChatRequest(BaseModel):
    question: str

class ChatResponse(BaseModel):
    answer: str

@app.get("/")
async def root():
    return {"message": "Legal Chatbot is running."}

@app.post("/chat", response_model=ChatResponse)
async def chat(request: ChatRequest):
    try:
        logger.debug(f"User question: {request.question}")
        result = qa.invoke(input=request.question)
        answer = result.get("answer", "The chatbot could not generate a response.")
        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 if run directly
if __name__ == "__main__":
    uvicorn.run("main:app", host="0.0.0.0", port=7860)