Research Agent Memory

1"""Research Agent Memory: durable, searchable memory for an AI research assistant, in Pixeltable tables.
2
3Sources are the papers and pages the assistant has read (with an optional cover image). Memories are
4facts, preferences and episodes per user, embedded with OpenAI text-embedding-3-small into an
5EmbeddingIndex. recall_memories (@pxt.query) is semantic recall scoped to one user; lookup_source
6(@pxt.query) is an exact-key lookup the LLM can call as a tool. Turns is the agent table: each
7message recalls memories, calls gpt-4o-mini with tools, and proposes a memory to write back.
8run.py also builds pxt.retrieval_udf lookups on the live tables and consolidates memories with the
9salience and memory_digest aggregates from udfs.py.
10
11Layout: UDFs and UDAs live in udfs.py; this file holds table models and @pxt.query functions.
12"""
13
14from __future__ import annotations
15
16import pixeltable as pxt
17import pixeltable.functions as pxtf
18from pixeltable.functions import openai
19
20from udfs import (
21    CHAT_MODEL,
22    EMBED_MODEL,
23    agent_messages,
24    answer_messages,
25    extract_memory,
26    memory_block,
27    pick_answer,
28    reading_tier,
29)
30
31TableModel = pxt.model_base()
32
33
34class Sources(TableModel, name='sources', has_default_idxs=False):
35    """A paper, web page or note the assistant has read."""
36
37    source_id = pxt.Column(type=pxt.String, primary_key=True)
38    title: pxt.String
39    kind: pxt.String
40    year: pxt.Int
41    url: pxt.String | None
42    cover: pxt.Image | None
43
44    cover_width = cover.width  # type: ignore[attr-defined]
45
46
47class Memories(TableModel, name='memories', has_default_idxs=False):
48    """Long-term memory: one fact, preference, episode or summary per row, per user."""
49
50    memory_id = pxt.Column(value=pxtf.uuid.uuid7(), primary_key=True)
51    user_id: pxt.String
52    kind: pxt.String
53    content: pxt.String
54    importance: pxt.Int
55    created_at: pxt.Timestamp
56    source_id: pxt.String | None
57
58    __indexes__ = [
59        pxt.EmbeddingIndex(
60            content,  # type: ignore[name-defined]  # noqa: F821
61            string_embed=openai.embeddings.using(model=EMBED_MODEL),
62            metric='cosine',
63            name='memory_embed',
64        ),
65        pxt.BtreeIndex(user_id),  # type: ignore[name-defined]  # noqa: F821
66    ]
67
68
69@pxt.query
70def recall_memories(user_id: str, context: str, k: int = 4):
71    """Top-k memories of one user that are semantically closest to the context."""
72    sim = Memories.content.similarity(string=context, idx='memory_embed')
73    return (
74        Memories.where(Memories.user_id == user_id)
75        .order_by(sim, asc=False)
76        .limit(k)
77        .select(Memories.kind, Memories.content, Memories.source_id, score=sim)
78    )
79
80
81@pxt.query
82def lookup_source(source_id: str):
83    """Look up one source (title, kind, year, url) by its id, e.g. S-02."""
84    return Sources.where(Sources.source_id == source_id).limit(1).select(
85        Sources.source_id, Sources.title, Sources.kind, Sources.year, Sources.url
86    )
87
88
89@pxt.query
90def recent_memories(user_id: str, n: int = 5):
91    """The newest n memories of one user."""
92    return (
93        Memories.where(Memories.user_id == user_id)
94        .order_by(Memories.created_at, asc=False)
95        .limit(n)
96        .select(Memories.kind, Memories.content, Memories.created_at)
97    )
98
99
100# Keep at least one @pxt.udf in the toolset: with only @pxt.query tools, invoke_tools waits
101# forever on a response that has no tool calls (seen on pixeltable 0.7.16).
102TOOLS = pxt.tools(lookup_source, reading_tier)
103
104
105class Turns(TableModel, name='turns', has_default_idxs=False):
106    """Agent table: message -> recall -> tool-calling LLM -> answer + a proposed memory write-back."""
107
108    turn_id = pxt.Column(value=pxtf.uuid.uuid7(), primary_key=True)
109    user_id: pxt.String
110    message: pxt.String
111
112    recalled = recall_memories(user_id, message, 4)
113    memory_text = memory_block(recalled)
114    response = openai.chat_completions(
115        agent_messages(user_id, message, memory_text),
116        model=CHAT_MODEL,
117        model_kwargs={'max_tokens': 150, 'temperature': 0},
118        tools=TOOLS,
119    )
120    tool_output = openai.invoke_tools(TOOLS, response)
121    final = openai.chat_completions(
122        answer_messages(message, memory_text, tool_output),
123        model=CHAT_MODEL,
124        model_kwargs={'max_tokens': 160, 'temperature': 0},
125    )
126    answer = pick_answer(response.choices[0].message.content, final.choices[0].message.content, tool_output)
127    remember = extract_memory(message)