Chapter 2: Memory

// Storage: text describing the user's attitude and its cause
const userContext = {
  userId: 'alice',
  attitudes: [
    {
      topic: 'cloud infrastructure',
      sentiment: 'skeptical',
      reason: 'Experienced significant downtime during AWS outage in Q2 2024, affecting client deliverables'
    }
  ]
};

contextStore.set('alice', userContext);

// Later, when building context, you serialize it:
const context = contextStore.get('alice');
const systemPrompt = `User context: ${JSON.stringify(context, null, 2)}`;

// The model reads this string and interprets its meaning
const messages = [
  { role: 'system', content: systemPrompt },
  { role: 'user', content: 'Should we migrate to Cloudflare Workers?' }
];

const reply = await llm.complete(messages);
// reply: "Given your past experience with cloud outages, let me address
// reliability concerns first. Cloudflare has a different architecture..."

// Traditional approach: explicit logic interprets the data
if (user.units === 'metric') {
  return `${distanceInKm} km`;
} else {
  return `${distanceInMiles} miles`;
}

const profile = {
  name: 'Alice',
  company: 'Acme Corp',
  preferences: { units: 'metric' },
  projects: ['European expansion']
};

// Serialized directly to string for context:
const systemPrompt = `User profile: ${JSON.stringify(profile, null, 2)}`;
// User profile: {
//   "name": "Alice",
//   "company": "Acme Corp",
//   "preferences": { "units": "metric" },
//   "projects": ["European expansion"]
// }

// Unstructured representation: natural language
const profile = `
Alice works at Acme Corp, a company with a budget of $50k.
She prefers metric units and is working on expanding their
European market presence.
`;

// Structured representation: subject-predicate-object triplets
const facts = [
  { subject: 'Alice', predicate: 'works_at', object: 'Acme Corp' },
  { subject: 'Acme Corp', predicate: 'budget', object: '$50k' },
  { subject: 'Alice', predicate: 'prefers', object: 'metric units' }
];

// When building context, these become strings:
const factStrings = facts.map(f =>
  `${f.subject} ${f.predicate.replace('_', ' ')} ${f.object}`
);

const systemPrompt = `Known facts:\n${factStrings.join('\n')}`;
// Known facts:
// Alice works at Acme Corp
// Acme Corp budget $50k
// Alice prefers metric units

type Preferences = { name?: string; units?: 'metric' | 'imperial' };

const preferencesStore = new Map<string, Preferences>();

async function chat(userId: string, input: string) {
  // 1. Load memory from external storage
  const prefs = preferencesStore.get(userId) ?? {};

  // 2. Build context that *includes* remembered data
  const systemPrompt = [
    prefs.name && `The user's name is ${prefs.name}.`,
    prefs.units && `The user prefers ${prefs.units} units.`,
  ]
    .filter(Boolean)
    .join(' ');

  const messages = [
    { role: 'system', content: systemPrompt },
    { role: 'user', content: input },
  ];

  // 3. Call the stateless model with this constructed context
  const reply = await llm.complete(messages);

  // 4. Optionally *update* memory based on this turn
  const updatedPrefs = extractPreferencesFromTurn(input, reply, prefs);
  preferencesStore.set(userId, updatedPrefs);

  return reply;
}

// Structured memory: easy to update specific facts
const profile = memoryStore.get('alice');
profile.attitudes[0].sentiment = 'neutral'; // User's skepticism has faded
profile.attitudes[0].reason += '. Recent experience with Cloudflare has been positive.';
memoryStore.set('alice', profile);

// Prose memory: harder to update specific facts
let profile = `
Alice is skeptical of cloud infrastructure after a bad AWS outage.
She prefers on-premise solutions.
`;

// To update, you either:
// 1. Rewrite the entire string manually
profile = `
Alice was initially skeptical of cloud infrastructure after a bad AWS outage,
but recent positive experience with Cloudflare has made her more open to it.
She still values reliability over convenience.
`;

// 2. Or ask the model to update it
const updated = await llm.complete([
  { role: 'system', content: 'Update this user profile to reflect that their cloud skepticism has faded due to positive recent experience.' },
  { role: 'user', content: profile }
]);
profile = updated;

type Message = { role: 'user' | 'assistant'; content: string };

const messageLog: Message[] = [];

function storeTurn(user: string, assistant: string) {
  messageLog.push({ role: 'user', content: user });
  messageLog.push({ role: 'assistant', content: assistant });
}

function recentMessages(limit: number): Message[] {
  return messageLog.slice(-limit);
}

type Profile = { company?: string; role?: string };

const profiles = new Map<string, Profile>();

function rememberProfile(userId: string, profile: Profile) {
  profiles.set(userId, profile);
}

function getProfile(userId: string): Profile | undefined {
  return profiles.get(userId);
}

type KnowledgeItem = { id: string; content: string };

const knowledge = new VectorStore<KnowledgeItem>();

async function rememberFact(item: KnowledgeItem) {
  await knowledge.add(item.content, item);
}

async function recallRelevant(query: string, k = 5): Promise<KnowledgeItem[]> {
  return knowledge.search(query, { limit: k });
}

type Event = {
  id: string;
  content: string;
  timestamp: number;
  important?: boolean;
};

// Episodic store: append-only log
const events: Event[] = [];

// Semantic index: meaning-based retrieval
const eventIndex = new VectorStore<Event>();

async function rememberEvent(event: Event) {
  events.push(event);
  await eventIndex.add(event.content, event);
}

async function buildContextFor(input: string, tokenBudget: number) {
  const now = Date.now();

  // 1. Recency: take the latest N events
  const recent = events.slice(-20);

  // 2. Semantic: events similar to the current input
  const similar = await eventIndex.search(input, { limit: 20 });

  // 3. Importance: all events marked important, regardless of age
  const important = events.filter(e => e.important);

  // 4. Merge with simple scoring
  const scored = new Map<string, { event: Event; score: number }>();

  function bump(e: Event, delta: number) {
    const existing = scored.get(e.id) ?? { event: e, score: 0 };
    existing.score += delta;
    scored.set(e.id, existing);
  }

  for (const e of recent) bump(e, 1.0);
  for (const e of similar) bump(e, 2.0);
  for (const e of important) bump(e, 5.0);

  // 5. Sort by score (and slight bias toward more recent within ties)
  const ranked = Array.from(scored.values())
    .sort(
      (a, b) =>
        b.score - a.score ||
        b.event.timestamp - a.event.timestamp,
    )
    .map(s => s.event);

  // 6. Pack into context until you hit the token budget
  const selected: Event[] = [];
  let used = 0;

  for (const e of ranked) {
    const cost = estimateTokens(e.content);
    if (used + cost > tokenBudget) break;
    selected.push(e);
    used += cost;
  }

  return selected.map(e => ({ role: 'system', content: e.content }));
}

type Turn = { role: 'user' | 'assistant'; content: string; timestamp: number };

let turns: Turn[] = [];
let summary: string | null = null;

async function maybeCompressHistory() {
  // Only compress when we have many turns
  if (turns.length < 40) return;

  const old = turns.slice(0, -10);
  const recent = turns.slice(-10);

  const prompt = [
    {
      role: 'system',
      content:
        'Summarize this interaction. Preserve goals, decisions, and enduring preferences. Omit greetings and small talk.',
    },
    { role: 'user', content: JSON.stringify(old) },
  ];

  const summaryText = await llm.complete(prompt);

  summary = summary
    ? `${summary}\n\nAdditional context:\n${summaryText}`
    : summaryText;

  // Replace old turns with a single synthetic system message
  turns = [
    { role: 'system', content: `Summary of earlier conversation:\n${summary}`, timestamp: Date.now() },
    ...recent,
  ];
}

const ONE_WEEK = 7 * 24 * 60 * 60 * 1000;

function purgeEphemeralEvents(now: number = Date.now()) {
  // Keep only recent events or those explicitly marked important
  events = events.filter(
    e => e.important || now - e.timestamp < ONE_WEEK,
  );
}

function markImportant(id: string) {
  const event = events.find(e => e.id === id);
  if (event) event.important = true;
}

function pruneUnimportant(maxSize: number) {
  // Keep all important events, plus the most recent unimportant ones
  const important = events.filter(e => e.important);
  const unimportant = events.filter(e => !e.important).slice(-maxSize);
  events = [...important, ...unimportant];
}

export class Assistant {
  // Episodic: full turns plus summary
  private turns: Turn[] = [];
  private summary: string | null = null;

  // Structured: preferences and project state
  private preferences: Preferences = {};
  private project: { name: string; objective: string } | null = null;

  // Semantic: knowledge across sessions
  private knowledge = new VectorStore<{ content: string; timestamp: number }>();

  async chat(userId: string, input: string) {
    // 1. Update memory with the new user turn (episodic)
    this.turns.push({ role: 'user', content: input, timestamp: Date.now() });

    // 2. Retrieve from memory for this specific context
    const context = await this.buildContext(input);

    // 3. Call the model
    const reply = await llm.complete(context);

    // 4. Store assistant turn and maybe extract new persistent facts
    this.turns.push({ role: 'assistant', content: reply, timestamp: Date.now() });

    await this.updateStructuredMemory(input, reply);
    await this.updateSemanticMemory(input, reply);

    // 5. Manage growth
    await this.maybeCompressHistory();
    this.pruneUnimportantEvents();

    return reply;
  }

  private async buildContext(input: string) {
    const tokenBudget = 2000;

    // 1. Preferences: direct lookup
    const prefsText = this.renderPreferences();

    // 2. Project state: direct lookup
    const projectText = this.project
      ? `Current project: ${this.project.name}. Objective: ${this.project.objective}.`
      : '';

    // 3. Episodic + semantic hybrid retrieval
    const relevantEvents = await buildContextFor(input, tokenBudget / 2);

    // 4. Remaining budget goes to immediate recent turns
    const recentTurns = this.turns.slice(-10);

    return [
      { role: 'system', content: prefsText },
      projectText && { role: 'system', content: projectText },
      ...relevantEvents,
      ...recentTurns,
      { role: 'user', content: input },
    ].filter(Boolean) as { role: string; content: string }[];
  }

  private renderPreferences(): string {
    const { name, units } = this.preferences;
    const parts = [];
    if (name) parts.push(`The user's name is ${name}.`);
    if (units) parts.push(`The user prefers ${units} units.`);
    return parts.join(' ') || 'You are talking to a returning user.';
  }

  private async updateStructuredMemory(input: string, reply: string) {
    // Use the model (or rules) to extract durable preferences and project info
    const extractionPrompt = [
      {
        role: 'system',
        content:
          'From the following exchange, extract any stable preferences (like name, units) and project info.',
      },
      { role: 'user', content: input },
      { role: 'assistant', content: reply },
    ];
    const json = await llm.completeAsJson<{
      preferences?: Preferences;
      project?: { name?: string; objective?: string };
    }>(extractionPrompt);

    if (json.preferences) {
      this.preferences = { ...this.preferences, ...json.preferences };
    }
    if (json.project) {
      this.project = { ...(this.project ?? {}), ...json.project } as any;
    }
  }

  private async updateSemanticMemory(input: string, reply: string) {
    // Treat assistant replies as candidate knowledge
    await this.knowledge.add(reply, { content: reply, timestamp: Date.now() });
  }

  private async maybeCompressHistory() {
    // As in the earlier example
  }

  private pruneUnimportantEvents() {
    // As in the earlier example
  }
}