From 7b2bc93bc1994e2490d1527680f9ab6c406f790b Mon Sep 17 00:00:00 2001 From: Danil Nikolaev Date: Mon, 18 May 2026 21:27:41 +0300 Subject: [PATCH] Add stale RAG provider cleanup --- PIPELINE_TODO.md | 2 +- src/ai/rag-retention-planner.ts | 75 ++++++++++++++++++ src/ai/rag-retention.ts | 117 ++++++++++++++++++++++++++++ src/index.ts | 2 + test/rag-retention-planner.test.mjs | 53 +++++++++++++ 5 files changed, 248 insertions(+), 1 deletion(-) create mode 100644 src/ai/rag-retention-planner.ts create mode 100644 src/ai/rag-retention.ts create mode 100644 test/rag-retention-planner.test.mjs diff --git a/PIPELINE_TODO.md b/PIPELINE_TODO.md index cf50e78..0a9c716 100644 --- a/PIPELINE_TODO.md +++ b/PIPELINE_TODO.md @@ -136,7 +136,7 @@ ## 10. Operational cleanup and observability - [x] Add retention policy for `data/cache/internal-artifacts`. -- [ ] Add retention policy for stale RAG vector/library provider state. +- [x] Add retention policy for stale RAG vector/library provider state. - [x] Add command or admin view for recent `ai_requests`. - [x] Add command or admin view for request audit by message id. - [x] Add command to inspect artifacts for a message. diff --git a/src/ai/rag-retention-planner.ts b/src/ai/rag-retention-planner.ts new file mode 100644 index 0000000..fd0c509 --- /dev/null +++ b/src/ai/rag-retention-planner.ts @@ -0,0 +1,75 @@ +import type {RagArtifactPayload} from "./rag-artifact-payload"; + +export type ArtifactLike = { + id: string; + createdAt: string; + payload: string; +}; + +export type RagCleanupTarget = { + artifactId: string; + createdAt: string; + provider: RagArtifactPayload["providerState"]["provider"]; + vectorStoreIds?: string[]; + uploadedFileIds?: string[]; + libraryId?: string; +}; + +export type RagCleanupPlan = { + cutoffAt: string; + targets: RagCleanupTarget[]; +}; + +function parseRagArtifactPayload(payload: string): RagArtifactPayload | null { + try { + const parsed = JSON.parse(payload) as Partial; + if (!parsed || parsed.artifactKind !== "rag" || !parsed.providerState) return null; + return parsed as RagArtifactPayload; + } catch { + return null; + } +} + +export function buildStaleRagCleanupPlan( + artifacts: ArtifactLike[], + retentionDays = 14, + now = new Date(), +): RagCleanupPlan { + const cutoffAt = new Date(now.getTime() - retentionDays * 24 * 60 * 60 * 1000).toISOString(); + const targets: RagCleanupTarget[] = []; + + for (const artifact of artifacts) { + if (artifact.createdAt > cutoffAt) continue; + + const payload = parseRagArtifactPayload(artifact.payload); + if (!payload || payload.artifactKind !== "rag") continue; + + switch (payload.providerState.provider) { + case "OPENAI": + if (payload.providerState.vectorStoreIds.length || payload.providerState.uploadedFileIds.length) { + targets.push({ + artifactId: artifact.id, + createdAt: artifact.createdAt, + provider: payload.providerState.provider, + vectorStoreIds: [...payload.providerState.vectorStoreIds], + uploadedFileIds: [...payload.providerState.uploadedFileIds], + }); + } + break; + case "MISTRAL": + if (payload.providerState.libraryId) { + targets.push({ + artifactId: artifact.id, + createdAt: artifact.createdAt, + provider: payload.providerState.provider, + libraryId: payload.providerState.libraryId, + }); + } + break; + case "OLLAMA": + break; + } + } + + return {cutoffAt, targets}; +} diff --git a/src/ai/rag-retention.ts b/src/ai/rag-retention.ts new file mode 100644 index 0000000..2d9ea44 --- /dev/null +++ b/src/ai/rag-retention.ts @@ -0,0 +1,117 @@ +import {appLogger} from "../logging/logger.js"; +import {DatabaseManager} from "../db/database-manager.js"; +import {AiProvider} from "../model/ai-provider.js"; +import {createOpenAiClient, resolveAiRuntimeTarget} from "./ai-runtime-target.js"; +import {deleteMistralLibrary} from "./unified-ai-runner.shared.js"; +import {buildStaleRagCleanupPlan} from "./rag-retention-planner.js"; + +const logger = appLogger.child("rag-retention"); + +function unique(values: string[]): string[] { + return [...new Set(values.filter(Boolean))]; +} + +async function cleanupOpenAiRag(vectorStoreIds: string[], uploadedFileIds: string[]): Promise { + const target = resolveAiRuntimeTarget(AiProvider.OPENAI, "documents"); + const client = createOpenAiClient(target); + + for (const vectorStoreId of unique(vectorStoreIds)) { + const startedAt = Date.now(); + logger.info("openai.vector_store.cleanup.start", {vectorStoreId}); + try { + await client.vectorStores.delete(vectorStoreId); + logger.success("openai.vector_store.cleanup.done", {vectorStoreId, duration: `${Date.now() - startedAt}ms`}); + } catch (error) { + logger.warn("openai.vector_store.cleanup.failed", { + vectorStoreId, + duration: `${Date.now() - startedAt}ms`, + error: error instanceof Error ? error : String(error), + }); + } + } + + for (const fileId of unique(uploadedFileIds)) { + const startedAt = Date.now(); + logger.info("openai.file.cleanup.start", {fileId}); + try { + await client.files.delete(fileId); + logger.success("openai.file.cleanup.done", {fileId, duration: `${Date.now() - startedAt}ms`}); + } catch (error) { + logger.warn("openai.file.cleanup.failed", { + fileId, + duration: `${Date.now() - startedAt}ms`, + error: error instanceof Error ? error : String(error), + }); + } + } +} + +async function cleanupMistralRag(libraryId: string): Promise { + const target = resolveAiRuntimeTarget(AiProvider.MISTRAL, "documents"); + const startedAt = Date.now(); + logger.info("mistral.library.cleanup.start", {libraryId}); + try { + await deleteMistralLibrary(libraryId, target); + logger.success("mistral.library.cleanup.done", {libraryId, duration: `${Date.now() - startedAt}ms`}); + } catch (error) { + logger.warn("mistral.library.cleanup.failed", { + libraryId, + duration: `${Date.now() - startedAt}ms`, + error: error instanceof Error ? error : String(error), + }); + } +} + +export async function cleanupStaleRagProviderState(retentionDays = 14): Promise<{ + scannedArtifacts: number; + cleanupTargets: number; + openaiTargets: number; + mistralTargets: number; +}> { + const startedAt = Date.now(); + const artifacts = await DatabaseManager.getAllArtifacts().catch(() => []); + const plan = buildStaleRagCleanupPlan(artifacts, retentionDays); + + logger.info("cleanup.start", { + retentionDays, + scannedArtifacts: artifacts.length, + cleanupTargets: plan.targets.length, + cutoffAt: plan.cutoffAt, + }); + + let openaiTargets = 0; + let mistralTargets = 0; + + for (const target of plan.targets) { + switch (target.provider) { + case "OPENAI": + openaiTargets += 1; + await cleanupOpenAiRag(target.vectorStoreIds ?? [], target.uploadedFileIds ?? []); + break; + case "MISTRAL": + mistralTargets += 1; + if (target.libraryId) { + await cleanupMistralRag(target.libraryId); + } + break; + case "OLLAMA": + break; + } + } + + logger.success("cleanup.done", { + retentionDays, + scannedArtifacts: artifacts.length, + cleanupTargets: plan.targets.length, + openaiTargets, + mistralTargets, + duration: `${Date.now() - startedAt}ms`, + }); + + return { + scannedArtifacts: artifacts.length, + cleanupTargets: plan.targets.length, + openaiTargets, + mistralTargets, + }; +} diff --git a/src/index.ts b/src/index.ts index 8795588..bce4378 100644 --- a/src/index.ts +++ b/src/index.ts @@ -78,6 +78,7 @@ import {cleanupInternalArtifactCache} from "./ai/internal-artifact-store.js"; import {AIAudit} from "./commands/ai-audit.js"; import {AIMetrics} from "./commands/ai-metrics.js"; import {AIRequests} from "./commands/ai-requests.js"; +import {cleanupStaleRagProviderState} from "./ai/rag-retention.js"; process.setUncaughtExceptionCaptureCallback(logError); @@ -278,6 +279,7 @@ async function main() { }, () => ({notesRootFilePath})); await measureStartupStep("cleanup_internal_artifacts", () => cleanupInternalArtifactCache(), () => ({retentionDays: 14})); + await measureStartupStep("cleanup_stale_rag_provider_state", () => cleanupStaleRagProviderState(), () => ({retentionDays: 14})); await measureStartupStep("observability.snapshot", async () => { const [aiRequests, attachments, artifacts, requestAudits] = await Promise.all([ DatabaseManager.getAllAiRequests(), diff --git a/test/rag-retention-planner.test.mjs b/test/rag-retention-planner.test.mjs new file mode 100644 index 0000000..ae5e0ee --- /dev/null +++ b/test/rag-retention-planner.test.mjs @@ -0,0 +1,53 @@ +import test from "node:test"; +import assert from "node:assert/strict"; + +const {buildStaleRagCleanupPlan} = await import("../dist/ai/rag-retention-planner.js"); + +test("stale rag cleanup plan selects only older rag artifacts", () => { + const plan = buildStaleRagCleanupPlan([ + { + id: "recent-openai", + createdAt: "2026-05-18T00:00:00.000Z", + payload: JSON.stringify({ + artifactKind: "rag", + providerState: { + provider: "OPENAI", + vectorStoreIds: ["vs_1"], + uploadedFileIds: ["file_1"], + }, + }), + }, + { + id: "stale-openai", + createdAt: "2026-04-01T00:00:00.000Z", + payload: JSON.stringify({ + artifactKind: "rag", + providerState: { + provider: "OPENAI", + vectorStoreIds: ["vs_2"], + uploadedFileIds: ["file_2"], + }, + }), + }, + { + id: "stale-ollama", + createdAt: "2026-04-01T00:00:00.000Z", + payload: JSON.stringify({ + artifactKind: "rag", + providerState: { + provider: "OLLAMA", + prepared: true, + }, + }), + }, + ], 14, new Date("2026-05-18T00:00:00.000Z")); + + assert.equal(plan.targets.length, 1); + assert.deepEqual(plan.targets[0], { + artifactId: "stale-openai", + createdAt: "2026-04-01T00:00:00.000Z", + provider: "OPENAI", + vectorStoreIds: ["vs_2"], + uploadedFileIds: ["file_2"], + }); +});