generated from john/python-template
409 lines
13 KiB
Python
409 lines
13 KiB
Python
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import logging
|
|
from collections.abc import Sequence
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from uuid import UUID
|
|
from uuid import uuid4
|
|
|
|
from sqlmodel import select
|
|
from sqlmodel.ext.asyncio.session import AsyncSession
|
|
|
|
from transcription.config import Settings
|
|
from transcription.config import get_settings
|
|
from transcription.errors import AppError
|
|
from transcription.errors import ErrorCategory
|
|
from transcription.runtime_helpers import run_blocking
|
|
|
|
from ..db.models import Document
|
|
from ..db.models import Job
|
|
from ..db.models import JobSource
|
|
from ..db.models import JobSourceStatus
|
|
from ..db.models import Source
|
|
from ..db.session import SessionFactory
|
|
from ..db.session import session_scope
|
|
from .errors import TranscriptionError
|
|
from .media_storage import persist_named_media
|
|
from .normalization import normalize_orientation_async
|
|
from .sources import build_prompt_execution
|
|
from .sources import source_mime_type
|
|
from .sources import validate_source_content
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class SourceStorageError(AppError):
|
|
"""Raised when Source content cannot be validated or persisted safely."""
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class JobCreateResult:
|
|
"""Summary of explicit Job create records."""
|
|
|
|
document_id: UUID
|
|
job_id: UUID
|
|
source_ids: tuple[UUID, ...]
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class DocumentJobResult:
|
|
"""Summary of a Document, Source, and Job created together."""
|
|
|
|
document_id: UUID
|
|
job_id: UUID
|
|
stored_path: Path
|
|
original_filename: str
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class StoredSourceFile:
|
|
"""A persisted Source file and the identity of the bytes actually stored."""
|
|
|
|
path: Path
|
|
file_hash: str
|
|
file_size_bytes: int
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class PendingStoredSource:
|
|
"""Pre-staged Source artifact tied to a Source id."""
|
|
|
|
source_id: UUID
|
|
original_filename: str
|
|
stored_path: Path
|
|
file_hash: str
|
|
file_size_bytes: int
|
|
|
|
|
|
async def create_document_job(
|
|
*,
|
|
filename: str,
|
|
file_bytes: bytes,
|
|
session: AsyncSession | None = None,
|
|
session_factory: SessionFactory | None = None,
|
|
settings: Settings | None = None,
|
|
) -> DocumentJobResult:
|
|
"""Create a Document, its first Source, and a queued Job.
|
|
|
|
Owns its own session when the caller does not supply one, so UI callers
|
|
never have to import a session scope.
|
|
"""
|
|
runtime_settings = settings or get_settings()
|
|
prompt_execution = build_prompt_execution(settings=runtime_settings)
|
|
document_id = uuid4()
|
|
source_id = uuid4()
|
|
stored = await store_source_file(
|
|
filename=filename,
|
|
file_bytes=file_bytes,
|
|
settings=runtime_settings,
|
|
relative_directory=Path("documents") / str(document_id),
|
|
filename_stem=str(source_id),
|
|
)
|
|
stored_path = stored.path
|
|
try:
|
|
async with session_scope(
|
|
session_factory=session_factory,
|
|
session=session,
|
|
settings=runtime_settings,
|
|
) as _session:
|
|
document, job = await _create_document_job_records(
|
|
session=_session,
|
|
document_id=document_id,
|
|
source_id=source_id,
|
|
original_filename=filename,
|
|
stored_path=stored_path,
|
|
file_hash=stored.file_hash,
|
|
file_size_bytes=stored.file_size_bytes,
|
|
upload_dir=runtime_settings.upload_dir,
|
|
prompt_execution=prompt_execution,
|
|
)
|
|
except Exception as exc:
|
|
_best_effort_delete(stored_path)
|
|
raise SourceStorageError(
|
|
"Failed to create Document, Source, and Job records",
|
|
category=ErrorCategory.INFRA_PERSISTENT,
|
|
suggestion="Retry upload. If this keeps happening, verify database availability.",
|
|
) from exc
|
|
|
|
logger.info("Created document job document_id=%s job_id=%s", document.id, job.id)
|
|
return DocumentJobResult(
|
|
document_id=document.id,
|
|
job_id=job.id,
|
|
stored_path=stored_path,
|
|
original_filename=Path(filename).name,
|
|
)
|
|
|
|
|
|
async def create_job_for_document(
|
|
*,
|
|
document_id: UUID,
|
|
source_files: Sequence[tuple[str, bytes]],
|
|
session: AsyncSession | None = None,
|
|
session_factory: SessionFactory | None = None,
|
|
provider: str | None = None,
|
|
model: str | None = None,
|
|
settings: Settings | None = None,
|
|
) -> JobCreateResult:
|
|
"""Create a queued Job for an existing Document with one or more Sources.
|
|
|
|
Owns its own session when the caller does not supply one, so UI callers
|
|
never have to import a session scope.
|
|
"""
|
|
if not source_files:
|
|
raise SourceStorageError(
|
|
"At least one Source file is required to create a Job",
|
|
category=ErrorCategory.VALIDATION,
|
|
suggestion="Upload one or more files and try again.",
|
|
)
|
|
|
|
runtime_settings = settings or get_settings()
|
|
prompt_execution = build_prompt_execution(settings=runtime_settings)
|
|
sorted_source_files = sorted(source_files, key=lambda item: Path(item[0]).name.casefold())
|
|
stored_sources: list[PendingStoredSource] = []
|
|
for filename, file_bytes in sorted_source_files:
|
|
source_id = uuid4()
|
|
stored = await store_source_file(
|
|
filename=filename,
|
|
file_bytes=file_bytes,
|
|
settings=runtime_settings,
|
|
relative_directory=Path("documents") / str(document_id),
|
|
filename_stem=str(source_id),
|
|
)
|
|
stored_sources.append(
|
|
PendingStoredSource(
|
|
source_id=source_id,
|
|
original_filename=filename,
|
|
stored_path=stored.path,
|
|
file_hash=stored.file_hash,
|
|
file_size_bytes=stored.file_size_bytes,
|
|
)
|
|
)
|
|
|
|
try:
|
|
async with session_scope(
|
|
session_factory=session_factory,
|
|
session=session,
|
|
settings=runtime_settings,
|
|
) as _session:
|
|
job, source_ids = await _create_job_for_document_records(
|
|
session=_session,
|
|
document_id=document_id,
|
|
stored_sources=stored_sources,
|
|
provider=provider,
|
|
model=model,
|
|
upload_dir=runtime_settings.upload_dir,
|
|
prompt_execution=prompt_execution,
|
|
)
|
|
except Exception as exc:
|
|
for source in stored_sources:
|
|
_best_effort_delete(source.stored_path)
|
|
raise SourceStorageError(
|
|
"Failed to create Job records from Source files",
|
|
category=ErrorCategory.INFRA_PERSISTENT,
|
|
suggestion="Retry creation. If this keeps happening, verify database availability.",
|
|
) from exc
|
|
|
|
logger.info("Created explicit job document_id=%s job_id=%s sources=%s", document_id, job.id, len(source_ids))
|
|
return JobCreateResult(
|
|
document_id=document_id,
|
|
job_id=job.id,
|
|
source_ids=tuple(source_ids),
|
|
)
|
|
|
|
|
|
async def _create_document_job_records(
|
|
*,
|
|
session: AsyncSession,
|
|
document_id: UUID,
|
|
source_id: UUID,
|
|
original_filename: str,
|
|
stored_path: Path,
|
|
file_hash: str,
|
|
file_size_bytes: int,
|
|
upload_dir: Path,
|
|
prompt_execution,
|
|
) -> tuple[Document, Job]:
|
|
document = Document(
|
|
id=document_id,
|
|
name=Path(original_filename).name,
|
|
)
|
|
session.add(document)
|
|
await session.flush()
|
|
|
|
job = Job(
|
|
document_id=document.id,
|
|
prompt_name=prompt_execution.prompt_name,
|
|
prompt_hash=prompt_execution.prompt_hash,
|
|
system_prompt=prompt_execution.system_prompt,
|
|
user_prompt=prompt_execution.user_prompt,
|
|
temperature=prompt_execution.temperature,
|
|
top_p=prompt_execution.top_p,
|
|
)
|
|
session.add(job)
|
|
await session.flush()
|
|
|
|
source = Source(
|
|
id=source_id,
|
|
document_id=document.id,
|
|
page_number=1,
|
|
upload_name=Path(original_filename).name,
|
|
filename=stored_path.name,
|
|
file_path=_upload_relative_path(stored_path=stored_path, upload_dir=upload_dir),
|
|
file_hash=file_hash,
|
|
file_size_bytes=file_size_bytes,
|
|
)
|
|
session.add(source)
|
|
await session.flush()
|
|
|
|
session.add(
|
|
JobSource(
|
|
job_id=job.id,
|
|
source_id=source.id,
|
|
status=JobSourceStatus.PENDING,
|
|
)
|
|
)
|
|
|
|
await session.commit()
|
|
await session.refresh(document)
|
|
await session.refresh(job)
|
|
return document, job
|
|
|
|
|
|
async def _create_job_for_document_records(
|
|
*,
|
|
session: AsyncSession,
|
|
document_id: UUID,
|
|
stored_sources: Sequence[PendingStoredSource],
|
|
provider: str | None,
|
|
model: str | None,
|
|
upload_dir: Path,
|
|
prompt_execution,
|
|
) -> tuple[Job, list[UUID]]:
|
|
document = await session.get(Document, document_id)
|
|
if document is None:
|
|
raise SourceStorageError(
|
|
f"Document with id {document_id} not found",
|
|
category=ErrorCategory.NOT_FOUND,
|
|
suggestion="Select an existing document and retry.",
|
|
)
|
|
|
|
existing_sources = (await session.exec(select(Source).where(Source.document_id == document_id))).all()
|
|
next_page_number = max((source.page_number for source in existing_sources), default=0) + 1
|
|
|
|
job = Job(
|
|
document_id=document_id,
|
|
provider=(provider or None),
|
|
model=(model or None),
|
|
prompt_name=prompt_execution.prompt_name,
|
|
prompt_hash=prompt_execution.prompt_hash,
|
|
system_prompt=prompt_execution.system_prompt,
|
|
user_prompt=prompt_execution.user_prompt,
|
|
temperature=prompt_execution.temperature,
|
|
top_p=prompt_execution.top_p,
|
|
)
|
|
session.add(job)
|
|
await session.flush()
|
|
|
|
source_ids: list[UUID] = []
|
|
for page_offset, stored_source in enumerate(stored_sources):
|
|
source = Source(
|
|
id=stored_source.source_id,
|
|
document_id=document_id,
|
|
page_number=next_page_number + page_offset,
|
|
upload_name=Path(stored_source.original_filename).name,
|
|
filename=stored_source.stored_path.name,
|
|
file_path=_upload_relative_path(
|
|
stored_path=stored_source.stored_path,
|
|
upload_dir=upload_dir,
|
|
),
|
|
file_hash=stored_source.file_hash,
|
|
file_size_bytes=stored_source.file_size_bytes,
|
|
)
|
|
session.add(source)
|
|
await session.flush()
|
|
source_ids.append(source.id)
|
|
|
|
session.add(
|
|
JobSource(
|
|
job_id=job.id,
|
|
source_id=source.id,
|
|
status=JobSourceStatus.PENDING,
|
|
)
|
|
)
|
|
|
|
await session.commit()
|
|
await session.refresh(job)
|
|
return job, source_ids
|
|
|
|
|
|
def _best_effort_delete(path: Path) -> None:
|
|
try:
|
|
if path.exists():
|
|
path.unlink()
|
|
except OSError:
|
|
logger.warning("Failed to clean up Source file after database error: %s", path)
|
|
|
|
|
|
def _upload_relative_path(*, stored_path: Path, upload_dir: Path) -> str:
|
|
return stored_path.resolve().relative_to(upload_dir.resolve()).as_posix()
|
|
|
|
|
|
async def store_source_file(
|
|
*,
|
|
filename: str,
|
|
file_bytes: bytes,
|
|
settings: Settings | None = None,
|
|
relative_directory: Path | None = None,
|
|
filename_stem: str | None = None,
|
|
) -> StoredSourceFile:
|
|
"""Validate, orient, and persist a Source file to configured media storage.
|
|
|
|
Orientation is applied here, at the ingest boundary, so the stored bytes are
|
|
already upright and the hash and byte size recorded on the ``Source`` row
|
|
describe exactly what is on disk and exactly what a provider is later sent.
|
|
"""
|
|
runtime_settings = settings or get_settings()
|
|
try:
|
|
validate_source_content(filename=filename, content=file_bytes)
|
|
except TranscriptionError as exc:
|
|
raise SourceStorageError(
|
|
exc.message,
|
|
category=exc.category,
|
|
suggestion=exc.suggestion,
|
|
retriable=exc.retriable,
|
|
) from exc
|
|
|
|
normalized = await normalize_orientation_async(file_bytes, media_type=source_mime_type(filename))
|
|
if normalized is not None:
|
|
logger.info(
|
|
"Normalized Source orientation on ingest: %s (orientation=%s, rotation=%s)",
|
|
Path(filename).name,
|
|
normalized.original_orientation,
|
|
normalized.applied_rotation_degrees,
|
|
)
|
|
file_bytes = normalized.content
|
|
|
|
upload_dir = runtime_settings.upload_dir
|
|
stored_path = await persist_named_media(
|
|
root=upload_dir,
|
|
namespace=relative_directory,
|
|
filename=filename,
|
|
filename_stem=filename_stem,
|
|
file_bytes=file_bytes,
|
|
error=SourceStorageError,
|
|
failure_message="Failed to persist Source file",
|
|
failure_suggestion="Check upload directory permissions and available disk space, then retry.",
|
|
log_label="Source file",
|
|
)
|
|
return StoredSourceFile(
|
|
path=stored_path,
|
|
file_hash=await run_blocking(_sha256_hexdigest, file_bytes),
|
|
file_size_bytes=len(file_bytes),
|
|
)
|
|
|
|
|
|
def _sha256_hexdigest(data: bytes) -> str:
|
|
return hashlib.sha256(data).hexdigest()
|