@roadiehq/rag-ai-backend-retrieval-augmenter
Version:
142 lines (131 loc) • 5.95 kB
TypeScript
import { TokenManager, PluginEndpointDiscovery } from '@backstage/backend-common';
import { CatalogApi } from '@backstage/catalog-client';
import { LoggerService, AuthService, DiscoveryService } from '@backstage/backend-plugin-api';
import { RoadieVectorStore, AugmentationIndexer, EmbeddingsSource, EmbeddingDoc, EntityFilterShape, RetrievalPipeline, RetrievalRouter, AugmentationRetriever, AugmentationPostProcessor } from '@roadiehq/rag-ai-node';
import { Entity } from '@backstage/catalog-model';
import { Embeddings } from '@langchain/core/embeddings';
import { RecursiveCharacterTextSplitter } from 'langchain/text_splitter';
type AugmentationOptions = {
chunkSize?: number;
chunkOverlap?: number;
concurrencyLimit?: number;
};
interface RoadieEmbeddingsConfig {
logger: LoggerService;
tokenManager?: TokenManager;
auth?: AuthService;
vectorStore: RoadieVectorStore;
catalogApi: CatalogApi;
discovery: PluginEndpointDiscovery;
augmentationOptions?: AugmentationOptions;
}
type TechDocsDocument = {
text: string;
entity: Entity;
title: string;
location: string;
};
declare class DefaultVectorAugmentationIndexer implements AugmentationIndexer {
private readonly _vectorStore;
private readonly catalogApi;
private readonly logger;
private readonly auth;
private readonly discovery;
private readonly augmentationOptions?;
protected constructor({ vectorStore, catalogApi, logger, auth, tokenManager, embeddings, discovery, augmentationOptions, }: {
vectorStore: RoadieVectorStore;
catalogApi: CatalogApi;
logger: LoggerService;
auth?: AuthService;
tokenManager?: TokenManager;
embeddings: Embeddings;
discovery: DiscoveryService;
augmentationOptions?: AugmentationOptions;
});
get vectorStore(): RoadieVectorStore;
/**
* 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.
*/
protected getSplitter(): RecursiveCharacterTextSplitter;
protected constructCatalogEmbeddingDocuments(entities: Entity[], source: EmbeddingsSource): Promise<EmbeddingDoc[]>;
protected constructTechDocsEmbeddingDocuments(documents: TechDocsDocument[], source: EmbeddingsSource): Promise<EmbeddingDoc[]>;
protected getDocuments(source: EmbeddingsSource, filter?: EntityFilterShape): Promise<EmbeddingDoc[]>;
createEmbeddings(source: EmbeddingsSource, filter: EntityFilterShape): Promise<number>;
deleteEmbeddings(source: EmbeddingsSource, filter: EntityFilterShape): Promise<void>;
}
declare class DefaultRetrievalPipeline implements RetrievalPipeline {
private readonly routers;
private readonly retrievers;
private readonly postProcessors;
constructor({ routers, retrievers, postProcessors, }: {
routers?: RetrievalRouter[];
retrievers?: AugmentationRetriever[];
postProcessors?: AugmentationPostProcessor[];
});
retrieveAugmentationContext(query: string, source: EmbeddingsSource, filter?: EntityFilterShape): Promise<EmbeddingDoc[]>;
}
declare class CombiningPostProcessor implements AugmentationPostProcessor {
process(_query: string, _source: EmbeddingsSource, embeddingDocs: Map<string, EmbeddingDoc[]>): Promise<EmbeddingDoc[]>;
}
declare class VectorEmbeddingsRetriever implements AugmentationRetriever {
private readonly logger;
private readonly vectorStore;
constructor({ vectorStore, logger, }: {
vectorStore: RoadieVectorStore;
logger: LoggerService;
});
get id(): string;
retrieve(query: string, source: EmbeddingsSource, _filter?: EntityFilterShape): Promise<EmbeddingDoc[]>;
}
type SearchClientQuery = {
term: string;
source: EmbeddingsSource;
};
declare class SearchClient {
private readonly discoveryApi;
private readonly logger;
private readonly auth;
constructor(options: {
discoveryApi: PluginEndpointDiscovery;
logger: LoggerService;
auth?: AuthService;
tokenManager?: TokenManager;
});
query(query: SearchClientQuery): Promise<EmbeddingDoc[]>;
}
declare class SearchRetriever implements AugmentationRetriever {
private readonly searchClient;
private readonly logger;
constructor({ discovery, logger, searchClient, auth, tokenManager, }: {
discovery: PluginEndpointDiscovery;
logger: LoggerService;
searchClient?: SearchClient;
auth?: AuthService;
tokenManager?: TokenManager;
});
get id(): string;
retrieve(query: string, source: EmbeddingsSource): Promise<EmbeddingDoc[]>;
}
declare class SourceBasedRetrievalRouter implements RetrievalRouter {
private readonly logger;
private readonly retrievers;
constructor({ logger, retrievers, }: {
logger: LoggerService;
retrievers: Map<EmbeddingsSource, AugmentationRetriever[]>;
});
determineRetriever(_query: string, source: EmbeddingsSource): Promise<AugmentationRetriever[]>;
}
type DefaultRetrievalPipelineOptions = {
vectorStore: RoadieVectorStore;
logger: LoggerService;
discovery: PluginEndpointDiscovery;
auth?: AuthService;
tokenManager?: TokenManager;
};
declare const createDefaultRetrievalPipeline: ({ vectorStore, discovery, logger, auth, tokenManager, }: DefaultRetrievalPipelineOptions) => DefaultRetrievalPipeline;
export { type AugmentationOptions, CombiningPostProcessor, DefaultRetrievalPipeline, type DefaultRetrievalPipelineOptions, DefaultVectorAugmentationIndexer, type RoadieEmbeddingsConfig, SearchRetriever, SourceBasedRetrievalRouter, VectorEmbeddingsRetriever, createDefaultRetrievalPipeline };