Enhance Gradio UI and Card Display Logic

This commit introduces several improvements and bug fixes to the Gradio application, enhancing user experience and display reliability.

Key changes include:

- **Multiple Submissions:** Resolved an issue preventing consecutive submissions without clearing state. `GradioEventHandler` is now instantiated per submission.
- **OTS Content Reset:** Ensured the One-True-Shot (OTS) content area resets correctly to its default state for new queries.
- **UI Card Rendering:** Improved the reliability of dynamic UI card rendering (players, games, teams) by adding explicit yield statements.
- **Session Persistence:** Modified the "Clear State" functionality to preserve Zep and Freeplay session IDs, allowing users to continue their session.
- **Card Display Priority:** Corrected the logic in `on_tool_end` to ensure game and team cards are displayed correctly, not just player cards.
- **Semi-Final Game Cards:** Implemented a dedicated display for "semi-final" game cards, utilizing the existing game card format.

These changes address issues with submission flow, UI state management, and card display logic, while also introducing a new feature for specific game event cards.

Stagelink: https://huggingface.co/spaces/yamilsteven/ifx-huge-league-app

Makefile

@@ -31,7 +31,7 @@ clean-requirements:
 
 # Build the Docker image for the 'runtime' stage in api/Dockerfile, tagged as huge-ifx-api:prod
 build-prod:
-	cd api && docker build -f Dockerfile --target runtime -t huge-ifx-api:prod .
+	cd api && docker build --platform linux/amd64 -f Dockerfile --target runtime -t huge-ifx-api:prod .
 
 # Build the prod image and run it locally, mapping ports 7860 and 8000
 up-build-prod: build-prod
@@ -39,5 +39,5 @@ up-build-prod: build-prod
 
 # Push the prod image to GitHub Container Registry
 push-prod-ghcr:
-	docker tag huge-ifx-api:prod ghcr.io/rbalch/huge-ifx-api:prod
-	docker push ghcr.io/rbalch/huge-ifx-api:prod
+	docker tag huge-ifx-api:prod ghcr.io/ylassohugeinc/ifx-huge-league-api:prod
+	docker push ghcr.io/ylassohugeinc/ifx-huge-league-api:prod

api/README.md

@@ -109,7 +109,7 @@ To deploy:
 2. **Deployment Targets:**
     - GitHub: [aliss77777/IFX-sandbox](https://github.com/aliss77777/IFX-sandbox/tree/huge-league)
     - HuggingFace: [ryanbalch/IFX-huge-league](https://huggingface.co/spaces/ryanbalch/IFX-huge-league/tree/main)
-    - Docker images: `ghcr.io/rbalch/huge-ifx-api:prod`
+    - Docker images: `ghcr.io/ylassohugeinc/ifx-huge-league-api:prod`
 
 3. **Deployment Steps:**
 
@@ -196,7 +196,7 @@ The following `make` commands are available for development, build, and deployme
 | `make extract-lock`   | Extract the `poetry.lock` file from a built container to your local directory. **Note:** This is only needed if you've been deleting the lock file because build will not have access to local lock file. |
 | `make build-prod`     | Build the Docker image for the `runtime` stage in `api/Dockerfile`, tagged as `huge-ifx-api:prod`. Used for production deploys. |
 | `make up-build-prod`  | Build and run the production image locally, mapping ports 7860 and 8000, with `.env` and `DEV_MODE=true`. |
-| `make push-prod-ghcr` | Tag and push the production image to GitHub Container Registry at `ghcr.io/rbalch/huge-ifx-api:prod`. |
+| `make push-prod-ghcr` | Tag and push the production image to GitHub Container Registry at `ghcr.io/ylassohugeinc/ifx-huge-league-api:prod`. |
 
 **Typical workflow:**
 - Use `make build` and `make up` for local development.

api/components_bugfix_README.md

@@ -0,0 +1,80 @@
+# Technical Design Document: Gradio UI Enhancements & Card Display Logic
+
+**Version:** 1.0
+**Date:** 6/3/2025
+**Author:** @ylasso
+
+## 1. Overview
+
+This document outlines a series of enhancements and bug fixes applied to the Gradio-based application. The primary goals of these changes were to improve user experience by enabling seamless multiple submissions, ensuring correct UI state resets, enhancing the reliability of dynamic card displays (for players, games, teams), and introducing new features like session persistence on state clear and specific handling for semi-final game cards.
+
+## 2. Problems Addressed and Solutions Implemented
+
+### 2.1. Consecutive Submissions Requiring "Clear State"
+
+*   **Problem:** Users could not make a second successful submission without first clicking the "Clear State" button. The first submission would work, but subsequent ones would fail to produce the expected output or behavior.
+*   **Root Cause:** The `GradioEventHandler` instance, which manages a queue for streaming responses, was initialized globally once. Its internal queue was not reset between submissions. If a previous workflow ended by placing a `None` (end-of-stream marker) in the queue, subsequent submissions would immediately read this `None` and terminate prematurely.
+*   **Solution:**
+    *   The `GradioEventHandler` instantiation was moved from a global scope within the Gradio Blocks definition to occur locally within the `submit` and `user_query_submit` event handler functions in `api/server_gradio.py`.
+    *   This ensures that a fresh `GradioEventHandler` (and thus a new, empty queue) is created for each user submission, preventing interference from previous submission states.
+
+### 2.2. Persistent OTS (One-True-Shot) Content
+
+*   **Problem:** After a player card (or other specialized content in the "OTS" display area) was shown, it would remain visible even after a new, unrelated query was submitted. The OTS area was not resetting to its default landing state.
+*   **Root Cause:** The `ots_content` attribute in the `AppState` was not being explicitly reset at the beginning of a new submission.
+*   **Solution:**
+    *   A constant `INITIAL_OTS_IMAGE_HTML` was defined in `api/server_gradio.py` holding the HTML for the default landing image.
+    *   The `AppState` model's `ots_content` field now defaults to this initial image.
+    *   Crucially, at the beginning of the `submit_helper` function (called by both submit actions), `state.ots_content` is explicitly reset to `ots_default.format(content=INITIAL_OTS_IMAGE_HTML)`. This ensures the OTS display area starts fresh with the landing image for every new query.
+
+### 2.3. Intermittent Failure to Render UI Cards (OTS Content)
+
+*   **Problem:** Even when a tool was correctly invoked (indicated by a diagnostic message), the corresponding UI card (e.g., player card, game card) would sometimes fail to render, and only a text response would be shown. This behavior was inconsistent.
+*   **Root Cause:** The Gradio UI might not have been updated with the new `state.ots_content` if the "ots" token from the `handler.queue` was the last meaningful token before the stream ended, or if no further text tokens followed to trigger another `yield`. The state update with the card HTML might have been "missed."
+*   **Solution:**
+    *   In the `submit_helper` function in `api/server_gradio.py`, an explicit `yield state, result` was added immediately after an "ots" type token is processed and `state.ots_content` is updated. This forces Gradio to update with the card content as soon as it's ready.
+    *   An additional final `yield state, result` was added at the very end of the `submit_helper` function, after the main processing loop, to ensure the absolute final state (including any last-minute `ots_content` or the complete `result`) is always rendered.
+
+### 2.4. Session ID Reset on "Clear State"
+
+*   **Problem:** Clicking the "Clear State" button would reset all application state, including Zep and Freeplay session IDs. This meant subsequent interactions would start entirely new sessions, losing context from the prior interaction if the user only intended to clear the current query and history.
+*   **Desired Behavior:** Preserve Zep and Freeplay session IDs when "Clear State" is used, allowing the user to continue the same underlying session while clearing the immediate UI and chat history.
+*   **Root Cause:** The `clear_state` function created a completely new `AppState()` instance, which initializes session IDs to empty strings by default.
+*   **Solution:**
+    *   The `clear_state` function in `api/server_gradio.py` was modified to accept the current `state` (as `current_app_state: AppState`) as an input.
+    *   Before creating the `new_state`, it now extracts `existing_zep_session_id` and `existing_freeplay_session_id` from `current_app_state`.
+    *   These preserved session IDs are then passed explicitly when instantiating the `new_state = AppState(zep_session_id=..., freeplay_session_id=...)`. Other fields are reset to their defaults.
+
+### 2.5. Game and Team Cards Not Displaying (Player Card Priority Issue)
+
+*   **Problem:** Only player cards were being displayed. Queries that should have resulted in game cards or team cards were not showing these specific UIs.
+*   **Root Cause:** Within the `on_tool_end` method of `api/event_handlers/gradio_handler.py`, the logic for displaying player cards was checked first. If a player card was generated, the method would `return` immediately, preventing the subsequent `if/elif` blocks for game cards and team cards from being evaluated.
+*   **Solution:**
+    *   The conditional logic in `on_tool_end` was restructured to prioritize more specific card types. The new order of checks is:
+        1.  (New) Specific check for "Semi-Final" game cards (see section 3.1).
+        2.  Generic Game Cards (based on `output[0].metadata.get('type') == 'event'`).
+        3.  Team Cards (based on `output[0].metadata.get("show_team_card")`).
+        4.  Player Cards (iterating `output` for `doc.metadata.get("show_profile_card")`).
+    *   Each primary card-type block still `return`s after successfully queuing its card via `self.ots_box()`, assuming only one card is displayed per tool end.
+
+## 3. New Features
+
+### 3.1. Dedicated Semi-Final Game Card Display
+
+*   **Requirement:** Provide a distinct card display for queries related to "semi-finals," similar to the existing game card.
+*   **Implementation:**
+    *   An `if` condition was added at the beginning of the `on_tool_end` method in `api/event_handlers/gradio_handler.py`.
+    *   This condition specifically checks if `output[0].metadata.get('game_name')` (case-insensitive) contains the string "semi-final".
+    *   If true, the existing game card generation logic (`game_card_html` and associated data extraction) is used to render the semi-final information.
+    *   This ensures that semi-final games are explicitly identified and displayed using the game card format.
+    *   **Dependency:** This relies on the upstream tool providing `game_name` in `output[0].metadata` for semi-final related queries.
+
+## 4. Affected Files
+
+*   `api/server_gradio.py`:
+    *   Modified `submit` and `user_query_submit` event handlers.
+    *   Modified `submit_helper` function.
+    *   Modified `clear_state` function and its Gradio registration.
+    *   Added `INITIAL_OTS_IMAGE_HTML` constant and updated `AppState` default for `ots_content`.
+*   `api/event_handlers/gradio_handler.py`:
+    *   Significantly refactored the `on_tool_end` method's conditional logic for card display.

api/event_handlers/gradio_handler.py

@@ -115,52 +115,86 @@ class GradioEventHandler(AsyncCallbackHandler):
 
     async def on_tool_end(self, output: any, **kwargs):
         print(f"\n{Fore.CYAN}[TOOL END] {output}{Style.RESET_ALL}")
-        for doc in output:
-            if doc.metadata.get("show_profile_card"):
-                raw_player_name = doc.metadata.get("name", "unknown-player")
-                player_name = "Unknown Player"
-                if raw_player_name != "unknown-player":
-                    parts = raw_player_name.split('-')
-                    if len(parts) == 2:
-                        first_name = parts[0].capitalize()
-                        last_name = parts[1].capitalize()
-                        player_name = f"{last_name} {first_name}"
-                    else:
-                        player_name = raw_player_name.replace('-', ' ').title()
-                else:
-                    player_name = "Unknown Player"
 
-                team_slug = doc.metadata.get("team", "unknown-team")
-                formatted_team_name = team_slug.replace('-', ' ').title()
-                if team_slug == "unknown-team":
-                    formatted_team_name = "Unknown Team"
+        # Check for Semi-Final Game Card first, based on game_name
+        if output and isinstance(output, list) and len(output) > 0 and \
+           isinstance(output[0], object) and hasattr(output[0], 'metadata') and \
+           isinstance(output[0].metadata.get('game_name'), str) and \
+           "semi-final" in output[0].metadata.get('game_name', '').lower():
+            
+            # This is largely the same logic as the generic game card below
+            game_teams_set = set()
+            game_events_details = []
+            # Ensure raw_game_name uses the game_name from metadata if available
+            raw_game_name = output[0].metadata.get('game_name', "Unknown Semi-Final")
+
+            for i, doc in enumerate(output): # Iterate through all docs for game details
+                if hasattr(doc, 'metadata') and doc.metadata.get('team'):
+                    game_teams_set.add(doc.metadata.get('team'))
+                # game_name is taken from output[0] already
                 
-                player_number = doc.metadata.get("number", "N/A")
-                height = doc.metadata.get("height", "175")
-                weight = doc.metadata.get("weight", "150")
-                instagram_url = doc.metadata.get("instagram_url", raw_player_name)
-
-                img = image_base.format(
-                    filename=self.get_image_filename(doc),
-                    player_name=player_name,
-                    team_name=formatted_team_name,
-                    player_number=player_number,
-                    height=height,
-                    weight=weight,
-                    instagram_url=instagram_url
+                if hasattr(doc, 'metadata') and doc.metadata.get('event') == 'Goal':
+                    event_desc = doc.metadata.get('description', 'Goal scored')
+                    event_minute = doc.metadata.get('minute', '')
+                    event_team = doc.metadata.get('team', '')
+                    game_events_details.append(f"({event_minute}') {event_team}: {event_desc}")
+            
+            if len(game_teams_set) >= 1: 
+                team_list = list(game_teams_set)
+                team1_name_str = team_list[0]
+                team2_name_str = team_list[1] if len(team_list) > 1 else "(opponent not specified)"
+
+                score_team1 = 0
+                score_team2 = 0
+                for doc in output: 
+                    if hasattr(doc, 'metadata') and doc.metadata.get('event') == 'Goal':
+                        scoring_team = doc.metadata.get('team')
+                        if scoring_team == team1_name_str:
+                            score_team1 += 1
+                        elif scoring_team == team2_name_str:
+                            score_team2 += 1
+                
+                game_title_str = raw_game_name.replace('_', ' ').title()
+                highlights_html_str = "<ul style='list-style-type: none; padding-left: 0; text-align: left;'>"
+                for highlight in game_events_details[:3]: 
+                    highlights_html_str += f"<li style='margin-bottom: 5px; font-size: 14px;'>{highlight}</li>"
+                highlights_html_str += "</ul>"
+                if not game_events_details:
+                    highlights_html_str = "<p style='font-size: 14px;'>No goal highlights available.</p>"
+              
+                team1_logo_filename = GradioEventHandler.get_team_logo_slug(team1_name_str)
+                team2_logo_filename = GradioEventHandler.get_team_logo_slug(team2_name_str)
+
+                team1_logo_url = f"{TEAM_LOGO_BASE_URL}{team1_logo_filename}"
+                team2_logo_url = f"{TEAM_LOGO_BASE_URL}{team2_logo_filename}"
+                default_logo_url = f"{TEAM_LOGO_BASE_URL}default.png"
+
+                formatted_game_html = game_card_html.format(
+                    game_title=game_title_str,
+                    team_home=team1_name_str, 
+                    team_away=team2_name_str,
+                    team1_score=score_team1,
+                    team2_score=score_team2,
+                    highlights=highlights_html_str,
+                    team1_logo_url=team1_logo_url,
+                    team2_logo_url=team2_logo_url,
+                    default_logo_url=default_logo_url
                 )
-                self.ots_box(img)
+                self.ots_box(formatted_game_html)
                 return
+            else:
+                print(f"\n{Fore.RED}[TOOL END - SEMI-FINAL CARD] Not enough team data or structure not as expected for {raw_game_name}.{Style.RESET_ALL}")
 
-        if output and isinstance(output[0], object) and hasattr(output[0], 'metadata') and output[0].metadata.get('type') == 'event':
+        # Attempt to process Generic Game Card (e.g., type 'event') if not a Semi-Final
+        elif output and isinstance(output, list) and len(output) > 0 and isinstance(output[0], object) and hasattr(output[0], 'metadata') and output[0].metadata.get('type') == 'event':
             game_teams_set = set()
             game_events_details = []
             raw_game_name = "Unknown Game"
 
-            for i, doc in enumerate(output):
+            for i, doc in enumerate(output): 
                 if hasattr(doc, 'metadata') and doc.metadata.get('team'):
                     game_teams_set.add(doc.metadata.get('team'))
-                if i == 0:
+                if i == 0: 
                     raw_game_name = doc.metadata.get('game_name', raw_game_name)
                 
                 if hasattr(doc, 'metadata') and doc.metadata.get('event') == 'Goal':
@@ -169,14 +203,14 @@ class GradioEventHandler(AsyncCallbackHandler):
                     event_team = doc.metadata.get('team', '')
                     game_events_details.append(f"({event_minute}') {event_team}: {event_desc}")
             
-            if len(game_teams_set) >= 1:
+            if len(game_teams_set) >= 1: 
                 team_list = list(game_teams_set)
                 team1_name_str = team_list[0]
                 team2_name_str = team_list[1] if len(team_list) > 1 else "(opponent not specified)"
 
                 score_team1 = 0
                 score_team2 = 0
-                for doc in output:
+                for doc in output: 
                     if hasattr(doc, 'metadata') and doc.metadata.get('event') == 'Goal':
                         scoring_team = doc.metadata.get('team')
                         if scoring_team == team1_name_str:
@@ -191,8 +225,7 @@ class GradioEventHandler(AsyncCallbackHandler):
                 highlights_html_str += "</ul>"
                 if not game_events_details:
                     highlights_html_str = "<p style='font-size: 14px;'>No goal highlights available.</p>"
-
-                log_msg = f"Game: {game_title_str} | {team1_name_str} {score_team1} - {score_team2} {team2_name_str}"                
+              
                 team1_logo_filename = GradioEventHandler.get_team_logo_slug(team1_name_str)
                 team2_logo_filename = GradioEventHandler.get_team_logo_slug(team2_name_str)
 
@@ -212,34 +245,76 @@ class GradioEventHandler(AsyncCallbackHandler):
                     default_logo_url=default_logo_url
                 )
                 self.ots_box(formatted_game_html)
-                return
+                return 
             else:
-                print(f"\n{Fore.RED}[TOOL END - GAME CARD] Not enough team data found in events.{Style.RESET_ALL}")
-        
-        elif isinstance(output, list) and output:
+                print(f"\n{Fore.RED}[TOOL END - GAME CARD] Not enough team data found in events or output structure not as expected.{Style.RESET_ALL}")
+
+        # Attempt to process Team Card if no Game Card was processed
+        elif output and isinstance(output, list) and len(output) > 0 and isinstance(output[0], object) and hasattr(output[0], 'metadata') and output[0].metadata.get("show_team_card"):
             doc_for_team_card = output[0]
-            if hasattr(doc_for_team_card, 'metadata') and doc_for_team_card.metadata.get("show_team_card"):
-                team_name_from_meta = doc_for_team_card.metadata.get("team_name", "Unknown Team")
-                city_raw = doc_for_team_card.metadata.get("city", "N/A")
-                display_team_name = str(team_name_from_meta).replace('-', ' ').replace('_', ' ').title()
-                city_display = city_raw.title() if city_raw != "N/A" else "Location N/A"
-                logo_filename = GradioEventHandler.get_team_logo_slug(team_name_from_meta)
-                team_specific_logo_url = f"{TEAM_LOGO_BASE_URL}{logo_filename}"
-                current_default_logo_url = f"{TEAM_LOGO_BASE_URL}default.png"
-                team_id_slug = doc_for_team_card.metadata.get("team_id", "")
-                team_page_url = f"https://www.team.com/{team_id_slug}" if team_id_slug else "#" 
-                team_page_cta_html = f'''<a href="{team_page_url}" target="_blank" style="background-color: #007bff; color: white; padding: 10px 20px; text-decoration: none; border-radius: 5px; display: inline-block; font-size: 14px; margin-top: 10px;">Visit Team Page</a>'''
-
-                formatted_team_card_html = team_info_card_html.format(
-                    team_logo_url=team_specific_logo_url,
-                    team_display_name=display_team_name,
-                    city_display=city_display,
-                    default_logo_url=current_default_logo_url,
-                    team_page_cta_html=team_page_cta_html
-                )
-                
-                self.ots_box(formatted_team_card_html)
-                return
+            team_name_from_meta = doc_for_team_card.metadata.get("team_name", "Unknown Team")
+            city_raw = doc_for_team_card.metadata.get("city", "N/A")
+            display_team_name = str(team_name_from_meta).replace('-', ' ').replace('_', ' ').title()
+            city_display = city_raw.title() if city_raw != "N/A" else "Location N/A"
+            logo_filename = GradioEventHandler.get_team_logo_slug(team_name_from_meta)
+            team_specific_logo_url = f"{TEAM_LOGO_BASE_URL}{logo_filename}"
+            current_default_logo_url = f"{TEAM_LOGO_BASE_URL}default.png"
+            team_id_slug = doc_for_team_card.metadata.get("team_id", "")
+            team_page_url = f"https://www.team.com/{team_id_slug}" if team_id_slug else "#" 
+            team_page_cta_html = f'''<a href="{team_page_url}" target="_blank" style="background-color: #007bff; color: white; padding: 10px 20px; text-decoration: none; border-radius: 5px; display: inline-block; font-size: 14px; margin-top: 10px;">Visit Team Page</a>'''
+
+            formatted_team_card_html = team_info_card_html.format(
+                team_logo_url=team_specific_logo_url,
+                team_display_name=display_team_name,
+                city_display=city_display,
+                default_logo_url=current_default_logo_url,
+                team_page_cta_html=team_page_cta_html
+            )
+            self.ots_box(formatted_team_card_html)
+            return
+
+        # Attempt to process Player Card if no Game or Team Card was processed
+        # This iterates through all docs in output, so it's more general
+        elif output and isinstance(output, list):
+            for doc in output:
+                if hasattr(doc, 'metadata') and doc.metadata.get("show_profile_card"):
+                    raw_player_name = doc.metadata.get("name", "unknown-player")
+                    player_name = "Unknown Player"
+                    if raw_player_name != "unknown-player":
+                        parts = raw_player_name.split('-')
+                        if len(parts) == 2:
+                            first_name = parts[0].capitalize()
+                            last_name = parts[1].capitalize()
+                            player_name = f"{last_name} {first_name}"
+                        else:
+                            player_name = raw_player_name.replace('-', ' ').title()
+                    else:
+                        player_name = "Unknown Player"
+
+                    team_slug = doc.metadata.get("team", "unknown-team")
+                    formatted_team_name = team_slug.replace('-', ' ').title()
+                    if team_slug == "unknown-team":
+                        formatted_team_name = "Unknown Team"
+                    
+                    player_number = doc.metadata.get("number", "N/A")
+                    height = doc.metadata.get("height", "175")
+                    weight = doc.metadata.get("weight", "150")
+                    instagram_url = doc.metadata.get("instagram_url", raw_player_name)
+
+                    img = image_base.format(
+                        filename=self.get_image_filename(doc),
+                        player_name=player_name,
+                        team_name=formatted_team_name,
+                        player_number=player_number,
+                        height=height,
+                        weight=weight,
+                        instagram_url=instagram_url
+                    )
+                    self.ots_box(img)
+                    return # Return after the first player card is found
+
+        # Fallback if no specific card type was identified
+        print(f"\n{Fore.YELLOW}[TOOL END] No specific card type processed for OTS.{Style.RESET_ALL}")
 
     async def on_tool_start(self, input: any, *args, **kwargs):
         self.info_box(input.get("name", "[TOOL START]"))

api/server_gradio.py

@@ -18,6 +18,12 @@ MESSAGE_TYPE_MAP = {
     "ai": AIMessage,
     # Add other message types as needed
 }
+INITIAL_OTS_IMAGE_HTML = """
+    <img 
+        src="https://huggingface.co/spaces/ryanbalch/IFX-huge-league/resolve/main/assets/huge_landing.png"
+        style="max-width: 100%; max-height: 100%; object-fit: contain; display: block; margin: 0 auto;"
+    />
+    """
 ots_default = """
 <div style="display: flex; justify-content: center; align-items: center; width: 100%; max-width: 727px; height: 363px; margin: 0 auto;">
     {content}
@@ -34,12 +40,7 @@ class AppState(BaseModel):
     history: list = []
     zep_session_id: str = ""
     freeplay_session_id: str = ""
-    ots_content: str = ots_default.format(content="""
-    <img 
-        src="https://huggingface.co/spaces/ryanbalch/IFX-huge-league/resolve/main/assets/huge_landing.png"
-        style="max-width: 100%; max-height: 100%; object-fit: contain; display: block; margin: 0 auto;"
-    />
-    """)
+    ots_content: str = ots_default.format(content=INITIAL_OTS_IMAGE_HTML)
     
     def ensure_sessions(self):
         if not self.zep_session_id:
@@ -67,6 +68,7 @@ class AppState(BaseModel):
 ### Helpers ###
 
 def submit_helper(state, handler, user_query):
+    state.ots_content = ots_default.format(content=INITIAL_OTS_IMAGE_HTML)
     state.count += 1
     state.ensure_sessions()
     message = HumanMessage(content=user_query)
@@ -117,17 +119,18 @@ def submit_helper(state, handler, user_query):
                 print('OTS: ' + token["message"])
                 state.ots_content = ots_default.format(content=token["message"])
                 state = AppState(**state.model_dump())
+                yield state, result
                 continue
         result += token
         yield state, result
     
     state.history.append(AIMessage(content=result))
+    yield state, result
 
 ### Interface ###
 
 with gr.Blocks(theme="soft") as demo:
     state = gr.State(AppState())
-    handler = GradioEventHandler()
 
     gr.Markdown("# Huge League Soccer")
     with gr.Row():
@@ -207,21 +210,33 @@ with gr.Blocks(theme="soft") as demo:
     def state_change(state):
         return state.count, state.persona, state.zep_session_id, state.freeplay_session_id, "", state.ots_content
     
-    @clear_state_btn.click(outputs=[state, llm_response, persona, user_query, email, first_name, last_name])
-    def clear_state():
-        new_state = AppState()
+    @clear_state_btn.click(inputs=[state], outputs=[state, llm_response, persona, user_query, email, first_name, last_name])
+    def clear_state(current_app_state: AppState):
+        # Preserve session IDs
+        existing_zep_session_id = current_app_state.zep_session_id
+        existing_freeplay_session_id = current_app_state.freeplay_session_id
+
+        # Create new AppState, passing preserved session IDs
+        # Other fields will get their default values from AppState definition
+        new_state = AppState(
+            zep_session_id=existing_zep_session_id,
+            freeplay_session_id=existing_freeplay_session_id
+        )
+        # The email, first_name, last_name, persona will be the defaults from new_state
         return new_state, "", new_state.persona, "", new_state.email, new_state.first_name, new_state.last_name
 
     @submit_btn.click(inputs=[state, user_query], outputs=[state, llm_response])
     def submit(state, user_query):
         # user_query = user_query or "tell me about some players in everglade fc"
         user_query = user_query or "tell me about Ryan Martinez of everglade fc"
+        handler = GradioEventHandler()
         yield from submit_helper(state, handler, user_query)
         
     @user_query.submit(inputs=[state, user_query], outputs=[state, llm_response])
     def user_query_change(state, user_query):
         # user_query = user_query or "tell me about some players in everglade fc"
         user_query = user_query or "tell me about Ryan Martinez of everglade fc"
+        handler = GradioEventHandler()
         yield from submit_helper(state, handler, user_query)
 
     @persona.change(inputs=[persona, state], outputs=[persona_disp])