stack/apps/api/src/findings/findings.service.ts

import { BadRequestException, Injectable, Logger, NotFoundException } from "@nestjs/common";
import { Prisma } from "@prisma/client";
import { PrismaService } from "../prisma/prisma.service";
import { EmbeddingService } from "../knowledge/services/embedding.service";
import type { CreateFindingDto, QueryFindingsDto, SearchFindingsDto } from "./dto";

const findingSelect = {
  id: true,
  workspaceId: true,
  taskId: true,
  agentId: true,
  type: true,
  title: true,
  data: true,
  summary: true,
  createdAt: true,
  updatedAt: true,
} satisfies Prisma.FindingSelect;

type FindingRecord = Prisma.FindingGetPayload<{ select: typeof findingSelect }>;

interface RawFindingSearchResult {
  id: string;
  workspace_id: string;
  task_id: string | null;
  agent_id: string;
  type: string;
  title: string;
  data: Prisma.JsonValue;
  summary: string;
  created_at: Date;
  updated_at: Date;
  score: number;
}

export interface FindingSearchResult extends FindingRecord {
  score: number;
}

interface PaginatedMeta {
  total: number;
  page: number;
  limit: number;
  totalPages: number;
}

export interface PaginatedFindingsResponse {
  data: FindingRecord[];
  meta: PaginatedMeta;
}

export interface FindingsSearchResponse {
  data: FindingSearchResult[];
  meta: PaginatedMeta;
  query: string;
  similarityThreshold: number;
}

/**
 * Service for managing structured findings with vector search support
 */
@Injectable()
export class FindingsService {
  private readonly logger = new Logger(FindingsService.name);
  private readonly defaultSimilarityThreshold: number;

  constructor(
    private readonly prisma: PrismaService,
    private readonly embeddingService: EmbeddingService
  ) {
    const parsedThreshold = Number.parseFloat(process.env.FINDINGS_SIMILARITY_THRESHOLD ?? "0.5");

    this.defaultSimilarityThreshold =
      Number.isFinite(parsedThreshold) && parsedThreshold >= 0 && parsedThreshold <= 1
        ? parsedThreshold
        : 0.5;
  }

  /**
   * Create a finding and generate its embedding from the summary when available
   */
  async create(workspaceId: string, createFindingDto: CreateFindingDto): Promise<FindingRecord> {
    if (createFindingDto.taskId) {
      const task = await this.prisma.agentTask.findUnique({
        where: {
          id: createFindingDto.taskId,
          workspaceId,
        },
        select: { id: true },
      });

      if (!task) {
        throw new NotFoundException(`Agent task with ID ${createFindingDto.taskId} not found`);
      }
    }

    const createInput: Prisma.FindingUncheckedCreateInput = {
      workspaceId,
      agentId: createFindingDto.agentId,
      type: createFindingDto.type,
      title: createFindingDto.title,
      data: createFindingDto.data as Prisma.InputJsonValue,
      summary: createFindingDto.summary,
    };

    if (createFindingDto.taskId) {
      createInput.taskId = createFindingDto.taskId;
    }

    const finding = await this.prisma.finding.create({
      data: createInput,
      select: findingSelect,
    });

    await this.generateAndStoreEmbedding(finding.id, workspaceId, finding.summary);

    if (this.embeddingService.isConfigured()) {
      return this.findOne(finding.id, workspaceId);
    }

    return finding;
  }

  /**
   * Get paginated findings with optional filters
   */
  async findAll(workspaceId: string, query: QueryFindingsDto): Promise<PaginatedFindingsResponse> {
    const page = query.page ?? 1;
    const limit = query.limit ?? 50;
    const skip = (page - 1) * limit;

    const where: Prisma.FindingWhereInput = {
      workspaceId,
    };

    if (query.agentId) {
      where.agentId = query.agentId;
    }

    if (query.type) {
      where.type = query.type;
    }

    if (query.taskId) {
      where.taskId = query.taskId;
    }

    const [data, total] = await Promise.all([
      this.prisma.finding.findMany({
        where,
        select: findingSelect,
        orderBy: {
          createdAt: "desc",
        },
        skip,
        take: limit,
      }),
      this.prisma.finding.count({ where }),
    ]);

    return {
      data,
      meta: {
        total,
        page,
        limit,
        totalPages: Math.ceil(total / limit),
      },
    };
  }

  /**
   * Get a single finding by ID
   */
  async findOne(id: string, workspaceId: string): Promise<FindingRecord> {
    const finding = await this.prisma.finding.findUnique({
      where: {
        id,
        workspaceId,
      },
      select: findingSelect,
    });

    if (!finding) {
      throw new NotFoundException(`Finding with ID ${id} not found`);
    }

    return finding;
  }

  /**
   * Semantic search findings using vector similarity
   */
  async search(workspaceId: string, searchDto: SearchFindingsDto): Promise<FindingsSearchResponse> {
    if (!this.embeddingService.isConfigured()) {
      throw new BadRequestException(
        "Finding vector search requires OPENAI_API_KEY to be configured"
      );
    }

    const page = searchDto.page ?? 1;
    const limit = searchDto.limit ?? 20;
    const offset = (page - 1) * limit;
    const similarityThreshold = searchDto.similarityThreshold ?? this.defaultSimilarityThreshold;
    const distanceThreshold = 1 - similarityThreshold;

    const queryEmbedding = await this.embeddingService.generateEmbedding(searchDto.query);
    const embeddingString = `[${queryEmbedding.join(",")}]`;

    const agentFilter = searchDto.agentId
      ? Prisma.sql`AND f.agent_id = ${searchDto.agentId}`
      : Prisma.sql``;
    const typeFilter = searchDto.type ? Prisma.sql`AND f.type = ${searchDto.type}` : Prisma.sql``;
    const taskFilter = searchDto.taskId
      ? Prisma.sql`AND f.task_id = ${searchDto.taskId}::uuid`
      : Prisma.sql``;

    const searchResults = await this.prisma.$queryRaw<RawFindingSearchResult[]>`
      SELECT
        f.id,
        f.workspace_id,
        f.task_id,
        f.agent_id,
        f.type,
        f.title,
        f.data,
        f.summary,
        f.created_at,
        f.updated_at,
        (1 - (f.embedding <=> ${embeddingString}::vector)) AS score
      FROM findings f
      WHERE f.workspace_id = ${workspaceId}::uuid
        AND f.embedding IS NOT NULL
        ${agentFilter}
        ${typeFilter}
        ${taskFilter}
        AND (f.embedding <=> ${embeddingString}::vector) <= ${distanceThreshold}
      ORDER BY f.embedding <=> ${embeddingString}::vector
      LIMIT ${limit}
      OFFSET ${offset}
    `;

    const countResult = await this.prisma.$queryRaw<[{ count: bigint }]>`
      SELECT COUNT(*) as count
      FROM findings f
      WHERE f.workspace_id = ${workspaceId}::uuid
        AND f.embedding IS NOT NULL
        ${agentFilter}
        ${typeFilter}
        ${taskFilter}
        AND (f.embedding <=> ${embeddingString}::vector) <= ${distanceThreshold}
    `;

    const total = Number(countResult[0].count);

    const data: FindingSearchResult[] = searchResults.map((row) => ({
      id: row.id,
      workspaceId: row.workspace_id,
      taskId: row.task_id,
      agentId: row.agent_id,
      type: row.type,
      title: row.title,
      data: row.data,
      summary: row.summary,
      createdAt: row.created_at,
      updatedAt: row.updated_at,
      score: row.score,
    }));

    return {
      data,
      meta: {
        total,
        page,
        limit,
        totalPages: Math.ceil(total / limit),
      },
      query: searchDto.query,
      similarityThreshold,
    };
  }

  /**
   * Delete a finding
   */
  async remove(id: string, workspaceId: string): Promise<{ message: string }> {
    const existingFinding = await this.prisma.finding.findUnique({
      where: {
        id,
        workspaceId,
      },
      select: { id: true },
    });

    if (!existingFinding) {
      throw new NotFoundException(`Finding with ID ${id} not found`);
    }

    await this.prisma.finding.delete({
      where: {
        id,
        workspaceId,
      },
    });

    return { message: "Finding deleted successfully" };
  }

  /**
   * Generate and persist embedding for a finding summary
   */
  private async generateAndStoreEmbedding(
    findingId: string,
    workspaceId: string,
    summary: string
  ): Promise<void> {
    if (!this.embeddingService.isConfigured()) {
      return;
    }

    try {
      const embedding = await this.embeddingService.generateEmbedding(summary);
      const embeddingString = `[${embedding.join(",")}]`;

      await this.prisma.$executeRaw`
        UPDATE findings
        SET embedding = ${embeddingString}::vector,
            updated_at = NOW()
        WHERE id = ${findingId}::uuid
          AND workspace_id = ${workspaceId}::uuid
      `;
    } catch (error) {
      const message = error instanceof Error ? error.message : String(error);
      this.logger.warn(`Failed to generate embedding for finding ${findingId}: ${message}`);
    }
  }
}