YoVinchen
ecdb10c37a
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
1205 lines
42 KiB
Python
1205 lines
42 KiB
Python
"""Short-transaction PostgreSQL adapter for lease-fenced document indexing."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import math
|
|
import re
|
|
import uuid
|
|
from collections.abc import Callable, Iterator, Mapping, Sequence
|
|
from contextlib import contextmanager
|
|
from typing import Any, cast
|
|
|
|
import psycopg
|
|
from pgvector.psycopg import register_vector
|
|
from pgvector.vector import Vector
|
|
from psycopg import Connection
|
|
from psycopg.rows import dict_row
|
|
from psycopg.types.json import Jsonb
|
|
|
|
from app.core.config import Settings
|
|
from app.core.secrets import SecretFileError
|
|
from app.persistence.job_queue import JobLease, LeaseLostError
|
|
from app.persistence.retrieval import ActiveEmbeddingProfile
|
|
from app.ports.model_providers import ProviderUsage
|
|
from app.services.indexing import (
|
|
ApprovedIndexingPlan,
|
|
AssignmentProgress,
|
|
AssignmentStatus,
|
|
CachedEmbedding,
|
|
EmbeddingCacheLookup,
|
|
EmbeddingWrite,
|
|
IndexingItem,
|
|
InvocationStatus,
|
|
embedding_cache_key,
|
|
)
|
|
|
|
type IndexingRow = dict[str, Any]
|
|
type ConnectionFactory = Callable[[str, int], Connection[IndexingRow]]
|
|
type VectorRegistrar = Callable[[Connection[Any]], None]
|
|
|
|
_HASH = re.compile(r"^[0-9a-f]{64}$")
|
|
_SAFE_IDENTIFIER = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._:/-]{0,511}$")
|
|
_ERROR_CODE = re.compile(r"^[A-Z][A-Z0-9_]{0,127}$")
|
|
|
|
|
|
class IndexingPersistenceError(RuntimeError):
|
|
"""Sanitized infrastructure failure with no SQL, DSN, text, or credential data."""
|
|
|
|
def __init__(self) -> None:
|
|
super().__init__("indexing persistence unavailable")
|
|
|
|
|
|
class IndexingPersistenceConflict(RuntimeError):
|
|
"""A governed profile, manifest, invocation, or count contract did not match."""
|
|
|
|
def __init__(self) -> None:
|
|
super().__init__("indexing state does not satisfy the governed contract")
|
|
|
|
|
|
LEASE_FENCE_SQL = """
|
|
SELECT job.resource_id
|
|
FROM rag.background_jobs AS job
|
|
WHERE job.id = %s
|
|
AND job.status = 'RUNNING'
|
|
AND job.job_type = 'EMBED_DOCUMENT'
|
|
AND job.required_capability = 'embedding'
|
|
AND job.resource_type = 'document_version'
|
|
AND job.lease_owner = %s
|
|
AND job.lease_token = %s
|
|
AND job.lease_until >= now()
|
|
FOR UPDATE
|
|
"""
|
|
|
|
LOAD_PLAN_SQL = """
|
|
SELECT
|
|
document.knowledge_base_id,
|
|
version.id AS document_version_id,
|
|
version.review_state,
|
|
version.outbound_manifest_sha256,
|
|
version.expected_chunk_count,
|
|
profile.profile_hash,
|
|
profile.model,
|
|
profile.dimension,
|
|
profile.synthetic,
|
|
(
|
|
SELECT count(*)::integer
|
|
FROM rag.outbound_manifest_items AS item
|
|
WHERE item.document_version_id = version.id
|
|
AND item.outbound_manifest_sha256 = version.outbound_manifest_sha256
|
|
) AS manifest_count,
|
|
(
|
|
SELECT count(*)::integer
|
|
FROM rag.chunks AS chunk
|
|
JOIN rag.outbound_manifest_items AS item
|
|
ON item.document_version_id = chunk.document_version_id
|
|
AND item.ordinal = chunk.ordinal
|
|
AND item.outbound_manifest_sha256 = chunk.outbound_manifest_sha256
|
|
AND item.cloud_text_sha256 = chunk.cloud_text_sha256
|
|
AND item.embedding_text_sha256 = chunk.embedding_text_sha256
|
|
WHERE chunk.document_version_id = version.id
|
|
AND chunk.approval_status = 'CLOUD_APPROVED'
|
|
AND chunk.outbound_manifest_sha256 = version.outbound_manifest_sha256
|
|
AND chunk.embedding_profile_hash = profile.profile_hash
|
|
AND chunk.embedding_model = profile.model
|
|
AND chunk.embedding_dimension = 1024
|
|
AND chunk.deleted_at IS NULL
|
|
) AS eligible_chunk_count
|
|
FROM rag.document_versions AS version
|
|
JOIN rag.documents AS document
|
|
ON document.id = version.document_id
|
|
AND document.deleted_at IS NULL
|
|
JOIN rag.knowledge_bases AS knowledge_base
|
|
ON knowledge_base.id = document.knowledge_base_id
|
|
JOIN rag.model_profiles AS profile
|
|
ON profile.profile_hash = version.embedding_profile_hash
|
|
AND profile.kind = 'embedding'
|
|
AND profile.enabled IS TRUE
|
|
AND profile.dimension = 1024
|
|
WHERE version.id = %s
|
|
AND version.review_state = 'CLOUD_APPROVED'
|
|
AND version.status IN ('PENDING', 'PROCESSING', 'READY')
|
|
AND version.outbound_manifest_sha256 IS NOT NULL
|
|
AND version.expected_chunk_count IS NOT NULL
|
|
AND knowledge_base.active_embedding_profile_kind = 'embedding'
|
|
AND knowledge_base.active_embedding_profile_hash = profile.profile_hash
|
|
AND document.status IN (
|
|
'CLOUD_APPROVED', 'EMBEDDING', 'INDEXING', 'VALIDATING', 'READY'
|
|
)
|
|
FOR SHARE OF version, document, knowledge_base, profile
|
|
"""
|
|
|
|
LOAD_PLAN_ITEMS_SQL = """
|
|
SELECT
|
|
chunk.id AS chunk_id,
|
|
chunk.ordinal,
|
|
chunk.embedding_text,
|
|
chunk.embedding_text_sha256,
|
|
CASE
|
|
WHEN assignment.status = 'READY'
|
|
AND assignment.embedding_text_sha256 = chunk.embedding_text_sha256
|
|
AND assignment.cache_profile_hash = profile.profile_hash
|
|
AND assignment.cache_embedding_text_sha256 = chunk.embedding_text_sha256
|
|
AND cache.resolved_model = profile.model
|
|
AND vector_dims(cache.embedding) = 1024
|
|
AND vector_norm(cache.embedding) > 0
|
|
AND chunk.index_status = 'READY'
|
|
AND chunk.embedded_text_sha256 = chunk.embedding_text_sha256
|
|
AND chunk.embedding = cache.embedding
|
|
THEN 'READY'
|
|
WHEN assignment.status = 'READY' THEN 'STALE'
|
|
ELSE COALESCE(assignment.status, 'PENDING')
|
|
END AS assignment_status
|
|
FROM rag.chunks AS chunk
|
|
JOIN rag.document_versions AS version
|
|
ON version.id = chunk.document_version_id
|
|
JOIN rag.documents AS document
|
|
ON document.id = chunk.document_id
|
|
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 = version.embedding_profile_hash
|
|
AND profile.kind = 'embedding'
|
|
AND profile.enabled IS TRUE
|
|
AND profile.dimension = 1024
|
|
JOIN rag.outbound_manifest_items AS item
|
|
ON item.document_version_id = chunk.document_version_id
|
|
AND item.ordinal = chunk.ordinal
|
|
AND item.outbound_manifest_sha256 = chunk.outbound_manifest_sha256
|
|
AND item.cloud_text_sha256 = chunk.cloud_text_sha256
|
|
AND item.embedding_text_sha256 = chunk.embedding_text_sha256
|
|
LEFT JOIN rag.chunk_embedding_assignments AS assignment
|
|
ON assignment.chunk_id = chunk.id
|
|
AND assignment.profile_hash = profile.profile_hash
|
|
LEFT JOIN rag.embedding_cache AS cache
|
|
ON cache.profile_hash = assignment.cache_profile_hash
|
|
AND cache.embedding_text_sha256 = assignment.cache_embedding_text_sha256
|
|
WHERE version.id = %s
|
|
AND version.review_state = 'CLOUD_APPROVED'
|
|
AND chunk.approval_status = 'CLOUD_APPROVED'
|
|
AND chunk.outbound_manifest_sha256 = version.outbound_manifest_sha256
|
|
AND chunk.embedding_profile_hash = profile.profile_hash
|
|
AND chunk.embedding_model = profile.model
|
|
AND chunk.embedding_dimension = 1024
|
|
AND chunk.deleted_at IS NULL
|
|
AND document.deleted_at IS NULL
|
|
AND knowledge_base.active_embedding_profile_hash = profile.profile_hash
|
|
AND knowledge_base.active_embedding_profile_kind = 'embedding'
|
|
ORDER BY chunk.ordinal, chunk.id
|
|
FOR SHARE OF chunk
|
|
"""
|
|
|
|
CACHE_LOOKUP_SQL = """
|
|
SELECT
|
|
cache.profile_hash,
|
|
cache.embedding_text_sha256,
|
|
cache.resolved_model,
|
|
vector_dims(cache.embedding)::integer AS dimension
|
|
FROM rag.embedding_cache AS cache
|
|
JOIN rag.model_profiles AS profile
|
|
ON profile.profile_hash = cache.profile_hash
|
|
AND profile.kind = 'embedding'
|
|
AND profile.enabled IS TRUE
|
|
AND profile.dimension = 1024
|
|
JOIN rag.document_versions AS version
|
|
ON version.id = %s
|
|
AND version.review_state = 'CLOUD_APPROVED'
|
|
AND version.embedding_profile_hash = profile.profile_hash
|
|
JOIN rag.documents AS document
|
|
ON document.id = version.document_id
|
|
AND document.deleted_at IS NULL
|
|
JOIN rag.knowledge_bases AS knowledge_base
|
|
ON knowledge_base.id = document.knowledge_base_id
|
|
AND knowledge_base.active_embedding_profile_hash = profile.profile_hash
|
|
AND knowledge_base.active_embedding_profile_kind = 'embedding'
|
|
WHERE cache.profile_hash = %s
|
|
AND cache.embedding_text_sha256 = %s
|
|
AND cache.resolved_model = profile.model
|
|
AND vector_dims(cache.embedding) = 1024
|
|
AND vector_norm(cache.embedding) > 0
|
|
FOR SHARE OF cache
|
|
"""
|
|
|
|
BEGIN_INVOCATION_SQL = """
|
|
INSERT INTO rag.model_invocations (
|
|
trace_id,
|
|
caller,
|
|
operation,
|
|
profile_hash,
|
|
model,
|
|
status,
|
|
item_count
|
|
)
|
|
SELECT
|
|
%s,
|
|
'worker',
|
|
'embedding',
|
|
profile.profile_hash,
|
|
profile.model,
|
|
'STARTED',
|
|
%s
|
|
FROM rag.document_versions AS version
|
|
JOIN rag.documents AS document
|
|
ON document.id = version.document_id
|
|
AND document.deleted_at IS NULL
|
|
JOIN rag.knowledge_bases AS knowledge_base
|
|
ON knowledge_base.id = document.knowledge_base_id
|
|
JOIN rag.model_profiles AS profile
|
|
ON profile.profile_hash = version.embedding_profile_hash
|
|
AND profile.kind = 'embedding'
|
|
AND profile.enabled IS TRUE
|
|
AND profile.dimension = 1024
|
|
WHERE version.id = %s
|
|
AND version.review_state = 'CLOUD_APPROVED'
|
|
AND profile.profile_hash = %s
|
|
AND profile.model = %s
|
|
AND knowledge_base.active_embedding_profile_hash = profile.profile_hash
|
|
AND knowledge_base.active_embedding_profile_kind = 'embedding'
|
|
RETURNING id
|
|
"""
|
|
|
|
FINISH_INVOCATION_SQL = """
|
|
UPDATE rag.model_invocations AS invocation
|
|
SET status = %s,
|
|
provider_request_id = %s,
|
|
prompt_tokens = %s,
|
|
completion_tokens = %s,
|
|
total_tokens = %s,
|
|
elapsed_ms = %s,
|
|
error_code = %s,
|
|
finished_at = now()
|
|
FROM rag.document_versions AS version
|
|
JOIN rag.documents AS document
|
|
ON document.id = version.document_id
|
|
AND document.deleted_at IS NULL
|
|
JOIN rag.knowledge_bases AS knowledge_base
|
|
ON knowledge_base.id = document.knowledge_base_id
|
|
JOIN rag.model_profiles AS profile
|
|
ON profile.profile_hash = version.embedding_profile_hash
|
|
AND profile.kind = 'embedding'
|
|
AND profile.enabled IS TRUE
|
|
WHERE invocation.id = %s
|
|
AND invocation.trace_id = %s
|
|
AND invocation.operation = 'embedding'
|
|
AND invocation.status = 'STARTED'
|
|
AND invocation.profile_hash = profile.profile_hash
|
|
AND invocation.model = profile.model
|
|
AND version.id = %s
|
|
AND version.review_state = 'CLOUD_APPROVED'
|
|
AND knowledge_base.active_embedding_profile_hash = profile.profile_hash
|
|
RETURNING invocation.id
|
|
"""
|
|
|
|
INSERT_CACHE_SQL = """
|
|
INSERT INTO rag.embedding_cache (
|
|
profile_hash,
|
|
embedding_text_sha256,
|
|
embedding,
|
|
resolved_model,
|
|
provider_request_id,
|
|
usage,
|
|
elapsed_ms
|
|
)
|
|
SELECT %s, %s, %s, %s, %s, %s, %s
|
|
FROM rag.document_versions AS version
|
|
JOIN rag.documents AS document
|
|
ON document.id = version.document_id
|
|
AND document.deleted_at IS NULL
|
|
JOIN rag.knowledge_bases AS knowledge_base
|
|
ON knowledge_base.id = document.knowledge_base_id
|
|
JOIN rag.model_profiles AS profile
|
|
ON profile.profile_hash = version.embedding_profile_hash
|
|
AND profile.kind = 'embedding'
|
|
AND profile.enabled IS TRUE
|
|
AND profile.dimension = 1024
|
|
WHERE version.id = %s
|
|
AND version.review_state = 'CLOUD_APPROVED'
|
|
AND version.embedding_profile_hash = %s
|
|
AND profile.model = %s
|
|
AND knowledge_base.active_embedding_profile_hash = profile.profile_hash
|
|
AND vector_dims(%s::vector) = 1024
|
|
AND vector_norm(%s::vector) > 0
|
|
ON CONFLICT (profile_hash, embedding_text_sha256) DO NOTHING
|
|
RETURNING profile_hash
|
|
"""
|
|
|
|
VERIFY_CACHE_FOR_WRITE_SQL = """
|
|
SELECT 1 AS available
|
|
FROM rag.embedding_cache AS cache
|
|
JOIN rag.model_profiles AS profile
|
|
ON profile.profile_hash = cache.profile_hash
|
|
AND profile.kind = 'embedding'
|
|
AND profile.enabled IS TRUE
|
|
AND profile.dimension = 1024
|
|
JOIN rag.document_versions AS version
|
|
ON version.id = %s
|
|
AND version.review_state = 'CLOUD_APPROVED'
|
|
AND version.embedding_profile_hash = profile.profile_hash
|
|
JOIN rag.documents AS document
|
|
ON document.id = version.document_id
|
|
AND document.deleted_at IS NULL
|
|
JOIN rag.knowledge_bases AS knowledge_base
|
|
ON knowledge_base.id = document.knowledge_base_id
|
|
AND knowledge_base.active_embedding_profile_hash = profile.profile_hash
|
|
WHERE cache.profile_hash = %s
|
|
AND cache.embedding_text_sha256 = %s
|
|
AND cache.resolved_model = %s
|
|
AND vector_dims(cache.embedding) = 1024
|
|
AND vector_norm(cache.embedding) > 0
|
|
"""
|
|
|
|
UPSERT_ASSIGNMENT_SQL = """
|
|
INSERT INTO rag.chunk_embedding_assignments (
|
|
chunk_id,
|
|
profile_hash,
|
|
embedding_text_sha256,
|
|
cache_profile_hash,
|
|
cache_embedding_text_sha256,
|
|
status,
|
|
error_code,
|
|
completed_at
|
|
)
|
|
SELECT
|
|
chunk.id,
|
|
profile.profile_hash,
|
|
chunk.embedding_text_sha256,
|
|
cache.profile_hash,
|
|
cache.embedding_text_sha256,
|
|
'READY',
|
|
NULL,
|
|
now()
|
|
FROM rag.chunks AS chunk
|
|
JOIN rag.document_versions AS version
|
|
ON version.id = chunk.document_version_id
|
|
JOIN rag.documents AS document
|
|
ON document.id = chunk.document_id
|
|
AND document.deleted_at IS NULL
|
|
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 = version.embedding_profile_hash
|
|
AND profile.kind = 'embedding'
|
|
AND profile.enabled IS TRUE
|
|
AND profile.dimension = 1024
|
|
JOIN rag.embedding_cache AS cache
|
|
ON cache.profile_hash = profile.profile_hash
|
|
AND cache.embedding_text_sha256 = chunk.embedding_text_sha256
|
|
AND cache.resolved_model = profile.model
|
|
WHERE chunk.id = %s
|
|
AND chunk.document_version_id = %s
|
|
AND chunk.approval_status = 'CLOUD_APPROVED'
|
|
AND chunk.outbound_manifest_sha256 = version.outbound_manifest_sha256
|
|
AND chunk.embedding_profile_hash = profile.profile_hash
|
|
AND chunk.embedding_model = profile.model
|
|
AND chunk.embedding_dimension = 1024
|
|
AND chunk.embedding_text_sha256 = %s
|
|
AND chunk.deleted_at IS NULL
|
|
AND knowledge_base.active_embedding_profile_hash = profile.profile_hash
|
|
AND profile.profile_hash = %s
|
|
AND vector_dims(cache.embedding) = 1024
|
|
AND vector_norm(cache.embedding) > 0
|
|
ON CONFLICT (chunk_id, profile_hash) DO UPDATE
|
|
SET cache_profile_hash = EXCLUDED.cache_profile_hash,
|
|
cache_embedding_text_sha256 = EXCLUDED.cache_embedding_text_sha256,
|
|
status = 'READY',
|
|
error_code = NULL,
|
|
completed_at = COALESCE(rag.chunk_embedding_assignments.completed_at, now()),
|
|
updated_at = now()
|
|
WHERE rag.chunk_embedding_assignments.embedding_text_sha256 = EXCLUDED.embedding_text_sha256
|
|
RETURNING chunk_id
|
|
"""
|
|
|
|
UPDATE_CHUNK_FROM_CACHE_SQL = """
|
|
UPDATE rag.chunks AS chunk
|
|
SET embedding = cache.embedding,
|
|
embedded_text_sha256 = chunk.embedding_text_sha256,
|
|
embedding_profile_hash = profile.profile_hash,
|
|
embedding_model = profile.model,
|
|
embedding_dimension = 1024,
|
|
index_status = 'READY',
|
|
searchable = false,
|
|
updated_at = now()
|
|
FROM rag.document_versions AS version
|
|
JOIN rag.documents AS document
|
|
ON document.id = version.document_id
|
|
AND document.deleted_at IS NULL
|
|
JOIN rag.knowledge_bases AS knowledge_base
|
|
ON knowledge_base.id = document.knowledge_base_id
|
|
JOIN rag.model_profiles AS profile
|
|
ON profile.profile_hash = version.embedding_profile_hash
|
|
AND profile.kind = 'embedding'
|
|
AND profile.enabled IS TRUE
|
|
AND profile.dimension = 1024
|
|
JOIN rag.embedding_cache AS cache
|
|
ON cache.profile_hash = profile.profile_hash
|
|
WHERE chunk.id = %s
|
|
AND chunk.document_version_id = version.id
|
|
AND version.id = %s
|
|
AND version.review_state = 'CLOUD_APPROVED'
|
|
AND chunk.approval_status = 'CLOUD_APPROVED'
|
|
AND chunk.outbound_manifest_sha256 = version.outbound_manifest_sha256
|
|
AND chunk.embedding_profile_hash = profile.profile_hash
|
|
AND chunk.embedding_text_sha256 = %s
|
|
AND cache.embedding_text_sha256 = chunk.embedding_text_sha256
|
|
AND cache.resolved_model = profile.model
|
|
AND knowledge_base.active_embedding_profile_hash = profile.profile_hash
|
|
AND chunk.deleted_at IS NULL
|
|
AND chunk.index_status <> 'READY'
|
|
AND vector_dims(cache.embedding) = 1024
|
|
AND vector_norm(cache.embedding) > 0
|
|
RETURNING chunk.id
|
|
"""
|
|
|
|
PREPARE_STALE_CHUNK_SQL = """
|
|
UPDATE rag.chunks AS chunk
|
|
SET searchable = false,
|
|
index_status = 'EMBEDDING',
|
|
updated_at = now()
|
|
FROM rag.embedding_cache AS cache
|
|
WHERE chunk.id = %s
|
|
AND chunk.document_version_id = %s
|
|
AND chunk.embedding_profile_hash = %s
|
|
AND chunk.embedding_text_sha256 = %s
|
|
AND cache.profile_hash = chunk.embedding_profile_hash
|
|
AND cache.embedding_text_sha256 = chunk.embedding_text_sha256
|
|
AND cache.resolved_model = chunk.embedding_model
|
|
AND chunk.index_status = 'READY'
|
|
AND (
|
|
chunk.embedded_text_sha256 IS DISTINCT FROM chunk.embedding_text_sha256
|
|
OR chunk.embedding IS DISTINCT FROM cache.embedding
|
|
OR vector_dims(chunk.embedding) IS DISTINCT FROM 1024
|
|
OR vector_norm(chunk.embedding) <= 0
|
|
)
|
|
RETURNING chunk.id
|
|
"""
|
|
|
|
VERIFY_READY_CHUNK_SQL = """
|
|
SELECT chunk.id
|
|
FROM rag.chunks AS chunk
|
|
JOIN rag.embedding_cache AS cache
|
|
ON cache.profile_hash = chunk.embedding_profile_hash
|
|
AND cache.embedding_text_sha256 = chunk.embedding_text_sha256
|
|
WHERE chunk.id = %s
|
|
AND chunk.document_version_id = %s
|
|
AND chunk.embedding_profile_hash = %s
|
|
AND chunk.embedding_text_sha256 = %s
|
|
AND chunk.embedding_model = cache.resolved_model
|
|
AND chunk.embedding_dimension = 1024
|
|
AND chunk.index_status = 'READY'
|
|
AND chunk.embedded_text_sha256 = chunk.embedding_text_sha256
|
|
AND chunk.embedding = cache.embedding
|
|
AND chunk.approval_status = 'CLOUD_APPROVED'
|
|
AND chunk.deleted_at IS NULL
|
|
AND vector_dims(chunk.embedding) = 1024
|
|
AND vector_norm(chunk.embedding) > 0
|
|
"""
|
|
|
|
PROGRESS_SQL = """
|
|
SELECT
|
|
version.expected_chunk_count AS expected_count,
|
|
(
|
|
SELECT count(*)::integer
|
|
FROM rag.chunks AS chunk
|
|
WHERE chunk.document_version_id = version.id
|
|
AND chunk.approval_status = 'CLOUD_APPROVED'
|
|
AND chunk.outbound_manifest_sha256 = version.outbound_manifest_sha256
|
|
AND chunk.embedding_profile_hash = profile.profile_hash
|
|
AND chunk.embedding_model = profile.model
|
|
AND chunk.embedding_dimension = 1024
|
|
AND chunk.deleted_at IS NULL
|
|
) AS chunk_count,
|
|
(
|
|
SELECT count(*)::integer
|
|
FROM rag.chunk_embedding_assignments AS assignment
|
|
JOIN rag.chunks AS chunk ON chunk.id = assignment.chunk_id
|
|
WHERE chunk.document_version_id = version.id
|
|
AND assignment.profile_hash = profile.profile_hash
|
|
) AS assignment_count,
|
|
(
|
|
SELECT count(*)::integer
|
|
FROM rag.chunks AS chunk
|
|
JOIN rag.chunk_embedding_assignments AS assignment
|
|
ON assignment.chunk_id = chunk.id
|
|
AND assignment.profile_hash = profile.profile_hash
|
|
AND assignment.status = 'READY'
|
|
AND assignment.embedding_text_sha256 = chunk.embedding_text_sha256
|
|
AND assignment.cache_profile_hash = profile.profile_hash
|
|
AND assignment.cache_embedding_text_sha256 = chunk.embedding_text_sha256
|
|
JOIN rag.embedding_cache AS cache
|
|
ON cache.profile_hash = assignment.cache_profile_hash
|
|
AND cache.embedding_text_sha256 = assignment.cache_embedding_text_sha256
|
|
AND cache.resolved_model = profile.model
|
|
WHERE chunk.document_version_id = version.id
|
|
AND chunk.approval_status = 'CLOUD_APPROVED'
|
|
AND chunk.outbound_manifest_sha256 = version.outbound_manifest_sha256
|
|
AND chunk.embedding_profile_hash = profile.profile_hash
|
|
AND chunk.embedding_model = profile.model
|
|
AND chunk.embedding_dimension = 1024
|
|
AND chunk.index_status = 'READY'
|
|
AND chunk.embedded_text_sha256 = chunk.embedding_text_sha256
|
|
AND chunk.embedding = cache.embedding
|
|
AND chunk.deleted_at IS NULL
|
|
AND vector_dims(cache.embedding) = 1024
|
|
AND vector_norm(cache.embedding) > 0
|
|
) AS ready_count
|
|
FROM rag.document_versions AS version
|
|
JOIN rag.documents AS document
|
|
ON document.id = version.document_id
|
|
AND document.deleted_at IS NULL
|
|
JOIN rag.knowledge_bases AS knowledge_base
|
|
ON knowledge_base.id = document.knowledge_base_id
|
|
JOIN rag.model_profiles AS profile
|
|
ON profile.profile_hash = version.embedding_profile_hash
|
|
AND profile.kind = 'embedding'
|
|
AND profile.enabled IS TRUE
|
|
AND profile.dimension = 1024
|
|
WHERE version.id = %s
|
|
AND version.review_state = 'CLOUD_APPROVED'
|
|
AND version.expected_chunk_count IS NOT NULL
|
|
AND profile.profile_hash = %s
|
|
AND knowledge_base.active_embedding_profile_hash = profile.profile_hash
|
|
AND knowledge_base.active_embedding_profile_kind = 'embedding'
|
|
FOR SHARE OF version, document, knowledge_base, profile
|
|
"""
|
|
|
|
MARK_VERSION_READY_SQL = """
|
|
UPDATE rag.document_versions AS version
|
|
SET status = 'READY',
|
|
error_code = NULL,
|
|
completed_at = COALESCE(version.completed_at, now())
|
|
FROM rag.documents AS document
|
|
JOIN rag.knowledge_bases AS knowledge_base
|
|
ON knowledge_base.id = document.knowledge_base_id
|
|
JOIN rag.model_profiles AS profile
|
|
ON profile.profile_hash = knowledge_base.active_embedding_profile_hash
|
|
AND profile.kind = 'embedding'
|
|
AND profile.enabled IS TRUE
|
|
WHERE version.id = %s
|
|
AND version.document_id = document.id
|
|
AND version.review_state = 'CLOUD_APPROVED'
|
|
AND version.embedding_profile_hash = %s
|
|
AND profile.profile_hash = version.embedding_profile_hash
|
|
AND document.deleted_at IS NULL
|
|
RETURNING version.document_id
|
|
"""
|
|
|
|
MARK_DOCUMENT_ACTIVE_SQL = """
|
|
UPDATE rag.documents AS document
|
|
SET active_version_id = %s,
|
|
status = 'READY',
|
|
updated_at = now()
|
|
FROM rag.document_versions AS version
|
|
WHERE document.id = version.document_id
|
|
AND version.id = %s
|
|
AND version.status = 'READY'
|
|
AND version.review_state = 'CLOUD_APPROVED'
|
|
AND document.deleted_at IS NULL
|
|
RETURNING document.id
|
|
"""
|
|
|
|
DEACTIVATE_OLD_CHUNKS_SQL = """
|
|
UPDATE rag.chunks
|
|
SET searchable = false,
|
|
updated_at = now()
|
|
WHERE document_id = %s
|
|
AND document_version_id <> %s
|
|
AND searchable IS TRUE
|
|
"""
|
|
|
|
ACTIVATE_CURRENT_CHUNKS_SQL = """
|
|
UPDATE rag.chunks AS chunk
|
|
SET searchable = true,
|
|
updated_at = now()
|
|
FROM rag.document_versions AS version
|
|
JOIN rag.documents AS document
|
|
ON document.id = version.document_id
|
|
JOIN rag.knowledge_bases AS knowledge_base
|
|
ON knowledge_base.id = document.knowledge_base_id
|
|
WHERE chunk.document_version_id = version.id
|
|
AND version.id = %s
|
|
AND version.status = 'READY'
|
|
AND version.review_state = 'CLOUD_APPROVED'
|
|
AND version.embedding_profile_hash = %s
|
|
AND document.active_version_id = version.id
|
|
AND document.status = 'READY'
|
|
AND document.deleted_at IS NULL
|
|
AND knowledge_base.active_embedding_profile_hash = version.embedding_profile_hash
|
|
AND chunk.approval_status = 'CLOUD_APPROVED'
|
|
AND chunk.outbound_manifest_sha256 = version.outbound_manifest_sha256
|
|
AND chunk.embedding_profile_hash = version.embedding_profile_hash
|
|
AND chunk.index_status = 'READY'
|
|
AND chunk.embedded_text_sha256 = chunk.embedding_text_sha256
|
|
AND chunk.embedding IS NOT NULL
|
|
AND chunk.deleted_at IS NULL
|
|
RETURNING chunk.id
|
|
"""
|
|
|
|
|
|
def _default_connection_factory(dsn: str, timeout: int) -> Connection[IndexingRow]:
|
|
return psycopg.connect(
|
|
dsn,
|
|
connect_timeout=timeout,
|
|
row_factory=dict_row,
|
|
application_name="geological-rag-indexing-worker",
|
|
)
|
|
|
|
|
|
class PostgresIndexingRepository:
|
|
"""Implement every protocol method as one independently committed transaction."""
|
|
|
|
def __init__(
|
|
self,
|
|
settings: Settings,
|
|
*,
|
|
connect_timeout: int = 5,
|
|
connection_factory: ConnectionFactory = _default_connection_factory,
|
|
vector_registrar: VectorRegistrar = register_vector,
|
|
) -> None:
|
|
if isinstance(connect_timeout, bool) or not 1 <= connect_timeout <= 60:
|
|
raise ValueError("connect_timeout must be between 1 and 60")
|
|
self._settings = settings
|
|
self._connect_timeout = connect_timeout
|
|
self._connection_factory = connection_factory
|
|
self._vector_registrar = vector_registrar
|
|
|
|
def _dsn(self) -> str:
|
|
return (
|
|
self._settings.database_url()
|
|
.set(drivername="postgresql")
|
|
.render_as_string(hide_password=False)
|
|
)
|
|
|
|
@contextmanager
|
|
def _transaction(self) -> Iterator[Connection[IndexingRow]]:
|
|
try:
|
|
with self._connection_factory(self._dsn(), self._connect_timeout) as connection:
|
|
with connection.transaction():
|
|
self._vector_registrar(cast(Connection[Any], connection))
|
|
yield connection
|
|
except (LeaseLostError, IndexingPersistenceConflict):
|
|
raise
|
|
except (OSError, SecretFileError, psycopg.Error, KeyError, TypeError, ValueError):
|
|
raise IndexingPersistenceError from None
|
|
|
|
@staticmethod
|
|
def _fence(connection: Connection[IndexingRow], lease: JobLease) -> uuid.UUID:
|
|
row = connection.execute(
|
|
LEASE_FENCE_SQL,
|
|
(lease.job_id, lease.worker_id, lease.lease_token),
|
|
).fetchone()
|
|
if row is None:
|
|
raise LeaseLostError("indexing job lease is no longer owned")
|
|
resource_id = row.get("resource_id")
|
|
if not isinstance(resource_id, uuid.UUID):
|
|
raise IndexingPersistenceConflict
|
|
return resource_id
|
|
|
|
def load_approved_plan(
|
|
self,
|
|
*,
|
|
lease: JobLease,
|
|
document_version_id: uuid.UUID,
|
|
) -> ApprovedIndexingPlan:
|
|
with self._transaction() as connection:
|
|
resource_id = self._fence(connection, lease)
|
|
if resource_id != document_version_id:
|
|
raise IndexingPersistenceConflict
|
|
row = connection.execute(LOAD_PLAN_SQL, (document_version_id,)).fetchone()
|
|
if row is None:
|
|
raise IndexingPersistenceConflict
|
|
item_rows = connection.execute(
|
|
LOAD_PLAN_ITEMS_SQL,
|
|
(document_version_id,),
|
|
).fetchall()
|
|
return self._plan_from_rows(row, item_rows)
|
|
|
|
def lookup_cache(
|
|
self,
|
|
*,
|
|
lease: JobLease,
|
|
lookups: Sequence[EmbeddingCacheLookup],
|
|
) -> Mapping[str, CachedEmbedding]:
|
|
if not lookups or len(lookups) > 10:
|
|
raise IndexingPersistenceConflict
|
|
with self._transaction() as connection:
|
|
resource_id = self._fence(connection, lease)
|
|
found: dict[str, CachedEmbedding] = {}
|
|
for lookup in lookups:
|
|
if (
|
|
embedding_cache_key(
|
|
lookup.embedding_text_sha256,
|
|
lookup.profile_hash,
|
|
)
|
|
!= lookup.cache_key
|
|
):
|
|
raise IndexingPersistenceConflict
|
|
row = connection.execute(
|
|
CACHE_LOOKUP_SQL,
|
|
(resource_id, lookup.profile_hash, lookup.embedding_text_sha256),
|
|
).fetchone()
|
|
if row is not None:
|
|
record = self._cache_from_row(lookup.cache_key, row)
|
|
found[lookup.cache_key] = record
|
|
return found
|
|
|
|
def begin_model_invocation(
|
|
self,
|
|
*,
|
|
lease: JobLease,
|
|
trace_id: uuid.UUID,
|
|
profile_hash: str,
|
|
model: str,
|
|
item_count: int,
|
|
) -> uuid.UUID:
|
|
if (
|
|
trace_id != lease.job_id
|
|
or not _valid_hash(profile_hash)
|
|
or not _valid_model(model)
|
|
or isinstance(item_count, bool)
|
|
or not 1 <= item_count <= 10
|
|
):
|
|
raise IndexingPersistenceConflict
|
|
with self._transaction() as connection:
|
|
resource_id = self._fence(connection, lease)
|
|
row = connection.execute(
|
|
BEGIN_INVOCATION_SQL,
|
|
(trace_id, item_count, resource_id, profile_hash, model),
|
|
).fetchone()
|
|
if row is None or not isinstance(row.get("id"), uuid.UUID):
|
|
raise IndexingPersistenceConflict
|
|
return cast(uuid.UUID, row["id"])
|
|
|
|
def finish_model_invocation(
|
|
self,
|
|
*,
|
|
lease: JobLease,
|
|
invocation_id: uuid.UUID,
|
|
status: InvocationStatus,
|
|
provider_request_id: str | None,
|
|
usage: ProviderUsage,
|
|
elapsed_ms: float,
|
|
error_code: str | None,
|
|
) -> None:
|
|
_validate_invocation_finish(
|
|
status=status,
|
|
provider_request_id=provider_request_id,
|
|
elapsed_ms=elapsed_ms,
|
|
error_code=error_code,
|
|
)
|
|
prompt_tokens, completion_tokens, total_tokens = _invocation_counts(usage)
|
|
elapsed = _elapsed_integer(elapsed_ms)
|
|
with self._transaction() as connection:
|
|
resource_id = self._fence(connection, lease)
|
|
row = connection.execute(
|
|
FINISH_INVOCATION_SQL,
|
|
(
|
|
status,
|
|
provider_request_id,
|
|
prompt_tokens,
|
|
completion_tokens,
|
|
total_tokens,
|
|
elapsed,
|
|
error_code,
|
|
invocation_id,
|
|
lease.job_id,
|
|
resource_id,
|
|
),
|
|
).fetchone()
|
|
if row is None:
|
|
raise IndexingPersistenceConflict
|
|
|
|
def fenced_persist_batch(
|
|
self,
|
|
*,
|
|
lease: JobLease,
|
|
document_version_id: uuid.UUID,
|
|
profile_hash: str,
|
|
writes: Sequence[EmbeddingWrite],
|
|
) -> AssignmentProgress:
|
|
if not writes or len(writes) > 10 or not _valid_hash(profile_hash):
|
|
raise IndexingPersistenceConflict
|
|
with self._transaction() as connection:
|
|
resource_id = self._fence(connection, lease)
|
|
if resource_id != document_version_id:
|
|
raise IndexingPersistenceConflict
|
|
for write in writes:
|
|
self._persist_write(
|
|
connection,
|
|
document_version_id=document_version_id,
|
|
profile_hash=profile_hash,
|
|
write=write,
|
|
)
|
|
return self._read_progress(
|
|
connection,
|
|
document_version_id=document_version_id,
|
|
profile_hash=profile_hash,
|
|
)
|
|
|
|
def fenced_activate(
|
|
self,
|
|
*,
|
|
lease: JobLease,
|
|
document_version_id: uuid.UUID,
|
|
profile_hash: str,
|
|
expected_count: int,
|
|
) -> bool:
|
|
if not _valid_hash(profile_hash) or isinstance(expected_count, bool) or expected_count < 0:
|
|
raise IndexingPersistenceConflict
|
|
with self._transaction() as connection:
|
|
resource_id = self._fence(connection, lease)
|
|
if resource_id != document_version_id:
|
|
raise IndexingPersistenceConflict
|
|
progress, chunk_count, assignment_count = self._read_progress_details(
|
|
connection,
|
|
document_version_id=document_version_id,
|
|
profile_hash=profile_hash,
|
|
)
|
|
if not (
|
|
progress.expected_count
|
|
== expected_count
|
|
== progress.ready_count
|
|
== chunk_count
|
|
== assignment_count
|
|
):
|
|
return False
|
|
|
|
version_row = connection.execute(
|
|
MARK_VERSION_READY_SQL,
|
|
(document_version_id, profile_hash),
|
|
).fetchone()
|
|
if version_row is None or not isinstance(version_row.get("document_id"), uuid.UUID):
|
|
raise IndexingPersistenceConflict
|
|
document_id = cast(uuid.UUID, version_row["document_id"])
|
|
document_row = connection.execute(
|
|
MARK_DOCUMENT_ACTIVE_SQL,
|
|
(document_version_id, document_version_id),
|
|
).fetchone()
|
|
if document_row is None:
|
|
raise IndexingPersistenceConflict
|
|
connection.execute(
|
|
DEACTIVATE_OLD_CHUNKS_SQL,
|
|
(document_id, document_version_id),
|
|
)
|
|
activated_rows = connection.execute(
|
|
ACTIVATE_CURRENT_CHUNKS_SQL,
|
|
(document_version_id, profile_hash),
|
|
).fetchall()
|
|
if len(activated_rows) != expected_count:
|
|
raise IndexingPersistenceConflict
|
|
return True
|
|
|
|
@staticmethod
|
|
def _plan_from_rows(
|
|
row: IndexingRow,
|
|
item_rows: Sequence[IndexingRow],
|
|
) -> ApprovedIndexingPlan:
|
|
expected_count = _required_integer(row, "expected_chunk_count")
|
|
manifest_count = _required_integer(row, "manifest_count")
|
|
eligible_count = _required_integer(row, "eligible_chunk_count")
|
|
if expected_count < 0 or not expected_count == manifest_count == eligible_count == len(
|
|
item_rows
|
|
):
|
|
raise IndexingPersistenceConflict
|
|
profile = ActiveEmbeddingProfile(
|
|
profile_hash=_required_hash(row, "profile_hash"),
|
|
model=_required_text(row, "model"),
|
|
dimension=_required_integer(row, "dimension"),
|
|
synthetic=_required_boolean(row, "synthetic"),
|
|
)
|
|
if profile.dimension != 1024:
|
|
raise IndexingPersistenceConflict
|
|
items = tuple(PostgresIndexingRepository._item_from_row(item) for item in item_rows)
|
|
return ApprovedIndexingPlan(
|
|
knowledge_base_id=_required_uuid(row, "knowledge_base_id"),
|
|
document_version_id=_required_uuid(row, "document_version_id"),
|
|
review_state=_required_text(row, "review_state"),
|
|
outbound_manifest_sha256=_required_hash(row, "outbound_manifest_sha256"),
|
|
expected_count=expected_count,
|
|
profile=profile,
|
|
items=items,
|
|
)
|
|
|
|
@staticmethod
|
|
def _item_from_row(row: IndexingRow) -> IndexingItem:
|
|
status = _required_text(row, "assignment_status")
|
|
if status not in {"PENDING", "EMBEDDING", "READY", "FAILED", "STALE"}:
|
|
raise IndexingPersistenceConflict
|
|
return IndexingItem(
|
|
chunk_id=_required_uuid(row, "chunk_id"),
|
|
ordinal=_required_integer(row, "ordinal"),
|
|
embedding_text=_required_text(row, "embedding_text", strip=False),
|
|
embedding_text_sha256=_required_hash(row, "embedding_text_sha256"),
|
|
assignment_status=cast(AssignmentStatus, status),
|
|
)
|
|
|
|
@staticmethod
|
|
def _cache_from_row(cache_key: str, row: IndexingRow) -> CachedEmbedding:
|
|
dimension = _required_integer(row, "dimension")
|
|
if dimension != 1024:
|
|
raise IndexingPersistenceConflict
|
|
return CachedEmbedding(
|
|
cache_key=cache_key,
|
|
profile_hash=_required_hash(row, "profile_hash"),
|
|
embedding_text_sha256=_required_hash(row, "embedding_text_sha256"),
|
|
resolved_model=_required_text(row, "resolved_model"),
|
|
dimension=dimension,
|
|
)
|
|
|
|
@staticmethod
|
|
def _persist_write(
|
|
connection: Connection[IndexingRow],
|
|
*,
|
|
document_version_id: uuid.UUID,
|
|
profile_hash: str,
|
|
write: EmbeddingWrite,
|
|
) -> None:
|
|
_validate_write(write, profile_hash)
|
|
if write.source == "provider":
|
|
vector = Vector(list(cast(tuple[float, ...], write.embedding)))
|
|
usage = Jsonb(_usage_json(write.usage))
|
|
elapsed = _elapsed_integer(write.elapsed_ms)
|
|
connection.execute(
|
|
INSERT_CACHE_SQL,
|
|
(
|
|
profile_hash,
|
|
write.embedding_text_sha256,
|
|
vector,
|
|
write.resolved_model,
|
|
write.provider_request_id,
|
|
usage,
|
|
elapsed,
|
|
document_version_id,
|
|
profile_hash,
|
|
write.resolved_model,
|
|
vector,
|
|
vector,
|
|
),
|
|
)
|
|
cache_row = connection.execute(
|
|
VERIFY_CACHE_FOR_WRITE_SQL,
|
|
(
|
|
document_version_id,
|
|
profile_hash,
|
|
write.embedding_text_sha256,
|
|
write.resolved_model,
|
|
),
|
|
).fetchone()
|
|
if cache_row is None:
|
|
raise IndexingPersistenceConflict
|
|
assignment_row = connection.execute(
|
|
UPSERT_ASSIGNMENT_SQL,
|
|
(
|
|
write.chunk_id,
|
|
document_version_id,
|
|
write.embedding_text_sha256,
|
|
profile_hash,
|
|
),
|
|
).fetchone()
|
|
if assignment_row is None:
|
|
raise IndexingPersistenceConflict
|
|
connection.execute(
|
|
PREPARE_STALE_CHUNK_SQL,
|
|
(
|
|
write.chunk_id,
|
|
document_version_id,
|
|
profile_hash,
|
|
write.embedding_text_sha256,
|
|
),
|
|
)
|
|
chunk_row = connection.execute(
|
|
UPDATE_CHUNK_FROM_CACHE_SQL,
|
|
(
|
|
write.chunk_id,
|
|
document_version_id,
|
|
write.embedding_text_sha256,
|
|
),
|
|
).fetchone()
|
|
if chunk_row is None:
|
|
chunk_row = connection.execute(
|
|
VERIFY_READY_CHUNK_SQL,
|
|
(
|
|
write.chunk_id,
|
|
document_version_id,
|
|
profile_hash,
|
|
write.embedding_text_sha256,
|
|
),
|
|
).fetchone()
|
|
if chunk_row is None:
|
|
raise IndexingPersistenceConflict
|
|
|
|
@staticmethod
|
|
def _read_progress(
|
|
connection: Connection[IndexingRow],
|
|
*,
|
|
document_version_id: uuid.UUID,
|
|
profile_hash: str,
|
|
) -> AssignmentProgress:
|
|
progress, _, _ = PostgresIndexingRepository._read_progress_details(
|
|
connection,
|
|
document_version_id=document_version_id,
|
|
profile_hash=profile_hash,
|
|
)
|
|
return progress
|
|
|
|
@staticmethod
|
|
def _read_progress_details(
|
|
connection: Connection[IndexingRow],
|
|
*,
|
|
document_version_id: uuid.UUID,
|
|
profile_hash: str,
|
|
) -> tuple[AssignmentProgress, int, int]:
|
|
row = connection.execute(
|
|
PROGRESS_SQL,
|
|
(document_version_id, profile_hash),
|
|
).fetchone()
|
|
if row is None:
|
|
raise IndexingPersistenceConflict
|
|
expected_count = _required_integer(row, "expected_count")
|
|
chunk_count = _required_integer(row, "chunk_count")
|
|
assignment_count = _required_integer(row, "assignment_count")
|
|
ready_count = _required_integer(row, "ready_count")
|
|
if not 0 <= ready_count <= assignment_count <= chunk_count <= expected_count:
|
|
raise IndexingPersistenceConflict
|
|
return AssignmentProgress(expected_count, ready_count), chunk_count, assignment_count
|
|
|
|
|
|
def _validate_write(write: EmbeddingWrite, profile_hash: str) -> None:
|
|
if (
|
|
write.profile_hash != profile_hash
|
|
or not _valid_hash(write.embedding_text_sha256)
|
|
or embedding_cache_key(write.embedding_text_sha256, profile_hash) != write.cache_key
|
|
or not _valid_model(write.resolved_model)
|
|
or isinstance(write.batch_index, bool)
|
|
or not 0 <= write.batch_index < 10
|
|
or isinstance(write.elapsed_ms, bool)
|
|
or not isinstance(write.elapsed_ms, (int, float))
|
|
or not math.isfinite(float(write.elapsed_ms))
|
|
or write.elapsed_ms < 0
|
|
or (
|
|
write.provider_request_id is not None
|
|
and not _valid_identifier(write.provider_request_id)
|
|
)
|
|
):
|
|
raise IndexingPersistenceConflict
|
|
_invocation_counts(write.usage)
|
|
if write.source == "cache":
|
|
if write.embedding is not None:
|
|
raise IndexingPersistenceConflict
|
|
return
|
|
if write.source != "provider" or write.embedding is None:
|
|
raise IndexingPersistenceConflict
|
|
if len(write.embedding) != 1024:
|
|
raise IndexingPersistenceConflict
|
|
norm = 0.0
|
|
for component in write.embedding:
|
|
if (
|
|
isinstance(component, bool)
|
|
or not isinstance(component, (int, float))
|
|
or not math.isfinite(float(component))
|
|
):
|
|
raise IndexingPersistenceConflict
|
|
norm += float(component) ** 2
|
|
if norm <= 0:
|
|
raise IndexingPersistenceConflict
|
|
|
|
|
|
def _validate_invocation_finish(
|
|
*,
|
|
status: InvocationStatus,
|
|
provider_request_id: str | None,
|
|
elapsed_ms: float,
|
|
error_code: str | None,
|
|
) -> None:
|
|
if status not in {"SUCCEEDED", "FAILED", "UNKNOWN"}:
|
|
raise IndexingPersistenceConflict
|
|
if (status == "SUCCEEDED") != (error_code is None):
|
|
raise IndexingPersistenceConflict
|
|
if error_code is not None and _ERROR_CODE.fullmatch(error_code) is None:
|
|
raise IndexingPersistenceConflict
|
|
if provider_request_id is not None and not _valid_identifier(provider_request_id):
|
|
raise IndexingPersistenceConflict
|
|
_elapsed_integer(elapsed_ms)
|
|
|
|
|
|
def _invocation_counts(usage: ProviderUsage) -> tuple[int, int, int]:
|
|
values = (usage.input_tokens, usage.output_tokens, usage.total_tokens)
|
|
if any(
|
|
isinstance(value, bool) or (value is not None and (not isinstance(value, int) or value < 0))
|
|
for value in values
|
|
):
|
|
raise IndexingPersistenceConflict
|
|
prompt = usage.input_tokens
|
|
completion = usage.output_tokens or 0
|
|
if prompt is None:
|
|
prompt = max(0, (usage.total_tokens or 0) - completion)
|
|
total = prompt + completion
|
|
if usage.total_tokens is not None and usage.total_tokens != total:
|
|
raise IndexingPersistenceConflict
|
|
return prompt, completion, total
|
|
|
|
|
|
def _usage_json(usage: ProviderUsage) -> dict[str, int]:
|
|
prompt, completion, total = _invocation_counts(usage)
|
|
return {
|
|
"input_tokens": prompt,
|
|
"output_tokens": completion,
|
|
"total_tokens": total,
|
|
}
|
|
|
|
|
|
def _elapsed_integer(value: float) -> int:
|
|
if (
|
|
isinstance(value, bool)
|
|
or not isinstance(value, (int, float))
|
|
or not math.isfinite(float(value))
|
|
or not 0 <= value <= 2_147_483_647
|
|
):
|
|
raise IndexingPersistenceConflict
|
|
return int(round(value))
|
|
|
|
|
|
def _valid_hash(value: str) -> bool:
|
|
return isinstance(value, str) and _HASH.fullmatch(value) is not None
|
|
|
|
|
|
def _valid_model(value: str) -> bool:
|
|
return isinstance(value, str) and value == value.strip() and _valid_identifier(value)
|
|
|
|
|
|
def _valid_identifier(value: str) -> bool:
|
|
return isinstance(value, str) and _SAFE_IDENTIFIER.fullmatch(value) is not None
|
|
|
|
|
|
def _required_uuid(row: Mapping[str, object], key: str) -> uuid.UUID:
|
|
value = row.get(key)
|
|
if not isinstance(value, uuid.UUID):
|
|
raise IndexingPersistenceConflict
|
|
return value
|
|
|
|
|
|
def _required_text(row: Mapping[str, object], key: str, *, strip: bool = True) -> str:
|
|
value = row.get(key)
|
|
if not isinstance(value, str) or not value or (strip and value != value.strip()):
|
|
raise IndexingPersistenceConflict
|
|
return value
|
|
|
|
|
|
def _required_hash(row: Mapping[str, object], key: str) -> str:
|
|
value = _required_text(row, key)
|
|
if not _valid_hash(value):
|
|
raise IndexingPersistenceConflict
|
|
return value
|
|
|
|
|
|
def _required_integer(row: Mapping[str, object], key: str) -> int:
|
|
value = row.get(key)
|
|
if isinstance(value, bool) or not isinstance(value, int):
|
|
raise IndexingPersistenceConflict
|
|
return value
|
|
|
|
|
|
def _required_boolean(row: Mapping[str, object], key: str) -> bool:
|
|
value = row.get(key)
|
|
if not isinstance(value, bool):
|
|
raise IndexingPersistenceConflict
|
|
return value
|