Merge pull request #1629 from kuangp/feat/rag-ai/new-auth

feat(rag-retrieval): support new auth system

.changeset/lucky-maps-confess.md

@@ -0,0 +1,5 @@
+---
+'@roadiehq/rag-ai-backend-retrieval-augmenter': major
+---
+
+Support new auth system and replace winston with `LoggerService`

plugins/backend/rag-ai-backend-retrieval-augmenter/package.json

@@ -38,13 +38,13 @@
     "@backstage/catalog-client": "^1.6.6",
     "@backstage/catalog-model": "^1.6.0",
     "@backstage/config": "^1.2.0",
+    "@backstage/errors": "^1.2.4",
     "@backstage/plugin-search-common": "^1.2.14",
     "@langchain/core": "^0.2.27",
     "@roadiehq/rag-ai-node": "^0.1.6",
     "langchain": "^0.1.21",
     "node-fetch": "^2.6.7",
     "p-limit": "^3.0.2",
-    "winston": "^3.11.0",
     "yn": "^4.0.0"
   },
   "devDependencies": {

plugins/backend/rag-ai-backend-retrieval-augmenter/src/defaultInitializer.ts

@@ -21,23 +21,25 @@ import {
   VectorEmbeddingsRetriever,
 } from './retrieval';
 import { RoadieVectorStore } from '@roadiehq/rag-ai-node';
-import { Logger } from 'winston';
 import {
   PluginEndpointDiscovery,
   TokenManager,
 } from '@backstage/backend-common';
+import { AuthService, LoggerService } from '@backstage/backend-plugin-api';
 
 export type DefaultRetrievalPipelineOptions = {
   vectorStore: RoadieVectorStore;
-  logger: Logger;
+  logger: LoggerService;
   discovery: PluginEndpointDiscovery;
-  tokenManager: TokenManager;
+  auth?: AuthService;
+  tokenManager?: TokenManager;
 };
 
 export const createDefaultRetrievalPipeline = ({
   vectorStore,
   discovery,
   logger,
+  auth,
   tokenManager,
 }: DefaultRetrievalPipelineOptions) => {
   const vectorEmbeddingsRetriever = new VectorEmbeddingsRetriever({
@@ -48,6 +50,7 @@ export const createDefaultRetrievalPipeline = ({
   const searchRetriever = new SearchRetriever({
     discovery,
     logger,
+    auth,
     tokenManager,
   });
 

plugins/backend/rag-ai-backend-retrieval-augmenter/src/indexing/DefaultVectorAugmentationIndexer.ts

@@ -14,9 +14,11 @@
  * limitations under the License.
  */
 
-import { TokenManager } from '@backstage/backend-common';
+import {
+  createLegacyAuthAdapters,
+  TokenManager,
+} from '@backstage/backend-common';
 import { CATALOG_FILTER_EXISTS, CatalogApi } from '@backstage/catalog-client';
-import { Logger } from 'winston';
 import { SearchIndex, AugmentationOptions, TechDocsDocument } from './types';
 import { Embeddings } from '@langchain/core/embeddings';
 import { RecursiveCharacterTextSplitter } from 'langchain/text_splitter';
@@ -28,14 +30,18 @@ import {
   EntityFilterShape,
   RoadieVectorStore,
 } from '@roadiehq/rag-ai-node';
-import { DiscoveryService } from '@backstage/backend-plugin-api';
+import {
+  AuthService,
+  DiscoveryService,
+  LoggerService,
+} from '@backstage/backend-plugin-api';
 import pLimit from 'p-limit';
 
 export class DefaultVectorAugmentationIndexer implements AugmentationIndexer {
   private readonly _vectorStore: RoadieVectorStore;
   private readonly catalogApi: CatalogApi;
-  private readonly logger: Logger;
-  private readonly tokenManager: TokenManager;
+  private readonly logger: LoggerService;
+  private readonly auth: AuthService;
   private readonly discovery: DiscoveryService;
 
   private readonly augmentationOptions?: AugmentationOptions;
@@ -44,15 +50,17 @@ export class DefaultVectorAugmentationIndexer implements AugmentationIndexer {
     vectorStore,
     catalogApi,
     logger,
+    auth,
     tokenManager,
     embeddings,
     discovery,
     augmentationOptions,
   }: {
     vectorStore: RoadieVectorStore;
     catalogApi: CatalogApi;
-    logger: Logger;
-    tokenManager: TokenManager;
+    logger: LoggerService;
+    auth?: AuthService;
+    tokenManager?: TokenManager;
     embeddings: Embeddings;
     discovery: DiscoveryService;
     augmentationOptions?: AugmentationOptions;
@@ -62,7 +70,11 @@ export class DefaultVectorAugmentationIndexer implements AugmentationIndexer {
     this.augmentationOptions = augmentationOptions;
     this.catalogApi = catalogApi;
     this.logger = logger;
-    this.tokenManager = tokenManager;
+    this.auth = createLegacyAuthAdapters({
+      auth,
+      discovery,
+      tokenManager,
+    }).auth;
     this.discovery = discovery;
   }
 
@@ -144,7 +156,10 @@ export class DefaultVectorAugmentationIndexer implements AugmentationIndexer {
 
     switch (source) {
       case 'catalog': {
-        const { token } = await this.tokenManager.getToken();
+        const { token } = await this.auth.getPluginRequestToken({
+          onBehalfOf: await this.auth.getOwnServiceCredentials(),
+          targetPluginId: 'catalog',
+        });
 
         const entitiesResponse = await this.catalogApi.getEntities(
           { filter },
@@ -162,7 +177,10 @@ export class DefaultVectorAugmentationIndexer implements AugmentationIndexer {
         return constructCatalogEmbeddingDocuments;
       }
       case 'tech-docs': {
-        const { token } = await this.tokenManager.getToken();
+        const { token } = await this.auth.getPluginRequestToken({
+          onBehalfOf: await this.auth.getOwnServiceCredentials(),
+          targetPluginId: 'techdocs',
+        });
 
         const entitiesResponse = await this.catalogApi.getEntities(
           {
@@ -206,7 +224,7 @@ export class DefaultVectorAugmentationIndexer implements AugmentationIndexer {
             } catch (e) {
               this.logger.debug(
                 `Failed to retrieve tech docs search index for entity ${namespace}/${kind}/${name}`,
-                e,
+                e as Error,
               );
               return [];
             }
@@ -244,7 +262,10 @@ export class DefaultVectorAugmentationIndexer implements AugmentationIndexer {
     source: EmbeddingsSource,
     filter: EntityFilterShape,
   ): Promise<void> {
-    const { token } = await this.tokenManager.getToken();
+    const { token } = await this.auth.getPluginRequestToken({
+      onBehalfOf: await this.auth.getOwnServiceCredentials(),
+      targetPluginId: 'catalog',
+    });
 
     const entities = (
       await this.catalogApi.getEntities({ filter }, { token })

plugins/backend/rag-ai-backend-retrieval-augmenter/src/indexing/types.ts

@@ -14,8 +14,8 @@
  * limitations under the License.
  */
 import { TokenManager } from '@backstage/backend-common';
-import { Logger } from 'winston';
 import { CatalogApi } from '@backstage/catalog-client';
+import { LoggerService } from '@backstage/backend-plugin-api';
 import { PluginEndpointDiscovery } from '@backstage/backend-common';
 import { RoadieVectorStore } from '@roadiehq/rag-ai-node';
 import { Entity } from '@backstage/catalog-model';
@@ -27,7 +27,7 @@ export type AugmentationOptions = {
 };
 
 export interface RoadieEmbeddingsConfig {
-  logger: Logger;
+  logger: LoggerService;
   tokenManager: TokenManager;
   vectorStore: RoadieVectorStore;
   catalogApi: CatalogApi;

plugins/backend/rag-ai-backend-retrieval-augmenter/src/retrieval/retrievers/SearchClient.ts

@@ -14,12 +14,14 @@
  * limitations under the License.
  */
 import {
+  createLegacyAuthAdapters,
   PluginEndpointDiscovery,
   TokenManager,
 } from '@backstage/backend-common';
+import { AuthService, LoggerService } from '@backstage/backend-plugin-api';
+import { ResponseError } from '@backstage/errors';
 import { EmbeddingDoc, EmbeddingsSource } from '@roadiehq/rag-ai-node';
 import { SearchResultSet } from '@backstage/plugin-search-common';
-import { Logger } from 'winston';
 
 export type SearchClientQuery = {
   term: string;
@@ -39,24 +41,31 @@ const embeddingsSourceToBackstageSearchType = (source: EmbeddingsSource) => {
 
 export class SearchClient {
   private readonly discoveryApi: PluginEndpointDiscovery;
-  private readonly logger: Logger;
-  private readonly tokenManager: TokenManager;
+  private readonly logger: LoggerService;
+  private readonly auth: AuthService;
 
   constructor(options: {
     discoveryApi: PluginEndpointDiscovery;
-    logger: Logger;
-    tokenManager: TokenManager;
+    logger: LoggerService;
+    auth?: AuthService;
+    tokenManager?: TokenManager;
   }) {
     this.discoveryApi = options.discoveryApi;
     this.logger = options.logger;
-    this.tokenManager = options.tokenManager;
+    this.auth = createLegacyAuthAdapters({
+      ...options,
+      discovery: options.discoveryApi,
+    }).auth;
   }
 
   async query(query: SearchClientQuery): Promise<EmbeddingDoc[]> {
     const url = `${await this.discoveryApi.getBaseUrl('search')}/query?term=${
       query.term
     }&types[0]=${embeddingsSourceToBackstageSearchType(query.source)}`;
-    const { token } = await this.tokenManager.getToken();
+    const { token } = await this.auth.getPluginRequestToken({
+      onBehalfOf: await this.auth.getOwnServiceCredentials(),
+      targetPluginId: 'search',
+    });
     const response = await fetch(url, {
       method: 'GET',
       headers: {
@@ -67,7 +76,7 @@ export class SearchClient {
     if (!response.ok) {
       this.logger.warn(
         'Unable to query Backstage search API for embeddable results.',
-        await response.text(),
+        await ResponseError.fromResponse(response),
       );
       return [];
     }

plugins/backend/rag-ai-backend-retrieval-augmenter/src/retrieval/retrievers/SearchRetriever.ts

@@ -19,33 +19,36 @@ import {
   EmbeddingDoc,
   EmbeddingsSource,
 } from '@roadiehq/rag-ai-node';
-import { Logger } from 'winston';
 import { SearchClient } from './SearchClient';
 import {
   PluginEndpointDiscovery,
   TokenManager,
 } from '@backstage/backend-common';
+import { AuthService, LoggerService } from '@backstage/backend-plugin-api';
 
 export class SearchRetriever implements AugmentationRetriever {
   private readonly searchClient: SearchClient;
-  private readonly logger: Logger;
+  private readonly logger: LoggerService;
 
   constructor({
     discovery,
     logger,
     searchClient,
+    auth,
     tokenManager,
   }: {
     discovery: PluginEndpointDiscovery;
-    logger: Logger;
+    logger: LoggerService;
     searchClient?: SearchClient;
-    tokenManager: TokenManager;
+    auth?: AuthService;
+    tokenManager?: TokenManager;
   }) {
     this.searchClient =
       searchClient ??
       new SearchClient({
         discoveryApi: discovery,
         logger: logger.child({ label: 'rag-ai-searchclient' }),
+        auth,
         tokenManager,
       });
     this.logger = logger;

plugins/backend/rag-ai-backend-retrieval-augmenter/src/retrieval/retrievers/VectorEmbeddingsRetriever.ts

@@ -21,18 +21,18 @@ import {
   EntityFilterShape,
   RoadieVectorStore,
 } from '@roadiehq/rag-ai-node';
-import { Logger } from 'winston';
+import { LoggerService } from '@backstage/backend-plugin-api';
 
 export class VectorEmbeddingsRetriever implements AugmentationRetriever {
-  private readonly logger: Logger;
+  private readonly logger: LoggerService;
   private readonly vectorStore: RoadieVectorStore;
 
   constructor({
     vectorStore,
     logger,
   }: {
     vectorStore: RoadieVectorStore;
-    logger: Logger;
+    logger: LoggerService;
   }) {
     this.vectorStore = vectorStore;
     this.logger = logger;

plugins/backend/rag-ai-backend-retrieval-augmenter/src/retrieval/routers/SourceBasedRetrievalRouter.ts

@@ -19,17 +19,17 @@ import {
   EmbeddingsSource,
   RetrievalRouter,
 } from '@roadiehq/rag-ai-node';
-import { Logger } from 'winston';
+import { LoggerService } from '@backstage/backend-plugin-api';
 
 export class SourceBasedRetrievalRouter implements RetrievalRouter {
-  private readonly logger: Logger;
+  private readonly logger: LoggerService;
   private readonly retrievers: Map<EmbeddingsSource, AugmentationRetriever[]>;
 
   constructor({
     logger,
     retrievers,
   }: {
-    logger: Logger;
+    logger: LoggerService;
     retrievers: Map<EmbeddingsSource, AugmentationRetriever[]>;
   }) {
     this.retrievers = retrievers;