generated from kgod/ai-review-template
224 lines
8.9 KiB
Python
224 lines
8.9 KiB
Python
from __future__ import annotations
|
|
|
|
from datetime import UTC, datetime
|
|
from typing import Any
|
|
from uuid import uuid4
|
|
|
|
from sqlalchemy import Engine, func, insert, select, update
|
|
from sqlalchemy.exc import IntegrityError
|
|
|
|
from .schema import build_session_tables
|
|
|
|
|
|
class RevisionConflict(Exception):
|
|
"""The caller attempted to replace a stale resume document."""
|
|
|
|
|
|
def _now() -> datetime:
|
|
return datetime.now(UTC)
|
|
|
|
|
|
class PostgresSessionRepository:
|
|
"""Core PostgreSQL store, shaped for incremental replacement of SQLite."""
|
|
|
|
def __init__(self, engine: Engine, *, schema: str = "resume_agent") -> None:
|
|
self.engine = engine
|
|
self.schema = schema
|
|
self.metadata, self.tables = build_session_tables(schema)
|
|
|
|
def initialize(self) -> None:
|
|
with self.engine.begin() as connection:
|
|
connection.exec_driver_sql(f'CREATE SCHEMA IF NOT EXISTS "{self.schema}"')
|
|
self.metadata.create_all(connection)
|
|
|
|
def create_session(self, session_id: str, stage: str, profile: dict[str, Any]) -> dict[str, Any]:
|
|
now = _now()
|
|
values = {
|
|
"id": session_id,
|
|
"stage": stage,
|
|
"revision": 0,
|
|
"profile": profile,
|
|
"created_at": now,
|
|
"updated_at": now,
|
|
}
|
|
with self.engine.begin() as connection:
|
|
connection.execute(insert(self.tables["sessions"]).values(**values))
|
|
return self.get_session(session_id) or values
|
|
|
|
def get_session(self, session_id: str) -> dict[str, Any] | None:
|
|
with self.engine.connect() as connection:
|
|
row = connection.execute(
|
|
select(self.tables["sessions"]).where(self.tables["sessions"].c.id == session_id)
|
|
).mappings().first()
|
|
return dict(row) if row else None
|
|
|
|
def insert_turn(
|
|
self, session_id: str, *, role: str, content: str | None, composer_mode: str, blocks: list[dict[str, Any]]
|
|
) -> dict[str, Any]:
|
|
turns, block_table = self.tables["turns"], self.tables["blocks"]
|
|
now, turn_id = _now(), f"turn_{uuid4().hex}"
|
|
with self.engine.begin() as connection:
|
|
maximum = connection.execute(
|
|
select(func.coalesce(func.max(turns.c.sequence), 0)).where(turns.c.session_id == session_id)
|
|
).scalar_one()
|
|
sequence = maximum + 1
|
|
connection.execute(insert(turns).values(
|
|
id=turn_id, session_id=session_id, sequence=sequence, role=role,
|
|
content=content, composer_mode=composer_mode, created_at=now,
|
|
))
|
|
for index, block in enumerate(blocks):
|
|
connection.execute(insert(block_table).values(
|
|
id=block.get("id", f"block_{uuid4().hex}"), session_id=session_id,
|
|
turn_id=turn_id, block_index=index, type=block["type"],
|
|
lifecycle=block.get("lifecycle", "active"), data=block.get("data", {}),
|
|
version=1, created_at=now, updated_at=now,
|
|
))
|
|
return self._turn(turn_id)
|
|
|
|
def list_turns(self, session_id: str) -> list[dict[str, Any]]:
|
|
turns = self.tables["turns"]
|
|
with self.engine.connect() as connection:
|
|
ids = connection.execute(
|
|
select(turns.c.id).where(turns.c.session_id == session_id).order_by(turns.c.sequence)
|
|
).scalars().all()
|
|
return [self._turn(turn_id) for turn_id in ids]
|
|
|
|
def _turn(self, turn_id: str) -> dict[str, Any]:
|
|
turns, blocks = self.tables["turns"], self.tables["blocks"]
|
|
with self.engine.connect() as connection:
|
|
turn = connection.execute(select(turns).where(turns.c.id == turn_id)).mappings().one()
|
|
rows = connection.execute(
|
|
select(blocks).where(blocks.c.turn_id == turn_id).order_by(blocks.c.block_index)
|
|
).mappings().all()
|
|
result = dict(turn)
|
|
result["blocks"] = [dict(row) for row in rows]
|
|
return result
|
|
|
|
def create_resume(
|
|
self, session_id: str, resume_id: str, idempotency_key: str | None, content: dict[str, Any]
|
|
) -> dict[str, Any]:
|
|
existing = self.get_resume(session_id)
|
|
if existing:
|
|
return existing
|
|
now = _now()
|
|
try:
|
|
with self.engine.begin() as connection:
|
|
connection.execute(insert(self.tables["resumes"]).values(
|
|
id=resume_id, session_id=session_id, idempotency_key=idempotency_key,
|
|
revision=1, content=content, created_at=now, updated_at=now,
|
|
))
|
|
except IntegrityError:
|
|
existing = self.get_resume(session_id)
|
|
if existing:
|
|
return existing
|
|
raise
|
|
return self.get_resume(session_id) # type: ignore[return-value]
|
|
|
|
def get_resume(self, session_id: str) -> dict[str, Any] | None:
|
|
with self.engine.connect() as connection:
|
|
row = connection.execute(
|
|
select(self.tables["resumes"]).where(self.tables["resumes"].c.session_id == session_id)
|
|
).mappings().first()
|
|
return dict(row) if row else None
|
|
|
|
def update_resume(self, session_id: str, content: dict[str, Any], *, expected_revision: int) -> dict[str, Any]:
|
|
resumes = self.tables["resumes"]
|
|
with self.engine.begin() as connection:
|
|
result = connection.execute(update(resumes).where(
|
|
resumes.c.session_id == session_id, resumes.c.revision == expected_revision
|
|
).values(content=content, revision=expected_revision + 1, updated_at=_now()))
|
|
if result.rowcount != 1:
|
|
raise RevisionConflict(session_id)
|
|
return self.get_resume(session_id) # type: ignore[return-value]
|
|
|
|
def create_resume_import(
|
|
self,
|
|
session_id: str,
|
|
*,
|
|
import_id: str,
|
|
file_name: str,
|
|
mime_type: str,
|
|
size_bytes: int,
|
|
sha256: str,
|
|
object_key: str,
|
|
document: dict[str, Any] | None,
|
|
field_reviews: list[dict[str, Any]],
|
|
status: str = "awaiting_review",
|
|
error_code: str | None = None,
|
|
) -> dict[str, Any]:
|
|
existing = self.find_resume_import_by_sha256(session_id, sha256)
|
|
if existing:
|
|
return existing
|
|
now = _now()
|
|
values = {
|
|
"id": import_id,
|
|
"session_id": session_id,
|
|
"file_name": file_name,
|
|
"mime_type": mime_type,
|
|
"size_bytes": size_bytes,
|
|
"sha256": sha256,
|
|
"object_key": object_key,
|
|
"status": status,
|
|
"document": document,
|
|
"field_reviews": field_reviews,
|
|
"error_code": error_code,
|
|
"created_at": now,
|
|
"updated_at": now,
|
|
}
|
|
try:
|
|
with self.engine.begin() as connection:
|
|
connection.execute(insert(self.tables["resume_imports"]).values(**values))
|
|
except IntegrityError:
|
|
existing = self.find_resume_import_by_sha256(session_id, sha256)
|
|
if existing:
|
|
return existing
|
|
raise
|
|
return self.get_resume_import(session_id, import_id) # type: ignore[return-value]
|
|
|
|
def get_resume_import(self, session_id: str, import_id: str) -> dict[str, Any] | None:
|
|
imports = self.tables["resume_imports"]
|
|
with self.engine.connect() as connection:
|
|
row = connection.execute(
|
|
select(imports).where(
|
|
imports.c.id == import_id,
|
|
imports.c.session_id == session_id,
|
|
)
|
|
).mappings().first()
|
|
return dict(row) if row else None
|
|
|
|
def find_resume_import_by_sha256(self, session_id: str, sha256: str) -> dict[str, Any] | None:
|
|
imports = self.tables["resume_imports"]
|
|
with self.engine.connect() as connection:
|
|
row = connection.execute(
|
|
select(imports).where(
|
|
imports.c.session_id == session_id,
|
|
imports.c.sha256 == sha256,
|
|
)
|
|
).mappings().first()
|
|
return dict(row) if row else None
|
|
|
|
def update_resume_import_status(
|
|
self,
|
|
session_id: str,
|
|
import_id: str,
|
|
status: str,
|
|
*,
|
|
error_code: str | None = None,
|
|
) -> dict[str, Any]:
|
|
imports = self.tables["resume_imports"]
|
|
with self.engine.begin() as connection:
|
|
result = connection.execute(
|
|
update(imports)
|
|
.where(imports.c.id == import_id, imports.c.session_id == session_id)
|
|
.values(status=status, error_code=error_code, updated_at=_now())
|
|
)
|
|
if result.rowcount != 1:
|
|
raise KeyError(import_id)
|
|
return self.get_resume_import(session_id, import_id) # type: ignore[return-value]
|
|
def delete_session(self, session_id: str) -> bool:
|
|
with self.engine.begin() as connection:
|
|
result = connection.execute(
|
|
self.tables["sessions"].delete().where(self.tables["sessions"].c.id == session_id)
|
|
)
|
|
return result.rowcount == 1
|