embedding-storage.service.ts

/**
 * Embedding Storage Service
 * Handles persistence of embedding metadata to PostgreSQL
 */

import { ChunkEmbedding } from '../types';
import { logger } from '../utils/logger';
import { db } from '../database/connection';

export class EmbeddingStorageService {
  /**
   * Save embedding metadata to database
   */
  async saveEmbeddings(embeddings: ChunkEmbedding[]): Promise<void> {
    if (embeddings.length === 0) {
      logger.warn('No embeddings to save');
      return;
    }

    logger.info('Saving embedding metadata to database', {
      count: embeddings.length,
    });

    await db.transaction(async (client) => {
      for (const embedding of embeddings) {
        await client.query(
          `INSERT INTO chunk_embeddings (chunk_id, model, dimensions, generated_at)
           VALUES ($1, $2, $3, $4)
           ON CONFLICT (chunk_id, model)
           DO UPDATE SET
             dimensions = EXCLUDED.dimensions,
             generated_at = EXCLUDED.generated_at`,
          [
            embedding.chunkId,
            embedding.model,
            embedding.dimensions,
            embedding.generatedAt,
          ]
        );
      }
    });

    logger.info('Embedding metadata saved successfully', {
      count: embeddings.length,
    });
  }

  /**
   * Get embedding metadata by chunk ID
   */
  async getEmbeddingByChunk(chunkId: string): Promise<{
    model: string;
    dimensions: number;
    generatedAt: Date;
  } | null> {
    const result = await db.query(
      `SELECT model, dimensions, generated_at
       FROM chunk_embeddings
       WHERE chunk_id = $1
       LIMIT 1`,
      [chunkId]
    );

    if (result.rows.length === 0) {
      return null;
    }

    const row = result.rows[0];
    return {
      model: row.model,
      dimensions: row.dimensions,
      generatedAt: row.generated_at,
    };
  }

  /**
   * Check if chunk has embedding
   */
  async hasEmbedding(chunkId: string): Promise<boolean> {
    const result = await db.query(
      'SELECT EXISTS(SELECT 1 FROM chunk_embeddings WHERE chunk_id = $1) as exists',
      [chunkId]
    );

    return result.rows[0].exists;
  }

  /**
   * Delete embeddings for specific chunks
   */
  async deleteEmbeddings(chunkIds: string[]): Promise<void> {
    if (chunkIds.length === 0) {
      return;
    }

    logger.info('Deleting embedding metadata', { count: chunkIds.length });

    await db.query(
      'DELETE FROM chunk_embeddings WHERE chunk_id = ANY($1)',
      [chunkIds]
    );

    logger.info('Embedding metadata deleted', { count: chunkIds.length });
  }

  /**
   * Get embedding count for document
   */
  async getEmbeddingCountForDocument(documentId: string): Promise<number> {
    const result = await db.query(
      `SELECT COUNT(*) as count
       FROM chunk_embeddings ce
       JOIN text_chunks tc ON ce.chunk_id = tc.id
       WHERE tc.document_id = $1`,
      [documentId]
    );

    return parseInt(result.rows[0].count, 10);
  }
}