update

app.py

@@ -457,18 +457,28 @@ def process_query(gen_config: dict, ret_config: dict, query: str):
     return response
 
 
-def display_retrieved_chunks(response):
+def display_retrieved_chunks(docs: list[Document], tag: str|None=None):
     with st.container(border=True):
-        doc_grps = group_docs(response["docs"])
-        st.write(
-            "Retrieved Chunks (note that you may need to 'right click' on links in the expanders to follow them)"
-        )
+        doc_grps = group_docs(docs)
+        if tag is None:
+            st.write(
+                "Retrieved Chunks\n\nleft click to expand, right click to follow links"
+            )
+        else:
+            st.write(
+                f"Retrieved Chunks ({tag})\n\nleft click to expand, right click to follow links"
+            )
         for legis_id, doc_grp in doc_grps:
             render_doc_grp(legis_id, doc_grp)
 
 
 def display_response(
-    response, model_info, provider, should_escape_markdown, should_add_legis_urls
+    response,
+    model_info: dict,
+    provider: str,
+    should_escape_markdown: bool,
+    should_add_legis_urls: bool,
+    tag: str|None=None
 ):
     out_display = response["aimessage"].content
     if should_escape_markdown:
@@ -477,11 +487,14 @@ def display_response(
         out_display = replace_legis_ids_with_urls(out_display)
 
     with st.container(border=True):
-        st.write("Response")
+        if tag is None:
+            st.write("Response")
+        else:
+            st.write(f"Response ({tag})")
         st.info(out_display)
 
-    usage.display_api_usage(response, model_info, provider)
-    display_retrieved_chunks(response)
+    usage.display_api_usage(response, model_info, provider, tag=tag)
+    display_retrieved_chunks(response["docs"], tag=tag)
 
 
 def render_query_rag_tab():
@@ -588,6 +601,7 @@ def render_query_rag_sbs_tab():
                     gen_configs[post_key_prefix]["provider"],
                     gen_configs[post_key_prefix]["should_escape_markdown"],
                     gen_configs[post_key_prefix]["should_add_legis_urls"],
+                    tag = grp_names[post_key_prefix],
                 )
 
 

usage.py

@@ -55,9 +55,12 @@ def get_token_usage(metadata: dict, model_info: dict, provider: str):
             raise ValueError()
 
 
-def display_api_usage(response, model_info, provider: str):
+def display_api_usage(response, model_info, provider: str, tag: str|None=None):
     with st.container(border=True):
-        st.write("API Usage")
+        if tag is None:
+            st.write("API Usage")
+        else:
+            st.write(f"API Usage ({tag})")
         token_usage = get_token_usage(
             response["aimessage"].response_metadata, model_info, provider
         )