Feat/google adk polish

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()

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)

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",
 ]

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,18 +196,14 @@ 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)
 
         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)

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)

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