Files
RAG/backend/app/persistence/document_workflows.py
YoVinchen ecdb10c37a
Some checks failed
verify / verify (push) Has been cancelled
Make the governed RAG evidence path executable end to end
Separate local parsing from model indexing, bind review decisions to immutable manifests, persist vectors behind active profiles, and expose retrieval, chat, evaluation, and document workflows through the React workbench.

Constraint: Live Bailian authentication currently fails for all three configured capabilities

Rejected: Direct upload-to-embedding flow | bypasses local review and manifest binding

Confidence: high

Scope-risk: broad

Directive: Keep private-data deployment blocked until authentication, RBAC, and separate database roles land

Tested: make verify; fresh and replay Docker document smoke; worker recovery smoke; frozen synthetic evaluation; migration 0003-0004 roundtrip

Not-tested: Successful live Bailian calls, OCR, real multi-user authorization
2026-07-13 05:58:11 +08:00

1122 lines
36 KiB
Python

"""Lease-fenced persistence for the PARSE_DOCUMENT workflow."""
from __future__ import annotations
import hashlib
import json
import re
import uuid
from collections.abc import Callable, Mapping
from dataclasses import dataclass
from typing import Any, Protocol, cast
import psycopg
from psycopg import Connection
from psycopg.rows import dict_row
from app.core.config import Settings
from app.core.secrets import SecretFileError
from app.persistence.job_queue import BackgroundJob, JobLease, LeaseLostError
from app.services.document_ingestion import (
CloudTextPolicy,
IngestionArtifact,
IngestionStatus,
)
type WorkflowRow = dict[str, Any]
type ConnectionFactory = Callable[[str, int], Connection[WorkflowRow]]
_DB_ID_NAMESPACE = uuid.UUID("07fb2c87-b4c9-5dd8-b965-fb5d8e53e350")
_HASH = re.compile(r"^[0-9a-f]{64}$")
_ERROR_CODE = re.compile(r"^[A-Z][A-Z0-9_]{0,127}$")
_NORMALIZATION_PROFILE_HASH = hashlib.sha256(b"document-ingestion-normalized-text-v1").hexdigest()
@dataclass(frozen=True, slots=True)
class DocumentSource:
upload_id: uuid.UUID
document_id: uuid.UUID
knowledge_base_id: uuid.UUID
access_scope_id: uuid.UUID
filename: str
mime_type: str
storage_key: uuid.UUID
byte_size: int
raw_sha256: str
@dataclass(frozen=True, slots=True)
class PlannedPage:
id: uuid.UUID
ordinal: int
page_number: int | None
text: str
text_sha256: str
line_start: int
line_end: int
@dataclass(frozen=True, slots=True)
class PlannedBlock:
id: uuid.UUID
page_id: uuid.UUID
ordinal: int
kind: str
text: str
text_sha256: str
section_path: tuple[str, ...]
anchor_id: str
normalized_text_sha256: str
char_start: int
char_end: int
line_start: int
line_end: int
page_start: int | None
page_end: int | None
@dataclass(frozen=True, slots=True)
class PlannedChunk:
id: uuid.UUID
ordinal: int
display_text: str
cloud_text: str
cloud_text_sha256: str
embedding_prefix: str
embedding_text: str
embedding_text_sha256: str
token_count: int
page_start: int | None
page_end: int | None
section_path: tuple[str, ...]
metadata: Mapping[str, object]
@dataclass(frozen=True, slots=True)
class ArtifactPlan:
version_id: uuid.UUID
parser_profile_hash: str
normalization_profile_hash: str
chunk_profile_hash: str
cloud_policy_id: str
outbound_manifest_sha256: str | None
review_state: str
version_status: str
version_error_code: str | None
expected_chunk_count: int
document_status: str
job_stage: str
pages: tuple[PlannedPage, ...]
blocks: tuple[PlannedBlock, ...]
chunks: tuple[PlannedChunk, ...]
class InvalidDocumentJobError(RuntimeError):
def __init__(self) -> None:
super().__init__("document job contract is invalid")
class DocumentWorkflowPersistenceError(RuntimeError):
def __init__(self) -> None:
super().__init__("document workflow persistence unavailable")
class ArtifactConflictError(RuntimeError):
def __init__(self) -> None:
super().__init__("stored ingestion artifact conflicts with deterministic output")
class DocumentWorkflowRepository(Protocol):
def load_source(self, job: BackgroundJob) -> DocumentSource: ...
def record_terminal_parse_failure(
self,
*,
lease: JobLease,
source: DocumentSource,
error_code: str,
) -> None: ...
def persist_artifact(
self,
*,
lease: JobLease,
source: DocumentSource,
artifact: IngestionArtifact,
cloud_policy: CloudTextPolicy,
embedding_model: str,
embedding_dimension: int,
) -> ArtifactPlan: ...
LOAD_SOURCE_SQL = """
SELECT
upload.id AS upload_id,
upload.document_id,
upload.knowledge_base_id,
upload.access_scope_id,
upload.original_filename,
upload.declared_mime_type,
upload.storage_key,
upload.actual_size,
upload.actual_sha256,
document.filename AS document_filename,
document.mime_type AS document_mime_type,
document.raw_sha256 AS document_raw_sha256,
document.storage_key AS document_storage_key
FROM rag.background_jobs AS job
JOIN rag.document_uploads AS upload
ON upload.parse_job_id = job.id
JOIN rag.documents AS document
ON document.id = upload.document_id
AND document.knowledge_base_id = upload.knowledge_base_id
AND document.access_scope_id = upload.access_scope_id
WHERE job.id = %s
AND job.status = 'RUNNING'
AND job.lease_owner = %s
AND job.lease_token = %s
AND job.lease_until >= now()
AND job.job_type = 'PARSE_DOCUMENT'
AND job.required_capability = 'document_parse'
AND job.resource_type = 'document'
AND job.resource_id = %s
AND upload.id = %s
AND upload.document_id = %s
AND upload.status = 'COMPLETED'
AND upload.actual_size = upload.expected_size
AND upload.actual_sha256 = upload.expected_sha256
AND upload.actual_sha256 = document.raw_sha256
AND upload.storage_key::text = document.storage_key
AND upload.original_filename = document.filename
AND upload.declared_mime_type = document.mime_type
AND document.deleted_at IS NULL
LIMIT 1
"""
FAIL_PARSE_SQL = """
WITH fenced AS (
SELECT job.id AS job_id, document.id AS document_id
FROM rag.background_jobs AS job
JOIN rag.document_uploads AS upload ON upload.parse_job_id = job.id
JOIN rag.documents AS document
ON document.id = upload.document_id
AND document.knowledge_base_id = upload.knowledge_base_id
AND document.access_scope_id = upload.access_scope_id
WHERE job.id = %s
AND job.status = 'RUNNING'
AND job.lease_owner = %s
AND job.lease_token = %s
AND job.lease_until >= now()
AND job.job_type = 'PARSE_DOCUMENT'
AND job.required_capability = 'document_parse'
AND job.resource_type = 'document'
AND job.resource_id = %s
AND upload.id = %s
AND upload.document_id = %s
AND upload.knowledge_base_id = %s
AND upload.access_scope_id = %s
AND upload.actual_size = %s
AND upload.actual_sha256 = %s
AND upload.storage_key = %s
AND upload.status = 'COMPLETED'
AND document.raw_sha256 = %s
AND document.storage_key = %s
AND document.deleted_at IS NULL
FOR UPDATE OF job, document
), updated_document AS (
UPDATE rag.documents AS document
SET status = 'FAILED', updated_at = now()
FROM fenced
WHERE document.id = fenced.document_id
RETURNING document.id
)
UPDATE rag.background_jobs AS job
SET stage = 'PARSE_REJECTED',
progress = 100,
last_error_code = %s,
last_error_message = 'Document parsing was rejected by a deterministic local policy.',
updated_at = now()
FROM fenced
JOIN updated_document ON updated_document.id = fenced.document_id
WHERE job.id = fenced.job_id
RETURNING job.id
"""
INSERT_VERSION_SQL = """
INSERT INTO rag.document_versions (
id,
document_id,
parser_profile_hash,
ocr_profile_hash,
normalization_profile_hash,
chunk_profile_hash,
cloud_policy_id,
outbound_manifest_sha256,
review_state,
embedding_profile_hash,
status,
expected_chunk_count,
error_code
) VALUES (%s, %s, %s, NULL, %s, %s, %s, %s, %s, NULL, %s, %s, %s)
ON CONFLICT DO NOTHING
"""
INSERT_PAGE_SQL = """
INSERT INTO rag.document_pages (
id, document_version_id, ordinal, page_number, display_text, text_sha256,
line_start, line_end
) VALUES (%s, %s, %s, %s, %s, %s, %s, %s)
ON CONFLICT DO NOTHING
"""
INSERT_BLOCK_SQL = """
INSERT INTO rag.document_blocks (
id, document_version_id, page_id, ordinal, block_kind, display_text,
text_sha256, section_path, anchor_id, normalized_text_sha256,
char_start, char_end, line_start, line_end, page_start, page_end
) VALUES (%s, %s, %s, %s, %s, %s, %s, %s::jsonb, %s, %s, %s, %s, %s, %s, %s, %s)
ON CONFLICT DO NOTHING
"""
INSERT_MANIFEST_ITEM_SQL = """
INSERT INTO rag.outbound_manifest_items (
document_version_id,
ordinal,
outbound_manifest_sha256,
cloud_text_sha256,
embedding_text_sha256
) VALUES (%s, %s, %s, %s, %s)
ON CONFLICT DO NOTHING
"""
INSERT_CHUNK_SQL = """
INSERT INTO rag.chunks (
id,
knowledge_base_id,
document_id,
document_version_id,
access_scope_id,
ordinal,
display_text,
cloud_text,
cloud_text_sha256,
embedding_prefix,
embedding_text,
embedding_text_sha256,
embedded_text_sha256,
embedding_profile_hash,
outbound_manifest_sha256,
token_count,
page_start,
page_end,
section_path,
metadata,
embedding_model,
embedding_dimension,
embedding,
approval_status,
index_status,
searchable
) VALUES (
%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s,
NULL, NULL, %s, %s, %s, %s, %s::jsonb, %s::jsonb, %s, %s, NULL,
'LOCAL_PARSED_PENDING_CLOUD_REVIEW', 'PENDING', false
)
ON CONFLICT DO NOTHING
"""
SELECT_VERSION_SQL = """
SELECT
id, document_id, parser_profile_hash, normalization_profile_hash,
chunk_profile_hash, cloud_policy_id, outbound_manifest_sha256,
expected_chunk_count
FROM rag.document_versions
WHERE id = %s AND document_id = %s
"""
SELECT_PAGES_SQL = """
SELECT id, ordinal, page_number, display_text, text_sha256, line_start, line_end
FROM rag.document_pages
WHERE document_version_id = %s
ORDER BY ordinal
"""
SELECT_BLOCKS_SQL = """
SELECT
id, page_id, ordinal, block_kind, display_text, text_sha256, section_path,
anchor_id, normalized_text_sha256, char_start, char_end, line_start, line_end,
page_start, page_end
FROM rag.document_blocks
WHERE document_version_id = %s
ORDER BY ordinal
"""
SELECT_MANIFEST_ITEMS_SQL = """
SELECT ordinal, outbound_manifest_sha256, cloud_text_sha256, embedding_text_sha256
FROM rag.outbound_manifest_items
WHERE document_version_id = %s
ORDER BY ordinal
"""
SELECT_CHUNKS_SQL = """
SELECT
id, ordinal, display_text, cloud_text, cloud_text_sha256, embedding_prefix,
embedding_text, embedding_text_sha256, token_count, page_start, page_end,
section_path, metadata
FROM rag.chunks
WHERE document_version_id = %s AND deleted_at IS NULL
ORDER BY ordinal
"""
UPDATE_DOCUMENT_SQL = """
UPDATE rag.documents
SET status = CASE WHEN active_version_id IS NULL THEN %s ELSE status END,
updated_at = now()
WHERE id = %s
AND knowledge_base_id = %s
AND access_scope_id = %s
AND raw_sha256 = %s
AND storage_key = %s
AND deleted_at IS NULL
RETURNING id
"""
FINALIZE_JOB_SQL = """
UPDATE rag.background_jobs AS job
SET stage = %s,
progress = 100,
last_error_code = NULL,
last_error_message = NULL,
updated_at = now()
WHERE job.id = %s
AND job.status = 'RUNNING'
AND job.lease_owner = %s
AND job.lease_token = %s
AND job.lease_until >= now()
AND job.job_type = 'PARSE_DOCUMENT'
AND job.required_capability = 'document_parse'
AND job.resource_type = 'document'
AND job.resource_id = %s
AND EXISTS (
SELECT 1
FROM rag.document_uploads AS upload
JOIN rag.documents AS document
ON document.id = upload.document_id
AND document.knowledge_base_id = upload.knowledge_base_id
AND document.access_scope_id = upload.access_scope_id
WHERE upload.parse_job_id = job.id
AND upload.id = %s
AND upload.document_id = %s
AND upload.knowledge_base_id = %s
AND upload.access_scope_id = %s
AND upload.actual_size = %s
AND upload.actual_sha256 = %s
AND upload.storage_key = %s
AND upload.status = 'COMPLETED'
AND document.raw_sha256 = %s
AND document.storage_key = %s
AND document.deleted_at IS NULL
)
RETURNING job.id
"""
class PostgresDocumentWorkflowRepository:
def __init__(
self,
settings: Settings,
*,
connect_timeout: int = 5,
connection_factory: ConnectionFactory | None = None,
) -> None:
self._settings = settings
self._connect_timeout = connect_timeout
self._connection_factory = connection_factory or _default_connection_factory
def _dsn(self) -> str:
return (
self._settings.database_url()
.set(drivername="postgresql")
.render_as_string(hide_password=False)
)
def load_source(self, job: BackgroundJob) -> DocumentSource:
upload_id, document_id = _validate_job_contract(job)
parameters = (
job.id,
job.lease.worker_id,
job.lease.lease_token,
document_id,
upload_id,
document_id,
)
try:
with self._connection_factory(self._dsn(), self._connect_timeout) as connection:
with connection.transaction():
row = connection.execute(LOAD_SOURCE_SQL, parameters).fetchone()
except (OSError, SecretFileError, psycopg.Error):
raise DocumentWorkflowPersistenceError from None
if row is None:
raise LeaseLostError("document source fence is no longer valid")
source = _source_from_row(row)
if source.upload_id != upload_id or source.document_id != document_id:
raise ArtifactConflictError
return source
def record_terminal_parse_failure(
self,
*,
lease: JobLease,
source: DocumentSource,
error_code: str,
) -> None:
if _ERROR_CODE.fullmatch(error_code) is None:
raise ValueError("error_code must be a stable uppercase identifier")
parameters = _fence_parameters(lease, source) + (error_code,)
try:
with self._connection_factory(self._dsn(), self._connect_timeout) as connection:
with connection.transaction():
row = connection.execute(FAIL_PARSE_SQL, parameters).fetchone()
if row is None:
raise LeaseLostError("document failure fence is no longer valid")
except (OSError, SecretFileError, psycopg.Error):
raise DocumentWorkflowPersistenceError from None
def persist_artifact(
self,
*,
lease: JobLease,
source: DocumentSource,
artifact: IngestionArtifact,
cloud_policy: CloudTextPolicy,
embedding_model: str,
embedding_dimension: int,
) -> ArtifactPlan:
plan = plan_artifact(
source,
artifact,
cloud_policy=cloud_policy,
)
if not embedding_model.strip() or embedding_dimension != 1024:
raise ArtifactConflictError
try:
with self._connection_factory(self._dsn(), self._connect_timeout) as connection:
with connection.transaction():
self._insert_plan(
connection,
source,
plan,
embedding_model=embedding_model,
embedding_dimension=embedding_dimension,
)
self._verify_plan(connection, source, plan)
document = connection.execute(
UPDATE_DOCUMENT_SQL,
(
plan.document_status,
source.document_id,
source.knowledge_base_id,
source.access_scope_id,
source.raw_sha256,
str(source.storage_key),
),
).fetchone()
if document is None:
raise ArtifactConflictError
finalized = connection.execute(
FINALIZE_JOB_SQL,
(plan.job_stage, *_fence_parameters(lease, source)),
).fetchone()
if finalized is None:
raise LeaseLostError("document artifact fence is no longer valid")
except (OSError, SecretFileError, psycopg.Error):
raise DocumentWorkflowPersistenceError from None
return plan
@staticmethod
def _insert_plan(
connection: Connection[WorkflowRow],
source: DocumentSource,
plan: ArtifactPlan,
*,
embedding_model: str,
embedding_dimension: int,
) -> None:
connection.execute(
INSERT_VERSION_SQL,
(
plan.version_id,
source.document_id,
plan.parser_profile_hash,
plan.normalization_profile_hash,
plan.chunk_profile_hash,
plan.cloud_policy_id,
plan.outbound_manifest_sha256,
plan.review_state,
plan.version_status,
plan.expected_chunk_count,
plan.version_error_code,
),
)
with connection.cursor() as cursor:
pages = [
(
page.id,
plan.version_id,
page.ordinal,
page.page_number,
page.text,
page.text_sha256,
page.line_start,
page.line_end,
)
for page in plan.pages
]
if pages:
cursor.executemany(INSERT_PAGE_SQL, pages)
blocks = [
(
block.id,
plan.version_id,
block.page_id,
block.ordinal,
block.kind,
block.text,
block.text_sha256,
_json(block.section_path),
block.anchor_id,
block.normalized_text_sha256,
block.char_start,
block.char_end,
block.line_start,
block.line_end,
block.page_start,
block.page_end,
)
for block in plan.blocks
]
if blocks:
cursor.executemany(INSERT_BLOCK_SQL, blocks)
if plan.outbound_manifest_sha256 is not None and plan.chunks:
cursor.executemany(
INSERT_MANIFEST_ITEM_SQL,
[
(
plan.version_id,
chunk.ordinal,
plan.outbound_manifest_sha256,
chunk.cloud_text_sha256,
chunk.embedding_text_sha256,
)
for chunk in plan.chunks
],
)
chunks = [
(
chunk.id,
source.knowledge_base_id,
source.document_id,
plan.version_id,
source.access_scope_id,
chunk.ordinal,
chunk.display_text,
chunk.cloud_text,
chunk.cloud_text_sha256,
chunk.embedding_prefix,
chunk.embedding_text,
chunk.embedding_text_sha256,
plan.outbound_manifest_sha256,
chunk.token_count,
chunk.page_start,
chunk.page_end,
_json(chunk.section_path),
_json(chunk.metadata),
embedding_model,
embedding_dimension,
)
for chunk in plan.chunks
]
if chunks:
cursor.executemany(INSERT_CHUNK_SQL, chunks)
@staticmethod
def _verify_plan(
connection: Connection[WorkflowRow],
source: DocumentSource,
plan: ArtifactPlan,
) -> None:
version = connection.execute(
SELECT_VERSION_SQL, (plan.version_id, source.document_id)
).fetchone()
expected_version = (
plan.version_id,
source.document_id,
plan.parser_profile_hash,
plan.normalization_profile_hash,
plan.chunk_profile_hash,
plan.cloud_policy_id,
plan.outbound_manifest_sha256,
plan.expected_chunk_count,
)
if version is None or _version_tuple(version) != expected_version:
raise ArtifactConflictError
pages = list(connection.execute(SELECT_PAGES_SQL, (plan.version_id,)).fetchall())
blocks = list(connection.execute(SELECT_BLOCKS_SQL, (plan.version_id,)).fetchall())
manifests = list(
connection.execute(SELECT_MANIFEST_ITEMS_SQL, (plan.version_id,)).fetchall()
)
chunks = list(connection.execute(SELECT_CHUNKS_SQL, (plan.version_id,)).fetchall())
if tuple(_page_tuple(row) for row in pages) != tuple(
_planned_page_tuple(value) for value in plan.pages
):
raise ArtifactConflictError
if tuple(_block_tuple(row) for row in blocks) != tuple(
_planned_block_tuple(value) for value in plan.blocks
):
raise ArtifactConflictError
expected_manifests = (
tuple(
(
chunk.ordinal,
plan.outbound_manifest_sha256,
chunk.cloud_text_sha256,
chunk.embedding_text_sha256,
)
for chunk in plan.chunks
)
if plan.outbound_manifest_sha256 is not None
else ()
)
if tuple(_manifest_tuple(row) for row in manifests) != expected_manifests:
raise ArtifactConflictError
if tuple(_chunk_tuple(row) for row in chunks) != tuple(
_planned_chunk_tuple(value) for value in plan.chunks
):
raise ArtifactConflictError
def plan_artifact(
source: DocumentSource,
artifact: IngestionArtifact,
*,
cloud_policy: CloudTextPolicy,
) -> ArtifactPlan:
if artifact.raw_sha256 != source.raw_sha256:
raise ArtifactConflictError
try:
source_version_id = uuid.UUID(artifact.document_version_id)
except ValueError:
raise ArtifactConflictError from None
version_id = _namespaced_id(
"version",
source.knowledge_base_id,
source.document_id,
source_version_id,
)
if artifact.status is IngestionStatus.OCR_REQUIRED:
if artifact.pages or artifact.blocks or artifact.chunks or artifact.manifest is not None:
raise ArtifactConflictError
error_code = artifact.limitations[0] if artifact.limitations else "OCR_REQUIRED"
if _ERROR_CODE.fullmatch(error_code) is None:
raise ArtifactConflictError
return ArtifactPlan(
version_id=version_id,
parser_profile_hash=_checked_hash(artifact.parser_profile_hash),
normalization_profile_hash=_NORMALIZATION_PROFILE_HASH,
chunk_profile_hash=_checked_hash(artifact.chunk_profile_hash),
cloud_policy_id=cloud_policy.policy_id,
outbound_manifest_sha256=None,
review_state="QUARANTINED_LOCAL_REVIEW",
version_status="PENDING",
version_error_code=error_code,
expected_chunk_count=0,
document_status="LOCAL_OCR_REQUIRED",
job_stage="OCR_REQUIRED",
pages=(),
blocks=(),
chunks=(),
)
if (
artifact.status is not IngestionStatus.READY_FOR_LOCAL_REVIEW
or artifact.manifest is None
or artifact.normalized_text_sha256 is None
or not artifact.pages
or not artifact.blocks
or not artifact.chunks
or artifact.manifest.cloud_policy_id != cloud_policy.policy_id
or artifact.manifest.raw_sha256 != source.raw_sha256
or artifact.manifest.normalized_text_sha256 != artifact.normalized_text_sha256
):
raise ArtifactConflictError
page_number_ids: dict[int | None, uuid.UUID] = {}
pages: list[PlannedPage] = []
for page in artifact.pages:
page_id = _namespaced_id("page", source.knowledge_base_id, version_id, page.page_id)
if page.page_number in page_number_ids:
raise ArtifactConflictError
page_number_ids[page.page_number] = page_id
pages.append(
PlannedPage(
id=page_id,
ordinal=page.ordinal,
page_number=page.page_number,
text=page.text,
text_sha256=_checked_hash(page.text_sha256),
line_start=page.line_start,
line_end=page.line_end,
)
)
block_ids: dict[str, uuid.UUID] = {}
blocks: list[PlannedBlock] = []
for block in artifact.blocks:
block_id = _namespaced_id("block", source.knowledge_base_id, version_id, block.block_id)
block_ids[block.block_id] = block_id
block_page_id = page_number_ids.get(block.anchor.page_start)
if block_page_id is None:
raise ArtifactConflictError
blocks.append(
PlannedBlock(
id=block_id,
page_id=block_page_id,
ordinal=block.ordinal,
kind=block.kind.value,
text=block.text,
text_sha256=_checked_hash(block.text_sha256),
section_path=block.section_path,
anchor_id=_checked_hash(block.anchor.anchor_id),
normalized_text_sha256=_checked_hash(block.anchor.normalized_text_sha256),
char_start=block.anchor.char_start,
char_end=block.anchor.char_end,
line_start=block.anchor.line_start,
line_end=block.anchor.line_end,
page_start=block.anchor.page_start,
page_end=block.anchor.page_end,
)
)
chunks: list[PlannedChunk] = []
for chunk in artifact.chunks:
try:
remapped_blocks = tuple(str(block_ids[value]) for value in chunk.anchor.block_ids)
except KeyError:
raise ArtifactConflictError from None
chunk_id = _namespaced_id("chunk", source.knowledge_base_id, version_id, chunk.chunk_id)
chunks.append(
PlannedChunk(
id=chunk_id,
ordinal=chunk.ordinal,
display_text=chunk.display_text,
cloud_text=chunk.cloud_text,
cloud_text_sha256=_checked_hash(chunk.cloud_text_sha256),
embedding_prefix=chunk.embedding_prefix,
embedding_text=chunk.embedding_text,
embedding_text_sha256=_checked_hash(chunk.embedding_text_sha256),
token_count=chunk.token_count,
page_start=chunk.anchor.page_start,
page_end=chunk.anchor.page_end,
section_path=chunk.section_path,
metadata={
"source_anchor": {
"anchor_id": _checked_hash(chunk.anchor.anchor_id),
"normalized_text_sha256": _checked_hash(
chunk.anchor.normalized_text_sha256
),
"char_start": chunk.anchor.char_start,
"char_end": chunk.anchor.char_end,
"line_start": chunk.anchor.line_start,
"line_end": chunk.anchor.line_end,
"page_start": chunk.anchor.page_start,
"page_end": chunk.anchor.page_end,
"block_ids": remapped_blocks,
},
"ingestion_profile": "document-ingestion-v1",
},
)
)
if [item.ordinal for item in artifact.manifest.items] != [chunk.ordinal for chunk in chunks]:
raise ArtifactConflictError
for item, source_chunk, planned_chunk in zip(
artifact.manifest.items, artifact.chunks, chunks, strict=True
):
if (
item.chunk_id != source_chunk.chunk_id
or item.anchor_id != source_chunk.anchor.anchor_id
or item.cloud_text_sha256 != planned_chunk.cloud_text_sha256
or item.embedding_text_sha256 != planned_chunk.embedding_text_sha256
):
raise ArtifactConflictError
return ArtifactPlan(
version_id=version_id,
parser_profile_hash=_checked_hash(artifact.parser_profile_hash),
normalization_profile_hash=_NORMALIZATION_PROFILE_HASH,
chunk_profile_hash=_checked_hash(artifact.chunk_profile_hash),
cloud_policy_id=cloud_policy.policy_id,
outbound_manifest_sha256=_checked_hash(artifact.manifest.manifest_sha256),
review_state="LOCAL_PARSED_PENDING_CLOUD_REVIEW",
version_status="PROCESSING",
version_error_code=None,
expected_chunk_count=len(chunks),
document_status="LOCAL_PARSED_PENDING_CLOUD_REVIEW",
job_stage="LOCAL_PARSED_PENDING_CLOUD_REVIEW",
pages=tuple(pages),
blocks=tuple(blocks),
chunks=tuple(chunks),
)
def _validate_job_contract(job: BackgroundJob) -> tuple[uuid.UUID, uuid.UUID]:
if (
job.lease.job_id != job.id
or not job.lease.worker_id.strip()
or job.job_type != "PARSE_DOCUMENT"
or job.required_capability != "document_parse"
or job.resource_type != "document"
or set(job.payload) != {"upload_id", "document_id"}
):
raise InvalidDocumentJobError
try:
upload_id = uuid.UUID(cast(str, job.payload["upload_id"]))
document_id = uuid.UUID(cast(str, job.payload["document_id"]))
except (ValueError, TypeError, KeyError):
raise InvalidDocumentJobError from None
if document_id != job.resource_id:
raise InvalidDocumentJobError
return upload_id, document_id
def _source_from_row(row: Mapping[str, object]) -> DocumentSource:
source = DocumentSource(
upload_id=_uuid(row, "upload_id"),
document_id=_uuid(row, "document_id"),
knowledge_base_id=_uuid(row, "knowledge_base_id"),
access_scope_id=_uuid(row, "access_scope_id"),
filename=_text(row, "original_filename"),
mime_type=_text(row, "declared_mime_type"),
storage_key=_uuid(row, "storage_key"),
byte_size=_integer(row, "actual_size"),
raw_sha256=_checked_hash(_text(row, "actual_sha256")),
)
if (
_text(row, "document_filename") != source.filename
or _text(row, "document_mime_type") != source.mime_type
or _checked_hash(_text(row, "document_raw_sha256")) != source.raw_sha256
or _text(row, "document_storage_key") != str(source.storage_key)
):
raise ArtifactConflictError
return source
def _fence_parameters(lease: JobLease, source: DocumentSource) -> tuple[object, ...]:
return (
lease.job_id,
lease.worker_id,
lease.lease_token,
source.document_id,
source.upload_id,
source.document_id,
source.knowledge_base_id,
source.access_scope_id,
source.byte_size,
source.raw_sha256,
source.storage_key,
source.raw_sha256,
str(source.storage_key),
)
def _namespaced_id(kind: str, *parts: object) -> uuid.UUID:
return uuid.uuid5(_DB_ID_NAMESPACE, "\x1f".join((kind, *(str(value) for value in parts))))
def _checked_hash(value: str) -> str:
if _HASH.fullmatch(value) is None:
raise ArtifactConflictError
return value
def _json(value: object) -> str:
return json.dumps(value, ensure_ascii=False, sort_keys=True, separators=(",", ":"))
def _default_connection_factory(dsn: str, timeout: int) -> Connection[WorkflowRow]:
return psycopg.connect(
dsn,
connect_timeout=timeout,
row_factory=dict_row,
application_name="geological-rag-document-workflow",
)
def _version_tuple(row: Mapping[str, object]) -> tuple[object, ...]:
return (
_uuid(row, "id"),
_uuid(row, "document_id"),
_checked_hash(_text(row, "parser_profile_hash")),
_checked_hash(_text(row, "normalization_profile_hash")),
_checked_hash(_text(row, "chunk_profile_hash")),
_text(row, "cloud_policy_id"),
_optional_hash(row, "outbound_manifest_sha256"),
_integer(row, "expected_chunk_count"),
)
def _planned_page_tuple(value: PlannedPage) -> tuple[object, ...]:
return (
value.id,
value.ordinal,
value.page_number,
value.text,
value.text_sha256,
value.line_start,
value.line_end,
)
def _page_tuple(row: Mapping[str, object]) -> tuple[object, ...]:
return (
_uuid(row, "id"),
_integer(row, "ordinal"),
_optional_integer(row, "page_number"),
_text(row, "display_text"),
_checked_hash(_text(row, "text_sha256")),
_integer(row, "line_start"),
_integer(row, "line_end"),
)
def _planned_block_tuple(value: PlannedBlock) -> tuple[object, ...]:
return (
value.id,
value.page_id,
value.ordinal,
value.kind,
value.text,
value.text_sha256,
value.section_path,
value.anchor_id,
value.normalized_text_sha256,
value.char_start,
value.char_end,
value.line_start,
value.line_end,
value.page_start,
value.page_end,
)
def _block_tuple(row: Mapping[str, object]) -> tuple[object, ...]:
return (
_uuid(row, "id"),
_uuid(row, "page_id"),
_integer(row, "ordinal"),
_text(row, "block_kind"),
_text(row, "display_text"),
_checked_hash(_text(row, "text_sha256")),
_string_tuple(row, "section_path"),
_checked_hash(_text(row, "anchor_id")),
_checked_hash(_text(row, "normalized_text_sha256")),
_integer(row, "char_start"),
_integer(row, "char_end"),
_integer(row, "line_start"),
_integer(row, "line_end"),
_optional_integer(row, "page_start"),
_optional_integer(row, "page_end"),
)
def _manifest_tuple(row: Mapping[str, object]) -> tuple[object, ...]:
return (
_integer(row, "ordinal"),
_checked_hash(_text(row, "outbound_manifest_sha256")),
_checked_hash(_text(row, "cloud_text_sha256")),
_checked_hash(_text(row, "embedding_text_sha256")),
)
def _planned_chunk_tuple(value: PlannedChunk) -> tuple[object, ...]:
return (
value.id,
value.ordinal,
value.display_text,
value.cloud_text,
value.cloud_text_sha256,
value.embedding_prefix,
value.embedding_text,
value.embedding_text_sha256,
value.token_count,
value.page_start,
value.page_end,
value.section_path,
_canonical_json(value.metadata),
)
def _chunk_tuple(row: Mapping[str, object]) -> tuple[object, ...]:
metadata = row.get("metadata")
if not isinstance(metadata, dict):
raise ArtifactConflictError
return (
_uuid(row, "id"),
_integer(row, "ordinal"),
_text(row, "display_text"),
_text(row, "cloud_text"),
_checked_hash(_text(row, "cloud_text_sha256")),
_text(row, "embedding_prefix"),
_text(row, "embedding_text"),
_checked_hash(_text(row, "embedding_text_sha256")),
_integer(row, "token_count"),
_optional_integer(row, "page_start"),
_optional_integer(row, "page_end"),
_string_tuple(row, "section_path"),
metadata,
)
def _canonical_json(value: object) -> object:
"""Match PostgreSQL jsonb's array representation for tuple-backed plans."""
return json.loads(_json(value))
def _uuid(row: Mapping[str, object], name: str) -> uuid.UUID:
value = row.get(name)
if isinstance(value, uuid.UUID):
return value
if isinstance(value, str):
try:
return uuid.UUID(value)
except ValueError:
pass
raise ArtifactConflictError
def _text(row: Mapping[str, object], name: str) -> str:
value = row.get(name)
if not isinstance(value, str) or not value.strip():
raise ArtifactConflictError
return value
def _integer(row: Mapping[str, object], name: str) -> int:
value = row.get(name)
if isinstance(value, bool) or not isinstance(value, int):
raise ArtifactConflictError
return value
def _optional_integer(row: Mapping[str, object], name: str) -> int | None:
return None if row.get(name) is None else _integer(row, name)
def _optional_hash(row: Mapping[str, object], name: str) -> str | None:
return None if row.get(name) is None else _checked_hash(_text(row, name))
def _string_tuple(row: Mapping[str, object], name: str) -> tuple[str, ...]:
value = row.get(name)
if not isinstance(value, list) or not all(isinstance(item, str) for item in value):
raise ArtifactConflictError
return tuple(value)