generated from kgod/ai-review-template
406 lines
14 KiB
Python
406 lines
14 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import sqlite3
|
|
from contextlib import contextmanager
|
|
from datetime import UTC, datetime
|
|
from pathlib import Path
|
|
from typing import Any, Iterator
|
|
from uuid import uuid4
|
|
|
|
from .models import (
|
|
BusinessResume,
|
|
ComponentBlock,
|
|
ConversationTurn,
|
|
SessionView,
|
|
)
|
|
|
|
|
|
def utc_now() -> str:
|
|
return datetime.now(UTC).isoformat()
|
|
|
|
|
|
class Database:
|
|
def __init__(self, path: str | Path) -> None:
|
|
self.path = str(path)
|
|
if self.path != ":memory:":
|
|
Path(self.path).parent.mkdir(parents=True, exist_ok=True)
|
|
|
|
def connect(self) -> sqlite3.Connection:
|
|
connection = sqlite3.connect(self.path, timeout=10, isolation_level=None)
|
|
connection.row_factory = sqlite3.Row
|
|
connection.execute("PRAGMA foreign_keys = ON")
|
|
connection.execute("PRAGMA busy_timeout = 10000")
|
|
if self.path != ":memory:":
|
|
connection.execute("PRAGMA journal_mode = WAL")
|
|
return connection
|
|
|
|
@contextmanager
|
|
def transaction(self, *, immediate: bool = False) -> Iterator[sqlite3.Connection]:
|
|
connection = self.connect()
|
|
try:
|
|
connection.execute("BEGIN IMMEDIATE" if immediate else "BEGIN")
|
|
yield connection
|
|
connection.commit()
|
|
except Exception:
|
|
connection.rollback()
|
|
raise
|
|
finally:
|
|
connection.close()
|
|
|
|
def initialize(self) -> None:
|
|
with self.transaction(immediate=True) as connection:
|
|
connection.executescript(
|
|
"""
|
|
CREATE TABLE IF NOT EXISTS sessions (
|
|
id TEXT PRIMARY KEY,
|
|
stage TEXT NOT NULL,
|
|
revision INTEGER NOT NULL DEFAULT 0,
|
|
profile_json TEXT NOT NULL,
|
|
draft_id TEXT,
|
|
resume_id TEXT,
|
|
created_at TEXT NOT NULL,
|
|
updated_at TEXT NOT NULL
|
|
);
|
|
|
|
CREATE TABLE IF NOT EXISTS turns (
|
|
id TEXT PRIMARY KEY,
|
|
session_id TEXT NOT NULL REFERENCES sessions(id) ON DELETE CASCADE,
|
|
sequence INTEGER NOT NULL,
|
|
role TEXT NOT NULL,
|
|
content TEXT,
|
|
composer_mode TEXT NOT NULL,
|
|
created_at TEXT NOT NULL,
|
|
UNIQUE(session_id, sequence)
|
|
);
|
|
|
|
CREATE TABLE IF NOT EXISTS blocks (
|
|
id TEXT PRIMARY KEY,
|
|
session_id TEXT NOT NULL REFERENCES sessions(id) ON DELETE CASCADE,
|
|
turn_id TEXT NOT NULL REFERENCES turns(id) ON DELETE CASCADE,
|
|
block_index INTEGER NOT NULL,
|
|
type TEXT NOT NULL,
|
|
lifecycle TEXT NOT NULL,
|
|
data_json TEXT NOT NULL,
|
|
version INTEGER NOT NULL DEFAULT 1,
|
|
created_at TEXT NOT NULL,
|
|
updated_at TEXT NOT NULL,
|
|
UNIQUE(turn_id, block_index)
|
|
);
|
|
|
|
CREATE TABLE IF NOT EXISTS resumes (
|
|
id TEXT PRIMARY KEY,
|
|
session_id TEXT NOT NULL UNIQUE REFERENCES sessions(id) ON DELETE CASCADE,
|
|
idempotency_key TEXT,
|
|
revision INTEGER NOT NULL,
|
|
content_json TEXT NOT NULL,
|
|
created_at TEXT NOT NULL,
|
|
updated_at TEXT NOT NULL
|
|
);
|
|
|
|
CREATE INDEX IF NOT EXISTS idx_turns_session
|
|
ON turns(session_id, sequence);
|
|
CREATE INDEX IF NOT EXISTS idx_blocks_session
|
|
ON blocks(session_id, turn_id, block_index);
|
|
"""
|
|
)
|
|
|
|
def create_session(
|
|
self,
|
|
session_id: str,
|
|
stage: str,
|
|
profile: dict[str, Any],
|
|
initial_turn: dict[str, Any],
|
|
) -> None:
|
|
now = utc_now()
|
|
with self.transaction(immediate=True) as connection:
|
|
connection.execute(
|
|
"""INSERT INTO sessions
|
|
(id, stage, revision, profile_json, created_at, updated_at)
|
|
VALUES (?, ?, 0, ?, ?, ?)""",
|
|
(session_id, stage, json.dumps(profile, ensure_ascii=False), now, now),
|
|
)
|
|
self.insert_turn(connection, session_id=session_id, **initial_turn)
|
|
|
|
def fetch_session(
|
|
self, connection: sqlite3.Connection, session_id: str
|
|
) -> dict[str, Any] | None:
|
|
row = connection.execute(
|
|
"SELECT * FROM sessions WHERE id = ?", (session_id,)
|
|
).fetchone()
|
|
if row is None:
|
|
return None
|
|
result = dict(row)
|
|
result["profile"] = json.loads(result.pop("profile_json"))
|
|
return result
|
|
|
|
def get_session(self, session_id: str) -> dict[str, Any] | None:
|
|
with self.transaction() as connection:
|
|
return self.fetch_session(connection, session_id)
|
|
|
|
def update_session(
|
|
self,
|
|
connection: sqlite3.Connection,
|
|
session_id: str,
|
|
*,
|
|
stage: str,
|
|
profile: dict[str, Any],
|
|
draft_id: str | None = None,
|
|
resume_id: str | None = None,
|
|
increment_revision: bool = True,
|
|
) -> dict[str, Any]:
|
|
current = self.fetch_session(connection, session_id)
|
|
if current is None:
|
|
raise KeyError(session_id)
|
|
revision = current["revision"] + (1 if increment_revision else 0)
|
|
draft_value = draft_id if draft_id is not None else current["draft_id"]
|
|
resume_value = resume_id if resume_id is not None else current["resume_id"]
|
|
connection.execute(
|
|
"""UPDATE sessions
|
|
SET stage = ?, revision = ?, profile_json = ?, draft_id = ?,
|
|
resume_id = ?, updated_at = ?
|
|
WHERE id = ?""",
|
|
(
|
|
stage,
|
|
revision,
|
|
json.dumps(profile, ensure_ascii=False),
|
|
draft_value,
|
|
resume_value,
|
|
utc_now(),
|
|
session_id,
|
|
),
|
|
)
|
|
updated = self.fetch_session(connection, session_id)
|
|
assert updated is not None
|
|
return updated
|
|
|
|
def insert_turn(
|
|
self,
|
|
connection: sqlite3.Connection,
|
|
*,
|
|
session_id: str,
|
|
role: str,
|
|
content: str | None,
|
|
composer_mode: str,
|
|
blocks: list[dict[str, Any]],
|
|
) -> str:
|
|
turn_id = f"turn_{uuid4().hex}"
|
|
sequence = connection.execute(
|
|
"SELECT COALESCE(MAX(sequence), 0) + 1 FROM turns WHERE session_id = ?",
|
|
(session_id,),
|
|
).fetchone()[0]
|
|
now = utc_now()
|
|
connection.execute(
|
|
"""INSERT INTO turns
|
|
(id, session_id, sequence, role, content, composer_mode, created_at)
|
|
VALUES (?, ?, ?, ?, ?, ?, ?)""",
|
|
(turn_id, session_id, sequence, role, content, composer_mode, now),
|
|
)
|
|
for index, block in enumerate(blocks):
|
|
connection.execute(
|
|
"""INSERT INTO blocks
|
|
(id, session_id, turn_id, block_index, type, lifecycle,
|
|
data_json, version, created_at, updated_at)
|
|
VALUES (?, ?, ?, ?, ?, ?, ?, 1, ?, ?)""",
|
|
(
|
|
block.get("id", f"block_{uuid4().hex}"),
|
|
session_id,
|
|
turn_id,
|
|
index,
|
|
block["type"],
|
|
block.get("lifecycle", "active"),
|
|
json.dumps(block.get("data", {}), ensure_ascii=False),
|
|
now,
|
|
now,
|
|
),
|
|
)
|
|
return turn_id
|
|
|
|
def fetch_block(
|
|
self, connection: sqlite3.Connection, session_id: str, block_id: str
|
|
) -> dict[str, Any] | None:
|
|
row = connection.execute(
|
|
"SELECT * FROM blocks WHERE id = ? AND session_id = ?",
|
|
(block_id, session_id),
|
|
).fetchone()
|
|
if row is None:
|
|
return None
|
|
result = dict(row)
|
|
result["data"] = json.loads(result.pop("data_json"))
|
|
return result
|
|
|
|
def update_block(
|
|
self,
|
|
connection: sqlite3.Connection,
|
|
block_id: str,
|
|
*,
|
|
lifecycle: str,
|
|
data: dict[str, Any] | None = None,
|
|
) -> None:
|
|
row = connection.execute(
|
|
"SELECT data_json FROM blocks WHERE id = ?", (block_id,)
|
|
).fetchone()
|
|
if row is None:
|
|
raise KeyError(block_id)
|
|
serialized = row["data_json"] if data is None else json.dumps(data, ensure_ascii=False)
|
|
connection.execute(
|
|
"""UPDATE blocks
|
|
SET lifecycle = ?, data_json = ?, version = version + 1, updated_at = ?
|
|
WHERE id = ?""",
|
|
(lifecycle, serialized, utc_now(), block_id),
|
|
)
|
|
|
|
def supersede_active_components(
|
|
self,
|
|
connection: sqlite3.Connection,
|
|
session_id: str,
|
|
) -> None:
|
|
"""Make component submissions single-use when chat or create advances the flow."""
|
|
connection.execute(
|
|
"""UPDATE blocks
|
|
SET lifecycle = 'superseded', version = version + 1, updated_at = ?
|
|
WHERE session_id = ? AND type = 'component' AND lifecycle = 'active'""",
|
|
(utc_now(), session_id),
|
|
)
|
|
|
|
def get_turn(self, turn_id: str) -> ConversationTurn:
|
|
with self.transaction() as connection:
|
|
return self.fetch_turn(connection, turn_id)
|
|
|
|
def fetch_turn(
|
|
self,
|
|
connection: sqlite3.Connection,
|
|
turn_id: str,
|
|
) -> ConversationTurn:
|
|
row = connection.execute("SELECT * FROM turns WHERE id = ?", (turn_id,)).fetchone()
|
|
if row is None:
|
|
raise KeyError(turn_id)
|
|
return self._turn_from_row(connection, row)
|
|
|
|
def list_turns(self, session_id: str) -> list[ConversationTurn]:
|
|
with self.transaction() as connection:
|
|
rows = connection.execute(
|
|
"SELECT * FROM turns WHERE session_id = ? ORDER BY sequence",
|
|
(session_id,),
|
|
).fetchall()
|
|
return [self._turn_from_row(connection, row) for row in rows]
|
|
|
|
def _turn_from_row(
|
|
self, connection: sqlite3.Connection, row: sqlite3.Row
|
|
) -> ConversationTurn:
|
|
block_rows = connection.execute(
|
|
"SELECT * FROM blocks WHERE turn_id = ? ORDER BY block_index", (row["id"],)
|
|
).fetchall()
|
|
blocks = [
|
|
ComponentBlock(
|
|
id=block["id"],
|
|
type=block["type"],
|
|
lifecycle=block["lifecycle"],
|
|
data=json.loads(block["data_json"]),
|
|
version=block["version"],
|
|
created_at=block["created_at"],
|
|
updated_at=block["updated_at"],
|
|
)
|
|
for block in block_rows
|
|
]
|
|
return ConversationTurn(
|
|
id=row["id"],
|
|
sequence=row["sequence"],
|
|
role=row["role"],
|
|
content=row["content"],
|
|
composer_mode=row["composer_mode"],
|
|
blocks=blocks,
|
|
created_at=row["created_at"],
|
|
)
|
|
|
|
def session_view(self, session: dict[str, Any]) -> SessionView:
|
|
profile = session["profile"]
|
|
phone = profile.get("phone")
|
|
masked_phone = f"{phone[:3]}****{phone[-4:]}" if phone else None
|
|
return SessionView(
|
|
id=session["id"],
|
|
stage=session["stage"],
|
|
revision=session["revision"],
|
|
job_type=profile.get("job_type"),
|
|
anchor_type=profile.get("anchor_type"),
|
|
masked_phone=masked_phone,
|
|
phone_source=profile.get("phone_source"),
|
|
name=profile.get("name"),
|
|
draft_id=session.get("draft_id"),
|
|
resume_id=session.get("resume_id"),
|
|
created_at=session["created_at"],
|
|
updated_at=session["updated_at"],
|
|
)
|
|
|
|
def fetch_resume(
|
|
self, connection: sqlite3.Connection, session_id: str
|
|
) -> dict[str, Any] | None:
|
|
row = connection.execute(
|
|
"SELECT * FROM resumes WHERE session_id = ?", (session_id,)
|
|
).fetchone()
|
|
if row is None:
|
|
return None
|
|
result = dict(row)
|
|
result["content"] = json.loads(result.pop("content_json"))
|
|
return result
|
|
|
|
def insert_resume(
|
|
self,
|
|
connection: sqlite3.Connection,
|
|
*,
|
|
resume_id: str,
|
|
session_id: str,
|
|
idempotency_key: str | None,
|
|
content: dict[str, Any],
|
|
) -> dict[str, Any]:
|
|
now = utc_now()
|
|
connection.execute(
|
|
"""INSERT INTO resumes
|
|
(id, session_id, idempotency_key, revision, content_json, created_at, updated_at)
|
|
VALUES (?, ?, ?, 1, ?, ?, ?)""",
|
|
(
|
|
resume_id,
|
|
session_id,
|
|
idempotency_key,
|
|
json.dumps(content, ensure_ascii=False),
|
|
now,
|
|
now,
|
|
),
|
|
)
|
|
result = self.fetch_resume(connection, session_id)
|
|
assert result is not None
|
|
return result
|
|
|
|
def update_resume(
|
|
self,
|
|
connection: sqlite3.Connection,
|
|
session_id: str,
|
|
content: dict[str, Any],
|
|
) -> dict[str, Any]:
|
|
connection.execute(
|
|
"""UPDATE resumes
|
|
SET revision = revision + 1, content_json = ?, updated_at = ?
|
|
WHERE session_id = ?""",
|
|
(json.dumps(content, ensure_ascii=False), utc_now(), session_id),
|
|
)
|
|
result = self.fetch_resume(connection, session_id)
|
|
if result is None:
|
|
raise KeyError(session_id)
|
|
return result
|
|
|
|
@staticmethod
|
|
def resume_view(resume: dict[str, Any]) -> BusinessResume:
|
|
return BusinessResume(
|
|
id=resume["id"],
|
|
session_id=resume["session_id"],
|
|
revision=resume["revision"],
|
|
content=resume["content"],
|
|
created_at=resume["created_at"],
|
|
updated_at=resume["updated_at"],
|
|
)
|
|
|
|
def delete_session(self, session_id: str) -> bool:
|
|
with self.transaction(immediate=True) as connection:
|
|
cursor = connection.execute("DELETE FROM sessions WHERE id = ?", (session_id,))
|
|
return cursor.rowcount > 0
|