feat(connector): inject only user-selected connectors per request

- Read ext_info.connector_ids (list) at start of _react_agent_stream;
  fall back to legacy ext_info.connector_id (string) for older clients
- Filter pack injection to selected ids via get_connector_tools()
- Empty/missing selection no longer leaks all active connectors into prompt
- Warn (not error) on stale ids so users can recover gracefully
This commit is contained in:
alan.cl
2026-05-26 12:06:12 +08:00
parent beb9e6244b
commit babd975465
2 changed files with 202 additions and 13 deletions

View File

@@ -8,7 +8,7 @@ import tempfile
import uuid
import zipfile
from pathlib import Path
from typing import TYPE_CHECKING, Any, AsyncGenerator, Dict, List, Optional
from typing import TYPE_CHECKING, Any, AsyncGenerator, Dict, List, Optional, Tuple
from urllib.parse import urlparse
from fastapi import APIRouter, Body, Depends, File, Query, Request, UploadFile
@@ -35,6 +35,7 @@ logger = logging.getLogger(__name__)
if TYPE_CHECKING:
from dbgpt.agent.core.memory.gpts import GptsMemory
from dbgpt.agent.resource.connector.manager import ConnectorManager
REACT_AGENT_MEMORY_CACHE: Dict[str, "GptsMemory"] = {}
@@ -68,6 +69,54 @@ def _extract_auto_data_markers(text: str) -> tuple[str, Dict[str, str]]:
return cleaned, extracted
def _parse_connector_ids(ext_info: Optional[Dict[str, Any]]) -> List[str]:
"""Extract connector IDs from ``ext_info``.
Supports two shapes:
- ``ext_info.connector_ids`` — a list of UUID strings (preferred).
- ``ext_info.connector_id`` — a single UUID string (legacy back-compat).
Returns an empty list when no valid IDs are found.
"""
if not ext_info or not isinstance(ext_info, dict):
return []
raw_ids = ext_info.get("connector_ids")
if isinstance(raw_ids, list):
return [cid for cid in raw_ids if isinstance(cid, str) and cid]
legacy = ext_info.get("connector_id")
if isinstance(legacy, str) and legacy:
return [legacy]
return []
def _select_connector_packs(
connector_ids: List[str],
connector_manager: Optional["ConnectorManager"],
) -> Tuple[List["MCPToolPack"], List[str]]:
"""Return tool packs for the requested *connector_ids*.
Args:
connector_ids: IDs the user selected in the frontend.
connector_manager: A :class:`ConnectorManager` instance (or ``None``).
Returns:
A tuple of ``(packs, missing_ids)`` where *packs* are the resolved
:class:`MCPToolPack` objects and *missing_ids* lists IDs that could
not be resolved (e.g. deleted mid-conversation).
"""
packs: List["MCPToolPack"] = []
missing_ids: List[str] = []
if connector_manager is None or not connector_ids:
return packs, missing_ids
for cid in connector_ids:
pack = connector_manager.get_connector_tools(cid)
if pack is None:
missing_ids.append(cid)
continue
packs.append(pack)
return packs, missing_ids
async def _execute_skill_script_impl(
skill_name: str, script_name: str, args: dict
) -> str:
@@ -856,6 +905,9 @@ async def _react_agent_stream(
)
database_name = dialogue.ext_info.get("database_name")
# Connector selection (Task C): only inject user-selected connectors.
connector_ids: List[str] = _parse_connector_ids(dialogue.ext_info)
def infer_phase(action: str) -> str:
"""根据 action 名称推断所属阶段,返回自由文本描述"""
if not action:
@@ -2652,7 +2704,7 @@ print(json.dumps(summary, ensure_ascii=False))
is_skill_mode = pre_matched_skill is not None
_skill_name = pre_matched_skill.metadata.name if pre_matched_skill else "skill"
# Inject connector tools (safe: no-op if manager not registered or no active connectors)
# Inject connector tools — only the ones the user explicitly selected.
connector_tool_extras: List[Any] = []
try:
from dbgpt.agent.resource.connector.manager import (
@@ -2662,13 +2714,21 @@ print(json.dumps(summary, ensure_ascii=False))
_connector_manager = CFG.SYSTEM_APP.get_component(
"connector_manager", _ConnectorManager, default_component=None
)
if _connector_manager is not None:
_connector_packs = _connector_manager.get_all_tools()
if _connector_packs:
connector_tool_extras.extend(_connector_packs)
if _connector_manager is not None and connector_ids:
connector_tool_extras, _missing = _select_connector_packs(
connector_ids, _connector_manager
)
for _mid in _missing:
logger.warning(
"_react_agent_stream: connector_id %s not active, skipping",
_mid,
)
if connector_tool_extras:
logger.info(
"_react_agent_stream: injected %d connector tool pack(s)",
len(_connector_packs),
"_react_agent_stream: injected %d connector tool pack(s) "
"(selected: %d)",
len(connector_tool_extras),
len(connector_ids),
)
except Exception:
pass # graceful degradation — connector tools are optional
@@ -2933,14 +2993,20 @@ Action Input: The JSON format of tool parameters
_cm = CFG.SYSTEM_APP.get_component(
"connector_manager", _ConnectorManager, default_component=None
)
if _cm is not None:
if _cm is not None and connector_ids:
_active = _cm.list_active()
if _active:
# Only describe connectors the user explicitly selected.
# Iterate connector_ids (not _active) so prompt order matches
# user selection order and stays consistent with _select_connector_packs.
_active_map = {c["connector_id"]: c for c in _active if isinstance(c, dict)}
_selected = [
_active_map[cid] for cid in connector_ids if cid in _active_map
]
if _selected:
_connector_lines = []
for _c in _active:
for _c in _selected:
_tool_names = ", ".join(
t.name
for t in (_c.get("tools", []) if isinstance(_c, dict) else [])
t.get("name", "unknown") for t in _c.get("tools", [])
)
_connector_lines.append(
f"- {_c.get('name', 'unknown')} ({_c.get('connector_type', 'unknown')}): "

View File

@@ -0,0 +1,123 @@
"""Unit tests for connector selection helpers in agentic_data_api.
Tests ``_parse_connector_ids`` and ``_select_connector_packs`` which are
extracted helpers that ``_react_agent_stream`` delegates to.
"""
from typing import Any, Dict
from unittest.mock import MagicMock
from dbgpt_app.openapi.api_v1.agentic_data_api import (
_parse_connector_ids,
_select_connector_packs,
)
# ---------------------------------------------------------------------------
# _parse_connector_ids
# ---------------------------------------------------------------------------
class TestParseConnectorIds:
"""Verify ext_info parsing for connector_ids / connector_id."""
def test_none_ext_info(self):
assert _parse_connector_ids(None) == []
def test_empty_dict(self):
assert _parse_connector_ids({}) == []
def test_not_a_dict(self):
assert _parse_connector_ids("not-a-dict") == []
def test_connector_ids_list(self):
ext = {"connector_ids": ["id1", "id2"]}
assert _parse_connector_ids(ext) == ["id1", "id2"]
def test_connector_ids_filters_non_strings(self):
ext = {"connector_ids": ["id1", 123, "", None, "id2"]}
assert _parse_connector_ids(ext) == ["id1", "id2"]
def test_connector_ids_empty_list(self):
ext = {"connector_ids": []}
assert _parse_connector_ids(ext) == []
def test_legacy_connector_id_string(self):
ext = {"connector_id": "legacy-id"}
assert _parse_connector_ids(ext) == ["legacy-id"]
def test_legacy_connector_id_empty_string(self):
ext = {"connector_id": ""}
assert _parse_connector_ids(ext) == []
def test_legacy_connector_id_not_string(self):
ext = {"connector_id": 42}
assert _parse_connector_ids(ext) == []
def test_connector_ids_takes_precedence_over_legacy(self):
ext = {"connector_ids": ["new-id"], "connector_id": "old-id"}
assert _parse_connector_ids(ext) == ["new-id"]
def test_connector_ids_empty_falls_through_to_legacy(self):
"""Empty list means user explicitly chose nothing -> no fallback."""
ext = {"connector_ids": [], "connector_id": "old-id"}
# connector_ids is present (empty list), so we return empty -- no
# fallback to legacy. This matches the code: isinstance([], list)
# is True, so the list comprehension runs and returns [].
assert _parse_connector_ids(ext) == []
# ---------------------------------------------------------------------------
# _select_connector_packs
# ---------------------------------------------------------------------------
def _make_manager(packs: Dict[str, Any]) -> MagicMock:
"""Create a mock ConnectorManager with the given id->pack mapping."""
mgr = MagicMock()
mgr.get_connector_tools.side_effect = lambda cid: packs.get(cid)
return mgr
class TestSelectConnectorPacks:
"""Verify filtered connector pack selection."""
def test_empty_ids_returns_empty(self):
mgr = _make_manager({"id1": "pack1"})
packs, missing = _select_connector_packs([], mgr)
assert packs == []
assert missing == []
def test_none_manager_returns_empty(self):
packs, missing = _select_connector_packs(["id1"], None)
assert packs == []
assert missing == []
def test_both_active(self):
mgr = _make_manager({"id1": "pack1", "id2": "pack2"})
packs, missing = _select_connector_packs(["id1", "id2"], mgr)
assert packs == ["pack1", "pack2"]
assert missing == []
def test_one_missing(self):
mgr = _make_manager({"id1": "pack1"})
packs, missing = _select_connector_packs(["id1", "missing"], mgr)
assert packs == ["pack1"]
assert missing == ["missing"]
def test_all_missing(self):
mgr = _make_manager({})
packs, missing = _select_connector_packs(["a", "b"], mgr)
assert packs == []
assert missing == ["a", "b"]
def test_single_id(self):
mgr = _make_manager({"solo": "solo-pack"})
packs, missing = _select_connector_packs(["solo"], mgr)
assert packs == ["solo-pack"]
assert missing == []
def test_preserves_order(self):
mgr = _make_manager({"c": "pc", "a": "pa", "b": "pb"})
packs, missing = _select_connector_packs(["a", "b", "c"], mgr)
assert packs == ["pa", "pb", "pc"]
assert missing == []