Use OpenAI API and improve interactivity

Signed-off-by: Aivin V. Solatorio <avsolatorio@gmail.com>

mcp_openai_client.py

@@ -22,7 +22,9 @@ from openai.types.responses import (
     ResponseAudioDeltaEvent,
     ResponseMcpCallCompletedEvent,
     ResponseOutputItemDoneEvent,
+    ResponseOutputItemAddedEvent,
 )
+import ast
 
 load_dotenv()
 
@@ -67,10 +69,17 @@ When responding you must always plan the steps and enumerate all the tools that
    - Summarize the data in a table format with clear column names and values.
    - If the data is not available, respond by clearly stating that you do not have access to that information.
 
+7. **Tool Use**:
+   - Fetch each indicator data using independent tool calls.
+   - Provide some brief explanation between tool calls.
+
 Stay strictly within these boundaries while maintaining a helpful and respectful tone."""
 
 
 LLM_MODEL = "claude-3-5-haiku-20241022"
+OPENAI_MODEL = "gpt-4.1"
+# OPENAI_MODEL = "gpt-4.1-mini"
+# OPENAI_MODEL = "gpt-4.1-nano"
 # What is the military spending of bangladesh in 2014?
 # When a tool is needed for any step, ensure to add the token `TOOL_USE`.
 
@@ -202,9 +211,10 @@ class MCPClientWrapper:
                 else:
                     is_delta = False
                     messages.extend(partial)
+                    print(partial)
 
                 yield messages, gr.Textbox(value="")
-                await asyncio.sleep(0.05)
+                await asyncio.sleep(0.01)
 
                 if (
                     messages[-1]["role"] == "assistant"
@@ -220,7 +230,7 @@ class MCPClientWrapper:
         self, message: str, history: List[Union[Dict[str, Any], ChatMessage]]
     ):
         response = self.openai.responses.create(
-            model="gpt-4.1-mini",
+            model=OPENAI_MODEL,
             tools=[
                 {
                     "type": "mcp",
@@ -241,24 +251,100 @@ class MCPClientWrapper:
         )
 
         is_tool_call = False
+        tool_name = None
+        tool_args = None
         for event in response:
-            if isinstance(event, ResponseMcpCallInProgressEvent):
+            if (
+                isinstance(event, ResponseOutputItemAddedEvent)
+                and event.item.type == "mcp_call"
+            ):
                 is_tool_call = True
-                yield [
-                    {
-                        "role": "assistant",
-                        "content": "I'll use the tool to help answer your question.",
-                    }
-                ]
-            elif isinstance(event, ResponseOutputItemDoneEvent):
-                if is_tool_call:
+                tool_name = event.item.name
+            # if isinstance(event, ResponseMcpCallInProgressEvent):
+            #     is_tool_call = True
+            #     yield [
+            #         {
+            #             "role": "assistant",
+            #             "content": "I'll use the tool to help answer your question.",
+            #         }
+            #     ]
+            if is_tool_call:
+                if (
+                    isinstance(event, ResponseAudioDeltaEvent)
+                    and event.type == "response.mcp_call_arguments.done"
+                ):
+                    tool_args = event.arguments
+
+                    try:
+                        tool_args = json.dumps(
+                            json.loads(tool_args), ensure_ascii=True, indent=2
+                        )
+                    except:
+                        pass
+
+                    yield [
+                        {
+                            "role": "assistant",
+                            "content": f"I'll use the {tool_name} tool to help answer your question.",
+                            "metadata": {
+                                "title": f"Using tool: {tool_name.replace('avsolatorio_test_data_mcp_server', '')}",
+                                "log": f"Parameters: {tool_args}",
+                                # "status": "pending",
+                                "status": "done",
+                                "id": f"tool_call_{tool_name}",
+                            },
+                        }
+                    ]
+
                     yield [
                         {
                             "role": "assistant",
-                            "content": "I've used the tool to help answer your question.",
+                            "content": "```json\n" + tool_args + "\n```",
+                            "metadata": {
+                                "parent_id": f"tool_call_{tool_name}",
+                                "id": f"params_{tool_name}",
+                                "title": "Tool Parameters",
+                            },
                         }
                     ]
-                    is_tool_call = False
+
+                elif isinstance(event, ResponseOutputItemDoneEvent):
+                    if event.item.type == "mcp_call":
+                        yield [
+                            {
+                                "role": "assistant",
+                                "content": "Here are the results from the tool:",
+                                "metadata": {
+                                    "title": f"Tool Result for {tool_name.replace('avsolatorio_test_data_mcp_server', '')}",
+                                    "status": "done",
+                                    "id": f"result_{tool_name}",
+                                },
+                            }
+                        ]
+
+                        result_content = event.item.output
+                        if result_content.startswith("root="):
+                            result_content = result_content[5:]
+                            try:
+                                result_content = ast.literal_eval(result_content)
+                                result_content = json.dumps(result_content, indent=2)
+                            except:
+                                pass
+
+                        yield [
+                            {
+                                "role": "assistant",
+                                "content": "```\n" + result_content + "\n```",
+                                "metadata": {
+                                    "parent_id": f"result_{tool_name}",
+                                    "id": f"raw_result_{tool_name}",
+                                    "title": "Raw Output",
+                                },
+                            }
+                        ]
+                        is_tool_call = False
+                        tool_name = None
+                        tool_args = None
             elif isinstance(event, ResponseTextDeltaEvent):
                 yield [{"role": "assistant", "content": None, "delta": event.delta}]
 
@@ -509,7 +595,7 @@ def gradio_interface(
     # Disable auto-dark mode by setting theme to None
     with gr.Blocks(title="WDI MCP Client", css=custom_css, theme=None) as demo:
         try:
-            gr.Markdown("# Development Data Chat")
+            gr.Markdown("# Data360 Chat [Prototype]")
             # gr.Markdown("Connect to the WDI MCP server and chat with the assistant")
 
             with gr.Accordion(