Coverage for src/lite_agent/message_transfers.py: 96%

1""" 

2Predefined message transfer functions for lite-agent. 

3 

4This module provides common message transfer functions that can be used 

5with agents to preprocess messages before sending them to the API. 

6""" 

7 

8from lite_agent.types import NewUserMessage, RunnerMessages, UserTextContent 

9 

10 

11def consolidate_history_transfer(messages: RunnerMessages) -> RunnerMessages: 

12 """Consolidate all message history into a single user message with XML format. 

13 

14 This message transfer function converts all message history into XML format 

15 and creates a single user message asking what to do next. This is useful when 

16 you want to summarize the entire conversation context in a single prompt. 

17 

18 Args: 

19 messages: The original messages to be processed 

20 

21 Returns: 

22 A single user message containing the consolidated history in XML format 

23 

24 Example: 

25 >>> agent = Agent( 

26 ... model="gpt-4", 

27 ... name="HistoryAgent", 

28 ... instructions="You are a helpful assistant.", 

29 ... message_transfer=consolidate_history_transfer 

30 ... ) 

31 """ 

32 if not messages: 

33 return messages 

34 

35 # Convert messages to XML format 

36 xml_content = ["<conversation_history>"] 

37 

38 for message in messages: 

39 xml_content.extend(_process_message_to_xml(message)) 

40 

41 xml_content.append("</conversation_history>") 

42 

43 # Create the consolidated message 

44 consolidated_content = "以下是目前发生的所有交互:\n\n" + "\n".join(xml_content) + "\n\n接下来该做什么?" 

45 

46 # Return a single user message using NewMessage format 

47 return [NewUserMessage(content=[UserTextContent(text=consolidated_content)])] 

48 

49 

50def _process_message_to_xml(message: dict | object) -> list[str]: 

51 """Process a single message and convert it to XML format. 

52 

53 Args: 

54 message: A single message to process 

55 

56 Returns: 

57 List of XML strings representing the message 

58 """ 

59 xml_lines = [] 

60 

61 if isinstance(message, dict): 

62 xml_lines.extend(_process_dict_message(message)) 

63 elif hasattr(message, "role"): 

64 # Handle Pydantic model format messages 

65 role = getattr(message, "role", "unknown") 

66 content = getattr(message, "content", "") 

67 

68 # Handle new message format where content is a list 

69 if isinstance(content, list): 

70 # Extract text from content items 

71 text_parts = [item.text for item in content if (hasattr(item, "type") and item.type == "text") or hasattr(item, "text")] 

72 content_text = " ".join(text_parts) 

73 if content_text: 

74 xml_lines.append(f" <message role='{role}'>{content_text}</message>") 

75 elif isinstance(content, str): 

76 xml_lines.append(f" <message role='{role}'>{content}</message>") 

77 elif hasattr(message, "type"): 

78 # Handle function call messages 

79 xml_lines.extend(_process_function_message(message)) 

80 

81 return xml_lines 

82 

83 

84def _process_dict_message(message: dict) -> list[str]: 

85 """Process dictionary format message to XML.""" 

86 xml_lines = [] 

87 role = message.get("role", "unknown") 

88 content = message.get("content", "") 

89 message_type = message.get("type") 

90 

91 if message_type == "function_call": 

92 name = message.get("name", "unknown") 

93 arguments = message.get("arguments", "") 

94 xml_lines.append(f" <function_call name='{name}' arguments='{arguments}' />") 

95 elif message_type == "function_call_output": 

96 call_id = message.get("call_id", "unknown") 

97 output = message.get("output", "") 

98 xml_lines.append(f" <function_result call_id='{call_id}'>{output}</function_result>") 

99 elif role in ["user", "assistant", "system"]: 

100 xml_lines.append(f" <message role='{role}'>{content}</message>") 

101 

102 return xml_lines 

103 

104 

105def _process_function_message(message: dict | object) -> list[str]: 

106 """Process function call message to XML.""" 

107 xml_lines = [] 

108 message_type = getattr(message, "type", "unknown") 

109 

110 if message_type == "function_call": 

111 name = getattr(message, "name", "unknown") 

112 arguments = getattr(message, "arguments", "") 

113 xml_lines.append(f" <function_call name='{name}' arguments='{arguments}' />") 

114 elif message_type == "function_call_output": 

115 call_id = getattr(message, "call_id", "unknown") 

116 output = getattr(message, "output", "") 

117 xml_lines.append(f" <function_result call_id='{call_id}'>{output}</function_result>") 

118 

119 return xml_lines