From 91168222739435d40ddfd9a5d5769735878e0eb8 Mon Sep 17 00:00:00 2001 From: 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 Date: Thu, 6 Aug 2026 19:01:16 +0300 Subject: [PATCH 1/5] added some tests and polishing of google adk integration --- src/tealtiger/cost/pricing.py | 29 ++++++++++++++++++++++++ src/tealtiger/integrations/__init__.py | 2 ++ src/tealtiger/integrations/google_adk.py | 19 +++++++++++----- tests/test_google_adk_integration.py | 29 ++++++++++++++++++++++++ 4 files changed, 73 insertions(+), 6 deletions(-) create mode 100644 tests/test_google_adk_integration.py diff --git a/src/tealtiger/cost/pricing.py b/src/tealtiger/cost/pricing.py index dce4342..e27eb8c 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", diff --git a/src/tealtiger/integrations/__init__.py b/src/tealtiger/integrations/__init__.py index 7013b2e..f7aad35 100644 --- a/src/tealtiger/integrations/__init__.py +++ b/src/tealtiger/integrations/__init__.py @@ -2,6 +2,7 @@ from tealtiger.integrations.langfuse import LangfuseGovernanceExporter from tealtiger.integrations.agentops import AgentOpsGovernanceReporter +from tealtiger.integrations.google_adk import TealTigerCallback from tealtiger.integrations.opik import ( GovernanceAccuracyMetric, PIIDetectionMetric, @@ -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..5fa95e0 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, @@ -58,6 +58,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,11 +68,15 @@ 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 @@ -157,6 +163,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.cost_per_tool_call if action == "ALLOW" else 0.0 + if action == "ALLOW": + self._cumulative_cost += cost # Record decision decision = { @@ -169,7 +179,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 +187,7 @@ 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/test_google_adk_integration.py b/tests/test_google_adk_integration.py new file mode 100644 index 0000000..5569f7d --- /dev/null +++ b/tests/test_google_adk_integration.py @@ -0,0 +1,29 @@ +import pytest +from tealtiger.integrations.google_adk import TealTigerCallback + +def test_allow_tracks_cost(): + g = TealTigerCallback(cost_per_tool_call=0.01, mode="ENFORCE") + result = g.before_tool(None, "search",{}) + assert result is not None + assert g.total_cost == 0.01 + assert g.decisions[0]["action"] == "ALLOW" + assert g.decisions[0]["cumulative_cost"] == 0.01 + +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 \ No newline at end of file From aa01d60fd74e5b7f63864b2ba2237ca9fd4c3f5f Mon Sep 17 00:00:00 2001 From: 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 Date: Fri, 7 Aug 2026 14:42:40 +0300 Subject: [PATCH 2/5] feat: add Google ADK example and integration tests --- examples/google_adk_governance.py | 52 ++++++++++++++++++++++++++++ tests/test_google_adk_integration.py | 2 +- 2 files changed, 53 insertions(+), 1 deletion(-) create mode 100644 examples/google_adk_governance.py diff --git a/examples/google_adk_governance.py b/examples/google_adk_governance.py new file mode 100644 index 0000000..e0922ae --- /dev/null +++ b/examples/google_adk_governance.py @@ -0,0 +1,52 @@ +"""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 ===") + 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() \ No newline at end of file diff --git a/tests/test_google_adk_integration.py b/tests/test_google_adk_integration.py index 5569f7d..dba731e 100644 --- a/tests/test_google_adk_integration.py +++ b/tests/test_google_adk_integration.py @@ -4,7 +4,7 @@ def test_allow_tracks_cost(): g = TealTigerCallback(cost_per_tool_call=0.01, mode="ENFORCE") result = g.before_tool(None, "search",{}) - assert result is not None + assert result is None assert g.total_cost == 0.01 assert g.decisions[0]["action"] == "ALLOW" assert g.decisions[0]["cumulative_cost"] == 0.01 From 8a87a1b36a2e0277cb43c8814b5716fec15052b6 Mon Sep 17 00:00:00 2001 From: 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 Date: Fri, 7 Aug 2026 14:53:24 +0300 Subject: [PATCH 3/5] test: add Gemini pricing lookup coverage --- tests/cost/test_pricing.py | 25 +++++++++++++++++++++++++ 1 file changed, 25 insertions(+) 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) From c85832a205ad18c30ad751f601d9a485a0bfb189 Mon Sep 17 00:00:00 2001 From: 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 Date: Sat, 8 Aug 2026 15:11:19 +0300 Subject: [PATCH 4/5] feat: wire Gemini pricing into Google ADK cost estimates --- examples/google_adk_governance.py | 1 - src/tealtiger/integrations/google_adk.py | 18 +++++++++++++++++- tests/test_google_adk_integration.py | 22 +++++++++++++++++----- 3 files changed, 34 insertions(+), 7 deletions(-) diff --git a/examples/google_adk_governance.py b/examples/google_adk_governance.py index e0922ae..e57e99d 100644 --- a/examples/google_adk_governance.py +++ b/examples/google_adk_governance.py @@ -37,7 +37,6 @@ def main(): ) print("block result:", result) print("\n=== 4) freeze / unfreeze ===") - print("\n=== 4) freeze / unfreeze ===") governance.freeze() result = governance.before_tool(None, "search", {"query": "x"}) print("frozen block:", result) # dict — blocked while frozen diff --git a/src/tealtiger/integrations/google_adk.py b/src/tealtiger/integrations/google_adk.py index 5fa95e0..902ef8a 100644 --- a/src/tealtiger/integrations/google_adk.py +++ b/src/tealtiger/integrations/google_adk.py @@ -30,6 +30,7 @@ import uuid import time from typing import Any, Dict, List +from tealtiger.cost.pricing import get_model_pricing # PII patterns _PII_PATTERNS = { @@ -81,6 +82,21 @@ def __init__( self._cumulative_cost: float = 0.0 self._frozen: bool = False + 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): """Before-tool callback for Google ADK. @@ -164,7 +180,7 @@ def before_tool(self, callback_context, tool, args, tool_context=None): eval_time = (time.perf_counter() - start_time) * 1000 # Track cost for allowed actions - cost = self.cost_per_tool_call if action == "ALLOW" else 0.0 + cost = self._estimate_tool_cost() if action == "ALLOW" else 0.0 if action == "ALLOW": self._cumulative_cost += cost diff --git a/tests/test_google_adk_integration.py b/tests/test_google_adk_integration.py index dba731e..2e6aad7 100644 --- a/tests/test_google_adk_integration.py +++ b/tests/test_google_adk_integration.py @@ -1,13 +1,25 @@ import pytest from tealtiger.integrations.google_adk import TealTigerCallback -def test_allow_tracks_cost(): - g = TealTigerCallback(cost_per_tool_call=0.01, mode="ENFORCE") - result = g.before_tool(None, "search",{}) +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 == 0.01 + assert g.total_cost == pytest.approx(0.01) assert g.decisions[0]["action"] == "ALLOW" - assert g.decisions[0]["cumulative_cost"] == 0.01 + 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( From c93eac23a8a65b69182e2697ba6cac521e4c7ea7 Mon Sep 17 00:00:00 2001 From: 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 Date: Sat, 8 Aug 2026 15:43:44 +0300 Subject: [PATCH 5/5] feature: polish governance callbacks for google agent development kit --- examples/google_adk_governance.py | 4 +++- src/tealtiger/cost/pricing.py | 2 +- src/tealtiger/integrations/__init__.py | 6 +++--- src/tealtiger/integrations/google_adk.py | 12 ++++++------ tests/test_google_adk_integration.py | 13 +++++++++---- 5 files changed, 22 insertions(+), 15 deletions(-) diff --git a/examples/google_adk_governance.py b/examples/google_adk_governance.py index e57e99d..64d5637 100644 --- a/examples/google_adk_governance.py +++ b/examples/google_adk_governance.py @@ -47,5 +47,7 @@ def main(): print("total_cost:", governance.total_cost) print("decisions:", len(governance.decisions)) print("denies:", governance.deny_count) + + if __name__ == "__main__": - main() \ No newline at end of file + main() diff --git a/src/tealtiger/cost/pricing.py b/src/tealtiger/cost/pricing.py index e27eb8c..cf005b2 100644 --- a/src/tealtiger/cost/pricing.py +++ b/src/tealtiger/cost/pricing.py @@ -454,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 f7aad35..6aca697 100644 --- a/src/tealtiger/integrations/__init__.py +++ b/src/tealtiger/integrations/__init__.py @@ -1,14 +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__ = [ diff --git a/src/tealtiger/integrations/google_adk.py b/src/tealtiger/integrations/google_adk.py index 902ef8a..60a1d2d 100644 --- a/src/tealtiger/integrations/google_adk.py +++ b/src/tealtiger/integrations/google_adk.py @@ -27,9 +27,10 @@ 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 @@ -89,15 +90,15 @@ def _estimate_tool_cost(self) -> float: 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 + 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): + 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. @@ -203,7 +204,6 @@ def before_tool(self, callback_context, tool, args, tool_context=None): if self.on_decision: self.on_decision(decision) - # Mode-based behavior if self.mode == "ENFORCE" and action == "DENY": # Return a dict to block execution (ADK pattern) diff --git a/tests/test_google_adk_integration.py b/tests/test_google_adk_integration.py index 2e6aad7..beacadd 100644 --- a/tests/test_google_adk_integration.py +++ b/tests/test_google_adk_integration.py @@ -1,6 +1,8 @@ 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( @@ -14,6 +16,7 @@ def test_allow_tracks_cost_fallback(): 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") @@ -21,21 +24,23 @@ def test_allow_uses_model_pricing(): 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"]}], + policies=[{"type": "tool_allowlist", "allowed": ["search"]}], mode="ENFORCE", ) - result = g.before_tool(None, "delete_all",{}) + 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", {}) + 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 \ No newline at end of file + assert g.total_cost == 0.0