minor

app.py

@@ -1,5 +1,3 @@
-# app.py
-
 import gradio as gr
 import json
 from rag.rag_pipeline import RAGPipeline
@@ -31,12 +29,6 @@ def query_rag(study_name: str, question: str, prompt_type: str) -> str:
     # Extract study information using RAG
     study_info = rag.extract_study_info()
 
-    # Prepare a dictionary with all possible prompt parameters
-    prompt_params = {
-        **study_info,  # Unpack the extracted study info
-        "query_str": question,  # Add the question to the prompt parameters
-    }
-
     if prompt_type == "Highlight":
         prompt = highlight_prompt
     elif prompt_type == "Evidence-based":
@@ -48,8 +40,17 @@ def query_rag(study_name: str, question: str, prompt_type: str) -> str:
     else:
         prompt = None
 
-    # Use the prompt_params in the query
-    response = rag.query(question, prompt, **prompt_params)
+    # Prepare the context with study info
+    context = "Study Information:\n"
+    for key, value in study_info.items():
+        context += f"{key}: {value}\n"
+    context += "\n"
+
+    # Add the question to the context
+    context += f"Question: {question}\n"
+
+    # Use the prepared context in the query
+    response = rag.query(context, prompt_template=prompt)
 
     # Format the response as Markdown
     formatted_response = f"## Question\n\n{question}\n\n## Answer\n\n{response['answer']}\n\n## Sources\n\n"

rag/rag_pipeline.py

@@ -1,5 +1,3 @@
-# rag/rag_pipeline.py
-
 import json
 from typing import Dict, Any
 from llama_index.core import Document, VectorStoreIndex
@@ -94,7 +92,7 @@ class RAGPipeline:
         return extracted_info
 
     def query(
-        self, question: str, prompt_template: PromptTemplate = None, **kwargs
+        self, context: str, prompt_template: PromptTemplate = None
     ) -> Dict[str, Any]:
         if prompt_template is None:
             prompt_template = PromptTemplate(
@@ -102,7 +100,7 @@ class RAGPipeline:
                 "---------------------\n"
                 "{context_str}\n"
                 "---------------------\n"
-                "Given this information, please answer the question: {query_str}\n"
+                "Given this information, please answer the question provided in the context. "
                 "Include all relevant information from the provided context. "
                 "If information comes from multiple sources, please mention all of them. "
                 "If the information is not available in the context, please state that clearly. "
@@ -113,11 +111,9 @@ class RAGPipeline:
             text_qa_template=prompt_template, similarity_top_k=5
         )
 
-        # Use kwargs to pass additional parameters to the query
-        response = query_engine.query(question, **kwargs)
+        response = query_engine.query(context)
 
         return {
-            "question": question,
             "answer": response.response,
             "sources": [node.metadata for node in response.source_nodes],
         }