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:
@@ -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")),
|
||||
)
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -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.*
|
||||
"""
|
||||
|
||||
|
||||
@@ -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