mubit-adk

Google ADK BaseMemoryService backed by MuBit.

This adapter now uses the canonical Python SDK transport internally instead of raw httpx, while preserving the ADK memory-service contract and adding MuBit-specific helpers for checkpointing, observability, and multi-agent coordination.

Install

pip install mubit-adk[adk]

Basic usage

from google.adk.agents import Agent
from google.adk.tools.preload_memory_tool import PreloadMemoryTool

from mubit_adk import (
    MubitMemoryService,
    make_bigquery_tool_callback,
    make_session_memory_callback,
)

memory = MubitMemoryService(
    endpoint="http://127.0.0.1:3000",
    api_key="mbt_...",
)

agent = Agent(
    name="bigquery_agent",
    description="Read-only BigQuery agent with MuBit memory.",
    model="gemini-2.0-flash",
    instruction="Use BigQuery tools and prior MuBit lessons to answer safely.",
    tools=[
        PreloadMemoryTool(),  # ADK calls memory.search_memory() before each turn.
        bigquery_toolset,     # Your ADK BigQueryToolset.
    ],
    after_agent_callback=make_session_memory_callback(memory),
    after_tool_callback=make_bigquery_tool_callback(memory),
)

For Vertex AI Agent Engine, pass the same memory service through AdkApp:

from vertexai.agent_engines import AdkApp

adk_app = AdkApp(
    agent=agent,
    memory_service_builder=lambda: memory,
    enable_tracing=True,
)

ADK to MuBit mapping

Closing the attribution loop (recall -> act -> record_outcome)

search_memory() accepts an opt-in with_entry_ids=True that attaches response.mubit_entry_ids (the recalled reference IDs) and response.mubit_citations (grounding indices) to the standard SearchMemoryResponse. For a self-contained handle, search_with_attribution() returns the response plus entry_ids, citations, and the raw recall dict. Feed those IDs straight back into record_outcome(entry_ids=...) so MuBit can credit every entry that contributed to the result.

result = await memory.search_with_attribution(
    app_name="bigquery_agent", user_id="u1", query="how do I dry-run?"
)
# ... agent acts on result["response"].memories ...
await memory.record_outcome(
    app_name="bigquery_agent",
    user_id="u1",
    outcome="success",
    entry_ids=result["entry_ids"],     # credit the recalled entries
    verified_in_production=True,        # boost lessons confirmed in live use
)

The module also re-exports extract_entry_ids(result, cited_only=...) and extract_citations(result) for pulling those handles off any recall() result.

All ingest paths (add_session_to_memory, add_events_to_memory, add_memory) take an optional idempotency_key= to override the default per-session token, so a retried write is deduplicated server-side instead of double-counting reinforcement.

MuBit extension methods

Existing MAS helpers:

• checkpoint()
• record_outcome() — accepts entry_ids= and verified_in_production= for v0.7.0 attribution
• record_step_outcome()
• surface_strategies()

Current observability and coordination helpers:

• get_context()
• memory_health()
• diagnose()
• archive()
• dereference()
• register_agent()
• list_agents()
• handoff()
• feedback()
• make_session_memory_callback() for fail-open ADK session ingestion
• make_bigquery_tool_callback() for redacted BigQuery tool lessons and step outcomes

Config

For tests or advanced embedding, you can inject mubit_client directly.

Development

PYTHONPATH=sdk/python/mubit-sdk/src:integrations/python/mubit_adk:integrations/python/mubit_integration_base \
python3 -m unittest integrations.python.mubit_adk.tests.test_memory_service integrations.python.mubit_adk.tests.test_mubit_client -v

License

Apache-2.0