Factor out LLM chat rendering so that it persists even when the submit button isn't active.

app.py

@@ -516,6 +516,206 @@ def get_selected_models_to_streamlit_column_map(st_columns, selected_models):
     return selected_models_to_streamlit_column_map
 
 
+def get_aggregator_key(llm_aggregator):
+    return "agg__" + llm_aggregator
+
+
+def st_render_responses(user_prompt):
+    """Renders the responses from the LLMs.
+
+    Uses cached responses from the session state, if available.
+    Otherwise, streams the responses anew.
+
+    Assumes that the session state has already been set up with selected models and selected aggregator.
+    """
+    st.markdown("#### Responses")
+
+    response_columns = st.columns(3)
+    selected_models_to_streamlit_column_map = (
+        get_selected_models_to_streamlit_column_map(
+            response_columns, st.session_state.selected_models
+        )
+    )
+    for response_model in st.session_state.selected_models:
+        st_column = selected_models_to_streamlit_column_map.get(
+            response_model, response_columns[0]
+        )
+
+        with st_column.chat_message(
+            response_model,
+            avatar=get_llm_avatar(response_model),
+        ):
+            st.write(get_ui_friendly_name(response_model))
+            if response_model in st.session_state.responses:
+                # Use the cached response from session state.
+                st.write(st.session_state.responses[response_model])
+            else:
+                # Stream the response from the LLM.
+                message_placeholder = st.empty()
+                stream = get_llm_response_stream(response_model, user_prompt)
+                st.session_state.responses[response_model] = (
+                    message_placeholder.write_stream(stream)
+                )
+
+    # Render the aggregator response.
+    aggregator_prompt = get_default_aggregator_prompt(
+        user_prompt=user_prompt, llms=st.session_state.selected_models
+    )
+
+    # Streaming response from the aggregator.
+    with st.chat_message(
+        get_aggregator_key(st.session_state.selected_aggregator),
+        avatar="img/council_icon.png",
+    ):
+        st.write(
+            f"{get_ui_friendly_name(get_aggregator_key(st.session_state.selected_aggregator))}"
+        )
+        if (
+            get_aggregator_key(st.session_state.selected_aggregator)
+            in st.session_state.responses
+        ):
+            st.write(
+                st.session_state.responses[
+                    get_aggregator_key(st.session_state.selected_aggregator)
+                ]
+            )
+        else:
+            message_placeholder = st.empty()
+            aggregator_stream = get_llm_response_stream(
+                selected_aggregator, aggregator_prompt
+            )
+            if aggregator_stream:
+                st.session_state.responses[get_aggregator_key(selected_aggregator)] = (
+                    message_placeholder.write_stream(aggregator_stream)
+                )
+
+    st.session_state.responses_collected = True
+
+
+def st_direct_assessment_results(user_prompt, direct_assessment_prompt, criteria_list):
+    """Renders the direct assessment results block.
+
+    Uses session state to render results from LLMs. If the session state isn't set, then fetches the
+    responses from the LLMs services from scratch (and sets the session state).
+
+    Assumes that the session state has already been set up with responses.
+    """
+    responses_for_judging = st.session_state.responses
+
+    # Get judging responses.
+    response_judging_columns = st.columns(3)
+    responses_for_judging_to_streamlit_column_map = (
+        get_selected_models_to_streamlit_column_map(
+            response_judging_columns, responses_for_judging.keys()
+        )
+    )
+
+    for response_model, response in responses_for_judging.items():
+        st_column = responses_for_judging_to_streamlit_column_map[response_model]
+
+        with st_column:
+            st.write(f"Judging for {get_ui_friendly_name(response_model)}")
+            judging_prompt = get_direct_assessment_prompt(
+                direct_assessment_prompt=direct_assessment_prompt,
+                user_prompt=user_prompt,
+                response=response,
+                criteria_list=criteria_list,
+                options=SEVEN_POINT_DIRECT_ASSESSMENT_OPTIONS,
+            )
+
+            with st.expander("Final Judging Prompt"):
+                st.code(judging_prompt)
+
+            for judging_model in st.session_state.selected_models:
+                with st.expander(get_ui_friendly_name(judging_model), expanded=True):
+                    with st.chat_message(
+                        judging_model,
+                        avatar=PROVIDER_TO_AVATAR_MAP[judging_model],
+                    ):
+                        if (
+                            judging_model
+                            in st.session_state.direct_assessment_judging_responses[
+                                response_model
+                            ]
+                        ):
+                            # Use the session state cached response.
+                            st.write(
+                                st.session_state.direct_assessment_judging_responses[
+                                    response_model
+                                ][judging_model]
+                            )
+                        else:
+                            message_placeholder = st.empty()
+                            # Get the judging response from the LLM.
+                            judging_stream = get_llm_response_stream(
+                                judging_model, judging_prompt
+                            )
+                            st.session_state.direct_assessment_judging_responses[
+                                response_model
+                            ][judging_model] = message_placeholder.write_stream(
+                                judging_stream
+                            )
+
+            # Extract actual scores from open-ended responses using structured outputs.
+            # Since we're extracting structured data for the first time, we can save the dataframe
+            # to the session state so that it's cached.
+            if response_model not in st.session_state.direct_assessment_judging_df:
+                judging_responses = (
+                    st.session_state.direct_assessment_judging_responses[response_model]
+                )
+                parse_judging_response_prompt = (
+                    get_parse_judging_response_for_direct_assessment_prompt(
+                        judging_responses,
+                        criteria_list,
+                        SEVEN_POINT_DIRECT_ASSESSMENT_OPTIONS,
+                    )
+                )
+                parsed_judging_responses = parse_judging_responses(
+                    parse_judging_response_prompt, judging_responses
+                )
+                st.session_state.direct_assessment_judging_df[response_model] = (
+                    create_dataframe_for_direct_assessment_judging_response(
+                        parsed_judging_responses
+                    )
+                )
+
+            # Uses the session state to plot the criteria scores and graphs for a given response
+            # model.
+            plot_criteria_scores(
+                st.session_state.direct_assessment_judging_df[response_model]
+            )
+
+            plot_per_judge_overall_scores(
+                st.session_state.direct_assessment_judging_df[response_model]
+            )
+
+            grouped = (
+                st.session_state.direct_assessment_judging_df[response_model]
+                .groupby(["judging_model"])
+                .agg({"score": ["mean"]})
+                .reset_index()
+            )
+            grouped.columns = ["judging_model", "overall_score"]
+
+            # Save the overall scores to the session state if it's not already there.
+            for record in grouped.to_dict(orient="records"):
+                if (
+                    response_model
+                    not in st.session_state.direct_assessment_overall_scores
+                ):
+                    st.session_state.direct_assessment_overall_scores[response_model][
+                        record["judging_model"]
+                    ] = record["overall_score"]
+
+            overall_score = grouped["overall_score"].mean()
+            controversy = grouped["overall_score"].std()
+            st.write(f"Overall Score: {overall_score:.2f}")
+            st.write(f"Controversy: {controversy:.2f}")
+
+    # Mark judging as complete.
+    st.session_state.judging_status = "complete"
+
+
 # Main Streamlit App
 def main():
     st.set_page_config(
@@ -632,71 +832,11 @@ def main():
             st.session_state.selected_aggregator = selected_aggregator
 
             # Render the chats.
-            response_columns = st.columns(3)
-
-            selected_models_to_streamlit_column_map = (
-                get_selected_models_to_streamlit_column_map(
-                    response_columns, selected_models
-                )
-            )
+            st_render_responses(user_prompt)
 
-            # Fetching and streaming responses from each selected model
-            for selected_model in st.session_state.selected_models:
-                with selected_models_to_streamlit_column_map[selected_model]:
-                    st.write(get_ui_friendly_name(selected_model))
-                    with st.chat_message(
-                        selected_model,
-                        avatar=PROVIDER_TO_AVATAR_MAP[selected_model],
-                    ):
-                        message_placeholder = st.empty()
-                        stream = get_llm_response_stream(selected_model, user_prompt)
-                        if stream:
-                            st.session_state["responses"][selected_model] = (
-                                message_placeholder.write_stream(stream)
-                            )
-
-            # Get the aggregator prompt.
-            aggregator_prompt = get_default_aggregator_prompt(
-                user_prompt=user_prompt, llms=selected_models
-            )
-
-            # Fetching and streaming response from the aggregator
-            st.write(f"{get_ui_friendly_name(selected_aggregator)}")
-            with st.chat_message(
-                selected_aggregator,
-                avatar="img/council_icon.png",
-            ):
-                message_placeholder = st.empty()
-                aggregator_stream = get_llm_response_stream(
-                    selected_aggregator, aggregator_prompt
-                )
-                if aggregator_stream:
-                    st.session_state.responses["agg__" + selected_aggregator] = (
-                        message_placeholder.write_stream(aggregator_stream)
-                    )
-
-            st.session_state.responses_collected = True
-
-        # Render chats generally?
-        if st.session_state.responses and not submit_button:
-            st.markdown("#### Responses")
-
-            response_columns = st.columns(3)
-            selected_models_to_streamlit_column_map = (
-                get_selected_models_to_streamlit_column_map(
-                    response_columns, st.session_state.selected_models
-                )
-            )
-            for response_model, response in st.session_state.responses.items():
-                st_column = selected_models_to_streamlit_column_map.get(
-                    response_model, response_columns[0]
-                )
-                with st_column.chat_message(
-                    response_model,
-                    avatar=get_llm_avatar(response_model),
-                ):
-                    st.write(get_ui_friendly_name(response_model))
-                    st.write(response)
+        # Render chats generally even they are available, if the submit button isn't clicked.
+        elif st.session_state.responses:
+            st_render_responses(user_prompt)
 
         # Judging.
         if st.session_state.responses_collected:
@@ -727,228 +867,41 @@ def main():
                     # TODO: Add option to edit criteria list with a basic text field.
                     criteria_list = DEFAULT_DIRECT_ASSESSMENT_CRITERIA_LIST
 
+                with center_column:
                     judging_submit_button = st.form_submit_button(
                         "Submit Judging", use_container_width=True
                     )
 
             if judging_submit_button:
+                # Update session state.
                 st.session_state.assessment_type = assessment_type
-                st.session_state.direct_assessment_config = {
-                    "prompt": direct_assessment_prompt,
-                    "criteria_list": criteria_list,
-                }
-
-                responses_for_judging = st.session_state.responses
-
-                # Get judging responses.
-                response_judging_columns = st.columns(3)
-                responses_for_judging_to_streamlit_column_map = (
-                    get_selected_models_to_streamlit_column_map(
-                        response_judging_columns, responses_for_judging.keys()
-                    )
-                )
-
                 if st.session_state.assessment_type == "Direct Assessment":
-                    for response_model, response in responses_for_judging.items():
-                        st_column = responses_for_judging_to_streamlit_column_map[
-                            response_model
-                        ]
-
-                        with st_column:
-                            st.write(
-                                f"Judging for {get_ui_friendly_name(response_model)}"
-                            )
-                            judging_prompt = get_direct_assessment_prompt(
-                                direct_assessment_prompt=direct_assessment_prompt,
-                                user_prompt=user_prompt,
-                                response=response,
-                                criteria_list=criteria_list,
-                                options=SEVEN_POINT_DIRECT_ASSESSMENT_OPTIONS,
-                            )
-
-                            with st.expander("Final Judging Prompt"):
-                                st.code(judging_prompt)
-
-                            for judging_model in selected_models:
-                                with st.expander(
-                                    get_ui_friendly_name(judging_model), expanded=True
-                                ):
-                                    with st.chat_message(
-                                        judging_model,
-                                        avatar=PROVIDER_TO_AVATAR_MAP[judging_model],
-                                    ):
-                                        message_placeholder = st.empty()
-                                        judging_stream = get_llm_response_stream(
-                                            judging_model, judging_prompt
-                                        )
-                                        st.session_state[
-                                            "direct_assessment_judging_responses"
-                                        ][response_model][
-                                            judging_model
-                                        ] = message_placeholder.write_stream(
-                                            judging_stream
-                                        )
-                            # When all of the judging is finished for the given response, get the actual
-                            # values, parsed.
-                            judging_responses = st.session_state[
-                                "direct_assessment_judging_responses"
-                            ][response_model]
-
-                            if not judging_responses:
-                                st.error(f"No judging responses for {response_model}")
-                                quit()
-                            parse_judging_response_prompt = (
-                                get_parse_judging_response_for_direct_assessment_prompt(
-                                    judging_responses,
-                                    criteria_list,
-                                    SEVEN_POINT_DIRECT_ASSESSMENT_OPTIONS,
-                                )
-                            )
-                            # Issue the prompt to openai mini with structured outputs
-                            parsed_judging_responses = parse_judging_responses(
-                                parse_judging_response_prompt, judging_responses
-                            )
-
-                            st.session_state["direct_assessment_judging_df"][
-                                response_model
-                            ] = create_dataframe_for_direct_assessment_judging_response(
-                                parsed_judging_responses
-                            )
-
-                            plot_criteria_scores(
-                                st.session_state["direct_assessment_judging_df"][
-                                    response_model
-                                ]
-                            )
-
-                            # Find the overall score by finding the overall score for each judge, and then averaging
-                            # over all judges.
-                            plot_per_judge_overall_scores(
-                                st.session_state["direct_assessment_judging_df"][
-                                    response_model
-                                ]
-                            )
-
-                            grouped = (
-                                st.session_state["direct_assessment_judging_df"][
-                                    response_model
-                                ]
-                                .groupby(["judging_model"])
-                                .agg({"score": ["mean"]})
-                                .reset_index()
-                            )
-                            grouped.columns = ["judging_model", "overall_score"]
-
-                            # Save the overall scores to the session state.
-                            for record in grouped.to_dict(orient="records"):
-                                st.session_state["direct_assessment_overall_scores"][
-                                    response_model
-                                ][record["judging_model"]] = record["overall_score"]
-
-                            overall_score = grouped["overall_score"].mean()
-                            controversy = grouped["overall_score"].std()
-                            st.write(f"Overall Score: {overall_score:.2f}")
-                            st.write(f"Controversy: {controversy:.2f}")
-
-                    st.session_state.judging_status = "complete"
+                    st.session_state.direct_assessment_config = {
+                        "prompt": direct_assessment_prompt,
+                        "criteria_list": criteria_list,
+                    }
+                    st_direct_assessment_results(
+                        user_prompt=st.session_state.user_prompt,
+                        direct_assessment_prompt=direct_assessment_prompt,
+                        criteria_list=criteria_list,
+                    )
             # If judging is complete, but the submit button is cleared, still render the results.
             elif st.session_state.judging_status == "complete":
                 if st.session_state.assessment_type == "Direct Assessment":
-                    responses_for_judging = st.session_state.responses
-
-                    # Get judging responses.
-                    response_judging_columns = st.columns(3)
-                    responses_for_judging_to_streamlit_column_map = (
-                        get_selected_models_to_streamlit_column_map(
-                            response_judging_columns, responses_for_judging.keys()
-                        )
+                    st_direct_assessment_results(
+                        user_prompt=st.session_state.user_prompt,
+                        direct_assessment_prompt=direct_assessment_prompt,
+                        criteria_list=criteria_list,
                     )
 
-                    for response_model, response in responses_for_judging.items():
-                        st_column = responses_for_judging_to_streamlit_column_map[
-                            response_model
-                        ]
-
-                        with st_column:
-                            st.write(
-                                f"Judging for {get_ui_friendly_name(response_model)}"
-                            )
-                            judging_prompt = get_direct_assessment_prompt(
-                                direct_assessment_prompt=direct_assessment_prompt,
-                                user_prompt=user_prompt,
-                                response=response,
-                                criteria_list=criteria_list,
-                                options=SEVEN_POINT_DIRECT_ASSESSMENT_OPTIONS,
-                            )
-
-                            with st.expander("Final Judging Prompt"):
-                                st.code(judging_prompt)
-
-                            for judging_model in selected_models:
-                                with st.expander(
-                                    get_ui_friendly_name(judging_model), expanded=True
-                                ):
-                                    with st.chat_message(
-                                        judging_model,
-                                        avatar=PROVIDER_TO_AVATAR_MAP[judging_model],
-                                    ):
-                                        st.write(
-                                            st.session_state.direct_assessment_judging_responses[
-                                                response_model
-                                            ][
-                                                judging_model
-                                            ]
-                                        )
-                            # When all of the judging is finished for the given response, get the actual
-                            # values, parsed.
-                            judging_responses = (
-                                st.session_state.direct_assessment_judging_responses[
-                                    response_model
-                                ]
-                            )
-
-                            parse_judging_response_prompt = (
-                                get_parse_judging_response_for_direct_assessment_prompt(
-                                    judging_responses,
-                                    criteria_list,
-                                    SEVEN_POINT_DIRECT_ASSESSMENT_OPTIONS,
-                                )
-                            )
-
-                            plot_criteria_scores(
-                                st.session_state.direct_assessment_judging_df[
-                                    response_model
-                                ]
-                            )
-
-                            plot_per_judge_overall_scores(
-                                st.session_state.direct_assessment_judging_df[
-                                    response_model
-                                ]
-                            )
-
-                            grouped = (
-                                st.session_state.direct_assessment_judging_df[
-                                    response_model
-                                ]
-                                .groupby(["judging_model"])
-                                .agg({"score": ["mean"]})
-                                .reset_index()
-                            )
-                            grouped.columns = ["judging_model", "overall_score"]
-
-                            overall_score = grouped["overall_score"].mean()
-                            controversy = grouped["overall_score"].std()
-                            st.write(f"Overall Score: {overall_score:.2f}")
-                            st.write(f"Controversy: {controversy:.2f}")
-
-            # Judging is complete, stuff that would be rendered that's not stream-specific.
+            # Judging is complete.
+            # Render stuff that would be rendered that's not stream-specific.
             # The session state now contains the overall scores for each response from each judge.
             if st.session_state.judging_status == "complete":
                 st.write("#### Results")
 
                 overall_scores_df_raw = pd.DataFrame(
-                    st.session_state["direct_assessment_overall_scores"]
+                    st.session_state.direct_assessment_overall_scores
                 ).reset_index()
 
                 overall_scores_df = pd.melt(