Gemini
Track the google-genai SDK on Vertex AI or AI Studio.
Guard supports the raw google-genai SDK in sync and async code, streaming and non-streaming, on both Vertex AI and AI Studio.
Install
pip install "impraint-guard[gemini]"This installs google-genai (1.53.0 or newer) and langchain-google-genai (4.0.0 or newer, which Vertex AI needs when used through LangChain).
Option 1: patch the SDK (process-wide)
Call patch_gemini(guard) once at startup, before any Gemini call. Every google-genai client in the process is tracked.
import os
from google import genai
from impraint_guard import Guard
from impraint_guard.integrations.gemini import patch_gemini
guard = Guard(api_key=os.environ["IMPRAINT_API_KEY"], agent_id="gemini-agent")
patch_gemini(guard)
client = genai.Client(api_key=os.environ["GEMINI_API_KEY"])
@guard.track
def ask(question: str) -> str:
resp = client.models.generate_content(
model="gemini-2.5-flash",
contents=question,
)
return resp.textPatching is process-wide
Calling patch_gemini again with the same Guard does nothing. If you create several Guard
instances in one process, only the first one that patches receives the Gemini data.
Option 2: wrap a single client
client = genai.Client(api_key=os.environ["GEMINI_API_KEY"])
client = guard.wrap_gemini(client)Authentication modes
client = genai.Client(api_key=os.environ["GEMINI_API_KEY"])What Guard records
- Model name, input tokens and output tokens
- Gemini 2.5 thinking tokens, which are included in output tokens and cost
- Cached input tokens, when a context cache is used
- Cost, using Guard's pricing table (see Tokens and cost)
Only inside @guard.track
Gemini calls are recorded as steps of the trace that is currently running. Outside a function
decorated with @guard.track, nothing is recorded.
Using LangChain too?
If a call goes through LangChain's ChatGoogleGenerativeAI while the LangChain handler is
recording, Guard skips the raw patch for that call so tokens and cost are not counted twice.
The SDK reports the model as models/gemini-2.5-flash. Guard removes the models/ prefix, so it shows as gemini-2.5-flash and matches your pricing table.
Prompt checks
By default, Guard also checks the text going into the model and reports what it finds (on_input="report"). To change this, pass on_input to patch_gemini or wrap_gemini:
patch_gemini(guard, on_input="block") # stop the call if the prompt is found unsafe
patch_gemini(guard, on_input="off") # no prompt checksSee Content safety for what is checked and what each mode does.