Spaces:
Running
Running
File size: 4,157 Bytes
b1f90a5 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 |
import asyncio
import pdb
import sys
import time
sys.path.append(".")
from dotenv import load_dotenv
load_dotenv()
async def test_mcp_client():
from src.utils.mcp_client import setup_mcp_client_and_tools, create_tool_param_model
test_server_config = {
"mcpServers": {
# "markitdown": {
# "command": "docker",
# "args": [
# "run",
# "--rm",
# "-i",
# "markitdown-mcp:latest"
# ]
# },
"desktop-commander": {
"command": "npx",
"args": [
"-y",
"@wonderwhy-er/desktop-commander"
]
},
# "filesystem": {
# "command": "npx",
# "args": [
# "-y",
# "@modelcontextprotocol/server-filesystem",
# "/Users/xxx/ai_workspace",
# ]
# },
}
}
mcp_tools, mcp_client = await setup_mcp_client_and_tools(test_server_config)
for tool in mcp_tools:
tool_param_model = create_tool_param_model(tool)
print(tool.name)
print(tool.description)
print(tool_param_model.model_json_schema())
pdb.set_trace()
async def test_controller_with_mcp():
import os
from src.controller.custom_controller import CustomController
from browser_use.controller.registry.views import ActionModel
mcp_server_config = {
"mcpServers": {
# "markitdown": {
# "command": "docker",
# "args": [
# "run",
# "--rm",
# "-i",
# "markitdown-mcp:latest"
# ]
# },
"desktop-commander": {
"command": "npx",
"args": [
"-y",
"@wonderwhy-er/desktop-commander"
]
},
# "filesystem": {
# "command": "npx",
# "args": [
# "-y",
# "@modelcontextprotocol/server-filesystem",
# "/Users/xxx/ai_workspace",
# ]
# },
}
}
controller = CustomController()
await controller.setup_mcp_client(mcp_server_config)
action_name = "mcp.desktop-commander.execute_command"
action_info = controller.registry.registry.actions[action_name]
param_model = action_info.param_model
print(param_model.model_json_schema())
params = {"command": f"python ./tmp/test.py"
}
validated_params = param_model(**params)
ActionModel_ = controller.registry.create_action_model()
# Create ActionModel instance with the validated parameters
action_model = ActionModel_(**{action_name: validated_params})
result = await controller.act(action_model)
result = result.extracted_content
print(result)
if result and "Command is still running. Use read_output to get more output." in result and "PID" in \
result.split("\n")[0]:
pid = int(result.split("\n")[0].split("PID")[-1].strip())
action_name = "mcp.desktop-commander.read_output"
action_info = controller.registry.registry.actions[action_name]
param_model = action_info.param_model
print(param_model.model_json_schema())
params = {"pid": pid}
validated_params = param_model(**params)
action_model = ActionModel_(**{action_name: validated_params})
output_result = ""
while True:
time.sleep(1)
result = await controller.act(action_model)
result = result.extracted_content
if result:
pdb.set_trace()
output_result = result
break
print(output_result)
pdb.set_trace()
await controller.close_mcp_client()
pdb.set_trace()
if __name__ == '__main__':
# asyncio.run(test_mcp_client())
asyncio.run(test_controller_with_mcp())
|