""" server_ui.py - Multi-Agent Clinical Decision Support Chatbot Stack: Streamlit, LangGraph, LlamaIndex/ChromaDB, DeepSeek V4, CrossEncoder """ # ============================================================================= # 1. Imports & Initializations # ============================================================================= import os import time os.environ["CUDA_VISIBLE_DEVICES"] = "" import json import uuid import logging from datetime import date from typing import TypedDict from rich.logging import RichHandler from rich.console import Console from rich.theme import Theme console = Console(theme=Theme({ "logging.level.info": "bright_cyan", "logging.level.warning": "bright_yellow", "logging.level.error": "bright_red", "log.time": "bright_black", "log.message": "white", })) logging.basicConfig( level=logging.INFO, format="%(message)s", datefmt="%Y-%m-%d %H:%M:%S", handlers=[RichHandler( show_time=True, show_level=True, show_path=False, omit_repeated_times=False, console=console, )], ) for noisy in ("huggingface_hub", "sentence_transformers", "urllib3", "httpx"): logging.getLogger(noisy).setLevel(logging.WARNING) logger = logging.getLogger(__name__) import streamlit as st import openai from openai import OpenAI from langgraph.graph import StateGraph, END import chromadb from llama_index.core import VectorStoreIndex, StorageContext from llama_index.vector_stores.chroma import ChromaVectorStore from llama_index.embeddings.huggingface import HuggingFaceEmbedding from llama_index.core.vector_stores import ( MetadataFilter, MetadataFilters, FilterCondition, FilterOperator, ) from sentence_transformers import CrossEncoder import transformers transformers.logging.set_verbosity_error() CHROMA_DB_PATH = "./chroma_db" COLLECTION_NAME = "clinical_guidelines" EMBED_MODEL_NAME = "BAAI/bge-small-en-v1.5" RERANKER_MODEL_NAME = "BAAI/bge-reranker-base" DEEPSEEK_API_KEY = os.environ.get("DEEPSEEK_API_KEY", "") DEEPSEEK_BASE_URL = os.environ.get( "DEEPSEEK_BASE_URL", "https://api.deepseek.com/v1" ) DEEPSEEK_MODEL = os.environ.get("DEEPSEEK_MODEL", "deepseek-chat") CURRENT_YEAR = date.today().year AVAILABLE_DOCS = { "2026 Guideline for the Early Management of Patients With Acute Ischemic Stroke: A Guideline From the American Heart Association/American Stroke Association": [ 2026 ], "Guidelines for the Early Management of Patients With Acute Ischemic Stroke: 2019 Update to the 2018 Guidelines for the Early Management of Acute Ischemic Stroke: A Guideline for Healthcare Professionals From the American Heart Association/American Stroke Association": [ 2019 ], } DOC_URLS = { "2026 Guideline for the Early Management of Patients With Acute Ischemic Stroke: A Guideline From the American Heart Association/American Stroke Association": "https://www.ahajournals.org/doi/10.1161/STR.0000000000000513", "Guidelines for the Early Management of Patients With Acute Ischemic Stroke: 2019 Update to the 2018 Guidelines for the Early Management of Acute Ischemic Stroke: A Guideline for Healthcare Professionals From the American Heart Association/American Stroke Association": "https://www.ahajournals.org/doi/full/10.1161/STR.0000000000000211", } deepseek_client: OpenAI | None = None def call_llm( system_prompt: str, user_prompt: str, thinking: bool = False, temperature: float = 0.1, ) -> str: logger.info("[LLM] → Calling DeepSeek (thinking=%s, temperature=%s)", thinking, temperature) t0 = time.time() global deepseek_client if deepseek_client is None: if not DEEPSEEK_API_KEY: raise ValueError("API key not configured") deepseek_client = OpenAI( api_key=DEEPSEEK_API_KEY, base_url=DEEPSEEK_BASE_URL, ) messages = [ {"role": "system", "content": system_prompt}, {"role": "user", "content": user_prompt}, ] kwargs = { "model": DEEPSEEK_MODEL, "messages": messages, "temperature": temperature, } if thinking: kwargs["extra_body"] = {"thinking": {"type": "enabled"}, "reasoning_effort": "high"} response = deepseek_client.chat.completions.create(**kwargs) result = response.choices[0].message.content.strip() elapsed = time.time() - t0 logger.info("[LLM] ← Response in %.1fs (%d chars)", elapsed, len(result)) return result def format_conversation_history(history: list) -> str: if not history: return "" lines = ["## Previous Conversation"] for m in history: role_label = "User" if m["role"] == "user" else "Assistant" lines.append(f"{role_label}: {m['content']}") return "\n\n".join(lines) + "\n\n" @st.cache_resource def init_vector_index(): db = chromadb.PersistentClient(path=CHROMA_DB_PATH) chroma_collection = db.get_collection(COLLECTION_NAME) vector_store = ChromaVectorStore(chroma_collection=chroma_collection) embed_model = HuggingFaceEmbedding(model_name=EMBED_MODEL_NAME) return VectorStoreIndex.from_vector_store( vector_store, embed_model=embed_model ) @st.cache_resource def load_cross_encoder(): return CrossEncoder(RERANKER_MODEL_NAME) @st.cache_data def get_available_docs() -> list[tuple[str, int]]: db = chromadb.PersistentClient(path=CHROMA_DB_PATH) chroma_collection = db.get_collection(COLLECTION_NAME) result = chroma_collection.get(include=["metadatas"]) seen = set() docs = [] for meta in result["metadatas"]: key = (meta["doc_name"], meta["doc_year"]) if key not in seen: seen.add(key) docs.append((meta["doc_name"], meta["doc_year"])) return docs def query_chromadb( query_text: str, doc_name: str, doc_year: int, k: int = 5 ) -> list: logger.info("[DB] → query_chromadb: %.60s… [%s (%d)]", query_text, doc_name, doc_year) index = init_vector_index() filters = MetadataFilters( filters=[ MetadataFilter( key="doc_name", value=doc_name, operator=FilterOperator.EQ, ), MetadataFilter( key="doc_year", value=doc_year, operator=FilterOperator.EQ, ), ], condition=FilterCondition.AND, ) retriever = index.as_retriever(similarity_top_k=k, filters=filters) results = [] for node in retriever.retrieve(query_text): results.append( { "text": node.text, "doc_name": node.metadata.get("doc_name", ""), "doc_year": node.metadata.get("doc_year", 0), "score": getattr(node, "score", 0.0), } ) logger.info("[DB] ← %d result(s)", len(results)) return results # ============================================================================= # 2. LangGraph State & Nodes # ============================================================================= # # The graph has 5 nodes executing in this order: # # guard_agent (triages query relevance) # -> [if relevant] retrieval_agent (decomposes query, searches ChromaDB) # -> re_ranker (CrossEncoder, keeps top 3) # -> main_agent (synthesizes final answer) # -> [if irrelevant] end_inappropriate (returns out-of-scope message) # class GraphState(TypedDict): """State object passed between LangGraph nodes at each step. - user_query: The original question from the doctor. - is_appropriate: Set by GuardAgent — whether the query is clinically relevant. - fetched_chunks: Accumulated raw chunks from ChromaDB; later filtered by re-ranker. - final_response: The final answer produced by MainAgent (or the "can't answer" message). - debug_system_prompt: System prompt sent to MainAgent (for debug dialog). - debug_user_prompt: User prompt sent to MainAgent (for debug dialog). - retrieval_system_prompt: System prompt sent to RetrievalAgent (for debug). - retrieval_user_prompt: User prompt sent to RetrievalAgent (for debug). - conversation_history: Prior turns (list of {"role", "content"}) for context. - top_k_retrieval: Number of chunks to fetch from ChromaDB per sub-query. - top_k_rerank: Number of best chunks to keep after re-ranking. - temperature: LLM temperature for main_agent synthesis. """ user_query: str is_appropriate: bool fetched_chunks: list final_response: str debug_system_prompt: str debug_user_prompt: str retrieval_system_prompt: str retrieval_user_prompt: str retrieval_raw_response: str conversation_history: list top_k_retrieval: int top_k_rerank: int temperature: float def main_agent(state: GraphState) -> dict: """Synthesis node. Builds final answer from re-ranked chunks. Passes through if a final response was already set upstream.""" logger.info("[AGENT] main_agent → start") if state.get("final_response"): logger.info("[AGENT] main_agent ← skipped (final_response already set)") return {} chunks_text = "\n\n---\n\n".join( [ f"[Document: {c['doc_name']} ({c['doc_year']})]\n{c['text']}" for c in state["fetched_chunks"] ] ) system_prompt = ( "You are a clinical decision support AI assistant for doctors. " "Answer **exclusively** from the retrieved guideline chunks below. " "Do not use your own medical knowledge or training data — " "only the information provided in the chunks.\n\n" "If the chunks do not contain enough information to fully answer " "the question, clearly state what is missing rather than making up an answer.\n\n" "Each chunk is structured as:\n" " Document: (use this for citation)\n" " Section: (use this for section reference)\n" " \n\n" "Cite the specific guideline year and relevant section.\n\n" "Be concise by default — doctors need key facts quickly. " "However, if the user's query explicitly asks for more detail, " "elaboration, or a comprehensive explanation, provide a thorough answer." ) history_text = format_conversation_history(state.get("conversation_history", [])) user_prompt = ( f"{history_text}" f"## User Query\n\n{state['user_query']}\n\n" f"## Retrieved Guideline Chunks\n\n{chunks_text}\n\n" "Based on the above retrieved information, provide a medical answer for the doctor." ) # Thinking mode ON for careful synthesis temp = state.get("temperature", 0.1) response = call_llm(system_prompt, user_prompt, thinking=True, temperature=temp) logger.info("[AGENT] main_agent ← done (%d chars)", len(response)) return { "final_response": response, "debug_system_prompt": system_prompt, "debug_user_prompt": user_prompt, } def guard_agent(state: GraphState) -> dict: """Guardrail node. Uses a fast LLM call (thinking OFF) to decide if the query belongs to clinical medicine. Sets `is_appropriate` in state.""" logger.info("[AGENT] guard_agent → start") system_prompt = ( "You are a medical triage guard for a clinical decision support " "system used by doctors. Determine if a user query is relevant to " "clinical medicine, patient management, or medical guidelines.\n\n" "Ignore years, document names, or specific guideline versions — " "that is checked downstream. Only assess medical relevance.\n\n" 'Respond with ONLY "YES" if clinically relevant, or "NO" if not. ' "YES: blood pressure, medication, patient care, " "symptoms, diagnosis, treatment, rehabilitation, imaging.\n" "NO: sports, politics, entertainment, programming, weather, " "history, geography, general knowledge." ) history_text = format_conversation_history(state.get("conversation_history", [])) user_prompt = ( f"{history_text}" f"Query: {state['user_query']}\n\n" "Is this query relevant to clinical medicine? Answer YES or NO only." ) # Thinking mode OFF = minimum latency for this simple classification response = call_llm( system_prompt, user_prompt, thinking=False, temperature=0.0 ) is_appropriate = response.strip().upper().startswith("YES") logger.info("[AGENT] guard_agent ← done (appropriate=%s)", is_appropriate) return {"is_appropriate": is_appropriate} def end_inappropriate(state: GraphState) -> dict: """Terminal node for out-of-scope queries. Sets a polite refusal message.""" logger.info("[AGENT] end_inappropriate → start") logger.info("[AGENT] end_inappropriate ← done") return { "final_response": ( "I can't answer this type of question. Please ask a " "clinical/medical question related to the available healthcare " "guidelines." ) } def retrieval_agent(state: GraphState) -> dict: """Query-decomposition node. LLM (thinking ON) analyses the user query and outputs a JSON list of sub-queries. Each sub-query specifies keywords, doc_name, and doc_year for a `query_chromadb` call. Handles temporal reasoning: explicit years, relative terms like "latest", or defaults to the newest version, and validates that every sub-query targets an available document. """ logger.info("[AGENT] retrieval_agent → start") docs_info = "\n".join( [f'- "{name}" (years: {years})' for name, years in AVAILABLE_DOCS.items()] ) system_prompt = ( "You are a medical retrieval specialist for healthcare guidelines. " "Conversation history is provided below for context. " "Your ONLY task is to decompose the LATEST user query into vector " "database searches. Ignore the assistant responses in history.\n\n" f"Current year: {CURRENT_YEAR}\n\n" f"Available documents:\n{docs_info}\n\n" "Output a JSON object with a 'sub_queries' array. Each item:\n" '- "query_text": clinical keywords for semantic search\n' '- "doc_name": exact document name string\n' '- "doc_year": integer year\n\n' "Rules:\n" "1. If user specifies years (e.g., '2019 vs 2026'), create separate " "sub_queries for each year.\n" f"2. If user says 'latest'/'current', use {CURRENT_YEAR}.\n" f"3. If no year specified, use the latest ({CURRENT_YEAR}).\n" "4. Extract concise clinical keywords optimized for vector search.\n" "5. If the query does not match any available document, return " '{"sub_queries": []}.\n' "6. If the user specifies a year that is not in the available years " "for a document, return {\"sub_queries\": []} — do not substitute " "a different year." ) history_text = format_conversation_history(state.get("conversation_history", [])) user_prompt = ( f"{history_text}" f"## Latest User Query\n\n{state['user_query']}\n\n" "Based on the conversation history and the latest user query above, " "output a JSON object with a 'sub_queries' array for the latest query." ) # Thinking mode ON ensures accurate extraction of search arguments response = call_llm(system_prompt, user_prompt, thinking=True, temperature=0.0) def extract_json(text: str) -> dict | None: """Extract the first JSON object from text, even if surrounded by conversational wrapper text.""" text = text.strip() # Strip markdown code fences if "```json" in text: text = text.split("```json")[1].split("```")[0].strip() elif "```" in text: text = text.split("```")[1].split("```")[0].strip() # Find the outermost { … } start = text.find("{") end = text.rfind("}") if start == -1 or end == -1 or end <= start: return None try: return json.loads(text[start : end + 1]) except json.JSONDecodeError: return None parsed = extract_json(response) try: if parsed is None: raise ValueError sub_queries = parsed.get("sub_queries", []) if not isinstance(sub_queries, list) or not sub_queries: raise ValueError except (ValueError, TypeError, AttributeError): logger.info("[AGENT] retrieval_agent ← no sub_queries could be extracted") return { "final_response": "Sorry, I am unable to utilize my knowledge base.", "fetched_chunks": [], "retrieval_system_prompt": system_prompt, "retrieval_user_prompt": user_prompt, "retrieval_raw_response": response, } # Validate every sub-query against the available documents fallback_response = { "final_response": "Sorry, that is outside my current knowledge base.", "fetched_chunks": [], "retrieval_system_prompt": system_prompt, "retrieval_user_prompt": user_prompt, "retrieval_raw_response": response, } try: for sq in sub_queries: # If sq isn't a dict, .get() raises an AttributeError doc_name = sq.get("doc_name", "") doc_year = sq.get("doc_year", 0) # Just to check if query_text exists. query_text = sq.get("query_text") # If doc_name or doc_year don't exist, it raises a KeyError if doc_year not in AVAILABLE_DOCS[doc_name]: logger.info( "[AGENT] retrieval_agent ← year %d not available for doc '%s'", doc_year, doc_name, ) return fallback_response except (AttributeError, KeyError): logger.info("[AGENT] retrieval_agent ← sub_query validation failed (unknown doc/field)") return fallback_response all_chunks = [] for sq_idx, sq in enumerate(sub_queries): doc_name = sq.get("doc_name") doc_year = sq.get("doc_year") query_text = sq.get("query_text") results = query_chromadb(query_text, doc_name, doc_year, k=state.get("top_k_retrieval", 5)) for r in results: r["sq_idx"] = sq_idx all_chunks.append(r) logger.info("[AGENT] retrieval_agent ← done (%d sub_queries, %d chunks)", len(sub_queries), len(all_chunks)) return { "fetched_chunks": all_chunks, "retrieval_system_prompt": system_prompt, "retrieval_user_prompt": user_prompt, "retrieval_raw_response": response, } def re_ranker(state: GraphState) -> dict: """Pure-Python re-ranker node (no LLM call). Uses a CrossEncoder to score every fetched chunk against the original query, keeps only the top 3 most relevant results. Passes through if a final response was already set upstream.""" logger.info("[AGENT] re_ranker → start (input: %d chunks)", len(state.get("fetched_chunks", []))) if state.get("final_response"): logger.info("[AGENT] re_ranker ← skipped (final_response already set)") return {} cross_encoder = load_cross_encoder() grouped: dict[int, list] = {} for chunk in state["fetched_chunks"]: grouped.setdefault(chunk["sq_idx"], []).append(chunk) selected = [] for group in grouped.values(): passages = [c["text"] for c in group] if not passages: continue scores = cross_encoder.predict([(state["user_query"], p) for p in passages]) scored = list(zip(scores, group)) scored.sort(key=lambda x: x[0], reverse=True) k_rerank = state.get("top_k_rerank", 3) selected.extend(c for _, c in scored[:k_rerank]) logger.info("[AGENT] re_ranker ← done (%d chunks after filtering)", len(selected)) return {"fetched_chunks": selected} # ============================================================================= # 3. Workflow Compilation # ============================================================================= # # LangGraph flow: # # START # | # v # [guard_agent] # | # ├──[if relevant]──→ [retrieval_agent] → [re_ranker] → [main_agent] → END # └──[if irrelevant]─→ [end_inappropriate] → END @st.cache_resource def build_graph(): """Build and compile the LangGraph state machine. Cached so the graph topology is only constructed once.""" builder = StateGraph(GraphState) # Register all 5 nodes builder.add_node("main_agent", main_agent) builder.add_node("guard_agent", guard_agent) builder.add_node("end_inappropriate", end_inappropriate) builder.add_node("retrieval_agent", retrieval_agent) builder.add_node("re_ranker", re_ranker) # First node to execute builder.set_entry_point("guard_agent") # -- guard_agent routes -- # If query is clinically relevant → proceed to retrieval_agent. # Otherwise → skip retrieval and go straight to end_inappropriate. def route_guard(state): return "retrieval_agent" if state["is_appropriate"] else "end_inappropriate" builder.add_conditional_edges( "guard_agent", route_guard, { "retrieval_agent": "retrieval_agent", "end_inappropriate": "end_inappropriate", }, ) # -- deterministic edges -- builder.add_edge("retrieval_agent", "re_ranker") # retrieve → re-rank builder.add_edge("re_ranker", "main_agent") # re-rank → synthesise builder.add_edge("main_agent", END) # synthesis → end builder.add_edge("end_inappropriate", END) # terminal node return builder.compile() # ============================================================================= # 4. Streamlit Application # ============================================================================= def init_session_state(): if "thread_id" not in st.session_state: st.session_state.thread_id = str(uuid.uuid4()) if "messages" not in st.session_state: st.session_state.messages = [] if "processing" not in st.session_state: st.session_state.processing = False if "debug_transcript" not in st.session_state: st.session_state.debug_transcript = [] if "quota_exceeded" not in st.session_state: st.session_state.quota_exceeded = False if "_quota_error_detected" not in st.session_state: st.session_state._quota_error_detected = False if "_text_to_copy" not in st.session_state: st.session_state._text_to_copy = "" if "top_k_retrieval" not in st.session_state: st.session_state.top_k_retrieval = 16 if "top_k_rerank" not in st.session_state: st.session_state.top_k_rerank = 8 if "temperature" not in st.session_state: st.session_state.temperature = 0.1 def _request_copy(text: str): st.session_state._text_to_copy = text def _copy_conversation(): lines = [] for msg in st.session_state.messages: role = "You" if msg["role"] == "user" else "Assistant" lines.append(f"{role}: {msg['content']}") _request_copy("\n\n---\n\n".join(lines)) def _render_copy_script(): text = st.session_state.pop("_text_to_copy", "") if not text.strip(): return import html srcdoc = f"""""" st.markdown( f'', unsafe_allow_html=True, ) st.toast("Copied!", icon="📋") def reset_session(): prefs = { "top_k_retrieval": st.session_state.get("top_k_retrieval", 16), "top_k_rerank": st.session_state.get("top_k_rerank", 8), "temperature": st.session_state.get("temperature", 0.1), } st.session_state.thread_id = str(uuid.uuid4()) st.session_state.messages = [] st.session_state.processing = False st.session_state.debug_transcript = [] st.session_state.quota_exceeded = False st.session_state._quota_error_detected = False st.session_state.top_k_retrieval = prefs["top_k_retrieval"] st.session_state.top_k_rerank = prefs["top_k_rerank"] st.session_state.temperature = prefs["temperature"] @st.dialog("📜 Debug Transcript", width="large") def show_debug_transcript(): for entry in st.session_state.debug_transcript: st.markdown(f"### Turn {entry['turn']}: {entry['user_query']}") guard_icon = "✅" if entry["guard_decision"] else "❌" st.markdown(f"**Guard:** {guard_icon} {'Relevant' if entry['guard_decision'] else 'Irrelevant'} | **Chunks:** {entry['num_chunks']}") st.markdown("**Retrieval Agent — System Prompt:**") st.code(entry["retrieval_system_prompt"]) st.markdown("**Retrieval Agent — User Prompt:**") st.code(entry["retrieval_user_prompt"]) st.markdown("**Retrieval Agent — Decomposed Sub-queries:**") st.code(entry["retrieval_raw_response"]) st.markdown("**Main Agent — System Prompt:**") st.code(entry["main_system_prompt"]) st.markdown("**Main Agent — User Prompt:**") st.code(entry["main_user_prompt"]) st.markdown(f"**Response:** {entry['final_response']}") st.divider() if st.button("Close"): st.rerun() @st.dialog("🏗️ Architecture", width="large") def show_architecture(): left, right = st.columns(2) with left: try: st.image("assets/architecture.png", width=350) except Exception: st.caption("Unable to render graph diagram.") with right: st.markdown("### ChromaDB Schema") st.markdown( "| Column | Type | Description |\n" "|---|---|---|\n" "| `id` | string | Auto-generated UUID |\n" "| `embedding` | `List[float]` | 384-dim vector from `BAAI/bge-small-en-v1.5` |\n" "| `document` | string | Enriched chunk text |\n" "| `metadata` | dict | `{\"doc_year\": int, \"doc_name\": string}` |\n" ) if st.button("Close"): st.rerun() SAMPLES = [ { "desc": "Can't answer an out-of-scope question (guard fails)", "question": "Why did President Trump put tariff on everyone?", }, { "desc": "Can't answer a clinically relevant question if it doesn't match available docs (retrieval fails)", "question": "Tell me about Antiplatelet Treatment for patients with minor Acute Ischemic Stroke in Guidelines 2015.", }, { "desc": "Answer using one call to ChromaDB (vector database), filtering with doc_year=2019", "question": "Tell me about Antiplatelet Treatment for patients with minor Acute Ischemic Stroke in Guidelines 2019.", }, { "desc": "Test follow-up question: 'Does the treatment have side effects?'", "question": "Tell me about Antiplatelet Treatment for patients with minor Acute Ischemic Stroke.", }, { "desc": "Compare treatment between two medical documents (2019, 2026), filtering with doc_year=2019 and then doc_year=2026", "question": "How did blood pressure management recommendations change after thrombolysis or thrombectomy between 2026 and 2019 guidelines?", }, ] @st.dialog("❓ Help & Examples", width="large") def show_help_dialog(): st.markdown( "Each sample below demonstrates a different type of question " "the system can answer. Click **Use this question** to try it." ) st.markdown( "📖 For full documentation, visit the " "[README](https://huggingface.co/spaces/tnt306/agentic_retrieval/blob/main/README.md)." ) st.markdown( "", unsafe_allow_html=True, ) st.markdown('
', unsafe_allow_html=True) for sample in SAMPLES: st.markdown( f"**{sample['question']}**
" f'👉 ' f'{sample["desc"]}.', unsafe_allow_html=True, ) if st.button("Use this question", key=f"sample_{hash(sample['question'])}"): reset_session() st.session_state.chat_input = sample["question"] st.rerun() st.markdown('
', unsafe_allow_html=True) if st.button("Close"): st.rerun() def main(): st.set_page_config( page_title="Clinical Decision Support - Healthcare Guidelines", page_icon="⚡", layout="wide", ) st.markdown( "", unsafe_allow_html=True, ) init_session_state() with st.sidebar: st.title("⚡ Agentic Retrieval") st.markdown("**Clinical Decision Support System**") st.divider() if st.button("🗑️ New Session", use_container_width=True, type="primary"): reset_session() st.rerun() col1, col2 = st.columns(2) with col1: if st.button("📜 Debug", use_container_width=True, disabled=len(st.session_state.debug_transcript) == 0): show_debug_transcript() with col2: if st.button("📋 Copy", use_container_width=True): _copy_conversation() col3, col4 = st.columns(2) with col3: if st.button("🏗️ Architecture", use_container_width=True): show_architecture() with col4: if st.button("❓ Examples", use_container_width=True): show_help_dialog() st.checkbox( "⚠️ Simulate quota exceeded", value=st.session_state.get("quota_exceeded", False), key="quota_exceeded", help="Show the quota-exceeded warning without actually hitting the API.", ) with st.expander("⚙️ Settings"): st.number_input("Chunks to retrieve", 1, 20, key="top_k_retrieval") st.number_input("Chunks after re-ranking", min_value=1, max_value=st.session_state.top_k_retrieval, key="top_k_rerank") st.slider("LLM temperature", 0.0, 1.0, key="temperature", step=0.05) if not DEEPSEEK_API_KEY: st.warning( "⚠️ DEEPSEEK_API_KEY not set. Set it in your environment.", icon="⚠️", ) st.markdown("### 🏥 Healthcare Guideline Assistant") st.caption( "Multi-Agent RAG | LangGraph · ChromaDB · BGE Reranker · DeepSeek V4" ) with st.expander("📚 Medical Documents Covered"): docs = get_available_docs() for doc_name, doc_year in sorted(docs, key=lambda x: -x[1]): url = DOC_URLS.get(doc_name, "") if url: st.markdown(f"- [**{doc_name}** ({doc_year})]({url})") else: st.markdown(f"- **{doc_name}** ({doc_year})") for i, msg in enumerate(st.session_state.messages): with st.chat_message(msg["role"]): st.markdown(msg["content"]) if msg["role"] == "assistant": label = f"📋 Copy ⏱ {msg['response_time']:.1f}s" if msg.get("response_time") else "📋 Copy" st.button(label, key=f"copy_{i}", on_click=_request_copy, args=(msg["content"],)) prompt = st.chat_input( "Ask a clinical question about healthcare guidelines...", key="chat_input", disabled=st.session_state.processing or st.session_state.get("quota_exceeded") or st.session_state.get("_quota_error_detected"), ) if not DEEPSEEK_API_KEY: st.markdown( '
' '' "⚠️ The API key is not configured." "
", unsafe_allow_html=True, ) elif st.session_state.get("quota_exceeded") or st.session_state.get("_quota_error_detected"): st.markdown( '
' '' "⚠️ The LLM is currently unavailable due to quota limits." "
", unsafe_allow_html=True, ) if prompt: st.session_state.messages.append({"role": "user", "content": prompt}) with st.chat_message("user"): st.markdown(prompt) st.session_state.processing = True t_start = time.time() graph = build_graph() initial_state: GraphState = { "user_query": prompt, "conversation_history": [ {"role": m["role"], "content": m["content"]} for m in st.session_state.messages[:-1] ], "is_appropriate": False, "fetched_chunks": [], "final_response": "", "debug_system_prompt": "", "debug_user_prompt": "", "retrieval_system_prompt": "", "retrieval_user_prompt": "", "retrieval_raw_response": "", "top_k_retrieval": st.session_state.top_k_retrieval, "top_k_rerank": st.session_state.top_k_rerank, "temperature": st.session_state.temperature, } with st.chat_message("assistant"): response = None error_msg = None with st.status("Starting...", expanded=False) as status: try: final_state = {} for step in graph.stream(initial_state): for node_name, node_data in step.items(): if node_data is None: continue if node_name == "guard_agent": is_appropriate = node_data.get("is_appropriate", False) if is_appropriate: status.update(label="✅ Query is clinically relevant", state="complete") else: status.update(label="❌ Query is out of scope", state="error") elif node_name == "retrieval_agent": chunks = node_data.get("fetched_chunks", []) status.update(label=f"🔍 Retrieved {len(chunks)} guideline chunks", state="complete") elif node_name == "re_ranker": chunks = node_data.get("fetched_chunks", []) status.update(label=f"🎯 Selected top {len(chunks)} most relevant chunks", state="complete") elif node_name == "main_agent": status.update(label="🧠 Synthesizing clinical answer...", state="running") elif node_name == "end_inappropriate": status.update(label="❌ Cannot answer this query", state="error") final_state.update(node_data) st.session_state.debug_transcript.append({ "turn": len(st.session_state.debug_transcript) + 1, "user_query": prompt, "guard_decision": final_state.get("is_appropriate"), "num_chunks": len(final_state.get("fetched_chunks", [])), "main_system_prompt": final_state.get("debug_system_prompt", ""), "main_user_prompt": final_state.get("debug_user_prompt", ""), "retrieval_system_prompt": final_state.get("retrieval_system_prompt", ""), "retrieval_user_prompt": final_state.get("retrieval_user_prompt", ""), "retrieval_raw_response": final_state.get("retrieval_raw_response", ""), "final_response": final_state.get("final_response", ""), }) response = final_state.get( "final_response", "I'm sorry, I couldn't process that query.", ) status.update(label="✅ Complete", state="complete") except openai.RateLimitError: st.session_state._quota_error_detected = True error_msg = "The LLM is currently unavailable due to quota limits." status.update(label="❌ Quota exceeded", state="error") except Exception as e: error_msg = ( "I encountered an error processing your request: " f"{str(e)}" ) status.update(label="❌ Error", state="error") elapsed = time.time() - t_start if response: def word_stream(text): words = text.split(" ") for i, word in enumerate(words): yield word + (" " if i < len(words) - 1 else "") time.sleep(0.02) st.write_stream(word_stream(response)) st.session_state.messages.append( {"role": "assistant", "content": response, "response_time": elapsed} ) elif error_msg: st.error(error_msg) st.session_state.messages.append( {"role": "assistant", "content": error_msg} ) if response: st.button(f"📋 Copy ⏱ {elapsed:.1f}s", key="copy_new", on_click=_request_copy, args=(response,)) st.session_state.processing = False st.rerun() _render_copy_script() if __name__ == "__main__": main()