Example: an agent's memory

The project

import Lake
open Lake DSL

package «agent-memory»

require ladybug from git "https://github.com/Volland/ladybug-lean" @ "v0.1.0"

lean_lib AgentMemory

@[default_target]
lean_exe «agent-memory» where
  root := `Main

lake build AgentMemory     # compiles the schema and runs every check
lake exe agent-memory      # runs on a LadybugDB file, memory.lbug

1. The ontology

The whole schema is five structures and five declarations. Nothing else describes the database: the DDL, the decoders and the typed query variables are derived from these.

structure Entity where
  name : String
  kind : String

structure Memory where
  id : Int64
  text : String
  importance : Int32
  happenedAt : Timestamp
  embedding : Embedding 4

structure Mentions where
  weight : Option Float

structure Relates where
  label : String

structure Next

node_table Entity (key := name)
node_table Memory
rel_table Mentions (from := Memory) (to := Entity)
rel_table Relates (from := Entity) (to := Entity)
rel_table Next (from := Memory) (to := Memory)

def schema : Schema :=
  [.ofNode Entity, .ofNode Memory, .ofRel Mentions, .ofRel Relates, .ofRel Next]

2. Remembering and recalling

Every function is written once against MonadGraph, so it runs on LadybugDB and on the in-memory graph. Notice the mix: the repository for writes, the builder for a vector search and for paths, and plain Cypher for a lookup, all checked against the schema.

variable {m : Type → Type} [MonadGraph m]

/-- Store a memory, link it to the entities it mentions and to the previous memory. -/
def remember (mem : Memory) (mentions : List (String × Float)) (previous : Option Int64 := none) : m Unit :=
  withTx do
    Repo.insert mem
    for (name, w) in mentions do
      Repo.link mem.id (Mentions.mk (some w)) name
    if let some prev := previous then
      Repo.link prev Next.mk mem.id

/-- The `k` memories closest in meaning to `q`, nearest first. -/
def recallSimilar (q : Embedding 4) (k : Nat) : m (Array (String × Float)) :=
  runQueryM do
    let (mem, dist) ← vectorSearch Memory "memory_vec" q k
    orderBy dist
    return_ (mem.text, dist)

/-- The memories that mention an entity, most important first. -/
def about (entity : String) : m (Array (String × Int32)) :=
  runQuery (cypher%
    MATCH (mem:Memory)-[:Mentions]->(e:Entity)
    WHERE e.name = $entity
    RETURN mem.text, mem.importance ORDER BY mem.importance DESC)

/-- Entities reachable from one entity in one or two `Relates` hops. -/
def neighbourhood (entity : String) : m (Array String) :=
  runQueryM do
    let e ← match_ (node Entity)
    let (_, other) ← matchPaths e Relates (node Entity) (max := 2) (dir := .both)
    where_ (e.name ==. lit entity &&. other.name !=. lit entity)
    return_ other.name (distinct := true)

/-- How many steps of the timeline separate two memories. -/
def distance (a b : Int64) : m (Option Int64) := do
  let rows ← runQueryM do
    let x ← match_ (node Memory)
    let (path, y) ← shortestPath x Next (node Memory) (max := 30)
    where_ (x.id ==. lit a &&. y.id ==. lit b)
    return_ (pathLength path)
  return rows[0]?

/-- Entities no memory mentions any more. -/
def forgotten : m (Array String) :=
  runQueryM do
    let e ← match_ (node Entity)
    where_ (Expr.not (← exists_ (discard (match_ (inc e Mentions (node Memory))))))
    return_ e.name

3. Checked while compiling

This file has no main. It is checked every time it builds: #guard runs the memory functions on the in-memory backend and compares the results, and #guard_msgs pins the error for each query the schema must reject. Change a function, a query or the schema in a way that changes behaviour, and the build fails.

def onSeed {α : Type} (prog : MemM α) : Except Error α := (do seed; prog).run' {}

#guard (onSeed (about "Ann")).toOption == some #[("Ann started the Ladybug bindings", 5),
  ("Meetup in Kyiv about graph memory", 4), ("Bob reviewed Ann's first query", 3)]

#guard (onSeed (distance 1 4)).toOption == some (some 3)

#guard (onSeed forgotten).toOption == some #["Kyiv"]

/-- error: unknown property 'importanse' of table Memory; did you mean 'importance'? -/
#guard_msgs in
example := cypher% MATCH (mem:Memory) RETURN mem.importanse

-- `Mentions` goes from a memory to an entity, not the other way round.
/-- error: failed to synthesize instance of type class Matchable (PatOut Entity.Var Mentions (NodeRef Memory)) -/
#guard_msgs (whitespace := lax, substring := true) in
example : Query String := query do
  let e ← match_ (node Entity)
  let (_, mem) ← match_ (out e Mentions (node Memory))
  return_ mem.text

-- Importance is an `Int32`: comparing it with a string is a type error.
/-- error: failed to synthesize instance of type class Expr.LitFor Int32 String -/
#guard_msgs (whitespace := lax, substring := true) in
example : Query String := query do
  let mem ← match_ (node Memory)
  where_ (mem.importance >. lit "high")
  return_ mem.text

-- An optional match makes every property nullable, and the result type says so.
example : Query (String × Option String) :=
  cypher% MATCH (mem:Memory) OPTIONAL MATCH (mem)-[:Mentions]->(e:Entity) RETURN mem.text, e.kind

4. Running on LadybugDB

The executable runs the same functions on a database file. Its output agrees with what the checks assert for the in-memory backend:

def main : IO Unit :=
  Ladybug.withDatabase "memory.lbug" fun db => db.withConnection do
    seed
    IO.println s!"about Ann:        {← about "Ann"}"
    IO.println s!"near Ann:         {← neighbourhood "Ann"}"
    IO.println s!"memories 1 → 4:   {(← distance 1 4).getD 0} steps"
    IO.println s!"like memory 1:    {← recallSimilar (vec [1, 0, 0, 0]) 2}"
    IO.println s!"nobody mentions:  {← forgotten}"

about Ann: #[(Ann started the Ladybug bindings, 5), (Meetup in Kyiv about graph memory, 4), (Bob reviewed Ann's first query, 3)] near Ann: #[Ladybug, Bob, Lean] memories 1 → 4: 3 steps like memory 1: #[(Ann started the Ladybug bindings, 0.000000), (Bob reviewed Ann's first query, 0.029858)] nobody mentions: #[Kyiv]