Update ui.py

ui.py

@@ -70,6 +70,8 @@ def create_ui(llm_model, solr_client):
                         "Qualitative URL will appear here...", visible=False)
                     qualitative_data_display = gr.Markdown(
                         "Example data will appear here...", visible=False)
+                with gr.Accordion("Token Usage", open=False):
+                    token_summary_box = gr.Markdown(visible=False)
                 plot_display = gr.Image(
                     label="Visualization", type="filepath", visible=False)
                 report_display = gr.Markdown(
@@ -79,25 +81,28 @@ def create_ui(llm_model, solr_client):
             """
             Manages the conversation and yields UI updates.
             """
+            analysis_plan_input_token_count = analysis_plan_output_token_count = analysis_plan_total_token_count = None
+            enriched_report_input_token_count = enriched_report_output_token_count =  enriched_report_total_token_count = None
+            visualization_input_token_count = visualization_output_token_count = visualization_total_token_count = None
             if state is None:
                 state = {'query_count': 0, 'last_suggestions': []}
             if history is None:
                 history = []
 
             # Reset all displays at the beginning of a new flow
-            yield (history, state, gr.update(value=None, visible=False), gr.update(value=None, visible=False), gr.update(value=None, visible=False), gr.update(value=None, visible=False), gr.update(value=None, visible=False), gr.update(value=None, visible=False), gr.update(value=None, visible=False), gr.update(value="Suggestions from the external API will appear here...", visible=False))
+            yield (history, state, gr.update(value=None, visible=False), gr.update(value=None, visible=False), gr.update(value=None, visible=False), gr.update(value=None, visible=False), gr.update(value=None, visible=False), gr.update(value=None, visible=False), gr.update(value=None, visible=False), gr.update(value=None, visible=False), gr.update(value="Suggestions from the external API will appear here...", visible=False))
 
             query_context = user_input.strip()
             if not query_context:
                 history.append((user_input, "Please enter a question to analyze."))
-                yield (history, state, None, None, None, None, None, None, None, None)
+                yield (history, state, None, None, None, None, None, None, None, None, None)
                 return
 
             history.append((user_input, f"Analyzing: '{query_context}'\n\n*Generating analysis plan...*"))
-            yield (history, state, None, None, None, None, None, None, None, None)
+            yield (history, state, None, None, None, None, None, None, None, None, None)
 
             # Generate plan, get search field suggestions, and intent.
-            analysis_plan, mapped_search_fields, core_name, intent = llm_generate_analysis_plan_with_history(llm_model, query_context, history)
+            analysis_plan, mapped_search_fields, core_name, intent, analysis_plan_input_token_count, analysis_plan_output_token_count, analysis_plan_total_token_count = llm_generate_analysis_plan_with_history(llm_model, query_context, history)
 
             # Update and display search field suggestions in its own accordion
             if mapped_search_fields:
@@ -112,7 +117,7 @@ def create_ui(llm_model, solr_client):
                 else:
                     message = "I'm sorry, I couldn't generate a valid analysis plan. Please try rephrasing your question."
                 history.append((None, message))
-                yield (history, state, None, None, None, None, None, None, None, suggestions_display_update)
+                yield (history, state, None, None, None, None, None, None, None, None, suggestions_display_update)
                 return
 
             history.append((None, f"✅ Analysis plan generated for core: **`{core_name}`**"))
@@ -123,10 +128,10 @@ def create_ui(llm_model, solr_client):
         """
             history.append((None, plan_summary))
             formatted_plan = f"**Full Analysis Plan (Core: `{core_name}`):**\n```json\n{json.dumps(analysis_plan, indent=2)}\n```"
-            yield (history, state, None, None, gr.update(value=formatted_plan, visible=True), None, None, None, None, suggestions_display_update)
+            yield (history, state, None, None, gr.update(value=formatted_plan, visible=True), None, None, None, None, None, suggestions_display_update)
 
             history.append((None, "*Executing queries for aggregates and examples...*"))
-            yield (history, state, None, None, gr.update(value=formatted_plan, visible=True), None, None, None, None, suggestions_display_update)
+            yield (history, state, None, None, gr.update(value=formatted_plan, visible=True), None, None, None, None, None, suggestions_display_update)
 
             # --- DYNAMIC CORE SWITCH (Thread-safe) ---
             with solr_lock:
@@ -153,7 +158,7 @@ def create_ui(llm_model, solr_client):
 
             if not aggregate_data or aggregate_data.get('count', 0) == 0:
                 history.append((None, f"No data was found for your query in the '{core_name}' core. Please try a different question."))
-                yield (history, state, None, None, gr.update(value=formatted_plan, visible=True), None, None, None, None, suggestions_display_update)
+                yield (history, state, None, None, gr.update(value=formatted_plan, visible=True), None, None, None, None, None, suggestions_display_update)
                 return
 
             # Display retrieved data
@@ -162,36 +167,66 @@ def create_ui(llm_model, solr_client):
             formatted_agg_data = f"**Quantitative (Aggregate) Data:**\n```json\n{json.dumps(aggregate_data, indent=2)}\n```"
             formatted_qual_data = f"**Qualitative (Example) Data:**\n```json\n{json.dumps(example_data, indent=2)}\n```"
             qual_data_display_update = gr.update(value=formatted_qual_data, visible=True)
-            yield (history, state, None, None, gr.update(value=formatted_plan, visible=True), quantitative_url_update, gr.update(value=formatted_agg_data, visible=True), qualitative_url_update, qual_data_display_update, suggestions_display_update)
+            yield (history, state, None, None, gr.update(value=formatted_plan, visible=True), quantitative_url_update, gr.update(value=formatted_agg_data, visible=True), qualitative_url_update, qual_data_display_update, None, suggestions_display_update)
 
             history.append((None, "✅ Data retrieved. Generating visualization and final report..."))
-            yield (history, state, None, None, gr.update(value=formatted_plan, visible=True), quantitative_url_update, gr.update(value=formatted_agg_data, visible=True), qualitative_url_update, qual_data_display_update, suggestions_display_update)
-
+            yield (history, state, None, None, gr.update(value=formatted_plan, visible=True), quantitative_url_update, gr.update(value=formatted_agg_data, visible=True), qualitative_url_update, qual_data_display_update, None, suggestions_display_update)
             # Generate viz and report
             with concurrent.futures.ThreadPoolExecutor() as executor:
                 viz_future = executor.submit(llm_generate_visualization_code, llm_model, query_context, aggregate_data)
-
+                viz_code, visualization_input_token_count, visualization_output_token_count, visualization_total_token_count = viz_future.result()
                 report_text = ""
                 stream_history = history[:]
-                for chunk in llm_synthesize_enriched_report_stream(llm_model, query_context, aggregate_data, example_data, analysis_plan):
-                    report_text += chunk
-                    yield (stream_history, state, None, gr.update(value=report_text, visible=True), gr.update(value=formatted_plan, visible=True), quantitative_url_update, gr.update(value=formatted_agg_data, visible=True), qualitative_url_update, qual_data_display_update, suggestions_display_update)
-
+                report_stream = llm_synthesize_enriched_report_stream(llm_model, query_context, aggregate_data, example_data, analysis_plan)
+                for item in report_stream:
+                    if item["text"] is not None:
+                        report_text += item["text"]
+                        yield (stream_history, state, None, gr.update(value=report_text, visible=True), gr.update(value=formatted_plan, visible=True), quantitative_url_update, gr.update(value=formatted_agg_data, visible=True), qualitative_url_update, qual_data_display_update, None, suggestions_display_update)
+                    elif item["tokens"] is not None:
+                        enriched_report_input_token_count = item["tokens"]["input"]
+                        enriched_report_output_token_count = item["tokens"]["output"]
+                        enriched_report_total_token_count = item["tokens"]["total"]
                 history.append((None, report_text))
 
-                viz_code = viz_future.result()
                 plot_path = execute_viz_code_and_get_path(viz_code, aggregate_data)
                 output_plot = gr.update(value=plot_path, visible=True) if plot_path else gr.update(visible=False)
                 if not plot_path:
                     history.append((None, "*I was unable to generate a plot for this data.*\n"))
-
-                yield (history, state, output_plot, gr.update(value=report_text), gr.update(value=formatted_plan, visible=True), quantitative_url_update, gr.update(value=formatted_agg_data, visible=True), qualitative_url_update, qual_data_display_update, suggestions_display_update)
+                
+                cumulative_tokens = sum(filter(None, [
+                    analysis_plan_total_token_count,
+                    enriched_report_total_token_count,
+                    visualization_total_token_count
+                ]))
+
+                total_input = sum(filter(None, [
+                    analysis_plan_input_token_count,
+                    enriched_report_input_token_count,
+                    visualization_input_token_count
+                ]))
+                total_output = sum(filter(None, [
+                    analysis_plan_output_token_count,
+                    enriched_report_output_token_count,
+                    visualization_output_token_count
+                ]))
+                expected_cost = round((total_input*0.3+total_output*2.5)/1000000, 3)
+
+                token_summary_box_update = gr.update(
+                    value=f"""**Analysis Plan Tokens** → Prompt: `{analysis_plan_input_token_count or '-'}`,  Output: `{analysis_plan_output_token_count or '-'}`,  Total: `{analysis_plan_total_token_count or '-'}`  
+                    **Report Tokens** → Prompt: `{enriched_report_input_token_count or '-'}`,  Output: `{enriched_report_output_token_count or '-'}`,  Total: `{enriched_report_total_token_count or '-'}`  
+                    **Visualization Tokens** → Prompt: `{visualization_input_token_count or '-'}`,  Output: `{visualization_output_token_count or '-'}`,  Total: `{visualization_total_token_count or '-'}` 
+                      
+                    **Cumulative Tokens** → `{cumulative_tokens}`  
+                    **Expected Cost** → `{expected_cost}$`""",
+                    visible=True
+                    )
+                yield (history, state, output_plot, gr.update(value=report_text), gr.update(value=formatted_plan, visible=True), quantitative_url_update, gr.update(value=formatted_agg_data, visible=True), qualitative_url_update, qual_data_display_update, token_summary_box_update, suggestions_display_update)        
 
             state['query_count'] += 1
             state['last_suggestions'] = parse_suggestions_from_report(report_text)
             next_prompt = "Analysis complete. What would you like to explore next?"
             history.append((None, next_prompt))
-            yield (history, state, output_plot, gr.update(value=report_text), gr.update(value=formatted_plan, visible=True), quantitative_url_update, gr.update(value=formatted_agg_data, visible=True), qualitative_url_update, qual_data_display_update, suggestions_display_update)
+            yield (history, state, output_plot, gr.update(value=report_text), gr.update(value=formatted_plan, visible=True), quantitative_url_update, gr.update(value=formatted_agg_data, visible=True), qualitative_url_update, qual_data_display_update, token_summary_box_update, suggestions_display_update)
 
         def reset_all():
             """Resets the entire UI for a new analysis session."""
@@ -206,6 +241,7 @@ def create_ui(llm_model, solr_client):
                 gr.update(value=None, visible=False),
                 gr.update(value=None, visible=False),
                 gr.update(value=None, visible=False),
+                gr.update(value=None, visible=False),
                 gr.update(value=None, visible=False)
             )
 
@@ -213,7 +249,7 @@ def create_ui(llm_model, solr_client):
             fn=process_analysis_flow,
             inputs=[msg_textbox, chatbot, state],
             outputs=[chatbot, state, plot_display, report_display, plan_display, quantitative_url_display,
-                     quantitative_data_display, qualitative_url_display, qualitative_data_display, suggestions_display],
+                     quantitative_data_display, qualitative_url_display, qualitative_data_display, token_summary_box, suggestions_display],
         ).then(
             lambda: gr.update(value=""),
             None,
@@ -225,7 +261,7 @@ def create_ui(llm_model, solr_client):
             fn=reset_all,
             inputs=None,
             outputs=[chatbot, state, msg_textbox, plot_display, report_display, plan_display, quantitative_url_display,
-                     quantitative_data_display, qualitative_url_display, qualitative_data_display, suggestions_display],
+                     quantitative_data_display, qualitative_url_display, qualitative_data_display, token_summary_box, suggestions_display],
             queue=False
         )