From 45fbcc540303e7f9fa00c1820413b0eec40e9eb5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Luk=C3=A1=C5=A1=20Jane=C4=8Dek?= Date: Wed, 11 Dec 2024 14:45:34 +0100 Subject: [PATCH] fixup! feat(log): add basic tracking MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Lukáš Janeček --- src/common/log.entity.ts | 38 ++++++++------- src/common/log.subscriber.ts | 94 ++++++++++++++++++++---------------- 2 files changed, 73 insertions(+), 59 deletions(-) diff --git a/src/common/log.entity.ts b/src/common/log.entity.ts index 0b3fb01..217717d 100644 --- a/src/common/log.entity.ts +++ b/src/common/log.entity.ts @@ -14,12 +14,11 @@ * limitations under the License. */ -import { ChangeSetType, Entity, ManyToOne, PrimaryKey, Property, ref, Ref } from '@mikro-orm/core'; -import { User } from '@zilliz/milvus2-sdk-node'; -import { requestContext } from '@fastify/request-context'; +import { ChangeSetType, Entity, PrimaryKey, Property } from '@mikro-orm/core'; +import { requestContext, RequestContextData } from '@fastify/request-context'; import { generatePrefixedObjectId } from '@/utils/id.js'; -import { ProjectPrincipal } from '@/administration/entities/project-principal.entity'; +import { jobLocalStorage } from '@/context'; @Entity() export class Log { @@ -29,11 +28,8 @@ export class Log { @Property() createdAt: Date = new Date(); - @ManyToOne() - projectPrincipal?: Ref; - - @ManyToOne() - user?: Ref; + @Property() + requestContext: any = {}; @Property() entity?: string; @@ -48,16 +44,24 @@ export class Log { change?: any; @Property() - additionalData: any; + additionalData?: any; constructor({ entity, entityId, type, change, additionalData }: LogInput) { - const user = requestContext.get('user'); - if (user) { - this.user = ref(user); - } - const projectPrincipal = requestContext.get('projectPrincipal'); - if (projectPrincipal) { - this.projectPrincipal = ref(projectPrincipal); + const contextData: (keyof RequestContextData)[] = [ + 'user', + 'organizationUser', + 'apiKey', + 'projectPrincipal', + 'artifact' + ]; + + contextData.forEach((key: keyof RequestContextData) => { + const value = requestContext.get(key); + if (value) Object.assign(this.requestContext, { [key]: value.id }); + }); + const job = jobLocalStorage.getStore()?.job.name; + if (job) { + Object.assign(this.requestContext, { job }); } this.entity = entity; this.entityId = entityId; diff --git a/src/common/log.subscriber.ts b/src/common/log.subscriber.ts index 634e7c5..fde4386 100644 --- a/src/common/log.subscriber.ts +++ b/src/common/log.subscriber.ts @@ -4,82 +4,92 @@ import { Log } from './log.entity'; import { BaseEntity } from './base.entity'; import { inJob, inSeeder } from '@/context'; +import { getLogger } from '@/logger'; const loggedEntities: { [key: string]: { types?: ChangeSetType[]; entities: string[] } } = { - Assistant: { + assistant: { entities: ['project', 'agent'] }, - Artifact: { + artifact: { entities: ['project', 'thread', 'name'] }, - Chat: { + chat: { types: [ChangeSetType.CREATE], - entities: ['artifact'] + entities: [] }, - Message: { + message: { entities: ['project'] }, - Thread: { + thread: { entities: ['project'] }, - Tool: { + tool: { entities: ['project', 'name'] }, - VectorStore: { + 'vector-store': { entities: ['project', 'name'] }, - VectorStoreFile: { + 'vector-store-file': { entities: ['project', 'file'] }, - File: { + file: { entities: ['project', 'filename'] }, - Run: { + run: { entities: ['project', 'assistant', 'status'] }, - User: { + user: { entities: ['email'] }, - Organization: { + organization: { entities: ['name'] }, - Project: { + project: { entities: ['name', 'organization'] }, - ApiKey: { + 'project-api-key': { entities: ['project'] } }; export class LogSubscriber implements EventSubscriber { onFlush(args: FlushEventArgs): void | Promise { - args.uow.getChangeSets().forEach((cs) => { - if ( - loggedEntities[cs.name] && - (loggedEntities[cs.name].types?.includes(cs.type) ?? true) && - inSeeder() === false - ) { - if (cs.type === ChangeSetType.DELETE && inJob()) return; + try { + args.uow.getChangeSets().forEach((cs) => { + if ( + loggedEntities[cs.collection] && + (loggedEntities[cs.collection].types?.includes(cs.type) ?? true) && + inSeeder() === false + ) { + // Do not log cleanup deletes from jobs + if (cs.type === ChangeSetType.DELETE && inJob()) return; - const log = new Log({ - entity: cs.name, - entityId: cs.entity?.id, - type: cs.type, - change: cs.type === ChangeSetType.UPDATE ? cs.payload : undefined, - additionalData: - loggedEntities[cs.name].entities.reduce( - (acc, name) => ({ - ...acc, - [name]: - cs.entity[name].entity instanceof BaseEntity - ? cs.entity[name].id - : cs.entity[name] - }), - {} - ) ?? undefined - }); - args.uow.computeChangeSet(log); - } - }); + const log = new Log({ + entity: cs.name, + entityId: cs.entity?.id, + type: cs.type, + // save changes only for update operation + change: cs.type === ChangeSetType.UPDATE ? cs.payload : undefined, + ...(loggedEntities[cs.collection].entities.length > 0 + ? { + additionalData: loggedEntities[cs.collection].entities.reduce( + (acc, name) => ({ + ...acc, + [name]: + cs.entity[name]?.entity instanceof BaseEntity + ? cs.entity[name].id + : cs.entity[name] + }), + {} + ) + } + : {}) + }); + args.uow.computeChangeSet(log); + } + }); + } catch (e) { + getLogger().warn(e, 'Error during database logging.'); + } } }