← All posts
Use case

Typed memory for AI agents with TypeDB and Lean

An agent's memory is a database, so it should have a schema. This post models observations, beliefs, evidence, time, confidence and contradictions in TypeDB, and lets Lean check the code that reads and writes them.

An agent that runs for more than one conversation needs to remember things. The usual first attempts are to keep the chat history in the context window, to embed text chunks into a vector store, or to save a JSON document per session. All three are good at one question: find me text that looks like this.

Agents soon have to answer other questions, though:

Those are questions about relationships between facts, not about text similarity. They need a model with types: what kinds of things exist, how they relate, and what each relationship carries. This post builds such a model with TypeDB, a database whose schema is a type system, and typedb-lean, which makes the Lean compiler check every query and transaction against it.

What memory needs

Strip away the language model, and an agent's memory has six requirements:

  1. Facts are things. An observation has an identity, so it can be referenced, supported, contradicted or retracted. A sentence in a text blob can't be.
  2. Provenance. Every belief points at the observation it came from. Later, other observations can count as evidence for it.
  3. Time. When something was observed, and since when an agent has held it. "The bridge is closed" was true on Tuesday.
  4. Confidence, exactly. Thresholds such as "act only on beliefs at 0.8 or above" have to compare correctly. A 0.1 stored as a binary float is 0.1000000000000000055….
  5. Relations between facts. Support and contradiction relate observations and beliefs, not strings.
  6. Many agents. Several agents share one memory, and each holds its own beliefs.

Why a typed graph

Vector store or JSON blob

  • Facts are text chunks, with no identity to point at
  • Provenance and time live in free-form metadata, if anywhere
  • "Contradicts" and "supports" are not representable, only nearby
  • Confidence is a float in a dictionary
  • The shape of a record is whatever the last writer wrote

Typed graph in TypeDB

  • Observations are entities with a key
  • Beliefs are relations that carry the time they were adopted
  • Support and contradiction are relations, and a relation can play a role in another one
  • Confidence is an exact decimal
  • The server rejects data that does not fit the schema

This doesn't replace embeddings. Fuzzy recall ("anything about bridges?") is what vectors are for, and an embedding index can store observation keys next to its vectors. The typed graph is where the agent keeps what it knows, why it knows it, and how the pieces fit together.

The model

We have three entity types and four relation types. Entities are drawn as circles and relations as diamonds, as in an entity-relationship diagram.

holder content source topic supported evidence claim claim agent name @key observation id @key · text observed-at · confidence topic name @key belief since mention support weight contradiction
Three entity types and four relation types. support takes a belief, itself a relation, as its supported player.

The model in Lean

With typedb-lean you write the model once, as Lean structures. Markers in the field types say how each field maps: Key is a key attribute, Role "…" (Ref T) is a role player, and every other field is an owned attribute whose TypeDB value type follows from its Lean type.

import TypeDB
open TypeDB Std.Time

@[typedb "agent"]
structure Agent where
  name : Key String
  deriving Repr, Entity

@[typedb "topic"]
structure Topic where
  name : Key String
  deriving Repr, Entity

@[typedb "observation"]
structure Observation where
  id : Key String
  text : String
  observedAt : PlainDateTime
  confidence : Decimal
  deriving Repr, Entity

/-- An observation mentions a topic. -/
structure Mention where
  source : Role "source" (Ref Observation)
  topic : Role "topic" (Ref Topic)
  deriving Repr, Relation

/-- An agent holds an observation as one of its beliefs. -/
structure Belief where
  holder : Role "holder" (Ref Agent)
  content : Role "content" (Ref Observation)
  since : PlainDateTime
  deriving Repr, Relation

/-- Two observations that cannot both be true. -/
structure Contradiction where
  a : Role "claim" (Ref Observation)
  b : Role "claim" (Ref Observation)
  deriving Repr, Relation

/-- Evidence for a belief. The `supported` player is itself a relation. -/
structure Support where
  supported : Role "supported" (Ref Belief)
  evidence : Role "evidence" (Ref Observation)
  weight : Decimal
  deriving Repr, Relation

def schema : List TypeSpec :=
  [TypeSpec.of Agent, TypeSpec.of Topic, TypeSpec.of Observation,
   TypeSpec.of Mention, TypeSpec.of Belief, TypeSpec.of Contradiction, TypeSpec.of Support]

Note the types. confidence is a Decimal, exact to 19 fractional digits. observedAt is a Std.Time.PlainDateTime, not a string. Support.supported is a Ref Belief: the compiler won't let you put an observation there.

The schema it generates

Schema.defineText schema prints the TypeQL that the structures stand for. This is the real output:

define
  attribute name, value string;
  attribute id, value string;
  attribute text, value string;
  attribute observed-at, value datetime;
  attribute confidence, value decimal;
  attribute since, value datetime;
  attribute weight, value decimal;
  entity agent, owns name @key;
  entity topic, owns name @key;
  entity observation, owns id @key, owns text @card(1..1), owns observed-at @card(1..1), owns confidence @card(1..1);
  relation mention, relates source, relates topic;
  relation belief, relates holder, relates content, owns since @card(1..1);
  relation contradiction, relates claim @card(0..2);
  relation support, relates supported, relates evidence, owns weight @card(1..1);
  observation plays mention:source;
  topic plays mention:topic;
  agent plays belief:holder;
  observation plays belief:content;
  observation plays contradiction:claim;
  belief plays support:supported;
  observation plays support:evidence;

Field names become kebab-case labels (observedAt → observed-at). Required fields get @card(1..1), so the server rejects an observation without a text. Two fields that share the role claim become @card(0..2). The last line, belief plays support:supported, is where a relation plays a role.

Remember: one atomic write

Remembering something takes several inserts: the observation, a mention per topic (creating topics the agent hasn't seen yet), and the agent's belief. They belong together. Either the agent knows the whole thing or nothing at all, so they run in one write transaction:

def remember (agent id text : String) (at_ : PlainDateTime) (confidence : Decimal)
    (topics : List String) : Tx .write (Ref Belief) := do
  let obs ← Entity.insert { id, text, observedAt := at_, confidence : Observation }
  for t in topics do
    let topic := Ref.key (α := Topic) t
    if (← Entity.get? topic).isNone then
      discard <| Entity.insert { name := t : Topic }
    discard <| Relation.insert { source := obs, topic : Mention }
  Relation.insert { holder := Ref.key (α := Agent) agent, content := obs, since := at_ : Belief }

The type Tx .write (Ref Belief) says it all: it needs a write transaction, and it returns a reference to the new belief. If any step fails (an unknown agent key, a duplicate observation id), Driver.runSync closes the transaction without committing. No observation is left without its belief.

Several agents may write at once. TypeDB detects conflicting writes at commit, and a retry policy re-runs the whole body when, and only when, the commit loses a conflict:

let belief ← d.runSync "memory" .write (retry := .conflicts 5 (.exponential 20))
  (remember "scout" "obs-1" "The bridge on Route 9 is closed for repairs." now (pct 90) ["route-9", "bridge"])

Here pct 90 is a two-line helper, Decimal.ofScaled (n * 10 ^ 17), that builds the exact decimal 0.9.

Recall: typed queries

Recall is a TypeQL query whose rows decode straight into a structure. deriving FromRow reads each field from the query variable of the same name:

structure Recalled where
  text : String
  confidence : Decimal
  deriving Repr, FromRow

def recallAbout (topic : String) : Tx .read (Array Recalled) :=
  query typeql!"match
    $t isa topic, has name {topic};
    $m isa mention, links (source: $o, topic: $t);
    $o has text $text, has confidence $confidence;
    sort $confidence desc;"

The {topic} comes from the user, or from the language model, which is worse. typeql! renders every interpolated value as a quoted, escaped TypeQL literal, so the value can't change the query's structure. Here is what a hostile topic turns into:

def topic := "bridge\"; match $x isa agent; delete $x;"
#eval IO.println (typeql!"match $t isa topic, has name {topic};").text
match $t isa topic, has name "bridge\"; match $x isa agent; delete $x;";

That's one string literal. Run against our test database, the query found no topic with that odd name and returned zero rows, and both agents were still there afterwards. Also note the query is a Tx .read: even if text did slip through, a read transaction cannot delete.

Questions a vector store can't answer

With the relationships in the schema, the questions from the start of this post become short queries.

Which beliefs clear a confidence bar?

def beliefs (agent : String) (atLeast : Decimal) : Tx .read (Array Recalled) :=
  query typeql!"match
    $a isa agent, has name {agent};
    $b isa belief, links (holder: $a, content: $o);
    $o has text $text, has confidence $confidence;
    $confidence >= {atLeast};"

{atLeast} with pct 80 renders as 0.8dec, a TypeQL decimal literal, so the comparison is exact on both sides.

What do I believe that something contradicts?

structure Conflict where
  mine : String
  theirs : String
  deriving Repr, FromRow

def conflicts (agent : String) : Tx .read (Array Conflict) :=
  query typeql!"match
    $a isa agent, has name {agent};
    $b isa belief, links (holder: $a, content: $o1);
    $c isa contradiction, links (claim: $o1, claim: $o2);
    $o1 has text $mine;
    $o2 has text $theirs;"

Which beliefs have no evidence?

TypeQL's not matches the absence of a pattern. A belief with no support relation pointing at it is unsupported:

def unsupported (agent : String) : Tx .read (Array (Col "text" String)) :=
  query typeql!"match
    $a isa agent, has name {agent};
    $b isa belief, links (holder: $a, content: $o);
    not \{ $s isa support, links (supported: $b); };
    $o has text $text;"

In typeql!, TypeQL's own braces are written \{ and \}, because a plain {…} is interpolation.

Everything known about a topic

Questions you ask often can live in the schema as TypeQL functions. The server type-checks them when they're defined:

define
  fun about($t: topic) -> { observation }:
    match
      $m isa mention, links (source: $o, topic: $t);
    return { $o };
match
  $t isa topic, has name "bridge";
  let $o in about($t);
  $o has text $text;

Where do two agents disagree?

Because beliefs are per agent, contradictions between agents are one pattern. A team supervisor can run this before agents act on shared plans:

match
  $a isa agent;
  $b isa belief, links (holder: $a, content: $mine);
  $c isa contradiction, links (claim: $mine, claim: $other);
  $b2 isa belief, links (holder: $a2, content: $other);
  $a has name $who; $a2 has name $whom;

Show the evidence, as JSON for the prompt

When the agent builds its next prompt, a fetch query returns documents ready to paste in, with each claim next to its evidence:

match
  $b isa belief, links (holder: $a, content: $o);
  $s isa support, links (supported: $b, evidence: $e);
  $a has name $agent; $o has text $claim; $e has text $evidence;
fetch { "agent": $agent, "claim": $claim, "evidence": $evidence };

In Lean, fetch (α := T) decodes those documents into any type with a Lean.FromJson instance.

Many agents, one memory

Each Driver.run is a Std.Async computation on its own thread, and one driver connection multiplexes them all. An agent that recalls two topics before planning can do it concurrently:

def recallBoth (d : Driver) (t1 t2 : String) : EIO TypeDBError (Array Recalled × Array Recalled) :=
  (Std.Async.EAsync.concurrently
    (d.run "memory" .read (recallAbout t1))
    (d.run "memory" .read (recallAbout t2))).block

If either transaction fails, the whole computation fails with that transaction's typed error.

Start-up checks and an evolving memory

Memory outlives code. The agent you deploy next month will open a database written by this month's version. Two features keep them in step.

Verification. At start-up, Schema.verify! compares the Lean structures with the live schema. It fails with TypeDBError.schemaMismatch and a list of issues, such as type 'observation' does not own 'confidence', instead of letting the first recall fail an hour later.

Migrations. Suppose observations now need an origin: a sensor, a web page, a user. Migrations are versioned steps, and each step runs in its own schema transaction together with its history record:

def migrations : List Migration := [
  Migration.define 1 "agent memory" schema,
  Migration.tql 2 "observation origin"
    "define attribute origin, value string; observation owns origin;"
    (down := some "undefine owns origin from observation; origin;"),
  { version := 3, name := "default origin", fingerprint := "origin := unknown, v1"
    up := exec typeql!"match $o isa observation; not \{ $o has origin $x; }; insert $o has origin \"unknown\";"
    down := .step (exec typeql!"match $o isa observation, has origin $x; delete has $x of $o;") }
]

Step 3 is Lean code that backfills existing observations. Lean code can't be hashed, so it carries a fingerprint. If an applied step's checksum later changes, migrate stops before running anything. Your agents never run against a history they didn't write.

Mistakes that don't compile

A memory layer is written once and called from everywhere, which is where small mistakes hide. Two of the most common ones never reach the database. Writing in a read transaction:

def forget : Tx .read Unit :=
  exec typeql!"match $o isa observation; insert $o has text \"x\";"
error: failed to synthesize instance of type class
  TxMode.read.Allows TxMode.write

Interpolating something that isn't a value, such as a whole request object:

structure UserInput where
  raw : String

def lookup (u : UserInput) : Tx .read (Array Recalled) :=
  query typeql!"match $t isa topic, has name {u}; …"
error: failed to synthesize instance of type class
  ToTypeQL UserInput

More of these are in Five bugs the compiler catches for you.

Putting it together

def main : IO Unit := do
  let r ← (Driver.withDriver { address := "127.0.0.1:1729" } fun d => do
    d.ensureDatabase "memory"
    discard <| d.migrateSync "memory" migrations
    d.runSync "memory" .read (Schema.verify! schema)
    d.runSync "memory" .write do
      discard <| Entity.insert { name := "scout" : Agent }
    let now ← PlainDateTime.now
    let b ← d.runSync "memory" .write (retry := .conflicts 5 (.exponential 20))
      (remember "scout" "obs-1" "The bridge on Route 9 is closed for repairs." now (pct 90) ["route-9", "bridge"])
    let ev ← d.runSync "memory" .write
      (Entity.insert { id := "obs-2", text := "A road sign says: bridge closed until Friday.",
                       observedAt := now, confidence := pct 70 : Observation })
    d.runSync "memory" .write (support b ev (pct 60))
    for r in ← d.runSync "memory" .read (recallAbout "bridge") do
      IO.println s!"{r.confidence}  {r.text}"
    : EIO TypeDBError Unit).toBaseIO
  if let .error e := r then IO.eprintln e

Here support is a one-line helper around Relation.insert { supported := belief, evidence, weight : Support }. Every Lean snippet in this post compiles against typedb-lean, and every TypeQL query in it ran against TypeDB 3.14.0 on a database with this schema.

Where to go next

TypeDB for Edge AI AgentsVolodymyr Pavlyshyn
Go deeper: TypeDB for Edge AI Agents

This post is a small slice of a larger subject. The book covers TypeDB's type system and TypeQL in depth, and how to use a typed knowledge graph as the memory of AI agents running at the edge.

Get it on Leanpub →