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 bea774315..f67f982f0 100644 --- a/backend/openapi/generateOpenApiDoc.ts +++ b/backend/openapi/generateOpenApiDoc.ts @@ -40,6 +40,8 @@ delete ModelInfo.Schemas.Reference.$id; delete ModelInfo.Schemas.POST.Response.Payload.$id; delete ModelInfo.Schemas.POST.Request.Payload.$id; delete ModelInfo.Schemas.GET.Response.Payload.$id; +delete ModelInfo.ModelInfo.EmbeddingProcessStates.POST.Schemas.Request.Payload.$id; +delete ModelInfo.ModelInfo.EmbeddingProcessStates.POST.Schemas.Response.Payload.$id; delete Locale.Schemas.Payload.$id; delete Presigned.Schemas.GET.Response.Payload.$id; delete Import.Schemas.POST.Request.Payload.$id; @@ -435,11 +437,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, LocaleSchema: Locale.Schemas.Payload, ImportSchema: Import.Schemas.POST.Request.Payload, ExportSchema: Export.Schemas.POST.Request.Payload, diff --git a/backend/package.json b/backend/package.json index 86549a590..236ea2b1f 100644 --- a/backend/package.json +++ b/backend/package.json @@ -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/common/validations/validateRequest.test.ts b/backend/src/common/validations/validateRequest.test.ts new file mode 100644 index 000000000..8832f17ca --- /dev/null +++ b/backend/src/common/validations/validateRequest.test.ts @@ -0,0 +1,138 @@ +// silence chatty console +import "_test_utilities/consoleMock"; + +import { APIGatewayProxyEvent } from "aws-lambda"; +import { SchemaObject } from "ajv"; +import ErrorAPISpecs from "api-specifications/error"; +import { validateEvent, validateSchema } from "./validateRequest"; +import { StatusCodes } from "server/httpUtils"; + +const givenSchema: SchemaObject = { + $id: "test://validateRequest/testPayloadSchema", + type: "object", + properties: { + name: { type: "string" }, + }, + required: ["name"], + additionalProperties: false, +}; + +const givenMaxPayloadLength = 100; + +function getEvent(options: { body?: string; contentType?: string | null }): APIGatewayProxyEvent { + const headers: Record = {}; + if (options.contentType !== null) { + headers["Content-Type"] = options.contentType ?? "application/json"; + } + return { + body: options.body ?? JSON.stringify({ name: "foo" }), + headers, + } as never; +} + +describe("Test the validateSchema function", () => { + test("should return the payload when it conforms to the schema", () => { + // GIVEN a payload that conforms to the schema + const givenPayload = { name: "foo" }; + + // WHEN validating the payload against the schema + const actualResult = validateSchema(givenSchema, givenPayload); + + // THEN expect the payload to be returned without an error response + expect(actualResult.payload).toEqual(givenPayload); + expect(actualResult.errorResponse).toBeUndefined(); + }); + + test("should return the INVALID_JSON_SCHEMA error response when the payload does not conform to the schema", () => { + // GIVEN a payload that does not conform to the schema + const givenPayload = { foo: "bar" }; + + // WHEN validating the payload against the schema + const actualResult = validateSchema(givenSchema, givenPayload); + + // THEN expect an error response with the BAD_REQUEST status code and the INVALID_JSON_SCHEMA error code + expect(actualResult.payload).toBeUndefined(); + expect(actualResult.errorResponse!.statusCode).toEqual(StatusCodes.BAD_REQUEST); + expect(JSON.parse(actualResult.errorResponse!.body).errorCode).toEqual( + ErrorAPISpecs.Constants.ErrorCodes.INVALID_JSON_SCHEMA + ); + }); +}); + +describe("Test the validateEvent function", () => { + test("should return the parsed payload when the event is valid", () => { + // GIVEN a valid event + const givenPayload = { name: "foo" }; + const givenEvent = getEvent({ body: JSON.stringify(givenPayload) }); + + // WHEN validating the event + const actualResult = validateEvent(givenEvent, givenSchema, givenMaxPayloadLength); + + // THEN expect the parsed payload to be returned without an error response + expect(actualResult.payload).toEqual(givenPayload); + expect(actualResult.errorResponse).toBeUndefined(); + }); + + test.each([ + ["is not json", "text/html"], + ["is missing", null], + ])( + "should return the UNSUPPORTED_MEDIA_TYPE error response when the content type %s", + (_description, givenContentType) => { + // GIVEN an event with an invalid content type + const givenEvent = getEvent({ contentType: givenContentType }); + + // WHEN validating the event + const actualResult = validateEvent(givenEvent, givenSchema, givenMaxPayloadLength); + + // THEN expect an error response with the UNSUPPORTED_MEDIA_TYPE status code + expect(actualResult.payload).toBeUndefined(); + expect(actualResult.errorResponse!.statusCode).toEqual(StatusCodes.UNSUPPORTED_MEDIA_TYPE); + } + ); + + test("should return the TOO_LARGE_PAYLOAD error response when the body is longer than the maximum payload length", () => { + // GIVEN an event with a body that is longer than the maximum payload length + const givenEvent = getEvent({ body: "a".repeat(givenMaxPayloadLength + 1) }); + + // WHEN validating the event + const actualResult = validateEvent(givenEvent, givenSchema, givenMaxPayloadLength); + + // THEN expect an error response with the TOO_LARGE_PAYLOAD status code + expect(actualResult.payload).toBeUndefined(); + expect(actualResult.errorResponse!.statusCode).toEqual(StatusCodes.TOO_LARGE_PAYLOAD); + expect(JSON.parse(actualResult.errorResponse!.body).errorCode).toEqual( + ErrorAPISpecs.Constants.ErrorCodes.TOO_LARGE_PAYLOAD + ); + }); + + test("should return the MALFORMED_BODY error response when the body is not valid json", () => { + // GIVEN an event with a body that is not valid json + const givenEvent = getEvent({ body: "{ not json" }); + + // WHEN validating the event + const actualResult = validateEvent(givenEvent, givenSchema, givenMaxPayloadLength); + + // THEN expect an error response with the BAD_REQUEST status code and the MALFORMED_BODY error code + expect(actualResult.payload).toBeUndefined(); + expect(actualResult.errorResponse!.statusCode).toEqual(StatusCodes.BAD_REQUEST); + expect(JSON.parse(actualResult.errorResponse!.body).errorCode).toEqual( + ErrorAPISpecs.Constants.ErrorCodes.MALFORMED_BODY + ); + }); + + test("should return the INVALID_JSON_SCHEMA error response when the body does not conform to the schema", () => { + // GIVEN an event with a body that does not conform to the schema + const givenEvent = getEvent({ body: JSON.stringify({ foo: "bar" }) }); + + // WHEN validating the event + const actualResult = validateEvent(givenEvent, givenSchema, givenMaxPayloadLength); + + // THEN expect an error response with the BAD_REQUEST status code and the INVALID_JSON_SCHEMA error code + expect(actualResult.payload).toBeUndefined(); + expect(actualResult.errorResponse!.statusCode).toEqual(StatusCodes.BAD_REQUEST); + expect(JSON.parse(actualResult.errorResponse!.body).errorCode).toEqual( + ErrorAPISpecs.Constants.ErrorCodes.INVALID_JSON_SCHEMA + ); + }); +}); diff --git a/backend/src/common/validations/validateRequest.ts b/backend/src/common/validations/validateRequest.ts new file mode 100644 index 000000000..ba90ebbd5 --- /dev/null +++ b/backend/src/common/validations/validateRequest.ts @@ -0,0 +1,64 @@ +import { APIGatewayProxyEvent, APIGatewayProxyResult } from "aws-lambda"; +import { SchemaObject, ValidateFunction } from "ajv"; +import { ajvInstance, ParseValidationError } from "validator"; +import { STD_ERRORS_RESPONSES } from "server/httpUtils"; + +/** + * The result of validating a request: either the typed, validated payload, + * or the error response that the handler should return to the caller. + */ +export type RequestValidationResult = + | { payload: T; errorResponse?: never } + | { payload?: never; errorResponse: APIGatewayProxyResult }; + +/** + * Validates a payload against the given JSON schema. + * + * @param schema - The JSON schema to validate against. If it is already registered in the ajv instance it is reused, otherwise it is compiled. + * @param payload - The payload to validate. + * @return The typed payload when it is valid, otherwise the INVALID_JSON_SCHEMA error response to return to the caller. + */ +export function validateSchema(schema: SchemaObject, payload: unknown): RequestValidationResult { + const validateFunction = (ajvInstance.getSchema(schema.$id as string) ?? + ajvInstance.compile(schema)) as ValidateFunction; + const isValid = validateFunction(payload); + if (!isValid) { + const errorDetail = ParseValidationError(validateFunction.errors); + return { errorResponse: STD_ERRORS_RESPONSES.INVALID_JSON_SCHEMA_ERROR(errorDetail) }; + } + return { payload: payload as T }; +} + +/** + * Validates a JSON request event: the content type, the payload length, that the body is valid JSON, + * and that it conforms to the given JSON schema. + * + * @param event - The API Gateway event to validate. + * @param schema - The JSON schema the body must conform to. + * @param maxPayloadLength - The maximum allowed length of the body. + * @return The parsed and validated payload, otherwise the error response to return to the caller. + */ +export function validateEvent( + event: APIGatewayProxyEvent, + schema: SchemaObject, + maxPayloadLength: number +): RequestValidationResult { + if (!event.headers["Content-Type"]?.includes("application/json")) { + return { errorResponse: STD_ERRORS_RESPONSES.UNSUPPORTED_MEDIA_TYPE_ERROR }; + } + + if (event.body?.length && event.body.length > maxPayloadLength) { + return { + errorResponse: STD_ERRORS_RESPONSES.TOO_LARGE_PAYLOAD_ERROR(`Expected maximum length is ${maxPayloadLength}`), + }; + } + + let payload: unknown; + try { + payload = JSON.parse(event.body as string); + } catch (error: unknown) { + return { errorResponse: STD_ERRORS_RESPONSES.MALFORMED_BODY_ERROR((error as Error).message) }; + } + + return validateSchema(schema, payload); +} 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..d6da167ca --- /dev/null +++ b/backend/src/embeddings/embeddingProcess/embeddingProcess.service.test.ts @@ -0,0 +1,332 @@ +// silence chatty console +import "_test_utilities/consoleMock"; + +import { Readable } from "node:stream"; +import ModelInfoApiSpecs from "api-specifications/modelInfo"; +import { EmbeddingProcessService, TASKS_FLUSH_SIZE } 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()), + deleteById: jest.fn().mockResolvedValue(undefined), + }; + 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 = { pushTasksToQueue: 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 modelId is not a valid object id", async () => { + // GIVEN a modelId that is not a valid object id + const givenInvalidModelId = "not-a-valid-object-id"; + const { service, modelRepository, embeddingProcessStateRepository } = setupService({ model: getMockModel(true) }); + + // WHEN triggering the embedding process with the invalid modelId + const actualPromise = service.triggerEmbeddingProcess(givenInvalidModelId, givenEmbeddingServiceId); + + // THEN expect it to reject with a ModelNotFoundError + await expect(actualPromise).rejects.toThrow(ModelNotFoundError); + // AND expect the model to not be queried + expect(modelRepository.getModelById).not.toHaveBeenCalled(); + // AND expect no embedding process state to be created + expect(embeddingProcessStateRepository.create).not.toHaveBeenCalled(); + }); + + 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 batch of tasks to be pushed per entity type + expect(embeddingClient.pushTasksToQueue).toHaveBeenCalledTimes(4); + // AND expect the tasks for the skills to be pushed with the skill fields + expect(embeddingClient.pushTasksToQueue).toHaveBeenCalledWith( + givenSkills.map((skill) => ({ + modelId: givenModelId, + entityId: skill.id, + entityType: EmbeddableEntityType.Skill, + fields: [EmbeddableField.preferredLabel, EmbeddableField.description, EmbeddableField.altLabels], + })) + ); + // AND expect the tasks for the skill groups to be pushed with the skill group fields + expect(embeddingClient.pushTasksToQueue).toHaveBeenCalledWith([ + { + modelId: givenModelId, + entityId: givenSkillGroups[0].id, + entityType: EmbeddableEntityType.SkillGroup, + fields: [ + EmbeddableField.preferredLabel, + EmbeddableField.description, + EmbeddableField.altLabels, + EmbeddableField.scopeNote, + ], + }, + ]); + // AND expect the tasks for the occupations to be pushed + expect(embeddingClient.pushTasksToQueue).toHaveBeenCalledWith([ + { + modelId: givenModelId, + entityId: givenOccupations[0].id, + entityType: EmbeddableEntityType.Occupation, + fields: [EmbeddableField.preferredLabel, EmbeddableField.description, EmbeddableField.altLabels], + }, + ]); + // AND expect the tasks for the occupation groups to be pushed + expect(embeddingClient.pushTasksToQueue).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); + }); + + test("should flush the tasks to the queue in chunks when an entity type has more entities than the flush size", async () => { + // GIVEN a released model with no unfinished embedding process + // AND more skills than fit in a single flush + const givenSkillCount = TASKS_FLUSH_SIZE + 2; + const givenSkills = Array.from({ length: givenSkillCount }, (_, index) => ({ id: getMockStringId(index + 100) })); + const { service, embeddingClient, embeddingProcessStateRepository } = setupService({ + model: getMockModel(true), + pendingProcess: null, + skills: givenSkills, + }); + + // WHEN triggering the embedding process + await service.triggerEmbeddingProcess(givenModelId, givenEmbeddingServiceId); + + // THEN expect the tasks for the skills to be flushed in chunks of at most the flush size + const actualPushedBatches = (embeddingClient.pushTasksToQueue as jest.Mock).mock.calls.map((call) => call[0]); + actualPushedBatches.forEach((actualBatch) => { + expect(actualBatch.length).toBeLessThanOrEqual(TASKS_FLUSH_SIZE); + }); + // AND expect all the skills to be pushed exactly once, in order + const actualPushedEntityIds = actualPushedBatches.flat().map((task) => task.entityId); + const expectedEntityIds = givenSkills.map((skill) => skill.id); + expect(actualPushedEntityIds).toEqual(expectedEntityIds); + // AND expect the process state to be updated with the total number of skills + expect(embeddingProcessStateRepository.update).toHaveBeenCalledWith(expect.any(String), { + status: ModelInfoApiSpecs.ModelInfo.EmbeddingProcessStates.Enums.Status.IN_PROGRESS, + totalDocuments: givenSkillCount, + }); + }); + + test("should delete the created process state when pushing the tasks to the queue fails", async () => { + // GIVEN a released model with no unfinished embedding process and some skills + const givenSkills = [{ id: getMockStringId(101) }]; + // AND a created process state + const givenCreatedProcessState = getMockEmbeddingProcessState({ id: getMockStringId(10) }); + const { service, embeddingClient, embeddingProcessStateRepository } = setupService({ + model: getMockModel(true), + pendingProcess: null, + skills: givenSkills, + createdProcessState: givenCreatedProcessState, + }); + // AND pushing the tasks to the queue will fail + const givenCause = new Error("SQS unavailable"); + (embeddingClient.pushTasksToQueue as jest.Mock).mockRejectedValue(givenCause); + + // WHEN triggering the embedding process + const actualPromise = service.triggerEmbeddingProcess(givenModelId, givenEmbeddingServiceId); + + // THEN expect it to reject with a wrapped error + await expect(actualPromise).rejects.toThrow( + expect.toMatchErrorWithCause(`Failed to trigger embedding process for model ${givenModelId}.`, givenCause.message) + ); + // AND expect the created process state to be deleted so that the model is not blocked by an orphaned record + expect(embeddingProcessStateRepository.deleteById).toHaveBeenCalledWith(givenCreatedProcessState.id); + }); + + test("should still reject with the original error when the cleanup of the created process state fails", async () => { + // GIVEN a released model with no unfinished embedding process and some skills + const givenSkills = [{ id: getMockStringId(101) }]; + // AND a created process state + const givenCreatedProcessState = getMockEmbeddingProcessState({ id: getMockStringId(10) }); + const { service, embeddingClient, embeddingProcessStateRepository } = setupService({ + model: getMockModel(true), + pendingProcess: null, + skills: givenSkills, + createdProcessState: givenCreatedProcessState, + }); + // AND pushing the tasks to the queue will fail + const givenCause = new Error("SQS unavailable"); + (embeddingClient.pushTasksToQueue as jest.Mock).mockRejectedValue(givenCause); + // AND deleting the created process state will also fail + (embeddingProcessStateRepository.deleteById as jest.Mock).mockRejectedValue(new Error("delete failed")); + + // WHEN triggering the embedding process + const actualPromise = service.triggerEmbeddingProcess(givenModelId, givenEmbeddingServiceId); + + // THEN expect it to reject with the original wrapped error + await expect(actualPromise).rejects.toThrow( + expect.toMatchErrorWithCause(`Failed to trigger embedding process for model ${givenModelId}.`, givenCause.message) + ); + // AND expect the cleanup to have been attempted + expect(embeddingProcessStateRepository.deleteById).toHaveBeenCalledWith(givenCreatedProcessState.id); + }); +}); diff --git a/backend/src/embeddings/embeddingProcess/embeddingProcess.service.ts b/backend/src/embeddings/embeddingProcess/embeddingProcess.service.ts new file mode 100644 index 000000000..18325c3eb --- /dev/null +++ b/backend/src/embeddings/embeddingProcess/embeddingProcess.service.ts @@ -0,0 +1,181 @@ +import { Readable } from "node:stream"; +import mongoose from "mongoose"; +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, IGenerateEmbeddingTask } 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, + ], +}; + +/** + * Number of tasks to accumulate from the entity stream before flushing them to the queue + * in a single pushTasksToQueue call, so that the entire entity set is never held in memory at once. + */ +export const TASKS_FLUSH_SIZE = 1000; + +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 { + // An invalid object id can never reference an existing model, + // and would cause a CastError in the queries below. + if (!mongoose.Types.ObjectId.isValid(modelId)) { + throw new ModelNotFoundError(modelId); + } + + 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); + } + + let embeddingProcessState: IEmbeddingProcessState; + try { + // 3. Create a new embedding process state. + embeddingProcessState = await this.embeddingProcessStateRepository.create({ + modelId, + status: ModelInfoApiSpecs.ModelInfo.EmbeddingProcessStates.Enums.Status.PENDING, + embeddingServiceId, + totalDocuments: 0, + errorCounts: 0, + warningCounts: 0, + completedDocuments: 0, + }); + } catch (e) { + throw new Error(`Failed to trigger embedding process for model ${modelId}.`, { cause: e }); + } + + try { + // 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 await this.embeddingProcessStateRepository.update(embeddingProcessState.id, { + status: ModelInfoApiSpecs.ModelInfo.EmbeddingProcessStates.Enums.Status.IN_PROGRESS, + totalDocuments, + }); + } catch (e) { + // Clean up the created process state so that a failed attempt does not leave an orphaned + // unfinished record behind, which would block any future embedding process for the model. + try { + await this.embeddingProcessStateRepository.deleteById(embeddingProcessState.id); + } catch (cleanupError) { + console.error( + new Error( + `Failed to clean up the embedding process state ${embeddingProcessState.id} after a failed trigger for model ${modelId}.`, + { cause: cleanupError } + ) + ); + } + 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; + let tasksBuffer: IGenerateEmbeddingTask[] = []; + for await (const entity of stream) { + tasksBuffer.push({ + modelId, + entityId: (entity as { id: string }).id, + entityType, + fields: EMBEDDABLE_FIELDS_BY_ENTITY_TYPE[entityType], + }); + count++; + if (tasksBuffer.length >= TASKS_FLUSH_SIZE) { + await this.embeddingClient.pushTasksToQueue(tasksBuffer); + tasksBuffer = []; + } + } + if (tasksBuffer.length > 0) { + await this.embeddingClient.pushTasksToQueue(tasksBuffer); + } + 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..a3e292188 --- /dev/null +++ b/backend/src/embeddings/embeddingProcessState/embeddingProcessStateRepository.test.ts @@ -0,0 +1,265 @@ +// 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)); + }); + }); + + describe("Test deleteById() EmbeddingProcessState", () => { + test("should successfully delete an EmbeddingProcessState by id", async () => { + // GIVEN an EmbeddingProcessState in the database + const givenNewEmbeddingProcessStateSpec = getNewEmbeddingProcessStateSpec(); + const createdEmbeddingProcessState = await repository.create(givenNewEmbeddingProcessStateSpec); + + // WHEN deleting the EmbeddingProcessState by id + await repository.deleteById(createdEmbeddingProcessState.id); + + // THEN expect the EmbeddingProcessState to no longer be found + const actualFoundEmbeddingProcessState = await repository.findById(createdEmbeddingProcessState.id); + expect(actualFoundEmbeddingProcessState).toBeNull(); + }); + + test("should resolve without an error when deleting an EmbeddingProcessState that does not exist", async () => { + // GIVEN an id of an EmbeddingProcessState that does not exist + const givenId = getMockStringId(1); + + // WHEN deleting the EmbeddingProcessState by id + const actualDeletePromise = repository.deleteById(givenId); + + // THEN expect it to resolve without an error + await expect(actualDeletePromise).resolves.toBeUndefined(); + }); + + TestDBConnectionFailureNoSetup((repositoryRegistry) => { + return repositoryRegistry.embeddingProcessState.deleteById(getMockStringId(1)); + }); + }); +}); diff --git a/backend/src/embeddings/embeddingProcessState/embeddingProcessStateRepository.ts b/backend/src/embeddings/embeddingProcessState/embeddingProcessStateRepository.ts new file mode 100644 index 000000000..b20b3bfdf --- /dev/null +++ b/backend/src/embeddings/embeddingProcessState/embeddingProcessStateRepository.ts @@ -0,0 +1,155 @@ +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; + + /** + * Deletes the EmbeddingProcessState entry with the given ID. + * + * @param {string} id - The unique ID of the EmbeddingProcessState entry. + * @return {Promise} - A Promise that resolves when the entry has been deleted or does not exist. + * Rejects with an error if the operation fails. + */ + deleteById(id: 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; + } + } + + async deleteById(id: string): Promise { + try { + await this.Model.findByIdAndDelete(id).exec(); + } catch (e: unknown) { + const err = new Error("EmbeddingProcessStateRepository.deleteById: deleteById 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 index 2be624fc9..45858c6b4 100644 --- a/backend/src/embeddings/handler/index.ts +++ b/backend/src/embeddings/handler/index.ts @@ -5,6 +5,7 @@ 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); @@ -45,6 +46,19 @@ async function handleRecord(embeddingService: EmbeddingService, record: SQSRecor 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/service/client.test.ts b/backend/src/embeddings/service/client.test.ts new file mode 100644 index 000000000..167d5bfda --- /dev/null +++ b/backend/src/embeddings/service/client.test.ts @@ -0,0 +1,199 @@ +// silence chatty console +import "_test_utilities/consoleMock"; + +import { SendMessageBatchCommand, SendMessageCommand, SQSClient } from "@aws-sdk/client-sqs"; +import { EmbeddingClient, SQS_MAX_BATCH_SIZE } 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) + ); + }); + + describe("pushTasksToQueue", () => { + function getValidTasks(count: number): IGenerateEmbeddingTask[] { + return Array.from({ length: count }, (_, index) => ({ + ...getValidTask(), + entityId: getMockStringId(index + 10), + })); + } + + test("should push multiple tasks to the queue in batches of the maximum SQS batch size", async () => { + // GIVEN a mock SQS client + const givenSqsClient = getMockSqsClient(jest.fn().mockResolvedValue({ Failed: [] })); + // AND an EmbeddingClient that uses it + const givenEmbeddingClient = new EmbeddingClient(givenSqsClient); + // AND more valid tasks than fit in a single SQS batch + const givenTaskCount = SQS_MAX_BATCH_SIZE * 2 + 5; + const givenTasks = getValidTasks(givenTaskCount); + + // WHEN pushing the tasks to the queue + await givenEmbeddingClient.pushTasksToQueue(givenTasks); + + // THEN expect the SQS client to send one batch command per chunk of the maximum batch size + const expectedBatchCount = Math.ceil(givenTaskCount / SQS_MAX_BATCH_SIZE); + expect(givenSqsClient.send).toHaveBeenCalledTimes(expectedBatchCount); + // AND expect every batch to target the queue url and carry the tasks as message bodies + const actualCommands = (givenSqsClient.send as jest.Mock).mock.calls.map( + (call) => call[0] as SendMessageBatchCommand + ); + actualCommands.forEach((actualCommand) => { + expect(actualCommand).toBeInstanceOf(SendMessageBatchCommand); + expect(actualCommand.input.QueueUrl).toEqual(givenQueueUrl); + expect(actualCommand.input.Entries!.length).toBeLessThanOrEqual(SQS_MAX_BATCH_SIZE); + }); + const actualMessageBodies = actualCommands.flatMap((command) => + command.input.Entries!.map((entry) => entry.MessageBody) + ); + const expectedMessageBodies = givenTasks.map((task) => JSON.stringify(task)); + expect(actualMessageBodies).toEqual(expectedMessageBodies); + }); + + test("should not push any task to the queue and should throw when one of the tasks is invalid", async () => { + // GIVEN a mock SQS client + const givenSqsClient = getMockSqsClient(); + // AND an EmbeddingClient that uses it + const givenEmbeddingClient = new EmbeddingClient(givenSqsClient); + // AND a list of tasks where one is invalid (missing the required fields property) + const givenInvalidTask = { ...getValidTask(), fields: [] }; + const givenTasks = [...getValidTasks(3), givenInvalidTask as IGenerateEmbeddingTask]; + + // WHEN pushing the tasks to the queue + const actualPromise = givenEmbeddingClient.pushTasksToQueue(givenTasks); + + // 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 a batch", 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 valid tasks + const givenTasks = getValidTasks(3); + + // WHEN pushing the tasks to the queue + const actualPromise = givenEmbeddingClient.pushTasksToQueue(givenTasks); + + // THEN expect it to reject with a wrapped error + await expect(actualPromise).rejects.toThrow( + expect.toMatchErrorWithCause( + "EmbeddingClient.pushTasksToQueue: failed to push tasks to queue", + givenCause.message + ) + ); + }); + + test("should reject with an error when some messages in a batch fail to be sent", async () => { + // GIVEN a mock SQS client that reports a partial batch failure + const givenFailedEntry = { Id: "0", SenderFault: true, Code: "InternalError", Message: "some error" }; + const givenSqsClient = getMockSqsClient(jest.fn().mockResolvedValue({ Failed: [givenFailedEntry] })); + // AND an EmbeddingClient that uses it + const givenEmbeddingClient = new EmbeddingClient(givenSqsClient); + // AND valid tasks + const givenTasks = getValidTasks(3); + + // WHEN pushing the tasks to the queue + const actualPromise = givenEmbeddingClient.pushTasksToQueue(givenTasks); + + // THEN expect it to reject with a wrapped error that reports the failed messages + await expect(actualPromise).rejects.toThrow( + expect.toMatchErrorWithCause( + "EmbeddingClient.pushTasksToQueue: failed to push tasks to queue", + /some messages failed to be sent/ + ) + ); + }); + }); +}); diff --git a/backend/src/embeddings/service/client.ts b/backend/src/embeddings/service/client.ts index b0394abba..757d5e070 100644 --- a/backend/src/embeddings/service/client.ts +++ b/backend/src/embeddings/service/client.ts @@ -1,12 +1,87 @@ +import { SQSClient, SendMessageCommand, SendMessageBatchCommand } from "@aws-sdk/client-sqs"; import { IGenerateEmbeddingTask } from "./types"; +import { getEmbeddingsQueueUrl } from "server/config/config"; +import { validateEmbeddingQueueJob } from "embeddings/specs/queueJob.schema"; + +// SQS SendMessageBatch accepts at most 10 messages per request (AWS hard limit). +// reference: https://docs.aws.amazon.com/AWSJavaScriptSDK/v3/latest/client/sqs/command/SendMessageBatchCommand/ +export const SQS_MAX_BATCH_SIZE = 10; export interface IEmbeddingClient { pushTaskToQueue(task: IGenerateEmbeddingTask): Promise; + pushTasksToQueue(tasks: IGenerateEmbeddingTask[]): Promise; } export class EmbeddingClient implements IEmbeddingClient { - pushTaskToQueue(task: IGenerateEmbeddingTask): Promise { - console.debug("EmbeddingClient.pushTaskToQueue", task); - throw new Error("Method not implemented."); + 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; + } + } + + async pushTasksToQueue(tasks: IGenerateEmbeddingTask[]): Promise { + // Validate all the tasks against the queue job schema before pushing any of them, + // so that a malformed job never ends up on the queue. + for (const task of tasks) { + const isValid = validateEmbeddingQueueJob(task); + if (!isValid) { + const err = new Error( + `EmbeddingClient.pushTasksToQueue: a task does not conform to the queue job schema: ${JSON.stringify( + validateEmbeddingQueueJob.errors + )}` + ); + console.error(err); + throw err; + } + } + + try { + for (let i = 0; i < tasks.length; i += SQS_MAX_BATCH_SIZE) { + const batch = tasks.slice(i, i + SQS_MAX_BATCH_SIZE); + const response = await this.sqsClient.send( + new SendMessageBatchCommand({ + QueueUrl: getEmbeddingsQueueUrl(), + Entries: batch.map((task, index) => ({ + // The ID only needs to be unique within a single batch request, + // it is used to correlate the Successful/Failed entries of the response. + Id: `${i + index}`, + MessageBody: JSON.stringify(task), + })), + }) + ); + // SendMessageBatch can partially fail without the send() call throwing. + if (response.Failed && response.Failed.length > 0) { + throw new Error(`some messages failed to be sent: ${JSON.stringify(response.Failed)}`); + } + } + } catch (e: unknown) { + const err = new Error("EmbeddingClient.pushTasksToQueue: failed to push tasks 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/types.ts b/backend/src/embeddings/service/types.ts index 099c1f450..39673c575 100644 --- a/backend/src/embeddings/service/types.ts +++ b/backend/src/embeddings/service/types.ts @@ -16,5 +16,5 @@ export interface IGenerateEmbeddingTask { modelId: string; entityId: string; entityType: EmbeddableEntityType; - fields: EmbeddableField; + 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 9de79ceb6..7784948b5 100644 --- a/backend/src/index.ts +++ b/backend/src/index.ts @@ -66,7 +66,7 @@ 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)) { + } else if (pathToRegexp([Routes.MODEL_EMBEDDING_PROCESSES_ROUTE]).regexp.test(path)) { return ModelEmbeddingProcessesHandler(event); } else if ( pathToRegexp([ 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..041d8f6a6 --- /dev/null +++ b/backend/src/modelInfo/embeddingProcesses/POST/index.integration.test.ts @@ -0,0 +1,306 @@ +// 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 (SendMessageBatchCommand, 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 { SendMessageBatchCommand } 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 { SQS_MAX_BATCH_SIZE } from "embeddings/service/client"; +import { validateEmbeddingQueueJob } from "embeddings/specs/queueJob.schema"; + +const givenEntityCountPerType = 100; +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 valid SQS message batches covering 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 SendMessageBatchCommands of at most the maximum SQS batch size, + // one chunk per entity type since each entity type is flushed to the queue separately + const expectedBatchCount = 4 * Math.ceil(givenEntityCountPerType / SQS_MAX_BATCH_SIZE); + expect(mockSQSClientSend).toHaveBeenCalledTimes(expectedBatchCount); + const actualCommands = mockSQSClientSend.mock.calls.map(([command]) => command); + for (const actualCommand of actualCommands) { + expect(actualCommand).toBeInstanceOf(SendMessageBatchCommand); + expect(actualCommand.input.QueueUrl).toEqual(getEmbeddingsQueueUrl()); + expect(actualCommand.input.Entries.length).toBeLessThanOrEqual(SQS_MAX_BATCH_SIZE); + } + // AND every message body of every batch to be a valid embedding queue job + const actualTasks: IGenerateEmbeddingTask[] = actualCommands.flatMap((command) => + command.input.Entries.map((entry: { MessageBody: string }) => JSON.parse(entry.MessageBody)) + ); + 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..3de61a8d5 --- /dev/null +++ b/backend/src/modelInfo/embeddingProcesses/POST/index.test.ts @@ -0,0 +1,239 @@ +// 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 MODEL_NOT_FOUND_BY_ID error code as documented in the API specification + expect(actualResponse.statusCode).toEqual(StatusCodes.NOT_FOUND); + expect(JSON.parse(actualResponse.body).errorCode).toEqual( + ModelInfoApiSpecs.ModelInfo.EmbeddingProcessStates.POST.Enums.Response.Status404.ErrorCodes.MODEL_NOT_FOUND_BY_ID + ); + }); + + 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 index d66e963a4..2355d227e 100644 --- a/backend/src/modelInfo/embeddingProcesses/POST/index.ts +++ b/backend/src/modelInfo/embeddingProcesses/POST/index.ts @@ -1,8 +1,150 @@ import { APIGatewayProxyEvent, APIGatewayProxyResult } from "aws-lambda"; +import ModelInfoApiSpecs from "api-specifications/modelInfo"; +import ErrorAPISpecs from "api-specifications/error"; +import AuthAPISpecs from "api-specifications/auth"; +import { errorResponse, responseJSON, StatusCodes } from "server/httpUtils"; +import { validateEvent } from "common/validations/validateRequest"; +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 { - async handle(_event: APIGatewayProxyEvent): Promise { - console.log(_event); - throw new Error("Not implemented yet"); + @RoleRequired(AuthAPISpecs.Enums.TabiyaRoles.MODEL_MANAGER) // Applying role-based access control + async handle(event: APIGatewayProxyEvent): Promise { + const validationResult = + validateEvent( + event, + ModelInfoApiSpecs.ModelInfo.EmbeddingProcessStates.POST.Schemas.Request.Payload, + ModelInfoApiSpecs.ModelInfo.EmbeddingProcessStates.POST.Constants.MAX_PAYLOAD_LENGTH + ); + if (validationResult.errorResponse) { + return validationResult.errorResponse; + } + const payload = validationResult.payload; + + 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 }) + ); + } + + 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, + ModelInfoApiSpecs.ModelInfo.EmbeddingProcessStates.POST.Enums.Response.Status404.ErrorCodes + .MODEL_NOT_FOUND_BY_ID, + "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/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 c26323619..12afd8213 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 02b29ac8f..bdfe3be8b 100644 --- a/backend/src/validator.ts +++ b/backend/src/validator.ts @@ -235,6 +235,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"