mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-13 00:14:02 +00:00
Add an endpoint to fix a connection tool's parameters
PUT /api/connections/<id>/tools/<tool_id>/parameters takes an action and a map of parameter to value (always use it) or null (let the model decide), checked against the action's schema; only the tool's owner can call it. Connection tools now list each action's parameters and whether they are fixed.
This commit is contained in:
1 parent
54bea5a1d8
commit
28b56e8b03
4 files changed
+171
-1
No files matched your search
@@ -518,6 +518,36 @@ class ConnectionToolPermissions(Resource):
|
||||
return make_response(jsonify({"success": True, "tool": payload}), 200)
|
||||
|
||||
|
||||
@connections_ns.route("/connections/<string:connection_id>/tools/<string:tool_id>/parameters")
|
||||
class ConnectionToolParameters(Resource):
|
||||
@api.doc(
|
||||
description=(
|
||||
"Fix or release an action's parameters: {action, parameters: {name: value | null}}. "
|
||||
"A value is sent on every call and hidden from the model; null lets the model decide. Owner only."
|
||||
)
|
||||
)
|
||||
def put(self, connection_id: str, tool_id: str):
|
||||
user_id = _user_id()
|
||||
if not user_id:
|
||||
return _unauthorized()
|
||||
body = _json_body()
|
||||
action, pins = body.get("action"), body.get("parameters")
|
||||
if not isinstance(action, str) or not isinstance(pins, dict) or not pins:
|
||||
return _error("Send the action and a map of parameters to values or null", 400)
|
||||
with db_session() as conn:
|
||||
row = _owned(conn, connection_id, user_id)
|
||||
if row is None:
|
||||
return _not_found()
|
||||
try:
|
||||
tool = service.set_tool_parameters(conn, user_id, connection_id, tool_id, action, pins)
|
||||
except ValueError as err:
|
||||
return _error(str(err), 400)
|
||||
if tool is None:
|
||||
return _not_found()
|
||||
payload = service.serialize_tool(tool)
|
||||
return make_response(jsonify({"success": True, "tool": payload}), 200)
|
||||
|
||||
|
||||
@connections_ns.route("/connections/<string:connection_id>/refresh-tools")
|
||||
class ConnectionRefreshTools(Resource):
|
||||
@api.doc(description="MCP: re-scan the server's actions and return what was added and removed")
|
||||
|
||||
@@ -151,8 +151,34 @@ def serialize_source(row: dict, connection_status: Optional[str] = None) -> dict
|
||||
}
|
||||
|
||||
|
||||
def serialize_parameters(action: dict) -> list[dict]:
|
||||
"""An action's parameters as the connection drawer shows them.
|
||||
|
||||
``fixed`` parameters are sent with ``value`` on every call and hidden
|
||||
from the model; the others are left to it.
|
||||
"""
|
||||
from docsgpt.agents.tool_pins import is_pinned, iter_parameters
|
||||
|
||||
schema = action.get("parameters") if isinstance(action.get("parameters"), dict) else {}
|
||||
required = set(schema.get("required") or []) if isinstance(schema.get("required"), list) else set()
|
||||
parameters = []
|
||||
for _section, name, details in iter_parameters(action):
|
||||
fixed = is_pinned(details)
|
||||
parameters.append(
|
||||
{
|
||||
"name": name,
|
||||
"description": details.get("description") or "",
|
||||
"type": details.get("type") if isinstance(details.get("type"), str) else "string",
|
||||
"required": bool(details.get("required")) or name in required,
|
||||
"fixed": fixed,
|
||||
"value": details.get("value") if fixed else None,
|
||||
}
|
||||
)
|
||||
return parameters
|
||||
|
||||
|
||||
def serialize_tool(row: dict) -> dict:
|
||||
"""A tool linked to a connection, with its actions and their permissions."""
|
||||
"""A tool linked to a connection, with its actions, permissions and parameters."""
|
||||
from docsgpt.connectors.permissions import action_access, action_permission
|
||||
|
||||
actions = []
|
||||
@@ -166,6 +192,7 @@ def serialize_tool(row: dict) -> dict:
|
||||
"description": action.get("description", ""),
|
||||
"access": access,
|
||||
"permission": action_permission(action),
|
||||
"parameters": serialize_parameters(action),
|
||||
}
|
||||
)
|
||||
return {
|
||||
@@ -1277,6 +1304,46 @@ def set_tool_permissions(
|
||||
return tools.get_any(str(tool["id"]), user_id)
|
||||
|
||||
|
||||
def set_tool_parameters(
|
||||
conn, user_id: str, connection_id: str, tool_id: str, action_name: str, pins: dict,
|
||||
) -> Optional[dict]:
|
||||
"""Fix or release parameters of one action of a tool the user owns.
|
||||
|
||||
Args:
|
||||
conn: Open database connection.
|
||||
user_id: The caller, who must own the tool.
|
||||
connection_id: The connection the tool must belong to.
|
||||
tool_id: The tool to update.
|
||||
action_name: The action whose parameters change.
|
||||
pins: Parameter name to the value to always use, or None to let the
|
||||
model decide.
|
||||
|
||||
Returns:
|
||||
The updated tool, or None (nothing written) when the tool is not the
|
||||
user's or belongs to another connection.
|
||||
|
||||
Raises:
|
||||
ValueError: The action or a parameter does not exist, or a value does
|
||||
not fit its parameter.
|
||||
"""
|
||||
from docsgpt.agents.tool_pins import set_pins
|
||||
from docsgpt.storage.db.repositories.user_tools import UserToolsRepository
|
||||
|
||||
tools = UserToolsRepository(conn)
|
||||
tool = tools.get_any(tool_id, user_id)
|
||||
if tool is None or tool.get("user_id") != user_id or str(tool.get("connection_id")) != connection_id:
|
||||
return None
|
||||
actions = _json(tool.get("actions")) or []
|
||||
if not any(isinstance(a, dict) and a.get("name") == action_name for a in actions):
|
||||
raise ValueError(f"Unknown action: {action_name}")
|
||||
actions = [
|
||||
set_pins(action, pins) if isinstance(action, dict) and action.get("name") == action_name else action
|
||||
for action in actions
|
||||
]
|
||||
tools.update(str(tool["id"]), user_id, {"actions": actions})
|
||||
return tools.get_any(str(tool["id"]), user_id)
|
||||
|
||||
|
||||
def reencrypt_all(batch_size: int = 500) -> dict:
|
||||
"""Rewrite every connection's credentials with the current key.
|
||||
|
||||
|
||||
@@ -320,6 +320,78 @@ class TestPermissions:
|
||||
assert resp.status_code == 400
|
||||
|
||||
|
||||
class TestParameters:
|
||||
@staticmethod
|
||||
def _tool(pg_conn, user="alice"):
|
||||
cid = _connection(pg_conn, user=user, secrets={"credentials": {"token": "t"}})
|
||||
tool = service.ensure_connection_tools(pg_conn, user, service.ConnectorSessionsRepository(pg_conn).get(cid))[0]
|
||||
return cid, str(tool["id"])
|
||||
|
||||
@staticmethod
|
||||
def _put(app, cid, tool_id, body, user="alice"):
|
||||
from docsgpt.api.connector.connections import ConnectionToolParameters
|
||||
|
||||
return _call(app, ConnectionToolParameters, "put", f"/api/connections/{cid}/tools/{tool_id}/parameters",
|
||||
user=user, body=body, args=[cid, tool_id])
|
||||
|
||||
@staticmethod
|
||||
def _parameters(payload, action="telegram_send_message"):
|
||||
actions = {a["name"]: a for a in payload["tool"]["actions"]}
|
||||
return {p["name"]: p for p in actions[action]["parameters"]}
|
||||
|
||||
def test_serialized_actions_list_their_parameters(self, pg_conn):
|
||||
cid, _ = self._tool(pg_conn)
|
||||
detail = service.connection_detail(pg_conn, service.ConnectorSessionsRepository(pg_conn).get(cid))
|
||||
action = next(a for a in detail["tools"][0]["actions"] if a["name"] == "telegram_send_message")
|
||||
names = [p["name"] for p in action["parameters"]]
|
||||
assert names == ["text", "chat_id"]
|
||||
chat_id = action["parameters"][1]
|
||||
assert chat_id["fixed"] is False and chat_id["value"] is None
|
||||
assert chat_id["type"] == "string"
|
||||
assert chat_id["description"]
|
||||
|
||||
def test_owner_fixes_and_releases_a_parameter(self, app, pg_conn):
|
||||
cid, tool_id = self._tool(pg_conn)
|
||||
with _db(pg_conn):
|
||||
fixed = self._put(app, cid, tool_id, {"action": "telegram_send_message",
|
||||
"parameters": {"chat_id": "-1001"}})
|
||||
assert fixed.status_code == 200
|
||||
chat_id = self._parameters(fixed.get_json())["chat_id"]
|
||||
assert chat_id["fixed"] is True and chat_id["value"] == "-1001"
|
||||
released = self._put(app, cid, tool_id, {"action": "telegram_send_message",
|
||||
"parameters": {"chat_id": None}})
|
||||
chat_id = self._parameters(released.get_json())["chat_id"]
|
||||
assert chat_id["fixed"] is False and chat_id["value"] is None
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"body",
|
||||
[
|
||||
{"action": "telegram_send_message", "parameters": {"token": "x"}},
|
||||
{"action": "not_an_action", "parameters": {"chat_id": "1"}},
|
||||
{"action": "telegram_send_message", "parameters": {"chat_id": ""}},
|
||||
{"action": "telegram_send_message", "parameters": "chat_id"},
|
||||
{"parameters": {"chat_id": "1"}},
|
||||
],
|
||||
)
|
||||
def test_rejects_what_the_action_does_not_have(self, app, pg_conn, body):
|
||||
cid, tool_id = self._tool(pg_conn)
|
||||
with _db(pg_conn):
|
||||
resp = self._put(app, cid, tool_id, body)
|
||||
assert resp.status_code == 400
|
||||
|
||||
def test_another_user_cannot_fix_values(self, app, pg_conn):
|
||||
cid, tool_id = self._tool(pg_conn)
|
||||
with _db(pg_conn):
|
||||
resp = self._put(app, cid, tool_id, {"action": "telegram_send_message",
|
||||
"parameters": {"chat_id": "666"}}, user="mallory")
|
||||
assert resp.status_code == 404
|
||||
stored = pg_conn.execute(
|
||||
text("SELECT actions FROM user_tools WHERE id = CAST(:i AS uuid)"), {"i": tool_id}
|
||||
).scalar()
|
||||
chat_id = stored[0]["parameters"]["properties"]["chat_id"]
|
||||
assert chat_id["filled_by_llm"] is True
|
||||
|
||||
|
||||
class TestDelete:
|
||||
def test_delete_with_source_removal(self, app, pg_conn):
|
||||
from docsgpt.api.connector.connections import ConnectionDetail
|
||||
|
||||
@@ -108,6 +108,7 @@ class TestListing:
|
||||
actions = {a["name"]: a for a in detail["tools"][0]["actions"]}
|
||||
assert actions["search_pages"] == {
|
||||
"name": "search_pages", "description": "Search", "access": "read", "permission": "always",
|
||||
"parameters": [],
|
||||
}
|
||||
assert actions["create_page"]["access"] == "write"
|
||||
assert actions["create_page"]["permission"] == "ask"
|
||||
|
||||
Reference in new issue
Block a user