added a cache

app.py

@@ -1,22 +1,25 @@
-# app.py
-
 import gradio as gr
 import json
 from rag.rag_pipeline import RAGPipeline
 from utils.prompts import highlight_prompt, evidence_based_prompt
 from config import STUDY_FILES
 
+# Cache for RAG pipelines
+rag_cache = {}
 
-def load_rag_pipeline(study_name):
-    study_file = STUDY_FILES.get(study_name)
-    if study_file:
-        return RAGPipeline(study_file)
-    else:
-        raise ValueError(f"Invalid study name: {study_name}")
+
+def get_rag_pipeline(study_name):
+    if study_name not in rag_cache:
+        study_file = STUDY_FILES.get(study_name)
+        if study_file:
+            rag_cache[study_name] = RAGPipeline(study_file)
+        else:
+            raise ValueError(f"Invalid study name: {study_name}")
+    return rag_cache[study_name]
 
 
 def query_rag(study_name, question, prompt_type):
-    rag = load_rag_pipeline(study_name)
+    rag = get_rag_pipeline(study_name)
 
     if prompt_type == "Highlight":
         prompt = highlight_prompt

rag/rag_pipeline.py

@@ -4,59 +4,62 @@ import json
 from llama_index.core import Document, VectorStoreIndex
 from llama_index.core.node_parser import SentenceWindowNodeParser, SentenceSplitter
 from llama_index.core import PromptTemplate
-from typing import List
 
 
 class RAGPipeline:
     def __init__(self, study_json, use_semantic_splitter=False):
         self.study_json = study_json
-        self.index = None
         self.use_semantic_splitter = use_semantic_splitter
-        self.load_documents()
-        self.build_index()
+        self.documents = None
+        self.index = None
 
     def load_documents(self):
-        with open(self.study_json, "r") as f:
-            self.data = json.load(f)
+        if self.documents is None:
+            with open(self.study_json, "r") as f:
+                self.data = json.load(f)
 
-        self.documents = []
+            self.documents = []
 
-        for index, doc_data in enumerate(self.data):
-            doc_content = (
-                f"Title: {doc_data['title']}\n"
-                f"Authors: {', '.join(doc_data['authors'])}\n"
-                f"Full Text: {doc_data['full_text']}"
-            )
+            for index, doc_data in enumerate(self.data):
+                doc_content = (
+                    f"Title: {doc_data['title']}\n"
+                    f"Authors: {', '.join(doc_data['authors'])}\n"
+                    f"Full Text: {doc_data['full_text']}"
+                )
 
-            metadata = {
-                "title": doc_data.get("title"),
-                "abstract": doc_data.get("abstract"),
-                "authors": doc_data.get("authors", []),
-                "year": doc_data.get("year"),
-                "doi": doc_data.get("doi"),
-            }
+                metadata = {
+                    "title": doc_data.get("title"),
+                    "abstract": doc_data.get("abstract"),
+                    "authors": doc_data.get("authors", []),
+                    "year": doc_data.get("year"),
+                    "doi": doc_data.get("doi"),
+                }
 
-            self.documents.append(
-                Document(text=doc_content, id_=f"doc_{index}", metadata=metadata)
-            )
+                self.documents.append(
+                    Document(text=doc_content, id_=f"doc_{index}", metadata=metadata)
+                )
 
     def build_index(self):
-        sentence_splitter = SentenceSplitter(chunk_size=128, chunk_overlap=13)
+        if self.index is None:
+            self.load_documents()
+            sentence_splitter = SentenceSplitter(chunk_size=128, chunk_overlap=13)
 
-        def _split(text: str) -> List[str]:
-            return sentence_splitter.split_text(text)
+            def _split(text: str) -> List[str]:
+                return sentence_splitter.split_text(text)
 
-        node_parser = SentenceWindowNodeParser.from_defaults(
-            sentence_splitter=_split,
-            window_size=3,
-            window_metadata_key="window",
-            original_text_metadata_key="original_text",
-        )
+            node_parser = SentenceWindowNodeParser.from_defaults(
+                sentence_splitter=_split,
+                window_size=3,
+                window_metadata_key="window",
+                original_text_metadata_key="original_text",
+            )
 
-        nodes = node_parser.get_nodes_from_documents(self.documents)
-        self.index = VectorStoreIndex(nodes)
+            nodes = node_parser.get_nodes_from_documents(self.documents)
+            self.index = VectorStoreIndex(nodes)
 
     def query(self, question, prompt_template=None):
+        self.build_index()  # This will only build the index if it hasn't been built yet
+
         if prompt_template is None:
             prompt_template = PromptTemplate(
                 "Context information is below.\n"