diff --git a/api-specifications/src/embeddings/__snapshots__/constants.test.ts.snap b/api-specifications/src/embeddings/__snapshots__/constants.test.ts.snap new file mode 100644 index 000000000..a1e3e30f6 --- /dev/null +++ b/api-specifications/src/embeddings/__snapshots__/constants.test.ts.snap @@ -0,0 +1,11 @@ +// Jest Snapshot v1, https://goo.gl/fbAQLP + +exports[`Test the Embeddings Constants the embedding model 'models/gemini-embedding-2' should match the snapshot 1`] = ` +{ + "enabled": false, + "id": "77bb8ff3-a6b0-460b-bcaa-00631a907852", + "modelName": "models/gemini-embedding-2", + "modelProvider": "gemini", + "numberOfDimensions": 768, +} +`; diff --git a/api-specifications/src/embeddings/__snapshots__/index.test.ts.snap b/api-specifications/src/embeddings/__snapshots__/index.test.ts.snap new file mode 100644 index 000000000..44343ef77 --- /dev/null +++ b/api-specifications/src/embeddings/__snapshots__/index.test.ts.snap @@ -0,0 +1,71 @@ +// Jest Snapshot v1, https://goo.gl/fbAQLP + +exports[`Test the embeddings module The embeddings module matches the snapshot 1`] = ` +{ + "Constants": { + "EEmbeddingModelProvider": { + "GEMINI": "gemini", + }, + "EmbeddingServiceIds": [ + "77bb8ff3-a6b0-460b-bcaa-00631a907852", + ], + "EmbeddingServices": [ + { + "enabled": false, + "id": "77bb8ff3-a6b0-460b-bcaa-00631a907852", + "modelName": "models/gemini-embedding-2", + "modelProvider": "gemini", + "numberOfDimensions": 768, + }, + ], + "MAX_EMBEDDING_MODEL_ID_LENGTH": 36, + "MAX_NUMBER_OF_DIMENSIONS": 100000, + "MIN_NUMBER_OF_DIMENSIONS": 1, + "MODEL_NAME_MAX_LENGTH": 256, + }, + "Schemas": { + "EmbeddingModel": { + "$id": "/components/schemas/EmbeddingModelSchema", + "additionalProperties": false, + "properties": { + "enabled": { + "description": "Whether the embedding model is enabled and can be used to trigger an embedding process.", + "type": "boolean", + }, + "id": { + "description": "The UUID of the embedding model.", + "pattern": "^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$", + "type": "string", + }, + "modelName": { + "description": "The provider specific name of the embedding model.", + "maxLength": 256, + "pattern": "\\S", + "type": "string", + }, + "modelProvider": { + "description": "The provider of the embedding model.", + "enum": [ + "gemini", + ], + "type": "string", + }, + "numberOfDimensions": { + "description": "The number of dimensions of the vectors produced by the embedding model.", + "maximum": 100000, + "minimum": 1, + "type": "integer", + }, + }, + "required": [ + "id", + "modelProvider", + "modelName", + "numberOfDimensions", + "enabled", + ], + "type": "object", + }, + }, +} +`; diff --git a/api-specifications/src/embeddings/constants.test.ts b/api-specifications/src/embeddings/constants.test.ts new file mode 100644 index 000000000..a28147d78 --- /dev/null +++ b/api-specifications/src/embeddings/constants.test.ts @@ -0,0 +1,34 @@ +import EmbeddingsConstants from "./constants"; + +describe("Test the Embeddings Constants", () => { + describe("Test the EmbeddingModels registry", () => { + test("every embedding model should have a unique id in the entire registry", () => { + // GIVEN the registry of embedding models + const givenEmbeddingModels = EmbeddingsConstants.EmbeddingServices; + + // WHEN the unique ids are extracted from the registry + const actualUniqueIds = new Set(givenEmbeddingModels.map((model) => model.id)); + + // THEN expect as many unique ids as there are registered embedding models + expect(actualUniqueIds.size).toBe(givenEmbeddingModels.length); + }); + + test("the EmbeddingServiceIds should have one id for every registered embedding model", () => { + // GIVEN the registry of embedding models + // WHEN the supported embedding model ids are derived from the registry + // THEN expect the number of ids to be the same as the number of registered embedding models + expect(EmbeddingsConstants.EmbeddingServiceIds).toHaveLength(EmbeddingsConstants.EmbeddingServices.length); + }); + }); + + test.each(EmbeddingsConstants.EmbeddingServices.map((model) => [model.modelName, model] as const))( + "the embedding model '%s' should match the snapshot", + (_modelName, givenEmbeddingModel) => { + // GIVEN an embedding model from the registry + + // WHEN the embedding model is inspected + // THEN expect the embedding model to match the snapshot + expect(givenEmbeddingModel).toMatchSnapshot(); + } + ); +}); diff --git a/api-specifications/src/embeddings/constants.ts b/api-specifications/src/embeddings/constants.ts new file mode 100644 index 000000000..8f7352586 --- /dev/null +++ b/api-specifications/src/embeddings/constants.ts @@ -0,0 +1,41 @@ +namespace EmbeddingsConstants { + export enum EEmbeddingModelProvider { + GEMINI = "gemini", + } + + export interface IEmbeddingService { + id: string; + modelProvider: EEmbeddingModelProvider; + modelName: string; + numberOfDimensions: number; + enabled: boolean; + } + + export const EmbeddingServices: IEmbeddingService[] = [ + { + id: `77bb8ff3-a6b0-460b-bcaa-00631a907852`, + modelProvider: EEmbeddingModelProvider.GEMINI, + modelName: "models/gemini-embedding-2", + numberOfDimensions: 768, + enabled: false, + }, + ]; + + /** + * The identifiers of the embedding models that are supported. + * These are the only values accepted as the `embeddingServiceId` when triggering an embedding process. + */ + export const EmbeddingServiceIds = Array.from(new Set(EmbeddingServices.map((model) => model.id))); + + // Since the embedding model id is a UUID, it is 36 characters long. + export const MAX_EMBEDDING_MODEL_ID_LENGTH = 36; + + // The maximum length of the human/provider readable name of an embedding model (e.g. "models/gemini-embedding-2"). + export const MODEL_NAME_MAX_LENGTH = 256; + + // The bounds for the number of dimensions produced by an embedding model. + export const MIN_NUMBER_OF_DIMENSIONS = 1; + export const MAX_NUMBER_OF_DIMENSIONS = 100000; +} + +export default EmbeddingsConstants; diff --git a/api-specifications/src/embeddings/index.test.ts b/api-specifications/src/embeddings/index.test.ts new file mode 100644 index 000000000..cd2d92369 --- /dev/null +++ b/api-specifications/src/embeddings/index.test.ts @@ -0,0 +1,25 @@ +describe("Test the embeddings module", () => { + test("The embeddings module can be required via the index", async () => { + // GIVEN the module + // WHEN the module is required via the index + expect(async () => await import("./")).not.toThrow(); // We check that it doesn't throw an error instead of simply letting it fail on import because we want an easier error message + const embeddingsModule = await import("./"); + + // THEN Check if the Constants are defined in it + expect(embeddingsModule.default.Constants.EmbeddingServices).toBeDefined(); + expect(embeddingsModule.default.Constants.EmbeddingServiceIds).toBeDefined(); + expect(embeddingsModule.default.Constants.MAX_EMBEDDING_MODEL_ID_LENGTH).toBeDefined(); + }); + + test("The embeddings module matches the snapshot", () => { + // GIVEN the module + // WHEN the module is required via the index + expect(async () => { + // THEN Check if the module can be required without error + const embeddingsModule = await import("./"); + expect(embeddingsModule.default).toMatchSnapshot(); + }).not.toThrowError(); + }); +}); + +export {}; diff --git a/api-specifications/src/embeddings/index.ts b/api-specifications/src/embeddings/index.ts new file mode 100644 index 000000000..77659e746 --- /dev/null +++ b/api-specifications/src/embeddings/index.ts @@ -0,0 +1,28 @@ +// we need to disable the eslint rule here because this is a top level export +/* eslint-disable @typescript-eslint/no-unused-vars */ + +import EmbeddingsConstants from "./constants"; +import EmbeddingModelSchema from "./schema"; + +/** + * This file should be imported in the following way + + import EmbeddingsAPISpecs from "api-specifications/embeddings"; + + * And the general pattern is EmbeddingsAPISpecs.{Constants/Schemas}.{...} + * + * NOTE: The API contract for /models/{modelId}/embedding-process-states (schemas, types, enums) + * lives under modelInfo — see ModelInfoAPISpecs.ModelInfo.EmbeddingProcessStates. + * Only the embedding-model constants remain here. + */ + +namespace EmbeddingsSchemas { + export const EmbeddingModel = EmbeddingModelSchema; +} + +namespace EmbeddingsAPISpecs { + export import Constants = EmbeddingsConstants; + export import Schemas = EmbeddingsSchemas; +} + +export default EmbeddingsAPISpecs; diff --git a/api-specifications/src/embeddings/package.json b/api-specifications/src/embeddings/package.json new file mode 100644 index 000000000..14ab704d8 --- /dev/null +++ b/api-specifications/src/embeddings/package.json @@ -0,0 +1,3 @@ +{ + "main": "index.js" +} diff --git a/api-specifications/src/embeddings/schema.test.ts b/api-specifications/src/embeddings/schema.test.ts new file mode 100644 index 000000000..2a370b1bf --- /dev/null +++ b/api-specifications/src/embeddings/schema.test.ts @@ -0,0 +1,145 @@ +import { randomUUID } from "crypto"; +import Ajv from "ajv"; +import addFormats from "ajv-formats"; +import { getTestString } from "_test_utilities/specialCharacters"; +import { + testBooleanField, + testEnumField, + testNonEmptyStringField, + testSchemaWithAdditionalProperties, + testSchemaWithValidObject, + testUUIDField, + testValidSchema, +} from "_test_utilities/stdSchemaTests"; +import { assertCaseForProperty, CaseType, constructSchemaError } from "_test_utilities/assertCaseForProperty"; +import EmbeddingsAPISpecs from "./index"; +import EmbeddingsConstants from "./constants"; + +describe("Test the EmbeddingModel Schema", () => { + // GIVEN the EmbeddingsAPISpecs.Schemas.EmbeddingModel schema + // WHEN the schema is validated + // THEN expect the schema to be valid + testValidSchema("EmbeddingsAPISpecs.Schemas.EmbeddingModel", EmbeddingsAPISpecs.Schemas.EmbeddingModel); + + describe("Validate JSON against the EmbeddingModel Schema", () => { + // GIVEN a valid EmbeddingModel object + const givenValidEmbeddingModel: EmbeddingsConstants.IEmbeddingService = { + id: randomUUID(), + modelProvider: EmbeddingsConstants.EEmbeddingModelProvider.GEMINI, + modelName: getTestString(EmbeddingsConstants.MODEL_NAME_MAX_LENGTH), + numberOfDimensions: 768, + enabled: true, + }; + + // WHEN the object is validated + // THEN expect the object to validate successfully + testSchemaWithValidObject( + "EmbeddingsAPISpecs.Schemas.EmbeddingModel", + EmbeddingsAPISpecs.Schemas.EmbeddingModel, + givenValidEmbeddingModel + ); + + // AND WHEN the object has additional properties + // THEN expect the object to not validate + testSchemaWithAdditionalProperties( + "EmbeddingsAPISpecs.Schemas.EmbeddingModel", + EmbeddingsAPISpecs.Schemas.EmbeddingModel, + givenValidEmbeddingModel + ); + }); + + describe("validate EmbeddingsAPISpecs.Schemas.EmbeddingModel fields", () => { + describe("Test validation of 'id'", () => { + testUUIDField("id", EmbeddingsAPISpecs.Schemas.EmbeddingModel); + }); + + describe("Test validation of 'modelProvider'", () => { + testEnumField( + "modelProvider", + EmbeddingsAPISpecs.Schemas.EmbeddingModel, + Object.values(EmbeddingsConstants.EEmbeddingModelProvider) + ); + }); + + describe("Test validation of 'modelName'", () => { + testNonEmptyStringField( + "modelName", + EmbeddingsConstants.MODEL_NAME_MAX_LENGTH, + EmbeddingsAPISpecs.Schemas.EmbeddingModel + ); + }); + + describe("Test validation of 'numberOfDimensions'", () => { + test.each([ + [ + CaseType.Failure, + "undefined", + undefined, + constructSchemaError("", "required", "must have required property 'numberOfDimensions'"), + ], + [CaseType.Failure, "null", null, constructSchemaError("/numberOfDimensions", "type", "must be integer")], + [CaseType.Failure, "a string", "768", constructSchemaError("/numberOfDimensions", "type", "must be integer")], + [CaseType.Failure, "a float", 768.5, constructSchemaError("/numberOfDimensions", "type", "must be integer")], + [ + CaseType.Failure, + "less than the minimum", + EmbeddingsConstants.MIN_NUMBER_OF_DIMENSIONS - 1, + constructSchemaError( + "/numberOfDimensions", + "minimum", + `must be >= ${EmbeddingsConstants.MIN_NUMBER_OF_DIMENSIONS}` + ), + ], + [ + CaseType.Failure, + "greater than the maximum", + EmbeddingsConstants.MAX_NUMBER_OF_DIMENSIONS + 1, + constructSchemaError( + "/numberOfDimensions", + "maximum", + `must be <= ${EmbeddingsConstants.MAX_NUMBER_OF_DIMENSIONS}` + ), + ], + [CaseType.Success, "the minimum", EmbeddingsConstants.MIN_NUMBER_OF_DIMENSIONS, undefined], + [CaseType.Success, "the maximum", EmbeddingsConstants.MAX_NUMBER_OF_DIMENSIONS, undefined], + [CaseType.Success, "a valid value", 768, undefined], + ])("(%s) Validate 'numberOfDimensions' when it is %s", (caseType, _description, givenValue, failureMessages) => { + // GIVEN an object with the given value + const givenObject = { numberOfDimensions: givenValue }; + // THEN expect the object to validate accordingly + assertCaseForProperty( + "numberOfDimensions", + givenObject, + EmbeddingsAPISpecs.Schemas.EmbeddingModel, + caseType, + failureMessages + ); + }); + }); + + describe("Test validation of 'enabled'", () => { + testBooleanField("enabled", EmbeddingsAPISpecs.Schemas.EmbeddingModel); + }); + }); + + describe("Validate the registered EmbeddingModels against the EmbeddingModel Schema", () => { + // GIVEN an AJV instance compiled with the EmbeddingModel schema + const givenAjv = new Ajv({ validateSchema: true, allErrors: true, strict: true }); + addFormats(givenAjv); + const givenValidateFunction = givenAjv.compile(EmbeddingsAPISpecs.Schemas.EmbeddingModel); + + // AND every embedding model registered in the constants registry + test.each(EmbeddingsConstants.EmbeddingServices.map((model) => [model.modelName, model] as const))( + "the registered embedding model '%s' validates against the EmbeddingModel schema", + (_modelName, givenEmbeddingModel) => { + // WHEN the registered embedding model is validated against the schema + const actualIsValid = givenValidateFunction(givenEmbeddingModel); + + // THEN expect the embedding model to validate successfully + expect(actualIsValid).toBe(true); + // AND no errors to be present + expect(givenValidateFunction.errors).toBeNull(); + } + ); + }); +}); diff --git a/api-specifications/src/embeddings/schema.ts b/api-specifications/src/embeddings/schema.ts new file mode 100644 index 000000000..648b78dff --- /dev/null +++ b/api-specifications/src/embeddings/schema.ts @@ -0,0 +1,40 @@ +import { SchemaObject } from "ajv"; +import { RegExp_Str_NotEmptyString, RegExp_Str_UUIDv4 } from "../regex"; +import EmbeddingsConstants from "./constants"; + +const EmbeddingModelSchema: SchemaObject = { + $id: "/components/schemas/EmbeddingModelSchema", + type: "object", + additionalProperties: false, + properties: { + id: { + description: "The UUID of the embedding model.", + type: "string", + pattern: RegExp_Str_UUIDv4, + }, + modelProvider: { + description: "The provider of the embedding model.", + type: "string", + enum: Object.values(EmbeddingsConstants.EEmbeddingModelProvider), + }, + modelName: { + description: "The provider specific name of the embedding model.", + type: "string", + pattern: RegExp_Str_NotEmptyString, + maxLength: EmbeddingsConstants.MODEL_NAME_MAX_LENGTH, + }, + numberOfDimensions: { + description: "The number of dimensions of the vectors produced by the embedding model.", + type: "integer", + minimum: EmbeddingsConstants.MIN_NUMBER_OF_DIMENSIONS, + maximum: EmbeddingsConstants.MAX_NUMBER_OF_DIMENSIONS, + }, + enabled: { + description: "Whether the embedding model is enabled and can be used to trigger an embedding process.", + type: "boolean", + }, + }, + required: ["id", "modelProvider", "modelName", "numberOfDimensions", "enabled"], +}; + +export default EmbeddingModelSchema; diff --git a/api-specifications/src/error/__snapshots__/index.test.ts.snap b/api-specifications/src/error/__snapshots__/index.test.ts.snap index 3aa2c3254..feb5ecac4 100644 --- a/api-specifications/src/error/__snapshots__/index.test.ts.snap +++ b/api-specifications/src/error/__snapshots__/index.test.ts.snap @@ -281,6 +281,10 @@ exports[`Test the error module The export module matches the snapshot 1`] = ` "SKILL_GROUP_PARENT_NOT_FOUND", "DB_FAILED_TO_RETRIEVE_SKILL_GROUP_PARENT", "DB_FAILED_TO_RETRIEVE_SKILL_GROUP_CHILDREN", + "MODEL_NOT_RELEASED", + "MODEL_NOT_FOUND_BY_ID", + "EMBEDDING_PROCESS_ALREADY_RUNNING", + "FAILED_TO_TRIGGER_EMBEDDING_PROCESS", ], "pattern": "\\S", "type": "string", diff --git a/api-specifications/src/error/schema.test.ts b/api-specifications/src/error/schema.test.ts index 916b09fd8..41791d511 100644 --- a/api-specifications/src/error/schema.test.ts +++ b/api-specifications/src/error/schema.test.ts @@ -91,6 +91,10 @@ describe("Validate JSON against the APIError Schema", () => { ...Object.values(SkillsAPI.GET.Errors.Status400.ErrorCodes), ...Object.values(SkillsAPI.GET.Errors.Status404.ErrorCodes), ...Object.values(SkillsAPI.GET.Errors.Status500.ErrorCodes), + ...Object.values(ModelInfoAPI.ModelInfo.EmbeddingProcessStates.POST.Enums.Response.Status400.ErrorCodes), + ...Object.values(ModelInfoAPI.ModelInfo.EmbeddingProcessStates.POST.Enums.Response.Status404.ErrorCodes), + ...Object.values(ModelInfoAPI.ModelInfo.EmbeddingProcessStates.POST.Enums.Response.Status409.ErrorCodes), + ...Object.values(ModelInfoAPI.ModelInfo.EmbeddingProcessStates.POST.Enums.Response.Status500.ErrorCodes), ]; test.each([ diff --git a/api-specifications/src/error/schema.ts b/api-specifications/src/error/schema.ts index 28b526e10..4c5a789a9 100644 --- a/api-specifications/src/error/schema.ts +++ b/api-specifications/src/error/schema.ts @@ -87,6 +87,10 @@ export const ErrorSchema: SchemaObject = { ...Object.values(SkillGroupAPI.SkillGroup.Children.GET.Enums.Response.Status400.ErrorCodes), ...Object.values(SkillGroupAPI.SkillGroup.Children.GET.Enums.Response.Status404.ErrorCodes), ...Object.values(SkillGroupAPI.SkillGroup.Children.GET.Enums.Response.Status500.ErrorCodes), + ...Object.values(ModelInfoAPI.ModelInfo.EmbeddingProcessStates.POST.Enums.Response.Status400.ErrorCodes), + ...Object.values(ModelInfoAPI.ModelInfo.EmbeddingProcessStates.POST.Enums.Response.Status404.ErrorCodes), + ...Object.values(ModelInfoAPI.ModelInfo.EmbeddingProcessStates.POST.Enums.Response.Status409.ErrorCodes), + ...Object.values(ModelInfoAPI.ModelInfo.EmbeddingProcessStates.POST.Enums.Response.Status500.ErrorCodes), ]) ), pattern: RegExp_Str_NotEmptyString, diff --git a/api-specifications/src/modelInfo/[id]/embeddingProcessStates/POST/constants.ts b/api-specifications/src/modelInfo/[id]/embeddingProcessStates/POST/constants.ts new file mode 100644 index 000000000..5e961fc32 --- /dev/null +++ b/api-specifications/src/modelInfo/[id]/embeddingProcessStates/POST/constants.ts @@ -0,0 +1,7 @@ +import EmbeddingsConstants from "../../../../embeddings/constants"; + +namespace EmbeddingProcessStatesConstants { + export const MAX_PAYLOAD_LENGTH = EmbeddingsConstants.MAX_EMBEDDING_MODEL_ID_LENGTH + 1000; +} + +export default EmbeddingProcessStatesConstants; diff --git a/api-specifications/src/modelInfo/[id]/embeddingProcessStates/POST/enums.ts b/api-specifications/src/modelInfo/[id]/embeddingProcessStates/POST/enums.ts new file mode 100644 index 000000000..f34f9f35e --- /dev/null +++ b/api-specifications/src/modelInfo/[id]/embeddingProcessStates/POST/enums.ts @@ -0,0 +1,35 @@ +namespace EmbeddingProcessStatesEnums { + /** + * The status of an embedding process of a model. + */ + export enum Status { + PENDING = "pending", + IN_PROGRESS = "in_progress", + COMPLETED = "completed", + } + + export namespace Response { + export namespace Status400 { + export enum ErrorCodes { + MODEL_NOT_RELEASED = "MODEL_NOT_RELEASED", + } + } + export namespace Status404 { + export enum ErrorCodes { + MODEL_NOT_FOUND_BY_ID = "MODEL_NOT_FOUND_BY_ID", + } + } + export namespace Status409 { + export enum ErrorCodes { + EMBEDDING_PROCESS_ALREADY_RUNNING = "EMBEDDING_PROCESS_ALREADY_RUNNING", + } + } + export namespace Status500 { + export enum ErrorCodes { + FAILED_TO_TRIGGER_EMBEDDING_PROCESS = "FAILED_TO_TRIGGER_EMBEDDING_PROCESS", + } + } + } +} + +export default EmbeddingProcessStatesEnums; diff --git a/api-specifications/src/modelInfo/[id]/embeddingProcessStates/POST/index.ts b/api-specifications/src/modelInfo/[id]/embeddingProcessStates/POST/index.ts new file mode 100644 index 000000000..b741222ec --- /dev/null +++ b/api-specifications/src/modelInfo/[id]/embeddingProcessStates/POST/index.ts @@ -0,0 +1,30 @@ +import POSTEmbeddingProcessStatesEnums from "./enums"; +import POSTEmbeddingProcessStatesTypes from "./types"; +import EmbeddingProcessStatesConstants from "./constants"; + +import SchemaPOSTRequest from "./schema.request"; +import SchemaPOSTResponse from "./schema.response"; + +// ─── POST /models/{modelId}/embedding-process-states ─── +namespace POSTEmbeddingProcessStatesOperation { + export namespace Schemas { + export namespace Request { + export const Payload = SchemaPOSTRequest; + } + export namespace Response { + export const Payload = SchemaPOSTResponse; + } + } + export namespace Types { + export namespace Request { + export type Payload = POSTEmbeddingProcessStatesTypes.POST.Request.Payload; + } + export namespace Response { + export type Payload = POSTEmbeddingProcessStatesTypes.POST.Response.Payload; + } + } + export import Enums = POSTEmbeddingProcessStatesEnums; + export import Constants = EmbeddingProcessStatesConstants; +} + +export default POSTEmbeddingProcessStatesOperation; diff --git a/api-specifications/src/modelInfo/[id]/embeddingProcessStates/POST/schema.request.test.ts b/api-specifications/src/modelInfo/[id]/embeddingProcessStates/POST/schema.request.test.ts new file mode 100644 index 000000000..e37dd5850 --- /dev/null +++ b/api-specifications/src/modelInfo/[id]/embeddingProcessStates/POST/schema.request.test.ts @@ -0,0 +1,51 @@ +import EmbeddingProcessStatesAPISpecs from "../index"; +import EmbeddingsAPISpecs from "embeddings"; +import { + testSchemaWithAdditionalProperties, + testSchemaWithValidObject, + testValidSchema, + testEnumField, +} from "_test_utilities/stdSchemaTests"; + +describe("Test the EmbeddingProcessStates POST Request Schema", () => { + // GIVEN the EmbeddingProcessStatesAPISpecs.POST.Schemas.Request.Payload schema + // WHEN the schema is validated + // THEN expect the schema to be valid + testValidSchema( + "EmbeddingProcessStatesAPISpecs.POST.Schemas.Request.Payload", + EmbeddingProcessStatesAPISpecs.POST.Schemas.Request.Payload + ); +}); + +describe("Validate JSON against the EmbeddingProcessStates POST Request Schema", () => { + // GIVEN a valid EmbeddingProcessStates POST Request object + const givenValidEmbeddingProcessStatesRequestPayload: EmbeddingProcessStatesAPISpecs.POST.Types.Request.Payload = { + embeddingServiceId: EmbeddingsAPISpecs.Constants.EmbeddingServiceIds[0], + }; + + // WHEN the object is validated + // THEN expect the object to validate successfully + testSchemaWithValidObject( + "EmbeddingProcessStatesAPISpecs.POST.Schemas.Request.Payload", + EmbeddingProcessStatesAPISpecs.POST.Schemas.Request.Payload, + givenValidEmbeddingProcessStatesRequestPayload + ); + + // AND WHEN the object has additional properties + // THEN expect the object to not validate + testSchemaWithAdditionalProperties( + "EmbeddingProcessStatesAPISpecs.POST.Schemas.Request.Payload", + EmbeddingProcessStatesAPISpecs.POST.Schemas.Request.Payload, + givenValidEmbeddingProcessStatesRequestPayload + ); + + describe("Validate the EmbeddingProcessStates POST Request fields", () => { + describe("Test validation of 'embeddingServiceId'", () => { + testEnumField( + "embeddingServiceId", + EmbeddingProcessStatesAPISpecs.POST.Schemas.Request.Payload, + EmbeddingsAPISpecs.Constants.EmbeddingServiceIds + ); + }); + }); +}); diff --git a/api-specifications/src/modelInfo/[id]/embeddingProcessStates/POST/schema.request.ts b/api-specifications/src/modelInfo/[id]/embeddingProcessStates/POST/schema.request.ts new file mode 100644 index 000000000..822ba930b --- /dev/null +++ b/api-specifications/src/modelInfo/[id]/embeddingProcessStates/POST/schema.request.ts @@ -0,0 +1,18 @@ +import { SchemaObject } from "ajv"; +import EmbeddingsConstants from "../../../../embeddings/constants"; + +export const SchemaPOSTRequest: SchemaObject = { + $id: "/components/schemas/EmbeddingsRequestSchemaPOST", + type: "object", + additionalProperties: false, + properties: { + embeddingServiceId: { + description: "The identifier of the embedding model to use for generating the embeddings of the model.", + type: "string", + enum: EmbeddingsConstants.EmbeddingServiceIds, + }, + }, + required: ["embeddingServiceId"], +}; + +export default SchemaPOSTRequest; diff --git a/api-specifications/src/modelInfo/[id]/embeddingProcessStates/POST/schema.response.test.ts b/api-specifications/src/modelInfo/[id]/embeddingProcessStates/POST/schema.response.test.ts new file mode 100644 index 000000000..138c99f99 --- /dev/null +++ b/api-specifications/src/modelInfo/[id]/embeddingProcessStates/POST/schema.response.test.ts @@ -0,0 +1,134 @@ +import EmbeddingProcessStatesAPISpecs from "../index"; +import EmbeddingsAPISpecs from "embeddings"; +import { getMockId } from "_test_utilities/mockMongoId"; +import { + testObjectIdField, + testSchemaWithAdditionalProperties, + testSchemaWithValidObject, + testTimestampField, + testValidSchema, + testEnumField, +} from "_test_utilities/stdSchemaTests"; +import { assertCaseForProperty, CaseType, constructSchemaError } from "_test_utilities/assertCaseForProperty"; +import { WHITESPACE } from "_test_utilities/specialCharacters"; + +describe("Test the EmbeddingProcessStates POST Response Schema", () => { + // GIVEN the EmbeddingProcessStatesAPISpecs.POST.Schemas.Response.Payload schema + // WHEN the schema is validated + // THEN expect the schema to be valid + testValidSchema( + "EmbeddingProcessStatesAPISpecs.POST.Schemas.Response.Payload", + EmbeddingProcessStatesAPISpecs.POST.Schemas.Response.Payload + ); +}); + +describe("Validate JSON against the EmbeddingProcessStates POST Response Schema", () => { + // GIVEN a valid EmbeddingProcessStates POST Response object + const givenValidEmbeddingProcessStatesResponsePayload: EmbeddingProcessStatesAPISpecs.POST.Types.Response.Payload = { + id: getMockId(1), + modelId: getMockId(2), + status: EmbeddingProcessStatesAPISpecs.Enums.Status.PENDING, + embeddingServiceId: EmbeddingsAPISpecs.Constants.EmbeddingServiceIds[0], + totalDocuments: 10, + errorCounts: 0, + warningCounts: 0, + completedDocuments: 0, + createdAt: new Date().toISOString(), + updatedAt: new Date().toISOString(), + }; + + // WHEN the object is validated + // THEN expect the object to validate successfully + testSchemaWithValidObject( + "EmbeddingProcessStatesAPISpecs.POST.Schemas.Response.Payload", + EmbeddingProcessStatesAPISpecs.POST.Schemas.Response.Payload, + givenValidEmbeddingProcessStatesResponsePayload + ); + + // AND WHEN the object has additional properties + // THEN expect the object to not validate + testSchemaWithAdditionalProperties( + "EmbeddingProcessStatesAPISpecs.POST.Schemas.Response.Payload", + EmbeddingProcessStatesAPISpecs.POST.Schemas.Response.Payload, + givenValidEmbeddingProcessStatesResponsePayload + ); + + describe("Validate the EmbeddingProcessStates POST Response fields", () => { + describe("Test validation of 'id'", () => { + testObjectIdField("id", EmbeddingProcessStatesAPISpecs.POST.Schemas.Response.Payload); + }); + + describe("Test validation of 'modelId'", () => { + testObjectIdField("modelId", EmbeddingProcessStatesAPISpecs.POST.Schemas.Response.Payload); + }); + + describe("Test validation of 'status'", () => { + testEnumField( + "status", + EmbeddingProcessStatesAPISpecs.POST.Schemas.Response.Payload, + Object.values(EmbeddingProcessStatesAPISpecs.Enums.Status) + ); + }); + + describe("Test validation of 'embeddingServiceId'", () => { + testEnumField( + "embeddingServiceId", + EmbeddingProcessStatesAPISpecs.POST.Schemas.Response.Payload, + EmbeddingsAPISpecs.Constants.EmbeddingServiceIds + ); + }); + + describe.each([["totalDocuments"], ["errorCounts"], ["warningCounts"], ["completedDocuments"]])( + "Test validation of '%s'", + (fieldName) => { + test.each([ + [ + CaseType.Failure, + "undefined", + undefined, + constructSchemaError("", "required", `must have required property '${fieldName}'`), + ], + [CaseType.Failure, "null", null, constructSchemaError(`/${fieldName}`, "type", "must be integer")], + [CaseType.Failure, "a string", "foo", constructSchemaError(`/${fieldName}`, "type", "must be integer")], + [ + CaseType.Failure, + "only whitespace characters", + WHITESPACE, + constructSchemaError(`/${fieldName}`, "type", "must be integer"), + ], + [CaseType.Failure, "a float", 1.5, constructSchemaError(`/${fieldName}`, "type", "must be integer")], + [CaseType.Failure, "a negative number", -1, constructSchemaError(`/${fieldName}`, "minimum", "must be >= 0")], + [CaseType.Success, "zero", 0, undefined], + [CaseType.Success, "a positive integer", 42, undefined], + ])(`(%s) Validate '${fieldName}' when it is %s`, (caseType, _description, givenValue, failureMessages) => { + // GIVEN an object with the given value + const givenObject = { + [fieldName]: givenValue, + }; + // THEN expect the object to validate accordingly + assertCaseForProperty( + fieldName, + givenObject, + EmbeddingProcessStatesAPISpecs.POST.Schemas.Response.Payload, + caseType, + failureMessages + ); + }); + } + ); + + describe("Test validation of 'createdAt'", () => { + testTimestampField( + "createdAt", + EmbeddingProcessStatesAPISpecs.POST.Schemas.Response.Payload + ); + }); + + describe("Test validation of 'updatedAt'", () => { + testTimestampField( + "updatedAt", + EmbeddingProcessStatesAPISpecs.POST.Schemas.Response.Payload + ); + }); + }); +}); diff --git a/api-specifications/src/modelInfo/[id]/embeddingProcessStates/POST/schema.response.ts b/api-specifications/src/modelInfo/[id]/embeddingProcessStates/POST/schema.response.ts new file mode 100644 index 000000000..d6ece9fcd --- /dev/null +++ b/api-specifications/src/modelInfo/[id]/embeddingProcessStates/POST/schema.response.ts @@ -0,0 +1,70 @@ +import { SchemaObject } from "ajv"; +import { RegExp_Str_ID } from "../../../../regex"; +import EmbeddingProcessStatesEnums from "./enums"; +import EmbeddingsConstants from "../../../../embeddings/constants"; + +const SchemaPOSTResponse: SchemaObject = { + description: + "The state of the embedding process of the model. Since the embedding process is asynchronous, use the status to check if the embedding process has completed and the counts to check its progress.", + $id: "/components/schemas/EmbeddingsResponseSchemaPOST", + type: "object", + additionalProperties: false, + properties: { + id: { + description: "The identifier of the specific embedding process state.", + type: "string", + pattern: RegExp_Str_ID, + }, + modelId: { + description: "The identifier of the model whose entities are being embedded.", + type: "string", + pattern: RegExp_Str_ID, + }, + status: { + description: "The status of the embedding process of the model.", + type: "string", + enum: Object.values(EmbeddingProcessStatesEnums.Status), + }, + embeddingServiceId: { + description: "The identifier of the embedding model that is used to generate the embeddings.", + type: "string", + enum: EmbeddingsConstants.EmbeddingServiceIds, + }, + totalDocuments: { + description: "The total number of documents (entities) that will be embedded.", + type: "integer", + minimum: 0, + }, + errorCounts: { + description: "The number of documents that failed to be embedded.", + type: "integer", + minimum: 0, + }, + warningCounts: { + description: "The number of documents that were embedded with warnings.", + type: "integer", + minimum: 0, + }, + completedDocuments: { + description: "The number of documents that have been successfully embedded so far.", + type: "integer", + minimum: 0, + }, + createdAt: { type: "string", format: "date-time" }, + updatedAt: { type: "string", format: "date-time" }, + }, + required: [ + "id", + "modelId", + "status", + "embeddingServiceId", + "totalDocuments", + "errorCounts", + "warningCounts", + "completedDocuments", + "createdAt", + "updatedAt", + ], +}; + +export default SchemaPOSTResponse; diff --git a/api-specifications/src/modelInfo/[id]/embeddingProcessStates/POST/types.ts b/api-specifications/src/modelInfo/[id]/embeddingProcessStates/POST/types.ts new file mode 100644 index 000000000..d0b1a3a83 --- /dev/null +++ b/api-specifications/src/modelInfo/[id]/embeddingProcessStates/POST/types.ts @@ -0,0 +1,28 @@ +import EmbeddingProcessStatesEnums from "../enums"; + +namespace POSTEmbeddingProcessStatesTypes { + export namespace POST { + export namespace Request { + export interface Payload { + embeddingServiceId: string; + } + } + + export namespace Response { + export interface Payload { + id: string; + modelId: string; + status: EmbeddingProcessStatesEnums.Status; + embeddingServiceId: string; + totalDocuments: number; + errorCounts: number; + warningCounts: number; + completedDocuments: number; + createdAt: string; + updatedAt: string; + } + } + } +} + +export default POSTEmbeddingProcessStatesTypes; diff --git a/api-specifications/src/modelInfo/[id]/embeddingProcessStates/__snapshots__/index.test.ts.snap b/api-specifications/src/modelInfo/[id]/embeddingProcessStates/__snapshots__/index.test.ts.snap new file mode 100644 index 000000000..953ef6556 --- /dev/null +++ b/api-specifications/src/modelInfo/[id]/embeddingProcessStates/__snapshots__/index.test.ts.snap @@ -0,0 +1,144 @@ +// Jest Snapshot v1, https://goo.gl/fbAQLP + +exports[`Test the EmbeddingProcessStates module The EmbeddingProcessStates module matches the snapshot 1`] = ` +{ + "Enums": { + "Status": { + "COMPLETED": "completed", + "IN_PROGRESS": "in_progress", + "PENDING": "pending", + }, + }, + "POST": { + "Constants": { + "MAX_PAYLOAD_LENGTH": 1036, + }, + "Enums": { + "Response": { + "Status400": { + "ErrorCodes": { + "MODEL_NOT_RELEASED": "MODEL_NOT_RELEASED", + }, + }, + "Status404": { + "ErrorCodes": { + "MODEL_NOT_FOUND_BY_ID": "MODEL_NOT_FOUND_BY_ID", + }, + }, + "Status409": { + "ErrorCodes": { + "EMBEDDING_PROCESS_ALREADY_RUNNING": "EMBEDDING_PROCESS_ALREADY_RUNNING", + }, + }, + "Status500": { + "ErrorCodes": { + "FAILED_TO_TRIGGER_EMBEDDING_PROCESS": "FAILED_TO_TRIGGER_EMBEDDING_PROCESS", + }, + }, + }, + "Status": { + "COMPLETED": "completed", + "IN_PROGRESS": "in_progress", + "PENDING": "pending", + }, + }, + "Schemas": { + "Request": { + "Payload": { + "$id": "/components/schemas/EmbeddingsRequestSchemaPOST", + "additionalProperties": false, + "properties": { + "embeddingServiceId": { + "description": "The identifier of the embedding model to use for generating the embeddings of the model.", + "enum": [ + "77bb8ff3-a6b0-460b-bcaa-00631a907852", + ], + "type": "string", + }, + }, + "required": [ + "embeddingServiceId", + ], + "type": "object", + }, + }, + "Response": { + "Payload": { + "$id": "/components/schemas/EmbeddingsResponseSchemaPOST", + "additionalProperties": false, + "description": "The state of the embedding process of the model. Since the embedding process is asynchronous, use the status to check if the embedding process has completed and the counts to check its progress.", + "properties": { + "completedDocuments": { + "description": "The number of documents that have been successfully embedded so far.", + "minimum": 0, + "type": "integer", + }, + "createdAt": { + "format": "date-time", + "type": "string", + }, + "embeddingServiceId": { + "description": "The identifier of the embedding model that is used to generate the embeddings.", + "enum": [ + "77bb8ff3-a6b0-460b-bcaa-00631a907852", + ], + "type": "string", + }, + "errorCounts": { + "description": "The number of documents that failed to be embedded.", + "minimum": 0, + "type": "integer", + }, + "id": { + "description": "The identifier of the specific embedding process state.", + "pattern": "^[0-9a-f]{24}$", + "type": "string", + }, + "modelId": { + "description": "The identifier of the model whose entities are being embedded.", + "pattern": "^[0-9a-f]{24}$", + "type": "string", + }, + "status": { + "description": "The status of the embedding process of the model.", + "enum": [ + "pending", + "in_progress", + "completed", + ], + "type": "string", + }, + "totalDocuments": { + "description": "The total number of documents (entities) that will be embedded.", + "minimum": 0, + "type": "integer", + }, + "updatedAt": { + "format": "date-time", + "type": "string", + }, + "warningCounts": { + "description": "The number of documents that were embedded with warnings.", + "minimum": 0, + "type": "integer", + }, + }, + "required": [ + "id", + "modelId", + "status", + "embeddingServiceId", + "totalDocuments", + "errorCounts", + "warningCounts", + "completedDocuments", + "createdAt", + "updatedAt", + ], + "type": "object", + }, + }, + }, + }, +} +`; diff --git a/api-specifications/src/modelInfo/[id]/embeddingProcessStates/enums.ts b/api-specifications/src/modelInfo/[id]/embeddingProcessStates/enums.ts new file mode 100644 index 000000000..eaa958f90 --- /dev/null +++ b/api-specifications/src/modelInfo/[id]/embeddingProcessStates/enums.ts @@ -0,0 +1,12 @@ +namespace EmbeddingProcessStatesEnums { + /** + * The status of an embedding process of a model. + */ + export enum Status { + PENDING = "pending", + IN_PROGRESS = "in_progress", + COMPLETED = "completed", + } +} + +export default EmbeddingProcessStatesEnums; diff --git a/api-specifications/src/modelInfo/[id]/embeddingProcessStates/index.test.ts b/api-specifications/src/modelInfo/[id]/embeddingProcessStates/index.test.ts new file mode 100644 index 000000000..89206b82b --- /dev/null +++ b/api-specifications/src/modelInfo/[id]/embeddingProcessStates/index.test.ts @@ -0,0 +1,24 @@ +describe("Test the EmbeddingProcessStates module", () => { + test("The EmbeddingProcessStates module can be required via the index", async () => { + // GIVEN the module + // WHEN the module is required via the index + expect(async () => await import("./")).not.toThrow(); // We check that it doesn't throw an error instead of simply letting it fail on import because we want an easier error message + const embeddingProcessStatesModule = await import("./"); + + // THEN Check if the POST operation Schemas are defined in it + expect(embeddingProcessStatesModule.default.POST.Schemas.Request.Payload).toBeDefined(); + expect(embeddingProcessStatesModule.default.POST.Schemas.Response.Payload).toBeDefined(); + // AND Check if the various exports are defined in it + expect(embeddingProcessStatesModule.default.Enums.Status).toBeDefined(); + }); + + test("The EmbeddingProcessStates module matches the snapshot", async () => { + // GIVEN the module + // WHEN the module is required via the index + const embeddingProcessStatesModule = await import("./"); + // THEN expect the module to match the snapshot + expect(embeddingProcessStatesModule.default).toMatchSnapshot(); + }); +}); + +export {}; diff --git a/api-specifications/src/modelInfo/[id]/embeddingProcessStates/index.ts b/api-specifications/src/modelInfo/[id]/embeddingProcessStates/index.ts new file mode 100644 index 000000000..32b1e5746 --- /dev/null +++ b/api-specifications/src/modelInfo/[id]/embeddingProcessStates/index.ts @@ -0,0 +1,14 @@ +// we need to disable the eslint rule here because this is a top level export +/* eslint-disable @typescript-eslint/no-unused-vars */ + +import EmbeddingProcessStatesEnums from "./enums"; +import POSTEmbeddingProcessStatesOperation from "./POST"; + +// Concept aggregator for the embedding process states of a model +namespace EmbeddingProcessStatesAPISpecs { + export import Enums = EmbeddingProcessStatesEnums; + + export import POST = POSTEmbeddingProcessStatesOperation; +} + +export default EmbeddingProcessStatesAPISpecs; diff --git a/api-specifications/src/modelInfo/[id]/index.ts b/api-specifications/src/modelInfo/[id]/index.ts new file mode 100644 index 000000000..76eb49609 --- /dev/null +++ b/api-specifications/src/modelInfo/[id]/index.ts @@ -0,0 +1,12 @@ +// we need to disable the eslint rule here because this is a top level export +/* eslint-disable @typescript-eslint/no-unused-vars */ + +import EmbeddingProcessStatesAPISpecs from "./embeddingProcessStates"; + +// Concept aggregator for [id] (ModelInfo Instance) +namespace ModelInfoInstanceAPISpecs { + // Child API Paths + export import EmbeddingProcessStates = EmbeddingProcessStatesAPISpecs; +} + +export default ModelInfoInstanceAPISpecs; diff --git a/api-specifications/src/modelInfo/__snapshots__/index.test.ts.snap b/api-specifications/src/modelInfo/__snapshots__/index.test.ts.snap index f45cc967e..f07f68c0a 100644 --- a/api-specifications/src/modelInfo/__snapshots__/index.test.ts.snap +++ b/api-specifications/src/modelInfo/__snapshots__/index.test.ts.snap @@ -28,6 +28,148 @@ exports[`Test the modelInfo module The export module matches the snapshot 1`] = }, }, }, + "ModelInfo": { + "EmbeddingProcessStates": { + "Enums": { + "Status": { + "COMPLETED": "completed", + "IN_PROGRESS": "in_progress", + "PENDING": "pending", + }, + }, + "POST": { + "Constants": { + "MAX_PAYLOAD_LENGTH": 1036, + }, + "Enums": { + "Response": { + "Status400": { + "ErrorCodes": { + "MODEL_NOT_RELEASED": "MODEL_NOT_RELEASED", + }, + }, + "Status404": { + "ErrorCodes": { + "MODEL_NOT_FOUND_BY_ID": "MODEL_NOT_FOUND_BY_ID", + }, + }, + "Status409": { + "ErrorCodes": { + "EMBEDDING_PROCESS_ALREADY_RUNNING": "EMBEDDING_PROCESS_ALREADY_RUNNING", + }, + }, + "Status500": { + "ErrorCodes": { + "FAILED_TO_TRIGGER_EMBEDDING_PROCESS": "FAILED_TO_TRIGGER_EMBEDDING_PROCESS", + }, + }, + }, + "Status": { + "COMPLETED": "completed", + "IN_PROGRESS": "in_progress", + "PENDING": "pending", + }, + }, + "Schemas": { + "Request": { + "Payload": { + "$id": "/components/schemas/EmbeddingsRequestSchemaPOST", + "additionalProperties": false, + "properties": { + "embeddingServiceId": { + "description": "The identifier of the embedding model to use for generating the embeddings of the model.", + "enum": [ + "77bb8ff3-a6b0-460b-bcaa-00631a907852", + ], + "type": "string", + }, + }, + "required": [ + "embeddingServiceId", + ], + "type": "object", + }, + }, + "Response": { + "Payload": { + "$id": "/components/schemas/EmbeddingsResponseSchemaPOST", + "additionalProperties": false, + "description": "The state of the embedding process of the model. Since the embedding process is asynchronous, use the status to check if the embedding process has completed and the counts to check its progress.", + "properties": { + "completedDocuments": { + "description": "The number of documents that have been successfully embedded so far.", + "minimum": 0, + "type": "integer", + }, + "createdAt": { + "format": "date-time", + "type": "string", + }, + "embeddingServiceId": { + "description": "The identifier of the embedding model that is used to generate the embeddings.", + "enum": [ + "77bb8ff3-a6b0-460b-bcaa-00631a907852", + ], + "type": "string", + }, + "errorCounts": { + "description": "The number of documents that failed to be embedded.", + "minimum": 0, + "type": "integer", + }, + "id": { + "description": "The identifier of the specific embedding process state.", + "pattern": "^[0-9a-f]{24}$", + "type": "string", + }, + "modelId": { + "description": "The identifier of the model whose entities are being embedded.", + "pattern": "^[0-9a-f]{24}$", + "type": "string", + }, + "status": { + "description": "The status of the embedding process of the model.", + "enum": [ + "pending", + "in_progress", + "completed", + ], + "type": "string", + }, + "totalDocuments": { + "description": "The total number of documents (entities) that will be embedded.", + "minimum": 0, + "type": "integer", + }, + "updatedAt": { + "format": "date-time", + "type": "string", + }, + "warningCounts": { + "description": "The number of documents that were embedded with warnings.", + "minimum": 0, + "type": "integer", + }, + }, + "required": [ + "id", + "modelId", + "status", + "embeddingServiceId", + "totalDocuments", + "errorCounts", + "warningCounts", + "completedDocuments", + "createdAt", + "updatedAt", + ], + "type": "object", + }, + }, + }, + }, + }, + }, "Schemas": { "GET": { "Response": { diff --git a/api-specifications/src/modelInfo/index.ts b/api-specifications/src/modelInfo/index.ts index b89f407fc..fe5971ab6 100644 --- a/api-specifications/src/modelInfo/index.ts +++ b/api-specifications/src/modelInfo/index.ts @@ -9,6 +9,7 @@ import SchemaPOSTRequest from "./schema.POST.request"; import SchemaPOSTResponse from "./schema.POST.response"; import SchemaModelInfoReference from "./schema.reference"; import ModelInfoEnums from "./enums"; +import ModelInfoInstanceAPISpecs from "./[id]"; /** * This file should be imported in the following way @@ -40,6 +41,9 @@ namespace ModelInfoAPISpecs { export import Constants = ModelInfoConstants; export import Types = ModelInfoTypes; export import Schemas = ModelInfoSchemas; + + // Instance-level operations (via [id] concept) + export import ModelInfo = ModelInfoInstanceAPISpecs; } export default ModelInfoAPISpecs; diff --git a/backend/openapi/generateOpenApiDoc.ts b/backend/openapi/generateOpenApiDoc.ts index 28b3e3ead..494ea4f81 100644 --- a/backend/openapi/generateOpenApiDoc.ts +++ b/backend/openapi/generateOpenApiDoc.ts @@ -41,6 +41,8 @@ delete ModelInfo.Schemas.POST.Response.Payload.$id; delete ModelInfo.Schemas.POST.Request.Payload.$id; delete ModelInfo.Schemas.GET.Response.Payload.$id; delete ModelInfo.Schemas.Reference.$id; +delete ModelInfo.ModelInfo.EmbeddingProcessStates.POST.Schemas.Request.Payload.$id; +delete ModelInfo.ModelInfo.EmbeddingProcessStates.POST.Schemas.Response.Payload.$id; delete Occupation.Schemas.Reference.$id; delete OccupationGroup.Schemas.Reference.$id; delete Skill.Schemas.Reference.$id; @@ -408,11 +410,32 @@ function getOpenAPISpecification( 404, Object.values(Skill.GET.Errors.Status404.RelatedSkills.ErrorCodes) ), + // EmbeddingProcess-specific error schemas + POSTEmbeddingProcess400ErrorSchema: APIError.Schemas.getPayload( + "POST", + "EmbeddingProcess", + 400, + Object.values(ModelInfo.ModelInfo.EmbeddingProcessStates.POST.Enums.Response.Status400.ErrorCodes) + ), + POSTEmbeddingProcess404ErrorSchema: APIError.Schemas.getPayload( + "POST", + "EmbeddingProcess", + 404, + Object.values(ModelInfo.ModelInfo.EmbeddingProcessStates.POST.Enums.Response.Status404.ErrorCodes) + ), + POSTEmbeddingProcess409ErrorSchema: APIError.Schemas.getPayload( + "POST", + "EmbeddingProcess", + 409, + Object.values(ModelInfo.ModelInfo.EmbeddingProcessStates.POST.Enums.Response.Status409.ErrorCodes) + ), PresignedSchema: Presigned.Schemas.GET.Response.Payload, ModelInfoResponseSchemaPOST: ModelInfo.Schemas.POST.Response.Payload, ModelInfoRequestSchemaPOST: ModelInfo.Schemas.POST.Request.Payload, ModelInfoResponseSchemaGET: ModelInfo.Schemas.GET.Response.Payload, ModelInfoReferenceSchema: ModelInfo.Schemas.Reference, + EmbeddingsRequestSchemaPOST: ModelInfo.ModelInfo.EmbeddingProcessStates.POST.Schemas.Request.Payload, + EmbeddingsResponseSchemaPOST: ModelInfo.ModelInfo.EmbeddingProcessStates.POST.Schemas.Response.Payload, OccupationReferenceSchema: Occupation.Schemas.Reference, OccupationGroupReferenceSchema: OccupationGroup.Schemas.Reference, SkillReferenceSchema: Skill.Schemas.Reference, diff --git a/backend/package.json b/backend/package.json index 87d000e9e..236ea2b1f 100644 --- a/backend/package.json +++ b/backend/package.json @@ -8,8 +8,8 @@ "clean": "rimraf -I build coverage deploy", "compile": "tsc --build tsconfig.prod.json", "prebuild": "npm run clean && npm run compile", - "build": "esbuild ./src/index.ts ./src/import/async/index.ts ./src/export/async/index.ts ./src/auth/authenticator/index.ts --bundle --tsconfig=tsconfig.prod.json --sourcemap --minify --platform=node --target=node16 --outdir=build --main-fields=module,main", - "postbuild": "cp ./package.json ./build/ && (cd ./build; npm pkg delete devDependencies; npm pkg delete dependencies; npm pkg delete scripts) && cp ./build/package.json ./build/import/async && mkdir build/rest && mv build/index.js build/index.js.map build/package.json ./build/rest/", + "build": "esbuild ./src/index.ts ./src/import/async/index.ts ./src/export/async/index.ts ./src/embeddings/handler/index.ts ./src/auth/authenticator/index.ts --bundle --tsconfig=tsconfig.prod.json --sourcemap --minify --platform=node --target=node16 --outdir=build --main-fields=module,main", + "postbuild": "cp ./package.json ./build/ && (cd ./build; npm pkg delete devDependencies; npm pkg delete dependencies; npm pkg delete scripts) && cp ./build/package.json ./build/import/async && cp ./build/package.json ./build/embeddings/handler && mkdir build/rest && mv build/index.js build/index.js.map build/package.json ./build/rest/", "test": "jest --coverage", "test:integration": "jest --config jest.integration.config.js --runInBand", "test:smoke": "jest --config jest.smoke.config.js --runInBand test/smoke/*.test.ts", @@ -59,6 +59,7 @@ "@aws-sdk/client-cognito-identity-provider": "^3.567.0", "@aws-sdk/client-lambda": "~3.468.0", "@aws-sdk/client-s3": "~3.468.0", + "@aws-sdk/client-sqs": "~3.468.0", "@aws-sdk/lib-storage": "~3.468.0", "@aws-sdk/s3-presigned-post": "~3.468.0", "@aws-sdk/s3-request-presigner": "~3.468.0", diff --git a/backend/src/_test_utilities/getTestConfiguration.ts b/backend/src/_test_utilities/getTestConfiguration.ts index 642b9e6e3..b0387bc07 100644 --- a/backend/src/_test_utilities/getTestConfiguration.ts +++ b/backend/src/_test_utilities/getTestConfiguration.ts @@ -13,5 +13,9 @@ export function getTestConfiguration(dbname: string): IConfiguration { asyncImportLambdaFunctionArn: "arn:aws:lambda:foo:bar:baz:import", asyncExportLambdaFunctionArn: "arn:aws:lambda:foo:bar:baz:export", asyncLambdaFunctionRegion: "foo", + geminiApiKey: "test-gemini-api-key", + geminiEmbeddingModel: "text-embedding-004", + embeddingsQueueUrl: "https://sqs.foo.amazonaws.com/123456789012/test-embeddings-queue", + embeddingsQueueRegion: "foo", }; } diff --git a/backend/src/common/lambda.types.ts b/backend/src/common/lambda.types.ts index d11cb646c..019516f5f 100644 --- a/backend/src/common/lambda.types.ts +++ b/backend/src/common/lambda.types.ts @@ -6,4 +6,5 @@ export enum Lambdas { API = "API", IMPORT = "IMPORT", EXPORT = "EXPORT", + EMBEDDING = "EMBEDDING", } diff --git a/backend/src/embeddings/embeddingProcess/embeddingProcess.service.test.ts b/backend/src/embeddings/embeddingProcess/embeddingProcess.service.test.ts new file mode 100644 index 000000000..f75fd9aad --- /dev/null +++ b/backend/src/embeddings/embeddingProcess/embeddingProcess.service.test.ts @@ -0,0 +1,223 @@ +// silence chatty console +import "_test_utilities/consoleMock"; + +import { Readable } from "node:stream"; +import ModelInfoApiSpecs from "api-specifications/modelInfo"; +import { EmbeddingProcessService } from "./embeddingProcess.service"; +import { EmbeddingProcessAlreadyRunningError, ModelNotFoundError, ModelNotReleasedError } from "./errors"; +import { EmbeddableEntityType, EmbeddableField } from "embeddings/service/types"; +import { IEmbeddingProcessState } from "embeddings/embeddingProcessState/embeddingProcessState.types"; +import { getMockStringId } from "_test_utilities/mockMongoId"; + +function getMockModel(released: boolean) { + return { id: getMockStringId(1), released }; +} + +function getMockEmbeddingProcessState(overrides: Partial = {}): IEmbeddingProcessState { + return { + id: getMockStringId(10), + modelId: getMockStringId(1), + status: ModelInfoApiSpecs.ModelInfo.EmbeddingProcessStates.Enums.Status.PENDING, + embeddingServiceId: "gemini$$models/gemini-embedding-2", + totalDocuments: 0, + errorCounts: 0, + warningCounts: 0, + completedDocuments: 0, + createdAt: new Date(), + updatedAt: new Date(), + ...overrides, + }; +} + +function setupService(mocks: { + model?: { id: string; released: boolean } | null; + pendingProcess?: IEmbeddingProcessState | null; + skills?: object[]; + skillGroups?: object[]; + occupations?: object[]; + occupationGroups?: object[]; + createdProcessState?: IEmbeddingProcessState; + updatedProcessState?: IEmbeddingProcessState; +}) { + const modelRepository = { + getModelById: jest.fn().mockResolvedValue(mocks.model), + }; + const embeddingProcessStateRepository = { + findPendingByModelId: jest.fn().mockResolvedValue(mocks.pendingProcess ?? null), + create: jest.fn().mockResolvedValue(mocks.createdProcessState ?? getMockEmbeddingProcessState()), + update: jest.fn().mockResolvedValue(mocks.updatedProcessState ?? getMockEmbeddingProcessState()), + }; + const skillRepository = { findAll: jest.fn().mockReturnValue(Readable.from(mocks.skills ?? [])) }; + const skillGroupRepository = { findAll: jest.fn().mockReturnValue(Readable.from(mocks.skillGroups ?? [])) }; + const occupationRepository = { findAll: jest.fn().mockReturnValue(Readable.from(mocks.occupations ?? [])) }; + const occupationGroupRepository = { + findAll: jest.fn().mockReturnValue(Readable.from(mocks.occupationGroups ?? [])), + }; + const embeddingClient = { pushTaskToQueue: jest.fn().mockResolvedValue(undefined) }; + + const service = new EmbeddingProcessService( + // @ts-ignore + modelRepository, + // @ts-ignore + embeddingProcessStateRepository, + // @ts-ignore + skillRepository, + // @ts-ignore + skillGroupRepository, + // @ts-ignore + occupationRepository, + // @ts-ignore + occupationGroupRepository, + // @ts-ignore + embeddingClient + ); + + return { + service, + modelRepository, + embeddingProcessStateRepository, + skillRepository, + skillGroupRepository, + occupationRepository, + occupationGroupRepository, + embeddingClient, + }; +} + +describe("Test the EmbeddingProcessService", () => { + const givenModelId = getMockStringId(1); + const givenEmbeddingServiceId = "gemini$$models/gemini-embedding-2"; + + beforeEach(() => { + jest.clearAllMocks(); + }); + + test("should throw a ModelNotFoundError when the model does not exist", async () => { + // GIVEN a model that does not exist + const { service, embeddingProcessStateRepository } = setupService({ model: null }); + + // WHEN triggering the embedding process + const actualPromise = service.triggerEmbeddingProcess(givenModelId, givenEmbeddingServiceId); + + // THEN expect it to reject with a ModelNotFoundError + await expect(actualPromise).rejects.toThrow(ModelNotFoundError); + // AND expect no embedding process state to be created + expect(embeddingProcessStateRepository.create).not.toHaveBeenCalled(); + }); + + test("should throw a ModelNotReleasedError when the model is not released", async () => { + // GIVEN a model that is not released + const { service, embeddingProcessStateRepository } = setupService({ model: getMockModel(false) }); + + // WHEN triggering the embedding process + const actualPromise = service.triggerEmbeddingProcess(givenModelId, givenEmbeddingServiceId); + + // THEN expect it to reject with a ModelNotReleasedError + await expect(actualPromise).rejects.toThrow(ModelNotReleasedError); + // AND expect no embedding process state to be created + expect(embeddingProcessStateRepository.create).not.toHaveBeenCalled(); + }); + + test("should throw an EmbeddingProcessAlreadyRunningError when there is an unfinished process for the model", async () => { + // GIVEN a released model AND an unfinished embedding process for the model + const { service, embeddingProcessStateRepository } = setupService({ + model: getMockModel(true), + pendingProcess: getMockEmbeddingProcessState(), + }); + + // WHEN triggering the embedding process + const actualPromise = service.triggerEmbeddingProcess(givenModelId, givenEmbeddingServiceId); + + // THEN expect it to reject with an EmbeddingProcessAlreadyRunningError + await expect(actualPromise).rejects.toThrow(EmbeddingProcessAlreadyRunningError); + // AND expect no embedding process state to be created + expect(embeddingProcessStateRepository.create).not.toHaveBeenCalled(); + }); + + test("should create a process state, push all the entities to the queue and update the process state", async () => { + // GIVEN a released model with no unfinished embedding process + // AND some entities of each type + const givenSkills = [{ id: getMockStringId(101) }, { id: getMockStringId(102) }]; + const givenSkillGroups = [{ id: getMockStringId(201) }]; + const givenOccupations = [{ id: getMockStringId(301) }]; + const givenOccupationGroups = [{ id: getMockStringId(401) }]; + const givenTotalDocuments = + givenSkills.length + givenSkillGroups.length + givenOccupations.length + givenOccupationGroups.length; + // AND a created process state + const givenCreatedProcessState = getMockEmbeddingProcessState({ id: getMockStringId(10) }); + // AND an updated process state + const givenUpdatedProcessState = getMockEmbeddingProcessState({ + id: getMockStringId(10), + status: ModelInfoApiSpecs.ModelInfo.EmbeddingProcessStates.Enums.Status.IN_PROGRESS, + totalDocuments: givenTotalDocuments, + }); + const { service, modelRepository, embeddingProcessStateRepository, embeddingClient } = setupService({ + model: getMockModel(true), + pendingProcess: null, + skills: givenSkills, + skillGroups: givenSkillGroups, + occupations: givenOccupations, + occupationGroups: givenOccupationGroups, + createdProcessState: givenCreatedProcessState, + updatedProcessState: givenUpdatedProcessState, + }); + + // WHEN triggering the embedding process + const actualProcessState = await service.triggerEmbeddingProcess(givenModelId, givenEmbeddingServiceId); + + // THEN expect the model to be fetched + expect(modelRepository.getModelById).toHaveBeenCalledWith(givenModelId); + // AND expect a new process state to be created with the given embedding model id + expect(embeddingProcessStateRepository.create).toHaveBeenCalledWith({ + modelId: givenModelId, + status: ModelInfoApiSpecs.ModelInfo.EmbeddingProcessStates.Enums.Status.PENDING, + embeddingServiceId: givenEmbeddingServiceId, + totalDocuments: 0, + errorCounts: 0, + warningCounts: 0, + completedDocuments: 0, + }); + // AND expect one task to be pushed per entity + expect(embeddingClient.pushTaskToQueue).toHaveBeenCalledTimes(givenTotalDocuments); + // AND expect the tasks for the skills to be pushed with the skill fields + expect(embeddingClient.pushTaskToQueue).toHaveBeenCalledWith({ + modelId: givenModelId, + entityId: givenSkills[0].id, + entityType: EmbeddableEntityType.Skill, + fields: [EmbeddableField.preferredLabel, EmbeddableField.description, EmbeddableField.altLabels], + }); + // AND expect the task for the skill group to be pushed with the skill group fields + expect(embeddingClient.pushTaskToQueue).toHaveBeenCalledWith({ + modelId: givenModelId, + entityId: givenSkillGroups[0].id, + entityType: EmbeddableEntityType.SkillGroup, + fields: [ + EmbeddableField.preferredLabel, + EmbeddableField.description, + EmbeddableField.altLabels, + EmbeddableField.scopeNote, + ], + }); + // AND expect the task for the occupation to be pushed + expect(embeddingClient.pushTaskToQueue).toHaveBeenCalledWith({ + modelId: givenModelId, + entityId: givenOccupations[0].id, + entityType: EmbeddableEntityType.Occupation, + fields: [EmbeddableField.preferredLabel, EmbeddableField.description, EmbeddableField.altLabels], + }); + // AND expect the task for the occupation group to be pushed + expect(embeddingClient.pushTaskToQueue).toHaveBeenCalledWith({ + modelId: givenModelId, + entityId: givenOccupationGroups[0].id, + entityType: EmbeddableEntityType.OccupationGroup, + fields: [EmbeddableField.preferredLabel, EmbeddableField.description, EmbeddableField.altLabels], + }); + // AND expect the process state to be updated with the total documents and the in-progress status + expect(embeddingProcessStateRepository.update).toHaveBeenCalledWith(givenCreatedProcessState.id, { + status: ModelInfoApiSpecs.ModelInfo.EmbeddingProcessStates.Enums.Status.IN_PROGRESS, + totalDocuments: givenTotalDocuments, + }); + // AND expect the returned process state to be the updated one + expect(actualProcessState).toEqual(givenUpdatedProcessState); + }); +}); diff --git a/backend/src/embeddings/embeddingProcess/embeddingProcess.service.ts b/backend/src/embeddings/embeddingProcess/embeddingProcess.service.ts new file mode 100644 index 000000000..8681fa602 --- /dev/null +++ b/backend/src/embeddings/embeddingProcess/embeddingProcess.service.ts @@ -0,0 +1,143 @@ +import { Readable } from "node:stream"; +import ModelInfoApiSpecs from "api-specifications/modelInfo"; +import { IModelRepository } from "modelInfo/modelInfoRepository"; +import { IEmbeddingProcessStateRepository } from "embeddings/embeddingProcessState/embeddingProcessStateRepository"; +import { IEmbeddingProcessState } from "embeddings/embeddingProcessState/embeddingProcessState.types"; +import { ISkillRepository } from "esco/skill/repository/skill.repository"; +import { ISkillGroupRepository } from "esco/skillGroup/repository/SkillGroup.repository"; +import { IOccupationRepository } from "esco/occupations/repository/occupation.repository"; +import { IOccupationGroupRepository } from "esco/occupationGroup/repository/OccupationGroup.repository"; +import { IEmbeddingClient } from "embeddings/service/client"; +import { EmbeddableEntityType, EmbeddableField } from "embeddings/service/types"; +import { IEmbeddingProcessService } from "./embeddingProcess.service.types"; +import { + DatabaseError, + EmbeddingProcessAlreadyRunningError, + ModelNotFoundError, + ModelNotReleasedError, +} from "./errors"; + +/** + * The fields of each entity type that should be embedded. + */ +const EMBEDDABLE_FIELDS_BY_ENTITY_TYPE: Record = { + [EmbeddableEntityType.Skill]: [ + EmbeddableField.preferredLabel, + EmbeddableField.description, + EmbeddableField.altLabels, + ], + [EmbeddableEntityType.SkillGroup]: [ + EmbeddableField.preferredLabel, + EmbeddableField.description, + EmbeddableField.altLabels, + EmbeddableField.scopeNote, + ], + [EmbeddableEntityType.Occupation]: [ + EmbeddableField.preferredLabel, + EmbeddableField.description, + EmbeddableField.altLabels, + ], + [EmbeddableEntityType.OccupationGroup]: [ + EmbeddableField.preferredLabel, + EmbeddableField.description, + EmbeddableField.altLabels, + ], +}; + +export class EmbeddingProcessService implements IEmbeddingProcessService { + constructor( + private readonly modelRepository: IModelRepository, + private readonly embeddingProcessStateRepository: IEmbeddingProcessStateRepository, + private readonly skillRepository: ISkillRepository, + private readonly skillGroupRepository: ISkillGroupRepository, + private readonly occupationRepository: IOccupationRepository, + private readonly occupationGroupRepository: IOccupationGroupRepository, + private readonly embeddingClient: IEmbeddingClient + ) {} + + async triggerEmbeddingProcess(modelId: string, embeddingServiceId: string): Promise { + let model = null; + try { + // 1. Check that the model exists and has been released. + model = await this.modelRepository.getModelById(modelId); + } catch (e) { + throw new DatabaseError("Failed to find the taxonomy model by id"); + } + + if (model === null) { + throw new ModelNotFoundError(modelId); + } + + if (!model.released) { + throw new ModelNotReleasedError(modelId); + } + + let pendingProcess; + try { + // 2. Check that there is no other unfinished embedding process for the same model. + pendingProcess = await this.embeddingProcessStateRepository.findPendingByModelId(modelId); + } catch (e) { + throw new DatabaseError("Failed to find the embedding process state by model id"); + } + + if (pendingProcess !== null) { + throw new EmbeddingProcessAlreadyRunningError(modelId); + } + + try { + // 3. Create a new embedding process state. + const embeddingProcessState = await this.embeddingProcessStateRepository.create({ + modelId, + status: ModelInfoApiSpecs.ModelInfo.EmbeddingProcessStates.Enums.Status.PENDING, + embeddingServiceId, + totalDocuments: 0, + errorCounts: 0, + warningCounts: 0, + completedDocuments: 0, + }); + + // 4. Loop through all the entities of the model and push each of them to the embeddings queue. + // The order follows the requirement: skills, skill groups, occupations, occupation groups. + const entitySources: { entityType: EmbeddableEntityType; findAll: () => Readable }[] = [ + { entityType: EmbeddableEntityType.Skill, findAll: () => this.skillRepository.findAll(modelId) }, + { entityType: EmbeddableEntityType.SkillGroup, findAll: () => this.skillGroupRepository.findAll(modelId) }, + { entityType: EmbeddableEntityType.Occupation, findAll: () => this.occupationRepository.findAll(modelId) }, + { + entityType: EmbeddableEntityType.OccupationGroup, + findAll: () => this.occupationGroupRepository.findAll(modelId), + }, + ]; + + let totalDocuments = 0; + for (const entitySource of entitySources) { + totalDocuments += await this.pushEntitiesToQueue(modelId, entitySource.entityType, entitySource.findAll()); + } + + // 5. Update the embedding process state with the total number of documents that were pushed to the queue. + return this.embeddingProcessStateRepository.update(embeddingProcessState.id, { + status: ModelInfoApiSpecs.ModelInfo.EmbeddingProcessStates.Enums.Status.IN_PROGRESS, + totalDocuments, + }); + } catch (e) { + throw new Error(`Failed to trigger embedding process for model ${modelId}.`, { cause: e }); + } + } + + private async pushEntitiesToQueue( + modelId: string, + entityType: EmbeddableEntityType, + stream: Readable + ): Promise { + let count = 0; + for await (const entity of stream) { + await this.embeddingClient.pushTaskToQueue({ + modelId, + entityId: (entity as { id: string }).id, + entityType, + fields: EMBEDDABLE_FIELDS_BY_ENTITY_TYPE[entityType], + }); + count++; + } + return count; + } +} diff --git a/backend/src/embeddings/embeddingProcess/embeddingProcess.service.types.ts b/backend/src/embeddings/embeddingProcess/embeddingProcess.service.types.ts new file mode 100644 index 000000000..d5948bc14 --- /dev/null +++ b/backend/src/embeddings/embeddingProcess/embeddingProcess.service.types.ts @@ -0,0 +1,23 @@ +import { IEmbeddingProcessState } from "embeddings/embeddingProcessState/embeddingProcessState.types"; + +export interface IEmbeddingProcessService { + /** + * Triggers the generation of the embeddings for all the entities of a model. + * + * The process: + * 1. checks that the model exists and has been released (embeddings can only be generated for released models), + * 2. checks that there is no other unfinished embedding process for the same model, + * 3. creates a new embedding process state, + * 4. loops through all the entities of the model (skills, skill groups, occupations, occupation groups) and + * pushes each of them to the embeddings queue, + * 5. updates the embedding process state with the total number of documents that were pushed to the queue. + * + * @param {string} modelId - The id of the model to generate the embeddings for. + * @param {string} embeddingServiceId - The id of the embedding model service to use. + * @return {Promise} - A Promise that resolves to the created embedding process state. + * @throws {ModelNotFoundError} - If the model does not exist. + * @throws {ModelNotReleasedError} - If the model has not been released. + * @throws {EmbeddingProcessAlreadyRunningError} - If there is already an unfinished embedding process for the model. + */ + triggerEmbeddingProcess(modelId: string, embeddingServiceId: string): Promise; +} diff --git a/backend/src/embeddings/embeddingProcess/errors.ts b/backend/src/embeddings/embeddingProcess/errors.ts new file mode 100644 index 000000000..fa11891d6 --- /dev/null +++ b/backend/src/embeddings/embeddingProcess/errors.ts @@ -0,0 +1,41 @@ +/** + * Thrown when an embedding process is triggered for a model that does not exist. + */ +export class ModelNotFoundError extends Error { + constructor(modelId: string) { + super(`Model with id ${modelId} was not found`); + this.name = "ModelNotFoundError"; + } +} + +/** + * Thrown when an embedding process is triggered for a model that has not been released. + * Embeddings can only be generated for released models. + */ +export class ModelNotReleasedError extends Error { + constructor(modelId: string) { + super(`Model with id ${modelId} is not released; embeddings can only be generated for released models`); + this.name = "ModelNotReleasedError"; + } +} + +/** + * Thrown when an embedding process is triggered for a model that already has an unfinished embedding process. + * The new process should wait for the previous one to complete. + */ +export class EmbeddingProcessAlreadyRunningError extends Error { + constructor(modelId: string) { + super(`An embedding process is already running for model with id ${modelId}`); + this.name = "EmbeddingProcessAlreadyRunningError"; + } +} + +/** + * Thrown when an error occurs while interacting with the database. + */ +export class DatabaseError extends Error { + constructor(message: string) { + super(message); + this.name = "DatabaseError"; + } +} diff --git a/backend/src/embeddings/embeddingProcessState/embeddingProcessState.types.ts b/backend/src/embeddings/embeddingProcessState/embeddingProcessState.types.ts new file mode 100644 index 000000000..72ab07e76 --- /dev/null +++ b/backend/src/embeddings/embeddingProcessState/embeddingProcessState.types.ts @@ -0,0 +1,35 @@ +import mongoose from "mongoose"; +import ModelInfoApiSpecs from "api-specifications/modelInfo"; + +/** + * Describes how an embedding process state is saved in the database. + */ +export interface IEmbeddingProcessStateDoc { + modelId: mongoose.Types.ObjectId; + status: ModelInfoApiSpecs.ModelInfo.EmbeddingProcessStates.Enums.Status; + embeddingServiceId: string; + totalDocuments: number; + errorCounts: number; + warningCounts: number; + completedDocuments: number; + createdAt: Date; + updatedAt: Date; +} + +/** + * Describes how an embedding process state is returned from the API. + */ +export interface IEmbeddingProcessState extends Omit { + id: string; + modelId: string; +} + +/** + * Describes how a new embedding process state is created. + */ +export type INewEmbeddingProcessStateSpec = Omit; + +/** + * Describes how an embedding process state is updated. + */ +export type IUpdateEmbeddingProcessStateSpec = Partial>; diff --git a/backend/src/embeddings/embeddingProcessState/embeddingProcessStateModel.test.ts b/backend/src/embeddings/embeddingProcessState/embeddingProcessStateModel.test.ts new file mode 100644 index 000000000..bbfbfade6 --- /dev/null +++ b/backend/src/embeddings/embeddingProcessState/embeddingProcessStateModel.test.ts @@ -0,0 +1,149 @@ +// suppress chatty log output when testing +import "_test_utilities/consoleMock"; + +import mongoose from "mongoose"; +import ModelInfoApiSpecs from "api-specifications/modelInfo"; +import { IEmbeddingProcessStateDoc } from "./embeddingProcessState.types"; +import { getTestConfiguration } from "_test_utilities/getTestConfiguration"; +import { getNewConnection } from "server/connection/newConnection"; +import { initializeSchemaAndModel } from "./embeddingProcessStateModel"; +import { testObjectIdField } from "esco/_test_utilities/modelSchemaTestFunctions"; +import { assertCaseForProperty, CaseType } from "_test_utilities/dataModel"; +import { WHITESPACE } from "_test_utilities/getMockRandomData"; + +describe("Test the definition of EmbeddingProcessState Model", () => { + let dbConnection: mongoose.Connection; + let model: mongoose.Model; + + beforeAll(async () => { + const config = getTestConfiguration("EmbeddingProcessStateModelTestDB"); + dbConnection = await getNewConnection(config.dbURI); + model = initializeSchemaAndModel(dbConnection); + }); + + afterAll(async () => { + if (dbConnection) { + await dbConnection.dropDatabase(); + await dbConnection.close(); + } + }); + + test("should successfully validate with mandatory fields", async () => { + // GIVEN an object with the mandatory fields + const givenObject: IEmbeddingProcessStateDoc = { + modelId: new mongoose.Types.ObjectId(), + status: ModelInfoApiSpecs.ModelInfo.EmbeddingProcessStates.Enums.Status.PENDING, + embeddingServiceId: "gemini$$models/gemini-embedding-2", + totalDocuments: 10, + errorCounts: 0, + warningCounts: 0, + completedDocuments: 0, + createdAt: new Date(), + updatedAt: new Date(), + }; + // AND an EmbeddingProcessState document based on the given object + const givenEmbeddingStateDocument = new model(givenObject); + + // WHEN validating that given document + const actualValidationErrors = givenEmbeddingStateDocument.validateSync(); + + // THEN expect it to validate without any error + expect(actualValidationErrors).toBeUndefined(); + // AND the document to be saved successfully + await givenEmbeddingStateDocument.save(); + // AND the toObject() transformation to return the correct properties + expect(givenEmbeddingStateDocument.toObject()).toEqual({ + ...givenObject, + modelId: givenObject.modelId.toString(), + id: givenEmbeddingStateDocument._id.toString(), + createdAt: expect.any(Date), + updatedAt: expect.any(Date), + }); + }); + + describe("Validate EmbeddingProcessState fields", () => { + testObjectIdField(() => model, "modelId"); + + describe("Test validation of 'status'", () => { + test.each([ + [CaseType.Failure, "undefined", undefined, "Path `{0}` is required."], + [CaseType.Failure, "null", null, "Path `{0}` is required."], + [CaseType.Failure, "only whitespace characters", WHITESPACE, ` is not a valid enum value for path \`{0}\`.`], + [CaseType.Failure, "string", "foo", `\`foo\` is not a valid enum value for path \`{0}\`.`], + [ + CaseType.Success, + ModelInfoApiSpecs.ModelInfo.EmbeddingProcessStates.Enums.Status.PENDING, + ModelInfoApiSpecs.ModelInfo.EmbeddingProcessStates.Enums.Status.PENDING, + undefined, + ], + [ + CaseType.Success, + ModelInfoApiSpecs.ModelInfo.EmbeddingProcessStates.Enums.Status.IN_PROGRESS, + ModelInfoApiSpecs.ModelInfo.EmbeddingProcessStates.Enums.Status.IN_PROGRESS, + undefined, + ], + [ + CaseType.Success, + ModelInfoApiSpecs.ModelInfo.EmbeddingProcessStates.Enums.Status.COMPLETED, + ModelInfoApiSpecs.ModelInfo.EmbeddingProcessStates.Enums.Status.COMPLETED, + undefined, + ], + ])( + `(%s) Validate 'status' when it is %s`, + (caseType: CaseType, caseDescription, value, expectedFailureMessage) => { + assertCaseForProperty({ + model, + propertyNames: "status", + caseType, + testValue: value, + expectedFailureMessage, + }); + } + ); + }); + + describe("Test validation of 'embeddingServiceId'", () => { + test.each([ + [CaseType.Failure, "undefined", undefined, "Path `{0}` is required."], + [CaseType.Failure, "null", null, "Path `{0}` is required."], + [CaseType.Success, "a valid string", "gemini$$models/gemini-embedding-2", undefined], + ])( + `(%s) Validate 'embeddingServiceId' when it is %s`, + (caseType: CaseType, caseDescription, value, expectedFailureMessage) => { + assertCaseForProperty({ + model, + propertyNames: "embeddingServiceId", + caseType, + testValue: value, + expectedFailureMessage, + }); + } + ); + }); + + describe.each([["totalDocuments"], ["errorCounts"], ["warningCounts"], ["completedDocuments"]])( + "Test validation of '%s'", + (fieldName) => { + test.each([ + [CaseType.Failure, "undefined", undefined, "Path `{0}` is required."], + [CaseType.Failure, "null", null, "Path `{0}` is required."], + [CaseType.Failure, "not a number", "foo", "Cast to Number failed"], + [CaseType.Failure, "a negative number", -1, "is less than minimum allowed value"], + [CaseType.Success, "zero", 0, undefined], + [CaseType.Success, "a positive integer", 42, undefined], + ])( + `(%s) Validate '${fieldName}' when it is %s`, + (caseType: CaseType, caseDescription, value, expectedFailureMessage) => { + assertCaseForProperty({ + model, + propertyNames: fieldName, + caseType, + testValue: value, + expectedFailureMessage, + }); + } + ); + } + ); + }); +}); diff --git a/backend/src/embeddings/embeddingProcessState/embeddingProcessStateModel.ts b/backend/src/embeddings/embeddingProcessState/embeddingProcessStateModel.ts new file mode 100644 index 000000000..614988784 --- /dev/null +++ b/backend/src/embeddings/embeddingProcessState/embeddingProcessStateModel.ts @@ -0,0 +1,55 @@ +import mongoose from "mongoose"; +import ModelInfoAPISpecs from "api-specifications/modelInfo"; +import { IEmbeddingProcessStateDoc } from "./embeddingProcessState.types"; +import { getGlobalTransformOptions } from "server/repositoryRegistry/globalTransform"; + +export const ModelName = "EmbeddingProcessStateModel"; + +export const EmbeddingProcessStateModelPaths = { + modelId: "modelId", +}; + +export function initializeSchemaAndModel(dbConnection: mongoose.Connection): mongoose.Model { + const schema = new mongoose.Schema( + { + [EmbeddingProcessStateModelPaths.modelId]: { type: mongoose.Schema.Types.ObjectId, required: true }, + status: { + type: String, + required: true, + enum: Object.values(ModelInfoAPISpecs.ModelInfo.EmbeddingProcessStates.Enums.Status), + }, + embeddingServiceId: { + type: String, + required: true, + }, + totalDocuments: { + type: Number, + required: true, + min: 0, + }, + errorCounts: { + type: Number, + required: true, + min: 0, + }, + warningCounts: { + type: Number, + required: true, + min: 0, + }, + completedDocuments: { + type: Number, + required: true, + min: 0, + }, + }, + { + timestamps: true, + strict: "throw", + toObject: getGlobalTransformOptions(), + toJSON: getGlobalTransformOptions(), + } + ); + + return dbConnection.model(ModelName, schema); +} diff --git a/backend/src/embeddings/embeddingProcessState/embeddingProcessStateRepository.test.ts b/backend/src/embeddings/embeddingProcessState/embeddingProcessStateRepository.test.ts new file mode 100644 index 000000000..8ad145ded --- /dev/null +++ b/backend/src/embeddings/embeddingProcessState/embeddingProcessStateRepository.test.ts @@ -0,0 +1,235 @@ +// suppress chatty log output when testing +import "_test_utilities/consoleMock"; + +import { Connection } from "mongoose"; +import ModelInfoApiSpecs from "api-specifications/modelInfo"; +import { IEmbeddingProcessStateRepository } from "./embeddingProcessStateRepository"; +import { getTestConfiguration } from "_test_utilities/getTestConfiguration"; +import { getNewConnection } from "server/connection/newConnection"; +import { getRepositoryRegistry, RepositoryRegistry } from "server/repositoryRegistry/repositoryRegistry"; +import { initOnce } from "server/init"; +import { getConnectionManager } from "server/connection/connectionManager"; +import { INewEmbeddingProcessStateSpec } from "./embeddingProcessState.types"; +import { getMockStringId } from "_test_utilities/mockMongoId"; +import { TestDBConnectionFailureNoSetup } from "_test_utilities/testDBConnectionFaillure"; + +function getNewEmbeddingProcessStateSpec(): INewEmbeddingProcessStateSpec { + return { + modelId: getMockStringId(2), + status: ModelInfoApiSpecs.ModelInfo.EmbeddingProcessStates.Enums.Status.PENDING, + embeddingServiceId: "gemini$$models/gemini-embedding-2", + totalDocuments: 10, + errorCounts: 0, + warningCounts: 0, + completedDocuments: 0, + }; +} + +function expectedFromGivenSpec(givenSpec: INewEmbeddingProcessStateSpec) { + return { + ...givenSpec, + id: expect.any(String), + createdAt: expect.any(Date), + updatedAt: expect.any(Date), + }; +} + +describe("Test EmbeddingProcessState Repository with an in-memory mongodb", () => { + let dbConnection: Connection; + let repository: IEmbeddingProcessStateRepository; + + beforeAll(async () => { + const config = getTestConfiguration("EmbeddingProcessStateRepositoryTestDB"); + dbConnection = await getNewConnection(config.dbURI); + const repositoryRegistry = new RepositoryRegistry(); + await repositoryRegistry.initialize(dbConnection); + repository = repositoryRegistry.embeddingProcessState; + }); + + afterAll(async () => { + if (dbConnection) { + await dbConnection.dropDatabase(); + await dbConnection.close(false); + } + }); + + beforeEach(async () => { + await repository.Model.deleteMany({}).exec(); + }); + + afterEach(async () => { + await repository.Model.deleteMany({}).exec(); + }); + + test("should return the model", async () => { + expect(repository.Model).toBeDefined(); + }); + + test("initOnce has registered the repository", async () => { + // GIVEN that the mongodb uri should be set + expect(process.env.MONGODB_URI).toBeDefined(); + + // WHEN initOnce has been called + await initOnce(); + + // THEN expect the repository to be defined + expect(getRepositoryRegistry().embeddingProcessState).toBeDefined(); + + // Clean up + await getConnectionManager().getCurrentDBConnection()!.close(false); + }); + + describe("Test create() EmbeddingProcessState", () => { + test("should successfully create a new EmbeddingProcessState", async () => { + // GIVEN a valid newEmbeddingProcessStateSpec + const givenNewEmbeddingProcessStateSpec = getNewEmbeddingProcessStateSpec(); + + // WHEN creating a new EmbeddingProcessState + const actualEmbeddingProcessState = await repository.create(givenNewEmbeddingProcessStateSpec); + + // THEN expect the created EmbeddingProcessState to match the given spec + expect(actualEmbeddingProcessState).toEqual(expectedFromGivenSpec(givenNewEmbeddingProcessStateSpec)); + }); + + TestDBConnectionFailureNoSetup((repositoryRegistry) => { + return repositoryRegistry.embeddingProcessState.create(getNewEmbeddingProcessStateSpec()); + }); + }); + + describe("Test update() EmbeddingProcessState", () => { + test("should successfully update all the mutable fields of an EmbeddingProcessState", async () => { + // GIVEN an EmbeddingProcessState in the database + const givenNewEmbeddingProcessStateSpec = getNewEmbeddingProcessStateSpec(); + const createdEmbeddingProcessState = await repository.create(givenNewEmbeddingProcessStateSpec); + // AND new values for all the mutable fields + const givenUpdateSpec = { + status: ModelInfoApiSpecs.ModelInfo.EmbeddingProcessStates.Enums.Status.COMPLETED, + embeddingServiceId: "gemini$$models/gemini-embedding-2", + totalDocuments: 20, + errorCounts: 1, + warningCounts: 2, + completedDocuments: 17, + }; + + // WHEN updating the EmbeddingProcessState + const actualUpdatedEmbeddingProcessState = await repository.update( + createdEmbeddingProcessState.id, + givenUpdateSpec + ); + + // THEN expect the updated EmbeddingProcessState to match the given spec + expect(actualUpdatedEmbeddingProcessState).toEqual({ + ...expectedFromGivenSpec(givenNewEmbeddingProcessStateSpec), + ...givenUpdateSpec, + }); + }); + + test("should reject with an error when updating an EmbeddingProcessState that does not exist", async () => { + // GIVEN an id of an EmbeddingProcessState that does not exist + const givenId = getMockStringId(1); + // AND valid updateSpecs + const givenUpdateSpecs = { + status: ModelInfoApiSpecs.ModelInfo.EmbeddingProcessStates.Enums.Status.COMPLETED, + }; + + // WHEN updating the EmbeddingProcessState with an id that does not exist + const actualUpdatedEmbeddingProcessStatePromise = repository.update(givenId, givenUpdateSpecs); + + // THEN expect to reject with an error + await expect(actualUpdatedEmbeddingProcessStatePromise).rejects.toThrow( + expect.toMatchErrorWithCause( + "EmbeddingProcessStateRepository.update: update failed", + `Update failed to find embedding process with id: ${givenId}` + ) + ); + }); + + TestDBConnectionFailureNoSetup((repositoryRegistry) => { + return repositoryRegistry.embeddingProcessState.update(getMockStringId(1), { + status: ModelInfoApiSpecs.ModelInfo.EmbeddingProcessStates.Enums.Status.COMPLETED, + }); + }); + }); + + describe("Test findById() EmbeddingProcessState", () => { + test("should successfully find an EmbeddingProcessState by id", async () => { + // GIVEN an EmbeddingProcessState in the database + const givenNewEmbeddingProcessStateSpec = getNewEmbeddingProcessStateSpec(); + const createdEmbeddingProcessState = await repository.create(givenNewEmbeddingProcessStateSpec); + + // WHEN finding the EmbeddingProcessState by id + const actualFoundEmbeddingProcessState = await repository.findById(createdEmbeddingProcessState.id); + + // THEN expect the found EmbeddingProcessState to match the created EmbeddingProcessState + expect(actualFoundEmbeddingProcessState).toEqual(createdEmbeddingProcessState); + }); + + test("should return null when finding an EmbeddingProcessState by id that does not exist", async () => { + // GIVEN an id of an EmbeddingProcessState that does not exist + const givenId = getMockStringId(1); + + // WHEN finding the EmbeddingProcessState by id + const actualFoundEmbeddingProcessState = await repository.findById(givenId); + + // THEN expect the found EmbeddingProcessState to be null + expect(actualFoundEmbeddingProcessState).toBeNull(); + }); + + TestDBConnectionFailureNoSetup((repositoryRegistry) => { + return repositoryRegistry.embeddingProcessState.findById(getMockStringId(1)); + }); + }); + + describe("Test findPendingByModelId() EmbeddingProcessState", () => { + test.each([ + [ModelInfoApiSpecs.ModelInfo.EmbeddingProcessStates.Enums.Status.PENDING], + [ModelInfoApiSpecs.ModelInfo.EmbeddingProcessStates.Enums.Status.IN_PROGRESS], + ])("should find an unfinished EmbeddingProcessState with status '%s'", async (givenStatus) => { + // GIVEN an unfinished EmbeddingProcessState in the database for a model + const givenModelId = getMockStringId(3); + const givenNewEmbeddingProcessStateSpec = { + ...getNewEmbeddingProcessStateSpec(), + modelId: givenModelId, + status: givenStatus, + }; + const createdEmbeddingProcessState = await repository.create(givenNewEmbeddingProcessStateSpec); + + // WHEN finding a pending EmbeddingProcessState by the model id + const actualFoundEmbeddingProcessState = await repository.findPendingByModelId(givenModelId); + + // THEN expect the found EmbeddingProcessState to match the created one + expect(actualFoundEmbeddingProcessState).toEqual(createdEmbeddingProcessState); + }); + + test("should return null when the only EmbeddingProcessState for the model is completed", async () => { + // GIVEN a completed EmbeddingProcessState in the database for a model + const givenModelId = getMockStringId(3); + await repository.create({ + ...getNewEmbeddingProcessStateSpec(), + modelId: givenModelId, + status: ModelInfoApiSpecs.ModelInfo.EmbeddingProcessStates.Enums.Status.COMPLETED, + }); + + // WHEN finding a pending EmbeddingProcessState by the model id + const actualFoundEmbeddingProcessState = await repository.findPendingByModelId(givenModelId); + + // THEN expect the found EmbeddingProcessState to be null + expect(actualFoundEmbeddingProcessState).toBeNull(); + }); + + test("should return null when there is no EmbeddingProcessState for the model", async () => { + // GIVEN no EmbeddingProcessState in the database for the model + const givenModelId = getMockStringId(4); + + // WHEN finding a pending EmbeddingProcessState by the model id + const actualFoundEmbeddingProcessState = await repository.findPendingByModelId(givenModelId); + + // THEN expect the found EmbeddingProcessState to be null + expect(actualFoundEmbeddingProcessState).toBeNull(); + }); + + TestDBConnectionFailureNoSetup((repositoryRegistry) => { + return repositoryRegistry.embeddingProcessState.findPendingByModelId(getMockStringId(1)); + }); + }); +}); diff --git a/backend/src/embeddings/embeddingProcessState/embeddingProcessStateRepository.ts b/backend/src/embeddings/embeddingProcessState/embeddingProcessStateRepository.ts new file mode 100644 index 000000000..fa05bafe2 --- /dev/null +++ b/backend/src/embeddings/embeddingProcessState/embeddingProcessStateRepository.ts @@ -0,0 +1,136 @@ +import mongoose from "mongoose"; +import ModelInfoApiSpecs from "api-specifications/modelInfo"; +import { + IEmbeddingProcessState, + IEmbeddingProcessStateDoc, + INewEmbeddingProcessStateSpec, + IUpdateEmbeddingProcessStateSpec, +} from "./embeddingProcessState.types"; + +export interface IEmbeddingProcessStateRepository { + readonly Model: mongoose.Model; + + /** + * Creates a new EmbeddingProcessState entry. + * + * @param {INewEmbeddingProcessStateSpec} newSpecs - The specification for the new EmbeddingProcessState entry. + * @return {Promise} - A Promise that resolves to the newly created EmbeddingProcessState entry. + * Rejects with an error if the EmbeddingProcessState entry cannot be created. + */ + create(newSpecs: INewEmbeddingProcessStateSpec): Promise; + + /** + * Updates the EmbeddingProcessState entry with the given ID. + * + * @param {string} id - The unique ID of the EmbeddingProcessState entry. + * @param {IUpdateEmbeddingProcessStateSpec} updateSpecs - The specification to update with + * @return {Promise} - A Promise that resolves to the updated EmbeddingProcessState entry. + * Rejects with an error if the EmbeddingProcessState entry cannot be updated. + */ + update(id: string, updateSpecs: IUpdateEmbeddingProcessStateSpec): Promise; + + /** + * Finds an EmbeddingProcessState entry by its ID. + * + * @param {string} id - The unique ID of the EmbeddingProcessState entry. + * @return {Promise} - A Promise that resolves to the found EmbeddingProcessState entry or null if not found. + * Rejects with an error if the operation fails. + */ + findById(id: string): Promise; + + /** + * Finds an unfinished (pending or in progress) EmbeddingProcessState entry for the given model. + * + * @param {string} modelId - The unique ID of the model. + * @return {Promise} - A Promise that resolves to the found EmbeddingProcessState entry or null if there is none. + * Rejects with an error if the operation fails. + */ + findPendingByModelId(modelId: string): Promise; +} + +export class EmbeddingProcessStateRepository implements IEmbeddingProcessStateRepository { + public readonly Model: mongoose.Model; + + constructor(model: mongoose.Model) { + this.Model = model; + } + + async create(newSpecs: INewEmbeddingProcessStateSpec): Promise { + try { + const newDoc = new this.Model({ + ...newSpecs, + }); + await newDoc.save(); + return newDoc.toObject(); + } catch (e: unknown) { + const err = new Error("EmbeddingProcessStateRepository.create: create failed", { cause: e }); + console.error(err); + throw err; + } + } + + async update(id: string, updateSpecs: IUpdateEmbeddingProcessStateSpec): Promise { + try { + const doc = await this.Model.findById(id).exec(); + if (doc === null) { + throw new Error("Update failed to find embedding process with id: " + id); + } + if (updateSpecs.status !== undefined) { + doc.status = updateSpecs.status; + } + if (updateSpecs.embeddingServiceId !== undefined) { + doc.embeddingServiceId = updateSpecs.embeddingServiceId; + } + if (updateSpecs.totalDocuments !== undefined) { + doc.totalDocuments = updateSpecs.totalDocuments; + } + if (updateSpecs.errorCounts !== undefined) { + doc.errorCounts = updateSpecs.errorCounts; + } + if (updateSpecs.warningCounts !== undefined) { + doc.warningCounts = updateSpecs.warningCounts; + } + if (updateSpecs.completedDocuments !== undefined) { + doc.completedDocuments = updateSpecs.completedDocuments; + } + await doc.save(); + return doc.toObject(); + } catch (e: unknown) { + const err = new Error("EmbeddingProcessStateRepository.update: update failed", { cause: e }); + console.error(err); + throw err; + } + } + + async findById(id: string): Promise { + try { + const embeddingProcessState = (await this.Model.findById(id)) as mongoose.Document; + return embeddingProcessState !== null ? embeddingProcessState.toObject() : null; + } catch (e: unknown) { + const err = new Error("EmbeddingProcessStateRepository.findById: findById failed", { cause: e }); + console.error(err); + throw err; + } + } + + async findPendingByModelId(modelId: string): Promise { + try { + const embeddingProcessState = await this.Model.findOne({ + modelId: { $eq: modelId }, + // Pass a bare array (not an explicit { $in: [...] }): mongoose applies $in automatically, and unlike an + // operator object this is not rewritten by the connection's sanitizeFilter=true. + status: [ + ModelInfoApiSpecs.ModelInfo.EmbeddingProcessStates.Enums.Status.PENDING, + ModelInfoApiSpecs.ModelInfo.EmbeddingProcessStates.Enums.Status.IN_PROGRESS, + ], + }).exec(); + return embeddingProcessState !== null ? embeddingProcessState.toObject() : null; + } catch (e: unknown) { + const err = new Error("EmbeddingProcessStateRepository.findPendingByModelId: findPendingByModelId failed", { + cause: e, + }); + console.error(err); + throw err; + } + } +} diff --git a/backend/src/embeddings/handler/index.test.ts b/backend/src/embeddings/handler/index.test.ts new file mode 100644 index 000000000..6aae23324 --- /dev/null +++ b/backend/src/embeddings/handler/index.test.ts @@ -0,0 +1,101 @@ +// mute console +import "_test_utilities/consoleMock"; +import "_test_utilities/mockSentry"; + +// ############## +// Mock the server init function +jest.mock("server/init", () => ({ + initOnce: jest.fn(), +})); + +// Mock the EmbeddingService so we can observe the processing of the tasks +const mockProcessTask = jest.fn(); +jest.mock("embeddings/service/service", () => ({ + EmbeddingService: jest.fn().mockImplementation(() => ({ processTask: mockProcessTask })), +})); + +// ############## +import { handler } from "./index"; +import { SQSEvent } from "aws-lambda"; +import { EmbeddableEntityType, EmbeddableField, IGenerateEmbeddingTask } from "embeddings/service/types"; +import { getMockStringId } from "_test_utilities/mockMongoId"; + +function getValidTask(index: number): IGenerateEmbeddingTask { + return { + modelId: getMockStringId(1), + entityId: getMockStringId(index), + entityType: EmbeddableEntityType.Skill, + fields: [EmbeddableField.preferredLabel], + }; +} + +function getSQSEvent(bodies: string[]): SQSEvent { + return { + Records: bodies.map((body, index) => ({ body, messageId: `message-${index}` })), + } as never; +} + +describe("Test the embeddings lambda handler", () => { + beforeEach(() => { + jest.clearAllMocks(); + }); + + test("should process each valid task in the event", async () => { + // GIVEN an event with two valid embedding tasks + const givenTask1 = getValidTask(2); + const givenTask2 = getValidTask(3); + const givenEvent = getSQSEvent([JSON.stringify(givenTask1), JSON.stringify(givenTask2)]); + + // WHEN the handler is invoked with the given event + await handler(givenEvent, {} as never, {} as never); + + // THEN expect each task to be processed + expect(mockProcessTask).toHaveBeenCalledTimes(2); + expect(mockProcessTask).toHaveBeenCalledWith(givenTask1); + expect(mockProcessTask).toHaveBeenCalledWith(givenTask2); + }); + + test("should skip an unparseable record without throwing", async () => { + // GIVEN an event with an unparseable record and a valid record + const givenValidTask = getValidTask(2); + const givenEvent = getSQSEvent(["{ not json", JSON.stringify(givenValidTask)]); + + // WHEN the handler is invoked with the given event + const actualPromise = handler(givenEvent, {} as never, {} as never); + + // THEN expect it to resolve without throwing + await expect(actualPromise).resolves.toBeUndefined(); + // AND expect only the valid task to be processed + expect(mockProcessTask).toHaveBeenCalledTimes(1); + expect(mockProcessTask).toHaveBeenCalledWith(givenValidTask); + }); + + test("should skip a record that does not conform to the queue job schema", async () => { + // GIVEN an event with a record that does not conform to the schema and a valid record + const givenInvalidTask = { ...getValidTask(2), fields: [] }; + const givenValidTask = getValidTask(3); + const givenEvent = getSQSEvent([JSON.stringify(givenInvalidTask), JSON.stringify(givenValidTask)]); + + // WHEN the handler is invoked with the given event + const actualPromise = handler(givenEvent, {} as never, {} as never); + + // THEN expect it to resolve without throwing + await expect(actualPromise).resolves.toBeUndefined(); + // AND expect only the valid task to be processed + expect(mockProcessTask).toHaveBeenCalledTimes(1); + expect(mockProcessTask).toHaveBeenCalledWith(givenValidTask); + }); + + test("should handle an event with no records", async () => { + // GIVEN an event with no records + const givenEvent = {} as SQSEvent; + + // WHEN the handler is invoked with the given event + const actualPromise = handler(givenEvent, {} as never, {} as never); + + // THEN expect it to resolve without throwing + await expect(actualPromise).resolves.toBeUndefined(); + // AND expect no task to be processed + expect(mockProcessTask).not.toHaveBeenCalled(); + }); +}); diff --git a/backend/src/embeddings/handler/index.ts b/backend/src/embeddings/handler/index.ts new file mode 100644 index 000000000..45858c6b4 --- /dev/null +++ b/backend/src/embeddings/handler/index.ts @@ -0,0 +1,64 @@ +import { SQSEvent, SQSRecord } from "aws-lambda"; +import { initOnce as serverInitOnce } from "server/init"; +import { initializeSentry } from "initializeSentry"; +import * as Sentry from "@sentry/aws-serverless"; +import { Lambdas } from "common/lambda.types"; +import { IGenerateEmbeddingTask } from "embeddings/service/types"; +import { EmbeddingService } from "embeddings/service/service"; +import { validateEmbeddingQueueJob } from "embeddings/specs/queueJob.schema"; + +initializeSentry(Lambdas.EMBEDDING); + +async function initOnce() { + await serverInitOnce(); + return new EmbeddingService(); +} + +/** + * Entry point for the embeddings lambda function. + * + * This lambda is triggered by the embeddings SQS queue via an event source mapping: when N messages + * are available on the queue, the lambda is invoked with up to N records in `event.Records`, and this + * handler generates the embeddings for each of them in the background. + * + * If a record throws, we rethrow so that AWS retries the message and, after the configured number of + * attempts, moves it to the dead-letter queue. + */ +export const handler = Sentry.wrapHandler(async (event: SQSEvent): Promise => { + const records = event.Records ?? []; + console.info(`Embeddings lambda triggered with ${records.length} record(s)`); + + // Initialize the connection to the database (and registries). If it fails, rethrow so the lambda is retried. + const embeddingService = await initOnce(); + + for (const record of records) { + await handleRecord(embeddingService, record); + } +}); + +async function handleRecord(embeddingService: EmbeddingService, record: SQSRecord): Promise { + let task: IGenerateEmbeddingTask; + try { + task = JSON.parse(record.body) as IGenerateEmbeddingTask; + } catch (e: unknown) { + // A malformed message will never succeed on retry, so log and skip it rather than throwing. + console.error(new Error(`Embeddings lambda: skipping unparseable SQS record ${record.messageId}`, { cause: e })); + return; + } + + // Validate the task against the queue job schema. An invalid job will never succeed on retry, + // so log and skip it rather than throwing. + if (!validateEmbeddingQueueJob(task)) { + console.error( + new Error( + `Embeddings lambda: skipping SQS record ${ + record.messageId + } that does not conform to the queue job schema: ${JSON.stringify(validateEmbeddingQueueJob.errors)}` + ) + ); + return; + } + + console.info("Processing embedding task", { messageId: record.messageId, task }); + await embeddingService.processTask(task); +} diff --git a/backend/src/embeddings/models/gemini/geminiService.test.ts b/backend/src/embeddings/models/gemini/geminiService.test.ts new file mode 100644 index 000000000..0c8dc16a8 --- /dev/null +++ b/backend/src/embeddings/models/gemini/geminiService.test.ts @@ -0,0 +1,54 @@ +// silence chatty console +import "_test_utilities/consoleMock"; + +import { GeminiService } from "./geminiService"; + +describe("Test the GeminiService", () => { + const givenApiKey = "some-api-key"; + const givenModel = "models/gemini-embedding-2"; + + test("should throw when the api key is not configured", () => { + // GIVEN no api key + // WHEN constructing a GeminiService + // THEN expect it to throw + expect(() => new GeminiService("", givenModel)).toThrow("GeminiService: GEMINI_API_KEY is not configured"); + }); + + test("should throw when the model is not configured", () => { + // GIVEN an api key but no model + // WHEN constructing a GeminiService + // THEN expect it to throw + expect(() => new GeminiService(givenApiKey, "")).toThrow("GeminiService: GEMINI_MODEL_NAME is not configured"); + }); + + test("should construct successfully with an api key and a model", () => { + // GIVEN an api key and a model + // WHEN constructing a GeminiService + const actualService = new GeminiService(givenApiKey, givenModel); + + // THEN expect the service to be constructed + expect(actualService).toBeInstanceOf(GeminiService); + }); + + test("should return a batch of embeddings", async () => { + // GIVEN a GeminiService + const givenService = new GeminiService(givenApiKey, givenModel); + + // WHEN generating a batch of embeddings + const actualEmbeddings = await givenService.generateEmbeddingBatch(["foo", "bar"]); + + // THEN expect an array of embeddings to be returned + expect(actualEmbeddings).toEqual([[]]); + }); + + test("should return a single embedding", async () => { + // GIVEN a GeminiService + const givenService = new GeminiService(givenApiKey, givenModel); + + // WHEN generating a single embedding + const actualEmbedding = await givenService.generateEmbedding("foo"); + + // THEN expect the first embedding of the batch to be returned + expect(actualEmbedding).toEqual([]); + }); +}); diff --git a/backend/src/embeddings/models/gemini/geminiService.ts b/backend/src/embeddings/models/gemini/geminiService.ts new file mode 100644 index 000000000..7bb366df7 --- /dev/null +++ b/backend/src/embeddings/models/gemini/geminiService.ts @@ -0,0 +1,43 @@ +import { IEmbeddingModelService } from "embeddings/models/modelsServiceTypes"; + +/** + * Base URL of the Gemini (Generative Language) REST API. + */ +// const GEMINI_API_BASE_URL = "https://generativelanguage.googleapis.com/v1beta"; + +/** + * GeminiService talks to the Gemini REST API directly via `fetch` + * (no SDK dependency, so it bundles cleanly with esbuild). + * + * See: + * - https://ai.google.dev/api/embeddings#method:-models.embedcontent + * - https://ai.google.dev/api/embeddings#method:-models.batchembedcontents + */ +export class GeminiService implements IEmbeddingModelService { + private readonly apiKey: string; + private readonly model: string; + + constructor(apiKey: string, model: string) { + if (!apiKey) { + throw new Error("GeminiService: GEMINI_API_KEY is not configured"); + } + this.apiKey = apiKey; + + if (!model) + throw new Error( + "GeminiService: GEMINI_MODEL_NAME is not configured. Please set it in the environment variables or in the config file." + ); + this.model = model; + } + + async generateEmbeddingBatch(texts: string[]): Promise { + console.debug("GeminiService: generateEmbeddingBatch", { texts }); + console.warn("Not implemented yet"); + return [[]]; + } + + async generateEmbedding(text: string): Promise { + const [vector] = await this.generateEmbeddingBatch([text]); + return vector; + } +} diff --git a/backend/src/embeddings/models/modelsServiceTypes.ts b/backend/src/embeddings/models/modelsServiceTypes.ts new file mode 100644 index 000000000..440f8988d --- /dev/null +++ b/backend/src/embeddings/models/modelsServiceTypes.ts @@ -0,0 +1,16 @@ +export interface IEmbeddingModelService { + /** + * Generate an embedding vector for the given text. + * + * @param text - The text to embed. + * @returns The embedding vector as an array of numbers. + * @throws If the API key is not configured, or the API returns an error / empty response. + */ + generateEmbedding(text: string): Promise; + + /** + * Generate embeddings for a batch of texts. + * @param texts + */ + generateEmbeddingBatch(texts: string[]): Promise; +} diff --git a/backend/src/embeddings/service/client.test.ts b/backend/src/embeddings/service/client.test.ts new file mode 100644 index 000000000..002fb98cd --- /dev/null +++ b/backend/src/embeddings/service/client.test.ts @@ -0,0 +1,99 @@ +// silence chatty console +import "_test_utilities/consoleMock"; + +import { SendMessageCommand, SQSClient } from "@aws-sdk/client-sqs"; +import { EmbeddingClient } from "./client"; +import { EmbeddableEntityType, EmbeddableField, IGenerateEmbeddingTask } from "./types"; +import { getMockStringId } from "_test_utilities/mockMongoId"; + +const givenQueueUrl = "https://sqs.test.amazonaws.com/123456789012/test-embeddings-queue"; + +jest.mock("server/config/config", () => ({ + getEmbeddingsQueueUrl: jest.fn().mockReturnValue("https://sqs.test.amazonaws.com/123456789012/test-embeddings-queue"), +})); + +function getValidTask(): IGenerateEmbeddingTask { + return { + modelId: getMockStringId(1), + entityId: getMockStringId(2), + entityType: EmbeddableEntityType.Skill, + fields: [EmbeddableField.preferredLabel], + }; +} + +function getMockSqsClient(send: jest.Mock = jest.fn().mockResolvedValue(undefined)): SQSClient { + return { send } as unknown as SQSClient; +} + +describe("Test the EmbeddingClient", () => { + beforeEach(() => { + jest.clearAllMocks(); + }); + + test("should construct with an injected SQSClient", () => { + // GIVEN an SQSClient + const givenSqsClient = getMockSqsClient(); + + // WHEN constructing an EmbeddingClient with the injected SQSClient + const actualClient = new EmbeddingClient(givenSqsClient); + + // THEN expect the client to be constructed + expect(actualClient).toBeInstanceOf(EmbeddingClient); + }); + + test("should push a valid task to the queue", async () => { + // GIVEN a mock SQS client + const givenSqsClient = getMockSqsClient(); + // AND an EmbeddingClient that uses it + const givenEmbeddingClient = new EmbeddingClient(givenSqsClient); + // AND a valid task + const givenTask = getValidTask(); + + // WHEN pushing the task to the queue + await givenEmbeddingClient.pushTaskToQueue(givenTask); + + // THEN expect the SQS client to send a message with the queue url and the task as the body + expect(givenSqsClient.send).toHaveBeenCalledTimes(1); + const actualCommand = (givenSqsClient.send as jest.Mock).mock.calls[0][0] as SendMessageCommand; + expect(actualCommand).toBeInstanceOf(SendMessageCommand); + expect(actualCommand.input).toEqual({ + QueueUrl: givenQueueUrl, + MessageBody: JSON.stringify(givenTask), + }); + }); + + test("should not push an invalid task to the queue and should throw", async () => { + // GIVEN a mock SQS client + const givenSqsClient = getMockSqsClient(); + // AND an EmbeddingClient that uses it + const givenEmbeddingClient = new EmbeddingClient(givenSqsClient); + // AND an invalid task (missing the required fields property) + const givenInvalidTask = { ...getValidTask(), fields: [] }; + + // WHEN pushing the invalid task to the queue + const actualPromise = givenEmbeddingClient.pushTaskToQueue(givenInvalidTask as IGenerateEmbeddingTask); + + // THEN expect it to reject with an error + await expect(actualPromise).rejects.toThrow("does not conform to the queue job schema"); + // AND expect the SQS client to not send any message + expect(givenSqsClient.send).not.toHaveBeenCalled(); + }); + + test("should reject with an error when the SQS client fails to send the message", async () => { + // GIVEN a mock SQS client that fails to send messages + const givenCause = new Error("SQS unavailable"); + const givenSqsClient = getMockSqsClient(jest.fn().mockRejectedValue(givenCause)); + // AND an EmbeddingClient that uses it + const givenEmbeddingClient = new EmbeddingClient(givenSqsClient); + // AND a valid task + const givenTask = getValidTask(); + + // WHEN pushing the task to the queue + const actualPromise = givenEmbeddingClient.pushTaskToQueue(givenTask); + + // THEN expect it to reject with a wrapped error + await expect(actualPromise).rejects.toThrow( + expect.toMatchErrorWithCause("EmbeddingClient.pushTaskToQueue: failed to push task to queue", givenCause.message) + ); + }); +}); diff --git a/backend/src/embeddings/service/client.ts b/backend/src/embeddings/service/client.ts new file mode 100644 index 000000000..3a943fe44 --- /dev/null +++ b/backend/src/embeddings/service/client.ts @@ -0,0 +1,40 @@ +import { SQSClient, SendMessageCommand } from "@aws-sdk/client-sqs"; +import { IGenerateEmbeddingTask } from "./types"; +import { getEmbeddingsQueueUrl } from "server/config/config"; +import { validateEmbeddingQueueJob } from "embeddings/specs/queueJob.schema"; + +export interface IEmbeddingClient { + pushTaskToQueue(task: IGenerateEmbeddingTask): Promise; +} + +export class EmbeddingClient implements IEmbeddingClient { + constructor(private sqsClient: SQSClient) {} + + async pushTaskToQueue(task: IGenerateEmbeddingTask): Promise { + // Validate the task against the queue job schema before pushing it, + // so that a malformed job never ends up on the queue. + const isValid = validateEmbeddingQueueJob(task); + if (!isValid) { + const err = new Error( + `EmbeddingClient.pushTaskToQueue: the task does not conform to the queue job schema: ${JSON.stringify( + validateEmbeddingQueueJob.errors + )}` + ); + console.error(err); + throw err; + } + + try { + await this.sqsClient.send( + new SendMessageCommand({ + QueueUrl: getEmbeddingsQueueUrl(), + MessageBody: JSON.stringify(task), + }) + ); + } catch (e: unknown) { + const err = new Error("EmbeddingClient.pushTaskToQueue: failed to push task to queue", { cause: e }); + console.error(err); + throw err; + } + } +} diff --git a/backend/src/embeddings/service/service.test.ts b/backend/src/embeddings/service/service.test.ts new file mode 100644 index 000000000..7073530bc --- /dev/null +++ b/backend/src/embeddings/service/service.test.ts @@ -0,0 +1,34 @@ +// silence chatty console +import "_test_utilities/consoleMock"; + +import { EmbeddingService } from "./service"; +import { EmbeddableEntityType, EmbeddableField, IGenerateEmbeddingTask } from "./types"; +import { getMockStringId } from "_test_utilities/mockMongoId"; + +describe("Test the EmbeddingService", () => { + test("should process a task without throwing", async () => { + // GIVEN an embedding service + const givenService = new EmbeddingService(); + // AND a task + const givenTask: IGenerateEmbeddingTask = { + modelId: getMockStringId(1), + entityId: getMockStringId(2), + entityType: EmbeddableEntityType.Skill, + fields: [EmbeddableField.preferredLabel], + }; + // AND the console.info is spied on + const consoleInfoSpy = jest.spyOn(console, "info"); + + // WHEN processing the task + const actualPromise = givenService.processTask(givenTask); + + // THEN expect it to resolve + await expect(actualPromise).resolves.toBeUndefined(); + // AND expect the task to be logged + expect(consoleInfoSpy).toHaveBeenCalledWith("Generated embedding", { + modelId: givenTask.modelId, + entityType: givenTask.entityType, + entityId: givenTask.entityId, + }); + }); +}); diff --git a/backend/src/embeddings/service/service.ts b/backend/src/embeddings/service/service.ts new file mode 100644 index 000000000..8373b9abd --- /dev/null +++ b/backend/src/embeddings/service/service.ts @@ -0,0 +1,17 @@ +import { IGenerateEmbeddingTask } from "./types"; + +export interface IEmbeddingService { + processTask(task: IGenerateEmbeddingTask): Promise; +} + +export class EmbeddingService implements IEmbeddingService { + constructor() {} + + async processTask(task: IGenerateEmbeddingTask): Promise { + console.info("Generated embedding", { + modelId: task.modelId, + entityType: task.entityType, + entityId: task.entityId, + }); + } +} diff --git a/backend/src/embeddings/service/types.ts b/backend/src/embeddings/service/types.ts new file mode 100644 index 000000000..39673c575 --- /dev/null +++ b/backend/src/embeddings/service/types.ts @@ -0,0 +1,20 @@ +export enum EmbeddableEntityType { + Skill = "Skill", + Occupation = "Occupation", + OccupationGroup = "OccupationGroup", + SkillGroup = "SkillGroup", +} + +export enum EmbeddableField { + preferredLabel = "preferredLabel", + description = "description", + altLabels = "altLabels", + scopeNote = "scopeNote", +} + +export interface IGenerateEmbeddingTask { + modelId: string; + entityId: string; + entityType: EmbeddableEntityType; + fields: EmbeddableField[]; +} diff --git a/backend/src/embeddings/specs/queueJob.schema.test.ts b/backend/src/embeddings/specs/queueJob.schema.test.ts new file mode 100644 index 000000000..8f96c1772 --- /dev/null +++ b/backend/src/embeddings/specs/queueJob.schema.test.ts @@ -0,0 +1,124 @@ +// silence chatty console +import "_test_utilities/consoleMock"; + +import { EmbeddingQueueJobSchema, validateEmbeddingQueueJob } from "./queueJob.schema"; +import { EmbeddableEntityType, EmbeddableField, IGenerateEmbeddingTask } from "embeddings/service/types"; +import { getMockStringId } from "_test_utilities/mockMongoId"; + +function getValidEmbeddingQueueJob(): IGenerateEmbeddingTask { + return { + modelId: getMockStringId(1), + entityId: getMockStringId(2), + entityType: EmbeddableEntityType.Skill, + fields: [EmbeddableField.preferredLabel, EmbeddableField.description], + }; +} + +describe("Test the EmbeddingQueueJobSchema", () => { + test("the schema should have a valid $id", () => { + // GIVEN the EmbeddingQueueJobSchema + // THEN expect the schema to have a $id + expect(EmbeddingQueueJobSchema.$id).toBe("/components/schemas/EmbeddingQueueJobSchema"); + }); + + test("should validate a valid embedding queue job", () => { + // GIVEN a valid embedding queue job + const givenJob = getValidEmbeddingQueueJob(); + + // WHEN the job is validated + const actualIsValid = validateEmbeddingQueueJob(givenJob); + + // THEN expect the job to be valid + expect(actualIsValid).toBe(true); + // AND expect no validation errors + expect(validateEmbeddingQueueJob.errors).toBeNull(); + }); + + test("should not validate a job with additional properties", () => { + // GIVEN a job with additional properties + const givenJob = { ...getValidEmbeddingQueueJob(), foo: "bar" }; + + // WHEN the job is validated + const actualIsValid = validateEmbeddingQueueJob(givenJob); + + // THEN expect the job to be invalid + expect(actualIsValid).toBe(false); + }); + + describe.each([["modelId"], ["entityId"], ["entityType"], ["fields"]])( + "should not validate a job that is missing the required property '%s'", + (propertyName) => { + test(`missing '${propertyName}'`, () => { + // GIVEN a job that is missing the given required property + const givenJob: Partial = getValidEmbeddingQueueJob(); + delete givenJob[propertyName as keyof IGenerateEmbeddingTask]; + + // WHEN the job is validated + const actualIsValid = validateEmbeddingQueueJob(givenJob); + + // THEN expect the job to be invalid + expect(actualIsValid).toBe(false); + }); + } + ); + + describe.each([["modelId"], ["entityId"]])("Test validation of the id field '%s'", (propertyName) => { + test.each([ + ["not a valid object id", "foo", false], + ["an empty string", "", false], + ["a valid object id", getMockStringId(3), true], + ])(`(%s) should validate '${propertyName}' accordingly`, (_description, givenValue, expectedIsValid) => { + // GIVEN a job with the given value for the id field + const givenJob = { ...getValidEmbeddingQueueJob(), [propertyName]: givenValue }; + + // WHEN the job is validated + const actualIsValid = validateEmbeddingQueueJob(givenJob); + + // THEN expect the job to validate accordingly + expect(actualIsValid).toBe(expectedIsValid); + }); + }); + + describe("Test validation of 'entityType'", () => { + test("should not validate a job with an unknown entityType", () => { + // GIVEN a job with an unknown entityType + const givenJob = { ...getValidEmbeddingQueueJob(), entityType: "UnknownEntityType" }; + + // WHEN the job is validated + const actualIsValid = validateEmbeddingQueueJob(givenJob); + + // THEN expect the job to be invalid + expect(actualIsValid).toBe(false); + }); + + test.each(Object.values(EmbeddableEntityType))("should validate a job with entityType '%s'", (givenEntityType) => { + // GIVEN a job with the given valid entityType + const givenJob = { ...getValidEmbeddingQueueJob(), entityType: givenEntityType }; + + // WHEN the job is validated + const actualIsValid = validateEmbeddingQueueJob(givenJob); + + // THEN expect the job to be valid + expect(actualIsValid).toBe(true); + }); + }); + + describe("Test validation of 'fields'", () => { + test.each([ + ["an empty array", [], false], + ["an array with an unknown field", ["unknownField"], false], + ["an array with duplicate fields", [EmbeddableField.description, EmbeddableField.description], false], + ["not an array", EmbeddableField.description, false], + ["a valid array of fields", [EmbeddableField.preferredLabel, EmbeddableField.altLabels], true], + ])("(%s) should validate 'fields' accordingly", (_description, givenValue, expectedIsValid) => { + // GIVEN a job with the given value for the fields property + const givenJob = { ...getValidEmbeddingQueueJob(), fields: givenValue }; + + // WHEN the job is validated + const actualIsValid = validateEmbeddingQueueJob(givenJob); + + // THEN expect the job to validate accordingly + expect(actualIsValid).toBe(expectedIsValid); + }); + }); +}); diff --git a/backend/src/embeddings/specs/queueJob.schema.ts b/backend/src/embeddings/specs/queueJob.schema.ts new file mode 100644 index 000000000..5d8924e68 --- /dev/null +++ b/backend/src/embeddings/specs/queueJob.schema.ts @@ -0,0 +1,55 @@ +import Ajv, { SchemaObject, ValidateFunction } from "ajv"; +import addFormats from "ajv-formats"; +import { RegExp_Str_ID } from "server/regex"; +import { EmbeddableEntityType, EmbeddableField } from "embeddings/service/types"; + +/** + * The AJV schema of a single embedding generation job that travels through the embeddings SQS queue. + * + * This schema is backend-only (the frontend never produces or consumes queue jobs) and is used on both + * ends of the queue: + * - the producer (EmbeddingClient) validates a job before pushing it to the queue, and + * - the consumer (embeddings lambda) validates a job after receiving it from the queue. + */ +export const EmbeddingQueueJobSchema: SchemaObject = { + $id: "/components/schemas/EmbeddingQueueJobSchema", + type: "object", + additionalProperties: false, + properties: { + modelId: { + description: "The identifier of the model that the entity belongs to.", + type: "string", + pattern: RegExp_Str_ID, + }, + entityId: { + description: "The identifier of the entity to generate the embedding for.", + type: "string", + pattern: RegExp_Str_ID, + }, + entityType: { + description: "The type of the entity to generate the embedding for.", + type: "string", + enum: Object.values(EmbeddableEntityType), + }, + fields: { + description: "The fields of the entity that should be embedded.", + type: "array", + minItems: 1, + uniqueItems: true, + items: { + type: "string", + enum: Object.values(EmbeddableField), + }, + }, + }, + required: ["modelId", "entityId", "entityType", "fields"], +}; + +// A dedicated Ajv instance keeps the embeddings lambda bundle self-contained, +// so it does not need to pull in the full API validator with every api-specification schema. +const ajvInstance = new Ajv({ validateSchema: true, allErrors: true, strict: true }); +addFormats(ajvInstance); + +export const validateEmbeddingQueueJob: ValidateFunction = ajvInstance.compile(EmbeddingQueueJobSchema); + +export default EmbeddingQueueJobSchema; diff --git a/backend/src/index.ts b/backend/src/index.ts index 631d9ff6a..7784948b5 100644 --- a/backend/src/index.ts +++ b/backend/src/index.ts @@ -1,6 +1,7 @@ import { APIGatewayProxyEvent, Handler } from "aws-lambda"; import { handler as InfoHandler } from "applicationInfo"; import { handler as ModelHandler } from "modelInfo"; +import { handler as ModelEmbeddingProcessesHandler } from "modelInfo/embeddingProcesses"; import { handler as ImportHandler } from "import"; import { handler as OccupationGroupHandler } from "esco/occupationGroup/index"; import { handler as OccupationHandler } from "esco/occupations"; @@ -65,6 +66,8 @@ export const handleRouteEvent = async (event: APIGatewayProxyEvent) => { return ImportHandler(event); } else if (path === Routes.EXPORT_ROUTE) { return ExportHandler(event); + } else if (pathToRegexp([Routes.MODEL_EMBEDDING_PROCESSES_ROUTE]).regexp.test(path)) { + return ModelEmbeddingProcessesHandler(event); } else if ( pathToRegexp([ Routes.OCCUPATION_GROUPS_ROUTE, diff --git a/backend/src/modelInfo/embeddingProcesses/POST/index.integration.test.ts b/backend/src/modelInfo/embeddingProcesses/POST/index.integration.test.ts new file mode 100644 index 000000000..32faf4585 --- /dev/null +++ b/backend/src/modelInfo/embeddingProcesses/POST/index.integration.test.ts @@ -0,0 +1,302 @@ +// Suppress chatty console during the tests +import "_test_utilities/consoleMock"; + +// Mock only the AWS SQS client so that no message is ever pushed to a real queue. +// The rest of the module (SendMessageCommand, etc.) stays real so that the actual commands can be asserted. +const mockSQSClientSend = jest.fn().mockResolvedValue({}); +jest.mock("@aws-sdk/client-sqs", () => { + const actual = jest.requireActual("@aws-sdk/client-sqs"); + return { + ...actual, + SQSClient: jest.fn().mockImplementation(() => ({ + send: mockSQSClientSend, + })), + }; +}); + +import Ajv, { ValidateFunction } from "ajv"; +import addFormats from "ajv-formats"; +import { randomUUID } from "crypto"; +import { Connection } from "mongoose"; +import { SendMessageCommand } from "@aws-sdk/client-sqs"; + +import LocaleAPISpecs from "api-specifications/locale"; +import ModelInfoAPISpecs from "api-specifications/modelInfo"; +import EmbeddingsAPISpecs from "api-specifications/embeddings"; + +import { handler as embeddingProcessesHandler } from "modelInfo/embeddingProcesses"; +import { HTTP_VERBS, StatusCodes } from "server/httpUtils"; +import { initOnce } from "server/init"; +import { getConnectionManager } from "server/connection/connectionManager"; +import { getEmbeddingsQueueUrl } from "server/config/config"; +import { getRepositoryRegistry } from "server/repositoryRegistry/repositoryRegistry"; +import { getTestConfiguration } from "_test_utilities/getTestConfiguration"; +import { getTestString } from "_test_utilities/getMockRandomData"; +import { usersRequestContext } from "_test_utilities/dataModel"; +import { + getSimpleNewESCOOccupationSpec, + getSimpleNewISCOGroupSpec, + getSimpleNewSkillGroupSpec, + getSimpleNewSkillSpec, +} from "esco/_test_utilities/getNewSpecs"; +import { MongooseModelName } from "esco/common/mongooseModelNames"; +import { ModelName as ModelInfoModelName } from "modelInfo/modelInfoModel"; +import { ModelName as EmbeddingProcessStateModelName } from "embeddings/embeddingProcessState/embeddingProcessStateModel"; +import { IModelInfo } from "modelInfo/modelInfo.types"; +import { EmbeddableEntityType, EmbeddableField, IGenerateEmbeddingTask } from "embeddings/service/types"; +import { validateEmbeddingQueueJob } from "embeddings/specs/queueJob.schema"; + +const givenEntityCountPerType = 10; +const givenEmbeddingServiceId = EmbeddingsAPISpecs.Constants.EmbeddingServiceIds[0]; + +/** + * The fields that the embedding process is expected to embed for each entity type. + */ +const expectedFieldsByEntityType: Record = { + [EmbeddableEntityType.Skill]: [ + EmbeddableField.preferredLabel, + EmbeddableField.description, + EmbeddableField.altLabels, + ], + [EmbeddableEntityType.SkillGroup]: [ + EmbeddableField.preferredLabel, + EmbeddableField.description, + EmbeddableField.altLabels, + EmbeddableField.scopeNote, + ], + [EmbeddableEntityType.Occupation]: [ + EmbeddableField.preferredLabel, + EmbeddableField.description, + EmbeddableField.altLabels, + ], + [EmbeddableEntityType.OccupationGroup]: [ + EmbeddableField.preferredLabel, + EmbeddableField.description, + EmbeddableField.altLabels, + ], +}; + +async function createModelInDB(released: boolean): Promise { + const modelInfoRepository = getRepositoryRegistry().modelInfo; + const newModel = await modelInfoRepository.create({ + name: getTestString(ModelInfoAPISpecs.Constants.NAME_MAX_LENGTH), + locale: { + UUID: randomUUID(), + name: getTestString(LocaleAPISpecs.Constants.NAME_MAX_LENGTH), + shortCode: getTestString(LocaleAPISpecs.Constants.LOCALE_SHORTCODE_MAX_LENGTH), + }, + description: getTestString(ModelInfoAPISpecs.Constants.DESCRIPTION_MAX_LENGTH), + license: getTestString(ModelInfoAPISpecs.Constants.LICENSE_MAX_LENGTH), + UUIDHistory: [randomUUID()], + }); + if (released) { + // the repository always creates models as not released, so release it directly in the DB + await modelInfoRepository.Model.findByIdAndUpdate(newModel.id, { released: true }); + } + return { ...newModel, released }; +} + +async function createEntitiesInDB(modelId: string) { + const repositoryRegistry = getRepositoryRegistry(); + const skills = await repositoryRegistry.skill.createMany( + Array.from({ length: givenEntityCountPerType }, (_, index) => getSimpleNewSkillSpec(modelId, `Skill ${index + 1}`)) + ); + const skillGroups = await repositoryRegistry.skillGroup.createMany( + Array.from({ length: givenEntityCountPerType }, (_, index) => + getSimpleNewSkillGroupSpec(modelId, `Skill Group ${index + 1}`) + ) + ); + const occupations = await repositoryRegistry.occupation.createMany( + Array.from({ length: givenEntityCountPerType }, (_, index) => + getSimpleNewESCOOccupationSpec(modelId, `Occupation ${index + 1}`) + ) + ); + const occupationGroups = await repositoryRegistry.OccupationGroup.createMany( + Array.from({ length: givenEntityCountPerType }, (_, index) => + getSimpleNewISCOGroupSpec(modelId, `Occupation Group ${index + 1}`) + ) + ); + return { skills, skillGroups, occupations, occupationGroups }; +} + +function getEvent(modelId: string, requestContext: object) { + return { + httpMethod: HTTP_VERBS.POST, + path: `/models/${modelId}/embedding-processes`, + body: JSON.stringify({ embeddingServiceId: givenEmbeddingServiceId }), + headers: { + "Content-Type": "application/json", + }, + requestContext, + }; +} + +describe("Test for the POST model embedding processes handler with a DB", () => { + // set up the ajv validate POST response function + const ajv = new Ajv({ + validateSchema: true, + strict: true, + allErrors: true, + }); + addFormats(ajv); + ajv.addSchema(ModelInfoAPISpecs.ModelInfo.EmbeddingProcessStates.POST.Schemas.Response.Payload); + const validatePOSTResponse: ValidateFunction = ajv.getSchema( + ModelInfoAPISpecs.ModelInfo.EmbeddingProcessStates.POST.Schemas.Response.Payload.$id as string + ) as ValidateFunction; + // --- + + let dbConnection: Connection | undefined; + beforeAll(async () => { + // using the in-memory mongodb instance that is started up with @shelf/jest-mongodb + const config = getTestConfiguration("ModelEmbeddingProcessesHandlerTestDB"); + const configModule = await import("server/config/config"); + jest.spyOn(configModule, "readEnvironmentConfiguration").mockReturnValue(config); + await initOnce(); + dbConnection = getConnectionManager().getCurrentDBConnection(); + }); + + afterAll(async () => { + if (dbConnection) { + await dbConnection.dropDatabase(); + await dbConnection.close(); + } + }); + + beforeEach(async () => { + mockSQSClientSend.mockClear(); + if (dbConnection) { + // delete all documents in the DB + await Promise.all( + [ + ModelInfoModelName, + EmbeddingProcessStateModelName, + MongooseModelName.Skill, + MongooseModelName.SkillGroup, + MongooseModelName.Occupation, + MongooseModelName.OccupationGroup, + ].map((modelName) => dbConnection!.models[modelName].deleteMany({})) + ); + } + }); + + test("POST should respond with the ACCEPTED status code and push one valid SQS message for every entity of the model", async () => { + // GIVEN a released model in the DB + const givenModel = await createModelInDB(true); + // AND the model has 10 skills, 10 skill groups, 10 occupations and 10 occupation groups in the DB + const givenEntities = await createEntitiesInDB(givenModel.id); + // guard to ensure that all the entities were actually created in the DB + expect(givenEntities.skills).toHaveLength(givenEntityCountPerType); + expect(givenEntities.skillGroups).toHaveLength(givenEntityCountPerType); + expect(givenEntities.occupations).toHaveLength(givenEntityCountPerType); + expect(givenEntities.occupationGroups).toHaveLength(givenEntityCountPerType); + // AND a valid request from a model manager to trigger the embedding process of the model + const givenEvent = getEvent(givenModel.id, usersRequestContext.MODEL_MANAGER); + + // WHEN the handler is invoked with the given event + // @ts-ignore + const actualResponse = await embeddingProcessesHandler(givenEvent); + + // THEN expect the handler to respond with the ACCEPTED status code + expect(actualResponse.statusCode).toEqual(StatusCodes.ACCEPTED); + // AND the response passes the JSON Schema validation + const actualPayload = JSON.parse(actualResponse.body); + validatePOSTResponse(actualPayload); + expect(validatePOSTResponse.errors).toBeNull(); + // AND the response to refer to an IN_PROGRESS embedding process for the given model that counts all the entities + const expectedTotalDocuments = 4 * givenEntityCountPerType; + expect(actualPayload).toMatchObject({ + modelId: givenModel.id, + status: ModelInfoAPISpecs.ModelInfo.EmbeddingProcessStates.Enums.Status.IN_PROGRESS, + embeddingServiceId: givenEmbeddingServiceId, + totalDocuments: expectedTotalDocuments, + }); + + // AND expect the SQS client to have been called with a SendMessageCommand for every entity of the model + expect(mockSQSClientSend).toHaveBeenCalledTimes(expectedTotalDocuments); + const actualCommands = mockSQSClientSend.mock.calls.map(([command]) => command); + for (const actualCommand of actualCommands) { + expect(actualCommand).toBeInstanceOf(SendMessageCommand); + expect(actualCommand.input.QueueUrl).toEqual(getEmbeddingsQueueUrl()); + } + // AND every message body to be a valid embedding queue job + const actualTasks: IGenerateEmbeddingTask[] = actualCommands.map((command) => + JSON.parse(command.input.MessageBody as string) + ); + for (const actualTask of actualTasks) { + expect(validateEmbeddingQueueJob(actualTask)).toBe(true); + } + // AND the queued jobs to reference exactly the entities in the DB with the expected fields for each entity type + const expectedTasks: IGenerateEmbeddingTask[] = [ + ...givenEntities.skills.map((skill) => ({ + modelId: givenModel.id, + entityId: skill.id, + entityType: EmbeddableEntityType.Skill, + fields: expectedFieldsByEntityType[EmbeddableEntityType.Skill], + })), + ...givenEntities.skillGroups.map((skillGroup) => ({ + modelId: givenModel.id, + entityId: skillGroup.id, + entityType: EmbeddableEntityType.SkillGroup, + fields: expectedFieldsByEntityType[EmbeddableEntityType.SkillGroup], + })), + ...givenEntities.occupations.map((occupation) => ({ + modelId: givenModel.id, + entityId: occupation.id, + entityType: EmbeddableEntityType.Occupation, + fields: expectedFieldsByEntityType[EmbeddableEntityType.Occupation], + })), + ...givenEntities.occupationGroups.map((occupationGroup) => ({ + modelId: givenModel.id, + entityId: occupationGroup.id, + entityType: EmbeddableEntityType.OccupationGroup, + fields: expectedFieldsByEntityType[EmbeddableEntityType.OccupationGroup], + })), + ]; + expect(actualTasks).toHaveLength(expectedTasks.length); + expect(actualTasks).toEqual(expect.arrayContaining(expectedTasks)); + + // AND the embedding process state to have been persisted in the DB + const actualProcessState = await getRepositoryRegistry().embeddingProcessState.findById(actualPayload.id); + expect(actualProcessState).toMatchObject({ + modelId: givenModel.id, + status: ModelInfoAPISpecs.ModelInfo.EmbeddingProcessStates.Enums.Status.IN_PROGRESS, + embeddingServiceId: givenEmbeddingServiceId, + totalDocuments: expectedTotalDocuments, + }); + }); + + test("POST should respond with the FORBIDDEN status code if the user is not a model manager", async () => { + // GIVEN a released model in the DB + const givenModel = await createModelInDB(true); + // AND a valid request from a user that is not a model manager + const givenEvent = getEvent(givenModel.id, usersRequestContext.REGISTED_USER); + + // WHEN the handler is invoked with the given event + // @ts-ignore + const actualResponse = await embeddingProcessesHandler(givenEvent); + + // THEN expect the handler to respond with the FORBIDDEN status code + expect(actualResponse.statusCode).toEqual(StatusCodes.FORBIDDEN); + // AND expect no message to have been pushed to the SQS queue + expect(mockSQSClientSend).not.toHaveBeenCalled(); + }); + + test("POST should respond with the BAD_REQUEST status code and not push any SQS message when the model is not released", async () => { + // GIVEN a model in the DB that is not released + const givenModel = await createModelInDB(false); + // AND a valid request from a model manager to trigger the embedding process of the model + const givenEvent = getEvent(givenModel.id, usersRequestContext.MODEL_MANAGER); + + // WHEN the handler is invoked with the given event + // @ts-ignore + const actualResponse = await embeddingProcessesHandler(givenEvent); + + // THEN expect the handler to respond with the BAD_REQUEST status code and the MODEL_NOT_RELEASED error code + expect(actualResponse.statusCode).toEqual(StatusCodes.BAD_REQUEST); + expect(JSON.parse(actualResponse.body).errorCode).toEqual( + ModelInfoAPISpecs.ModelInfo.EmbeddingProcessStates.POST.Enums.Response.Status400.ErrorCodes.MODEL_NOT_RELEASED + ); + // AND expect no message to have been pushed to the SQS queue + expect(mockSQSClientSend).not.toHaveBeenCalled(); + }); +}); diff --git a/backend/src/modelInfo/embeddingProcesses/POST/index.test.ts b/backend/src/modelInfo/embeddingProcesses/POST/index.test.ts new file mode 100644 index 000000000..3e951c3db --- /dev/null +++ b/backend/src/modelInfo/embeddingProcesses/POST/index.test.ts @@ -0,0 +1,237 @@ +// silence chatty console +import "_test_utilities/consoleMock"; + +// Mock the authorizer so the role check is a no-op +jest.mock("auth/authorizer", () => ({ + checkRole: jest.fn().mockImplementation(() => true), + RoleRequired: jest.fn().mockImplementation(() => { + return (target: unknown, propertyKey: string, descriptor: PropertyDescriptor) => { + return descriptor; + }; + }), +})); + +// Mock the service registry so we can control the embedding process service +const mockTriggerEmbeddingProcess = jest.fn(); +jest.mock("server/serviceRegistry/serviceRegistry", () => ({ + getServiceRegistry: jest.fn().mockReturnValue({ + embeddingProcess: { triggerEmbeddingProcess: mockTriggerEmbeddingProcess }, + }), +})); + +import { POSTModelEmbeddingProcessesHandler } from "./index"; +import { transformEmbeddingProcessState } from "./transform"; +import ModelInfoApiSpecs from "api-specifications/modelInfo"; +import EmbeddingsAPISpecs from "api-specifications/embeddings"; +import ErrorAPISpecs from "api-specifications/error"; +import { HTTP_VERBS, StatusCodes } from "server/httpUtils"; +import { APIGatewayProxyEvent } from "aws-lambda"; +import { getMockStringId } from "_test_utilities/mockMongoId"; +import { IEmbeddingProcessState } from "embeddings/embeddingProcessState/embeddingProcessState.types"; +import { + EmbeddingProcessAlreadyRunningError, + ModelNotFoundError, + ModelNotReleasedError, +} from "embeddings/embeddingProcess/errors"; + +const givenModelId = getMockStringId(1); +const givenEmbeddingServiceId = EmbeddingsAPISpecs.Constants.EmbeddingServiceIds[0]; + +function getEvent(options: { body?: string; modelId?: string | null; contentType?: string }): APIGatewayProxyEvent { + const modelId = options.modelId === undefined ? givenModelId : options.modelId; + const path = modelId === null ? `/models/embedding-processes` : `/models/${modelId}/embedding-processes`; + return { + httpMethod: HTTP_VERBS.POST, + path, + body: options.body ?? JSON.stringify({ embeddingServiceId: givenEmbeddingServiceId }), + headers: { + "Content-Type": options.contentType ?? "application/json", + }, + } as never; +} + +function getMockEmbeddingProcessState(): IEmbeddingProcessState { + return { + id: getMockStringId(10), + modelId: givenModelId, + status: ModelInfoApiSpecs.ModelInfo.EmbeddingProcessStates.Enums.Status.IN_PROGRESS, + embeddingServiceId: givenEmbeddingServiceId, + totalDocuments: 5, + errorCounts: 0, + warningCounts: 0, + completedDocuments: 0, + createdAt: new Date(), + updatedAt: new Date(), + }; +} + +describe("Test the POSTModelEmbeddingProcessesHandler", () => { + beforeEach(() => { + jest.clearAllMocks(); + }); + + test("should respond with ACCEPTED and the created embedding process state on success", async () => { + // GIVEN a valid event + const givenEvent = getEvent({}); + // AND the service will trigger the embedding process successfully + const givenProcessState = getMockEmbeddingProcessState(); + mockTriggerEmbeddingProcess.mockResolvedValueOnce(givenProcessState); + + // WHEN the handler is invoked with the given event + const actualResponse = await new POSTModelEmbeddingProcessesHandler().handle(givenEvent); + + // THEN expect the service to be called with the model id and the embedding model id + expect(mockTriggerEmbeddingProcess).toHaveBeenCalledWith(givenModelId, givenEmbeddingServiceId); + // AND expect a 202 response with the transformed embedding process state + expect(actualResponse.statusCode).toEqual(StatusCodes.ACCEPTED); + expect(JSON.parse(actualResponse.body)).toEqual(transformEmbeddingProcessState(givenProcessState)); + }); + + test("should respond with UNSUPPORTED_MEDIA_TYPE when the content type is not json", async () => { + // GIVEN an event with an invalid content type + const givenEvent = getEvent({ contentType: "text/html" }); + + // WHEN the handler is invoked with the given event + const actualResponse = await new POSTModelEmbeddingProcessesHandler().handle(givenEvent); + + // THEN expect a 415 response + expect(actualResponse.statusCode).toEqual(StatusCodes.UNSUPPORTED_MEDIA_TYPE); + // AND expect the service to not be called + expect(mockTriggerEmbeddingProcess).not.toHaveBeenCalled(); + }); + + test("should respond with TOO_LARGE_PAYLOAD when the payload is too large", async () => { + // GIVEN an event with a payload that is too large + const givenEvent = getEvent({ + body: "a".repeat(ModelInfoApiSpecs.ModelInfo.EmbeddingProcessStates.POST.Constants.MAX_PAYLOAD_LENGTH + 1), + }); + + // WHEN the handler is invoked with the given event + const actualResponse = await new POSTModelEmbeddingProcessesHandler().handle(givenEvent); + + // THEN expect a 413 response + expect(actualResponse.statusCode).toEqual(StatusCodes.TOO_LARGE_PAYLOAD); + // AND expect the service to not be called + expect(mockTriggerEmbeddingProcess).not.toHaveBeenCalled(); + }); + + test("should respond with BAD_REQUEST when the modelId is missing in the path", async () => { + // GIVEN an event whose path does not contain a modelId + const givenEvent = getEvent({ modelId: null }); + + // WHEN the handler is invoked with the given event + const actualResponse = await new POSTModelEmbeddingProcessesHandler().handle(givenEvent); + + // THEN expect a 400 response with the INVALID_JSON_SCHEMA error code + expect(actualResponse.statusCode).toEqual(StatusCodes.BAD_REQUEST); + expect(JSON.parse(actualResponse.body).errorCode).toEqual(ErrorAPISpecs.Constants.ErrorCodes.INVALID_JSON_SCHEMA); + // AND expect the service to not be called + expect(mockTriggerEmbeddingProcess).not.toHaveBeenCalled(); + }); + + test("should respond with BAD_REQUEST when the body is malformed", async () => { + // GIVEN an event with a malformed body + const givenEvent = getEvent({ body: "{ not json" }); + + // WHEN the handler is invoked with the given event + const actualResponse = await new POSTModelEmbeddingProcessesHandler().handle(givenEvent); + + // THEN expect a 400 response with the MALFORMED_BODY error code + expect(actualResponse.statusCode).toEqual(StatusCodes.BAD_REQUEST); + expect(JSON.parse(actualResponse.body).errorCode).toEqual(ErrorAPISpecs.Constants.ErrorCodes.MALFORMED_BODY); + // AND expect the service to not be called + expect(mockTriggerEmbeddingProcess).not.toHaveBeenCalled(); + }); + + test("should respond with BAD_REQUEST when the body does not conform to the schema", async () => { + // GIVEN an event with a body that does not conform to the schema + const givenEvent = getEvent({ body: JSON.stringify({ foo: "bar" }) }); + + // WHEN the handler is invoked with the given event + const actualResponse = await new POSTModelEmbeddingProcessesHandler().handle(givenEvent); + + // THEN expect a 400 response with the INVALID_JSON_SCHEMA error code + expect(actualResponse.statusCode).toEqual(StatusCodes.BAD_REQUEST); + expect(JSON.parse(actualResponse.body).errorCode).toEqual(ErrorAPISpecs.Constants.ErrorCodes.INVALID_JSON_SCHEMA); + // AND expect the service to not be called + expect(mockTriggerEmbeddingProcess).not.toHaveBeenCalled(); + }); + + test("should respond with BAD_REQUEST when the embeddingServiceId is not supported", async () => { + // GIVEN an event with an unsupported embeddingServiceId + const givenEvent = getEvent({ body: JSON.stringify({ embeddingServiceId: "unsupported-model" }) }); + + // WHEN the handler is invoked with the given event + const actualResponse = await new POSTModelEmbeddingProcessesHandler().handle(givenEvent); + + // THEN expect a 400 response with the INVALID_JSON_SCHEMA error code + expect(actualResponse.statusCode).toEqual(StatusCodes.BAD_REQUEST); + expect(JSON.parse(actualResponse.body).errorCode).toEqual(ErrorAPISpecs.Constants.ErrorCodes.INVALID_JSON_SCHEMA); + // AND expect the service to not be called + expect(mockTriggerEmbeddingProcess).not.toHaveBeenCalled(); + }); + + test("should respond with NOT_FOUND when the model does not exist", async () => { + // GIVEN a valid event + const givenEvent = getEvent({}); + // AND the service throws a ModelNotFoundError + mockTriggerEmbeddingProcess.mockRejectedValueOnce(new ModelNotFoundError(givenModelId)); + + // WHEN the handler is invoked with the given event + const actualResponse = await new POSTModelEmbeddingProcessesHandler().handle(givenEvent); + + // THEN expect a 404 response with the NOT_FOUND error code + expect(actualResponse.statusCode).toEqual(StatusCodes.NOT_FOUND); + expect(JSON.parse(actualResponse.body).errorCode).toEqual(ErrorAPISpecs.Constants.ErrorCodes.NOT_FOUND); + }); + + test("should respond with BAD_REQUEST when the model is not released", async () => { + // GIVEN a valid event + const givenEvent = getEvent({}); + // AND the service throws a ModelNotReleasedError + mockTriggerEmbeddingProcess.mockRejectedValueOnce(new ModelNotReleasedError(givenModelId)); + + // WHEN the handler is invoked with the given event + const actualResponse = await new POSTModelEmbeddingProcessesHandler().handle(givenEvent); + + // THEN expect a 400 response with the MODEL_NOT_RELEASED error code + expect(actualResponse.statusCode).toEqual(StatusCodes.BAD_REQUEST); + expect(JSON.parse(actualResponse.body).errorCode).toEqual( + ModelInfoApiSpecs.ModelInfo.EmbeddingProcessStates.POST.Enums.Response.Status400.ErrorCodes.MODEL_NOT_RELEASED + ); + }); + + test("should respond with CONFLICT when an embedding process is already running", async () => { + // GIVEN a valid event + const givenEvent = getEvent({}); + // AND the service throws an EmbeddingProcessAlreadyRunningError + mockTriggerEmbeddingProcess.mockRejectedValueOnce(new EmbeddingProcessAlreadyRunningError(givenModelId)); + + // WHEN the handler is invoked with the given event + const actualResponse = await new POSTModelEmbeddingProcessesHandler().handle(givenEvent); + + // THEN expect a 409 response with the EMBEDDING_PROCESS_ALREADY_RUNNING error code + expect(actualResponse.statusCode).toEqual(StatusCodes.CONFLICT); + expect(JSON.parse(actualResponse.body).errorCode).toEqual( + ModelInfoApiSpecs.ModelInfo.EmbeddingProcessStates.POST.Enums.Response.Status409.ErrorCodes + .EMBEDDING_PROCESS_ALREADY_RUNNING + ); + }); + + test("should respond with INTERNAL_SERVER_ERROR when the service fails unexpectedly", async () => { + // GIVEN a valid event + const givenEvent = getEvent({}); + // AND the service throws an unexpected error + mockTriggerEmbeddingProcess.mockRejectedValueOnce(new Error("unexpected")); + + // WHEN the handler is invoked with the given event + const actualResponse = await new POSTModelEmbeddingProcessesHandler().handle(givenEvent); + + // THEN expect a 500 response with the FAILED_TO_TRIGGER_EMBEDDING_PROCESS error code + expect(actualResponse.statusCode).toEqual(StatusCodes.INTERNAL_SERVER_ERROR); + expect(JSON.parse(actualResponse.body).errorCode).toEqual( + ModelInfoApiSpecs.ModelInfo.EmbeddingProcessStates.POST.Enums.Response.Status500.ErrorCodes + .FAILED_TO_TRIGGER_EMBEDDING_PROCESS + ); + }); +}); diff --git a/backend/src/modelInfo/embeddingProcesses/POST/index.ts b/backend/src/modelInfo/embeddingProcesses/POST/index.ts new file mode 100644 index 000000000..b98860cd4 --- /dev/null +++ b/backend/src/modelInfo/embeddingProcesses/POST/index.ts @@ -0,0 +1,173 @@ +import { APIGatewayProxyEvent, APIGatewayProxyResult } from "aws-lambda"; +import { ValidateFunction } from "ajv"; +import ModelInfoApiSpecs from "api-specifications/modelInfo"; +import ErrorAPISpecs from "api-specifications/error"; +import AuthAPISpecs from "api-specifications/auth"; +import { errorResponse, responseJSON, StatusCodes, STD_ERRORS_RESPONSES } from "server/httpUtils"; +import { ajvInstance, ParseValidationError } from "validator"; +import { RoleRequired } from "auth/authorizer"; +import { getServiceRegistry } from "server/serviceRegistry/serviceRegistry"; +import { parsePath } from "common/parsePath/parsePath"; +import { Routes } from "routes.constant"; +import { + EmbeddingProcessAlreadyRunningError, + ModelNotFoundError, + ModelNotReleasedError, +} from "embeddings/embeddingProcess/errors"; +import { transformEmbeddingProcessState } from "./transform"; + +/** + * @openapi + * + * /models/{modelId}/embedding-processes: + * post: + * operationId: PostModelEmbeddingProcess + * tags: + * - embeddings + * - model + * summary: Trigger the generation of the embeddings of a model. + * description: | + * Asynchronously trigger the generation of the embeddings for all the entities of a released model. + * The entities are pushed to the embeddings queue and processed in the background. + * security: + * - api_key: [] + * - jwt_auth: [] + * parameters: + * - in: path + * name: modelId + * required: true + * schema: + * type: string + * requestBody: + * content: + * application/json: + * schema: + * $ref: '#/components/schemas/EmbeddingsRequestSchemaPOST' + * required: true + * responses: + * '202': + * description: The embedding process was successfully triggered. + * content: + * application/json: + * schema: + * $ref: '#/components/schemas/EmbeddingsResponseSchemaPOST' + * '400': + * description: | + * Failed to trigger the embedding process because the model is not released. Additional information can be found in the response body. + * content: + * application/json: + * schema: + * $ref: '#/components/schemas/POSTEmbeddingProcess400ErrorSchema' + * '401': + * $ref: '#/components/responses/UnAuthorizedResponse' + * '403': + * $ref: '#/components/responses/ForbiddenResponse' + * '404': + * description: The model was not found. + * content: + * application/json: + * schema: + * $ref: '#/components/schemas/POSTEmbeddingProcess404ErrorSchema' + * '409': + * description: An embedding process is already running for the model. + * content: + * application/json: + * schema: + * $ref: '#/components/schemas/POSTEmbeddingProcess409ErrorSchema' + * '415': + * $ref: '#/components/responses/AcceptOnlyJSONResponse' + * '500': + * $ref: '#/components/responses/InternalServerErrorResponse' + */ +export class POSTModelEmbeddingProcessesHandler { + @RoleRequired(AuthAPISpecs.Enums.TabiyaRoles.MODEL_MANAGER) // Applying role-based access control + async handle(event: APIGatewayProxyEvent): Promise { + if (!event.headers["Content-Type"]?.includes("application/json")) { + return STD_ERRORS_RESPONSES.UNSUPPORTED_MEDIA_TYPE_ERROR; + } + + if ( + event.body?.length && + event.body.length > ModelInfoApiSpecs.ModelInfo.EmbeddingProcessStates.POST.Constants.MAX_PAYLOAD_LENGTH + ) { + return STD_ERRORS_RESPONSES.TOO_LARGE_PAYLOAD_ERROR( + `Expected maximum length is ${ModelInfoApiSpecs.ModelInfo.EmbeddingProcessStates.POST.Constants.MAX_PAYLOAD_LENGTH}` + ); + } + + const { modelId } = parsePath<{ modelId?: string }>(Routes.MODEL_EMBEDDING_PROCESSES_ROUTE, event.path); + if (!modelId) { + return errorResponse( + StatusCodes.BAD_REQUEST, + ErrorAPISpecs.Constants.ErrorCodes.INVALID_JSON_SCHEMA, + "modelId is missing in the path", + JSON.stringify({ path: event.path }) + ); + } + + let payload: ModelInfoApiSpecs.ModelInfo.EmbeddingProcessStates.POST.Types.Request.Payload; + try { + payload = JSON.parse(event.body as string); + } catch (error: unknown) { + return errorResponse( + StatusCodes.BAD_REQUEST, + ErrorAPISpecs.Constants.ErrorCodes.MALFORMED_BODY, + "Payload is malformed, it should be a valid json", + (error as Error).message + ); + } + + const validateFunction = ajvInstance.getSchema( + ModelInfoApiSpecs.ModelInfo.EmbeddingProcessStates.POST.Schemas.Request.Payload.$id as string + ) as ValidateFunction; + const isValid = validateFunction(payload); + if (!isValid) { + const errorDetail = ParseValidationError(validateFunction.errors); + return STD_ERRORS_RESPONSES.INVALID_JSON_SCHEMA_ERROR(errorDetail); + } + + try { + const embeddingProcessState = await getServiceRegistry().embeddingProcess.triggerEmbeddingProcess( + modelId, + payload.embeddingServiceId + ); + return responseJSON(StatusCodes.ACCEPTED, transformEmbeddingProcessState(embeddingProcessState)); + } catch (error: unknown) { + if (error instanceof ModelNotFoundError) { + return errorResponse( + StatusCodes.NOT_FOUND, + ErrorAPISpecs.Constants.ErrorCodes.NOT_FOUND, + "Model not found", + `No model found with id: ${modelId}` + ); + } + if (error instanceof ModelNotReleasedError) { + return errorResponse( + StatusCodes.BAD_REQUEST, + ModelInfoApiSpecs.ModelInfo.EmbeddingProcessStates.POST.Enums.Response.Status400.ErrorCodes + .MODEL_NOT_RELEASED, + "The model is not released", + "Embeddings can only be generated for released models" + ); + } + if (error instanceof EmbeddingProcessAlreadyRunningError) { + return errorResponse( + StatusCodes.CONFLICT, + ModelInfoApiSpecs.ModelInfo.EmbeddingProcessStates.POST.Enums.Response.Status409.ErrorCodes + .EMBEDDING_PROCESS_ALREADY_RUNNING, + "An embedding process is already running for this model", + "Please wait for the current embedding process to complete before triggering a new one" + ); + } + console.error(new Error("Failed to trigger the embedding process", { cause: error })); + // Do not surface the error message to the user as it can contain sensitive information. + return errorResponse( + StatusCodes.INTERNAL_SERVER_ERROR, + ModelInfoApiSpecs.ModelInfo.EmbeddingProcessStates.POST.Enums.Response.Status500.ErrorCodes + .FAILED_TO_TRIGGER_EMBEDDING_PROCESS, + "Failed to trigger the embedding process", + "" + ); + } + } +} diff --git a/backend/src/modelInfo/embeddingProcesses/POST/transform.test.ts b/backend/src/modelInfo/embeddingProcesses/POST/transform.test.ts new file mode 100644 index 000000000..0ffb9d4c4 --- /dev/null +++ b/backend/src/modelInfo/embeddingProcesses/POST/transform.test.ts @@ -0,0 +1,72 @@ +// silence chatty console +import "_test_utilities/consoleMock"; + +import ModelInfoApiSpecs from "api-specifications/modelInfo"; +import EmbeddingsAPISpecs from "api-specifications/embeddings"; +import { transformEmbeddingProcessState } from "./transform"; +import { IEmbeddingProcessState } from "embeddings/embeddingProcessState/embeddingProcessState.types"; +import { getMockStringId } from "_test_utilities/mockMongoId"; + +describe("Test transformEmbeddingProcessState", () => { + test("should transform an embedding process state into the API response payload", () => { + // GIVEN an embedding process state + const givenCreatedAt = new Date("2024-01-01T00:00:00.000Z"); + const givenUpdatedAt = new Date("2024-01-02T00:00:00.000Z"); + const givenEmbeddingProcessState: IEmbeddingProcessState = { + id: getMockStringId(1), + modelId: getMockStringId(2), + status: ModelInfoApiSpecs.ModelInfo.EmbeddingProcessStates.Enums.Status.IN_PROGRESS, + embeddingServiceId: "gemini$$models/gemini-embedding-2", + totalDocuments: 100, + errorCounts: 1, + warningCounts: 2, + completedDocuments: 50, + createdAt: givenCreatedAt, + updatedAt: givenUpdatedAt, + }; + + // WHEN transforming the embedding process state + const actualPayload = transformEmbeddingProcessState(givenEmbeddingProcessState); + + // THEN expect the payload to match the embedding process state with the dates as ISO strings + const expectedPayload: ModelInfoApiSpecs.ModelInfo.EmbeddingProcessStates.POST.Types.Response.Payload = { + id: givenEmbeddingProcessState.id, + modelId: givenEmbeddingProcessState.modelId, + status: givenEmbeddingProcessState.status, + embeddingServiceId: givenEmbeddingProcessState.embeddingServiceId, + totalDocuments: givenEmbeddingProcessState.totalDocuments, + errorCounts: givenEmbeddingProcessState.errorCounts, + warningCounts: givenEmbeddingProcessState.warningCounts, + completedDocuments: givenEmbeddingProcessState.completedDocuments, + createdAt: givenCreatedAt.toISOString(), + updatedAt: givenUpdatedAt.toISOString(), + }; + expect(actualPayload).toEqual(expectedPayload); + }); + + test("the transformed payload should conform to the response schema", () => { + // GIVEN an embedding process state + const givenEmbeddingProcessState: IEmbeddingProcessState = { + id: getMockStringId(1), + modelId: getMockStringId(2), + status: ModelInfoApiSpecs.ModelInfo.EmbeddingProcessStates.Enums.Status.PENDING, + embeddingServiceId: EmbeddingsAPISpecs.Constants.EmbeddingServiceIds[0], + totalDocuments: 0, + errorCounts: 0, + warningCounts: 0, + completedDocuments: 0, + createdAt: new Date(), + updatedAt: new Date(), + }; + + // WHEN transforming the embedding process state + const actualPayload = transformEmbeddingProcessState(givenEmbeddingProcessState); + + // THEN expect the payload to have all the required properties + const requiredProperties = ModelInfoApiSpecs.ModelInfo.EmbeddingProcessStates.POST.Schemas.Response.Payload + .required as string[]; + requiredProperties.forEach((property) => { + expect(actualPayload).toHaveProperty(property); + }); + }); +}); diff --git a/backend/src/modelInfo/embeddingProcesses/POST/transform.ts b/backend/src/modelInfo/embeddingProcesses/POST/transform.ts new file mode 100644 index 000000000..af8c8dcb9 --- /dev/null +++ b/backend/src/modelInfo/embeddingProcesses/POST/transform.ts @@ -0,0 +1,22 @@ +import ModelInfoApiSpecs from "api-specifications/modelInfo"; +import { IEmbeddingProcessState } from "embeddings/embeddingProcessState/embeddingProcessState.types"; + +/** + * Transforms an EmbeddingProcessState (as returned by the repository) into the API response payload. + */ +export function transformEmbeddingProcessState( + embeddingProcessState: IEmbeddingProcessState +): ModelInfoApiSpecs.ModelInfo.EmbeddingProcessStates.POST.Types.Response.Payload { + return { + id: embeddingProcessState.id, + modelId: embeddingProcessState.modelId, + status: embeddingProcessState.status, + embeddingServiceId: embeddingProcessState.embeddingServiceId, + totalDocuments: embeddingProcessState.totalDocuments, + errorCounts: embeddingProcessState.errorCounts, + warningCounts: embeddingProcessState.warningCounts, + completedDocuments: embeddingProcessState.completedDocuments, + createdAt: embeddingProcessState.createdAt.toISOString(), + updatedAt: embeddingProcessState.updatedAt.toISOString(), + }; +} diff --git a/backend/src/modelInfo/embeddingProcesses/index.test.ts b/backend/src/modelInfo/embeddingProcesses/index.test.ts new file mode 100644 index 000000000..3580d7a8b --- /dev/null +++ b/backend/src/modelInfo/embeddingProcesses/index.test.ts @@ -0,0 +1,53 @@ +// silence chatty console +import "_test_utilities/consoleMock"; + +// Mock the POST handler so we can observe the dispatch +const mockHandle = jest.fn(); +jest.mock("./POST", () => ({ + POSTModelEmbeddingProcessesHandler: jest.fn().mockImplementation(() => ({ handle: mockHandle })), +})); + +import * as ModelEmbeddingProcessesHandler from "./index"; +import { HTTP_VERBS, response, StatusCodes, STD_ERRORS_RESPONSES } from "server/httpUtils"; +import { APIGatewayProxyEvent } from "aws-lambda"; +import { testMethodsNotAllowed } from "_test_utilities/stdRESTHandlerTests"; + +describe("Test the ModelEmbeddingProcesses handler", () => { + beforeEach(() => { + jest.clearAllMocks(); + }); + + test("POST should delegate to the POSTModelEmbeddingProcessesHandler", async () => { + // GIVEN a POST event + const givenEvent = { httpMethod: HTTP_VERBS.POST } as APIGatewayProxyEvent; + // AND the POST handler will return a response + const givenResponse = response(StatusCodes.ACCEPTED, {}); + mockHandle.mockResolvedValueOnce(givenResponse); + + // WHEN the handler is invoked with the given event + const actualResponse = await ModelEmbeddingProcessesHandler.handler(givenEvent); + + // THEN expect the POST handler to be invoked with the given event + expect(mockHandle).toHaveBeenCalledWith(givenEvent); + // AND the handler to return the response from the POST handler + expect(actualResponse).toEqual(givenResponse); + }); + + testMethodsNotAllowed( + [HTTP_VERBS.PUT, HTTP_VERBS.DELETE, HTTP_VERBS.OPTIONS, HTTP_VERBS.PATCH, HTTP_VERBS.GET], + ModelEmbeddingProcessesHandler.handler + ); + + test("should not delegate to the POST handler for non-POST methods", async () => { + // GIVEN a GET event + const givenEvent = { httpMethod: HTTP_VERBS.GET } as APIGatewayProxyEvent; + + // WHEN the handler is invoked with the given event + const actualResponse = await ModelEmbeddingProcessesHandler.handler(givenEvent); + + // THEN expect the POST handler to not be invoked + expect(mockHandle).not.toHaveBeenCalled(); + // AND expect a METHOD_NOT_ALLOWED response + expect(actualResponse).toEqual(STD_ERRORS_RESPONSES.METHOD_NOT_ALLOWED); + }); +}); diff --git a/backend/src/modelInfo/embeddingProcesses/index.ts b/backend/src/modelInfo/embeddingProcesses/index.ts new file mode 100644 index 000000000..acc2ff051 --- /dev/null +++ b/backend/src/modelInfo/embeddingProcesses/index.ts @@ -0,0 +1,15 @@ +import { APIGatewayProxyEvent } from "aws-lambda"; +import { APIGatewayProxyResult } from "aws-lambda/trigger/api-gateway-proxy"; +import { HTTP_VERBS, STD_ERRORS_RESPONSES } from "server/httpUtils"; +import { POSTModelEmbeddingProcessesHandler } from "./POST"; + +export const handler: ( + event: APIGatewayProxyEvent /*, context: Context, callback: Callback*/ +) => Promise = async ( + event: APIGatewayProxyEvent /*, context: Context, callback: Callback*/ +) => { + if (event?.httpMethod === HTTP_VERBS.POST) { + return new POSTModelEmbeddingProcessesHandler().handle(event); + } + return STD_ERRORS_RESPONSES.METHOD_NOT_ALLOWED; +}; diff --git a/backend/src/routes.constant.ts b/backend/src/routes.constant.ts index 3bd08eeff..57b089a05 100644 --- a/backend/src/routes.constant.ts +++ b/backend/src/routes.constant.ts @@ -27,4 +27,5 @@ export const Routes = { SKILL_OCCUPATIONS_ROUTE: "/models/:modelId/skills/:id/occupations", SKILL_RELATED_ROUTE: "/models/:modelId/skills/:id/related", SKILL_HISTORY_ROUTE: "/models/:modelId/skills/:id/history", + MODEL_EMBEDDING_PROCESSES_ROUTE: "/models/:modelId/embedding-processes", }; diff --git a/backend/src/server/config/config.test.ts b/backend/src/server/config/config.test.ts index 816214c65..438cf82e2 100644 --- a/backend/src/server/config/config.test.ts +++ b/backend/src/server/config/config.test.ts @@ -47,6 +47,10 @@ describe("Test read Configuration()", () => { process.env.ASYNC_IMPORT_LAMBDA_FUNCTION_ARN = getRandomString(10); process.env.ASYNC_EXPORT_LAMBDA_FUNCTION_ARN = getRandomString(10); process.env.ASYNC_LAMBDA_FUNCTION_REGION = getRandomString(10); + process.env.GEMINI_API_KEY = getRandomString(10); + process.env.GEMINI_EMBEDDING_MODEL = getRandomString(10); + process.env.EMBEDDINGS_QUEUE_URL = getRandomString(10); + process.env.EMBEDDINGS_QUEUE_REGION = getRandomString(10); // WHEN reading the configuration from the environment const actualConfig = readEnvironmentConfiguration(); @@ -63,6 +67,10 @@ describe("Test read Configuration()", () => { asyncImportLambdaFunctionArn: process.env.ASYNC_IMPORT_LAMBDA_FUNCTION_ARN, asyncExportLambdaFunctionArn: process.env.ASYNC_EXPORT_LAMBDA_FUNCTION_ARN, asyncLambdaFunctionRegion: process.env.ASYNC_LAMBDA_FUNCTION_REGION, + geminiApiKey: process.env.GEMINI_API_KEY, + geminiEmbeddingModel: process.env.GEMINI_EMBEDDING_MODEL, + embeddingsQueueUrl: process.env.EMBEDDINGS_QUEUE_URL, + embeddingsQueueRegion: process.env.EMBEDDINGS_QUEUE_REGION, }); }); @@ -77,6 +85,10 @@ describe("Test read Configuration()", () => { delete process.env.DOWNLOAD_BUCKET_REGION; delete process.env.ASYNC_LAMBDA_FUNCTION_ARN; delete process.env.ASYNC_LAMBDA_FUNCTION_REGION; + delete process.env.GEMINI_API_KEY; + delete process.env.GEMINI_EMBEDDING_MODEL; + delete process.env.EMBEDDINGS_QUEUE_URL; + delete process.env.EMBEDDINGS_QUEUE_REGION; // WHEN reading the configuration from the environment const config = readEnvironmentConfiguration(); @@ -93,6 +105,10 @@ describe("Test read Configuration()", () => { asyncImportLambdaFunctionArn: "", asyncExportLambdaFunctionArn: "", asyncLambdaFunctionRegion: "", + geminiApiKey: "", + geminiEmbeddingModel: "", + embeddingsQueueUrl: "", + embeddingsQueueRegion: "", }); }); }); @@ -189,5 +205,9 @@ function getMockConfig(): IConfiguration { asyncImportLambdaFunctionArn: getTestString(10), asyncExportLambdaFunctionArn: getTestString(10), asyncLambdaFunctionRegion: getTestString(10), + geminiApiKey: getTestString(10), + geminiEmbeddingModel: getTestString(10), + embeddingsQueueRegion: getTestString(10), + embeddingsQueueUrl: getTestString(10), }; } diff --git a/backend/src/server/config/config.ts b/backend/src/server/config/config.ts index c77ea8463..cf1c031c5 100644 --- a/backend/src/server/config/config.ts +++ b/backend/src/server/config/config.ts @@ -11,6 +11,10 @@ export const ENV_VAR_NAMES = { ASYNC_IMPORT_LAMBDA_FUNCTION_ARN: "ASYNC_IMPORT_LAMBDA_FUNCTION_ARN", ASYNC_EXPORT_LAMBDA_FUNCTION_ARN: "ASYNC_EXPORT_LAMBDA_FUNCTION_ARN", ASYNC_LAMBDA_FUNCTION_REGION: "ASYNC_LAMBDA_FUNCTION_REGION", + GEMINI_API_KEY: "GEMINI_API_KEY", + GEMINI_EMBEDDING_MODEL: "GEMINI_EMBEDDING_MODEL", + EMBEDDINGS_QUEUE_URL: "EMBEDDINGS_QUEUE_URL", + EMBEDDINGS_QUEUE_REGION: "EMBEDDINGS_QUEUE_REGION", }; export interface IConfiguration { @@ -24,6 +28,10 @@ export interface IConfiguration { asyncImportLambdaFunctionArn: string; asyncExportLambdaFunctionArn: string; asyncLambdaFunctionRegion: string; + geminiApiKey: string; + geminiEmbeddingModel: string; + embeddingsQueueUrl: string; + embeddingsQueueRegion: string; } export function readEnvironmentConfiguration(): IConfiguration { return { @@ -37,6 +45,10 @@ export function readEnvironmentConfiguration(): IConfiguration { asyncImportLambdaFunctionArn: process.env[ENV_VAR_NAMES.ASYNC_IMPORT_LAMBDA_FUNCTION_ARN] ?? "", asyncExportLambdaFunctionArn: process.env[ENV_VAR_NAMES.ASYNC_EXPORT_LAMBDA_FUNCTION_ARN] ?? "", asyncLambdaFunctionRegion: process.env[ENV_VAR_NAMES.ASYNC_LAMBDA_FUNCTION_REGION] ?? "", + geminiApiKey: process.env[ENV_VAR_NAMES.GEMINI_API_KEY] ?? "", + geminiEmbeddingModel: process.env[ENV_VAR_NAMES.GEMINI_EMBEDDING_MODEL] ?? "", + embeddingsQueueUrl: process.env[ENV_VAR_NAMES.EMBEDDINGS_QUEUE_URL] ?? "", + embeddingsQueueRegion: process.env[ENV_VAR_NAMES.EMBEDDINGS_QUEUE_REGION] ?? "", }; } @@ -88,3 +100,19 @@ export function getAsyncLambdaFunctionRegion() { export function getDomainName() { return _configuration?.domainName ?? ""; } + +export function getGeminiApiKey() { + return _configuration?.geminiApiKey ?? ""; +} + +export function getGeminiEmbeddingModel() { + return _configuration?.geminiEmbeddingModel ?? ""; +} + +export function getEmbeddingsQueueUrl() { + return _configuration?.embeddingsQueueUrl ?? ""; +} + +export function getEmbeddingsQueueRegion() { + return _configuration?.embeddingsQueueRegion ?? ""; +} diff --git a/backend/src/server/httpUtils.ts b/backend/src/server/httpUtils.ts index 552426133..41099e91f 100644 --- a/backend/src/server/httpUtils.ts +++ b/backend/src/server/httpUtils.ts @@ -21,6 +21,7 @@ export enum StatusCodes { FORBIDDEN = 403, NOT_FOUND = 404, METHOD_NOT_ALLOWED = 405, + CONFLICT = 409, TOO_LARGE_PAYLOAD = 413, UNSUPPORTED_MEDIA_TYPE = 415, INTERNAL_SERVER_ERROR = 500, diff --git a/backend/src/server/repositoryRegistry/repositoryRegistry.test.ts b/backend/src/server/repositoryRegistry/repositoryRegistry.test.ts index f952496af..32b395c32 100644 --- a/backend/src/server/repositoryRegistry/repositoryRegistry.test.ts +++ b/backend/src/server/repositoryRegistry/repositoryRegistry.test.ts @@ -46,6 +46,7 @@ describe("test the RepositoryRegistry", () => { expect(repositoryRegistry.skillHierarchy).toBeDefined(); expect(repositoryRegistry.skillToSkillRelation).toBeDefined(); expect(repositoryRegistry.occupationToSkillRelation).toBeDefined(); + expect(repositoryRegistry.embeddingProcessState).toBeDefined(); }); test("should reject the connection is not defined", async () => { diff --git a/backend/src/server/repositoryRegistry/repositoryRegistry.ts b/backend/src/server/repositoryRegistry/repositoryRegistry.ts index 804b64967..a9f896f4c 100644 --- a/backend/src/server/repositoryRegistry/repositoryRegistry.ts +++ b/backend/src/server/repositoryRegistry/repositoryRegistry.ts @@ -11,6 +11,7 @@ import * as skillToSkillRelationModel from "esco/skillToSkillRelation/skillToSki import * as occupationToSkillRelationModel from "esco/occupationToSkillRelation/occupationToSkillRelationModel"; import * as importStateModel from "import/ImportProcessState/importProcessStateModel"; import * as exportProcessStateModel from "export/exportProcessState/exportProcessStateModel"; +import * as embeddingProcessStateModel from "embeddings/embeddingProcessState/embeddingProcessStateModel"; import { IOccupationGroupRepository, @@ -40,6 +41,10 @@ import { ExportProcessStateRepository, IExportProcessStateRepository, } from "export/exportProcessState/exportProcessStateRepository"; +import { + EmbeddingProcessStateRepository, + IEmbeddingProcessStateRepository, +} from "embeddings/embeddingProcessState/embeddingProcessStateRepository"; export class RepositoryRegistry { // eslint-disable-next-line @@ -133,6 +138,14 @@ export class RepositoryRegistry { this._repositories.set("IExportProcessStateRepository", repository); } + public get embeddingProcessState(): IEmbeddingProcessStateRepository { + return this._repositories.get("IEmbeddingProcessStateRepository"); + } + + public set embeddingProcessState(repository: IEmbeddingProcessStateRepository) { + this._repositories.set("IEmbeddingProcessStateRepository", repository); + } + async initialize(connection: Connection | undefined) { if (!connection) throw new Error("Connection is undefined"); @@ -173,6 +186,9 @@ export class RepositoryRegistry { this.exportProcessState = new ExportProcessStateRepository( exportProcessStateModel.initializeSchemaAndModel(connection) ); + this.embeddingProcessState = new EmbeddingProcessStateRepository( + embeddingProcessStateModel.initializeSchemaAndModel(connection) + ); // Set up the indexes // This is done here because the autoIndex is turned off in production @@ -191,6 +207,7 @@ export class RepositoryRegistry { await this.occupationToSkillRelation.relationModel.createIndexes(); await this.importProcessState.Model.createIndexes(); await this.exportProcessState.Model.createIndexes(); + await this.embeddingProcessState.Model.createIndexes(); } } diff --git a/backend/src/server/serviceRegistry/serviceRegistry.test.ts b/backend/src/server/serviceRegistry/serviceRegistry.test.ts index f232e64d7..54e854097 100644 --- a/backend/src/server/serviceRegistry/serviceRegistry.test.ts +++ b/backend/src/server/serviceRegistry/serviceRegistry.test.ts @@ -1,6 +1,8 @@ // Suppress chatty console during the tests import "_test_utilities/consoleMock"; +import { SQSClient } from "@aws-sdk/client-sqs"; +import { getEmbeddingsQueueRegion } from "server/config/config"; import { getServiceRegistry, ServiceRegistry } from "./serviceRegistry"; import { SkillService } from "esco/skill/services/skill.service"; import { OccupationService } from "esco/occupations/services/occupation.service"; @@ -10,8 +12,22 @@ import { OccupationHierarchyService } from "esco/occupationHierarchy/occupationH import { OccupationToSkillRelationService } from "esco/occupationToSkillRelation/occupationToSkillRelation.service"; import { SkillHierarchyService } from "esco/skillHierarchy/skillHierarchy.service"; import { SkillToSkillRelationService } from "esco/skillToSkillRelation/skillToSkillRelation.service"; +import { EmbeddingProcessService } from "embeddings/embeddingProcess/embeddingProcess.service"; + +jest.mock("@aws-sdk/client-sqs", () => ({ + SQSClient: jest.fn(), +})); + +jest.mock("server/config/config", () => ({ + ...jest.requireActual("server/config/config"), + getEmbeddingsQueueRegion: jest.fn(), +})); describe("test the ServiceRegistry", () => { + beforeEach(() => { + jest.clearAllMocks(); + }); + test("should return a singleton ServiceRegistry", () => { // WHEN trying to get the ServiceRegistry const serviceRegistry = getServiceRegistry(); @@ -25,10 +41,17 @@ describe("test the ServiceRegistry", () => { }); test("should initialize and set services successfully", async () => { + // GIVEN the embeddings queue region is configured + const givenEmbeddingsQueueRegion = "eu-west-1"; + (getEmbeddingsQueueRegion as jest.Mock).mockReturnValue(givenEmbeddingsQueueRegion); + // WHEN trying to initialize the ServiceRegistry const serviceRegistry = new ServiceRegistry(); await serviceRegistry.initialize(); + // THEN the SQSClient should be constructed with the configured region + expect(SQSClient).toHaveBeenCalledWith({ region: givenEmbeddingsQueueRegion }); + // THEN the services should be initialized expect(serviceRegistry.occupation).toBeDefined(); expect(serviceRegistry.occupation).toBeInstanceOf(OccupationService); @@ -46,5 +69,7 @@ describe("test the ServiceRegistry", () => { expect(serviceRegistry.skillHierarchy).toBeInstanceOf(SkillHierarchyService); expect(serviceRegistry.skillToSkillRelation).toBeDefined(); expect(serviceRegistry.skillToSkillRelation).toBeInstanceOf(SkillToSkillRelationService); + expect(serviceRegistry.embeddingProcess).toBeDefined(); + expect(serviceRegistry.embeddingProcess).toBeInstanceOf(EmbeddingProcessService); }); }); diff --git a/backend/src/server/serviceRegistry/serviceRegistry.ts b/backend/src/server/serviceRegistry/serviceRegistry.ts index aa4c66542..c2c72d23c 100644 --- a/backend/src/server/serviceRegistry/serviceRegistry.ts +++ b/backend/src/server/serviceRegistry/serviceRegistry.ts @@ -1,3 +1,5 @@ +import { SQSClient } from "@aws-sdk/client-sqs"; + import { IOccupationGroupService } from "esco/occupationGroup/services/occupationGroup.service.type"; import { getRepositoryRegistry } from "../repositoryRegistry/repositoryRegistry"; import { OccupationService } from "esco/occupations/services/occupation.service"; @@ -15,6 +17,10 @@ import { ISkillHierarchyService } from "esco/skillHierarchy/skillHierarchy.servi import { SkillHierarchyService } from "esco/skillHierarchy/skillHierarchy.service"; import { ISkillToSkillRelationService } from "esco/skillToSkillRelation/skillToSkillRelation.service.types"; import { SkillToSkillRelationService } from "esco/skillToSkillRelation/skillToSkillRelation.service"; +import { IEmbeddingProcessService } from "embeddings/embeddingProcess/embeddingProcess.service.types"; +import { EmbeddingProcessService } from "embeddings/embeddingProcess/embeddingProcess.service"; +import { EmbeddingClient } from "embeddings/service/client"; +import { getEmbeddingsQueueRegion } from "server/config/config"; export class ServiceRegistry { // eslint-disable-next-line @typescript-eslint/no-explicit-any @@ -44,6 +50,9 @@ export class ServiceRegistry { public get skillToSkillRelation(): ISkillToSkillRelationService { return this._services.get("SkillToSkillRelationService"); } + public get embeddingProcess(): IEmbeddingProcessService { + return this._services.get("EmbeddingProcessService"); + } public set occupation(service: IOccupationService) { this._services.set("OccupationService", service); @@ -70,8 +79,12 @@ export class ServiceRegistry { public set skillToSkillRelation(service: ISkillToSkillRelationService) { this._services.set("SkillToSkillRelationService", service); } + public set embeddingProcess(service: IEmbeddingProcessService) { + this._services.set("EmbeddingProcessService", service); + } async initialize() { + const awsSQSClient = new SQSClient({ region: getEmbeddingsQueueRegion() }); const repositoryRegistry = getRepositoryRegistry(); this.occupation = new OccupationService(repositoryRegistry.occupation, repositoryRegistry.modelInfo); this.occupationGroup = new OccupationGroupService( @@ -99,6 +112,16 @@ export class ServiceRegistry { repositoryRegistry.skill, repositoryRegistry.skillToSkillRelation ); + const embeddingClient = new EmbeddingClient(awsSQSClient); + this.embeddingProcess = new EmbeddingProcessService( + repositoryRegistry.modelInfo, + repositoryRegistry.embeddingProcessState, + repositoryRegistry.skill, + repositoryRegistry.skillGroup, + repositoryRegistry.occupation, + repositoryRegistry.OccupationGroup, + embeddingClient + ); } } diff --git a/backend/src/validator.ts b/backend/src/validator.ts index 8078de68f..1420ec706 100644 --- a/backend/src/validator.ts +++ b/backend/src/validator.ts @@ -211,6 +211,10 @@ ajvInstance.addSchema( SkillAPISpecs.Skill.RelatedSkills.POST.Schemas.Request.Payload, SkillAPISpecs.Skill.RelatedSkills.POST.Schemas.Request.Payload.$id ); +ajvInstance.addSchema( + ModelInfoAPISpecs.ModelInfo.EmbeddingProcessStates.POST.Schemas.Request.Payload, + ModelInfoAPISpecs.ModelInfo.EmbeddingProcessStates.POST.Schemas.Request.Payload.$id +); /** * Turn the errors from ajv and turn into a string that consumers can read. diff --git a/backend/yarn.lock b/backend/yarn.lock index c214e8ca0..f8244238f 100644 --- a/backend/yarn.lock +++ b/backend/yarn.lock @@ -314,6 +314,53 @@ fast-xml-parser "4.2.5" tslib "^2.5.0" +"@aws-sdk/client-sqs@~3.468.0": + version "3.468.0" + resolved "https://registry.yarnpkg.com/@aws-sdk/client-sqs/-/client-sqs-3.468.0.tgz#c3636f1d26e9f8641fe39576c9b45851e9ca8d89" + integrity sha512-1Hubnp2e1yRMLHkAKgR+U3ptqBt7uan64/A88EIFTCzzMHEUmgDsIcpPGVpd14YXqxG2cCeU7qA/Z6uiHGWK0g== + dependencies: + "@aws-crypto/sha256-browser" "3.0.0" + "@aws-crypto/sha256-js" "3.0.0" + "@aws-sdk/client-sts" "3.468.0" + "@aws-sdk/core" "3.468.0" + "@aws-sdk/credential-provider-node" "3.468.0" + "@aws-sdk/middleware-host-header" "3.468.0" + "@aws-sdk/middleware-logger" "3.468.0" + "@aws-sdk/middleware-recursion-detection" "3.468.0" + "@aws-sdk/middleware-sdk-sqs" "3.468.0" + "@aws-sdk/middleware-signing" "3.468.0" + "@aws-sdk/middleware-user-agent" "3.468.0" + "@aws-sdk/region-config-resolver" "3.468.0" + "@aws-sdk/types" "3.468.0" + "@aws-sdk/util-endpoints" "3.468.0" + "@aws-sdk/util-user-agent-browser" "3.468.0" + "@aws-sdk/util-user-agent-node" "3.468.0" + "@smithy/config-resolver" "^2.0.20" + "@smithy/fetch-http-handler" "^2.3.1" + "@smithy/hash-node" "^2.0.17" + "@smithy/invalid-dependency" "^2.0.15" + "@smithy/md5-js" "^2.0.17" + "@smithy/middleware-content-length" "^2.0.17" + "@smithy/middleware-endpoint" "^2.2.2" + "@smithy/middleware-retry" "^2.0.23" + "@smithy/middleware-serde" "^2.0.15" + "@smithy/middleware-stack" "^2.0.9" + "@smithy/node-config-provider" "^2.1.7" + "@smithy/node-http-handler" "^2.2.1" + "@smithy/protocol-http" "^3.0.11" + "@smithy/smithy-client" "^2.1.18" + "@smithy/types" "^2.7.0" + "@smithy/url-parser" "^2.0.15" + "@smithy/util-base64" "^2.0.1" + "@smithy/util-body-length-browser" "^2.0.1" + "@smithy/util-body-length-node" "^2.1.0" + "@smithy/util-defaults-mode-browser" "^2.0.22" + "@smithy/util-defaults-mode-node" "^2.0.28" + "@smithy/util-endpoints" "^1.0.6" + "@smithy/util-retry" "^2.0.8" + "@smithy/util-utf8" "^2.0.2" + tslib "^2.5.0" + "@aws-sdk/client-sso@3.468.0": version "3.468.0" resolved "https://registry.npmjs.org/@aws-sdk/client-sso/-/client-sso-3.468.0.tgz" @@ -769,6 +816,17 @@ "@smithy/util-config-provider" "^2.0.0" tslib "^2.5.0" +"@aws-sdk/middleware-sdk-sqs@3.468.0": + version "3.468.0" + resolved "https://registry.yarnpkg.com/@aws-sdk/middleware-sdk-sqs/-/middleware-sdk-sqs-3.468.0.tgz#7887dcb89553ab643d8806fa44d5529fa3510021" + integrity sha512-ehDQHKimed1DrdXYG12LPKWzHlG9iXihGEok07nOJR/2RpOuSpAcavc4eeuOVeEHh58Eu8TXvvVAJ906Vtm6YA== + dependencies: + "@aws-sdk/types" "3.468.0" + "@smithy/types" "^2.7.0" + "@smithy/util-hex-encoding" "^2.0.0" + "@smithy/util-utf8" "^2.0.2" + tslib "^2.5.0" + "@aws-sdk/middleware-sdk-sts@3.468.0": version "3.468.0" resolved "https://registry.npmjs.org/@aws-sdk/middleware-sdk-sts/-/middleware-sdk-sts-3.468.0.tgz" diff --git a/iac/backend/src/embeddings.ts b/iac/backend/src/embeddings.ts new file mode 100644 index 000000000..1a61c9ae6 --- /dev/null +++ b/iac/backend/src/embeddings.ts @@ -0,0 +1,98 @@ +import * as aws from "@pulumi/aws"; +import * as pulumi from "@pulumi/pulumi"; +import {asset} from "@pulumi/pulumi"; + +// The embedding lambda is bundled by esbuild from backend/src/embeddings/handler/index.ts +const buildFolderPath = "../../backend/build/embeddings/handler"; + +const LOG_RETENTION_IN_DAYS = 7; + +const LAMBDA_TIMEOUT_IN_SECONDS = 300; + +const LAMBDA_MEMORY_IN_MB = 1024; + +const LAMBDA_MAXIMUM_CONCURRENT_EXECUTIONS = 2; + +interface SetupEmbeddingsFnConfig { + mongodb_uri: string, + resourcesBaseUrl: string, + gemini_api_key: string, + gemini_embedding_model: string, + sentry_backend_dsn: string +} + +interface SetupEmbeddingsFnOutput { + embeddingsLambdaRole: aws.iam.Role, + embeddingsLambdaFunction: aws.lambda.Function +} + +export function setupEmbeddingsFn(environment: string, config: SetupEmbeddingsFnConfig): SetupEmbeddingsFnOutput { + /** + * Lambda for embedding generation (triggered by the embeddings SQS queue) + */ + + // Create a new IAM role for the Lambda function + const embeddingsLambdaRole = new aws.iam.Role("embeddings-function-role", { + assumeRolePolicy: JSON.stringify({ + Version: "2012-10-17", + Statement: [ + { + Action: "sts:AssumeRole", + Effect: "Allow", + Principal: { + Service: "lambda.amazonaws.com", + }, + }, + ], + }), + }); + + // Attach the necessary policies to the IAM role + const cloudwatchPolicy = aws.iam.getPolicy({ + arn: "arn:aws:iam::aws:policy/CloudWatchLogsFullAccess", + }); + + const embeddingsLambdaPolicy = new aws.iam.Policy("embeddings-function-policy", { + policy: cloudwatchPolicy.then((cp) => cp.policy), + }); + + new aws.iam.RolePolicyAttachment("embeddings-function-role-policy-attachment", { + policyArn: embeddingsLambdaPolicy.arn, + role: embeddingsLambdaRole.name, + }); + + const embeddingsFileArchive = new asset.FileArchive(buildFolderPath); + + // Create a new AWS Lambda function + const embeddingsLambdaFunction = new aws.lambda.Function("embeddings-function", { + role: embeddingsLambdaRole.arn, + code: embeddingsFileArchive, + handler: "index.handler", + runtime: 'nodejs20.x', + timeout: LAMBDA_TIMEOUT_IN_SECONDS, + memorySize: LAMBDA_MEMORY_IN_MB, + reservedConcurrentExecutions: LAMBDA_MAXIMUM_CONCURRENT_EXECUTIONS, + environment: { + variables: { + NODE_OPTIONS: '--enable-source-maps', + RESOURCES_BASE_URL: config.resourcesBaseUrl, + MONGODB_URI: config.mongodb_uri, + GEMINI_API_KEY: config.gemini_api_key, + GEMINI_EMBEDDING_MODEL: config.gemini_embedding_model, + SENTRY_BACKEND_DSN: config.sentry_backend_dsn, + TARGET_ENVIRONMENT: environment, + } + } + }); + + // Create log group with retention of days, + // log group is assigned to the lambda function via the name of the log group (see https://docs.aws.amazon.com/lambda/latest/dg/monitoring-cloudwatchlogs.html) + new aws.cloudwatch.LogGroup("embeddings-log-group", { + name: pulumi.interpolate`/aws/lambda/${embeddingsLambdaFunction.name}`, + retentionInDays: LOG_RETENTION_IN_DAYS + }); + + return {embeddingsLambdaRole, embeddingsLambdaFunction}; +} + +export const EMBEDDINGS_LAMBDA_TIMEOUT_IN_SECONDS = LAMBDA_TIMEOUT_IN_SECONDS; diff --git a/iac/backend/src/embeddingsQueue.ts b/iac/backend/src/embeddingsQueue.ts new file mode 100644 index 000000000..62b6f0e42 --- /dev/null +++ b/iac/backend/src/embeddingsQueue.ts @@ -0,0 +1,91 @@ +import * as aws from "@pulumi/aws"; +import {EMBEDDINGS_LAMBDA_TIMEOUT_IN_SECONDS} from "./embeddings"; + +// How many times a message is retried before it is moved to the dead-letter queue. +const MAX_RECEIVE_COUNT = 5; + +// The visibility timeout must be greater than or equal to the consuming lambda's timeout, otherwise a +// message could become visible again (and be reprocessed) while the lambda is still handling it. +const QUEUE_VISIBILITY_TIMEOUT_IN_SECONDS = EMBEDDINGS_LAMBDA_TIMEOUT_IN_SECONDS + 60; + +// AWS maximum retention (14 days) for messages that end up in the dead-letter queue. +const DEAD_LETTER_QUEUE_RETENTION_IN_SECONDS = 1209600; + +// Number of records delivered to the lambda per invocation. If N messages are available, the event +// source mapping invokes the lambda with up to BATCH_SIZE records at a time. +const BATCH_SIZE = 10; + +/** + * Set up the embeddings SQS queue and its dead-letter queue and connect the queue to the embedding + * lambda via an event source mapping so that messages pushed onto the queue automatically trigger the + * lambda to process them. + */ +export function setupEmbeddingsQueue(config: { + embeddingsLambdaRole: aws.iam.Role, + embeddingsLambdaFunction: aws.lambda.Function, +}): { + embeddingsQueue: aws.sqs.Queue, + embeddingsDeadLetterQueue: aws.sqs.Queue, +} { + /** + * Dead-letter queue: receives messages that fail to process MAX_RECEIVE_COUNT times. + */ + const embeddingsDeadLetterQueue = new aws.sqs.Queue("embeddings-dead-letter-queue", { + messageRetentionSeconds: DEAD_LETTER_QUEUE_RETENTION_IN_SECONDS, + }); + + /** + * Main embeddings queue, with a redrive policy pointing at the dead-letter queue. + */ + const embeddingsQueue = new aws.sqs.Queue("embeddings-queue", { + visibilityTimeoutSeconds: QUEUE_VISIBILITY_TIMEOUT_IN_SECONDS, + redrivePolicy: embeddingsDeadLetterQueue.arn.apply((deadLetterTargetArn) => + JSON.stringify({ + deadLetterTargetArn, + maxReceiveCount: MAX_RECEIVE_COUNT, + }) + ), + }); + + /** + * Allow the embedding lambda to consume from the queue. + */ + const consumePolicy = new aws.iam.Policy("embeddings-queue-consume-policy", { + policy: embeddingsQueue.arn.apply((queueArn) => + JSON.stringify({ + Version: "2012-10-17", + Statement: [ + { + Effect: "Allow", + Action: [ + "sqs:ReceiveMessage", + "sqs:DeleteMessage", + "sqs:GetQueueAttributes", + ], + Resource: queueArn, + }, + ], + }) + ), + }); + + new aws.iam.RolePolicyAttachment("embeddings-queue-consume-policy-attachment", { + policyArn: consumePolicy.arn, + role: config.embeddingsLambdaRole.name, + }); + + /** + * Event source mapping: SQS queue -> embeddings lambda. + * + * When N messages are available on the queue, the lambda is automatically invoked with up to + * BATCH_SIZE records per invocation until the queue is drained. + */ + new aws.lambda.EventSourceMapping("embeddings-queue-event-source-mapping", { + eventSourceArn: embeddingsQueue.arn, + functionName: config.embeddingsLambdaFunction.arn, + batchSize: BATCH_SIZE, + enabled: true, + }); + + return {embeddingsQueue, embeddingsDeadLetterQueue}; +} diff --git a/iac/backend/src/index.ts b/iac/backend/src/index.ts index 9f890510a..f76246d8b 100644 --- a/iac/backend/src/index.ts +++ b/iac/backend/src/index.ts @@ -7,6 +7,8 @@ import {setupRedocBucket, setupSwaggerBucket} from "./openapiBuckets"; import {setupDownloadBucket, setupDownloadBucketWritePolicy} from "./downloadBucket"; import {setupAsyncExportApi} from "./asyncExport"; import {setupAuthorizer} from "./authorizer"; +import {setupEmbeddingsFn} from "./embeddings"; +import {setupEmbeddingsQueue} from "./embeddingsQueue"; export const environment = pulumi.getStack(); export const domainName = process.env.DOMAIN_NAME!; @@ -26,6 +28,13 @@ if(!authDatabaseURI) throw new Error("environment variable AUTH_DATABASE_URI is const sentryBackendDSN = process.env.SENTRY_BACKEND_DSN ?? ""; +// Gemini configuration (used by both the core REST backend and the embedding lambda to generate vectors). +const geminiApiKey = process.env.GEMINI_API_KEY; +if(!geminiApiKey) throw new Error("environment variable GEMINI_API_KEY is required") + +const geminiEmbeddingModel = process.env.GEMINI_EMBEDDING_MODEL; +if(!geminiEmbeddingModel) throw new Error("environment variable GEMINI_EMBEDDING_MODEL is required") + const userPoolId = process.env.USER_POOL_ID; if(!userPoolId) throw new Error("environment variable USER_POOL_ID is required"); @@ -92,6 +101,27 @@ const {asyncImportLambdaRole, asyncImportLambdaFunction} = setupAsyncImportApi(e sentry_backend_dsn: sentryBackendDSN, }); +/** + * Setup Embeddings lambda and queue + * + * The embedding lambda generates vectors in the background. It is triggered by the embeddings SQS + * queue via an event source mapping (see setupEmbeddingsQueue), which also wires the dead-letter queue. + */ +const {embeddingsLambdaRole, embeddingsLambdaFunction} = setupEmbeddingsFn(environment, { + mongodb_uri: mongoDbUri, + resourcesBaseUrl, + gemini_api_key: geminiApiKey, + gemini_embedding_model: geminiEmbeddingModel, + sentry_backend_dsn: sentryBackendDSN, +}); + +const {embeddingsQueue} = setupEmbeddingsQueue({ + embeddingsLambdaRole, + embeddingsLambdaFunction, +}); + +export const embeddingsQueueUrl = embeddingsQueue.url; + /** * Setup Authorizer Lambda */ @@ -117,6 +147,11 @@ const {restApi, stage, restApiLambdaRole} = setupBackendRESTApi(environment, { async_lambda_function_region: currentRegion, authorizer_lambda_function_invoke_arn: authorizerLambdaFunction.invokeArn, authorizer_lambda_function_name: authorizerLambdaFunction.name, + embeddings_queue_url: embeddingsQueue.url, + embeddings_queue_arn: embeddingsQueue.arn, + embeddings_queue_region: currentRegion, + gemini_api_key: geminiApiKey, + gemini_embedding_model: geminiEmbeddingModel, sentry_backend_dsn: sentryBackendDSN, }); diff --git a/iac/backend/src/restApi.ts b/iac/backend/src/restApi.ts index c5762706d..4f77b0464 100644 --- a/iac/backend/src/restApi.ts +++ b/iac/backend/src/restApi.ts @@ -14,7 +14,7 @@ const LAMBDA_MEMORY_IN_MB = 512; const LAMBDA_MAXIMUM_CONCURRENT_EXECUTIONS = 10; -export function setupBackendRESTApi(environment: string, config: { +interface SetupBackendRESTAPIConfig { mongodb_uri: string, auth_database_uri: string, resourcesBaseUrl: string, @@ -27,8 +27,21 @@ export function setupBackendRESTApi(environment: string, config: { async_lambda_function_region: Output, authorizer_lambda_function_invoke_arn: Output authorizer_lambda_function_name: Output, + embeddings_queue_url: Output, + embeddings_queue_arn: Output, + embeddings_queue_region: Output, + gemini_api_key: string, + gemini_embedding_model: string, sentry_backend_dsn: string -}): { restApi: RestApi, stage: Stage, restApiLambdaRole: aws.iam.Role } { +} + +interface SetupBackendRESTAPIOutput { + restApi: RestApi, + stage: Stage, + restApiLambdaRole: aws.iam.Role +} + +export function setupBackendRESTApi(environment: string, config: SetupBackendRESTAPIConfig): SetupBackendRESTAPIOutput { /** * Lambda for api */ @@ -104,6 +117,30 @@ export function setupBackendRESTApi(environment: string, config: { role: lambdaRole.name, }); + // Send message to the embedding queue policy + // The REST API is the producer that pushes entities onto the embedding queue. + const embeddingsQueueSendMessagePolicy = new aws.iam.Policy("model-api-function-embeddings-queue-send-policy", { + policy: config.embeddings_queue_arn.apply((queueArn) => JSON.stringify({ + Version: "2012-10-17", + Statement: [ + { + Action: [ + "sqs:SendMessage", + "sqs:SendMessageBatch", + "sqs:GetQueueAttributes", + ], + Effect: "Allow", + Resource: queueArn, + } + ] + })) + }); + + new aws.iam.RolePolicyAttachment("model-api-function-role-embeddings-queue-send-policy-attachment", { + policyArn: embeddingsQueueSendMessagePolicy.arn, + role: lambdaRole.name, + }); + // Build the source code archive let fileArchive = new asset.FileArchive(buildFolderPath); @@ -129,6 +166,10 @@ export function setupBackendRESTApi(environment: string, config: { ASYNC_IMPORT_LAMBDA_FUNCTION_ARN: config.async_import_lambda_function_arn, ASYNC_EXPORT_LAMBDA_FUNCTION_ARN: config.async_export_lambda_function_arn, ASYNC_LAMBDA_FUNCTION_REGION: config.async_lambda_function_region, + EMBEDDINGS_QUEUE_URL: config.embeddings_queue_url, + EMBEDDINGS_QUEUE_REGION: config.embeddings_queue_region, + GEMINI_API_KEY: config.gemini_api_key, + GEMINI_EMBEDDING_MODEL: config.gemini_embedding_model, SENTRY_BACKEND_DSN: config.sentry_backend_dsn, TARGET_ENVIRONMENT: environment, }