[Tagar]

Summary

Adds memory_store() to databricks_agentkit.langgraph. It returns DatabricksMemoryStore, a LangGraph BaseStore on Agent Bricks managed memory. LangGraph and LangMem code that takes store=, such as create_agent(..., store=memory_store()), then gets memory that lasts across conversations, with no Postgres connection, schema or embedding endpoint to run.

Why: AgentKit's LangGraph helpers already cover sessions (checkpointer()) and memory tools (memory_tools()). LangGraph's own long-term memory interface is BaseStore, though, and LangMem and many LangGraph agents are written against it. On Databricks today that means databricks_langchain.DatabricksStore, a PostgresStore on a Lakebase database the team owns and operates. This PR puts the same interface on the managed memory service that AgentKit already uses.

In short, the store should:

• Serve every user of an app from one instance. It takes the user from each run's config (thread_config(session_id, actor)), and raises rather than mixes users when an explicit actor= and the run disagree.
• Keep its entries apart from memory_tools, sessions and automatic memory pipeline runs by storing them under the storage actor <actor>:langgraph.
• Cover the whole store interface, sync and async, with LangGraph's filter semantics.
• Refuse what the service can't store before making any call: values that aren't JSON objects, NaN, and oversized requests or paths.
• Support TTL on the client side, documented where it differs from LangGraph's Postgres store.
• Fall back to an in-process stand-in only for local runs. A deployed App with no store bound raises.

Changes

• langgraph/store.py (new): DatabricksMemoryStore and memory_store().

  • Storage: each item is one cross-session memory entry at /langgraph/<namespace>/<key>, with each segment percent-encoded. The entry holds the value in a small JSON envelope ({"lg": 1, "v": ...}, plus "x", the expiry, for TTL items) and a plain-text description.
  • Search: the query goes to the service; the namespace prefix, filter and offset apply client-side.
  • TTL: counts from the last write, and reads never refresh it. sweep_ttl(actor=...) deletes one user's expired items.
  • Attribution: the filter helpers are adapted from LangGraph's InMemoryStore (MIT) and credited in the source and in NOTICE.

• Reserved actor suffix (_actors.py, new): these now refuse actor ids ending in :langgraph, so no session or memory-tools actor can reach another user's store entries:

  • thread_config;
  • the LangGraph checkpointer's config parsing;
  • memory_tools;
  • the OpenAI session_store and memory_tools.

  This is a behaviour change for any existing caller whose actor or session id ends with that suffix. Other ids, including non-string ones, pass through unchanged.

• Docs:

  • A README section: setup, where entries live, identity, LangMem, limits, TTL, a check for semantic matching, and access.
  • A migration page, docs/migrate-from-databricks-store.md.
  • The README's memory section no longer describes search as BM25-only, to match the managed memory docs, and the keying section says where the store's entries live.

• Tests:

  • Unit tests run against a wire-level fake of the memory API. The real AgentKit client runs against it, and the fake rejects the same requests the service rejects.
  • A gated live suite, tests/suites/test_langgraph_store.py, run with RUN_LANGGRAPH_STORE_TESTS=1.

It lives in databricks_agentkit.langgraph, next to checkpointer() and memory_tools(), because it's built on AgentKit's memory client, and databricks-langchain doesn't depend on databricks-agentbricks. DatabricksStore is unchanged, for teams that want their own Postgres and embeddings.

Known limitations

• Ranking: ranking is the service's. The response doesn't say whether semantic matching contributed, and scores are rank-based, not similarities. The unit tests use a keyword fake and don't cover ranking.
• Service limits:
  • At most 100 results per search, with filters applied after ranking.
  • put is last-writer-wins, and overwriting an existing key takes three API calls.
• Pipeline separation:
  • Automatic memory pipeline runs find actors through sessions, so they skip <actor>:langgraph. That relies on current pipeline behaviour and isn't tested here against a live pipeline.
  • A run that names its actors explicitly can still reach the entries. The README recommends a memory store with no pipeline for the strongest separation.
• Templates: this PR doesn't wire the store into the templates. The bundled agent-langgraph template takes input.actor from the request body, and the README says not to add the store until the actor comes from the forwarded identity.
• Existing client issue: MemoryStore._put_memory doesn't retry an update that hits NOT_FOUND. A strict xfail marks it.

Testing

• Unit tests:
  • The four store test files pass on Python 3.12 (120 passed, 1 strict xfail) and on 3.10 (117 passed, 3 skipped, 1 xfail).
  • The full agentbricks unit suite passes.
  • The store modules skip cleanly without the langgraph extra, as in CI's 3.10 unit job.
• Static checks: ruff is clean. ty check reports only the existing model_upgrades/optimization.py diagnostic.
• Live, against a managed memory store on a Databricks workspace:
  • The gated suite passes: put, get, overwrite and delete; two users kept apart; entries under the storage actor; TTL expiry and sweep; and the semantic-search check.
  • A short benchmark of about 2,100 calls had no errors or throttling. Median latencies were about 100 ms for a put of a new key, 300 ms for an existing key (three calls), 95 ms for a get, and 280 ms for a search.
  • Eight concurrent writers to one new key left exactly one entry.
  • LangMem's tools, memory manager and ReflectionExecutor ran against a live store.
  • A tutorial notebook ran end to end on serverless compute.

Requesting review from @nihit, @qinxinw-db, @elainewang-db and @junchoi-db.

This pull request and its description were written by Isaac.

[Tagar]

/trigger all