1616from streamlit import session_state as state
1717
1818from client .content .config .tabs .models import get_models
19- from client .utils import st_common , api_call , client
19+ from client .utils import st_common , api_call , client , vs_options
2020from client .utils .st_footer import render_chat_footer
2121from common import logging_config
2222
2626#############################################################################
2727# Functions
2828#############################################################################
29- def show_vector_search_refs (context ):
29+ def show_vector_search_refs (context , vs_metadata = None ):
3030 """When Vector Search Content Found, show the references"""
3131 st .markdown ("**References:**" )
3232 ref_src = set ()
3333 ref_cols = st .columns ([3 , 3 , 3 ])
3434 # Create a button in each column
35- for i , (ref_col , chunk ) in enumerate (zip (ref_cols , context [0 ])):
35+ for i , (ref_col , chunk ) in enumerate (zip (ref_cols , context ["documents" ])):
3636 with ref_col .popover (f"Reference: { i + 1 } " ):
37- chunk = context [0 ][i ]
37+ chunk = context ["documents" ][i ]
3838 logger .debug ("Chunk Content: %s" , chunk )
3939 st .subheader ("Reference Text" , divider = "red" )
4040 st .markdown (chunk ["page_content" ])
@@ -46,9 +46,32 @@ def show_vector_search_refs(context):
4646 except KeyError :
4747 logger .error ("Chunk Metadata NOT FOUND!!" )
4848
49- for link in ref_src :
50- st .markdown ("- " + link )
51- st .markdown (f"**Notes:** Vector Search Query - { context [1 ]} " )
49+ # Display Vector Search details in expander
50+ if vs_metadata or ref_src :
51+ with st .expander ("Vector Search Details" , expanded = False ):
52+ if ref_src :
53+ st .markdown ("**Source Documents:**" )
54+ for link in ref_src :
55+ st .markdown (f"- { link } " )
56+
57+ if vs_metadata and vs_metadata .get ("searched_tables" ):
58+ st .markdown ("**Tables Searched:**" )
59+ for table in vs_metadata ["searched_tables" ]:
60+ st .markdown (f"- { table } " )
61+
62+ if vs_metadata and vs_metadata .get ("context_input" ):
63+ st .markdown (f"**Search Query:** { vs_metadata .get ('context_input' )} " )
64+ elif context .get ("context_input" ):
65+ st .markdown (f"**Search Query:** { context .get ('context_input' )} " )
66+
67+
68+ def show_token_usage (token_usage ):
69+ """Display token usage for AI responses using caption"""
70+ if token_usage :
71+ prompt_tokens = token_usage .get ("prompt_tokens" , 0 )
72+ completion_tokens = token_usage .get ("completion_tokens" , 0 )
73+ total_tokens = token_usage .get ("total_tokens" , 0 )
74+ st .caption (f"Token usage: { prompt_tokens } prompt + { completion_tokens } completion = { total_tokens } total" )
5275
5376
5477def setup_sidebar ():
@@ -62,7 +85,7 @@ def setup_sidebar():
6285 st_common .tools_sidebar ()
6386 st_common .history_sidebar ()
6487 st_common .ll_sidebar ()
65- st_common .vector_search_sidebar ()
88+ vs_options .vector_search_sidebar ()
6689
6790 if not state .enable_client :
6891 st .stop ()
@@ -80,22 +103,33 @@ def create_client():
80103
81104
82105def display_chat_history (history ):
83- """Display chat history messages"""
106+ """Display chat history messages with metadata """
84107 st .chat_message ("ai" ).write ("Hello, how can I help you?" )
85108 vector_search_refs = []
86109
87110 for message in history or []:
88111 if not message ["content" ]:
89112 continue
90113
91- if message ["role" ] == "tool" and message ["name" ] == "oraclevs_tool " :
114+ if message ["role" ] == "tool" and message ["name" ] == "optimizer_vs-retriever " :
92115 vector_search_refs = json .loads (message ["content" ])
93116
94117 elif message ["role" ] in ("ai" , "assistant" ):
95118 with st .chat_message ("ai" ):
96119 st .markdown (message ["content" ])
120+
121+ # Extract metadata from response_metadata
122+ response_metadata = message .get ("response_metadata" , {})
123+ vs_metadata = response_metadata .get ("vs_metadata" , {})
124+ token_usage = response_metadata .get ("token_usage" , {})
125+
126+ # Show token usage immediately after message
127+ if token_usage :
128+ show_token_usage (token_usage )
129+
130+ # Show vector search references if available
97131 if vector_search_refs :
98- show_vector_search_refs (vector_search_refs )
132+ show_vector_search_refs (vector_search_refs , vs_metadata )
99133 vector_search_refs = []
100134
101135 elif message ["role" ] in ("human" , "user" ):
@@ -131,9 +165,32 @@ async def handle_chat_input(user_client):
131165 try :
132166 message_placeholder = st .chat_message ("ai" ).empty ()
133167 full_answer = ""
134- async for chunk in user_client .stream (message = human_request .text , image_b64 = file_b64 ):
135- full_answer += chunk
136- message_placeholder .markdown (full_answer )
168+
169+ # Animated thinking indicator
170+ async def animate_thinking ():
171+ """Animate the thinking indicator with increasing dots"""
172+ dots = 0
173+ while True :
174+ message_placeholder .markdown (f"🤔 Thinking{ '.' * (dots % 4 )} " )
175+ dots += 1
176+ await asyncio .sleep (0.5 ) # Update every 500ms
177+
178+ # Start the thinking animation
179+ thinking_task = asyncio .create_task (animate_thinking ())
180+
181+ try :
182+ async for chunk in user_client .stream (message = human_request .text , image_b64 = file_b64 ):
183+ # Cancel thinking animation on first chunk
184+ if thinking_task and not thinking_task .done ():
185+ thinking_task .cancel ()
186+ thinking_task = None
187+ full_answer += chunk
188+ message_placeholder .markdown (full_answer )
189+ finally :
190+ # Ensure thinking task is cancelled
191+ if thinking_task and not thinking_task .done ():
192+ thinking_task .cancel ()
193+
137194 st .rerun ()
138195 except (ConnectionError , TimeoutError , api_call .ApiError ) as ex :
139196 logger .exception ("Error during chat streaming: %s" , ex )
0 commit comments