281 lines
11 KiB
TypeScript
281 lines
11 KiB
TypeScript
// The index job's body: what the worker runs for one document. Read the file's bytes, hand them to
|
|
// the search service, and record the outcome on the durable row.
|
|
// Transient failures (the service hiccups) throw, and the queue retries the job; the row the
|
|
// retry resumes from makes the retry safe. Permanent failures are recorded on the row with a
|
|
// stable code, so a person sees why a document failed and a retried request does not pay twice.
|
|
// The automatic PDF intake (core-11) runs the same lifecycle with another send: the PDF is read into
|
|
// passages and each passage is indexed on its own (`passageSend`).
|
|
|
|
import { internal } from "@conusai-kit/core"
|
|
import { downloadFile, getFile, readAll, sha256Hex } from "@conusai-kit/files"
|
|
import {
|
|
type DocumentView,
|
|
documents,
|
|
INDEX_FAILURES,
|
|
type IndexFailure,
|
|
MAX_AUTO_INDEX_BYTES,
|
|
MAX_INDEX_BYTES,
|
|
passageSets,
|
|
} from "../contract/index.js"
|
|
import { documentById, documentView } from "./documents.js"
|
|
import { type KnowledgeDeps, type KnowledgePassage, type KnowledgeScope, knowledgeEnabled } from "./ports.js"
|
|
|
|
type DocumentRow = Awaited<ReturnType<typeof documentById>>
|
|
|
|
/** What one run sends between the row's "indexing" and its outcome; it throws to fail or to retry. */
|
|
type Send = (row: DocumentRow, signal: AbortSignal | undefined) => Promise<void>
|
|
|
|
export async function indexDocument(
|
|
deps: KnowledgeDeps,
|
|
scope: KnowledgeScope,
|
|
documentId: string,
|
|
options: { signal?: AbortSignal; finalAttempt?: boolean } = {},
|
|
): Promise<DocumentView> {
|
|
return runIndex(deps, scope, documentId, () => textSend(deps, scope), options)
|
|
}
|
|
|
|
/** The automatic PDF intake's run of one document: the PDF read into passages, each indexed on its own. */
|
|
export async function indexPassages(
|
|
deps: KnowledgeDeps,
|
|
scope: KnowledgeScope,
|
|
documentId: string,
|
|
options: { signal?: AbortSignal; finalAttempt?: boolean } = {},
|
|
): Promise<DocumentView> {
|
|
return runIndex(deps, scope, documentId, () => passageSend(deps, scope), options)
|
|
}
|
|
|
|
/** The whole file as text, to the search service. */
|
|
function textSend(deps: KnowledgeDeps, scope: KnowledgeScope): Send {
|
|
const service = deps.search
|
|
if (!service) throw internal("the knowledge search service is not bound in this process")
|
|
return async (row, signal) => {
|
|
const bytes = await readFileBytes(deps, scope, row.fileId, MAX_INDEX_BYTES)
|
|
const text = decodeText(bytes)
|
|
await deps.beforeIndex?.(row.requestedBy)
|
|
signal?.throwIfAborted()
|
|
await service.index(scope, {
|
|
collectionName: row.collectionName,
|
|
fileId: row.fileId,
|
|
filename: row.filename,
|
|
bytes: new TextEncoder().encode(text),
|
|
...(signal ? { signal } : {}),
|
|
})
|
|
}
|
|
}
|
|
|
|
/**
|
|
* A PDF, read by the platform's document reader into passages, each indexed as one item. What is about to
|
|
* be indexed is recorded before any of it is sent, so a run the service refuses partway can still be
|
|
* removed passage by passage. The read is keyed by the file and its digest: a retry reads the same job.
|
|
*/
|
|
function passageSend(deps: KnowledgeDeps, scope: KnowledgeScope): Send {
|
|
const { reader, passages, enablement } = deps
|
|
if (!reader || !passages) throw internal("the document reader and the passage index are not bound in this process")
|
|
if (!enablement) throw internal("the enablement platform is not bound in this process")
|
|
return async (row, signal) => {
|
|
const file = await getFile(deps.files, scope, row.fileId)
|
|
if (file.size > MAX_AUTO_INDEX_BYTES)
|
|
throw Object.assign(new Error("the file is too large to read"), { permanent: "too_large" })
|
|
// The organization's eligibility and quota, as the enablement flow checks them, before the read is paid
|
|
// for: a rollout kill switch or a spent quota stops the next PDF (root's core-11 decision). The service
|
|
// still refuses past its quota at call time; this check only stops earlier.
|
|
const eligibility = await enablement.checkEligibility(scope)
|
|
if (!eligibility.ok)
|
|
throw Object.assign(new Error(eligibility.detail ?? eligibility.reason), { permanent: eligibility.reason })
|
|
if ((await enablement.indexingQuota(scope)) <= 0)
|
|
throw Object.assign(new Error("the indexing volume for today is spent"), { permanent: "quota_exhausted" })
|
|
signal?.throwIfAborted()
|
|
const bytes = await readFileBytes(deps, scope, row.fileId, MAX_AUTO_INDEX_BYTES)
|
|
await deps.beforeIndex?.(row.requestedBy)
|
|
signal?.throwIfAborted()
|
|
const read = await reader.read(bytes, {
|
|
filename: row.filename,
|
|
idempotencyKey: `knowledge.autoindex.${row.fileId}.${file.sha256}`,
|
|
...(signal ? { signal } : {}),
|
|
})
|
|
const documentRef = await passages.reference(scope, row.fileId)
|
|
await recordPassages(deps, scope, row, documentRef, read)
|
|
await deps.beforeIndex?.(row.requestedBy)
|
|
signal?.throwIfAborted()
|
|
await passages.index(scope, {
|
|
collectionName: row.collectionName,
|
|
documentRef,
|
|
fileId: row.fileId,
|
|
passages: read,
|
|
...(signal ? { signal } : {}),
|
|
})
|
|
}
|
|
}
|
|
|
|
/** The passages per page, comma-separated from page 1: what `passagesOf` turns back into references. */
|
|
const pagesOf = (read: readonly KnowledgePassage[]): string => {
|
|
const counts: number[] = []
|
|
for (const passage of read) {
|
|
while (counts.length < passage.page) counts.push(0)
|
|
counts[passage.page - 1] = (counts[passage.page - 1] ?? 0) + 1
|
|
}
|
|
return counts.join(",")
|
|
}
|
|
|
|
const passagesOf = (pages: string): Array<{ page: number; index: number }> =>
|
|
pages === ""
|
|
? []
|
|
: pages
|
|
.split(",")
|
|
.flatMap((count, at) => Array.from({ length: Number(count) }, (_, index) => ({ page: at + 1, index })))
|
|
|
|
async function recordPassages(
|
|
deps: KnowledgeDeps,
|
|
scope: KnowledgeScope,
|
|
row: DocumentRow,
|
|
documentRef: string,
|
|
read: readonly KnowledgePassage[],
|
|
): Promise<void> {
|
|
const fields = { documentRef, pages: pagesOf(read) }
|
|
await deps.data.unitOfWork.run(scope, async (tx) => {
|
|
const existing = (await tx.of(passageSets).list(passageSetOf(row))).items[0]
|
|
if (existing) await tx.of(passageSets).update(existing.id, fields, existing.version)
|
|
else await tx.of(passageSets).insert({ collectionName: row.collectionName, fileId: row.fileId, ...fields })
|
|
})
|
|
}
|
|
|
|
/** The query of a document's passage record: one per (collection, file), as the document row. */
|
|
const passageSetOf = (row: { collectionName: string; fileId: string }) => ({
|
|
where: [
|
|
{ field: "collectionName" as const, op: "eq" as const, value: row.collectionName },
|
|
{ field: "fileId" as const, op: "eq" as const, value: row.fileId },
|
|
],
|
|
limit: 1,
|
|
})
|
|
|
|
async function runIndex(
|
|
deps: KnowledgeDeps,
|
|
scope: KnowledgeScope,
|
|
documentId: string,
|
|
sendOf: () => Send,
|
|
options: { signal?: AbortSignal; finalAttempt?: boolean },
|
|
): Promise<DocumentView> {
|
|
const row = await documentById(deps, scope, documentId)
|
|
if (row.status === "indexed" || row.status === "cancelled" || row.status === "failed") {
|
|
// A redelivered job or a crashed retry: the row says where the work stands.
|
|
return {
|
|
documentId: row.id,
|
|
collectionName: row.collectionName,
|
|
fileId: row.fileId,
|
|
filename: row.filename,
|
|
status: row.status,
|
|
error: row.error,
|
|
indexedAt: row.indexedAt,
|
|
requestedBy: row.requestedBy,
|
|
}
|
|
}
|
|
const send = sendOf()
|
|
const indexing = await deps.data
|
|
.forScope(scope)
|
|
.of(documents)
|
|
.update(row.id, { status: "indexing", error: null }, row.version)
|
|
|
|
try {
|
|
options.signal?.throwIfAborted()
|
|
if (!(await knowledgeEnabled(deps, scope)))
|
|
throw Object.assign(new Error("processing disabled"), { permanent: "feature_disabled" })
|
|
await deps.beforeIndex?.(row.requestedBy)
|
|
options.signal?.throwIfAborted()
|
|
await send(indexing, options.signal)
|
|
options.signal?.throwIfAborted()
|
|
const current = await deps.data.forScope(scope).of(documents).get(indexing.id)
|
|
if (!current || current.status === "cancelled") return documentView({ ...indexing, status: "cancelled" })
|
|
const done = await deps.data
|
|
.forScope(scope)
|
|
.of(documents)
|
|
.update(indexing.id, { status: "indexed", indexedAt: new Date().toISOString(), error: null }, indexing.version)
|
|
return {
|
|
documentId: done.id,
|
|
collectionName: done.collectionName,
|
|
fileId: done.fileId,
|
|
filename: done.filename,
|
|
status: done.status as DocumentView["status"],
|
|
error: done.error,
|
|
indexedAt: done.indexedAt,
|
|
requestedBy: done.requestedBy,
|
|
}
|
|
} catch (error) {
|
|
if (options.signal?.aborted) throw error
|
|
const permanent = permanentFailure(error) ?? (options.finalAttempt === true ? "processing_exhausted" : null)
|
|
if (permanent === null) throw error
|
|
const current = await deps.data.forScope(scope).of(documents).get(indexing.id)
|
|
if (!current || current.status === "cancelled") return documentView({ ...indexing, status: "cancelled" })
|
|
const failed = await deps.data
|
|
.forScope(scope)
|
|
.of(documents)
|
|
.update(indexing.id, { status: "failed", error: permanent }, indexing.version)
|
|
return {
|
|
documentId: failed.id,
|
|
collectionName: failed.collectionName,
|
|
fileId: failed.fileId,
|
|
filename: failed.filename,
|
|
status: "failed",
|
|
error: failed.error,
|
|
indexedAt: failed.indexedAt,
|
|
requestedBy: failed.requestedBy,
|
|
}
|
|
}
|
|
}
|
|
|
|
export async function removeDocument(
|
|
deps: KnowledgeDeps,
|
|
scope: KnowledgeScope,
|
|
documentId: string,
|
|
storageKey?: string,
|
|
): Promise<void> {
|
|
const row = await deps.data.forScope(scope).of(documents).get(documentId)
|
|
if (!row) return
|
|
const set = (await deps.data.forScope(scope).of(passageSets).list(passageSetOf(row))).items[0]
|
|
if (set) {
|
|
if (!deps.passages) throw internal("the passage index is not bound in this process")
|
|
await deps.passages.remove(scope, { documentRef: set.documentRef, passages: passagesOf(set.pages) })
|
|
await deps.data.unitOfWork.run(scope, async (tx) => {
|
|
await tx.of(documents).delete(row.id)
|
|
await tx.of(passageSets).delete(set.id)
|
|
})
|
|
return
|
|
}
|
|
const service = deps.search
|
|
if (!service) throw internal("the knowledge search service is not bound in this process")
|
|
await service.remove(scope, {
|
|
collectionName: row.collectionName,
|
|
fileId: row.fileId,
|
|
...(storageKey === undefined ? {} : { storageKey }),
|
|
})
|
|
await deps.data.forScope(scope).of(documents).delete(row.id)
|
|
}
|
|
|
|
async function readFileBytes(
|
|
deps: KnowledgeDeps,
|
|
scope: KnowledgeScope,
|
|
fileId: string,
|
|
maxBytes: number,
|
|
): Promise<Uint8Array> {
|
|
const file = await getFile(deps.files, scope, fileId)
|
|
const response = await downloadFile(deps.files, scope, fileId)
|
|
if (!response.body) throw Object.assign(new Error("source missing"), { permanent: "not_indexable" })
|
|
const bytes = await readAll(response.body, maxBytes)
|
|
if (!file.sha256 || (await sha256Hex(bytes)) !== file.sha256)
|
|
throw Object.assign(new Error("source changed"), { permanent: "file_changed" })
|
|
return bytes
|
|
}
|
|
|
|
function decodeText(bytes: Uint8Array): string {
|
|
try {
|
|
return new TextDecoder("utf-8", { fatal: true }).decode(bytes)
|
|
} catch {
|
|
throw Object.assign(new Error("the file is not valid UTF-8 text"), { permanent: "not_text" })
|
|
}
|
|
}
|
|
|
|
/** Maps an indexing failure to the stable code recorded on the row; null when it is transient. */
|
|
function permanentFailure(error: unknown): string | null {
|
|
const code = (error as { permanent?: unknown }).permanent
|
|
return (INDEX_FAILURES as readonly unknown[]).includes(code) ? (code as IndexFailure) : null
|
|
}
|