gadget/gadget-code/src/services/vector-store.ts

// src/services/vector-store.ts
// Copyright (C) 2026 Robert Colbert <rob.colbert@openplatform.us>
// All Rights Reserved

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

import env from "../config/env.js";
import { DtpService } from "../lib/service.js";
import { ChatTurnStatus } from "@gadget/api";
import {
  createAiApi,
  IAiEnvironment,
  IAiProvider as IAiApiProvider,
  AiApi,
} from "@gadget/ai";
import ChatTurn from "../models/chat-turn.js";
import AiProvider from "../models/ai-provider.js";
import { ChatSessionService } from "./index.js";

export interface ISearchFilters {
  userId?: string;
  projectId?: string;
  sessionId?: string;
  turnId?: string;
}

export interface ISearchResult {
  id: string;
  content: string;
  score: number;
  userId: string;
  projectId: string;
  sessionId: string;
  turnId: string;
  role: string;
  createdAt: string;
}

const aiEnv: IAiEnvironment = {
  NODE_ENV: env.NODE_ENV || "develop",
  services: {
    google: {
      cse: {
        apiKey: env.google.cse.apiKey,
        engineId: env.google.cse.engineId,
      },
    },
  },
};

class VectorStoreService extends DtpService {
  private client!: QdrantClient;
  private aiApi!: AiApi;
  private splitter!: RecursiveCharacterTextSplitter;
  private _initialized = false;
  private _dimensionMismatch = false;

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

  async start(): Promise<void> {
    if (!env.qdrant.providerId) {
      this.log.warn(
        "qdrant.providerId is not configured — vector store service will not start. " +
          "Set providerId in gadget-code.yaml under qdrant to enable semantic search.",
      );
      return;
    }

    this.log.info("initializing Qdrant client", {
      host: env.qdrant.host,
      port: env.qdrant.port,
      collection: env.qdrant.collection,
      providerId: env.qdrant.providerId,
      embeddingModel: env.qdrant.embeddingModel,
      vectorSize: env.qdrant.vectorSize,
    });

    // Initialize Qdrant client
    this.client = new QdrantClient({
      url: `http://${env.qdrant.host}:${env.qdrant.port}`,
      apiKey: env.qdrant.apiKey || undefined,
    });

    // Load the configured AI provider for embeddings
    const providerDoc = await AiProvider.findById(env.qdrant.providerId);
    if (!providerDoc) {
      this.log.error(
        `Qdrant providerId "${env.qdrant.providerId}" not found in database — ` +
          "vector store service will not start.",
      );
      return;
    }

    const aiProvider: IAiApiProvider = {
      _id: providerDoc._id,
      name: providerDoc.name,
      sdk: providerDoc.apiType, // map apiType → sdk
      baseUrl: providerDoc.baseUrl,
      apiKey: providerDoc.apiKey,
    };
    this.aiApi = createAiApi(aiEnv, aiProvider, this.log);

    // Initialize text splitter
    this.splitter = new RecursiveCharacterTextSplitter({
      chunkSize: 1000,
      chunkOverlap: 200,
    });

    // Ensure the Qdrant collection exists and validate dimensions
    await this.ensureCollection();

    // Validate that the embedding model produces vectors of the configured size
    await this.validateEmbeddingDimensions();

    if (this._dimensionMismatch) {
      this.log.error(
        "VectorStoreService started with DIMENSION MISMATCH — searches and ingestion will fail. " +
          "See previous error logs for fix instructions.",
      );
    }

    this._initialized = true;
    this.log.info("started", {
      host: env.qdrant.host,
      port: env.qdrant.port,
      collection: env.qdrant.collection,
      providerId: env.qdrant.providerId,
      embeddingModel: env.qdrant.embeddingModel,
      vectorSize: env.qdrant.vectorSize,
      dimensionMismatch: this._dimensionMismatch,
    });
  }

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

  /**
   * Check if the service is initialized and ready.
   */
  get isReady(): boolean {
    return this._initialized;
  }

  /**
   * Create the Qdrant collection if it doesn't already exist.
   * Validates existing collection dimensions against the configured vectorSize.
   */
  private async ensureCollection(): Promise<void> {
    const collections = await this.client.getCollections();
    const exists = collections.collections.some(
      (c) => c.name === env.qdrant.collection,
    );

    if (!exists) {
      await this.client.createCollection(env.qdrant.collection, {
        vectors: {
          size: env.qdrant.vectorSize,
          distance: "Cosine" as const,
        },
      });
      this.log.info(`created Qdrant collection "${env.qdrant.collection}"`, {
        vectorSize: env.qdrant.vectorSize,
        distance: "Cosine",
      });
    } else {
      // Validate existing collection dimensions against config
      try {
        const collectionInfo = await this.client.getCollection(
          env.qdrant.collection,
        );
        const vectorsConfig = collectionInfo.config?.params?.vectors;
        // Handle both named and unnamed vector configurations
        // Unnamed: { size: N, distance: "Cosine" } — Named: { "default": { size: N, distance: "Cosine" } }
        const actualSize =
          (vectorsConfig as Record<string, any>)?.size ??
          (
            Object.values(
              (vectorsConfig as Record<string, any>) || {},
            )[0] as Record<string, any>
          )?.size;

        if (actualSize && actualSize !== env.qdrant.vectorSize) {
          this._dimensionMismatch = true;
          this.log.error(
            "QDRANT COLLECTION DIMENSION MISMATCH — searches and ingestion will fail!",
            {
              collectionName: env.qdrant.collection,
              configuredVectorSize: env.qdrant.vectorSize,
              actualCollectionVectorSize: actualSize,
              fix: `Either (1) delete the collection and restart to recreate it, or (2) update qdrant.vectorSize in your config to ${actualSize}`,
            },
          );
        } else {
          this.log.info(
            `Qdrant collection "${env.qdrant.collection}" already exists`,
            {
              vectorSize: actualSize || env.qdrant.vectorSize,
            },
          );
        }
      } catch (err) {
        this.log.warn("Could not validate Qdrant collection dimensions", {
          error: (err as Error).message,
        });
      }
    }
  }

  /**
   * Validate that the embedding model produces vectors matching the configured vectorSize.
   * Sends a test embedding and compares its length against env.qdrant.vectorSize.
   */
  private async validateEmbeddingDimensions(): Promise<void> {
    try {
      const testEmbedding = await this.getEmbedding("test");
      if (testEmbedding.length !== env.qdrant.vectorSize) {
        this._dimensionMismatch = true;
        this.log.error(
          "EMBEDDING MODEL DIMENSION MISMATCH — the configured vectorSize does not match the model output!",
          {
            configuredVectorSize: env.qdrant.vectorSize,
            actualModelDimensions: testEmbedding.length,
            embeddingModel: env.qdrant.embeddingModel,
            fix: `Update qdrant.vectorSize in your config to ${testEmbedding.length}`,
          },
        );
      } else {
        this.log.info("embedding dimension validation passed", {
          vectorSize: testEmbedding.length,
          model: env.qdrant.embeddingModel,
        });
      }
    } catch (err) {
      this.log.warn("Could not validate embedding dimensions at startup", {
        error: (err as Error).message,
      });
    }
  }

  /**
   * Generate an embedding vector for the given text.
   */
  private async getEmbedding(text: string): Promise<number[]> {
    const response = await this.aiApi.embeddings(
      env.qdrant.embeddingModel,
      text,
    );
    return response.embedding;
  }

  /**
   * Ingest a chat turn into the vector store.
   * Fetches the turn by ID, extracts content, chunks, embeds, and upserts to Qdrant.
   * Fire-and-forget safe — logs errors but does not throw.
   */
  async ingestTurn(turnId: string): Promise<void> {
    if (!this._initialized) {
      this.log.warn(
        "ingestTurn called but service is not initialized — skipping",
        { turnId },
      );
      return;
    }

    if (this._dimensionMismatch) {
      this.log.error(
        "ingestTurn skipped — vector dimension mismatch detected at startup. Fix config and restart.",
        {
          turnId,
          configuredVectorSize: env.qdrant.vectorSize,
        },
      );
      return;
    }

    try {
      // Fetch and populate the turn
      let turn = await ChatTurn.findById(turnId);
      if (!turn) {
        this.log.warn("ingestTurn: turn not found", { turnId });
        return;
      }

      // Only ingest finished turns
      if (turn.status !== ChatTurnStatus.Finished) {
        this.log.debug("ingestTurn: skipping non-finished turn", {
          turnId,
          status: turn.status,
        });
        return;
      }

      // Populate references for metadata
      turn = await ChatTurn.populate(turn, ChatSessionService.populateChatTurn);

      // Extract content
      const userPrompt = turn.prompts.user || "";

      // Extract agent response: concatenate all 'responding' mode blocks
      let agentResponse = "";
      for (const block of turn.blocks) {
        if (block.mode === "responding" && typeof block.content === "string") {
          agentResponse += block.content + "\n";
        }
      }

      // Combine user + agent text for chunking
      const combinedText = [
        userPrompt ? `User: ${userPrompt}` : "",
        agentResponse ? `Agent: ${agentResponse.trim()}` : "",
      ]
        .filter(Boolean)
        .join("\n\n");

      if (!combinedText.trim()) {
        this.log.debug("ingestTurn: no content to ingest", { turnId });
        return;
      }

      // Chunk the text
      const chunks = await this.splitter.splitText(combinedText);

      // Determine role for each chunk based on content
      const points = [];
      for (let i = 0; i < chunks.length; i++) {
        const chunk = chunks[i]!;
        let role: string;
        if (chunk.startsWith("User:") && !chunk.includes("Agent:")) {
          role = "user";
        } else if (chunk.startsWith("Agent:") && !chunk.includes("User:")) {
          role = "agent";
        } else {
          role = "both";
        }

        const embedding = await this.getEmbedding(chunk);

        // Validate embedding dimensions before upsert
        if (embedding.length !== env.qdrant.vectorSize) {
          this.log.error(
            "embedding dimension mismatch during ingest — skipping chunk",
            {
              expected: env.qdrant.vectorSize,
              actual: embedding.length,
              turnId,
              chunkIndex: i,
              fix: `Update qdrant.vectorSize in your config to ${embedding.length}`,
            },
          );
          continue;
        }

        points.push({
          id: `${turnId}:${i}`,
          vector: embedding,
          payload: {
            content: chunk,
            userId: String(turn.user),
            projectId: turn.project ? String(turn.project) : "",
            sessionId: String(turn.session),
            turnId: turn._id,
            role,
            createdAt: turn.createdAt.toISOString(),
          },
        });
      }

      // Upsert all points
      await this.client.upsert(env.qdrant.collection, {
        wait: false,
        points,
      });

      this.log.info("ingested turn to vector store", {
        turnId,
        chunkCount: points.length,
      });
    } catch (error) {
      this.log.error("ingestTurn failed", {
        turnId,
        error,
      });
      // Do not rethrow — this is designed to be fire-and-forget
    }
  }

  /**
   * Search the vector store for relevant chunks.
   */
  async search(
    query: string,
    filters?: ISearchFilters,
    topK: number = 10,
  ): Promise<ISearchResult[]> {
    if (!this._initialized) {
      throw new Error("VectorStoreService is not initialized");
    }

    if (this._dimensionMismatch) {
      throw new Error(
        `Vector dimension mismatch: the Qdrant collection dimensions do not match the ` +
          `configured vectorSize (${env.qdrant.vectorSize}). Either delete the collection ` +
          `and restart, or update qdrant.vectorSize in your config.`,
      );
    }

    const queryVector = await this.getEmbedding(query);

    // Validate embedding dimensions before searching
    if (queryVector.length !== env.qdrant.vectorSize) {
      throw new Error(
        `Embedding dimension mismatch: model produced ${queryVector.length} dimensions, ` +
          `but collection expects ${env.qdrant.vectorSize}. ` +
          `Update qdrant.vectorSize in your config to ${queryVector.length}.`,
      );
    }

    // Build Qdrant filter from provided filters
    const must: Array<{
      key: string;
      match: { value: string };
    }> = [];

    if (filters?.userId) {
      must.push({ key: "userId", match: { value: filters.userId } });
    }
    if (filters?.projectId) {
      must.push({ key: "projectId", match: { value: filters.projectId } });
    }
    if (filters?.sessionId) {
      must.push({ key: "sessionId", match: { value: filters.sessionId } });
    }
    if (filters?.turnId) {
      must.push({ key: "turnId", match: { value: filters.turnId } });
    }

    const searchResults = await this.client.search(env.qdrant.collection, {
      vector: queryVector,
      limit: topK,
      filter: must.length > 0 ? { must } : undefined,
      with_payload: true,
    });

    return searchResults.map((point) => ({
      id: point.id as string,
      content: (point.payload as any)?.content || "",
      score: point.score ?? 0,
      userId: (point.payload as any)?.userId || "",
      projectId: (point.payload as any)?.projectId || "",
      sessionId: (point.payload as any)?.sessionId || "",
      turnId: (point.payload as any)?.turnId || "",
      role: (point.payload as any)?.role || "",
      createdAt: (point.payload as any)?.createdAt || "",
    }));
  }

  /**
   * Remove all vector points for a given turn.
   * Used for cleanup when a turn is deleted.
   */
  async removeTurnPoints(turnId: string): Promise<void> {
    if (!this._initialized) {
      this.log.warn(
        "removeTurnPoints called but service is not initialized — skipping",
        { turnId },
      );
      return;
    }

    try {
      await this.client.delete(env.qdrant.collection, {
        filter: {
          must: [{ key: "turnId", match: { value: turnId } }],
        },
      });
      this.log.info("removed vector points for turn", { turnId });
    } catch (error) {
      this.log.error("removeTurnPoints failed", {
        turnId,
        error,
      });
    }
  }
}

export default new VectorStoreService();