Kokorone/apps/api/src/memory/memory-retriever.service.ts

import { Injectable } from "@nestjs/common";
import type { EpisodicMemory } from "@kokorone/shared";
import { PrismaService } from "../prisma/prisma.service.js";

export interface RetrieveOptions {
  limit?: number;
  // E-5 關係加權檢索接點：與此使用者相關的記憶會被加權排到前面。
  relatedUserId?: string;
  // K-3 隱私邊界：群聊等有旁人在場的情境傳 true，標記為私密的記憶不會被檢索到。
  excludePrivate?: boolean;
}

export interface MemoryRetriever {
  retrieve(characterId: string, query: string, options?: RetrieveOptions): Promise<EpisodicMemory[]>;
}

const WEIGHTS = {
  keyword: 0.4,
  recency: 0.25,
  emotion: 0.15,
  relationship: 0.2,
};

const DEFAULT_LIMIT = 5;

function keywordRelevance(content: string, keywords: string[]): number {
  if (keywords.length === 0) {
    return 0;
  }
  const hits = keywords.filter((keyword) => content.includes(keyword)).length;
  return hits / keywords.length;
}

// C-6 記憶檢索器（關鍵字版本）：排序權重＝關鍵字相關度＋時間近因＋情緒權重。
// R-2 會加上 pgvector 語意檢索實作，取代這裡的關鍵字比對。
@Injectable()
export class KeywordMemoryRetriever implements MemoryRetriever {
  constructor(private readonly prisma: PrismaService) {}

  async retrieve(characterId: string, query: string, options: RetrieveOptions = {}): Promise<EpisodicMemory[]> {
    const limit = options.limit ?? DEFAULT_LIMIT;
    const rows = await this.prisma.client.episodicMemory.findMany({
      where: { characterId, ...(options.excludePrivate ? { isPrivate: false } : {}) },
    });
    const now = Date.now();
    const keywords = query.split(/\s+/).filter(Boolean);

    const scored = rows.map((row) => {
      const keywordScore = keywordRelevance(row.content, keywords);
      const daysSince = (now - row.occurredAt.getTime()) / (1000 * 60 * 60 * 24);
      const recencyScore = 1 / (1 + Math.max(0, daysSince));
      const emotionScore = row.emotionIntensity;
      const relationshipScore = options.relatedUserId && row.relatedUserId === options.relatedUserId ? 1 : 0;
      const score =
        keywordScore * WEIGHTS.keyword +
        recencyScore * WEIGHTS.recency +
        emotionScore * WEIGHTS.emotion +
        relationshipScore * WEIGHTS.relationship;
      return { row, score };
    });

    scored.sort((a, b) => b.score - a.score);
    const top = scored.slice(0, limit).map((s) => s.row);

    if (top.length > 0) {
      await this.prisma.client.episodicMemory.updateMany({
        where: { id: { in: top.map((row) => row.id) } },
        data: { lastRetrievedAt: new Date(now), retrievalCount: { increment: 1 } },
      });
    }

    return top.map((row) => ({
      id: row.id,
      characterId: row.characterId,
      content: row.content,
      occurredAt: row.occurredAt.toISOString(),
      emotionTag: row.emotionTag,
      emotionIntensity: row.emotionIntensity,
      retrievalCount: row.retrievalCount + 1,
      lastRetrievedAt: new Date(now).toISOString(),
      source: row.source,
      weight: row.weight,
      isPrivate: row.isPrivate,
    }));
  }
}