Update app.py

app.py

@@ -172,18 +172,20 @@ def generate(
                     call_idx = output[generation_start_idx:].find("<call>")
                 if response_idx == -1:
                     response_idx = output[generation_start_idx:].find("<response>")
+                    if response_idx != -1:
+                        response_idx += len("<response>")
                 if submit_idx == -1:
                     submit_idx = output[generation_start_idx:].find("<submit>")
                 # I am sorry about the code
                 print("-------", generation_start_idx, request_idx, call_idx, response_idx)
                 highlighted_output = [
-                    (prompt, ""),
-                    (output[generation_start_idx:], "") if request_idx == -1 else ("", ""),
-                    (output[generation_start_idx:generation_start_idx+request_idx], ""),
-                    (output[generation_start_idx+request_idx:], "") if call_idx == -1 else  ("", ""),
-                    (output[generation_start_idx+request_idx:generation_start_idx+call_idx], "request"),
-                    (output[generation_start_idx+call_idx:generation_start_idx+response_idx], "call"),
-                    (output[generation_start_idx+response_idx:], "submit") if submit_idx != -1 else ("", ""),
+                    (prompt, "query"),
+                    (output[generation_start_idx:], "model") if request_idx == -1 else ("", ""),
+                    (output[generation_start_idx:generation_start_idx+request_idx], "model"),
+                    (output[generation_start_idx+request_idx:], "model") if call_idx == -1 else  ("", ""),
+                    (output[generation_start_idx+request_idx:generation_start_idx+call_idx], "tool request"),
+                    (output[generation_start_idx+call_idx:generation_start_idx+response_idx], "tool call"),
+                    (output[generation_start_idx+response_idx:], "model") if submit_idx != -1 else ("", ""),
                     # (output[generation_start_idx:generation_start_idx+request_idx], ""),
                     # (output[generation_start_idx+request_idx:generation_start_idx+call_idx], "request"),
                     # (output[generation_start_idx+call_idx:], "call"),
@@ -288,7 +290,7 @@ with gr.Blocks(theme=theme, analytics_enabled=False, css=css) as demo:
                 
                 output = gr.HighlightedText(
                     label="Output",
-                    color_map={"query": "red", "call": "green", "response": "blue", "submit": "yellow", "model": "pink"},
+                    color_map={"query": "red", "tool call": "green", "tool response": "blue", "model": "pink"},
                 )
 
                 gr_examples = gr.Examples(