oclaw/platform/llm/transports/google_gemini_sse.py
oliver 4a23b715a2 重构仓库目录为统一的 runtime 分层并清理历史 openclaw 残留。
本次迁移将网关/通道/工具/技能/脚本与协议资源集中到新结构,统一路径常量与脚本转发机制,减少顶层噪音并保证运行与测试行为一致。

Made-with: Cursor
2026-04-25 01:24:23 +08:00

265 lines
11 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

from __future__ import annotations
import json
import os
import uuid
from typing import Any, Optional
from collections.abc import Callable
from oclaw.platform.llm.transports.base import ChatModel, LLMResponse, LLMToolCall
class GoogleGeminiChatModel(ChatModel):
"""Native Google Gemini streaming transport (SSE), modeled after Oclaw's google transport."""
def __init__(self, *, model: str, api_key: str, base_url: str | None = None):
self.model = (model or "").strip() or "gemini-2.5-pro"
self.api_key = (api_key or "").strip()
self.base_url = (base_url or "").strip() or "https://generativelanguage.googleapis.com/v1beta"
if not self.api_key:
raise RuntimeError("未设置 GOOGLE_API_KEY,无法使用 Gemini 原生接口")
def _tool_decls_from_openai_tools(self, tools: list[dict[str, Any]]) -> list[dict[str, Any]]:
decls: list[dict[str, Any]] = []
for t in tools or []:
if not isinstance(t, dict) or str(t.get("type") or "") != "function":
continue
fn = t.get("function")
if not isinstance(fn, dict):
continue
name = str(fn.get("name") or "").strip()
if not name:
continue
decls.append(
{
"name": name,
"description": str(fn.get("description") or ""),
# Oclaw uses `parametersJsonSchema`
"parametersJsonSchema": fn.get("parameters") if isinstance(fn.get("parameters"), dict) else {"type": "object"},
}
)
return decls
@staticmethod
def _thinking_config_from_env() -> dict[str, Any] | None:
enabled = str(os.getenv("AIA_GEMINI_THINKING") or "").strip().lower()
if enabled in ("0", "false", "no", "off"):
return {"includeThoughts": False}
if enabled in ("1", "true", "yes", "on"):
cfg: dict[str, Any] = {"includeThoughts": True}
lvl = str(os.getenv("AIA_GEMINI_THINKING_LEVEL") or "").strip()
if lvl:
cfg["thinkingLevel"] = lvl
bud = str(os.getenv("AIA_GEMINI_THINKING_BUDGET") or "").strip()
if bud.isdigit():
cfg["thinkingBudget"] = int(bud)
return cfg
return None
@staticmethod
def _parse_tool_result_payload(raw: Any) -> Any:
if raw is None:
return {}
if isinstance(raw, (dict, list)):
return raw
s = str(raw or "").strip()
if not s:
return {}
try:
return json.loads(s)
except Exception:
return {"raw": s}
@staticmethod
def _openai_messages_to_gemini_contents(messages: list[dict[str, Any]]) -> tuple[dict[str, Any] | None, list[dict[str, Any]]]:
system_parts: list[dict[str, Any]] = []
contents: list[dict[str, Any]] = []
def _append(role: str, parts: list[dict[str, Any]]) -> None:
if not parts:
return
contents.append({"role": role, "parts": parts})
for m in messages or []:
if not isinstance(m, dict):
continue
role = str(m.get("role") or "")
content = m.get("content")
if role == "system":
if isinstance(content, str) and content.strip():
system_parts.append({"text": content})
continue
if role == "user":
parts: list[dict[str, Any]] = []
if isinstance(content, str):
if content.strip():
parts.append({"text": content})
elif isinstance(content, list):
for item in content:
if not isinstance(item, dict):
continue
if item.get("type") == "text" and str(item.get("text") or "").strip():
parts.append({"text": str(item.get("text") or "")})
_append("user", parts)
continue
if role == "assistant":
parts: list[dict[str, Any]] = []
if isinstance(content, str) and content.strip():
parts.append({"text": content})
_append("model", parts)
continue
if role == "tool":
name = str(m.get("name") or "").strip() or str(m.get("tool_name") or m.get("toolName") or "").strip()
payload = GoogleGeminiChatModel._parse_tool_result_payload(content)
_append("user", [{"functionResponse": {"name": name or "tool", "response": payload}}])
continue
system_instruction = {"parts": system_parts} if system_parts else None
return system_instruction, contents
def _build_url(self) -> str:
base = self._normalize_base_url_for_native_sse(self.base_url).rstrip("/")
model_path = self.model
if not model_path.startswith("models/") and not model_path.startswith("tunedModels/"):
model_path = f"models/{model_path}"
# Prefer header auth (Oclaw-style); avoid forcing key into query string.
return f"{base}/{model_path}:streamGenerateContent?alt=sse"
@staticmethod
def _normalize_base_url_for_native_sse(base_url: str) -> str:
"""Normalize unified gateway base_url for Gemini native SSE.
Many unified gateways expose OpenAI-compat under `/v1/*`, but Gemini native endpoints
are usually implemented under `/v1beta/*` semantics.
"""
raw = str(base_url or "").strip()
if not raw:
return raw
# Strip trailing slashes first.
while raw.endswith("/"):
raw = raw[:-1]
# Common unified gateway layout: https://host/v1 -> https://host/v1beta
if raw.lower().endswith("/v1"):
return raw[:-3] + "/v1beta"
# Some gateways keep an `/openai` suffix for compat endpoints.
if raw.lower().endswith("/openai"):
raw = raw[:-7]
if raw.lower().endswith("/v1"):
return raw[:-3] + "/v1beta"
return raw
return raw
def _auth_headers(self) -> dict[str, str]:
k = str(self.api_key or "").strip()
if not k:
return {}
# If user stored a bearer token in the profile secret.
if k.lower().startswith("bearer "):
return {"authorization": k}
# Gemini API supports x-goog-api-key; unified gateways often accept this too.
return {"x-goog-api-key": k}
def chat(
self,
messages: list[dict[str, Any]],
tools: list[dict[str, Any]],
*,
on_token: Optional[Callable[[str], None]] = None,
) -> LLMResponse:
system_instruction, contents = self._openai_messages_to_gemini_contents(messages)
decls = self._tool_decls_from_openai_tools(tools)
req: dict[str, Any] = {"contents": contents}
if system_instruction is not None:
req["systemInstruction"] = system_instruction
if decls:
req["tools"] = [{"functionDeclarations": decls}]
thinking_cfg = self._thinking_config_from_env()
if thinking_cfg is not None:
req["generationConfig"] = {"thinkingConfig": thinking_cfg}
url = self._build_url()
try:
import httpx
except Exception as exc:
raise RuntimeError("missing dependency: httpx") from exc
thinking_parts: list[str] = []
thinking_signature: str | None = None
out_text: list[str] = []
tool_calls: list[LLMToolCall] = []
def _stream_once(stream_url: str) -> None:
with httpx.Client(timeout=httpx.Timeout(60.0, connect=10.0)) as client:
with client.stream(
"POST",
stream_url,
json=req,
headers={
"Content-Type": "application/json",
"accept": "text/event-stream",
**self._auth_headers(),
},
) as r:
r.raise_for_status()
for line in r.iter_lines():
if not line or not str(line).startswith("data:"):
continue
data = str(line)[len("data:") :].strip()
if not data or data == "[DONE]":
continue
try:
chunk = json.loads(data)
except Exception:
continue
cands = chunk.get("candidates") if isinstance(chunk, dict) else None
if not isinstance(cands, list) or not cands:
continue
c0 = cands[0] if isinstance(cands[0], dict) else {}
content0 = c0.get("content") if isinstance(c0.get("content"), dict) else {}
parts = content0.get("parts") if isinstance(content0.get("parts"), list) else []
for p in parts:
if not isinstance(p, dict):
continue
if p.get("thought") is True:
txt = str(p.get("text") or "")
if txt:
thinking_parts.append(txt)
sig = p.get("thoughtSignature")
if isinstance(sig, str) and sig:
thinking_signature = sig
continue
fc = p.get("functionCall") if isinstance(p.get("functionCall"), dict) else None
if fc is not None:
name = str(fc.get("name") or "").strip()
args = fc.get("args") if isinstance(fc.get("args"), dict) else {}
tid = str(fc.get("id") or "") or f"call_{uuid.uuid4().hex}"
tool_calls.append(
LLMToolCall(id=tid, name=name, arguments=args, thought_signature=thinking_signature)
)
continue
txt = str(p.get("text") or "")
if txt:
out_text.append(txt)
if on_token:
on_token(txt)
try:
_stream_once(url)
except Exception as exc:
# If gateway returns 500 on /v1* SSE path, try /v1beta* fallback once.
msg = str(exc or "").lower()
if "500" in msg and "/v1/" in str(url).lower():
alt = str(url).replace("/v1/", "/v1beta/")
_stream_once(alt)
else:
raise
return LLMResponse(
content="".join(out_text),
tool_calls=tool_calls,
reasoning_content="".join(thinking_parts).strip(),
)
__all__ = ["GoogleGeminiChatModel"]