Update prompt

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

mcp_client.py

@@ -16,6 +16,18 @@ 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.
+
+You must not provide answers beyond what the tools provide.
+
+Do not make up data or information and never simulate the `get_wdi_data` tool. Instead, you must always call the `get_wdi_data` tool when the user asks for data.
+
+You can use multiple tools if needed.
+
+If the user asks for any information beyond what the tools available to you provide, you must say that you do not have that information."""
+
+LLM_MODEL = "claude-3-5-haiku-20241022"
+
 
 class MCPClientWrapper:
     def __init__(self):
@@ -98,7 +110,9 @@ class MCPClientWrapper:
         claude_messages.append({"role": "user", "content": message})
 
         response = self.anthropic.messages.create(
-            model="claude-3-5-sonnet-20241022",
+            # model="claude-3-5-sonnet-20241022",
+            model=LLM_MODEL,
+            system=SYSTEM_PROMPT,
             max_tokens=1000,
             messages=claude_messages,
             tools=self.tools,
@@ -106,11 +120,14 @@ class MCPClientWrapper:
 
         result_messages = []
 
+        print(response.content)
+
         for content in response.content:
             if content.type == "text":
                 result_messages.append({"role": "assistant", "content": content.text})
 
             elif content.type == "tool_use":
+                tool_id = content.id
                 tool_name = content.name
                 tool_args = content.input
 
@@ -141,6 +158,7 @@ class MCPClientWrapper:
                     }
                 )
 
+                print(f"Calling tool: {tool_name} with args: {tool_args}")
                 result = await self.session.call_tool(tool_name, tool_args)
 
                 if result_messages and "metadata" in result_messages[-2]:
@@ -159,8 +177,10 @@ 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)
+                    print("result_content", result_content)
 
                 try:
                     result_json = json.loads(result_content)
@@ -207,18 +227,36 @@ class MCPClientWrapper:
                         }
                     )
 
+                # claude_messages.append(
+                #     {
+                #         "role": "user",
+                #         "content": f"Tool result for {tool_name}: {result_content}",
+                #     }
+                # )
+                claude_messages.append(
+                    {"role": "assistant", "content": [content.model_dump()]}
+                )
                 claude_messages.append(
                     {
                         "role": "user",
-                        "content": f"Tool result for {tool_name}: {result_content}",
+                        "content": [
+                            {
+                                "type": "tool_result",
+                                "tool_use_id": tool_id,
+                                "content": result_content,
+                            }
+                        ],
                     }
                 )
                 next_response = self.anthropic.messages.create(
-                    model="claude-3-5-sonnet-20241022",
+                    model=LLM_MODEL,
+                    system=SYSTEM_PROMPT,
                     max_tokens=1000,
                     messages=claude_messages,
                 )
 
+                print("next_response", next_response.content)
+
                 if next_response.content and next_response.content[0].type == "text":
                     result_messages.append(
                         {"role": "assistant", "content": next_response.content[0].text}

wdi_mcp_server.py

@@ -157,7 +157,7 @@ async def get_wdi_data(
     indicator_id: str,
     country_codes: str | list[str],
     date: Optional[str] = None,
-    per_page: Optional[int] = 100,
+    per_page: Optional[int] = 5,
 ) -> dict[str, list[dict[str, Any]] | str]:
     """Fetches indicator data for a given indicator id (idno) from the World Bank's World Development Indicators (WDI) API. The LLM must exclusively use this tool when the user asks for data. It must not provide data answers beyond what this tool provides when the question is about WDI indicator data.
 
@@ -171,7 +171,7 @@ async def get_wdi_data(
         A dictionary with keys `data` and `note`. The `data` key contains a list of indicator data entries requested. The `note` key contains a note about the data returned.
     """
     print("Hello...")
-    MAX_INFO = 100
+    MAX_INFO = 20
     note = ""
 
     if isinstance(country_codes, str):