mirror of
https://github.com/hansjone/oclaw.git
synced 2026-10-11 06:53:22 +08:00
oliver
398d85ba06
- Forward reasoning_content and non-stream reasoning to on_token in OpenAI chat completions; Gemini thought parts to on_token. - Stream bubble: always render session.tool panels; format tool_use_call payloads and title as call · tool_name. - Emit tool_use_call via on_tool_ui before tool execution so WS shows invoke cards before tool_use_result. Co-authored-by: Cursor <cursoragent@cursor.com>
267 lines
11 KiB
Python
267 lines
11 KiB
Python
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)
|
||
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"]
|
||
|