index.ts

import { EMBEDDING_PATH, VECTOR_INDEX } from "./config.js";
import { extractFacts } from "./extract.js";
import { isReady, memories, messages } from "./store.js";
import type { MemoryScope, MessageDoc, RecalledMemory, Role } from "./types.js";

export { attach, detach, isReady } from "./store.js";
export type { MemoryScope, RecalledMemory } from "./types.js";

/**
 * Write a message down, and keep any durable facts it contains.
 *
 * Returns the facts that were stored so the caller can show them. Seeing the
 * agent say "I will remember that you are vegetarian" is much more convincing
 * than trusting that something happened in the background.
 */
export async function remember(
  scope: MemoryScope,
  role: Role,
  content: string,
  options: { extract?: boolean } = {}
): Promise<string[]> {
  if (!isReady()) return [];

  const now = new Date();

  await messages().insertOne({
    threadId: scope.threadId,
    userId: scope.userId,
    role,
    content,
    createdAt: now,
  });

  // Only user messages are worth mining. The assistant repeating your own
  // preferences back at you is not new information.
  const shouldExtract = options.extract ?? role === "user";
  if (!shouldExtract) return [];

  const facts = await extractFacts(content);
  if (facts.length === 0) return [];

  const stored: string[] = [];

  for (const { subject, fact } of facts) {
    // No embedding call here. We write the fact as plain text and Atlas
    // generates the vector for it, because the index declares this field
    // with the "autoEmbed" type.
    //
    // Upsert on subject. If the person told us last week that they love spicy
    // food and today that spicy food destroys them, we keep today's version
    // instead of storing both and confusing the agent.
    await memories().updateOne(
      { threadId: scope.threadId, subject },
      {
        $set: {
          userId: scope.userId,
          fact,
          sourceText: content,
          updatedAt: now,
        },
        $setOnInsert: { threadId: scope.threadId, subject, createdAt: now },
      },
      { upsert: true }
    );

    stored.push(fact);
  }

  return stored;
}

/**
 * Find the memories that matter for what the person just said.
 *
 * The filter is the important line. It keeps this search inside one thread.
 * Delete it and the agent will happily answer one person using another
 * person's memories.
 */
export async function recall(
  scope: MemoryScope,
  query: string,
  limit = 5
): Promise<RecalledMemory[]> {
  if (!isReady()) return [];

  return memories()
    .aggregate<RecalledMemory>([
      {
        $vectorSearch: {
          index: VECTOR_INDEX,
          path: EMBEDDING_PATH,
          // Plain text, not a vector. Atlas embeds this with the same model the
          // index declares. The shape is { text }, not a bare string.
          query: { text: query },
          numCandidates: limit * 15,
          limit,
          filter: {
            threadId: { $eq: scope.threadId },
            userId: { $eq: scope.userId },
          },
        },
      },
      {
        $project: {
          _id: 0,
          subject: 1,
          fact: 1,
          updatedAt: 1,
          score: { $meta: "vectorSearchScore" },
        },
      },
    ])
    .toArray();
}

/** The last few lines of the transcript, oldest first. No vectors involved. */
export async function history(
  scope: MemoryScope,
  limit = 10
): Promise<Array<Pick<MessageDoc, "role" | "content">>> {
  if (!isReady()) return [];

  const recent = await messages()
    .find(
      { threadId: scope.threadId, userId: scope.userId },
      { projection: { _id: 0, role: 1, content: 1 } }
    )
    .sort({ createdAt: -1 })
    .limit(limit)
    .toArray();

  return recent.reverse();
}

/** Delete a single subject, or wipe the whole thread. */
export async function forget(
  scope: MemoryScope,
  subject?: string
): Promise<number> {
  if (!isReady()) return 0;

  if (subject) {
    const result = await memories().deleteOne({
      threadId: scope.threadId,
      subject: subject.toLowerCase(),
    });
    return result.deletedCount;
  }

  const result = await memories().deleteMany({ threadId: scope.threadId });
  await messages().deleteMany({ threadId: scope.threadId });
  return result.deletedCount;
}

/** List everything stored for a thread. Used by the inspector. */
export async function listMemories(scope: MemoryScope) {
  if (!isReady()) return [];

  return memories()
    .find(
      { threadId: scope.threadId },
      { projection: { _id: 0, subject: 1, fact: 1, updatedAt: 1 } }
    )
    .sort({ updatedAt: -1 })
    .toArray();
}