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 threading

11from collections.abc import Coroutine, Sequence

12from dataclasses import dataclass

13from importlib import metadata

14from typing import TYPE_CHECKING, 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

35if TYPE_CHECKING:

36 from pydantic_ai.tools import Tool as PydanticAITool

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

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

43class SearchConfig(TypedDict, total=False):

44 """Search configuration for the StackOneToolSet constructor.

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

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

48 Per-call options override these defaults.

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

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

52 """

54 method: SearchMode

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

56 top_k: int

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

58 min_similarity: float

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

62class ExecuteToolsConfig(TypedDict, total=False):

63 """Execution configuration for the StackOneToolSet constructor.

65 Controls default account scoping and timeout for tool execution.

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

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

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

70 """

72 account_ids: list[str]

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

75 timeout: float

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

77 constructor param which takes precedence."""

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

82try:

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

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

85 _SDK_VERSION = "dev"

86_RPC_PARAMETER_LOCATIONS = {

87 "action": ParameterLocation.BODY,

88 "body": ParameterLocation.BODY,

89 "headers": ParameterLocation.BODY,

90 "path": ParameterLocation.BODY,

91 "query": ParameterLocation.BODY,

92}

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

96# --- Internal tool_search + tool_execute ---

99class _SearchInput(BaseModel):

100 """Input validation for tool_search."""

101

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

103 connector: str | None = None

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

105

106 @field_validator("query")

107 @classmethod

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

109 trimmed = v.strip()

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

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

112 return trimmed

113

114

115class _SearchTool(StackOneTool):

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

117

118 _toolset: Any = PrivateAttr(default=None)

119

120 def execute(

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

122 ) -> JsonDict:

123 try:

124 if isinstance(arguments, str):

125 raw_params = json.loads(arguments)

126 else:

127 raw_params = arguments or {}

128

129 parsed = _SearchInput(**raw_params)

130

131 search_config = self._toolset._search_config or {}

132 results = self._toolset.search_tools(

133 parsed.query,

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

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

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

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

138 account_ids=self._toolset._account_ids,

139 )

140

141 return {

142 "tools": [

143 {

144 "name": t.name,

145 "description": t.description,

146 "parameters": t.parameters.properties,

147 }

148 for t in results

149 ],

150 "total": len(results),

151 "query": parsed.query,

152 }

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

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

155

156

157class _ExecuteInput(BaseModel):

158 """Input validation for tool_execute."""

159

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

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

162

163 @field_validator("tool_name")

164 @classmethod

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

166 trimmed = v.strip()

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

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

169 return trimmed

170

171

172class _ExecuteTool(StackOneTool):

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

174

175 _toolset: Any = PrivateAttr(default=None)

176

177 def execute(

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

179 ) -> JsonDict:

180 tool_name = "unknown"

181 try:

182 if isinstance(arguments, str):

183 raw_params = json.loads(arguments)

184 else:

185 raw_params = arguments or {}

186

187 parsed = _ExecuteInput(**raw_params)

188 tool_name = parsed.tool_name

189

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

191 target = tools.get_tool(parsed.tool_name)

192

193 if target is None:

194 return {

195 "error": (

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

197 ),

198 }

199

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

201 except StackOneAPIError as exc:

202 return {

203 "error": str(exc),

204 "status_code": exc.status_code,

205 "response_body": exc.response_body,

206 "tool_name": tool_name,

207 }

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

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

210

211

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

213 name = "tool_search"

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

215 description = (

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

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

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

219 f"when calling tool_execute.{connector_line}"

220 )

221 parameters = ToolParameters(

222 type="object",

223 properties={

224 "query": {

225 "type": "string",

226 "description": (

227 "Natural language description of what you need "

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

229 ),

230 },

231 "connector": {

232 "type": "string",

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

234 "nullable": True,

235 },

236 "top_k": {

237 "type": "integer",

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

239 "minimum": 1,

240 "maximum": 50,

241 "nullable": True,

242 },

243 },

244 )

245 execute_config = ExecuteConfig(

246 name=name,

247 method="POST",

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

249 parameter_locations={

250 "query": ParameterLocation.BODY,

251 "connector": ParameterLocation.BODY,

252 "top_k": ParameterLocation.BODY,

253 },

254 )

255

256 tool = _SearchTool.__new__(_SearchTool)

257 StackOneTool.__init__(

258 tool,

259 description=description,

260 parameters=parameters,

261 _execute_config=execute_config,

262 _api_key=api_key,

263 )

264 return tool

265

266

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

268 name = "tool_execute"

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

270 description = (

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

272 "Use tool_search first to find available tools. "

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

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

275 )

276 parameters = ToolParameters(

277 type="object",

278 properties={

279 "tool_name": {

280 "type": "string",

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

282 },

283 "parameters": {

284 "type": "object",

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

286 "nullable": True,

287 },

288 },

289 )

290 execute_config = ExecuteConfig(

291 name=name,

292 method="POST",

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

294 parameter_locations={

295 "tool_name": ParameterLocation.BODY,

296 "parameters": ParameterLocation.BODY,

297 },

298 )

299

300 tool = _ExecuteTool.__new__(_ExecuteTool)

301 StackOneTool.__init__(

302 tool,

303 description=description,

304 parameters=parameters,

305 _execute_config=execute_config,

306 _api_key=api_key,

307 )

308 return tool

311T = TypeVar("T")

314@dataclass

315class _McpToolDefinition:

316 name: str

317 description: str | None

318 input_schema: dict[str, Any]

319

320

321class ToolsetError(Exception):

322 """Base exception for toolset errors"""

323

324 pass

325

326

327class ToolsetConfigError(ToolsetError):

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

329

330 pass

331

332

333class ToolsetLoadError(ToolsetError):

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

335

336 pass

337

338

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

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

341

342 try:

343 asyncio.get_running_loop()

344 except RuntimeError:

345 return asyncio.run(awaitable)

346

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

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

349

350 def runner() -> None:

351 try:

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

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

354 error["error"] = exc

355

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

357 thread.start()

358 thread.join()

359

360 if "error" in error:

361 raise error["error"]

362

363 return result["value"]

364

365

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

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

368 return f"Basic {token}"

369

370

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

372 try:

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

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

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

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

377 raise ToolsetConfigError(

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

379 ) from exc

380

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

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

383 session = ClientSession(

384 read_stream,

385 write_stream,

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

387 )

388 async with session:

389 await session.initialize()

390 cursor: str | None = None

391 collected: list[_McpToolDefinition] = []

392 while True:

393 result = await session.list_tools(cursor)

394 for tool in result.tools:

395 input_schema = tool.inputSchema or {}

396 collected.append(

397 _McpToolDefinition(

398 name=tool.name,

399 description=tool.description,

400 input_schema=dict(input_schema),

401 )

402 )

403 cursor = result.nextCursor

404 if cursor is None:

405 break

406 return collected

407

408 return _run_async(_list())

409

410

411class _StackOneRpcTool(StackOneTool):

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

413

414 def __init__(

415 self,

416 *,

417 name: str,

418 description: str,

419 parameters: ToolParameters,

420 api_key: str,

421 base_url: str,

422 account_id: str | None,

423 timeout: float = 60.0,

424 ) -> None:

425 execute_config = ExecuteConfig(

426 method="POST",

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

428 name=name,

429 headers={},

430 body_type="json",

431 parameter_locations=dict(_RPC_PARAMETER_LOCATIONS),

432 timeout=timeout,

433 )

434 super().__init__(

435 description=description,

436 parameters=parameters,

437 _execute_config=execute_config,

438 _api_key=api_key,

439 _account_id=account_id,

440 )

441

442 def execute(

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

444 ) -> dict[str, Any]:

445 parsed_arguments = self._parse_arguments(arguments)

446

447 body_payload = self._extract_record(parsed_arguments.pop("body", None))

448 headers_payload = self._extract_record(parsed_arguments.pop("headers", None))

449 path_payload = self._extract_record(parsed_arguments.pop("path", None))

450 query_payload = self._extract_record(parsed_arguments.pop("query", None))

451

452 rpc_body: dict[str, Any] = dict(body_payload or {})

453 for key, value in parsed_arguments.items():

454 rpc_body[key] = value

455

456 payload: dict[str, Any] = {

457 "action": self.name,

458 "body": rpc_body,

459 "headers": self._build_action_headers(headers_payload),

460 }

461 if path_payload:

462 payload["path"] = path_payload

463 if query_payload:

464 payload["query"] = query_payload

465

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

467

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

469 if arguments is None:

470 return {}

471 if isinstance(arguments, str):

472 parsed = json.loads(arguments)

473 else:

474 parsed = arguments

475 if not isinstance(parsed, dict):

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

477 return dict(parsed)

478

479 @staticmethod

480 def _extract_record(value: Any) -> dict[str, Any] | None:

481 if isinstance(value, dict):

482 return dict(value)

483 return None

484

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

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

487 account_id = self.get_account_id()

488 if account_id:

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

490

491 if additional_headers:

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

493 if value is None:

494 continue

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

496

497 headers.pop("Authorization", None)

498 return headers

499

500

501class SearchTool:

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

503

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

505

506 Example::

507

508 toolset = StackOneToolSet()

509 search_tool = toolset.get_search_tool()

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

511 """

512

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

514 self._toolset = toolset

515 self._config: SearchConfig = config or {}

516

517 def __call__(

518 self,

519 query: str,

520 *,

521 connector: str | None = None,

522 top_k: int | None = None,

523 min_similarity: float | None = None,

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

525 search: SearchMode | None = None,

526 ) -> Tools:

527 """Search for tools using natural language.

528

529 Args:

530 query: Natural language description of needed functionality

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

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

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

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

535 search: Override the default search mode for this call

536

537 Returns:

538 Tools collection with matched tools

539 """

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

541 effective_min_sim = (

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

543 )

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

545 return self._toolset.search_tools(

546 query,

547 connector=connector,

548 top_k=effective_top_k,

549 min_similarity=effective_min_sim,

550 account_ids=account_ids,

551 search=effective_search,

552 )

553

554

555class StackOneToolSet:

556 """Main class for accessing StackOne tools"""

557

558 def __init__(

559 self,

560 api_key: str | None = None,

561 account_id: str | None = None,

562 base_url: str | None = None,

563 search: SearchConfig | None = None,

564 execute: ExecuteToolsConfig | None = None,

565 timeout: float | None = None,

566 ) -> None:

567 """Initialize StackOne tools with authentication

568

569 Args:

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

571 account_id: Optional account ID

572 base_url: Optional base URL override for API requests

573 search: Search configuration. Controls default search behavior.

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

575 will return all regular tools.

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

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

578 Per-call options always override these defaults.

579 execute: Execution configuration. Controls default account scoping

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

581 tools to specific accounts.

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

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

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

585

586 Raises:

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

588 """

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

590 if not api_key_value:

591 raise ToolsetConfigError(

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

593 "STACKONE_API_KEY environment variable"

594 )

595 self.api_key: str = api_key_value

596 self.account_id = account_id

597 self.base_url = base_url or DEFAULT_BASE_URL

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

599 self._semantic_client: SemanticSearchClient | None = None

600 self._search_config: SearchConfig | None = search

601 self._execute_config: ExecuteToolsConfig | None = execute

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

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

604 self._tools_cache: Tools | None = None

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

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

607

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

609 """Set account IDs for filtering tools

610

611 Args:

612 account_ids: List of account IDs to filter tools by

613

614 Returns:

615 This toolset instance for chaining

616 """

617 self._account_ids = account_ids

618 self.clear_catalog_cache()

619 return self

620

621 def clear_catalog_cache(self) -> None:

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

623

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

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

626 """

627 self._catalog_cache.clear()

628 self._tool_index_cache = None

629

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

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

632

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

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

635 :class:`Tools`.

636

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

638

639 Args:

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

641 the constructor's search config.

642

643 Returns:

644 SearchTool instance

645

646 Example::

647

648 toolset = StackOneToolSet()

649 search_tool = toolset.get_search_tool()

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

651 """

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

653 raise ToolsetConfigError(

654 "Search is disabled. Initialize StackOneToolSet with a search config to enable."

655 )

656

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

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

659 config["method"] = search

660

661 return SearchTool(self, config=config)

662

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

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

665 if self._search_config is None:

666 raise ToolsetConfigError(

667 "Search is disabled. Initialize StackOneToolSet with a search config to enable."

668 )

669

670 if account_ids:

671 self._account_ids = account_ids

672

673 # Discover available connectors for dynamic descriptions

674 connectors_str = ""

675 try:

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

677 connectors = sorted(all_tools.get_connectors())

678 if connectors:

679 connectors_str = ", ".join(connectors)

680 except Exception:

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

682

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

684 search_tool._toolset = self

685

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

687 execute_tool._toolset = self

688

689 return Tools([search_tool, execute_tool])

690

691 def openai(

692 self,

693 *,

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

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

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

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

698

699 Args:

700 mode: Tool mode.

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

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

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

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

705 config from the constructor.

706

707 Returns:

708 List of tool definitions in OpenAI function format.

709

710 Examples::

711

712 # All tools

713 toolset = StackOneToolSet()

714 tools = toolset.openai()

715

716 # Meta tools for agent-driven discovery

717 toolset = StackOneToolSet()

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

719 """

720 effective_account_ids = account_ids or (

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

722 )

723

724 if mode == "search_and_execute":

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

726

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

728

729 def langchain(

730 self,

731 *,

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

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

734 ) -> Sequence[Any]:

735 """Get tools in LangChain format.

736

737 Args:

738 mode: Tool mode.

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

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

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

742 The framework handles tool execution automatically.

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

744 config from the constructor.

745

746 Returns:

747 List of LangChain tool objects.

748 """

749 effective_account_ids = account_ids or (

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

751 )

752

753 if mode == "search_and_execute":

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

755

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

757

758 def pydantic_ai(

759 self,

760 *,

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

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

763 ) -> list[PydanticAITool]:

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

765

766 Args:

767 mode: Tool mode.

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

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

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

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

772 config from the constructor.

773

774 Returns:

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

776

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

778

779 Examples::

780

781 # All tools

782 toolset = StackOneToolSet()

783 tools = toolset.pydantic_ai()

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

785

786 # Meta tools for agent-driven discovery

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

788 """

789 effective_account_ids = account_ids or (

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

791 )

792

793 if mode == "search_and_execute":

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

795

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

797

798 def execute(

799 self,

800 tool_name: str,

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

802 ) -> dict[str, Any]:

803 """Execute a tool by name.

804

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

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

807

808 Tools are cached after the first call.

809

810 Args:

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

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

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

814 as a JSON string or dict.

815

816 Returns:

817 Tool execution result as a dict.

818 """

819 if self._tools_cache is None:

820 self._tools_cache = self._build_tools()

821

822 tool = self._tools_cache.get_tool(tool_name)

823 if tool is None:

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

825 return tool.execute(arguments)

826

827 @property

828 def semantic_client(self) -> SemanticSearchClient:

829 """Lazy initialization of semantic search client.

830

831 Returns:

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

833 """

834 if self._semantic_client is None:

835 self._semantic_client = SemanticSearchClient(

836 api_key=self.api_key,

837 base_url=self.base_url,

838 )

839 return self._semantic_client

840

841 def _local_search(

842 self,

843 query: str,

844 all_tools: Tools,

845 *,

846 connector: str | None = None,

847 top_k: int | None = None,

848 min_similarity: float | None = None,

849 ) -> Tools:

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

851 from stackone_ai.local_search import ToolIndex

852

853 available_connectors = all_tools.get_connectors()

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

855 return Tools([])

856

857 cache_key = id(all_tools)

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

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

860 index = self._tool_index_cache[1]

861 results = index.search(

862 query,

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

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

865 )

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

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

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

869 matched_tools = [

870 tool_map[name]

871 for name in matched_names

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

873 ]

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

875

876 def search_tools(

877 self,

878 query: str,

879 *,

880 connector: str | None = None,

881 top_k: int | None = None,

882 min_similarity: float | None = None,

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

884 search: SearchMode | None = None,

885 ) -> Tools:

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

887

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

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

890

891 Args:

892 query: Natural language description of needed functionality

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

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

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

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

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

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

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

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

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

902 ``SemanticSearchError`` on failure.

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

904 semantic search endpoint).

905

906 Returns:

907 Tools collection with matched tools from linked accounts

908

909 Raises:

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

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

912

913 Examples:

914 # Semantic search (default with local fallback)

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

916

917 # Explicit semantic search

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

919

920 # Local BM25+TF-IDF search

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

922

923 # Filter by connector

924 tools = toolset.search_tools(

925 "create time off request",

926 connector="bamboohr",

927 search="semantic",

928 )

929 """

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

931 raise ToolsetConfigError(

932 "Search is disabled. Initialize StackOneToolSet with a search config to enable."

933 )

934

935 # Merge constructor defaults with per-call overrides

936 effective_search: SearchMode = (

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

938 )

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

940 effective_min_sim = (

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

942 )

943

944 all_tools = self.fetch_tools(account_ids=account_ids)

945 available_connectors = all_tools.get_connectors()

946

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

948 return Tools([])

949

950 # Local-only search — skip semantic API entirely

951 if effective_search == "local":

952 return self._local_search(

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

954 )

955

956 try:

957 # Determine which connectors to search

958 if connector:

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

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

961 return Tools([])

962 else:

963 connectors_to_search = available_connectors

964

965 # Search each connector in parallel

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

967 resp = self.semantic_client.search(

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

969 )

970 return list(resp.results)

971

972 all_results: list[SemanticSearchResult] = []

973 last_error: SemanticSearchError | None = None

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

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

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

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

978 try:

979 all_results.extend(future.result())

980 except SemanticSearchError as e:

981 last_error = e

982

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

984 if not all_results and last_error is not None:

985 raise last_error

986

987 # Sort by score, apply top_k

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

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

990 all_results = all_results[:effective_top_k]

991

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

993 return Tools([])

994

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

996 seen_names: set[str] = set()

997 action_names: list[str] = []

998 for result in all_results:

999 name = _normalize_action_name(result.id)

1000 if name in seen_names:

1001 continue

1002 seen_names.add(name)

1003 action_names.append(name)

1004

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

1006 return Tools([])

1007

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

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

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

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

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

1013

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

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

1016 logger.warning(

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

1018 "falling back to local search",

1019 len(all_results),

1020 )

1021 return self._local_search(

1022 query,

1023 all_tools,

1024 connector=connector,

1025 top_k=effective_top_k,

1026 min_similarity=effective_min_sim,

1027 )

1028

1029 return Tools(matched_tools)

1030

1031 except SemanticSearchError as e:

1032 if effective_search == "semantic":

1033 raise

1034

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

1036 return self._local_search(

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

1038 )

1039

1040 def search_action_names(

1041 self,

1042 query: str,

1043 *,

1044 connector: str | None = None,

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

1046 top_k: int | None = None,

1047 min_similarity: float | None = None,

1048 ) -> list[SemanticSearchResult]:

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

1050

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

1052 or when building custom filtering logic.

1053

1054 Args:

1055 query: Natural language description of needed functionality

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

1057 account_ids: Optional account IDs to scope results to connectors

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

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

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

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

1062 the server uses its default.

1063

1064 Returns:

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

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

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

1068 with their individual scores.

1069

1070 Examples:

1071 # Lightweight: inspect results before fetching

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

1073 for r in results:

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

1075

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

1077 results = toolset.search_action_names(

1078 "create employee",

1079 account_ids=["acc-123"],

1080 top_k=5

1081 )

1082 """

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

1084 raise ToolsetConfigError(

1085 "Search is disabled. Initialize StackOneToolSet with search config to enable."

1086 )

1087

1088 # Merge constructor defaults with per-call overrides

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

1090 effective_min_sim = (

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

1092 )

1093

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

1095 available_connectors: set[str] | None = None

1096 effective_account_ids = account_ids or self._account_ids

1097 if effective_account_ids:

1098 all_tools = self.fetch_tools(account_ids=effective_account_ids)

1099 available_connectors = all_tools.get_connectors()

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

1101 return []

1102

1103 try:

1104 if available_connectors:

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

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

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

1108 else:

1109 connectors_to_search = available_connectors

1110

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

1112 try:

1113 resp = self.semantic_client.search(

1114 query=query,

1115 connector=c,

1116 top_k=effective_top_k,

1117 min_similarity=effective_min_sim,

1118 )

1119 return list(resp.results)

1120 except SemanticSearchError:

1121 return []

1122

1123 all_results: list[SemanticSearchResult] = []

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

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

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

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

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

1129 all_results.extend(future.result())

1130 else:

1131 # No account filtering — single global search

1132 response = self.semantic_client.search(

1133 query=query,

1134 connector=connector,

1135 top_k=effective_top_k,

1136 min_similarity=effective_min_sim,

1137 )

1138 all_results = list(response.results)

1139

1140 except SemanticSearchError as e:

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

1142 return []

1143

1144 # Sort by score

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

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

1147

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

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

1150

1151 Args:

1152 tool_name: Name of the tool to check

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

1154

1155 Returns:

1156 True if the tool matches any provider, False otherwise

1157 """

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

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

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

1161 return provider in provider_set

1162

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

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

1165

1166 Args:

1167 tool_name: Name of the tool to check

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

1169

1170 Returns:

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

1172 """

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

1174

1175 def fetch_tools(

1176 self,

1177 *,

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

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

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

1181 ) -> Tools:

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

1183

1184 Args:

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

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

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

1188 Case-insensitive matching.

1189 actions: Optional list of action patterns with glob support

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

1191

1192 Returns:

1193 Collection of tools matching the filter criteria

1194

1195 Raises:

1196 ToolsetLoadError: If there is an error loading the tools

1197

1198 Examples:

1199 # Filter by account IDs

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

1201

1202 # Filter by providers

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

1204

1205 # Filter by actions with glob patterns

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

1207

1208 # Combine filters

1209 tools = toolset.fetch_tools(

1210 account_ids=['123'],

1211 providers=['hibob'],

1212 actions=['*_list_*']

1213 )

1214

1215 # Use set_accounts() for account filtering

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

1217 tools = toolset.fetch_tools()

1218 """

1219 try:

1220 effective_account_ids = account_ids or self._account_ids

1221 if not effective_account_ids and self.account_id:

1222 effective_account_ids = [self.account_id]

1223

1224 if effective_account_ids:

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

1226 else:

1227 account_scope = [None]

1228

1229 cache_key = (

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

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

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

1233 )

1234 cached = self._catalog_cache.get(cache_key)

1235 if cached is not None:

1236 return cached

1237

1238 endpoint = f"{self.base_url.rstrip('/')}/mcp"

1239

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

1241 headers = self._build_mcp_headers(account)

1242 catalog = _fetch_mcp_tools(endpoint, headers)

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

1244

1245 all_tools: list[StackOneTool] = []

1246 if len(account_scope) == 1:

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

1248 else:

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

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

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

1252 for future in futures:

1253 all_tools.extend(future.result())

1254

1255 if providers:

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

1257

1258 if actions:

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

1260

1261 result = Tools(all_tools)

1262 self._catalog_cache[cache_key] = result

1263 return result

1264

1265 except ToolsetError:

1266 raise

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

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

1269

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

1271 headers = {

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

1273 "User-Agent": _USER_AGENT,

1274 }

1275 if account_id:

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

1277 return headers

1278

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

1280 schema = tool_def.input_schema or {}

1281 parameters = ToolParameters(

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

1283 properties=self._normalize_schema_properties(schema),

1284 )

1285 return _StackOneRpcTool(

1286 name=tool_def.name,

1287 description=tool_def.description or "",

1288 parameters=parameters,

1289 api_key=self.api_key,

1290 base_url=self.base_url,

1291 account_id=account_id,

1292 timeout=self._timeout,

1293 )

1294

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

1296 properties = schema.get("properties")

1297 if not isinstance(properties, dict):

1298 return {}

1299

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

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

1302

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

1304 if isinstance(details, dict):

1305 prop = dict(details)

1306 else:

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

1308

1309 if name in required_fields:

1310 prop.setdefault("nullable", False)

1311 else:

1312 prop.setdefault("nullable", True)

1313

1314 normalized[str(name)] = prop

1315

1316 return normalized

Coverage for stackone_ai / toolset.py: 91%

467 statements