UNPKG

@roadiehq/rag-ai-backend-retrieval-augmenter

Version:

142 lines (131 loc) 5.95 kB
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 };