Coverage for src/integry/resources/functions/types.py: 52%
92 statements
« prev ^ index » next coverage.py v7.6.1, created at 2025-03-26 01:26 +0500
1from pydantic import BaseModel, Field
2from typing import (
3 Any,
4 Awaitable,
5 Callable,
6 Dict,
7 List,
8 Union,
9 Literal,
10 Optional,
11 Annotated,
12 TYPE_CHECKING,
13 cast,
14)
16from integry.utils.common import generate_docstring_from_schema_for_smolagent
17from integry.utils.pydantic import get_pydantic_model_from_json_schema
19if TYPE_CHECKING: 19 ↛ 20line 19 didn't jump to line 20 because the condition on line 19 was never true
20 from integry.resources.functions.api import Functions as FunctionsResource
23class StringSchema(BaseModel):
24 type: Literal["string"]
27class NumberSchema(BaseModel):
28 type: Literal["number"]
31class BooleanSchema(BaseModel):
32 type: Literal["boolean"]
35class NullSchema(BaseModel):
36 type: Literal["null"]
39class ObjectSchema(BaseModel):
40 type: Literal["object"]
41 properties: Dict[str, "JSONSchemaType"] = Field(default_factory=dict)
42 required: List[str] = []
43 additionalProperties: Union["JSONSchemaType", bool] = True
46class ArraySchema(BaseModel):
47 type: Literal["array"]
48 items: Union["JSONSchemaType", List["JSONSchemaType"], None] = None
51JSONSchemaType = Annotated[
52 StringSchema
53 | NumberSchema
54 | BooleanSchema
55 | NullSchema
56 | ObjectSchema
57 | ArraySchema,
58 Field(discriminator="type"),
59]
62class FunctionCallOutput(BaseModel):
63 network_code: int
64 output: Any
67class PaginatedFunctionCallOutput(FunctionCallOutput):
68 cursor: str = Field(alias="_cursor")
71class Function(BaseModel):
72 name: str
73 description: str
74 parameters: JSONSchemaType
75 arguments: dict[str, Any] = Field(default_factory=dict)
77 _json_schema: dict[str, Any]
78 _resource: "FunctionsResource"
80 def __init__(self, **data: Any):
81 super().__init__(**data)
83 self._resource = data.pop("_resource")
85 self._json_schema = data
87 def get_json_schema(self) -> dict[str, Any]:
88 """
89 Returns the JSON schema of the function which can be passed directly to an LLM.
91 Returns:
92 The JSON schema.
93 """
94 return self._json_schema
96 def get_langchain_tool[T](
97 self,
98 from_function: Callable[..., T],
99 user_id: str,
100 variables: Optional[dict[str, Any]] = None,
101 ) -> T:
102 """
103 Returns a LangChain tool for the function.
105 Args:
106 from_function: This should be LangChain's `StructuredTool.from_function` method.
107 user_id: The user ID of the user on whose behalf the function will be called.
108 variables: The variables to use for mapping the arguments, if applicable.
110 Returns:
111 The LangChain tool.
112 """
113 argument_schema = get_pydantic_model_from_json_schema(
114 json_schema=self.get_json_schema()["parameters"],
115 )
117 tool = from_function(
118 coroutine=self._get_callable(user_id, variables),
119 func=self._get_sync_callable(user_id, variables),
120 name=self.name,
121 description=self.description,
122 args_schema=argument_schema,
123 )
124 return tool
126 def get_haystack_tool[T](
127 self,
128 haystack_tool: Callable[..., T],
129 user_id: str,
130 variables: Optional[dict[str, Any]] = None,
131 ) -> T:
132 """
133 Returns a Haystack tool for the function.
135 Args:
136 haystack_tool: This should be Haystack Tool constructor (`from haystack.tools import Tool`).
137 user_id: The user ID for authentication.
139 Returns:
140 The Haystack tool.
141 """
143 schema = self.get_json_schema()
145 return haystack_tool(
146 name=schema["name"],
147 description=schema["description"],
148 function=self._get_sync_callable(user_id=user_id, variables=variables),
149 parameters=schema["parameters"],
150 )
152 def get_litellm_tool[T](
153 self,
154 ) -> Dict[str, Any]:
155 """
156 Returns a Litellm tool for the function.
158 Generates a Litellm tool based on the function's JSON schema.
160 Returns:
161 The Litellm tool.
162 """
164 schema = self.get_json_schema()
166 return {
167 "type": "function",
168 "function": {
169 "name": schema["name"],
170 "description": schema["description"],
171 "parameters": schema["parameters"],
172 },
173 }
175 def get_mistralai_tool[T](
176 self,
177 ) -> Dict[str, Any]:
178 """
179 Returns a Mistral AI tool for the function.
181 Generates a Mistral AI tool based on the function's JSON schema.
183 Returns:
184 The Mistral AI tool.
185 """
187 schema = self.get_json_schema()
189 return {
190 "type": "function",
191 "function": {
192 "name": schema["name"],
193 "description": schema["description"],
194 "parameters": schema["parameters"],
195 },
196 }
198 def get_llamaindex_tool[T](
199 self,
200 tool_from_defaults: Callable[..., T],
201 tools_metadata: Callable[..., Any],
202 user_id: str,
203 variables: Optional[dict[str, Any]] = None,
204 ) -> T:
205 """
206 Returns a LlamaIndex tool for the function.
208 Args:
209 tool_from_defaults: This should be LlamaIndex's `FunctionTool.from_defaults` method.
210 tools_metadata: This should be LlamaIndex's `ToolMetadata` class.
211 user_id: The user ID for authentication.
213 Returns:
214 The LlamaIndex tool.
215 """
217 function_schema = get_pydantic_model_from_json_schema(
218 json_schema=self.get_json_schema()["parameters"],
219 )
221 metadata = tools_metadata(
222 name=self.name,
223 description=self.description,
224 fn_schema=function_schema,
225 )
227 return tool_from_defaults(
228 async_fn=self._get_callable(user_id=user_id, variables=variables),
229 tool_metadata=metadata,
230 )
232 def register_with_autogen_agents(
233 self,
234 register_function: Callable[..., None],
235 caller: Any,
236 executor: Any,
237 user_id: str,
238 variables: Optional[dict[str, Any]] = None,
239 ):
240 """
241 Registers the function as a tool with AutoGen caller and executor agents.
243 Args:
244 register_function: This should be AutoGen's `register_function` function (`from autogen import register_function`).
245 caller: The caller agent.
246 executor: The executor agent.
247 user_id: The ID of the user on whose behalf the function will be called.
248 variables: The variables to use for mapping the arguments, if applicable.
250 """
251 argument_schema = get_pydantic_model_from_json_schema(
252 json_schema=self.get_json_schema()["parameters"],
253 )
255 async def autogen_function(input: Annotated[argument_schema, f"Input to the {self.name}."]) -> FunctionCallOutput: # type: ignore
256 args = cast(BaseModel, input).model_dump(by_alias=True, exclude_unset=True)
257 return await self._resource.call(
258 self.name,
259 args,
260 user_id,
261 variables,
262 )
264 register_function(
265 autogen_function,
266 caller=caller,
267 executor=executor,
268 name=self.name,
269 description=self.description,
270 )
272 def get_smolagent_tool[T](
273 self,
274 newTool: Callable[..., T],
275 user_id: str,
276 variables: Optional[dict[str, Any]] = None,
277 ) -> T:
278 """
279 Returns a SmolAgent tool for the function.
281 Args:
282 newTool: This should be SmolAgent tool.
283 user_id: The user ID for authentication
284 variables: The variables to use for mapping the arguments, if applicable.
286 Returns:
287 The SmolAgent tool.
288 """
290 def add_docstring(func: Callable[..., Any]) -> Callable[..., Any]:
291 schema = self.get_json_schema()
292 exec_docstring = generate_docstring_from_schema_for_smolagent(schema)
293 func.__doc__ = exec_docstring
294 return func
296 @newTool
297 @add_docstring
298 def execute_function(**kwargs: dict[str, Any]) -> dict[str, Any]:
299 callable_function = self._get_sync_callable(
300 user_id=user_id, variables=variables
301 )
302 result = callable_function(**kwargs)
303 return cast(dict[str, Any], result)
305 return execute_function
307 async def __call__(
308 self,
309 user_id: str,
310 arguments: dict[str, Any],
311 variables: Optional[dict[str, Any]] = None,
312 connected_account_id: Optional[int] = None,
313 ) -> FunctionCallOutput:
314 return await self._resource.call(self.name, arguments, user_id, variables, connected_account_id)
316 def _get_callable(
317 self, user_id: str, variables: Optional[dict[str, Any]] = None
318 ) -> Callable[..., Awaitable[FunctionCallOutput]]:
320 async def callable(**arguments: dict[str, Any]) -> FunctionCallOutput:
321 return await self._resource.call(self.name, arguments, user_id, variables)
323 return callable
325 def _get_sync_callable(
326 self, user_id: str, variables: Optional[dict[str, Any]] = None
327 ):
328 def sync_callable(**arguments: dict[str, Any]) -> FunctionCallOutput:
329 return self._resource.call_sync(self.name, arguments, user_id, variables)
331 return sync_callable
334class FunctionsPage(BaseModel):
335 functions: list[Function]
336 cursor: str
339IncludeOptions = list[Literal["meta"]]
341FunctionType = Literal["ACTION", "QUERY"]