convex/memory.ts

/**
 * AI Town — Memory System (Core of the Agent Architecture)
 * ==========================================================
 * Implements the Stanford "Generative Agents" memory stream:
 * - Store observations with importance + embedding
 * - Retrieve by: recency × importance × relevance (cosine similarity)
 * - Trigger reflection when cumulative importance > threshold
 */

import { v } from "convex/values";
import { mutation, query, action, internalAction } from "./_generated/server";
import { internal } from "./_generated/api";

// ═══ STORE NEW MEMORY ═══

export const addMemory = mutation({
  args: {
    characterId: v.id("characters"),
    description: v.string(),
    type: v.string(),
    importance: v.number(),
    embedding: v.array(v.float64()),
    participants: v.optional(v.array(v.string())),
  },
  handler: async (ctx, args) => {
    const now = Date.now();
    return await ctx.db.insert("memories", {
      characterId: args.characterId,
      description: args.description,
      type: args.type,
      importance: args.importance,
      embedding: args.embedding,
      participants: args.participants,
      createdAt: now,
      lastAccessed: now,
    });
  },
});

// ═══ RETRIEVE RELEVANT MEMORIES (RAG) ═══

export const retrieveMemories = action({
  args: {
    characterId: v.id("characters"),
    queryEmbedding: v.array(v.float64()),
    limit: v.optional(v.number()),
  },
  handler: async (ctx, args) => {
    const limit = args.limit || 10;

    // Vector search for semantically relevant memories
    const results = await ctx.vectorSearch("memories", "by_embedding", {
      vector: args.queryEmbedding,
      limit: limit * 2, // Get extra for scoring
      filter: (q) => q.eq("characterId", args.characterId),
    });

    // Score each memory: recency × importance × relevance
    const now = Date.now();
    const scored = results.map((result) => {
      const hoursSinceAccess = (now - result.lastAccessed) / (1000 * 60 * 60);
      const recencyScore = Math.pow(0.995, hoursSinceAccess);
      const importanceScore = result.importance / 10;
      const relevanceScore = result._score; // cosine similarity from vector search

      return {
        ...result,
        finalScore: recencyScore + importanceScore + relevanceScore,
      };
    });

    // Sort by combined score and return top N
    scored.sort((a, b) => b.finalScore - a.finalScore);
    const topMemories = scored.slice(0, limit);

    // Update lastAccessed for retrieved memories
    for (const mem of topMemories) {
      await ctx.runMutation(internal.memory.touchMemory, { memoryId: mem._id });
    }

    return topMemories;
  },
});

// Internal: update lastAccessed
export const touchMemory = mutation({
  args: { memoryId: v.id("memories") },
  handler: async (ctx, { memoryId }) => {
    await ctx.db.patch(memoryId, { lastAccessed: Date.now() });
  },
});

// ═══ GET RECENT MEMORIES (short-term buffer) ═══

export const getRecentMemories = query({
  args: {
    characterId: v.id("characters"),
    limit: v.optional(v.number()),
  },
  handler: async (ctx, args) => {
    return await ctx.db
      .query("memories")
      .withIndex("by_character", (q) => q.eq("characterId", args.characterId))
      .order("desc")
      .take(args.limit || 20);
  },
});

// ═══ CHECK REFLECTION THRESHOLD ═══

export const shouldReflect = query({
  args: { characterId: v.id("characters") },
  handler: async (ctx, { characterId }) => {
    // Get memories since last reflection
    const lastReflection = await ctx.db
      .query("memories")
      .withIndex("by_character_type", (q) =>
        q.eq("characterId", characterId).eq("type", "reflection")
      )
      .order("desc")
      .first();

    const since = lastReflection?.createdAt || 0;

    const recentMemories = await ctx.db
      .query("memories")
      .withIndex("by_character", (q) => q.eq("characterId", characterId))
      .order("desc")
      .take(100);

    const newMemories = recentMemories.filter((m) => m.createdAt > since);
    const importanceSum = newMemories.reduce((sum, m) => sum + m.importance, 0);

    // Threshold: 150 (from Stanford paper)
    return { shouldReflect: importanceSum >= 150, importanceSum, memoryCount: newMemories.length };
  },
});

// ═══ GET MEMORY COUNT ═══

export const getMemoryCount = query({
  args: { characterId: v.id("characters") },
  handler: async (ctx, { characterId }) => {
    const memories = await ctx.db
      .query("memories")
      .withIndex("by_character", (q) => q.eq("characterId", characterId))
      .collect();
    return memories.length;
  },
});