app.py

import os
import platform
import re
import shutil
import tempfile
import subprocess
import time
import urllib.error
import urllib.request
from pathlib import Path
from typing import Final
import chromadb
import ollama
import streamlit as st
from chromadb.utils.embedding_functions.ollama_embedding_function import (
    OllamaEmbeddingFunction,
)
from langchain_community.document_loaders import PyMuPDFLoader
from langchain_core.documents import Document
from langchain_text_splitters import RecursiveCharacterTextSplitter
from sentence_transformers import CrossEncoder
from streamlit.runtime.uploaded_file_manager import UploadedFile


default_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:"

To answer the question:
1. Thoroughly analyze the context, identifying key information relevant to the question.
2. Organize your thoughts and plan your response to ensure a logical flow of information.
3. Formulate a detailed answer that directly addresses the question, using only the information provided in the context.
4. Ensure your answer is comprehensive, covering all relevant aspects found in the context.
5. If the context doesn't contain sufficient information to fully answer the question, state this clearly in your response.

Format your response as follows:
1. Return the PDB code in the format PDB:XXXX for example. Add it at the start of the response. Search in tables, references and details. If you don't know return PDB:XXXX .
2. Use clear, concise language.
3. Organize your answer into paragraphs for readability.
4. Use bullet points or numbered lists where appropriate to break down complex information.
5. If relevant, include any headings or subheadings to structure your response.
6. Ensure proper grammar, punctuation, and spelling throughout your answer.


Important: Base your entire response solely on the information provided in the context. 
Do not include any external knowledge or assumptions not present in the given text.
"""

computational_biologist_system_prompt = """
You are an AI assistant for structural biology and computational biology question answering.

Your task is to answer the user's question using ONLY the provided context.
You must be strictly evidence-grounded.

INPUT FORMAT
- context will be passed as "Context:"
- user question will be passed as "Question:"

PRIMARY RULES
1. Use only information explicitly present in the Context.
2. Do NOT use outside knowledge, domain priors, common biological assumptions, or likely completions.
3. Do NOT guess. Do NOT fill gaps. Do NOT infer unstated details unless the question explicitly asks for a limited inference and the inference is unavoidable and clearly marked as an inference.
4. If the Context is insufficient, ambiguous, contradictory, or incomplete, say so plainly.
5. It is better to return a partial but correct answer than a complete but speculative one.
6. If a detail is not directly supported by the Context, state: "Not stated in the provided context."
7. If multiple interpretations are possible, list them separately and explain why the Context does not allow disambiguation.
8. Never merge entities that may be distinct, including proteins, genes, isoforms, species, constructs, chains, domains, ligands, mutants, conformations, assemblies, or PDB entries.

DOMAIN-SPECIFIC RULES FOR STRUCTURAL / COMPUTATIONAL BIOLOGY
1. Preserve technical precision. Use exact identifiers, residue numbers, chain IDs, construct names, species names, ligand names, mutations, methods, and experimental qualifiers exactly as stated.
2. Distinguish carefully between:
   - experimental structure vs predicted model
   - apo vs holo
   - wild type vs mutant
   - monomer vs complex / assembly
   - chain-level vs protein-level statements
   - structure determination method vs downstream computational analysis
3. Do not generalize biological function beyond what is explicitly stated.
4. Do not claim binding, catalysis, mechanism, conformational change, interface details, or residue involvement unless explicitly supported by the Context.
5. If a PDB code, method, resolution, chain, residue range, mutation, ligand, organism, or dataset identifier is absent or uncertain, say so explicitly.

PDB EXTRACTION RULES
1. First, search the entire Context for candidate PDB identifiers, including tables, references, captions, supplementary-style text, and structured metadata.
2. Return exactly one leading line in this format:
   PDB:XXXX
3. If no PDB identifier is explicitly stated, return:
   PDB:XXXX
4. If multiple PDB identifiers are present and the relevant one cannot be determined from the Question and Context alone, return:
   PDB:AMBIGUOUS
   Then explain the candidate PDB entries in the body.
5. Never invent or normalize a PDB code that is not explicitly present.

RESPONSE CONSTRUCTION PROCESS
Follow this procedure before writing the answer:
1. Identify the exact entity or entities being asked about.
2. Extract only facts directly supported by the Context that are relevant to the Question.
3. Identify missing, ambiguous, or conflicting information.
4. Compose a conservative answer that includes only supported claims.
5. If the Context does not fully answer the Question, explicitly state what is missing.

OUTPUT FORMAT
- First line: PDB:XXXX (or PDB:AMBIGUOUS if applicable)
- Then include the following sections when relevant:

Answer:
A concise, evidence-grounded answer to the Question.

Supported details:
- Bullet points containing only facts explicitly stated in the Context.

Limitations / uncertainty:
- Bullet points listing missing, ambiguous, conflicting, or insufficiently supported details.

STYLE REQUIREMENTS
1. Write for an expert audience of structural biologists and computational biologists.
2. Be precise, restrained, and technically specific.
3. Prefer exact wording over broad paraphrase when technical details matter.
4. Do not add background explanation unless it is necessary to answer the Question and is explicitly supported by the Context.
5. Do not use promotional, conversational, or speculative language.

FINAL CHECK BEFORE RESPONDING
Before producing the final answer, verify:
- Every substantive claim is explicitly supported by the Context.
- No external knowledge has been introduced.
- No ambiguity has been silently resolved.
- No missing detail has been guessed.
- The answer is useful even if partial.
"""

OLLAMA_BASE_URL = os.getenv("OLLAMA_BASE_URL", "http://localhost:11434")
OLLAMA_CHAT_MODEL = os.getenv("OLLAMA_CHAT_MODEL", "llama3.2:3b")
OLLAMA_EMBED_MODEL = os.getenv("OLLAMA_EMBED_MODEL", "nomic-embed-text:latest")
OLLAMA_EMBEDDINGS_URL = f"{OLLAMA_BASE_URL.rstrip('/')}/api/embeddings"
CHROMA_CLIENT_MODE: Final[str] = os.getenv("CHROMA_CLIENT_MODE", "persistent").lower()
CHROMA_HOST = os.getenv("CHROMA_HOST", "localhost")
CHROMA_PORT = int(os.getenv("CHROMA_PORT", "8000"))
CHROMA_SSL = os.getenv("CHROMA_SSL", "false").lower() == "true"
CHROMA_STARTUP_RETRIES = int(os.getenv("CHROMA_STARTUP_RETRIES", "30"))
CHROMA_STARTUP_DELAY_SECONDS = float(os.getenv("CHROMA_STARTUP_DELAY_SECONDS", "2"))
APP_NAME = "RAGU"
APP_TITLE = "Local-first RAG for structural biology PDFs"
APP_LOGO_PATH = Path("public/images/logo.png")
VECTOR_STORE_DIR = "./ragu-data"
VECTOR_COLLECTION_NAME = "ragu_app"
RERANKER_MODEL_NAME = "cross-encoder/ms-marco-MiniLM-L-6-v2"
PIPER_COMMAND = os.getenv("PIPER_COMMAND", "piper")
PIPER_MODEL_PATH = Path(
    os.getenv("PIPER_MODEL_PATH", "/app/piper/models/en_US-lessac-medium.onnx")
)
PIPER_CONFIG_PATH = Path(
    os.getenv("PIPER_CONFIG_PATH", f"{PIPER_MODEL_PATH}.json")
)

def search_pdb_code(response: str):
    """
    Extract the first PDB identifier from a response string.

    This function searches for occurrences of PDB identifiers formatted as
    "XXXX" (where XXXX is a 4-character alphanumeric code) within
    the given input string and returns the first match found.

    The search is case-insensitive.

    Args:
        response (str): Input text that may contain PDB identifiers.

    Returns:
        str | None:
            - The first matching PDB identifier (e.g., "1ABC") if found.
            - None if no match is found or if the input is not a string.
    """

    if isinstance(response, str):

        result = re.search(r"PDB(\s+)?:(\s+)?[A-Z0-9]{4}", response, flags=re.IGNORECASE)

        if result:

            tmp = re.sub(r"PDB(\s+)?:(\s+)?", "", result.group(0))

            if tmp != "XXXX":
            
                return tmp
        
    return None

def process_document(uploaded_file: UploadedFile) -> list[Document]:
    """Processes an uploaded PDF file by converting it to text chunks.

    Takes an uploaded PDF file, saves it temporarily, loads and splits the content
    into text chunks using recursive character splitting.

    Args:
        uploaded_file: A Streamlit UploadedFile object containing the PDF file

    Returns:
        A list of Document objects containing the chunked text from the PDF

    Raises:
        IOError: If there are issues reading/writing the temporary file
    """
    # Store uploaded file as a temp file
    temp_file = tempfile.NamedTemporaryFile("wb", suffix=".pdf", delete=False)
    temp_file.write(uploaded_file.read())

    loader = PyMuPDFLoader(temp_file.name)
    docs = loader.load()
    # os.unlink(temp_file.name) # On windows doesn't work

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


@st.cache_resource(show_spinner=False)
def get_ollama_client() -> ollama.Client:
    return ollama.Client(host=OLLAMA_BASE_URL)


@st.cache_resource(show_spinner=False)
def get_chroma_client():
    if CHROMA_CLIENT_MODE == "http":
        scheme = "https" if CHROMA_SSL else "http"
        heartbeat_url = f"{scheme}://{CHROMA_HOST}:{CHROMA_PORT}/api/v2/heartbeat"

        last_error: Exception | None = None
        for _ in range(CHROMA_STARTUP_RETRIES):
            try:
                with urllib.request.urlopen(heartbeat_url, timeout=5) as response:
                    if response.status == 200:
                        break
            except (urllib.error.URLError, TimeoutError) as exc:
                last_error = exc
                time.sleep(CHROMA_STARTUP_DELAY_SECONDS)
        else:
            raise RuntimeError(
                f"Chroma server is unavailable at {heartbeat_url}"
            ) from last_error

        return chromadb.HttpClient(
            host=CHROMA_HOST,
            port=CHROMA_PORT,
            ssl=CHROMA_SSL,
        )

    return chromadb.PersistentClient(path=VECTOR_STORE_DIR)


@st.cache_resource(show_spinner=False)
def get_vector_collection() -> chromadb.Collection:
    """Gets or creates a ChromaDB collection for vector storage.

    Creates an Ollama embedding function using the nomic-embed-text model and initializes
    a persistent ChromaDB client. Returns a collection that can be used to store and
    query document embeddings.

    Returns:
        chromadb.Collection: A ChromaDB collection configured with the Ollama embedding
            function and cosine similarity space.
    """
    ollama_ef = OllamaEmbeddingFunction(
        url=OLLAMA_EMBEDDINGS_URL,
        model_name=OLLAMA_EMBED_MODEL,
    )

    chroma_client = get_chroma_client()
    return chroma_client.get_or_create_collection(
        name=VECTOR_COLLECTION_NAME,
        embedding_function=ollama_ef,
        metadata={"hnsw:space": "cosine"},
    )


def delete_vector_collection() -> None:
    chroma_client = get_chroma_client()
    chroma_client.delete_collection(name=VECTOR_COLLECTION_NAME)
    get_vector_collection.clear()


def add_to_vector_collection(all_splits: list[Document], file_name: str):
    """Adds document splits to a vector collection for semantic search.

    Takes a list of document splits and adds them to a ChromaDB vector collection
    along with their metadata and unique IDs based on the filename.

    Args:
        all_splits: List of Document objects containing text chunks and metadata
        file_name: String identifier used to generate unique IDs for the chunks

    Returns:
        None. Displays a success message via Streamlit when complete.

    Raises:
        ChromaDBError: If there are issues upserting documents to the collection
    """
    collection = get_vector_collection()
    documents, metadatas, ids = [], [], []

    for idx, split in enumerate(all_splits):
        documents.append(split.page_content)
        metadatas.append(split.metadata)
        ids.append(f"{file_name}_{idx}")

    collection.upsert(
        documents=documents,
        metadatas=metadatas,
        ids=ids,
    )
    st.success("Data added to the vector store!")


def query_collection(prompt: str, n_results: int = 10):
    """Queries the vector collection with a given prompt to retrieve relevant documents.

    Args:
        prompt: The search query text to find relevant documents.
        n_results: Maximum number of results to return. Defaults to 10.

    Returns:
        dict: Query results containing documents, distances and metadata from the collection.

    Raises:
        ChromaDBError: If there are issues querying the collection.
    """
    collection = get_vector_collection()
    results = collection.query(query_texts=[prompt], n_results=n_results)
    return results


def call_llm(context: str, prompt: str):
    """Calls the language model with context and prompt to generate a response.

    Uses Ollama to stream responses from a language model by providing context and a
    question prompt. The model uses a system prompt to format and ground its responses appropriately.

    Args:
        context: String containing the relevant context for answering the question
        prompt: String containing the user's question

    Yields:
        String chunks of the generated response as they become available from the model

    Raises:
        OllamaError: If there are issues communicating with the Ollama API
    """

    client = get_ollama_client()
    response = client.chat(
        model=OLLAMA_CHAT_MODEL,
        stream=True,
        messages=[
            {
                "role": "system",
                "content": computational_biologist_system_prompt,
            },
            {
                "role": "user",
                "content": f"Context: {context}, Question: {prompt}",
            },
        ],
    )
    for chunk in response:
        if chunk["done"] is False:
            yield chunk["message"]["content"]
        else:
            break


@st.cache_resource(show_spinner=False)
def get_reranker() -> CrossEncoder:
    return CrossEncoder(RERANKER_MODEL_NAME)


def re_rank_cross_encoders(documents: list[str]) -> tuple[str, list[int]]:
    """Re-ranks documents using a cross-encoder model for more accurate relevance scoring.

    Uses the MS MARCO MiniLM cross-encoder model to re-rank the input documents based on
    their relevance to the query prompt. Returns the concatenated text of the top 3 most
    relevant documents along with their indices.

    Args:
        documents: List of document strings to be re-ranked.

    Returns:
        tuple: A tuple containing:
            - relevant_text (str): Concatenated text from the top 3 ranked documents
            - relevant_text_ids (list[int]): List of indices for the top ranked documents

    Raises:
        ValueError: If documents list is empty
        RuntimeError: If cross-encoder model fails to load or rank documents
    """
    relevant_text = ""
    relevant_text_ids = []

    encoder_model = get_reranker()
    ranks = encoder_model.rank(prompt, documents, top_k=3)
    for rank in ranks:
        relevant_text += documents[rank["corpus_id"]]
        relevant_text_ids.append(rank["corpus_id"])

    return relevant_text, relevant_text_ids


def get_piper_command() -> str | None:
    return shutil.which(PIPER_COMMAND)


def get_piper_status() -> tuple[bool, str]:
    piper_command = get_piper_command()
    current_arch = platform.machine().lower()

    if piper_command is None:
        if current_arch in {"aarch64", "arm64"}:
            return (
                False,
                "Text-to-speech unavailable in this ARM container build: Piper is missing or failed to install.",
            )
        return False, "Text-to-speech unavailable: Piper binary not found."

    if not PIPER_MODEL_PATH.exists():
        return False, f"Text-to-speech unavailable: Piper model not found at {PIPER_MODEL_PATH}."

    if not PIPER_CONFIG_PATH.exists():
        return False, f"Text-to-speech unavailable: Piper config not found at {PIPER_CONFIG_PATH}."

    return True, "Text-to-speech available."


def piper_is_available() -> bool:
    available, _ = get_piper_status()
    return available


def synthesize_speech(text: str) -> bytes | None:
    piper_command = get_piper_command()
    if piper_command is None or not PIPER_MODEL_PATH.exists() or not PIPER_CONFIG_PATH.exists():
        return None

    with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as output_file:
        output_path = Path(output_file.name)

    try:
        try:
            subprocess.run(
                [
                    piper_command,
                    "--model",
                    str(PIPER_MODEL_PATH),
                    "--output_file",
                    str(output_path),
                ],
                input=text.encode("utf-8"),
                check=True,
                stdout=subprocess.DEVNULL,
                stderr=subprocess.PIPE,
            )
        except (OSError, subprocess.CalledProcessError):
            return None

        return output_path.read_bytes()
    finally:
        output_path.unlink(missing_ok=True)


if __name__ == "__main__":
    page_config = {"page_title": APP_NAME}
    if APP_LOGO_PATH.exists():
        page_config["page_icon"] = str(APP_LOGO_PATH)
    st.set_page_config(**page_config)
    if APP_LOGO_PATH.exists():
        st.logo(str(APP_LOGO_PATH))

    # Document Upload Area
    with st.sidebar:
        uploaded_file_list = st.file_uploader(
            "**📑 Upload PDF files for QnA**", type=["pdf"], accept_multiple_files=True
        )

        process = st.button(
            "⚡️ Process",
        )
        if uploaded_file_list and process:

            for uploaded_file in uploaded_file_list:
                
                normalize_uploaded_file_name = uploaded_file.name.translate(
                    str.maketrans({"-": "_", ".": "_", " ": "_"})
                )
                all_splits = process_document(uploaded_file)
                add_to_vector_collection(all_splits, normalize_uploaded_file_name)

        delete = st.button("💣 Delete Embeddings")

        if delete:

            delete_vector_collection()

    st.write("Session ID:", id(st.session_state))

    # Question and Answer Area
    st.header(f"🗣️ {APP_NAME}: {APP_TITLE}")
    piper_available, piper_status_message = get_piper_status()
    if piper_available:
        st.caption(piper_status_message)
    else:
        st.info(piper_status_message)
    prompt = st.text_area("**Ask a question related to your document:**")
    ask = st.button("🔥 Ask",)
    
    # Using .mvsj files in Javascript
    #
    # const sourceUrl = 'https://raw.githubusercontent.com/molstar/molstar/master/examples/mvs/1h9t_domain_labels.mvsj';
    # molstar.Viewer
    #        .create('entry-structure-viewer', { layoutIsExpanded: false, layoutShowControls: false })
    #        .then(viewer => viewer.loadMvsFromUrl(sourceUrl, 'mvsj'));

    if ask and prompt:
        
        pdb = None
        
        results = query_collection(prompt)
        context = results.get("documents")[0]
        relevant_text, relevant_text_ids = re_rank_cross_encoders(context)
        response = call_llm(context=relevant_text, prompt=prompt)

        def stream_and_collect(gen):
            full_text = ""
            for chunk in gen:
                text = chunk if isinstance(chunk, str) else chunk.get("content", "")
                full_text += text
                yield text
            st.session_state["last_response"] = full_text  # store it

        st.write_stream(stream_and_collect(response))

        full_text = st.session_state.get("last_response", "")

        pdb = search_pdb_code(full_text)

        if isinstance(pdb, str):

            st.components.v1.html("""
                <!-- Replace "latest" by the specific version you want to use, e.g. "4.0.0" -->
                <script src="https://cdn.jsdelivr.net/npm/molstar@latest/build/viewer/molstar.js"></script>
                <!-- Replace "latest" by the specific version you want to use, e.g. "4.0.0" -->
                <link rel="stylesheet" type="text/css" href="https://cdn.jsdelivr.net/npm/molstar@latest/build/viewer/molstar.css"/>
                <style>
                    .molstar-viewer {
                        position: relative;
                        width: 100%;
                        height: 100%;
                        overflow: hidden;
                        background: var(--bg);
                    }
                    .molstar-viewer .msp-plugin,
                    .molstar-viewer .msp-plugin-content,
                    .molstar-viewer .msp-layout-state,
                    .molstar-viewer .msp-layout-root {
                        width: 100%;
                        height: 100%;
                    }
                    .viewer-placeholder {
                        display: grid;
                        gap: 0.65rem;
                        min-height: 18rem;
                        place-content: center;
                        text-align: center;
                        border: 1px dashed var(--border-strong);
                        border-radius: var(--radius-sm);
                        background:
                            linear-gradient(180deg, rgba(10, 77, 122, 0.04), rgba(17, 167, 163, 0.01)),
                            var(--surface-muted);
                        padding: 1.5rem;
                    }
                    .structure-viewer-shell {
                        height: 500px;
                    }
                </style>
                <div class="structure-viewer-shell">
                    <div id="entry-structure-viewer" class="molstar-viewer" aria-label="3D protein structure viewer"></div>
                </div>
                <script>
                    const sourceUrl = 'https://files.rcsb.org/download/"""+pdb+""".cif';
                    molstar.Viewer
                    .create('entry-structure-viewer', { 
                        layoutIsExpanded: false, 
                        layoutShowControls: false 
                    })
                    .then(viewer => viewer.loadStructureFromUrl(sourceUrl, 'mmcif'));
                </script>
                """, height=500)

        with st.expander("See retrieved documents"):
            st.write(results)

        with st.expander("See most relevant document ids"):
            st.write(relevant_text_ids)
            st.write(relevant_text)

        if piper_available:
            audio_bytes = synthesize_speech(full_text)
            if audio_bytes is not None:
                st.subheader("Listen")
                st.audio(audio_bytes, format="audio/wav")
            else:
                st.warning("Text-to-speech is configured, but Piper failed to synthesize audio for this answer.")