evOo-local/packages/memory/src/index.ts

import type {
  AdapterContext,
  MemoryCommitRequest,
  MemoryExportRequest,
  MemoryReadRequest,
  MemoryRevision,
  MemorySearchRequest,
  MemorySearchResult,
  MemorySnapshot,
  MemoryStore,
  PortableFile,
} from "@rakazo/adapter-kit";
import { MEMORY_DOCUMENT_MAX_BYTES, MemoryDocumentTooLargeError } from "@rakazo/adapter-kit";
import type { PrismaClient } from "@rakazo/db";
import { Prisma, withTransactionRetry } from "@rakazo/db";

export class MarkdownMemoryStore implements MemoryStore {
  constructor(private readonly prisma: PrismaClient) {}

  describe() {
    return {
      id: "markdown",
      contractVersion: "1",
      adapterVersion: "0.1.0",
      capabilities: { search: true, revisions: true, markdownPortable: true },
    };
  }

  async read(request: MemoryReadRequest, context: AdapterContext): Promise<MemorySnapshot> {
    const documents = await this.prisma.memoryDocument.findMany({
      where: {
        spaceId: context.spaceId,
        userId: context.userId,
        scope: request.scope,
        ...(request.botId ? { botId: request.botId } : {}),
        ...(request.path ? { path: request.path } : {}),
      },
      orderBy: [{ updatedAt: "desc" }, { path: "asc" }],
    });
    return {
      documents: documents.map((doc) => ({
        id: doc.id,
        path: doc.path,
        content: doc.content,
        revision: doc.revision,
        updatedAt: doc.updatedAt.toISOString(),
      })),
    };
  }

  async search(
    request: MemorySearchRequest,
    context: AdapterContext,
  ): Promise<MemorySearchResult[]> {
    const documents = await this.prisma.memoryDocument.findMany({
      where: {
        spaceId: context.spaceId,
        userId: context.userId,
        ...(request.scope === "all" ? {} : { scope: request.scope }),
        ...(request.botId ? { botId: request.botId } : {}),
      },
    });
    return rankMemoryDocuments(documents, request.query);
  }

  async commit(request: MemoryCommitRequest, context: AdapterContext): Promise<MemoryRevision> {
    // A save replaces (or appends to) the current content and appends its history atomically.
    // Conflicts retry the whole transaction so revision numbers and appends use fresh state.
    return withTransactionRetry(() =>
      this.prisma.$transaction(
        async (tx) => {
          const existing = await tx.memoryDocument.findFirst({
            where: {
              spaceId: context.spaceId,
              userId: context.userId,
              scope: request.scope,
              botId: request.botId ?? null,
              path: request.path,
            },
          });
          const content =
            request.mode === "append"
              ? appendMarkdown(existing?.content ?? "", request.content)
              : request.content;
          if (Buffer.byteLength(content, "utf8") > MEMORY_DOCUMENT_MAX_BYTES) {
            throw new MemoryDocumentTooLargeError(request.path);
          }
          const doc = existing
            ? await tx.memoryDocument.update({
                where: { id: existing.id },
                data: { content, revision: existing.revision + 1 },
              })
            : await tx.memoryDocument.create({
                data: {
                  spaceId: context.spaceId,
                  userId: context.userId,
                  botId: request.botId,
                  scope: request.scope,
                  path: request.path,
                  content,
                },
              });
          await tx.memoryRevision.create({
            data: {
              documentId: doc.id,
              revision: doc.revision,
              content,
              sourceRunId: request.sourceRunId,
              sourceThreadId: request.sourceThreadId,
            },
          });
          return { id: doc.id, path: doc.path, revision: doc.revision, content: doc.content };
        },
        { isolationLevel: Prisma.TransactionIsolationLevel.Serializable },
      ),
    );
  }

  async *exportMarkdown(
    request: MemoryExportRequest,
    context: AdapterContext,
  ): AsyncIterable<PortableFile> {
    const snapshot = await this.read(
      { scope: request.scope === "all" ? "user" : request.scope, botId: request.botId },
      context,
    );
    for (const doc of snapshot.documents) {
      yield { path: doc.path, content: new TextEncoder().encode(doc.content) };
    }
  }

  async importMarkdown(
    files: AsyncIterable<PortableFile>,
    context: AdapterContext,
  ): Promise<MemoryRevision> {
    let last: MemoryRevision | undefined;
    for await (const file of files) {
      last = await this.commit(
        {
          scope: "user",
          path: file.path,
          content: new TextDecoder().decode(file.content),
        },
        context,
      );
    }
    if (!last) throw new Error("No memory files to import");
    return last;
  }
}

/**
 * `addition` after the existing content, ending with a newline: right below a list, so notes stay
 * one list, and after a blank line below anything else, such as a heading.
 */
export function appendMarkdown(existing: string, addition: string): string {
  const base = existing.trimEnd();
  const next = addition.trim();
  if (!next) return existing;
  if (!base) return `${next}\n`;
  const existingLines = base.split("\n").map((line) => line.trim());
  if (!next.includes("\n") && existingLines.includes(next)) {
    return `${base}\n`;
  }
  const lastLine = base.slice(base.lastIndexOf("\n") + 1);
  const inList = /^\s*(?:[-*+]|\d+\.)\s|^\s{2,}\S/.test(lastLine);
  return `${base}${inList ? "\n" : "\n\n"}${next}\n`;
}

export function tokenizeMemoryQuery(query: string): string[] {
  const seen = new Set<string>();
  const tokens: string[] = [];
  for (const match of query.toLowerCase().matchAll(/[\p{L}\p{N}_-]+/gu)) {
    const token = match[0];
    if (token && !seen.has(token)) {
      seen.add(token);
      tokens.push(token);
    }
  }
  return tokens;
}

function countOccurrences(haystack: string, needle: string): number {
  if (!needle) return 0;
  let count = 0;
  let pos = 0;
  while (true) {
    const idx = haystack.indexOf(needle, pos);
    if (idx < 0) break;
    count += 1;
    pos = idx + needle.length;
  }
  return count;
}

function extractHeadingText(markdown: string): string {
  return markdown
    .split("\n")
    .filter((line) => /^\s*#{1,6}\s+/.test(line))
    .join("\n")
    .toLowerCase();
}

export function rankMemoryDocuments(
  documents: ReadonlyArray<{ path: string; content: string; updatedAt?: Date }>,
  query: string,
): MemorySearchResult[] {
  const normalizedQuery = query.trim().toLowerCase();
  if (!normalizedQuery) {
    return documents.map((doc) => ({
      path: doc.path,
      snippet: snippet(doc.content, ""),
      score: 1,
    }));
  }

  const tokens = tokenizeMemoryQuery(normalizedQuery);
  const scored: Array<MemorySearchResult & { rawScore: number; index: number }> = [];

  documents.forEach((doc, index) => {
    const lowerPath = doc.path.toLowerCase();
    const lowerContent = doc.content.toLowerCase();
    const headings = extractHeadingText(doc.content);

    const exactPhraseInPath = lowerPath.includes(normalizedQuery);
    const exactPhraseInContent = lowerContent.includes(normalizedQuery);

    let matchedTokens = 0;
    let tokenScore = 0;
    let bestAnchor = exactPhraseInContent ? normalizedQuery : "";

    for (const token of tokens) {
      const inPath = lowerPath.includes(token);
      const inHeading = headings.includes(token);
      const bodyOccurrences = countOccurrences(lowerContent, token);
      if (inPath || inHeading || bodyOccurrences > 0) {
        matchedTokens += 1;
        if (!bestAnchor && bodyOccurrences > 0) bestAnchor = token;
        const tf = bodyOccurrences / (bodyOccurrences + 1.2);
        tokenScore += (inPath ? 1.5 : 0) + (inHeading ? 1.0 : 0) + tf;
      }
    }

    if (!exactPhraseInPath && !exactPhraseInContent && matchedTokens === 0) {
      return;
    }

    const coverage = tokens.length > 0 ? matchedTokens / tokens.length : 1;
    const phraseBoost = (exactPhraseInPath ? 3.0 : 0) + (exactPhraseInContent ? 2.0 : 0);
    const rawScore = (phraseBoost + tokenScore) * (0.5 + 0.5 * coverage);

    scored.push({
      path: doc.path,
      snippet: snippet(doc.content, bestAnchor || normalizedQuery),
      score: 0,
      rawScore,
      index,
    });
  });

  if (scored.length === 0) return [];

  scored.sort((a, b) => b.rawScore - a.rawScore || a.index - b.index);
  const maxRaw = scored[0]?.rawScore || 1;

  return scored.map(({ path, snippet: textSnippet, rawScore }) => ({
    path,
    snippet: textSnippet,
    score: Math.max(0.01, Math.round((rawScore / maxRaw) * 10000) / 10000),
  }));
}

function snippet(content: string, q: string): string {
  if (!q) return content.slice(0, 140);
  const idx = content.toLowerCase().indexOf(q);
  if (idx < 0) return content.slice(0, 140);
  return content.slice(Math.max(0, idx - 40), idx + q.length + 80);
}