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(search={"method": "auto"})

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. Pass search={} (or search={'method': 'auto'}) to "

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

656 )

657

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

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

660 config["method"] = search

661

662 return SearchTool(self, config=config)

663

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

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

666 if self._search_config is None:

667 raise ToolsetConfigError(

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

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

670 )

671

672 if account_ids:

673 self._account_ids = account_ids

674

675 # Discover available connectors for dynamic descriptions

676 connectors_str = ""

677 try:

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

679 connectors = sorted(all_tools.get_connectors())

680 if connectors:

681 connectors_str = ", ".join(connectors)

682 except Exception:

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

684

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

686 search_tool._toolset = self

687

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

689 execute_tool._toolset = self

690

691 return Tools([search_tool, execute_tool])

692

693 def openai(

694 self,

695 *,

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

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

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

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

700

701 Args:

702 mode: Tool mode.

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

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

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

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

707 config from the constructor.

708

709 Returns:

710 List of tool definitions in OpenAI function format.

711

712 Examples::

713

714 # All tools

715 toolset = StackOneToolSet()

716 tools = toolset.openai()

717

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

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

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

721 """

722 effective_account_ids = account_ids or (

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

724 )

725

726 if mode == "search_and_execute":

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

728

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

730

731 def langchain(

732 self,

733 *,

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

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

736 ) -> Sequence[Any]:

737 """Get tools in LangChain format.

738

739 Args:

740 mode: Tool mode.

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

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

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

744 The framework handles tool execution automatically.

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

746 config from the constructor.

747

748 Returns:

749 List of LangChain tool objects.

750 """

751 effective_account_ids = account_ids or (

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

753 )

754

755 if mode == "search_and_execute":

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

757

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

759

760 def pydantic_ai(

761 self,

762 *,

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

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

765 ) -> list[PydanticAITool]:

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

767

768 Args:

769 mode: Tool mode.

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

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

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

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

774 config from the constructor.

775

776 Returns:

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

778

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

780

781 Examples::

782

783 # All tools

784 toolset = StackOneToolSet()

785 tools = toolset.pydantic_ai()

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

787

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

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

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

791 """

792 effective_account_ids = account_ids or (

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

794 )

795

796 if mode == "search_and_execute":

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

798

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

800

801 def execute(

802 self,

803 tool_name: str,

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

805 ) -> dict[str, Any]:

806 """Execute a tool by name.

807

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

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

810

811 Tools are cached after the first call.

812

813 Args:

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

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

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

817 as a JSON string or dict.

818

819 Returns:

820 Tool execution result as a dict.

821 """

822 if self._tools_cache is None:

823 self._tools_cache = self._build_tools()

824

825 tool = self._tools_cache.get_tool(tool_name)

826 if tool is None:

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

828 return tool.execute(arguments)

829

830 @property

831 def semantic_client(self) -> SemanticSearchClient:

832 """Lazy initialization of semantic search client.

833

834 Returns:

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

836 """

837 if self._semantic_client is None:

838 self._semantic_client = SemanticSearchClient(

839 api_key=self.api_key,

840 base_url=self.base_url,

841 )

842 return self._semantic_client

843

844 def _local_search(

845 self,

846 query: str,

847 all_tools: Tools,

848 *,

849 connector: str | None = None,

850 top_k: int | None = None,

851 min_similarity: float | None = None,

852 ) -> Tools:

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

854 from stackone_ai.local_search import ToolIndex

855

856 available_connectors = all_tools.get_connectors()

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

858 return Tools([])

859

860 cache_key = id(all_tools)

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

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

863 index = self._tool_index_cache[1]

864 results = index.search(

865 query,

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

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

868 )

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

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

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

872 matched_tools = [

873 tool_map[name]

874 for name in matched_names

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

876 ]

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

878

879 def search_tools(

880 self,

881 query: str,

882 *,

883 connector: str | None = None,

884 top_k: int | None = None,

885 min_similarity: float | None = None,

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

887 search: SearchMode | None = None,

888 ) -> Tools:

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

890

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

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

893

894 Args:

895 query: Natural language description of needed functionality

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

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

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

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

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

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

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

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

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

905 ``SemanticSearchError`` on failure.

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

907 semantic search endpoint).

908

909 Returns:

910 Tools collection with matched tools from linked accounts

911

912 Raises:

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

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

915

916 Examples:

917 # Semantic search (default with local fallback)

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

919

920 # Explicit semantic search

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

922

923 # Local BM25+TF-IDF search

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

925

926 # Filter by connector

927 tools = toolset.search_tools(

928 "create time off request",

929 connector="bamboohr",

930 search="semantic",

931 )

932 """

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

934 raise ToolsetConfigError(

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

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

937 )

938

939 # Merge constructor defaults with per-call overrides

940 effective_search: SearchMode = (

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

942 )

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

944 effective_min_sim = (

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

946 )

947

948 all_tools = self.fetch_tools(account_ids=account_ids)

949 available_connectors = all_tools.get_connectors()

950

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

952 return Tools([])

953

954 # Local-only search — skip semantic API entirely

955 if effective_search == "local":

956 return self._local_search(

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

958 )

959

960 try:

961 # Determine which connectors to search

962 if connector:

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

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

965 return Tools([])

966 else:

967 connectors_to_search = available_connectors

968

969 # Search each connector in parallel

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

971 resp = self.semantic_client.search(

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

973 )

974 return list(resp.results)

975

976 all_results: list[SemanticSearchResult] = []

977 last_error: SemanticSearchError | None = None

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

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

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

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

982 try:

983 all_results.extend(future.result())

984 except SemanticSearchError as e:

985 last_error = e

986

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

988 if not all_results and last_error is not None:

989 raise last_error

990

991 # Sort by score, apply top_k

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

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

994 all_results = all_results[:effective_top_k]

995

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

997 return Tools([])

998

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

1000 seen_names: set[str] = set()

1001 action_names: list[str] = []

1002 for result in all_results:

1003 name = _normalize_action_name(result.id)

1004 if name in seen_names:

1005 continue

1006 seen_names.add(name)

1007 action_names.append(name)

1008

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

1010 return Tools([])

1011

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

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

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

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

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

1017

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

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

1020 logger.warning(

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

1022 "falling back to local search",

1023 len(all_results),

1024 )

1025 return self._local_search(

1026 query,

1027 all_tools,

1028 connector=connector,

1029 top_k=effective_top_k,

1030 min_similarity=effective_min_sim,

1031 )

1032

1033 return Tools(matched_tools)

1034

1035 except SemanticSearchError as e:

1036 if effective_search == "semantic":

1037 raise

1038

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

1040 return self._local_search(

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

1042 )

1043

1044 def search_action_names(

1045 self,

1046 query: str,

1047 *,

1048 connector: str | None = None,

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

1050 top_k: int | None = None,

1051 min_similarity: float | None = None,

1052 ) -> list[SemanticSearchResult]:

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

1054

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

1056 or when building custom filtering logic.

1057

1058 Args:

1059 query: Natural language description of needed functionality

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

1061 account_ids: Optional account IDs to scope results to connectors

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

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

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

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

1066 the server uses its default.

1067

1068 Returns:

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

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

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

1072 with their individual scores.

1073

1074 Examples:

1075 # Lightweight: inspect results before fetching

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

1077 for r in results:

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

1079

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

1081 results = toolset.search_action_names(

1082 "create employee",

1083 account_ids=["acc-123"],

1084 top_k=5

1085 )

1086 """

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

1088 raise ToolsetConfigError(

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

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

1091 )

1092

1093 # Merge constructor defaults with per-call overrides

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

1095 effective_min_sim = (

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

1097 )

1098

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

1100 available_connectors: set[str] | None = None

1101 effective_account_ids = account_ids or self._account_ids

1102 if effective_account_ids:

1103 all_tools = self.fetch_tools(account_ids=effective_account_ids)

1104 available_connectors = all_tools.get_connectors()

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

1106 return []

1107

1108 try:

1109 if available_connectors:

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

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

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

1113 else:

1114 connectors_to_search = available_connectors

1115

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

1117 try:

1118 resp = self.semantic_client.search(

1119 query=query,

1120 connector=c,

1121 top_k=effective_top_k,

1122 min_similarity=effective_min_sim,

1123 )

1124 return list(resp.results)

1125 except SemanticSearchError:

1126 return []

1127

1128 all_results: list[SemanticSearchResult] = []

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

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

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

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

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

1134 all_results.extend(future.result())

1135 else:

1136 # No account filtering — single global search

1137 response = self.semantic_client.search(

1138 query=query,

1139 connector=connector,

1140 top_k=effective_top_k,

1141 min_similarity=effective_min_sim,

1142 )

1143 all_results = list(response.results)

1144

1145 except SemanticSearchError as e:

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

1147 return []

1148

1149 # Sort by score

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

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

1152

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

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

1155

1156 Args:

1157 tool_name: Name of the tool to check

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

1159

1160 Returns:

1161 True if the tool matches any provider, False otherwise

1162 """

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

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

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

1166 return provider in provider_set

1167

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

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

1170

1171 Args:

1172 tool_name: Name of the tool to check

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

1174

1175 Returns:

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

1177 """

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

1179

1180 def fetch_tools(

1181 self,

1182 *,

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

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

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

1186 ) -> Tools:

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

1188

1189 Args:

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

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

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

1193 Case-insensitive matching.

1194 actions: Optional list of action patterns with glob support

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

1196

1197 Returns:

1198 Collection of tools matching the filter criteria

1199

1200 Raises:

1201 ToolsetLoadError: If there is an error loading the tools

1202

1203 Examples:

1204 # Filter by account IDs

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

1206

1207 # Filter by providers

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

1209

1210 # Filter by actions with glob patterns

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

1212

1213 # Combine filters

1214 tools = toolset.fetch_tools(

1215 account_ids=['123'],

1216 providers=['hibob'],

1217 actions=['*_list_*']

1218 )

1219

1220 # Use set_accounts() for account filtering

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

1222 tools = toolset.fetch_tools()

1223 """

1224 try:

1225 effective_account_ids = account_ids or self._account_ids

1226 if not effective_account_ids and self.account_id:

1227 effective_account_ids = [self.account_id]

1228

1229 if effective_account_ids:

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

1231 else:

1232 account_scope = [None]

1233

1234 cache_key = (

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

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

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

1238 )

1239 cached = self._catalog_cache.get(cache_key)

1240 if cached is not None:

1241 return cached

1242

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

1244

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

1246 headers = self._build_mcp_headers(account)

1247 catalog = _fetch_mcp_tools(endpoint, headers)

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

1249

1250 all_tools: list[StackOneTool] = []

1251 if len(account_scope) == 1:

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

1253 else:

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

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

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

1257 for future in futures:

1258 all_tools.extend(future.result())

1259

1260 if providers:

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

1262

1263 if actions:

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

1265

1266 result = Tools(all_tools)

1267 self._catalog_cache[cache_key] = result

1268 return result

1269

1270 except ToolsetError:

1271 raise

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

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

1274

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

1276 headers = {

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

1278 "User-Agent": _USER_AGENT,

1279 }

1280 if account_id:

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

1282 return headers

1283

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

1285 schema = tool_def.input_schema or {}

1286 parameters = ToolParameters(

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

1288 properties=self._normalize_schema_properties(schema),

1289 )

1290 return _StackOneRpcTool(

1291 name=tool_def.name,

1292 description=tool_def.description or "",

1293 parameters=parameters,

1294 api_key=self.api_key,

1295 base_url=self.base_url,

1296 account_id=account_id,

1297 timeout=self._timeout,

1298 )

1299

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

1301 properties = schema.get("properties")

1302 if not isinstance(properties, dict):

1303 return {}

1304

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

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

1307

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

1309 if isinstance(details, dict):

1310 prop = dict(details)

1311 else:

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

1313

1314 if name in required_fields:

1315 prop.setdefault("nullable", False)

1316 else:

1317 prop.setdefault("nullable", True)

1318

1319 normalized[str(name)] = prop

1320

1321 return normalized

Coverage for stackone_ai / toolset.py: 91%

467 statements