Updating package file structure

app.py

@@ -1,351 +1,388 @@
 import os
-import sys
-import subprocess
-import re
+from typing import List, Dict, Tuple, Optional
 
+# Disable telemetry for LangChain and Chroma by default
 os.environ.setdefault("LANGCHAIN_TELEMETRY_ENABLED", "false")
 os.environ.setdefault("LANGCHAIN_DISABLE_TELEMETRY", "true")
 os.environ.setdefault("CHROMA_TELEMETRY_ENABLED", "false")
 
 import streamlit as st
-from src.vectorstore import get_retriever
-from src.qa_chain import make_conversational_chain
 
-st.set_page_config(page_title="Abalone RAG Chatbot", page_icon="🐚")
-
-st.title("Abalone RAG Chatbot")
-st.write(
-    "Ask natural-language questions about abalone studies and data. "
-    "The app uses a local Chroma vectorstore and OpenAI to retrieve and answer."
+from src.utils.rag_runtime import (
+    run_ingest_cli,
+    build_or_load_retriever_cached,
+    get_chain_cached,
 )
+from src.utils.metrics import compute_quality_scores
+from src.utils.formatting import format_source_label
+from src.utils.env import ensure_openai_key
 
-# ---------------- Sidebar ----------------
-
-st.sidebar.header("Model Settings")
 
-model_name = st.sidebar.selectbox(
-    "Model",
-    options=["gpt-3.5-turbo", "gpt-4"],
-    index=0,
-)
+class AbaloneRAGApp:
+    """Main application class for the Abalone RAG Chatbot."""
 
-st.sidebar.markdown("---")
+    def __init__(self) -> None:
+        """Initialize the Streamlit page and application state."""
+        st.set_page_config(page_title="Abalone RAG Chatbot", page_icon="🐚")
 
-st.sidebar.header("Retrieval Configuration")
+        st.title("Abalone RAG Chatbot")
+        st.write(
+            "Ask natural-language questions about abalone biology, ecology, "
+            "and research datasets. The app uses a local Chroma vectorstore "
+            "and OpenAI to retrieve and answer questions accurately."
+        )
 
-top_k = st.sidebar.slider(
-    "Number of retrieved chunks (k)",
-    min_value=2,
-    max_value=10,
-    value=4,
-)
+        # Data and vectorstore locations
+        self.data_dir = "./data"
+        self.persist_dir = "./vectorstore"
+
+        # Initialize session state
+        st.session_state.setdefault("chat_history", [])
+        st.session_state.setdefault("rebuild_pending", False)
+        self.chat_history: List[Dict] = st.session_state["chat_history"]
+
+        # Sidebar configuration
+        (
+            self.model_name,
+            self.top_k,
+            self.retrieval_mode,
+            self.temperature,
+            self.answer_length,
+            self.style_instruction,
+            self.rebuild_clicked,
+        ) = self._build_sidebar()
+
+        # QA chain instance (loaded lazily)
+        self.chain: Optional[object] = None
+
+    # ------------------------------------------------------------------
+    # Sidebar configuration
+    # ------------------------------------------------------------------
+
+    def _build_sidebar(self) -> Tuple[str, int, str, float, str, str, bool]:
+        """Render all sidebar controls and return model configuration.
+
+        Returns:
+            Tuple containing:
+                - model_name: Which LLM to use.
+                - top_k: Number of chunks to retrieve.
+                - retrieval_mode: Strategy (mmr, similarity, hybrid).
+                - temperature: LLM temperature.
+                - answer_length: Short/Medium/Long preference.
+                - style_instruction: Natural-language style directive.
+                - rebuild_clicked: Whether "Rebuild vectorstore" was pressed.
+        """
+        st.sidebar.header("Model Settings")
 
-retrieval_mode_label = st.sidebar.selectbox(
-    "Retrieval mode",
-    ["MMR (diverse)", "Similarity", "Hybrid (dense + MMR)"],
-    index=0,
-)
+        model_name = st.sidebar.selectbox(
+            "Model",
+            options=["gpt-3.5-turbo", "gpt-4"],
+            index=0,
+        )
 
-retrieval_mode_map = {
-    "MMR (diverse)": "mmr",
-    "Similarity": "similarity",
-    "Hybrid (dense + MMR)": "hybrid",
-}
-retrieval_mode = retrieval_mode_map[retrieval_mode_label]
+        st.sidebar.markdown("---")
 
-st.sidebar.markdown("---")
+        # Retrieval configuration
+        st.sidebar.header("Retrieval Configuration")
 
-st.sidebar.header("Answer Style")
+        top_k = st.sidebar.slider(
+            "Number of retrieved chunks (k)",
+            min_value=2,
+            max_value=10,
+            value=4,
+        )
 
-temperature = st.sidebar.slider(
-    "Temperature",
-    min_value=0.0,
-    max_value=1.0,
-    value=0.2,
-    step=0.05,
-)
+        retrieval_mode_label = st.sidebar.selectbox(
+            "Retrieval mode",
+            ["MMR (diverse)", "Similarity", "Hybrid (dense + MMR)"],
+            index=0,
+        )
+        retrieval_mode_map = {
+            "MMR (diverse)": "mmr",
+            "Similarity": "similarity",
+            "Hybrid (dense + MMR)": "hybrid",
+        }
+        retrieval_mode = retrieval_mode_map[retrieval_mode_label]
 
-answer_length = st.sidebar.selectbox(
-    "Answer length",
-    ["Short", "Medium", "Long"],
-    index=1,
-)
+        st.sidebar.markdown("---")
 
-st.sidebar.markdown("---")
+        # Answer style
+        st.sidebar.header("Answer Style")
 
-st.sidebar.header("Vectorstore Controls")
+        temperature = st.sidebar.slider(
+            "Temperature",
+            min_value=0.0,
+            max_value=1.0,
+            value=0.2,
+            step=0.05,
+        )
 
-rebuild_clicked = st.sidebar.button("Rebuild vectorstore", use_container_width=True)
+        answer_length = st.sidebar.selectbox(
+            "Answer length",
+            ["Short", "Medium", "Long"],
+            index=1,
+        )
 
-st.sidebar.markdown(
-    "<small>Use this when you add or modify files in <code>./data</code>.</small>",
-    unsafe_allow_html=True,
-)
+        st.sidebar.markdown("---")
 
-# -------------- Core config ----------------
-
-length_instruction_map = {
-    "Short": "Answer in 1–3 sentences.",
-    "Medium": "Answer in 1–2 paragraphs.",
-    "Long": "Provide a detailed, multi-paragraph explanation.",
-}
-length_instruction = length_instruction_map[answer_length]
-style_instruction = (
-        length_instruction
-        + f" Use a response style appropriate for a temperature of {temperature:.2f}, "
-          "where lower values are more factual and higher values are more exploratory."
-)
+        # Vectorstore controls
+        st.sidebar.header("Vectorstore Controls")
 
-data_dir = "./data"
-persist_dir = "./vectorstore"
-
-if "chat_history" not in st.session_state:
-    st.session_state["chat_history"] = []
-
-if "rebuild_pending" not in st.session_state:
-    st.session_state["rebuild_pending"] = False
-
-# -------------- Helpers ----------------
-
-def ensure_openai_key() -> bool:
-    if not os.environ.get("OPENAI_API_KEY"):
-        st.error("OPENAI_API_KEY is not set.")
-        return False
-    return True
-
-def run_ingest_cli(data_dir: str, persist_dir: str):
-    cmd = [
-        sys.executable,
-        "-m",
-        "src.ingest",
-        "--data-dir",
-        data_dir,
-        "--persist-dir",
-        persist_dir,
-    ]
-    subprocess.run(cmd, check=True)
-
-@st.cache_resource(show_spinner=False)
-def build_or_load_retriever_cached(
-        data_dir: str,
-        persist_dir: str,
-        top_k: int,
-        retrieval_mode: str,
-):
-    try:
-        return get_retriever(
-            persist_dir=persist_dir,
-            top_k=top_k,
-            retrieval_mode=retrieval_mode,
+        rebuild_clicked = st.sidebar.button(
+            "Rebuild vectorstore",
+            use_container_width=True,
         )
-    except Exception:
-        run_ingest_cli(data_dir=data_dir, persist_dir=persist_dir)
-        return get_retriever(
-            persist_dir=persist_dir,
-            top_k=top_k,
-            retrieval_mode=retrieval_mode,
+
+        st.sidebar.markdown(
+            "<small>Use this when you add or modify files in <code>./data</code>.</small>",
+            unsafe_allow_html=True,
         )
 
-@st.cache_resource(show_spinner=False)
-def get_chain(model_name: str, top_k: int, retrieval_mode: str):
-    retriever = build_or_load_retriever_cached(
-        data_dir=data_dir,
-        persist_dir=persist_dir,
-        top_k=top_k,
-        retrieval_mode=retrieval_mode,
-    )
-    return make_conversational_chain(retriever, model_name=model_name)
-
-def format_source_label(meta: dict, index: int) -> str:
-    source = (
-            meta.get("source")
-            or meta.get("file_path")
-            or meta.get("path")
-            or meta.get("document_id")
-            or "Unknown source"
-    )
-    return f"[{index}] {source}"
-
-def tokenize(text: str):
-    return [w.lower() for w in re.findall(r"\w+", text) if len(w) > 3]
-
-def compute_quality_scores(question: str, answer: str, sources: list):
-    all_chunk_text = " ".join(s.get("content", "") for s in sources)
-    q_tokens = tokenize(question)
-    a_tokens = tokenize(answer)
-    c_tokens = set(tokenize(all_chunk_text))
-    if not c_tokens:
-        return 0.0, 0.0
-    if not q_tokens:
-        coverage = 0.0
-    else:
-        coverage = sum(1 for t in q_tokens if t in c_tokens) / len(q_tokens)
-    if not a_tokens:
-        grounding = 0.0
-    else:
-        grounding = sum(1 for t in a_tokens if t in c_tokens) / len(a_tokens)
-    return coverage, grounding
-
-if not ensure_openai_key():
-    st.stop()
-
-# -------------- Rebuild confirmation + chain init ----------------
-
-if rebuild_clicked:
-    st.session_state["rebuild_pending"] = True
-
-chain = None
-
-if st.session_state["rebuild_pending"]:
-    st.warning(
-        "Rebuild the vectorstore from the current contents of ./data? "
-        "This will overwrite existing embeddings."
-    )
-
-    col_left, col_center, col_right = st.columns([1, 2, 1])
-
-    with col_center:
-        confirm_rebuild = st.button(
-            "Yes, rebuild",
-            key="confirm_rebuild",
-            use_container_width=True,
+        # Build style instruction for the LLM
+        length_instruction_map = {
+            "Short": "Answer in 1–3 sentences.",
+            "Medium": "Answer in 1–2 paragraphs.",
+            "Long": "Provide a detailed, multi-paragraph explanation.",
+        }
+        length_instruction = length_instruction_map[answer_length]
+        style_instruction = (
+                length_instruction
+                + f" Use a response style appropriate for a temperature of {temperature:.2f}, "
+                  "where lower values are more factual and higher values are more exploratory."
         )
-        cancel_rebuild = st.button(
-            "Cancel",
-            key="cancel_rebuild",
-            use_container_width=True,
+
+        return (
+            model_name,
+            top_k,
+            retrieval_mode,
+            temperature,
+            answer_length,
+            style_instruction,
+            rebuild_clicked,
         )
 
-    st.markdown(
+    # ------------------------------------------------------------------
+    # Vectorstore rebuild workflow
+    # ------------------------------------------------------------------
+
+    def handle_rebuild(self) -> None:
+        """Render rebuild confirmation dialog and rebuild if confirmed.
+
+        This manages the 2-step rebuild process:
+
+        1. User clicks "Rebuild vectorstore".
+        2. A confirmation dialog appears with "Yes, rebuild" and "Cancel".
+
+        If confirmed, the vectorstore is regenerated and caches are cleared.
         """
-        <style>
-            div[data-testid="column"] div:has(> button[aria-label="Yes, rebuild"]) button {
-                background-color: #27ae60 !important;
-                color: white !important;
-            }
-            div[data-testid="column"] div:has(> button[aria-label="Cancel"]) button {
-                background-color: #c0392b !important;
-                color: white !important;
-            }
-        </style>
-        """,
-        unsafe_allow_html=True,
-    )
-
-    if confirm_rebuild:
-        with st.spinner("Rebuilding vectorstore..."):
-            run_ingest_cli(data_dir=data_dir, persist_dir=persist_dir)
-            build_or_load_retriever_cached.clear()
-            get_chain.clear()
-            chain = get_chain(
-                model_name=model_name,
-                top_k=top_k,
-                retrieval_mode=retrieval_mode,
-            )
-        st.session_state["rebuild_pending"] = False
-        st.success("Vectorstore rebuilt successfully.")
-
-    elif cancel_rebuild:
-        st.session_state["rebuild_pending"] = False
-        st.info("Rebuild canceled.")
-
-if chain is None and not st.session_state["rebuild_pending"]:
-    with st.spinner("Initializing knowledge base and chat model..."):
-        chain = get_chain(
-            model_name=model_name,
-            top_k=top_k,
-            retrieval_mode=retrieval_mode,
+        if self.rebuild_clicked:
+            st.session_state["rebuild_pending"] = True
+
+        if not st.session_state["rebuild_pending"]:
+            return
+
+        st.warning(
+            "Rebuild the vectorstore from the current contents of ./data? "
+            "This will overwrite existing embeddings."
         )
-    st.success("Knowledge base and model are ready.")
-elif chain is not None and not st.session_state["rebuild_pending"]:
-    st.success("Knowledge base and model are ready.")
 
-# -------------- Render chat history ----------------
+        col_left, col_center, col_right = st.columns([1, 2, 1])
 
-if st.session_state["chat_history"]:
-    for turn in st.session_state["chat_history"]:
-        with st.chat_message("user"):
-            st.markdown(turn["question"])
-        answer_text = turn["answer"]
-        with st.chat_message("assistant"):
-            st.markdown(answer_text)
+        with col_center:
+            confirm = st.button(
+                "Yes, rebuild",
+                key="confirm_rebuild",
+                use_container_width=True,
+            )
+            cancel = st.button(
+                "Cancel",
+                key="cancel_rebuild",
+                use_container_width=True,
+            )
 
-# -------------- New user input ----------------
+        # Centered green (confirm) and red (cancel) buttons
+        st.markdown(
+            """
+            <style>
+                div[data-testid="column"] div:has(> button[aria-label="Yes, rebuild"]) button {
+                    background-color: #27ae60 !important;
+                    color: white !important;
+                }
+                div[data-testid="column"] div:has(> button[aria-label="Cancel"]) button {
+                    background-color: #c0392b !important;
+                    color: white !important;
+                }
+            </style>
+            """,
+            unsafe_allow_html=True,
+        )
 
-user_input = st.chat_input("Ask a question about abalone (biology, data, methodology, etc.)")
+        if confirm:
+            with st.spinner("Rebuilding vectorstore..."):
+                run_ingest_cli(data_dir=self.data_dir, persist_dir=self.persist_dir)
+                build_or_load_retriever_cached.clear()
+                get_chain_cached.clear()
+
+                self.chain = get_chain_cached(
+                    model_name=self.model_name,
+                    top_k=self.top_k,
+                    retrieval_mode=self.retrieval_mode,
+                    data_dir=self.data_dir,
+                    persist_dir=self.persist_dir,
+                )
+
+            st.session_state["rebuild_pending"] = False
+            st.success("Vectorstore rebuilt successfully.")
+
+        elif cancel:
+            st.session_state["rebuild_pending"] = False
+            st.info("Rebuild canceled.")
+
+    # ------------------------------------------------------------------
+    # Chain loading
+    # ------------------------------------------------------------------
+
+    def ensure_chain_ready(self) -> None:
+        """Load or create the QA chain unless a rebuild is still pending."""
+        if st.session_state["rebuild_pending"]:
+            return
+
+        if self.chain is None:
+            with st.spinner("Initializing knowledge base and chat model..."):
+                self.chain = get_chain_cached(
+                    model_name=self.model_name,
+                    top_k=self.top_k,
+                    retrieval_mode=self.retrieval_mode,
+                    data_dir=self.data_dir,
+                    persist_dir=self.persist_dir,
+                )
+            st.success("Knowledge base and model are ready.")
+        else:
+            st.success("Knowledge base and model are ready.")
+
+    # ------------------------------------------------------------------
+    # Chat UI
+    # ------------------------------------------------------------------
+
+    def render_chat_history(self) -> None:
+        """Render previous user and assistant messages."""
+        for turn in self.chat_history:
+            with st.chat_message("user"):
+                st.markdown(turn["question"])
+            with st.chat_message("assistant"):
+                st.markdown(turn["answer"])
+
+    def handle_user_input(self) -> None:
+        """Process new user queries, run RAG, compute metrics, and display results."""
+        if st.session_state["rebuild_pending"] or self.chain is None:
+            return
+
+        user_input = st.chat_input(
+            "Ask a question about abalone (biology, data, methodology, etc.)"
+        )
+        if not user_input:
+            return
 
-if user_input and chain is not None and not st.session_state["rebuild_pending"]:
-    with st.chat_message("user"):
-        st.markdown(user_input)
+        # Render user message
+        with st.chat_message("user"):
+            st.markdown(user_input)
 
-    with st.spinner("Thinking..."):
-        prior_history = [
-            (h.get("question"), h.get("answer", "")) for h in st.session_state["chat_history"]
-        ]
+        # Run inference
+        with st.spinner("Thinking..."):
+            prior_history: List[Tuple[str, str]] = [
+                (h.get("question"), h.get("answer", ""))
+                for h in self.chat_history
+            ]
 
-        styled_question = style_instruction + "\n\nQuestion: " + user_input
+            styled_question = self.style_instruction + "\n\nQuestion: " + user_input
 
-        result = chain(
-            {"question": styled_question, "chat_history": prior_history}
+            result = self.chain(
+                {"question": styled_question, "chat_history": prior_history}
+            )
+
+            answer = (
+                    result.get("answer")
+                    or result.get("result")
+                    or result.get("output_text")
+                    or ""
+            )
+            source_docs = result.get("source_documents") or []
+
+            # Normalize retrieved docs for UI and metrics
+            formatted_sources: List[Dict] = []
+            for idx, sd in enumerate(source_docs, start=1):
+                if isinstance(sd, dict):
+                    meta = sd.get("metadata", {}) or {}
+                    text = (
+                            sd.get("page_content")
+                            or sd.get("content")
+                            or sd.get("text", "")
+                            or ""
+                    )
+                else:
+                    meta = getattr(sd, "metadata", {}) or {}
+                    text = (
+                            getattr(sd, "page_content", None)
+                            or getattr(sd, "content", "")
+                            or ""
+                    )
+
+                formatted_sources.append(
+                    {"index": idx, "metadata": meta, "content": str(text)}
+                )
+
+        # Compute simple retrieval quality metrics
+        coverage, grounding = compute_quality_scores(
+            user_input, answer, formatted_sources
         )
+        coverage_pct = int(round(coverage * 100))
+        grounding_pct = int(round(grounding * 100))
 
-        answer = (
-                result.get("answer")
-                or result.get("result")
-                or result.get("output_text")
-                or ""
+        # Render assistant message + debug block
+        with st.chat_message("assistant"):
+            st.markdown(answer)
+
+            with st.expander("Retrieval Metrics and Sources"):
+                st.markdown(f"- Retrieval mode: `{self.retrieval_mode}`")
+                st.markdown(f"- k: `{self.top_k}`")
+                st.markdown(
+                    f"- Coverage score (question vs sources): **{coverage_pct}%**"
+                )
+                st.markdown(
+                    f"- Grounding score (answer vs sources): **{grounding_pct}%**"
+                )
+
+                if formatted_sources:
+                    st.markdown("**Retrieved chunks:**")
+                    for src in formatted_sources:
+                        label = format_source_label(src["metadata"], src["index"])
+                        snippet = src["content"][:200].replace("\n", " ")
+                        st.markdown(f"**[{src['index']}] {label}**")
+                        st.code(snippet + "...")
+
+        # Persist turn in chat history
+        self.chat_history.append(
+            {
+                "question": user_input,
+                "answer": answer,
+                "sources": formatted_sources,
+            }
         )
-        source_docs = result.get("source_documents") or []
-
-        sources_for_ui = []
-        for idx, sd in enumerate(source_docs, start=1):
-            if isinstance(sd, dict):
-                meta = sd.get("metadata", {}) or {}
-                content_full = sd.get("page_content") or sd.get("content") or sd.get("text", "")
-            else:
-                meta = getattr(sd, "metadata", {}) or {}
-                content_full = getattr(sd, "page_content", None)
-                if content_full is None:
-                    content_full = getattr(sd, "content", "")
-            if content_full is None:
-                content_full = ""
-            sources_for_ui.append(
-                {
-                    "index": idx,
-                    "metadata": meta,
-                    "content": str(content_full),
-                }
-            )
+        st.session_state["chat_history"] = self.chat_history
 
-    coverage, grounding = compute_quality_scores(user_input, answer, sources_for_ui)
-    coverage_pct = int(round(coverage * 100))
-    grounding_pct = int(round(grounding * 100))
-    answer_text = answer
-
-    with st.chat_message("assistant"):
-        st.markdown(answer_text)
-
-        with st.expander("Retrieval Metrics and Sources"):
-            st.markdown(f"- Retrieval mode: `{retrieval_mode}`")
-            st.markdown(f"- k: `{top_k}`")
-            st.markdown(f"- Coverage score (question vs sources): **{coverage_pct}%**")
-            st.markdown(f"- Grounding score (answer vs sources): **{grounding_pct}%**")
-
-            if sources_for_ui:
-                st.markdown("**Retrieved chunks:**")
-                for src in sources_for_ui:
-                    idx = src.get("index", 0)
-                    meta = src.get("metadata", {}) or {}
-                    label = format_source_label(meta, idx)
-                    chunk_text = src.get("content", "")
-                    snippet = chunk_text[:200].replace("\n", " ")
-                    st.markdown(f"**[{idx}] {label}**")
-                    st.code(snippet + "...")
-
-    st.session_state["chat_history"].append(
-        {
-            "question": user_input,
-            "answer": answer,
-            "sources": sources_for_ui,
-        }
-    )
+
+def main() -> None:
+    """Main entry point for running the Abalone RAG Chatbot app."""
+    app = AbaloneRAGApp()
+
+    if not ensure_openai_key():
+        st.stop()
+
+    app.handle_rebuild()
+    app.ensure_chain_ready()
+    app.render_chat_history()
+    app.handle_user_input()
+
+
+if __name__ == "__main__":
+    main()

src/utils/env.py

@@ -0,0 +1,10 @@
+import os
+import streamlit as st
+
+
+def ensure_openai_key(env_var: str = "OPENAI_API_KEY") -> bool:
+    """Ensure the specified OpenAI API key environment variable is present."""
+    if not os.environ.get(env_var):
+        st.error(f"{env_var} is not set.")
+        return False
+    return True

src/utils/formatting.py

@@ -0,0 +1,14 @@
+from typing import Dict
+
+
+def format_source_label(meta: Dict, index: int) -> str:
+    """Create a readable label for a retrieved chunk."""
+    source = (
+            meta.get("source")
+            or meta.get("file_path")
+            or meta.get("path")
+            or meta.get("document_id")
+            or "Unknown source"
+    )
+    return f"[{index}] {source}"
+

src/utils/metrics.py

@@ -0,0 +1,44 @@
+import re
+from typing import List, Dict, Tuple
+
+
+def tokenize(text: str) -> List[str]:
+    """Tokenize a string into lowercase words >3 chars."""
+    return [w.lower() for w in re.findall(r"\w+", text) if len(w) > 3]
+
+
+def compute_quality_scores(
+        question: str,
+        answer: str,
+        sources: List[Dict],
+) -> Tuple[float, float]:
+    """Compute retrieval quality metrics (coverage & grounding).
+
+    Args:
+        question: User's question text.
+        answer: Model-generated answer text.
+        sources: Retrieved documents/chunks, each with a 'content' field.
+
+    Returns:
+        (coverage, grounding) as floats in [0.0, 1.0].
+    """
+    all_chunk_text = " ".join(s.get("content", "") for s in sources)
+    q_tokens = tokenize(question)
+    a_tokens = tokenize(answer)
+    c_tokens = set(tokenize(all_chunk_text))
+
+    if not c_tokens:
+        return 0.0, 0.0
+
+    coverage = (
+        sum(1 for t in q_tokens if t in c_tokens) / len(q_tokens)
+        if q_tokens
+        else 0.0
+    )
+    grounding = (
+        sum(1 for t in a_tokens if t in c_tokens) / len(a_tokens)
+        if a_tokens
+        else 0.0
+    )
+
+    return coverage, grounding

src/utils/rag_runtime.py

@@ -0,0 +1,95 @@
+import sys
+import subprocess
+from typing import Any
+
+import streamlit as st
+
+from src.vectorstore import get_retriever
+from src.qa_chain import make_conversational_chain
+
+
+def run_ingest_cli(data_dir: str, persist_dir: str) -> None:
+    """Run the ingestion module to rebuild the vectorstore.
+
+    Args:
+        data_dir: Directory containing the raw text files.
+        persist_dir: Directory where embeddings and Chroma DB should be stored.
+
+    Raises:
+        CalledProcessError: If the underlying subprocess fails.
+    """
+    cmd = [
+        sys.executable,
+        "-m",
+        "src.ingest",
+        "--data-dir",
+        data_dir,
+        "--persist-dir",
+        persist_dir,
+    ]
+    subprocess.run(cmd, check=True)
+
+
+@st.cache_resource(show_spinner=False)
+def build_or_load_retriever_cached(
+        data_dir: str,
+        persist_dir: str,
+        top_k: int,
+        retrieval_mode: str,
+) -> Any:
+    """Load a retriever from the persisted vectorstore or build a new one.
+
+    If loading fails—usually because the vectorstore doesn't exist—this
+    function triggers ingestion and retries loading.
+
+    Args:
+        data_dir: Directory containing input documents.
+        persist_dir: Directory where the Chroma vectorstore is stored.
+        top_k: Number of chunks to retrieve for queries.
+        retrieval_mode: Retrieval strategy (mmr, similarity, hybrid).
+
+    Returns:
+        An initialized retriever instance.
+    """
+    try:
+        return get_retriever(
+            persist_dir=persist_dir,
+            top_k=top_k,
+            retrieval_mode=retrieval_mode,
+        )
+    except Exception:
+        run_ingest_cli(data_dir=data_dir, persist_dir=persist_dir)
+        return get_retriever(
+            persist_dir=persist_dir,
+            top_k=top_k,
+            retrieval_mode=retrieval_mode,
+        )
+
+
+@st.cache_resource(show_spinner=False)
+def get_chain_cached(
+        model_name: str,
+        top_k: int,
+        retrieval_mode: str,
+        data_dir: str,
+        persist_dir: str,
+) -> Any:
+    """Create or load a cached conversational QA chain.
+
+    Args:
+        model_name: The OpenAI model to use (gpt-3.5-turbo, gpt-4).
+        top_k: Number of chunks to retrieve.
+        retrieval_mode: Retrieval mode for the retriever.
+        data_dir: Path to data directory.
+        persist_dir: Path to vectorstore directory.
+
+    Returns:
+        A fully configured conversational QA chain.
+    """
+    retriever = build_or_load_retriever_cached(
+        data_dir=data_dir,
+        persist_dir=persist_dir,
+        top_k=top_k,
+        retrieval_mode=retrieval_mode,
+    )
+    return make_conversational_chain(retriever, model_name=model_name)