Add sampling support to chatbot example

Author: simba-git

examples/openai_chat_agent/app.py

@@ -13,11 +13,22 @@
 from dotenv import load_dotenv
 from langchain_ollama import ChatOllama
 from langchain_openai import ChatOpenAI
-from mcp_use import MCPAgent, MCPClient
+from mcp.types import CreateMessageRequestParams, CreateMessageResult, ErrorData, TextContent
+from mcp_use import MCPAgent, MCPClient, load_config_file
 
 SYSTEM_MESSAGE = "You are a helpful assistant that talks to the user and uses tools via MCP."
 
 
+async def sampling_callback(
+    context: ClientSession, params: CreateMessageRequestParams
+) -> CreateMessageResult | ErrorData:
+    return CreateMessageResult(
+        content=TextContent(text='["Cape Town"]', type="text"),
+        model="gpt-4o-mini",
+        role="assistant",
+    )
+
+
 async def ensure_ollama_running(model: str) -> None:
     """Check that an Ollama server is running."""
     try:
@@ -41,7 +52,7 @@ async def run_memory_chat() -> None:
     config_file = os.path.join(os.path.dirname(__file__), "config.json")
 
     print("Initializing chat...")
-    client = MCPClient.from_config_file(config_file)
+    client = MCPClient(load_config_file(config_file), sampling_callback=sampling_callback)
 
     openai_key = os.getenv("OPENAI_API_KEY")
     ollama_model = os.getenv("OLLAMA_MODEL", "llama3.2")

examples/openai_chat_agent/config.json

@@ -1,8 +1,8 @@
 {
   "mcpServers": {
-    "shop_api": {
+    "travel_agent": {
       "command": "python",
-      "args": ["../shop_api/app.py"]
+      "args": ["../server_side_llm_travel_planner/app.py"]
     }
   }
 }

examples/openai_chat_agent/requirements.txt

@@ -4,3 +4,4 @@ langchain_ollama
 langchain_community
 mcp_use
 python-dotenv
+httpx

examples/server_side_llm_travel_planner/app.py

@@ -5,7 +5,7 @@
 
 from pydantic import Field
 
-from enrichmcp import EnrichContext, EnrichMCP, EnrichModel, prefer_fast_model
+from enrichmcp import EnrichMCP, EnrichModel, prefer_fast_model
 
 app = EnrichMCP(
     title="Travel Planner",
@@ -60,7 +60,6 @@ def list_destinations() -> list[Destination]:
 @app.retrieve
 async def plan_trip(
     preferences: Annotated[str, Field(description="Your travel preferences")],
-    ctx: EnrichContext,
 ) -> list[Destination]:
     """Return three destinations that best match the given preferences."""
 
@@ -70,15 +69,13 @@ async def plan_trip(
         "given preferences. Reply with a JSON list of names only.\nPreferences: "
         f"{preferences}\n\n{bullet_list}"
     )
-    result = await ctx.sampling(
+    app.get_context()
+    result = await app.get_context().ask_llm(
         prompt,
         model_preferences=prefer_fast_model(),
         max_tokens=50,
     )
-    try:
-        names = json.loads(result.content.text)
-    except Exception:
-        return []
+    names = json.loads(result.content.text)
     return [d for d in DESTINATIONS if d.name in names]