Coverage for stackone_ai/models.py: 98%

1from __future__ import annotations

3import base64

4import json

5import logging

6import re

7from collections.abc import Sequence

8from datetime import datetime, timezone

9from enum import Enum

10from typing import TYPE_CHECKING, Annotated, Any, ClassVar, TypeAlias, cast

11from urllib.parse import quote, unquote

13import httpx

14from langchain_core.tools import BaseTool

15from pydantic import BaseModel, BeforeValidator, Field, PrivateAttr

17if TYPE_CHECKING:

18 from pydantic_ai.tools import Tool as PydanticAITool

20# Type aliases for common types

21JsonDict: TypeAlias = dict[str, Any]

22Headers: TypeAlias = dict[str, str]

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

28class StackOneError(Exception):

29 """Base exception for StackOne errors"""

31 pass

34class StackOneAPIError(StackOneError):

35 """Raised when the StackOne API returns an error"""

37 def __init__(self, message: str, status_code: int, response_body: Any) -> None:

38 super().__init__(message)

39 self.status_code = status_code

40 self.response_body = response_body

43class ParameterLocation(str, Enum):

44 """Valid locations for parameters in requests"""

46 HEADER = "header"

47 QUERY = "query"

48 PATH = "path"

49 BODY = "body"

50 FILE = "file" # For file uploads

53def validate_method(v: str) -> str:

54 """Validate HTTP method is uppercase and supported"""

55 method = v.upper()

56 if method not in {"GET", "POST", "PUT", "DELETE", "PATCH"}:

57 raise ValueError(f"Unsupported HTTP method: {method}")

58 return method

61def _is_json_content_type(content_type: str) -> bool:

62 """Whether a response body should be parsed as JSON based on its Content-Type.

64 Only genuine JSON media types are parsed (``application/json`` and structured

65 suffixes such as ``application/problem+json``). Anything else - including a

66 missing Content-Type - is treated as opaque content (a file download), so the

67 raw bytes are returned instead of being force-decoded as UTF-8/JSON. This mirrors

68 how the StackOne generated SDKs default unknown bodies to ``application/octet-stream``.

69 """

70 media_type = content_type.split(";", 1)[0].strip().lower()

71 return media_type == "application/json" or media_type.endswith("+json")

74def _filename_from_content_disposition(value: str | None) -> str | None:

75 """Extract the filename from a Content-Disposition header value, if present.

77 Handles both the plain ``filename="example.pdf"`` form and the RFC 5987 extended

78 ``filename*=UTF-8''example%20file.pdf`` form (which takes precedence when present).

79 The extended form is percent-decoded using its declared charset (RFC 5987 permits

80 both ``UTF-8`` and ``ISO-8859-1``); an unknown or empty charset falls back to UTF-8.

81 """

82 if not value:

83 return None

84 extended = re.search(r"filename\*\s*=\s*([^']*)'[^']*'([^;]+)", value, re.IGNORECASE)

85 if extended:

86 charset = extended.group(1).strip() or "utf-8"

87 encoded = extended.group(2).strip().strip('"')

88 try:

89 return unquote(encoded, encoding=charset, errors="replace") or None

90 except LookupError:

91 # Unrecognised charset label - decode as UTF-8 rather than failing.

92 return unquote(encoded, encoding="utf-8", errors="replace") or None

93 quoted = re.search(r'filename\s*=\s*"([^"]*)"', value, re.IGNORECASE)

94 if quoted:

95 return quoted.group(1).strip() or None

96 bare = re.search(r"filename\s*=\s*([^;]+)", value, re.IGNORECASE)

97 if bare:

98 return bare.group(1).strip().strip('"') or None

99 return None

100

101

102class ExecuteConfig(BaseModel):

103 """Configuration for executing a tool against an API endpoint"""

104

105 headers: Headers = Field(default_factory=dict, description="HTTP headers to include in the request")

106 method: Annotated[str, BeforeValidator(validate_method)] = Field(description="HTTP method to use")

107 url: str = Field(description="API endpoint URL")

108 name: str = Field(description="Tool name")

109 body_type: str | None = Field(default=None, description="Content type for request body")

110 parameter_locations: dict[str, ParameterLocation] = Field(

111 default_factory=dict, description="Maps parameter names to their location in the request"

112 )

113 timeout: float = Field(default=60.0, description="Request timeout in seconds")

114

115

116class ToolParameters(BaseModel):

117 """Schema definition for tool parameters"""

118

119 type: str = Field(description="JSON Schema type")

120 properties: JsonDict = Field(description="JSON Schema properties")

121

122

123class ToolDefinition(BaseModel):

124 """Complete definition of a tool including its schema and execution config"""

125

126 description: str = Field(description="Tool description")

127 parameters: ToolParameters = Field(description="Tool parameter schema")

128 execute: ExecuteConfig = Field(description="Tool execution configuration")

129

130

131class StackOneTool(BaseModel):

132 """Base class for all StackOne tools. Provides functionality for executing API calls

133 and converting to various formats (OpenAI, LangChain)."""

134

135 name: str = Field(description="Tool name")

136 description: str = Field(description="Tool description")

137 parameters: ToolParameters = Field(description="Tool parameters")

138 _execute_config: ExecuteConfig = PrivateAttr()

139 _api_key: str = PrivateAttr()

140 _account_id: str | None = PrivateAttr(default=None)

141 _FEEDBACK_OPTION_KEYS: ClassVar[set[str]] = {

142 "feedback_session_id",

143 "feedback_user_id",

144 "feedback_metadata",

145 }

146

147 @property

148 def connector(self) -> str:

149 """Extract connector from tool name.

150

151 Tool names follow the format: {connector}_{action}_{entity}

152 e.g., 'bamboohr_create_employee' -> 'bamboohr'

153

154 Returns:

155 Connector name in lowercase

156 """

157 return self.name.split("_")[0].lower()

158

159 def __init__(

160 self,

161 description: str,

162 parameters: ToolParameters,

163 _execute_config: ExecuteConfig,

164 _api_key: str,

165 _account_id: str | None = None,

166 ) -> None:

167 super().__init__(

168 name=_execute_config.name,

169 description=description,

170 parameters=parameters,

171 )

172 self._execute_config = _execute_config

173 self._api_key = _api_key

174 self._account_id = _account_id

175

176 @classmethod

177 def _split_feedback_options(cls, params: JsonDict, options: JsonDict | None) -> tuple[JsonDict, JsonDict]:

178 merged_params = dict(params)

179 feedback_options = dict(options or {})

180 for key in cls._FEEDBACK_OPTION_KEYS:

181 if key in merged_params and key not in feedback_options:

182 feedback_options[key] = merged_params.pop(key)

183 return merged_params, feedback_options

184

185 def _prepare_headers(self) -> Headers:

186 """Prepare headers for the API request

187

188 Returns:

189 Headers to use in the request

190 """

191 auth_string = base64.b64encode(f"{self._api_key}:".encode()).decode()

192 headers: Headers = {

193 "Authorization": f"Basic {auth_string}",

194 "User-Agent": "stackone-python/1.0.0",

195 }

196

197 if self._account_id:

198 headers["x-account-id"] = self._account_id

199

200 # Add predefined headers

201 headers.update(self._execute_config.headers)

202 return headers

203

204 def _prepare_request_params(self, kwargs: JsonDict) -> tuple[str, JsonDict, JsonDict]:

205 """Prepare URL and parameters for the API request

206

207 Args:

208 kwargs: Arguments to process

209

210 Returns:

211 Tuple of (url, body_params, query_params)

212 """

213 url = self._execute_config.url

214 body_params: JsonDict = {}

215 query_params: JsonDict = {}

216

217 for key, value in kwargs.items():

218 param_location = self._execute_config.parameter_locations.get(key)

219

220 if param_location == ParameterLocation.PATH:

221 # Safely encode path parameters to prevent SSRF attacks

222 encoded_value = quote(str(value), safe="")

223 url = url.replace(f"{{{key}}}", encoded_value)

224 elif param_location == ParameterLocation.QUERY:

225 query_params[key] = value

226 elif param_location in (ParameterLocation.BODY, ParameterLocation.FILE):

227 body_params[key] = value

228 else:

229 # Default behavior

230 if f"{{{key}}}" in url:

231 # Safely encode path parameters to prevent SSRF attacks

232 encoded_value = quote(str(value), safe="")

233 url = url.replace(f"{{{key}}}", encoded_value)

234 elif self._execute_config.method in {"GET", "DELETE"}:

235 query_params[key] = value

236 else:

237 body_params[key] = value

238

239 return url, body_params, query_params

240

241 def execute(

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

243 ) -> JsonDict:

244 """Execute the tool with the given parameters

245

246 Args:

247 arguments: Tool arguments as string or dict

248 options: Execution options (e.g. feedback metadata)

249

250 Returns:

251 For JSON responses, the parsed API response as a dict.

252

253 For file downloads (any non-JSON Content-Type, e.g. a

254 ``documents_download_file`` action), a dict describing the file:

255 ``{"content": <bytes>, "content_type": str, "status_code": int,

256 "headers": dict, "file_name": str | None}``. Note ``content`` holds

257 the raw bytes and is therefore not JSON-serializable - callers that

258 re-serialize tool results (e.g. for an LLM) should handle this key.

259

260 Raises:

261 StackOneAPIError: If the API request fails

262 ValueError: If the arguments are invalid

263 """

264 datetime.now(timezone.utc)

265 feedback_options: JsonDict = {}

266 result_payload: JsonDict | None = None

267 response_status: int | None = None

268 error_message: str | None = None

269 status = "success"

270 url_used = self._execute_config.url

271

272 try:

273 if isinstance(arguments, str):

274 parsed_arguments = json.loads(arguments)

275 else:

276 parsed_arguments = arguments or {}

277

278 if not isinstance(parsed_arguments, dict):

279 status = "error"

280 error_message = "Tool arguments must be a JSON object"

281 raise ValueError(error_message)

282

283 kwargs = parsed_arguments

284 dict(kwargs)

285

286 headers = self._prepare_headers()

287 url_used, body_params, query_params = self._prepare_request_params(kwargs)

288

289 request_kwargs: dict[str, Any] = {

290 "method": self._execute_config.method,

291 "url": url_used,

292 "headers": headers,

293 }

294

295 if body_params:

296 body_type = self._execute_config.body_type or "json"

297 if body_type == "json":

298 request_kwargs["json"] = body_params

299 elif body_type == "form": 299 ↛ 302line 299 didn't jump to line 302 because the condition on line 299 was always true

300 request_kwargs["data"] = body_params

301

302 if query_params:

303 request_kwargs["params"] = query_params

304

305 response = httpx.request(**request_kwargs, timeout=self._execute_config.timeout)

306 response_status = response.status_code

307 response.raise_for_status()

308

309 content_type = response.headers.get("content-type", "")

310 if _is_json_content_type(content_type):

311 result = response.json()

312 result_payload = cast(JsonDict, result) if isinstance(result, dict) else {"result": result}

313 return result_payload

314

315 # Non-JSON bodies are file downloads (e.g. documents_download_file), which the

316 # API serves as raw binary with the file's own MIME type and a Content-Disposition

317 # header. Return the bytes plus metadata rather than forcing a JSON/UTF-8 decode.

318 # The shape mirrors the StackOne generated SDKs' download response.

319 return {

320 "content": response.content,

321 "content_type": content_type or "application/octet-stream",

322 "status_code": response.status_code,

323 "headers": dict(response.headers),

324 "file_name": _filename_from_content_disposition(response.headers.get("content-disposition")),

325 }

326

327 except json.JSONDecodeError as exc:

328 status = "error"

329 error_message = f"Invalid JSON in arguments: {exc}"

330 raise ValueError(error_message) from exc

331 except httpx.HTTPStatusError as exc:

332 status = "error"

333 response_body = None

334 if exc.response.text: 334 ↛ 339line 334 didn't jump to line 339 because the condition on line 334 was always true

335 try:

336 response_body = exc.response.json()

337 except json.JSONDecodeError:

338 response_body = exc.response.text

339 raise StackOneAPIError(

340 str(exc),

341 exc.response.status_code,

342 response_body,

343 ) from exc

344 except httpx.RequestError as exc:

345 status = "error"

346 raise StackOneError(f"Request failed: {exc}") from exc

347 finally:

348 datetime.now(timezone.utc)

349 metadata: JsonDict = {

350 "http_method": self._execute_config.method,

351 "url": url_used,

352 "status_code": response_status,

353 "status": status,

354 }

355

356 feedback_metadata = feedback_options.get("feedback_metadata")

357 if isinstance(feedback_metadata, dict): 357 ↛ 358line 357 didn't jump to line 358 because the condition on line 357 was never true

358 metadata["feedback_metadata"] = feedback_metadata

359

360 if feedback_options:

361 metadata["feedback_options"] = {

362 key: value

363 for key, value in feedback_options.items()

364 if key in {"feedback_session_id", "feedback_user_id"} and value is not None

365 }

366

367 # Implicit feedback removed - just API calls

368

369 def call(self, *args: Any, options: JsonDict | None = None, **kwargs: Any) -> JsonDict:

370 """Call the tool with the given arguments

371

372 This method provides a more intuitive way to execute tools directly.

373

374 Args:

375 *args: If a single argument is provided, it's treated as the full arguments dict/string

376 **kwargs: Keyword arguments to pass to the tool

377 options: Optional execution options

378

379 Returns:

380 API response as dict

381

382 Raises:

383 StackOneAPIError: If the API request fails

384 ValueError: If the arguments are invalid

385

386 Examples:

387 >>> tool.call({"name": "John", "email": "john@example.com"})

388 >>> tool.call(name="John", email="john@example.com")

389 """

390 if args and kwargs:

391 raise ValueError("Cannot provide both positional and keyword arguments")

392

393 if args:

394 if len(args) > 1:

395 raise ValueError("Only one positional argument is allowed")

396 return self.execute(args[0])

397

398 return self.execute(kwargs if kwargs else None)

399

400 def __call__(self, *args: Any, options: JsonDict | None = None, **kwargs: Any) -> JsonDict:

401 """Make the tool directly callable.

402

403 Alias for :meth:`call` so that ``tool(query="…")`` works.

404 """

405 return self.call(*args, options=options, **kwargs)

406

407 def to_openai_function(self) -> JsonDict:

408 """Convert this tool to OpenAI's function format

409

410 Returns:

411 Tool definition in OpenAI function format

412 """

413 # Clean properties and handle special types

414 properties = {}

415 required = []

416

417 for name, prop in self.parameters.properties.items():

418 if isinstance(prop, dict):

419 # Only keep standard JSON Schema properties

420 cleaned_prop = {}

421

422 # Copy basic properties

423 if "type" in prop:

424 cleaned_prop["type"] = prop["type"]

425 if "description" in prop:

426 cleaned_prop["description"] = prop["description"]

427 if "enum" in prop:

428 cleaned_prop["enum"] = prop["enum"]

429

430 # Handle array types

431 if cleaned_prop.get("type") == "array" and "items" in prop:

432 if isinstance(prop["items"], dict): 432 ↛ 438line 432 didn't jump to line 438 because the condition on line 432 was always true

433 cleaned_prop["items"] = {

434 k: v for k, v in prop["items"].items() if k in ("type", "description", "enum")

435 }

436

437 # Handle object types

438 if cleaned_prop.get("type") == "object" and "properties" in prop:

439 cleaned_prop["properties"] = {

440 k: {sk: sv for sk, sv in v.items() if sk in ("type", "description", "enum")}

441 for k, v in prop["properties"].items()

442 }

443

444 # Handle required fields - if not explicitly nullable

445 if not prop.get("nullable", False):

446 required.append(name)

447

448 properties[name] = cleaned_prop

449 else:

450 properties[name] = {"type": "string"}

451 required.append(name)

452

453 # Create the OpenAI function schema

454 parameters = {

455 "type": "object",

456 "properties": properties,

457 }

458

459 # Only include required if there are required fields

460 if required:

461 parameters["required"] = required

462

463 return {

464 "type": "function",

465 "function": {

466 "name": self.name,

467 "description": self.description,

468 "parameters": parameters,

469 },

470 }

471

472 def to_langchain(self) -> BaseTool:

473 """Convert this tool to LangChain format

474

475 Returns:

476 Tool in LangChain format

477 """

478 # Create properly annotated schema for the tool

479 schema_props: dict[str, Any] = {}

480 annotations: dict[str, Any] = {}

481

482 for name, details in self.parameters.properties.items():

483 python_type: type = str # Default to str

484 is_nullable = False

485 if isinstance(details, dict):

486 type_str = details.get("type", "string")

487 is_nullable = details.get("nullable", False)

488 if type_str == "number":

489 python_type = float

490 elif type_str == "integer":

491 python_type = int

492 elif type_str == "boolean":

493 python_type = bool

494 elif type_str == "object":

495 python_type = dict

496 elif type_str == "array": 496 ↛ 497line 496 didn't jump to line 497 because the condition on line 496 was never true

497 python_type = list

498

499 if is_nullable:

500 field = Field(default=None, description=details.get("description", ""))

501 else:

502 field = Field(description=details.get("description", ""))

503 else:

504 field = Field(description="")

505

506 schema_props[name] = field

507 if is_nullable:

508 annotations[name] = python_type | None

509 else:

510 annotations[name] = python_type

511

512 # Create the schema class with proper annotations

513 schema_class = type(

514 f"{self.name.title()}Args",

515 (BaseModel,),

516 {

517 "__annotations__": annotations,

518 "__module__": __name__,

519 **schema_props,

520 },

521 )

522

523 parent_tool = self

524

525 class StackOneLangChainTool(BaseTool):

526 name: str = parent_tool.name

527 description: str = parent_tool.description

528 args_schema: type[BaseModel] = schema_class # ty: ignore[invalid-assignment]

529 func = staticmethod(parent_tool.execute) # Required by CrewAI

530

531 def _run(self, **kwargs: Any) -> Any:

532 return parent_tool.execute(kwargs)

533

534 return StackOneLangChainTool()

535

536 def to_pydantic_ai_tool(self) -> PydanticAITool:

537 """Convert this tool to a Pydantic AI ``Tool``.

538

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

540

541 Returns:

542 A ``pydantic_ai.tools.Tool`` ready to pass to ``Agent(tools=[...])``.

543 """

544 try:

545 from pydantic_ai.tools import Tool

546 except ImportError as e:

547 raise ImportError(

548 "Install `pydantic-ai-slim` (or `stackone-ai[pydantic-ai]`) "

549 "to use the Pydantic AI integration."

550 ) from e

551

552 openai_function = self.to_openai_function()

553 json_schema = openai_function["function"]["parameters"]

554 parent_tool = self

555

556 def implementation(**kwargs: Any) -> Any:

557 return parent_tool.execute(kwargs)

558

559 return Tool.from_schema(

560 function=implementation,

561 name=self.name,

562 description=self.description,

563 json_schema=json_schema,

564 )

565

566 def set_account_id(self, account_id: str | None) -> None:

567 """Set the account ID for this tool

568

569 Args:

570 account_id: The account ID to use, or None to clear it

571 """

572 self._account_id = account_id

573

574 def get_account_id(self) -> str | None:

575 """Get the current account ID for this tool

576

577 Returns:

578 Current account ID or None if not set

579 """

580 return self._account_id

581

582

583class Tools:

584 """Container for Tool instances with lookup capabilities"""

585

586 def __init__(

587 self,

588 tools: list[StackOneTool],

589 ) -> None:

590 """Initialize Tools container

591

592 Args:

593 tools: List of Tool instances to manage

594 """

595 self.tools = tools

596 self._tool_map = {tool.name: tool for tool in tools}

597

598 def __getitem__(self, index: int) -> StackOneTool:

599 return self.tools[index]

600

601 def __len__(self) -> int:

602 return len(self.tools)

603

604 def __iter__(self) -> Any:

605 """Make Tools iterable"""

606 return iter(self.tools)

607

608 def to_list(self) -> list[StackOneTool]:

609 """Convert to list of tools

610

611 Returns:

612 List of StackOneTool instances

613 """

614 return list(self.tools)

615

616 def get_tool(self, name: str) -> StackOneTool | None:

617 """Get a tool by its name

618

619 Args:

620 name: Name of the tool to retrieve

621

622 Returns:

623 The tool if found, None otherwise

624 """

625 return self._tool_map.get(name)

626

627 def set_account_id(self, account_id: str | None) -> None:

628 """Set the account ID for all tools in this collection

629

630 Args:

631 account_id: The account ID to use, or None to clear it

632 """

633 for tool in self.tools:

634 tool.set_account_id(account_id)

635

636 def get_account_id(self) -> str | None:

637 """Get the current account ID for this collection

638

639 Returns:

640 The first non-None account ID found, or None if none set

641 """

642 for tool in self.tools:

643 account_id = tool.get_account_id()

644 if isinstance(account_id, str):

645 return account_id

646 return None

647

648 def get_connectors(self) -> set[str]:

649 """Get unique connector names from all tools.

650

651 Returns:

652 Set of connector names (lowercase)

653

654 Example:

655 tools = toolset.fetch_tools()

656 connectors = tools.get_connectors()

657 # {'bamboohr', 'hibob', 'slack', ...}

658 """

659 return {tool.connector for tool in self.tools}

660

661 def to_openai(self) -> list[JsonDict]:

662 """Convert all tools to OpenAI function format

663

664 Returns:

665 List of tools in OpenAI function format

666 """

667 return [tool.to_openai_function() for tool in self.tools]

668

669 def to_langchain(self) -> Sequence[BaseTool]:

670 """Convert all tools to LangChain format

671

672 Returns:

673 Sequence of tools in LangChain format

674 """

675 return [tool.to_langchain() for tool in self.tools]

676

677 def to_pydantic_ai(self) -> list[PydanticAITool]:

678 """Convert all tools to Pydantic AI ``Tool`` instances.

679

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

681

682 Returns:

683 List of ``pydantic_ai.tools.Tool`` ready to pass to ``Agent(tools=[...])``.

684 """

685 return [tool.to_pydantic_ai_tool() for tool in self.tools]

Coverage for stackone_ai / models.py: 98%

301 statements