@roadiehq/rag-ai-backend-retrieval-augmenter/dist/indexing/DefaultVectorAugmentationIndexer.cjs.js

'use strict';

var backendCommon = require('@backstage/backend-common');
var catalogClient = require('@backstage/catalog-client');
var text_splitter = require('langchain/text_splitter');
var catalogModel = require('@backstage/catalog-model');
var pLimit = require('p-limit');

function _interopDefaultCompat (e) { return e && typeof e === 'object' && 'default' in e ? e : { default: e }; }

var pLimit__default = /*#__PURE__*/_interopDefaultCompat(pLimit);

const TECHDOCS_ENTITY_FILTER = {
  "metadata.annotations.backstage.io/techdocs-ref": catalogClient.CATALOG_FILTER_EXISTS
};
class DefaultVectorAugmentationIndexer {
  _vectorStore;
  catalogApi;
  logger;
  auth;
  discovery;
  augmentationOptions;
  constructor({
    vectorStore,
    catalogApi,
    logger,
    auth,
    tokenManager,
    embeddings,
    discovery,
    augmentationOptions
  }) {
    vectorStore.connectEmbeddings(embeddings);
    this._vectorStore = vectorStore;
    this.augmentationOptions = augmentationOptions;
    this.catalogApi = catalogApi;
    this.logger = logger;
    this.auth = backendCommon.createLegacyAuthAdapters({
      auth,
      discovery,
      tokenManager
    }).auth;
    this.discovery = discovery;
  }
  get vectorStore() {
    return this._vectorStore;
  }
  /**
   * Returns the splitter object. Default implementation is using a naive RecursiveCharacterTextSplitter
   * which is likely not the best candidate for structured data splitting.
   *
   * It is recommended that this method is overwritten with more applicable implementation
   *
   * @returns {RecursiveCharacterTextSplitter} The splitter object.
   */
  getSplitter() {
    return new text_splitter.RecursiveCharacterTextSplitter({
      chunkSize: this.augmentationOptions?.chunkSize,
      chunkOverlap: this.augmentationOptions?.chunkOverlap
    });
  }
  async constructCatalogEmbeddingDocuments(entities, source) {
    const splitter = this.getSplitter();
    let docs = [];
    for (const entity of entities) {
      const splits = await splitter.splitText(JSON.stringify(entity));
      docs = docs.concat(
        splits.map((text, idx) => ({
          content: text,
          metadata: {
            splitId: idx.toString(),
            source,
            entityRef: catalogModel.stringifyEntityRef(entity),
            kind: entity.kind
          }
        }))
      );
    }
    return docs;
  }
  async constructTechDocsEmbeddingDocuments(documents, source) {
    const splitter = this.getSplitter();
    let docs = [];
    for (const { text, entity, title, location } of documents) {
      const splits = await splitter.splitText(text);
      docs = docs.concat(
        splits.map((content, idx) => ({
          content,
          metadata: {
            splitId: idx.toString(),
            source,
            entityRef: catalogModel.stringifyEntityRef(entity),
            kind: entity.kind,
            title,
            location
          }
        }))
      );
    }
    return docs;
  }
  async getDocuments(source, filter) {
    const limit = pLimit__default.default(this.augmentationOptions?.concurrencyLimit ?? 10);
    switch (source) {
      case "catalog": {
        const { token } = await this.auth.getPluginRequestToken({
          onBehalfOf: await this.auth.getOwnServiceCredentials(),
          targetPluginId: "catalog"
        });
        const entitiesResponse = await this.catalogApi.getEntities(
          { filter },
          { token }
        );
        const constructCatalogEmbeddingDocuments = await this.constructCatalogEmbeddingDocuments(
          entitiesResponse.items,
          source
        );
        this.logger.info(
          `Constructed ${constructCatalogEmbeddingDocuments.length} embedding documents for ${entitiesResponse.items.length} catalog items.`
        );
        return constructCatalogEmbeddingDocuments;
      }
      case "tech-docs": {
        const { token } = await this.auth.getPluginRequestToken({
          onBehalfOf: await this.auth.getOwnServiceCredentials(),
          targetPluginId: "catalog"
        });
        const entitiesResponse = await this.catalogApi.getEntities(
          {
            filter: { ...TECHDOCS_ENTITY_FILTER, ...filter }
          },
          { token }
        );
        const techDocsBaseUrl = await this.discovery.getBaseUrl("techdocs");
        const documentsPromises = entitiesResponse.items.map(
          (entity) => limit(async () => {
            const { kind } = entity;
            const { namespace = "default", name } = entity.metadata;
            const searchIndexUrl = `${techDocsBaseUrl}/static/docs/${namespace}/${kind}/${name}/search/search_index.json`;
            try {
              const { token: techDocsToken } = await this.auth.getPluginRequestToken({
                onBehalfOf: await this.auth.getOwnServiceCredentials(),
                targetPluginId: "techdocs"
              });
              const searchIndexResponse = await fetch(searchIndexUrl, {
                headers: {
                  Authorization: `Bearer ${techDocsToken}`
                }
              });
              const searchIndex = await searchIndexResponse.json();
              return searchIndex.docs.reduce((acc, doc) => {
                if (doc.location.includes("#") && doc.text)
                  acc.push({
                    ...doc,
                    entity
                  });
                return acc;
              }, []);
            } catch (e) {
              this.logger.debug(
                `Failed to retrieve tech docs search index for entity ${namespace}/${kind}/${name}`,
                e
              );
              return [];
            }
          })
        );
        const documents = (await Promise.all(documentsPromises)).flat();
        const constructTechDocsEmbeddingDocuments = await this.constructTechDocsEmbeddingDocuments(documents, source);
        this.logger.info(
          `Constructed ${constructTechDocsEmbeddingDocuments.length} embedding documents for ${entitiesResponse.items.length} TechDocs.`
        );
        return constructTechDocsEmbeddingDocuments;
      }
      default:
        throw new Error(
          `Attempting to create embeddings for a source not implemented yet: ${source} `
        );
    }
  }
  async createEmbeddings(source, filter) {
    await this.deleteEmbeddings(source, filter);
    const documents = await this.getDocuments(source, filter);
    await this._vectorStore.addDocuments(documents);
    return documents.length;
  }
  async deleteEmbeddings(source, filter) {
    const { token } = await this.auth.getPluginRequestToken({
      onBehalfOf: await this.auth.getOwnServiceCredentials(),
      targetPluginId: "catalog"
    });
    const entityFilter = source === "tech-docs" ? { ...TECHDOCS_ENTITY_FILTER, ...filter } : filter;
    const entities = (await this.catalogApi.getEntities({ filter: entityFilter }, { token })).items.map(catalogModel.stringifyEntityRef);
    for (const entityRef of entities) {
      await this._vectorStore.deleteDocuments({
        filter: { source, entityRef }
      });
    }
  }
}

exports.DefaultVectorAugmentationIndexer = DefaultVectorAugmentationIndexer;
//# sourceMappingURL=DefaultVectorAugmentationIndexer.cjs.js.map