prime-sdk 2.0.5 → 2.1.2

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "prime-sdk",
-  "version": "2.0.5",
+  "version": "2.1.2",
   "description": "swiss army knife",
   "main": "dist/index.js",
   "types": "dist/index.d.ts",
@@ -23,10 +23,12 @@
   },
   "dependencies": {
     "@aws-sdk/client-kms": "^3.606.0",
+    "@aws-sdk/client-bedrock-runtime": "^3.620.0",
     "@aws-sdk/client-s3": "^3.572.0",
     "@aws-sdk/client-ses": "^3.572.0",
     "@aws-sdk/client-sns": "^3.572.0",
     "@aws-sdk/client-sqs": "^3.569.0",
+    "@aws-sdk/client-textract": "^3.620.0",
     "@aws-sdk/s3-request-presigner": "^3.572.0",
     "@nestjs/common": "^11.1.5",
     "@nestjs/config": "^4.0.2",

package/src/app.module.ts

@@ -9,6 +9,7 @@ import { CacheModule } from './lib/cache/cache.module';
 import { TypeRedis } from './external/redis/types/ioredis.types';
 import { ICache } from './lib/cache/interfaces/cache.interface';
 import { Cache } from './lib/cache/cache';
+import { DocumentsModule } from './lib/documents/documents.module';
 
 //Example
 @Module({
@@ -53,6 +54,16 @@ import { Cache } from './lib/cache/cache';
       endpoint: 'http://localhost:4566/',
       region: 'us-east-1',
     }),
+    DocumentsModule.forRoot({
+      bedrock: {
+        region: 'us-east-1',
+        endpoint: 'http://localhost:4566/',
+        credentials: {
+          accessKeyId: '',
+          secretAccessKey: '',
+        },
+      },
+    }),
     CacheModule.forRoot({
       type: TypeRedis.SIMPLE,
       connection: {

package/src/external/aws/bedrock/bedrock.provider.ts

@@ -0,0 +1,39 @@
+import { BedrockRuntimeClient } from '@aws-sdk/client-bedrock-runtime';
+import { DynamicModule } from '@nestjs/common';
+
+export type BedrockSetup = {
+  region: string;
+  endpoint?: string;
+  credentials: {
+    accessKeyId: string;
+    secretAccessKey: string;
+  };
+};
+
+const BedrockProvider = (params: BedrockSetup) => {
+  return {
+    provide: 'CONNECTION_BEDROCK_RUNTIME',
+    useFactory: () => {
+      return new BedrockRuntimeClient({
+        credentials: {
+          accessKeyId: params.credentials.accessKeyId,
+          secretAccessKey: params.credentials.secretAccessKey,
+        },
+        ...(params.endpoint && { endpoint: params.endpoint }),
+        region: params.region,
+      });
+    },
+  };
+};
+
+export class BedrockModule {
+  static forRoot(params: BedrockSetup): DynamicModule {
+    const provider = BedrockProvider(params);
+
+    return {
+      module: BedrockModule,
+      providers: [provider],
+      exports: [provider],
+    };
+  }
+}

package/src/external/aws/bedrock/bedrock.service.ts

@@ -0,0 +1,29 @@
+import { Inject, Injectable } from '@nestjs/common';
+import {
+  BedrockRuntimeClient,
+  InvokeModelCommand,
+  InvokeModelCommandOutput,
+} from '@aws-sdk/client-bedrock-runtime';
+import { IBedrock } from '../../../lib/documents/interfaces/documents.interface';
+import { InvokeModelPayload } from '../../../lib/documents/types/bedrock.type';
+
+@Injectable()
+export class Bedrock implements IBedrock {
+  constructor(
+    @Inject('CONNECTION_BEDROCK_RUNTIME')
+    private readonly client: BedrockRuntimeClient,
+  ) {}
+
+  async invokeModel(
+    payload: InvokeModelPayload,
+  ): Promise<InvokeModelCommandOutput> {
+    const command = new InvokeModelCommand({
+      body: payload.body,
+      modelId: payload.modelId,
+      contentType: payload.contentType || 'application/json',
+      accept: payload.accept || 'application/json',
+    });
+
+    return await this.client.send(command);
+  }
+}

package/src/external/aws/textract/textract.provider.ts

@@ -0,0 +1,39 @@
+import { TextractClient } from '@aws-sdk/client-textract';
+import { DynamicModule } from '@nestjs/common';
+
+export type TextractSetup = {
+  region: string;
+  endpoint?: string;
+  credentials: {
+    accessKeyId: string;
+    secretAccessKey: string;
+  };
+};
+
+const TextractProvider = (params: TextractSetup) => {
+  return {
+    provide: 'CONNECTION_TEXTRACT',
+    useFactory: () => {
+      return new TextractClient({
+        credentials: {
+          accessKeyId: params.credentials.accessKeyId,
+          secretAccessKey: params.credentials.secretAccessKey,
+        },
+        ...(params.endpoint && { endpoint: params.endpoint }),
+        region: params.region,
+      });
+    },
+  };
+};
+
+export class TextractModule {
+  static forRoot(params: TextractSetup): DynamicModule {
+    const provider = TextractProvider(params);
+
+    return {
+      module: TextractModule,
+      providers: [provider],
+      exports: [provider],
+    };
+  }
+}

package/src/external/aws/textract/textract.service.ts

@@ -0,0 +1,167 @@
+import { Inject, Injectable } from '@nestjs/common';
+import {
+  AnalyzeDocumentCommand,
+  AnalyzeDocumentCommandOutput,
+  DetectDocumentTextCommand,
+  DetectDocumentTextCommandOutput,
+  StartDocumentAnalysisCommand,
+  StartDocumentAnalysisCommandOutput,
+  GetDocumentAnalysisCommand,
+  GetDocumentAnalysisCommandOutput,
+  StartDocumentTextDetectionCommand,
+  StartDocumentTextDetectionCommandOutput,
+  GetDocumentTextDetectionCommand,
+  GetDocumentTextDetectionCommandOutput,
+  Document,
+  TextractClient,
+} from '@aws-sdk/client-textract';
+import { ITextract } from '../../../lib/documents/interfaces/documents.interface';
+import {
+  AnalyzeDocumentInput,
+  DetectDocumentTextBytesInput,
+  DetectDocumentTextS3Input,
+  StartDocumentAnalysisInput,
+  GetDocumentAnalysisInput,
+  StartDocumentTextDetectionInput,
+  GetDocumentTextDetectionInput,
+} from '../../../lib/documents/types/textract.type';
+
+@Injectable()
+export class Textract implements ITextract {
+  constructor(
+    @Inject('CONNECTION_TEXTRACT') private readonly client: TextractClient,
+  ) {}
+
+  async detectDocumentTextBytes(
+    input: DetectDocumentTextBytesInput,
+  ): Promise<DetectDocumentTextCommandOutput> {
+    const document: Document = { Bytes: input.bytes };
+    const command = new DetectDocumentTextCommand({ Document: document });
+    return await this.client.send(command);
+  }
+
+  async detectDocumentTextS3(
+    input: DetectDocumentTextS3Input,
+  ): Promise<DetectDocumentTextCommandOutput> {
+    const document: Document = {
+      S3Object: { Bucket: input.bucket, Name: input.name },
+    };
+    const command = new DetectDocumentTextCommand({ Document: document });
+    return await this.client.send(command);
+  }
+
+  async analyzeDocument(
+    input: AnalyzeDocumentInput,
+  ): Promise<AnalyzeDocumentCommandOutput> {
+    const document: Document = input.bytes
+      ? { Bytes: input.bytes }
+      : { S3Object: { Bucket: input.bucket, Name: input.name } };
+
+    const command = new AnalyzeDocumentCommand({
+      Document: document,
+      FeatureTypes: input.featureTypes,
+      ...(input.queriesConfig && { QueriesConfig: input.queriesConfig }),
+      ...(input.adaptersConfig && { AdaptersConfig: input.adaptersConfig }),
+      ...(input.humanLoopConfig && { HumanLoopConfig: input.humanLoopConfig }),
+    });
+
+    return await this.client.send(command);
+  }
+
+  async startDocumentAnalysis(
+    input: StartDocumentAnalysisInput,
+  ): Promise<StartDocumentAnalysisCommandOutput> {
+    const command = new StartDocumentAnalysisCommand({
+      DocumentLocation: {
+        S3Object: {
+          Bucket: input.bucket,
+          Name: input.name,
+          ...(input.version && { Version: input.version }),
+        },
+      },
+      FeatureTypes: input.featureTypes,
+      ...(input.queriesConfig && { QueriesConfig: input.queriesConfig }),
+      ...(input.adaptersConfig && { AdaptersConfig: input.adaptersConfig }),
+      ...(input.clientRequestToken && {
+        ClientRequestToken: input.clientRequestToken,
+      }),
+      ...(input.jobTag && { JobTag: input.jobTag }),
+      ...(input.notificationChannel && {
+        NotificationChannel: {
+          SNSTopicArn: input.notificationChannel.snsTopicArn,
+          RoleArn: input.notificationChannel.roleArn,
+        },
+      }),
+      ...(input.outputConfig && {
+        OutputConfig: {
+          S3Bucket: input.outputConfig.s3Bucket,
+          ...(input.outputConfig.s3Prefix && {
+            S3Prefix: input.outputConfig.s3Prefix,
+          }),
+        },
+      }),
+      ...(input.kmsKeyId && { KMSKeyId: input.kmsKeyId }),
+    });
+
+    return await this.client.send(command);
+  }
+
+  async getDocumentAnalysis(
+    input: GetDocumentAnalysisInput,
+  ): Promise<GetDocumentAnalysisCommandOutput> {
+    const command = new GetDocumentAnalysisCommand({
+      JobId: input.jobId,
+      ...(input.maxResults && { MaxResults: input.maxResults }),
+      ...(input.nextToken && { NextToken: input.nextToken }),
+    });
+
+    return await this.client.send(command);
+  }
+
+  async startDocumentTextDetection(
+    input: StartDocumentTextDetectionInput,
+  ): Promise<StartDocumentTextDetectionCommandOutput> {
+    const command = new StartDocumentTextDetectionCommand({
+      DocumentLocation: {
+        S3Object: {
+          Bucket: input.bucket,
+          Name: input.name,
+          ...(input.version && { Version: input.version }),
+        },
+      },
+      ...(input.clientRequestToken && {
+        ClientRequestToken: input.clientRequestToken,
+      }),
+      ...(input.jobTag && { JobTag: input.jobTag }),
+      ...(input.notificationChannel && {
+        NotificationChannel: {
+          SNSTopicArn: input.notificationChannel.snsTopicArn,
+          RoleArn: input.notificationChannel.roleArn,
+        },
+      }),
+      ...(input.outputConfig && {
+        OutputConfig: {
+          S3Bucket: input.outputConfig.s3Bucket,
+          ...(input.outputConfig.s3Prefix && {
+            S3Prefix: input.outputConfig.s3Prefix,
+          }),
+        },
+      }),
+      ...(input.kmsKeyId && { KMSKeyId: input.kmsKeyId }),
+    });
+
+    return await this.client.send(command);
+  }
+
+  async getDocumentTextDetection(
+    input: GetDocumentTextDetectionInput,
+  ): Promise<GetDocumentTextDetectionCommandOutput> {
+    const command = new GetDocumentTextDetectionCommand({
+      JobId: input.jobId,
+      ...(input.maxResults && { MaxResults: input.maxResults }),
+      ...(input.nextToken && { NextToken: input.nextToken }),
+    });
+
+    return await this.client.send(command);
+  }
+}

package/src/external/redis/interfaces/redis.interface.ts

@@ -1,3 +1,5 @@
+import { PubSubRedis } from '../types/ioredis.types';
+
 export interface IRedisService {
   setRedisKey(
     key: string,
@@ -8,6 +10,8 @@ export interface IRedisService {
   getRedisKey(key: string): Promise<string>;
 
   deleteRedisKey(key: string): Promise<void>;
+
+  createPubSubClient(): Promise<PubSubRedis>;
 }
 
 export const IRedisService = Symbol('IRedisService');

package/src/external/redis/redis.service.ts

@@ -1,6 +1,7 @@
 import { Inject, Injectable } from '@nestjs/common';
 import { Cluster, Redis } from 'ioredis';
 import { IRedisService } from './interfaces/redis.interface';
+import { PubSubRedis } from './types/ioredis.types';
 
 @Injectable()
 export class RedisService implements IRedisService {
@@ -29,4 +30,11 @@ export class RedisService implements IRedisService {
   async deleteRedisKey(key: string): Promise<void> {
     await this.client.del(key);
   }
+
+  async createPubSubClient(): Promise<PubSubRedis> {
+    return {
+      pub: this.client,
+      sub: this.client.duplicate(),
+    };
+  }
 }

package/src/external/redis/types/ioredis.types.ts

@@ -1,4 +1,9 @@
-import { ClusterOptions, RedisOptions } from 'ioredis';
+import {
+  type Cluster,
+  type Redis,
+  ClusterOptions,
+  RedisOptions,
+} from 'ioredis';
 
 export enum TypeRedis {
   SIMPLE = 'simple',
@@ -21,3 +26,8 @@ export type ClusterRedisSetup = {
   nodes: Connection[];
   clusterOptions?: ClusterOptions;
 };
+
+export type PubSubRedis = {
+  pub: Redis | Cluster;
+  sub: Redis | Cluster;
+};

package/src/index.ts

@@ -21,10 +21,17 @@ import { ICrypto } from './lib/crypto/interfaces/crypto.interface';
 import { CacheModule } from './lib/cache/cache.module';
 import { Cache } from './lib/cache/cache';
 import { ICache } from './lib/cache/interfaces/cache.interface';
+import { DocumentsModule } from './lib/documents/documents.module';
+import { Documents } from './lib/documents/documents';
+import {
+  IBedrock,
+  ITextract,
+} from './lib/documents/interfaces/documents.interface';
 import {
   SimpleRedisSetup,
   ClusterRedisSetup,
   TypeRedis,
+  PubSubRedis,
 } from './external/redis/types/ioredis.types';
 
 export {
@@ -46,7 +53,12 @@ export {
   CacheModule,
   Cache,
   ICache,
+  DocumentsModule,
+  Documents,
+  IBedrock,
+  ITextract,
   SimpleRedisSetup,
   ClusterRedisSetup,
   TypeRedis,
+  PubSubRedis,
 };

package/src/lib/cache/cache.ts

@@ -1,6 +1,7 @@
 import { Inject, Injectable } from '@nestjs/common';
 import { ICache } from './interfaces/cache.interface';
 import { IRedisService } from 'src/external/redis/interfaces/redis.interface';
+import { PubSubRedis } from 'src/external/redis/types/ioredis.types';
 
 @Injectable()
 export class Cache implements ICache {
@@ -19,4 +20,8 @@ export class Cache implements ICache {
   async delete(key: string): Promise<void> {
     await this.client.deleteRedisKey(key);
   }
+
+  async createPubSubClient(): Promise<PubSubRedis> {
+    return await this.client.createPubSubClient();
+  }
 }

package/src/lib/cache/interfaces/cache.interface.ts

@@ -1,3 +1,5 @@
+import { PubSubRedis } from 'src/external/redis/types/ioredis.types';
+
 export interface ICache {
   set(
     key: string,
@@ -8,6 +10,8 @@ export interface ICache {
   get(key: string): Promise<string>;
 
   delete(key: string): Promise<void>;
+
+  createPubSubClient(): Promise<PubSubRedis>;
 }
 
 export const ICache = Symbol('ICache');

package/src/lib/documents/documents.module.ts

@@ -0,0 +1,26 @@
+import { DynamicModule, Module } from '@nestjs/common';
+import {
+  BedrockModule as BedrockProviderModule,
+  BedrockSetup,
+} from '../../external/aws/bedrock/bedrock.provider';
+import { Bedrock as BedrockService } from '../../external/aws/bedrock/bedrock.service';
+import { TextractSetup } from '../../external/aws/textract/textract.provider';
+import { Documents } from './documents';
+import { IBedrock } from './interfaces/documents.interface';
+
+export type DocumentsSetup = {
+  bedrock: BedrockSetup;
+  textract?: TextractSetup;
+};
+
+@Module({})
+export class DocumentsModule {
+  static forRoot(params: DocumentsSetup): DynamicModule {
+    return {
+      module: DocumentsModule,
+      imports: [BedrockProviderModule.forRoot(params.bedrock)],
+      providers: [Documents, { provide: IBedrock, useClass: BedrockService }],
+      exports: [Documents, { provide: IBedrock, useClass: BedrockService }],
+    };
+  }
+}

package/src/lib/documents/documents.ts

@@ -0,0 +1,83 @@
+import { Inject, Injectable, Optional } from '@nestjs/common';
+import { IBedrock, ITextract } from './interfaces/documents.interface';
+import { InvokeModelCommandOutput } from '@aws-sdk/client-bedrock-runtime';
+import {
+  AnalyzeDocumentCommandOutput,
+  DetectDocumentTextCommandOutput,
+  StartDocumentAnalysisCommandOutput,
+  GetDocumentAnalysisCommandOutput,
+  StartDocumentTextDetectionCommandOutput,
+  GetDocumentTextDetectionCommandOutput,
+} from '@aws-sdk/client-textract';
+import { InvokeModelPayload } from './types/bedrock.type';
+import {
+  AnalyzeDocumentInput,
+  DetectDocumentTextBytesInput,
+  DetectDocumentTextS3Input,
+  StartDocumentAnalysisInput,
+  GetDocumentAnalysisInput,
+  StartDocumentTextDetectionInput,
+  GetDocumentTextDetectionInput,
+} from './types/textract.type';
+
+@Injectable()
+export class Documents {
+  constructor(
+    @Optional()
+    @Inject(IBedrock)
+    private readonly bedrock: IBedrock,
+    @Optional()
+    @Inject(ITextract)
+    private readonly textract: ITextract,
+  ) {}
+
+  // bedrock
+  async invokeBedrockModel(
+    payload: InvokeModelPayload,
+  ): Promise<InvokeModelCommandOutput> {
+    return await this.bedrock.invokeModel(payload);
+  }
+
+  // textract
+  async detectDocumentTextBytes(
+    input: DetectDocumentTextBytesInput,
+  ): Promise<DetectDocumentTextCommandOutput> {
+    return await this.textract.detectDocumentTextBytes(input);
+  }
+
+  async detectDocumentTextS3(
+    input: DetectDocumentTextS3Input,
+  ): Promise<DetectDocumentTextCommandOutput> {
+    return await this.textract.detectDocumentTextS3(input);
+  }
+
+  async analyzeDocument(
+    input: AnalyzeDocumentInput,
+  ): Promise<AnalyzeDocumentCommandOutput> {
+    return await this.textract.analyzeDocument(input);
+  }
+
+  async startDocumentAnalysis(
+    input: StartDocumentAnalysisInput,
+  ): Promise<StartDocumentAnalysisCommandOutput> {
+    return await this.textract.startDocumentAnalysis(input);
+  }
+
+  async getDocumentAnalysis(
+    input: GetDocumentAnalysisInput,
+  ): Promise<GetDocumentAnalysisCommandOutput> {
+    return await this.textract.getDocumentAnalysis(input);
+  }
+
+  async startDocumentTextDetection(
+    input: StartDocumentTextDetectionInput,
+  ): Promise<StartDocumentTextDetectionCommandOutput> {
+    return await this.textract.startDocumentTextDetection(input);
+  }
+
+  async getDocumentTextDetection(
+    input: GetDocumentTextDetectionInput,
+  ): Promise<GetDocumentTextDetectionCommandOutput> {
+    return await this.textract.getDocumentTextDetection(input);
+  }
+}

package/src/lib/documents/interfaces/documents.interface.ts

@@ -0,0 +1,51 @@
+import { InvokeModelCommandOutput } from '@aws-sdk/client-bedrock-runtime';
+import {
+  AnalyzeDocumentCommandOutput,
+  DetectDocumentTextCommandOutput,
+  StartDocumentAnalysisCommandOutput,
+  GetDocumentAnalysisCommandOutput,
+  StartDocumentTextDetectionCommandOutput,
+  GetDocumentTextDetectionCommandOutput,
+} from '@aws-sdk/client-textract';
+import { InvokeModelPayload } from '../types/bedrock.type';
+import {
+  AnalyzeDocumentInput,
+  DetectDocumentTextBytesInput,
+  DetectDocumentTextS3Input,
+  StartDocumentAnalysisInput,
+  GetDocumentAnalysisInput,
+  StartDocumentTextDetectionInput,
+  GetDocumentTextDetectionInput,
+} from '../types/textract.type';
+
+export interface IBedrock {
+  invokeModel(payload: InvokeModelPayload): Promise<InvokeModelCommandOutput>;
+}
+
+export const IBedrock = Symbol('IBedrock');
+
+export interface ITextract {
+  detectDocumentTextBytes(
+    input: DetectDocumentTextBytesInput,
+  ): Promise<DetectDocumentTextCommandOutput>;
+  detectDocumentTextS3(
+    input: DetectDocumentTextS3Input,
+  ): Promise<DetectDocumentTextCommandOutput>;
+  analyzeDocument(
+    input: AnalyzeDocumentInput,
+  ): Promise<AnalyzeDocumentCommandOutput>;
+  startDocumentAnalysis(
+    input: StartDocumentAnalysisInput,
+  ): Promise<StartDocumentAnalysisCommandOutput>;
+  getDocumentAnalysis(
+    input: GetDocumentAnalysisInput,
+  ): Promise<GetDocumentAnalysisCommandOutput>;
+  startDocumentTextDetection(
+    input: StartDocumentTextDetectionInput,
+  ): Promise<StartDocumentTextDetectionCommandOutput>;
+  getDocumentTextDetection(
+    input: GetDocumentTextDetectionInput,
+  ): Promise<GetDocumentTextDetectionCommandOutput>;
+}
+
+export const ITextract = Symbol('ITextract');

package/src/lib/documents/types/bedrock.type.ts

@@ -0,0 +1,6 @@
+export type InvokeModelPayload = {
+  modelId: string;
+  body: Uint8Array | Buffer | string;
+  contentType?: string;
+  accept?: string;
+};

package/src/lib/documents/types/textract.type.ts

@@ -0,0 +1,78 @@
+import {
+  FeatureType,
+  QueriesConfig,
+  AdaptersConfig,
+  HumanLoopConfig,
+} from '@aws-sdk/client-textract';
+
+export type DetectDocumentTextBytesInput = {
+  bytes: Uint8Array;
+};
+
+export type DetectDocumentTextS3Input = {
+  bucket: string;
+  name: string;
+};
+
+export type AnalyzeDocumentInput =
+  | ({ bytes: Uint8Array; bucket?: never; name?: never } & {
+      featureTypes: FeatureType[];
+      queriesConfig?: QueriesConfig;
+      adaptersConfig?: AdaptersConfig;
+      humanLoopConfig?: HumanLoopConfig;
+    })
+  | ({ bytes?: never; bucket: string; name: string } & {
+      featureTypes: FeatureType[];
+      queriesConfig?: QueriesConfig;
+      adaptersConfig?: AdaptersConfig;
+      humanLoopConfig?: HumanLoopConfig;
+    });
+
+export type StartDocumentAnalysisInput = {
+  bucket: string;
+  name: string;
+  version?: string;
+  featureTypes: FeatureType[];
+  queriesConfig?: QueriesConfig;
+  adaptersConfig?: AdaptersConfig;
+  clientRequestToken?: string;
+  jobTag?: string;
+  notificationChannel?: {
+    snsTopicArn: string;
+    roleArn: string;
+  };
+  outputConfig?: {
+    s3Bucket: string;
+    s3Prefix?: string;
+  };
+  kmsKeyId?: string;
+};
+
+export type GetDocumentAnalysisInput = {
+  jobId: string;
+  maxResults?: number;
+  nextToken?: string;
+};
+
+export type StartDocumentTextDetectionInput = {
+  bucket: string;
+  name: string;
+  version?: string;
+  clientRequestToken?: string;
+  jobTag?: string;
+  notificationChannel?: {
+    snsTopicArn: string;
+    roleArn: string;
+  };
+  outputConfig?: {
+    s3Bucket: string;
+    s3Prefix?: string;
+  };
+  kmsKeyId?: string;
+};
+
+export type GetDocumentTextDetectionInput = {
+  jobId: string;
+  maxResults?: number;
+  nextToken?: string;
+};