feat(indexing): contentType check instead of GLOB match

aws-samples · Nov 2, 2023 · b3032ab · b3032ab
1 parent e6262a1
commit b3032ab
Show file tree

Hide file tree

Showing 4 changed files with 65 additions and 17 deletions.
diff --git a/demo/corpus/logic/src/env.ts b/demo/corpus/logic/src/env.ts
@@ -34,6 +34,10 @@ export interface IProcessEnv {
  PROCESSING_INPUT_LOCAL_PATH?: string;
  /** Glob pattern(s) to index in source bucket (CSV) */
  INDEXING_GLOB?: string;
+ /**
+ * Comma-separated list of supported 'Content-Type' values to index from source bucket.
+ */
+ INDEXING_SUPPORTED_CONTENT_TYPES?: string;
  /** Indicates if delta check is skipped, which is the per file last indexed checking */
  INDEXING_SKIP_DELTA_CHECK?: string;
 
@@ -76,7 +80,8 @@ export namespace ENV {
  export const CHUNK_OVERLAP = parseInt(process.env.CHUNK_OVERLAP || '200');
  export const VECTOR_INDEX_LISTS = parseInt(process.env.VECTOR_INDEX_LISTS || '1000');
  export const PROCESSING_INPUT_LOCAL_PATH = process.env.PROCESSING_INPUT_LOCAL_PATH || '/opt/ml/processing/input_data';
- export const INDEXING_GLOB = process.env.INDEXING_GLOB || '**/*.txt';
+ export const INDEXING_GLOB = process.env.INDEXING_GLOB || '**/*.*';
+ export const INDEXING_SUPPORTED_CONTENT_TYPES = process.env.INDEXING_SUPPORTED_CONTENT_TYPES || 'text/plain';
 
  export const EMBEDDING_TABLENAME = normalizePostgresTableName(
  `${EMBEDDING_SENTENCE_TRANSFORMER_MODEL}_${VECTOR_SIZE}`,

diff --git a/demo/corpus/logic/src/indexing/datastore.ts b/demo/corpus/logic/src/indexing/datastore.ts
@@ -14,6 +14,7 @@ import {
 } from '@aws-sdk/lib-dynamodb';
 import * as async from 'async';
 import { chunkArray } from './utils';
+import { ENV } from '../env';
 import { measurable } from '../metrics';
 
 const logger = getLogger('indexing/datastore');
@@ -59,6 +60,8 @@ export class IndexingCache {
  protected lastIndexedMap: Map<string, Date> = new Map();
  protected modelLastExecuted?: Date;
 
+ private supportedContentTypes: string[];
+
  get entityCount(): number {
  return this.entities.size;
  }
@@ -72,6 +75,9 @@ export class IndexingCache {
  this.ddbClient = new DynamoDBClient(ddbClientConfig ?? {});
  this.ddbDocClient = DynamoDBDocumentClient.from(this.ddbClient);
  this.s3Client = new S3Client(s3ClientConfig ?? {});
+
+ // remove spaces and split by comma
+ this.supportedContentTypes = ENV.INDEXING_SUPPORTED_CONTENT_TYPES.replace(' ', '').split(',');
  }
 
  formatSourceLocation(key: string): string {
@@ -218,17 +224,22 @@ export class IndexingCache {
  }),
  );
 
- const entity: IndexEntity = {
- objectKey: _objectKey,
- localPath: path.join(this.baseLocalPath, _objectKey),
- sourceLocation,
- metadata: normalizeMetadata(response.Metadata),
- lastModified: response.LastModified ?? new Date(),
- };
-
- this.entities.set(sourceLocation, entity);
-
- return entity;
+ if (response.ContentType && this.supportedContentTypes.includes(response.ContentType)) {
+ const entity: IndexEntity = {
+ objectKey: _objectKey,
+ localPath: path.join(this.baseLocalPath, _objectKey),
+ sourceLocation,
+ metadata: normalizeMetadata(response.Metadata),
+ lastModified: response.LastModified ?? new Date(),
+ };
+
+ this.entities.set(sourceLocation, entity);
+
+ return;
+ } else {
+ logger.warn(`${sourceLocation}'s Content-Type (${response.ContentType} not supported. Skipping...)`);
+ return;
+ }
  } catch (error) {
  logger.warn(`Failed to resolve S3 object key: "${sourceLocation}"`, error as Error);
  throw error;

diff --git a/demo/corpus/logic/test/datastore.test.ts b/demo/corpus/logic/test/datastore.test.ts
@@ -28,13 +28,12 @@ describe('metadata', () => {
  // console.log(base64EncodedValue);
 
  const metadata = {
- 'key': 'value',
- 'json-base64':
- 'eyJjaMOsYSBraMOzYSI6Imdpw6EgdHLhu4siLCJnacOhIHRy4buLIjoiY2jDrGEga2jDs2EifQ==', // base64EncodedValue,
+ key: 'value',
+ 'json-base64': 'eyJjaMOsYSBraMOzYSI6Imdpw6EgdHLhu4siLCJnacOhIHRy4buLIjoiY2jDrGEga2jDs2EifQ==', // base64EncodedValue,
  };
 
  const expectedResult = {
- 'key': 'value',
+ key: 'value',
  'chìa khóa': 'giá trị',
  'giá trị': 'chìa khóa',
  };
@@ -45,7 +44,7 @@ describe('metadata', () => {
 
  test('normalizeMetadata - with wrong base64 field', () => {
  const metadata = {
- 'key': 'value',
+ key: 'value',
  'json-base64': 'wrongly-encoded-value',
  };
 

diff --git a/demo/corpus/logic/test/indexing.test.ts b/demo/corpus/logic/test/indexing.test.ts
@@ -73,6 +73,7 @@ describe('indexing', () => {
  a: 'a',
  b: 'b',
  },
+ ContentType: 'text/plain',
  } as HeadObjectCommandOutput);
 
  dynamoDBDocumentMock.on(GetCommand).resolves({});
@@ -102,6 +103,7 @@ describe('indexing', () => {
  a: 'a',
  b: 'b',
  },
+ ContentType: 'text/plain',
  } as HeadObjectCommandOutput);
 
  // model last indexed
@@ -148,6 +150,7 @@ describe('indexing', () => {
  a: 'a',
  b: 'b',
  },
+ ContentType: 'text/plain',
  } as HeadObjectCommandOutput);
  });
 
@@ -178,4 +181,34 @@ describe('indexing', () => {
  const count = await require('../src/indexing').main();
  expect(count).toBe(Object.keys(inputFiles).length / 2);
  });
+
+ test('content-type filtered', async () => {
+ const _s3Head = s3Mock.on(HeadObjectCommand);
+ Object.keys(inputFiles).forEach((_, i) => {
+ _s3Head.resolvesOnce({
+ $metadata: {},
+ LastModified: new Date(),
+ Metadata: {
+ a: 'a',
+ b: 'b',
+ },
+ ContentType: i % 2 == 0 ? 'text/plain' : 'application/octet-stream',
+ } as HeadObjectCommandOutput);
+ });
+
+ dynamoDBDocumentMock.on(GetCommand).resolves({});
+ dynamoDBDocumentMock.on(PutCommand).resolves({});
+
+ dynamoDBDocumentMock.on(BatchGetCommand).resolves({
+ $metadata: {} as any,
+ Responses: {
+ [INDEXING_CACHE_TABLE]: [],
+ },
+ } as BatchGetCommandOutput);
+
+ dynamoDBDocumentMock.on(BatchWriteCommand).resolves({});
+
+ const count = await require('../src/indexing').main();
+ expect(count).toBe(Object.keys(inputFiles).length / 2);
+ });
 });