mirror of
https://github.com/HKUDS/OpenSpace.git
synced 2026-08-28 05:15:00 +00:00
520 lines
17 KiB
Python
520 lines
17 KiB
Python
from __future__ import annotations
|
|
|
|
import re
|
|
from typing import Any
|
|
|
|
from openspace.grounding.core.tool import BaseTool
|
|
from openspace.grounding.core.types import BackendType, ToolResult, ToolStatus
|
|
|
|
|
|
TOOL_DISCOVERY_TOOL_NAME = "tool_search"
|
|
TOOL_SEARCH_TOOL_NAME = TOOL_DISCOVERY_TOOL_NAME
|
|
LOW_CONFIDENCE_MIN_SCORE = 6
|
|
LARGE_CANDIDATE_THRESHOLD = 80
|
|
PRESELECTOR_MIN_CANDIDATES = 20
|
|
PRESELECTOR_CANDIDATE_MULTIPLIER = 4
|
|
|
|
|
|
class ToolSearchTool(BaseTool):
|
|
"""Discover deferred tools and load their schemas on the next model turn."""
|
|
|
|
_name = TOOL_DISCOVERY_TOOL_NAME
|
|
_description = (
|
|
"Search deferred tools by name or capability and load matching tool "
|
|
"schemas for the next model turn."
|
|
)
|
|
backend_type = BackendType.META
|
|
_is_read_only = True
|
|
_is_concurrency_safe = True
|
|
always_load = True
|
|
search_hint = "discover deferred tools by name or capability"
|
|
parameter_descriptions = {
|
|
"query": (
|
|
"Search query. Use select:<tool_name> for exact tool loading, "
|
|
"or keywords such as 'github issue create'."
|
|
),
|
|
"max_results": "Maximum number of deferred tools to load.",
|
|
}
|
|
|
|
def __init__(self, all_tools: list[BaseTool] | None = None):
|
|
self._all_tools = list(all_tools or [])
|
|
self._current_context: Any | None = None
|
|
super().__init__()
|
|
|
|
def set_context(self, context: Any) -> None:
|
|
self._current_context = context
|
|
|
|
def get_prompt(self, context: Any | None = None) -> str:
|
|
return (
|
|
"Search deferred tools and make matching tool schemas available on "
|
|
"the next model turn. Use this when a needed MCP/GUI/domain tool is "
|
|
"listed as deferred but is not currently callable. Prefer "
|
|
"`select:<exact_tool_name>` when you know the tool name."
|
|
)
|
|
|
|
async def _arun(self, query: str, max_results: int = 5) -> ToolResult:
|
|
ctx = self._current_context
|
|
all_tools = list(getattr(ctx, "all_tools", None) or self._all_tools)
|
|
active_names = {tool.name for tool in getattr(ctx, "tools", []) or []}
|
|
discovered_names = set(getattr(ctx, "discovered_tool_names", set()) or set())
|
|
deferred_names = set(getattr(ctx, "deferred_tool_names", set()) or set())
|
|
deferred_tools = [
|
|
tool for tool in all_tools
|
|
if (tool.is_deferred or tool.name in deferred_names)
|
|
and tool.name not in discovered_names
|
|
]
|
|
|
|
ranked_matches, search_metadata = await _rank_with_optional_fallback(
|
|
query,
|
|
deferred_tools,
|
|
max_results=max_results,
|
|
context=ctx,
|
|
)
|
|
matches = [tool for tool, _ in ranked_matches]
|
|
if query.lower().strip().startswith("select:") and not matches:
|
|
# If the requested tool is already active, return a helpful no-op.
|
|
requested = _parse_select_names(query)
|
|
already_loaded = sorted(name for name in requested if name in active_names)
|
|
if already_loaded:
|
|
return ToolResult(
|
|
status=ToolStatus.SUCCESS,
|
|
content=(
|
|
"Requested tool(s) are already loaded: "
|
|
+ ", ".join(already_loaded)
|
|
),
|
|
metadata={"matches": already_loaded, "already_loaded": True},
|
|
)
|
|
|
|
matched_names = [tool.name for tool in matches]
|
|
if ctx is not None and hasattr(ctx, "mark_tools_discovered"):
|
|
ctx.mark_tools_discovered(matched_names)
|
|
|
|
payload = {
|
|
"query": query,
|
|
"matches": matched_names,
|
|
"loaded_next_turn": matched_names,
|
|
"total_deferred_tools": len(deferred_tools),
|
|
**search_metadata,
|
|
}
|
|
if not matched_names:
|
|
pending_mcp_servers = _get_pending_mcp_servers(ctx)
|
|
if pending_mcp_servers:
|
|
payload["pending_mcp_servers"] = pending_mcp_servers
|
|
content = "No matching deferred tools found"
|
|
if pending_mcp_servers:
|
|
content += (
|
|
". Some MCP servers are still connecting: "
|
|
+ ", ".join(pending_mcp_servers)
|
|
+ ". Their tools will become available shortly - try searching again."
|
|
)
|
|
else:
|
|
content = (
|
|
"Loaded deferred tool schema(s) for the next model turn:\n"
|
|
+ "\n".join(f"- {name}" for name in matched_names)
|
|
+ "\n\nYou can call these tools after the next assistant turn begins."
|
|
)
|
|
return ToolResult(
|
|
status=ToolStatus.SUCCESS,
|
|
content=content,
|
|
metadata=payload,
|
|
)
|
|
|
|
|
|
def search_tools_by_keyword(
|
|
query: str,
|
|
deferred_tools: list[BaseTool],
|
|
*,
|
|
max_results: int = 5,
|
|
) -> list[BaseTool]:
|
|
"""OpenSpace keyword search over tool names, descriptions and hints."""
|
|
|
|
return [tool for tool, _ in rank_tools_by_keyword(
|
|
query,
|
|
deferred_tools,
|
|
max_results=max_results,
|
|
)]
|
|
|
|
|
|
async def _rank_with_optional_fallback(
|
|
query: str,
|
|
deferred_tools: list[BaseTool],
|
|
*,
|
|
max_results: int,
|
|
context: Any | None,
|
|
) -> tuple[list[tuple[BaseTool, int]], dict[str, Any]]:
|
|
"""Rank deferred tools, falling back to the system-side preselector.
|
|
|
|
``tool_search`` is primarily a schema-hydration handshake. The OpenSpace
|
|
keyword scorer should win when it has a clear lexical match. The older
|
|
OpenSpace preselector stays as a recall enhancer for two cases:
|
|
1. very large candidate sets, where it can narrow the pool; and
|
|
2. no/low-confidence keyword results.
|
|
"""
|
|
|
|
query_kind = _query_kind(query, deferred_tools)
|
|
ranked_all = rank_tools_by_keyword(
|
|
query,
|
|
deferred_tools,
|
|
max_results=max_results,
|
|
)
|
|
ranked = ranked_all
|
|
confidence = _keyword_confidence(ranked_all)
|
|
metadata: dict[str, Any] = {
|
|
"selection_method": query_kind,
|
|
"keyword_top_score": confidence["top_score"],
|
|
"keyword_confidence": confidence,
|
|
"preselector_used": False,
|
|
"preselector_reason": None,
|
|
}
|
|
|
|
if query_kind != "keyword":
|
|
return ranked, metadata
|
|
|
|
large_candidate_set = len(deferred_tools) > LARGE_CANDIDATE_THRESHOLD
|
|
if large_candidate_set:
|
|
narrowed_limit = max(
|
|
PRESELECTOR_MIN_CANDIDATES,
|
|
max_results * PRESELECTOR_CANDIDATE_MULTIPLIER,
|
|
)
|
|
narrowed = await _preselect_deferred_tools(
|
|
query,
|
|
deferred_tools,
|
|
max_results=narrowed_limit,
|
|
context=context,
|
|
)
|
|
if narrowed:
|
|
ranked_narrowed = _rank_preselected_tools(
|
|
query,
|
|
narrowed,
|
|
max_results=max_results,
|
|
)
|
|
if confidence["low_confidence"]:
|
|
ranked = _merge_ranked(ranked_narrowed, ranked_all, max_results)
|
|
else:
|
|
ranked = _merge_keyword_with_preselector(
|
|
ranked_all,
|
|
ranked_narrowed,
|
|
max_results,
|
|
)
|
|
metadata.update({
|
|
"selection_method": "preselector_narrowed_keyword",
|
|
"preselector_used": True,
|
|
"preselector_reason": "large_candidate_set",
|
|
"preselector_candidate_count": len(narrowed),
|
|
})
|
|
|
|
if confidence["low_confidence"] and not metadata["preselector_used"]:
|
|
fallback = await _preselect_deferred_tools(
|
|
query,
|
|
deferred_tools,
|
|
max_results=max_results,
|
|
context=context,
|
|
)
|
|
if fallback:
|
|
fallback_ranked = _rank_preselected_tools(
|
|
query,
|
|
fallback,
|
|
max_results=max_results,
|
|
)
|
|
ranked = _merge_ranked(fallback_ranked, ranked, max_results)
|
|
metadata.update({
|
|
"selection_method": "preselector_fallback",
|
|
"preselector_used": True,
|
|
"preselector_reason": confidence["reason"],
|
|
"preselector_candidate_count": len(fallback),
|
|
})
|
|
|
|
return ranked, metadata
|
|
|
|
|
|
async def _preselect_deferred_tools(
|
|
query: str,
|
|
deferred_tools: list[BaseTool],
|
|
*,
|
|
max_results: int,
|
|
context: Any | None,
|
|
) -> list[BaseTool]:
|
|
if not deferred_tools:
|
|
return []
|
|
|
|
try:
|
|
from .search_tools import ToolPreselector
|
|
except Exception:
|
|
return []
|
|
|
|
llm_client = getattr(context, "llm_client", None)
|
|
quality_manager = getattr(context, "quality_manager", None)
|
|
|
|
try:
|
|
preselector = ToolPreselector(
|
|
max_tools=max_results,
|
|
llm=llm_client,
|
|
quality_manager=quality_manager,
|
|
)
|
|
if len(deferred_tools) <= max_results:
|
|
return _rank_with_preselector_without_return_all(
|
|
preselector,
|
|
query,
|
|
deferred_tools,
|
|
max_results=max_results,
|
|
)
|
|
return await preselector._arun(
|
|
task_prompt=query,
|
|
candidate_tools=deferred_tools,
|
|
max_tools=max_results,
|
|
)
|
|
except Exception:
|
|
return []
|
|
|
|
|
|
def _rank_with_preselector_without_return_all(
|
|
preselector: Any,
|
|
query: str,
|
|
deferred_tools: list[BaseTool],
|
|
*,
|
|
max_results: int,
|
|
) -> list[BaseTool]:
|
|
mode = getattr(preselector, "_default_mode", "keyword")
|
|
use_embedding_ranker = mode in {"semantic", "hybrid"}
|
|
ranked = preselector._rank_tools(
|
|
query,
|
|
deferred_tools,
|
|
top_k=max_results,
|
|
mode=mode,
|
|
use_embedding_ranker=use_embedding_ranker,
|
|
)
|
|
return [tool for tool, score in ranked if score > 0][:max_results]
|
|
|
|
|
|
def _query_kind(query: str, tools: list[BaseTool]) -> str:
|
|
query_lower = (query or "").lower().strip()
|
|
if query_lower.startswith("select:"):
|
|
return "select"
|
|
if any(tool.name.lower() == query_lower for tool in tools):
|
|
return "exact"
|
|
if query_lower.startswith("mcp__") and len(query_lower) > len("mcp__"):
|
|
return "server_prefix"
|
|
return "keyword"
|
|
|
|
|
|
def _keyword_confidence(
|
|
ranked: list[tuple[BaseTool, int]],
|
|
) -> dict[str, Any]:
|
|
if not ranked:
|
|
return {
|
|
"low_confidence": True,
|
|
"reason": "no_keyword_matches",
|
|
"top_score": 0,
|
|
"min_score": LOW_CONFIDENCE_MIN_SCORE,
|
|
}
|
|
|
|
top_score = ranked[0][1]
|
|
low_confidence = top_score < LOW_CONFIDENCE_MIN_SCORE
|
|
return {
|
|
"low_confidence": low_confidence,
|
|
"reason": (
|
|
"top_score_below_threshold"
|
|
if low_confidence
|
|
else "keyword_match"
|
|
),
|
|
"top_score": top_score,
|
|
"min_score": LOW_CONFIDENCE_MIN_SCORE,
|
|
}
|
|
|
|
|
|
def _merge_ranked(
|
|
primary: list[tuple[BaseTool, int]],
|
|
secondary: list[tuple[BaseTool, int]],
|
|
max_results: int,
|
|
) -> list[tuple[BaseTool, int]]:
|
|
merged: list[tuple[BaseTool, int]] = []
|
|
seen: set[str] = set()
|
|
for tool, score in [*primary, *secondary]:
|
|
if tool.name in seen:
|
|
continue
|
|
seen.add(tool.name)
|
|
merged.append((tool, score))
|
|
if len(merged) >= max_results:
|
|
break
|
|
return merged
|
|
|
|
|
|
def _rank_preselected_tools(
|
|
query: str,
|
|
tools: list[BaseTool],
|
|
*,
|
|
max_results: int,
|
|
) -> list[tuple[BaseTool, int]]:
|
|
keyword_scores = {
|
|
tool.name: score
|
|
for tool, score in rank_tools_by_keyword(
|
|
query,
|
|
tools,
|
|
max_results=len(tools),
|
|
)
|
|
}
|
|
fallback_score = max(1, LOW_CONFIDENCE_MIN_SCORE - 1)
|
|
ranked: list[tuple[BaseTool, int]] = []
|
|
for tool in tools[:max_results]:
|
|
ranked.append((tool, keyword_scores.get(tool.name, fallback_score)))
|
|
return ranked
|
|
|
|
|
|
def _merge_keyword_with_preselector(
|
|
keyword_ranked: list[tuple[BaseTool, int]],
|
|
preselector_ranked: list[tuple[BaseTool, int]],
|
|
max_results: int,
|
|
) -> list[tuple[BaseTool, int]]:
|
|
if not keyword_ranked:
|
|
return preselector_ranked[:max_results]
|
|
if max_results <= 1:
|
|
return keyword_ranked[:1]
|
|
|
|
reserved_count = min(len(preselector_ranked), max(1, max_results // 2))
|
|
primary = [
|
|
*keyword_ranked[:1],
|
|
*preselector_ranked[:reserved_count],
|
|
]
|
|
secondary = [
|
|
*keyword_ranked[1:],
|
|
*preselector_ranked[reserved_count:],
|
|
]
|
|
return _merge_ranked(primary, secondary, max_results)
|
|
|
|
|
|
def rank_tools_by_keyword(
|
|
query: str,
|
|
tools: list[BaseTool],
|
|
*,
|
|
max_results: int = 5,
|
|
) -> list[tuple[BaseTool, int]]:
|
|
"""Rank tools with the lightweight OpenSpace ToolSearch keyword scoring rules."""
|
|
|
|
query_lower = (query or "").lower().strip()
|
|
if not query_lower or not tools:
|
|
return []
|
|
|
|
# Fast path: exact name.
|
|
for tool in tools:
|
|
if tool.name.lower() == query_lower:
|
|
return [(tool, 10_000)]
|
|
|
|
# Fast path: mcp__server prefix.
|
|
if query_lower.startswith("mcp__") and len(query_lower) > len("mcp__"):
|
|
matches = [
|
|
tool
|
|
for tool in tools
|
|
if tool.name.lower().startswith(query_lower)
|
|
]
|
|
if matches:
|
|
return [(tool, 9_000 - index) for index, tool in enumerate(matches[:max_results])]
|
|
|
|
# Fast path: select:tool_a,tool_b.
|
|
if query_lower.startswith("select:"):
|
|
requested = _parse_select_names(query)
|
|
by_name = {tool.name: tool for tool in tools}
|
|
return [(by_name[name], 10_000) for name in requested if name in by_name][:max_results]
|
|
|
|
raw_terms = [term for term in re.split(r"\s+", query_lower) if term]
|
|
required_terms = [term[1:] for term in raw_terms if term.startswith("+") and len(term) > 1]
|
|
optional_terms = [term for term in raw_terms if not term.startswith("+")]
|
|
scoring_terms = required_terms + optional_terms if required_terms else raw_terms
|
|
if not scoring_terms:
|
|
return []
|
|
|
|
term_patterns = {
|
|
term: re.compile(r"\b" + re.escape(term) + r"\b")
|
|
for term in scoring_terms
|
|
}
|
|
|
|
candidates = tools
|
|
if required_terms:
|
|
candidates = [
|
|
tool for tool in tools
|
|
if _matches_all_required_terms(tool, required_terms, term_patterns)
|
|
]
|
|
|
|
scored: list[tuple[BaseTool, int]] = []
|
|
for tool in candidates:
|
|
score = _score_tool(tool, scoring_terms, term_patterns)
|
|
if score > 0:
|
|
scored.append((tool, score))
|
|
scored.sort(key=lambda item: (-item[1], item[0].name))
|
|
return scored[:max_results]
|
|
|
|
|
|
def _parse_select_names(query: str) -> list[str]:
|
|
raw = query.split(":", 1)[1] if ":" in query else query
|
|
return [name.strip() for name in raw.split(",") if name.strip()]
|
|
|
|
|
|
def _get_pending_mcp_servers(context: Any | None) -> list[str]:
|
|
if context is None:
|
|
return []
|
|
for attr in ("pending_mcp_servers", "pending_mcp_server_names"):
|
|
value = getattr(context, attr, None)
|
|
if isinstance(value, (list, tuple, set)):
|
|
return sorted(str(item) for item in value if item)
|
|
return []
|
|
|
|
|
|
def _matches_all_required_terms(
|
|
tool: BaseTool,
|
|
required_terms: list[str],
|
|
term_patterns: dict[str, re.Pattern[str]],
|
|
) -> bool:
|
|
parsed = _parse_tool_name(tool.name)
|
|
desc = (tool.description or "").lower()
|
|
hint = (getattr(tool, "search_hint", "") or "").lower()
|
|
return all(
|
|
term in parsed["parts"]
|
|
or any(term in part for part in parsed["parts"])
|
|
or term_patterns[term].search(desc)
|
|
or (hint and term_patterns[term].search(hint))
|
|
for term in required_terms
|
|
)
|
|
|
|
|
|
def _score_tool(
|
|
tool: BaseTool,
|
|
query_terms: list[str],
|
|
term_patterns: dict[str, re.Pattern[str]],
|
|
) -> int:
|
|
parsed = _parse_tool_name(tool.name)
|
|
desc = (tool.description or "").lower()
|
|
hint = (getattr(tool, "search_hint", "") or "").lower()
|
|
score = 0
|
|
for term in query_terms:
|
|
pattern = term_patterns[term]
|
|
name_matched = False
|
|
if term in parsed["parts"]:
|
|
score += 12 if parsed["is_mcp"] else 10
|
|
name_matched = True
|
|
elif any(term in part for part in parsed["parts"]):
|
|
score += 6 if parsed["is_mcp"] else 5
|
|
name_matched = True
|
|
if not name_matched and term in parsed["full"]:
|
|
score += 3
|
|
if hint and pattern.search(hint):
|
|
score += 4
|
|
if pattern.search(desc):
|
|
score += 2
|
|
return score
|
|
|
|
|
|
def _parse_tool_name(name: str) -> dict[str, Any]:
|
|
if name.startswith("mcp__"):
|
|
without_prefix = name[5:].lower()
|
|
parts: list[str] = []
|
|
for segment in without_prefix.split("__"):
|
|
parts.extend(segment.split("_"))
|
|
parts = [part for part in parts if part]
|
|
return {
|
|
"parts": parts,
|
|
"full": without_prefix.replace("__", " ").replace("_", " "),
|
|
"is_mcp": True,
|
|
}
|
|
spaced = re.sub(r"([a-z])([A-Z])", r"\1 \2", name)
|
|
spaced = spaced.replace("_", " ").lower()
|
|
parts = [part for part in spaced.split() if part]
|
|
return {"parts": parts, "full": " ".join(parts), "is_mcp": False}
|