avsolatorio commited on
Commit
feb936d
·
verified ·
1 Parent(s): 11faf9a

Use OpenAI service for chat and other fixes (#1)

Browse files

- Duplicate code (8cd6fde1b446eedcc7fb93e454909012aa44a3e1)
- Update chat name (0c92aa0549c831e54b6ff5dda16f6a13b98dc5be)
- Bootstrap openai client (c8120990b49188805a48dd9b35cd250722ea3308)
- Use OpenAI API and improve interactivity (48042197b62244ebb703b1931d1f2778cda2694f)
- Update requirements (33f4b8c7893308049ff86b0da3a5376c31ae0268)
- Add history (8362096a1d78aa658553bfe16002ab76fa0b132a)
- Implement proper buffering of the output (e8809f1d039d7a252ca230042d59a3eb268a6482)
- Update language handling and store logs (0e2092543ad57d4d1863a677fec5a761e935a76e)
- Handle case where the data for the country is empty. (5eb37d7a551259ccf4d20781f703d981ab8cee8f)

Files changed (5) hide show
  1. mcp_openai_client.py +758 -0
  2. mcp_remote_client.py +1 -1
  3. pyproject.toml +1 -0
  4. services.py +8 -0
  5. uv.lock +21 -0
mcp_openai_client.py ADDED
@@ -0,0 +1,758 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import asyncio
2
+ import os
3
+ import json
4
+ from typing import List, Dict, Any, Union
5
+ from contextlib import AsyncExitStack
6
+ from datetime import datetime
7
+ import gradio as gr
8
+ from gradio.components.chatbot import ChatMessage
9
+ from mcp import ClientSession, StdioServerParameters
10
+ from mcp.client.stdio import stdio_client
11
+ from mcp.client.sse import sse_client
12
+ from anthropic import Anthropic
13
+ from anthropic._exceptions import OverloadedError
14
+ from dotenv import load_dotenv
15
+ from openai import OpenAI
16
+ import openai
17
+ from openai.types.responses import (
18
+ ResponseTextDeltaEvent,
19
+ ResponseContentPartAddedEvent,
20
+ ResponseContentPartDoneEvent,
21
+ ResponseTextDoneEvent,
22
+ ResponseMcpCallInProgressEvent,
23
+ ResponseAudioDeltaEvent,
24
+ ResponseMcpCallCompletedEvent,
25
+ ResponseOutputItemDoneEvent,
26
+ ResponseOutputItemAddedEvent,
27
+ ResponseCompletedEvent,
28
+ )
29
+ import ast
30
+
31
+ load_dotenv()
32
+
33
+ # LLM_PROVIDER = "anthropic"
34
+ LLM_PROVIDER = "openai"
35
+
36
+ SYSTEM_PROMPT = f"""You are a helpful assistant. Today is {datetime.now().strftime("%Y-%m-%d")}.
37
+
38
+ 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.
39
+
40
+ Detect the language of the user's query and use that language for your response, unless the user specifies otherwise.
41
+
42
+ When responding you must always plan the steps and enumerate all the tools that you plan to use to answer the user's query.
43
+
44
+ ### Your Instructions:
45
+
46
+ 1. **Tool Use Only**:
47
+ - You must not provide any answers based on prior knowledge or assumptions.
48
+ - You must **not** fabricate data or simulate the behavior of the `get_wdi_data` tool.
49
+ - You cannot use the `get_wdi_data` tool without using the `search_relevant_indicators` tool first.
50
+ - If the user requests WDI data, you **MUST ALWAYS** first call the `search_relevant_indicators` tool to see if there's any relevant data.
51
+ - If relevant data exists, call the `get_wdi_data` tool to get the data.
52
+
53
+ 2. **Tool Invocation**:
54
+ - Use any relevant tools provided to you to answer the user's question.
55
+ - You may call multiple tools if needed, and you should do so in a logical sequence to minimize unnecessary user interaction.
56
+ - Do not hesitate to invoke tools as soon as they are relevant.
57
+
58
+ 3. **Limitations**:
59
+ - If a user request cannot be fulfilled using the tools available, respond by clearly stating that you do not have access to that information.
60
+
61
+ 4. **Ethical Guidelines**:
62
+ - Do not make or endorse statements based on stereotypes, bias, or assumptions.
63
+ - Ensure all claims and explanations are grounded in the data or factual evidence retrieved via tools.
64
+ - Politely refuse to respond to requests that involve stereotypes or unfounded generalizations.
65
+
66
+ 5. **Communication Style**:
67
+ - Present the data in clear, user-friendly language.
68
+ - You may summarize or explain the data retrieved, but do **not** elaborate based on outside or implicit knowledge.
69
+ - You may describe the data in a way that is easy to understand but you MUST NOT elaborate based on external knowledge.
70
+ - Provide summary of the answer in the last step describing some observations and insights solely based on the data.
71
+
72
+ 6. **Presentation**:
73
+ - Present the data in a way that is easy to understand.
74
+ - Summarize the data in a table format with clear column names and values.
75
+ - If the data is not available, respond by clearly stating that you do not have access to that information.
76
+
77
+ 7. **Tool Use**:
78
+ - Fetch each indicator data using independent tool calls.
79
+ - Provide some brief explanation between tool calls.
80
+
81
+ Stay strictly within these boundaries while maintaining a helpful and respectful tone."""
82
+
83
+
84
+ LLM_MODEL = "claude-3-5-haiku-20241022"
85
+ OPENAI_MODEL = "gpt-4.1"
86
+ # OPENAI_MODEL = "gpt-4.1-mini"
87
+ # OPENAI_MODEL = "gpt-4.1-nano"
88
+ # What is the military spending of bangladesh in 2014?
89
+ # When a tool is needed for any step, ensure to add the token `TOOL_USE`.
90
+
91
+
92
+ loop = asyncio.new_event_loop()
93
+ asyncio.set_event_loop(loop)
94
+
95
+
96
+ class MCPClientWrapper:
97
+ def __init__(self):
98
+ self.session = None
99
+ self.exit_stack = None
100
+ self.anthropic = Anthropic()
101
+ self.openai = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
102
+ self.tools = []
103
+
104
+ async def connect(self, server_path_or_url: str) -> str:
105
+ try:
106
+ # If there's an existing session, close it
107
+ if self.exit_stack:
108
+ return "Already connected to an MCP server. Please disconnect first."
109
+ # await self.exit_stack.aclose()
110
+
111
+ self.exit_stack = AsyncExitStack()
112
+
113
+ if server_path_or_url.endswith(".py"):
114
+ command = "python"
115
+
116
+ server_params = StdioServerParameters(
117
+ command=command,
118
+ args=[server_path_or_url],
119
+ env={"PYTHONIOENCODING": "utf-8", "PYTHONUNBUFFERED": "1"},
120
+ )
121
+
122
+ print(
123
+ f"Starting MCP server with command: {command} {server_path_or_url}"
124
+ )
125
+ # Launch MCP subprocess and bind streams on the current running loop
126
+ stdio_transport = await self.exit_stack.enter_async_context(
127
+ stdio_client(server_params)
128
+ )
129
+ self.stdio, self.write = stdio_transport
130
+ else:
131
+ print(f"Connecting to MCP server at: {server_path_or_url}")
132
+ sse_transport = await self.exit_stack.enter_async_context(
133
+ sse_client(
134
+ server_path_or_url,
135
+ headers={"Authorization": f"Bearer {os.getenv('HF_TOKEN')}"},
136
+ )
137
+ )
138
+ self.stdio, self.write = sse_transport
139
+
140
+ print("Creating MCP client session...")
141
+ # Create ClientSession on this same loop
142
+ self.session = await self.exit_stack.enter_async_context(
143
+ ClientSession(self.stdio, self.write)
144
+ )
145
+ await self.session.initialize()
146
+ print("MCP session initialized successfully")
147
+
148
+ response = await self.session.list_tools()
149
+ self.tools = [
150
+ {
151
+ "name": tool.name,
152
+ "description": tool.description,
153
+ "input_schema": tool.inputSchema,
154
+ }
155
+ for tool in response.tools
156
+ ]
157
+
158
+ print("Available tools:", self.tools)
159
+ tool_names = [tool["name"] for tool in self.tools]
160
+ return f"Connected to MCP server. Available tools: {', '.join(tool_names)}"
161
+ except Exception as e:
162
+ error_msg = f"Failed to connect to MCP server: {str(e)}"
163
+ print(error_msg)
164
+ # Clean up on error
165
+ if self.exit_stack:
166
+ await self.exit_stack.aclose()
167
+ self.exit_stack = None
168
+ self.session = None
169
+ return error_msg
170
+
171
+ async def disconnect(self):
172
+ if self.exit_stack:
173
+ print("Disconnecting from MCP server...")
174
+ await self.exit_stack.aclose()
175
+ self.exit_stack = None
176
+ self.session = None
177
+
178
+ async def process_message(
179
+ self,
180
+ message: str,
181
+ history: List[Union[Dict[str, Any], ChatMessage]],
182
+ previous_response_id: str = None,
183
+ ):
184
+ if not self.session and LLM_PROVIDER == "anthropic":
185
+ messages = history + [
186
+ {"role": "user", "content": message},
187
+ {
188
+ "role": "assistant",
189
+ "content": "Please connect to an MCP server first by reloading the page.",
190
+ },
191
+ ]
192
+ yield messages, gr.Textbox(value=""), gr.Textbox(value=previous_response_id)
193
+ else:
194
+ messages = history + [
195
+ {"role": "user", "content": message},
196
+ {
197
+ "role": "assistant",
198
+ "content": "Ok, let me think about your query 🤔...",
199
+ },
200
+ ]
201
+
202
+ yield messages, gr.Textbox(value=""), gr.Textbox(value=previous_response_id)
203
+ # simulate thinking with asyncio.sleep
204
+ await asyncio.sleep(0.1)
205
+ messages.pop(-1)
206
+
207
+ is_delta = False
208
+ async for partial in self._process_query(
209
+ message, history, previous_response_id
210
+ ):
211
+ if partial[-1].get("delta"):
212
+ if not is_delta:
213
+ is_delta = True
214
+ messages.append(
215
+ {
216
+ "role": "assistant",
217
+ "content": "",
218
+ }
219
+ )
220
+ messages[-1]["content"] += partial[-1]["delta"]
221
+ if partial[-1].get("status") == "done":
222
+ await asyncio.sleep(0.05)
223
+ else:
224
+ is_delta = False
225
+ if partial[-1].get("response_id"):
226
+ previous_response_id = partial[-1]["response_id"]
227
+ yield (
228
+ messages,
229
+ gr.Textbox(value=""),
230
+ gr.Textbox(value=previous_response_id),
231
+ )
232
+ await asyncio.sleep(0.01)
233
+ continue
234
+ else:
235
+ messages.extend(partial)
236
+ print(partial)
237
+
238
+ yield (
239
+ messages,
240
+ gr.Textbox(value=""),
241
+ gr.Textbox(value=previous_response_id),
242
+ )
243
+ await asyncio.sleep(0.01)
244
+
245
+ if (
246
+ messages[-1]["role"] == "assistant"
247
+ and messages[-1]["content"]
248
+ == "The LLM API is overloaded now, try again later..."
249
+ ):
250
+ break
251
+
252
+ with open("messages.log.jsonl", "a+") as fl:
253
+ fl.write(json.dumps(dict(time=f"{datetime.now()}", messages=messages)))
254
+
255
+ async def _process_query_openai(
256
+ self,
257
+ message: str,
258
+ history: List[Union[Dict[str, Any], ChatMessage]],
259
+ previous_response_id: str = None,
260
+ ):
261
+ response = self.openai.responses.create(
262
+ model=OPENAI_MODEL,
263
+ tools=[
264
+ {
265
+ "type": "mcp",
266
+ "server_label": "wdi_mcp",
267
+ "server_url": "https://avsolatorio-test-data-mcp-server.hf.space/gradio_api/mcp/sse",
268
+ "require_approval": "never",
269
+ "headers": {"Authorization": f"Bearer {os.getenv('HF_TOKEN')}"},
270
+ # "server_token": userdata.get('MCP_HF_TOKEN'),
271
+ },
272
+ ],
273
+ # input="What transport protocols are supported in the 2025-03-26 version of the MCP spec?",
274
+ instructions=SYSTEM_PROMPT,
275
+ # input="What is the gdp of india in 2020?",
276
+ input=message,
277
+ parallel_tool_calls=False,
278
+ stream=True,
279
+ max_output_tokens=32768,
280
+ temperature=0,
281
+ previous_response_id=previous_response_id
282
+ if previous_response_id.strip()
283
+ else None,
284
+ store=True, # Store the response in the OpenAIlogs
285
+ )
286
+
287
+ is_tool_call = False
288
+ tool_name = None
289
+ tool_args = None
290
+ for event in response:
291
+ if isinstance(event, ResponseCompletedEvent):
292
+ yield [
293
+ {
294
+ "response_id": event.response.id,
295
+ }
296
+ ]
297
+ elif (
298
+ isinstance(event, ResponseOutputItemAddedEvent)
299
+ and event.item.type == "mcp_call"
300
+ ):
301
+ is_tool_call = True
302
+ tool_name = event.item.name
303
+ # if isinstance(event, ResponseMcpCallInProgressEvent):
304
+ # is_tool_call = True
305
+ # yield [
306
+ # {
307
+ # "role": "assistant",
308
+ # "content": "I'll use the tool to help answer your question.",
309
+ # }
310
+ # ]
311
+ if is_tool_call:
312
+ if (
313
+ isinstance(event, ResponseAudioDeltaEvent)
314
+ and event.type == "response.mcp_call_arguments.done"
315
+ ):
316
+ tool_args = event.arguments
317
+
318
+ try:
319
+ tool_args = json.dumps(
320
+ json.loads(tool_args), ensure_ascii=True, indent=2
321
+ )
322
+ except:
323
+ pass
324
+
325
+ yield [
326
+ {
327
+ "role": "assistant",
328
+ "content": f"I'll use the {tool_name} tool to help answer your question.",
329
+ "metadata": {
330
+ "title": f"Using tool: {tool_name.replace('avsolatorio_test_data_mcp_server', '')}",
331
+ "log": f"Parameters: {tool_args}",
332
+ # "status": "pending",
333
+ "status": "done",
334
+ "id": f"tool_call_{tool_name}",
335
+ },
336
+ }
337
+ ]
338
+
339
+ yield [
340
+ {
341
+ "role": "assistant",
342
+ "content": "```json\n" + tool_args + "\n```",
343
+ "metadata": {
344
+ "parent_id": f"tool_call_{tool_name}",
345
+ "id": f"params_{tool_name}",
346
+ "title": "Tool Parameters",
347
+ },
348
+ }
349
+ ]
350
+
351
+ elif isinstance(event, ResponseOutputItemDoneEvent):
352
+ if event.item.type == "mcp_call":
353
+ yield [
354
+ {
355
+ "role": "assistant",
356
+ "content": "Here are the results from the tool:",
357
+ "metadata": {
358
+ "title": f"Tool Result for {tool_name.replace('avsolatorio_test_data_mcp_server', '')}",
359
+ "status": "done",
360
+ "id": f"result_{tool_name}",
361
+ },
362
+ }
363
+ ]
364
+
365
+ result_content = event.item.output
366
+ if result_content.startswith("root="):
367
+ result_content = result_content[5:]
368
+ try:
369
+ result_content = ast.literal_eval(result_content)
370
+ result_content = json.dumps(result_content, indent=2)
371
+ except:
372
+ pass
373
+
374
+ yield [
375
+ {
376
+ "role": "assistant",
377
+ "content": "```\n" + result_content + "\n```",
378
+ "metadata": {
379
+ "parent_id": f"result_{tool_name}",
380
+ "id": f"raw_result_{tool_name}",
381
+ "title": "Raw Output",
382
+ },
383
+ }
384
+ ]
385
+ is_tool_call = False
386
+ tool_name = None
387
+ tool_args = None
388
+
389
+ elif (
390
+ isinstance(event, ResponseContentPartDoneEvent)
391
+ and event.type == "response.content_part.done"
392
+ ):
393
+ yield [
394
+ {
395
+ "role": "assistant",
396
+ "content": "",
397
+ "delta": "",
398
+ "status": "done",
399
+ }
400
+ ]
401
+ elif isinstance(event, ResponseTextDeltaEvent):
402
+ yield [{"role": "assistant", "content": None, "delta": event.delta}]
403
+
404
+ async def _process_query_anthropic(
405
+ self, message: str, history: List[Union[Dict[str, Any], ChatMessage]]
406
+ ):
407
+ claude_messages = []
408
+ for msg in history:
409
+ if isinstance(msg, ChatMessage):
410
+ role, content = msg.role, msg.content
411
+ else:
412
+ role, content = msg.get("role"), msg.get("content")
413
+
414
+ if role in ["user", "assistant", "system"]:
415
+ claude_messages.append({"role": role, "content": content})
416
+
417
+ claude_messages.append({"role": "user", "content": message})
418
+
419
+ try:
420
+ response = self.anthropic.messages.create(
421
+ # model="claude-3-5-sonnet-20241022",
422
+ model=LLM_MODEL,
423
+ system=SYSTEM_PROMPT,
424
+ max_tokens=1000,
425
+ messages=claude_messages,
426
+ tools=self.tools,
427
+ )
428
+ except OverloadedError:
429
+ yield [
430
+ {
431
+ "role": "assistant",
432
+ "content": "The LLM API is overloaded now, try again later...",
433
+ }
434
+ ]
435
+ # TODO: Add a retry mechanism
436
+
437
+ result_messages = []
438
+ partial_messages = []
439
+
440
+ print(response.content)
441
+ contents = response.content
442
+
443
+ MAX_CALLS = 10
444
+ auto_calls = 0
445
+
446
+ while len(contents) > 0 and auto_calls < MAX_CALLS:
447
+ content = contents.pop(0)
448
+
449
+ if content.type == "text":
450
+ result_messages.append({"role": "assistant", "content": content.text})
451
+ claude_messages.append({"role": "assistant", "content": content.text})
452
+ partial_messages.append(result_messages[-1])
453
+ yield [result_messages[-1]]
454
+ partial_messages = []
455
+
456
+ elif content.type == "tool_use":
457
+ tool_id = content.id
458
+ tool_name = content.name
459
+ tool_args = content.input
460
+
461
+ result_messages.append(
462
+ {
463
+ "role": "assistant",
464
+ "content": f"I'll use the {tool_name} tool to help answer your question.",
465
+ "metadata": {
466
+ "title": f"Using tool: {tool_name.replace('avsolatorio_test_data_mcp_server', '')}",
467
+ "log": f"Parameters: {json.dumps(tool_args, ensure_ascii=True)}",
468
+ # "status": "pending",
469
+ "status": "done",
470
+ "id": f"tool_call_{tool_name}",
471
+ },
472
+ }
473
+ )
474
+ partial_messages.append(result_messages[-1])
475
+ yield [result_messages[-1]]
476
+
477
+ result_messages.append(
478
+ {
479
+ "role": "assistant",
480
+ "content": "```json\n"
481
+ + json.dumps(tool_args, indent=2, ensure_ascii=True)
482
+ + "\n```",
483
+ "metadata": {
484
+ "parent_id": f"tool_call_{tool_name}",
485
+ "id": f"params_{tool_name}",
486
+ "title": "Tool Parameters",
487
+ },
488
+ }
489
+ )
490
+ partial_messages.append(result_messages[-1])
491
+ yield [result_messages[-1]]
492
+
493
+ print(f"Calling tool: {tool_name} with args: {tool_args}")
494
+ try:
495
+ # Check if session is still valid
496
+ if not self.session or not self.stdio or not self.write:
497
+ raise Exception(
498
+ "MCP session is not connected or has been closed"
499
+ )
500
+
501
+ result = await self.session.call_tool(tool_name, tool_args)
502
+ except Exception as e:
503
+ error_msg = f"Error calling tool {tool_name}: {str(e)}"
504
+ print(error_msg)
505
+ result_messages.append(
506
+ {
507
+ "role": "assistant",
508
+ "content": f"Sorry, I encountered an error while calling the tool: {error_msg}. Please try again or reload the page.",
509
+ "metadata": {
510
+ "title": f"Tool Error for {tool_name.replace('avsolatorio_test_data_mcp_server', '')}",
511
+ "status": "done",
512
+ "id": f"error_{tool_name}",
513
+ },
514
+ }
515
+ )
516
+ partial_messages.append(result_messages[-1])
517
+ yield [result_messages[-1]]
518
+ partial_messages = []
519
+ continue
520
+
521
+ if result_messages and "metadata" in result_messages[-2]:
522
+ result_messages[-2]["metadata"]["status"] = "done"
523
+
524
+ result_messages.append(
525
+ {
526
+ "role": "assistant",
527
+ "content": "Here are the results from the tool:",
528
+ "metadata": {
529
+ "title": f"Tool Result for {tool_name.replace('avsolatorio_test_data_mcp_server', '')}",
530
+ "status": "done",
531
+ "id": f"result_{tool_name}",
532
+ },
533
+ }
534
+ )
535
+ partial_messages.append(result_messages[-1])
536
+ yield [result_messages[-1]]
537
+ partial_messages = []
538
+
539
+ result_content = result.content
540
+ print(result_content)
541
+ if isinstance(result_content, list):
542
+ result_content = [r.model_dump() for r in result_content]
543
+
544
+ for r in result_content:
545
+ # Remove annotations field from each item if it exists
546
+ r.pop("annotations", None)
547
+ try:
548
+ r["text"] = json.loads(r["text"])
549
+ except:
550
+ pass
551
+
552
+ print("result_content", result_content)
553
+
554
+ result_messages.append(
555
+ {
556
+ "role": "assistant",
557
+ "content": "```\n"
558
+ + json.dumps(result_content, indent=2)
559
+ + "\n```",
560
+ "metadata": {
561
+ "parent_id": f"result_{tool_name}",
562
+ "id": f"raw_result_{tool_name}",
563
+ "title": "Raw Output",
564
+ },
565
+ }
566
+ )
567
+ partial_messages.append(result_messages[-1])
568
+ yield [result_messages[-1]]
569
+ partial_messages = []
570
+
571
+ claude_messages.append(
572
+ {"role": "assistant", "content": [content.model_dump()]}
573
+ )
574
+ claude_messages.append(
575
+ {
576
+ "role": "user",
577
+ "content": [
578
+ {
579
+ "type": "tool_result",
580
+ "tool_use_id": tool_id,
581
+ "content": json.dumps(result_content, indent=2),
582
+ }
583
+ ],
584
+ }
585
+ )
586
+
587
+ try:
588
+ next_response = self.anthropic.messages.create(
589
+ model=LLM_MODEL,
590
+ system=SYSTEM_PROMPT,
591
+ max_tokens=1000,
592
+ messages=claude_messages,
593
+ tools=self.tools,
594
+ )
595
+ auto_calls += 1
596
+ except OverloadedError:
597
+ yield [
598
+ {
599
+ "role": "assistant",
600
+ "content": "The LLM API is overloaded now, try again later...",
601
+ }
602
+ ]
603
+
604
+ print("next_response", next_response.content)
605
+
606
+ contents.extend(next_response.content)
607
+
608
+ async def _process_query(
609
+ self,
610
+ message: str,
611
+ history: List[Union[Dict[Any, Any], ChatMessage]],
612
+ previous_response_id: str = None,
613
+ ):
614
+ if LLM_PROVIDER == "anthropic":
615
+ async for partial in self._process_query_anthropic(message, history):
616
+ yield partial
617
+ elif LLM_PROVIDER == "openai":
618
+ try:
619
+ async for partial in self._process_query_openai(
620
+ message, history, previous_response_id
621
+ ):
622
+ yield partial
623
+ except openai.APIError as e:
624
+ print(e)
625
+ yield [
626
+ {
627
+ "role": "assistant",
628
+ "content": "The LLM encountered an error. Please try again or reload the page.",
629
+ }
630
+ ]
631
+ except Exception as e:
632
+ print(e)
633
+ yield [
634
+ {
635
+ "role": "assistant",
636
+ "content": f"Sorry, I encountered an unexpected error: `{e}`. Please try again or reload the page.",
637
+ }
638
+ ]
639
+
640
+
641
+ def gradio_interface(
642
+ server_path_or_url: str = "https://avsolatorio-test-data-mcp-server.hf.space/gradio_api/mcp/sse",
643
+ ):
644
+ # server_path_or_url = "https://avsolatorio-test-data-mcp-server.hf.space/gradio_api/mcp/sse"
645
+ # server_path_or_url = "wdi_mcp_server.py"
646
+
647
+ client = MCPClientWrapper()
648
+ custom_css = """
649
+ .gradio-container {
650
+ background-color: #fff !important;
651
+ }
652
+ .message-row.panel.bot-row {
653
+ background-color: #fff !important;
654
+ }
655
+ .message-row.panel.user-row {
656
+ background-color: #fff !important;
657
+ }
658
+ .user {
659
+ background-color: #f1f6ff !important;
660
+ }
661
+ .bot {
662
+ background-color: #fff !important;
663
+ }
664
+ .role {
665
+ margin-left: 10px !important;
666
+ }
667
+ footer{display:none !important}
668
+ """
669
+
670
+ # Disable auto-dark mode by setting theme to None
671
+ with gr.Blocks(title="WDI MCP Client", css=custom_css, theme=None) as demo:
672
+ try:
673
+ gr.Markdown("# Data360 Chat [Prototype]")
674
+ # gr.Markdown("Connect to the WDI MCP server and chat with the assistant")
675
+
676
+ with gr.Accordion(
677
+ "Connect to the WDI MCP server and chat with the assistant",
678
+ open=False,
679
+ visible=server_path_or_url.endswith(".py"),
680
+ ):
681
+ with gr.Row(equal_height=True):
682
+ with gr.Column(scale=4):
683
+ server_path = gr.Textbox(
684
+ label="Server Script Path",
685
+ placeholder="Enter path to server script (e.g., wdi_mcp_server.py)",
686
+ value=server_path_or_url,
687
+ )
688
+ with gr.Column(scale=1):
689
+ connect_btn = gr.Button("Connect")
690
+
691
+ status = gr.Textbox(label="Connection Status", interactive=False)
692
+
693
+ chatbot = gr.Chatbot(
694
+ value=[],
695
+ height="81vh",
696
+ type="messages",
697
+ show_copy_button=False,
698
+ avatar_images=("img/small-user.png", "img/small-robot.png"),
699
+ autoscroll=True,
700
+ layout="panel",
701
+ placeholder="Ask development data questions!",
702
+ )
703
+ previous_response_id = gr.Textbox(
704
+ label="Previous Response ID", interactive=False, visible=False
705
+ )
706
+
707
+ with gr.Row(equal_height=True):
708
+ msg = gr.Textbox(
709
+ label=None,
710
+ placeholder="Ask about what indicators are available for a specific topic (e.g., What's the definition of GDP?)",
711
+ scale=4,
712
+ show_label=False,
713
+ )
714
+ # clear_btn = gr.Button("Clear Chat", scale=1)
715
+
716
+ # connect_btn.click(client.connect, inputs=server_path, outputs=status)
717
+ # Automatically call client.connect(...) as soon as the interface loads
718
+ if LLM_PROVIDER == "anthropic":
719
+ demo.load(
720
+ fn=client.connect,
721
+ inputs=server_path,
722
+ outputs=status,
723
+ show_progress="full",
724
+ )
725
+
726
+ msg.submit(
727
+ client.process_message,
728
+ [msg, chatbot, previous_response_id],
729
+ [chatbot, msg, previous_response_id],
730
+ concurrency_limit=10,
731
+ )
732
+ # clear_btn.click(lambda: [], None, chatbot)
733
+
734
+ except KeyboardInterrupt:
735
+ if LLM_PROVIDER == "anthropic":
736
+ print("Keyboard interrupt received. Disconnecting from MCP server...")
737
+ asyncio.run(client.disconnect())
738
+ raise KeyboardInterrupt
739
+ # demo.unload(client.disconnect)
740
+
741
+ return demo
742
+
743
+
744
+ if __name__ == "__main__":
745
+ if not os.getenv("ANTHROPIC_API_KEY"):
746
+ print(
747
+ "Warning: ANTHROPIC_API_KEY not found in environment. Please set it in your .env file."
748
+ )
749
+
750
+ # interface = gradio_interface(server_path_or_url="wdi_mcp_server.py")
751
+ interface = gradio_interface(
752
+ server_path_or_url="https://avsolatorio-test-data-mcp-server.hf.space/gradio_api/mcp/sse"
753
+ )
754
+ interface.launch(
755
+ server_name=os.getenv("SERVER_NAME", "127.0.0.1"),
756
+ server_port=os.getenv("SERVER_PORT", 7860),
757
+ debug=True,
758
+ )
mcp_remote_client.py CHANGED
@@ -425,7 +425,7 @@ def gradio_interface(
425
  # Disable auto-dark mode by setting theme to None
426
  with gr.Blocks(title="WDI MCP Client", css=custom_css, theme=None) as demo:
427
  try:
428
- gr.Markdown("# Development Data Chat")
429
  # gr.Markdown("Connect to the WDI MCP server and chat with the assistant")
430
 
431
  with gr.Accordion(
 
425
  # Disable auto-dark mode by setting theme to None
426
  with gr.Blocks(title="WDI MCP Client", css=custom_css, theme=None) as demo:
427
  try:
428
+ gr.Markdown("# Data360 Chat [Prototype]")
429
  # gr.Markdown("Connect to the WDI MCP server and chat with the assistant")
430
 
431
  with gr.Accordion(
pyproject.toml CHANGED
@@ -11,6 +11,7 @@ dependencies = [
11
  "httpx>=0.28.1",
12
  "mcp>=1.10.0",
13
  "numpy>=2.2.6",
 
14
  "python-ulid>=3.0.0",
15
  "scikit-learn>=1.6.1",
16
  "sentence-transformers>=4.1.0",
 
11
  "httpx>=0.28.1",
12
  "mcp>=1.10.0",
13
  "numpy>=2.2.6",
14
+ "openai>=1.93.1",
15
  "python-ulid>=3.0.0",
16
  "scikit-learn>=1.6.1",
17
  "sentence-transformers>=4.1.0",
services.py CHANGED
@@ -202,6 +202,14 @@ def get_wdi_data(
202
  break
203
 
204
  metadata, data_page = json_response
 
 
 
 
 
 
 
 
205
  all_data.extend(data_page)
206
 
207
  if len(all_data) >= MAX_INFO:
 
202
  break
203
 
204
  metadata, data_page = json_response
205
+
206
+ if data_page is None:
207
+ if metadata.get("total") == 0:
208
+ note = "IMPORTANT: Let the user know that the indicator data is not available for the given country and date."
209
+ else:
210
+ note = "ERROR: The API response is invalid or empty."
211
+ break
212
+
213
  all_data.extend(data_page)
214
 
215
  if len(all_data) >= MAX_INFO:
uv.lock CHANGED
@@ -1215,6 +1215,25 @@ wheels = [
1215
  { url = "https://files.pythonhosted.org/packages/9e/4e/0d0c945463719429b7bd21dece907ad0bde437a2ff12b9b12fee94722ab0/nvidia_nvtx_cu12-12.6.77-py3-none-manylinux2014_x86_64.whl", hash = "sha256:6574241a3ec5fdc9334353ab8c479fe75841dbe8f4532a8fc97ce63503330ba1", size = 89265, upload-time = "2024-10-01T17:00:38.172Z" },
1216
  ]
1217
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1218
  [[package]]
1219
  name = "orjson"
1220
  version = "3.10.18"
@@ -2298,6 +2317,7 @@ dependencies = [
2298
  { name = "httpx" },
2299
  { name = "mcp" },
2300
  { name = "numpy" },
 
2301
  { name = "python-ulid" },
2302
  { name = "scikit-learn" },
2303
  { name = "sentence-transformers" },
@@ -2316,6 +2336,7 @@ requires-dist = [
2316
  { name = "httpx", specifier = ">=0.28.1" },
2317
  { name = "mcp", specifier = ">=1.10.0" },
2318
  { name = "numpy", specifier = ">=2.2.6" },
 
2319
  { name = "python-ulid", specifier = ">=3.0.0" },
2320
  { name = "scikit-learn", specifier = ">=1.6.1" },
2321
  { name = "sentence-transformers", specifier = ">=4.1.0" },
 
1215
  { url = "https://files.pythonhosted.org/packages/9e/4e/0d0c945463719429b7bd21dece907ad0bde437a2ff12b9b12fee94722ab0/nvidia_nvtx_cu12-12.6.77-py3-none-manylinux2014_x86_64.whl", hash = "sha256:6574241a3ec5fdc9334353ab8c479fe75841dbe8f4532a8fc97ce63503330ba1", size = 89265, upload-time = "2024-10-01T17:00:38.172Z" },
1216
  ]
1217
 
1218
+ [[package]]
1219
+ name = "openai"
1220
+ version = "1.93.1"
1221
+ source = { registry = "https://pypi.org/simple" }
1222
+ dependencies = [
1223
+ { name = "anyio" },
1224
+ { name = "distro" },
1225
+ { name = "httpx" },
1226
+ { name = "jiter" },
1227
+ { name = "pydantic" },
1228
+ { name = "sniffio" },
1229
+ { name = "tqdm" },
1230
+ { name = "typing-extensions" },
1231
+ ]
1232
+ sdist = { url = "https://files.pythonhosted.org/packages/5e/a8/e4427729da048cb33bda15e70f09f7520bdf3577bafc546b135ecb36af7d/openai-1.93.1.tar.gz", hash = "sha256:11eb8932965d0f79ecc4cb38a60a0c4cef4bcd5fcf08b99fc9a399fa5f1e50ab", size = 487124, upload-time = "2025-07-07T16:40:38.389Z" }
1233
+ wheels = [
1234
+ { url = "https://files.pythonhosted.org/packages/64/4f/875e5af1fb4e5ed4ea9e4a88f482d9ca2e48932105605b6c516e9a14de25/openai-1.93.1-py3-none-any.whl", hash = "sha256:a2c2946c4f21346d4902311a7440381fd8a33466ee7ca688133d1cad29a9357c", size = 755081, upload-time = "2025-07-07T16:40:36.585Z" },
1235
+ ]
1236
+
1237
  [[package]]
1238
  name = "orjson"
1239
  version = "3.10.18"
 
2317
  { name = "httpx" },
2318
  { name = "mcp" },
2319
  { name = "numpy" },
2320
+ { name = "openai" },
2321
  { name = "python-ulid" },
2322
  { name = "scikit-learn" },
2323
  { name = "sentence-transformers" },
 
2336
  { name = "httpx", specifier = ">=0.28.1" },
2337
  { name = "mcp", specifier = ">=1.10.0" },
2338
  { name = "numpy", specifier = ">=2.2.6" },
2339
+ { name = "openai", specifier = ">=1.93.1" },
2340
  { name = "python-ulid", specifier = ">=3.0.0" },
2341
  { name = "scikit-learn", specifier = ">=1.6.1" },
2342
  { name = "sentence-transformers", specifier = ">=4.1.0" },