Coverage for stackone_ai / toolset.py: 92%
429 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-04-02 08:51 +0000
1from __future__ import annotations
3import asyncio
4import base64
5import concurrent.futures
6import fnmatch
7import json
8import logging
9import os
10import threading
11from collections.abc import Coroutine, Sequence
12from dataclasses import dataclass
13from importlib import metadata
14from typing import Any, Literal, TypedDict, TypeVar
16from pydantic import BaseModel, Field, PrivateAttr, ValidationError, field_validator
18from stackone_ai.constants import DEFAULT_BASE_URL
19from stackone_ai.models import (
20 ExecuteConfig,
21 JsonDict,
22 ParameterLocation,
23 StackOneAPIError,
24 StackOneTool,
25 ToolParameters,
26 Tools,
27)
28from stackone_ai.semantic_search import (
29 SemanticSearchClient,
30 SemanticSearchError,
31 SemanticSearchResult,
32)
33from stackone_ai.utils.normalize import _normalize_action_name
35logger = logging.getLogger("stackone.tools")
37SearchMode = Literal["auto", "semantic", "local"]
40class SearchConfig(TypedDict, total=False):
41 """Search configuration for the StackOneToolSet constructor.
43 When provided as a dict, sets default search options that flow through
44 to ``search_tools()``, ``get_search_tool()``, and ``search_action_names()``.
45 Per-call options override these defaults.
47 When set to ``None``, search is disabled entirely.
48 When omitted, defaults to ``{"method": "auto"}``.
49 """
51 method: SearchMode
52 """Search backend to use. Defaults to ``"auto"``."""
53 top_k: int
54 """Maximum number of tools to return."""
55 min_similarity: float
56 """Minimum similarity score threshold 0-1."""
59class ExecuteToolsConfig(TypedDict, total=False):
60 """Execution configuration for the StackOneToolSet constructor.
62 Controls default account scoping for tool execution.
64 When set to ``None`` (default), no account scoping is applied.
65 When provided, ``account_ids`` flow through to ``openai(mode="search_and_execute")``
66 and ``fetch_tools()`` as defaults.
67 """
69 account_ids: list[str]
70 """Account IDs to scope tool discovery and execution."""
73_SEARCH_DEFAULT: SearchConfig = {"method": "auto"}
75try:
76 _SDK_VERSION = metadata.version("stackone-ai")
77except metadata.PackageNotFoundError: # pragma: no cover - best-effort fallback when running from source
78 _SDK_VERSION = "dev"
79_RPC_PARAMETER_LOCATIONS = {
80 "action": ParameterLocation.BODY,
81 "body": ParameterLocation.BODY,
82 "headers": ParameterLocation.BODY,
83 "path": ParameterLocation.BODY,
84 "query": ParameterLocation.BODY,
85}
86_USER_AGENT = f"stackone-ai-python/{_SDK_VERSION}"
89# --- Internal tool_search + tool_execute ---
92class _SearchInput(BaseModel):
93 """Input validation for tool_search."""
95 query: str = Field(..., min_length=1)
96 connector: str | None = None
97 top_k: int | None = Field(default=None, ge=1, le=50)
99 @field_validator("query")
100 @classmethod
101 def validate_query(cls, v: str) -> str:
102 trimmed = v.strip()
103 if not trimmed: 103 ↛ 104line 103 didn't jump to line 104 because the condition on line 103 was never true
104 raise ValueError("query must be a non-empty string")
105 return trimmed
108class _SearchTool(StackOneTool):
109 """LLM-callable tool that searches for available StackOne tools."""
111 _toolset: Any = PrivateAttr(default=None)
113 def execute(
114 self, arguments: str | JsonDict | None = None, *, options: JsonDict | None = None
115 ) -> JsonDict:
116 try:
117 if isinstance(arguments, str):
118 raw_params = json.loads(arguments)
119 else:
120 raw_params = arguments or {}
122 parsed = _SearchInput(**raw_params)
124 search_config = self._toolset._search_config or {}
125 results = self._toolset.search_tools(
126 parsed.query,
127 connector=parsed.connector or search_config.get("connector"),
128 top_k=parsed.top_k or search_config.get("top_k") or 5,
129 min_similarity=search_config.get("min_similarity"),
130 search=search_config.get("method"),
131 account_ids=self._toolset._account_ids,
132 )
134 return {
135 "tools": [
136 {
137 "name": t.name,
138 "description": t.description,
139 "parameters": t.parameters.properties,
140 }
141 for t in results
142 ],
143 "total": len(results),
144 "query": parsed.query,
145 }
146 except (json.JSONDecodeError, ValidationError) as exc:
147 return {"error": f"Invalid input: {exc}", "query": raw_params if "raw_params" in dir() else None}
150class _ExecuteInput(BaseModel):
151 """Input validation for tool_execute."""
153 tool_name: str = Field(..., min_length=1)
154 parameters: dict[str, Any] = Field(default_factory=dict)
156 @field_validator("tool_name")
157 @classmethod
158 def validate_tool_name(cls, v: str) -> str:
159 trimmed = v.strip()
160 if not trimmed: 160 ↛ 161line 160 didn't jump to line 161 because the condition on line 160 was never true
161 raise ValueError("tool_name must be a non-empty string")
162 return trimmed
165class _ExecuteTool(StackOneTool):
166 """LLM-callable tool that executes a StackOne tool by name."""
168 _toolset: Any = PrivateAttr(default=None)
169 _cached_tools: Any = PrivateAttr(default=None)
171 def execute(
172 self, arguments: str | JsonDict | None = None, *, options: JsonDict | None = None
173 ) -> JsonDict:
174 tool_name = "unknown"
175 try:
176 if isinstance(arguments, str):
177 raw_params = json.loads(arguments)
178 else:
179 raw_params = arguments or {}
181 parsed = _ExecuteInput(**raw_params)
182 tool_name = parsed.tool_name
184 if self._cached_tools is None:
185 self._cached_tools = self._toolset.fetch_tools(account_ids=self._toolset._account_ids)
187 target = self._cached_tools.get_tool(parsed.tool_name)
189 if target is None:
190 return {
191 "error": (
192 f'Tool "{parsed.tool_name}" not found. Use tool_search to find available tools.'
193 ),
194 }
196 return target.execute(parsed.parameters, options=options)
197 except StackOneAPIError as exc:
198 return {
199 "error": str(exc),
200 "status_code": exc.status_code,
201 "response_body": exc.response_body,
202 "tool_name": tool_name,
203 }
204 except (json.JSONDecodeError, ValidationError) as exc:
205 return {"error": f"Invalid input: {exc}", "tool_name": tool_name}
208def _create_search_tool(api_key: str) -> _SearchTool:
209 name = "tool_search"
210 description = (
211 "Search for available tools by describing what you need. "
212 "Returns matching tool names, descriptions, and parameter schemas. "
213 "Use the returned parameter schemas to know exactly what to pass "
214 "when calling tool_execute."
215 )
216 parameters = ToolParameters(
217 type="object",
218 properties={
219 "query": {
220 "type": "string",
221 "description": (
222 "Natural language description of what you need "
223 '(e.g. "create an employee", "list time off requests")'
224 ),
225 },
226 "connector": {
227 "type": "string",
228 "description": 'Optional connector filter (e.g. "bamboohr")',
229 "nullable": True,
230 },
231 "top_k": {
232 "type": "integer",
233 "description": "Max results to return (1-50, default 5)",
234 "minimum": 1,
235 "maximum": 50,
236 "nullable": True,
237 },
238 },
239 )
240 execute_config = ExecuteConfig(
241 name=name,
242 method="POST",
243 url="local://meta/search",
244 parameter_locations={
245 "query": ParameterLocation.BODY,
246 "connector": ParameterLocation.BODY,
247 "top_k": ParameterLocation.BODY,
248 },
249 )
251 tool = _SearchTool.__new__(_SearchTool)
252 StackOneTool.__init__(
253 tool,
254 description=description,
255 parameters=parameters,
256 _execute_config=execute_config,
257 _api_key=api_key,
258 )
259 return tool
262def _create_execute_tool(api_key: str) -> _ExecuteTool:
263 name = "tool_execute"
264 description = (
265 "Execute a tool by name with the given parameters. "
266 "Use tool_search first to find available tools. "
267 "The parameters field must match the parameter schema returned "
268 "by tool_search. Pass parameters as a nested object matching "
269 "the schema structure."
270 )
271 parameters = ToolParameters(
272 type="object",
273 properties={
274 "tool_name": {
275 "type": "string",
276 "description": "Exact tool name from tool_search results",
277 },
278 "parameters": {
279 "type": "object",
280 "description": "Parameters for the tool, matching the schema from tool_search.",
281 "nullable": True,
282 },
283 },
284 )
285 execute_config = ExecuteConfig(
286 name=name,
287 method="POST",
288 url="local://meta/execute",
289 parameter_locations={
290 "tool_name": ParameterLocation.BODY,
291 "parameters": ParameterLocation.BODY,
292 },
293 )
295 tool = _ExecuteTool.__new__(_ExecuteTool)
296 StackOneTool.__init__(
297 tool,
298 description=description,
299 parameters=parameters,
300 _execute_config=execute_config,
301 _api_key=api_key,
302 )
303 return tool
306T = TypeVar("T")
309@dataclass
310class _McpToolDefinition:
311 name: str
312 description: str | None
313 input_schema: dict[str, Any]
316class ToolsetError(Exception):
317 """Base exception for toolset errors"""
319 pass
322class ToolsetConfigError(ToolsetError):
323 """Raised when there is an error in the toolset configuration"""
325 pass
328class ToolsetLoadError(ToolsetError):
329 """Raised when there is an error loading tools"""
331 pass
334def _run_async(awaitable: Coroutine[Any, Any, T]) -> T:
335 """Run a coroutine, even when called from an existing event loop."""
337 try:
338 asyncio.get_running_loop()
339 except RuntimeError:
340 return asyncio.run(awaitable)
342 result: dict[str, T] = {}
343 error: dict[str, BaseException] = {}
345 def runner() -> None:
346 try:
347 result["value"] = asyncio.run(awaitable)
348 except BaseException as exc: # pragma: no cover - surfaced in caller context
349 error["error"] = exc
351 thread = threading.Thread(target=runner, daemon=True)
352 thread.start()
353 thread.join()
355 if "error" in error:
356 raise error["error"]
358 return result["value"]
361def _build_auth_header(api_key: str) -> str:
362 token = base64.b64encode(f"{api_key}:".encode()).decode()
363 return f"Basic {token}"
366def _fetch_mcp_tools(endpoint: str, headers: dict[str, str]) -> list[_McpToolDefinition]:
367 try:
368 from mcp import types as mcp_types # ty: ignore[unresolved-import]
369 from mcp.client.session import ClientSession # ty: ignore[unresolved-import]
370 from mcp.client.streamable_http import streamablehttp_client # ty: ignore[unresolved-import]
371 except ImportError as exc: # pragma: no cover - depends on optional extra
372 raise ToolsetConfigError(
373 "MCP dependencies are required for fetch_tools. Install with 'pip install \"stackone-ai[mcp]\"'."
374 ) from exc
376 async def _list() -> list[_McpToolDefinition]:
377 async with streamablehttp_client(endpoint, headers=headers) as (read_stream, write_stream, _):
378 session = ClientSession(
379 read_stream,
380 write_stream,
381 client_info=mcp_types.Implementation(name="stackone-ai-python", version=_SDK_VERSION),
382 )
383 async with session:
384 await session.initialize()
385 cursor: str | None = None
386 collected: list[_McpToolDefinition] = []
387 while True:
388 result = await session.list_tools(cursor)
389 for tool in result.tools:
390 input_schema = tool.inputSchema or {}
391 collected.append(
392 _McpToolDefinition(
393 name=tool.name,
394 description=tool.description,
395 input_schema=dict(input_schema),
396 )
397 )
398 cursor = result.nextCursor
399 if cursor is None:
400 break
401 return collected
403 return _run_async(_list())
406class _StackOneRpcTool(StackOneTool):
407 """RPC-backed tool wired to the StackOne actions RPC endpoint."""
409 def __init__(
410 self,
411 *,
412 name: str,
413 description: str,
414 parameters: ToolParameters,
415 api_key: str,
416 base_url: str,
417 account_id: str | None,
418 ) -> None:
419 execute_config = ExecuteConfig(
420 method="POST",
421 url=f"{base_url.rstrip('/')}/actions/rpc",
422 name=name,
423 headers={},
424 body_type="json",
425 parameter_locations=dict(_RPC_PARAMETER_LOCATIONS),
426 )
427 super().__init__(
428 description=description,
429 parameters=parameters,
430 _execute_config=execute_config,
431 _api_key=api_key,
432 _account_id=account_id,
433 )
435 def execute(
436 self, arguments: str | dict[str, Any] | None = None, *, options: dict[str, Any] | None = None
437 ) -> dict[str, Any]:
438 parsed_arguments = self._parse_arguments(arguments)
440 body_payload = self._extract_record(parsed_arguments.pop("body", None))
441 headers_payload = self._extract_record(parsed_arguments.pop("headers", None))
442 path_payload = self._extract_record(parsed_arguments.pop("path", None))
443 query_payload = self._extract_record(parsed_arguments.pop("query", None))
445 rpc_body: dict[str, Any] = dict(body_payload or {})
446 for key, value in parsed_arguments.items():
447 rpc_body[key] = value
449 payload: dict[str, Any] = {
450 "action": self.name,
451 "body": rpc_body,
452 "headers": self._build_action_headers(headers_payload),
453 }
454 if path_payload:
455 payload["path"] = path_payload
456 if query_payload:
457 payload["query"] = query_payload
459 return super().execute(payload, options=options)
461 def _parse_arguments(self, arguments: str | dict[str, Any] | None) -> dict[str, Any]:
462 if arguments is None:
463 return {}
464 if isinstance(arguments, str):
465 parsed = json.loads(arguments)
466 else:
467 parsed = arguments
468 if not isinstance(parsed, dict):
469 raise ValueError("Tool arguments must be a JSON object")
470 return dict(parsed)
472 @staticmethod
473 def _extract_record(value: Any) -> dict[str, Any] | None:
474 if isinstance(value, dict):
475 return dict(value)
476 return None
478 def _build_action_headers(self, additional_headers: dict[str, Any] | None) -> dict[str, str]:
479 headers: dict[str, str] = {}
480 account_id = self.get_account_id()
481 if account_id:
482 headers["x-account-id"] = account_id
484 if additional_headers:
485 for key, value in additional_headers.items():
486 if value is None:
487 continue
488 headers[str(key)] = str(value)
490 headers.pop("Authorization", None)
491 return headers
494class SearchTool:
495 """Callable search tool that wraps StackOneToolSet.search_tools().
497 Designed for agent loops — call it with a query to get Tools back.
499 Example::
501 toolset = StackOneToolSet()
502 search_tool = toolset.get_search_tool()
503 tools = search_tool("manage employee records", account_ids=["acc-123"])
504 """
506 def __init__(self, toolset: StackOneToolSet, config: SearchConfig | None = None) -> None:
507 self._toolset = toolset
508 self._config: SearchConfig = config or {}
510 def __call__(
511 self,
512 query: str,
513 *,
514 connector: str | None = None,
515 top_k: int | None = None,
516 min_similarity: float | None = None,
517 account_ids: list[str] | None = None,
518 search: SearchMode | None = None,
519 ) -> Tools:
520 """Search for tools using natural language.
522 Args:
523 query: Natural language description of needed functionality
524 connector: Optional provider/connector filter (e.g., "bamboohr", "slack")
525 top_k: Maximum number of tools to return. Overrides constructor default.
526 min_similarity: Minimum similarity score threshold 0-1. Overrides constructor default.
527 account_ids: Optional account IDs (uses set_accounts() if not provided)
528 search: Override the default search mode for this call
530 Returns:
531 Tools collection with matched tools
532 """
533 effective_top_k = top_k if top_k is not None else self._config.get("top_k")
534 effective_min_sim = (
535 min_similarity if min_similarity is not None else self._config.get("min_similarity")
536 )
537 effective_search = search if search is not None else self._config.get("method", "auto")
538 return self._toolset.search_tools(
539 query,
540 connector=connector,
541 top_k=effective_top_k,
542 min_similarity=effective_min_sim,
543 account_ids=account_ids,
544 search=effective_search,
545 )
548class StackOneToolSet:
549 """Main class for accessing StackOne tools"""
551 def __init__(
552 self,
553 api_key: str | None = None,
554 account_id: str | None = None,
555 base_url: str | None = None,
556 search: SearchConfig | None = None,
557 execute: ExecuteToolsConfig | None = None,
558 ) -> None:
559 """Initialize StackOne tools with authentication
561 Args:
562 api_key: Optional API key. If not provided, will try to get from STACKONE_API_KEY env var
563 account_id: Optional account ID
564 base_url: Optional base URL override for API requests
565 search: Search configuration. Controls default search behavior.
566 Pass ``None`` (default) to disable search — ``toolset.openai()``
567 will return all regular tools.
568 Pass ``{}`` or ``{"method": "auto"}`` to enable search with defaults.
569 Pass ``{"method": "semantic", "top_k": 5}`` for custom defaults.
570 Per-call options always override these defaults.
571 execute: Execution configuration. Controls default account scoping
572 for tool execution. Pass ``{"account_ids": ["acc-1"]}`` to scope
573 meta tools to specific accounts.
575 Raises:
576 ToolsetConfigError: If no API key is provided or found in environment
577 """
578 api_key_value = api_key or os.getenv("STACKONE_API_KEY")
579 if not api_key_value:
580 raise ToolsetConfigError(
581 "API key must be provided either through api_key parameter or "
582 "STACKONE_API_KEY environment variable"
583 )
584 self.api_key: str = api_key_value
585 self.account_id = account_id
586 self.base_url = base_url or DEFAULT_BASE_URL
587 self._account_ids: list[str] = []
588 self._semantic_client: SemanticSearchClient | None = None
589 self._search_config: SearchConfig | None = search
590 self._execute_config: ExecuteToolsConfig | None = execute
591 self._tools_cache: Tools | None = None
593 def set_accounts(self, account_ids: list[str]) -> StackOneToolSet:
594 """Set account IDs for filtering tools
596 Args:
597 account_ids: List of account IDs to filter tools by
599 Returns:
600 This toolset instance for chaining
601 """
602 self._account_ids = account_ids
603 return self
605 def get_search_tool(self, *, search: SearchMode | None = None) -> SearchTool:
606 """Get a callable search tool that returns Tools collections.
608 Returns a callable that wraps :meth:`search_tools` for use in agent loops.
609 The returned tool is directly callable: ``search_tool("query")`` returns
610 :class:`Tools`.
612 Uses the constructor's search config as defaults. Per-call options override.
614 Args:
615 search: Override the default search mode. If not provided, uses
616 the constructor's search config.
618 Returns:
619 SearchTool instance
621 Example::
623 toolset = StackOneToolSet()
624 search_tool = toolset.get_search_tool()
625 tools = search_tool("manage employee records", account_ids=["acc-123"])
626 """
627 if self._search_config is None: 627 ↛ 628line 627 didn't jump to line 628 because the condition on line 627 was never true
628 raise ToolsetConfigError(
629 "Search is disabled. Initialize StackOneToolSet with a search config to enable."
630 )
632 config: SearchConfig = {**self._search_config}
633 if search is not None: 633 ↛ 636line 633 didn't jump to line 636 because the condition on line 633 was always true
634 config["method"] = search
636 return SearchTool(self, config=config)
638 def _build_tools(self, account_ids: list[str] | None = None) -> Tools:
639 """Build tool_search + tool_execute tools scoped to this toolset."""
640 if self._search_config is None:
641 raise ToolsetConfigError(
642 "Search is disabled. Initialize StackOneToolSet with a search config to enable."
643 )
645 if account_ids:
646 self._account_ids = account_ids
648 search_tool = _create_search_tool(self.api_key)
649 search_tool._toolset = self
651 execute_tool = _create_execute_tool(self.api_key)
652 execute_tool._toolset = self
654 return Tools([search_tool, execute_tool])
656 def openai(
657 self,
658 *,
659 mode: Literal["search_and_execute"] | None = None,
660 account_ids: list[str] | None = None,
661 ) -> list[dict[str, Any]]:
662 """Get tools in OpenAI function calling format.
664 Args:
665 mode: Tool mode.
666 ``None`` (default): fetch all tools and convert to OpenAI format.
667 ``"search_and_execute"``: return two meta tools (tool_search + tool_execute)
668 that let the LLM discover and execute tools on-demand.
669 account_ids: Account IDs to scope tools. Overrides the ``execute``
670 config from the constructor.
672 Returns:
673 List of tool definitions in OpenAI function format.
675 Examples::
677 # All tools
678 toolset = StackOneToolSet()
679 tools = toolset.openai()
681 # Meta tools for agent-driven discovery
682 toolset = StackOneToolSet()
683 tools = toolset.openai(mode="search_and_execute")
684 """
685 effective_account_ids = account_ids or (
686 self._execute_config.get("account_ids") if self._execute_config else None
687 )
689 if mode == "search_and_execute":
690 return self._build_tools(account_ids=effective_account_ids).to_openai()
692 return self.fetch_tools(account_ids=effective_account_ids).to_openai()
694 def langchain(
695 self,
696 *,
697 mode: Literal["search_and_execute"] | None = None,
698 account_ids: list[str] | None = None,
699 ) -> Sequence[Any]:
700 """Get tools in LangChain format.
702 Args:
703 mode: Tool mode.
704 ``None`` (default): fetch all tools and convert to LangChain format.
705 ``"search_and_execute"``: return two tools (tool_search + tool_execute)
706 that let the LLM discover and execute tools on-demand.
707 The framework handles tool execution automatically.
708 account_ids: Account IDs to scope tools. Overrides the ``execute``
709 config from the constructor.
711 Returns:
712 List of LangChain tool objects.
713 """
714 effective_account_ids = account_ids or (
715 self._execute_config.get("account_ids") if self._execute_config else None
716 )
718 if mode == "search_and_execute":
719 return self._build_tools(account_ids=effective_account_ids).to_langchain()
721 return self.fetch_tools(account_ids=effective_account_ids).to_langchain()
723 def execute(
724 self,
725 tool_name: str,
726 arguments: str | dict[str, Any] | None = None,
727 ) -> dict[str, Any]:
728 """Execute a tool by name.
730 Use with ``openai(mode="search_and_execute")`` in manual agent loops —
731 pass the tool name and arguments from the LLM's tool call directly.
733 Tools are cached after the first call.
735 Args:
736 tool_name: The tool name from the LLM's tool call
737 (e.g. ``"tool_search"`` or ``"tool_execute"``).
738 arguments: The arguments from the LLM's tool call,
739 as a JSON string or dict.
741 Returns:
742 Tool execution result as a dict.
743 """
744 if self._tools_cache is None:
745 self._tools_cache = self._build_tools()
747 tool = self._tools_cache.get_tool(tool_name)
748 if tool is None:
749 return {"error": f'Tool "{tool_name}" not found.'}
750 return tool.execute(arguments)
752 @property
753 def semantic_client(self) -> SemanticSearchClient:
754 """Lazy initialization of semantic search client.
756 Returns:
757 SemanticSearchClient instance configured with the toolset's API key and base URL
758 """
759 if self._semantic_client is None:
760 self._semantic_client = SemanticSearchClient(
761 api_key=self.api_key,
762 base_url=self.base_url,
763 )
764 return self._semantic_client
766 def _local_search(
767 self,
768 query: str,
769 all_tools: Tools,
770 *,
771 connector: str | None = None,
772 top_k: int | None = None,
773 min_similarity: float | None = None,
774 ) -> Tools:
775 """Run local BM25+TF-IDF search over already-fetched tools."""
776 from stackone_ai.local_search import ToolIndex
778 available_connectors = all_tools.get_connectors()
779 if not available_connectors: 779 ↛ 780line 779 didn't jump to line 780 because the condition on line 779 was never true
780 return Tools([])
782 index = ToolIndex(list(all_tools))
783 results = index.search(
784 query,
785 limit=top_k if top_k is not None else 5,
786 min_score=min_similarity if min_similarity is not None else 0.0,
787 )
788 matched_names = [r.name for r in results]
789 tool_map = {t.name: t for t in all_tools}
790 filter_connectors = {connector.lower()} if connector else available_connectors
791 matched_tools = [
792 tool_map[name]
793 for name in matched_names
794 if name in tool_map and name.split("_")[0].lower() in filter_connectors
795 ]
796 return Tools(matched_tools[:top_k] if top_k is not None else matched_tools)
798 def search_tools(
799 self,
800 query: str,
801 *,
802 connector: str | None = None,
803 top_k: int | None = None,
804 min_similarity: float | None = None,
805 account_ids: list[str] | None = None,
806 search: SearchMode | None = None,
807 ) -> Tools:
808 """Search for and fetch tools using semantic or local search.
810 This method discovers relevant tools based on natural language queries.
811 Constructor search config provides defaults; per-call args override.
813 Args:
814 query: Natural language description of needed functionality
815 (e.g., "create employee", "send a message")
816 connector: Optional provider/connector filter (e.g., "bamboohr", "slack")
817 top_k: Maximum number of tools to return. Overrides constructor default.
818 min_similarity: Minimum similarity score threshold 0-1. Overrides constructor default.
819 account_ids: Optional account IDs (uses set_accounts() if not provided)
820 search: Search backend to use. Overrides constructor default.
821 - ``"auto"`` (default): try semantic search first, fall back to local
822 BM25+TF-IDF if the API is unavailable.
823 - ``"semantic"``: use only the semantic search API; raises
824 ``SemanticSearchError`` on failure.
825 - ``"local"``: use only local BM25+TF-IDF search (no API call to the
826 semantic search endpoint).
828 Returns:
829 Tools collection with matched tools from linked accounts
831 Raises:
832 ToolsetConfigError: If search is disabled (``search=None`` in constructor)
833 SemanticSearchError: If the API call fails and search is ``"semantic"``
835 Examples:
836 # Semantic search (default with local fallback)
837 tools = toolset.search_tools("manage employee records", top_k=5)
839 # Explicit semantic search
840 tools = toolset.search_tools("manage employees", search="semantic")
842 # Local BM25+TF-IDF search
843 tools = toolset.search_tools("manage employees", search="local")
845 # Filter by connector
846 tools = toolset.search_tools(
847 "create time off request",
848 connector="bamboohr",
849 search="semantic",
850 )
851 """
852 if self._search_config is None: 852 ↛ 853line 852 didn't jump to line 853 because the condition on line 852 was never true
853 raise ToolsetConfigError(
854 "Search is disabled. Initialize StackOneToolSet with a search config to enable."
855 )
857 # Merge constructor defaults with per-call overrides
858 effective_search: SearchMode = (
859 search if search is not None else self._search_config.get("method", "auto")
860 )
861 effective_top_k = top_k if top_k is not None else self._search_config.get("top_k")
862 effective_min_sim = (
863 min_similarity if min_similarity is not None else self._search_config.get("min_similarity")
864 )
866 all_tools = self.fetch_tools(account_ids=account_ids)
867 available_connectors = all_tools.get_connectors()
869 if not available_connectors: 869 ↛ 870line 869 didn't jump to line 870 because the condition on line 869 was never true
870 return Tools([])
872 # Local-only search — skip semantic API entirely
873 if effective_search == "local":
874 return self._local_search(
875 query, all_tools, connector=connector, top_k=effective_top_k, min_similarity=effective_min_sim
876 )
878 try:
879 # Determine which connectors to search
880 if connector:
881 connectors_to_search = {connector.lower()} & available_connectors
882 if not connectors_to_search: 882 ↛ 883line 882 didn't jump to line 883 because the condition on line 882 was never true
883 return Tools([])
884 else:
885 connectors_to_search = available_connectors
887 # Search each connector in parallel
888 def _search_one(c: str) -> list[SemanticSearchResult]:
889 resp = self.semantic_client.search(
890 query=query, connector=c, top_k=effective_top_k, min_similarity=effective_min_sim
891 )
892 return list(resp.results)
894 all_results: list[SemanticSearchResult] = []
895 last_error: SemanticSearchError | None = None
896 max_workers = min(len(connectors_to_search), 10)
897 with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as pool:
898 futures = {pool.submit(_search_one, c): c for c in connectors_to_search}
899 for future in concurrent.futures.as_completed(futures):
900 try:
901 all_results.extend(future.result())
902 except SemanticSearchError as e:
903 last_error = e
905 # If ALL connector searches failed, re-raise to trigger fallback
906 if not all_results and last_error is not None:
907 raise last_error
909 # Sort by score, apply top_k
910 all_results.sort(key=lambda r: r.similarity_score, reverse=True)
911 if effective_top_k is not None: 911 ↛ 914line 911 didn't jump to line 914 because the condition on line 911 was always true
912 all_results = all_results[:effective_top_k]
914 if not all_results: 914 ↛ 915line 914 didn't jump to line 915 because the condition on line 914 was never true
915 return Tools([])
917 # 1. Parse composite IDs to MCP-format action names, deduplicate
918 seen_names: set[str] = set()
919 action_names: list[str] = []
920 for result in all_results:
921 name = _normalize_action_name(result.id)
922 if name in seen_names:
923 continue
924 seen_names.add(name)
925 action_names.append(name)
927 if not action_names: 927 ↛ 928line 927 didn't jump to line 928 because the condition on line 927 was never true
928 return Tools([])
930 # 2. Use MCP tools (already fetched) — schemas come from the source of truth
931 # 3. Filter to only the tools search found, preserving search relevance order
932 action_order = {name: i for i, name in enumerate(action_names)}
933 matched_tools = [t for t in all_tools if t.name in seen_names]
934 matched_tools.sort(key=lambda t: action_order.get(t.name, float("inf")))
936 # Auto mode: if semantic returned results but none matched MCP tools, fall back to local
937 if effective_search == "auto" and len(matched_tools) == 0:
938 logger.warning(
939 "Semantic search returned %d results but none matched MCP tools, "
940 "falling back to local search",
941 len(all_results),
942 )
943 return self._local_search(
944 query,
945 all_tools,
946 connector=connector,
947 top_k=effective_top_k,
948 min_similarity=effective_min_sim,
949 )
951 return Tools(matched_tools)
953 except SemanticSearchError as e:
954 if effective_search == "semantic":
955 raise
957 logger.warning("Semantic search failed (%s), falling back to local BM25+TF-IDF search", e)
958 return self._local_search(
959 query, all_tools, connector=connector, top_k=effective_top_k, min_similarity=effective_min_sim
960 )
962 def search_action_names(
963 self,
964 query: str,
965 *,
966 connector: str | None = None,
967 account_ids: list[str] | None = None,
968 top_k: int | None = None,
969 min_similarity: float | None = None,
970 ) -> list[SemanticSearchResult]:
971 """Search for action names without fetching tools.
973 Useful when you need to inspect search results before fetching,
974 or when building custom filtering logic.
976 Args:
977 query: Natural language description of needed functionality
978 connector: Optional provider/connector filter (single connector)
979 account_ids: Optional account IDs to scope results to connectors
980 available in those accounts (uses set_accounts() if not provided).
981 When provided, results are filtered to only matching connectors.
982 top_k: Maximum number of results. If None, uses the backend default.
983 min_similarity: Minimum similarity score threshold 0-1. If not provided,
984 the server uses its default.
986 Returns:
987 List of SemanticSearchResult with action names, scores, and metadata.
988 Versioned API names are normalized to MCP format but results are NOT
989 deduplicated — multiple API versions of the same action may appear
990 with their individual scores.
992 Examples:
993 # Lightweight: inspect results before fetching
994 results = toolset.search_action_names("manage employees")
995 for r in results:
996 print(f"{r.id}: {r.similarity_score:.2f}")
998 # Account-scoped: only results for connectors in linked accounts
999 results = toolset.search_action_names(
1000 "create employee",
1001 account_ids=["acc-123"],
1002 top_k=5
1003 )
1004 """
1005 if self._search_config is None: 1005 ↛ 1006line 1005 didn't jump to line 1006 because the condition on line 1005 was never true
1006 raise ToolsetConfigError(
1007 "Search is disabled. Initialize StackOneToolSet with search config to enable."
1008 )
1010 # Merge constructor defaults with per-call overrides
1011 effective_top_k = top_k if top_k is not None else self._search_config.get("top_k")
1012 effective_min_sim = (
1013 min_similarity if min_similarity is not None else self._search_config.get("min_similarity")
1014 )
1016 # Resolve available connectors from account_ids (same pattern as search_tools)
1017 available_connectors: set[str] | None = None
1018 effective_account_ids = account_ids or self._account_ids
1019 if effective_account_ids:
1020 all_tools = self.fetch_tools(account_ids=effective_account_ids)
1021 available_connectors = all_tools.get_connectors()
1022 if not available_connectors: 1022 ↛ 1023line 1022 didn't jump to line 1023 because the condition on line 1022 was never true
1023 return []
1025 try:
1026 if available_connectors:
1027 # Parallel per-connector search (only user's connectors)
1028 if connector: 1028 ↛ 1029line 1028 didn't jump to line 1029 because the condition on line 1028 was never true
1029 connectors_to_search = {connector.lower()} & available_connectors
1030 else:
1031 connectors_to_search = available_connectors
1033 def _search_one(c: str) -> list[SemanticSearchResult]:
1034 try:
1035 resp = self.semantic_client.search(
1036 query=query,
1037 connector=c,
1038 top_k=effective_top_k,
1039 min_similarity=effective_min_sim,
1040 )
1041 return list(resp.results)
1042 except SemanticSearchError:
1043 return []
1045 all_results: list[SemanticSearchResult] = []
1046 if connectors_to_search: 1046 ↛ 1067line 1046 didn't jump to line 1067 because the condition on line 1046 was always true
1047 max_workers = min(len(connectors_to_search), 10)
1048 with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as pool:
1049 futures = [pool.submit(_search_one, c) for c in connectors_to_search]
1050 for future in concurrent.futures.as_completed(futures):
1051 all_results.extend(future.result())
1052 else:
1053 # No account filtering — single global search
1054 response = self.semantic_client.search(
1055 query=query,
1056 connector=connector,
1057 top_k=effective_top_k,
1058 min_similarity=effective_min_sim,
1059 )
1060 all_results = list(response.results)
1062 except SemanticSearchError as e:
1063 logger.warning("Semantic search failed: %s", e)
1064 return []
1066 # Sort by score
1067 all_results.sort(key=lambda r: r.similarity_score, reverse=True)
1068 return all_results[:effective_top_k] if effective_top_k is not None else all_results
1070 def _filter_by_provider(self, tool_name: str, providers: list[str]) -> bool:
1071 """Check if a tool name matches any of the provider filters
1073 Args:
1074 tool_name: Name of the tool to check
1075 providers: List of provider names (case-insensitive)
1077 Returns:
1078 True if the tool matches any provider, False otherwise
1079 """
1080 # Extract provider from tool name (assuming format: provider_action)
1081 provider = tool_name.split("_")[0].lower()
1082 provider_set = {p.lower() for p in providers}
1083 return provider in provider_set
1085 def _filter_by_action(self, tool_name: str, actions: list[str]) -> bool:
1086 """Check if a tool name matches any of the action patterns
1088 Args:
1089 tool_name: Name of the tool to check
1090 actions: List of action patterns (supports glob patterns)
1092 Returns:
1093 True if the tool matches any action pattern, False otherwise
1094 """
1095 return any(fnmatch.fnmatch(tool_name, pattern) for pattern in actions)
1097 def fetch_tools(
1098 self,
1099 *,
1100 account_ids: list[str] | None = None,
1101 providers: list[str] | None = None,
1102 actions: list[str] | None = None,
1103 ) -> Tools:
1104 """Fetch tools with optional filtering by account IDs, providers, and actions
1106 Args:
1107 account_ids: Optional list of account IDs to filter by.
1108 If not provided, uses accounts set via set_accounts()
1109 providers: Optional list of provider names (e.g., ['hibob', 'bamboohr']).
1110 Case-insensitive matching.
1111 actions: Optional list of action patterns with glob support
1112 (e.g., ['*_list_employees', 'hibob_create_employees'])
1114 Returns:
1115 Collection of tools matching the filter criteria
1117 Raises:
1118 ToolsetLoadError: If there is an error loading the tools
1120 Examples:
1121 # Filter by account IDs
1122 tools = toolset.fetch_tools(account_ids=['123', '456'])
1124 # Filter by providers
1125 tools = toolset.fetch_tools(providers=['hibob', 'bamboohr'])
1127 # Filter by actions with glob patterns
1128 tools = toolset.fetch_tools(actions=['*_list_employees'])
1130 # Combine filters
1131 tools = toolset.fetch_tools(
1132 account_ids=['123'],
1133 providers=['hibob'],
1134 actions=['*_list_*']
1135 )
1137 # Use set_accounts() for account filtering
1138 toolset.set_accounts(['123', '456'])
1139 tools = toolset.fetch_tools()
1140 """
1141 try:
1142 effective_account_ids = account_ids or self._account_ids
1143 if not effective_account_ids and self.account_id:
1144 effective_account_ids = [self.account_id]
1146 if effective_account_ids:
1147 account_scope: list[str | None] = list(dict.fromkeys(effective_account_ids))
1148 else:
1149 account_scope = [None]
1151 endpoint = f"{self.base_url.rstrip('/')}/mcp"
1152 all_tools: list[StackOneTool] = []
1154 for account in account_scope:
1155 headers = self._build_mcp_headers(account)
1156 catalog = _fetch_mcp_tools(endpoint, headers)
1157 for tool_def in catalog:
1158 all_tools.append(self._create_rpc_tool(tool_def, account))
1160 if providers:
1161 all_tools = [tool for tool in all_tools if self._filter_by_provider(tool.name, providers)]
1163 if actions:
1164 all_tools = [tool for tool in all_tools if self._filter_by_action(tool.name, actions)]
1166 return Tools(all_tools)
1168 except ToolsetError:
1169 raise
1170 except Exception as exc: # pragma: no cover - unexpected runtime errors
1171 raise ToolsetLoadError(f"Error fetching tools: {exc}") from exc
1173 def _build_mcp_headers(self, account_id: str | None) -> dict[str, str]:
1174 headers = {
1175 "Authorization": _build_auth_header(self.api_key),
1176 "User-Agent": _USER_AGENT,
1177 }
1178 if account_id:
1179 headers["x-account-id"] = account_id
1180 return headers
1182 def _create_rpc_tool(self, tool_def: _McpToolDefinition, account_id: str | None) -> StackOneTool:
1183 schema = tool_def.input_schema or {}
1184 parameters = ToolParameters(
1185 type=str(schema.get("type") or "object"),
1186 properties=self._normalize_schema_properties(schema),
1187 )
1188 return _StackOneRpcTool(
1189 name=tool_def.name,
1190 description=tool_def.description or "",
1191 parameters=parameters,
1192 api_key=self.api_key,
1193 base_url=self.base_url,
1194 account_id=account_id,
1195 )
1197 def _normalize_schema_properties(self, schema: dict[str, Any]) -> dict[str, Any]:
1198 properties = schema.get("properties")
1199 if not isinstance(properties, dict):
1200 return {}
1202 required_fields = {str(name) for name in schema.get("required", [])}
1203 normalized: dict[str, Any] = {}
1205 for name, details in properties.items():
1206 if isinstance(details, dict):
1207 prop = dict(details)
1208 else:
1209 prop = {"description": str(details)}
1211 if name in required_fields:
1212 prop.setdefault("nullable", False)
1213 else:
1214 prop.setdefault("nullable", True)
1216 normalized[str(name)] = prop
1218 return normalized