Coverage for src / agent / providers / gemini / types.py: 97%

1# Copyright 2025-2026 Microsoft Corporation 

2# 

3# Licensed under the Apache License, Version 2.0 (the "License"); 

4# you may not use this file except in compliance with the License. 

5# You may obtain a copy of the License at 

6# 

7# http://www.apache.org/licenses/LICENSE-2.0 

8# 

9# Unless required by applicable law or agreed to in writing, software 

10# distributed under the License is distributed on an "AS IS" BASIS, 

11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 

12# See the License for the specific language governing permissions and 

13# limitations under the License. 

14 

15""" 

16Type definitions and conversion utilities for Gemini provider. 

17 

18This module provides functions to convert between agent-framework's message 

19formats and Google Gemini's expected formats. 

20""" 

21 

22from typing import Any 

23 

24from agent_framework import ( 

25 AIFunction, 

26 ChatMessage, 

27 FunctionCallContent, 

28 FunctionResultContent, 

29 TextContent, 

30) 

31 

32 

33def to_gemini_message( 

34 message: ChatMessage, call_id_to_name: dict[str, str] | None = None 

35) -> dict[str, Any]: 

36 """Convert agent-framework ChatMessage to Gemini message format. 

37 

38 Args: 

39 message: ChatMessage from agent-framework 

40 

41 Returns: 

42 Dictionary in Gemini's message format 

43 

44 Example: 

45 >>> msg = ChatMessage(role="user", contents=[TextContent(text="Hello")]) 

46 >>> to_gemini_message(msg) 

47 {"role": "user", "parts": [{"text": "Hello"}]} 

48 """ 

49 # Map roles (agent-framework -> Gemini) 

50 role_mapping = { 

51 "user": "user", 

52 "assistant": "model", 

53 "system": "user", # Gemini doesn't have system role, treat as user 

54 "tool": "user", # Tool results are supplied as user role in Gemini 

55 } 

56 # Handle Role enum or string 

57 role_str = message.role.value if hasattr(message.role, "value") else str(message.role) 

58 gemini_role = role_mapping.get(role_str, "user") 

59 

60 # Convert contents to parts 

61 parts = [] 

62 

63 # Get contents from message (agent-framework uses 'contents' attribute) 

64 contents = message.contents if hasattr(message, "contents") else [] 

65 

66 # Handle different content types with role-aware rules 

67 for item in contents: 

68 if isinstance(item, TextContent): 

69 # Text content is allowed in any turn 

70 parts.append({"text": item.text}) 

71 continue 

72 

73 if isinstance(item, FunctionCallContent): 

74 # Function call should only appear in a model (assistant) turn 

75 if gemini_role == "model": 

76 parts.append({"function_call": {"name": item.name, "args": item.arguments}}) # type: ignore[dict-item] 

77 # Skip function_call in non-model turns 

78 continue 

79 

80 if isinstance(item, FunctionResultContent): 

81 # Function result should appear only in a user/tool turn 

82 if gemini_role == "user": 

83 name = None 

84 if call_id_to_name and item.call_id in call_id_to_name: 

85 name = call_id_to_name[item.call_id] 

86 # If we can't resolve the name, skip to avoid invalid request 

87 if not name: 

88 continue 

89 response_dict: dict[str, Any] = { 

90 "function_response": { 

91 "name": name, 

92 "response": {"result": item.result}, 

93 } 

94 } 

95 parts.append(response_dict) 

96 # Skip function_result in model turns 

97 continue 

98 

99 # If no parts were added, attempt to use message.text; otherwise add empty text part 

100 if not parts: 

101 # Some ChatMessage instances may expose text; use it if present 

102 text_value = getattr(message, "text", "") 

103 parts.append({"text": text_value}) 

104 

105 return {"role": gemini_role, "parts": parts} 

106 

107 

108def from_gemini_message(gemini_response: Any) -> ChatMessage: 

109 """Convert Gemini response to agent-framework ChatMessage. 

110 

111 Args: 

112 gemini_response: Response object from Gemini API 

113 

114 Returns: 

115 ChatMessage for agent-framework 

116 

117 Example: 

118 >>> response = gemini_client.generate_content("Hello") 

119 >>> msg = from_gemini_message(response) 

120 >>> msg.role 

121 'assistant' 

122 """ 

123 content_items: list[TextContent | FunctionCallContent] = [] 

124 

125 # Extract text and function calls from response 

126 if hasattr(gemini_response, "candidates") and gemini_response.candidates: 

127 candidate = gemini_response.candidates[0] 

128 if hasattr(candidate, "content") and candidate.content: 

129 for part in candidate.content.parts: 

130 # Text content 

131 if hasattr(part, "text") and part.text: 

132 content_items.append(TextContent(text=part.text)) 

133 

134 # Function call 

135 if hasattr(part, "function_call") and part.function_call: 

136 fc = part.function_call 

137 # Generate a call_id if not provided by Gemini 

138 call_id = getattr(fc, "id", None) or f"{fc.name}_{id(fc)}" 

139 content_items.append( 

140 FunctionCallContent( 

141 call_id=call_id, 

142 name=fc.name, 

143 arguments=dict(fc.args) if hasattr(fc, "args") else {}, 

144 ) 

145 ) 

146 

147 # Return ChatMessage with contents list 

148 # Note: agent-framework uses 'contents' (plural) not 'content' 

149 if not content_items: 

150 content_items = [TextContent(text="")] 

151 

152 return ChatMessage(role="assistant", contents=content_items) 

153 

154 

155def to_gemini_tools(tools: list[AIFunction]) -> list[dict[str, Any]]: 

156 """Convert agent-framework AIFunction tools to Gemini function declarations. 

157 

158 Args: 

159 tools: List of AIFunction objects from agent-framework 

160 

161 Returns: 

162 List of Gemini function declaration dictionaries 

163 

164 Example: 

165 >>> tools = [AIFunction(name="get_weather", description="Get weather")] 

166 >>> gemini_tools = to_gemini_tools(tools) 

167 >>> gemini_tools[0]["name"] 

168 'get_weather' 

169 """ 

170 function_declarations = [] 

171 

172 for tool in tools: 

173 # Extract function schema 

174 function_declaration = { 

175 "name": tool.name, 

176 "description": tool.description or "", 

177 } 

178 

179 # Add parameters if available 

180 # Note: tool.parameters might be a method, so check if callable 

181 if hasattr(tool, "parameters"): 

182 params = tool.parameters() if callable(tool.parameters) else tool.parameters 

183 if params: 

184 # Convert parameters to Gemini format 

185 function_declaration["parameters"] = _convert_parameters(params) # type: ignore[assignment] 

186 

187 function_declarations.append(function_declaration) 

188 

189 # Return a single tools object with all function declarations 

190 return [{"function_declarations": function_declarations}] if function_declarations else [] 

191 

192 

193def _convert_parameters(parameters: dict[str, Any]) -> dict[str, Any]: 

194 """Convert agent-framework parameter schema to Gemini format. 

195 

196 Args: 

197 parameters: Parameter schema from AIFunction 

198 

199 Returns: 

200 Gemini-compatible parameter schema 

201 """ 

202 # Gemini expects OpenAPI-like schema 

203 gemini_params = { 

204 "type": "object", 

205 "properties": {}, 

206 } 

207 

208 # Copy properties if they exist 

209 if "properties" in parameters: 

210 gemini_params["properties"] = parameters["properties"] 

211 

212 # Copy required fields if they exist 

213 if "required" in parameters: 

214 gemini_params["required"] = parameters["required"] 

215 

216 return gemini_params 

217 

218 

219def extract_usage_metadata(gemini_response: Any) -> dict[str, Any]: 

220 """Extract token usage information from Gemini response. 

221 

222 Args: 

223 gemini_response: Response object from Gemini API 

224 

225 Returns: 

226 Dictionary with usage metadata 

227 

228 Example: 

229 >>> usage = extract_usage_metadata(response) 

230 >>> usage["prompt_tokens"] 

231 15 

232 """ 

233 usage = {} 

234 

235 if hasattr(gemini_response, "usage_metadata"): 

236 metadata = gemini_response.usage_metadata 

237 usage["prompt_tokens"] = getattr(metadata, "prompt_token_count", 0) 

238 usage["completion_tokens"] = getattr(metadata, "candidates_token_count", 0) 

239 usage["total_tokens"] = getattr(metadata, "total_token_count", 0) 

240 

241 return usage