reducing chunks retrieved and adding retrieved chunks to display

app.py

@@ -31,6 +31,7 @@ def rag_retrieve_tool(query):
     result = rag_retrieve(query, llm=llm)
     return {
         "rag_result": result.get("rag_result", ""),
+        "rag_sources": result.get("rag_sources", []),
         "last_tool": "rag_retrieve",
     }
 
@@ -164,7 +165,7 @@ def chat_with_patient(question: str, patient_id: str, sitecode: str, thread_id:
 
     assistant_message = output_state["messages"][-1].content
 
-    return assistant_message, thread_id
+    return assistant_message, thread_id, output_state.get("rag_sources", "")
 
 def init_session():
     return str(uuid.uuid4())
@@ -195,16 +196,19 @@ with gr.Blocks() as app:
             label="Sitecode",
         )
 
+    gr.Markdown("### Ask a Clinical Question")
     question_input = gr.Textbox(label="Question")
     thread_id_state = gr.State(init_session())
     output_chat = gr.Textbox(label="Assistant Response")
 
+    retrieved_sources_display = gr.HTML(label="Retrieved Sources (if applicable)")
+
     submit_btn = gr.Button("Ask")
 
     submit_btn.click(  # pylint: disable=no-member
         chat_with_patient,
         inputs=[question_input, id_selected, sitecode_selection, thread_id_state],
-        outputs=[output_chat, thread_id_state],
+        outputs=[output_chat, thread_id_state, retrieved_sources_display],
     )
 
 app.launch(

chatlib/assistant_node.py

@@ -39,6 +39,7 @@ def assistant(state: AppState, sys_msg, llm, llm_with_tools) -> AppState:
     state.setdefault("pk_hash", "")
     state.setdefault("sitecode", "")
     state.setdefault("rag_result", "")
+    state.setdefault("rag_sources", "")
     state.setdefault("answer", "")
     state.setdefault("last_answer", None)
     state.setdefault("last_user_message", None)
@@ -176,9 +177,14 @@ def assistant(state: AppState, sys_msg, llm, llm_with_tools) -> AppState:
     elif state.get("rag_result"):
         # Use conversation history + a system message to inject RAG guidance
         rag_msg = SystemMessage(
-            content="The following clinical guidelines may help answer the user's question:\n\n"
-            f"{state['rag_result']}\n\n"
-            "Use this information when responding."
+            content = (
+                "Based on the following clinical guideline excerpts, answer the clinician's question as precisely as possible.\n\n"
+                "Focus only on information that directly addresses the question.\n"
+                "Do not include background or general recommendations unless they are explicitly relevant.\n\n"
+                "Guideline excerpts:\n"
+                f"{state['rag_result']}\n\n"
+                "Respond with a focused summary tailored to the question about advanced HIV disease."
+            )
         )
         messages_with_rag = messages + [rag_msg]
         llm_response = llm.invoke(messages_with_rag)

chatlib/guidlines_rag_agent_li.py

@@ -18,7 +18,7 @@ embedding_model = OpenAIEmbedding()
 llm_llama = OpenAI(model="gpt-4o", temperature=0.0)
 
 # Create LLM reranker
-reranker = LLMRerank(llm=llm_llama, top_n=5)
+reranker = LLMRerank(llm=llm_llama, top_n=3)
 
 # Define a prompt template for query expansion
 query_expansion_prompt = ChatPromptTemplate.from_messages([
@@ -49,6 +49,20 @@ def cosine_similarity_numpy(query_vec: np.ndarray, matrix: np.ndarray) -> np.nda
     # Dot product gives cosine similarity
     return matrix_norm @ query_norm
 
+def format_sources_for_html(sources):
+    html_blocks = []
+    for i, source in enumerate(sources):
+        text = source.text.replace("\n", "<br>").strip()
+        block = f"""
+        <details style='margin-bottom: 1em;'>
+            <summary><strong>Source {i+1}</strong></summary>
+            <div style='margin-top: 0.5em; font-family: monospace;'>{text}</div>
+        </details>
+        """
+        html_blocks.append(block)
+    return "\n".join(html_blocks)
+
+
 def rag_retrieve(query: str, llm) -> AppState:
     """Perform RAG search of repository containing authoritative information on HIV/AIDS in Kenya."""
     
@@ -59,7 +73,7 @@ def rag_retrieve(query: str, llm) -> AppState:
     # Embed the expanded query and find similar summaries
     query_embedding = embedding_model.get_text_embedding(expanded_query)
     similarities = cosine_similarity_numpy(query_embedding, embeddings)
-    top_indices = similarities.argsort()[-5:][::-1]
+    top_indices = similarities.argsort()[-3:][::-1]
     selected_paths = df.loc[top_indices, "vectorestore_path"].tolist()
     print(f"Selected paths for retrieval: {selected_paths}")
 
@@ -81,6 +95,7 @@ def rag_retrieve(query: str, llm) -> AppState:
             "rag_result": "No relevant information found in the sources. Please try rephrasing your question.",
             "last_tool": "rag_retrieve"
         }
+    # Format the retrieved sources for the response (and remove lengthy white space or repeated dashes)
     retrieved_text = "\n\n".join([
         f"Source {i+1}: {source.text}" for i, source in enumerate(sources)
     ])
@@ -97,4 +112,7 @@ def rag_retrieve(query: str, llm) -> AppState:
     print("Prompt length in characters:", len(summarization_prompt))
     summary_response = llm.invoke(summarization_prompt)
 
-    return {"rag_result": summary_response.content, "last_tool": "rag_retrieve"}  # type: ignore
+    return {"rag_result": summary_response.content,
+            "rag_sources": format_sources_for_html(sources),
+            "last_tool": "rag_retrieve"
+        }  # type: ignore

chatlib/state_types.py

@@ -11,6 +11,7 @@ class AppState(TypedDict):
     pk_hash: str
     sitecode: str
     rag_result: str
+    rag_sources: Optional[str]  # Added to store retrieved sources
     answer: str
     last_answer: Optional[str] = None
     last_user_message: Optional[str] = None