Use Agent Platform models, sessions and Memory Bank

"""The docs assistant on Gemini Enterprise Agent Platform: Gemini through Vertex, recall through Memory Bank."""
import os
from pathlib import Path

from google.adk import Agent
from google.adk.models.google_llm import Gemini
from google.adk.tools import load_memory
from google.genai import types

DOCS = Path(__file__).resolve().parent / "docs"
PAGES = sorted(p.stem for p in DOCS.glob("*.md"))
# The recorded run uses global for this model. Sessions and Memory Bank stay regional
# (GOOGLE_CLOUD_LOCATION), so the model client gets its own location.
MODEL_LOCATION = os.environ.get("DOCS_MODEL_LOCATION", "global")


def read_page(page: str) -> dict:
    """Read one page of the docuctl documentation: config, deploy, rollback or tokens."""
    if page not in PAGES:
        return {"error": f"No page {page!r}. Pages: {', '.join(PAGES)}."}
    return {"page": page, "text": (DOCS / f"{page}.md").read_text(encoding="utf-8")}


INSTRUCTION = """You answer questions from developers about docuctl, the team's deploy CLI.
Read the relevant page with read_page before you answer. Answers depend on the developer's
docuctl major version (1.x or 2.x). If the developer has not said it in this conversation,
call load_memory to look for it in earlier conversations before you ask. Keep the answer
under 80 words."""

root_agent = Agent(
    name="docs_helper",
    model=Gemini(model="gemini-3.8-flash", client_kwargs={"location": MODEL_LOCATION}),
    instruction=INSTRUCTION,
    tools=[read_page, load_memory],
    generate_content_config=types.GenerateContentConfig(
        temperature=0.2, thinking_config=types.ThinkingConfig(thinking_level="low")),
)

"""Create one Agent Runtime instance with no code: it holds Sessions and Memory Bank for the docs assistant."""
import os
import sys

import vertexai

LABEL = "use-agent-platform-models-sessions-and-memory-bank"


def memory_config() -> dict:
    return {"display_name": "docs-helper-memory", "labels": {"aiml-lesson": LABEL},
            "context_spec": {"memory_bank_config": {
                "ttl_config": {"default_ttl": "86400s"},
                "customization_configs": [{"memory_topics": [
                    {"managed_memory_topic": {"managed_topic_enum": "USER_PREFERENCES"}},
                    {"custom_memory_topic": {"label": "tool_versions",
                        "description": "Which version of a CLI, SDK or service the developer runs, for example 'runs docuctl 1.x'."}},
                ]}],
            }}}


def main() -> None:
    client = vertexai.Client(project=os.environ["GOOGLE_CLOUD_PROJECT"],
                             location=os.environ.get("GOOGLE_CLOUD_LOCATION", "us-central1"))
    engine = client.agent_engines.create(config=memory_config())
    print("created", engine.api_resource.name)
    print("agent_engine_id", engine.api_resource.name.split("/")[-1], file=sys.stderr)


if __name__ == "__main__":
    main()

"""Run the docs assistant on Agent Platform Sessions and Memory Bank, and record what happens.

Usage: AGENT_ENGINE_ID=<id> python record.py
Needs ADC (gcloud auth application-default login), GOOGLE_CLOUD_PROJECT and GOOGLE_CLOUD_LOCATION.
"""
import asyncio
import json
import os
import re
import time
from pathlib import Path

from google.adk import Agent, Runner
from google.adk.memory import VertexAiMemoryBankService
from google.adk.models.google_llm import Gemini
from google.adk.sessions import InMemorySessionService, VertexAiSessionService
from google.genai import types

from docs_helper.agent import root_agent

HERE = Path(__file__).resolve().parent
PROJECT = os.environ["GOOGLE_CLOUD_PROJECT"]
LOCATION = os.environ["GOOGLE_CLOUD_LOCATION"]
ENGINE = os.environ["AGENT_ENGINE_ID"]
APP = ENGINE  # with Agent Platform Sessions, the app name is the Agent Runtime instance
USER = os.environ.get("DOCS_USER", "dana")
out, record = [], {"model": "gemini-3.8-flash", "location": LOCATION, "failures": {}}


def redact(text: str) -> str:
    return re.sub(r"projects/\d+/", "projects/PROJECT_NUMBER/", str(text))


def log(line: str = "") -> None:
    out.append(redact(line))
    print(redact(line))


def summary(events) -> list[str]:
    lines = []
    for event in events:
        lines += [f"call   {c.name} {json.dumps(c.args)}" for c in event.get_function_calls()]
        for r in event.get_function_responses():
            body = r.response if r.name == "load_memory" else sorted(r.response)
            text = json.dumps(body, default=lambda o: o.model_dump(mode="json", exclude_none=True))
            lines.append(f"result {r.name} {redact(text)[:300]}")
        if event.is_final_response() and event.content:
            lines.append("answer " + "".join(p.text or "" for p in event.content.parts if not p.thought).strip())
    return lines


async def turn(runner: Runner, session_id: str, text: str) -> list:
    message = types.Content(role="user", parts=[types.Part(text=text)])
    return [e async for e in runner.run_async(user_id=USER, session_id=session_id, new_message=message)]


async def failures() -> None:
    # 1. Record the failure from using the session region for the model request.
    regional = Agent(name="docs_helper", model=Gemini(model="gemini-3.8-flash", client_kwargs={"location": LOCATION}),
                     instruction="Answer in one word.")
    runner = Runner(app_name="docs", agent=regional, session_service=InMemorySessionService())
    session = await runner.session_service.create_session(app_name="docs", user_id=USER)
    try:
        await turn(runner, session.id, "Say ok.")
    except Exception as error:  # noqa: BLE001 - record whatever the API raised
        record["failures"]["model_region"] = f"{type(error).__name__}: {redact(error)[:400]}"
    # 2. A Sessions service with no instance id and an ADK-style app name.
    try:
        await VertexAiSessionService(project=PROJECT, location=LOCATION).create_session(app_name="docs_helper", user_id=USER)
    except ValueError as error:
        record["failures"]["app_name"] = f"ValueError: {error}"
    for name, text in record["failures"].items():
        log(f"# failure {name}\n{text}\n")


async def require_empty() -> None:
    """Refuse an already-used instance instead of deleting its memories or sessions."""
    import vertexai
    client = vertexai.Client(project=PROJECT, location=LOCATION)
    name = f"projects/{PROJECT}/locations/{LOCATION}/reasoningEngines/{ENGINE}"
    if any(client.agent_engines.memories.list(name=name)):
        raise RuntimeError("Use a fresh lesson instance; existing memories will not be deleted.")
    sessions = VertexAiSessionService(PROJECT, LOCATION, ENGINE)
    if (await sessions.list_sessions(app_name=APP, user_id=USER)).sessions:
        raise RuntimeError("Use a fresh lesson instance; existing sessions will not be deleted.")


async def main() -> None:
    await require_empty()
    await failures()
    sessions = VertexAiSessionService(PROJECT, LOCATION, ENGINE)
    memory = VertexAiMemoryBankService(PROJECT, LOCATION, ENGINE)
    runner = Runner(app_name=APP, agent=root_agent, session_service=sessions, memory_service=memory)

    # Conversation 1: Dana says which version she runs.
    s1 = await sessions.create_session(app_name=APP, user_id=USER)
    log(f"# conversation 1 · session {s1.id}")
    log("> I run docuctl 1.x. How do I undo release r42 in production?")
    for line in summary(await turn(runner, s1.id, "I run docuctl 1.x. How do I undo release r42 in production?")):
        log(line)
    stored = await sessions.get_session(app_name=APP, user_id=USER, session_id=s1.id)
    record["session_1"] = {"id": s1.id, "stored_events": len(stored.events),
                           "authors": [e.author for e in stored.events]}

    # IngestEvents returns at once; Memory Bank extracts memories in the background.
    # force_flush asks it to process this conversation now instead of on its own trigger.
    started = time.monotonic()
    await memory.add_events_to_memory(app_name=APP, user_id=USER, events=stored.events, session_id=s1.id,
                                      custom_metadata={"force_flush": True})
    record["ingest_return_seconds"] = round(time.monotonic() - started, 2)
    early = await memory.search_memory(app_name=APP, user_id=USER, query="docuctl version")
    record["memories_right_after_ingest"] = len(early.memories)
    log(f"\n# ingest returned in {record['ingest_return_seconds']} s; search right after: {len(early.memories)} memories")
    found = early
    while not found.memories and time.monotonic() - started < 180:
        await asyncio.sleep(3)
        found = await memory.search_memory(app_name=APP, user_id=USER, query="docuctl version")
    record["searchable_after_seconds"] = round(time.monotonic() - started)
    record["memories"] = ["".join(p.text or "" for p in m.content.parts) for m in found.memories]
    log(f"# searchable after about {record['searchable_after_seconds']} s:")
    for text in record["memories"]:
        log(f"  - {text}")

    # Conversation 2: a new session; the version is not in it.
    s2 = await sessions.create_session(app_name=APP, user_id=USER)
    log(f"\n# conversation 2 · session {s2.id}")
    log("> How do I undo the last production release?")
    lines = summary(await turn(runner, s2.id, "How do I undo the last production release?"))
    for line in lines:
        log(line)
    record["session_2"] = {"id": s2.id, "lines": lines}
    listed = await sessions.list_sessions(app_name=APP, user_id=USER)
    record["sessions_listed"] = len(listed.sessions)
    log(f"\n# list_sessions for {USER}: {len(listed.sessions)} sessions")
    await runner.close()

    (HERE / "run-output.txt").write_text("\n".join(out) + "\n", encoding="utf-8")
    (HERE / "recorded-runs.json").write_text(redact(json.dumps(record, indent=2)) + "\n", encoding="utf-8")


if __name__ == "__main__":
    asyncio.run(main())