generated from kgod/ai-review-template
feat: initialize resume agent with OfferPai sync
This commit is contained in:
Codex
@@ -0,0 +1,223 @@
|
||||
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
|
||||
Reference in New Issue
Block a user