mirror of
https://github.com/x1xhlol/system-prompts-and-models-of-ai-tools.git
synced 2026-08-16 21:24:09 +00:00
488 lines
19 KiB
Python
488 lines
19 KiB
Python
"""
|
|
Gemini client — the Gemini side of llm.py's LLM_BACKEND switch.
|
|
|
|
Activated when LLM_BACKEND=gemini (auto-selected if GCP_API_KEY is set).
|
|
Routes Claude-shaped calls through the google-genai SDK with an API key. Exposes:
|
|
|
|
- GeminiBedrockClient: .invoke_model() drop-in for boto3 bedrock-runtime.
|
|
Translates Anthropic-style request body to Gemini and shapes the response
|
|
back to the Anthropic JSON the rest of the repo expects.
|
|
- GeminiChatModel: langchain ChatBedrock-style .invoke().
|
|
- GeminiConverseClient: converse-API wrapper for RLM / tearsheet loops.
|
|
- gemini_model_call / gemini_invoke_model: functional helpers used by llm.py.
|
|
"""
|
|
|
|
import base64
|
|
import json
|
|
import logging
|
|
import os
|
|
import uuid
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
try:
|
|
from utils.token_tracker import record_usage as _record_token_usage
|
|
except ImportError:
|
|
try:
|
|
from task.utils.token_tracker import record_usage as _record_token_usage
|
|
except ImportError:
|
|
def _record_token_usage(_model, _input_tokens, _output_tokens): # type: ignore
|
|
pass
|
|
|
|
|
|
GEMINI_MODEL_ID = os.environ.get("GEMINI_MODEL_ID", "gemini-3.5-flash")
|
|
GCP_API_KEY = os.environ.get("GCP_API_KEY") or os.environ.get("GEMINI_API_KEY", "")
|
|
|
|
|
|
_client_init_logged = False
|
|
_gemini_client_singleton = None
|
|
|
|
|
|
def _get_gemini_client():
|
|
"""Return a google-genai Client authenticated with GCP_API_KEY.
|
|
|
|
Lazily constructed and cached: the genai client (and the GCP_API_KEY
|
|
requirement) is only resolved on the FIRST LLM call, not at module import.
|
|
This keeps a missing GCP_API_KEY from crashing the whole worker at boot —
|
|
it surfaces only on the specific request that needs the model.
|
|
"""
|
|
global _gemini_client_singleton
|
|
if _gemini_client_singleton is not None:
|
|
return _gemini_client_singleton
|
|
from google import genai
|
|
if not GCP_API_KEY:
|
|
raise RuntimeError(
|
|
"GCP_API_KEY (or GEMINI_API_KEY) is not set. "
|
|
"Either set it or switch LLM_BACKEND."
|
|
)
|
|
global _client_init_logged
|
|
if not _client_init_logged:
|
|
_client_init_logged = True
|
|
logger.info("[gemini_client] init: model=%s (api-key auth)", GEMINI_MODEL_ID)
|
|
print(f"[gemini_client] init: model={GEMINI_MODEL_ID} (api-key auth)", flush=True)
|
|
_gemini_client_singleton = genai.Client(api_key=GCP_API_KEY)
|
|
return _gemini_client_singleton
|
|
|
|
|
|
def _to_gemini_contents(messages: list[dict]):
|
|
"""Convert Anthropic-style messages to google-genai Content list.
|
|
|
|
Anthropic role "assistant" maps to Gemini role "model". Image blocks
|
|
in Anthropic format ({"type": "image", "source": {"type": "base64",
|
|
"media_type": ..., "data": ...}}) become inline Parts.
|
|
"""
|
|
from google.genai import types
|
|
|
|
contents = []
|
|
for msg in messages:
|
|
role = "model" if msg.get("role") == "assistant" else "user"
|
|
content = msg.get("content")
|
|
parts = []
|
|
if isinstance(content, str):
|
|
parts.append(types.Part.from_text(text=content))
|
|
elif isinstance(content, list):
|
|
for block in content:
|
|
btype = block.get("type")
|
|
if btype == "text":
|
|
parts.append(types.Part.from_text(text=block.get("text", "")))
|
|
elif btype == "image":
|
|
src = block.get("source", {}) or {}
|
|
if src.get("type") == "base64":
|
|
parts.append(types.Part.from_bytes(
|
|
data=base64.b64decode(src.get("data", "")),
|
|
mime_type=src.get("media_type", "image/png"),
|
|
))
|
|
else:
|
|
# Tearsheet/Bedrock converse-style block — {"text": ...} or {"image": ...}.
|
|
if "text" in block:
|
|
parts.append(types.Part.from_text(text=block["text"]))
|
|
elif "image" in block:
|
|
img = block["image"]
|
|
fmt = img.get("format", "png")
|
|
data = img.get("source", {}).get("bytes", b"")
|
|
parts.append(types.Part.from_bytes(
|
|
data=data, mime_type=f"image/{fmt}",
|
|
))
|
|
else:
|
|
parts.append(types.Part.from_text(text=str(content)))
|
|
contents.append(types.Content(role=role, parts=parts))
|
|
return contents
|
|
|
|
|
|
def _build_generate_config(*, system=None, max_tokens=4096, temperature=0,
|
|
use_thinking=False, thinking_budget=None,
|
|
json_mode=False, response_schema=None,
|
|
use_search=False):
|
|
"""Assemble a GenerateContentConfig with optional thinking + system prompt.
|
|
|
|
json_mode=True forces structured JSON output via the Gemini decoder
|
|
(`response_mime_type="application/json"`) — the model cannot emit
|
|
prose / markdown / table rows even if the prompt is sloppy.
|
|
|
|
response_schema (optional) binds the output to a specific shape. Pass
|
|
a Python type (`list[str]`, a TypedDict, a pydantic BaseModel) or a
|
|
`types.Schema` dict. Implies json_mode. The decoder will refuse to
|
|
emit anything that doesn't conform.
|
|
|
|
use_search=True attaches Gemini's built-in Google Search grounding tool
|
|
so the model searches the live web before answering. OFF by default —
|
|
callers opt in per request (it adds latency + search cost and makes
|
|
output web-dependent). Verified on gemini-2.5-pro to coexist
|
|
with json_mode / response_schema and with thinking, so it composes with
|
|
the other flags here.
|
|
"""
|
|
from google.genai import types
|
|
|
|
cfg_kwargs = {
|
|
"temperature": temperature,
|
|
"max_output_tokens": max_tokens,
|
|
}
|
|
if system:
|
|
cfg_kwargs["system_instruction"] = system
|
|
if use_thinking or (thinking_budget and thinking_budget > 0):
|
|
budget = thinking_budget if (thinking_budget and thinking_budget > 0) else 5000
|
|
cfg_kwargs["thinking_config"] = types.ThinkingConfig(thinking_budget=budget)
|
|
if json_mode or response_schema is not None:
|
|
cfg_kwargs["response_mime_type"] = "application/json"
|
|
if response_schema is not None:
|
|
cfg_kwargs["response_schema"] = response_schema
|
|
if use_search:
|
|
# Gemini 2.0+ grounding uses GoogleSearch(); serializes to the same
|
|
# {"google_search": {}} tool the vendor-research REST path sends.
|
|
cfg_kwargs["tools"] = [types.Tool(google_search=types.GoogleSearch())]
|
|
return types.GenerateContentConfig(**cfg_kwargs)
|
|
|
|
|
|
def _extract_text(response) -> str:
|
|
"""Pull the concatenated text from a google-genai GenerateContentResponse."""
|
|
if getattr(response, "text", None):
|
|
return response.text
|
|
out = []
|
|
for cand in getattr(response, "candidates", []) or []:
|
|
content = getattr(cand, "content", None)
|
|
if not content:
|
|
continue
|
|
for part in getattr(content, "parts", []) or []:
|
|
if getattr(part, "text", None):
|
|
out.append(part.text)
|
|
return "".join(out)
|
|
|
|
|
|
def _extract_usage(response) -> tuple[int, int]:
|
|
"""Pull token counts from response.usage_metadata per Gemini API docs.
|
|
|
|
Documented fields (https://ai.google.dev/gemini-api/docs/tokens):
|
|
- prompt_token_count — input tokens
|
|
- candidates_token_count — output tokens (excluding thinking)
|
|
- thoughts_token_count — thinking tokens (billed as OUTPUT per
|
|
https://ai.google.dev/gemini-api/docs/pricing
|
|
"Output price (including thinking tokens)")
|
|
- cached_content_token_count, total_token_count — informational
|
|
|
|
Returns (input_tokens, output_tokens) using:
|
|
input = prompt_token_count
|
|
output = candidates_token_count + thoughts_token_count
|
|
"""
|
|
um = getattr(response, "usage_metadata", None)
|
|
if um is None:
|
|
return 0, 0
|
|
in_toks = int(getattr(um, "prompt_token_count", 0) or 0)
|
|
out_toks = (
|
|
int(getattr(um, "candidates_token_count", 0) or 0)
|
|
+ int(getattr(um, "thoughts_token_count", 0) or 0)
|
|
)
|
|
return in_toks, out_toks
|
|
|
|
|
|
def _record_usage(response, model_id: str) -> tuple[int, int]:
|
|
"""Manual token tracking — read documented fields, record once."""
|
|
in_toks, out_toks = _extract_usage(response)
|
|
if in_toks or out_toks:
|
|
_record_token_usage(model_id, in_toks, out_toks)
|
|
return in_toks, out_toks
|
|
|
|
|
|
def _stream_generate(client, *, model: str, contents, config):
|
|
"""Run generate_content_stream and return (full_text, last_chunk).
|
|
|
|
Streaming avoids the 10-minute timeout on long requests. The last chunk
|
|
carries the cumulative ``usage_metadata`` for token tracking.
|
|
"""
|
|
text_parts: list[str] = []
|
|
last_chunk = None
|
|
for chunk in client.models.generate_content_stream(
|
|
model=model, contents=contents, config=config,
|
|
):
|
|
last_chunk = chunk
|
|
chunk_text = _extract_text(chunk)
|
|
if chunk_text:
|
|
text_parts.append(chunk_text)
|
|
return "".join(text_parts), last_chunk
|
|
|
|
|
|
class _StreamBody:
|
|
"""Mimic boto3 StreamingBody so callers can do response['body'].read()."""
|
|
|
|
def __init__(self, data: bytes):
|
|
self._data = data
|
|
|
|
def read(self):
|
|
return self._data
|
|
|
|
|
|
class GeminiBedrockClient:
|
|
"""Drop-in replacement for boto3 bedrock-runtime client.
|
|
|
|
Translates an Anthropic-shaped request body to Gemini, then shapes the
|
|
response back to Anthropic JSON so the rest of the repo (which parses
|
|
response_body['content'][i]['text'] / 'usage') keeps working.
|
|
|
|
Token tracking — pure shim: returns the real ``usage`` in the response
|
|
and does NOT record itself. The caller
|
|
parses ``response_body['usage']`` and records via the documented
|
|
``record_usage`` API (``task/utils/token_tracker.py:218``). Token counts
|
|
are computed from the documented ``usage_metadata`` fields per the
|
|
Gemini docs (``_extract_usage``).
|
|
"""
|
|
|
|
def __init__(self):
|
|
# Client built lazily on first invoke_model() so a missing GCP_API_KEY
|
|
# does not crash the worker at import/construction time.
|
|
self._model = GEMINI_MODEL_ID
|
|
|
|
@property
|
|
def _client(self):
|
|
return _get_gemini_client()
|
|
|
|
def invoke_model(self, modelId: str, body: str, **_kwargs):
|
|
logger.info("[gemini_client.GeminiBedrockClient.invoke_model] model=%s (caller passed modelId=%s)", self._model, modelId)
|
|
request = json.loads(body) if isinstance(body, str) else body
|
|
messages = request.get("messages", [])
|
|
max_tokens = request.get("max_tokens", 16384)
|
|
temperature = request.get("temperature", 0)
|
|
system = request.get("system")
|
|
thinking = request.get("thinking") or {}
|
|
|
|
use_thinking = thinking.get("type") == "enabled"
|
|
budget = thinking.get("budget_tokens") if use_thinking else None
|
|
# Optional Google Search grounding — set "use_search": true in the
|
|
# request body (non-Anthropic field; absent → off, behavior unchanged).
|
|
use_search = bool(request.get("use_search", False))
|
|
# Unlike Anthropic, Gemini accepts thinking WITH temperature=0 — keep
|
|
# the caller's temperature (default 0) so thinking calls stay
|
|
# deterministic. Never force temperature=1 here.
|
|
|
|
contents = _to_gemini_contents(messages)
|
|
config = _build_generate_config(
|
|
system=system, max_tokens=max_tokens, temperature=temperature,
|
|
use_thinking=use_thinking, thinking_budget=budget,
|
|
use_search=use_search,
|
|
)
|
|
|
|
text, last_chunk = _stream_generate(
|
|
self._client, model=self._model, contents=contents, config=config,
|
|
)
|
|
|
|
in_toks, out_toks = _extract_usage(last_chunk) if last_chunk else (0, 0)
|
|
|
|
response_dict = {
|
|
"id": f"gemini-{uuid.uuid4().hex[:24]}",
|
|
"type": "message",
|
|
"role": "assistant",
|
|
"content": [{"type": "text", "text": text}],
|
|
"model": self._model,
|
|
"stop_reason": "end_turn",
|
|
"usage": {
|
|
"input_tokens": in_toks,
|
|
"output_tokens": out_toks,
|
|
},
|
|
}
|
|
return {"body": _StreamBody(json.dumps(response_dict).encode("utf-8"))}
|
|
|
|
|
|
class _InvokeResult:
|
|
def __init__(self, content: str):
|
|
self.content = content
|
|
|
|
|
|
class GeminiChatModel:
|
|
"""Drop-in replacement for LangChain ChatBedrock.
|
|
|
|
The ``model_id`` argument is accepted for signature compatibility with
|
|
callers that hardcode a Bedrock id (e.g. ``us.anthropic.claude-sonnet-4-6``)
|
|
but is **ignored** — under ``LLM_BACKEND=gemini`` we always call
|
|
``GEMINI_MODEL_ID``.
|
|
"""
|
|
|
|
def __init__(self, model_id: str = None, model_kwargs: dict = None):
|
|
del model_id # accepted for signature compat; we always use GEMINI_MODEL_ID
|
|
# Client built lazily on first invoke() so a missing GCP_API_KEY does
|
|
# not crash the worker at import/construction time.
|
|
self._model = GEMINI_MODEL_ID
|
|
self._kwargs = model_kwargs or {}
|
|
|
|
@property
|
|
def _client(self):
|
|
return _get_gemini_client()
|
|
|
|
def invoke(self, prompt):
|
|
logger.info("[gemini_client.GeminiChatModel.invoke] model=%s", self._model)
|
|
if isinstance(prompt, str):
|
|
messages = [{"role": "user", "content": prompt}]
|
|
else:
|
|
messages = prompt
|
|
|
|
config = _build_generate_config(
|
|
max_tokens=self._kwargs.get("max_tokens", 16384),
|
|
temperature=self._kwargs.get("temperature", 0),
|
|
use_search=bool(self._kwargs.get("use_search", False)),
|
|
)
|
|
text, last_chunk = _stream_generate(
|
|
self._client, model=self._model,
|
|
contents=_to_gemini_contents(messages), config=config,
|
|
)
|
|
if last_chunk is not None:
|
|
_record_usage(last_chunk, self._model)
|
|
return _InvokeResult(text)
|
|
|
|
|
|
class GeminiConverseClient:
|
|
"""Drop-in replacement for pdf_to_xlsx / tearsheets BedrockClient.
|
|
|
|
The ``model_id`` and ``region`` arguments are accepted for signature
|
|
compatibility (RLMBedrockClient callers hardcode Bedrock model ids and
|
|
AWS regions) but are **ignored** — under ``LLM_BACKEND=gemini`` we always
|
|
call ``GEMINI_MODEL_ID``.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
model_id: str = None,
|
|
region: str = None,
|
|
max_tokens: int = 16384,
|
|
thinking_budget: int = 4096,
|
|
use_search: bool = False,
|
|
):
|
|
del model_id, region # accepted for signature compat; always GEMINI_MODEL_ID
|
|
self.model_id = GEMINI_MODEL_ID
|
|
self.max_tokens = max_tokens
|
|
self.thinking_budget = thinking_budget if (thinking_budget and thinking_budget > 0) else 0
|
|
self.use_search = use_search
|
|
# Client built lazily on first completion() so a missing GCP_API_KEY
|
|
# does not crash the worker at import/construction time.
|
|
self.total_input_tokens = 0
|
|
self.total_output_tokens = 0
|
|
self.call_count = 0
|
|
|
|
@property
|
|
def _client(self):
|
|
return _get_gemini_client()
|
|
|
|
def _generate(self, messages: list[dict], system: str | None) -> str:
|
|
logger.info(
|
|
"[gemini_client.GeminiConverseClient] model=%s thinking=%s call#%d",
|
|
self.model_id, bool(self.thinking_budget), self.call_count + 1,
|
|
)
|
|
contents = _to_gemini_contents(messages)
|
|
config = _build_generate_config(
|
|
system=system,
|
|
max_tokens=self.max_tokens,
|
|
temperature=0, # Gemini allows thinking at temp 0 — deterministic
|
|
use_thinking=bool(self.thinking_budget),
|
|
thinking_budget=self.thinking_budget or None,
|
|
use_search=self.use_search,
|
|
)
|
|
text, last_chunk = _stream_generate(
|
|
self._client, model=self.model_id, contents=contents, config=config,
|
|
)
|
|
if last_chunk is not None:
|
|
in_toks, out_toks = _record_usage(last_chunk, self.model_id)
|
|
self.total_input_tokens += in_toks
|
|
self.total_output_tokens += out_toks
|
|
self.call_count += 1
|
|
return text
|
|
|
|
def completion(self, messages: list[dict], system: str | None = None) -> str:
|
|
return self._generate(messages, system)
|
|
|
|
def completion_with_image(
|
|
self,
|
|
image_bytes: bytes,
|
|
mime_type: str,
|
|
prompt: str,
|
|
system: str | None = None,
|
|
) -> str:
|
|
media_type = mime_type if "/" in mime_type else f"image/{mime_type}"
|
|
b64 = base64.standard_b64encode(image_bytes).decode("utf-8")
|
|
messages = [{
|
|
"role": "user",
|
|
"content": [
|
|
{"type": "image", "source": {
|
|
"type": "base64", "media_type": media_type, "data": b64,
|
|
}},
|
|
{"type": "text", "text": prompt},
|
|
],
|
|
}]
|
|
return self._generate(messages, system)
|
|
|
|
def usage_summary(self) -> str:
|
|
return (
|
|
f"Calls: {self.call_count} | "
|
|
f"Input: {self.total_input_tokens:,} | "
|
|
f"Output: {self.total_output_tokens:,}"
|
|
)
|
|
|
|
|
|
def gemini_model_call(prompt, max_tokens=16384, use_thinking=False, system=None,
|
|
json_mode=False, response_schema=None, use_search=False):
|
|
"""Single-shot Gemini call. Returns the text response.
|
|
|
|
use_search=True grounds the call with Google Search (off by default)."""
|
|
logger.info("[gemini_client.gemini_model_call] model=%s thinking=%s json_mode=%s schema=%s search=%s", GEMINI_MODEL_ID, use_thinking, json_mode, response_schema is not None, use_search)
|
|
client = _get_gemini_client()
|
|
config = _build_generate_config(
|
|
system=system, max_tokens=max_tokens,
|
|
temperature=0, # Gemini allows thinking at temp 0 — deterministic
|
|
use_thinking=use_thinking,
|
|
json_mode=json_mode,
|
|
response_schema=response_schema,
|
|
use_search=use_search,
|
|
)
|
|
contents = _to_gemini_contents([{"role": "user", "content": prompt}])
|
|
try:
|
|
text, last_chunk = _stream_generate(
|
|
client, model=GEMINI_MODEL_ID, contents=contents, config=config,
|
|
)
|
|
if last_chunk is not None:
|
|
_record_usage(last_chunk, GEMINI_MODEL_ID)
|
|
return text
|
|
except Exception as e:
|
|
logger.error(f"Gemini model call failed: {e}")
|
|
raise
|
|
|
|
|
|
def gemini_invoke_model(messages, system=None, max_tokens=16384, temperature=0,
|
|
use_thinking=False, use_search=False):
|
|
"""Multi-turn Gemini call. Returns the text response.
|
|
|
|
use_search=True grounds the call with Google Search (off by default)."""
|
|
logger.info("[gemini_client.gemini_invoke_model] model=%s msgs=%d thinking=%s search=%s", GEMINI_MODEL_ID, len(messages), use_thinking, use_search)
|
|
client = _get_gemini_client()
|
|
config = _build_generate_config(
|
|
system=system, max_tokens=max_tokens,
|
|
temperature=temperature, # Gemini allows thinking at temp 0 — no forcing
|
|
use_thinking=use_thinking,
|
|
use_search=use_search,
|
|
)
|
|
try:
|
|
text, last_chunk = _stream_generate(
|
|
client, model=GEMINI_MODEL_ID,
|
|
contents=_to_gemini_contents(messages), config=config,
|
|
)
|
|
if last_chunk is not None:
|
|
_record_usage(last_chunk, GEMINI_MODEL_ID)
|
|
return text
|
|
except Exception as e:
|
|
logger.error(f"Gemini invoke_model failed: {e}")
|
|
raise
|