Interactive MCP Client Test Script.
This script allows you to test arbitrary commands with the MCP server
using an LLM agent that can decide which tools to call.
Supports elicitation - when tools require user input (like authentication tokens),
the script will prompt you interactively.
async
test_mcp_interactive() -> None
Test the MCP server interactively with LLM agent.
Source code in datarobot_genai/drmcp/test_utils/test_interactive.py
| async def test_mcp_interactive() -> None:
"""Test the MCP server interactively with LLM agent."""
# Check for required environment variables
datarobot_api_token = os.environ.get("DATAROBOT_API_TOKEN")
if not datarobot_api_token:
print("❌ Error: DATAROBOT_API_TOKEN environment variable is required")
print("Please set it in your .env file or export it")
return
# Optional DataRobot settings
datarobot_endpoint = os.environ.get("DATAROBOT_ENDPOINT")
dr_llm_gateway_model = os.environ.get("DR_LLM_GATEWAY_MODEL")
llm_temperature = os.environ.get("LLM_TEMPERATURE")
print("🤖 Initializing LLM MCP Client...")
# Initialize the LLM client with elicitation handler
config = {
"datarobot_api_token": datarobot_api_token,
"save_llm_responses": False,
}
if datarobot_endpoint:
config["datarobot_endpoint"] = datarobot_endpoint
if dr_llm_gateway_model:
config["model"] = dr_llm_gateway_model
if llm_temperature is not None:
config["temperature"] = llm_temperature
if not config.get("model"):
print("❌ Error: no LLM model configured")
print("Set DR_LLM_GATEWAY_MODEL in your .env (same as ETE tests).")
return
llm_client = DRLLMGatewayMCPClient(str(config))
# Get MCP server URL
mcp_server_url = get_dr_mcp_server_url()
if not mcp_server_url:
print("❌ Error: MCP server URL is not configured")
print("Please set DR_MCP_SERVER_URL environment variable or run: task test-interactive")
return
print(f"🔗 Connecting to MCP server at: {mcp_server_url}")
# Elicitation handler: prompt user for required values
async def elicitation_handler(
context: RequestContext[ClientSession, Any], params: ElicitRequestParams
) -> ElicitResult:
print(f"\n📋 Elicitation Request: {params.message}")
if params.requestedSchema:
print(f" Schema: {params.requestedSchema}")
while True:
try:
response = input(" Enter value (or 'decline'/'cancel'): ").strip()
except (EOFError, KeyboardInterrupt):
return ElicitResult(action="cancel")
if response.lower() == "decline":
return ElicitResult(action="decline")
if response.lower() == "cancel":
return ElicitResult(action="cancel")
if response:
return ElicitResult(action="accept", content={"value": response})
print(" Please enter a value or 'decline'/'cancel'")
try:
async with streamablehttp_client(
url=mcp_server_url,
headers=get_headers(),
) as (read_stream, write_stream, _):
async with ClientSession(
read_stream,
write_stream,
elicitation_callback=elicitation_handler,
) as session:
await session.initialize()
print("✅ Connected to MCP server!")
print("📋 Available tools:")
tools_result = await session.list_tools()
for i, tool in enumerate(tools_result.tools, 1):
print(f" {i}. {tool.name}: {tool.description}")
print("\n" + "=" * 60)
print("🎯 Interactive Testing Mode")
print("=" * 60)
print("Type your questions/commands. The AI will decide which tools to use.")
print("If a tool requires additional information, you will be prompted.")
print("Type 'quit' or 'exit' to stop.")
print()
while True:
try:
user_input = input("🤔 You: ").strip()
if user_input.lower() in ["quit", "exit", "q"]:
print("👋 Goodbye!")
break
if not user_input:
continue
except (EOFError, KeyboardInterrupt):
print("\n👋 Goodbye!")
break
print("🤖 AI is thinking...")
response = await llm_client.process_prompt_with_mcp_support(
prompt=user_input,
mcp_session=session,
)
print("\n🤖 AI Response:")
print("-" * 40)
print(response.content)
if response.tool_calls:
print("\n🔧 Tools Used:")
for i, tool_call in enumerate(response.tool_calls, 1):
print(f" {i}. {tool_call.tool_name}")
print(f" Parameters: {tool_call.parameters}")
print(f" Reasoning: {tool_call.reasoning}")
if i <= len(response.tool_results):
result = response.tool_results[i - 1]
try:
result_data = json.loads(result)
if result_data.get("status") == "error":
error_msg = result_data.get("error", "Unknown error")
print(f" ❌ Error: {error_msg}")
elif result_data.get("status") == "success":
print(" ✅ Success")
except json.JSONDecodeError:
if len(result) > 100:
print(f" Result: {result[:100]}...")
else:
print(f" Result: {result}")
print("\n" + "=" * 60)
except Exception as e:
print(f"❌ Connection Error: {e}")
print(f" Server URL: {mcp_server_url}")
traceback.print_exc()
return
|