WeHub Mirror
6bf8bebf51
CI / Test and Build (push) Failing after 1s
CI / Migrate Dev DB (push) Has been skipped
CI / Migrate DB (push) Has been skipped
CodeQL / Analyze actions (push) Has been cancelled
CodeQL / Analyze javascript-typescript (push) Has been cancelled
CI / Detect Version (push) Has been cancelled
CI / Detect Desktop Changes (push) Has been cancelled
CI / Build AMD64 (blacksmith-2vcpu-ubuntu-2404, ./docker/cron.Dockerfile, ubuntu-latest, ghcr.io/simstudioai/cron) (push) Has been cancelled
CI / Build AMD64 (blacksmith-2vcpu-ubuntu-2404, ./docker/db.Dockerfile, ECR_MIGRATIONS, ubuntu-latest, ghcr.io/simstudioai/migrations) (push) Has been cancelled
CI / Build AMD64 (blacksmith-4vcpu-ubuntu-2404, ./docker/pii.Dockerfile, ECR_PII, ubuntu-latest, ghcr.io/simstudioai/pii) (push) Has been cancelled
CI / Build AMD64 (blacksmith-4vcpu-ubuntu-2404, ./docker/realtime.Dockerfile, ECR_REALTIME, ubuntu-latest, ghcr.io/simstudioai/realtime) (push) Has been cancelled
CI / Build AMD64 (blacksmith-8vcpu-ubuntu-2404, ./docker/app.Dockerfile, ECR_APP, linux-x64-8-core, ghcr.io/simstudioai/simstudio) (push) Has been cancelled
CI / Build ARM64 (GHCR Only) (blacksmith-4vcpu-ubuntu-2404-arm, ./docker/cron.Dockerfile, ubuntu-24.04-arm, ghcr.io/simstudioai/cron) (push) Has been cancelled
CI / Build ARM64 (GHCR Only) (blacksmith-4vcpu-ubuntu-2404-arm, ./docker/db.Dockerfile, ubuntu-24.04-arm, ghcr.io/simstudioai/migrations) (push) Has been cancelled
CI / Build ARM64 (GHCR Only) (blacksmith-4vcpu-ubuntu-2404-arm, ./docker/pii.Dockerfile, ubuntu-24.04-arm, ghcr.io/simstudioai/pii) (push) Has been cancelled
CI / Build ARM64 (GHCR Only) (blacksmith-4vcpu-ubuntu-2404-arm, ./docker/realtime.Dockerfile, ubuntu-24.04-arm, ghcr.io/simstudioai/realtime) (push) Has been cancelled
CI / Build ARM64 (GHCR Only) (blacksmith-8vcpu-ubuntu-2404-arm, ./docker/app.Dockerfile, linux-arm64-8-core, ghcr.io/simstudioai/simstudio) (push) Has been cancelled
CI / Check Docs Changes (push) Has been cancelled
Publish CLI Package / publish-npm (push) Has been cancelled
Publish Python SDK / publish-pypi (push) Has been cancelled
CI / Deploy Trigger.dev (Dev) (push) Has been cancelled
Helm Chart / Lint, test, and validate chart (push) Has been cancelled
Helm Chart / Chart version bumped (push) Has been cancelled
Publish TypeScript SDK / publish-npm (push) Has been cancelled
CI / Build Dev ECR (blacksmith-8vcpu-ubuntu-2404, ./docker/app.Dockerfile, ECR_APP, linux-x64-8-core) (push) Has been cancelled
CI / Promote Images (push) Has been cancelled
CI / Create GHCR Manifests (ghcr.io/simstudioai/cron) (push) Has been cancelled
CI / Create GHCR Manifests (ghcr.io/simstudioai/migrations) (push) Has been cancelled
CI / Create GHCR Manifests (ghcr.io/simstudioai/pii) (push) Has been cancelled
CI / Create GHCR Manifests (ghcr.io/simstudioai/realtime) (push) Has been cancelled
CI / Build Dev ECR (blacksmith-2vcpu-ubuntu-2404, ./docker/db.Dockerfile, ECR_MIGRATIONS, ubuntu-latest) (push) Has been cancelled
CI / Build Dev ECR (blacksmith-4vcpu-ubuntu-2404, ./docker/pii.Dockerfile, ECR_PII, ubuntu-latest) (push) Has been cancelled
CI / Build Dev ECR (blacksmith-4vcpu-ubuntu-2404, ./docker/realtime.Dockerfile, ECR_REALTIME, ubuntu-latest) (push) Has been cancelled
CI / Create GHCR Manifests (ghcr.io/simstudioai/simstudio) (push) Has been cancelled
CI / Process Docs (push) Has been cancelled
CI / Create GitHub Release (push) Has been cancelled
CI / Check Desktop Signing Secrets (push) Has been cancelled
CI / Desktop Release (push) Has been cancelled
CI / Create Desktop Prerelease (push) Has been cancelled
CI / Desktop Prerelease Build (push) Has been cancelled
CI / Publish Desktop Prerelease (push) Has been cancelled
CI / Prune Desktop Prereleases (push) Has been cancelled
Helm Chart / Install on kind and run helm test (push) Has been cancelled
2891 lines
89 KiB
TypeScript
2891 lines
89 KiB
TypeScript
import { db } from '@sim/db'
|
|
import {
|
|
document,
|
|
documentSecretProvenance,
|
|
embedding,
|
|
embeddingSecretProvenance,
|
|
knowledgeBase,
|
|
knowledgeBaseTagDefinitions,
|
|
knowledgeConnector,
|
|
workspace as workspaceTable,
|
|
} from '@sim/db/schema'
|
|
import { createLogger } from '@sim/logger'
|
|
import { sha256Hex } from '@sim/security/hash'
|
|
import { getErrorMessage, toError } from '@sim/utils/errors'
|
|
import { generateId } from '@sim/utils/id'
|
|
import { tasks } from '@trigger.dev/sdk'
|
|
import { and, asc, desc, eq, inArray, isNotNull, isNull, type SQL, sql } from 'drizzle-orm'
|
|
import { checkActorUsageLimits } from '@/lib/billing/calculations/usage-monitor'
|
|
import {
|
|
assertBillingAttributionSnapshot,
|
|
type BillingAttributionSnapshot,
|
|
checkAttributedUsageLimits,
|
|
toBillingContext,
|
|
} from '@/lib/billing/core/billing-attribution'
|
|
import type { HighestPrioritySubscription } from '@/lib/billing/core/plan'
|
|
import { getHighestPrioritySubscription } from '@/lib/billing/core/subscription'
|
|
import { recordUsage } from '@/lib/billing/core/usage-log'
|
|
import {
|
|
applyStorageUsageDeltasInTx,
|
|
checkAndIncrementStorageUsageInTx,
|
|
checkStorageQuota,
|
|
checkStorageQuotaForBillingContext,
|
|
incrementStorageUsageForBillingContextInTx,
|
|
maybeNotifyStorageLimitForBillingContext,
|
|
resolveStorageBillingContext,
|
|
type StorageBillingContext,
|
|
} from '@/lib/billing/storage'
|
|
import {
|
|
checkAndBillOverageThreshold,
|
|
checkAndBillPayerOverageThreshold,
|
|
} from '@/lib/billing/threshold-billing'
|
|
import type { ChunkingStrategy, StrategyOptions } from '@/lib/chunkers/types'
|
|
import { resolveTriggerRegion } from '@/lib/core/async-jobs/region'
|
|
import { env, envNumber } from '@/lib/core/config/env'
|
|
import { getCostMultiplier, isTriggerDevEnabled } from '@/lib/core/config/env-flags'
|
|
import { mapWithConcurrency } from '@/lib/core/utils/concurrency'
|
|
import {
|
|
type DurableSecretProvenance,
|
|
durableSecretProvenanceFromRegistry,
|
|
EXACT_EMPTY_DURABLE_SECRET_PROVENANCE,
|
|
mergeDurableSecretProvenance,
|
|
} from '@/lib/execution/durable-secret-provenance'
|
|
import { processDocument } from '@/lib/knowledge/documents/document-processor'
|
|
import {
|
|
assertDocumentProcessingBillingContext,
|
|
createDocumentProcessingPayload,
|
|
createNonWorkspaceDocumentProcessingBillingContext,
|
|
createWorkspaceDocumentProcessingBillingContext,
|
|
type DocumentProcessingBillingContext,
|
|
type DocumentProcessingPayload,
|
|
hasDocumentProcessingBillingScope,
|
|
} from '@/lib/knowledge/documents/processing-payload'
|
|
import {
|
|
buildTagFilterCondition,
|
|
type TagFilterCondition,
|
|
} from '@/lib/knowledge/documents/tag-filter'
|
|
import type { DocumentSortField, SortOrder } from '@/lib/knowledge/documents/types'
|
|
import { getEmbeddingModelInfo } from '@/lib/knowledge/embedding-models'
|
|
import { generateEmbeddings } from '@/lib/knowledge/embeddings'
|
|
import { runWithKnowledgeModelInputProvenance } from '@/lib/knowledge/model-input-provenance'
|
|
import {
|
|
bindKnowledgeDocumentFieldSecretProvenance,
|
|
createKnowledgeDocumentSourceValue,
|
|
type KnowledgeDocumentMetadataField,
|
|
type KnowledgeDocumentWriteSecretProvenance,
|
|
loadKnowledgeDocumentSecretRegistry,
|
|
readBoundKnowledgeDocumentSecretProvenance,
|
|
rebindKnowledgeDocumentSecretProvenance,
|
|
replaceKnowledgeDocumentSecretProvenanceInTx,
|
|
} from '@/lib/knowledge/secret-provenance'
|
|
import {
|
|
buildUndefinedTagsError,
|
|
parseBooleanValue,
|
|
parseDateValue,
|
|
parseNumberValue,
|
|
validateTagValue,
|
|
} from '@/lib/knowledge/tags/utils'
|
|
import type { ProcessedDocumentTags } from '@/lib/knowledge/types'
|
|
import { estimateTokenCount } from '@/lib/tokenization/estimators'
|
|
import {
|
|
getBoundWorkspaceFileSecretProvenanceByMetadata,
|
|
type WorkspaceFileSecretProvenance,
|
|
} from '@/lib/uploads/contexts/workspace/workspace-file-secret-provenance'
|
|
import { deleteFile } from '@/lib/uploads/core/storage-service'
|
|
import {
|
|
deleteFileMetadataByIdentity,
|
|
type FileMetadataRecord,
|
|
getFileMetadataByKeys,
|
|
} from '@/lib/uploads/server/metadata'
|
|
import { extractStorageKey } from '@/lib/uploads/utils/file-utils'
|
|
import type { processDocument as processDocumentTask } from '@/background/knowledge-processing'
|
|
import { calculateCost } from '@/providers/utils'
|
|
|
|
const logger = createLogger('DocumentService')
|
|
|
|
/**
|
|
* Thrown when a knowledge-base document's `fileUrl` references an internal
|
|
* knowledge-base storage object not owned by the target knowledge base's workspace.
|
|
* Routes map this to a 403.
|
|
*/
|
|
export class KnowledgeBaseFileOwnershipError extends Error {
|
|
constructor(public readonly storageKey: string) {
|
|
super('Document file is not owned by this knowledge base')
|
|
this.name = 'KnowledgeBaseFileOwnershipError'
|
|
}
|
|
}
|
|
|
|
/**
|
|
* Guard document `fileUrl`s at creation time. When a URL points at an internal
|
|
* knowledge-base storage object, require that the target knowledge base owns the object,
|
|
* resolved from the trusted `workspace_files` binding:
|
|
*
|
|
* - Workspace KB (`kbWorkspaceId` set): the binding's `workspaceId` must match.
|
|
* - Personal KB (`kbWorkspaceId` null): the binding's `userId` must be the KB
|
|
* owner. A key bound to another tenant is rejected; an unbound key (legacy /
|
|
* never reserved) passes since it carries no cross-tenant ownership.
|
|
*
|
|
* External `http(s)`/`data:` URLs (ingestion sources) and other internal keys
|
|
* pass through unchanged. This blocks a user from asserting ownership of another
|
|
* tenant's object via a planted `fileUrl` — including in a personal KB, which
|
|
* otherwise could be moved into a workspace to launder the binding. All
|
|
* referenced bindings are resolved in one query (no N+1 inside the `FOR UPDATE`
|
|
* window). Single-document callers pass a one-element array.
|
|
*/
|
|
function isKnowledgeBaseOwnedStorageKey(key: string): boolean {
|
|
return key.startsWith('kb/') || key.startsWith('knowledge-base/')
|
|
}
|
|
|
|
function getKnowledgeBaseStorageKeys(fileUrls: readonly string[]): string[] {
|
|
return [
|
|
...new Set(
|
|
fileUrls
|
|
.map((url) => getKnowledgeBaseStorageKey(url))
|
|
.filter(
|
|
(key): key is string => typeof key === 'string' && isKnowledgeBaseOwnedStorageKey(key)
|
|
)
|
|
),
|
|
]
|
|
}
|
|
|
|
function getWorkspaceSourceStorageKeys(fileUrls: readonly string[]): string[] {
|
|
return [
|
|
...new Set(
|
|
fileUrls
|
|
.map((url) => getKnowledgeBaseStorageKey(url))
|
|
.filter((key): key is string => typeof key === 'string' && key.startsWith('workspace/'))
|
|
),
|
|
]
|
|
}
|
|
|
|
async function loadKnowledgeBaseFileBindings(
|
|
fileUrls: readonly string[],
|
|
executor: DbExecutor = db
|
|
): Promise<Map<string, FileMetadataRecord>> {
|
|
const keys = getKnowledgeBaseStorageKeys(fileUrls)
|
|
const bindings =
|
|
keys.length > 0 ? await getFileMetadataByKeys(keys, 'knowledge-base', executor) : []
|
|
|
|
return new Map(bindings.map((binding) => [binding.key, binding]))
|
|
}
|
|
|
|
async function loadWorkspaceSourceFileBindings(
|
|
fileUrls: readonly string[],
|
|
executor: DbExecutor = db
|
|
): Promise<Map<string, FileMetadataRecord>> {
|
|
const keys = getWorkspaceSourceStorageKeys(fileUrls)
|
|
if (keys.length === 0) return new Map()
|
|
|
|
const workspaceBindings = await getFileMetadataByKeys(keys, 'workspace', executor)
|
|
const mothershipBindings = await getFileMetadataByKeys(keys, 'mothership', executor)
|
|
|
|
return new Map(
|
|
[...workspaceBindings, ...mothershipBindings].map((binding) => [binding.key, binding])
|
|
)
|
|
}
|
|
|
|
async function assertKnowledgeBaseFileUrlsOwnership(
|
|
fileUrls: string[],
|
|
kbWorkspaceId: string | null,
|
|
kbUserId: string,
|
|
requestId: string,
|
|
executor: DbExecutor = db
|
|
): Promise<Map<string, FileMetadataRecord>> {
|
|
const keys = getKnowledgeBaseStorageKeys(fileUrls)
|
|
if (keys.length === 0) {
|
|
return new Map()
|
|
}
|
|
|
|
const bindingByKey = await loadKnowledgeBaseFileBindings(fileUrls, executor)
|
|
|
|
for (const key of keys) {
|
|
const binding = bindingByKey.get(key)
|
|
|
|
if (kbWorkspaceId) {
|
|
if (!binding || binding.workspaceId !== kbWorkspaceId) {
|
|
logger.warn(`[${requestId}] Rejected document referencing unowned knowledge-base file`, {
|
|
storageKey: key,
|
|
kbWorkspaceId,
|
|
bindingWorkspaceId: binding?.workspaceId ?? null,
|
|
})
|
|
throw new KnowledgeBaseFileOwnershipError(key)
|
|
}
|
|
continue
|
|
}
|
|
|
|
// Personal KB: reject a key whose binding belongs to a different user. An
|
|
// unbound key carries no ownership and is allowed (legacy personal files).
|
|
if (binding && binding.userId !== kbUserId) {
|
|
logger.warn(
|
|
`[${requestId}] Rejected personal-KB document referencing another tenant's file`,
|
|
{
|
|
storageKey: key,
|
|
kbUserId,
|
|
bindingUserId: binding.userId,
|
|
bindingWorkspaceId: binding.workspaceId ?? null,
|
|
}
|
|
)
|
|
throw new KnowledgeBaseFileOwnershipError(key)
|
|
}
|
|
}
|
|
|
|
return bindingByKey
|
|
}
|
|
|
|
async function loadCurrentWorkspaceSourceFileSecretProvenance(options: {
|
|
fileUrl: string
|
|
}): Promise<DurableSecretProvenance | undefined> {
|
|
const storageKey = getKnowledgeBaseStorageKey(options.fileUrl)
|
|
if (!storageKey?.startsWith('workspace/')) return undefined
|
|
|
|
const bindingByKey = await loadWorkspaceSourceFileBindings([options.fileUrl])
|
|
const binding = bindingByKey.get(storageKey)
|
|
if (!binding) return undefined
|
|
|
|
const provenanceById = await getBoundWorkspaceFileSecretProvenanceByMetadata(db, [binding])
|
|
const provenance = provenanceById.get(binding.id) ?? { status: 'unknown' as const }
|
|
return durableSecretProvenanceFromWorkspaceFile(provenance, binding)
|
|
}
|
|
|
|
const TIMEOUTS = {
|
|
OVERALL_PROCESSING: envNumber(env.KB_CONFIG_MAX_DURATION, 600) * 1000,
|
|
} as const
|
|
|
|
const LARGE_DOC_CONFIG = {
|
|
MAX_CHUNKS_PER_BATCH: 500,
|
|
MAX_EMBEDDING_BATCH: envNumber(env.KB_CONFIG_BATCH_SIZE, 2000),
|
|
MAX_FILE_SIZE: 100 * 1024 * 1024,
|
|
MAX_CHUNKS_PER_DOCUMENT: 100000,
|
|
}
|
|
|
|
const HARD_DELETE_DOCUMENT_BATCH_SIZE = 250
|
|
|
|
function withTimeout<T>(
|
|
promise: Promise<T>,
|
|
timeoutMs: number,
|
|
operation = 'Operation'
|
|
): Promise<T> {
|
|
return Promise.race([
|
|
promise,
|
|
new Promise<never>((_, reject) =>
|
|
setTimeout(() => reject(new Error(`${operation} timed out after ${timeoutMs}ms`)), timeoutMs)
|
|
),
|
|
])
|
|
}
|
|
|
|
const PROCESSING_CONFIG = {
|
|
maxConcurrentDocuments:
|
|
Math.max(1, Math.floor(envNumber(env.KB_CONFIG_CONCURRENCY_LIMIT, 20) / 5)) || 4,
|
|
batchSize: Math.max(1, Math.floor(envNumber(env.KB_CONFIG_BATCH_SIZE, 20) / 2)) || 10,
|
|
delayBetweenBatches: envNumber(env.KB_CONFIG_DELAY_BETWEEN_BATCHES, 100) * 2,
|
|
delayBetweenDocuments: envNumber(env.KB_CONFIG_DELAY_BETWEEN_DOCUMENTS, 50) * 2,
|
|
}
|
|
|
|
export function getProcessingConfig() {
|
|
return PROCESSING_CONFIG
|
|
}
|
|
|
|
export interface DocumentData {
|
|
documentId: string
|
|
filename: string
|
|
fileUrl: string
|
|
fileSize: number
|
|
mimeType: string
|
|
}
|
|
|
|
export interface ProcessingOptions {
|
|
recipe?: string
|
|
lang?: string
|
|
}
|
|
|
|
interface DocumentTagData {
|
|
tagName: string
|
|
fieldType: string
|
|
value: string
|
|
}
|
|
|
|
type TagDefinition = typeof knowledgeBaseTagDefinitions.$inferSelect
|
|
type TagDefinitionsByName = Map<string, TagDefinition>
|
|
type DbExecutor = Pick<typeof db, 'select'>
|
|
|
|
async function loadTagDefinitions(
|
|
knowledgeBaseId: string,
|
|
executor: DbExecutor = db
|
|
): Promise<TagDefinitionsByName> {
|
|
const defs = await executor
|
|
.select()
|
|
.from(knowledgeBaseTagDefinitions)
|
|
.where(eq(knowledgeBaseTagDefinitions.knowledgeBaseId, knowledgeBaseId))
|
|
return new Map(defs.map((def) => [def.displayName, def]))
|
|
}
|
|
|
|
function resolveDocumentTags(
|
|
tagData: DocumentTagData[],
|
|
tagDefinitions: TagDefinitionsByName,
|
|
requestId: string
|
|
): ProcessedDocumentTags {
|
|
const setTagValue = (
|
|
tags: ProcessedDocumentTags,
|
|
slot: string,
|
|
value: string | number | Date | boolean | null
|
|
): void => {
|
|
switch (slot) {
|
|
case 'tag1':
|
|
tags.tag1 = value as string | null
|
|
break
|
|
case 'tag2':
|
|
tags.tag2 = value as string | null
|
|
break
|
|
case 'tag3':
|
|
tags.tag3 = value as string | null
|
|
break
|
|
case 'tag4':
|
|
tags.tag4 = value as string | null
|
|
break
|
|
case 'tag5':
|
|
tags.tag5 = value as string | null
|
|
break
|
|
case 'tag6':
|
|
tags.tag6 = value as string | null
|
|
break
|
|
case 'tag7':
|
|
tags.tag7 = value as string | null
|
|
break
|
|
case 'number1':
|
|
tags.number1 = value as number | null
|
|
break
|
|
case 'number2':
|
|
tags.number2 = value as number | null
|
|
break
|
|
case 'number3':
|
|
tags.number3 = value as number | null
|
|
break
|
|
case 'number4':
|
|
tags.number4 = value as number | null
|
|
break
|
|
case 'number5':
|
|
tags.number5 = value as number | null
|
|
break
|
|
case 'date1':
|
|
tags.date1 = value as Date | null
|
|
break
|
|
case 'date2':
|
|
tags.date2 = value as Date | null
|
|
break
|
|
case 'boolean1':
|
|
tags.boolean1 = value as boolean | null
|
|
break
|
|
case 'boolean2':
|
|
tags.boolean2 = value as boolean | null
|
|
break
|
|
case 'boolean3':
|
|
tags.boolean3 = value as boolean | null
|
|
break
|
|
}
|
|
}
|
|
|
|
const result: ProcessedDocumentTags = {
|
|
tag1: null,
|
|
tag2: null,
|
|
tag3: null,
|
|
tag4: null,
|
|
tag5: null,
|
|
tag6: null,
|
|
tag7: null,
|
|
number1: null,
|
|
number2: null,
|
|
number3: null,
|
|
number4: null,
|
|
number5: null,
|
|
date1: null,
|
|
date2: null,
|
|
boolean1: null,
|
|
boolean2: null,
|
|
boolean3: null,
|
|
}
|
|
|
|
if (!Array.isArray(tagData) || tagData.length === 0) {
|
|
return result
|
|
}
|
|
|
|
const undefinedTags: string[] = []
|
|
const typeErrors: string[] = []
|
|
|
|
for (const tag of tagData) {
|
|
if (!tag.tagName?.trim()) continue
|
|
|
|
const tagName = tag.tagName.trim()
|
|
const fieldType = tag.fieldType || 'text'
|
|
|
|
const hasValue =
|
|
fieldType === 'boolean'
|
|
? tag.value !== undefined && tag.value !== null && tag.value !== ''
|
|
: tag.value?.trim && tag.value.trim().length > 0
|
|
|
|
if (!hasValue) continue
|
|
|
|
const existingDef = tagDefinitions.get(tagName)
|
|
if (!existingDef) {
|
|
undefinedTags.push(tagName)
|
|
continue
|
|
}
|
|
|
|
const rawValue = typeof tag.value === 'string' ? tag.value.trim() : tag.value
|
|
const actualFieldType = existingDef.fieldType || fieldType
|
|
const validationError = validateTagValue(tagName, String(rawValue), actualFieldType)
|
|
if (validationError) {
|
|
typeErrors.push(validationError)
|
|
}
|
|
}
|
|
|
|
if (undefinedTags.length > 0 || typeErrors.length > 0) {
|
|
const errorParts: string[] = []
|
|
|
|
if (undefinedTags.length > 0) {
|
|
errorParts.push(buildUndefinedTagsError(undefinedTags))
|
|
}
|
|
|
|
if (typeErrors.length > 0) {
|
|
errorParts.push(...typeErrors)
|
|
}
|
|
|
|
throw new Error(errorParts.join('\n'))
|
|
}
|
|
|
|
for (const tag of tagData) {
|
|
if (!tag.tagName?.trim()) continue
|
|
|
|
const tagName = tag.tagName.trim()
|
|
const fieldType = tag.fieldType || 'text'
|
|
|
|
const hasValue =
|
|
fieldType === 'boolean'
|
|
? tag.value !== undefined && tag.value !== null && tag.value !== ''
|
|
: tag.value?.trim && tag.value.trim().length > 0
|
|
|
|
if (!hasValue) continue
|
|
|
|
const existingDef = tagDefinitions.get(tagName)
|
|
if (!existingDef) continue
|
|
|
|
const targetSlot = existingDef.tagSlot
|
|
const actualFieldType = existingDef.fieldType || fieldType
|
|
const rawValue = typeof tag.value === 'string' ? tag.value.trim() : tag.value
|
|
const stringValue = String(rawValue).trim()
|
|
|
|
if (actualFieldType === 'boolean') {
|
|
setTagValue(result, targetSlot, parseBooleanValue(stringValue) ?? false)
|
|
} else if (actualFieldType === 'number') {
|
|
setTagValue(result, targetSlot, parseNumberValue(stringValue))
|
|
} else if (actualFieldType === 'date') {
|
|
setTagValue(result, targetSlot, parseDateValue(stringValue))
|
|
} else {
|
|
setTagValue(result, targetSlot, stringValue)
|
|
}
|
|
|
|
logger.info(`[${requestId}] Set tag ${tagName} (${targetSlot})`, {
|
|
fieldType: actualFieldType,
|
|
})
|
|
}
|
|
|
|
return result
|
|
}
|
|
|
|
const KNOWLEDGE_DOCUMENT_TAG_FIELDS = new Set<KnowledgeDocumentMetadataField>([
|
|
'tag1',
|
|
'tag2',
|
|
'tag3',
|
|
'tag4',
|
|
'tag5',
|
|
'tag6',
|
|
'tag7',
|
|
'number1',
|
|
'number2',
|
|
'number3',
|
|
'number4',
|
|
'number5',
|
|
'date1',
|
|
'date2',
|
|
'boolean1',
|
|
'boolean2',
|
|
'boolean3',
|
|
])
|
|
|
|
function durableSecretProvenanceFromWorkspaceFile(
|
|
provenance: WorkspaceFileSecretProvenance,
|
|
binding: FileMetadataRecord
|
|
): DurableSecretProvenance {
|
|
if (provenance.status === 'unknown') return provenance
|
|
return {
|
|
status: 'exact',
|
|
entries: provenance.entries.map((entry) => ({
|
|
...entry,
|
|
sourceUserId: binding.userId,
|
|
...(binding.workspaceId ? { sourceWorkspaceId: binding.workspaceId } : {}),
|
|
})),
|
|
}
|
|
}
|
|
|
|
function bindKnowledgeDocumentWriteSecretProvenance(options: {
|
|
source: ReturnType<typeof createKnowledgeDocumentSourceValue>
|
|
provenance?: KnowledgeDocumentWriteSecretProvenance
|
|
tagDefinitions: TagDefinitionsByName
|
|
boundFile?: {
|
|
binding: FileMetadataRecord
|
|
provenance: WorkspaceFileSecretProvenance
|
|
}
|
|
}): DurableSecretProvenance | undefined {
|
|
const values: DurableSecretProvenance[] = []
|
|
if (options.provenance) {
|
|
values.push(
|
|
bindKnowledgeDocumentFieldSecretProvenance(
|
|
options.provenance.filename,
|
|
'filename',
|
|
options.source.filename
|
|
),
|
|
bindKnowledgeDocumentFieldSecretProvenance(options.provenance.content, 'content', {
|
|
fileUrl: options.source.fileUrl,
|
|
contentHash: options.source.contentHash,
|
|
})
|
|
)
|
|
for (const tag of options.provenance.tags) {
|
|
const tagField = options.tagDefinitions.get(tag.tagName)?.tagSlot
|
|
if (
|
|
!tagField ||
|
|
!KNOWLEDGE_DOCUMENT_TAG_FIELDS.has(tagField as KnowledgeDocumentMetadataField)
|
|
) {
|
|
return { status: 'unknown' }
|
|
}
|
|
const field = tagField as KnowledgeDocumentMetadataField
|
|
values.push(
|
|
bindKnowledgeDocumentFieldSecretProvenance(tag.provenance, field, options.source[field])
|
|
)
|
|
}
|
|
}
|
|
if (options.boundFile) {
|
|
values.push(
|
|
bindKnowledgeDocumentFieldSecretProvenance(
|
|
durableSecretProvenanceFromWorkspaceFile(
|
|
options.boundFile.provenance,
|
|
options.boundFile.binding
|
|
),
|
|
'content',
|
|
{ fileUrl: options.source.fileUrl, contentHash: options.source.contentHash }
|
|
)
|
|
)
|
|
}
|
|
return values.length > 0 ? mergeDurableSecretProvenance(...values) : undefined
|
|
}
|
|
|
|
/** Per-call cap for `tasks.batchTrigger` on Trigger.dev SDK 4.3.1+. */
|
|
const TRIGGER_BATCH_SIZE = 1000
|
|
|
|
function buildJobPayload(
|
|
doc: DocumentData,
|
|
knowledgeBaseId: string,
|
|
processingOptions: ProcessingOptions,
|
|
requestId: string,
|
|
billingContext: DocumentProcessingBillingContext
|
|
): DocumentProcessingPayload {
|
|
return createDocumentProcessingPayload(
|
|
{
|
|
knowledgeBaseId,
|
|
documentId: doc.documentId,
|
|
docData: {
|
|
filename: doc.filename,
|
|
fileUrl: doc.fileUrl,
|
|
fileSize: doc.fileSize,
|
|
mimeType: doc.mimeType,
|
|
},
|
|
processingOptions,
|
|
requestId,
|
|
},
|
|
billingContext
|
|
)
|
|
}
|
|
|
|
async function resolveDocumentProcessingBillingContext(
|
|
knowledgeBaseId: string,
|
|
providedBillingAttribution: BillingAttributionSnapshot | undefined
|
|
): Promise<DocumentProcessingBillingContext> {
|
|
const [knowledgeBaseContext] = await db
|
|
.select({
|
|
userId: knowledgeBase.userId,
|
|
workspaceId: knowledgeBase.workspaceId,
|
|
})
|
|
.from(knowledgeBase)
|
|
.where(and(eq(knowledgeBase.id, knowledgeBaseId), isNull(knowledgeBase.deletedAt)))
|
|
.limit(1)
|
|
|
|
if (!knowledgeBaseContext) {
|
|
throw new Error(`Knowledge base ${knowledgeBaseId} not found for document processing`)
|
|
}
|
|
|
|
if (knowledgeBaseContext.workspaceId) {
|
|
if (!providedBillingAttribution) {
|
|
throw new Error('Workspace document processing requires a billing attribution snapshot')
|
|
}
|
|
const billingContext = createWorkspaceDocumentProcessingBillingContext(
|
|
providedBillingAttribution
|
|
)
|
|
if (billingContext.workspaceId !== knowledgeBaseContext.workspaceId) {
|
|
throw new Error('Document processing workspace does not match billing attribution')
|
|
}
|
|
return billingContext
|
|
}
|
|
|
|
if (providedBillingAttribution !== undefined) {
|
|
throw new Error('Non-workspace document processing cannot include billing attribution')
|
|
}
|
|
|
|
return createNonWorkspaceDocumentProcessingBillingContext(knowledgeBaseContext.userId)
|
|
}
|
|
|
|
/**
|
|
* Dispatches document processing jobs via Trigger.dev's `batchTrigger` when
|
|
* available, or in-process otherwise. Throws only when every dispatch fails;
|
|
* partial failures are logged and recovered by the next sync's stuck-doc pass.
|
|
*/
|
|
export async function processDocumentsWithQueue(
|
|
createdDocuments: DocumentData[],
|
|
knowledgeBaseId: string,
|
|
processingOptions: ProcessingOptions,
|
|
requestId: string,
|
|
billingAttribution: BillingAttributionSnapshot | undefined
|
|
): Promise<void> {
|
|
if (createdDocuments.length === 0) return
|
|
|
|
const billingContext = await resolveDocumentProcessingBillingContext(
|
|
knowledgeBaseId,
|
|
billingAttribution
|
|
)
|
|
const jobPayloads = createdDocuments.map((doc) =>
|
|
buildJobPayload(doc, knowledgeBaseId, processingOptions, requestId, billingContext)
|
|
)
|
|
|
|
const useTrigger = isTriggerAvailable()
|
|
logger.info(
|
|
`[${requestId}] Dispatching background processing for ${jobPayloads.length} documents`,
|
|
{ backend: useTrigger ? 'trigger-dev' : 'direct' }
|
|
)
|
|
|
|
const dispatched = useTrigger
|
|
? await dispatchViaBatchTrigger(jobPayloads, requestId)
|
|
: await dispatchInProcess(jobPayloads, requestId)
|
|
|
|
logger.info(
|
|
`[${requestId}] Document dispatch complete: ${dispatched}/${jobPayloads.length} succeeded`
|
|
)
|
|
|
|
if (dispatched === 0) {
|
|
throw new Error(`All ${jobPayloads.length} document processing dispatches failed`)
|
|
}
|
|
}
|
|
|
|
async function dispatchViaBatchTrigger(
|
|
jobPayloads: DocumentProcessingPayload[],
|
|
requestId: string
|
|
): Promise<number> {
|
|
let dispatched = 0
|
|
const batchIds: string[] = []
|
|
const region = await resolveTriggerRegion()
|
|
for (let i = 0; i < jobPayloads.length; i += TRIGGER_BATCH_SIZE) {
|
|
const chunk = jobPayloads.slice(i, i + TRIGGER_BATCH_SIZE)
|
|
try {
|
|
const result = await tasks.batchTrigger<typeof processDocumentTask>(
|
|
'knowledge-process-document',
|
|
chunk.map((payload) => ({
|
|
payload,
|
|
options: {
|
|
// Scoped to (documentId, requestId): blocks intra-dispatch retries
|
|
// from double-enqueuing; later syncs use a fresh requestId.
|
|
idempotencyKey: `doc-process-${payload.documentId}-${requestId}`,
|
|
tags: [
|
|
`knowledgeBaseId:${payload.knowledgeBaseId}`,
|
|
`documentId:${payload.documentId}`,
|
|
],
|
|
region,
|
|
},
|
|
}))
|
|
)
|
|
batchIds.push(result.batchId)
|
|
dispatched += chunk.length
|
|
} catch (error) {
|
|
logger.error(`[${requestId}] Failed to batchTrigger ${chunk.length} document jobs`, {
|
|
error: getErrorMessage(error),
|
|
})
|
|
}
|
|
}
|
|
if (batchIds.length > 0) {
|
|
logger.info(`[${requestId}] Trigger.dev batches dispatched`, { batchIds })
|
|
}
|
|
return dispatched
|
|
}
|
|
|
|
/** Each in-process job runs chunking + embedding + many DB inserts. */
|
|
const IN_PROCESS_DISPATCH_CONCURRENCY = 5
|
|
|
|
async function dispatchInProcess(
|
|
jobPayloads: DocumentProcessingPayload[],
|
|
requestId: string
|
|
): Promise<number> {
|
|
const results = await mapWithConcurrency(
|
|
jobPayloads,
|
|
IN_PROCESS_DISPATCH_CONCURRENCY,
|
|
async (p) => {
|
|
try {
|
|
await processDocumentAsync(
|
|
p.knowledgeBaseId,
|
|
p.documentId,
|
|
p.docData,
|
|
p.processingOptions,
|
|
p
|
|
)
|
|
return true
|
|
} catch (error) {
|
|
logger.error(`[${requestId}] Document dispatch failed`, { error: getErrorMessage(error) })
|
|
return false
|
|
}
|
|
}
|
|
)
|
|
return results.filter(Boolean).length
|
|
}
|
|
|
|
export async function processDocumentAsync(
|
|
knowledgeBaseId: string,
|
|
documentId: string,
|
|
docData: {
|
|
filename: string
|
|
fileUrl: string
|
|
fileSize: number
|
|
mimeType: string
|
|
},
|
|
processingOptions: ProcessingOptions = {},
|
|
providedBillingContext?: BillingAttributionSnapshot | DocumentProcessingBillingContext
|
|
): Promise<void> {
|
|
const startTime = Date.now()
|
|
try {
|
|
logger.info(`[${documentId}] Starting document processing`, {
|
|
knowledgeBaseId,
|
|
mimeType: docData.mimeType,
|
|
fileSize: docData.fileSize,
|
|
})
|
|
|
|
// KB config + workspace billing + doc tags in one JOIN (was 3 SELECTs).
|
|
const contextRows = await db
|
|
.select({
|
|
workspaceId: knowledgeBase.workspaceId,
|
|
knowledgeBaseUserId: knowledgeBase.userId,
|
|
chunkingConfig: knowledgeBase.chunkingConfig,
|
|
embeddingModel: knowledgeBase.embeddingModel,
|
|
billedAccountUserId: workspaceTable.billedAccountUserId,
|
|
uploadedBy: document.uploadedBy,
|
|
filename: document.filename,
|
|
fileUrl: document.fileUrl,
|
|
fileSize: document.fileSize,
|
|
mimeType: document.mimeType,
|
|
tag1: document.tag1,
|
|
tag2: document.tag2,
|
|
tag3: document.tag3,
|
|
tag4: document.tag4,
|
|
tag5: document.tag5,
|
|
tag6: document.tag6,
|
|
tag7: document.tag7,
|
|
number1: document.number1,
|
|
number2: document.number2,
|
|
number3: document.number3,
|
|
number4: document.number4,
|
|
number5: document.number5,
|
|
date1: document.date1,
|
|
date2: document.date2,
|
|
boolean1: document.boolean1,
|
|
boolean2: document.boolean2,
|
|
boolean3: document.boolean3,
|
|
})
|
|
.from(document)
|
|
.innerJoin(knowledgeBase, eq(knowledgeBase.id, document.knowledgeBaseId))
|
|
.leftJoin(
|
|
workspaceTable,
|
|
and(eq(workspaceTable.id, knowledgeBase.workspaceId), isNull(workspaceTable.archivedAt))
|
|
)
|
|
.where(
|
|
and(
|
|
eq(document.id, documentId),
|
|
eq(knowledgeBase.id, knowledgeBaseId),
|
|
isNull(document.archivedAt),
|
|
isNull(document.deletedAt),
|
|
isNull(knowledgeBase.deletedAt)
|
|
)
|
|
)
|
|
.limit(1)
|
|
|
|
if (contextRows.length === 0) {
|
|
logger.warn(
|
|
`[${documentId}] Skipping document processing: document or knowledge base ${knowledgeBaseId} no longer exists`
|
|
)
|
|
await db
|
|
.update(document)
|
|
.set({
|
|
processingStatus: 'failed',
|
|
processingError: 'Document or knowledge base no longer exists',
|
|
processingCompletedAt: new Date(),
|
|
})
|
|
.where(eq(document.id, documentId))
|
|
return
|
|
}
|
|
|
|
const ctx = contextRows[0]
|
|
const persistedDocData = {
|
|
filename: ctx.filename,
|
|
fileUrl: ctx.fileUrl,
|
|
fileSize: ctx.fileSize,
|
|
mimeType: ctx.mimeType,
|
|
}
|
|
|
|
await db
|
|
.update(document)
|
|
.set({
|
|
processingStatus: 'processing',
|
|
processingStartedAt: new Date(),
|
|
processingCompletedAt: null,
|
|
processingError: null,
|
|
})
|
|
.where(
|
|
and(eq(document.id, documentId), isNull(document.archivedAt), isNull(document.deletedAt))
|
|
)
|
|
|
|
logger.info(`[${documentId}] Status updated to 'processing', starting document processor`)
|
|
|
|
const rawConfig = ctx.chunkingConfig as {
|
|
maxSize?: number
|
|
minSize?: number
|
|
overlap?: number
|
|
strategy?: ChunkingStrategy
|
|
strategyOptions?: StrategyOptions
|
|
} | null
|
|
const kbConfig = {
|
|
maxSize: rawConfig?.maxSize ?? 1024,
|
|
minSize: rawConfig?.minSize ?? 100,
|
|
overlap: rawConfig?.overlap ?? 200,
|
|
}
|
|
|
|
const kbEmbeddingModel = ctx.embeddingModel
|
|
const queuedBillingContext = hasDocumentProcessingBillingScope(providedBillingContext)
|
|
? assertDocumentProcessingBillingContext(providedBillingContext)
|
|
: undefined
|
|
const restoredBillingAttribution =
|
|
queuedBillingContext?.billingScope === 'workspace'
|
|
? queuedBillingContext.billingAttribution
|
|
: providedBillingContext && !queuedBillingContext
|
|
? assertBillingAttributionSnapshot(providedBillingContext)
|
|
: undefined
|
|
const documentActorUserId =
|
|
queuedBillingContext?.actorUserId ??
|
|
restoredBillingAttribution?.actorUserId ??
|
|
ctx.uploadedBy ??
|
|
ctx.billedAccountUserId ??
|
|
ctx.knowledgeBaseUserId
|
|
let billingAttribution: BillingAttributionSnapshot | undefined
|
|
if (ctx.workspaceId) {
|
|
if (queuedBillingContext?.billingScope === 'non-workspace') {
|
|
throw new Error('Document processing billing scope does not match knowledge base workspace')
|
|
}
|
|
if (!restoredBillingAttribution) {
|
|
throw new Error('Billing attribution is required for queued document processing')
|
|
}
|
|
billingAttribution = restoredBillingAttribution
|
|
if (
|
|
billingAttribution.actorUserId !== documentActorUserId ||
|
|
billingAttribution.workspaceId !== ctx.workspaceId
|
|
) {
|
|
throw new Error('Document billing attribution does not match its actor and workspace')
|
|
}
|
|
} else if (restoredBillingAttribution || queuedBillingContext?.billingScope === 'workspace') {
|
|
throw new Error('Workspace-less document processing cannot use workspace billing attribution')
|
|
}
|
|
|
|
/**
|
|
* Authoritative gate covering every indexing path. Workspace-less legacy
|
|
* knowledge bases retain account-only enforcement.
|
|
*/
|
|
const usageGate = billingAttribution
|
|
? await checkAttributedUsageLimits(billingAttribution)
|
|
: await checkActorUsageLimits(documentActorUserId)
|
|
if (usageGate.isExceeded) {
|
|
logger.warn(`[${documentId}] Usage limit reached — skipping document indexing`)
|
|
await db
|
|
.update(document)
|
|
.set({
|
|
processingStatus: 'failed',
|
|
processingError:
|
|
usageGate.message ?? 'Usage limit exceeded. Please upgrade your plan to continue.',
|
|
processingCompletedAt: new Date(),
|
|
})
|
|
.where(eq(document.id, documentId))
|
|
return
|
|
}
|
|
let totalEmbeddingTokens = 0
|
|
let embeddingIsBYOK = false
|
|
let embeddingModelName = kbEmbeddingModel
|
|
let embeddingPricingId = kbEmbeddingModel
|
|
|
|
const currentSourceFileProvenance = await loadCurrentWorkspaceSourceFileSecretProvenance({
|
|
fileUrl: persistedDocData.fileUrl,
|
|
})
|
|
const documentSecretContext = await loadKnowledgeDocumentSecretRegistry(
|
|
documentId,
|
|
{
|
|
userId: documentActorUserId,
|
|
...(ctx.workspaceId ? { workspaceId: ctx.workspaceId } : {}),
|
|
},
|
|
currentSourceFileProvenance
|
|
)
|
|
|
|
await withTimeout(
|
|
runWithKnowledgeModelInputProvenance(
|
|
documentSecretContext.registry,
|
|
async () => {
|
|
const processed = await processDocument(
|
|
persistedDocData.fileUrl,
|
|
persistedDocData.filename,
|
|
persistedDocData.mimeType,
|
|
kbConfig.maxSize,
|
|
kbConfig.overlap,
|
|
kbConfig.minSize,
|
|
/**
|
|
* Authorize source-file processing as the actor, not the payer. Using
|
|
* the KB owner would let a writer ingest an internal file that only the
|
|
* owner can read.
|
|
*/
|
|
documentActorUserId,
|
|
ctx.workspaceId,
|
|
rawConfig?.strategy,
|
|
rawConfig?.strategyOptions
|
|
)
|
|
|
|
if (processed.chunks.length > LARGE_DOC_CONFIG.MAX_CHUNKS_PER_DOCUMENT) {
|
|
throw new Error(
|
|
`Document has ${processed.chunks.length.toLocaleString()} chunks, exceeding maximum of ${LARGE_DOC_CONFIG.MAX_CHUNKS_PER_DOCUMENT.toLocaleString()}. ` +
|
|
`This document is unusually large and may need to be split into multiple files or preprocessed to reduce content.`
|
|
)
|
|
}
|
|
|
|
const now = new Date()
|
|
|
|
logger.info(
|
|
`[${documentId}] Document parsed successfully, generating embeddings for ${processed.chunks.length} chunks`
|
|
)
|
|
|
|
const chunkTexts = processed.chunks.map((chunk) => chunk.text)
|
|
const embeddings: number[][] = []
|
|
|
|
if (chunkTexts.length > 0) {
|
|
const batchSize = LARGE_DOC_CONFIG.MAX_EMBEDDING_BATCH
|
|
const totalBatches = Math.ceil(chunkTexts.length / batchSize)
|
|
|
|
logger.info(`[${documentId}] Generating embeddings in ${totalBatches} batches`)
|
|
|
|
for (let i = 0; i < chunkTexts.length; i += batchSize) {
|
|
const batch = chunkTexts.slice(i, i + batchSize)
|
|
const batchNum = Math.floor(i / batchSize) + 1
|
|
|
|
logger.info(`[${documentId}] Processing embedding batch ${batchNum}/${totalBatches}`)
|
|
const {
|
|
embeddings: batchEmbeddings,
|
|
totalTokens: batchTokens,
|
|
isBYOK,
|
|
modelName,
|
|
pricingId,
|
|
} = await generateEmbeddings(batch, kbEmbeddingModel, ctx.workspaceId)
|
|
for (const emb of batchEmbeddings) {
|
|
embeddings.push(emb)
|
|
}
|
|
totalEmbeddingTokens += batchTokens
|
|
if (i === 0) {
|
|
embeddingIsBYOK = isBYOK
|
|
embeddingModelName = modelName
|
|
embeddingPricingId = pricingId
|
|
}
|
|
}
|
|
}
|
|
|
|
// Tag values prefetched above; reuse for the embedding rows.
|
|
const documentTags = ctx
|
|
|
|
logger.info(`[${documentId}] Embeddings generated, creating embedding records with tags`)
|
|
|
|
const tokenizerProvider = getEmbeddingModelInfo(kbEmbeddingModel).tokenizerProvider
|
|
|
|
const chunkProvenances = processed.chunks.map((chunk) =>
|
|
documentSecretContext.tracked
|
|
? documentSecretContext.registry
|
|
? durableSecretProvenanceFromRegistry(documentSecretContext.registry, chunk.text)
|
|
: EXACT_EMPTY_DURABLE_SECRET_PROVENANCE
|
|
: undefined
|
|
)
|
|
const embeddingRecords = processed.chunks.map((chunk, chunkIndex) => ({
|
|
id: generateId(),
|
|
knowledgeBaseId,
|
|
documentId,
|
|
chunkIndex,
|
|
chunkHash: sha256Hex(chunk.text),
|
|
content: chunk.text,
|
|
secretProvenanceVersion: chunkProvenances[chunkIndex] ? 1 : null,
|
|
contentLength: chunk.text.length,
|
|
tokenCount: estimateTokenCount(chunk.text, tokenizerProvider).count,
|
|
embedding: embeddings[chunkIndex] || null,
|
|
embeddingModel: kbEmbeddingModel,
|
|
startOffset: chunk.metadata.startIndex,
|
|
endOffset: chunk.metadata.endIndex,
|
|
tag1: documentTags.tag1,
|
|
tag2: documentTags.tag2,
|
|
tag3: documentTags.tag3,
|
|
tag4: documentTags.tag4,
|
|
tag5: documentTags.tag5,
|
|
tag6: documentTags.tag6,
|
|
tag7: documentTags.tag7,
|
|
number1: documentTags.number1,
|
|
number2: documentTags.number2,
|
|
number3: documentTags.number3,
|
|
number4: documentTags.number4,
|
|
number5: documentTags.number5,
|
|
date1: documentTags.date1,
|
|
date2: documentTags.date2,
|
|
boolean1: documentTags.boolean1,
|
|
boolean2: documentTags.boolean2,
|
|
boolean3: documentTags.boolean3,
|
|
createdAt: now,
|
|
updatedAt: now,
|
|
}))
|
|
|
|
await db.transaction(async (tx) => {
|
|
const activeDocument = await tx
|
|
.select({ id: document.id })
|
|
.from(document)
|
|
.innerJoin(knowledgeBase, eq(document.knowledgeBaseId, knowledgeBase.id))
|
|
.where(
|
|
and(
|
|
eq(document.id, documentId),
|
|
isNull(document.archivedAt),
|
|
isNull(document.deletedAt),
|
|
isNull(knowledgeBase.deletedAt)
|
|
)
|
|
)
|
|
.limit(1)
|
|
|
|
if (activeDocument.length === 0) {
|
|
return
|
|
}
|
|
|
|
if (embeddingRecords.length > 0) {
|
|
await tx.delete(embedding).where(eq(embedding.documentId, documentId))
|
|
|
|
const insertBatchSize = LARGE_DOC_CONFIG.MAX_CHUNKS_PER_BATCH
|
|
const batches: (typeof embeddingRecords)[] = []
|
|
for (let i = 0; i < embeddingRecords.length; i += insertBatchSize) {
|
|
batches.push(embeddingRecords.slice(i, i + insertBatchSize))
|
|
}
|
|
|
|
logger.info(`[${documentId}] Inserting ${embeddingRecords.length} embeddings`)
|
|
for (const batch of batches) {
|
|
await tx.insert(embedding).values(batch)
|
|
}
|
|
const provenanceRecords = embeddingRecords.flatMap((record, index) => {
|
|
const provenance = chunkProvenances[index]
|
|
if (!provenance) return []
|
|
return [
|
|
{
|
|
embeddingId: record.id,
|
|
contentHash: record.chunkHash,
|
|
status: provenance.status,
|
|
entries: provenance.status === 'exact' ? [...provenance.entries] : [],
|
|
updatedAt: now,
|
|
},
|
|
]
|
|
})
|
|
for (let i = 0; i < provenanceRecords.length; i += insertBatchSize) {
|
|
await tx
|
|
.insert(embeddingSecretProvenance)
|
|
.values(provenanceRecords.slice(i, i + insertBatchSize))
|
|
}
|
|
}
|
|
|
|
await tx
|
|
.update(document)
|
|
.set({
|
|
chunkCount: processed.metadata.chunkCount,
|
|
tokenCount: processed.metadata.tokenCount,
|
|
characterCount: processed.metadata.characterCount,
|
|
processingStatus: 'completed',
|
|
processingCompletedAt: now,
|
|
processingError: null,
|
|
})
|
|
.where(eq(document.id, documentId))
|
|
})
|
|
},
|
|
{
|
|
opaqueInputSafe:
|
|
documentSecretContext.provenance.status === 'exact' &&
|
|
documentSecretContext.provenance.entries.length === 0,
|
|
}
|
|
),
|
|
TIMEOUTS.OVERALL_PROCESSING,
|
|
'Document processing'
|
|
)
|
|
|
|
const processingTime = Date.now() - startTime
|
|
logger.info(`[${documentId}] Successfully processed document in ${processingTime}ms`)
|
|
|
|
if (!embeddingIsBYOK && totalEmbeddingTokens > 0) {
|
|
try {
|
|
const costMultiplier = getCostMultiplier()
|
|
const { total: cost } = calculateCost(
|
|
embeddingPricingId,
|
|
totalEmbeddingTokens,
|
|
0,
|
|
false,
|
|
costMultiplier
|
|
)
|
|
if (cost > 0) {
|
|
await recordUsage({
|
|
userId: documentActorUserId,
|
|
workspaceId: ctx.workspaceId ?? undefined,
|
|
...(billingAttribution ? toBillingContext(billingAttribution) : {}),
|
|
entries: [
|
|
{
|
|
category: 'model',
|
|
source: 'knowledge-base',
|
|
description: embeddingModelName,
|
|
cost,
|
|
sourceReference: `knowledge-document:${documentId}:${startTime}`,
|
|
metadata: { inputTokens: totalEmbeddingTokens, outputTokens: 0 },
|
|
},
|
|
],
|
|
})
|
|
if (billingAttribution) {
|
|
await checkAndBillPayerOverageThreshold(billingAttribution.billingEntity)
|
|
} else {
|
|
await checkAndBillOverageThreshold(documentActorUserId)
|
|
}
|
|
} else {
|
|
logger.warn(
|
|
`[${documentId}] Embedding model "${embeddingModelName}" has no pricing entry — billing skipped`,
|
|
{ totalEmbeddingTokens, embeddingModelName }
|
|
)
|
|
}
|
|
} catch (billingError) {
|
|
logger.error(`[${documentId}] Failed to record embedding usage`, { error: billingError })
|
|
}
|
|
}
|
|
} catch (error) {
|
|
const processingTime = Date.now() - startTime
|
|
const errorMessage = getErrorMessage(error, 'Unknown error')
|
|
logger.error(`[${documentId}] Failed to process document after ${processingTime}ms:`, {
|
|
errorType: toError(error).name,
|
|
knowledgeBaseId,
|
|
mimeType: docData.mimeType,
|
|
fileSize: docData.fileSize,
|
|
})
|
|
|
|
await db
|
|
.update(document)
|
|
.set({
|
|
processingStatus: 'failed',
|
|
processingError: errorMessage,
|
|
processingCompletedAt: new Date(),
|
|
})
|
|
.where(eq(document.id, documentId))
|
|
|
|
throw error
|
|
}
|
|
}
|
|
|
|
export function isTriggerAvailable(): boolean {
|
|
return Boolean(env.TRIGGER_SECRET_KEY) && isTriggerDevEnabled
|
|
}
|
|
|
|
type DocumentStorageBilling =
|
|
| {
|
|
readonly context: StorageBillingContext
|
|
readonly bytes: number
|
|
}
|
|
| {
|
|
readonly userId: string
|
|
readonly bytes: number
|
|
readonly sub: HighestPrioritySubscription | null
|
|
}
|
|
|
|
interface DocumentStorageNotification {
|
|
readonly context: StorageBillingContext
|
|
readonly updatedUsage: number
|
|
}
|
|
|
|
interface DocumentStorageAdmission {
|
|
readonly workspaceId: string | null
|
|
readonly knowledgeBaseUserId: string
|
|
readonly billing?: DocumentStorageBilling
|
|
}
|
|
|
|
/**
|
|
* Uses trusted file metadata for a KB object size when that metadata already
|
|
* exists. External/data URLs and legacy unbound objects retain the caller's
|
|
* size; this path deliberately does not add provider HEAD requests.
|
|
*/
|
|
function getServerKnownDocumentSize(
|
|
fileUrl: string,
|
|
fallbackSize: number,
|
|
bindingByKey: ReadonlyMap<string, FileMetadataRecord>
|
|
): number {
|
|
const storageKey = getKnowledgeBaseStorageKey(fileUrl)
|
|
return storageKey ? (bindingByKey.get(storageKey)?.size ?? fallbackSize) : fallbackSize
|
|
}
|
|
|
|
/**
|
|
* Resolves server-known KB object sizes outside the insertion transaction.
|
|
*/
|
|
async function resolveServerKnownDocumentSizes<
|
|
T extends { readonly fileUrl: string; readonly fileSize: number },
|
|
>(documents: readonly T[]): Promise<Array<T & { fileSize: number }>> {
|
|
const bindingByKey = await loadKnowledgeBaseFileBindings(
|
|
documents.map((document) => document.fileUrl)
|
|
)
|
|
return documents.map((docData) => ({
|
|
...docData,
|
|
fileSize: getServerKnownDocumentSize(docData.fileUrl, docData.fileSize, bindingByKey),
|
|
}))
|
|
}
|
|
|
|
/**
|
|
* Resolves storage admission before opening the short document transaction.
|
|
* The transaction revalidates the locked KB ownership snapshot and the
|
|
* workspace helper atomically rechecks quota while applying both ledgers.
|
|
*/
|
|
async function resolveDocumentStorageAdmission(
|
|
knowledgeBaseId: string,
|
|
uploadedBy: string | null,
|
|
bytes: number
|
|
): Promise<DocumentStorageAdmission> {
|
|
const [kb] = await db
|
|
.select({
|
|
workspaceId: knowledgeBase.workspaceId,
|
|
userId: knowledgeBase.userId,
|
|
})
|
|
.from(knowledgeBase)
|
|
.where(and(eq(knowledgeBase.id, knowledgeBaseId), isNull(knowledgeBase.deletedAt)))
|
|
.limit(1)
|
|
if (!kb) {
|
|
throw new Error('Knowledge base not found')
|
|
}
|
|
|
|
if (bytes <= 0) {
|
|
return { workspaceId: kb.workspaceId, knowledgeBaseUserId: kb.userId }
|
|
}
|
|
|
|
const billedUserId = uploadedBy ?? kb.userId
|
|
if (kb.workspaceId) {
|
|
const context = await resolveStorageBillingContext(kb.workspaceId)
|
|
const quotaCheck = await checkStorageQuotaForBillingContext(context, bytes)
|
|
if (!quotaCheck.allowed) {
|
|
throw new Error(quotaCheck.error || 'Storage limit exceeded')
|
|
}
|
|
return {
|
|
workspaceId: kb.workspaceId,
|
|
knowledgeBaseUserId: kb.userId,
|
|
billing: { context, bytes },
|
|
}
|
|
}
|
|
|
|
const [quotaCheck, sub] = await Promise.all([
|
|
checkStorageQuota(billedUserId, bytes),
|
|
getHighestPrioritySubscription(billedUserId),
|
|
])
|
|
if (!quotaCheck.allowed) {
|
|
throw new Error(quotaCheck.error || 'Storage limit exceeded')
|
|
}
|
|
return {
|
|
workspaceId: null,
|
|
knowledgeBaseUserId: kb.userId,
|
|
billing: { userId: billedUserId, bytes, sub },
|
|
}
|
|
}
|
|
|
|
export async function createDocumentRecords(
|
|
documents: Array<{
|
|
filename: string
|
|
fileUrl: string
|
|
fileSize: number
|
|
mimeType: string
|
|
documentTagsData?: string
|
|
tag1?: string
|
|
tag2?: string
|
|
tag3?: string
|
|
tag4?: string
|
|
tag5?: string
|
|
tag6?: string
|
|
tag7?: string
|
|
}>,
|
|
knowledgeBaseId: string,
|
|
requestId: string,
|
|
uploadedBy: string | null = null,
|
|
secretProvenances?: readonly KnowledgeDocumentWriteSecretProvenance[]
|
|
): Promise<DocumentData[]> {
|
|
if (secretProvenances && secretProvenances.length !== documents.length) {
|
|
throw new Error('Knowledge document secret provenance count does not match the request')
|
|
}
|
|
const resolvedDocuments = await resolveServerKnownDocumentSizes(documents)
|
|
const totalBytes = resolvedDocuments.reduce((sum, docData) => sum + (docData.fileSize || 0), 0)
|
|
const admission = await resolveDocumentStorageAdmission(knowledgeBaseId, uploadedBy, totalBytes)
|
|
const { returnData, storageNotification } = await db.transaction(async (tx) => {
|
|
let storageNotification: DocumentStorageNotification | null = null
|
|
|
|
await tx.execute(sql`SELECT 1 FROM knowledge_base WHERE id = ${knowledgeBaseId} FOR UPDATE`)
|
|
|
|
const kb = await tx
|
|
.select({
|
|
id: knowledgeBase.id,
|
|
workspaceId: knowledgeBase.workspaceId,
|
|
userId: knowledgeBase.userId,
|
|
})
|
|
.from(knowledgeBase)
|
|
.where(and(eq(knowledgeBase.id, knowledgeBaseId), isNull(knowledgeBase.deletedAt)))
|
|
.limit(1)
|
|
|
|
if (kb.length === 0) {
|
|
throw new Error('Knowledge base not found')
|
|
}
|
|
|
|
if (
|
|
kb[0].workspaceId !== admission.workspaceId ||
|
|
kb[0].userId !== admission.knowledgeBaseUserId
|
|
) {
|
|
throw new Error(
|
|
'Knowledge base storage ownership changed; retry with fresh storage admission'
|
|
)
|
|
}
|
|
|
|
const kbWorkspaceId = kb[0].workspaceId
|
|
const bindingByKey = await assertKnowledgeBaseFileUrlsOwnership(
|
|
resolvedDocuments.map((docData) => docData.fileUrl),
|
|
kbWorkspaceId,
|
|
kb[0].userId,
|
|
requestId,
|
|
tx
|
|
)
|
|
const sourceBindingByKey = await loadWorkspaceSourceFileBindings(
|
|
resolvedDocuments.map((docData) => docData.fileUrl),
|
|
tx
|
|
)
|
|
const trackedBindings = [
|
|
...[...bindingByKey.values()].filter((binding) => binding.secretProvenanceVersion !== null),
|
|
...sourceBindingByKey.values(),
|
|
]
|
|
const boundFileProvenanceById = await getBoundWorkspaceFileSecretProvenanceByMetadata(
|
|
tx,
|
|
trackedBindings
|
|
)
|
|
for (const [documentIndex, docData] of resolvedDocuments.entries()) {
|
|
const currentSize = getServerKnownDocumentSize(
|
|
docData.fileUrl,
|
|
docData.fileSize,
|
|
bindingByKey
|
|
)
|
|
if (currentSize !== docData.fileSize) {
|
|
throw new Error('Knowledge base file metadata changed; retry document insertion')
|
|
}
|
|
}
|
|
|
|
if (admission.billing) {
|
|
const preparedBilling = admission.billing
|
|
if ('context' in preparedBilling) {
|
|
const updatedUsage = await incrementStorageUsageForBillingContextInTx(
|
|
tx,
|
|
preparedBilling.context,
|
|
preparedBilling.bytes
|
|
)
|
|
if (updatedUsage !== undefined) {
|
|
storageNotification = { context: preparedBilling.context, updatedUsage }
|
|
}
|
|
} else {
|
|
const quotaCheck = await checkAndIncrementStorageUsageInTx(
|
|
tx,
|
|
preparedBilling.sub,
|
|
preparedBilling.userId,
|
|
preparedBilling.bytes
|
|
)
|
|
if (!quotaCheck.allowed) {
|
|
throw new Error(quotaCheck.error || 'Storage limit exceeded')
|
|
}
|
|
}
|
|
}
|
|
|
|
// One load per batch (was N+1); skip entirely if no doc carries tags.
|
|
const hasTaggedDocs = resolvedDocuments.some((d) => d.documentTagsData)
|
|
const tagDefinitions = hasTaggedDocs
|
|
? await loadTagDefinitions(knowledgeBaseId, tx)
|
|
: (new Map() as TagDefinitionsByName)
|
|
|
|
const now = new Date()
|
|
const documentRecords = []
|
|
const documentProvenances: (DurableSecretProvenance | undefined)[] = []
|
|
const returnData: DocumentData[] = []
|
|
|
|
for (const [documentIndex, docData] of resolvedDocuments.entries()) {
|
|
const documentId = generateId()
|
|
|
|
let processedTags: Partial<ProcessedDocumentTags> = {}
|
|
|
|
if (docData.documentTagsData) {
|
|
try {
|
|
const tagData = JSON.parse(docData.documentTagsData)
|
|
if (Array.isArray(tagData)) {
|
|
processedTags = resolveDocumentTags(tagData, tagDefinitions, requestId)
|
|
}
|
|
} catch (error) {
|
|
if (error instanceof SyntaxError) {
|
|
logger.warn(`[${requestId}] Failed to parse documentTagsData for bulk document:`, error)
|
|
} else {
|
|
throw error
|
|
}
|
|
}
|
|
}
|
|
|
|
const storageKey = getKnowledgeBaseStorageKey(docData.fileUrl)
|
|
const baseDocument = {
|
|
id: documentId,
|
|
knowledgeBaseId,
|
|
filename: docData.filename,
|
|
fileUrl: docData.fileUrl,
|
|
storageKey,
|
|
contentHash: null,
|
|
fileSize: docData.fileSize,
|
|
mimeType: docData.mimeType,
|
|
chunkCount: 0,
|
|
tokenCount: 0,
|
|
characterCount: 0,
|
|
processingStatus: 'pending' as const,
|
|
enabled: true,
|
|
uploadedAt: now,
|
|
uploadedBy,
|
|
tag1: processedTags.tag1 ?? docData.tag1 ?? null,
|
|
tag2: processedTags.tag2 ?? docData.tag2 ?? null,
|
|
tag3: processedTags.tag3 ?? docData.tag3 ?? null,
|
|
tag4: processedTags.tag4 ?? docData.tag4 ?? null,
|
|
tag5: processedTags.tag5 ?? docData.tag5 ?? null,
|
|
tag6: processedTags.tag6 ?? docData.tag6 ?? null,
|
|
tag7: processedTags.tag7 ?? docData.tag7 ?? null,
|
|
number1: processedTags.number1 ?? null,
|
|
number2: processedTags.number2 ?? null,
|
|
number3: processedTags.number3 ?? null,
|
|
number4: processedTags.number4 ?? null,
|
|
number5: processedTags.number5 ?? null,
|
|
date1: processedTags.date1 ?? null,
|
|
date2: processedTags.date2 ?? null,
|
|
boolean1: processedTags.boolean1 ?? null,
|
|
boolean2: processedTags.boolean2 ?? null,
|
|
boolean3: processedTags.boolean3 ?? null,
|
|
}
|
|
const source = createKnowledgeDocumentSourceValue(baseDocument)
|
|
const binding = storageKey
|
|
? (bindingByKey.get(storageKey) ?? sourceBindingByKey.get(storageKey))
|
|
: undefined
|
|
const provenanceBinding =
|
|
binding && (binding.secretProvenanceVersion !== null || sourceBindingByKey.has(binding.key))
|
|
? binding
|
|
: undefined
|
|
const boundFileProvenance = provenanceBinding
|
|
? (boundFileProvenanceById.get(provenanceBinding.id) ?? { status: 'unknown' as const })
|
|
: undefined
|
|
const provenance = bindKnowledgeDocumentWriteSecretProvenance({
|
|
source,
|
|
provenance: secretProvenances?.[documentIndex],
|
|
tagDefinitions,
|
|
...(provenanceBinding && boundFileProvenance
|
|
? { boundFile: { binding: provenanceBinding, provenance: boundFileProvenance } }
|
|
: {}),
|
|
})
|
|
const newDocument = {
|
|
...baseDocument,
|
|
secretProvenanceVersion: provenance ? 1 : null,
|
|
}
|
|
|
|
documentRecords.push(newDocument)
|
|
documentProvenances.push(provenance)
|
|
returnData.push({
|
|
documentId,
|
|
filename: docData.filename,
|
|
fileUrl: docData.fileUrl,
|
|
fileSize: docData.fileSize,
|
|
mimeType: docData.mimeType,
|
|
})
|
|
}
|
|
|
|
if (documentRecords.length > 0) {
|
|
await tx.insert(document).values(documentRecords)
|
|
for (const [documentIndex, record] of documentRecords.entries()) {
|
|
const provenance = documentProvenances[documentIndex]
|
|
if (!provenance) continue
|
|
await replaceKnowledgeDocumentSecretProvenanceInTx(
|
|
tx,
|
|
record.id,
|
|
createKnowledgeDocumentSourceValue(record),
|
|
provenance
|
|
)
|
|
}
|
|
logger.info(
|
|
`[${requestId}] Bulk created ${documentRecords.length} document records in knowledge base ${knowledgeBaseId}`
|
|
)
|
|
|
|
await tx
|
|
.update(knowledgeBase)
|
|
.set({ updatedAt: now })
|
|
.where(eq(knowledgeBase.id, knowledgeBaseId))
|
|
}
|
|
|
|
return { returnData, storageNotification }
|
|
})
|
|
|
|
if (storageNotification) {
|
|
void maybeNotifyStorageLimitForBillingContext(
|
|
storageNotification.context,
|
|
storageNotification.updatedUsage
|
|
)
|
|
}
|
|
|
|
return returnData
|
|
}
|
|
|
|
export async function getDocuments(
|
|
knowledgeBaseId: string,
|
|
options: {
|
|
enabledFilter?: 'all' | 'enabled' | 'disabled'
|
|
search?: string
|
|
limit?: number
|
|
offset?: number
|
|
sortBy?: DocumentSortField
|
|
sortOrder?: SortOrder
|
|
tagFilters?: TagFilterCondition[]
|
|
},
|
|
requestId: string
|
|
): Promise<{
|
|
documents: Array<{
|
|
id: string
|
|
knowledgeBaseId: string
|
|
filename: string
|
|
fileUrl: string
|
|
fileSize: number
|
|
mimeType: string
|
|
chunkCount: number
|
|
tokenCount: number
|
|
characterCount: number
|
|
processingStatus: 'pending' | 'processing' | 'completed' | 'failed'
|
|
processingStartedAt: Date | null
|
|
processingCompletedAt: Date | null
|
|
processingError: string | null
|
|
enabled: boolean
|
|
uploadedAt: Date
|
|
tag1: string | null
|
|
tag2: string | null
|
|
tag3: string | null
|
|
tag4: string | null
|
|
tag5: string | null
|
|
tag6: string | null
|
|
tag7: string | null
|
|
number1: number | null
|
|
number2: number | null
|
|
number3: number | null
|
|
number4: number | null
|
|
number5: number | null
|
|
date1: Date | null
|
|
date2: Date | null
|
|
boolean1: boolean | null
|
|
boolean2: boolean | null
|
|
boolean3: boolean | null
|
|
connectorId: string | null
|
|
connectorType: string | null
|
|
sourceUrl: string | null
|
|
}>
|
|
pagination: {
|
|
total: number
|
|
limit: number
|
|
offset: number
|
|
hasMore: boolean
|
|
}
|
|
}> {
|
|
const {
|
|
enabledFilter = 'all',
|
|
search,
|
|
limit = 50,
|
|
offset = 0,
|
|
sortBy = 'filename',
|
|
sortOrder = 'asc',
|
|
tagFilters,
|
|
} = options
|
|
|
|
const whereConditions: (SQL | undefined)[] = [
|
|
eq(document.knowledgeBaseId, knowledgeBaseId),
|
|
eq(document.userExcluded, false),
|
|
isNull(document.archivedAt),
|
|
isNull(document.deletedAt),
|
|
]
|
|
|
|
if (enabledFilter === 'enabled') {
|
|
whereConditions.push(eq(document.enabled, true))
|
|
} else if (enabledFilter === 'disabled') {
|
|
whereConditions.push(eq(document.enabled, false))
|
|
}
|
|
|
|
if (search) {
|
|
whereConditions.push(sql`LOWER(${document.filename}) LIKE LOWER(${`%${search}%`})`)
|
|
}
|
|
|
|
if (tagFilters && tagFilters.length > 0) {
|
|
for (const filter of tagFilters) {
|
|
const condition = buildTagFilterCondition(filter)
|
|
if (condition) {
|
|
whereConditions.push(condition)
|
|
}
|
|
}
|
|
}
|
|
|
|
const totalResult = await db
|
|
.select({ count: sql<number>`COUNT(*)` })
|
|
.from(document)
|
|
.where(and(...whereConditions))
|
|
|
|
const total = Number(totalResult[0]?.count ?? 0)
|
|
const hasMore = offset + limit < total
|
|
|
|
const getOrderByColumn = () => {
|
|
switch (sortBy) {
|
|
case 'filename':
|
|
return document.filename
|
|
case 'fileSize':
|
|
return document.fileSize
|
|
case 'tokenCount':
|
|
return document.tokenCount
|
|
case 'chunkCount':
|
|
return document.chunkCount
|
|
case 'uploadedAt':
|
|
return document.uploadedAt
|
|
case 'processingStatus':
|
|
return document.processingStatus
|
|
case 'enabled':
|
|
return document.enabled
|
|
default:
|
|
return document.uploadedAt
|
|
}
|
|
}
|
|
|
|
const primaryOrderBy = sortOrder === 'asc' ? asc(getOrderByColumn()) : desc(getOrderByColumn())
|
|
const secondaryOrderBy =
|
|
sortBy === 'filename' ? desc(document.uploadedAt) : asc(document.filename)
|
|
|
|
const documents = await db
|
|
.select({
|
|
id: document.id,
|
|
knowledgeBaseId: document.knowledgeBaseId,
|
|
filename: document.filename,
|
|
fileUrl: document.fileUrl,
|
|
fileSize: document.fileSize,
|
|
mimeType: document.mimeType,
|
|
chunkCount: document.chunkCount,
|
|
tokenCount: document.tokenCount,
|
|
characterCount: document.characterCount,
|
|
processingStatus: document.processingStatus,
|
|
processingStartedAt: document.processingStartedAt,
|
|
processingCompletedAt: document.processingCompletedAt,
|
|
processingError: document.processingError,
|
|
enabled: document.enabled,
|
|
uploadedAt: document.uploadedAt,
|
|
tag1: document.tag1,
|
|
tag2: document.tag2,
|
|
tag3: document.tag3,
|
|
tag4: document.tag4,
|
|
tag5: document.tag5,
|
|
tag6: document.tag6,
|
|
tag7: document.tag7,
|
|
number1: document.number1,
|
|
number2: document.number2,
|
|
number3: document.number3,
|
|
number4: document.number4,
|
|
number5: document.number5,
|
|
date1: document.date1,
|
|
date2: document.date2,
|
|
boolean1: document.boolean1,
|
|
boolean2: document.boolean2,
|
|
boolean3: document.boolean3,
|
|
connectorId: document.connectorId,
|
|
connectorType: knowledgeConnector.connectorType,
|
|
sourceUrl: document.sourceUrl,
|
|
})
|
|
.from(document)
|
|
.leftJoin(knowledgeConnector, eq(document.connectorId, knowledgeConnector.id))
|
|
.where(and(...whereConditions))
|
|
.orderBy(primaryOrderBy, secondaryOrderBy)
|
|
.limit(limit)
|
|
.offset(offset)
|
|
|
|
logger.info(
|
|
`[${requestId}] Retrieved ${documents.length} documents (${offset}-${offset + documents.length} of ${total}) for knowledge base ${knowledgeBaseId}`
|
|
)
|
|
|
|
return {
|
|
documents: documents.map((doc) => ({
|
|
id: doc.id,
|
|
knowledgeBaseId: doc.knowledgeBaseId,
|
|
filename: doc.filename,
|
|
fileUrl: doc.fileUrl,
|
|
fileSize: doc.fileSize,
|
|
mimeType: doc.mimeType,
|
|
chunkCount: doc.chunkCount,
|
|
tokenCount: doc.tokenCount,
|
|
characterCount: doc.characterCount,
|
|
processingStatus: doc.processingStatus as 'pending' | 'processing' | 'completed' | 'failed',
|
|
processingStartedAt: doc.processingStartedAt,
|
|
processingCompletedAt: doc.processingCompletedAt,
|
|
processingError: doc.processingError,
|
|
enabled: doc.enabled,
|
|
uploadedAt: doc.uploadedAt,
|
|
tag1: doc.tag1,
|
|
tag2: doc.tag2,
|
|
tag3: doc.tag3,
|
|
tag4: doc.tag4,
|
|
tag5: doc.tag5,
|
|
tag6: doc.tag6,
|
|
tag7: doc.tag7,
|
|
number1: doc.number1,
|
|
number2: doc.number2,
|
|
number3: doc.number3,
|
|
number4: doc.number4,
|
|
number5: doc.number5,
|
|
date1: doc.date1,
|
|
date2: doc.date2,
|
|
boolean1: doc.boolean1,
|
|
boolean2: doc.boolean2,
|
|
boolean3: doc.boolean3,
|
|
connectorId: doc.connectorId,
|
|
connectorType: doc.connectorType ?? null,
|
|
sourceUrl: doc.sourceUrl,
|
|
})),
|
|
pagination: {
|
|
total,
|
|
limit,
|
|
offset,
|
|
hasMore,
|
|
},
|
|
}
|
|
}
|
|
|
|
export async function createSingleDocument(
|
|
documentData: {
|
|
filename: string
|
|
fileUrl: string
|
|
fileSize: number
|
|
mimeType: string
|
|
documentTagsData?: string
|
|
tag1?: string
|
|
tag2?: string
|
|
tag3?: string
|
|
tag4?: string
|
|
tag5?: string
|
|
tag6?: string
|
|
tag7?: string
|
|
},
|
|
knowledgeBaseId: string,
|
|
requestId: string,
|
|
uploadedBy: string | null = null,
|
|
secretProvenance?: KnowledgeDocumentWriteSecretProvenance
|
|
): Promise<{
|
|
id: string
|
|
knowledgeBaseId: string
|
|
filename: string
|
|
fileUrl: string
|
|
fileSize: number
|
|
mimeType: string
|
|
chunkCount: number
|
|
tokenCount: number
|
|
characterCount: number
|
|
enabled: boolean
|
|
uploadedAt: Date
|
|
tag1: string | null
|
|
tag2: string | null
|
|
tag3: string | null
|
|
tag4: string | null
|
|
tag5: string | null
|
|
tag6: string | null
|
|
tag7: string | null
|
|
}> {
|
|
const documentId = generateId()
|
|
const now = new Date()
|
|
const [resolvedDocumentData] = await resolveServerKnownDocumentSizes([documentData])
|
|
const admission = await resolveDocumentStorageAdmission(
|
|
knowledgeBaseId,
|
|
uploadedBy,
|
|
resolvedDocumentData.fileSize
|
|
)
|
|
let processedTags: ProcessedDocumentTags = {
|
|
tag1: documentData.tag1 ?? null,
|
|
tag2: documentData.tag2 ?? null,
|
|
tag3: documentData.tag3 ?? null,
|
|
tag4: documentData.tag4 ?? null,
|
|
tag5: documentData.tag5 ?? null,
|
|
tag6: documentData.tag6 ?? null,
|
|
tag7: documentData.tag7 ?? null,
|
|
number1: null,
|
|
number2: null,
|
|
number3: null,
|
|
number4: null,
|
|
number5: null,
|
|
date1: null,
|
|
date2: null,
|
|
boolean1: null,
|
|
boolean2: null,
|
|
boolean3: null,
|
|
}
|
|
let tagDefinitions: TagDefinitionsByName = new Map()
|
|
|
|
if (documentData.documentTagsData) {
|
|
try {
|
|
const tagData = JSON.parse(documentData.documentTagsData)
|
|
if (Array.isArray(tagData)) {
|
|
tagDefinitions = await loadTagDefinitions(knowledgeBaseId)
|
|
processedTags = resolveDocumentTags(tagData, tagDefinitions, requestId)
|
|
}
|
|
} catch (error) {
|
|
if (error instanceof SyntaxError) {
|
|
logger.warn(`[${requestId}] Failed to parse documentTagsData:`, error)
|
|
} else {
|
|
throw error
|
|
}
|
|
}
|
|
}
|
|
|
|
const newDocument = {
|
|
id: documentId,
|
|
knowledgeBaseId,
|
|
filename: resolvedDocumentData.filename,
|
|
fileUrl: resolvedDocumentData.fileUrl,
|
|
storageKey: getKnowledgeBaseStorageKey(resolvedDocumentData.fileUrl),
|
|
fileSize: resolvedDocumentData.fileSize,
|
|
mimeType: resolvedDocumentData.mimeType,
|
|
chunkCount: 0,
|
|
tokenCount: 0,
|
|
characterCount: 0,
|
|
enabled: true,
|
|
uploadedAt: now,
|
|
uploadedBy,
|
|
...processedTags,
|
|
}
|
|
|
|
const storageNotification = await db.transaction(async (tx) => {
|
|
let storageNotification: DocumentStorageNotification | null = null
|
|
|
|
await tx.execute(sql`SELECT 1 FROM knowledge_base WHERE id = ${knowledgeBaseId} FOR UPDATE`)
|
|
|
|
const kb = await tx
|
|
.select({
|
|
id: knowledgeBase.id,
|
|
workspaceId: knowledgeBase.workspaceId,
|
|
userId: knowledgeBase.userId,
|
|
})
|
|
.from(knowledgeBase)
|
|
.where(and(eq(knowledgeBase.id, knowledgeBaseId), isNull(knowledgeBase.deletedAt)))
|
|
.limit(1)
|
|
|
|
if (kb.length === 0) {
|
|
throw new Error('Knowledge base not found')
|
|
}
|
|
|
|
if (
|
|
kb[0].workspaceId !== admission.workspaceId ||
|
|
kb[0].userId !== admission.knowledgeBaseUserId
|
|
) {
|
|
throw new Error(
|
|
'Knowledge base storage ownership changed; retry with fresh storage admission'
|
|
)
|
|
}
|
|
|
|
const bindingByKey = await assertKnowledgeBaseFileUrlsOwnership(
|
|
[resolvedDocumentData.fileUrl],
|
|
kb[0].workspaceId,
|
|
kb[0].userId,
|
|
requestId,
|
|
tx
|
|
)
|
|
const sourceBindingByKey = await loadWorkspaceSourceFileBindings(
|
|
[resolvedDocumentData.fileUrl],
|
|
tx
|
|
)
|
|
const storageKey = getKnowledgeBaseStorageKey(resolvedDocumentData.fileUrl)
|
|
const binding = storageKey
|
|
? (bindingByKey.get(storageKey) ?? sourceBindingByKey.get(storageKey))
|
|
: undefined
|
|
const provenanceBinding =
|
|
binding && (binding.secretProvenanceVersion !== null || sourceBindingByKey.has(binding.key))
|
|
? binding
|
|
: undefined
|
|
const boundFileProvenanceById = await getBoundWorkspaceFileSecretProvenanceByMetadata(
|
|
tx,
|
|
provenanceBinding ? [provenanceBinding] : []
|
|
)
|
|
const currentSize = getServerKnownDocumentSize(
|
|
resolvedDocumentData.fileUrl,
|
|
resolvedDocumentData.fileSize,
|
|
bindingByKey
|
|
)
|
|
if (currentSize !== resolvedDocumentData.fileSize) {
|
|
throw new Error('Knowledge base file metadata changed; retry document insertion')
|
|
}
|
|
|
|
if (admission.billing) {
|
|
const preparedBilling = admission.billing
|
|
if ('context' in preparedBilling) {
|
|
const updatedUsage = await incrementStorageUsageForBillingContextInTx(
|
|
tx,
|
|
preparedBilling.context,
|
|
preparedBilling.bytes
|
|
)
|
|
if (updatedUsage !== undefined) {
|
|
storageNotification = { context: preparedBilling.context, updatedUsage }
|
|
}
|
|
} else {
|
|
const quotaCheck = await checkAndIncrementStorageUsageInTx(
|
|
tx,
|
|
preparedBilling.sub,
|
|
preparedBilling.userId,
|
|
preparedBilling.bytes
|
|
)
|
|
if (!quotaCheck.allowed) {
|
|
throw new Error(quotaCheck.error || 'Storage limit exceeded')
|
|
}
|
|
}
|
|
}
|
|
|
|
const source = createKnowledgeDocumentSourceValue(newDocument)
|
|
const boundFileProvenance = provenanceBinding
|
|
? (boundFileProvenanceById.get(provenanceBinding.id) ?? { status: 'unknown' as const })
|
|
: undefined
|
|
const documentProvenance = bindKnowledgeDocumentWriteSecretProvenance({
|
|
source,
|
|
provenance: secretProvenance,
|
|
tagDefinitions,
|
|
...(provenanceBinding && boundFileProvenance
|
|
? { boundFile: { binding: provenanceBinding, provenance: boundFileProvenance } }
|
|
: {}),
|
|
})
|
|
await tx.insert(document).values({
|
|
...newDocument,
|
|
secretProvenanceVersion: documentProvenance ? 1 : null,
|
|
})
|
|
if (documentProvenance) {
|
|
await replaceKnowledgeDocumentSecretProvenanceInTx(tx, documentId, source, documentProvenance)
|
|
}
|
|
|
|
await tx
|
|
.update(knowledgeBase)
|
|
.set({ updatedAt: now })
|
|
.where(eq(knowledgeBase.id, knowledgeBaseId))
|
|
|
|
return storageNotification
|
|
})
|
|
|
|
if (storageNotification) {
|
|
void maybeNotifyStorageLimitForBillingContext(
|
|
storageNotification.context,
|
|
storageNotification.updatedUsage
|
|
)
|
|
}
|
|
|
|
logger.info(`[${requestId}] Document created: ${documentId} in knowledge base ${knowledgeBaseId}`)
|
|
|
|
return newDocument as {
|
|
id: string
|
|
knowledgeBaseId: string
|
|
filename: string
|
|
fileUrl: string
|
|
fileSize: number
|
|
mimeType: string
|
|
chunkCount: number
|
|
tokenCount: number
|
|
characterCount: number
|
|
enabled: boolean
|
|
uploadedAt: Date
|
|
tag1: string | null
|
|
tag2: string | null
|
|
tag3: string | null
|
|
tag4: string | null
|
|
tag5: string | null
|
|
tag6: string | null
|
|
tag7: string | null
|
|
}
|
|
}
|
|
|
|
export async function bulkDocumentOperation(
|
|
knowledgeBaseId: string,
|
|
operation: 'enable' | 'disable' | 'delete',
|
|
documentIds: string[],
|
|
requestId: string
|
|
): Promise<{
|
|
success: boolean
|
|
successCount: number
|
|
updatedDocuments: Array<{
|
|
id: string
|
|
enabled?: boolean
|
|
deletedAt?: Date | null
|
|
processingStatus?: string
|
|
}>
|
|
}> {
|
|
logger.info(
|
|
`[${requestId}] Starting bulk ${operation} operation on ${documentIds.length} documents in knowledge base ${knowledgeBaseId}`
|
|
)
|
|
|
|
const documentsToUpdate = await db
|
|
.select({
|
|
id: document.id,
|
|
enabled: document.enabled,
|
|
})
|
|
.from(document)
|
|
.where(
|
|
and(
|
|
eq(document.knowledgeBaseId, knowledgeBaseId),
|
|
inArray(document.id, documentIds),
|
|
eq(document.userExcluded, false),
|
|
isNull(document.archivedAt),
|
|
isNull(document.deletedAt)
|
|
)
|
|
)
|
|
|
|
if (documentsToUpdate.length === 0) {
|
|
throw new Error('No valid documents found to update')
|
|
}
|
|
|
|
if (documentsToUpdate.length !== documentIds.length) {
|
|
logger.warn(
|
|
`[${requestId}] Some documents not found or don't belong to knowledge base. Requested: ${documentIds.length}, Found: ${documentsToUpdate.length}`
|
|
)
|
|
}
|
|
|
|
let updateResult: Array<{
|
|
id: string
|
|
enabled?: boolean
|
|
deletedAt?: Date | null
|
|
processingStatus?: string
|
|
}>
|
|
|
|
if (operation === 'delete') {
|
|
const deletedIds = documentsToUpdate.map((doc) => doc.id)
|
|
const deletedCount = await deleteDocumentsByLifecyclePolicy(deletedIds, requestId)
|
|
updateResult = deletedIds.slice(0, deletedCount).map((id) => ({ id }))
|
|
} else {
|
|
const enabled = operation === 'enable'
|
|
|
|
updateResult = await db
|
|
.update(document)
|
|
.set({
|
|
enabled,
|
|
})
|
|
.where(
|
|
and(
|
|
eq(document.knowledgeBaseId, knowledgeBaseId),
|
|
inArray(document.id, documentIds),
|
|
eq(document.userExcluded, false),
|
|
isNull(document.archivedAt),
|
|
isNull(document.deletedAt)
|
|
)
|
|
)
|
|
.returning({ id: document.id, enabled: document.enabled })
|
|
}
|
|
|
|
const successCount = updateResult.length
|
|
|
|
logger.info(
|
|
`[${requestId}] Bulk ${operation} operation completed: ${successCount} documents updated in knowledge base ${knowledgeBaseId}`
|
|
)
|
|
|
|
return {
|
|
success: true,
|
|
successCount,
|
|
updatedDocuments: updateResult,
|
|
}
|
|
}
|
|
|
|
export async function bulkDocumentOperationByFilter(
|
|
knowledgeBaseId: string,
|
|
operation: 'enable' | 'disable' | 'delete',
|
|
enabledFilter: 'all' | 'enabled' | 'disabled' | undefined,
|
|
requestId: string
|
|
): Promise<{
|
|
success: boolean
|
|
successCount: number
|
|
updatedDocuments: Array<{
|
|
id: string
|
|
enabled?: boolean
|
|
deletedAt?: Date | null
|
|
}>
|
|
}> {
|
|
logger.info(
|
|
`[${requestId}] Starting bulk ${operation} operation on all documents (filter: ${enabledFilter || 'all'}) in knowledge base ${knowledgeBaseId}`
|
|
)
|
|
|
|
const whereConditions = [
|
|
eq(document.knowledgeBaseId, knowledgeBaseId),
|
|
eq(document.userExcluded, false),
|
|
isNull(document.archivedAt),
|
|
isNull(document.deletedAt),
|
|
]
|
|
|
|
if (enabledFilter === 'enabled') {
|
|
whereConditions.push(eq(document.enabled, true))
|
|
} else if (enabledFilter === 'disabled') {
|
|
whereConditions.push(eq(document.enabled, false))
|
|
}
|
|
|
|
let updateResult: Array<{
|
|
id: string
|
|
enabled?: boolean
|
|
deletedAt?: Date | null
|
|
}>
|
|
|
|
if (operation === 'delete') {
|
|
const matchingDocs = await db
|
|
.select({ id: document.id })
|
|
.from(document)
|
|
.where(and(...whereConditions))
|
|
|
|
const deletedIds = matchingDocs.map((doc) => doc.id)
|
|
const deletedCount = await deleteDocumentsByLifecyclePolicy(deletedIds, requestId)
|
|
updateResult = deletedIds.slice(0, deletedCount).map((id) => ({ id }))
|
|
} else {
|
|
const enabled = operation === 'enable'
|
|
|
|
updateResult = await db
|
|
.update(document)
|
|
.set({
|
|
enabled,
|
|
})
|
|
.where(and(...whereConditions))
|
|
.returning({ id: document.id, enabled: document.enabled })
|
|
}
|
|
|
|
const successCount = updateResult.length
|
|
|
|
logger.info(
|
|
`[${requestId}] Bulk ${operation} by filter completed: ${successCount} documents updated in knowledge base ${knowledgeBaseId}`
|
|
)
|
|
|
|
return {
|
|
success: true,
|
|
successCount,
|
|
updatedDocuments: updateResult,
|
|
}
|
|
}
|
|
|
|
export async function markDocumentAsFailedTimeout(
|
|
documentId: string,
|
|
processingStartedAt: Date,
|
|
requestId: string
|
|
): Promise<{ success: boolean; processingDuration: number }> {
|
|
const now = new Date()
|
|
const processingDuration = now.getTime() - processingStartedAt.getTime()
|
|
const DEAD_PROCESS_THRESHOLD_MS = 600 * 1000 // 10 minutes
|
|
|
|
if (processingDuration <= DEAD_PROCESS_THRESHOLD_MS) {
|
|
throw new Error('Document has not been processing long enough to be considered dead')
|
|
}
|
|
|
|
await db
|
|
.update(document)
|
|
.set({
|
|
processingStatus: 'failed',
|
|
processingError: 'Processing timed out. Please retry or re-sync the connector.',
|
|
processingCompletedAt: now,
|
|
})
|
|
.where(eq(document.id, documentId))
|
|
|
|
logger.info(
|
|
`[${requestId}] Marked document ${documentId} as failed due to dead process (processing time: ${Math.round(processingDuration / 1000)}s)`
|
|
)
|
|
|
|
return {
|
|
success: true,
|
|
processingDuration,
|
|
}
|
|
}
|
|
|
|
export async function retryDocumentProcessing(
|
|
knowledgeBaseId: string,
|
|
documentId: string,
|
|
docData: {
|
|
filename: string
|
|
fileUrl: string
|
|
fileSize: number
|
|
mimeType: string
|
|
},
|
|
requestId: string,
|
|
billingAttribution: BillingAttributionSnapshot | undefined
|
|
): Promise<{ success: boolean; status: string; message: string }> {
|
|
await db.transaction(async (tx) => {
|
|
await tx.delete(embedding).where(eq(embedding.documentId, documentId))
|
|
|
|
await tx
|
|
.update(document)
|
|
.set({
|
|
processingStatus: 'pending',
|
|
processingStartedAt: null,
|
|
processingCompletedAt: null,
|
|
processingError: null,
|
|
chunkCount: 0,
|
|
tokenCount: 0,
|
|
characterCount: 0,
|
|
})
|
|
.where(eq(document.id, documentId))
|
|
})
|
|
|
|
await processDocumentsWithQueue(
|
|
[
|
|
{
|
|
documentId,
|
|
filename: docData.filename,
|
|
fileUrl: docData.fileUrl,
|
|
fileSize: docData.fileSize,
|
|
mimeType: docData.mimeType,
|
|
},
|
|
],
|
|
knowledgeBaseId,
|
|
{},
|
|
requestId,
|
|
billingAttribution
|
|
)
|
|
|
|
logger.info(`[${requestId}] Document retry initiated: ${documentId}`)
|
|
|
|
return {
|
|
success: true,
|
|
status: 'pending',
|
|
message: 'Document retry processing started',
|
|
}
|
|
}
|
|
|
|
export async function updateDocument(
|
|
documentId: string,
|
|
updateData: {
|
|
filename?: string
|
|
enabled?: boolean
|
|
chunkCount?: number
|
|
tokenCount?: number
|
|
characterCount?: number
|
|
processingStatus?: 'pending' | 'processing' | 'completed' | 'failed'
|
|
processingError?: string
|
|
tag1?: string
|
|
tag2?: string
|
|
tag3?: string
|
|
tag4?: string
|
|
tag5?: string
|
|
tag6?: string
|
|
tag7?: string
|
|
number1?: string
|
|
number2?: string
|
|
number3?: string
|
|
number4?: string
|
|
number5?: string
|
|
date1?: string
|
|
date2?: string
|
|
boolean1?: string
|
|
boolean2?: string
|
|
boolean3?: string
|
|
},
|
|
requestId: string
|
|
): Promise<{
|
|
id: string
|
|
knowledgeBaseId: string
|
|
filename: string
|
|
fileUrl: string
|
|
fileSize: number
|
|
mimeType: string
|
|
chunkCount: number
|
|
tokenCount: number
|
|
characterCount: number
|
|
processingStatus: 'pending' | 'processing' | 'completed' | 'failed'
|
|
processingStartedAt: Date | null
|
|
processingCompletedAt: Date | null
|
|
processingError: string | null
|
|
enabled: boolean
|
|
uploadedAt: Date
|
|
tag1: string | null
|
|
tag2: string | null
|
|
tag3: string | null
|
|
tag4: string | null
|
|
tag5: string | null
|
|
tag6: string | null
|
|
tag7: string | null
|
|
number1: number | null
|
|
number2: number | null
|
|
number3: number | null
|
|
number4: number | null
|
|
number5: number | null
|
|
date1: Date | null
|
|
date2: Date | null
|
|
boolean1: boolean | null
|
|
boolean2: boolean | null
|
|
boolean3: boolean | null
|
|
deletedAt: Date | null
|
|
}> {
|
|
const dbUpdateData: Partial<{
|
|
filename: string
|
|
enabled: boolean
|
|
chunkCount: number
|
|
tokenCount: number
|
|
characterCount: number
|
|
processingStatus: 'pending' | 'processing' | 'completed' | 'failed'
|
|
processingError: string | null
|
|
processingStartedAt: Date | null
|
|
processingCompletedAt: Date | null
|
|
tag1: string | null
|
|
tag2: string | null
|
|
tag3: string | null
|
|
tag4: string | null
|
|
tag5: string | null
|
|
tag6: string | null
|
|
tag7: string | null
|
|
number1: number | null
|
|
number2: number | null
|
|
number3: number | null
|
|
number4: number | null
|
|
number5: number | null
|
|
date1: Date | null
|
|
date2: Date | null
|
|
boolean1: boolean | null
|
|
boolean2: boolean | null
|
|
boolean3: boolean | null
|
|
}> = {}
|
|
const ALL_TAG_SLOTS = [
|
|
'tag1',
|
|
'tag2',
|
|
'tag3',
|
|
'tag4',
|
|
'tag5',
|
|
'tag6',
|
|
'tag7',
|
|
'number1',
|
|
'number2',
|
|
'number3',
|
|
'number4',
|
|
'number5',
|
|
'date1',
|
|
'date2',
|
|
'boolean1',
|
|
'boolean2',
|
|
'boolean3',
|
|
] as const
|
|
type TagSlot = (typeof ALL_TAG_SLOTS)[number]
|
|
|
|
if (updateData.filename !== undefined) dbUpdateData.filename = updateData.filename
|
|
if (updateData.enabled !== undefined) dbUpdateData.enabled = updateData.enabled
|
|
if (updateData.chunkCount !== undefined) dbUpdateData.chunkCount = updateData.chunkCount
|
|
if (updateData.tokenCount !== undefined) dbUpdateData.tokenCount = updateData.tokenCount
|
|
if (updateData.characterCount !== undefined)
|
|
dbUpdateData.characterCount = updateData.characterCount
|
|
if (updateData.processingStatus !== undefined)
|
|
dbUpdateData.processingStatus = updateData.processingStatus
|
|
if (updateData.processingError !== undefined)
|
|
dbUpdateData.processingError = updateData.processingError
|
|
|
|
const convertTagValue = (
|
|
slot: string,
|
|
value: string | undefined
|
|
): string | number | Date | boolean | null => {
|
|
if (value === undefined || value === '') return null
|
|
|
|
if (slot.startsWith('number')) {
|
|
return parseNumberValue(value)
|
|
}
|
|
|
|
if (slot.startsWith('date')) {
|
|
return parseDateValue(value)
|
|
}
|
|
|
|
if (slot.startsWith('boolean')) {
|
|
return parseBooleanValue(value) ?? false
|
|
}
|
|
|
|
return value || null
|
|
}
|
|
|
|
type UpdateDataWithTags = typeof updateData & Record<TagSlot, string | undefined>
|
|
const typedUpdateData = updateData as UpdateDataWithTags
|
|
|
|
ALL_TAG_SLOTS.forEach((slot: TagSlot) => {
|
|
const updateValue = typedUpdateData[slot]
|
|
if (updateValue !== undefined) {
|
|
;(dbUpdateData as Record<TagSlot, string | number | Date | boolean | null>)[slot] =
|
|
convertTagValue(slot, updateValue)
|
|
}
|
|
})
|
|
|
|
const doc = await db.transaction(async (tx) => {
|
|
const hasTagUpdates = ALL_TAG_SLOTS.some((field) => typedUpdateData[field] !== undefined)
|
|
|
|
if (hasTagUpdates) {
|
|
const embeddingUpdateData: Partial<ProcessedDocumentTags> = {}
|
|
ALL_TAG_SLOTS.forEach((field) => {
|
|
if (typedUpdateData[field] !== undefined) {
|
|
;(embeddingUpdateData as Record<TagSlot, string | number | Date | boolean | null>)[
|
|
field
|
|
] = convertTagValue(field, typedUpdateData[field])
|
|
}
|
|
})
|
|
|
|
await tx
|
|
.update(embedding)
|
|
.set(embeddingUpdateData)
|
|
.where(eq(embedding.documentId, documentId))
|
|
}
|
|
|
|
const [current] = await tx
|
|
.select()
|
|
.from(document)
|
|
.where(eq(document.id, documentId))
|
|
.limit(1)
|
|
.for('update')
|
|
if (!current) return undefined
|
|
|
|
const [sidecar] =
|
|
current.secretProvenanceVersion === 1
|
|
? await tx
|
|
.select()
|
|
.from(documentSecretProvenance)
|
|
.where(eq(documentSecretProvenance.documentId, documentId))
|
|
.limit(1)
|
|
: []
|
|
const currentSource = createKnowledgeDocumentSourceValue(current)
|
|
const currentProvenance = readBoundKnowledgeDocumentSecretProvenance({
|
|
secretProvenanceVersion: current.secretProvenanceVersion,
|
|
source: currentSource,
|
|
provenanceSourceHash: sidecar?.sourceHash ?? null,
|
|
status: sidecar?.status ?? null,
|
|
entries: sidecar?.entries,
|
|
})
|
|
|
|
const [updated] = await tx
|
|
.update(document)
|
|
.set(dbUpdateData)
|
|
.where(eq(document.id, documentId))
|
|
.returning()
|
|
if (updated && current.secretProvenanceVersion === 1) {
|
|
const updatedSource = createKnowledgeDocumentSourceValue(updated)
|
|
await replaceKnowledgeDocumentSecretProvenanceInTx(
|
|
tx,
|
|
documentId,
|
|
updatedSource,
|
|
rebindKnowledgeDocumentSecretProvenance(currentProvenance, currentSource, updatedSource)
|
|
)
|
|
}
|
|
return updated
|
|
})
|
|
|
|
if (!doc) {
|
|
throw new Error(`Document ${documentId} not found`)
|
|
}
|
|
|
|
logger.info(`[${requestId}] Document updated: ${documentId}`)
|
|
|
|
return {
|
|
id: doc.id,
|
|
knowledgeBaseId: doc.knowledgeBaseId,
|
|
filename: doc.filename,
|
|
fileUrl: doc.fileUrl,
|
|
fileSize: doc.fileSize,
|
|
mimeType: doc.mimeType,
|
|
chunkCount: doc.chunkCount,
|
|
tokenCount: doc.tokenCount,
|
|
characterCount: doc.characterCount,
|
|
processingStatus: doc.processingStatus as 'pending' | 'processing' | 'completed' | 'failed',
|
|
processingStartedAt: doc.processingStartedAt,
|
|
processingCompletedAt: doc.processingCompletedAt,
|
|
processingError: doc.processingError,
|
|
enabled: doc.enabled,
|
|
uploadedAt: doc.uploadedAt,
|
|
tag1: doc.tag1,
|
|
tag2: doc.tag2,
|
|
tag3: doc.tag3,
|
|
tag4: doc.tag4,
|
|
tag5: doc.tag5,
|
|
tag6: doc.tag6,
|
|
tag7: doc.tag7,
|
|
number1: doc.number1,
|
|
number2: doc.number2,
|
|
number3: doc.number3,
|
|
number4: doc.number4,
|
|
number5: doc.number5,
|
|
date1: doc.date1,
|
|
date2: doc.date2,
|
|
boolean1: doc.boolean1,
|
|
boolean2: doc.boolean2,
|
|
boolean3: doc.boolean3,
|
|
deletedAt: doc.deletedAt,
|
|
}
|
|
}
|
|
|
|
function getKnowledgeBaseStorageKey(fileUrl: string | null): string | null {
|
|
if (!fileUrl) {
|
|
return null
|
|
}
|
|
|
|
try {
|
|
const urlPath = new URL(fileUrl, 'http://localhost').pathname
|
|
const storageKey = extractStorageKey(urlPath)
|
|
return storageKey !== urlPath ? storageKey : null
|
|
} catch {
|
|
return null
|
|
}
|
|
}
|
|
|
|
/** Each entry deletes a storage object plus its metadata row. */
|
|
const STORAGE_DELETE_CONCURRENCY = 10
|
|
|
|
export async function deleteDocumentStorageFiles(
|
|
documentsToDelete: Array<{ id: string; fileUrl: string | null; workspaceId?: string | null }>,
|
|
requestId: string
|
|
): Promise<void> {
|
|
const entries = documentsToDelete.map((doc) => ({
|
|
doc,
|
|
storageKey: getKnowledgeBaseStorageKey(doc.fileUrl),
|
|
}))
|
|
|
|
const storageKeys = [
|
|
...new Set(
|
|
entries
|
|
.map((entry) => entry.storageKey)
|
|
.filter(
|
|
(key): key is string => typeof key === 'string' && isKnowledgeBaseOwnedStorageKey(key)
|
|
)
|
|
),
|
|
]
|
|
const bindingByKey = new Map<string, FileMetadataRecord>()
|
|
if (storageKeys.length > 0) {
|
|
const bindings = await getFileMetadataByKeys(storageKeys, 'knowledge-base')
|
|
for (const binding of bindings) {
|
|
bindingByKey.set(binding.key, binding)
|
|
}
|
|
}
|
|
|
|
await mapWithConcurrency(entries, STORAGE_DELETE_CONCURRENCY, async ({ doc, storageKey }) => {
|
|
if (!storageKey) {
|
|
return
|
|
}
|
|
|
|
if (!isKnowledgeBaseOwnedStorageKey(storageKey)) {
|
|
return
|
|
}
|
|
|
|
const binding = bindingByKey.get(storageKey)
|
|
if (!binding?.workspaceId || binding.context !== 'knowledge-base') {
|
|
logger.warn(`[${requestId}] Skipping storage delete: no ownership binding for key`, {
|
|
documentId: doc.id,
|
|
storageKey,
|
|
})
|
|
return
|
|
}
|
|
if (!doc.workspaceId || binding.workspaceId !== doc.workspaceId) {
|
|
logger.warn(`[${requestId}] Skipping storage delete: ownership binding mismatch`, {
|
|
documentId: doc.id,
|
|
storageKey,
|
|
bindingWorkspaceId: binding.workspaceId,
|
|
documentWorkspaceId: doc.workspaceId ?? null,
|
|
})
|
|
return
|
|
}
|
|
|
|
try {
|
|
const metadataDeleted = await deleteFileMetadataByIdentity({
|
|
id: binding.id,
|
|
key: binding.key,
|
|
context: binding.context,
|
|
contentUpdatedAt: binding.contentUpdatedAt,
|
|
})
|
|
if (!metadataDeleted) {
|
|
logger.warn(`[${requestId}] Skipping storage delete: ownership binding changed`, {
|
|
documentId: doc.id,
|
|
storageKey,
|
|
})
|
|
return
|
|
}
|
|
await deleteFile({ key: storageKey, context: 'knowledge-base' })
|
|
} catch (error) {
|
|
logger.warn(`[${requestId}] Failed to delete document storage file`, {
|
|
documentId: doc.id,
|
|
error: toError(error).message,
|
|
})
|
|
}
|
|
})
|
|
}
|
|
|
|
async function excludeConnectorDocuments(
|
|
documentIds: string[],
|
|
requestId: string
|
|
): Promise<number> {
|
|
const ids = [...new Set(documentIds)]
|
|
if (ids.length === 0) {
|
|
return 0
|
|
}
|
|
|
|
const updated = await db
|
|
.update(document)
|
|
.set({
|
|
userExcluded: true,
|
|
enabled: false,
|
|
})
|
|
.where(and(inArray(document.id, ids), isNotNull(document.connectorId)))
|
|
.returning({ id: document.id })
|
|
|
|
if (updated.length > 0) {
|
|
logger.info(`[${requestId}] Excluded ${updated.length} connector-backed document(s)`, {
|
|
documentIds: updated.map((doc) => doc.id),
|
|
})
|
|
}
|
|
|
|
return updated.length
|
|
}
|
|
|
|
async function deleteDocumentsByLifecyclePolicy(
|
|
documentIds: string[],
|
|
requestId: string
|
|
): Promise<number> {
|
|
const ids = [...new Set(documentIds)]
|
|
if (ids.length === 0) {
|
|
return 0
|
|
}
|
|
|
|
const docs = await db
|
|
.select({
|
|
id: document.id,
|
|
connectorId: document.connectorId,
|
|
})
|
|
.from(document)
|
|
.where(inArray(document.id, ids))
|
|
|
|
const connectorBackedIds = docs.filter((doc) => doc.connectorId !== null).map((doc) => doc.id)
|
|
const hardDeleteIds = docs.filter((doc) => doc.connectorId === null).map((doc) => doc.id)
|
|
|
|
const [excludedCount, hardDeletedCount] = await Promise.all([
|
|
excludeConnectorDocuments(connectorBackedIds, requestId),
|
|
hardDeleteDocuments(hardDeleteIds, requestId),
|
|
])
|
|
|
|
return excludedCount + hardDeletedCount
|
|
}
|
|
|
|
export async function hardDeleteDocuments(
|
|
documentIds: string[],
|
|
requestId: string,
|
|
/**
|
|
* When provided, re-verifies each document's connectorId still matches at
|
|
* the moment of the actual delete query — not just the caller's earlier
|
|
* snapshot. A caller (e.g. connector sync reconciliation) can compute this
|
|
* ID list well before the delete runs; a concurrent request that detaches
|
|
* these same documents from the connector in between (e.g. "delete
|
|
* connector, keep documents") would otherwise still have them purged here
|
|
* despite no longer belonging to the connector the caller reasoned about.
|
|
*/
|
|
expectedConnectorId?: string
|
|
): Promise<number> {
|
|
const ids = [...new Set(documentIds)]
|
|
if (ids.length === 0) {
|
|
return 0
|
|
}
|
|
|
|
let deletedCount = 0
|
|
for (let offset = 0; offset < ids.length; offset += HARD_DELETE_DOCUMENT_BATCH_SIZE) {
|
|
deletedCount += await hardDeleteDocumentBatch(
|
|
ids.slice(offset, offset + HARD_DELETE_DOCUMENT_BATCH_SIZE),
|
|
requestId,
|
|
expectedConnectorId
|
|
)
|
|
}
|
|
return deletedCount
|
|
}
|
|
|
|
/**
|
|
* Hard-deletes one bounded metadata batch and applies every associated ledger
|
|
* delta atomically.
|
|
*/
|
|
async function hardDeleteDocumentBatch(
|
|
documentIds: string[],
|
|
requestId: string,
|
|
expectedConnectorId?: string
|
|
): Promise<number> {
|
|
const ids = [...new Set(documentIds)]
|
|
const documentsToDelete = await db
|
|
.select({
|
|
id: document.id,
|
|
knowledgeBaseId: document.knowledgeBaseId,
|
|
fileUrl: document.fileUrl,
|
|
fileSize: document.fileSize,
|
|
uploadedBy: document.uploadedBy,
|
|
connectorId: document.connectorId,
|
|
workspaceId: knowledgeBase.workspaceId,
|
|
kbUserId: knowledgeBase.userId,
|
|
})
|
|
.from(document)
|
|
.innerJoin(knowledgeBase, eq(document.knowledgeBaseId, knowledgeBase.id))
|
|
.where(
|
|
expectedConnectorId
|
|
? and(inArray(document.id, ids), eq(document.connectorId, expectedConnectorId))
|
|
: inArray(document.id, ids)
|
|
)
|
|
|
|
if (documentsToDelete.length === 0) {
|
|
return 0
|
|
}
|
|
|
|
const existingIds = documentsToDelete.map((doc) => doc.id)
|
|
|
|
/**
|
|
* Resolve immutable workspace payers and legacy account subscriptions before
|
|
* opening the deletion transaction. Connector documents were never metered.
|
|
*/
|
|
const storageContextByWorkspace = new Map<string, StorageBillingContext>()
|
|
const candidateUserIds = new Set<string>()
|
|
for (const doc of documentsToDelete) {
|
|
if (doc.connectorId || doc.fileSize <= 0) continue
|
|
if (doc.workspaceId) {
|
|
if (!storageContextByWorkspace.has(doc.workspaceId)) {
|
|
storageContextByWorkspace.set(
|
|
doc.workspaceId,
|
|
await resolveStorageBillingContext(doc.workspaceId)
|
|
)
|
|
}
|
|
continue
|
|
}
|
|
const billedUserId = doc.uploadedBy ?? doc.kbUserId
|
|
if (billedUserId) candidateUserIds.add(billedUserId)
|
|
}
|
|
const subByUser = new Map<string, HighestPrioritySubscription | null>()
|
|
for (const billedUserId of candidateUserIds) {
|
|
subByUser.set(billedUserId, await getHighestPrioritySubscription(billedUserId))
|
|
}
|
|
|
|
/**
|
|
* Key every decrement off rows this transaction actually deleted so
|
|
* concurrent deletion cannot double-decrement a payer.
|
|
*/
|
|
let deletedDocs: typeof documentsToDelete = []
|
|
await db.transaction(async (tx) => {
|
|
/**
|
|
* Lock every parent KB in stable ID order before deleting document rows.
|
|
* Normal inserts and KB moves take the same parent lock first, so the
|
|
* workspace snapshots used for accounting cannot change mid-delete.
|
|
*/
|
|
const knowledgeBaseIds = [
|
|
...new Set(documentsToDelete.map((doc) => doc.knowledgeBaseId)),
|
|
].sort()
|
|
const lockedKnowledgeBases = await tx
|
|
.select({
|
|
id: knowledgeBase.id,
|
|
workspaceId: knowledgeBase.workspaceId,
|
|
userId: knowledgeBase.userId,
|
|
})
|
|
.from(knowledgeBase)
|
|
.where(inArray(knowledgeBase.id, knowledgeBaseIds))
|
|
.orderBy(asc(knowledgeBase.id))
|
|
.for('update')
|
|
const lockedKnowledgeBaseById = new Map(lockedKnowledgeBases.map((kb) => [kb.id, kb]))
|
|
for (const doc of documentsToDelete) {
|
|
const lockedKb = lockedKnowledgeBaseById.get(doc.knowledgeBaseId)
|
|
if (
|
|
!lockedKb ||
|
|
lockedKb.workspaceId !== doc.workspaceId ||
|
|
lockedKb.userId !== doc.kbUserId
|
|
) {
|
|
throw new Error(
|
|
`Knowledge base ${doc.knowledgeBaseId} storage ownership changed; retry document deletion`
|
|
)
|
|
}
|
|
}
|
|
|
|
/**
|
|
* Re-verify `expectedConnectorId` here too, not only on the pre-transaction
|
|
* SELECT above — the billing lookups and KB locking between that SELECT
|
|
* and this delete are async and can span a concurrent "delete connector,
|
|
* keep documents" request that clears these rows' `connectorId` in
|
|
* between. Deleting a detached document's embeddings would corrupt its
|
|
* search index even if the document row itself were spared, so both the
|
|
* embedding delete and the document delete are scoped to this re-verified
|
|
* ID set rather than the stale `existingIds`.
|
|
*/
|
|
const stillTargetedIds = expectedConnectorId
|
|
? (
|
|
await tx
|
|
.select({ id: document.id })
|
|
.from(document)
|
|
.where(
|
|
and(inArray(document.id, existingIds), eq(document.connectorId, expectedConnectorId))
|
|
)
|
|
).map((d) => d.id)
|
|
: existingIds
|
|
|
|
await tx.delete(embedding).where(inArray(embedding.documentId, stillTargetedIds))
|
|
const deletedRows = await tx
|
|
.delete(document)
|
|
.where(inArray(document.id, stillTargetedIds))
|
|
.returning({ id: document.id })
|
|
|
|
const deletedIds = new Set(deletedRows.map((row) => row.id))
|
|
deletedDocs = documentsToDelete.filter((doc) => deletedIds.has(doc.id))
|
|
|
|
const bytesByWorkspace = new Map<string, number>()
|
|
const legacyBytesByUser = new Map<string, number>()
|
|
for (const doc of deletedDocs) {
|
|
if (doc.connectorId || doc.fileSize <= 0) continue
|
|
if (doc.workspaceId) {
|
|
bytesByWorkspace.set(
|
|
doc.workspaceId,
|
|
(bytesByWorkspace.get(doc.workspaceId) ?? 0) + doc.fileSize
|
|
)
|
|
continue
|
|
}
|
|
const billedUserId = doc.uploadedBy ?? doc.kbUserId
|
|
if (!billedUserId) continue
|
|
legacyBytesByUser.set(billedUserId, (legacyBytesByUser.get(billedUserId) ?? 0) + doc.fileSize)
|
|
}
|
|
await applyStorageUsageDeltasInTx(tx, {
|
|
workspaceDeltas: [...bytesByWorkspace.entries()]
|
|
.sort(([left], [right]) => left.localeCompare(right))
|
|
.map(([workspaceId, bytes]) => {
|
|
const context = storageContextByWorkspace.get(workspaceId)
|
|
if (!context) {
|
|
throw new Error(`Missing storage billing context for workspace ${workspaceId}`)
|
|
}
|
|
return { context, deltaBytes: -bytes }
|
|
}),
|
|
legacyDeltas: [...legacyBytesByUser.entries()]
|
|
.sort(([left], [right]) => left.localeCompare(right))
|
|
.map(([userId, bytes]) => ({
|
|
userId,
|
|
subscription: subByUser.get(userId) ?? null,
|
|
deltaBytes: -bytes,
|
|
})),
|
|
})
|
|
})
|
|
|
|
await deleteDocumentStorageFiles(deletedDocs, requestId)
|
|
|
|
logger.info(`[${requestId}] Hard deleted ${deletedDocs.length} documents`, {
|
|
documentIds: deletedDocs.map((doc) => doc.id),
|
|
})
|
|
|
|
return deletedDocs.length
|
|
}
|
|
|
|
export async function deleteDocument(
|
|
documentId: string,
|
|
requestId: string
|
|
): Promise<{ success: boolean; message: string }> {
|
|
await deleteDocumentsByLifecyclePolicy([documentId], requestId)
|
|
|
|
return {
|
|
success: true,
|
|
message: 'Document deleted successfully',
|
|
}
|
|
}
|