Improve prompt and handling of tool result

Signed-off-by: Aivin V. Solatorio <[email protected]>

mcp_client.py

@@ -3,7 +3,7 @@ import os
 import json
 from typing import List, Dict, Any, Union
 from contextlib import AsyncExitStack
-
+from datetime import datetime
 import gradio as gr
 from gradio.components.chatbot import ChatMessage
 from mcp import ClientSession, StdioServerParameters
@@ -16,7 +16,9 @@ load_dotenv()
 loop = asyncio.new_event_loop()
 asyncio.set_event_loop(loop)
 
-SYSTEM_PROMPT = """"You are a helpful assistant. You do not have any knowledge of the World Development Indicators (WDI) data. However, you can use the tools provided to answer questions.
+SYSTEM_PROMPT = f"""You are a helpful assistant and today is {datetime.now().strftime("%Y-%m-%d")}.
+
+You do not have any knowledge of the World Development Indicators (WDI) data. However, you can use the tools provided to answer questions.
 
 You must not provide answers beyond what the tools provide.
 
@@ -181,53 +183,28 @@ class MCPClientWrapper:
                 result_content = result.content
                 print(result_content)
                 if isinstance(result_content, list):
-                    result_content = "\n".join(str(item) for item in result_content)
+                    result_content = [r.model_dump() for r in result_content]
+
+                    for r in result_content:
+                        # Remove annotations field from each item if it exists
+                        r.pop("annotations", None)
+
+                    # result_content = "\n".join(str(item) for item in result_content)
                     print("result_content", result_content)
 
-                try:
-                    result_json = json.loads(result_content)
-                    if isinstance(result_json, dict) and "type" in result_json:
-                        if result_json["type"] == "image" and "url" in result_json:
-                            result_messages.append(
-                                {
-                                    "role": "assistant",
-                                    "content": {
-                                        "path": result_json["url"],
-                                        "alt_text": result_json.get(
-                                            "message", "Generated image"
-                                        ),
-                                    },
-                                    "metadata": {
-                                        "parent_id": f"result_{tool_name}",
-                                        "id": f"image_{tool_name}",
-                                        "title": "Generated Image",
-                                    },
-                                }
-                            )
-                        else:
-                            result_messages.append(
-                                {
-                                    "role": "assistant",
-                                    "content": "```\n" + result_content + "\n```",
-                                    "metadata": {
-                                        "parent_id": f"result_{tool_name}",
-                                        "id": f"raw_result_{tool_name}",
-                                        "title": "Raw Output",
-                                    },
-                                }
-                            )
-                except:
-                    result_messages.append(
-                        {
-                            "role": "assistant",
-                            "content": "```\n" + result_content + "\n```",
-                            "metadata": {
-                                "parent_id": f"result_{tool_name}",
-                                "id": f"raw_result_{tool_name}",
-                                "title": "Raw Output",
-                            },
-                        }
-                    )
+                result_messages.append(
+                    {
+                        "role": "assistant",
+                        "content": "```\n"
+                        + json.dumps(result_content, indent=2)
+                        + "\n```",
+                        "metadata": {
+                            "parent_id": f"result_{tool_name}",
+                            "id": f"raw_result_{tool_name}",
+                            "title": "Raw Output",
+                        },
+                    }
+                )
 
                 # claude_messages.append(
                 #     {