gadget/docs/archive/services/vector-store.ts

// src/services/vector-store.ts
// Copyright (C) 2025 DTP Technologies, LLC
// All Rights Reserved

import env from "../config/env.js";
import assert from "node:assert";

import { QdrantClient } from "@qdrant/js-client-rest";

import { DtpService } from "../lib/service.js";
import AiService from "./ai.js";

const VECTOR_DIMENSION = 768; //TODO: This should match the dimension of the embeddings from the AI provider

class VectorStoreService extends DtpService {
  private qdrant?: QdrantClient;

  get name(): string {
    return "VectorStoreService";
  }

  get slug(): string {
    return "vector-store";
  }

  async start(): Promise<void> {
    this.log.info("Initializing Qdrant client");
    this.qdrant = new QdrantClient({
      url: env.qdrant.host,
    });

    this.log.info("service started");
  }

  async stop(): Promise<void> {
    this.log.info("service stopped");
  }

  async ensureCollection(name: string): Promise<void> {
    assert(this.qdrant, "Qdrant client not initialized");
    const collections = await this.qdrant.getCollections();
    const exists = collections.collections.some((col) => col.name === name);

    if (!exists) {
      await this.qdrant.createCollection(name, {
        vectors: {
          size: VECTOR_DIMENSION,
          distance: "Cosine",
        },
      });
      this.log.info("Qdrant collection initialized", { name });
    } else {
      this.log.info("Qdrant collection already exists", { name });
    }
  }

  private async getEmbedding(
    userId: string,
    content: string,
  ): Promise<number[]> {
    const client = await AiService.getVectorClient(userId);
    const model = await AiService.getVectorModel(userId);
    const response = await client.embeddings(content, model);
    return response.embedding;
  }

  async addDocument(
    userId: string,
    collectionName: string,
    id: string,
    content: string,
    metadata?: any,
  ) {
    assert(this.qdrant, "Qdrant client not initialized");
    try {
      await this.ensureCollection(collectionName);

      const embedding = await this.getEmbedding(userId, content);

      await this.qdrant.upsert(collectionName, {
        points: [
          {
            id,
            vector: embedding,
            payload: {
              content,
              metadata: metadata || {},
              created_at: new Date().toISOString(),
            },
          },
        ],
      });

      this.log.info(`Document added to Qdrant`, {
        collectionName,
        id,
        content,
        metadata,
      });
    } catch (error) {
      const errorMessage =
        error instanceof Error ? error.message : String(error);
      this.log.error("Error adding document to Qdrant", {
        error: errorMessage,
        collectionName,
        id,
      });
      throw error;
    }
  }

  async search(
    userId: string,
    collectionName: string,
    query: string,
    topK: number = 5,
    filter?: Record<string, unknown>,
  ) {
    assert(this.qdrant, "Qdrant client not initialized");
    try {
      const queryVector = await this.getEmbedding(userId, query);

      const searchOptions: any = {
        vector: queryVector,
        limit: topK,
        with_payload: true,
        with_vector: false,
      };

      if (filter) {
        searchOptions.filter = {
          must: Object.entries(filter).map(([key, value]) => ({
            key,
            match: { value },
          })),
        };
      }

      const results = await this.qdrant.search(collectionName, searchOptions);

      this.log.debug("Vector search completed", {
        collectionName,
        query,
        resultCount: results.length,
      });

      return results.map((result: any) => ({
        id: result.id,
        content: result.payload?.content,
        metadata: result.payload?.metadata,
        score: result.score,
      }));
    } catch (error) {
      const errorMessage =
        error instanceof Error ? error.message : String(error);
      this.log.error("Error searching in Qdrant", {
        error: errorMessage,
        collectionName,
        query,
      });
      throw error;
    }
  }

  async searchWithContext(
    userId: string,
    collectionName: string,
    query: string,
    topK: number = 5,
  ): Promise<string> {
    const results = await this.search(
      userId,
      collectionName,
      query,
      topK,
      undefined,
    );
    return results
      .map(
        (result) => `Document ID: ${result.id}\nContent: ${result.content}\n`,
      )
      .join("\n---\n");
  }

  async removeDocument(collectionName: string, id: string) {
    assert(this.qdrant, "Qdrant client not initialized");
    await this.qdrant.delete(collectionName, {
      filter: {
        must: [
          {
            key: "id",
            match: {
              value: id,
            },
          },
        ],
      },
    });
    this.log.info("Document removed from Qdrant", {
      collectionName,
      id,
    });
  }
}

export default new VectorStoreService();