diff --git a/examples/google_adk_governance.py b/examples/google_adk_governance.py new file mode 100644 index 0000000..64d5637 --- /dev/null +++ b/examples/google_adk_governance.py @@ -0,0 +1,53 @@ +"""Example: TealTiger governance callbacks for Google ADK. +Calls before_tool / after_tool with mocks — no API key or google-adk required. +Run: + python examples/google_adk_governance.py +""" + +from tealtiger.integrations import TealTigerCallback + + +def main(): + governance = TealTigerCallback( + policies=[ + {"type": "tool_allowlist", "allowed": ["search", "lookup*"]}, + {"type": "pii_block", "categories": ["ssn", "email"]}, + {"type": "cost_limit", "max_per_session": 0.05}, + ], + mode="ENFORCE", + model="gemini-2.5-flash", + cost_per_tool_call=0.01, + agent_id="demo-adk-agent", + ) + print("=== 1) ALLOW: search ===") + result = governance.before_tool(None, "search", {"query": "weather"}) + print("block result:", result) # None = allowed + governance.after_tool(None, "search", {"query": "weather"}, result="ok") + print("cost:", governance.total_cost) + print("decision:", governance.decisions[-1]["action"], governance.decisions[-1]["reason_codes"]) + print("\n=== 2) DENY: tool not allowlisted ===") + result = governance.before_tool(None, "delete_all", {}) + print("block result:", result) # dict with content + print("deny_count:", governance.deny_count) + print("\n=== 3) DENY: PII in args ===") + result = governance.before_tool( + None, + "search", + {"note": "ssn 123-45-6789"}, + ) + print("block result:", result) + print("\n=== 4) freeze / unfreeze ===") + governance.freeze() + result = governance.before_tool(None, "search", {"query": "x"}) + print("frozen block:", result) # dict — blocked while frozen + governance.unfreeze() + result = governance.before_tool(None, "search", {"query": "x"}) + print("after unfreeze:", result) # None — allowed again + print("\n=== summary ===") + print("total_cost:", governance.total_cost) + print("decisions:", len(governance.decisions)) + print("denies:", governance.deny_count) + + +if __name__ == "__main__": + main() diff --git a/src/tealtiger/cost/pricing.py b/src/tealtiger/cost/pricing.py index dce4342..cf005b2 100644 --- a/src/tealtiger/cost/pricing.py +++ b/src/tealtiger/cost/pricing.py @@ -107,6 +107,7 @@ last_updated="2026-01-31", ), # Google PaLM/Gemini Models + # Google PaLM/Gemini Models from: https://ai.google.dev/gemini-api/docs/pricing?hl=en "gemini-pro": ModelPricing( model="gemini-pro", provider="google", @@ -143,6 +144,34 @@ output_cost_per_1k=0.0005, last_updated="2026-01-31", ), + "gemini-2.5-flash": ModelPricing( + model="gemini-2.5-flash", + provider="google", + input_cost_per_1k=0.0003, + output_cost_per_1k=0.0025, + last_updated="2026-03-06", + ), + "gemini-2.5-pro": ModelPricing( + model="gemini-2.5-pro", + provider="google", + input_cost_per_1k=0.00125, + output_cost_per_1k=0.01, + last_updated="2026-03-06", + ), + "gemini-3.5-flash": ModelPricing( + model="gemini-3.5-flash", + provider="google", + input_cost_per_1k=0.0015, + output_cost_per_1k=0.009, + last_updated="2026-03-06", + ), + "gemini-3.6-flash": ModelPricing( + model="gemini-3.6-flash", + provider="google", + input_cost_per_1k=0.0015, + output_cost_per_1k=0.0075, + last_updated="2026-03-06", + ), # Cohere Models "command": ModelPricing( model="command", @@ -425,5 +454,5 @@ def get_supported_providers() -> List[ModelProvider]: Returns: List of provider names """ - providers = set(p.provider for p in MODEL_PRICING.values()) + providers = {p.provider for p in MODEL_PRICING.values()} return list(providers) diff --git a/src/tealtiger/integrations/__init__.py b/src/tealtiger/integrations/__init__.py index 7013b2e..6aca697 100644 --- a/src/tealtiger/integrations/__init__.py +++ b/src/tealtiger/integrations/__init__.py @@ -1,13 +1,14 @@ """TealTiger integrations with external observability and monitoring platforms.""" -from tealtiger.integrations.langfuse import LangfuseGovernanceExporter from tealtiger.integrations.agentops import AgentOpsGovernanceReporter +from tealtiger.integrations.google_adk import TealTigerCallback +from tealtiger.integrations.langfuse import LangfuseGovernanceExporter from tealtiger.integrations.opik import ( - GovernanceAccuracyMetric, - PIIDetectionMetric, FalsePositiveRateMetric, + GovernanceAccuracyMetric, GovernanceLatencyMetric, GovernanceMultiMetric, + PIIDetectionMetric, ) __all__ = [ @@ -18,4 +19,5 @@ "FalsePositiveRateMetric", "GovernanceLatencyMetric", "GovernanceMultiMetric", + "TealTigerCallback", ] diff --git a/src/tealtiger/integrations/google_adk.py b/src/tealtiger/integrations/google_adk.py index 900a2a7..60a1d2d 100644 --- a/src/tealtiger/integrations/google_adk.py +++ b/src/tealtiger/integrations/google_adk.py @@ -17,7 +17,7 @@ ) agent = Agent( - model="gemini-2.0-flash", + model="gemini-3.6-flash", tools=[search_tool, code_tool], before_tool_callback=governance.before_tool, after_tool_callback=governance.after_tool, @@ -27,10 +27,12 @@ from __future__ import annotations import re -import uuid import time +import uuid from typing import Any, Dict, List +from tealtiger.cost.pricing import get_model_pricing + # PII patterns _PII_PATTERNS = { "ssn": re.compile(r"\b\d{3}-\d{2}-\d{4}\b"), @@ -58,6 +60,8 @@ class TealTigerCallback: mode: "OBSERVE", "MONITOR", or "ENFORCE". agent_id: Agent identifier for audit correlation. on_decision: Optional callback invoked with each governance decision. + model: Gemini 3.6 flash is default model. + cost_per_tool_call: Fallback USD cost when pricing/tokens is not available for the model. """ def __init__( @@ -66,16 +70,35 @@ def __init__( mode: str = "OBSERVE", agent_id: str = None, on_decision=None, + model: str = "gemini-3.6-flash", + cost_per_tool_call: float = 0.0015, ): self.policies = policies or [] self.mode = mode.upper() self.agent_id = agent_id or f"adk-agent-{str(uuid.uuid4())[:8]}" self.on_decision = on_decision + self.model = model + self.cost_per_tool_call = cost_per_tool_call self._decisions: List[Dict[str, Any]] = [] self._cumulative_cost: float = 0.0 self._frozen: bool = False - def before_tool(self, callback_context, tool, args, tool_context=None): + def _estimate_tool_cost(self) -> float: + """Estimate USD cost for one allowed tool call. + Used model pricing if available; otherwise cost_per_tool_call. + """ + pricing = get_model_pricing(self.model, provider="google") + if pricing is None: + return self.cost_per_tool_call + + estimated_input_tokens = 500 + estimated_output_tokens = 500 + + input_cost = (estimated_input_tokens / 1000) * pricing.input_cost_per_1k + output_cost = (estimated_output_tokens / 1000) * pricing.output_cost_per_1k + return input_cost + output_cost + + def before_tool(self, callback_context, tool, args, tool_context=None): # noqa: C901 """Before-tool callback for Google ADK. Evaluates governance policies before tool execution. @@ -157,6 +180,10 @@ def before_tool(self, callback_context, tool, args, tool_context=None): break eval_time = (time.perf_counter() - start_time) * 1000 + # Track cost for allowed actions + cost = self._estimate_tool_cost() if action == "ALLOW" else 0.0 + if action == "ALLOW": + self._cumulative_cost += cost # Record decision decision = { @@ -169,7 +196,7 @@ def before_tool(self, callback_context, tool, args, tool_context=None): "reason_codes": reason_codes or (["POLICY_ALLOW"] if action == "ALLOW" else []), "risk_score": risk_score, "evaluation_time_ms": eval_time, - "cost_tracked": 0.002 if action == "ALLOW" else 0.0, + "cost_tracked": cost, "cumulative_cost": self._cumulative_cost, } self._decisions.append(decision) @@ -177,10 +204,6 @@ def before_tool(self, callback_context, tool, args, tool_context=None): if self.on_decision: self.on_decision(decision) - # Track cost for allowed actions - if action == "ALLOW": - self._cumulative_cost += 0.002 - # Mode-based behavior if self.mode == "ENFORCE" and action == "DENY": # Return a dict to block execution (ADK pattern) diff --git a/tests/cost/test_pricing.py b/tests/cost/test_pricing.py index 5747d0f..21783ff 100644 --- a/tests/cost/test_pricing.py +++ b/tests/cost/test_pricing.py @@ -20,6 +20,12 @@ # ("open-mistral-nemo", 0.00015, 0.00015), # ("codestral-latest", 0.0003, 0.0009), ] +GEMINI_PRICING = [ + ("gemini-2.5-flash", 0.0003, 0.0025), + ("gemini-2.5-pro", 0.00125, 0.01), + ("gemini-3.5-flash", 0.0015, 0.009), + ("gemini-3.6-flash", 0.0015, 0.0075), +] @pytest.mark.parametrize("model,input_rate,output_rate", MISTRAL_PRICING) @@ -40,3 +46,22 @@ def test_mistral_large_latest_cost_calculation(): assert estimate.breakdown.input_cost == pytest.approx(0.0005) assert estimate.breakdown.output_cost == pytest.approx(0.0015) assert estimate.estimated_cost == pytest.approx(0.002) + + +@pytest.mark.parametrize("model,input_rate,output_rate", GEMINI_PRICING) +def test_gemini_pricing_lookup(model, input_rate, output_rate): + pricing = get_model_pricing(model) + assert pricing is not None + assert pricing.provider == "google" + assert pricing.input_cost_per_1k == input_rate + assert pricing.output_cost_per_1k == output_rate + + +def test_gemini_2_5_pro_cost_calculation(): + """Cost for 1000 in + 1000 out on gemini-2.5-pro.""" + tracker = CostTracker() + tokens = TokenUsage(input_tokens=1000, output_tokens=1000, total_tokens=2000) + estimate = tracker.estimate_cost("gemini-2.5-pro", tokens) + assert estimate.breakdown.input_cost == pytest.approx(0.00125) + assert estimate.breakdown.output_cost == pytest.approx(0.01) + assert estimate.estimated_cost == pytest.approx(0.01125) diff --git a/tests/test_google_adk_integration.py b/tests/test_google_adk_integration.py new file mode 100644 index 0000000..beacadd --- /dev/null +++ b/tests/test_google_adk_integration.py @@ -0,0 +1,46 @@ +import pytest + +from tealtiger.integrations.google_adk import TealTigerCallback + + +def test_allow_tracks_cost_fallback(): + """Fake model -> uses cost_per_tool_call (0.01).""" + g = TealTigerCallback( + model="not-a-real-model-xyz", + cost_per_tool_call=0.01, + mode="ENFORCE", + ) + result = g.before_tool(None, "search", {}) + assert result is None + assert g.total_cost == pytest.approx(0.01) + assert g.decisions[0]["action"] == "ALLOW" + assert g.decisions[0]["cumulative_cost"] == pytest.approx(0.01) + + +def test_allow_uses_model_pricing(): + """# Priced model -> (500/1000)*0.0015 + (500/1000)*0.0075 = 0.0045.""" + g = TealTigerCallback(model="gemini-3.6-flash", mode="ENFORCE") + result = g.before_tool(None, "search", {}) + assert result is None + assert g.total_cost == pytest.approx(0.0045) + + +def test_allowlist_deny_blocks_in_enforce_mode(): + g = TealTigerCallback( + policies=[{"type": "tool_allowlist", "allowed": ["search"]}], + mode="ENFORCE", + ) + result = g.before_tool(None, "delete_all", {}) + assert result is not None + assert "content" in result + assert g.decisions[0]["action"] == "DENY" + assert g.total_cost == 0.0 + + +def test_freeze_deni(): + g = TealTigerCallback(mode="ENFORCE") + g.freeze() + result = g.before_tool(None, "search", {}) + assert result is not None + assert "AGENT_FROZEN" in g.decisions[0]["reason_codes"] + assert g.total_cost == 0.0