Skip to content

Commit

Permalink
feat(rag-retrieval): support new auth system
Browse files Browse the repository at this point in the history
  • Loading branch information
kuangp committed Sep 20, 2024
1 parent 0f00d72 commit 8d209f4
Show file tree
Hide file tree
Showing 9 changed files with 77 additions and 36 deletions.
5 changes: 5 additions & 0 deletions .changeset/lucky-maps-confess.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
---
'@roadiehq/rag-ai-backend-retrieval-augmenter': major
---

Support new auth system and replace winston with `LoggerService`
Original file line number Diff line number Diff line change
Expand Up @@ -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": {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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({
Expand All @@ -48,6 +50,7 @@ export const createDefaultRetrievalPipeline = ({
const searchRetriever = new SearchRetriever({
discovery,
logger,
auth,
tokenManager,
});

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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';
Expand All @@ -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;
Expand All @@ -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;
Expand All @@ -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;
}

Expand Down Expand Up @@ -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 },
Expand All @@ -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(
{
Expand Down Expand Up @@ -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 [];
}
Expand Down Expand Up @@ -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 })
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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';
Expand All @@ -27,7 +27,7 @@ export type AugmentationOptions = {
};

export interface RoadieEmbeddingsConfig {
logger: Logger;
logger: LoggerService;
tokenManager: TokenManager;
vectorStore: RoadieVectorStore;
catalogApi: CatalogApi;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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: {
Expand All @@ -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 [];
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down

0 comments on commit 8d209f4

Please sign in to comment.