app.py

import os
from dotenv import load_dotenv
import streamlit as st
from pathlib import Path

load_dotenv()
groq_api_key = os.getenv('GROQ_API_KEY')
os.environ['HF_TOKEN'] = os.getenv('HF_TOKEN')

from langchain_groq import ChatGroq
llm = ChatGroq(model="llama-3.3-70b-versatile", groq_api_key=groq_api_key)

# # ---- Load Vector Store ----
from langchain.document_loaders import TextLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.vectorstores import FAISS

embeddings = HuggingFaceEmbeddings(model_name="BAAI/bge-large-en-v1.5",model_kwargs={"device": "cpu"})
loaded_db = FAISS.load_local("vectorstore/legal_db", embeddings, allow_dangerous_deserialization=True)
retriever = loaded_db.as_retriever()

# ---- Create Prompt & Chains ----
from langchain_core.prompts import ChatPromptTemplate
from langchain.chains.combine_documents import create_stuff_documents_chain
from langchain.chains import create_retrieval_chain

system_prompt = """
You are a highly knowledgeable legal assistant specializing in Pakistani law.
Your responses must use proper legal terminology and cite relevant sections of the law.
If asked about legal actions, suggest options like creating writs or filing appeals.
Be precise and avoid conjecture.

Context:
{context}
"""

prompt = ChatPromptTemplate.from_messages([
    ("system", system_prompt),
    ("human", "{input}"),
])

question_answer_chain = create_stuff_documents_chain(llm, prompt)
rag_chain = create_retrieval_chain(retriever, question_answer_chain)

# ---- Streamlit UI ----
st.set_page_config(page_title="Pakistan Legal Advisor", layout="wide")
st.title("📚 Pakistan Legal Advisor 🇵🇰")
st.write("Ask any legal question based on Pakistani law:")

query = st.text_input("Your legal question:")

if query:
    with st.spinner("Thinking..."):
        response = rag_chain.invoke({"input": query})
        st.subheader("📜 Legal Answer:")
        st.markdown(response['answer'])