Coverage for stackone_ai/toolset.py: 91%

1from __future__ import annotations

3import asyncio

4import base64

5import concurrent.futures

6import fnmatch

7import json

8import logging

9import os

10import re

11import threading

12from collections.abc import Coroutine, Sequence

13from dataclasses import dataclass

14from importlib import metadata

15from typing import TYPE_CHECKING, Any, Literal, TypedDict, TypeVar

17from pydantic import BaseModel, Field, PrivateAttr, ValidationError, field_validator

19from stackone_ai.constants import DEFAULT_BASE_URL

20from stackone_ai.models import (

21 ExecuteConfig,

22 JsonDict,

23 ParameterLocation,

24 StackOneAPIError,

25 StackOneTool,

26 ToolParameters,

27 Tools,

28)

29from stackone_ai.semantic_search import (

30 SemanticSearchClient,

31 SemanticSearchError,

32 SemanticSearchResult,

33)

34from stackone_ai.utils.normalize import _normalize_action_name

36if TYPE_CHECKING:

37 from pydantic_ai.tools import Tool as PydanticAITool

39logger = logging.getLogger("stackone.tools")

41SearchMode = Literal["auto", "semantic", "local"]

44class SearchConfig(TypedDict, total=False):

45 """Search configuration for the StackOneToolSet constructor.

47 When provided as a dict, sets default search options that flow through

48 to ``search_tools()``, ``get_search_tool()``, and ``search_action_names()``.

49 Per-call options override these defaults.

51 When set to ``None``, search is disabled entirely.

52 When omitted, defaults to ``{"method": "auto"}``.

53 """

55 method: SearchMode

56 """Search backend to use. Defaults to ``"auto"``."""

57 top_k: int

58 """Maximum number of tools to return."""

59 min_similarity: float

60 """Minimum similarity score threshold 0-1."""

63class ExecuteToolsConfig(TypedDict, total=False):

64 """Execution configuration for the StackOneToolSet constructor.

66 Controls default account scoping and timeout for tool execution.

68 When set to ``None`` (default), no account scoping is applied.

69 When provided, ``account_ids`` flow through to ``openai(mode="search_and_execute")``

70 and ``fetch_tools()`` as defaults.

71 """

73 account_ids: list[str]

74 """Account IDs to scope tool discovery and execution."""

76 timeout: float

77 """Request timeout in seconds. Default: 60. Can also be set as a top-level

78 constructor param which takes precedence."""

81_SEARCH_DEFAULT: SearchConfig = {"method": "auto"}

83try:

84 _SDK_VERSION = metadata.version("stackone-ai")

85except metadata.PackageNotFoundError: # pragma: no cover - best-effort fallback when running from source

86 _SDK_VERSION = "dev"

87_RPC_PARAMETER_LOCATIONS = {

88 "action": ParameterLocation.BODY,

89 "body": ParameterLocation.BODY,

90 "headers": ParameterLocation.BODY,

91 "path": ParameterLocation.BODY,

92 "query": ParameterLocation.BODY,

93}

94_USER_AGENT = f"stackone-ai-python/{_SDK_VERSION}"

96# Param-style pinned on the /mcp tool-listing URL. The MCP schema and the RPC-execution unwrap

97# (_split_envelope_params) must agree on this, so it is pinned rather than following the server

98# default — the server default is free to change without breaking the SDK.

99_MCP_PARAM_STYLE = "flat_prefixed"

100

101# Matches a flat_prefixed envelope key: `<location>_<field>` (e.g. `path_id`, `query_limit`).

102_FLAT_ENVELOPE_KEY_PATTERN = re.compile(r"^(path|query|body|headers)_(.+)$")

103

104

105# --- Internal tool_search + tool_execute ---

106

107

108class _SearchInput(BaseModel):

109 """Input validation for tool_search."""

110

111 query: str = Field(..., min_length=1)

112 connector: str | None = None

113 top_k: int | None = Field(default=None, ge=1, le=50)

114

115 @field_validator("query")

116 @classmethod

117 def validate_query(cls, v: str) -> str:

118 trimmed = v.strip()

119 if not trimmed: 119 ↛ 120line 119 didn't jump to line 120 because the condition on line 119 was never true

120 raise ValueError("query must be a non-empty string")

121 return trimmed

122

123

124class _SearchTool(StackOneTool):

125 """LLM-callable tool that searches for available StackOne tools."""

126

127 _toolset: Any = PrivateAttr(default=None)

128

129 def execute(

130 self, arguments: str | JsonDict | None = None, *, options: JsonDict | None = None

131 ) -> JsonDict:

132 try:

133 if isinstance(arguments, str):

134 raw_params = json.loads(arguments)

135 else:

136 raw_params = arguments or {}

137

138 parsed = _SearchInput(**raw_params)

139

140 search_config = self._toolset._search_config or {}

141 results = self._toolset.search_tools(

142 parsed.query,

143 connector=parsed.connector or search_config.get("connector"),

144 top_k=parsed.top_k or search_config.get("top_k") or 5,

145 min_similarity=search_config.get("min_similarity"),

146 search=search_config.get("method"),

147 account_ids=self._toolset._account_ids,

148 )

149

150 return {

151 "tools": [

152 {

153 "name": t.name,

154 "description": t.description,

155 "parameters": t.parameters.properties,

156 }

157 for t in results

158 ],

159 "total": len(results),

160 "query": parsed.query,

161 }

162 except (json.JSONDecodeError, ValidationError) as exc:

163 return {"error": f"Invalid input: {exc}", "query": raw_params if "raw_params" in dir() else None}

164

165

166class _ExecuteInput(BaseModel):

167 """Input validation for tool_execute."""

168

169 tool_name: str = Field(..., min_length=1)

170 parameters: dict[str, Any] = Field(default_factory=dict)

171

172 @field_validator("tool_name")

173 @classmethod

174 def validate_tool_name(cls, v: str) -> str:

175 trimmed = v.strip()

176 if not trimmed: 176 ↛ 177line 176 didn't jump to line 177 because the condition on line 176 was never true

177 raise ValueError("tool_name must be a non-empty string")

178 return trimmed

179

180

181class _ExecuteTool(StackOneTool):

182 """LLM-callable tool that executes a StackOne tool by name."""

183

184 _toolset: Any = PrivateAttr(default=None)

185

186 def execute(

187 self, arguments: str | JsonDict | None = None, *, options: JsonDict | None = None

188 ) -> JsonDict:

189 tool_name = "unknown"

190 try:

191 if isinstance(arguments, str):

192 raw_params = json.loads(arguments)

193 else:

194 raw_params = arguments or {}

195

196 parsed = _ExecuteInput(**raw_params)

197 tool_name = parsed.tool_name

198

199 tools = self._toolset.fetch_tools(account_ids=self._toolset._account_ids)

200 target = tools.get_tool(parsed.tool_name)

201

202 if target is None:

203 return {

204 "error": (

205 f'Tool "{parsed.tool_name}" not found. Use tool_search to find available tools.'

206 ),

207 }

208

209 return target.execute(parsed.parameters, options=options)

210 except StackOneAPIError as exc:

211 return {

212 "error": str(exc),

213 "status_code": exc.status_code,

214 "response_body": exc.response_body,

215 "tool_name": tool_name,

216 }

217 except (json.JSONDecodeError, ValidationError) as exc:

218 return {"error": f"Invalid input: {exc}", "tool_name": tool_name}

219

220

221def _create_search_tool(api_key: str, connectors: str = "") -> _SearchTool:

222 name = "tool_search"

223 connector_line = f" Available connectors: {connectors}." if connectors else ""

224 description = (

225 "Search for available tools by describing what you need. "

226 "Returns matching tool names, descriptions, and parameter schemas. "

227 "Use the returned parameter schemas to know exactly what to pass "

228 f"when calling tool_execute.{connector_line}"

229 )

230 parameters = ToolParameters(

231 type="object",

232 properties={

233 "query": {

234 "type": "string",

235 "description": (

236 "Natural language description of what you need "

237 '(e.g. "create an employee", "list time off requests")'

238 ),

239 },

240 "connector": {

241 "type": "string",

242 "description": 'Optional connector filter (e.g. "bamboohr")',

243 "nullable": True,

244 },

245 "top_k": {

246 "type": "integer",

247 "description": "Max results to return (1-50, default 5)",

248 "minimum": 1,

249 "maximum": 50,

250 "nullable": True,

251 },

252 },

253 )

254 execute_config = ExecuteConfig(

255 name=name,

256 method="POST",

257 url="local://meta/search",

258 parameter_locations={

259 "query": ParameterLocation.BODY,

260 "connector": ParameterLocation.BODY,

261 "top_k": ParameterLocation.BODY,

262 },

263 )

264

265 tool = _SearchTool.__new__(_SearchTool)

266 StackOneTool.__init__(

267 tool,

268 description=description,

269 parameters=parameters,

270 _execute_config=execute_config,

271 _api_key=api_key,

272 )

273 return tool

274

275

276def _create_execute_tool(api_key: str, connectors: str = "") -> _ExecuteTool:

277 name = "tool_execute"

278 connector_line = f" Available connectors: {connectors}." if connectors else ""

279 description = (

280 "Execute a tool by name with the given parameters. "

281 "Use tool_search first to find available tools. "

282 "The parameters field must match the parameter schema returned "

283 f"by tool_search. Pass parameters as a nested object matching the schema structure.{connector_line}"

284 )

285 parameters = ToolParameters(

286 type="object",

287 properties={

288 "tool_name": {

289 "type": "string",

290 "description": "Exact tool name from tool_search results",

291 },

292 "parameters": {

293 "type": "object",

294 "description": "Parameters for the tool, matching the schema from tool_search.",

295 "nullable": True,

296 },

297 },

298 )

299 execute_config = ExecuteConfig(

300 name=name,

301 method="POST",

302 url="local://meta/execute",

303 parameter_locations={

304 "tool_name": ParameterLocation.BODY,

305 "parameters": ParameterLocation.BODY,

306 },

307 )

308

309 tool = _ExecuteTool.__new__(_ExecuteTool)

310 StackOneTool.__init__(

311 tool,

312 description=description,

313 parameters=parameters,

314 _execute_config=execute_config,

315 _api_key=api_key,

316 )

317 return tool

320T = TypeVar("T")

323@dataclass

324class _McpToolDefinition:

325 name: str

326 description: str | None

327 input_schema: dict[str, Any]

328

329

330class ToolsetError(Exception):

331 """Base exception for toolset errors"""

332

333 pass

334

335

336class ToolsetConfigError(ToolsetError):

337 """Raised when there is an error in the toolset configuration"""

338

339 pass

340

341

342class ToolsetLoadError(ToolsetError):

343 """Raised when there is an error loading tools"""

344

345 pass

346

347

348def _run_async(awaitable: Coroutine[Any, Any, T]) -> T:

349 """Run a coroutine, even when called from an existing event loop."""

350

351 try:

352 asyncio.get_running_loop()

353 except RuntimeError:

354 return asyncio.run(awaitable)

355

356 result: dict[str, T] = {}

357 error: dict[str, BaseException] = {}

358

359 def runner() -> None:

360 try:

361 result["value"] = asyncio.run(awaitable)

362 except BaseException as exc: # pragma: no cover - surfaced in caller context

363 error["error"] = exc

364

365 thread = threading.Thread(target=runner, daemon=True)

366 thread.start()

367 thread.join()

368

369 if "error" in error:

370 raise error["error"]

371

372 return result["value"]

373

374

375def _build_auth_header(api_key: str) -> str:

376 token = base64.b64encode(f"{api_key}:".encode()).decode()

377 return f"Basic {token}"

378

379

380def _fetch_mcp_tools(endpoint: str, headers: dict[str, str]) -> list[_McpToolDefinition]:

381 try:

382 from mcp import types as mcp_types # ty: ignore[unresolved-import]

383 from mcp.client.session import ClientSession # ty: ignore[unresolved-import]

384 from mcp.client.streamable_http import streamablehttp_client # ty: ignore[unresolved-import]

385 except ImportError as exc: # pragma: no cover - depends on optional extra

386 raise ToolsetConfigError(

387 "MCP dependencies are required for fetch_tools. Install with 'pip install \"stackone-ai[mcp]\"'."

388 ) from exc

389

390 async def _list() -> list[_McpToolDefinition]:

391 async with streamablehttp_client(endpoint, headers=headers) as (read_stream, write_stream, _):

392 session = ClientSession(

393 read_stream,

394 write_stream,

395 client_info=mcp_types.Implementation(name="stackone-ai-python", version=_SDK_VERSION),

396 )

397 async with session:

398 await session.initialize()

399 cursor: str | None = None

400 collected: list[_McpToolDefinition] = []

401 while True:

402 result = await session.list_tools(cursor)

403 for tool in result.tools:

404 input_schema = tool.inputSchema or {}

405 collected.append(

406 _McpToolDefinition(

407 name=tool.name,

408 description=tool.description,

409 input_schema=dict(input_schema),

410 )

411 )

412 cursor = result.nextCursor

413 if cursor is None:

414 break

415 return collected

416

417 return _run_async(_list())

418

419

420class _StackOneRpcTool(StackOneTool):

421 """RPC-backed tool wired to the StackOne actions RPC endpoint."""

422

423 def __init__(

424 self,

425 *,

426 name: str,

427 description: str,

428 parameters: ToolParameters,

429 api_key: str,

430 base_url: str,

431 account_id: str | None,

432 timeout: float = 60.0,

433 ) -> None:

434 execute_config = ExecuteConfig(

435 method="POST",

436 url=f"{base_url.rstrip('/')}/actions/rpc",

437 name=name,

438 headers={},

439 body_type="json",

440 parameter_locations=dict(_RPC_PARAMETER_LOCATIONS),

441 timeout=timeout,

442 )

443 super().__init__(

444 description=description,

445 parameters=parameters,

446 _execute_config=execute_config,

447 _api_key=api_key,

448 _account_id=account_id,

449 )

450

451 def execute(

452 self, arguments: str | dict[str, Any] | None = None, *, options: dict[str, Any] | None = None

453 ) -> dict[str, Any]:

454 parsed_arguments = self._parse_arguments(arguments)

455

456 envelope = self._split_envelope_params(parsed_arguments)

457

458 payload: dict[str, Any] = {

459 "action": self.name,

460 "body": envelope["body"],

461 "headers": self._build_action_headers(envelope["headers"] or None),

462 }

463 if envelope["path"]:

464 payload["path"] = envelope["path"]

465 if envelope["query"]:

466 payload["query"] = envelope["query"]

467

468 return super().execute(payload, options=options)

469

470 def _parse_arguments(self, arguments: str | dict[str, Any] | None) -> dict[str, Any]:

471 if arguments is None:

472 return {}

473 if isinstance(arguments, str):

474 parsed = json.loads(arguments)

475 else:

476 parsed = arguments

477 if not isinstance(parsed, dict):

478 raise ValueError("Tool arguments must be a JSON object")

479 return dict(parsed)

480

481 @staticmethod

482 def _split_envelope_params(params: dict[str, Any]) -> dict[str, dict[str, Any]]:

483 """Split LLM-supplied tool arguments into the RPC envelope (path/query/headers/body).

484

485 Tools are listed with ``?param-style=flat_prefixed``, so keys arrive as

486 ``<location>_<field>`` (for example ``path_id``, ``query_limit``). The prefix carries

487 the parameter location, so the split needs no per-action schema. A bare dict-valued

488 ``path``/``query``/``headers``/``body`` key is still bucketed for clients holding a

489 cached nested schema, and any other key falls through to the body.

490 """

491 buckets: dict[str, dict[str, Any]] = {"path": {}, "query": {}, "headers": {}, "body": {}}

492 for key, value in params.items():

493 match = _FLAT_ENVELOPE_KEY_PATTERN.match(key)

494 if match:

495 location, field = match.group(1), match.group(2)

496 buckets[location].setdefault(field, value)

497 continue

498 if key in ("path", "query", "headers", "body") and isinstance(value, dict):

499 for field, field_value in value.items():

500 buckets[key].setdefault(field, field_value)

501 continue

502 buckets["body"][key] = value

503 return buckets

504

505 def _build_action_headers(self, additional_headers: dict[str, Any] | None) -> dict[str, str]:

506 headers: dict[str, str] = {}

507 account_id = self.get_account_id()

508 if account_id:

509 headers["x-account-id"] = account_id

510

511 if additional_headers:

512 for key, value in additional_headers.items():

513 if value is None:

514 continue

515 headers[str(key)] = str(value)

516

517 headers.pop("Authorization", None)

518 return headers

519

520

521class SearchTool:

522 """Callable search tool that wraps StackOneToolSet.search_tools().

523

524 Designed for agent loops — call it with a query to get Tools back.

525

526 Example::

527

528 toolset = StackOneToolSet()

529 search_tool = toolset.get_search_tool()

530 tools = search_tool("manage employee records", account_ids=["acc-123"])

531 """

532

533 def __init__(self, toolset: StackOneToolSet, config: SearchConfig | None = None) -> None:

534 self._toolset = toolset

535 self._config: SearchConfig = config or {}

536

537 def __call__(

538 self,

539 query: str,

540 *,

541 connector: str | None = None,

542 top_k: int | None = None,

543 min_similarity: float | None = None,

544 account_ids: list[str] | None = None,

545 search: SearchMode | None = None,

546 ) -> Tools:

547 """Search for tools using natural language.

548

549 Args:

550 query: Natural language description of needed functionality

551 connector: Optional provider/connector filter (e.g., "bamboohr", "slack")

552 top_k: Maximum number of tools to return. Overrides constructor default.

553 min_similarity: Minimum similarity score threshold 0-1. Overrides constructor default.

554 account_ids: Optional account IDs (uses set_accounts() if not provided)

555 search: Override the default search mode for this call

556

557 Returns:

558 Tools collection with matched tools

559 """

560 effective_top_k = top_k if top_k is not None else self._config.get("top_k")

561 effective_min_sim = (

562 min_similarity if min_similarity is not None else self._config.get("min_similarity")

563 )

564 effective_search = search if search is not None else self._config.get("method", "auto")

565 return self._toolset.search_tools(

566 query,

567 connector=connector,

568 top_k=effective_top_k,

569 min_similarity=effective_min_sim,

570 account_ids=account_ids,

571 search=effective_search,

572 )

573

574

575class StackOneToolSet:

576 """Main class for accessing StackOne tools"""

577

578 def __init__(

579 self,

580 api_key: str | None = None,

581 account_id: str | None = None,

582 base_url: str | None = None,

583 search: SearchConfig | None = None,

584 execute: ExecuteToolsConfig | None = None,

585 timeout: float | None = None,

586 ) -> None:

587 """Initialize StackOne tools with authentication

588

589 Args:

590 api_key: Optional API key. If not provided, will try to get from STACKONE_API_KEY env var

591 account_id: Optional account ID

592 base_url: Optional base URL override for API requests

593 search: Search configuration. Controls default search behavior.

594 Pass ``None`` (default) to disable search — ``toolset.openai()``

595 will return all regular tools.

596 Pass ``{}`` or ``{"method": "auto"}`` to enable search with defaults.

597 Pass ``{"method": "semantic", "top_k": 5}`` for custom defaults.

598 Per-call options always override these defaults.

599 execute: Execution configuration. Controls default account scoping

600 for tool execution. Pass ``{"account_ids": ["acc-1"]}`` to scope

601 tools to specific accounts.

602 timeout: Request timeout in seconds for tool execution HTTP calls.

603 Default: 60. Takes precedence over ``execute.timeout`` if set.

604 Increase for slow providers (e.g. Workday).

605

606 Raises:

607 ToolsetConfigError: If no API key is provided or found in environment

608 """

609 api_key_value = api_key or os.getenv("STACKONE_API_KEY")

610 if not api_key_value:

611 raise ToolsetConfigError(

612 "API key must be provided either through api_key parameter or "

613 "STACKONE_API_KEY environment variable"

614 )

615 self.api_key: str = api_key_value

616 self.account_id = account_id

617 self.base_url = base_url or DEFAULT_BASE_URL

618 self._account_ids: list[str] = execute.get("account_ids", []) if execute else []

619 self._semantic_client: SemanticSearchClient | None = None

620 self._search_config: SearchConfig | None = search

621 self._execute_config: ExecuteToolsConfig | None = execute

622 execute_timeout = execute.get("timeout") if execute else None

623 self._timeout: float = timeout if timeout is not None else (execute_timeout or 60.0)

624 self._tools_cache: Tools | None = None

625 self._catalog_cache: dict[tuple[Any, ...], Tools] = {}

626 self._tool_index_cache: tuple[int, Any] | None = None

627

628 def set_accounts(self, account_ids: list[str]) -> StackOneToolSet:

629 """Set account IDs for filtering tools

630

631 Args:

632 account_ids: List of account IDs to filter tools by

633

634 Returns:

635 This toolset instance for chaining

636 """

637 self._account_ids = account_ids

638 self.clear_catalog_cache()

639 return self

640

641 def clear_catalog_cache(self) -> None:

642 """Invalidate cached tool catalog and local search index.

643

644 Call when linked accounts change outside of ``set_accounts`` or when

645 you need to force a fresh fetch from the StackOne MCP endpoint.

646 """

647 self._catalog_cache.clear()

648 self._tool_index_cache = None

649

650 def get_search_tool(self, *, search: SearchMode | None = None) -> SearchTool:

651 """Get a callable search tool that returns Tools collections.

652

653 Returns a callable that wraps :meth:`search_tools` for use in agent loops.

654 The returned tool is directly callable: ``search_tool("query")`` returns

655 :class:`Tools`.

656

657 Uses the constructor's search config as defaults. Per-call options override.

658

659 Args:

660 search: Override the default search mode. If not provided, uses

661 the constructor's search config.

662

663 Returns:

664 SearchTool instance

665

666 Example::

667

668 toolset = StackOneToolSet(search={"method": "auto"})

669 search_tool = toolset.get_search_tool()

670 tools = search_tool("manage employee records", account_ids=["acc-123"])

671 """

672 if self._search_config is None: 672 ↛ 673line 672 didn't jump to line 673 because the condition on line 672 was never true

673 raise ToolsetConfigError(

674 "Search is disabled. Pass search={} (or search={'method': 'auto'}) to "

675 "StackOneToolSet(...) to enable. See README 'Search Tool' for options."

676 )

677

678 config: SearchConfig = {**self._search_config}

679 if search is not None: 679 ↛ 682line 679 didn't jump to line 682 because the condition on line 679 was always true

680 config["method"] = search

681

682 return SearchTool(self, config=config)

683

684 def _build_tools(self, account_ids: list[str] | None = None) -> Tools:

685 """Build tool_search + tool_execute tools scoped to this toolset."""

686 if self._search_config is None:

687 raise ToolsetConfigError(

688 "Search is disabled. Pass search={} (or search={'method': 'auto'}) to "

689 "StackOneToolSet(...) to enable. See README 'Search Tool' for options."

690 )

691

692 if account_ids:

693 self._account_ids = account_ids

694

695 # Discover available connectors for dynamic descriptions

696 connectors_str = ""

697 try:

698 all_tools = self.fetch_tools(account_ids=self._account_ids)

699 connectors = sorted(all_tools.get_connectors())

700 if connectors:

701 connectors_str = ", ".join(connectors)

702 except Exception:

703 logger.debug("Could not discover connectors for tool descriptions")

704

705 search_tool = _create_search_tool(self.api_key, connectors=connectors_str)

706 search_tool._toolset = self

707

708 execute_tool = _create_execute_tool(self.api_key, connectors=connectors_str)

709 execute_tool._toolset = self

710

711 return Tools([search_tool, execute_tool])

712

713 def openai(

714 self,

715 *,

716 mode: Literal["search_and_execute"] | None = None,

717 account_ids: list[str] | None = None,

718 ) -> list[dict[str, Any]]:

719 """Get tools in OpenAI function calling format.

720

721 Args:

722 mode: Tool mode.

723 ``None`` (default): fetch all tools and convert to OpenAI format.

724 ``"search_and_execute"``: return two meta tools (tool_search + tool_execute)

725 that let the LLM discover and execute tools on-demand.

726 account_ids: Account IDs to scope tools. Overrides the ``execute``

727 config from the constructor.

728

729 Returns:

730 List of tool definitions in OpenAI function format.

731

732 Examples::

733

734 # All tools

735 toolset = StackOneToolSet()

736 tools = toolset.openai()

737

738 # Meta tools for agent-driven discovery — search must be enabled

739 toolset = StackOneToolSet(search={"method": "auto"})

740 tools = toolset.openai(mode="search_and_execute")

741 """

742 effective_account_ids = account_ids or (

743 self._execute_config.get("account_ids") if self._execute_config else None

744 )

745

746 if mode == "search_and_execute":

747 return self._build_tools(account_ids=effective_account_ids).to_openai()

748

749 return self.fetch_tools(account_ids=effective_account_ids).to_openai()

750

751 def langchain(

752 self,

753 *,

754 mode: Literal["search_and_execute"] | None = None,

755 account_ids: list[str] | None = None,

756 ) -> Sequence[Any]:

757 """Get tools in LangChain format.

758

759 Args:

760 mode: Tool mode.

761 ``None`` (default): fetch all tools and convert to LangChain format.

762 ``"search_and_execute"``: return two tools (tool_search + tool_execute)

763 that let the LLM discover and execute tools on-demand.

764 The framework handles tool execution automatically.

765 account_ids: Account IDs to scope tools. Overrides the ``execute``

766 config from the constructor.

767

768 Returns:

769 List of LangChain tool objects.

770 """

771 effective_account_ids = account_ids or (

772 self._execute_config.get("account_ids") if self._execute_config else None

773 )

774

775 if mode == "search_and_execute":

776 return self._build_tools(account_ids=effective_account_ids).to_langchain()

777

778 return self.fetch_tools(account_ids=effective_account_ids).to_langchain()

779

780 def pydantic_ai(

781 self,

782 *,

783 mode: Literal["search_and_execute"] | None = None,

784 account_ids: list[str] | None = None,

785 ) -> list[PydanticAITool]:

786 """Get tools as Pydantic AI ``Tool`` instances.

787

788 Args:

789 mode: Tool mode.

790 ``None`` (default): fetch all tools and convert to Pydantic AI tools.

791 ``"search_and_execute"``: return two meta tools (tool_search + tool_execute)

792 that let the LLM discover and execute tools on-demand.

793 account_ids: Account IDs to scope tools. Overrides the ``execute``

794 config from the constructor.

795

796 Returns:

797 List of Pydantic AI ``Tool`` objects ready to pass to ``Agent(tools=...)``.

798

799 Requires ``stackone-ai[pydantic-ai]`` (installs ``pydantic-ai-slim``).

800

801 Examples::

802

803 # All tools

804 toolset = StackOneToolSet()

805 tools = toolset.pydantic_ai()

806 agent = Agent("openai:gpt-5.4", tools=tools)

807

808 # Meta tools for agent-driven discovery — search must be enabled

809 toolset = StackOneToolSet(search={"method": "auto"})

810 tools = toolset.pydantic_ai(mode="search_and_execute")

811 """

812 effective_account_ids = account_ids or (

813 self._execute_config.get("account_ids") if self._execute_config else None

814 )

815

816 if mode == "search_and_execute":

817 return self._build_tools(account_ids=effective_account_ids).to_pydantic_ai()

818

819 return self.fetch_tools(account_ids=effective_account_ids).to_pydantic_ai()

820

821 def execute(

822 self,

823 tool_name: str,

824 arguments: str | dict[str, Any] | None = None,

825 ) -> dict[str, Any]:

826 """Execute a tool by name.

827

828 Use with ``openai(mode="search_and_execute")`` in manual agent loops —

829 pass the tool name and arguments from the LLM's tool call directly.

830

831 Tools are cached after the first call.

832

833 Args:

834 tool_name: The tool name from the LLM's tool call

835 (e.g. ``"tool_search"`` or ``"tool_execute"``).

836 arguments: The arguments from the LLM's tool call,

837 as a JSON string or dict.

838

839 Returns:

840 Tool execution result as a dict.

841 """

842 if self._tools_cache is None:

843 self._tools_cache = self._build_tools()

844

845 tool = self._tools_cache.get_tool(tool_name)

846 if tool is None:

847 return {"error": f'Tool "{tool_name}" not found.'}

848 return tool.execute(arguments)

849

850 @property

851 def semantic_client(self) -> SemanticSearchClient:

852 """Lazy initialization of semantic search client.

853

854 Returns:

855 SemanticSearchClient instance configured with the toolset's API key and base URL

856 """

857 if self._semantic_client is None:

858 self._semantic_client = SemanticSearchClient(

859 api_key=self.api_key,

860 base_url=self.base_url,

861 )

862 return self._semantic_client

863

864 def _local_search(

865 self,

866 query: str,

867 all_tools: Tools,

868 *,

869 connector: str | None = None,

870 top_k: int | None = None,

871 min_similarity: float | None = None,

872 ) -> Tools:

873 """Run local BM25+TF-IDF search over already-fetched tools."""

874 from stackone_ai.local_search import ToolIndex

875

876 available_connectors = all_tools.get_connectors()

877 if not available_connectors: 877 ↛ 878line 877 didn't jump to line 878 because the condition on line 877 was never true

878 return Tools([])

879

880 cache_key = id(all_tools)

881 if self._tool_index_cache is None or self._tool_index_cache[0] != cache_key:

882 self._tool_index_cache = (cache_key, ToolIndex(list(all_tools)))

883 index = self._tool_index_cache[1]

884 results = index.search(

885 query,

886 limit=top_k if top_k is not None else 5,

887 min_score=min_similarity if min_similarity is not None else 0.0,

888 )

889 matched_names = [r.name for r in results]

890 tool_map = {t.name: t for t in all_tools}

891 filter_connectors = {connector.lower()} if connector else available_connectors

892 matched_tools = [

893 tool_map[name]

894 for name in matched_names

895 if name in tool_map and name.split("_")[0].lower() in filter_connectors

896 ]

897 return Tools(matched_tools[:top_k] if top_k is not None else matched_tools)

898

899 def search_tools(

900 self,

901 query: str,

902 *,

903 connector: str | None = None,

904 top_k: int | None = None,

905 min_similarity: float | None = None,

906 account_ids: list[str] | None = None,

907 search: SearchMode | None = None,

908 ) -> Tools:

909 """Search for and fetch tools using semantic or local search.

910

911 This method discovers relevant tools based on natural language queries.

912 Constructor search config provides defaults; per-call args override.

913

914 Args:

915 query: Natural language description of needed functionality

916 (e.g., "create employee", "send a message")

917 connector: Optional provider/connector filter (e.g., "bamboohr", "slack")

918 top_k: Maximum number of tools to return. Overrides constructor default.

919 min_similarity: Minimum similarity score threshold 0-1. Overrides constructor default.

920 account_ids: Optional account IDs (uses set_accounts() if not provided)

921 search: Search backend to use. Overrides constructor default.

922 - ``"auto"`` (default): try semantic search first, fall back to local

923 BM25+TF-IDF if the API is unavailable.

924 - ``"semantic"``: use only the semantic search API; raises

925 ``SemanticSearchError`` on failure.

926 - ``"local"``: use only local BM25+TF-IDF search (no API call to the

927 semantic search endpoint).

928

929 Returns:

930 Tools collection with matched tools from linked accounts

931

932 Raises:

933 ToolsetConfigError: If search is disabled (``search=None`` in constructor)

934 SemanticSearchError: If the API call fails and search is ``"semantic"``

935

936 Examples:

937 # Semantic search (default with local fallback)

938 tools = toolset.search_tools("manage employee records", top_k=5)

939

940 # Explicit semantic search

941 tools = toolset.search_tools("manage employees", search="semantic")

942

943 # Local BM25+TF-IDF search

944 tools = toolset.search_tools("manage employees", search="local")

945

946 # Filter by connector

947 tools = toolset.search_tools(

948 "create time off request",

949 connector="bamboohr",

950 search="semantic",

951 )

952 """

953 if self._search_config is None: 953 ↛ 954line 953 didn't jump to line 954 because the condition on line 953 was never true

954 raise ToolsetConfigError(

955 "Search is disabled. Pass search={} (or search={'method': 'auto'}) to "

956 "StackOneToolSet(...) to enable. See README 'Search Tool' for options."

957 )

958

959 # Merge constructor defaults with per-call overrides

960 effective_search: SearchMode = (

961 search if search is not None else self._search_config.get("method", "auto")

962 )

963 effective_top_k = top_k if top_k is not None else self._search_config.get("top_k")

964 effective_min_sim = (

965 min_similarity if min_similarity is not None else self._search_config.get("min_similarity")

966 )

967

968 all_tools = self.fetch_tools(account_ids=account_ids)

969 available_connectors = all_tools.get_connectors()

970

971 if not available_connectors: 971 ↛ 972line 971 didn't jump to line 972 because the condition on line 971 was never true

972 return Tools([])

973

974 # Local-only search — skip semantic API entirely

975 if effective_search == "local":

976 return self._local_search(

977 query, all_tools, connector=connector, top_k=effective_top_k, min_similarity=effective_min_sim

978 )

979

980 try:

981 # Determine which connectors to search

982 if connector:

983 connectors_to_search = {connector.lower()} & available_connectors

984 if not connectors_to_search: 984 ↛ 985line 984 didn't jump to line 985 because the condition on line 984 was never true

985 return Tools([])

986 else:

987 connectors_to_search = available_connectors

988

989 # Search each connector in parallel

990 def _search_one(c: str) -> list[SemanticSearchResult]:

991 resp = self.semantic_client.search(

992 query=query, connector=c, top_k=effective_top_k, min_similarity=effective_min_sim

993 )

994 return list(resp.results)

995

996 all_results: list[SemanticSearchResult] = []

997 last_error: SemanticSearchError | None = None

998 max_workers = min(len(connectors_to_search), 10)

999 with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as pool:

1000 futures = {pool.submit(_search_one, c): c for c in connectors_to_search}

1001 for future in concurrent.futures.as_completed(futures):

1002 try:

1003 all_results.extend(future.result())

1004 except SemanticSearchError as e:

1005 last_error = e

1006

1007 # If ALL connector searches failed, re-raise to trigger fallback

1008 if not all_results and last_error is not None:

1009 raise last_error

1010

1011 # Sort by score, apply top_k

1012 all_results.sort(key=lambda r: r.similarity_score, reverse=True)

1013 if effective_top_k is not None: 1013 ↛ 1016line 1013 didn't jump to line 1016 because the condition on line 1013 was always true

1014 all_results = all_results[:effective_top_k]

1015

1016 if not all_results: 1016 ↛ 1017line 1016 didn't jump to line 1017 because the condition on line 1016 was never true

1017 return Tools([])

1018

1019 # 1. Parse composite IDs to MCP-format action names, deduplicate

1020 seen_names: set[str] = set()

1021 action_names: list[str] = []

1022 for result in all_results:

1023 name = _normalize_action_name(result.id)

1024 if name in seen_names:

1025 continue

1026 seen_names.add(name)

1027 action_names.append(name)

1028

1029 if not action_names: 1029 ↛ 1030line 1029 didn't jump to line 1030 because the condition on line 1029 was never true

1030 return Tools([])

1031

1032 # 2. Use MCP tools (already fetched) — schemas come from the source of truth

1033 # 3. Filter to only the tools search found, preserving search relevance order

1034 action_order = {name: i for i, name in enumerate(action_names)}

1035 matched_tools = [t for t in all_tools if t.name in seen_names]

1036 matched_tools.sort(key=lambda t: action_order.get(t.name, float("inf")))

1037

1038 # Auto mode: if semantic returned results but none matched MCP tools, fall back to local

1039 if effective_search == "auto" and len(matched_tools) == 0:

1040 logger.warning(

1041 "Semantic search returned %d results but none matched MCP tools, "

1042 "falling back to local search",

1043 len(all_results),

1044 )

1045 return self._local_search(

1046 query,

1047 all_tools,

1048 connector=connector,

1049 top_k=effective_top_k,

1050 min_similarity=effective_min_sim,

1051 )

1052

1053 return Tools(matched_tools)

1054

1055 except SemanticSearchError as e:

1056 if effective_search == "semantic":

1057 raise

1058

1059 logger.warning("Semantic search failed (%s), falling back to local BM25+TF-IDF search", e)

1060 return self._local_search(

1061 query, all_tools, connector=connector, top_k=effective_top_k, min_similarity=effective_min_sim

1062 )

1063

1064 def search_action_names(

1065 self,

1066 query: str,

1067 *,

1068 connector: str | None = None,

1069 account_ids: list[str] | None = None,

1070 top_k: int | None = None,

1071 min_similarity: float | None = None,

1072 ) -> list[SemanticSearchResult]:

1073 """Search for action names without fetching tools.

1074

1075 Useful when you need to inspect search results before fetching,

1076 or when building custom filtering logic.

1077

1078 Args:

1079 query: Natural language description of needed functionality

1080 connector: Optional provider/connector filter (single connector)

1081 account_ids: Optional account IDs to scope results to connectors

1082 available in those accounts (uses set_accounts() if not provided).

1083 When provided, results are filtered to only matching connectors.

1084 top_k: Maximum number of results. If None, uses the backend default.

1085 min_similarity: Minimum similarity score threshold 0-1. If not provided,

1086 the server uses its default.

1087

1088 Returns:

1089 List of SemanticSearchResult with action names, scores, and metadata.

1090 Versioned API names are normalized to MCP format but results are NOT

1091 deduplicated — multiple API versions of the same action may appear

1092 with their individual scores.

1093

1094 Examples:

1095 # Lightweight: inspect results before fetching

1096 results = toolset.search_action_names("manage employees")

1097 for r in results:

1098 print(f"{r.id}: {r.similarity_score:.2f}")

1099

1100 # Account-scoped: only results for connectors in linked accounts

1101 results = toolset.search_action_names(

1102 "create employee",

1103 account_ids=["acc-123"],

1104 top_k=5

1105 )

1106 """

1107 if self._search_config is None: 1107 ↛ 1108line 1107 didn't jump to line 1108 because the condition on line 1107 was never true

1108 raise ToolsetConfigError(

1109 "Search is disabled. Pass search={} (or search={'method': 'auto'}) to "

1110 "StackOneToolSet(...) to enable. See README 'Search Tool' for options."

1111 )

1112

1113 # Merge constructor defaults with per-call overrides

1114 effective_top_k = top_k if top_k is not None else self._search_config.get("top_k")

1115 effective_min_sim = (

1116 min_similarity if min_similarity is not None else self._search_config.get("min_similarity")

1117 )

1118

1119 # Resolve available connectors from account_ids (same pattern as search_tools)

1120 available_connectors: set[str] | None = None

1121 effective_account_ids = account_ids or self._account_ids

1122 if effective_account_ids:

1123 all_tools = self.fetch_tools(account_ids=effective_account_ids)

1124 available_connectors = all_tools.get_connectors()

1125 if not available_connectors: 1125 ↛ 1126line 1125 didn't jump to line 1126 because the condition on line 1125 was never true

1126 return []

1127

1128 try:

1129 if available_connectors:

1130 # Parallel per-connector search (only user's connectors)

1131 if connector: 1131 ↛ 1132line 1131 didn't jump to line 1132 because the condition on line 1131 was never true

1132 connectors_to_search = {connector.lower()} & available_connectors

1133 else:

1134 connectors_to_search = available_connectors

1135

1136 def _search_one(c: str) -> list[SemanticSearchResult]:

1137 try:

1138 resp = self.semantic_client.search(

1139 query=query,

1140 connector=c,

1141 top_k=effective_top_k,

1142 min_similarity=effective_min_sim,

1143 )

1144 return list(resp.results)

1145 except SemanticSearchError:

1146 return []

1147

1148 all_results: list[SemanticSearchResult] = []

1149 if connectors_to_search: 1149 ↛ 1170line 1149 didn't jump to line 1170 because the condition on line 1149 was always true

1150 max_workers = min(len(connectors_to_search), 10)

1151 with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as pool:

1152 futures = [pool.submit(_search_one, c) for c in connectors_to_search]

1153 for future in concurrent.futures.as_completed(futures):

1154 all_results.extend(future.result())

1155 else:

1156 # No account filtering — single global search

1157 response = self.semantic_client.search(

1158 query=query,

1159 connector=connector,

1160 top_k=effective_top_k,

1161 min_similarity=effective_min_sim,

1162 )

1163 all_results = list(response.results)

1164

1165 except SemanticSearchError as e:

1166 logger.warning("Semantic search failed: %s", e)

1167 return []

1168

1169 # Sort by score

1170 all_results.sort(key=lambda r: r.similarity_score, reverse=True)

1171 return all_results[:effective_top_k] if effective_top_k is not None else all_results

1172

1173 def _filter_by_provider(self, tool_name: str, providers: list[str]) -> bool:

1174 """Check if a tool name matches any of the provider filters

1175

1176 Args:

1177 tool_name: Name of the tool to check

1178 providers: List of provider names (case-insensitive)

1179

1180 Returns:

1181 True if the tool matches any provider, False otherwise

1182 """

1183 # Extract provider from tool name (assuming format: provider_action)

1184 provider = tool_name.split("_")[0].lower()

1185 provider_set = {p.lower() for p in providers}

1186 return provider in provider_set

1187

1188 def _filter_by_action(self, tool_name: str, actions: list[str]) -> bool:

1189 """Check if a tool name matches any of the action patterns

1190

1191 Args:

1192 tool_name: Name of the tool to check

1193 actions: List of action patterns (supports glob patterns)

1194

1195 Returns:

1196 True if the tool matches any action pattern, False otherwise

1197 """

1198 return any(fnmatch.fnmatch(tool_name, pattern) for pattern in actions)

1199

1200 def fetch_tools(

1201 self,

1202 *,

1203 account_ids: list[str] | None = None,

1204 providers: list[str] | None = None,

1205 actions: list[str] | None = None,

1206 ) -> Tools:

1207 """Fetch tools with optional filtering by account IDs, providers, and actions

1208

1209 Args:

1210 account_ids: Optional list of account IDs to filter by.

1211 If not provided, uses accounts set via set_accounts()

1212 providers: Optional list of provider names (e.g., ['hibob', 'bamboohr']).

1213 Case-insensitive matching.

1214 actions: Optional list of action patterns with glob support

1215 (e.g., ['*_list_employees', 'hibob_create_employees'])

1216

1217 Returns:

1218 Collection of tools matching the filter criteria

1219

1220 Raises:

1221 ToolsetLoadError: If there is an error loading the tools

1222

1223 Examples:

1224 # Filter by account IDs

1225 tools = toolset.fetch_tools(account_ids=['123', '456'])

1226

1227 # Filter by providers

1228 tools = toolset.fetch_tools(providers=['hibob', 'bamboohr'])

1229

1230 # Filter by actions with glob patterns

1231 tools = toolset.fetch_tools(actions=['*_list_employees'])

1232

1233 # Combine filters

1234 tools = toolset.fetch_tools(

1235 account_ids=['123'],

1236 providers=['hibob'],

1237 actions=['*_list_*']

1238 )

1239

1240 # Use set_accounts() for account filtering

1241 toolset.set_accounts(['123', '456'])

1242 tools = toolset.fetch_tools()

1243 """

1244 try:

1245 effective_account_ids = account_ids or self._account_ids

1246 if not effective_account_ids and self.account_id:

1247 effective_account_ids = [self.account_id]

1248

1249 if effective_account_ids:

1250 account_scope: list[str | None] = list(dict.fromkeys(effective_account_ids))

1251 else:

1252 account_scope = [None]

1253

1254 cache_key = (

1255 tuple(sorted(account_scope, key=lambda a: (a is None, a))),

1256 tuple(sorted(p.lower() for p in providers)) if providers else None,

1257 tuple(sorted(actions)) if actions else None,

1258 )

1259 cached = self._catalog_cache.get(cache_key)

1260 if cached is not None:

1261 return cached

1262

1263 endpoint = f"{self.base_url.rstrip('/')}/mcp?param-style={_MCP_PARAM_STYLE}"

1264

1265 def _fetch_for_account(account: str | None) -> list[StackOneTool]:

1266 headers = self._build_mcp_headers(account)

1267 catalog = _fetch_mcp_tools(endpoint, headers)

1268 return [self._create_rpc_tool(tool_def, account) for tool_def in catalog]

1269

1270 all_tools: list[StackOneTool] = []

1271 if len(account_scope) == 1:

1272 all_tools.extend(_fetch_for_account(account_scope[0]))

1273 else:

1274 max_workers = min(len(account_scope), 10)

1275 with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as pool:

1276 futures = [pool.submit(_fetch_for_account, acc) for acc in account_scope]

1277 for future in futures:

1278 all_tools.extend(future.result())

1279

1280 if providers:

1281 all_tools = [tool for tool in all_tools if self._filter_by_provider(tool.name, providers)]

1282

1283 if actions:

1284 all_tools = [tool for tool in all_tools if self._filter_by_action(tool.name, actions)]

1285

1286 result = Tools(all_tools)

1287 self._catalog_cache[cache_key] = result

1288 return result

1289

1290 except ToolsetError:

1291 raise

1292 except Exception as exc: # pragma: no cover - unexpected runtime errors

1293 raise ToolsetLoadError(f"Error fetching tools: {exc}") from exc

1294

1295 def _build_mcp_headers(self, account_id: str | None) -> dict[str, str]:

1296 headers = {

1297 "Authorization": _build_auth_header(self.api_key),

1298 "User-Agent": _USER_AGENT,

1299 }

1300 if account_id:

1301 headers["x-account-id"] = account_id

1302 return headers

1303

1304 def _create_rpc_tool(self, tool_def: _McpToolDefinition, account_id: str | None) -> StackOneTool:

1305 schema = tool_def.input_schema or {}

1306 parameters = ToolParameters(

1307 type=str(schema.get("type") or "object"),

1308 properties=self._normalize_schema_properties(schema),

1309 )

1310 return _StackOneRpcTool(

1311 name=tool_def.name,

1312 description=tool_def.description or "",

1313 parameters=parameters,

1314 api_key=self.api_key,

1315 base_url=self.base_url,

1316 account_id=account_id,

1317 timeout=self._timeout,

1318 )

1319

1320 def _normalize_schema_properties(self, schema: dict[str, Any]) -> dict[str, Any]:

1321 properties = schema.get("properties")

1322 if not isinstance(properties, dict):

1323 return {}

1324

1325 required_fields = {str(name) for name in schema.get("required", [])}

1326 normalized: dict[str, Any] = {}

1327

1328 for name, details in properties.items():

1329 if isinstance(details, dict):

1330 prop = dict(details)

1331 else:

1332 prop = {"description": str(details)}

1333

1334 if name in required_fields:

1335 prop.setdefault("nullable", False)

1336 else:

1337 prop.setdefault("nullable", True)

1338

1339 normalized[str(name)] = prop

1340

1341 return normalized

Coverage for stackone_ai / toolset.py: 91%

474 statements