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