app.py

import os
import tempfile
import requests
import chromadb
import ollama
import sys
import streamlit as st
from chromadb.utils.embedding_functions import EmbeddingFunction
from langchain_community.document_loaders import PyMuPDFLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from sentence_transformers import CrossEncoder
from streamlit_extras.add_vertical_space import add_vertical_space
sys.modules["torch.classes"] = None
os.environ["STREAMLIT_SERVER_RUN_ON_SAVE"] = "false"
# ----- Custom Embedding Function -----
class CustomOllamaEmbeddingFunction(EmbeddingFunction):
    def __init__(self, url, model_name):
        self.url = url
        self.model_name = model_name

    def __call__(self, texts):
        embeddings = []
        for text in texts:
            try:
                response = requests.post(
                    self.url,
                    json={"model": self.model_name, "prompt": text},
                    timeout=10
                )
                response.raise_for_status()
                embeddings.append(response.json()["embedding"])
            except Exception as e:
                print(f"Embedding error: {e}")
                embeddings.append([0.0] * 768)
        return embeddings

# ----- System Prompt -----
system_prompt = """
You are an AI assistant tasked with providing detailed answers based solely on the given context. 
Your goal is to analyze the information provided and formulate a comprehensive, well-structured response to the question.

Context will be passed as "Context:"
User question will be passed as "Question:"

Format:
- Use clear, concise language.
- Organize into paragraphs or bullet points.
- If the context doesn't have enough info, state it clearly.
- Base your response **only** on the context.
"""

# ----- Utility Functions -----
def clean_text(text):
    return " ".join(text.split())

def process_document(uploaded_file):
    with tempfile.NamedTemporaryFile("wb", suffix=".pdf", delete=False) as temp_file:
        temp_file.write(uploaded_file.read())
        temp_file_path = temp_file.name

    try:
        loader = PyMuPDFLoader(temp_file_path)
        docs = loader.load()
    finally:
        os.unlink(temp_file_path)

    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=400,
        chunk_overlap=100,
        separators=["\n\n", "\n", ".", "?", "!", " ", ""]
    )
    return text_splitter.split_documents(docs)

def get_vector_collection():
    embedding_fn = CustomOllamaEmbeddingFunction(
        url="http://localhost:11434/api/embeddings",
        model_name="nomic-embed-text:latest"
    )
    chroma_client = chromadb.PersistentClient(path="./demo-rag-chroma")
    return chroma_client.get_or_create_collection(
        name="rag_app",
        embedding_function=embedding_fn,
        metadata={"hnsw:space": "cosine"}
    )

def reset_and_fill_collection(all_splits, file_name):
    collection = get_vector_collection()
    existing_ids = collection.get(include=["metadatas"])['ids']
    if existing_ids:
        collection.delete(ids=existing_ids)

    documents, metadatas, ids = [], [], []
    for idx, split in enumerate(all_splits[:50]):
        cleaned_content = clean_text(split.page_content)
        documents.append(cleaned_content)
        metadatas.append(split.metadata)
        ids.append(f"{file_name}_{idx}")

    collection.upsert(documents=documents, metadatas=metadatas, ids=ids)

def query_collection(prompt, n_results=10):
    collection = get_vector_collection()
    return collection.query(query_texts=[prompt], n_results=n_results)

def call_llm(context, prompt):
    response = ollama.chat(
        model="llama3.2:3b",
        stream=True,
        messages=[
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": f"Context: {context}\n\nQuestion: {prompt}"}
        ]
    )
    for chunk in response:
        if not chunk["done"]:
            yield chunk["message"]["content"]
        else:
            break

def re_rank_cross_encoders(prompt, documents):
    if not documents:
        return "", []
    encoder_model = CrossEncoder("./local_models/ms-marco-MiniLM-L-6-v2")
    sentence_pairs = [(prompt, doc) for doc in documents]
    scores = encoder_model.predict(sentence_pairs)
    top_indices = sorted(range(len(scores)), key=lambda i: scores[i], reverse=True)[:3]
    relevant_text = "\n\n".join([documents[idx] for idx in top_indices])
    return relevant_text, top_indices

# ----- Streamlit UI Theme -----
st.set_page_config(page_title="SmartDoc AI Q&A", layout="wide", page_icon="📖")

st.markdown("""
<style>
    body, .main {
        background-color: #1a1a1a !important;
        color: #e5e7eb !important;
        font-family: 'Segoe UI', sans-serif;
    }

    .stSidebar {
        background-color: #333 !important;
        color: #f1f1f1;
        border-right: 1px solid #444;
    }

    .stSidebar .stButton > button,
    .stButton > button {
        background-color: #00c896 !important;
        color: black !important;
        font-weight: 600;
        border-radius: 6px;
        border: none;
        padding: 0.5rem 1.2rem;
        transition: background-color 0.2s ease-in-out;
    }

    .stButton > button:hover {
        background-color: #00b088 !important;
    }

    .stTextArea textarea {
        font-size: 16px;
        background-color: #2a2a2a;
        color: #e5e7eb;
        border: 1px solid #444;
        border-radius: 6px;
    }

    .stMarkdown h1, .stMarkdown h2, .stMarkdown h3 {
        color: #f3f4f6;
    }

    .stFileUploader label {
        color: #ddd;
    }

    .stExpanderHeader {
        font-weight: 600;
        color: #e5e7eb;
    }

    .stAlert {
        display: none !important; /* Hide default green success/info boxes */
    }

    .custom-answer-box, .custom-generating-box {
        background-color: #333;
        color: #e5e7eb;
        padding: 1rem 1.5rem;
        border-radius: 6px;
        margin-top: 1rem;
        margin-bottom: 1rem;
        border-left: 6px solid #00c896;
        box-shadow: 0 2px 8px rgba(0, 0, 0, 0.3);
        position: relative;
        font-size: 16px;
        line-height: 1.6;
    }

    .custom-generating-box {
        display: block;
        font-style: italic;
        color: #999;
    }

    .stDownloadButton, .stDownloadButton button, .stDownloadButton label {
        color: white !important;
    }

    .stExpander .streamlit-expanderHeader:hover {
        color: #00c896 !important;
    }
</style>
""", unsafe_allow_html=True)

# ----- UI Components -----
st.title("SmartDoc AI Q&A Assistant")
st.markdown("Upload a PDF and ask context-aware questions. The assistant provides accurate responses based only on the document content.")

with st.sidebar:
    st.header("Upload Your Document")
    uploaded_file = st.file_uploader("Choose a PDF file:", type=["pdf"])
    if uploaded_file and st.button("Process Document"):
        try:
            normalized_file_name = uploaded_file.name.replace("-", "_").replace(".", "_").replace(" ", "_")
            all_splits = process_document(uploaded_file)
            reset_and_fill_collection(all_splits, normalized_file_name)
            st.success("Document processed and vector store updated.")
        except Exception as e:
            st.error(f"Failed to process document: {e}")

add_vertical_space(2)

st.subheader("Ask a Question")
prompt = st.text_area("Enter your question here:", height=100)

if st.button("Get Answer"):
    if not prompt.strip():
        st.warning("Please enter a question before submitting.")
    else:
        try:
            results = query_collection(prompt)
            documents = results.get("documents", [[]])[0]
            if not documents:
                st.error("No results found. Please upload a document and try again.")
            else:
                relevant_text, relevant_ids = re_rank_cross_encoders(prompt, documents)
                answer_stream = call_llm(context=relevant_text, prompt=prompt)

                # Custom styled box while generating
                # Display a placeholder while generating
                placeholder = st.empty()
                placeholder.markdown('<div class="custom-generating-box">Generating answer from document context...</div>', unsafe_allow_html=True)

                answer = ""
                for chunk in answer_stream:
                    answer += chunk

                # Clear the placeholder after answer is ready
                placeholder.empty()

                # Show final answer
                st.markdown(f'<div class="custom-answer-box">{answer}</div>', unsafe_allow_html=True)

                with st.expander("View Retrieved Passages"):
                    for idx, doc in enumerate(documents):
                        st.markdown(f"**Doc {idx+1}:** {doc}")

                with st.expander("Relevant Passages (Top 3)"):
                    st.markdown(relevant_text)

        except Exception as e:
            st.error(f"Error while generating answer: {e}")