import streamlit as st
import uuid
import os
import certifi
from pymongo import MongoClient
from dotenv import load_dotenv
import sys
import time
import html
st.set_page_config(
page_title="پاکستان قانونی AI",
layout="wide",
initial_sidebar_state="expanded"
)
# Space/local repo root
root_dir = os.path.dirname(os.path.abspath(__file__))
if root_dir not in sys.path:
sys.path.append(root_dir)
from rag_pipeline import RAGPipeline
# Optional for local laptop testing only
env_path = os.path.join(root_dir, ".env")
if os.path.exists(env_path):
load_dotenv(dotenv_path=env_path)
MONGO_URL = os.getenv("MONGO_URL")
# --- INITIALIZE RAG BACKEND ---
@st.cache_resource(show_spinner=False)
def init_rag():
rag = RAGPipeline(chunking_strategy="fixed")
# Correct path to your JSON in the scrapper folder
json_path = os.path.join(root_dir, "cleaned_ocr_output.json")
if os.path.exists(json_path):
rag.load_bm25_from_json(json_path)
else:
st.warning(f"JSON not found at: {json_path}")
return rag
# ... [Place the Import & Path Fix code from Step 1 here] ...
# --- SESSION STATE ---
if "current_page" not in st.session_state:
st.session_state.current_page = "chat"
# ... [Keep your existing session state logic for chat_sessions etc.] ...
# --- SIDEBAR ---
# --- MAIN CONTENT AREA ---
# if st.session_state.current_page == "dashboard":
# # --- PASTE CONTENT OF dashboard.py HERE ---
# st.title("📊 System Evaluation & Ablation Study")
# # ... (Tables, metrics from your dashboard.py) ...
# else:
# # --- YOUR ORIGINAL CHAT UI CODE ---
# # Welcome screen, Chips, Chat bubbles, and Chat Input
# # ...
# if prompt := st.chat_input("اپنا قانونی سوال یہاں لکھیں..."):
# # The logic we built to call rag_pipeline.query()
# with st.spinner("جواب تیار کیا جا رہا ہے..."):
# result = rag_pipeline.query(prompt, run_evaluation=False)
# # append result to messages and st.rerun()
# --- MONGODB CONNECTION ---
@st.cache_resource
def init_connection():
if not MONGO_URL:
# st.warning("MONGO_URL not set. Running without chat history.")
return None
try:
client = MongoClient(
MONGO_URL,
tlsCAFile=certifi.where(),
serverSelectionTimeoutMS=5000,
)
client.admin.command("ping")
return client
except Exception as e:
# st.warning(f"MongoDB unavailable. Running without chat history. ({e})")
return None
client = init_connection()
if client:
db = client["UrduLegalAI"]
collection = db["chat_histories"]
# --- HELPER FUNCTIONS FOR MEMORY ---
def get_all_past_chats():
if not client:
return {}
try:
cursor = collection.find({}, {"chat_id": 1, "messages": 1})
return {doc["chat_id"]: doc["messages"] for doc in cursor}
except Exception as e:
st.warning(f"Could not load chat history. ({e})")
return {}
def save_chat(chat_id, messages):
if not client: return
collection.update_one(
{"chat_id": chat_id},
{"$set": {"messages": messages}},
upsert=True
)
def render_user_bubble(text: str):
safe_text = html.escape(text).replace("\n", "
")
st.markdown(f"""
""", unsafe_allow_html=True)
def render_typing_indicator(container):
container.markdown("""
""", unsafe_allow_html=True)
def stream_assistant_bubble(container, full_text: str, delay: float = 0.012):
words = full_text.split()
shown = ""
for word in words:
shown = (shown + " " + word).strip()
safe_text = html.escape(shown).replace("\n", "
")
container.markdown(f"""
""", unsafe_allow_html=True)
time.sleep(delay)
# --- GLOBAL CSS (YOUR EXACT CSS) ---
st.markdown("""
""", unsafe_allow_html=True)
# --- INITIALIZE PIPELINE ---
rag_pipeline = init_rag()
# --- SESSION STATE ---
if "db_synced" not in st.session_state:
st.session_state.chat_sessions = get_all_past_chats()
st.session_state.db_synced = True
if "chat_sessions" not in st.session_state:
st.session_state.chat_sessions = {}
if "current_chat_id" not in st.session_state:
st.session_state.current_chat_id = None
if "chip_prompt" not in st.session_state:
st.session_state.chip_prompt = None
if not st.session_state.current_chat_id:
if len(st.session_state.chat_sessions) > 0:
st.session_state.current_chat_id = list(st.session_state.chat_sessions.keys())[-1]
else:
init_id = str(uuid.uuid4())
st.session_state.chat_sessions[init_id] = []
st.session_state.current_chat_id = init_id
# --- SIDEBAR ---
with st.sidebar:
st.markdown("""
""", unsafe_allow_html=True)
if st.button("➕ نیا سوال", use_container_width=True):
new_id = str(uuid.uuid4())
st.session_state.chat_sessions[new_id] = []
st.session_state.current_chat_id = new_id
st.rerun()
sessions_except_current = [
(cid, msgs) for cid, msgs in st.session_state.chat_sessions.items()
if cid != st.session_state.current_chat_id
]
if sessions_except_current:
st.markdown('Recent', unsafe_allow_html=True)
for chat_id, msgs in reversed(sessions_except_current):
label = next(
(m["content"][:30] + "…" for m in msgs if m["role"] == "user"),
f"گفتگو {chat_id[:6]}…"
)
if st.button(label, key=chat_id, use_container_width=True):
st.session_state.current_chat_id = chat_id
st.rerun()
# --- MAIN AREA ---
current_messages = st.session_state.chat_sessions[st.session_state.current_chat_id]
CHIPS = [
"پاکستان کا ریاستی مذہب کیا ہے؟",
"پاکستان میں قومی اسمبلی کی مدت کتنی ہے؟",
"اگر کسی شخص کو غلط طریقے سے گرفتار کیا جائے تو اسے کیا حقوق حاصل ہیں؟",
"انسانی عزت کے بارے میں پاکستان کے آئین میں کیا کہا گیا ہے؟",
"وزیر اعظم کو برطرف کرنے کے لیے کیا طریقہ کار ہے؟",
"پاکستان میں جنگ کی صورت میں کیا ہوتا ہے؟ حکومت کو کیا اختیارات مل جاتے ہیں؟",
"عدالت عظمیٰ براہ راست کوئی مقدمہ سن سکتی ہے یا پہلے نچلی عدالت میں جانا ضروری ہے؟",
"وفاقی شرعی عدالت کیا کام کرتی ہے؟",
"پاکستان کا ریاستی مذہب کیا ہے اور یہ آئین میں کہاں لکھا ہے؟",
"عدالت عظمیٰ کے جج کب ریٹائر ہوتے ہیں؟"
]
if len(current_messages) == 0:
st.markdown("""
Pakistan Legal AI • پاکستانی قانون
آپ کے قانونی سوالات کا جواب
"اپنے حقوق جاننا آپ کی طاقت ہے اور ان کی حفاظت آپ کی شہری ذمہ داری ہے"
""", unsafe_allow_html=True)
st.markdown('', unsafe_allow_html=True)
for i in range(0, len(CHIPS), 2):
cols = st.columns(2)
with cols[0]:
if st.button(CHIPS[i], key=f"chip_{i}", use_container_width=True):
st.session_state.chip_prompt = CHIPS[i]
st.rerun()
if i + 1 < len(CHIPS):
with cols[1]:
if st.button(CHIPS[i+1], key=f"chip_{i+1}", use_container_width=True):
st.session_state.chip_prompt = CHIPS[i+1]
st.rerun()
st.markdown("
", unsafe_allow_html=True)
else:
st.markdown('', unsafe_allow_html=True)
for message in current_messages:
if message["role"] == "user":
st.markdown(f"""
""", unsafe_allow_html=True)
else:
st.markdown(f"""
""", unsafe_allow_html=True)
# Show chunks and scores if they exist
if message.get("chunks"):
with st.expander("📊 Evaluation Scores & Retrieved Context (تفصیلات)", expanded=False):
f_score = message.get("faithfulness", "N/A")
r_score = message.get("relevancy", "N/A")
st.markdown(f"""
""", unsafe_allow_html=True)
for i, chunk in enumerate(message["chunks"]):
score = chunk.get("rerank_score", chunk.get("rrf_score", chunk.get("score", 0)))
chunk_text = html.escape(chunk["text"]).replace("\n", "
")
st.markdown(f"""
Document {i+1}
Score: {score:.3f}
{chunk_text}
""", unsafe_allow_html=True)
# --- EXECUTE RAG PIPELINE FUNCTION ---
def execute_rag(user_input):
current_messages.append({"role": "user", "content": user_input})
save_chat(st.session_state.current_chat_id, current_messages)
# show user message immediately
render_user_bubble(user_input)
# assistant typing placeholder
assistant_placeholder = st.empty()
render_typing_indicator(assistant_placeholder)
history_for_rag = [{"role": m["role"], "content": m["content"]} for m in current_messages[:-1]]
try:
result = rag_pipeline.query(
user_query=user_input,
conversation_history=history_for_rag,
run_evaluation=True
)
answer = result["answer"]
chunks = result.get("retrieved_chunks", [])
f_val = result.get("faithfulness", {}).get("score")
r_val = result.get("relevancy", {}).get("score")
faithfulness = f"{f_val:.2%}" if isinstance(f_val, float) else "N/A"
relevancy = f"{r_val:.2%}" if isinstance(r_val, float) else "N/A"
except Exception as e:
answer = f"معذرت، ایک تکنیکی خرابی پیش آ گئی ہے۔ برائے مہربانی دوبارہ کوشش کریں۔\n\n(System Error: {str(e)})"
chunks, faithfulness, relevancy = [], "N/A", "N/A"
# replace typing indicator with streamed answer
stream_assistant_bubble(assistant_placeholder, answer, delay=0.01)
current_messages.append({
"role": "assistant",
"content": answer,
"chunks": chunks,
"faithfulness": faithfulness,
"relevancy": relevancy
})
save_chat(st.session_state.current_chat_id, current_messages)
time.sleep(0.2)
st.rerun()
# --- HANDLE INPUTS ---
if st.session_state.chip_prompt:
prompt = st.session_state.chip_prompt
st.session_state.chip_prompt = None
execute_rag(prompt)
if prompt := st.chat_input("اپنا قانونی سوال یہاں لکھیں..."):
execute_rag(prompt) 