ai/gateway/src/clients/qdrant-client.ts

import { QdrantClient as QdrantRestClient } from '@qdrant/js-client-rest';
import type { FastifyBaseLogger } from 'fastify';

/**
 * Qdrant client configuration
 */
export interface QdrantConfig {
  url: string;
  apiKey?: string;
  collectionName?: string;
}

/**
 * Qdrant client wrapper for RAG vector storage
 *
 * Features:
 * - Global namespace (user_id = "0") for platform knowledge
 * - User-specific namespaces for personal memories
 * - Payload-indexed by user_id for GDPR compliance
 * - Cosine similarity search
 */
export class QdrantClient {
  private client: QdrantRestClient;
  private collectionName: string;
  private vectorDimension: number;
  private logger: FastifyBaseLogger;

  constructor(config: QdrantConfig, logger: FastifyBaseLogger, vectorDimension: number = 1536) {
    this.logger = logger;
    this.collectionName = config.collectionName || 'gateway_memory';
    this.vectorDimension = vectorDimension;

    // Initialize Qdrant REST client
    this.client = new QdrantRestClient({
      url: config.url,
      apiKey: config.apiKey,
    });

    this.logger.info({
      url: config.url,
      collection: this.collectionName,
      vectorDimension,
    }, 'Qdrant client initialized');
  }

  /**
   * Initialize collection with proper schema and indexes
   */
  async initialize(): Promise<void> {
    this.logger.info({ collection: this.collectionName }, 'Initializing Qdrant collection');

    try {
      // Check if collection exists
      const collections = await this.client.getCollections();
      const exists = collections.collections.some(c => c.name === this.collectionName);

      if (!exists) {
        this.logger.info({ collection: this.collectionName }, 'Creating new collection');

        // Create collection with vector configuration
        await this.client.createCollection(this.collectionName, {
          vectors: {
            size: this.vectorDimension,
            distance: 'Cosine',
          },
        });

        // Create payload indexes for efficient filtering
        await this.client.createPayloadIndex(this.collectionName, {
          field_name: 'user_id',
          field_schema: 'keyword',
        });

        await this.client.createPayloadIndex(this.collectionName, {
          field_name: 'session_id',
          field_schema: 'keyword',
        });

        await this.client.createPayloadIndex(this.collectionName, {
          field_name: 'timestamp',
          field_schema: 'integer',
        });

        this.logger.info({ collection: this.collectionName }, 'Collection created successfully');
      } else {
        this.logger.info({ collection: this.collectionName }, 'Collection already exists');
      }
    } catch (error) {
      this.logger.error({ error, collection: this.collectionName }, 'Failed to initialize collection');
      throw error;
    }
  }

  /**
   * Store a vector point with payload
   */
  async upsertPoint(
    id: string,
    vector: number[],
    payload: Record<string, any>
  ): Promise<void> {
    try {
      await this.client.upsert(this.collectionName, {
        wait: true,
        points: [{
          id,
          vector,
          payload,
        }],
      });
    } catch (error) {
      this.logger.error({ error, id }, 'Failed to upsert point');
      throw error;
    }
  }

  /**
   * Search for similar vectors
   * Queries both global (user_id="0") and user-specific vectors
   */
  async search(
    userId: string,
    queryVector: number[],
    options?: {
      limit?: number;
      scoreThreshold?: number;
      sessionId?: string;
      timeRange?: { start: number; end: number };
    }
  ): Promise<Array<{
    id: string;
    score: number;
    payload: Record<string, any>;
  }>> {
    const limit = options?.limit || 5;
    const scoreThreshold = options?.scoreThreshold || 0.7;

    try {
      // Build filter: (user_id = userId OR user_id = "0") AND other conditions
      const mustConditions: any[] = [];
      const shouldConditions: any[] = [
        { key: 'user_id', match: { value: userId } },
        { key: 'user_id', match: { value: '0' } }, // Global namespace
      ];

      // Add session filter if provided
      if (options?.sessionId) {
        mustConditions.push({
          key: 'session_id',
          match: { value: options.sessionId },
        });
      }

      // Add time range filter if provided
      if (options?.timeRange) {
        mustConditions.push({
          key: 'timestamp',
          range: {
            gte: options.timeRange.start,
            lte: options.timeRange.end,
          },
        });
      }

      // Perform search
      const results = await this.client.search(this.collectionName, {
        vector: queryVector,
        filter: {
          must: mustConditions.length > 0 ? mustConditions : undefined,
          should: shouldConditions,
        },
        limit,
        score_threshold: scoreThreshold,
        with_payload: true,
      });

      return results.map(r => ({
        id: r.id as string,
        score: r.score,
        payload: r.payload || {},
      }));
    } catch (error) {
      this.logger.error({ error, userId }, 'Search failed');
      throw error;
    }
  }

  /**
   * Get points by filter (without vector search)
   */
  async scroll(
    userId: string,
    options?: {
      limit?: number;
      sessionId?: string;
      offset?: string;
    }
  ): Promise<{
    points: Array<{ id: string; payload: Record<string, any> }>;
    nextOffset?: string;
  }> {
    try {
      const filter: any = {
        must: [
          { key: 'user_id', match: { value: userId } },
        ],
      };

      if (options?.sessionId) {
        filter.must.push({
          key: 'session_id',
          match: { value: options.sessionId },
        });
      }

      const result = await this.client.scroll(this.collectionName, {
        filter,
        limit: options?.limit || 10,
        offset: options?.offset,
        with_payload: true,
        with_vector: false,
      });

      return {
        points: result.points.map(p => ({
          id: p.id as string,
          payload: p.payload || {},
        })),
        nextOffset: result.next_page_offset as string | undefined,
      };
    } catch (error) {
      this.logger.error({ error, userId }, 'Scroll failed');
      throw error;
    }
  }

  /**
   * Delete all points for a user (GDPR compliance)
   */
  async deleteUserData(userId: string): Promise<void> {
    this.logger.info({ userId }, 'Deleting user vectors for GDPR compliance');

    try {
      await this.client.delete(this.collectionName, {
        wait: true,
        filter: {
          must: [
            { key: 'user_id', match: { value: userId } },
          ],
        },
      });

      this.logger.info({ userId }, 'User vectors deleted');
    } catch (error) {
      this.logger.error({ error, userId }, 'Failed to delete user data');
      throw error;
    }
  }

  /**
   * Delete points for a specific session
   */
  async deleteSession(userId: string, sessionId: string): Promise<void> {
    this.logger.info({ userId, sessionId }, 'Deleting session vectors');

    try {
      await this.client.delete(this.collectionName, {
        wait: true,
        filter: {
          must: [
            { key: 'user_id', match: { value: userId } },
            { key: 'session_id', match: { value: sessionId } },
          ],
        },
      });

      this.logger.info({ userId, sessionId }, 'Session vectors deleted');
    } catch (error) {
      this.logger.error({ error, userId, sessionId }, 'Failed to delete session');
      throw error;
    }
  }

  /**
   * Get collection info and statistics
   */
  async getCollectionInfo(): Promise<{
    vectorsCount: number;
    indexedVectorsCount: number;
    pointsCount: number;
  }> {
    try {
      const info = await this.client.getCollection(this.collectionName);

      return {
        vectorsCount: (info as any).vectors_count || 0,
        indexedVectorsCount: info.indexed_vectors_count || 0,
        pointsCount: info.points_count || 0,
      };
    } catch (error) {
      this.logger.error({ error }, 'Failed to get collection info');
      throw error;
    }
  }

  /**
   * Store global platform knowledge (user_id = "0")
   */
  async storeGlobalKnowledge(
    id: string,
    vector: number[],
    payload: Omit<Record<string, any>, 'user_id'>
  ): Promise<void> {
    return this.upsertPoint(id, vector, {
      ...payload,
      user_id: '0', // Global namespace
    });
  }
}