Make the governed RAG evidence path executable end to end
Some checks failed
verify / verify (push) Has been cancelled
Some checks failed
verify / verify (push) Has been cancelled
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
This commit is contained in:
517
backend/app/persistence/document_review.py
Normal file
517
backend/app/persistence/document_review.py
Normal file
@@ -0,0 +1,517 @@
|
||||
"""Optimistic, manifest-bound document review persistence."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import re
|
||||
import uuid
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime
|
||||
from typing import Literal
|
||||
|
||||
import psycopg
|
||||
from psycopg.rows import dict_row
|
||||
|
||||
from app.core.config import Settings
|
||||
from app.core.secrets import SecretFileError
|
||||
from app.persistence.documents import DocumentActor, SafeJob
|
||||
|
||||
type ReviewDecision = Literal["APPROVE", "REJECT"]
|
||||
type ReviewReason = Literal[
|
||||
"SYNTHETIC_REVIEW_APPROVED",
|
||||
"RIGHTS_NOT_VERIFIED",
|
||||
"CONTENT_QUALITY_REJECTED",
|
||||
"CLOUD_PROCESSING_REJECTED",
|
||||
]
|
||||
|
||||
_HASH = re.compile(r"^[0-9a-f]{64}$")
|
||||
LOGGER = logging.getLogger("geological_rag.document_review")
|
||||
|
||||
|
||||
class DocumentReviewError(RuntimeError):
|
||||
"""Base class for safe review persistence failures."""
|
||||
|
||||
|
||||
class DocumentReviewNotFoundError(DocumentReviewError):
|
||||
pass
|
||||
|
||||
|
||||
class DocumentReviewConflictError(DocumentReviewError):
|
||||
pass
|
||||
|
||||
|
||||
class DocumentReviewStateError(DocumentReviewError):
|
||||
pass
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class DocumentReviewResult:
|
||||
document_id: uuid.UUID
|
||||
document_version_id: uuid.UUID
|
||||
decision: ReviewDecision
|
||||
review_state: Literal["CLOUD_APPROVED", "REJECTED"]
|
||||
review_revision: int
|
||||
outbound_manifest_sha256: str | None
|
||||
embedding_profile_hash: str | None
|
||||
job: SafeJob | None
|
||||
|
||||
|
||||
_LOCK_REVIEW = """
|
||||
SELECT
|
||||
document.id AS document_id,
|
||||
version.id AS document_version_id,
|
||||
version.review_state,
|
||||
version.review_revision,
|
||||
version.status AS version_status,
|
||||
version.outbound_manifest_sha256,
|
||||
version.expected_chunk_count,
|
||||
knowledge_base.active_embedding_profile_hash,
|
||||
profile.model AS profile_model,
|
||||
profile.dimension AS profile_dimension,
|
||||
profile.enabled AS profile_enabled
|
||||
FROM rag.documents AS document
|
||||
JOIN rag.document_versions AS version
|
||||
ON version.id = (
|
||||
SELECT candidate.id
|
||||
FROM rag.document_versions AS candidate
|
||||
WHERE candidate.document_id = document.id
|
||||
ORDER BY candidate.created_at DESC, candidate.id DESC
|
||||
LIMIT 1
|
||||
)
|
||||
JOIN rag.knowledge_bases AS knowledge_base
|
||||
ON knowledge_base.id = document.knowledge_base_id
|
||||
LEFT JOIN rag.model_profiles AS profile
|
||||
ON profile.profile_hash = knowledge_base.active_embedding_profile_hash
|
||||
AND profile.kind = 'embedding'
|
||||
WHERE document.id = %s
|
||||
AND document.knowledge_base_id = %s
|
||||
AND document.access_scope_id = %s
|
||||
AND document.deleted_at IS NULL
|
||||
FOR UPDATE OF document, version
|
||||
"""
|
||||
|
||||
_APPROVE_VERSION = """
|
||||
UPDATE rag.document_versions
|
||||
SET review_state = 'CLOUD_APPROVED',
|
||||
embedding_profile_hash = %s,
|
||||
cloud_approved_at = now(),
|
||||
cloud_approved_by = %s,
|
||||
review_revision = review_revision + 1
|
||||
WHERE id = %s
|
||||
AND review_revision = %s
|
||||
AND review_state = 'LOCAL_PARSED_PENDING_CLOUD_REVIEW'
|
||||
AND status = 'PROCESSING'
|
||||
AND outbound_manifest_sha256 = %s
|
||||
RETURNING review_revision
|
||||
"""
|
||||
|
||||
_APPROVE_CHUNKS = """
|
||||
UPDATE rag.chunks
|
||||
SET approval_status = 'CLOUD_APPROVED',
|
||||
outbound_manifest_sha256 = %s,
|
||||
embedding_profile_hash = %s,
|
||||
embedding_model = %s,
|
||||
embedding_dimension = 1024,
|
||||
index_status = 'PENDING',
|
||||
searchable = false,
|
||||
updated_at = now()
|
||||
WHERE document_version_id = %s
|
||||
AND approval_status = 'LOCAL_PARSED_PENDING_CLOUD_REVIEW'
|
||||
RETURNING id, embedding_text_sha256
|
||||
"""
|
||||
|
||||
_REJECT_VERSION = """
|
||||
UPDATE rag.document_versions
|
||||
SET review_state = 'REJECTED',
|
||||
embedding_profile_hash = NULL,
|
||||
cloud_approved_at = NULL,
|
||||
cloud_approved_by = NULL,
|
||||
review_revision = review_revision + 1
|
||||
WHERE id = %s
|
||||
AND review_revision = %s
|
||||
AND review_state = 'LOCAL_PARSED_PENDING_CLOUD_REVIEW'
|
||||
AND status = 'PROCESSING'
|
||||
RETURNING review_revision
|
||||
"""
|
||||
|
||||
_ENQUEUE_EMBED_JOB = """
|
||||
INSERT INTO rag.background_jobs (
|
||||
job_type, required_capability, resource_type, resource_id,
|
||||
idempotency_key, payload, stage, status, max_attempts
|
||||
) VALUES (
|
||||
'EMBED_DOCUMENT', 'embedding', 'document_version', %s,
|
||||
%s, jsonb_build_object('document_version_id', %s::text),
|
||||
'PENDING', 'QUEUED', 3
|
||||
)
|
||||
ON CONFLICT (job_type, idempotency_key)
|
||||
DO UPDATE SET updated_at = rag.background_jobs.updated_at
|
||||
RETURNING id, job_type, stage, status, progress, attempt,
|
||||
max_attempts, last_error_code, created_at, updated_at, finished_at
|
||||
"""
|
||||
|
||||
|
||||
class PostgresDocumentReviewRepository:
|
||||
def __init__(self, settings: Settings, *, connect_timeout: int = 5) -> None:
|
||||
self._settings = settings
|
||||
self._connect_timeout = connect_timeout
|
||||
|
||||
def _dsn(self) -> str:
|
||||
return (
|
||||
self._settings.database_url()
|
||||
.set(drivername="postgresql")
|
||||
.render_as_string(hide_password=False)
|
||||
)
|
||||
|
||||
def apply_decision(
|
||||
self,
|
||||
*,
|
||||
actor: DocumentActor,
|
||||
document_id: uuid.UUID,
|
||||
decision: ReviewDecision,
|
||||
reason_code: ReviewReason,
|
||||
expected_revision: int,
|
||||
outbound_manifest_sha256: str | None,
|
||||
trace_id: uuid.UUID,
|
||||
) -> DocumentReviewResult:
|
||||
_validate_decision(
|
||||
decision=decision,
|
||||
reason_code=reason_code,
|
||||
expected_revision=expected_revision,
|
||||
outbound_manifest_sha256=outbound_manifest_sha256,
|
||||
)
|
||||
try:
|
||||
with psycopg.connect(
|
||||
self._dsn(),
|
||||
connect_timeout=self._connect_timeout,
|
||||
row_factory=dict_row,
|
||||
application_name="geological-rag-document-review",
|
||||
) as connection:
|
||||
with connection.transaction():
|
||||
row = connection.execute(
|
||||
_LOCK_REVIEW,
|
||||
(document_id, actor.knowledge_base_id, actor.access_scope_id),
|
||||
).fetchone()
|
||||
if row is None:
|
||||
raise DocumentReviewNotFoundError
|
||||
current_revision = int(row["review_revision"])
|
||||
if current_revision != expected_revision:
|
||||
raise DocumentReviewConflictError
|
||||
version_id = _uuid_value(row["document_version_id"])
|
||||
manifest = _optional_text(row["outbound_manifest_sha256"])
|
||||
if decision == "APPROVE":
|
||||
return self._approve(
|
||||
connection=connection,
|
||||
actor=actor,
|
||||
document_id=document_id,
|
||||
version_id=version_id,
|
||||
current=row,
|
||||
manifest=manifest,
|
||||
supplied_manifest=outbound_manifest_sha256,
|
||||
expected_revision=expected_revision,
|
||||
reason_code=reason_code,
|
||||
trace_id=trace_id,
|
||||
)
|
||||
return self._reject(
|
||||
connection=connection,
|
||||
actor=actor,
|
||||
document_id=document_id,
|
||||
version_id=version_id,
|
||||
manifest=manifest,
|
||||
expected_revision=expected_revision,
|
||||
reason_code=reason_code,
|
||||
trace_id=trace_id,
|
||||
)
|
||||
except DocumentReviewError:
|
||||
raise
|
||||
except psycopg.Error as exc:
|
||||
LOGGER.error(
|
||||
"document_review_database_error sqlstate=%s",
|
||||
exc.sqlstate or "UNKNOWN",
|
||||
)
|
||||
raise DocumentReviewError from None
|
||||
except (OSError, SecretFileError, KeyError, TypeError, ValueError):
|
||||
raise DocumentReviewError from None
|
||||
|
||||
def _approve(
|
||||
self,
|
||||
*,
|
||||
connection: psycopg.Connection[dict[str, object]],
|
||||
actor: DocumentActor,
|
||||
document_id: uuid.UUID,
|
||||
version_id: uuid.UUID,
|
||||
current: dict[str, object],
|
||||
manifest: str | None,
|
||||
supplied_manifest: str | None,
|
||||
expected_revision: int,
|
||||
reason_code: ReviewReason,
|
||||
trace_id: uuid.UUID,
|
||||
) -> DocumentReviewResult:
|
||||
profile_hash = _optional_text(current.get("active_embedding_profile_hash"))
|
||||
profile_model = _optional_text(current.get("profile_model"))
|
||||
expected_count = current.get("expected_chunk_count")
|
||||
if (
|
||||
current.get("review_state") != "LOCAL_PARSED_PENDING_CLOUD_REVIEW"
|
||||
or current.get("version_status") != "PROCESSING"
|
||||
or manifest is None
|
||||
or supplied_manifest != manifest
|
||||
or profile_hash is None
|
||||
or profile_model is None
|
||||
or current.get("profile_enabled") is not True
|
||||
or current.get("profile_dimension") != 1024
|
||||
or not isinstance(expected_count, int)
|
||||
or isinstance(expected_count, bool)
|
||||
or expected_count < 1
|
||||
):
|
||||
raise DocumentReviewStateError
|
||||
revision_row = connection.execute(
|
||||
_APPROVE_VERSION,
|
||||
(profile_hash, actor.subject, version_id, expected_revision, manifest),
|
||||
).fetchone()
|
||||
if revision_row is None:
|
||||
raise DocumentReviewConflictError
|
||||
chunks = list(
|
||||
connection.execute(
|
||||
_APPROVE_CHUNKS,
|
||||
(manifest, profile_hash, profile_model, version_id),
|
||||
).fetchall()
|
||||
)
|
||||
if len(chunks) != expected_count:
|
||||
raise DocumentReviewStateError
|
||||
connection.execute(
|
||||
"""
|
||||
INSERT INTO rag.chunk_embedding_assignments (
|
||||
chunk_id, profile_hash, embedding_text_sha256, status
|
||||
)
|
||||
SELECT id, %s, embedding_text_sha256, 'PENDING'
|
||||
FROM rag.chunks
|
||||
WHERE document_version_id = %s
|
||||
ON CONFLICT (chunk_id, profile_hash) DO NOTHING
|
||||
""",
|
||||
(profile_hash, version_id),
|
||||
)
|
||||
assignment_count = connection.execute(
|
||||
"""
|
||||
SELECT count(*)
|
||||
FROM rag.chunk_embedding_assignments AS assignment
|
||||
JOIN rag.chunks AS chunk ON chunk.id = assignment.chunk_id
|
||||
WHERE chunk.document_version_id = %s
|
||||
AND assignment.profile_hash = %s
|
||||
AND assignment.embedding_text_sha256 = chunk.embedding_text_sha256
|
||||
AND assignment.status = 'PENDING'
|
||||
""",
|
||||
(version_id, profile_hash),
|
||||
).fetchone()
|
||||
if assignment_count is None or assignment_count["count"] != expected_count:
|
||||
raise DocumentReviewStateError
|
||||
updated_document = connection.execute(
|
||||
"""
|
||||
UPDATE rag.documents
|
||||
SET status = 'CLOUD_APPROVED', updated_at = now()
|
||||
WHERE id = %s AND knowledge_base_id = %s AND access_scope_id = %s
|
||||
RETURNING id
|
||||
""",
|
||||
(document_id, actor.knowledge_base_id, actor.access_scope_id),
|
||||
).fetchone()
|
||||
if updated_document is None:
|
||||
raise DocumentReviewConflictError
|
||||
job = connection.execute(
|
||||
_ENQUEUE_EMBED_JOB,
|
||||
(
|
||||
version_id,
|
||||
f"embed-document:{version_id}:{profile_hash}",
|
||||
str(version_id),
|
||||
),
|
||||
).fetchone()
|
||||
if job is None:
|
||||
raise DocumentReviewError
|
||||
revision = _integer_value(revision_row["review_revision"])
|
||||
self._audit(
|
||||
connection=connection,
|
||||
document_id=document_id,
|
||||
version_id=version_id,
|
||||
actor=actor,
|
||||
decision="APPROVE",
|
||||
reason_code=reason_code,
|
||||
previous_revision=expected_revision,
|
||||
resulting_revision=revision,
|
||||
manifest=manifest,
|
||||
profile_hash=profile_hash,
|
||||
trace_id=trace_id,
|
||||
)
|
||||
return DocumentReviewResult(
|
||||
document_id=document_id,
|
||||
document_version_id=version_id,
|
||||
decision="APPROVE",
|
||||
review_state="CLOUD_APPROVED",
|
||||
review_revision=revision,
|
||||
outbound_manifest_sha256=manifest,
|
||||
embedding_profile_hash=profile_hash,
|
||||
job=_safe_job(job),
|
||||
)
|
||||
|
||||
def _reject(
|
||||
self,
|
||||
*,
|
||||
connection: psycopg.Connection[dict[str, object]],
|
||||
actor: DocumentActor,
|
||||
document_id: uuid.UUID,
|
||||
version_id: uuid.UUID,
|
||||
manifest: str | None,
|
||||
expected_revision: int,
|
||||
reason_code: ReviewReason,
|
||||
trace_id: uuid.UUID,
|
||||
) -> DocumentReviewResult:
|
||||
revision_row = connection.execute(
|
||||
_REJECT_VERSION,
|
||||
(version_id, expected_revision),
|
||||
).fetchone()
|
||||
if revision_row is None:
|
||||
raise DocumentReviewConflictError
|
||||
connection.execute(
|
||||
"""
|
||||
UPDATE rag.chunks
|
||||
SET approval_status = 'REJECTED', searchable = false,
|
||||
index_status = 'PENDING', embedding = NULL,
|
||||
embedded_text_sha256 = NULL, embedding_profile_hash = NULL,
|
||||
updated_at = now()
|
||||
WHERE document_version_id = %s
|
||||
""",
|
||||
(version_id,),
|
||||
)
|
||||
updated_document = connection.execute(
|
||||
"""
|
||||
UPDATE rag.documents
|
||||
SET status = 'REJECTED', active_version_id = NULL, updated_at = now()
|
||||
WHERE id = %s AND knowledge_base_id = %s AND access_scope_id = %s
|
||||
RETURNING id
|
||||
""",
|
||||
(document_id, actor.knowledge_base_id, actor.access_scope_id),
|
||||
).fetchone()
|
||||
if updated_document is None:
|
||||
raise DocumentReviewConflictError
|
||||
revision = _integer_value(revision_row["review_revision"])
|
||||
self._audit(
|
||||
connection=connection,
|
||||
document_id=document_id,
|
||||
version_id=version_id,
|
||||
actor=actor,
|
||||
decision="REJECT",
|
||||
reason_code=reason_code,
|
||||
previous_revision=expected_revision,
|
||||
resulting_revision=revision,
|
||||
manifest=manifest,
|
||||
profile_hash=None,
|
||||
trace_id=trace_id,
|
||||
)
|
||||
return DocumentReviewResult(
|
||||
document_id=document_id,
|
||||
document_version_id=version_id,
|
||||
decision="REJECT",
|
||||
review_state="REJECTED",
|
||||
review_revision=revision,
|
||||
outbound_manifest_sha256=manifest,
|
||||
embedding_profile_hash=None,
|
||||
job=None,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _audit(
|
||||
*,
|
||||
connection: psycopg.Connection[dict[str, object]],
|
||||
document_id: uuid.UUID,
|
||||
version_id: uuid.UUID,
|
||||
actor: DocumentActor,
|
||||
decision: ReviewDecision,
|
||||
reason_code: ReviewReason,
|
||||
previous_revision: int,
|
||||
resulting_revision: int,
|
||||
manifest: str | None,
|
||||
profile_hash: str | None,
|
||||
trace_id: uuid.UUID,
|
||||
) -> None:
|
||||
connection.execute(
|
||||
"""
|
||||
INSERT INTO rag.document_review_events (
|
||||
document_id, document_version_id, actor_subject, decision,
|
||||
reason_code, previous_revision, resulting_revision,
|
||||
outbound_manifest_sha256, embedding_profile_hash, trace_id
|
||||
) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
|
||||
""",
|
||||
(
|
||||
document_id,
|
||||
version_id,
|
||||
actor.subject,
|
||||
decision,
|
||||
reason_code,
|
||||
previous_revision,
|
||||
resulting_revision,
|
||||
manifest,
|
||||
profile_hash,
|
||||
trace_id,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _validate_decision(
|
||||
*,
|
||||
decision: ReviewDecision,
|
||||
reason_code: ReviewReason,
|
||||
expected_revision: int,
|
||||
outbound_manifest_sha256: str | None,
|
||||
) -> None:
|
||||
if isinstance(expected_revision, bool) or expected_revision < 0:
|
||||
raise ValueError("expected_revision must be non-negative")
|
||||
if decision == "APPROVE":
|
||||
if reason_code != "SYNTHETIC_REVIEW_APPROVED" or not (
|
||||
outbound_manifest_sha256 and _HASH.fullmatch(outbound_manifest_sha256)
|
||||
):
|
||||
raise ValueError("approval requires the reviewed manifest")
|
||||
elif decision == "REJECT":
|
||||
if reason_code == "SYNTHETIC_REVIEW_APPROVED":
|
||||
raise ValueError("rejection requires a rejection reason")
|
||||
else:
|
||||
raise ValueError("unsupported review decision")
|
||||
|
||||
|
||||
def _uuid_value(value: object) -> uuid.UUID:
|
||||
if not isinstance(value, uuid.UUID):
|
||||
raise DocumentReviewError
|
||||
return value
|
||||
|
||||
|
||||
def _optional_text(value: object) -> str | None:
|
||||
return value if isinstance(value, str) and value else None
|
||||
|
||||
|
||||
def _integer_value(value: object) -> int:
|
||||
if not isinstance(value, int) or isinstance(value, bool):
|
||||
raise DocumentReviewError
|
||||
return value
|
||||
|
||||
|
||||
def _datetime_value(value: object) -> datetime:
|
||||
if not isinstance(value, datetime):
|
||||
raise DocumentReviewError
|
||||
return value
|
||||
|
||||
|
||||
def _optional_datetime_value(value: object) -> datetime | None:
|
||||
if value is None:
|
||||
return None
|
||||
return _datetime_value(value)
|
||||
|
||||
|
||||
def _safe_job(row: dict[str, object]) -> SafeJob:
|
||||
return SafeJob(
|
||||
id=_uuid_value(row["id"]),
|
||||
job_type=str(row["job_type"]),
|
||||
stage=str(row["stage"]),
|
||||
status=str(row["status"]),
|
||||
progress=_integer_value(row["progress"]),
|
||||
attempt=_integer_value(row["attempt"]),
|
||||
max_attempts=_integer_value(row["max_attempts"]),
|
||||
last_error_code=_optional_text(row.get("last_error_code")),
|
||||
created_at=_datetime_value(row["created_at"]),
|
||||
updated_at=_datetime_value(row["updated_at"]),
|
||||
finished_at=_optional_datetime_value(row.get("finished_at")),
|
||||
)
|
||||
1121
backend/app/persistence/document_workflows.py
Normal file
1121
backend/app/persistence/document_workflows.py
Normal file
File diff suppressed because it is too large
Load Diff
1021
backend/app/persistence/documents.py
Normal file
1021
backend/app/persistence/documents.py
Normal file
File diff suppressed because it is too large
Load Diff
1204
backend/app/persistence/indexing.py
Normal file
1204
backend/app/persistence/indexing.py
Normal file
File diff suppressed because it is too large
Load Diff
345
backend/app/persistence/job_queue.py
Normal file
345
backend/app/persistence/job_queue.py
Normal file
@@ -0,0 +1,345 @@
|
||||
"""Transactional PostgreSQL runtime for the fenced background-job queue.
|
||||
|
||||
Every repository call owns a short database transaction. A claimed job is
|
||||
returned only after that transaction commits, so handlers never run while the
|
||||
claim row lock is held. Terminal mutations require the immutable worker/token
|
||||
lease pair and fail closed when the lease has moved to another worker.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
import uuid
|
||||
from collections.abc import Callable, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime
|
||||
from typing import Any
|
||||
|
||||
import psycopg
|
||||
from psycopg import Connection
|
||||
from psycopg.rows import dict_row
|
||||
|
||||
from app.persistence.job_queue_sql import (
|
||||
CLAIM_JOB_SQL,
|
||||
COMPLETE_JOB_SQL,
|
||||
FAIL_OR_RETRY_JOB_SQL,
|
||||
HEARTBEAT_JOB_SQL,
|
||||
REAP_EXPIRED_JOBS_SQL,
|
||||
)
|
||||
|
||||
type JobRow = dict[str, Any]
|
||||
type ConnectionFactory = Callable[[str, int], Connection[JobRow]]
|
||||
|
||||
_NAMED_BIND = re.compile(r"(?<!:):([a-zA-Z_][a-zA-Z0-9_]*)")
|
||||
_ERROR_CODE = re.compile(r"^[A-Z][A-Z0-9_]{0,127}$")
|
||||
|
||||
|
||||
def _psycopg_statement(statement: str) -> str:
|
||||
"""Translate the canonical named binds to psycopg's mapping syntax."""
|
||||
|
||||
return _NAMED_BIND.sub(r"%(\1)s", statement)
|
||||
|
||||
|
||||
_CLAIM = _psycopg_statement(CLAIM_JOB_SQL)
|
||||
_HEARTBEAT = _psycopg_statement(HEARTBEAT_JOB_SQL)
|
||||
_COMPLETE = _psycopg_statement(COMPLETE_JOB_SQL)
|
||||
_FAIL_OR_RETRY = _psycopg_statement(FAIL_OR_RETRY_JOB_SQL)
|
||||
_REAP_EXPIRED = _psycopg_statement(REAP_EXPIRED_JOBS_SQL)
|
||||
|
||||
|
||||
class JobQueueError(RuntimeError):
|
||||
"""Base class for safe, non-secret queue errors."""
|
||||
|
||||
|
||||
class LeaseLostError(JobQueueError):
|
||||
"""Raised when a fenced mutation no longer owns the job lease."""
|
||||
|
||||
|
||||
class InvalidJobRowError(JobQueueError):
|
||||
"""Raised when a claimed row violates the runtime shape contract."""
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class JobLease:
|
||||
"""The complete fencing identity required for mutations after claim."""
|
||||
|
||||
job_id: uuid.UUID
|
||||
worker_id: str
|
||||
lease_token: uuid.UUID
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class BackgroundJob:
|
||||
"""A claimed job and its immutable lease identity."""
|
||||
|
||||
id: uuid.UUID
|
||||
job_type: str
|
||||
required_capability: str
|
||||
resource_type: str
|
||||
resource_id: uuid.UUID
|
||||
idempotency_key: str
|
||||
payload: Mapping[str, object]
|
||||
stage: str
|
||||
progress: int
|
||||
priority: int
|
||||
attempt: int
|
||||
max_attempts: int
|
||||
run_after: datetime
|
||||
lease_until: datetime
|
||||
created_at: datetime
|
||||
updated_at: datetime
|
||||
lease: JobLease
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class LeaseHeartbeat:
|
||||
job_id: uuid.UUID
|
||||
lease_until: datetime
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class JobState:
|
||||
"""Metadata returned by a fenced terminal or reaper mutation."""
|
||||
|
||||
job_id: uuid.UUID
|
||||
status: str
|
||||
attempt: int
|
||||
max_attempts: int
|
||||
finished_at: datetime | None
|
||||
|
||||
|
||||
def _default_connection_factory(dsn: str, connect_timeout: int) -> Connection[JobRow]:
|
||||
return psycopg.connect(
|
||||
dsn,
|
||||
connect_timeout=connect_timeout,
|
||||
row_factory=dict_row,
|
||||
application_name="geological-rag-worker",
|
||||
)
|
||||
|
||||
|
||||
def _require_uuid(row: Mapping[str, object], name: str) -> uuid.UUID:
|
||||
value = row.get(name)
|
||||
if not isinstance(value, uuid.UUID):
|
||||
raise InvalidJobRowError(f"job row has invalid {name}")
|
||||
return value
|
||||
|
||||
|
||||
def _require_text(row: Mapping[str, object], name: str) -> str:
|
||||
value = row.get(name)
|
||||
if not isinstance(value, str) or not value.strip():
|
||||
raise InvalidJobRowError(f"job row has invalid {name}")
|
||||
return value
|
||||
|
||||
|
||||
def _require_integer(row: Mapping[str, object], name: str) -> int:
|
||||
value = row.get(name)
|
||||
if isinstance(value, bool) or not isinstance(value, int):
|
||||
raise InvalidJobRowError(f"job row has invalid {name}")
|
||||
return value
|
||||
|
||||
|
||||
def _require_datetime(row: Mapping[str, object], name: str) -> datetime:
|
||||
value = row.get(name)
|
||||
if not isinstance(value, datetime):
|
||||
raise InvalidJobRowError(f"job row has invalid {name}")
|
||||
return value
|
||||
|
||||
|
||||
def _job_from_row(row: JobRow, expected_worker_id: str) -> BackgroundJob:
|
||||
job_id = _require_uuid(row, "id")
|
||||
lease_owner = _require_text(row, "lease_owner")
|
||||
if lease_owner != expected_worker_id:
|
||||
raise InvalidJobRowError("claimed job has an unexpected lease owner")
|
||||
lease_token = _require_uuid(row, "lease_token")
|
||||
if row.get("status") != "RUNNING":
|
||||
raise InvalidJobRowError("claimed job is not running")
|
||||
|
||||
payload_value = row.get("payload")
|
||||
if not isinstance(payload_value, dict) or not all(
|
||||
isinstance(key, str) for key in payload_value
|
||||
):
|
||||
raise InvalidJobRowError("job row has invalid payload")
|
||||
|
||||
return BackgroundJob(
|
||||
id=job_id,
|
||||
job_type=_require_text(row, "job_type"),
|
||||
required_capability=_require_text(row, "required_capability"),
|
||||
resource_type=_require_text(row, "resource_type"),
|
||||
resource_id=_require_uuid(row, "resource_id"),
|
||||
idempotency_key=_require_text(row, "idempotency_key"),
|
||||
payload=dict(payload_value),
|
||||
stage=_require_text(row, "stage"),
|
||||
progress=_require_integer(row, "progress"),
|
||||
priority=_require_integer(row, "priority"),
|
||||
attempt=_require_integer(row, "attempt"),
|
||||
max_attempts=_require_integer(row, "max_attempts"),
|
||||
run_after=_require_datetime(row, "run_after"),
|
||||
lease_until=_require_datetime(row, "lease_until"),
|
||||
created_at=_require_datetime(row, "created_at"),
|
||||
updated_at=_require_datetime(row, "updated_at"),
|
||||
lease=JobLease(
|
||||
job_id=job_id,
|
||||
worker_id=lease_owner,
|
||||
lease_token=lease_token,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _state_from_row(row: JobRow) -> JobState:
|
||||
finished_at = row.get("finished_at")
|
||||
if finished_at is not None and not isinstance(finished_at, datetime):
|
||||
raise InvalidJobRowError("job row has invalid finished_at")
|
||||
return JobState(
|
||||
job_id=_require_uuid(row, "id"),
|
||||
status=_require_text(row, "status"),
|
||||
attempt=_require_integer(row, "attempt"),
|
||||
max_attempts=_require_integer(row, "max_attempts"),
|
||||
finished_at=finished_at,
|
||||
)
|
||||
|
||||
|
||||
def _validate_worker_id(worker_id: str) -> str:
|
||||
normalized = worker_id.strip()
|
||||
if not normalized or len(normalized) > 200:
|
||||
raise ValueError("worker_id must contain 1 to 200 characters")
|
||||
return normalized
|
||||
|
||||
|
||||
def _validate_lease_seconds(lease_seconds: int) -> int:
|
||||
if isinstance(lease_seconds, bool) or not 1 <= lease_seconds <= 86_400:
|
||||
raise ValueError("lease_seconds must be between 1 and 86400")
|
||||
return lease_seconds
|
||||
|
||||
|
||||
class PsycopgJobQueue:
|
||||
"""Short-transaction repository around the canonical queue statements."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dsn: str,
|
||||
*,
|
||||
connect_timeout: int = 5,
|
||||
connection_factory: ConnectionFactory = _default_connection_factory,
|
||||
) -> None:
|
||||
if not dsn.strip():
|
||||
raise ValueError("dsn must not be empty")
|
||||
if isinstance(connect_timeout, bool) or not 1 <= connect_timeout <= 60:
|
||||
raise ValueError("connect_timeout must be between 1 and 60")
|
||||
self._dsn = dsn
|
||||
self._connect_timeout = connect_timeout
|
||||
self._connection_factory = connection_factory
|
||||
|
||||
def _fetch_one(self, statement: str, parameters: Mapping[str, object]) -> JobRow | None:
|
||||
with self._connection_factory(self._dsn, self._connect_timeout) as connection:
|
||||
with connection.transaction():
|
||||
cursor = connection.execute(statement, parameters)
|
||||
return cursor.fetchone()
|
||||
|
||||
def _fetch_all(self, statement: str, parameters: Mapping[str, object]) -> list[JobRow]:
|
||||
with self._connection_factory(self._dsn, self._connect_timeout) as connection:
|
||||
with connection.transaction():
|
||||
cursor = connection.execute(statement, parameters)
|
||||
return list(cursor.fetchall())
|
||||
|
||||
def claim(
|
||||
self,
|
||||
*,
|
||||
worker_id: str,
|
||||
worker_capabilities: Sequence[str],
|
||||
lease_seconds: int,
|
||||
) -> BackgroundJob | None:
|
||||
owner = _validate_worker_id(worker_id)
|
||||
capabilities = tuple(
|
||||
capability.strip() for capability in worker_capabilities if capability.strip()
|
||||
)
|
||||
if not capabilities:
|
||||
raise ValueError("worker_capabilities must not be empty")
|
||||
if len(capabilities) != len(set(capabilities)):
|
||||
raise ValueError("worker_capabilities must not contain duplicates")
|
||||
|
||||
row = self._fetch_one(
|
||||
_CLAIM,
|
||||
{
|
||||
"worker_id": owner,
|
||||
"worker_capabilities": list(capabilities),
|
||||
"lease_seconds": _validate_lease_seconds(lease_seconds),
|
||||
},
|
||||
)
|
||||
if row is None:
|
||||
return None
|
||||
return _job_from_row(row, owner)
|
||||
|
||||
def heartbeat(self, lease: JobLease, *, lease_seconds: int) -> LeaseHeartbeat:
|
||||
row = self._fetch_one(
|
||||
_HEARTBEAT,
|
||||
{
|
||||
"job_id": lease.job_id,
|
||||
"worker_id": _validate_worker_id(lease.worker_id),
|
||||
"lease_token": lease.lease_token,
|
||||
"lease_seconds": _validate_lease_seconds(lease_seconds),
|
||||
},
|
||||
)
|
||||
if row is None:
|
||||
raise LeaseLostError("job lease is no longer owned")
|
||||
return LeaseHeartbeat(
|
||||
job_id=_require_uuid(row, "id"),
|
||||
lease_until=_require_datetime(row, "lease_until"),
|
||||
)
|
||||
|
||||
def complete(self, lease: JobLease) -> JobState:
|
||||
row = self._fetch_one(_COMPLETE, self._lease_parameters(lease))
|
||||
if row is None:
|
||||
raise LeaseLostError("job lease is no longer owned")
|
||||
return _state_from_row(row)
|
||||
|
||||
def fail_or_retry(
|
||||
self,
|
||||
lease: JobLease,
|
||||
*,
|
||||
error_code: str,
|
||||
error_message: str,
|
||||
retry_delay_seconds: int,
|
||||
) -> JobState:
|
||||
if not _ERROR_CODE.fullmatch(error_code):
|
||||
raise ValueError("error_code must be a stable uppercase identifier")
|
||||
if not error_message.strip():
|
||||
raise ValueError("error_message must not be empty")
|
||||
if isinstance(retry_delay_seconds, bool) or not 0 <= retry_delay_seconds <= 86_400:
|
||||
raise ValueError("retry_delay_seconds must be between 0 and 86400")
|
||||
|
||||
parameters = self._lease_parameters(lease)
|
||||
parameters.update(
|
||||
{
|
||||
"error_code": error_code,
|
||||
"error_message": error_message[:2000],
|
||||
"retry_delay_seconds": retry_delay_seconds,
|
||||
}
|
||||
)
|
||||
row = self._fetch_one(_FAIL_OR_RETRY, parameters)
|
||||
if row is None:
|
||||
raise LeaseLostError("job lease is no longer owned")
|
||||
return _state_from_row(row)
|
||||
|
||||
def reap_expired(
|
||||
self,
|
||||
*,
|
||||
lock_key: int,
|
||||
batch_size: int = 100,
|
||||
) -> tuple[JobState, ...]:
|
||||
if isinstance(lock_key, bool) or not -(2**63) <= lock_key < 2**63:
|
||||
raise ValueError("lock_key must be a signed 64-bit integer")
|
||||
if isinstance(batch_size, bool) or not 1 <= batch_size <= 1000:
|
||||
raise ValueError("batch_size must be between 1 and 1000")
|
||||
rows = self._fetch_all(
|
||||
_REAP_EXPIRED,
|
||||
{"lock_key": lock_key, "batch_size": batch_size},
|
||||
)
|
||||
return tuple(_state_from_row(row) for row in rows)
|
||||
|
||||
@staticmethod
|
||||
def _lease_parameters(lease: JobLease) -> dict[str, object]:
|
||||
return {
|
||||
"job_id": lease.job_id,
|
||||
"worker_id": _validate_worker_id(lease.worker_id),
|
||||
"lease_token": lease.lease_token,
|
||||
}
|
||||
@@ -40,6 +40,7 @@ WHERE job.id = :job_id
|
||||
AND job.status = 'RUNNING'
|
||||
AND job.lease_owner = :worker_id
|
||||
AND job.lease_token = :lease_token
|
||||
AND job.lease_until >= now()
|
||||
RETURNING job.id, job.lease_until
|
||||
"""
|
||||
|
||||
@@ -56,6 +57,7 @@ WHERE job.id = :job_id
|
||||
AND job.status = 'RUNNING'
|
||||
AND job.lease_owner = :worker_id
|
||||
AND job.lease_token = :lease_token
|
||||
AND job.lease_until >= now()
|
||||
RETURNING job.*
|
||||
"""
|
||||
|
||||
@@ -86,6 +88,7 @@ WHERE job.id = :job_id
|
||||
AND job.status = 'RUNNING'
|
||||
AND job.lease_owner = :worker_id
|
||||
AND job.lease_token = :lease_token
|
||||
AND job.lease_until >= now()
|
||||
RETURNING job.*
|
||||
"""
|
||||
|
||||
|
||||
313
backend/app/persistence/retrieval.py
Normal file
313
backend/app/persistence/retrieval.py
Normal file
@@ -0,0 +1,313 @@
|
||||
"""PostgreSQL/pgvector persistence boundary for formal retrieval.
|
||||
|
||||
Authorization, active-model selection, and document lifecycle checks belong in
|
||||
the candidate SQL. They must never be applied as an in-memory post-filter.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
import re
|
||||
import uuid
|
||||
from collections.abc import Sequence
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Protocol, cast
|
||||
|
||||
import psycopg
|
||||
from pgvector.psycopg import register_vector
|
||||
from pgvector.vector import Vector
|
||||
from psycopg.rows import dict_row
|
||||
|
||||
from app.core.config import Settings
|
||||
from app.core.secrets import SecretFileError
|
||||
|
||||
_PROFILE_HASH_PATTERN = re.compile(r"^[0-9a-f]{64}$")
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ActiveEmbeddingProfile:
|
||||
"""The enabled embedding profile selected by a knowledge base."""
|
||||
|
||||
profile_hash: str
|
||||
model: str
|
||||
dimension: int
|
||||
synthetic: bool = False
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class RetrievalCandidate:
|
||||
"""Approved and authorized projection returned by the vector query."""
|
||||
|
||||
citation_id: uuid.UUID
|
||||
document_id: uuid.UUID
|
||||
source_name: str
|
||||
cloud_text: str
|
||||
section_path: tuple[str, ...]
|
||||
page_start: int | None
|
||||
page_end: int | None
|
||||
vector_score: float
|
||||
|
||||
|
||||
class RetrievalPersistenceError(RuntimeError):
|
||||
"""Sanitized storage failure safe for translation at the service boundary."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
super().__init__("retrieval persistence unavailable")
|
||||
|
||||
|
||||
class RetrievalRepository(Protocol):
|
||||
"""Synchronous repository port; the async service runs calls in a thread."""
|
||||
|
||||
def resolve_active_profile(
|
||||
self,
|
||||
knowledge_base_id: uuid.UUID,
|
||||
*,
|
||||
allowed_scope_ids: Sequence[uuid.UUID],
|
||||
) -> ActiveEmbeddingProfile | None: ...
|
||||
|
||||
def search_candidates(
|
||||
self,
|
||||
knowledge_base_id: uuid.UUID,
|
||||
*,
|
||||
allowed_scope_ids: Sequence[uuid.UUID],
|
||||
profile_hash: str,
|
||||
query_vector: tuple[float, ...],
|
||||
limit: int,
|
||||
) -> list[RetrievalCandidate]: ...
|
||||
|
||||
|
||||
ACTIVE_PROFILE_SQL = """
|
||||
SELECT
|
||||
profile.profile_hash,
|
||||
profile.model,
|
||||
profile.dimension,
|
||||
profile.synthetic
|
||||
FROM rag.knowledge_bases AS knowledge_base
|
||||
JOIN rag.model_profiles AS profile
|
||||
ON profile.profile_hash = knowledge_base.active_embedding_profile_hash
|
||||
AND profile.kind = knowledge_base.active_embedding_profile_kind
|
||||
WHERE knowledge_base.id = %s
|
||||
AND profile.kind = 'embedding'
|
||||
AND profile.enabled IS TRUE
|
||||
AND profile.dimension = 1024
|
||||
AND EXISTS (
|
||||
SELECT 1
|
||||
FROM rag.access_scopes AS access_scope
|
||||
WHERE access_scope.knowledge_base_id = knowledge_base.id
|
||||
AND access_scope.id = ANY(%s::uuid[])
|
||||
)
|
||||
LIMIT 1
|
||||
"""
|
||||
|
||||
|
||||
CANDIDATE_SEARCH_SQL = """
|
||||
WITH query_input AS (
|
||||
SELECT %s::vector AS embedding
|
||||
)
|
||||
SELECT
|
||||
chunk.citation_id::text AS citation_id,
|
||||
document.id::text AS document_id,
|
||||
document.filename AS source_name,
|
||||
chunk.cloud_text,
|
||||
chunk.section_path,
|
||||
chunk.page_start,
|
||||
chunk.page_end,
|
||||
1 - (chunk.embedding <=> query_input.embedding) AS vector_score
|
||||
FROM rag.chunks AS chunk
|
||||
JOIN rag.knowledge_bases AS knowledge_base
|
||||
ON knowledge_base.id = chunk.knowledge_base_id
|
||||
JOIN rag.model_profiles AS profile
|
||||
ON profile.profile_hash = knowledge_base.active_embedding_profile_hash
|
||||
AND profile.kind = knowledge_base.active_embedding_profile_kind
|
||||
JOIN rag.access_scopes AS access_scope
|
||||
ON access_scope.id = chunk.access_scope_id
|
||||
AND access_scope.knowledge_base_id = chunk.knowledge_base_id
|
||||
JOIN rag.documents AS document
|
||||
ON document.id = chunk.document_id
|
||||
AND document.knowledge_base_id = chunk.knowledge_base_id
|
||||
AND document.access_scope_id = chunk.access_scope_id
|
||||
JOIN rag.document_versions AS document_version
|
||||
ON document_version.id = chunk.document_version_id
|
||||
AND document_version.document_id = chunk.document_id
|
||||
CROSS JOIN query_input
|
||||
WHERE chunk.knowledge_base_id = %s
|
||||
AND chunk.access_scope_id = ANY(%s::uuid[])
|
||||
AND knowledge_base.active_embedding_profile_hash = %s
|
||||
AND knowledge_base.active_embedding_profile_kind = 'embedding'
|
||||
AND profile.kind = 'embedding'
|
||||
AND profile.enabled IS TRUE
|
||||
AND profile.dimension = 1024
|
||||
AND chunk.embedding_profile_hash = knowledge_base.active_embedding_profile_hash
|
||||
AND chunk.embedding_model = profile.model
|
||||
AND chunk.embedding_dimension = profile.dimension
|
||||
AND chunk.embedding IS NOT NULL
|
||||
AND chunk.embedded_text_sha256 = chunk.embedding_text_sha256
|
||||
AND chunk.searchable IS TRUE
|
||||
AND chunk.index_status = 'READY'
|
||||
AND chunk.approval_status = 'CLOUD_APPROVED'
|
||||
AND chunk.deleted_at IS NULL
|
||||
AND document.status = 'READY'
|
||||
AND document.deleted_at IS NULL
|
||||
AND document.active_version_id = chunk.document_version_id
|
||||
AND document_version.status = 'READY'
|
||||
AND document_version.review_state = 'CLOUD_APPROVED'
|
||||
AND document_version.embedding_profile_hash = knowledge_base.active_embedding_profile_hash
|
||||
AND document_version.outbound_manifest_sha256 = chunk.outbound_manifest_sha256
|
||||
ORDER BY chunk.embedding <=> query_input.embedding, chunk.citation_id
|
||||
LIMIT %s
|
||||
"""
|
||||
|
||||
|
||||
class PostgresRetrievalRepository:
|
||||
"""Read-only PostgreSQL implementation with filtered HNSW candidate search."""
|
||||
|
||||
def __init__(self, settings: Settings) -> None:
|
||||
self._settings = settings
|
||||
|
||||
def _dsn(self) -> str:
|
||||
return (
|
||||
self._settings.database_url()
|
||||
.set(drivername="postgresql")
|
||||
.render_as_string(hide_password=False)
|
||||
)
|
||||
|
||||
def resolve_active_profile(
|
||||
self,
|
||||
knowledge_base_id: uuid.UUID,
|
||||
*,
|
||||
allowed_scope_ids: Sequence[uuid.UUID],
|
||||
) -> ActiveEmbeddingProfile | None:
|
||||
if not allowed_scope_ids:
|
||||
return None
|
||||
try:
|
||||
with psycopg.connect(
|
||||
self._dsn(),
|
||||
connect_timeout=2,
|
||||
row_factory=dict_row,
|
||||
) as connection:
|
||||
connection.execute("SET LOCAL statement_timeout = '3000ms'")
|
||||
row = connection.execute(
|
||||
ACTIVE_PROFILE_SQL,
|
||||
(knowledge_base_id, list(allowed_scope_ids)),
|
||||
).fetchone()
|
||||
except (OSError, SecretFileError, psycopg.Error) as exc:
|
||||
raise RetrievalPersistenceError from exc
|
||||
|
||||
if row is None:
|
||||
return None
|
||||
try:
|
||||
profile_hash = cast(str, row["profile_hash"])
|
||||
model = cast(str, row["model"])
|
||||
dimension = cast(int, row["dimension"])
|
||||
synthetic = cast(bool, row["synthetic"])
|
||||
if (
|
||||
_PROFILE_HASH_PATTERN.fullmatch(profile_hash) is None
|
||||
or not model
|
||||
or model != model.strip()
|
||||
or isinstance(dimension, bool)
|
||||
or dimension != 1024
|
||||
or not isinstance(synthetic, bool)
|
||||
):
|
||||
raise ValueError
|
||||
except (KeyError, TypeError, ValueError) as exc:
|
||||
raise RetrievalPersistenceError from exc
|
||||
return ActiveEmbeddingProfile(
|
||||
profile_hash=profile_hash,
|
||||
model=model,
|
||||
dimension=dimension,
|
||||
synthetic=synthetic,
|
||||
)
|
||||
|
||||
def search_candidates(
|
||||
self,
|
||||
knowledge_base_id: uuid.UUID,
|
||||
*,
|
||||
allowed_scope_ids: Sequence[uuid.UUID],
|
||||
profile_hash: str,
|
||||
query_vector: tuple[float, ...],
|
||||
limit: int,
|
||||
) -> list[RetrievalCandidate]:
|
||||
if (
|
||||
not allowed_scope_ids
|
||||
or _PROFILE_HASH_PATTERN.fullmatch(profile_hash) is None
|
||||
or len(query_vector) != 1024
|
||||
or isinstance(limit, bool)
|
||||
or not 1 <= limit <= 50
|
||||
):
|
||||
raise RetrievalPersistenceError
|
||||
|
||||
try:
|
||||
vector = Vector(list(query_vector))
|
||||
with psycopg.connect(
|
||||
self._dsn(),
|
||||
connect_timeout=2,
|
||||
row_factory=dict_row,
|
||||
) as connection:
|
||||
register_vector(connection)
|
||||
connection.execute("SET LOCAL statement_timeout = '3000ms'")
|
||||
connection.execute("SET LOCAL hnsw.iterative_scan = strict_order")
|
||||
connection.execute("SET LOCAL hnsw.ef_search = 100")
|
||||
rows = connection.execute(
|
||||
CANDIDATE_SEARCH_SQL,
|
||||
(
|
||||
vector,
|
||||
knowledge_base_id,
|
||||
list(allowed_scope_ids),
|
||||
profile_hash,
|
||||
limit,
|
||||
),
|
||||
).fetchall()
|
||||
except (OSError, SecretFileError, psycopg.Error) as exc:
|
||||
raise RetrievalPersistenceError from exc
|
||||
|
||||
try:
|
||||
return [self._candidate(row) for row in rows]
|
||||
except (KeyError, TypeError, ValueError) as exc:
|
||||
raise RetrievalPersistenceError from exc
|
||||
|
||||
@staticmethod
|
||||
def _candidate(row: dict[str, Any]) -> RetrievalCandidate:
|
||||
raw_section_path = row["section_path"]
|
||||
if not isinstance(raw_section_path, list) or any(
|
||||
not isinstance(part, str) or not part.strip() for part in raw_section_path
|
||||
):
|
||||
raise ValueError
|
||||
source_name = row["source_name"]
|
||||
cloud_text = row["cloud_text"]
|
||||
raw_score = row["vector_score"]
|
||||
if (
|
||||
not isinstance(source_name, str)
|
||||
or not source_name.strip()
|
||||
or not isinstance(cloud_text, str)
|
||||
or not cloud_text.strip()
|
||||
or isinstance(raw_score, bool)
|
||||
or not isinstance(raw_score, (int, float))
|
||||
or not math.isfinite(float(raw_score))
|
||||
):
|
||||
raise ValueError
|
||||
|
||||
page_start = row["page_start"]
|
||||
page_end = row["page_end"]
|
||||
if (page_start is None) != (page_end is None) or (
|
||||
page_start is not None
|
||||
and (
|
||||
isinstance(page_start, bool)
|
||||
or isinstance(page_end, bool)
|
||||
or not isinstance(page_start, int)
|
||||
or not isinstance(page_end, int)
|
||||
or page_start < 1
|
||||
or page_end < page_start
|
||||
)
|
||||
):
|
||||
raise ValueError
|
||||
|
||||
return RetrievalCandidate(
|
||||
citation_id=uuid.UUID(cast(str, row["citation_id"])),
|
||||
document_id=uuid.UUID(cast(str, row["document_id"])),
|
||||
source_name=source_name.strip(),
|
||||
cloud_text=cloud_text.strip(),
|
||||
section_path=tuple(part.strip() for part in raw_section_path),
|
||||
page_start=cast(int | None, page_start),
|
||||
page_end=cast(int | None, page_end),
|
||||
vector_score=float(raw_score),
|
||||
)
|
||||
Reference in New Issue
Block a user