app.py

import streamlit as st
import os
from langchain_community.vectorstores import Chroma
from langchain_google_genai import GoogleGenerativeAIEmbeddings
from langchain_groq import ChatGroq
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.runnables import RunnableLambda, RunnablePassthrough
from langchain_core.output_parsers import StrOutputParser
from langchain_core.messages import HumanMessage, AIMessage

# --- API Keys ---
if not os.getenv("GOOGLE_API_KEY") or not os.getenv("GROQ_API_KEY"):
    st.error("❌ Set GOOGLE_API_KEY and GROQ_API_KEY in your environment.")
    st.stop()

# --- Page Config ---
st.set_page_config(page_title="Jurisprudence Chatbot", layout="wide")
# CSS styling and JS
css = """
<style>
body {
    background: radial-gradient(circle at center, #1a1a2e 0%, #0f0f23 50%, #000000 100%);
    font-family: 'Arial', sans-serif;
    min-height: 100vh;
    margin: 0;
    position: relative;
    overflow-y: auto;
}
body::before {
    content: '';
    position: absolute;
    top: 0;
    left: 0;
    width: 100%;
    height: 100%;
    background: url('https://www.transparenttextures.com/patterns/dark-mosaic.png');
    opacity: 0.2;
    z-index: -1;
}
h1 {
    text-align: center;
    margin-top: 50px;
    background: linear-gradient(45deg, violet, purple, pink, white, black);
    -webkit-background-clip: text;
    background-clip: text;
    color: transparent;
    font-size: 2.5em;
    text-shadow: 0 0 10px rgba(255, 255, 255, 0.5), 0 0 20px rgba(255, 105, 180, 0.5);
    transition: text-shadow 0.3s ease, transform 0.3s ease;
    animation: slideIn 1s ease-out;
    background-color: transparent;
}
h1:hover {
    text-shadow: 0 0 15px rgba(255, 255, 255, 0.8), 0 0 30px rgba(255, 105, 180, 0.8);
    transform: scale(1.05);
}
.stChatInput {
    width: 90%;
    max-width: 800px;
    margin: 20px auto;
    z-index: 10;
}
.stChatInput.centered {
    position: fixed;
    top: 50%;
    left: 50%;
    transform: translate(-50%, -50%);
}
.stChatInput > div > div > input {
    background-color: black;
    color: white;
    font-size: 1.3em;
    height: 80px;
    border: 2px solid transparent;
    border-radius: 10px;
    padding: 10px;
    box-shadow: 0 0 0 2px linear-gradient(45deg, violet, purple, pink, white, black);
    background-clip: padding-box, border-box;
    background-origin: border-box;
    transition: box-shadow 0.3s ease, transform 0.3s ease;
}
.stChatInput > div > div > input:focus {
    box-shadow: 0 0 15px rgba(255, 105, 180, 0.8), 0 0 0 2px linear-gradient(45deg, violet, purple, pink, white, black);
    transform: scale(1.02);
    outline: none;
}
.stButton {
    width: 80%;
    max-width: 200px;
    margin: 20px auto;
    text-align: center;
    z-index: 10;
}
.stButton.centered {
    position: fixed;
    top: calc(50% + 100px);
    left: 50%;
    transform: translateX(-50%);
}
.stButton > button {
    background: linear-gradient(45deg, violet, purple, pink, white, black);
    -webkit-background-clip: text;
    background-clip: text;
    color: transparent;
    border: 2px solid transparent;
    border-radius: 10px;
    padding: 10px 20px;
    background-color: transparent;
    box-shadow: 0 0 10px rgba(255, 255, 255, 0.5), 0 0 20px rgba(255, 105, 180, 0.5);
    transition: box-shadow 0.3s ease, transform 0.3s ease;
    font-size: 1.2em;
    width: 100%;
}
.stButton > button:hover {
    box-shadow: 0 0 15px rgba(255, 255, 255, 0.8), 0 0 30px rgba(255, 105, 180, 0.8);
    transform: scale(1.05);
}
.stChatMessage {
    animation: fadeIn 0.5s ease-in;
    background-color: black;
    color: white;
    border-radius: 10px;
    padding: 15px;
    margin: 10px auto;
    width: 90%;
    max-width: 800px;
    box-shadow: 0 2px 4px rgba(0, 0, 0, 0.1);
}
.stChatMessage.user, .stChatMessage.assistant {
    background-color: black;
    color: white;
}
.stMarkdown p {
    font-size: 14px;
    color: white;
}
@keyframes slideIn {
    0% {
        opacity: 0;
        transform: translateY(-50px);
    }
    100% {
        opacity: 1;
        transform: translateY(0);
    }
}
@keyframes fadeIn {
    0% { opacity: 0; }
    100% { opacity: 1; }
}
</style>
<script>
(function() {
    let lastMessageCount = 0;
    const updateCentering = () => {
        const messages = document.querySelectorAll('.stChatMessage');
        const chatInput = document.querySelector('.stChatInput');
        const button = document.querySelector('.stButton');
        if (!chatInput || !button) return;
        if (messages.length === 0) {
            chatInput.classList.add('centered');
            button.classList.add('centered');
            lastMessageCount = 0;
        } else {
            chatInput.classList.remove('centered');
            button.classList.remove('centered');
            lastMessageCount = messages.length;
        }
    };
    setTimeout(updateCentering, 1000);
    setInterval(updateCentering, 300);
    const observer = new MutationObserver(updateCentering);
    observer.observe(document.body, { childList: true, subtree: true });
})();
</script>
"""
st.markdown(css, unsafe_allow_html=True)

# Title
st.markdown("""<h1 style='width:100%;background-color:transparent;'>📚 Jurisprudence Conversational Chatbot</h1>""", unsafe_allow_html=True)



# --- In-Memory Chat History (Thread-safe)
if "memory_dict" not in st.session_state:
    st.session_state.memory_dict = {"history": []}

memory_dict = st.session_state.memory_dict

# --- New Chat Button ---
if st.button("🆕 New Chat"):
    st.session_state.memory_dict = {"history": []}
    st.rerun()

# --- Display Chat History like ChatGPT ---
for msg in memory_dict["history"]:
    role = "user" if isinstance(msg, HumanMessage) else "assistant"
    with st.chat_message(role):
        st.markdown(msg.content)

# --- Chat Input ---
question = st.chat_input("💬 Ask a question from the Jurisprudence textbook...")

if question:
    # --- Show User Message ---
    memory_dict["history"].append(HumanMessage(content=question))
    with st.chat_message("user"):
        st.markdown(question)

    # --- Embeddings + Vector DB ---
    embeddings = GoogleGenerativeAIEmbeddings(model="models/embedding-001")
    vectordb = Chroma(persist_directory="chroma_db", embedding_function=embeddings)
    retriever = vectordb.as_retriever(search_type="mmr", search_kwargs={"k": 3})

    # --- Format Context from Docs ---
    def format_docs(docs):
        return "\n\n".join(doc.page_content for doc in docs)

    def get_context(question):
        docs = retriever.get_relevant_documents(question)
        return {
            "question": question,
            "context": format_docs(docs),
            "docs": docs
        }

    # --- Prompt & LLM ---
    prompt = ChatPromptTemplate.from_messages([
        ("system", """You are a legal assistant specialized in Jurisprudence, helping students understand complex legal topics from the textbook "Jurisprudence–I (Legal Method).

Instructions:
- Use ONLY the provided context to answer questions.
- If the user asks for a summary, generate a concise summary of the most relevant sections.
- If the user says "Explain like I’m five" or "ELI5", simplify the legal explanation as much as possible using analogies and plain language.
- Always provide citations in the format: [Section Title or Page Number].
- If you don’t know the answer, say you don’t know instead of making it up.

Context:
{context}"""),
        MessagesPlaceholder(variable_name="chat_history"),
        ("human", "{question}")
    ])

    llm = ChatGroq(model_name="llama3-70b-8192", api_key=os.getenv("GROQ_API_KEY"))
    parser = StrOutputParser()

    # --- Chain with memory injected manually (no st.session_state in threads) ---
    def chain_with_memory(input_dict):
        context_info = get_context(input_dict["input"])
        chain_input = {
            "question": context_info["question"],
            "context": context_info["context"],
            "chat_history": memory_dict["history"],
        }
        answer = prompt | llm | parser
        result = answer.invoke(chain_input)
        return {"result": result, "docs": context_info["docs"]}

    chain = RunnablePassthrough.assign() | RunnableLambda(chain_with_memory)

    # --- Get AI Response ---
    output = chain.invoke({"input": question})
    response = output["result"]

    memory_dict["history"].append(AIMessage(content=response))

    # --- Show AI Response ---
    with st.chat_message("assistant"):
        st.markdown(response)

    # # --- Show Sources (optional) ---
    # if output.get("docs"):
    #     with st.expander("📄 Source Documents"):
    #         for i, doc in enumerate(output["docs"], 1):
    #             st.markdown(f"**{i}.** {doc.page_content}")



st.markdown(
    """
    <style>
        .stApp {
            background-image: url("https://cdn-uploads.huggingface.co/production/uploads/67441c51a784a9d15cb12871/4iNiY-leNzRrY8SPCB3ZD.jpeg");
            background-size: cover;
            background-position: center;
            height: 100vh;
        }
        
        /* Semi-transparent overlay */
        .stApp::before {
            content: "";
            position: absolute;
            top: 0;
            left: 0;
            width: 100%;
            height: 100%;
            background: rgba(0, 0, 0, 0.4);  /* 40% transparency */
            z-index: -1;
        }
    </style>
    """, 
    unsafe_allow_html=True
)