vector-db.service.ts

/**
 * Vector Database Service
 * Handles vector storage and similarity search using Qdrant
 */

import axios, { AxiosInstance } from 'axios';
import { ChunkEmbedding, VectorSearchResult } from '../types';
import { logger } from '../utils/logger';
import config from '../config';

interface QdrantPoint {
  id: string;
  vector: number[];
  payload: Record<string, unknown>;
}

interface SearchRequest {
  vector: number[];
  limit: number;
  filter?: Record<string, unknown>;
  with_payload?: boolean;
  with_vector?: boolean;
}

interface SearchResponse {
  result: Array<{
    id: string;
    score: number;
    payload?: Record<string, unknown>;
    vector?: number[];
  }>;
}

export class VectorDbService {
  private client: AxiosInstance;
  private collectionName: string;
  private dimensions: number;

  constructor() {
    const { host, port } = config.vectorDb;
    this.collectionName = config.vectorDb.collection;
    this.dimensions = config.embedding.dimensions;

    this.client = axios.create({
      baseURL: `http://${host}:${port}`,
      timeout: 30000,
      headers: {
        'Content-Type': 'application/json',
      },
    });
  }

  /**
   * Initialize collection if it doesn't exist
   */
  async ensureCollection(): Promise<void> {
    try {
      // Check if collection exists
      const response = await this.client.get(`/collections/${this.collectionName}`);
      
      if (response.status === 200) {
        logger.info('Collection already exists', {
          collection: this.collectionName,
        });
        return;
      }
    } catch (error: any) {
      if (error.response?.status !== 404) {
        throw error;
      }
      // Collection doesn't exist, create it
    }

    logger.info('Creating vector collection', {
      collection: this.collectionName,
      dimensions: this.dimensions,
    });

    await this.client.put(`/collections/${this.collectionName}`, {
      vectors: {
        size: this.dimensions,
        distance: 'Cosine',
      },
    });

    logger.info('Collection created successfully', {
      collection: this.collectionName,
    });
  }

  /**
   * Store chunk embeddings in vector database
   */
  async storeEmbeddings(
    embeddings: ChunkEmbedding[],
    documentId?: string,
    dataset?: string
  ): Promise<void> {
    if (embeddings.length === 0) {
      logger.warn('No embeddings to store');
      return;
    }

    logger.info('Storing embeddings in vector database', {
      count: embeddings.length,
      collection: this.collectionName,
    });

    // Convert to Qdrant points
    const points: QdrantPoint[] = embeddings.map((emb) => ({
      id: emb.chunkId,
      vector: emb.vector,
      payload: {
        chunkId: emb.chunkId,
        documentId: documentId || emb.chunkId.split('-')[0],
        dataset: dataset || 'default', // Include dataset for filtering
        model: emb.model,
        dimensions: emb.dimensions,
        generatedAt: emb.generatedAt.toISOString(),
      },
    }));

    // Upsert points in batches
    const batchSize = 100;
    for (let i = 0; i < points.length; i += batchSize) {
      const batch = points.slice(i, i + batchSize);
      
      await this.client.put(`/collections/${this.collectionName}/points`, {
        points: batch,
      });

      logger.debug('Batch stored', {
        batchStart: i,
        batchSize: batch.length,
      });
    }

    logger.info('Embeddings stored successfully', {
      count: embeddings.length,
    });
  }

  /**
   * Search for similar vectors
   */
  async searchSimilar(
    queryVector: number[],
    limit: number = 10,
    filter?: Record<string, unknown>
  ): Promise<VectorSearchResult[]> {
    logger.info('Searching for similar vectors', {
      limit,
      collection: this.collectionName,
    });

    const request: SearchRequest = {
      vector: queryVector,
      limit,
      with_payload: true,
      with_vector: false,
    };

    if (filter) {
      request.filter = filter;
    }

    const response = await this.client.post<SearchResponse>(
      `/collections/${this.collectionName}/points/search`,
      request
    );

    const results: VectorSearchResult[] = response.data.result.map((hit) => ({
      chunkId: hit.id,
      score: hit.score,
      metadata: hit.payload,
      vector: hit.vector,
    }));

    logger.info('Search completed', {
      resultCount: results.length,
    });

    return results;
  }

  /**
   * Delete vectors for specific chunk IDs
   */
  async deleteVectors(chunkIds: string[]): Promise<void> {
    if (chunkIds.length === 0) {
      return;
    }

    logger.info('Deleting vectors', {
      count: chunkIds.length,
      collection: this.collectionName,
    });

    await this.client.post(`/collections/${this.collectionName}/points/delete`, {
      points: chunkIds,
    });

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

  /**
   * Delete all vectors for a document
   */
  async deleteDocumentVectors(documentId: string): Promise<void> {
    logger.info('Deleting document vectors', {
      documentId,
      collection: this.collectionName,
    });

    // Delete using filter (requires chunk metadata to include documentId)
    await this.client.post(`/collections/${this.collectionName}/points/delete`, {
      filter: {
        must: [
          {
            key: 'documentId',
            match: {
              value: documentId,
            },
          },
        ],
      },
    });

    logger.info('Document vectors deleted', { documentId });
  }

  /**
   * Get collection statistics
   */
  async getCollectionInfo(): Promise<{
    vectorsCount: number;
    status: string;
  }> {
    const response = await this.client.get(`/collections/${this.collectionName}`);

    return {
      vectorsCount: response.data.result.vectors_count || 0,
      status: response.data.result.status || 'unknown',
    };
  }

  /**
   * Health check
   */
  async healthCheck(): Promise<boolean> {
    try {
      const response = await this.client.get('/');
      return response.status === 200;
    } catch (error) {
      logger.error('Vector database health check failed', { error });
      return false;
    }
  }
}