stack/apps/api/src/agent-memory/agent-memory.service.ts

import { Injectable, NotFoundException } from "@nestjs/common";
import { PrismaService } from "../prisma/prisma.service";
import { Prisma } from "@prisma/client";
import type { UpsertAgentMemoryDto } from "./dto";

@Injectable()
export class AgentMemoryService {
  constructor(private readonly prisma: PrismaService) {}

  /**
   * Upsert a memory entry for an agent.
   */
  async upsert(workspaceId: string, agentId: string, key: string, dto: UpsertAgentMemoryDto) {
    return this.prisma.agentMemory.upsert({
      where: {
        workspaceId_agentId_key: { workspaceId, agentId, key },
      },
      create: {
        workspaceId,
        agentId,
        key,
        value: dto.value as Prisma.InputJsonValue,
      },
      update: {
        value: dto.value as Prisma.InputJsonValue,
      },
    });
  }

  /**
   * List all memory entries for an agent in a workspace.
   */
  async findAll(workspaceId: string, agentId: string) {
    return this.prisma.agentMemory.findMany({
      where: { workspaceId, agentId },
      orderBy: { key: "asc" },
    });
  }

  /**
   * Get a single memory entry by key.
   */
  async findOne(workspaceId: string, agentId: string, key: string) {
    const entry = await this.prisma.agentMemory.findUnique({
      where: {
        workspaceId_agentId_key: { workspaceId, agentId, key },
      },
    });

    if (!entry) {
      throw new NotFoundException(`Memory key "${key}" not found for agent "${agentId}"`);
    }

    return entry;
  }

  /**
   * Delete a memory entry by key.
   */
  async remove(workspaceId: string, agentId: string, key: string) {
    const entry = await this.prisma.agentMemory.findUnique({
      where: {
        workspaceId_agentId_key: { workspaceId, agentId, key },
      },
    });

    if (!entry) {
      throw new NotFoundException(`Memory key "${key}" not found for agent "${agentId}"`);
    }

    await this.prisma.agentMemory.delete({
      where: {
        workspaceId_agentId_key: { workspaceId, agentId, key },
      },
    });

    return { message: "Memory entry deleted successfully" };
  }
}