Use with Google ADK

"""Shared Aura connection through Bolt and readback for integration recipes."""

import os


def settings(embedding=None):
    from neo4j_agent_memory import BoltSettings
    from neo4j_agent_memory.llm import from_provider

    return BoltSettings(
        neo4j={
            "uri": os.environ["NEO4J_URI"],
            "username": os.environ["NEO4J_USERNAME"],
            "password": os.environ["NEO4J_PASSWORD"],
            "database": os.getenv("NEO4J_DATABASE", "neo4j"),
        },
        embedding=embedding or from_provider("openai/text-embedding-3-small", kind="embedding"),
        extraction={"extractor_type": "none"},
    )


async def verify_messages(client, session_id, expected):
    conversation = await client.short_term.get_conversation(session_id)
    contents = [message.content for message in conversation.messages]
    if not all(text in contents for text in expected):
        raise RuntimeError(f"Message readback failed for {session_id}")
    print(f"Verified stored messages; session={session_id}")


async def recorded_turn(client, session_id, prompt, respond):
    """Record a task outcome; this does not claim to capture hidden model reasoning."""
    trace = await client.reasoning.start_trace(session_id=session_id, task=prompt)
    try:
        await client.short_term.add_message(session_id, "user", prompt, extract_entities=False)
        context = await client.get_context(prompt, session_id=session_id)
        reply = str(await respond(context))
        if not reply.strip():
            raise RuntimeError("The framework returned no text")
        await client.short_term.add_message(session_id, "assistant", reply, extract_entities=False)
        await verify_messages(client, session_id, [prompt, reply])
    except Exception as exc:
        await client.reasoning.complete_trace(
            trace.id, success=False, outcome=f"Run failed: {type(exc).__name__}"
        )
        raise
    await client.reasoning.complete_trace(
        trace.id, success=True, outcome="Reply stored and read back"
    )
    stored = await client.reasoning.get_trace_with_steps(trace.id)
    if stored is None or stored.success is not True:
        raise RuntimeError(f"Trace readback failed for {trace.id}")
    print(f"Verified task trace: {trace.id}")
    return reply

"""Consume the ADK Runner stream, persist its session, then verify readback."""

import asyncio
import os

from common import settings, verify_messages


async def exercise(client):
    from google.adk.agents import LlmAgent
    from google.adk.runners import Runner
    from google.adk.sessions import InMemorySessionService
    from google.adk.tools import load_memory
    from google.genai import types

    from neo4j_agent_memory.integrations.google_adk import Neo4jMemoryService

    app_name, user_id = "docs-memory", "fictional-user"
    service = Neo4jMemoryService(client, user_id=user_id, extract_on_store=False)
    sessions = InMemorySessionService()
    agent = LlmAgent(
        name="checklist_agent",
        model=os.environ["ADK_MODEL"],
        instruction="Answer briefly. Use load_memory when prior context is needed.",
        tools=[load_memory],
    )
    runner = Runner(
        app_name=app_name, agent=agent, session_service=sessions, memory_service=service
    )
    session = await sessions.create_session(app_name=app_name, user_id=user_id)
    prompt = "The fictional project code is cedar-lantern-47. Acknowledge it briefly."
    async for event in runner.run_async(
        user_id=user_id,
        session_id=session.id,
        new_message=types.Content(role="user", parts=[types.Part(text=prompt)]),
    ):
        if event.error_code:
            raise RuntimeError(f"ADK run failed: {event.error_code}")
        if event.is_final_response() and event.content and event.content.parts:
            print("".join(part.text or "" for part in event.content.parts))
    stored_session = await sessions.get_session(
        app_name=app_name, user_id=user_id, session_id=session.id
    )
    if stored_session is None:
        raise RuntimeError("ADK session is missing after the run")
    await service.add_session_to_memory(stored_session)
    await verify_messages(client, session.id, [prompt])
    response = await service.search_memory(app_name=app_name, user_id=user_id, query="project code")
    print(f"ADK SearchMemoryResponse: {len(response.memories)} candidate(s)")
    return session.id


async def main():
    from neo4j_agent_memory import MemoryClient

    async with MemoryClient(settings()) as client:
        await exercise(client)


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