Coverage for stackone_ai / types.py: 100%

1"""Shared types, configuration and errors for the StackOne AI SDK.""" 

2 

3from __future__ import annotations 

4 

5from collections.abc import Sequence 

6from typing import Annotated, Any, Literal, TypeAlias, TypedDict 

7 

8from pydantic import ( 

9 BaseModel, 

10 ConfigDict, 

11 Field, 

12 PlainSerializer, 

13 SerializerFunctionWrapHandler, 

14 SkipValidation, 

15 model_serializer, 

16) 

17 

18JsonDict: TypeAlias = dict[str, Any] 

19Headers: TypeAlias = dict[str, str] 

20 

21# StackOne API base URL 

22DEFAULT_BASE_URL: str = "https://api.stackone.com" 

23 

24 

25ToolMode = Literal["individual", "search_execute"] 

26"""How the MCP endpoint lists tools. 

27 

28``"individual"`` (the server default) lists one tool per action — hundreds per 

29account. ``"search_execute"`` lists two meta tools per connector instead, a 

30``*_search_actions`` that ranks actions for a natural-language query and an 

31``*_execute_action`` that runs one by id. The catalog stays small regardless of 

32how many accounts are linked, which is what keeps it inside a model's context. 

33""" 

34 

35SUBMIT_FEEDBACK_TOOL_NAME: str = "stackone_submit_feedback" 

36"""The one global tool the MCP endpoint serves in every mode when feedback is enabled. 

37 

38It is not a connector action, and it is global rather than account-scoped, so the toolset 

39keeps one copy of it however many accounts list it. 

40""" 

41 

42FeedbackRating = Literal["positive", "negative", "neutral"] 

43FeedbackSource = Literal["model", "user", "system"] 

44FeedbackCategory = Literal["search", "execute", "defender", "connection", "general"] 

45 

46 

47class StackOneError(Exception): 

48 """Base exception for StackOne errors""" 

49 

50 pass 

51 

52 

53class StackOneAPIError(StackOneError): 

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

55 

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

57 super().__init__(message) 

58 self.status_code = status_code 

59 self.response_body = response_body 

60 

61 

62class ToolArgumentsError(StackOneError, ValueError): 

63 """Raised when a tool's arguments are unusable, before any request is made: 

64 

65 - a JSON string argument that does not parse, or parses to something other than 

66 an object 

67 - an argument that cannot be encoded as JSON (a set, bytes, NaN/Infinity, a lone 

68 surrogate, or a value that refers to itself, directly or through a cycle) 

69 - ``call()`` given both positional and keyword arguments, or more than one 

70 positional argument 

71 

72 A subclass of both StackOneError, so ``except StackOneError`` catches it, and 

73 ValueError, which is what these errors were before, so existing 

74 ``except ValueError`` clauses still do. 

75 """ 

76 

77 pass 

78 

79 

80class ToolsetError(StackOneError): 

81 """Base exception for toolset errors. 

82 

83 A subclass of StackOneError, so ``except StackOneError`` catches everything this 

84 SDK raises. The two used to be unrelated siblings, which meant the obvious 

85 catch-all silently missed ToolsetConfigError and ToolsetLoadError — the errors a 

86 user is most likely to hit on their very first call. Existing 

87 ``except ToolsetError`` clauses are unaffected. 

88 """ 

89 

90 pass 

91 

92 

93class ToolsetConfigError(ToolsetError): 

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

95 

96 pass 

97 

98 

99class ToolsetLoadError(ToolsetError): 

100 """Raised when there is an error loading tools. 

101 

102 ``failures`` holds each account's error when every account failed to list tools for 

103 differing reasons, in account order; otherwise it is empty. 

104 """ 

105 

106 def __init__(self, message: str, *, failures: Sequence[Exception] = ()) -> None: 

107 super().__init__(message) 

108 self.failures: list[Exception] = list(failures) 

109 

110 

111class ExecuteToolsConfig(TypedDict, total=False): 

112 """Execution configuration for the StackOneToolSet constructor. 

113 

114 Controls default account scoping and timeout for tool execution. 

115 """ 

116 

117 account_ids: list[str] 

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

119 

120 timeout: float 

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

122 constructor param which takes precedence.""" 

123 

124 

125class ExecuteConfig(BaseModel): 

126 """How a tool is called: its name, any extra request headers, and the timeout. 

127 

128 Every tool executes over MCP ``tools/call``, so there is no HTTP method, URL, body 

129 type or parameter location to configure. Extra fields are refused rather than ignored, 

130 so code still passing them finds out here instead of assuming they took effect. 

131 """ 

132 

133 model_config = ConfigDict(extra="forbid") 

134 

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

136 headers: Headers = Field( 

137 default_factory=dict, 

138 description="Extra HTTP headers for the MCP request. Authorization, x-account-id and " 

139 "User-Agent are always the SDK's own and cannot be replaced here.", 

140 ) 

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

142 

143 

144class ToolParameters(BaseModel): 

145 """Schema definition for tool parameters. 

146 

147 ``properties`` is a faithful mirror of the ``inputSchema`` the MCP server served, 

148 with the SDK's internal ``nullable`` marker added per property. Every other root 

149 keyword the server sent, ``required`` included, is kept verbatim as an extra field. 

150 Consumers that need the raw served schema (for example the ADK plugin) read this 

151 directly, so nothing here may be invented or dropped. 

152 """ 

153 

154 model_config = ConfigDict(extra="allow") 

155 

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

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

158 # Typed for callers building a tool by hand, but neither validated nor converted: a 

159 # served `required` is kept exactly as sent, even when malformed, and the adapters 

160 # drop what is not a list of names. 

161 required: Annotated[SkipValidation[list[str] | None], PlainSerializer(lambda value: value)] = Field( 

162 default=None, description="The names of the required properties, as served" 

163 ) 

164 

165 @model_serializer(mode="wrap") 

166 def _dump_as_served(self, handler: SerializerFunctionWrapHandler) -> JsonDict: 

167 # Omitted when it was never given, so a schema served without `required` dumps 

168 # without one, as it did when `required` was only an extra field. 

169 dumped = handler(self) 

170 if "required" not in self.model_fields_set: 

171 dumped.pop("required", None) 

172 return dumped