Coverage for src/integry/tool_handlers/lite_llm.py: 26%

1import json 

2from typing import Any, TypedDict, Optional 

3from integry.resources.functions.types import Function, FunctionCallOutput 

4 

5FunctionCall = TypedDict("FunctionCall", {"name": str, "arguments": str}) 

6 

7ToolCall = TypedDict("ToolCall", {"function": FunctionCall, "id": str, "type": str}) 

8 

9Message = TypedDict( 

10 "Message", 

11 {"role": str, "content": Optional[str], "tool_calls": Optional[list[ToolCall]]}, 

12) 

13 

14Choice = TypedDict( 

15 "Choice", 

16 {"index": int, "message": Message, "finish_reason": str}, 

17) 

18 

19LiteLLMResponse = TypedDict( 

20 "LiteLLMResponse", 

21 { 

22 "id": str, 

23 "created": int, 

24 "model": str, 

25 "object": str, 

26 "system_fingerprint": str, 

27 "choices": list[Choice], 

28 }, 

29) 

30 

31 

32async def handle_litellm_tool_calls( 

33 response: LiteLLMResponse, 

34 user_id: str, 

35 call_functions: list[Function], 

36 variables: Optional[dict[str, Any]] = None, 

37) -> list[FunctionCallOutput]: 

38 """ 

39 Processes multiple tool calls from LiteLLM's response and executes the corresponding functions. 

40 

41 Args: 

42 response: The LLM response possibly containing tool calls. 

43 user_id: The user ID on whose behalf the Integry function will be called. 

44 call_functions: A list of functions that can be called. 

45 variables: Additional variables passed to the callable function. 

46 

47 Returns: 

48 A list of results from executed tool functions. The order of results matches the order of tool calls in the response. 

49 """ 

50 choices = response["choices"] 

51 if not choices: 

52 return [] 

53 

54 tool_calls = choices[0]["message"].get("tool_calls") 

55 

56 if not tool_calls: 

57 return [] 

58 

59 results: list[FunctionCallOutput] = [] 

60 

61 for tool_call in tool_calls: 

62 function_name = tool_call["function"]["name"] 

63 function_args = json.loads(tool_call["function"]["arguments"]) 

64 

65 matching_function = next( 

66 ( 

67 func 

68 for func in call_functions 

69 if getattr(func, "name", None) == function_name 

70 ), 

71 None, 

72 ) 

73 

74 if matching_function is not None: 

75 result: FunctionCallOutput = await matching_function( 

76 user_id, function_args, variables 

77 ) 

78 results.append(result) 

79 

80 return results