Gemma on Amazon Bedrock — the super agent, streaming, clean logs

from pathlib import Path
import sys

ROOT = (
    Path.cwd().resolve().parent
    if Path.cwd().name == "notebooks"
    else Path.cwd().resolve()
)
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

from shipit_agent.llms import BedrockChatLLM, BedrockGemmaChatLLM

# Gemma 4 → routed to the OpenAI-compatible bedrock-mantle endpoint.
gemma = BedrockChatLLM(model="google.gemma-4-31b", region="us-east-1")
print("delegate type :", type(gemma._mantle_delegate).__name__)
print("delegate model:", gemma._mantle_delegate.model)
print("base_url      :", gemma._mantle_delegate.client_kwargs["base_url"])

assert isinstance(gemma._mantle_delegate, BedrockGemmaChatLLM)
assert (
    gemma._mantle_delegate.client_kwargs["base_url"]
    == "https://bedrock-mantle.us-east-1.api.aws/openai/v1"
)

claude = BedrockChatLLM(model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0")
gemma3 = BedrockChatLLM(model="bedrock/google.gemma-3-27b-it")
print("claude  ._mantle_delegate:", claude._mantle_delegate)
print("gemma-3 ._mantle_delegate:", gemma3._mantle_delegate)
assert claude._mantle_delegate is None
assert gemma3._mantle_delegate is None

# The region also comes from AWS_REGION_NAME / the region kwarg.
eu = BedrockChatLLM(model="google.gemma-4-26b-a4b", aws_region_name="eu-central-1")
print("eu-central-1 base_url    :", eu._mantle_delegate.client_kwargs["base_url"])

import sys
import types
from types import SimpleNamespace as ns


def install_fake_openai(*responses):
    """Inject a fake `openai` module returning the given responses in order."""
    calls = {"i": 0}

    class _Completions:
        def create(self, **_kwargs):
            # Clamp so extra turns keep returning the last (text) response —
            # the agent loop always terminates cleanly.
            r = responses[min(calls["i"], len(responses) - 1)]
            calls["i"] += 1
            return r

    class _OpenAI:
        def __init__(self, **_kwargs):
            self.chat = ns(completions=_Completions())

    fake = types.ModuleType("openai")
    fake.OpenAI = _OpenAI
    sys.modules["openai"] = fake


def tool_call_response(name, arguments):
    msg = ns(
        content="",
        tool_calls=[ns(id="c1", function=ns(name=name, arguments=arguments))],
        reasoning_content=None,
    )
    return ns(choices=[ns(message=msg)], usage=None)


def text_response(text):
    msg = ns(content=text, tool_calls=[], reasoning_content=None)
    return ns(choices=[ns(message=msg)], usage=None)

from typing import Any
from shipit_agent import Agent, FunctionTool

ran: list[str] = []


def add(a: int, b: int, **_: Any) -> str:
    """Add two numbers."""
    ran.append("add")
    return str(a + b)


# Turn 1: Gemma 4 emits a native tool call. Turn 2: it answers in text.
install_fake_openai(
    tool_call_response("add", '{"a": 2, "b": 3}'),
    text_response("The answer is 5."),
)

agent = Agent(
    llm=BedrockChatLLM(model="google.gemma-4-31b", region="us-east-1"),
    tools=[FunctionTool.from_callable(add, name="add")],
    auto_use_skills=False,
)

result = agent.run("What is 2 + 3?")

print("output    :", result.output)
print("tool ran  :", ran)          # ['add'] — the tool actually executed
assert ran == ["add"]
assert "answer is 5" in result.output

import json as _json
from shipit_agent import Agent, format_activity, format_event_line

# Script Gemma 4's two turns: a native build_document call, then the answer.
install_fake_openai(
    tool_call_response("build_document", _json.dumps({
        "kind": "xlsx",
        "title": "Q2 Close",
        "sheets": [{
            "name": "P&L",
            "headers": ["Item", "Amount"],
            "rows": [["Revenue", 124000], ["Costs", -78500], ["Net", "=B2+B3"]],
        }],
    })),
    text_response("Q2 close workbook is ready — net income formula included."),
)

analyst = Agent.for_role(
    "finance-analyst",
    llm=BedrockChatLLM(model="google.gemma-4-31b", region="us-east-1"),
)
print("role tools:", sorted(t.name for t in analyst.tools), "\n")

# ── real streaming: events render the moment they happen ──
events = []
for event in analyst.stream("Close Q2 and hand me the workbook."):
    events.append(event)
    line = format_event_line(event)
    if line:
        print(line)

# After the stream: the full run as one clean activity trace.
print(format_activity(events))

# The workbook is real — reopen it and verify the styling + live formula.
import openpyxl

done = next(e for e in events if e.type == "tool_completed")
path = done.payload["output"].split("→ ")[1].split(" (")[0]
ws = openpyxl.load_workbook(path)["P&L"]
print("header bold :", ws["A1"].font.bold)
print("frozen panes:", ws.freeze_panes)
print("live formula:", ws["B4"].value)
assert ws["B4"].value == "=B2+B3"

from shipit_agent import connect_mcp, list_mcp_catalog

for entry in list_mcp_catalog():
    need = f"  (needs {', '.join(entry.required_env)})" if entry.required_env else ""
    print(f"{entry.name:<14} {entry.description}{need}")

# Fail-fast demo — one clear message, not a subprocess stack trace:
import os
os.environ.pop("SLACK_BOT_TOKEN", None)
try:
    connect_mcp("slack")
except ValueError as e:
    print("\n", e)

# Live usage (with credentials set):
#   agent = Agent.for_role(
#       "finance-analyst",
#       llm=BedrockChatLLM(model="google.gemma-4-31b"),
#       mcps=[connect_mcp("github"), connect_mcp("filesystem", args=["."])],
#   )

import os

LIVE = bool(os.getenv("AWS_BEARER_TOKEN_BEDROCK"))
if not LIVE:
    try:  # generate a short-term bearer token from standard AWS credentials
        from aws_bedrock_token_generator import provide_token
        os.environ["AWS_BEARER_TOKEN_BEDROCK"] = provide_token(region="us-east-1")
        os.environ.setdefault("AWS_REGION_NAME", "us-east-1")
        LIVE = True
        print("✓ short-term Bedrock token generated from AWS credentials")
    except Exception as e:
        print(f"offline — no Bedrock credentials available ({type(e).__name__})")
else:
    print("✓ using AWS_BEARER_TOKEN_BEDROCK from the environment")

# Fresh, unfaked OpenAI client for the live calls
sys.modules.pop("openai", None)

from shipit_agent import Agent, FunctionTool

def get_time(city: str, **_):
    """Return the local time for a city."""
    return f"It's 3:00 PM in {city}."

if LIVE:
    live_agent = Agent(
        llm=BedrockChatLLM(model="google.gemma-4-31b", region="us-east-1"),
        tools=[FunctionTool.from_callable(get_time, name="get_time")],
        auto_use_skills=False,
    )
    for event in live_agent.stream("What time is it in Tokyo? Use the tool."):
        line = format_event_line(event)
        if line:
            print(line)
        if event.type == "run_completed":
            print("→", event.payload.get("output", ""))
else:
    print("skipped (offline)")

if LIVE:
    analyst_live = Agent.for_role(
        "finance-analyst",
        llm=BedrockChatLLM(model="google.gemma-4-31b", region="us-east-1"),
    )
    live_result = analyst_live.run(
        "Use the build_document tool to create an xlsx titled 'Live Q2 Close' with one "
        "sheet 'P&L': headers Item, Amount; rows Revenue 124000, Costs -78500, and a "
        "Net row with a formula =B2+B3. Then tell me it's done."
    )
    print(format_activity(live_result))
    print()
    print("final :", live_result.output)
    print("metrics:", live_result.summary())
else:
    print("skipped (offline)")

from shipit_agent import connect_mcp

if LIVE:
    fs = connect_mcp("filesystem", args=["."])
    print("MCP tools:", [t.name for t in fs.discover_tools()][:7], "…\n")

    mcp_agent = Agent(
        llm=BedrockChatLLM(model="google.gemma-4-31b", region="us-east-1"),
        mcps=[fs],
        auto_use_skills=False,
    )
    try:
        for event in mcp_agent.stream(
            "Use the list_directory tool to list the current directory, "
            "then name three notebooks you see."
        ):
            line = format_event_line(event)
            if line:
                print(line)
            if event.type == "run_completed":
                print("→", event.payload.get("output", "")[:300])
    finally:
        fs.close()
else:
    print("skipped (offline)")