oclaw/svc/llm/transports/google_gemini_sse.py
oliver 420abac9f1 refactor: root-package imports (svc/runtime/interfaces) and fix PYTHONPATH
- Rename platform/ to svc/ to avoid shadowing stdlib platform.
- Replace from oclaw.* with from svc/runtime/interfaces; update -m CLI paths.
- tests/conftest: prepend repo root to sys.path (no parent-folder package name).
- CI: paths and offline_eval script under repo root.
- Ops scripts: PYTHONPATH must be repo root for python -m runtime.* (fixes gateway/WhatsApp sidecar startup).
- Fix default oclaw.json path in tabular/file attachment limits; stabilize attachment test config.

Co-authored-by: Cursor <cursoragent@cursor.com>
2026-05-13 14:51:17 +08:00

267 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 svc.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)
if on_token:
on_token(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"]