Coverage for agentos/models/backends/gemini.py: 24%
153 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-07-06 23:17 +0800
1"""
2AgentOS v0.70 — Google Gemini Provider 全集成。
3基因来源: Google AI Studio SDK + Vertex AI
4支持: Gemini 2.5 Pro/Flash、Vision、System Instruction、Streaming、Token Counting、Safety Settings。
5"""
7from __future__ import annotations
9import json
10import os
11from collections.abc import AsyncIterator
12from dataclasses import dataclass, field
14import httpx
16from agentos.core.context import AgentContext
17from agentos.models.router import ModelResponse, ModelSpec
18from agentos.tools.base import ToolCall
20# ── Gemini Public API Endpoint ──────────────────
21GEMINI_API_BASE = "https://generativelanguage.googleapis.com/v1beta"
23# Prebuilt Gemini model specs
24GEMINI_MODELS: dict[str, ModelSpec] = {
25 "gemini-2.5-pro": ModelSpec(
26 provider="gemini",
27 model_id="gemini-2.5-pro-exp-03-25",
28 context_window=1_048_576,
29 cost_per_1m_input=1.25,
30 cost_per_1m_output=10.00,
31 ),
32 "gemini-2.5-flash": ModelSpec(
33 provider="gemini",
34 model_id="gemini-2.5-flash-preview-04-17",
35 context_window=1_048_576,
36 cost_per_1m_input=0.15,
37 cost_per_1m_output=0.60,
38 ),
39 "gemini-2.0-flash": ModelSpec(
40 provider="gemini",
41 model_id="gemini-2.0-flash",
42 context_window=1_048_576,
43 cost_per_1m_input=0.10,
44 cost_per_1m_output=0.40,
45 ),
46}
49@dataclass
50class GeminiSafetySetting:
51 """安全过滤配置。"""
53 category: str # HARM_CATEGORY_HARASSMENT | HATE_SPEECH | SEXUALLY_EXPLICIT | DANGEROUS_CONTENT
54 threshold: str = (
55 "BLOCK_ONLY_HIGH" # BLOCK_NONE | BLOCK_ONLY_HIGH | BLOCK_MEDIUM_AND_ABOVE | BLOCK_LOW_AND_ABOVE
56 )
59@dataclass
60class GeminiConfig:
61 """Gemini调用配置。"""
63 api_key: str = ""
64 temperature: float = 0.7
65 top_p: float = 0.95
66 top_k: int = 40
67 max_output_tokens: int = 8192
68 safety_settings: list[GeminiSafetySetting] = field(
69 default_factory=lambda: [
70 GeminiSafetySetting("HARM_CATEGORY_HARASSMENT", "BLOCK_ONLY_HIGH"),
71 GeminiSafetySetting("HARM_CATEGORY_HATE_SPEECH", "BLOCK_ONLY_HIGH"),
72 GeminiSafetySetting("HARM_CATEGORY_SEXUALLY_EXPLICIT", "BLOCK_ONLY_HIGH"),
73 GeminiSafetySetting("HARM_CATEGORY_DANGEROUS_CONTENT", "BLOCK_ONLY_HIGH"),
74 ]
75 )
78# ── Tool Declaration Helpers ─────────────────────
81def _convert_tools_to_gemini(openai_tools: list[dict]) -> list[dict]:
82 """将OpenAI格式的tools转换为Gemini functionDeclarations。"""
83 declarations = []
84 for tool in openai_tools:
85 if tool.get("type") != "function":
86 continue
87 func = tool.get("function", {})
88 declarations.append(
89 {
90 "name": func.get("name", ""),
91 "description": func.get("description", ""),
92 "parameters": func.get("parameters", {}),
93 }
94 )
95 return [{"function_declarations": declarations}] if declarations else []
98def _convert_gemini_tool_calls(parts: list[dict]) -> list[ToolCall]:
99 """将Gemini functionCall parts转为ToolCall列表。"""
100 tool_calls = []
101 for part in parts:
102 fc = part.get("functionCall")
103 if not fc:
104 continue
105 args = fc.get("args", {})
106 if isinstance(args, str):
107 try:
108 args = json.loads(args)
109 except json.JSONDecodeError:
110 args = {}
111 tool_calls.append(
112 ToolCall(
113 id=fc.get("name", "unknown"),
114 name=fc.get("name", "unknown"),
115 arguments=args,
116 )
117 )
118 return tool_calls
121# ── Core Gemini Client ───────────────────────────
124class GeminiClient:
125 """
126 Google Gemini API 客户端。
127 支持: chat/completions、Vision多模态、Streaming、System Instruction。
128 """
130 def __init__(
131 self,
132 config: GeminiConfig | None = None,
133 http_client: httpx.AsyncClient | None = None,
134 ):
135 self.config = config or GeminiConfig()
136 self._http = http_client or httpx.AsyncClient(timeout=180)
137 self._owned_http = http_client is None
139 @property
140 def api_key(self) -> str:
141 return self.config.api_key or os.environ.get("GEMINI_API_KEY", "")
143 async def close(self):
144 if self._owned_http:
145 await self._http.aclose()
147 async def call(
148 self,
149 spec: ModelSpec,
150 context: AgentContext,
151 ) -> ModelResponse:
152 """同步调用Gemini API。"""
153 contents, system_instruction = self._build_gemini_contents(context)
154 body = {
155 "contents": contents,
156 "generationConfig": {
157 "temperature": self.config.temperature,
158 "topP": self.config.top_p,
159 "topK": self.config.top_k,
160 "maxOutputTokens": self.config.max_output_tokens,
161 },
162 "safetySettings": [
163 {"category": s.category, "threshold": s.threshold}
164 for s in self.config.safety_settings
165 ],
166 }
167 if system_instruction:
168 body["systemInstruction"] = system_instruction
170 if context.tools:
171 body["tools"] = _convert_tools_to_gemini(context.tools)
173 url = f"{GEMINI_API_BASE}/models/{spec.model_id}:generateContent?key={self.api_key}"
174 resp = await self._http.post(url, json=body)
175 resp.raise_for_status()
176 data = resp.json()
178 return self._parse_response(data)
180 async def call_stream(
181 self,
182 spec: ModelSpec,
183 context: AgentContext,
184 ) -> AsyncIterator[dict]:
185 """流式调用Gemini API，逐个yield chunk。"""
186 contents, system_instruction = self._build_gemini_contents(context)
187 body = {
188 "contents": contents,
189 "generationConfig": {
190 "temperature": self.config.temperature,
191 "topP": self.config.top_p,
192 "topK": self.config.top_k,
193 "maxOutputTokens": self.config.max_output_tokens,
194 },
195 "safetySettings": [
196 {"category": s.category, "threshold": s.threshold}
197 for s in self.config.safety_settings
198 ],
199 }
200 if system_instruction:
201 body["systemInstruction"] = system_instruction
203 url = f"{GEMINI_API_BASE}/models/{spec.model_id}:streamGenerateContent?alt=sse&key={self.api_key}"
204 async with self._http.stream("POST", url, json=body) as resp:
205 resp.raise_for_status()
206 async for line in resp.aiter_lines():
207 line = line.strip()
208 if not line or not line.startswith("data: "):
209 continue
210 data_str = line[6:]
211 if data_str == "[DONE]":
212 break
213 try:
214 chunk = json.loads(data_str)
215 except json.JSONDecodeError:
216 continue
217 # skip safety / promptFeedback
218 if "candidates" not in chunk:
219 continue
220 yield chunk
222 async def call_with_image(
223 self,
224 spec: ModelSpec,
225 prompt: str,
226 image_data: bytes,
227 mime_type: str = "image/jpeg",
228 ) -> ModelResponse:
229 """Vision多模态调用。image_data为base64之前的内容。"""
230 import base64
232 b64 = base64.b64encode(image_data).decode()
233 contents = [
234 {
235 "role": "user",
236 "parts": [
237 {"text": prompt},
238 {"inlineData": {"mimeType": mime_type, "data": b64}},
239 ],
240 }
241 ]
242 body = {
243 "contents": contents,
244 "generationConfig": {
245 "temperature": self.config.temperature,
246 "maxOutputTokens": self.config.max_output_tokens,
247 },
248 }
249 url = f"{GEMINI_API_BASE}/models/{spec.model_id}:generateContent?key={self.api_key}"
250 resp = await self._http.post(url, json=body)
251 resp.raise_for_status()
252 return self._parse_response(resp.json())
254 async def count_tokens(self, spec: ModelSpec, context: AgentContext) -> dict:
255 """使用Gemini API统计输入/输出token数。"""
256 contents, _ = self._build_gemini_contents(context)
257 url = f"{GEMINI_API_BASE}/models/{spec.model_id}:countTokens?key={self.api_key}"
258 resp = await self._http.post(url, json={"contents": contents})
259 resp.raise_for_status()
260 data = resp.json()
261 return {
262 "total_tokens": data.get("totalTokens", 0),
263 "prompt_tokens": data.get("totalTokens", 0), # Gemini不区分输入输出
264 "model": spec.model_id,
265 }
267 # ── Internal helpers ──────────────────────────
269 def _build_gemini_contents(self, context: AgentContext) -> tuple[list[dict], dict | None]:
270 """将AgentContext转为Gemini contents格式。"""
271 contents = []
272 system_instruction = None
274 for msg in context.messages:
275 role = self._map_role(msg.role)
276 parts = []
278 # system prompt → systemInstruction
279 if msg.role == "system":
280 system_instruction = {"parts": [{"text": msg.content}]}
281 continue
283 # text content
284 if msg.content:
285 parts.append({"text": msg.content})
287 # tool calls from assistant
288 if msg.tool_calls:
289 for tc in msg.tool_calls:
290 parts.append(
291 {
292 "functionCall": {
293 "name": tc.name,
294 "args": tc.arguments,
295 }
296 }
297 )
299 # tool results
300 if msg.role == "tool" and msg.tool_call_id:
301 # Gemini uses functionResponse in user role
302 parts.append(
303 {
304 "functionResponse": {
305 "name": msg.tool_call_id,
306 "response": {"content": msg.content},
307 }
308 }
309 )
311 if parts:
312 contents.append({"role": role, "parts": parts})
314 # Ensure there's at least a user message
315 if not contents:
316 contents = [{"role": "user", "parts": [{"text": context.current_task or ""}]}]
318 return contents, system_instruction
320 def _map_role(self, role: str) -> str:
321 mapping = {
322 "user": "user",
323 "assistant": "model",
324 "system": "user", # handled separately via systemInstruction
325 "tool": "user", # functionResponse must be in user turn
326 }
327 return mapping.get(role, "user")
329 def _parse_response(self, data: dict) -> ModelResponse:
330 """解析Gemini API响应为ModelResponse。"""
331 candidates = data.get("candidates", [])
332 if not candidates:
333 # Safety blocked
334 block_reason = data.get("promptFeedback", {}).get("blockReason", "unknown")
335 return ModelResponse(content=f"[SAFETY_BLOCKED] {block_reason}")
337 candidate = candidates[0]
338 content = candidate.get("content", {})
339 parts = content.get("parts", [])
341 text_parts = []
342 tool_calls = []
344 for part in parts:
345 if "text" in part:
346 text_parts.append(part["text"])
347 if "functionCall" in part:
348 fc = part["functionCall"]
349 args = fc.get("args", {})
350 if isinstance(args, str):
351 try:
352 args = json.loads(args)
353 except json.JSONDecodeError:
354 args = {}
355 tool_calls.append(
356 ToolCall(
357 id=fc.get("name", "unknown"),
358 name=fc.get("name", "unknown"),
359 arguments=args,
360 )
361 )
363 return ModelResponse(
364 content="\n".join(text_parts),
365 tool_calls=tool_calls,
366 )