Google ADK integration

1
import os
2
from uuid import uuid4
3
4
from google.adk.agents import Agent
5
from google.adk.runners import Runner
6
from google.adk.sessions import InMemorySessionService
7
from google.genai import types
8
from zep_cloud.client import AsyncZep
9
from zep_adk import ZepContextTool, create_after_model_callback
10
11
zep = AsyncZep(api_key=os.environ["ZEP_API_KEY"])
12
13
# One shared agent definition — serves all users
14
agent = Agent(
15
    name="my_agent",
16
    model="gemini-2.5-flash",
17
    instruction="You are a helpful assistant with long-term memory.",
18
    tools=[
19
        your_existing_tool,                  # your tools stay as-is
20
        ZepContextTool(
21
            zep_client=zep,
22
            ignore_roles=["assistant"],      # optional — excludes assistant messages from graph ingestion
23
        ),
24
    ],
25
    after_model_callback=create_after_model_callback(
26
        zep_client=zep,
27
        assistant_name="my_agent",           # name shown in Zep (default: "Assistant")
28
        ignore_roles=["assistant"],          # optional — matches ZepContextTool setting
29
    ),
30
)
31
32
session_service = InMemorySessionService()
33
runner = Runner(agent=agent, app_name="my_app", session_service=session_service)
34
35
# Per-user session — identity maps automatically to Zep:
36
#   user_id    → Zep user ID
37
#   session_id → Zep thread ID (Zep's knowledge graph spans all threads for a user)
38
await session_service.create_session(
39
    app_name="my_app",
40
    user_id="user-123",
41
    session_id=f"session-{uuid4().hex[:8]}",
42
    state={
43
        **your_existing_state,               # your existing state stays as-is
44
        "zep_first_name": "Jane",            # recommended — anchors the identity node
45
        "zep_last_name": "Smith",            # optional (default: "User")
46
        "zep_email": "[email protected]",     # optional
47
    },
48
)
49
50
# Send a message — persistence and context injection happen automatically
51
content = types.Content(role="user", parts=[types.Part(text="Hi, I work at Acme Corp.")])
52
async for event in runner.run_async(
53
    user_id="user-123", session_id=session_id, new_message=content
54
):
55
    if event.is_final_response() and event.content:
56
        print(event.content.parts[0].text)

1
from pydantic import Field
2
from zep_cloud.client import AsyncZep
3
from zep_cloud.external_clients.ontology import EntityModel, EntityText
4
from zep_adk import ZepContextTool
5
6
class Company(EntityModel):
7
    """A company or organization the user is associated with."""
8
    industry: EntityText = Field(description="The company's industry", default=None)
9
10
async def setup_user(zep_client: AsyncZep, user_id: str) -> None:
11
    """Runs once when a new Zep user is created."""
12
    # Set a custom ontology for this user's knowledge graph
13
    await zep_client.graph.set_ontology(
14
        entities={"Company": Company},
15
        user_ids=[user_id],
16
    )
17
18
    # Add custom extraction instructions
19
    await zep_client.graph.add_custom_instructions(
20
        user_ids=[user_id],
21
        custom_instructions="Extract product preferences and purchase intent.",
22
    )
23
24
    # Configure how user summaries are generated
25
    await zep_client.user.add_user_summary_instructions(
26
        user_ids=[user_id],
27
        user_summary_instructions="Focus on the user's role, team, and active projects.",
28
    )
29
30
ZepContextTool(zep_client=zep, on_user_created=setup_user)

1
import asyncio
2
from zep_cloud.client import AsyncZep
3
from zep_adk import ZepContextTool, ContextBuilder
4
5
async def my_context_builder(
6
    zep_client: AsyncZep,
7
    user_id: str,
8
    thread_id: str,
9
    user_message: str,
10
) -> str | None:
11
    """Custom context: combine user context with a targeted graph search."""
12
    user_context, search_results = await asyncio.gather(
13
        zep_client.user.get_context(user_id),
14
        zep_client.graph.search(
15
            user_id=user_id,
16
            query=user_message,
17
            scope="edges",
18
            limit=10,
19
        ),
20
    )
21
22
    parts = []
23
    if user_context and user_context.context:
24
        parts.append(user_context.context)
25
    if search_results and search_results.edges:
26
        facts = [e.fact for e in search_results.edges if e.fact]
27
        if facts:
28
            parts.append("Additional facts:\n" + "\n".join(f"- {f}" for f in facts))
29
30
    return "\n\n".join(parts) if parts else None
31
32
# Pass it to ZepContextTool
33
tool = ZepContextTool(zep_client=zep, context_builder=my_context_builder)

1
ContextBuilder = Callable[[AsyncZep, str, str, str], Awaitable[str | None]]
2
#                          client   user  thread message
3
4
UserSetupHook = Callable[[AsyncZep, str], Awaitable[None]]
5
#                         client   user_id

1
from zep_adk import ZepContextTool, ZepGraphSearchTool, create_after_model_callback
2
3
agent = Agent(
4
    name="my_agent",
5
    model="gemini-2.5-flash",
6
    instruction="...",
7
    tools=[
8
        ZepContextTool(zep_client=zep),          # automatic context every turn
9
        ZepGraphSearchTool(zep_client=zep),      # on-demand search
10
    ],
11
    after_model_callback=create_after_model_callback(zep_client=zep),
12
)

1
# Model can only control query and scope — everything else is locked
2
ZepGraphSearchTool(
3
    zep_client=zep,
4
    reranker="cross_encoder",
5
    limit=5,
6
    search_filters={"node_labels": ["Person"]},
7
    bfs_origin_node_uuids=["node-uuid-1"],   # seed BFS traversal from specific nodes
8
)

1
agent = Agent(
2
    name="my_agent",
3
    model="gemini-2.5-flash",
4
    instruction="...",
5
    tools=[
6
        ZepContextTool(zep_client=zep),
7
        ZepGraphSearchTool(
8
            zep_client=zep,
9
            name="search_user_memory",
10
            description="Search the user's knowledge graph for information from previous conversations, known facts, or general context about the user.",
11
        ),
12
        ZepGraphSearchTool(
13
            zep_client=zep,
14
            name="search_docs",
15
            description="Search the shared documentation knowledge base.",
16
            graph_id="docs-graph-123",
17
        ),
18
    ],
19
    after_model_callback=create_after_model_callback(zep_client=zep),
20
)

1
import asyncio
2
from zep_cloud.client import AsyncZep
3
from zep_cloud import Message
4
5
zep = AsyncZep(api_key="your-zep-api-key")
6
7
async def backfill_user(
8
    user_id: str,                          # must match ADK create_session() user_id
9
    first_name: str,
10
    last_name: str,
11
    conversations: list[dict],             # list of {session_id, messages} dicts
12
):
13
    # 1. Create the user
14
    try:
15
        await zep.user.add(user_id=user_id, first_name=first_name, last_name=last_name)
16
    except Exception as e:
17
        if "already exists" not in str(e).lower():
18
            raise
19
20
    # 2. Load each conversation — use the original ADK session ID as the Zep thread ID
21
    for convo in conversations:
22
        thread_id = convo["session_id"]    # must match ADK session_id
23
        try:
24
            await zep.thread.create(thread_id=thread_id, user_id=user_id)
25
        except Exception as e:
26
            if "already exists" in str(e).lower():
27
                continue
28
            raise
29
30
        messages = [
31
            Message(
32
                role=msg["role"],
33
                content=msg["content"],
34
                name=f"{first_name} {last_name}" if msg["role"] == "user" else "Assistant",
35
            )
36
            for msg in convo["messages"]
37
        ]
38
        await zep.thread.add_messages(thread_id=thread_id, messages=messages)
39
40
    print(f"Backfilled {len(conversations)} conversations for {user_id}")
41
42
async def main():
43
    users = [
44
        {
45
            "user_id": "user-123",       # same ID used in ADK sessions
46
            "first_name": "Jane",
47
            "last_name": "Smith",
48
            "conversations": [
49
                {
50
                    "session_id": "session-abc",  # original ADK session ID
51
                    "messages": [
52
                        {"role": "user", "content": "I need help with my account settings."},
53
                        {"role": "assistant", "content": "I can help. What would you like to change?"},
54
                        {"role": "user", "content": "I want to enable two-factor authentication."},
55
                        {"role": "assistant", "content": "Go to Settings > Security > 2FA to enable it."},
56
                    ],
57
                },
58
            ],
59
        },
60
    ]
61
    for user in users:
62
        await backfill_user(**user)
63
64
asyncio.run(main())