Several UI improvements.

app.py

@@ -2,8 +2,6 @@ import json
 
 import gradio as gr
 
-from collections import UserList
-
 from flow import full_flow
 
 from utils import use_cost_tracker, get_costs, compute_hash
@@ -40,11 +38,18 @@ def parse_raw_messages(messages_raw):
 
 def main():
     with gr.Blocks() as demo:
-        gr.Markdown("### Agora Demo")
+        gr.Markdown("# Agora Demo")
         gr.Markdown("We will create a new Agora channel and offer it to Alice as a tool.")
 
-        chosen_task = gr.Dropdown(choices=list(SCHEMAS.keys()), label="Schema", value="weather_forecast")
-        custom_task = gr.Checkbox(label="Custom Task")
+        chosen_task = gr.Dropdown(choices=[
+            (v['display_name'], k) for k, v in SCHEMAS.items()
+        ], label="Schema", value="weather_forecast")
+
+        @gr.render(inputs=[chosen_task])
+        def render2(chosen_task):
+            gr.Markdown('**Description**: ' + SCHEMAS[chosen_task]["description"])
+
+        custom_task = gr.Checkbox(label="Override Demo Parameters")
 
         STATE_TRACKER = {}
         
@@ -52,13 +57,15 @@ def main():
         def render(chosen_task, custom_task):
             if STATE_TRACKER.get('chosen_task') != chosen_task:
                 STATE_TRACKER['chosen_task'] = chosen_task
-                for k, v in SCHEMAS[chosen_task].items():
+                for k, v in SCHEMAS[chosen_task]['schema'].items():
                     if isinstance(v, str):
                         STATE_TRACKER[k] = v
                     else:
                         STATE_TRACKER[k] = json.dumps(v, indent=2)
 
             if custom_task:
+                gr.Markdown('#### Custom Demo Parameters')
+                gr.Markdown('You can override the default parameters for the demo. Note: recommended for advanced users only.')
                 gr.Text(label="Description", value=STATE_TRACKER["description"], interactive=True).change(lambda x: STATE_TRACKER.update({'description': x}))
                 gr.TextArea(label="Input Schema", value=STATE_TRACKER["input"], interactive=True).change(lambda x: STATE_TRACKER.update({'input': x}))
                 gr.TextArea(label="Output Schema", value=STATE_TRACKER["output"], interactive=True).change(lambda x: STATE_TRACKER.update({'output': x}))
@@ -68,7 +75,8 @@ def main():
         model_options = [
             ('GPT 4o (Camel AI)', 'gpt-4o'),
             ('GPT 4o-mini (Camel AI)', 'gpt-4o-mini'),
-            ('Claude 3 Sonnet (LangChain)', 'claude-3-sonnet'),
+            ('Claude 3 Sonnet (LangChain)', 'claude-3-5-sonnet-latest'),
+            ('Claude 3 Haiku (LangChain)', 'claude-3-5-haiku-latest'),
             ('Gemini 1.5 Pro (Google GenAI)', 'gemini-1.5-pro'),
             ('Llama3 405B (Sambanova + LangChain)', 'llama3-405b')
         ]
@@ -90,59 +98,100 @@ def main():
             with gr.Column(scale=1):
                 bob_model_dd = gr.Dropdown(label="Bob Model", choices=model_options, value="gpt-4o")
         
-        button = gr.Button('Start', elem_id='start_button')
-        gr.Markdown('### Natural Language')
+        
 
         @gr.render(inputs=[alice_model_dd, bob_model_dd])
         def render_with_images(alice_model, bob_model):
+            button = gr.Button('Start', elem_id='start_button')
+            gr.Markdown('## Natural Language')
+
             avatar_images = [images.get(alice_model, fallback_image), images.get(bob_model, fallback_image)]
             chatbot_nl = gr.Chatbot(type="messages", avatar_images=avatar_images)
 
             with gr.Accordion(label="Raw Messages", open=False):
                 chatbot_nl_raw = gr.Chatbot(type="messages", avatar_images=avatar_images)
         
-            gr.Markdown('### Negotiation')
+            gr.Markdown('## Negotiation')
             chatbot_negotiation = gr.Chatbot(type="messages", avatar_images=avatar_images)
 
-            gr.Markdown('### Protocol')
+            gr.Markdown('## Protocol')
+            protocol_hash_result = gr.Text(interactive=False, label="Protocol Hash")
             protocol_result = gr.TextArea(interactive=False, label="Protocol")
 
-            gr.Markdown('### Implementation')
+            gr.Markdown('## Implementation')
             with gr.Row():
                 with gr.Column(scale=1):
                     alice_implementation = gr.TextArea(interactive=False, label="Alice Implementation")
                 with gr.Column(scale=1):
                     bob_implementation = gr.TextArea(interactive=False, label="Bob Implementation")
 
-            gr.Markdown('### Structured Communication')
+            gr.Markdown('## Structured Communication')
             structured_communication = gr.Chatbot(type="messages", avatar_images=avatar_images)
             
             with gr.Accordion(label="Raw Messages", open=False):
                 structured_communication_raw = gr.Chatbot(type="messages", avatar_images=avatar_images)
+            
+            gr.Markdown('## Cost')
+            cost_info = gr.State(value=None)
+            #cost_info = gr.TextArea(interactive=False, label="Cost")
+
+            query_slider = gr.Slider(label="Number of queries", minimum=1, maximum=10_000, step=1, value=50, interactive=True)
+            cost_display = gr.Markdown('')
+            
 
-            def respond(chosen_task, custom_task, alice_model, bob_model):
-                yield gr.update(interactive=False), gr.update(interactive=False), gr.update(interactive=False), \
-                    None, None, None, None, None, None, None, None
-
-                if custom_task:
-                    schema = dict(STATE_TRACKER)
-                    for k, v in schema.items():
-                        if isinstance(v, str):
-                            try:
-                                schema[k] = json.loads(v)
-                            except:
-                                pass
+            def render_info(query_count, cost_info):
+                if not cost_info:
+                    return ''
+                natural_cost = cost_info['conversation'] * query_count
+                agora_cost = cost_info['negotiation'] + cost_info['programming']
+
+                cost_message = f'Cost of one natural language conversation: {cost_info["conversation"]:.4f} USD\n\n'
+                cost_message += f'Cost of negotiating the protocol: {cost_info["negotiation"]:.4f} USD\n\n'
+                cost_message += f'Cost of implementing the protocol: {cost_info["programming"]:.4f} USD\n\n'
+                cost_message += f'Cost of {query_count} queries with natural language: {natural_cost:.4f} USD\n\n'
+                cost_message += f'Cost of {query_count} queries with Agora: {agora_cost:.4f} USD\n\n'
+
+                if natural_cost < agora_cost:
+                    factor = agora_cost / natural_cost
+                    cost_message += f'Natural language is {factor:.2f}x cheaper than Agora.'
                 else:
-                    schema = SCHEMAS[chosen_task]
+                    factor = natural_cost / agora_cost
+                    cost_message += f'Agora is {factor:.2f}x cheaper than natural language.'
+
+                return cost_message
+
+            cost_info.change(render_info, [query_slider, cost_info], [cost_display])
+            query_slider.change(render_info, [query_slider, cost_info], [cost_display])
+
+
+            def respond(chosen_task, custom_task, alice_model, bob_model, query_count):
+                with use_cost_tracker():
+                    yield gr.update(interactive=False), gr.update(interactive=False), gr.update(interactive=False), \
+                        None, None, None, None, None, None, None, None, None, None, None
+
+                    if custom_task:
+                        schema = dict(STATE_TRACKER)
+                        for k, v in schema.items():
+                            if isinstance(v, str):
+                                try:
+                                    schema[k] = json.loads(v)
+                                except:
+                                    pass
+                    else:
+                        schema = SCHEMAS[chosen_task]["schema"]
+
+                    for nl_messages_raw, negotiation_messages, structured_messages_raw, protocol, alice_implementation, bob_implementation in full_flow(schema, alice_model, bob_model):
+                        nl_messages_clean, nl_messages_agora = parse_raw_messages(nl_messages_raw)
+                        structured_messages_clean, structured_messages_agora = parse_raw_messages(structured_messages_raw)
+                        protocol_hash = compute_hash(protocol) if protocol else None
+                        yield gr.update(), gr.update(), gr.update(), None, None, nl_messages_clean, nl_messages_agora, negotiation_messages, structured_messages_clean, structured_messages_agora, protocol, protocol_hash, alice_implementation, bob_implementation
 
-                for nl_messages_raw, negotiation_messages, structured_messages_raw, protocol, alice_implementation, bob_implementation in full_flow(schema, alice_model, bob_model):
-                    nl_messages_clean, nl_messages_agora = parse_raw_messages(nl_messages_raw)
-                    structured_messages_clean, structured_messages_agora = parse_raw_messages(structured_messages_raw)
+                    #yield from full_flow(schema, alice_model, bob_model)
 
-                    yield gr.update(), gr.update(), gr.update(), nl_messages_clean, nl_messages_agora, negotiation_messages, structured_messages_clean, structured_messages_agora, protocol, alice_implementation, bob_implementation
+                    cost_data = get_costs()
+                    cost_data_formatted = render_info(query_count, cost_data)
 
-                #yield from full_flow(schema, alice_model, bob_model)
-                yield gr.update(interactive=True), gr.update(interactive=True), gr.update(interactive=True), gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), gr.update()
+                    yield gr.update(interactive=True), gr.update(interactive=True), gr.update(interactive=True), cost_data, cost_data_formatted, gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), gr.update()
 
             button.click(respond, [chosen_task, custom_task, alice_model_dd, bob_model_dd], [button, alice_model_dd, bob_model_dd, chatbot_nl, chatbot_nl_raw, chatbot_negotiation, structured_communication, structured_communication_raw, protocol_result, alice_implementation, bob_implementation])