Bootstrap openai client

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

mcp_openai_client.py

@@ -12,10 +12,23 @@ from mcp.client.sse import sse_client
 from anthropic import Anthropic
 from anthropic._exceptions import OverloadedError
 from dotenv import load_dotenv
-
+from openai import OpenAI
+from openai.types.responses import (
+    ResponseTextDeltaEvent,
+    ResponseContentPartAddedEvent,
+    ResponseContentPartDoneEvent,
+    ResponseTextDoneEvent,
+    ResponseMcpCallInProgressEvent,
+    ResponseAudioDeltaEvent,
+    ResponseMcpCallCompletedEvent,
+    ResponseOutputItemDoneEvent,
+)
 
 load_dotenv()
 
+# LLM_PROVIDER = "anthropic"
+LLM_PROVIDER = "openai"
+
 SYSTEM_PROMPT = f"""You are a helpful assistant. Today is {datetime.now().strftime("%Y-%m-%d")}.
 
 You **do not** have prior knowledge of the World Development Indicators (WDI) data. Instead, you must rely entirely on the tools available to you to answer the user's questions.
@@ -71,6 +84,7 @@ class MCPClientWrapper:
         self.session = None
         self.exit_stack = None
         self.anthropic = Anthropic()
+        self.openai = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
         self.tools = []
 
     async def connect(self, server_path_or_url: str) -> str:
@@ -150,7 +164,7 @@ class MCPClientWrapper:
     async def process_message(
         self, message: str, history: List[Union[Dict[str, Any], ChatMessage]]
     ):
-        if not self.session:
+        if not self.session and LLM_PROVIDER == "anthropic":
             messages = history + [
                 {"role": "user", "content": message},
                 {
@@ -173,8 +187,22 @@ class MCPClientWrapper:
             await asyncio.sleep(0.1)
             messages.pop(-1)
 
+            is_delta = False
             async for partial in self._process_query(message, history):
-                messages.extend(partial)
+                if partial[-1].get("delta"):
+                    if not is_delta:
+                        is_delta = True
+                        messages.append(
+                            {
+                                "role": "assistant",
+                                "content": "",
+                            }
+                        )
+                    messages[-1]["content"] += partial[-1]["delta"]
+                else:
+                    is_delta = False
+                    messages.extend(partial)
+
                 yield messages, gr.Textbox(value="")
                 await asyncio.sleep(0.05)
 
@@ -188,7 +216,53 @@ class MCPClientWrapper:
         with open("messages.log.jsonl", "a+") as fl:
             fl.write(json.dumps(dict(time=f"{datetime.now()}", messages=messages)))
 
-    async def _process_query(
+    async def _process_query_openai(
+        self, message: str, history: List[Union[Dict[str, Any], ChatMessage]]
+    ):
+        response = self.openai.responses.create(
+            model="gpt-4.1-mini",
+            tools=[
+                {
+                    "type": "mcp",
+                    "server_label": "wdi_mcp",
+                    "server_url": "https://avsolatorio-test-data-mcp-server.hf.space/gradio_api/mcp/sse",
+                    "require_approval": "never",
+                    "headers": {"Authorization": f"Bearer {os.getenv('HF_TOKEN')}"},
+                    # "server_token": userdata.get('MCP_HF_TOKEN'),
+                },
+            ],
+            # input="What transport protocols are supported in the 2025-03-26 version of the MCP spec?",
+            instructions=SYSTEM_PROMPT,
+            # input="What is the gdp of india in 2020?",
+            input=message,
+            parallel_tool_calls=False,
+            stream=True,
+            temperature=0,
+        )
+
+        is_tool_call = False
+        for event in response:
+            if isinstance(event, ResponseMcpCallInProgressEvent):
+                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:
+                    yield [
+                        {
+                            "role": "assistant",
+                            "content": "I've used the tool to help answer your question.",
+                        }
+                    ]
+                    is_tool_call = False
+            elif isinstance(event, ResponseTextDeltaEvent):
+                yield [{"role": "assistant", "content": None, "delta": event.delta}]
+
+    async def _process_query_anthropic(
         self, message: str, history: List[Union[Dict[str, Any], ChatMessage]]
     ):
         claude_messages = []
@@ -292,10 +366,10 @@ class MCPClientWrapper:
                     result_messages.append(
                         {
                             "role": "assistant",
-                            "content": f"Sorry, I encountered an error while calling the tool: {error_msg}. Please try again.",
+                            "content": f"Sorry, I encountered an error while calling the tool: {error_msg}. Please try again or reload the page.",
                             "metadata": {
                                 "title": f"Tool Error for {tool_name.replace('avsolatorio_test_data_mcp_server', '')}",
-                                "status": "error",
+                                "status": "done",
                                 "id": f"error_{tool_name}",
                             },
                         }
@@ -392,6 +466,16 @@ class MCPClientWrapper:
 
                 contents.extend(next_response.content)
 
+    async def _process_query(
+        self, message: str, history: List[Union[Dict[Any, Any], ChatMessage]]
+    ):
+        if LLM_PROVIDER == "anthropic":
+            async for partial in self._process_query_anthropic(message, history):
+                yield partial
+        elif LLM_PROVIDER == "openai":
+            async for partial in self._process_query_openai(message, history):
+                yield partial
+
 
 def gradio_interface(
     server_path_or_url: str = "https://avsolatorio-test-data-mcp-server.hf.space/gradio_api/mcp/sse",
@@ -467,12 +551,13 @@ def gradio_interface(
 
             # connect_btn.click(client.connect, inputs=server_path, outputs=status)
             # Automatically call client.connect(...) as soon as the interface loads
-            demo.load(
-                fn=client.connect,
-                inputs=server_path,
-                outputs=status,
-                show_progress="full",
-            )
+            if LLM_PROVIDER == "anthropic":
+                demo.load(
+                    fn=client.connect,
+                    inputs=server_path,
+                    outputs=status,
+                    show_progress="full",
+                )
 
             msg.submit(
                 client.process_message,
@@ -483,8 +568,9 @@ def gradio_interface(
             # clear_btn.click(lambda: [], None, chatbot)
 
         except KeyboardInterrupt:
-            print("Keyboard interrupt received. Disconnecting from MCP server...")
-            asyncio.run(client.disconnect())
+            if LLM_PROVIDER == "anthropic":
+                print("Keyboard interrupt received. Disconnecting from MCP server...")
+                asyncio.run(client.disconnect())
             raise KeyboardInterrupt
         # demo.unload(client.disconnect)