Coverage for src/lite_agent/utils/message_builder.py: 73%

1import json 

2from typing import Any 

3 

4from lite_agent.types import ( 

5 AssistantMessageContent, 

6 AssistantMessageMeta, 

7 AssistantTextContent, 

8 AssistantToolCall, 

9 AssistantToolCallResult, 

10 MessageMeta, 

11 NewAssistantMessage, 

12 NewSystemMessage, 

13 NewUserMessage, 

14 UserImageContent, 

15 UserMessageContent, 

16 UserTextContent, 

17) 

18 

19 

20class MessageBuilder: 

21 """Utility class for building and converting messages from various formats.""" 

22 

23 @staticmethod 

24 def build_user_message_from_dict(message: dict[str, Any]) -> NewUserMessage: 

25 """Build a NewUserMessage from a dictionary. 

26 

27 Args: 

28 message: Dictionary containing user message data 

29 

30 Returns: 

31 NewUserMessage instance 

32 """ 

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

34 

35 # Preserve meta information if present 

36 meta_data = message.get("meta", {}) 

37 meta = MessageMeta(**meta_data) if meta_data else MessageMeta() 

38 

39 if isinstance(content, str): 

40 return NewUserMessage(content=[UserTextContent(text=content)], meta=meta) 

41 

42 if isinstance(content, list): 

43 return NewUserMessage(content=MessageBuilder._build_user_content_items(content), meta=meta) 

44 

45 # Handle non-string, non-list content 

46 return NewUserMessage(content=[UserTextContent(text=str(content))], meta=meta) 

47 

48 @staticmethod 

49 def _build_user_content_items(content_list: list[Any]) -> list[UserMessageContent]: 

50 """Build user content items from a list of content data. 

51 

52 Args: 

53 content_list: List of content items (dicts or objects) 

54 

55 Returns: 

56 List of UserMessageContent items 

57 """ 

58 user_content_items: list[UserMessageContent] = [] 

59 

60 for item in content_list: 

61 if isinstance(item, dict): 

62 user_content_items.append(MessageBuilder._build_user_content_from_dict(item)) 

63 elif hasattr(item, "type"): 

64 user_content_items.append(MessageBuilder._build_user_content_from_object(item)) 

65 else: 

66 # Fallback: convert to text 

67 user_content_items.append(UserTextContent(text=str(item))) 

68 

69 return user_content_items 

70 

71 @staticmethod 

72 def _build_user_content_from_dict(item: dict[str, Any]) -> UserMessageContent: 

73 """Build user content from a dictionary item. 

74 

75 Args: 

76 item: Dictionary containing content item data 

77 

78 Returns: 

79 UserMessageContent instance 

80 """ 

81 item_type = item.get("type") 

82 

83 if item_type in {"input_text", "text"}: 

84 return UserTextContent(text=item.get("text", "")) 

85 

86 if item_type in {"input_image", "image_url"}: 

87 if item_type == "image_url": 

88 # Handle completion API format 

89 image_url_data = item.get("image_url", {}) 

90 url = image_url_data.get("url", "") if isinstance(image_url_data, dict) else str(image_url_data) 

91 return UserImageContent(image_url=url) 

92 

93 # Handle response API format 

94 return UserImageContent( 

95 image_url=item.get("image_url"), 

96 file_id=item.get("file_id"), 

97 detail=item.get("detail", "auto"), 

98 ) 

99 

100 # Fallback: treat as text 

101 return UserTextContent(text=str(item.get("text", item))) 

102 

103 @staticmethod 

104 def _build_user_content_from_object(item: Any) -> UserMessageContent: # noqa: ANN401 

105 """Build user content from an object with attributes. 

106 

107 Args: 

108 item: Object with type attribute and other properties 

109 

110 Returns: 

111 UserMessageContent instance 

112 """ 

113 if item.type == "input_text": 

114 return UserTextContent(text=item.text) 

115 

116 if item.type == "input_image": 

117 return UserImageContent( 

118 image_url=getattr(item, "image_url", None), 

119 file_id=getattr(item, "file_id", None), 

120 detail=getattr(item, "detail", "auto"), 

121 ) 

122 

123 # Fallback: convert to text 

124 return UserTextContent(text=str(item)) 

125 

126 @staticmethod 

127 def build_system_message_from_dict(message: dict[str, Any]) -> NewSystemMessage: 

128 """Build a NewSystemMessage from a dictionary. 

129 

130 Args: 

131 message: Dictionary containing system message data 

132 

133 Returns: 

134 NewSystemMessage instance 

135 """ 

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

137 

138 # Preserve meta information if present 

139 meta_data = message.get("meta", {}) 

140 meta = MessageMeta(**meta_data) if meta_data else MessageMeta() 

141 

142 return NewSystemMessage(content=str(content), meta=meta) 

143 

144 @staticmethod 

145 def build_assistant_message_from_dict(message: dict[str, Any]) -> NewAssistantMessage: 

146 """Build a NewAssistantMessage from a dictionary. 

147 

148 Args: 

149 message: Dictionary containing assistant message data 

150 

151 Returns: 

152 NewAssistantMessage instance 

153 """ 

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

155 assistant_content_items: list[AssistantMessageContent] = [] 

156 

157 if content: 

158 if isinstance(content, str): 

159 assistant_content_items = [AssistantTextContent(text=content)] 

160 elif isinstance(content, list): 

161 # Handle array content (from new format messages) 

162 for item in content: 

163 if isinstance(item, dict): 

164 item_type = item.get("type") 

165 if item_type == "text": 

166 assistant_content_items.append(AssistantTextContent(text=item.get("text", ""))) 

167 elif item_type == "tool_call": 

168 assistant_content_items.append( 

169 AssistantToolCall( 

170 call_id=item.get("call_id", ""), 

171 name=item.get("name", ""), 

172 arguments=item.get("arguments", "{}"), 

173 ), 

174 ) 

175 elif item_type == "tool_call_result": 

176 assistant_content_items.append( 

177 AssistantToolCallResult( 

178 call_id=item.get("call_id", ""), 

179 output=item.get("output", ""), 

180 execution_time_ms=item.get("execution_time_ms"), 

181 ), 

182 ) 

183 # Add more content types as needed 

184 else: 

185 # Fallback for unknown item format 

186 assistant_content_items.append(AssistantTextContent(text=str(item))) 

187 else: 

188 # Fallback for other content types 

189 assistant_content_items = [AssistantTextContent(text=str(content))] 

190 

191 # Handle tool calls if present (legacy format) 

192 if "tool_calls" in message: 

193 for tool_call in message.get("tool_calls", []): 

194 try: 

195 arguments = json.loads(tool_call["function"]["arguments"]) if isinstance(tool_call["function"]["arguments"], str) else tool_call["function"]["arguments"] 

196 except (json.JSONDecodeError, TypeError): 

197 arguments = tool_call["function"]["arguments"] 

198 

199 assistant_content_items.append( 

200 AssistantToolCall( 

201 call_id=tool_call["id"], 

202 name=tool_call["function"]["name"], 

203 arguments=arguments, 

204 ), 

205 ) 

206 

207 # Preserve meta information if present 

208 meta_data = message.get("meta", {}) 

209 meta = AssistantMessageMeta(**meta_data) if meta_data else AssistantMessageMeta() 

210 

211 return NewAssistantMessage(content=assistant_content_items, meta=meta)