Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
203bdfe89b | ||
|
|
edcacfcecc |
No files matched your search
@@ -1,30 +0,0 @@
|
||||
name: fix_eval_warnings
|
||||
on:
|
||||
workflow_run:
|
||||
workflows: ["build_systems"]
|
||||
types: [completed]
|
||||
|
||||
jobs:
|
||||
check-warnings:
|
||||
if: >-
|
||||
github.event.workflow_run.conclusion != 'cancelled' &&
|
||||
github.event.workflow_run.head_branch == 'main' &&
|
||||
(github.event.workflow_run.event == 'push' || github.event.workflow_run.event == 'schedule')
|
||||
runs-on: self-hosted
|
||||
permissions:
|
||||
contents: write
|
||||
pull-requests: write
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
|
||||
- name: Fix eval warnings
|
||||
env:
|
||||
GH_TOKEN: ${{ secrets.GH_TOKEN_FOR_UPDATES }}
|
||||
run: >-
|
||||
nix develop .#devShells.x86_64-linux.default -c
|
||||
python -m python.eval_warnings.main
|
||||
--run-id "${{ github.event.workflow_run.id }}"
|
||||
--repo "${{ github.repository }}"
|
||||
--ollama-url "${{ secrets.OLLAMA_URL }}"
|
||||
--run-url "${{ github.event.workflow_run.html_url }}"
|
||||
@@ -165,11 +165,3 @@ test.*
|
||||
|
||||
# syncthing
|
||||
.stfolder
|
||||
|
||||
# Frontend build output
|
||||
frontend/dist/
|
||||
frontend/node_modules/
|
||||
|
||||
# data dir for training, validation, and testing
|
||||
data/
|
||||
config.toml
|
||||
Vendored
+1
-6
@@ -77,7 +77,6 @@
|
||||
"esphome",
|
||||
"extest",
|
||||
"fadvise",
|
||||
"fastfetch",
|
||||
"fastforwardteam",
|
||||
"FASTFOX",
|
||||
"ffmpegthumbnailer",
|
||||
@@ -167,6 +166,7 @@
|
||||
"mypy",
|
||||
"ncdu",
|
||||
"nemo",
|
||||
"neofetch",
|
||||
"nerdfonts",
|
||||
"netdev",
|
||||
"netdevs",
|
||||
@@ -232,7 +232,6 @@
|
||||
"pyopenweathermap",
|
||||
"pyownet",
|
||||
"pytest",
|
||||
"qalculate",
|
||||
"quicksuggest",
|
||||
"radarr",
|
||||
"readahead",
|
||||
@@ -257,7 +256,6 @@
|
||||
"sessionmaker",
|
||||
"sessionstore",
|
||||
"shellcheck",
|
||||
"signalbot",
|
||||
"signon",
|
||||
"Signons",
|
||||
"skia",
|
||||
@@ -289,7 +287,6 @@
|
||||
"topstories",
|
||||
"treefmt",
|
||||
"twimg",
|
||||
"typedmonarchmoney",
|
||||
"typer",
|
||||
"uaccess",
|
||||
"ubiquiti",
|
||||
@@ -307,8 +304,6 @@
|
||||
"useragent",
|
||||
"usernamehw",
|
||||
"userprefs",
|
||||
"vaninventory",
|
||||
"vdev",
|
||||
"vfat",
|
||||
"victron",
|
||||
"virt",
|
||||
|
||||
@@ -3,10 +3,3 @@
|
||||
- use treefmt to format all files
|
||||
- make python code ruff compliant
|
||||
- use pytest to test python code
|
||||
- always use the minimum amount of complexity
|
||||
- if judgment calls are easy to reverse make them. if not ask me first
|
||||
- Match existing code style.
|
||||
- Use builtin helpers getenv() over os.environ.get.
|
||||
- Prefer single-purpose functions over “do everything” helpers.
|
||||
- Avoid compatibility branches like PG_USER and POSTGRESQL_URL unless requested.
|
||||
- Keep helpers only if reused or they simplify the code otherwise inline.
|
||||
@@ -33,8 +33,6 @@ in
|
||||
];
|
||||
warn-dirty = false;
|
||||
flake-registry = ""; # disable global flake registries
|
||||
connect-timeout = 10;
|
||||
fallback = true;
|
||||
};
|
||||
|
||||
# Add each flake input as a registry and nix_path
|
||||
|
||||
@@ -1,6 +0,0 @@
|
||||
{
|
||||
nix.settings = {
|
||||
trusted-substituters = [ "http://192.168.95.35:5000" ];
|
||||
substituters = [ "http://192.168.95.35:5000/?priority=1&want-mass-query=true" ];
|
||||
};
|
||||
}
|
||||
Generated
+18
-18
@@ -8,11 +8,11 @@
|
||||
},
|
||||
"locked": {
|
||||
"dir": "pkgs/firefox-addons",
|
||||
"lastModified": 1773979456,
|
||||
"narHash": "sha256-9kBMJ5IvxqNlkkj/swmE8uK1Sc7TL/LIRUI958m7uBM=",
|
||||
"lastModified": 1766762570,
|
||||
"narHash": "sha256-Nevsj5NYurwp3I6nSMeh3uirwoinVSbCldqOXu4smms=",
|
||||
"owner": "rycee",
|
||||
"repo": "nur-expressions",
|
||||
"rev": "81e28f47ac18d9e89513929c77e711e657b64851",
|
||||
"rev": "03d7d310ea91d6e4b47ed70aa86c781fcc5b38e1",
|
||||
"type": "gitlab"
|
||||
},
|
||||
"original": {
|
||||
@@ -29,11 +29,11 @@
|
||||
]
|
||||
},
|
||||
"locked": {
|
||||
"lastModified": 1774007980,
|
||||
"narHash": "sha256-FOnZjElEI8pqqCvB6K/1JRHTE8o4rer8driivTpq2uo=",
|
||||
"lastModified": 1766682973,
|
||||
"narHash": "sha256-GKO35onS711ThCxwWcfuvbIBKXwriahGqs+WZuJ3v9E=",
|
||||
"owner": "nix-community",
|
||||
"repo": "home-manager",
|
||||
"rev": "9670de2921812bc4e0452f6e3efd8c859696c183",
|
||||
"rev": "91cdb0e2d574c64fae80d221f4bf09d5592e9ec2",
|
||||
"type": "github"
|
||||
},
|
||||
"original": {
|
||||
@@ -44,11 +44,11 @@
|
||||
},
|
||||
"nixos-hardware": {
|
||||
"locked": {
|
||||
"lastModified": 1774018263,
|
||||
"narHash": "sha256-HHYEwK1A22aSaxv2ibhMMkKvrDGKGlA/qObG4smrSqc=",
|
||||
"lastModified": 1766568855,
|
||||
"narHash": "sha256-UXVtN77D7pzKmzOotFTStgZBqpOcf8cO95FcupWp4Zo=",
|
||||
"owner": "nixos",
|
||||
"repo": "nixos-hardware",
|
||||
"rev": "2d4b4717b2534fad5c715968c1cece04a172b365",
|
||||
"rev": "c5db9569ac9cc70929c268ac461f4003e3e5ca80",
|
||||
"type": "github"
|
||||
},
|
||||
"original": {
|
||||
@@ -60,11 +60,11 @@
|
||||
},
|
||||
"nixpkgs": {
|
||||
"locked": {
|
||||
"lastModified": 1773821835,
|
||||
"narHash": "sha256-TJ3lSQtW0E2JrznGVm8hOQGVpXjJyXY2guAxku2O9A4=",
|
||||
"lastModified": 1766651565,
|
||||
"narHash": "sha256-QEhk0eXgyIqTpJ/ehZKg9IKS7EtlWxF3N7DXy42zPfU=",
|
||||
"owner": "nixos",
|
||||
"repo": "nixpkgs",
|
||||
"rev": "b40629efe5d6ec48dd1efba650c797ddbd39ace0",
|
||||
"rev": "3e2499d5539c16d0d173ba53552a4ff8547f4539",
|
||||
"type": "github"
|
||||
},
|
||||
"original": {
|
||||
@@ -76,11 +76,11 @@
|
||||
},
|
||||
"nixpkgs-master": {
|
||||
"locked": {
|
||||
"lastModified": 1774051532,
|
||||
"narHash": "sha256-d3CGMweyYIcPuTj5BKq+1Lx4zwlgL31nVtN647tOZKo=",
|
||||
"lastModified": 1766794443,
|
||||
"narHash": "sha256-Q8IyTQ3Lu8vX/iqO3U+E4pjLbP1NsqFih6uElf8OYrQ=",
|
||||
"owner": "nixos",
|
||||
"repo": "nixpkgs",
|
||||
"rev": "8620c0b5cc8fbe76502442181be1d0514bc3a1b7",
|
||||
"rev": "088b069b8270ee36d83533c86b9f91d924d185d9",
|
||||
"type": "github"
|
||||
},
|
||||
"original": {
|
||||
@@ -125,11 +125,11 @@
|
||||
]
|
||||
},
|
||||
"locked": {
|
||||
"lastModified": 1773889674,
|
||||
"narHash": "sha256-+ycaiVAk3MEshJTg35cBTUa0MizGiS+bgpYw/f8ohkg=",
|
||||
"lastModified": 1766289575,
|
||||
"narHash": "sha256-BOKCwOQQIP4p9z8DasT5r+qjri3x7sPCOq+FTjY8Z+o=",
|
||||
"owner": "Mic92",
|
||||
"repo": "sops-nix",
|
||||
"rev": "29b6519f3e0780452bca0ac0be4584f04ac16cc5",
|
||||
"rev": "9836912e37aef546029e48c8749834735a6b9dad",
|
||||
"type": "github"
|
||||
},
|
||||
"original": {
|
||||
|
||||
@@ -1,24 +0,0 @@
|
||||
# Logs
|
||||
logs
|
||||
*.log
|
||||
npm-debug.log*
|
||||
yarn-debug.log*
|
||||
yarn-error.log*
|
||||
pnpm-debug.log*
|
||||
lerna-debug.log*
|
||||
|
||||
node_modules
|
||||
dist
|
||||
dist-ssr
|
||||
*.local
|
||||
|
||||
# Editor directories and files
|
||||
.vscode/*
|
||||
!.vscode/extensions.json
|
||||
.idea
|
||||
.DS_Store
|
||||
*.suo
|
||||
*.ntvs*
|
||||
*.njsproj
|
||||
*.sln
|
||||
*.sw?
|
||||
+3
-14
@@ -16,36 +16,25 @@
|
||||
};
|
||||
|
||||
python-env = final: _prev: {
|
||||
my_python = final.python314.withPackages (
|
||||
my_python = final.python313.withPackages (
|
||||
ps: with ps; [
|
||||
alembic
|
||||
apprise
|
||||
apscheduler
|
||||
fastapi
|
||||
fastapi-cli
|
||||
httpx
|
||||
huggingface-hub
|
||||
mypy
|
||||
orjson
|
||||
polars
|
||||
psycopg
|
||||
pydantic
|
||||
pyfakefs
|
||||
pytest
|
||||
pytest-cov
|
||||
pytest-mock
|
||||
pytest-xdist
|
||||
python-multipart
|
||||
requests
|
||||
ruff
|
||||
scalene
|
||||
sqlalchemy
|
||||
sqlalchemy
|
||||
tenacity
|
||||
textual
|
||||
tiktoken
|
||||
tinytuya
|
||||
typer
|
||||
websockets
|
||||
types-requests
|
||||
]
|
||||
);
|
||||
};
|
||||
|
||||
+5
-51
@@ -7,31 +7,7 @@ requires-python = "~=3.13.0"
|
||||
readme = "README.md"
|
||||
license = "MIT"
|
||||
# these dependencies are a best effort and aren't guaranteed to work
|
||||
# for up-to-date dependencies, see overlays/default.nix
|
||||
dependencies = [
|
||||
"alembic",
|
||||
"apprise",
|
||||
"apscheduler",
|
||||
"huggingface-hub",
|
||||
"httpx",
|
||||
"python-multipart",
|
||||
"polars",
|
||||
"psycopg[binary]",
|
||||
"pydantic",
|
||||
"pyyaml",
|
||||
"sqlalchemy",
|
||||
"typer",
|
||||
"websockets",
|
||||
]
|
||||
|
||||
[project.scripts]
|
||||
database = "python.database_cli:app"
|
||||
van-inventory = "python.van_inventory.main:serve"
|
||||
prompt-bench = "python.prompt_bench.main:cli"
|
||||
prompt-bench-download = "python.prompt_bench.downloader:cli"
|
||||
finetune = "python.prompt_bench.finetune:cli"
|
||||
finetune-container = "python.prompt_bench.finetune_container:cli"
|
||||
build-finetune-dataset = "python.prompt_bench.build_finetune_dataset:cli"
|
||||
dependencies = ["apprise", "apscheduler", "polars", "requests", "typer"]
|
||||
|
||||
[dependency-groups]
|
||||
dev = [
|
||||
@@ -42,6 +18,7 @@ dev = [
|
||||
"pytest-xdist",
|
||||
"pytest",
|
||||
"ruff",
|
||||
"types-requests",
|
||||
]
|
||||
|
||||
[tool.ruff]
|
||||
@@ -61,43 +38,21 @@ lint.ignore = [
|
||||
[tool.ruff.lint.per-file-ignores]
|
||||
|
||||
"tests/**" = [
|
||||
"ANN", # (perm) type annotations not needed in tests
|
||||
"D", # (perm) docstrings not needed in tests
|
||||
"PLR2004", # (perm) magic values are fine in test assertions
|
||||
"S101", # (perm) pytest needs asserts
|
||||
"S101", # (perm) pytest needs asserts
|
||||
]
|
||||
"python/stuff/**" = [
|
||||
"python/random/**" = [
|
||||
"T201", # (perm) I don't care about print statements dir
|
||||
]
|
||||
"python/testing/**" = [
|
||||
"T201", # (perm) I don't care about print statements dir
|
||||
"ERA001", # (perm) I don't care about print statements dir
|
||||
]
|
||||
|
||||
"python/splendor/**" = [
|
||||
"S311", # (perm) there is no security issue here
|
||||
"T201", # (perm) I don't care about print statements dir
|
||||
"PLR2004", # (temps) need to think about this
|
||||
]
|
||||
"python/orm/**" = [
|
||||
"TC003", # (perm) this creates issues because sqlalchemy uses these at runtime
|
||||
]
|
||||
"python/congress_tracker/**" = [
|
||||
"TC003", # (perm) this creates issues because sqlalchemy uses these at runtime
|
||||
]
|
||||
"python/eval_warnings/**" = [
|
||||
"S607", # (perm) gh and git are expected on PATH in the runner environment
|
||||
]
|
||||
"python/prompt_bench/**" = [
|
||||
"FBT002", # (perm) typer requires boolean defaults for --flag/--no-flag options
|
||||
"PLR0913", # (perm) typer CLIs naturally have many parameters
|
||||
"S607", # (perm) docker and nvidia-smi are expected on PATH
|
||||
]
|
||||
"python/alembic/**" = [
|
||||
"INP001", # (perm) this creates LSP issues for alembic
|
||||
]
|
||||
"python/signal_bot/**" = [
|
||||
"D107", # (perm) class docstrings cover __init__
|
||||
]
|
||||
|
||||
[tool.ruff.lint.pydocstyle]
|
||||
convention = "google"
|
||||
@@ -121,5 +76,4 @@ exclude_lines = [
|
||||
|
||||
[tool.pytest.ini_options]
|
||||
addopts = "-n auto -ra"
|
||||
testpaths = ["tests"]
|
||||
# --cov=system_tools --cov-report=term-missing --cov-report=xml --cov-report=html --cov-branch
|
||||
-1417
File diff suppressed because it is too large.
Load diff
-50
@@ -1,50 +0,0 @@
|
||||
"""adding FailedIngestion.
|
||||
|
||||
Revision ID: 2f43120e3ffc
|
||||
Revises: f99be864fe69
|
||||
Create Date: 2026-03-24 23:46:17.277897
|
||||
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import sqlalchemy as sa
|
||||
from alembic import op
|
||||
|
||||
from python.orm import DataScienceDevBase
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Sequence
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision: str = "2f43120e3ffc"
|
||||
down_revision: str | None = "f99be864fe69"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
schema = DataScienceDevBase.schema_name
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
"""Upgrade."""
|
||||
# ### commands auto generated by Alembic - please adjust! ###
|
||||
op.create_table(
|
||||
"failed_ingestion",
|
||||
sa.Column("raw_line", sa.Text(), nullable=False),
|
||||
sa.Column("error", sa.Text(), nullable=False),
|
||||
sa.Column("id", sa.Integer(), nullable=False),
|
||||
sa.Column("created", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False),
|
||||
sa.Column("updated", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False),
|
||||
sa.PrimaryKeyConstraint("id", name=op.f("pk_failed_ingestion")),
|
||||
schema=schema,
|
||||
)
|
||||
# ### end Alembic commands ###
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
"""Downgrade."""
|
||||
# ### commands auto generated by Alembic - please adjust! ###
|
||||
op.drop_table("failed_ingestion", schema=schema)
|
||||
# ### end Alembic commands ###
|
||||
-2770
File diff suppressed because it is too large.
Load diff
-1391
File diff suppressed because it is too large.
Load diff
@@ -1,72 +0,0 @@
|
||||
"""Attach all partition tables to the posts parent table.
|
||||
|
||||
Alembic autogenerate creates partition tables as standalone tables but does not
|
||||
emit the ALTER TABLE ... ATTACH PARTITION statements needed for PostgreSQL to
|
||||
route inserts to the correct partition.
|
||||
|
||||
Revision ID: a1b2c3d4e5f6
|
||||
Revises: 605b1794838f
|
||||
Create Date: 2026-03-25 10:00:00.000000
|
||||
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from alembic import op
|
||||
from sqlalchemy import text
|
||||
|
||||
from python.orm import DataScienceDevBase
|
||||
from python.orm.data_science_dev.posts.partitions import (
|
||||
PARTITION_END_YEAR,
|
||||
PARTITION_START_YEAR,
|
||||
iso_weeks_in_year,
|
||||
week_bounds,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Sequence
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision: str = "a1b2c3d4e5f6"
|
||||
down_revision: str | None = "605b1794838f"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
schema = DataScienceDevBase.schema_name
|
||||
|
||||
ALREADY_ATTACHED_QUERY = text("""
|
||||
SELECT inhrelid::regclass::text
|
||||
FROM pg_inherits
|
||||
WHERE inhparent = :parent::regclass
|
||||
""")
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
"""Attach all weekly partition tables to the posts parent table."""
|
||||
connection = op.get_bind()
|
||||
already_attached = {row[0] for row in connection.execute(ALREADY_ATTACHED_QUERY, {"parent": f"{schema}.posts"})}
|
||||
|
||||
for year in range(PARTITION_START_YEAR, PARTITION_END_YEAR + 1):
|
||||
for week in range(1, iso_weeks_in_year(year) + 1):
|
||||
table_name = f"posts_{year}_{week:02d}"
|
||||
qualified_name = f"{schema}.{table_name}"
|
||||
if qualified_name in already_attached:
|
||||
continue
|
||||
start, end = week_bounds(year, week)
|
||||
start_str = start.strftime("%Y-%m-%d %H:%M:%S")
|
||||
end_str = end.strftime("%Y-%m-%d %H:%M:%S")
|
||||
op.execute(
|
||||
f"ALTER TABLE {schema}.posts "
|
||||
f"ATTACH PARTITION {qualified_name} "
|
||||
f"FOR VALUES FROM ('{start_str}') TO ('{end_str}')"
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
"""Detach all weekly partition tables from the posts parent table."""
|
||||
for year in range(PARTITION_START_YEAR, PARTITION_END_YEAR + 1):
|
||||
for week in range(1, iso_weeks_in_year(year) + 1):
|
||||
table_name = f"posts_{year}_{week:02d}"
|
||||
op.execute(f"ALTER TABLE {schema}.posts DETACH PARTITION {schema}.{table_name}")
|
||||
-153
@@ -1,153 +0,0 @@
|
||||
"""adding congress data.
|
||||
|
||||
Revision ID: 83bfc8af92d8
|
||||
Revises: a1b2c3d4e5f6
|
||||
Create Date: 2026-03-27 10:43:02.324510
|
||||
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import sqlalchemy as sa
|
||||
from alembic import op
|
||||
|
||||
from python.orm import DataScienceDevBase
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Sequence
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision: str = "83bfc8af92d8"
|
||||
down_revision: str | None = "a1b2c3d4e5f6"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
schema = DataScienceDevBase.schema_name
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
"""Upgrade."""
|
||||
# ### commands auto generated by Alembic - please adjust! ###
|
||||
op.create_table(
|
||||
"bill",
|
||||
sa.Column("congress", sa.Integer(), nullable=False),
|
||||
sa.Column("bill_type", sa.String(), nullable=False),
|
||||
sa.Column("number", sa.Integer(), nullable=False),
|
||||
sa.Column("title", sa.String(), nullable=True),
|
||||
sa.Column("title_short", sa.String(), nullable=True),
|
||||
sa.Column("official_title", sa.String(), nullable=True),
|
||||
sa.Column("status", sa.String(), nullable=True),
|
||||
sa.Column("status_at", sa.Date(), nullable=True),
|
||||
sa.Column("sponsor_bioguide_id", sa.String(), nullable=True),
|
||||
sa.Column("subjects_top_term", sa.String(), nullable=True),
|
||||
sa.Column("id", sa.Integer(), nullable=False),
|
||||
sa.Column("created", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False),
|
||||
sa.Column("updated", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False),
|
||||
sa.PrimaryKeyConstraint("id", name=op.f("pk_bill")),
|
||||
sa.UniqueConstraint("congress", "bill_type", "number", name="uq_bill_congress_type_number"),
|
||||
schema=schema,
|
||||
)
|
||||
op.create_index("ix_bill_congress", "bill", ["congress"], unique=False, schema=schema)
|
||||
op.create_table(
|
||||
"legislator",
|
||||
sa.Column("bioguide_id", sa.Text(), nullable=False),
|
||||
sa.Column("thomas_id", sa.String(), nullable=True),
|
||||
sa.Column("lis_id", sa.String(), nullable=True),
|
||||
sa.Column("govtrack_id", sa.Integer(), nullable=True),
|
||||
sa.Column("opensecrets_id", sa.String(), nullable=True),
|
||||
sa.Column("fec_ids", sa.String(), nullable=True),
|
||||
sa.Column("first_name", sa.String(), nullable=False),
|
||||
sa.Column("last_name", sa.String(), nullable=False),
|
||||
sa.Column("official_full_name", sa.String(), nullable=True),
|
||||
sa.Column("nickname", sa.String(), nullable=True),
|
||||
sa.Column("birthday", sa.Date(), nullable=True),
|
||||
sa.Column("gender", sa.String(), nullable=True),
|
||||
sa.Column("current_party", sa.String(), nullable=True),
|
||||
sa.Column("current_state", sa.String(), nullable=True),
|
||||
sa.Column("current_district", sa.Integer(), nullable=True),
|
||||
sa.Column("current_chamber", sa.String(), nullable=True),
|
||||
sa.Column("id", sa.Integer(), nullable=False),
|
||||
sa.Column("created", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False),
|
||||
sa.Column("updated", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False),
|
||||
sa.PrimaryKeyConstraint("id", name=op.f("pk_legislator")),
|
||||
schema=schema,
|
||||
)
|
||||
op.create_index(op.f("ix_legislator_bioguide_id"), "legislator", ["bioguide_id"], unique=True, schema=schema)
|
||||
op.create_table(
|
||||
"bill_text",
|
||||
sa.Column("bill_id", sa.Integer(), nullable=False),
|
||||
sa.Column("version_code", sa.String(), nullable=False),
|
||||
sa.Column("version_name", sa.String(), nullable=True),
|
||||
sa.Column("text_content", sa.String(), nullable=True),
|
||||
sa.Column("date", sa.Date(), nullable=True),
|
||||
sa.Column("id", sa.Integer(), nullable=False),
|
||||
sa.Column("created", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False),
|
||||
sa.Column("updated", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False),
|
||||
sa.ForeignKeyConstraint(
|
||||
["bill_id"], [f"{schema}.bill.id"], name=op.f("fk_bill_text_bill_id_bill"), ondelete="CASCADE"
|
||||
),
|
||||
sa.PrimaryKeyConstraint("id", name=op.f("pk_bill_text")),
|
||||
sa.UniqueConstraint("bill_id", "version_code", name="uq_bill_text_bill_id_version_code"),
|
||||
schema=schema,
|
||||
)
|
||||
op.create_table(
|
||||
"vote",
|
||||
sa.Column("congress", sa.Integer(), nullable=False),
|
||||
sa.Column("chamber", sa.String(), nullable=False),
|
||||
sa.Column("session", sa.Integer(), nullable=False),
|
||||
sa.Column("number", sa.Integer(), nullable=False),
|
||||
sa.Column("vote_type", sa.String(), nullable=True),
|
||||
sa.Column("question", sa.String(), nullable=True),
|
||||
sa.Column("result", sa.String(), nullable=True),
|
||||
sa.Column("result_text", sa.String(), nullable=True),
|
||||
sa.Column("vote_date", sa.Date(), nullable=False),
|
||||
sa.Column("yea_count", sa.Integer(), nullable=True),
|
||||
sa.Column("nay_count", sa.Integer(), nullable=True),
|
||||
sa.Column("not_voting_count", sa.Integer(), nullable=True),
|
||||
sa.Column("present_count", sa.Integer(), nullable=True),
|
||||
sa.Column("bill_id", sa.Integer(), nullable=True),
|
||||
sa.Column("id", sa.Integer(), nullable=False),
|
||||
sa.Column("created", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False),
|
||||
sa.Column("updated", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False),
|
||||
sa.ForeignKeyConstraint(["bill_id"], [f"{schema}.bill.id"], name=op.f("fk_vote_bill_id_bill")),
|
||||
sa.PrimaryKeyConstraint("id", name=op.f("pk_vote")),
|
||||
sa.UniqueConstraint("congress", "chamber", "session", "number", name="uq_vote_congress_chamber_session_number"),
|
||||
schema=schema,
|
||||
)
|
||||
op.create_index("ix_vote_congress_chamber", "vote", ["congress", "chamber"], unique=False, schema=schema)
|
||||
op.create_index("ix_vote_date", "vote", ["vote_date"], unique=False, schema=schema)
|
||||
op.create_table(
|
||||
"vote_record",
|
||||
sa.Column("vote_id", sa.Integer(), nullable=False),
|
||||
sa.Column("legislator_id", sa.Integer(), nullable=False),
|
||||
sa.Column("position", sa.String(), nullable=False),
|
||||
sa.ForeignKeyConstraint(
|
||||
["legislator_id"],
|
||||
[f"{schema}.legislator.id"],
|
||||
name=op.f("fk_vote_record_legislator_id_legislator"),
|
||||
ondelete="CASCADE",
|
||||
),
|
||||
sa.ForeignKeyConstraint(
|
||||
["vote_id"], [f"{schema}.vote.id"], name=op.f("fk_vote_record_vote_id_vote"), ondelete="CASCADE"
|
||||
),
|
||||
sa.PrimaryKeyConstraint("vote_id", "legislator_id", name=op.f("pk_vote_record")),
|
||||
schema=schema,
|
||||
)
|
||||
# ### end Alembic commands ###
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
"""Downgrade."""
|
||||
# ### commands auto generated by Alembic - please adjust! ###
|
||||
op.drop_table("vote_record", schema=schema)
|
||||
op.drop_index("ix_vote_date", table_name="vote", schema=schema)
|
||||
op.drop_index("ix_vote_congress_chamber", table_name="vote", schema=schema)
|
||||
op.drop_table("vote", schema=schema)
|
||||
op.drop_table("bill_text", schema=schema)
|
||||
op.drop_index(op.f("ix_legislator_bioguide_id"), table_name="legislator", schema=schema)
|
||||
op.drop_table("legislator", schema=schema)
|
||||
op.drop_index("ix_bill_congress", table_name="bill", schema=schema)
|
||||
op.drop_table("bill", schema=schema)
|
||||
# ### end Alembic commands ###
|
||||
-58
@@ -1,58 +0,0 @@
|
||||
"""adding LegislatorSocialMedia.
|
||||
|
||||
Revision ID: 5cd7eee3549d
|
||||
Revises: 83bfc8af92d8
|
||||
Create Date: 2026-03-29 11:53:44.224799
|
||||
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import sqlalchemy as sa
|
||||
from alembic import op
|
||||
|
||||
from python.orm import DataScienceDevBase
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Sequence
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision: str = "5cd7eee3549d"
|
||||
down_revision: str | None = "83bfc8af92d8"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
schema = DataScienceDevBase.schema_name
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
"""Upgrade."""
|
||||
# ### commands auto generated by Alembic - please adjust! ###
|
||||
op.create_table(
|
||||
"legislator_social_media",
|
||||
sa.Column("legislator_id", sa.Integer(), nullable=False),
|
||||
sa.Column("platform", sa.String(), nullable=False),
|
||||
sa.Column("account_name", sa.String(), nullable=False),
|
||||
sa.Column("url", sa.String(), nullable=True),
|
||||
sa.Column("source", sa.String(), nullable=False),
|
||||
sa.Column("id", sa.Integer(), nullable=False),
|
||||
sa.Column("created", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False),
|
||||
sa.Column("updated", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False),
|
||||
sa.ForeignKeyConstraint(
|
||||
["legislator_id"],
|
||||
[f"{schema}.legislator.id"],
|
||||
name=op.f("fk_legislator_social_media_legislator_id_legislator"),
|
||||
),
|
||||
sa.PrimaryKeyConstraint("id", name=op.f("pk_legislator_social_media")),
|
||||
schema=schema,
|
||||
)
|
||||
# ### end Alembic commands ###
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
"""Downgrade."""
|
||||
# ### commands auto generated by Alembic - please adjust! ###
|
||||
op.drop_table("legislator_social_media", schema=schema)
|
||||
# ### end Alembic commands ###
|
||||
@@ -1,122 +0,0 @@
|
||||
"""Alembic."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Any, Literal
|
||||
|
||||
from alembic import context
|
||||
from alembic.script import write_hooks
|
||||
from sqlalchemy.schema import CreateSchema
|
||||
|
||||
from python.common import bash_wrapper
|
||||
from python.orm.common import get_postgres_engine
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import MutableMapping
|
||||
|
||||
from sqlalchemy.orm import DeclarativeBase
|
||||
|
||||
config = context.config
|
||||
|
||||
base_class: type[DeclarativeBase] = config.attributes.get("base")
|
||||
if base_class is None:
|
||||
error = "No base class provided. Use the database CLI to run alembic commands."
|
||||
raise RuntimeError(error)
|
||||
|
||||
target_metadata = base_class.metadata
|
||||
logging.basicConfig(
|
||||
level="DEBUG",
|
||||
datefmt="%Y-%m-%dT%H:%M:%S%z",
|
||||
format="%(asctime)s %(levelname)s %(filename)s:%(lineno)d - %(message)s",
|
||||
handlers=[logging.StreamHandler(sys.stdout)],
|
||||
)
|
||||
|
||||
|
||||
@write_hooks.register("dynamic_schema")
|
||||
def dynamic_schema(filename: str, _options: dict[Any, Any]) -> None:
|
||||
"""Dynamic schema."""
|
||||
original_file = Path(filename).read_text()
|
||||
schema_name = base_class.schema_name
|
||||
dynamic_schema_file_part1 = original_file.replace(f"schema='{schema_name}'", "schema=schema")
|
||||
dynamic_schema_file = dynamic_schema_file_part1.replace(f"'{schema_name}.", "f'{schema}.")
|
||||
Path(filename).write_text(dynamic_schema_file)
|
||||
|
||||
|
||||
@write_hooks.register("import_postgresql")
|
||||
def import_postgresql(filename: str, _options: dict[Any, Any]) -> None:
|
||||
"""Add postgresql dialect import when postgresql types are used."""
|
||||
content = Path(filename).read_text()
|
||||
if "postgresql." in content and "from sqlalchemy.dialects import postgresql" not in content:
|
||||
content = content.replace(
|
||||
"import sqlalchemy as sa\n",
|
||||
"import sqlalchemy as sa\nfrom sqlalchemy.dialects import postgresql\n",
|
||||
)
|
||||
Path(filename).write_text(content)
|
||||
|
||||
|
||||
@write_hooks.register("ruff")
|
||||
def ruff_check_and_format(filename: str, _options: dict[Any, Any]) -> None:
|
||||
"""Docstring for ruff_check_and_format."""
|
||||
bash_wrapper(f"ruff check --fix {filename}")
|
||||
bash_wrapper(f"ruff format {filename}")
|
||||
|
||||
|
||||
def include_name(
|
||||
name: str | None,
|
||||
type_: Literal["schema", "table", "column", "index", "unique_constraint", "foreign_key_constraint"],
|
||||
_parent_names: MutableMapping[Literal["schema_name", "table_name", "schema_qualified_table_name"], str | None],
|
||||
) -> bool:
|
||||
"""Filter tables to be included in the migration.
|
||||
|
||||
Args:
|
||||
name (str): The name of the table.
|
||||
type_ (str): The type of the table.
|
||||
_parent_names (MutableMapping): The names of the parent tables.
|
||||
|
||||
Returns:
|
||||
bool: True if the table should be included, False otherwise.
|
||||
|
||||
"""
|
||||
if type_ == "schema":
|
||||
# allows a database with multiple schemas to have separate alembic revisions
|
||||
return name == target_metadata.schema
|
||||
return True
|
||||
|
||||
|
||||
def run_migrations_online() -> None:
|
||||
"""Run migrations in 'online' mode.
|
||||
|
||||
In this scenario we need to create an Engine
|
||||
and associate a connection with the context.
|
||||
|
||||
"""
|
||||
env_prefix = config.attributes.get("env_prefix", "POSTGRES")
|
||||
connectable = get_postgres_engine(name=env_prefix)
|
||||
|
||||
with connectable.connect() as connection:
|
||||
schema = base_class.schema_name
|
||||
if not connectable.dialect.has_schema(connection, schema):
|
||||
answer = input(f"Schema {schema!r} does not exist. Create it? [y/N] ")
|
||||
if answer.lower() != "y":
|
||||
error = f"Schema {schema!r} does not exist. Exiting."
|
||||
raise SystemExit(error)
|
||||
connection.execute(CreateSchema(schema))
|
||||
connection.commit()
|
||||
|
||||
context.configure(
|
||||
connection=connection,
|
||||
target_metadata=target_metadata,
|
||||
include_schemas=True,
|
||||
version_table_schema=schema,
|
||||
include_name=include_name,
|
||||
)
|
||||
|
||||
with context.begin_transaction():
|
||||
context.run_migrations()
|
||||
connection.commit()
|
||||
|
||||
|
||||
run_migrations_online()
|
||||
@@ -1,113 +0,0 @@
|
||||
"""created contact api.
|
||||
|
||||
Revision ID: edd7dd61a3d2
|
||||
Revises:
|
||||
Create Date: 2026-01-11 15:45:59.909266
|
||||
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import sqlalchemy as sa
|
||||
from alembic import op
|
||||
|
||||
from python.orm import RichieBase
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Sequence
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision: str = "edd7dd61a3d2"
|
||||
down_revision: str | None = None
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
schema = RichieBase.schema_name
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
"""Upgrade."""
|
||||
# ### commands auto generated by Alembic - please adjust! ###
|
||||
op.create_table(
|
||||
"contact",
|
||||
sa.Column("name", sa.String(), nullable=False),
|
||||
sa.Column("age", sa.Integer(), nullable=True),
|
||||
sa.Column("bio", sa.String(), nullable=True),
|
||||
sa.Column("current_job", sa.String(), nullable=True),
|
||||
sa.Column("gender", sa.String(), nullable=True),
|
||||
sa.Column("goals", sa.String(), nullable=True),
|
||||
sa.Column("legal_name", sa.String(), nullable=True),
|
||||
sa.Column("profile_pic", sa.String(), nullable=True),
|
||||
sa.Column("safe_conversation_starters", sa.String(), nullable=True),
|
||||
sa.Column("self_sufficiency_score", sa.Integer(), nullable=True),
|
||||
sa.Column("social_structure_style", sa.String(), nullable=True),
|
||||
sa.Column("ssn", sa.String(), nullable=True),
|
||||
sa.Column("suffix", sa.String(), nullable=True),
|
||||
sa.Column("timezone", sa.String(), nullable=True),
|
||||
sa.Column("topics_to_avoid", sa.String(), nullable=True),
|
||||
sa.Column("id", sa.Integer(), nullable=False),
|
||||
sa.Column("created", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False),
|
||||
sa.Column("updated", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False),
|
||||
sa.PrimaryKeyConstraint("id", name=op.f("pk_contact")),
|
||||
schema=schema,
|
||||
)
|
||||
op.create_table(
|
||||
"need",
|
||||
sa.Column("name", sa.String(), nullable=False),
|
||||
sa.Column("description", sa.String(), nullable=True),
|
||||
sa.Column("id", sa.Integer(), nullable=False),
|
||||
sa.Column("created", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False),
|
||||
sa.Column("updated", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False),
|
||||
sa.PrimaryKeyConstraint("id", name=op.f("pk_need")),
|
||||
schema=schema,
|
||||
)
|
||||
op.create_table(
|
||||
"contact_need",
|
||||
sa.Column("contact_id", sa.Integer(), nullable=False),
|
||||
sa.Column("need_id", sa.Integer(), nullable=False),
|
||||
sa.ForeignKeyConstraint(
|
||||
["contact_id"],
|
||||
[f"{schema}.contact.id"],
|
||||
name=op.f("fk_contact_need_contact_id_contact"),
|
||||
ondelete="CASCADE",
|
||||
),
|
||||
sa.ForeignKeyConstraint(
|
||||
["need_id"], [f"{schema}.need.id"], name=op.f("fk_contact_need_need_id_need"), ondelete="CASCADE"
|
||||
),
|
||||
sa.PrimaryKeyConstraint("contact_id", "need_id", name=op.f("pk_contact_need")),
|
||||
schema=schema,
|
||||
)
|
||||
op.create_table(
|
||||
"contact_relationship",
|
||||
sa.Column("contact_id", sa.Integer(), nullable=False),
|
||||
sa.Column("related_contact_id", sa.Integer(), nullable=False),
|
||||
sa.Column("relationship_type", sa.String(length=100), nullable=False),
|
||||
sa.Column("closeness_weight", sa.Integer(), nullable=False),
|
||||
sa.ForeignKeyConstraint(
|
||||
["contact_id"],
|
||||
[f"{schema}.contact.id"],
|
||||
name=op.f("fk_contact_relationship_contact_id_contact"),
|
||||
ondelete="CASCADE",
|
||||
),
|
||||
sa.ForeignKeyConstraint(
|
||||
["related_contact_id"],
|
||||
[f"{schema}.contact.id"],
|
||||
name=op.f("fk_contact_relationship_related_contact_id_contact"),
|
||||
ondelete="CASCADE",
|
||||
),
|
||||
sa.PrimaryKeyConstraint("contact_id", "related_contact_id", name=op.f("pk_contact_relationship")),
|
||||
schema=schema,
|
||||
)
|
||||
# ### end Alembic commands ###
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
"""Downgrade."""
|
||||
# ### commands auto generated by Alembic - please adjust! ###
|
||||
op.drop_table("contact_relationship", schema=schema)
|
||||
op.drop_table("contact_need", schema=schema)
|
||||
op.drop_table("need", schema=schema)
|
||||
op.drop_table("contact", schema=schema)
|
||||
# ### end Alembic commands ###
|
||||
-135
@@ -1,135 +0,0 @@
|
||||
"""add congress tracker tables.
|
||||
|
||||
Revision ID: 3f71565e38de
|
||||
Revises: edd7dd61a3d2
|
||||
Create Date: 2026-02-12 16:36:09.457303
|
||||
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import sqlalchemy as sa
|
||||
from alembic import op
|
||||
|
||||
from python.orm import RichieBase
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Sequence
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision: str = "3f71565e38de"
|
||||
down_revision: str | None = "edd7dd61a3d2"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
schema = RichieBase.schema_name
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
"""Upgrade."""
|
||||
# ### commands auto generated by Alembic - please adjust! ###
|
||||
op.create_table(
|
||||
"bill",
|
||||
sa.Column("congress", sa.Integer(), nullable=False),
|
||||
sa.Column("bill_type", sa.String(), nullable=False),
|
||||
sa.Column("number", sa.Integer(), nullable=False),
|
||||
sa.Column("title", sa.String(), nullable=True),
|
||||
sa.Column("title_short", sa.String(), nullable=True),
|
||||
sa.Column("official_title", sa.String(), nullable=True),
|
||||
sa.Column("status", sa.String(), nullable=True),
|
||||
sa.Column("status_at", sa.Date(), nullable=True),
|
||||
sa.Column("sponsor_bioguide_id", sa.String(), nullable=True),
|
||||
sa.Column("subjects_top_term", sa.String(), nullable=True),
|
||||
sa.Column("id", sa.Integer(), nullable=False),
|
||||
sa.Column("created", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False),
|
||||
sa.Column("updated", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False),
|
||||
sa.PrimaryKeyConstraint("id", name=op.f("pk_bill")),
|
||||
sa.UniqueConstraint("congress", "bill_type", "number", name="uq_bill_congress_type_number"),
|
||||
schema=schema,
|
||||
)
|
||||
op.create_index("ix_bill_congress", "bill", ["congress"], unique=False, schema=schema)
|
||||
op.create_table(
|
||||
"legislator",
|
||||
sa.Column("bioguide_id", sa.Text(), nullable=False),
|
||||
sa.Column("thomas_id", sa.String(), nullable=True),
|
||||
sa.Column("lis_id", sa.String(), nullable=True),
|
||||
sa.Column("govtrack_id", sa.Integer(), nullable=True),
|
||||
sa.Column("opensecrets_id", sa.String(), nullable=True),
|
||||
sa.Column("fec_ids", sa.String(), nullable=True),
|
||||
sa.Column("first_name", sa.String(), nullable=False),
|
||||
sa.Column("last_name", sa.String(), nullable=False),
|
||||
sa.Column("official_full_name", sa.String(), nullable=True),
|
||||
sa.Column("nickname", sa.String(), nullable=True),
|
||||
sa.Column("birthday", sa.Date(), nullable=True),
|
||||
sa.Column("gender", sa.String(), nullable=True),
|
||||
sa.Column("current_party", sa.String(), nullable=True),
|
||||
sa.Column("current_state", sa.String(), nullable=True),
|
||||
sa.Column("current_district", sa.Integer(), nullable=True),
|
||||
sa.Column("current_chamber", sa.String(), nullable=True),
|
||||
sa.Column("id", sa.Integer(), nullable=False),
|
||||
sa.Column("created", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False),
|
||||
sa.Column("updated", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False),
|
||||
sa.PrimaryKeyConstraint("id", name=op.f("pk_legislator")),
|
||||
schema=schema,
|
||||
)
|
||||
op.create_index(op.f("ix_legislator_bioguide_id"), "legislator", ["bioguide_id"], unique=True, schema=schema)
|
||||
op.create_table(
|
||||
"vote",
|
||||
sa.Column("congress", sa.Integer(), nullable=False),
|
||||
sa.Column("chamber", sa.String(), nullable=False),
|
||||
sa.Column("session", sa.Integer(), nullable=False),
|
||||
sa.Column("number", sa.Integer(), nullable=False),
|
||||
sa.Column("vote_type", sa.String(), nullable=True),
|
||||
sa.Column("question", sa.String(), nullable=True),
|
||||
sa.Column("result", sa.String(), nullable=True),
|
||||
sa.Column("result_text", sa.String(), nullable=True),
|
||||
sa.Column("vote_date", sa.Date(), nullable=False),
|
||||
sa.Column("yea_count", sa.Integer(), nullable=True),
|
||||
sa.Column("nay_count", sa.Integer(), nullable=True),
|
||||
sa.Column("not_voting_count", sa.Integer(), nullable=True),
|
||||
sa.Column("present_count", sa.Integer(), nullable=True),
|
||||
sa.Column("bill_id", sa.Integer(), nullable=True),
|
||||
sa.Column("id", sa.Integer(), nullable=False),
|
||||
sa.Column("created", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False),
|
||||
sa.Column("updated", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False),
|
||||
sa.ForeignKeyConstraint(["bill_id"], [f"{schema}.bill.id"], name=op.f("fk_vote_bill_id_bill")),
|
||||
sa.PrimaryKeyConstraint("id", name=op.f("pk_vote")),
|
||||
sa.UniqueConstraint("congress", "chamber", "session", "number", name="uq_vote_congress_chamber_session_number"),
|
||||
schema=schema,
|
||||
)
|
||||
op.create_index("ix_vote_congress_chamber", "vote", ["congress", "chamber"], unique=False, schema=schema)
|
||||
op.create_index("ix_vote_date", "vote", ["vote_date"], unique=False, schema=schema)
|
||||
op.create_table(
|
||||
"vote_record",
|
||||
sa.Column("vote_id", sa.Integer(), nullable=False),
|
||||
sa.Column("legislator_id", sa.Integer(), nullable=False),
|
||||
sa.Column("position", sa.String(), nullable=False),
|
||||
sa.ForeignKeyConstraint(
|
||||
["legislator_id"],
|
||||
[f"{schema}.legislator.id"],
|
||||
name=op.f("fk_vote_record_legislator_id_legislator"),
|
||||
ondelete="CASCADE",
|
||||
),
|
||||
sa.ForeignKeyConstraint(
|
||||
["vote_id"], [f"{schema}.vote.id"], name=op.f("fk_vote_record_vote_id_vote"), ondelete="CASCADE"
|
||||
),
|
||||
sa.PrimaryKeyConstraint("vote_id", "legislator_id", name=op.f("pk_vote_record")),
|
||||
schema=schema,
|
||||
)
|
||||
# ### end Alembic commands ###
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
"""Downgrade."""
|
||||
# ### commands auto generated by Alembic - please adjust! ###
|
||||
op.drop_table("vote_record", schema=schema)
|
||||
op.drop_index("ix_vote_date", table_name="vote", schema=schema)
|
||||
op.drop_index("ix_vote_congress_chamber", table_name="vote", schema=schema)
|
||||
op.drop_table("vote", schema=schema)
|
||||
op.drop_index(op.f("ix_legislator_bioguide_id"), table_name="legislator", schema=schema)
|
||||
op.drop_table("legislator", schema=schema)
|
||||
op.drop_index("ix_bill_congress", table_name="bill", schema=schema)
|
||||
op.drop_table("bill", schema=schema)
|
||||
# ### end Alembic commands ###
|
||||
-58
@@ -1,58 +0,0 @@
|
||||
"""adding SignalDevice for DeviceRegistry for signal bot.
|
||||
|
||||
Revision ID: 4c410c16e39c
|
||||
Revises: 3f71565e38de
|
||||
Create Date: 2026-03-09 14:51:24.228976
|
||||
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import sqlalchemy as sa
|
||||
from alembic import op
|
||||
from sqlalchemy.dialects import postgresql
|
||||
|
||||
from python.orm import RichieBase
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Sequence
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision: str = "4c410c16e39c"
|
||||
down_revision: str | None = "3f71565e38de"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
schema = RichieBase.schema_name
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
"""Upgrade."""
|
||||
# ### commands auto generated by Alembic - please adjust! ###
|
||||
op.create_table(
|
||||
"signal_device",
|
||||
sa.Column("phone_number", sa.String(length=50), nullable=False),
|
||||
sa.Column("safety_number", sa.String(), nullable=False),
|
||||
sa.Column(
|
||||
"trust_level",
|
||||
postgresql.ENUM("VERIFIED", "UNVERIFIED", "BLOCKED", name="trust_level", schema=schema),
|
||||
nullable=False,
|
||||
),
|
||||
sa.Column("last_seen", sa.DateTime(timezone=True), nullable=False),
|
||||
sa.Column("id", sa.Integer(), nullable=False),
|
||||
sa.Column("created", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False),
|
||||
sa.Column("updated", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False),
|
||||
sa.PrimaryKeyConstraint("id", name=op.f("pk_signal_device")),
|
||||
sa.UniqueConstraint("phone_number", name=op.f("uq_signal_device_phone_number")),
|
||||
schema=schema,
|
||||
)
|
||||
# ### end Alembic commands ###
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
"""Downgrade."""
|
||||
# ### commands auto generated by Alembic - please adjust! ###
|
||||
op.drop_table("signal_device", schema=schema)
|
||||
# ### end Alembic commands ###
|
||||
@@ -1,41 +0,0 @@
|
||||
"""fixed safety number logic.
|
||||
|
||||
Revision ID: 99fec682516c
|
||||
Revises: 4c410c16e39c
|
||||
Create Date: 2026-03-09 16:25:25.085806
|
||||
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import sqlalchemy as sa
|
||||
from alembic import op
|
||||
|
||||
from python.orm import RichieBase
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Sequence
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision: str = "99fec682516c"
|
||||
down_revision: str | None = "4c410c16e39c"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
schema = RichieBase.schema_name
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
"""Upgrade."""
|
||||
# ### commands auto generated by Alembic - please adjust! ###
|
||||
op.alter_column("signal_device", "safety_number", existing_type=sa.VARCHAR(), nullable=True, schema=schema)
|
||||
# ### end Alembic commands ###
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
"""Downgrade."""
|
||||
# ### commands auto generated by Alembic - please adjust! ###
|
||||
op.alter_column("signal_device", "safety_number", existing_type=sa.VARCHAR(), nullable=False, schema=schema)
|
||||
# ### end Alembic commands ###
|
||||
-54
@@ -1,54 +0,0 @@
|
||||
"""add dead_letter_message table.
|
||||
|
||||
Revision ID: a1b2c3d4e5f6
|
||||
Revises: 99fec682516c
|
||||
Create Date: 2026-03-10 12:00:00.000000
|
||||
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import sqlalchemy as sa
|
||||
from alembic import op
|
||||
from sqlalchemy.dialects import postgresql
|
||||
|
||||
from python.orm import RichieBase
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Sequence
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision: str = "a1b2c3d4e5f6"
|
||||
down_revision: str | None = "99fec682516c"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
schema = RichieBase.schema_name
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
"""Upgrade."""
|
||||
op.create_table(
|
||||
"dead_letter_message",
|
||||
sa.Column("source", sa.String(), nullable=False),
|
||||
sa.Column("message", sa.Text(), nullable=False),
|
||||
sa.Column("received_at", sa.DateTime(timezone=True), nullable=False),
|
||||
sa.Column(
|
||||
"status",
|
||||
postgresql.ENUM("UNPROCESSED", "PROCESSED", name="message_status", schema=schema),
|
||||
nullable=False,
|
||||
),
|
||||
sa.Column("id", sa.Integer(), nullable=False),
|
||||
sa.Column("created", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False),
|
||||
sa.Column("updated", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False),
|
||||
sa.PrimaryKeyConstraint("id", name=op.f("pk_dead_letter_message")),
|
||||
schema=schema,
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
"""Downgrade."""
|
||||
op.drop_table("dead_letter_message", schema=schema)
|
||||
op.execute(sa.text(f"DROP TYPE IF EXISTS {schema}.message_status"))
|
||||
-66
@@ -1,66 +0,0 @@
|
||||
"""adding roles to signal devices.
|
||||
|
||||
Revision ID: 2ef7ba690159
|
||||
Revises: a1b2c3d4e5f6
|
||||
Create Date: 2026-03-16 19:22:38.020350
|
||||
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import sqlalchemy as sa
|
||||
from alembic import op
|
||||
|
||||
from python.orm import RichieBase
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Sequence
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision: str = "2ef7ba690159"
|
||||
down_revision: str | None = "a1b2c3d4e5f6"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
schema = RichieBase.schema_name
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
"""Upgrade."""
|
||||
# ### commands auto generated by Alembic - please adjust! ###
|
||||
op.create_table(
|
||||
"role",
|
||||
sa.Column("name", sa.String(length=50), nullable=False),
|
||||
sa.Column("id", sa.SmallInteger(), nullable=False),
|
||||
sa.Column("created", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False),
|
||||
sa.Column("updated", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False),
|
||||
sa.PrimaryKeyConstraint("id", name=op.f("pk_role")),
|
||||
sa.UniqueConstraint("name", name=op.f("uq_role_name")),
|
||||
schema=schema,
|
||||
)
|
||||
op.create_table(
|
||||
"device_role",
|
||||
sa.Column("device_id", sa.Integer(), nullable=False),
|
||||
sa.Column("role_id", sa.SmallInteger(), nullable=False),
|
||||
sa.Column("id", sa.Integer(), nullable=False),
|
||||
sa.Column("created", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False),
|
||||
sa.Column("updated", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False),
|
||||
sa.ForeignKeyConstraint(
|
||||
["device_id"], [f"{schema}.signal_device.id"], name=op.f("fk_device_role_device_id_signal_device")
|
||||
),
|
||||
sa.ForeignKeyConstraint(["role_id"], [f"{schema}.role.id"], name=op.f("fk_device_role_role_id_role")),
|
||||
sa.PrimaryKeyConstraint("id", name=op.f("pk_device_role")),
|
||||
sa.UniqueConstraint("device_id", "role_id", name="uq_device_role_device_role"),
|
||||
schema=schema,
|
||||
)
|
||||
# ### end Alembic commands ###
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
"""Downgrade."""
|
||||
# ### commands auto generated by Alembic - please adjust! ###
|
||||
op.drop_table("device_role", schema=schema)
|
||||
op.drop_table("role", schema=schema)
|
||||
# ### end Alembic commands ###
|
||||
-171
@@ -1,171 +0,0 @@
|
||||
"""seprating signal_bot database.
|
||||
|
||||
Revision ID: 6b275323f435
|
||||
Revises: 2ef7ba690159
|
||||
Create Date: 2026-03-18 08:34:28.785885
|
||||
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import sqlalchemy as sa
|
||||
from alembic import op
|
||||
from sqlalchemy.dialects import postgresql
|
||||
|
||||
from python.orm import RichieBase
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Sequence
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision: str = "6b275323f435"
|
||||
down_revision: str | None = "2ef7ba690159"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
schema = RichieBase.schema_name
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
"""Upgrade."""
|
||||
# ### commands auto generated by Alembic - please adjust! ###
|
||||
op.drop_table("device_role", schema=schema)
|
||||
op.drop_table("signal_device", schema=schema)
|
||||
op.drop_table("role", schema=schema)
|
||||
op.drop_table("dead_letter_message", schema=schema)
|
||||
# ### end Alembic commands ###
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
"""Downgrade."""
|
||||
# ### commands auto generated by Alembic - please adjust! ###
|
||||
op.create_table(
|
||||
"dead_letter_message",
|
||||
sa.Column("source", sa.VARCHAR(), autoincrement=False, nullable=False),
|
||||
sa.Column("message", sa.TEXT(), autoincrement=False, nullable=False),
|
||||
sa.Column("received_at", postgresql.TIMESTAMP(timezone=True), autoincrement=False, nullable=False),
|
||||
sa.Column(
|
||||
"status",
|
||||
postgresql.ENUM("UNPROCESSED", "PROCESSED", name="message_status", schema=schema),
|
||||
autoincrement=False,
|
||||
nullable=False,
|
||||
),
|
||||
sa.Column("id", sa.INTEGER(), autoincrement=True, nullable=False),
|
||||
sa.Column(
|
||||
"created",
|
||||
postgresql.TIMESTAMP(timezone=True),
|
||||
server_default=sa.text("now()"),
|
||||
autoincrement=False,
|
||||
nullable=False,
|
||||
),
|
||||
sa.Column(
|
||||
"updated",
|
||||
postgresql.TIMESTAMP(timezone=True),
|
||||
server_default=sa.text("now()"),
|
||||
autoincrement=False,
|
||||
nullable=False,
|
||||
),
|
||||
sa.PrimaryKeyConstraint("id", name=op.f("pk_dead_letter_message")),
|
||||
schema=schema,
|
||||
)
|
||||
op.create_table(
|
||||
"role",
|
||||
sa.Column("name", sa.VARCHAR(length=50), autoincrement=False, nullable=False),
|
||||
sa.Column(
|
||||
"id",
|
||||
sa.SMALLINT(),
|
||||
server_default=sa.text(f"nextval('{schema}.role_id_seq'::regclass)"),
|
||||
autoincrement=True,
|
||||
nullable=False,
|
||||
),
|
||||
sa.Column(
|
||||
"created",
|
||||
postgresql.TIMESTAMP(timezone=True),
|
||||
server_default=sa.text("now()"),
|
||||
autoincrement=False,
|
||||
nullable=False,
|
||||
),
|
||||
sa.Column(
|
||||
"updated",
|
||||
postgresql.TIMESTAMP(timezone=True),
|
||||
server_default=sa.text("now()"),
|
||||
autoincrement=False,
|
||||
nullable=False,
|
||||
),
|
||||
sa.PrimaryKeyConstraint("id", name=op.f("pk_role")),
|
||||
sa.UniqueConstraint(
|
||||
"name", name=op.f("uq_role_name"), postgresql_include=[], postgresql_nulls_not_distinct=False
|
||||
),
|
||||
schema=schema,
|
||||
)
|
||||
op.create_table(
|
||||
"signal_device",
|
||||
sa.Column("phone_number", sa.VARCHAR(length=50), autoincrement=False, nullable=False),
|
||||
sa.Column("safety_number", sa.VARCHAR(), autoincrement=False, nullable=True),
|
||||
sa.Column(
|
||||
"trust_level",
|
||||
postgresql.ENUM("VERIFIED", "UNVERIFIED", "BLOCKED", name="trust_level", schema=schema),
|
||||
autoincrement=False,
|
||||
nullable=False,
|
||||
),
|
||||
sa.Column("last_seen", postgresql.TIMESTAMP(timezone=True), autoincrement=False, nullable=False),
|
||||
sa.Column("id", sa.INTEGER(), autoincrement=True, nullable=False),
|
||||
sa.Column(
|
||||
"created",
|
||||
postgresql.TIMESTAMP(timezone=True),
|
||||
server_default=sa.text("now()"),
|
||||
autoincrement=False,
|
||||
nullable=False,
|
||||
),
|
||||
sa.Column(
|
||||
"updated",
|
||||
postgresql.TIMESTAMP(timezone=True),
|
||||
server_default=sa.text("now()"),
|
||||
autoincrement=False,
|
||||
nullable=False,
|
||||
),
|
||||
sa.PrimaryKeyConstraint("id", name=op.f("pk_signal_device")),
|
||||
sa.UniqueConstraint(
|
||||
"phone_number",
|
||||
name=op.f("uq_signal_device_phone_number"),
|
||||
postgresql_include=[],
|
||||
postgresql_nulls_not_distinct=False,
|
||||
),
|
||||
schema=schema,
|
||||
)
|
||||
op.create_table(
|
||||
"device_role",
|
||||
sa.Column("device_id", sa.INTEGER(), autoincrement=False, nullable=False),
|
||||
sa.Column("role_id", sa.SMALLINT(), autoincrement=False, nullable=False),
|
||||
sa.Column("id", sa.INTEGER(), autoincrement=True, nullable=False),
|
||||
sa.Column(
|
||||
"created",
|
||||
postgresql.TIMESTAMP(timezone=True),
|
||||
server_default=sa.text("now()"),
|
||||
autoincrement=False,
|
||||
nullable=False,
|
||||
),
|
||||
sa.Column(
|
||||
"updated",
|
||||
postgresql.TIMESTAMP(timezone=True),
|
||||
server_default=sa.text("now()"),
|
||||
autoincrement=False,
|
||||
nullable=False,
|
||||
),
|
||||
sa.ForeignKeyConstraint(
|
||||
["device_id"], [f"{schema}.signal_device.id"], name=op.f("fk_device_role_device_id_signal_device")
|
||||
),
|
||||
sa.ForeignKeyConstraint(["role_id"], [f"{schema}.role.id"], name=op.f("fk_device_role_role_id_role")),
|
||||
sa.PrimaryKeyConstraint("id", name=op.f("pk_device_role")),
|
||||
sa.UniqueConstraint(
|
||||
"device_id",
|
||||
"role_id",
|
||||
name=op.f("uq_device_role_device_role"),
|
||||
postgresql_include=[],
|
||||
postgresql_nulls_not_distinct=False,
|
||||
),
|
||||
schema=schema,
|
||||
)
|
||||
# ### end Alembic commands ###
|
||||
-187
@@ -1,187 +0,0 @@
|
||||
"""removed ds table from richie DB.
|
||||
|
||||
Revision ID: c8a794340928
|
||||
Revises: 6b275323f435
|
||||
Create Date: 2026-03-29 15:29:23.643146
|
||||
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import sqlalchemy as sa
|
||||
from alembic import op
|
||||
from sqlalchemy.dialects import postgresql
|
||||
|
||||
from python.orm import RichieBase
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Sequence
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision: str = "c8a794340928"
|
||||
down_revision: str | None = "6b275323f435"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
schema = RichieBase.schema_name
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
"""Upgrade."""
|
||||
# ### commands auto generated by Alembic - please adjust! ###
|
||||
op.drop_table("vote_record", schema=schema)
|
||||
op.drop_index(op.f("ix_vote_congress_chamber"), table_name="vote", schema=schema)
|
||||
op.drop_index(op.f("ix_vote_date"), table_name="vote", schema=schema)
|
||||
op.drop_index(op.f("ix_legislator_bioguide_id"), table_name="legislator", schema=schema)
|
||||
op.drop_table("legislator", schema=schema)
|
||||
op.drop_table("vote", schema=schema)
|
||||
op.drop_index(op.f("ix_bill_congress"), table_name="bill", schema=schema)
|
||||
op.drop_table("bill", schema=schema)
|
||||
# ### end Alembic commands ###
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
"""Downgrade."""
|
||||
# ### commands auto generated by Alembic - please adjust! ###
|
||||
op.create_table(
|
||||
"vote",
|
||||
sa.Column("congress", sa.INTEGER(), autoincrement=False, nullable=False),
|
||||
sa.Column("chamber", sa.VARCHAR(), autoincrement=False, nullable=False),
|
||||
sa.Column("session", sa.INTEGER(), autoincrement=False, nullable=False),
|
||||
sa.Column("number", sa.INTEGER(), autoincrement=False, nullable=False),
|
||||
sa.Column("vote_type", sa.VARCHAR(), autoincrement=False, nullable=True),
|
||||
sa.Column("question", sa.VARCHAR(), autoincrement=False, nullable=True),
|
||||
sa.Column("result", sa.VARCHAR(), autoincrement=False, nullable=True),
|
||||
sa.Column("result_text", sa.VARCHAR(), autoincrement=False, nullable=True),
|
||||
sa.Column("vote_date", sa.DATE(), autoincrement=False, nullable=False),
|
||||
sa.Column("yea_count", sa.INTEGER(), autoincrement=False, nullable=True),
|
||||
sa.Column("nay_count", sa.INTEGER(), autoincrement=False, nullable=True),
|
||||
sa.Column("not_voting_count", sa.INTEGER(), autoincrement=False, nullable=True),
|
||||
sa.Column("present_count", sa.INTEGER(), autoincrement=False, nullable=True),
|
||||
sa.Column("bill_id", sa.INTEGER(), autoincrement=False, nullable=True),
|
||||
sa.Column("id", sa.INTEGER(), autoincrement=True, nullable=False),
|
||||
sa.Column(
|
||||
"created",
|
||||
postgresql.TIMESTAMP(timezone=True),
|
||||
server_default=sa.text("now()"),
|
||||
autoincrement=False,
|
||||
nullable=False,
|
||||
),
|
||||
sa.Column(
|
||||
"updated",
|
||||
postgresql.TIMESTAMP(timezone=True),
|
||||
server_default=sa.text("now()"),
|
||||
autoincrement=False,
|
||||
nullable=False,
|
||||
),
|
||||
sa.ForeignKeyConstraint(["bill_id"], [f"{schema}.bill.id"], name=op.f("fk_vote_bill_id_bill")),
|
||||
sa.PrimaryKeyConstraint("id", name=op.f("pk_vote")),
|
||||
sa.UniqueConstraint(
|
||||
"congress",
|
||||
"chamber",
|
||||
"session",
|
||||
"number",
|
||||
name=op.f("uq_vote_congress_chamber_session_number"),
|
||||
postgresql_include=[],
|
||||
postgresql_nulls_not_distinct=False,
|
||||
),
|
||||
schema=schema,
|
||||
)
|
||||
op.create_index(op.f("ix_vote_date"), "vote", ["vote_date"], unique=False, schema=schema)
|
||||
op.create_index(op.f("ix_vote_congress_chamber"), "vote", ["congress", "chamber"], unique=False, schema=schema)
|
||||
op.create_table(
|
||||
"vote_record",
|
||||
sa.Column("vote_id", sa.INTEGER(), autoincrement=False, nullable=False),
|
||||
sa.Column("legislator_id", sa.INTEGER(), autoincrement=False, nullable=False),
|
||||
sa.Column("position", sa.VARCHAR(), autoincrement=False, nullable=False),
|
||||
sa.ForeignKeyConstraint(
|
||||
["legislator_id"],
|
||||
[f"{schema}.legislator.id"],
|
||||
name=op.f("fk_vote_record_legislator_id_legislator"),
|
||||
ondelete="CASCADE",
|
||||
),
|
||||
sa.ForeignKeyConstraint(
|
||||
["vote_id"], [f"{schema}.vote.id"], name=op.f("fk_vote_record_vote_id_vote"), ondelete="CASCADE"
|
||||
),
|
||||
sa.PrimaryKeyConstraint("vote_id", "legislator_id", name=op.f("pk_vote_record")),
|
||||
schema=schema,
|
||||
)
|
||||
op.create_table(
|
||||
"legislator",
|
||||
sa.Column("bioguide_id", sa.TEXT(), autoincrement=False, nullable=False),
|
||||
sa.Column("thomas_id", sa.VARCHAR(), autoincrement=False, nullable=True),
|
||||
sa.Column("lis_id", sa.VARCHAR(), autoincrement=False, nullable=True),
|
||||
sa.Column("govtrack_id", sa.INTEGER(), autoincrement=False, nullable=True),
|
||||
sa.Column("opensecrets_id", sa.VARCHAR(), autoincrement=False, nullable=True),
|
||||
sa.Column("fec_ids", sa.VARCHAR(), autoincrement=False, nullable=True),
|
||||
sa.Column("first_name", sa.VARCHAR(), autoincrement=False, nullable=False),
|
||||
sa.Column("last_name", sa.VARCHAR(), autoincrement=False, nullable=False),
|
||||
sa.Column("official_full_name", sa.VARCHAR(), autoincrement=False, nullable=True),
|
||||
sa.Column("nickname", sa.VARCHAR(), autoincrement=False, nullable=True),
|
||||
sa.Column("birthday", sa.DATE(), autoincrement=False, nullable=True),
|
||||
sa.Column("gender", sa.VARCHAR(), autoincrement=False, nullable=True),
|
||||
sa.Column("current_party", sa.VARCHAR(), autoincrement=False, nullable=True),
|
||||
sa.Column("current_state", sa.VARCHAR(), autoincrement=False, nullable=True),
|
||||
sa.Column("current_district", sa.INTEGER(), autoincrement=False, nullable=True),
|
||||
sa.Column("current_chamber", sa.VARCHAR(), autoincrement=False, nullable=True),
|
||||
sa.Column("id", sa.INTEGER(), autoincrement=True, nullable=False),
|
||||
sa.Column(
|
||||
"created",
|
||||
postgresql.TIMESTAMP(timezone=True),
|
||||
server_default=sa.text("now()"),
|
||||
autoincrement=False,
|
||||
nullable=False,
|
||||
),
|
||||
sa.Column(
|
||||
"updated",
|
||||
postgresql.TIMESTAMP(timezone=True),
|
||||
server_default=sa.text("now()"),
|
||||
autoincrement=False,
|
||||
nullable=False,
|
||||
),
|
||||
sa.PrimaryKeyConstraint("id", name=op.f("pk_legislator")),
|
||||
schema=schema,
|
||||
)
|
||||
op.create_index(op.f("ix_legislator_bioguide_id"), "legislator", ["bioguide_id"], unique=True, schema=schema)
|
||||
op.create_table(
|
||||
"bill",
|
||||
sa.Column("congress", sa.INTEGER(), autoincrement=False, nullable=False),
|
||||
sa.Column("bill_type", sa.VARCHAR(), autoincrement=False, nullable=False),
|
||||
sa.Column("number", sa.INTEGER(), autoincrement=False, nullable=False),
|
||||
sa.Column("title", sa.VARCHAR(), autoincrement=False, nullable=True),
|
||||
sa.Column("title_short", sa.VARCHAR(), autoincrement=False, nullable=True),
|
||||
sa.Column("official_title", sa.VARCHAR(), autoincrement=False, nullable=True),
|
||||
sa.Column("status", sa.VARCHAR(), autoincrement=False, nullable=True),
|
||||
sa.Column("status_at", sa.DATE(), autoincrement=False, nullable=True),
|
||||
sa.Column("sponsor_bioguide_id", sa.VARCHAR(), autoincrement=False, nullable=True),
|
||||
sa.Column("subjects_top_term", sa.VARCHAR(), autoincrement=False, nullable=True),
|
||||
sa.Column("id", sa.INTEGER(), autoincrement=True, nullable=False),
|
||||
sa.Column(
|
||||
"created",
|
||||
postgresql.TIMESTAMP(timezone=True),
|
||||
server_default=sa.text("now()"),
|
||||
autoincrement=False,
|
||||
nullable=False,
|
||||
),
|
||||
sa.Column(
|
||||
"updated",
|
||||
postgresql.TIMESTAMP(timezone=True),
|
||||
server_default=sa.text("now()"),
|
||||
autoincrement=False,
|
||||
nullable=False,
|
||||
),
|
||||
sa.PrimaryKeyConstraint("id", name=op.f("pk_bill")),
|
||||
sa.UniqueConstraint(
|
||||
"congress",
|
||||
"bill_type",
|
||||
"number",
|
||||
name=op.f("uq_bill_congress_type_number"),
|
||||
postgresql_include=[],
|
||||
postgresql_nulls_not_distinct=False,
|
||||
),
|
||||
schema=schema,
|
||||
)
|
||||
op.create_index(op.f("ix_bill_congress"), "bill", ["congress"], unique=False, schema=schema)
|
||||
# ### end Alembic commands ###
|
||||
@@ -1,36 +0,0 @@
|
||||
"""${message}.
|
||||
|
||||
Revision ID: ${up_revision}
|
||||
Revises: ${down_revision | comma,n}
|
||||
Create Date: ${create_date}
|
||||
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import sqlalchemy as sa
|
||||
|
||||
from alembic import op
|
||||
from python.orm import ${config.attributes["base"].__name__}
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Sequence
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision: str = ${repr(up_revision)}
|
||||
down_revision: str | None = ${repr(down_revision)}
|
||||
branch_labels: str | Sequence[str] | None = ${repr(branch_labels)}
|
||||
depends_on: str | Sequence[str] | None = ${repr(depends_on)}
|
||||
|
||||
schema=${config.attributes["base"].__name__}.schema_name
|
||||
|
||||
def upgrade() -> None:
|
||||
"""Upgrade."""
|
||||
${upgrades if upgrades else "pass"}
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
"""Downgrade."""
|
||||
${downgrades if downgrades else "pass"}
|
||||
-100
@@ -1,100 +0,0 @@
|
||||
"""seprating signal_bot database.
|
||||
|
||||
Revision ID: 6eaf696e07a5
|
||||
Revises:
|
||||
Create Date: 2026-03-17 21:35:37.612672
|
||||
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import sqlalchemy as sa
|
||||
from alembic import op
|
||||
from sqlalchemy.dialects import postgresql
|
||||
|
||||
from python.orm import SignalBotBase
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Sequence
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision: str = "6eaf696e07a5"
|
||||
down_revision: str | None = None
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
schema = SignalBotBase.schema_name
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
"""Upgrade."""
|
||||
# ### commands auto generated by Alembic - please adjust! ###
|
||||
op.create_table(
|
||||
"dead_letter_message",
|
||||
sa.Column("source", sa.String(), nullable=False),
|
||||
sa.Column("message", sa.Text(), nullable=False),
|
||||
sa.Column("received_at", sa.DateTime(timezone=True), nullable=False),
|
||||
sa.Column(
|
||||
"status", postgresql.ENUM("UNPROCESSED", "PROCESSED", name="message_status", schema=schema), nullable=False
|
||||
),
|
||||
sa.Column("id", sa.Integer(), nullable=False),
|
||||
sa.Column("created", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False),
|
||||
sa.Column("updated", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False),
|
||||
sa.PrimaryKeyConstraint("id", name=op.f("pk_dead_letter_message")),
|
||||
schema=schema,
|
||||
)
|
||||
op.create_table(
|
||||
"role",
|
||||
sa.Column("name", sa.String(length=50), nullable=False),
|
||||
sa.Column("id", sa.SmallInteger(), nullable=False),
|
||||
sa.Column("created", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False),
|
||||
sa.Column("updated", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False),
|
||||
sa.PrimaryKeyConstraint("id", name=op.f("pk_role")),
|
||||
sa.UniqueConstraint("name", name=op.f("uq_role_name")),
|
||||
schema=schema,
|
||||
)
|
||||
op.create_table(
|
||||
"signal_device",
|
||||
sa.Column("phone_number", sa.String(length=50), nullable=False),
|
||||
sa.Column("safety_number", sa.String(), nullable=True),
|
||||
sa.Column(
|
||||
"trust_level",
|
||||
postgresql.ENUM("VERIFIED", "UNVERIFIED", "BLOCKED", name="trust_level", schema=schema),
|
||||
nullable=False,
|
||||
),
|
||||
sa.Column("last_seen", sa.DateTime(timezone=True), nullable=False),
|
||||
sa.Column("id", sa.Integer(), nullable=False),
|
||||
sa.Column("created", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False),
|
||||
sa.Column("updated", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False),
|
||||
sa.PrimaryKeyConstraint("id", name=op.f("pk_signal_device")),
|
||||
sa.UniqueConstraint("phone_number", name=op.f("uq_signal_device_phone_number")),
|
||||
schema=schema,
|
||||
)
|
||||
op.create_table(
|
||||
"device_role",
|
||||
sa.Column("device_id", sa.Integer(), nullable=False),
|
||||
sa.Column("role_id", sa.SmallInteger(), nullable=False),
|
||||
sa.Column("id", sa.Integer(), nullable=False),
|
||||
sa.Column("created", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False),
|
||||
sa.Column("updated", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False),
|
||||
sa.ForeignKeyConstraint(
|
||||
["device_id"], [f"{schema}.signal_device.id"], name=op.f("fk_device_role_device_id_signal_device")
|
||||
),
|
||||
sa.ForeignKeyConstraint(["role_id"], [f"{schema}.role.id"], name=op.f("fk_device_role_role_id_role")),
|
||||
sa.PrimaryKeyConstraint("id", name=op.f("pk_device_role")),
|
||||
sa.UniqueConstraint("device_id", "role_id", name="uq_device_role_device_role"),
|
||||
schema=schema,
|
||||
)
|
||||
# ### end Alembic commands ###
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
"""Downgrade."""
|
||||
# ### commands auto generated by Alembic - please adjust! ###
|
||||
op.drop_table("device_role", schema=schema)
|
||||
op.drop_table("signal_device", schema=schema)
|
||||
op.drop_table("role", schema=schema)
|
||||
op.drop_table("dead_letter_message", schema=schema)
|
||||
# ### end Alembic commands ###
|
||||
@@ -1,72 +0,0 @@
|
||||
"""test.
|
||||
|
||||
Revision ID: 66bdd532bcab
|
||||
Revises: 6eaf696e07a5
|
||||
Create Date: 2026-03-18 19:21:14.561568
|
||||
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import sqlalchemy as sa
|
||||
from alembic import op
|
||||
from sqlalchemy.dialects import postgresql
|
||||
|
||||
from python.orm import SignalBotBase
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Sequence
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision: str = "66bdd532bcab"
|
||||
down_revision: str | None = "6eaf696e07a5"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
schema = SignalBotBase.schema_name
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
"""Upgrade."""
|
||||
# ### commands auto generated by Alembic - please adjust! ###
|
||||
op.alter_column(
|
||||
"dead_letter_message",
|
||||
"status",
|
||||
existing_type=postgresql.ENUM("UNPROCESSED", "PROCESSED", name="message_status", schema=schema),
|
||||
type_=sa.Enum("UNPROCESSED", "PROCESSED", name="message_status", native_enum=False),
|
||||
existing_nullable=False,
|
||||
schema=schema,
|
||||
)
|
||||
op.alter_column(
|
||||
"signal_device",
|
||||
"trust_level",
|
||||
existing_type=postgresql.ENUM("VERIFIED", "UNVERIFIED", "BLOCKED", name="trust_level", schema=schema),
|
||||
type_=sa.Enum("VERIFIED", "UNVERIFIED", "BLOCKED", name="trust_level", native_enum=False),
|
||||
existing_nullable=False,
|
||||
schema=schema,
|
||||
)
|
||||
# ### end Alembic commands ###
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
"""Downgrade."""
|
||||
# ### commands auto generated by Alembic - please adjust! ###
|
||||
op.alter_column(
|
||||
"signal_device",
|
||||
"trust_level",
|
||||
existing_type=sa.Enum("VERIFIED", "UNVERIFIED", "BLOCKED", name="trust_level", native_enum=False),
|
||||
type_=postgresql.ENUM("VERIFIED", "UNVERIFIED", "BLOCKED", name="trust_level", schema=schema),
|
||||
existing_nullable=False,
|
||||
schema=schema,
|
||||
)
|
||||
op.alter_column(
|
||||
"dead_letter_message",
|
||||
"status",
|
||||
existing_type=sa.Enum("UNPROCESSED", "PROCESSED", name="message_status", native_enum=False),
|
||||
type_=postgresql.ENUM("UNPROCESSED", "PROCESSED", name="message_status", schema=schema),
|
||||
existing_nullable=False,
|
||||
schema=schema,
|
||||
)
|
||||
# ### end Alembic commands ###
|
||||
-80
@@ -1,80 +0,0 @@
|
||||
"""starting van invintory.
|
||||
|
||||
Revision ID: 15e733499804
|
||||
Revises:
|
||||
Create Date: 2026-03-08 00:18:20.759720
|
||||
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import sqlalchemy as sa
|
||||
from alembic import op
|
||||
|
||||
from python.orm import VanInventoryBase
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Sequence
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision: str = "15e733499804"
|
||||
down_revision: str | None = None
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
schema = VanInventoryBase.schema_name
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
"""Upgrade."""
|
||||
# ### commands auto generated by Alembic - please adjust! ###
|
||||
op.create_table(
|
||||
"items",
|
||||
sa.Column("name", sa.String(), nullable=False),
|
||||
sa.Column("quantity", sa.Float(), nullable=False),
|
||||
sa.Column("unit", sa.String(), nullable=False),
|
||||
sa.Column("category", sa.String(), nullable=True),
|
||||
sa.Column("id", sa.Integer(), nullable=False),
|
||||
sa.Column("created", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False),
|
||||
sa.Column("updated", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False),
|
||||
sa.PrimaryKeyConstraint("id", name=op.f("pk_items")),
|
||||
sa.UniqueConstraint("name", name=op.f("uq_items_name")),
|
||||
schema=schema,
|
||||
)
|
||||
op.create_table(
|
||||
"meals",
|
||||
sa.Column("name", sa.String(), nullable=False),
|
||||
sa.Column("instructions", sa.String(), nullable=True),
|
||||
sa.Column("id", sa.Integer(), nullable=False),
|
||||
sa.Column("created", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False),
|
||||
sa.Column("updated", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False),
|
||||
sa.PrimaryKeyConstraint("id", name=op.f("pk_meals")),
|
||||
sa.UniqueConstraint("name", name=op.f("uq_meals_name")),
|
||||
schema=schema,
|
||||
)
|
||||
op.create_table(
|
||||
"meal_ingredients",
|
||||
sa.Column("meal_id", sa.Integer(), nullable=False),
|
||||
sa.Column("item_id", sa.Integer(), nullable=False),
|
||||
sa.Column("quantity_needed", sa.Float(), nullable=False),
|
||||
sa.Column("id", sa.Integer(), nullable=False),
|
||||
sa.Column("created", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False),
|
||||
sa.Column("updated", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False),
|
||||
sa.ForeignKeyConstraint(["item_id"], [f"{schema}.items.id"], name=op.f("fk_meal_ingredients_item_id_items")),
|
||||
sa.ForeignKeyConstraint(["meal_id"], [f"{schema}.meals.id"], name=op.f("fk_meal_ingredients_meal_id_meals")),
|
||||
sa.PrimaryKeyConstraint("id", name=op.f("pk_meal_ingredients")),
|
||||
sa.UniqueConstraint("meal_id", "item_id", name=op.f("uq_meal_ingredients_meal_id")),
|
||||
schema=schema,
|
||||
)
|
||||
# ### end Alembic commands ###
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
"""Downgrade."""
|
||||
# ### commands auto generated by Alembic - please adjust! ###
|
||||
op.drop_table("meal_ingredients", schema=schema)
|
||||
op.drop_table("meals", schema=schema)
|
||||
op.drop_table("items", schema=schema)
|
||||
# ### end Alembic commands ###
|
||||
@@ -1 +0,0 @@
|
||||
"""FastAPI applications."""
|
||||
@@ -1,16 +0,0 @@
|
||||
"""FastAPI dependencies."""
|
||||
|
||||
from collections.abc import Iterator
|
||||
from typing import Annotated
|
||||
|
||||
from fastapi import Depends, Request
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
|
||||
def get_db(request: Request) -> Iterator[Session]:
|
||||
"""Get database session from app state."""
|
||||
with Session(request.app.state.engine) as session:
|
||||
yield session
|
||||
|
||||
|
||||
DbSession = Annotated[Session, Depends(get_db)]
|
||||
@@ -1,52 +0,0 @@
|
||||
"""FastAPI interface for Contact database."""
|
||||
|
||||
import logging
|
||||
from collections.abc import AsyncIterator
|
||||
from contextlib import asynccontextmanager
|
||||
from typing import Annotated
|
||||
|
||||
import typer
|
||||
import uvicorn
|
||||
from fastapi import FastAPI
|
||||
|
||||
from python.api.middleware import ZstdMiddleware
|
||||
from python.api.routers import contact_router, views_router
|
||||
from python.common import configure_logger
|
||||
from python.orm.common import get_postgres_engine
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def create_app() -> FastAPI:
|
||||
"""Create and configure the FastAPI application."""
|
||||
|
||||
@asynccontextmanager
|
||||
async def lifespan(app: FastAPI) -> AsyncIterator[None]:
|
||||
"""Manage application lifespan."""
|
||||
app.state.engine = get_postgres_engine()
|
||||
yield
|
||||
app.state.engine.dispose()
|
||||
|
||||
app = FastAPI(title="Contact Database API", lifespan=lifespan)
|
||||
app.add_middleware(ZstdMiddleware)
|
||||
|
||||
app.include_router(contact_router)
|
||||
app.include_router(views_router)
|
||||
|
||||
return app
|
||||
|
||||
|
||||
def serve(
|
||||
host: Annotated[str, typer.Option("--host", "-h", help="Host to bind to")],
|
||||
port: Annotated[int, typer.Option("--port", "-p", help="Port to bind to")] = 8000,
|
||||
log_level: Annotated[str, typer.Option("--log-level", "-l", help="Log level")] = "INFO",
|
||||
) -> None:
|
||||
"""Start the Contact API server."""
|
||||
configure_logger(log_level)
|
||||
|
||||
app = create_app()
|
||||
uvicorn.run(app, host=host, port=port)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
typer.run(serve)
|
||||
@@ -1,49 +0,0 @@
|
||||
"""Middleware for the FastAPI application."""
|
||||
|
||||
from compression import zstd
|
||||
from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint
|
||||
from starlette.requests import Request
|
||||
from starlette.responses import Response
|
||||
|
||||
MINIMUM_RESPONSE_SIZE = 500
|
||||
|
||||
|
||||
class ZstdMiddleware(BaseHTTPMiddleware):
|
||||
"""Middleware that compresses responses with zstd when the client supports it."""
|
||||
|
||||
async def dispatch(self, request: Request, call_next: RequestResponseEndpoint) -> Response:
|
||||
"""Compress the response with zstd if the client accepts it."""
|
||||
accepted_encodings = request.headers.get("accept-encoding", "")
|
||||
if "zstd" not in accepted_encodings:
|
||||
return await call_next(request)
|
||||
|
||||
response = await call_next(request)
|
||||
|
||||
if response.headers.get("content-encoding") or "text/event-stream" in response.headers.get("content-type", ""):
|
||||
return response
|
||||
|
||||
body = b""
|
||||
async for chunk in response.body_iterator:
|
||||
body += chunk if isinstance(chunk, bytes) else chunk.encode()
|
||||
|
||||
if len(body) < MINIMUM_RESPONSE_SIZE:
|
||||
return Response(
|
||||
content=body,
|
||||
status_code=response.status_code,
|
||||
headers=dict(response.headers),
|
||||
media_type=response.media_type,
|
||||
)
|
||||
|
||||
compressed = zstd.compress(body)
|
||||
|
||||
headers = dict(response.headers)
|
||||
headers["content-encoding"] = "zstd"
|
||||
headers["content-length"] = str(len(compressed))
|
||||
headers.pop("transfer-encoding", None)
|
||||
|
||||
return Response(
|
||||
content=compressed,
|
||||
status_code=response.status_code,
|
||||
headers=headers,
|
||||
media_type=response.media_type,
|
||||
)
|
||||
@@ -1,6 +0,0 @@
|
||||
"""API routers."""
|
||||
|
||||
from python.api.routers.contact import router as contact_router
|
||||
from python.api.routers.views import router as views_router
|
||||
|
||||
__all__ = ["contact_router", "views_router"]
|
||||
@@ -1,481 +0,0 @@
|
||||
"""Contact API router."""
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
from fastapi import APIRouter, HTTPException, Request
|
||||
from fastapi.responses import HTMLResponse
|
||||
from fastapi.templating import Jinja2Templates
|
||||
from pydantic import BaseModel
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy.orm import selectinload
|
||||
|
||||
from python.api.dependencies import DbSession
|
||||
from python.orm.richie.contact import Contact, ContactRelationship, Need, RelationshipType
|
||||
|
||||
TEMPLATES_DIR = Path(__file__).parent.parent / "templates"
|
||||
templates = Jinja2Templates(directory=TEMPLATES_DIR)
|
||||
|
||||
|
||||
def _is_htmx(request: Request) -> bool:
|
||||
"""Check if the request is from HTMX."""
|
||||
return request.headers.get("HX-Request") == "true"
|
||||
|
||||
|
||||
class NeedBase(BaseModel):
|
||||
"""Base schema for Need."""
|
||||
|
||||
name: str
|
||||
description: str | None = None
|
||||
|
||||
|
||||
class NeedCreate(NeedBase):
|
||||
"""Schema for creating a Need."""
|
||||
|
||||
|
||||
class NeedResponse(NeedBase):
|
||||
"""Schema for Need response."""
|
||||
|
||||
id: int
|
||||
|
||||
model_config = {"from_attributes": True}
|
||||
|
||||
|
||||
class ContactRelationshipCreate(BaseModel):
|
||||
"""Schema for creating a contact relationship."""
|
||||
|
||||
related_contact_id: int
|
||||
relationship_type: RelationshipType
|
||||
closeness_weight: int | None = None
|
||||
|
||||
|
||||
class ContactRelationshipUpdate(BaseModel):
|
||||
"""Schema for updating a contact relationship."""
|
||||
|
||||
relationship_type: RelationshipType | None = None
|
||||
closeness_weight: int | None = None
|
||||
|
||||
|
||||
class ContactRelationshipResponse(BaseModel):
|
||||
"""Schema for contact relationship response."""
|
||||
|
||||
contact_id: int
|
||||
related_contact_id: int
|
||||
relationship_type: str
|
||||
closeness_weight: int
|
||||
|
||||
model_config = {"from_attributes": True}
|
||||
|
||||
|
||||
class RelationshipTypeInfo(BaseModel):
|
||||
"""Information about a relationship type."""
|
||||
|
||||
value: str
|
||||
display_name: str
|
||||
default_weight: int
|
||||
|
||||
|
||||
class GraphNode(BaseModel):
|
||||
"""Node in the relationship graph."""
|
||||
|
||||
id: int
|
||||
name: str
|
||||
current_job: str | None = None
|
||||
|
||||
|
||||
class GraphEdge(BaseModel):
|
||||
"""Edge in the relationship graph."""
|
||||
|
||||
source: int
|
||||
target: int
|
||||
relationship_type: str
|
||||
closeness_weight: int
|
||||
|
||||
|
||||
class GraphData(BaseModel):
|
||||
"""Complete graph data for visualization."""
|
||||
|
||||
nodes: list[GraphNode]
|
||||
edges: list[GraphEdge]
|
||||
|
||||
|
||||
class ContactBase(BaseModel):
|
||||
"""Base schema for Contact."""
|
||||
|
||||
name: str
|
||||
age: int | None = None
|
||||
bio: str | None = None
|
||||
current_job: str | None = None
|
||||
gender: str | None = None
|
||||
goals: str | None = None
|
||||
legal_name: str | None = None
|
||||
profile_pic: str | None = None
|
||||
safe_conversation_starters: str | None = None
|
||||
self_sufficiency_score: int | None = None
|
||||
social_structure_style: str | None = None
|
||||
ssn: str | None = None
|
||||
suffix: str | None = None
|
||||
timezone: str | None = None
|
||||
topics_to_avoid: str | None = None
|
||||
|
||||
|
||||
class ContactCreate(ContactBase):
|
||||
"""Schema for creating a Contact."""
|
||||
|
||||
need_ids: list[int] = []
|
||||
|
||||
|
||||
class ContactUpdate(BaseModel):
|
||||
"""Schema for updating a Contact."""
|
||||
|
||||
name: str | None = None
|
||||
age: int | None = None
|
||||
bio: str | None = None
|
||||
current_job: str | None = None
|
||||
gender: str | None = None
|
||||
goals: str | None = None
|
||||
legal_name: str | None = None
|
||||
profile_pic: str | None = None
|
||||
safe_conversation_starters: str | None = None
|
||||
self_sufficiency_score: int | None = None
|
||||
social_structure_style: str | None = None
|
||||
ssn: str | None = None
|
||||
suffix: str | None = None
|
||||
timezone: str | None = None
|
||||
topics_to_avoid: str | None = None
|
||||
need_ids: list[int] | None = None
|
||||
|
||||
|
||||
class ContactResponse(ContactBase):
|
||||
"""Schema for Contact response with relationships."""
|
||||
|
||||
id: int
|
||||
needs: list[NeedResponse] = []
|
||||
related_to: list[ContactRelationshipResponse] = []
|
||||
related_from: list[ContactRelationshipResponse] = []
|
||||
|
||||
model_config = {"from_attributes": True}
|
||||
|
||||
|
||||
class ContactListResponse(ContactBase):
|
||||
"""Schema for Contact list response."""
|
||||
|
||||
id: int
|
||||
|
||||
model_config = {"from_attributes": True}
|
||||
|
||||
|
||||
router = APIRouter(prefix="/api", tags=["contacts"])
|
||||
|
||||
|
||||
@router.post("/needs", response_model=NeedResponse)
|
||||
def create_need(need: NeedCreate, db: DbSession) -> Need:
|
||||
"""Create a new need."""
|
||||
db_need = Need(name=need.name, description=need.description)
|
||||
db.add(db_need)
|
||||
db.commit()
|
||||
db.refresh(db_need)
|
||||
return db_need
|
||||
|
||||
|
||||
@router.get("/needs", response_model=list[NeedResponse])
|
||||
def list_needs(db: DbSession) -> list[Need]:
|
||||
"""List all needs."""
|
||||
return list(db.scalars(select(Need)).all())
|
||||
|
||||
|
||||
@router.get("/needs/{need_id}", response_model=NeedResponse)
|
||||
def get_need(need_id: int, db: DbSession) -> Need:
|
||||
"""Get a need by ID."""
|
||||
need = db.get(Need, need_id)
|
||||
if not need:
|
||||
raise HTTPException(status_code=404, detail="Need not found")
|
||||
return need
|
||||
|
||||
|
||||
@router.delete("/needs/{need_id}", response_model=None)
|
||||
def delete_need(need_id: int, request: Request, db: DbSession) -> dict[str, bool] | HTMLResponse:
|
||||
"""Delete a need by ID."""
|
||||
need = db.get(Need, need_id)
|
||||
if not need:
|
||||
raise HTTPException(status_code=404, detail="Need not found")
|
||||
db.delete(need)
|
||||
db.commit()
|
||||
if _is_htmx(request):
|
||||
return HTMLResponse("")
|
||||
return {"deleted": True}
|
||||
|
||||
|
||||
@router.post("/contacts", response_model=ContactResponse)
|
||||
def create_contact(contact: ContactCreate, db: DbSession) -> Contact:
|
||||
"""Create a new contact."""
|
||||
need_ids = contact.need_ids
|
||||
contact_data = contact.model_dump(exclude={"need_ids"})
|
||||
db_contact = Contact(**contact_data)
|
||||
|
||||
if need_ids:
|
||||
needs = list(db.scalars(select(Need).where(Need.id.in_(need_ids))).all())
|
||||
db_contact.needs = needs
|
||||
|
||||
db.add(db_contact)
|
||||
db.commit()
|
||||
db.refresh(db_contact)
|
||||
return db_contact
|
||||
|
||||
|
||||
@router.get("/contacts", response_model=list[ContactListResponse])
|
||||
def list_contacts(
|
||||
db: DbSession,
|
||||
skip: int = 0,
|
||||
limit: int = 100,
|
||||
) -> list[Contact]:
|
||||
"""List all contacts with pagination."""
|
||||
return list(db.scalars(select(Contact).offset(skip).limit(limit)).all())
|
||||
|
||||
|
||||
@router.get("/contacts/{contact_id}", response_model=ContactResponse)
|
||||
def get_contact(contact_id: int, db: DbSession) -> Contact:
|
||||
"""Get a contact by ID with all relationships."""
|
||||
contact = db.scalar(
|
||||
select(Contact)
|
||||
.where(Contact.id == contact_id)
|
||||
.options(
|
||||
selectinload(Contact.needs),
|
||||
selectinload(Contact.related_to),
|
||||
selectinload(Contact.related_from),
|
||||
)
|
||||
)
|
||||
if not contact:
|
||||
raise HTTPException(status_code=404, detail="Contact not found")
|
||||
return contact
|
||||
|
||||
|
||||
@router.patch("/contacts/{contact_id}", response_model=ContactResponse)
|
||||
def update_contact(
|
||||
contact_id: int,
|
||||
contact: ContactUpdate,
|
||||
db: DbSession,
|
||||
) -> Contact:
|
||||
"""Update a contact by ID."""
|
||||
db_contact = db.get(Contact, contact_id)
|
||||
if not db_contact:
|
||||
raise HTTPException(status_code=404, detail="Contact not found")
|
||||
|
||||
update_data = contact.model_dump(exclude_unset=True)
|
||||
need_ids = update_data.pop("need_ids", None)
|
||||
|
||||
for key, value in update_data.items():
|
||||
setattr(db_contact, key, value)
|
||||
|
||||
if need_ids is not None:
|
||||
needs = list(db.scalars(select(Need).where(Need.id.in_(need_ids))).all())
|
||||
db_contact.needs = needs
|
||||
|
||||
db.commit()
|
||||
db.refresh(db_contact)
|
||||
return db_contact
|
||||
|
||||
|
||||
@router.delete("/contacts/{contact_id}", response_model=None)
|
||||
def delete_contact(contact_id: int, request: Request, db: DbSession) -> dict[str, bool] | HTMLResponse:
|
||||
"""Delete a contact by ID."""
|
||||
contact = db.get(Contact, contact_id)
|
||||
if not contact:
|
||||
raise HTTPException(status_code=404, detail="Contact not found")
|
||||
db.delete(contact)
|
||||
db.commit()
|
||||
if _is_htmx(request):
|
||||
return HTMLResponse("")
|
||||
return {"deleted": True}
|
||||
|
||||
|
||||
@router.post("/contacts/{contact_id}/needs/{need_id}")
|
||||
def add_need_to_contact(
|
||||
contact_id: int,
|
||||
need_id: int,
|
||||
db: DbSession,
|
||||
) -> dict[str, bool]:
|
||||
"""Add a need to a contact."""
|
||||
contact = db.get(Contact, contact_id)
|
||||
if not contact:
|
||||
raise HTTPException(status_code=404, detail="Contact not found")
|
||||
|
||||
need = db.get(Need, need_id)
|
||||
if not need:
|
||||
raise HTTPException(status_code=404, detail="Need not found")
|
||||
|
||||
if need not in contact.needs:
|
||||
contact.needs.append(need)
|
||||
db.commit()
|
||||
|
||||
return {"added": True}
|
||||
|
||||
|
||||
@router.delete("/contacts/{contact_id}/needs/{need_id}", response_model=None)
|
||||
def remove_need_from_contact(
|
||||
contact_id: int,
|
||||
need_id: int,
|
||||
request: Request,
|
||||
db: DbSession,
|
||||
) -> dict[str, bool] | HTMLResponse:
|
||||
"""Remove a need from a contact."""
|
||||
contact = db.get(Contact, contact_id)
|
||||
if not contact:
|
||||
raise HTTPException(status_code=404, detail="Contact not found")
|
||||
|
||||
need = db.get(Need, need_id)
|
||||
if not need:
|
||||
raise HTTPException(status_code=404, detail="Need not found")
|
||||
|
||||
if need in contact.needs:
|
||||
contact.needs.remove(need)
|
||||
db.commit()
|
||||
|
||||
if _is_htmx(request):
|
||||
return HTMLResponse("")
|
||||
return {"removed": True}
|
||||
|
||||
|
||||
@router.post(
|
||||
"/contacts/{contact_id}/relationships",
|
||||
response_model=ContactRelationshipResponse,
|
||||
)
|
||||
def add_contact_relationship(
|
||||
contact_id: int,
|
||||
relationship: ContactRelationshipCreate,
|
||||
db: DbSession,
|
||||
) -> ContactRelationship:
|
||||
"""Add a relationship between two contacts."""
|
||||
contact = db.get(Contact, contact_id)
|
||||
if not contact:
|
||||
raise HTTPException(status_code=404, detail="Contact not found")
|
||||
|
||||
related_contact = db.get(Contact, relationship.related_contact_id)
|
||||
if not related_contact:
|
||||
raise HTTPException(status_code=404, detail="Related contact not found")
|
||||
|
||||
if contact_id == relationship.related_contact_id:
|
||||
raise HTTPException(status_code=400, detail="Cannot relate contact to itself")
|
||||
|
||||
# Use provided weight or default from relationship type
|
||||
weight = relationship.closeness_weight
|
||||
if weight is None:
|
||||
weight = relationship.relationship_type.default_weight
|
||||
|
||||
db_relationship = ContactRelationship(
|
||||
contact_id=contact_id,
|
||||
related_contact_id=relationship.related_contact_id,
|
||||
relationship_type=relationship.relationship_type.value,
|
||||
closeness_weight=weight,
|
||||
)
|
||||
db.add(db_relationship)
|
||||
db.commit()
|
||||
db.refresh(db_relationship)
|
||||
return db_relationship
|
||||
|
||||
|
||||
@router.get(
|
||||
"/contacts/{contact_id}/relationships",
|
||||
response_model=list[ContactRelationshipResponse],
|
||||
)
|
||||
def get_contact_relationships(
|
||||
contact_id: int,
|
||||
db: DbSession,
|
||||
) -> list[ContactRelationship]:
|
||||
"""Get all relationships for a contact."""
|
||||
contact = db.get(Contact, contact_id)
|
||||
if not contact:
|
||||
raise HTTPException(status_code=404, detail="Contact not found")
|
||||
|
||||
outgoing = list(db.scalars(select(ContactRelationship).where(ContactRelationship.contact_id == contact_id)).all())
|
||||
incoming = list(
|
||||
db.scalars(select(ContactRelationship).where(ContactRelationship.related_contact_id == contact_id)).all()
|
||||
)
|
||||
return outgoing + incoming
|
||||
|
||||
|
||||
@router.patch(
|
||||
"/contacts/{contact_id}/relationships/{related_contact_id}",
|
||||
response_model=ContactRelationshipResponse,
|
||||
)
|
||||
def update_contact_relationship(
|
||||
contact_id: int,
|
||||
related_contact_id: int,
|
||||
update: ContactRelationshipUpdate,
|
||||
db: DbSession,
|
||||
) -> ContactRelationship:
|
||||
"""Update a relationship between two contacts."""
|
||||
relationship = db.scalar(
|
||||
select(ContactRelationship).where(
|
||||
ContactRelationship.contact_id == contact_id,
|
||||
ContactRelationship.related_contact_id == related_contact_id,
|
||||
)
|
||||
)
|
||||
if not relationship:
|
||||
raise HTTPException(status_code=404, detail="Relationship not found")
|
||||
|
||||
if update.relationship_type is not None:
|
||||
relationship.relationship_type = update.relationship_type.value
|
||||
if update.closeness_weight is not None:
|
||||
relationship.closeness_weight = update.closeness_weight
|
||||
|
||||
db.commit()
|
||||
db.refresh(relationship)
|
||||
return relationship
|
||||
|
||||
|
||||
@router.delete("/contacts/{contact_id}/relationships/{related_contact_id}", response_model=None)
|
||||
def remove_contact_relationship(
|
||||
contact_id: int,
|
||||
related_contact_id: int,
|
||||
request: Request,
|
||||
db: DbSession,
|
||||
) -> dict[str, bool] | HTMLResponse:
|
||||
"""Remove a relationship between two contacts."""
|
||||
relationship = db.scalar(
|
||||
select(ContactRelationship).where(
|
||||
ContactRelationship.contact_id == contact_id,
|
||||
ContactRelationship.related_contact_id == related_contact_id,
|
||||
)
|
||||
)
|
||||
if not relationship:
|
||||
raise HTTPException(status_code=404, detail="Relationship not found")
|
||||
|
||||
db.delete(relationship)
|
||||
db.commit()
|
||||
if _is_htmx(request):
|
||||
return HTMLResponse("")
|
||||
return {"deleted": True}
|
||||
|
||||
|
||||
@router.get("/relationship-types")
|
||||
def list_relationship_types() -> list[RelationshipTypeInfo]:
|
||||
"""List all available relationship types with their default weights."""
|
||||
return [
|
||||
RelationshipTypeInfo(
|
||||
value=rt.value,
|
||||
display_name=rt.display_name,
|
||||
default_weight=rt.default_weight,
|
||||
)
|
||||
for rt in RelationshipType
|
||||
]
|
||||
|
||||
|
||||
@router.get("/graph")
|
||||
def get_relationship_graph(db: DbSession) -> GraphData:
|
||||
"""Get all contacts and relationships as graph data for visualization."""
|
||||
contacts = list(db.scalars(select(Contact)).all())
|
||||
relationships = list(db.scalars(select(ContactRelationship)).all())
|
||||
|
||||
nodes = [GraphNode(id=c.id, name=c.name, current_job=c.current_job) for c in contacts]
|
||||
|
||||
edges = [
|
||||
GraphEdge(
|
||||
source=rel.contact_id,
|
||||
target=rel.related_contact_id,
|
||||
relationship_type=rel.relationship_type,
|
||||
closeness_weight=rel.closeness_weight,
|
||||
)
|
||||
for rel in relationships
|
||||
]
|
||||
|
||||
return GraphData(nodes=nodes, edges=edges)
|
||||
@@ -1,345 +0,0 @@
|
||||
"""HTMX server-rendered view router."""
|
||||
|
||||
from pathlib import Path
|
||||
from typing import Annotated, Any
|
||||
|
||||
from fastapi import APIRouter, Form, HTTPException, Request
|
||||
from fastapi.responses import HTMLResponse, RedirectResponse
|
||||
from fastapi.templating import Jinja2Templates
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy.orm import Session, selectinload
|
||||
|
||||
from python.api.dependencies import DbSession
|
||||
from python.orm.richie.contact import Contact, ContactRelationship, Need, RelationshipType
|
||||
|
||||
TEMPLATES_DIR = Path(__file__).parent.parent / "templates"
|
||||
templates = Jinja2Templates(directory=TEMPLATES_DIR)
|
||||
|
||||
router = APIRouter(tags=["views"])
|
||||
|
||||
FAMILIAL_TYPES = {
|
||||
"parent",
|
||||
"child",
|
||||
"sibling",
|
||||
"grandparent",
|
||||
"grandchild",
|
||||
"aunt_uncle",
|
||||
"niece_nephew",
|
||||
"cousin",
|
||||
"in_law",
|
||||
}
|
||||
FRIEND_TYPES = {"best_friend", "close_friend", "friend", "acquaintance", "neighbor"}
|
||||
PARTNER_TYPES = {"spouse", "partner"}
|
||||
PROFESSIONAL_TYPES = {"mentor", "mentee", "business_partner", "colleague", "manager", "direct_report", "client"}
|
||||
|
||||
CONTACT_STRING_FIELDS = (
|
||||
"name",
|
||||
"legal_name",
|
||||
"suffix",
|
||||
"gender",
|
||||
"current_job",
|
||||
"timezone",
|
||||
"profile_pic",
|
||||
"bio",
|
||||
"goals",
|
||||
"social_structure_style",
|
||||
"safe_conversation_starters",
|
||||
"topics_to_avoid",
|
||||
"ssn",
|
||||
)
|
||||
|
||||
CONTACT_INT_FIELDS = ("age", "self_sufficiency_score")
|
||||
|
||||
|
||||
def _group_relationships(relationships: list[ContactRelationship]) -> dict[str, list[ContactRelationship]]:
|
||||
"""Group relationships by category."""
|
||||
groups: dict[str, list[ContactRelationship]] = {
|
||||
"familial": [],
|
||||
"partners": [],
|
||||
"friends": [],
|
||||
"professional": [],
|
||||
"other": [],
|
||||
}
|
||||
for rel in relationships:
|
||||
if rel.relationship_type in FAMILIAL_TYPES:
|
||||
groups["familial"].append(rel)
|
||||
elif rel.relationship_type in PARTNER_TYPES:
|
||||
groups["partners"].append(rel)
|
||||
elif rel.relationship_type in FRIEND_TYPES:
|
||||
groups["friends"].append(rel)
|
||||
elif rel.relationship_type in PROFESSIONAL_TYPES:
|
||||
groups["professional"].append(rel)
|
||||
else:
|
||||
groups["other"].append(rel)
|
||||
return groups
|
||||
|
||||
|
||||
def _build_contact_name_map(database: Session, contact: Contact) -> dict[int, str]:
|
||||
"""Build a mapping of contact IDs to names for relationship display."""
|
||||
related_ids = {rel.related_contact_id for rel in contact.related_to}
|
||||
related_ids |= {rel.contact_id for rel in contact.related_from}
|
||||
related_ids.discard(contact.id)
|
||||
|
||||
if not related_ids:
|
||||
return {}
|
||||
|
||||
related_contacts = list(database.scalars(select(Contact).where(Contact.id.in_(related_ids))).all())
|
||||
return {related.id: related.name for related in related_contacts}
|
||||
|
||||
|
||||
def _get_relationship_type_display() -> dict[str, str]:
|
||||
"""Build a mapping of relationship type values to display names."""
|
||||
return {rel_type.value: rel_type.display_name for rel_type in RelationshipType}
|
||||
|
||||
|
||||
async def _parse_contact_form(request: Request) -> dict[str, Any]:
|
||||
"""Parse contact form data from a multipart/form request."""
|
||||
form_data = await request.form()
|
||||
result: dict[str, Any] = {}
|
||||
|
||||
for field in CONTACT_STRING_FIELDS:
|
||||
value = form_data.get(field, "")
|
||||
result[field] = str(value) if value else None
|
||||
|
||||
for field in CONTACT_INT_FIELDS:
|
||||
value = form_data.get(field, "")
|
||||
result[field] = int(value) if value else None
|
||||
|
||||
result["need_ids"] = [int(value) for value in form_data.getlist("need_ids")]
|
||||
return result
|
||||
|
||||
|
||||
def _save_contact_from_form(database: Session, contact: Contact, form_result: dict[str, Any]) -> None:
|
||||
"""Apply parsed form data to a Contact and save associated needs."""
|
||||
need_ids = form_result.pop("need_ids")
|
||||
|
||||
for key, value in form_result.items():
|
||||
setattr(contact, key, value)
|
||||
|
||||
if need_ids:
|
||||
contact.needs = list(database.scalars(select(Need).where(Need.id.in_(need_ids))).all())
|
||||
else:
|
||||
contact.needs = []
|
||||
|
||||
|
||||
@router.get("/", response_class=HTMLResponse)
|
||||
@router.get("/contacts", response_class=HTMLResponse)
|
||||
def contact_list_page(request: Request, database: DbSession) -> HTMLResponse:
|
||||
"""Render the contacts list page."""
|
||||
contacts = list(database.scalars(select(Contact)).all())
|
||||
return templates.TemplateResponse(request, "contact_list.html", {"contacts": contacts})
|
||||
|
||||
|
||||
@router.get("/contacts/new", response_class=HTMLResponse)
|
||||
def new_contact_page(request: Request, database: DbSession) -> HTMLResponse:
|
||||
"""Render the new contact form page."""
|
||||
all_needs = list(database.scalars(select(Need)).all())
|
||||
return templates.TemplateResponse(request, "contact_form.html", {"contact": None, "all_needs": all_needs})
|
||||
|
||||
|
||||
@router.post("/htmx/contacts/new")
|
||||
async def create_contact_form(request: Request, database: DbSession) -> RedirectResponse:
|
||||
"""Handle the create contact form submission."""
|
||||
form_result = await _parse_contact_form(request)
|
||||
contact = Contact()
|
||||
_save_contact_from_form(database, contact, form_result)
|
||||
|
||||
database.add(contact)
|
||||
database.commit()
|
||||
database.refresh(contact)
|
||||
return RedirectResponse(url=f"/contacts/{contact.id}", status_code=303)
|
||||
|
||||
|
||||
@router.get("/contacts/{contact_id}", response_class=HTMLResponse)
|
||||
def contact_detail_page(contact_id: int, request: Request, database: DbSession) -> HTMLResponse:
|
||||
"""Render the contact detail page."""
|
||||
contact = database.scalar(
|
||||
select(Contact)
|
||||
.where(Contact.id == contact_id)
|
||||
.options(
|
||||
selectinload(Contact.needs),
|
||||
selectinload(Contact.related_to),
|
||||
selectinload(Contact.related_from),
|
||||
)
|
||||
)
|
||||
if not contact:
|
||||
raise HTTPException(status_code=404, detail="Contact not found")
|
||||
|
||||
contact_names = _build_contact_name_map(database, contact)
|
||||
grouped_relationships = _group_relationships(contact.related_to)
|
||||
all_contacts = list(database.scalars(select(Contact)).all())
|
||||
all_needs = list(database.scalars(select(Need)).all())
|
||||
available_needs = [need for need in all_needs if need not in contact.needs]
|
||||
|
||||
return templates.TemplateResponse(
|
||||
request,
|
||||
"contact_detail.html",
|
||||
{
|
||||
"contact": contact,
|
||||
"contact_names": contact_names,
|
||||
"grouped_relationships": grouped_relationships,
|
||||
"all_contacts": all_contacts,
|
||||
"available_needs": available_needs,
|
||||
"relationship_types": list(RelationshipType),
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
@router.get("/contacts/{contact_id}/edit", response_class=HTMLResponse)
|
||||
def edit_contact_page(contact_id: int, request: Request, database: DbSession) -> HTMLResponse:
|
||||
"""Render the edit contact form page."""
|
||||
contact = database.scalar(select(Contact).where(Contact.id == contact_id).options(selectinload(Contact.needs)))
|
||||
if not contact:
|
||||
raise HTTPException(status_code=404, detail="Contact not found")
|
||||
|
||||
all_needs = list(database.scalars(select(Need)).all())
|
||||
return templates.TemplateResponse(request, "contact_form.html", {"contact": contact, "all_needs": all_needs})
|
||||
|
||||
|
||||
@router.post("/htmx/contacts/{contact_id}/edit")
|
||||
async def update_contact_form(contact_id: int, request: Request, database: DbSession) -> RedirectResponse:
|
||||
"""Handle the edit contact form submission."""
|
||||
contact = database.get(Contact, contact_id)
|
||||
if not contact:
|
||||
raise HTTPException(status_code=404, detail="Contact not found")
|
||||
|
||||
form_result = await _parse_contact_form(request)
|
||||
_save_contact_from_form(database, contact, form_result)
|
||||
|
||||
database.commit()
|
||||
return RedirectResponse(url=f"/contacts/{contact_id}", status_code=303)
|
||||
|
||||
|
||||
@router.post("/htmx/contacts/{contact_id}/add-need", response_class=HTMLResponse)
|
||||
def add_need_to_contact_htmx(
|
||||
contact_id: int,
|
||||
request: Request,
|
||||
database: DbSession,
|
||||
need_id: Annotated[int, Form()],
|
||||
) -> HTMLResponse:
|
||||
"""Add a need to a contact and return updated manage-needs partial."""
|
||||
contact = database.scalar(select(Contact).where(Contact.id == contact_id).options(selectinload(Contact.needs)))
|
||||
if not contact:
|
||||
raise HTTPException(status_code=404, detail="Contact not found")
|
||||
|
||||
need = database.get(Need, need_id)
|
||||
if not need:
|
||||
raise HTTPException(status_code=404, detail="Need not found")
|
||||
|
||||
if need not in contact.needs:
|
||||
contact.needs.append(need)
|
||||
database.commit()
|
||||
database.refresh(contact)
|
||||
|
||||
return templates.TemplateResponse(request, "partials/manage_needs.html", {"contact": contact})
|
||||
|
||||
|
||||
@router.post("/htmx/contacts/{contact_id}/add-relationship", response_class=HTMLResponse)
|
||||
def add_relationship_htmx(
|
||||
contact_id: int,
|
||||
request: Request,
|
||||
database: DbSession,
|
||||
related_contact_id: Annotated[int, Form()],
|
||||
relationship_type: Annotated[str, Form()],
|
||||
) -> HTMLResponse:
|
||||
"""Add a relationship and return updated manage-relationships partial."""
|
||||
contact = database.scalar(select(Contact).where(Contact.id == contact_id).options(selectinload(Contact.related_to)))
|
||||
if not contact:
|
||||
raise HTTPException(status_code=404, detail="Contact not found")
|
||||
|
||||
related_contact = database.get(Contact, related_contact_id)
|
||||
if not related_contact:
|
||||
raise HTTPException(status_code=404, detail="Related contact not found")
|
||||
|
||||
rel_type = RelationshipType(relationship_type)
|
||||
weight = rel_type.default_weight
|
||||
|
||||
relationship = ContactRelationship(
|
||||
contact_id=contact_id,
|
||||
related_contact_id=related_contact_id,
|
||||
relationship_type=relationship_type,
|
||||
closeness_weight=weight,
|
||||
)
|
||||
database.add(relationship)
|
||||
database.commit()
|
||||
database.refresh(contact)
|
||||
|
||||
contact_names = _build_contact_name_map(database, contact)
|
||||
return templates.TemplateResponse(
|
||||
request,
|
||||
"partials/manage_relationships.html",
|
||||
{"contact": contact, "contact_names": contact_names},
|
||||
)
|
||||
|
||||
|
||||
@router.post("/htmx/contacts/{contact_id}/relationships/{related_contact_id}/weight")
|
||||
def update_relationship_weight_htmx(
|
||||
contact_id: int,
|
||||
related_contact_id: int,
|
||||
database: DbSession,
|
||||
closeness_weight: Annotated[int, Form()],
|
||||
) -> HTMLResponse:
|
||||
"""Update a relationship's closeness weight from HTMX range input."""
|
||||
relationship = database.scalar(
|
||||
select(ContactRelationship).where(
|
||||
ContactRelationship.contact_id == contact_id,
|
||||
ContactRelationship.related_contact_id == related_contact_id,
|
||||
)
|
||||
)
|
||||
if not relationship:
|
||||
raise HTTPException(status_code=404, detail="Relationship not found")
|
||||
|
||||
relationship.closeness_weight = closeness_weight
|
||||
database.commit()
|
||||
return HTMLResponse("")
|
||||
|
||||
|
||||
@router.post("/htmx/needs", response_class=HTMLResponse)
|
||||
def create_need_htmx(
|
||||
request: Request,
|
||||
database: DbSession,
|
||||
name: Annotated[str, Form()],
|
||||
description: Annotated[str, Form()] = "",
|
||||
) -> HTMLResponse:
|
||||
"""Create a need via form data and return updated needs list."""
|
||||
need = Need(name=name, description=description or None)
|
||||
database.add(need)
|
||||
database.commit()
|
||||
needs = list(database.scalars(select(Need)).all())
|
||||
return templates.TemplateResponse(request, "partials/need_items.html", {"needs": needs})
|
||||
|
||||
|
||||
@router.get("/needs", response_class=HTMLResponse)
|
||||
def needs_page(request: Request, database: DbSession) -> HTMLResponse:
|
||||
"""Render the needs list page."""
|
||||
needs = list(database.scalars(select(Need)).all())
|
||||
return templates.TemplateResponse(request, "need_list.html", {"needs": needs})
|
||||
|
||||
|
||||
@router.get("/graph", response_class=HTMLResponse)
|
||||
def graph_page(request: Request, database: DbSession) -> HTMLResponse:
|
||||
"""Render the relationship graph page."""
|
||||
contacts = list(database.scalars(select(Contact)).all())
|
||||
relationships = list(database.scalars(select(ContactRelationship)).all())
|
||||
|
||||
graph_data = {
|
||||
"nodes": [{"id": contact.id, "name": contact.name, "current_job": contact.current_job} for contact in contacts],
|
||||
"edges": [
|
||||
{
|
||||
"source": rel.contact_id,
|
||||
"target": rel.related_contact_id,
|
||||
"relationship_type": rel.relationship_type,
|
||||
"closeness_weight": rel.closeness_weight,
|
||||
}
|
||||
for rel in relationships
|
||||
],
|
||||
}
|
||||
|
||||
return templates.TemplateResponse(
|
||||
request,
|
||||
"graph.html",
|
||||
{
|
||||
"graph_data": graph_data,
|
||||
"relationship_type_display": _get_relationship_type_display(),
|
||||
},
|
||||
)
|
||||
@@ -1,198 +0,0 @@
|
||||
<!DOCTYPE html>
|
||||
<html lang="en" data-theme="light">
|
||||
<head>
|
||||
<meta charset="UTF-8">
|
||||
<meta name="viewport" content="width=device-width, initial-scale=1.0">
|
||||
<title>{% block title %}Contact Database{% endblock %}</title>
|
||||
<script src="https://unpkg.com/htmx.org@2.0.4"></script>
|
||||
<style>
|
||||
:root {
|
||||
--color-bg: #f5f5f5;
|
||||
--color-bg-card: #ffffff;
|
||||
--color-bg-hover: #f0f0f0;
|
||||
--color-bg-muted: #f9f9f9;
|
||||
--color-bg-error: #ffe0e0;
|
||||
--color-text: #333333;
|
||||
--color-text-muted: #666666;
|
||||
--color-text-error: #cc0000;
|
||||
--color-border: #dddddd;
|
||||
--color-border-light: #eeeeee;
|
||||
--color-border-lighter: #f0f0f0;
|
||||
--color-primary: #0066cc;
|
||||
--color-primary-hover: #0055aa;
|
||||
--color-danger: #cc3333;
|
||||
--color-danger-hover: #aa2222;
|
||||
--color-tag-bg: #e0e0e0;
|
||||
--shadow: 0 1px 3px rgba(0, 0, 0, 0.1);
|
||||
font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, sans-serif;
|
||||
line-height: 1.5;
|
||||
color: var(--color-text);
|
||||
background-color: var(--color-bg);
|
||||
}
|
||||
[data-theme="dark"] {
|
||||
--color-bg: #1a1a1a;
|
||||
--color-bg-card: #2d2d2d;
|
||||
--color-bg-hover: #3d3d3d;
|
||||
--color-bg-muted: #252525;
|
||||
--color-bg-error: #4a2020;
|
||||
--color-text: #e0e0e0;
|
||||
--color-text-muted: #a0a0a0;
|
||||
--color-text-error: #ff6b6b;
|
||||
--color-border: #404040;
|
||||
--color-border-light: #353535;
|
||||
--color-border-lighter: #303030;
|
||||
--color-primary: #4da6ff;
|
||||
--color-primary-hover: #7dbfff;
|
||||
--color-danger: #ff6b6b;
|
||||
--color-danger-hover: #ff8a8a;
|
||||
--color-tag-bg: #404040;
|
||||
--shadow: 0 1px 3px rgba(0, 0, 0, 0.3);
|
||||
}
|
||||
* { box-sizing: border-box; }
|
||||
body { margin: 0; background: var(--color-bg); color: var(--color-text); }
|
||||
.app { max-width: 1000px; margin: 0 auto; padding: 20px; }
|
||||
nav { display: flex; align-items: center; gap: 20px; padding: 15px 0; border-bottom: 1px solid var(--color-border); margin-bottom: 20px; }
|
||||
nav a { color: var(--color-primary); text-decoration: none; font-weight: 500; }
|
||||
nav a:hover { text-decoration: underline; }
|
||||
.theme-toggle { margin-left: auto; }
|
||||
main { background: var(--color-bg-card); padding: 20px; border-radius: 8px; box-shadow: var(--shadow); }
|
||||
.header { display: flex; justify-content: space-between; align-items: center; margin-bottom: 20px; }
|
||||
.header h1 { margin: 0; }
|
||||
a { color: var(--color-primary); }
|
||||
a:hover { text-decoration: underline; }
|
||||
|
||||
.btn { display: inline-block; padding: 8px 16px; border: 1px solid var(--color-border); border-radius: 4px; background: var(--color-bg-card); color: var(--color-text); text-decoration: none; cursor: pointer; font-size: 14px; margin-left: 8px; }
|
||||
.btn:hover { background: var(--color-bg-hover); }
|
||||
.btn-primary { background: var(--color-primary); border-color: var(--color-primary); color: white; }
|
||||
.btn-primary:hover { background: var(--color-primary-hover); }
|
||||
.btn-danger { background: var(--color-danger); border-color: var(--color-danger); color: white; }
|
||||
.btn-danger:hover { background: var(--color-danger-hover); }
|
||||
.btn-small { padding: 4px 8px; font-size: 12px; }
|
||||
.btn:disabled { opacity: 0.6; cursor: not-allowed; }
|
||||
|
||||
table { width: 100%; border-collapse: collapse; }
|
||||
th, td { padding: 12px; text-align: left; border-bottom: 1px solid var(--color-border-light); }
|
||||
th { font-weight: 600; background: var(--color-bg-muted); }
|
||||
tr:hover { background: var(--color-bg-muted); }
|
||||
|
||||
.error { background: var(--color-bg-error); color: var(--color-text-error); padding: 10px; border-radius: 4px; margin-bottom: 20px; }
|
||||
.tag { display: inline-block; background: var(--color-tag-bg); padding: 2px 8px; border-radius: 12px; font-size: 12px; color: var(--color-text-muted); }
|
||||
|
||||
.add-form { display: flex; gap: 10px; margin-top: 15px; flex-wrap: wrap; }
|
||||
.add-form select, .add-form input { padding: 8px; border: 1px solid var(--color-border); border-radius: 4px; min-width: 200px; background: var(--color-bg-card); color: var(--color-text); }
|
||||
|
||||
.form-group { margin-bottom: 20px; }
|
||||
.form-group label { display: block; font-weight: 500; margin-bottom: 5px; }
|
||||
.form-group input, .form-group textarea, .form-group select { width: 100%; padding: 10px; border: 1px solid var(--color-border); border-radius: 4px; font-size: 14px; background: var(--color-bg-card); color: var(--color-text); }
|
||||
.form-group textarea { resize: vertical; }
|
||||
.form-row { display: grid; grid-template-columns: 1fr 1fr; gap: 20px; }
|
||||
.checkbox-group { display: flex; flex-wrap: wrap; gap: 15px; }
|
||||
.checkbox-label { display: flex; align-items: center; gap: 5px; cursor: pointer; }
|
||||
.form-actions { display: flex; gap: 10px; margin-top: 30px; padding-top: 20px; border-top: 1px solid var(--color-border-light); }
|
||||
|
||||
.need-form { background: var(--color-bg-muted); padding: 20px; border-radius: 4px; margin-bottom: 20px; }
|
||||
.need-items { list-style: none; padding: 0; }
|
||||
.need-items li { display: flex; justify-content: space-between; align-items: flex-start; padding: 15px; border: 1px solid var(--color-border-light); border-radius: 4px; margin-bottom: 10px; }
|
||||
.need-info p { margin: 5px 0 0; color: var(--color-text-muted); font-size: 14px; }
|
||||
|
||||
.graph-container { width: 100%; }
|
||||
.graph-hint { color: var(--color-text-muted); font-size: 14px; margin-bottom: 15px; }
|
||||
.selected-info { margin-top: 15px; padding: 15px; background: var(--color-bg-muted); border-radius: 8px; }
|
||||
.selected-info h3 { margin: 0 0 10px; }
|
||||
.selected-info p { margin: 5px 0; color: var(--color-text-muted); }
|
||||
.legend { margin-top: 20px; padding: 15px; background: var(--color-bg-muted); border-radius: 8px; }
|
||||
.legend h4 { margin: 0 0 10px; font-size: 14px; }
|
||||
.legend-items { display: flex; flex-wrap: wrap; gap: 15px; }
|
||||
.legend-item { display: flex; align-items: center; gap: 8px; font-size: 12px; color: var(--color-text-muted); }
|
||||
.legend-line { width: 30px; border-radius: 2px; }
|
||||
|
||||
.id-card { width: 100%; }
|
||||
.id-card-inner { background: linear-gradient(135deg, #0a0a0f 0%, #1a1a2e 50%, #0a0a0f 100%); background-image: radial-gradient(white 1px, transparent 1px), linear-gradient(135deg, #0a0a0f 0%, #1a1a2e 50%, #0a0a0f 100%); background-size: 50px 50px, 100% 100%; color: #fff; border-radius: 12px; padding: 25px; min-height: 500px; position: relative; overflow: hidden; }
|
||||
.id-card-header { display: flex; justify-content: space-between; align-items: flex-start; margin-bottom: 15px; }
|
||||
.id-card-header-left { flex: 1; }
|
||||
.id-card-header-right { display: flex; flex-direction: column; align-items: flex-end; gap: 10px; }
|
||||
.id-card-title { font-size: 2.5rem; font-weight: 700; margin: 0; color: #fff; text-shadow: 2px 2px 4px rgba(0,0,0,0.5); }
|
||||
.id-profile-pic { width: 80px; height: 80px; border-radius: 8px; object-fit: cover; border: 2px solid rgba(255,255,255,0.3); }
|
||||
.id-profile-placeholder { width: 80px; height: 80px; border-radius: 8px; background: linear-gradient(135deg, #4ecdc4 0%, #44a8a0 100%); display: flex; align-items: center; justify-content: center; border: 2px solid rgba(255,255,255,0.3); }
|
||||
.id-profile-placeholder span { font-size: 2rem; font-weight: 700; color: #fff; text-shadow: 1px 1px 2px rgba(0,0,0,0.3); }
|
||||
.id-card-actions { display: flex; gap: 8px; }
|
||||
.id-card-actions .btn { background: rgba(255,255,255,0.1); border-color: rgba(255,255,255,0.3); color: #fff; }
|
||||
.id-card-actions .btn:hover { background: rgba(255,255,255,0.2); }
|
||||
.id-card-body { display: grid; grid-template-columns: 1fr 1.5fr; gap: 30px; }
|
||||
.id-card-left { display: flex; flex-direction: column; gap: 8px; }
|
||||
.id-field { font-size: 1rem; line-height: 1.4; }
|
||||
.id-field-block { margin-top: 15px; font-size: 0.95rem; line-height: 1.5; }
|
||||
.id-label { color: #4ecdc4; font-weight: 500; }
|
||||
.id-card-right { display: flex; flex-direction: column; gap: 20px; }
|
||||
.id-bio { font-size: 0.9rem; line-height: 1.6; color: #e0e0e0; }
|
||||
.id-relationships { margin-top: 10px; }
|
||||
.id-section-title { font-size: 1.5rem; margin: 0 0 15px; color: #fff; border-bottom: 1px solid rgba(255,255,255,0.2); padding-bottom: 8px; }
|
||||
.id-rel-group { margin-bottom: 12px; font-size: 0.9rem; line-height: 1.6; }
|
||||
.id-rel-label { color: #a0a0a0; }
|
||||
.id-rel-group a { color: #4ecdc4; text-decoration: none; }
|
||||
.id-rel-group a:hover { text-decoration: underline; }
|
||||
.id-rel-type { color: #888; font-size: 0.85em; }
|
||||
.id-card-warnings { margin-top: 30px; padding-top: 20px; border-top: 1px solid rgba(255,255,255,0.2); display: flex; flex-wrap: wrap; gap: 20px; }
|
||||
.id-warning { display: flex; align-items: center; gap: 8px; font-size: 0.9rem; color: #ff6b6b; }
|
||||
.warning-dot { width: 8px; height: 8px; background: #ff6b6b; border-radius: 50%; flex-shrink: 0; }
|
||||
.warning-desc { color: #ccc; }
|
||||
|
||||
.id-card-manage { margin-top: 20px; background: var(--color-bg-muted); border-radius: 8px; padding: 15px; }
|
||||
.id-card-manage summary { cursor: pointer; font-weight: 600; font-size: 1.1rem; padding: 5px 0; }
|
||||
.id-card-manage[open] summary { margin-bottom: 15px; border-bottom: 1px solid var(--color-border-light); padding-bottom: 10px; }
|
||||
.manage-section { margin-bottom: 25px; }
|
||||
.manage-section h3 { margin: 0 0 15px; font-size: 1rem; }
|
||||
.manage-relationships { display: flex; flex-direction: column; gap: 10px; margin-bottom: 15px; }
|
||||
.manage-rel-item { display: flex; align-items: center; gap: 12px; padding: 10px; background: var(--color-bg-card); border-radius: 6px; flex-wrap: wrap; }
|
||||
.manage-rel-item a { font-weight: 500; min-width: 120px; }
|
||||
.weight-control { display: flex; align-items: center; gap: 8px; font-size: 12px; color: var(--color-text-muted); }
|
||||
.weight-control input[type="range"] { width: 80px; cursor: pointer; }
|
||||
.weight-value { min-width: 20px; text-align: center; font-weight: 600; }
|
||||
.manage-needs-list { list-style: none; padding: 0; margin: 0 0 15px; }
|
||||
.manage-needs-list li { display: flex; align-items: center; gap: 12px; padding: 10px; background: var(--color-bg-card); border-radius: 6px; margin-bottom: 8px; }
|
||||
.manage-needs-list li .btn { margin-left: auto; }
|
||||
|
||||
.htmx-indicator { display: none; }
|
||||
.htmx-request .htmx-indicator { display: inline; }
|
||||
.htmx-request.htmx-indicator { display: inline; }
|
||||
|
||||
@media (max-width: 768px) {
|
||||
.id-card-body { grid-template-columns: 1fr; }
|
||||
.id-card-title { font-size: 1.8rem; }
|
||||
.id-card-header { flex-direction: column; gap: 15px; }
|
||||
}
|
||||
</style>
|
||||
</head>
|
||||
<body>
|
||||
<div class="app">
|
||||
<nav>
|
||||
<a href="/contacts">Contacts</a>
|
||||
<a href="/graph">Graph</a>
|
||||
<a href="/needs">Needs</a>
|
||||
<button class="btn btn-small theme-toggle" onclick="toggleTheme()">
|
||||
<span id="theme-label">Dark</span>
|
||||
</button>
|
||||
</nav>
|
||||
|
||||
<main id="main-content">
|
||||
{% block content %}{% endblock %}
|
||||
</main>
|
||||
</div>
|
||||
|
||||
<script>
|
||||
function toggleTheme() {
|
||||
const html = document.documentElement;
|
||||
const current = html.getAttribute('data-theme');
|
||||
const next = current === 'light' ? 'dark' : 'light';
|
||||
html.setAttribute('data-theme', next);
|
||||
localStorage.setItem('theme', next);
|
||||
document.getElementById('theme-label').textContent = next === 'light' ? 'Dark' : 'Light';
|
||||
}
|
||||
(function() {
|
||||
const saved = localStorage.getItem('theme') || 'light';
|
||||
document.documentElement.setAttribute('data-theme', saved);
|
||||
document.getElementById('theme-label').textContent = saved === 'light' ? 'Dark' : 'Light';
|
||||
})();
|
||||
</script>
|
||||
</body>
|
||||
</html>
|
||||
@@ -1,204 +0,0 @@
|
||||
{% extends "base.html" %}
|
||||
{% block title %}{{ contact.name }}{% endblock %}
|
||||
{% block content %}
|
||||
<div class="id-card">
|
||||
<div class="id-card-inner">
|
||||
<div class="id-card-header">
|
||||
<div class="id-card-header-left">
|
||||
<h1 class="id-card-title">I.D.: {{ contact.name }}</h1>
|
||||
</div>
|
||||
<div class="id-card-header-right">
|
||||
{% if contact.profile_pic %}
|
||||
<img src="{{ contact.profile_pic }}" alt="{{ contact.name }}'s profile" class="id-profile-pic">
|
||||
{% else %}
|
||||
<div class="id-profile-placeholder">
|
||||
<span>{{ contact.name[0]|upper }}</span>
|
||||
</div>
|
||||
{% endif %}
|
||||
<div class="id-card-actions">
|
||||
<a href="/contacts/{{ contact.id }}/edit" class="btn btn-small">Edit</a>
|
||||
<a href="/contacts" class="btn btn-small">Back</a>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<div class="id-card-body">
|
||||
<div class="id-card-left">
|
||||
{% if contact.legal_name %}
|
||||
<div class="id-field">Legal name: {{ contact.legal_name }}</div>
|
||||
{% endif %}
|
||||
{% if contact.suffix %}
|
||||
<div class="id-field">Suffix: {{ contact.suffix }}</div>
|
||||
{% endif %}
|
||||
{% if contact.gender %}
|
||||
<div class="id-field">Gender: {{ contact.gender }}</div>
|
||||
{% endif %}
|
||||
{% if contact.age %}
|
||||
<div class="id-field">Age: {{ contact.age }}</div>
|
||||
{% endif %}
|
||||
{% if contact.current_job %}
|
||||
<div class="id-field">Job: {{ contact.current_job }}</div>
|
||||
{% endif %}
|
||||
{% if contact.social_structure_style %}
|
||||
<div class="id-field">Social style: {{ contact.social_structure_style }}</div>
|
||||
{% endif %}
|
||||
{% if contact.self_sufficiency_score is not none %}
|
||||
<div class="id-field">Self-Sufficiency: {{ contact.self_sufficiency_score }}</div>
|
||||
{% endif %}
|
||||
{% if contact.timezone %}
|
||||
<div class="id-field">Timezone: {{ contact.timezone }}</div>
|
||||
{% endif %}
|
||||
{% if contact.safe_conversation_starters %}
|
||||
<div class="id-field-block">
|
||||
<span class="id-label">Safe con starters:</span> {{ contact.safe_conversation_starters }}
|
||||
</div>
|
||||
{% endif %}
|
||||
{% if contact.topics_to_avoid %}
|
||||
<div class="id-field-block">
|
||||
<span class="id-label">Topics to avoid:</span> {{ contact.topics_to_avoid }}
|
||||
</div>
|
||||
{% endif %}
|
||||
{% if contact.goals %}
|
||||
<div class="id-field-block">
|
||||
<span class="id-label">Goals:</span> {{ contact.goals }}
|
||||
</div>
|
||||
{% endif %}
|
||||
</div>
|
||||
|
||||
<div class="id-card-right">
|
||||
{% if contact.bio %}
|
||||
<div class="id-bio">
|
||||
<span class="id-label">Bio:</span> {{ contact.bio }}
|
||||
</div>
|
||||
{% endif %}
|
||||
|
||||
<div class="id-relationships">
|
||||
<h2 class="id-section-title">Relationships</h2>
|
||||
|
||||
{% if grouped_relationships.familial %}
|
||||
<div class="id-rel-group">
|
||||
<span class="id-rel-label">Familial:</span>
|
||||
{% for rel in grouped_relationships.familial %}
|
||||
<a href="/contacts/{{ rel.related_contact_id }}">{{ contact_names[rel.related_contact_id] }}</a><span class="id-rel-type">({{ rel.relationship_type|replace("_", " ")|title }})</span>{% if not loop.last %}, {% endif %}
|
||||
{% endfor %}
|
||||
</div>
|
||||
{% endif %}
|
||||
|
||||
{% if grouped_relationships.partners %}
|
||||
<div class="id-rel-group">
|
||||
<span class="id-rel-label">Partners:</span>
|
||||
{% for rel in grouped_relationships.partners %}
|
||||
<a href="/contacts/{{ rel.related_contact_id }}">{{ contact_names[rel.related_contact_id] }}</a>{% if not loop.last %}, {% endif %}
|
||||
{% endfor %}
|
||||
</div>
|
||||
{% endif %}
|
||||
|
||||
{% if grouped_relationships.friends %}
|
||||
<div class="id-rel-group">
|
||||
<span class="id-rel-label">Friends:</span>
|
||||
{% for rel in grouped_relationships.friends %}
|
||||
<a href="/contacts/{{ rel.related_contact_id }}">{{ contact_names[rel.related_contact_id] }}</a>{% if not loop.last %}, {% endif %}
|
||||
{% endfor %}
|
||||
</div>
|
||||
{% endif %}
|
||||
|
||||
{% if grouped_relationships.professional %}
|
||||
<div class="id-rel-group">
|
||||
<span class="id-rel-label">Professional:</span>
|
||||
{% for rel in grouped_relationships.professional %}
|
||||
<a href="/contacts/{{ rel.related_contact_id }}">{{ contact_names[rel.related_contact_id] }}</a><span class="id-rel-type">({{ rel.relationship_type|replace("_", " ")|title }})</span>{% if not loop.last %}, {% endif %}
|
||||
{% endfor %}
|
||||
</div>
|
||||
{% endif %}
|
||||
|
||||
{% if grouped_relationships.other %}
|
||||
<div class="id-rel-group">
|
||||
<span class="id-rel-label">Other:</span>
|
||||
{% for rel in grouped_relationships.other %}
|
||||
<a href="/contacts/{{ rel.related_contact_id }}">{{ contact_names[rel.related_contact_id] }}</a><span class="id-rel-type">({{ rel.relationship_type|replace("_", " ")|title }})</span>{% if not loop.last %}, {% endif %}
|
||||
{% endfor %}
|
||||
</div>
|
||||
{% endif %}
|
||||
|
||||
{% if contact.related_from %}
|
||||
<div class="id-rel-group">
|
||||
<span class="id-rel-label">Known by:</span>
|
||||
{% for rel in contact.related_from %}
|
||||
<a href="/contacts/{{ rel.contact_id }}">{{ contact_names[rel.contact_id] }}</a>{% if not loop.last %}, {% endif %}
|
||||
{% endfor %}
|
||||
</div>
|
||||
{% endif %}
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
{% if contact.needs %}
|
||||
<div class="id-card-warnings">
|
||||
{% for need in contact.needs %}
|
||||
<div class="id-warning">
|
||||
<span class="warning-dot"></span>
|
||||
Warning: {{ need.name }}
|
||||
{% if need.description %}<span class="warning-desc"> - {{ need.description }}</span>{% endif %}
|
||||
</div>
|
||||
{% endfor %}
|
||||
</div>
|
||||
{% endif %}
|
||||
</div>
|
||||
|
||||
<details class="id-card-manage">
|
||||
<summary>Manage Contact</summary>
|
||||
|
||||
<div class="manage-section">
|
||||
<h3>Manage Relationships</h3>
|
||||
<div id="manage-relationships" class="manage-relationships">
|
||||
{% include "partials/manage_relationships.html" %}
|
||||
</div>
|
||||
|
||||
{% if all_contacts %}
|
||||
<form hx-post="/htmx/contacts/{{ contact.id }}/add-relationship"
|
||||
hx-target="#manage-relationships"
|
||||
hx-swap="innerHTML"
|
||||
class="add-form">
|
||||
<select name="related_contact_id" required>
|
||||
<option value="">Select contact...</option>
|
||||
{% for other in all_contacts %}
|
||||
{% if other.id != contact.id %}
|
||||
<option value="{{ other.id }}">{{ other.name }}</option>
|
||||
{% endif %}
|
||||
{% endfor %}
|
||||
</select>
|
||||
<select name="relationship_type" required>
|
||||
<option value="">Select relationship type...</option>
|
||||
{% for rel_type in relationship_types %}
|
||||
<option value="{{ rel_type.value }}">{{ rel_type.display_name }}</option>
|
||||
{% endfor %}
|
||||
</select>
|
||||
<button type="submit" class="btn btn-primary">Add Relationship</button>
|
||||
</form>
|
||||
{% endif %}
|
||||
</div>
|
||||
|
||||
<div class="manage-section">
|
||||
<h3>Manage Needs/Warnings</h3>
|
||||
<div id="manage-needs">
|
||||
{% include "partials/manage_needs.html" %}
|
||||
</div>
|
||||
|
||||
{% if available_needs %}
|
||||
<form hx-post="/htmx/contacts/{{ contact.id }}/add-need"
|
||||
hx-target="#manage-needs"
|
||||
hx-swap="innerHTML"
|
||||
class="add-form">
|
||||
<select name="need_id" required>
|
||||
<option value="">Select a need...</option>
|
||||
{% for need in available_needs %}
|
||||
<option value="{{ need.id }}">{{ need.name }}</option>
|
||||
{% endfor %}
|
||||
</select>
|
||||
<button type="submit" class="btn btn-primary">Add Need</button>
|
||||
</form>
|
||||
{% endif %}
|
||||
</div>
|
||||
</details>
|
||||
</div>
|
||||
{% endblock %}
|
||||
@@ -1,115 +0,0 @@
|
||||
{% extends "base.html" %}
|
||||
{% block title %}{{ "Edit " + contact.name if contact else "New Contact" }}{% endblock %}
|
||||
{% block content %}
|
||||
<div class="contact-form">
|
||||
<h1>{{ "Edit Contact" if contact else "New Contact" }}</h1>
|
||||
|
||||
{% if contact %}
|
||||
<form method="post" action="/htmx/contacts/{{ contact.id }}/edit">
|
||||
{% else %}
|
||||
<form method="post" action="/htmx/contacts/new">
|
||||
{% endif %}
|
||||
|
||||
<div class="form-group">
|
||||
<label for="name">Name *</label>
|
||||
<input id="name" name="name" type="text" value="{{ contact.name if contact else '' }}" required>
|
||||
</div>
|
||||
|
||||
<div class="form-row">
|
||||
<div class="form-group">
|
||||
<label for="legal_name">Legal Name</label>
|
||||
<input id="legal_name" name="legal_name" type="text" value="{{ contact.legal_name or '' }}">
|
||||
</div>
|
||||
<div class="form-group">
|
||||
<label for="suffix">Suffix</label>
|
||||
<input id="suffix" name="suffix" type="text" value="{{ contact.suffix or '' }}">
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<div class="form-row">
|
||||
<div class="form-group">
|
||||
<label for="age">Age</label>
|
||||
<input id="age" name="age" type="number" value="{{ contact.age if contact and contact.age is not none else '' }}">
|
||||
</div>
|
||||
<div class="form-group">
|
||||
<label for="gender">Gender</label>
|
||||
<input id="gender" name="gender" type="text" value="{{ contact.gender or '' }}">
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<div class="form-group">
|
||||
<label for="current_job">Current Job</label>
|
||||
<input id="current_job" name="current_job" type="text" value="{{ contact.current_job or '' }}">
|
||||
</div>
|
||||
|
||||
<div class="form-group">
|
||||
<label for="timezone">Timezone</label>
|
||||
<input id="timezone" name="timezone" type="text" value="{{ contact.timezone or '' }}">
|
||||
</div>
|
||||
|
||||
<div class="form-group">
|
||||
<label for="profile_pic">Profile Picture URL</label>
|
||||
<input id="profile_pic" name="profile_pic" type="url" placeholder="https://example.com/photo.jpg" value="{{ contact.profile_pic or '' }}">
|
||||
</div>
|
||||
|
||||
<div class="form-group">
|
||||
<label for="bio">Bio</label>
|
||||
<textarea id="bio" name="bio" rows="3">{{ contact.bio or '' }}</textarea>
|
||||
</div>
|
||||
|
||||
<div class="form-group">
|
||||
<label for="goals">Goals</label>
|
||||
<textarea id="goals" name="goals" rows="3">{{ contact.goals or '' }}</textarea>
|
||||
</div>
|
||||
|
||||
<div class="form-group">
|
||||
<label for="social_structure_style">Social Structure Style</label>
|
||||
<input id="social_structure_style" name="social_structure_style" type="text" value="{{ contact.social_structure_style or '' }}">
|
||||
</div>
|
||||
|
||||
<div class="form-group">
|
||||
<label for="self_sufficiency_score">Self-Sufficiency Score (1-10)</label>
|
||||
<input id="self_sufficiency_score" name="self_sufficiency_score" type="number" min="1" max="10" value="{{ contact.self_sufficiency_score if contact and contact.self_sufficiency_score is not none else '' }}">
|
||||
</div>
|
||||
|
||||
<div class="form-group">
|
||||
<label for="safe_conversation_starters">Safe Conversation Starters</label>
|
||||
<textarea id="safe_conversation_starters" name="safe_conversation_starters" rows="2">{{ contact.safe_conversation_starters or '' }}</textarea>
|
||||
</div>
|
||||
|
||||
<div class="form-group">
|
||||
<label for="topics_to_avoid">Topics to Avoid</label>
|
||||
<textarea id="topics_to_avoid" name="topics_to_avoid" rows="2">{{ contact.topics_to_avoid or '' }}</textarea>
|
||||
</div>
|
||||
|
||||
<div class="form-group">
|
||||
<label for="ssn">SSN</label>
|
||||
<input id="ssn" name="ssn" type="text" value="{{ contact.ssn or '' }}">
|
||||
</div>
|
||||
|
||||
{% if all_needs %}
|
||||
<div class="form-group">
|
||||
<label>Needs/Accommodations</label>
|
||||
<div class="checkbox-group">
|
||||
{% for need in all_needs %}
|
||||
<label class="checkbox-label">
|
||||
<input type="checkbox" name="need_ids" value="{{ need.id }}"
|
||||
{% if contact and need in contact.needs %}checked{% endif %}>
|
||||
{{ need.name }}
|
||||
</label>
|
||||
{% endfor %}
|
||||
</div>
|
||||
</div>
|
||||
{% endif %}
|
||||
|
||||
<div class="form-actions">
|
||||
<button type="submit" class="btn btn-primary">Save</button>
|
||||
{% if contact %}
|
||||
<a href="/contacts/{{ contact.id }}" class="btn">Cancel</a>
|
||||
{% else %}
|
||||
<a href="/contacts" class="btn">Cancel</a>
|
||||
{% endif %}
|
||||
</div>
|
||||
</form>
|
||||
</div>
|
||||
{% endblock %}
|
||||
@@ -1,14 +0,0 @@
|
||||
{% extends "base.html" %}
|
||||
{% block title %}Contacts{% endblock %}
|
||||
{% block content %}
|
||||
<div class="contact-list">
|
||||
<div class="header">
|
||||
<h1>Contacts</h1>
|
||||
<a href="/contacts/new" class="btn btn-primary">Add Contact</a>
|
||||
</div>
|
||||
|
||||
<div id="contact-table">
|
||||
{% include "partials/contact_table.html" %}
|
||||
</div>
|
||||
</div>
|
||||
{% endblock %}
|
||||
@@ -1,198 +0,0 @@
|
||||
{% extends "base.html" %}
|
||||
{% block title %}Relationship Graph{% endblock %}
|
||||
{% block content %}
|
||||
<div class="graph-container">
|
||||
<div class="header">
|
||||
<h1>Relationship Graph</h1>
|
||||
</div>
|
||||
<p class="graph-hint">Drag nodes to reposition. Closer relationships have shorter, darker edges.</p>
|
||||
<canvas id="graph-canvas" width="900" height="600"
|
||||
style="border: 1px solid var(--color-border); border-radius: 8px; background: var(--color-bg); cursor: grab;">
|
||||
</canvas>
|
||||
<div id="selected-info"></div>
|
||||
<div class="legend">
|
||||
<h4>Relationship Closeness (1-10)</h4>
|
||||
<div class="legend-items">
|
||||
<div class="legend-item">
|
||||
<span class="legend-line" style="background: hsl(220, 70%, 40%); height: 4px; display: inline-block;"></span>
|
||||
<span>10 - Very Close (Spouse, Partner)</span>
|
||||
</div>
|
||||
<div class="legend-item">
|
||||
<span class="legend-line" style="background: hsl(220, 70%, 52%); height: 3px; display: inline-block;"></span>
|
||||
<span>7 - Close (Family, Best Friend)</span>
|
||||
</div>
|
||||
<div class="legend-item">
|
||||
<span class="legend-line" style="background: hsl(220, 70%, 64%); height: 2px; display: inline-block;"></span>
|
||||
<span>4 - Moderate (Friend, Colleague)</span>
|
||||
</div>
|
||||
<div class="legend-item">
|
||||
<span class="legend-line" style="background: hsl(220, 70%, 72%); height: 1px; display: inline-block;"></span>
|
||||
<span>2 - Distant (Acquaintance)</span>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<script>
|
||||
(function() {
|
||||
const RELATIONSHIP_DISPLAY = {{ relationship_type_display|tojson }};
|
||||
const graphData = {{ graph_data|tojson }};
|
||||
|
||||
const canvas = document.getElementById('graph-canvas');
|
||||
const ctx = canvas.getContext('2d');
|
||||
const width = canvas.width;
|
||||
const height = canvas.height;
|
||||
const centerX = width / 2;
|
||||
const centerY = height / 2;
|
||||
|
||||
const nodes = graphData.nodes.map(function(node) {
|
||||
return Object.assign({}, node, {
|
||||
x: centerX + (Math.random() - 0.5) * 300,
|
||||
y: centerY + (Math.random() - 0.5) * 300,
|
||||
vx: 0,
|
||||
vy: 0
|
||||
});
|
||||
});
|
||||
|
||||
const nodeMap = new Map(nodes.map(function(node) { return [node.id, node]; }));
|
||||
|
||||
const edges = graphData.edges.map(function(edge) {
|
||||
const sourceNode = nodeMap.get(edge.source);
|
||||
const targetNode = nodeMap.get(edge.target);
|
||||
if (!sourceNode || !targetNode) return null;
|
||||
return Object.assign({}, edge, { sourceNode: sourceNode, targetNode: targetNode });
|
||||
}).filter(function(edge) { return edge !== null; });
|
||||
|
||||
let dragNode = null;
|
||||
let selectedNode = null;
|
||||
|
||||
const repulsion = 5000;
|
||||
const springStrength = 0.05;
|
||||
const baseSpringLength = 150;
|
||||
const damping = 0.9;
|
||||
const centerPull = 0.01;
|
||||
|
||||
function simulate() {
|
||||
for (const node of nodes) { node.vx = 0; node.vy = 0; }
|
||||
for (let i = 0; i < nodes.length; i++) {
|
||||
for (let j = i + 1; j < nodes.length; j++) {
|
||||
const dx = nodes[j].x - nodes[i].x;
|
||||
const dy = nodes[j].y - nodes[i].y;
|
||||
const dist = Math.sqrt(dx * dx + dy * dy) || 1;
|
||||
const force = repulsion / (dist * dist);
|
||||
const fx = (dx / dist) * force;
|
||||
const fy = (dy / dist) * force;
|
||||
nodes[i].vx -= fx; nodes[i].vy -= fy;
|
||||
nodes[j].vx += fx; nodes[j].vy += fy;
|
||||
}
|
||||
}
|
||||
for (const edge of edges) {
|
||||
const dx = edge.targetNode.x - edge.sourceNode.x;
|
||||
const dy = edge.targetNode.y - edge.sourceNode.y;
|
||||
const dist = Math.sqrt(dx * dx + dy * dy) || 1;
|
||||
const normalizedWeight = edge.closeness_weight / 10;
|
||||
const idealLength = baseSpringLength * (1.5 - normalizedWeight);
|
||||
const displacement = dist - idealLength;
|
||||
const force = springStrength * displacement;
|
||||
const fx = (dx / dist) * force;
|
||||
const fy = (dy / dist) * force;
|
||||
edge.sourceNode.vx += fx; edge.sourceNode.vy += fy;
|
||||
edge.targetNode.vx -= fx; edge.targetNode.vy -= fy;
|
||||
}
|
||||
for (const node of nodes) {
|
||||
node.vx += (centerX - node.x) * centerPull;
|
||||
node.vy += (centerY - node.y) * centerPull;
|
||||
}
|
||||
for (const node of nodes) {
|
||||
if (node === dragNode) continue;
|
||||
node.x += node.vx * damping;
|
||||
node.y += node.vy * damping;
|
||||
node.x = Math.max(30, Math.min(width - 30, node.x));
|
||||
node.y = Math.max(30, Math.min(height - 30, node.y));
|
||||
}
|
||||
}
|
||||
|
||||
function getEdgeColor(weight) {
|
||||
const normalized = weight / 10;
|
||||
return 'hsl(220, 70%, ' + (80 - normalized * 40) + '%)';
|
||||
}
|
||||
|
||||
function draw() {
|
||||
ctx.clearRect(0, 0, width, height);
|
||||
for (const edge of edges) {
|
||||
const lineWidth = 1 + (edge.closeness_weight / 10) * 3;
|
||||
ctx.strokeStyle = getEdgeColor(edge.closeness_weight);
|
||||
ctx.lineWidth = lineWidth;
|
||||
ctx.beginPath();
|
||||
ctx.moveTo(edge.sourceNode.x, edge.sourceNode.y);
|
||||
ctx.lineTo(edge.targetNode.x, edge.targetNode.y);
|
||||
ctx.stroke();
|
||||
const midX = (edge.sourceNode.x + edge.targetNode.x) / 2;
|
||||
const midY = (edge.sourceNode.y + edge.targetNode.y) / 2;
|
||||
ctx.fillStyle = '#666';
|
||||
ctx.font = '10px sans-serif';
|
||||
ctx.textAlign = 'center';
|
||||
const label = RELATIONSHIP_DISPLAY[edge.relationship_type] || edge.relationship_type;
|
||||
ctx.fillText(label, midX, midY - 5);
|
||||
}
|
||||
for (const node of nodes) {
|
||||
const isSelected = node === selectedNode;
|
||||
const radius = isSelected ? 25 : 20;
|
||||
ctx.beginPath();
|
||||
ctx.arc(node.x, node.y, radius, 0, Math.PI * 2);
|
||||
ctx.fillStyle = isSelected ? '#0066cc' : '#fff';
|
||||
ctx.fill();
|
||||
ctx.strokeStyle = '#0066cc';
|
||||
ctx.lineWidth = 2;
|
||||
ctx.stroke();
|
||||
ctx.fillStyle = isSelected ? '#fff' : '#333';
|
||||
ctx.font = '12px sans-serif';
|
||||
ctx.textAlign = 'center';
|
||||
ctx.textBaseline = 'middle';
|
||||
const name = node.name.length > 10 ? node.name.slice(0, 9) + '\u2026' : node.name;
|
||||
ctx.fillText(name, node.x, node.y);
|
||||
}
|
||||
}
|
||||
|
||||
function animate() {
|
||||
simulate();
|
||||
draw();
|
||||
requestAnimationFrame(animate);
|
||||
}
|
||||
animate();
|
||||
|
||||
function getNodeAt(x, y) {
|
||||
for (const node of nodes) {
|
||||
const dx = x - node.x;
|
||||
const dy = y - node.y;
|
||||
if (dx * dx + dy * dy < 400) return node;
|
||||
}
|
||||
return null;
|
||||
}
|
||||
|
||||
canvas.addEventListener('mousedown', function(event) {
|
||||
const rect = canvas.getBoundingClientRect();
|
||||
const node = getNodeAt(event.clientX - rect.left, event.clientY - rect.top);
|
||||
if (node) {
|
||||
dragNode = node;
|
||||
selectedNode = node;
|
||||
const infoDiv = document.getElementById('selected-info');
|
||||
let html = '<div class="selected-info"><h3>' + node.name + '</h3>';
|
||||
if (node.current_job) html += '<p>Job: ' + node.current_job + '</p>';
|
||||
html += '<a href="/contacts/' + node.id + '">View details</a></div>';
|
||||
infoDiv.innerHTML = html;
|
||||
}
|
||||
});
|
||||
|
||||
canvas.addEventListener('mousemove', function(event) {
|
||||
if (!dragNode) return;
|
||||
const rect = canvas.getBoundingClientRect();
|
||||
dragNode.x = event.clientX - rect.left;
|
||||
dragNode.y = event.clientY - rect.top;
|
||||
});
|
||||
|
||||
canvas.addEventListener('mouseup', function() { dragNode = null; });
|
||||
canvas.addEventListener('mouseleave', function() { dragNode = null; });
|
||||
})();
|
||||
</script>
|
||||
{% endblock %}
|
||||
@@ -1,31 +0,0 @@
|
||||
{% extends "base.html" %}
|
||||
{% block title %}Needs{% endblock %}
|
||||
{% block content %}
|
||||
<div class="need-list">
|
||||
<div class="header">
|
||||
<h1>Needs / Accommodations</h1>
|
||||
<button class="btn btn-primary" onclick="document.getElementById('need-form').toggleAttribute('hidden')">Add Need</button>
|
||||
</div>
|
||||
|
||||
<form id="need-form" hidden
|
||||
hx-post="/htmx/needs"
|
||||
hx-target="#need-items"
|
||||
hx-swap="innerHTML"
|
||||
hx-on::after-request="if(event.detail.successful) this.reset()"
|
||||
class="need-form">
|
||||
<div class="form-group">
|
||||
<label for="name">Name *</label>
|
||||
<input id="name" name="name" type="text" placeholder="e.g., Light Sensitive, ADHD" required>
|
||||
</div>
|
||||
<div class="form-group">
|
||||
<label for="description">Description</label>
|
||||
<textarea id="description" name="description" placeholder="Optional description..." rows="2"></textarea>
|
||||
</div>
|
||||
<button type="submit" class="btn btn-primary">Create</button>
|
||||
</form>
|
||||
|
||||
<div id="need-items">
|
||||
{% include "partials/need_items.html" %}
|
||||
</div>
|
||||
</div>
|
||||
{% endblock %}
|
||||
@@ -1,33 +0,0 @@
|
||||
{% if contacts %}
|
||||
<table>
|
||||
<thead>
|
||||
<tr>
|
||||
<th>Name</th>
|
||||
<th>Job</th>
|
||||
<th>Timezone</th>
|
||||
<th>Actions</th>
|
||||
</tr>
|
||||
</thead>
|
||||
<tbody>
|
||||
{% for contact in contacts %}
|
||||
<tr id="contact-row-{{ contact.id }}">
|
||||
<td><a href="/contacts/{{ contact.id }}">{{ contact.name }}</a></td>
|
||||
<td>{{ contact.current_job or "-" }}</td>
|
||||
<td>{{ contact.timezone or "-" }}</td>
|
||||
<td>
|
||||
<a href="/contacts/{{ contact.id }}/edit" class="btn">Edit</a>
|
||||
<button class="btn btn-danger"
|
||||
hx-delete="/api/contacts/{{ contact.id }}"
|
||||
hx-target="#contact-row-{{ contact.id }}"
|
||||
hx-swap="outerHTML"
|
||||
hx-confirm="Delete this contact?">
|
||||
Delete
|
||||
</button>
|
||||
</td>
|
||||
</tr>
|
||||
{% endfor %}
|
||||
</tbody>
|
||||
</table>
|
||||
{% else %}
|
||||
<p>No contacts yet.</p>
|
||||
{% endif %}
|
||||
@@ -1,14 +0,0 @@
|
||||
<ul class="manage-needs-list">
|
||||
{% for need in contact.needs %}
|
||||
<li id="contact-need-{{ need.id }}">
|
||||
<strong>{{ need.name }}</strong>
|
||||
{% if need.description %}<span> - {{ need.description }}</span>{% endif %}
|
||||
<button class="btn btn-small btn-danger"
|
||||
hx-delete="/api/contacts/{{ contact.id }}/needs/{{ need.id }}"
|
||||
hx-target="#contact-need-{{ need.id }}"
|
||||
hx-swap="outerHTML">
|
||||
Remove
|
||||
</button>
|
||||
</li>
|
||||
{% endfor %}
|
||||
</ul>
|
||||
@@ -1,23 +0,0 @@
|
||||
{% for rel in contact.related_to %}
|
||||
<div class="manage-rel-item" id="rel-{{ contact.id }}-{{ rel.related_contact_id }}">
|
||||
<a href="/contacts/{{ rel.related_contact_id }}">{{ contact_names[rel.related_contact_id] }}</a>
|
||||
<span class="tag">{{ rel.relationship_type|replace("_", " ")|title }}</span>
|
||||
<label class="weight-control">
|
||||
<span>Closeness:</span>
|
||||
<input type="range" min="1" max="10" value="{{ rel.closeness_weight }}"
|
||||
hx-post="/htmx/contacts/{{ contact.id }}/relationships/{{ rel.related_contact_id }}/weight"
|
||||
hx-trigger="change"
|
||||
hx-include="this"
|
||||
name="closeness_weight"
|
||||
hx-swap="none"
|
||||
oninput="this.nextElementSibling.textContent = this.value">
|
||||
<span class="weight-value">{{ rel.closeness_weight }}</span>
|
||||
</label>
|
||||
<button class="btn btn-small btn-danger"
|
||||
hx-delete="/api/contacts/{{ contact.id }}/relationships/{{ rel.related_contact_id }}"
|
||||
hx-target="#rel-{{ contact.id }}-{{ rel.related_contact_id }}"
|
||||
hx-swap="outerHTML">
|
||||
Remove
|
||||
</button>
|
||||
</div>
|
||||
{% endfor %}
|
||||
@@ -1,21 +0,0 @@
|
||||
{% if needs %}
|
||||
<ul class="need-items">
|
||||
{% for need in needs %}
|
||||
<li id="need-item-{{ need.id }}">
|
||||
<div class="need-info">
|
||||
<strong>{{ need.name }}</strong>
|
||||
{% if need.description %}<p>{{ need.description }}</p>{% endif %}
|
||||
</div>
|
||||
<button class="btn btn-danger"
|
||||
hx-delete="/api/needs/{{ need.id }}"
|
||||
hx-target="#need-item-{{ need.id }}"
|
||||
hx-swap="outerHTML"
|
||||
hx-confirm="Delete this need?">
|
||||
Delete
|
||||
</button>
|
||||
</li>
|
||||
{% endfor %}
|
||||
</ul>
|
||||
{% else %}
|
||||
<p>No needs defined yet.</p>
|
||||
{% endif %}
|
||||
@@ -1,3 +0,0 @@
|
||||
"""Data science CLI tools."""
|
||||
|
||||
from __future__ import annotations
|
||||
@@ -1,613 +0,0 @@
|
||||
"""Ingestion pipeline for loading congress data from unitedstates/congress JSON files.
|
||||
|
||||
Loads legislators, bills, votes, vote records, and bill text into the data_science_dev database.
|
||||
Expects the parent directory to contain congress-tracker/ and congress-legislators/ as siblings.
|
||||
|
||||
Usage:
|
||||
ingest-congress /path/to/parent/
|
||||
ingest-congress /path/to/parent/ --congress 118
|
||||
ingest-congress /path/to/parent/ --congress 118 --only bills
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from pathlib import Path # noqa: TC003 needed at runtime for typer CLI argument
|
||||
from typing import TYPE_CHECKING, Annotated
|
||||
|
||||
import orjson
|
||||
import typer
|
||||
import yaml
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from python.common import configure_logger
|
||||
from python.orm.common import get_postgres_engine
|
||||
from python.orm.data_science_dev.congress import Bill, BillText, Legislator, LegislatorSocialMedia, Vote, VoteRecord
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Iterator
|
||||
|
||||
from sqlalchemy.engine import Engine
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
BATCH_SIZE = 10_000
|
||||
|
||||
app = typer.Typer(help="Ingest unitedstates/congress data into data_science_dev.")
|
||||
|
||||
|
||||
@app.command()
|
||||
def main(
|
||||
parent_dir: Annotated[
|
||||
Path,
|
||||
typer.Argument(help="Parent directory containing congress-tracker/ and congress-legislators/"),
|
||||
],
|
||||
congress: Annotated[int | None, typer.Option(help="Only ingest a specific congress number")] = None,
|
||||
only: Annotated[
|
||||
str | None,
|
||||
typer.Option(help="Only run a specific step: legislators, social-media, bills, votes, bill-text"),
|
||||
] = None,
|
||||
) -> None:
|
||||
"""Ingest congress data from unitedstates/congress JSON files."""
|
||||
configure_logger(level="INFO")
|
||||
|
||||
data_dir = parent_dir / "congress-tracker/congress/data/"
|
||||
legislators_dir = parent_dir / "congress-legislators"
|
||||
|
||||
if not data_dir.is_dir():
|
||||
typer.echo(f"Expected congress-tracker/ directory: {data_dir}", err=True)
|
||||
raise typer.Exit(code=1)
|
||||
|
||||
if not legislators_dir.is_dir():
|
||||
typer.echo(f"Expected congress-legislators/ directory: {legislators_dir}", err=True)
|
||||
raise typer.Exit(code=1)
|
||||
|
||||
engine = get_postgres_engine(name="DATA_SCIENCE_DEV")
|
||||
|
||||
congress_dirs = _resolve_congress_dirs(data_dir, congress)
|
||||
if not congress_dirs:
|
||||
typer.echo("No congress directories found.", err=True)
|
||||
raise typer.Exit(code=1)
|
||||
|
||||
logger.info("Found %d congress directories to process", len(congress_dirs))
|
||||
|
||||
steps: dict[str, tuple] = {
|
||||
"legislators": (ingest_legislators, (engine, legislators_dir)),
|
||||
"legislators-social-media": (ingest_social_media, (engine, legislators_dir)),
|
||||
"bills": (ingest_bills, (engine, congress_dirs)),
|
||||
"votes": (ingest_votes, (engine, congress_dirs)),
|
||||
"bill-text": (ingest_bill_text, (engine, congress_dirs)),
|
||||
}
|
||||
|
||||
if only:
|
||||
if only not in steps:
|
||||
typer.echo(f"Unknown step: {only}. Choose from: {', '.join(steps)}", err=True)
|
||||
raise typer.Exit(code=1)
|
||||
steps = {only: steps[only]}
|
||||
|
||||
for step_name, (step_func, step_args) in steps.items():
|
||||
logger.info("=== Starting step: %s ===", step_name)
|
||||
step_func(*step_args)
|
||||
logger.info("=== Finished step: %s ===", step_name)
|
||||
|
||||
logger.info("ingest-congress done")
|
||||
|
||||
|
||||
def _resolve_congress_dirs(data_dir: Path, congress: int | None) -> list[Path]:
|
||||
"""Find congress number directories under data_dir."""
|
||||
if congress is not None:
|
||||
target = data_dir / str(congress)
|
||||
return [target] if target.is_dir() else []
|
||||
return sorted(path for path in data_dir.iterdir() if path.is_dir() and path.name.isdigit())
|
||||
|
||||
|
||||
def _flush_batch(session: Session, batch: list[object], label: str) -> int:
|
||||
"""Add a batch of ORM objects to the session and commit. Returns count added."""
|
||||
if not batch:
|
||||
return 0
|
||||
session.add_all(batch)
|
||||
session.commit()
|
||||
count = len(batch)
|
||||
logger.info("Committed %d %s", count, label)
|
||||
batch.clear()
|
||||
return count
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Legislators — loaded from congress-legislators YAML files
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def ingest_legislators(engine: Engine, legislators_dir: Path) -> None:
|
||||
"""Load legislators from congress-legislators YAML files."""
|
||||
legislators_data = _load_legislators_yaml(legislators_dir)
|
||||
logger.info("Loaded %d legislators from YAML files", len(legislators_data))
|
||||
|
||||
with Session(engine) as session:
|
||||
existing_legislators = {
|
||||
legislator.bioguide_id: legislator for legislator in session.scalars(select(Legislator)).all()
|
||||
}
|
||||
logger.info("Found %d existing legislators in DB", len(existing_legislators))
|
||||
|
||||
total_inserted = 0
|
||||
total_updated = 0
|
||||
for entry in legislators_data:
|
||||
bioguide_id = entry.get("id", {}).get("bioguide")
|
||||
if not bioguide_id:
|
||||
continue
|
||||
|
||||
fields = _parse_legislator(entry)
|
||||
if existing := existing_legislators.get(bioguide_id):
|
||||
changed = False
|
||||
for field, value in fields.items():
|
||||
if value is not None and getattr(existing, field) != value:
|
||||
setattr(existing, field, value)
|
||||
changed = True
|
||||
if changed:
|
||||
total_updated += 1
|
||||
else:
|
||||
session.add(Legislator(bioguide_id=bioguide_id, **fields))
|
||||
total_inserted += 1
|
||||
|
||||
session.commit()
|
||||
logger.info("Inserted %d new legislators, updated %d existing", total_inserted, total_updated)
|
||||
|
||||
|
||||
def _load_legislators_yaml(legislators_dir: Path) -> list[dict]:
|
||||
"""Load and combine legislators-current.yaml and legislators-historical.yaml."""
|
||||
legislators: list[dict] = []
|
||||
for filename in ("legislators-current.yaml", "legislators-historical.yaml"):
|
||||
path = legislators_dir / filename
|
||||
if not path.exists():
|
||||
logger.warning("Legislators file not found: %s", path)
|
||||
continue
|
||||
with path.open() as file:
|
||||
data = yaml.safe_load(file)
|
||||
if isinstance(data, list):
|
||||
legislators.extend(data)
|
||||
return legislators
|
||||
|
||||
|
||||
def _parse_legislator(entry: dict) -> dict:
|
||||
"""Extract Legislator fields from a congress-legislators YAML entry."""
|
||||
ids = entry.get("id", {})
|
||||
name = entry.get("name", {})
|
||||
bio = entry.get("bio", {})
|
||||
terms = entry.get("terms", [])
|
||||
latest_term = terms[-1] if terms else {}
|
||||
|
||||
fec_ids = ids.get("fec")
|
||||
fec_ids_joined = ",".join(fec_ids) if isinstance(fec_ids, list) else fec_ids
|
||||
|
||||
chamber = latest_term.get("type")
|
||||
chamber_normalized = {"rep": "House", "sen": "Senate"}.get(chamber, chamber)
|
||||
|
||||
return {
|
||||
"thomas_id": ids.get("thomas"),
|
||||
"lis_id": ids.get("lis"),
|
||||
"govtrack_id": ids.get("govtrack"),
|
||||
"opensecrets_id": ids.get("opensecrets"),
|
||||
"fec_ids": fec_ids_joined,
|
||||
"first_name": name.get("first"),
|
||||
"last_name": name.get("last"),
|
||||
"official_full_name": name.get("official_full"),
|
||||
"nickname": name.get("nickname"),
|
||||
"birthday": bio.get("birthday"),
|
||||
"gender": bio.get("gender"),
|
||||
"current_party": latest_term.get("party"),
|
||||
"current_state": latest_term.get("state"),
|
||||
"current_district": latest_term.get("district"),
|
||||
"current_chamber": chamber_normalized,
|
||||
}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Social Media — loaded from legislators-social-media.yaml
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
SOCIAL_MEDIA_PLATFORMS = {
|
||||
"twitter": "https://twitter.com/{account}",
|
||||
"facebook": "https://facebook.com/{account}",
|
||||
"youtube": "https://youtube.com/{account}",
|
||||
"instagram": "https://instagram.com/{account}",
|
||||
"mastodon": None,
|
||||
}
|
||||
|
||||
|
||||
def ingest_social_media(engine: Engine, legislators_dir: Path) -> None:
|
||||
"""Load social media accounts from legislators-social-media.yaml."""
|
||||
social_media_path = legislators_dir / "legislators-social-media.yaml"
|
||||
if not social_media_path.exists():
|
||||
logger.warning("Social media file not found: %s", social_media_path)
|
||||
return
|
||||
|
||||
with social_media_path.open() as file:
|
||||
social_media_data = yaml.safe_load(file)
|
||||
|
||||
if not isinstance(social_media_data, list):
|
||||
logger.warning("Unexpected format in %s", social_media_path)
|
||||
return
|
||||
|
||||
logger.info("Loaded %d entries from legislators-social-media.yaml", len(social_media_data))
|
||||
|
||||
with Session(engine) as session:
|
||||
legislator_map = _build_legislator_map(session)
|
||||
existing_accounts = {
|
||||
(account.legislator_id, account.platform)
|
||||
for account in session.scalars(select(LegislatorSocialMedia)).all()
|
||||
}
|
||||
logger.info("Found %d existing social media accounts in DB", len(existing_accounts))
|
||||
|
||||
total_inserted = 0
|
||||
total_updated = 0
|
||||
for entry in social_media_data:
|
||||
bioguide_id = entry.get("id", {}).get("bioguide")
|
||||
if not bioguide_id:
|
||||
continue
|
||||
|
||||
legislator_id = legislator_map.get(bioguide_id)
|
||||
if legislator_id is None:
|
||||
continue
|
||||
|
||||
social = entry.get("social", {})
|
||||
for platform, url_template in SOCIAL_MEDIA_PLATFORMS.items():
|
||||
account_name = social.get(platform)
|
||||
if not account_name:
|
||||
continue
|
||||
|
||||
url = url_template.format(account=account_name) if url_template else None
|
||||
|
||||
if (legislator_id, platform) in existing_accounts:
|
||||
total_updated += 1
|
||||
else:
|
||||
session.add(
|
||||
LegislatorSocialMedia(
|
||||
legislator_id=legislator_id,
|
||||
platform=platform,
|
||||
account_name=str(account_name),
|
||||
url=url,
|
||||
source="https://github.com/unitedstates/congress-legislators",
|
||||
)
|
||||
)
|
||||
existing_accounts.add((legislator_id, platform))
|
||||
total_inserted += 1
|
||||
|
||||
session.commit()
|
||||
logger.info("Inserted %d new social media accounts, updated %d existing", total_inserted, total_updated)
|
||||
|
||||
|
||||
def _iter_voters(position_group: object) -> Iterator[dict]:
|
||||
"""Yield voter dicts from a vote position group (handles list, single dict, or string)."""
|
||||
if isinstance(position_group, dict):
|
||||
yield position_group
|
||||
elif isinstance(position_group, list):
|
||||
for voter in position_group:
|
||||
if isinstance(voter, dict):
|
||||
yield voter
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Bills
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def ingest_bills(engine: Engine, congress_dirs: list[Path]) -> None:
|
||||
"""Load bill data.json files."""
|
||||
with Session(engine) as session:
|
||||
existing_bills = {(bill.congress, bill.bill_type, bill.number) for bill in session.scalars(select(Bill)).all()}
|
||||
logger.info("Found %d existing bills in DB", len(existing_bills))
|
||||
|
||||
total_inserted = 0
|
||||
batch: list[Bill] = []
|
||||
for congress_dir in congress_dirs:
|
||||
bills_dir = congress_dir / "bills"
|
||||
if not bills_dir.is_dir():
|
||||
continue
|
||||
logger.info("Scanning bills from %s", congress_dir.name)
|
||||
for bill_file in bills_dir.rglob("data.json"):
|
||||
data = _read_json(bill_file)
|
||||
if data is None:
|
||||
continue
|
||||
bill = _parse_bill(data, existing_bills)
|
||||
if bill is not None:
|
||||
batch.append(bill)
|
||||
if len(batch) >= BATCH_SIZE:
|
||||
total_inserted += _flush_batch(session, batch, "bills")
|
||||
|
||||
total_inserted += _flush_batch(session, batch, "bills")
|
||||
logger.info("Inserted %d new bills total", total_inserted)
|
||||
|
||||
|
||||
def _parse_bill(data: dict, existing_bills: set[tuple[int, str, int]]) -> Bill | None:
|
||||
"""Parse a bill data.json dict into a Bill ORM object, skipping existing."""
|
||||
raw_congress = data.get("congress")
|
||||
bill_type = data.get("bill_type")
|
||||
raw_number = data.get("number")
|
||||
if raw_congress is None or bill_type is None or raw_number is None:
|
||||
return None
|
||||
congress = int(raw_congress)
|
||||
number = int(raw_number)
|
||||
if (congress, bill_type, number) in existing_bills:
|
||||
return None
|
||||
|
||||
sponsor_bioguide = None
|
||||
sponsor = data.get("sponsor")
|
||||
if sponsor:
|
||||
sponsor_bioguide = sponsor.get("bioguide_id")
|
||||
|
||||
return Bill(
|
||||
congress=congress,
|
||||
bill_type=bill_type,
|
||||
number=number,
|
||||
title=data.get("short_title") or data.get("official_title"),
|
||||
title_short=data.get("short_title"),
|
||||
official_title=data.get("official_title"),
|
||||
status=data.get("status"),
|
||||
status_at=data.get("status_at"),
|
||||
sponsor_bioguide_id=sponsor_bioguide,
|
||||
subjects_top_term=data.get("subjects_top_term"),
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Votes (and vote records)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def ingest_votes(engine: Engine, congress_dirs: list[Path]) -> None:
|
||||
"""Load vote data.json files with their vote records."""
|
||||
with Session(engine) as session:
|
||||
legislator_map = _build_legislator_map(session)
|
||||
logger.info("Loaded %d legislators into lookup map", len(legislator_map))
|
||||
bill_map = _build_bill_map(session)
|
||||
logger.info("Loaded %d bills into lookup map", len(bill_map))
|
||||
existing_votes = {
|
||||
(vote.congress, vote.chamber, vote.session, vote.number) for vote in session.scalars(select(Vote)).all()
|
||||
}
|
||||
logger.info("Found %d existing votes in DB", len(existing_votes))
|
||||
|
||||
total_inserted = 0
|
||||
batch: list[Vote] = []
|
||||
for congress_dir in congress_dirs:
|
||||
votes_dir = congress_dir / "votes"
|
||||
if not votes_dir.is_dir():
|
||||
continue
|
||||
logger.info("Scanning votes from %s", congress_dir.name)
|
||||
for vote_file in votes_dir.rglob("data.json"):
|
||||
data = _read_json(vote_file)
|
||||
if data is None:
|
||||
continue
|
||||
vote = _parse_vote(data, legislator_map, bill_map, existing_votes)
|
||||
if vote is not None:
|
||||
batch.append(vote)
|
||||
if len(batch) >= BATCH_SIZE:
|
||||
total_inserted += _flush_batch(session, batch, "votes")
|
||||
|
||||
total_inserted += _flush_batch(session, batch, "votes")
|
||||
logger.info("Inserted %d new votes total", total_inserted)
|
||||
|
||||
|
||||
def _build_legislator_map(session: Session) -> dict[str, int]:
|
||||
"""Build a mapping of bioguide_id -> legislator.id."""
|
||||
return {legislator.bioguide_id: legislator.id for legislator in session.scalars(select(Legislator)).all()}
|
||||
|
||||
|
||||
def _build_bill_map(session: Session) -> dict[tuple[int, str, int], int]:
|
||||
"""Build a mapping of (congress, bill_type, number) -> bill.id."""
|
||||
return {(bill.congress, bill.bill_type, bill.number): bill.id for bill in session.scalars(select(Bill)).all()}
|
||||
|
||||
|
||||
def _parse_vote(
|
||||
data: dict,
|
||||
legislator_map: dict[str, int],
|
||||
bill_map: dict[tuple[int, str, int], int],
|
||||
existing_votes: set[tuple[int, str, int, int]],
|
||||
) -> Vote | None:
|
||||
"""Parse a vote data.json dict into a Vote ORM object with records."""
|
||||
raw_congress = data.get("congress")
|
||||
chamber = data.get("chamber")
|
||||
raw_number = data.get("number")
|
||||
vote_date = data.get("date")
|
||||
if raw_congress is None or chamber is None or raw_number is None or vote_date is None:
|
||||
return None
|
||||
|
||||
raw_session = data.get("session")
|
||||
if raw_session is None:
|
||||
return None
|
||||
|
||||
congress = int(raw_congress)
|
||||
number = int(raw_number)
|
||||
session_number = int(raw_session)
|
||||
|
||||
# Normalize chamber from "h"/"s" to "House"/"Senate"
|
||||
chamber_normalized = {"h": "House", "s": "Senate"}.get(chamber, chamber)
|
||||
|
||||
if (congress, chamber_normalized, session_number, number) in existing_votes:
|
||||
return None
|
||||
|
||||
# Resolve linked bill
|
||||
bill_id = None
|
||||
bill_ref = data.get("bill")
|
||||
if bill_ref:
|
||||
bill_key = (
|
||||
int(bill_ref.get("congress", congress)),
|
||||
bill_ref.get("type"),
|
||||
int(bill_ref.get("number", 0)),
|
||||
)
|
||||
bill_id = bill_map.get(bill_key)
|
||||
|
||||
raw_votes = data.get("votes", {})
|
||||
vote_counts = _count_votes(raw_votes)
|
||||
vote_records = _build_vote_records(raw_votes, legislator_map)
|
||||
|
||||
return Vote(
|
||||
congress=congress,
|
||||
chamber=chamber_normalized,
|
||||
session=session_number,
|
||||
number=number,
|
||||
vote_type=data.get("type"),
|
||||
question=data.get("question"),
|
||||
result=data.get("result"),
|
||||
result_text=data.get("result_text"),
|
||||
vote_date=vote_date[:10] if isinstance(vote_date, str) else vote_date,
|
||||
bill_id=bill_id,
|
||||
vote_records=vote_records,
|
||||
**vote_counts,
|
||||
)
|
||||
|
||||
|
||||
def _count_votes(raw_votes: dict) -> dict[str, int]:
|
||||
"""Count voters per position category, correctly handling dict and list formats."""
|
||||
yea_count = 0
|
||||
nay_count = 0
|
||||
not_voting_count = 0
|
||||
present_count = 0
|
||||
|
||||
for position, position_group in raw_votes.items():
|
||||
voter_count = sum(1 for _ in _iter_voters(position_group))
|
||||
if position in ("Yea", "Aye"):
|
||||
yea_count += voter_count
|
||||
elif position in ("Nay", "No"):
|
||||
nay_count += voter_count
|
||||
elif position == "Not Voting":
|
||||
not_voting_count += voter_count
|
||||
elif position == "Present":
|
||||
present_count += voter_count
|
||||
|
||||
return {
|
||||
"yea_count": yea_count,
|
||||
"nay_count": nay_count,
|
||||
"not_voting_count": not_voting_count,
|
||||
"present_count": present_count,
|
||||
}
|
||||
|
||||
|
||||
def _build_vote_records(raw_votes: dict, legislator_map: dict[str, int]) -> list[VoteRecord]:
|
||||
"""Build VoteRecord objects from raw vote data."""
|
||||
records: list[VoteRecord] = []
|
||||
for position, position_group in raw_votes.items():
|
||||
for voter in _iter_voters(position_group):
|
||||
bioguide_id = voter.get("id")
|
||||
if not bioguide_id:
|
||||
continue
|
||||
legislator_id = legislator_map.get(bioguide_id)
|
||||
if legislator_id is None:
|
||||
continue
|
||||
records.append(
|
||||
VoteRecord(
|
||||
legislator_id=legislator_id,
|
||||
position=position,
|
||||
)
|
||||
)
|
||||
return records
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Bill Text
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def ingest_bill_text(engine: Engine, congress_dirs: list[Path]) -> None:
|
||||
"""Load bill text from text-versions directories."""
|
||||
with Session(engine) as session:
|
||||
bill_map = _build_bill_map(session)
|
||||
logger.info("Loaded %d bills into lookup map", len(bill_map))
|
||||
existing_bill_texts = {
|
||||
(bill_text.bill_id, bill_text.version_code) for bill_text in session.scalars(select(BillText)).all()
|
||||
}
|
||||
logger.info("Found %d existing bill text versions in DB", len(existing_bill_texts))
|
||||
|
||||
total_inserted = 0
|
||||
batch: list[BillText] = []
|
||||
for congress_dir in congress_dirs:
|
||||
logger.info("Scanning bill texts from %s", congress_dir.name)
|
||||
for bill_text in _iter_bill_texts(congress_dir, bill_map, existing_bill_texts):
|
||||
batch.append(bill_text)
|
||||
if len(batch) >= BATCH_SIZE:
|
||||
total_inserted += _flush_batch(session, batch, "bill texts")
|
||||
|
||||
total_inserted += _flush_batch(session, batch, "bill texts")
|
||||
logger.info("Inserted %d new bill text versions total", total_inserted)
|
||||
|
||||
|
||||
def _iter_bill_texts(
|
||||
congress_dir: Path,
|
||||
bill_map: dict[tuple[int, str, int], int],
|
||||
existing_bill_texts: set[tuple[int, str]],
|
||||
) -> Iterator[BillText]:
|
||||
"""Yield BillText objects for a single congress directory, skipping existing."""
|
||||
bills_dir = congress_dir / "bills"
|
||||
if not bills_dir.is_dir():
|
||||
return
|
||||
|
||||
for bill_dir in bills_dir.rglob("text-versions"):
|
||||
if not bill_dir.is_dir():
|
||||
continue
|
||||
bill_key = _bill_key_from_dir(bill_dir.parent, congress_dir)
|
||||
if bill_key is None:
|
||||
continue
|
||||
bill_id = bill_map.get(bill_key)
|
||||
if bill_id is None:
|
||||
continue
|
||||
|
||||
for version_dir in sorted(bill_dir.iterdir()):
|
||||
if not version_dir.is_dir():
|
||||
continue
|
||||
if (bill_id, version_dir.name) in existing_bill_texts:
|
||||
continue
|
||||
text_content = _read_bill_text(version_dir)
|
||||
version_data = _read_json(version_dir / "data.json")
|
||||
yield BillText(
|
||||
bill_id=bill_id,
|
||||
version_code=version_dir.name,
|
||||
version_name=version_data.get("version_name") if version_data else None,
|
||||
date=version_data.get("issued_on") if version_data else None,
|
||||
text_content=text_content,
|
||||
)
|
||||
|
||||
|
||||
def _bill_key_from_dir(bill_dir: Path, congress_dir: Path) -> tuple[int, str, int] | None:
|
||||
"""Extract (congress, bill_type, number) from directory structure."""
|
||||
congress = int(congress_dir.name)
|
||||
bill_type = bill_dir.parent.name
|
||||
name = bill_dir.name
|
||||
# Directory name is like "hr3590" — strip the type prefix to get the number
|
||||
number_str = name[len(bill_type) :]
|
||||
if not number_str.isdigit():
|
||||
return None
|
||||
return (congress, bill_type, int(number_str))
|
||||
|
||||
|
||||
def _read_bill_text(version_dir: Path) -> str | None:
|
||||
"""Read bill text from a version directory, preferring .txt over .xml."""
|
||||
for extension in ("txt", "htm", "html", "xml"):
|
||||
candidates = list(version_dir.glob(f"document.{extension}"))
|
||||
if not candidates:
|
||||
candidates = list(version_dir.glob(f"*.{extension}"))
|
||||
if candidates:
|
||||
try:
|
||||
return candidates[0].read_text(encoding="utf-8")
|
||||
except Exception:
|
||||
logger.exception("Failed to read %s", candidates[0])
|
||||
return None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _read_json(path: Path) -> dict | None:
|
||||
"""Read and parse a JSON file, returning None on failure."""
|
||||
try:
|
||||
return orjson.loads(path.read_bytes())
|
||||
except FileNotFoundError:
|
||||
return None
|
||||
except Exception:
|
||||
logger.exception("Failed to parse %s", path)
|
||||
return None
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
app()
|
||||
@@ -1,247 +0,0 @@
|
||||
"""Ingestion pipeline for loading JSONL post files into the weekly-partitioned posts table.
|
||||
|
||||
Usage:
|
||||
ingest-posts /path/to/files/
|
||||
ingest-posts /path/to/single_file.jsonl
|
||||
ingest-posts /data/dir/ --workers 4 --batch-size 5000
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from datetime import UTC, datetime
|
||||
from pathlib import Path # noqa: TC003 this is needed for typer
|
||||
from typing import TYPE_CHECKING, Annotated
|
||||
|
||||
import orjson
|
||||
import psycopg
|
||||
import typer
|
||||
|
||||
from python.common import configure_logger
|
||||
from python.orm.common import get_connection_info
|
||||
from python.parallelize import parallelize_process
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Iterator
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
app = typer.Typer(help="Ingest JSONL post files into the partitioned posts table.")
|
||||
|
||||
|
||||
@app.command()
|
||||
def main(
|
||||
path: Annotated[Path, typer.Argument(help="Directory containing JSONL files, or a single JSONL file")],
|
||||
batch_size: Annotated[int, typer.Option(help="Rows per INSERT batch")] = 10000,
|
||||
workers: Annotated[int, typer.Option(help="Parallel workers for multi-file ingestion")] = 4,
|
||||
pattern: Annotated[str, typer.Option(help="Glob pattern for JSONL files")] = "*.jsonl",
|
||||
) -> None:
|
||||
"""Ingest JSONL post files into the weekly-partitioned posts table."""
|
||||
configure_logger(level="INFO")
|
||||
|
||||
logger.info("starting ingest-posts")
|
||||
logger.info("path=%s batch_size=%d workers=%d pattern=%s", path, batch_size, workers, pattern)
|
||||
if path.is_file():
|
||||
ingest_file(path, batch_size=batch_size)
|
||||
elif path.is_dir():
|
||||
ingest_directory(path, batch_size=batch_size, max_workers=workers, pattern=pattern)
|
||||
else:
|
||||
typer.echo(f"Path does not exist: {path}", err=True)
|
||||
raise typer.Exit(code=1)
|
||||
|
||||
logger.info("ingest-posts done")
|
||||
|
||||
|
||||
def ingest_directory(
|
||||
directory: Path,
|
||||
*,
|
||||
batch_size: int,
|
||||
max_workers: int,
|
||||
pattern: str = "*.jsonl",
|
||||
) -> None:
|
||||
"""Ingest all JSONL files in a directory using parallel workers."""
|
||||
files = sorted(directory.glob(pattern))
|
||||
if not files:
|
||||
logger.warning("No JSONL files found in %s", directory)
|
||||
return
|
||||
|
||||
logger.info("Found %d JSONL files to ingest", len(files))
|
||||
|
||||
kwargs_list = [{"path": fp, "batch_size": batch_size} for fp in files]
|
||||
parallelize_process(ingest_file, kwargs_list, max_workers=max_workers)
|
||||
|
||||
|
||||
SCHEMA = "main"
|
||||
|
||||
COLUMNS = (
|
||||
"post_id",
|
||||
"user_id",
|
||||
"instance",
|
||||
"date",
|
||||
"text",
|
||||
"langs",
|
||||
"like_count",
|
||||
"reply_count",
|
||||
"repost_count",
|
||||
"reply_to",
|
||||
"replied_author",
|
||||
"thread_root",
|
||||
"thread_root_author",
|
||||
"repost_from",
|
||||
"reposted_author",
|
||||
"quotes",
|
||||
"quoted_author",
|
||||
"labels",
|
||||
"sent_label",
|
||||
"sent_score",
|
||||
)
|
||||
|
||||
INSERT_FROM_STAGING = f"""
|
||||
INSERT INTO {SCHEMA}.posts ({", ".join(COLUMNS)})
|
||||
SELECT {", ".join(COLUMNS)} FROM pg_temp.staging
|
||||
ON CONFLICT (post_id, date) DO NOTHING
|
||||
""" # noqa: S608
|
||||
|
||||
FAILED_INSERT = f"""
|
||||
INSERT INTO {SCHEMA}.failed_ingestion (raw_line, error)
|
||||
VALUES (%(raw_line)s, %(error)s)
|
||||
""" # noqa: S608
|
||||
|
||||
|
||||
def get_psycopg_connection() -> psycopg.Connection:
|
||||
"""Create a raw psycopg3 connection from environment variables."""
|
||||
database, host, port, username, password = get_connection_info("DATA_SCIENCE_DEV")
|
||||
return psycopg.connect(
|
||||
dbname=database,
|
||||
host=host,
|
||||
port=int(port),
|
||||
user=username,
|
||||
password=password,
|
||||
autocommit=False,
|
||||
)
|
||||
|
||||
|
||||
def ingest_file(path: Path, *, batch_size: int) -> None:
|
||||
"""Ingest a single JSONL file into the posts table."""
|
||||
log_trigger = max(100_000 // batch_size, 1)
|
||||
failed_lines: list[dict] = []
|
||||
try:
|
||||
with get_psycopg_connection() as connection:
|
||||
for index, batch in enumerate(read_jsonl_batches(path, batch_size, failed_lines), 1):
|
||||
ingest_batch(connection, batch)
|
||||
if index % log_trigger == 0:
|
||||
logger.info("Ingested %d batches (%d rows) from %s", index, index * batch_size, path)
|
||||
|
||||
if failed_lines:
|
||||
logger.warning("Recording %d malformed lines from %s", len(failed_lines), path.name)
|
||||
with connection.cursor() as cursor:
|
||||
cursor.executemany(FAILED_INSERT, failed_lines)
|
||||
connection.commit()
|
||||
except Exception:
|
||||
logger.exception("Failed to ingest file: %s", path)
|
||||
raise
|
||||
|
||||
|
||||
def ingest_batch(connection: psycopg.Connection, batch: list[dict]) -> None:
|
||||
"""COPY batch into a temp staging table, then INSERT ... ON CONFLICT into posts."""
|
||||
if not batch:
|
||||
return
|
||||
|
||||
try:
|
||||
with connection.cursor() as cursor:
|
||||
cursor.execute(f"""
|
||||
CREATE TEMP TABLE IF NOT EXISTS staging
|
||||
(LIKE {SCHEMA}.posts INCLUDING DEFAULTS)
|
||||
ON COMMIT DELETE ROWS
|
||||
""")
|
||||
cursor.execute("TRUNCATE pg_temp.staging")
|
||||
|
||||
with cursor.copy(f"COPY pg_temp.staging ({', '.join(COLUMNS)}) FROM STDIN") as copy:
|
||||
for row in batch:
|
||||
copy.write_row(tuple(row.get(column) for column in COLUMNS))
|
||||
|
||||
cursor.execute(INSERT_FROM_STAGING)
|
||||
connection.commit()
|
||||
except Exception as error:
|
||||
connection.rollback()
|
||||
|
||||
if len(batch) == 1:
|
||||
logger.exception("Skipping bad row post_id=%s", batch[0].get("post_id"))
|
||||
with connection.cursor() as cursor:
|
||||
cursor.execute(
|
||||
FAILED_INSERT,
|
||||
{
|
||||
"raw_line": orjson.dumps(batch[0], default=str).decode(),
|
||||
"error": str(error),
|
||||
},
|
||||
)
|
||||
connection.commit()
|
||||
return
|
||||
|
||||
midpoint = len(batch) // 2
|
||||
ingest_batch(connection, batch[:midpoint])
|
||||
ingest_batch(connection, batch[midpoint:])
|
||||
|
||||
|
||||
def read_jsonl_batches(file_path: Path, batch_size: int, failed_lines: list[dict]) -> Iterator[list[dict]]:
|
||||
"""Stream a JSONL file and yield batches of transformed rows."""
|
||||
batch: list[dict] = []
|
||||
with file_path.open("r", encoding="utf-8") as handle:
|
||||
for raw_line in handle:
|
||||
line = raw_line.strip()
|
||||
if not line:
|
||||
continue
|
||||
batch.extend(parse_line(line, file_path, failed_lines))
|
||||
if len(batch) >= batch_size:
|
||||
yield batch
|
||||
batch = []
|
||||
if batch:
|
||||
yield batch
|
||||
|
||||
|
||||
def parse_line(line: str, file_path: Path, failed_lines: list[dict]) -> Iterator[dict]:
|
||||
"""Parse a JSONL line, handling concatenated JSON objects."""
|
||||
try:
|
||||
yield transform_row(orjson.loads(line))
|
||||
except orjson.JSONDecodeError:
|
||||
if "}{" not in line:
|
||||
logger.warning("Skipping malformed line in %s: %s", file_path.name, line[:120])
|
||||
failed_lines.append({"raw_line": line, "error": "malformed JSON"})
|
||||
return
|
||||
fragments = line.replace("}{", "}\n{").split("\n")
|
||||
for fragment in fragments:
|
||||
try:
|
||||
yield transform_row(orjson.loads(fragment))
|
||||
except (orjson.JSONDecodeError, KeyError, ValueError) as error:
|
||||
logger.warning("Skipping malformed fragment in %s: %s", file_path.name, fragment[:120])
|
||||
failed_lines.append({"raw_line": fragment, "error": str(error)})
|
||||
except Exception as error:
|
||||
logger.exception("Skipping bad row in %s: %s", file_path.name, line[:120])
|
||||
failed_lines.append({"raw_line": line, "error": str(error)})
|
||||
|
||||
|
||||
def transform_row(raw: dict) -> dict:
|
||||
"""Transform a raw JSONL row into a dict matching the Posts table columns."""
|
||||
raw["date"] = parse_date(raw["date"])
|
||||
if raw.get("langs") is not None:
|
||||
raw["langs"] = orjson.dumps(raw["langs"])
|
||||
if raw.get("text") is not None:
|
||||
raw["text"] = raw["text"].replace("\x00", "")
|
||||
return raw
|
||||
|
||||
|
||||
def parse_date(raw_date: int) -> datetime:
|
||||
"""Parse compact YYYYMMDDHHmm integer into a naive datetime (input is UTC by spec)."""
|
||||
return datetime(
|
||||
raw_date // 100000000,
|
||||
(raw_date // 1000000) % 100,
|
||||
(raw_date // 10000) % 100,
|
||||
(raw_date // 100) % 100,
|
||||
raw_date % 100,
|
||||
tzinfo=UTC,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
app()
|
||||
@@ -1,129 +0,0 @@
|
||||
"""CLI wrapper around alembic for multi-database support.
|
||||
|
||||
Usage:
|
||||
database <db_name> <command> [args...]
|
||||
|
||||
Examples:
|
||||
database van_inventory upgrade head
|
||||
database van_inventory downgrade head-1
|
||||
database van_inventory revision --autogenerate -m "add meals table"
|
||||
database van_inventory check
|
||||
database richie check
|
||||
database richie upgrade head
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from importlib import import_module
|
||||
from typing import TYPE_CHECKING, Annotated
|
||||
|
||||
import typer
|
||||
from alembic.config import CommandLine, Config
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from sqlalchemy.orm import DeclarativeBase
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class DatabaseConfig:
|
||||
"""Configuration for a database."""
|
||||
|
||||
env_prefix: str
|
||||
version_location: str
|
||||
base_module: str
|
||||
base_class_name: str
|
||||
models_module: str
|
||||
script_location: str = "python/alembic"
|
||||
file_template: str = "%%(year)d_%%(month).2d_%%(day).2d-%%(slug)s_%%(rev)s"
|
||||
|
||||
def get_base(self) -> type[DeclarativeBase]:
|
||||
"""Import and return the Base class."""
|
||||
module = import_module(self.base_module)
|
||||
return getattr(module, self.base_class_name)
|
||||
|
||||
def import_models(self) -> None:
|
||||
"""Import ORM models so alembic autogenerate can detect them."""
|
||||
import_module(self.models_module)
|
||||
|
||||
def alembic_config(self) -> Config:
|
||||
"""Build an alembic Config for this database."""
|
||||
# Runtime import needed — Config is in TYPE_CHECKING for the return type annotation
|
||||
from alembic.config import Config as AlembicConfig # noqa: PLC0415
|
||||
|
||||
cfg = AlembicConfig()
|
||||
cfg.set_main_option("script_location", self.script_location)
|
||||
cfg.set_main_option("file_template", self.file_template)
|
||||
cfg.set_main_option("prepend_sys_path", ".")
|
||||
cfg.set_main_option("version_path_separator", "os")
|
||||
cfg.set_main_option("version_locations", self.version_location)
|
||||
cfg.set_main_option("revision_environment", "true")
|
||||
cfg.set_section_option("post_write_hooks", "hooks", "dynamic_schema,import_postgresql,ruff")
|
||||
cfg.set_section_option("post_write_hooks", "dynamic_schema.type", "dynamic_schema")
|
||||
cfg.set_section_option("post_write_hooks", "import_postgresql.type", "import_postgresql")
|
||||
cfg.set_section_option("post_write_hooks", "ruff.type", "ruff")
|
||||
cfg.attributes["base"] = self.get_base()
|
||||
cfg.attributes["env_prefix"] = self.env_prefix
|
||||
self.import_models()
|
||||
return cfg
|
||||
|
||||
|
||||
DATABASES: dict[str, DatabaseConfig] = {
|
||||
"richie": DatabaseConfig(
|
||||
env_prefix="RICHIE",
|
||||
version_location="python/alembic/richie/versions",
|
||||
base_module="python.orm.richie.base",
|
||||
base_class_name="RichieBase",
|
||||
models_module="python.orm.richie",
|
||||
),
|
||||
"van_inventory": DatabaseConfig(
|
||||
env_prefix="VAN_INVENTORY",
|
||||
version_location="python/alembic/van_inventory/versions",
|
||||
base_module="python.orm.van_inventory.base",
|
||||
base_class_name="VanInventoryBase",
|
||||
models_module="python.orm.van_inventory.models",
|
||||
),
|
||||
"signal_bot": DatabaseConfig(
|
||||
env_prefix="SIGNALBOT",
|
||||
version_location="python/alembic/signal_bot/versions",
|
||||
base_module="python.orm.signal_bot.base",
|
||||
base_class_name="SignalBotBase",
|
||||
models_module="python.orm.signal_bot.models",
|
||||
),
|
||||
"data_science_dev": DatabaseConfig(
|
||||
env_prefix="DATA_SCIENCE_DEV",
|
||||
version_location="python/alembic/data_science_dev/versions",
|
||||
base_module="python.orm.data_science_dev.base",
|
||||
base_class_name="DataScienceDevBase",
|
||||
models_module="python.orm.data_science_dev.models",
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
app = typer.Typer(help="Multi-database alembic wrapper.")
|
||||
|
||||
|
||||
@app.command(
|
||||
context_settings={"allow_extra_args": True, "ignore_unknown_options": True},
|
||||
)
|
||||
def main(
|
||||
ctx: typer.Context,
|
||||
db_name: Annotated[str, typer.Argument(help=f"Database name. Options: {', '.join(DATABASES)}")],
|
||||
command: Annotated[str, typer.Argument(help="Alembic command (upgrade, downgrade, revision, check, etc.)")],
|
||||
) -> None:
|
||||
"""Run an alembic command against the specified database."""
|
||||
db_config = DATABASES.get(db_name)
|
||||
if not db_config:
|
||||
typer.echo(f"Unknown database: {db_name!r}. Available: {', '.join(DATABASES)}", err=True)
|
||||
raise typer.Exit(code=1)
|
||||
|
||||
alembic_cfg = db_config.alembic_config()
|
||||
|
||||
cmd_line = CommandLine()
|
||||
options = cmd_line.parser.parse_args([command, *ctx.args])
|
||||
cmd_line.run_cmd(alembic_cfg, options)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
app()
|
||||
|
||||
@@ -1 +0,0 @@
|
||||
"""Detect Nix evaluation warnings from build logs and create PRs with LLM-suggested fixes."""
|
||||
@@ -1,449 +0,0 @@
|
||||
"""Detect Nix evaluation warnings and create PRs with LLM-suggested fixes."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import logging
|
||||
import re
|
||||
import subprocess
|
||||
from dataclasses import dataclass
|
||||
from io import BytesIO
|
||||
from pathlib import Path
|
||||
from typing import Annotated
|
||||
from zipfile import ZipFile
|
||||
|
||||
import typer
|
||||
from httpx import HTTPError, post
|
||||
|
||||
from python.common import configure_logger
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class EvalWarning:
|
||||
"""A single Nix evaluation warning."""
|
||||
|
||||
system: str
|
||||
message: str
|
||||
|
||||
|
||||
@dataclass
|
||||
class FileChange:
|
||||
"""A file change suggested by the LLM."""
|
||||
|
||||
file_path: str
|
||||
original: str
|
||||
fixed: str
|
||||
|
||||
|
||||
def run_cmd(cmd: list[str], *, check: bool = True) -> subprocess.CompletedProcess[str]:
|
||||
"""Run a subprocess command and return the result.
|
||||
|
||||
Args:
|
||||
cmd: Command and arguments.
|
||||
check: Whether to raise on non-zero exit.
|
||||
|
||||
Returns:
|
||||
CompletedProcess with captured stdout/stderr.
|
||||
"""
|
||||
logger.debug("Running: %s", " ".join(cmd))
|
||||
return subprocess.run(cmd, capture_output=True, text=True, check=check)
|
||||
|
||||
|
||||
def download_logs(run_id: str, repo: str) -> dict[str, str]:
|
||||
"""Download build logs for a GitHub Actions run.
|
||||
|
||||
Args:
|
||||
run_id: The workflow run ID.
|
||||
repo: The GitHub repository (owner/repo).
|
||||
|
||||
Returns:
|
||||
Dict mapping zip entry names to their text content, filtered to build log files.
|
||||
|
||||
Raises:
|
||||
RuntimeError: If log download fails.
|
||||
"""
|
||||
result = subprocess.run(
|
||||
["gh", "api", f"repos/{repo}/actions/runs/{run_id}/logs"],
|
||||
capture_output=True,
|
||||
check=False,
|
||||
)
|
||||
if result.returncode != 0:
|
||||
msg = f"Failed to download logs: {result.stderr.decode(errors='replace')}"
|
||||
raise RuntimeError(msg)
|
||||
|
||||
logs: dict[str, str] = {}
|
||||
with ZipFile(BytesIO(result.stdout)) as zip_file:
|
||||
for name in zip_file.namelist():
|
||||
if name.startswith("build-") and name.endswith(".txt"):
|
||||
logs[name] = zip_file.read(name).decode(errors="replace")
|
||||
|
||||
return logs
|
||||
|
||||
|
||||
def parse_warnings(logs: dict[str, str]) -> set[EvalWarning]:
|
||||
"""Parse Nix evaluation warnings from build log contents.
|
||||
|
||||
Args:
|
||||
logs: Dict mapping zip entry names (e.g. "build-bob/2_Build.txt") to their text.
|
||||
|
||||
Returns:
|
||||
Deduplicated set of warnings.
|
||||
"""
|
||||
warnings: set[EvalWarning] = set()
|
||||
warning_pattern = re.compile(r"(?:^[\d\-T:.Z]+ )?(warning:|trace: warning:)")
|
||||
timestamp_prefix = re.compile(r"^[\d\-T:.Z]+ ")
|
||||
|
||||
for name, content in sorted(logs.items()):
|
||||
system = name.split("/")[0].removeprefix("build-")
|
||||
for line in content.splitlines():
|
||||
if warning_pattern.search(line):
|
||||
message = timestamp_prefix.sub("", line).strip()
|
||||
if message.startswith("warning: ignoring untrusted flake configuration setting"):
|
||||
continue
|
||||
logger.debug(f"Found warning: {line}")
|
||||
warnings.add(EvalWarning(system=system, message=message))
|
||||
|
||||
logger.info("Found %d unique warnings", len(warnings))
|
||||
return warnings
|
||||
|
||||
|
||||
def extract_referenced_files(warnings: set[EvalWarning]) -> dict[str, str]:
|
||||
"""Extract file paths referenced in warnings and read their contents.
|
||||
|
||||
Args:
|
||||
warnings: List of parsed warnings.
|
||||
|
||||
Returns:
|
||||
Dict mapping repo-relative file paths to their contents.
|
||||
"""
|
||||
paths: set[str] = set()
|
||||
warning_text = "\n".join(w.message for w in warnings)
|
||||
|
||||
nix_store_path = re.compile(r"/nix/store/[^/]+-source/([^:]+\.nix)")
|
||||
for match in nix_store_path.finditer(warning_text):
|
||||
paths.add(match.group(1))
|
||||
|
||||
repo_relative_path = re.compile(r"(?<![/\w])(systems|common|users|overlays)/[^:\s]+\.nix")
|
||||
for match in repo_relative_path.finditer(warning_text):
|
||||
paths.add(match.group(0))
|
||||
|
||||
files: dict[str, str] = {}
|
||||
for path_str in sorted(paths):
|
||||
path = Path(path_str)
|
||||
if path.is_file():
|
||||
files[path_str] = path.read_text()
|
||||
|
||||
if not files and Path("flake.nix").is_file():
|
||||
files["flake.nix"] = Path("flake.nix").read_text()
|
||||
|
||||
logger.info("Extracted %d referenced files", len(files))
|
||||
return files
|
||||
|
||||
|
||||
def compute_warning_hash(warnings: set[EvalWarning]) -> str:
|
||||
"""Compute a short hash of the warning set for deduplication.
|
||||
|
||||
Args:
|
||||
warnings: List of warnings.
|
||||
|
||||
Returns:
|
||||
8-character hex hash.
|
||||
"""
|
||||
text = "\n".join(sorted(f"[{w.system}] {w.message}" for w in warnings))
|
||||
return hashlib.sha256(text.encode()).hexdigest()[:8]
|
||||
|
||||
|
||||
def check_duplicate_pr(warning_hash: str) -> bool:
|
||||
"""Check if an open PR already exists for this warning hash.
|
||||
|
||||
Args:
|
||||
warning_hash: The hash to check.
|
||||
|
||||
Returns:
|
||||
True if a duplicate PR exists.
|
||||
|
||||
Raises:
|
||||
RuntimeError: If the gh CLI call fails.
|
||||
"""
|
||||
result = run_cmd(
|
||||
[
|
||||
"gh",
|
||||
"pr",
|
||||
"list",
|
||||
"--state",
|
||||
"open",
|
||||
"--label",
|
||||
"eval-warning-fix",
|
||||
"--json",
|
||||
"title",
|
||||
"--jq",
|
||||
".[].title",
|
||||
],
|
||||
check=False,
|
||||
)
|
||||
if result.returncode != 0:
|
||||
msg = f"Failed to check for duplicate PRs: {result.stderr}"
|
||||
raise RuntimeError(msg)
|
||||
|
||||
for title in result.stdout.splitlines():
|
||||
if warning_hash in title:
|
||||
logger.info("Duplicate PR found for hash %s", warning_hash)
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def query_ollama(
|
||||
warnings: set[EvalWarning],
|
||||
files: dict[str, str],
|
||||
ollama_url: str,
|
||||
) -> str | None:
|
||||
"""Query Ollama for a fix suggestion.
|
||||
|
||||
Args:
|
||||
warnings: List of warnings.
|
||||
files: Referenced file contents.
|
||||
ollama_url: Ollama API base URL.
|
||||
|
||||
Returns:
|
||||
LLM response text, or None on failure.
|
||||
"""
|
||||
warning_text = "\n".join(f"[{w.system}] {w.message}" for w in warnings)
|
||||
file_context = "\n".join(f"--- FILE: {path} ---\n{content}\n--- END FILE ---" for path, content in files.items())
|
||||
|
||||
prompt = f"""You are a NixOS configuration expert. \
|
||||
Analyze the following Nix evaluation warnings and suggest fixes.
|
||||
|
||||
## Warnings
|
||||
{warning_text}
|
||||
|
||||
## Referenced Files
|
||||
{file_context}
|
||||
|
||||
## Instructions
|
||||
- Identify the root cause of each warning
|
||||
- Provide the exact file changes needed to fix the warnings
|
||||
- Output your response in two clearly separated sections:
|
||||
1. **REASONING**: Brief explanation of what causes each warning and how to fix it
|
||||
2. **CHANGES**: For each file that needs changes, output a block like:
|
||||
FILE: path/to/file.nix
|
||||
<<<<<<< ORIGINAL
|
||||
the original lines to replace
|
||||
=======
|
||||
the replacement lines
|
||||
>>>>>>> FIXED
|
||||
- Only suggest changes for files that exist in the repository
|
||||
- Do not add unnecessary complexity
|
||||
- Preserve the existing code style
|
||||
- If a warning comes from upstream nixpkgs and cannot be fixed in this repo, \
|
||||
say so in REASONING and do not suggest changes"""
|
||||
|
||||
try:
|
||||
response = post(
|
||||
f"{ollama_url}/api/generate",
|
||||
json={
|
||||
"model": "qwen3-coder:30b",
|
||||
"prompt": prompt,
|
||||
"stream": False,
|
||||
"options": {"num_predict": 4096},
|
||||
},
|
||||
timeout=300,
|
||||
)
|
||||
response.raise_for_status()
|
||||
except HTTPError:
|
||||
logger.exception("Ollama request failed")
|
||||
return None
|
||||
|
||||
return response.json().get("response")
|
||||
|
||||
|
||||
def parse_changes(response: str) -> list[FileChange]:
|
||||
"""Parse file changes from the **CHANGES** section of the LLM response.
|
||||
|
||||
Expects blocks in the format:
|
||||
FILE: path/to/file.nix
|
||||
<<<<<<< ORIGINAL
|
||||
...
|
||||
=======
|
||||
...
|
||||
>>>>>>> FIXED
|
||||
|
||||
Args:
|
||||
response: Raw LLM response text.
|
||||
|
||||
Returns:
|
||||
List of parsed file changes.
|
||||
"""
|
||||
if "**CHANGES**" not in response:
|
||||
logger.warning("LLM response missing **CHANGES** section")
|
||||
return []
|
||||
|
||||
changes_section = response.split("**CHANGES**", 1)[1]
|
||||
|
||||
changes: list[FileChange] = []
|
||||
current_file = ""
|
||||
section: str | None = None
|
||||
original_lines: list[str] = []
|
||||
fixed_lines: list[str] = []
|
||||
|
||||
for line in changes_section.splitlines():
|
||||
stripped = line.strip()
|
||||
if stripped.startswith("FILE:"):
|
||||
current_file = stripped.removeprefix("FILE:").strip()
|
||||
elif stripped == "<<<<<<< ORIGINAL":
|
||||
section = "original"
|
||||
original_lines = []
|
||||
elif stripped == "=======" and section == "original":
|
||||
section = "fixed"
|
||||
fixed_lines = []
|
||||
elif stripped == ">>>>>>> FIXED" and section == "fixed":
|
||||
section = None
|
||||
if current_file:
|
||||
changes.append(FileChange(current_file, "\n".join(original_lines), "\n".join(fixed_lines)))
|
||||
elif section == "original":
|
||||
original_lines.append(line)
|
||||
elif section == "fixed":
|
||||
fixed_lines.append(line)
|
||||
|
||||
logger.info("Parsed %d file changes", len(changes))
|
||||
return changes
|
||||
|
||||
|
||||
def apply_changes(changes: list[FileChange]) -> int:
|
||||
"""Apply file changes to the working directory.
|
||||
|
||||
Args:
|
||||
changes: List of changes to apply.
|
||||
|
||||
Returns:
|
||||
Number of changes successfully applied.
|
||||
"""
|
||||
applied = 0
|
||||
cwd = Path.cwd().resolve()
|
||||
for change in changes:
|
||||
path = Path(change.file_path).resolve()
|
||||
if not path.is_relative_to(cwd):
|
||||
logger.warning("Path traversal blocked: %s", change.file_path)
|
||||
continue
|
||||
if not path.is_file():
|
||||
logger.warning("File not found: %s", change.file_path)
|
||||
continue
|
||||
|
||||
content = path.read_text()
|
||||
if change.original not in content:
|
||||
logger.warning("Original text not found in %s", change.file_path)
|
||||
continue
|
||||
|
||||
path.write_text(content.replace(change.original, change.fixed, 1))
|
||||
logger.info("Applied fix to %s", change.file_path)
|
||||
applied += 1
|
||||
|
||||
return applied
|
||||
|
||||
|
||||
def create_pr(
|
||||
warning_hash: str,
|
||||
warnings: set[EvalWarning],
|
||||
llm_response: str,
|
||||
run_url: str,
|
||||
) -> None:
|
||||
"""Create a git branch and PR with the applied fixes.
|
||||
|
||||
Args:
|
||||
warning_hash: Short hash for branch naming and deduplication.
|
||||
warnings: Original warnings for the PR body.
|
||||
llm_response: Full LLM response for extracting reasoning.
|
||||
run_url: URL to the triggering build run.
|
||||
"""
|
||||
branch = f"fix/eval-warning-{warning_hash}"
|
||||
warning_text = "\n".join(f"[{w.system}] {w.message}" for w in warnings)
|
||||
|
||||
if "**REASONING**" not in llm_response:
|
||||
logger.warning("LLM response missing **REASONING** section")
|
||||
reasoning = ""
|
||||
else:
|
||||
_, after = llm_response.split("**REASONING**", 1)
|
||||
reasoning = "\n".join(after.split("**CHANGES**", 1)[0].strip().splitlines()[:50])
|
||||
|
||||
run_cmd(["git", "config", "user.name", "github-actions[bot]"])
|
||||
run_cmd(["git", "config", "user.email", "github-actions[bot]@users.noreply.github.com"])
|
||||
run_cmd(["git", "checkout", "-b", branch])
|
||||
run_cmd(["git", "add", "-A"])
|
||||
|
||||
diff_result = run_cmd(["git", "diff", "--cached", "--quiet"], check=False)
|
||||
if diff_result.returncode == 0:
|
||||
logger.info("No file changes to commit")
|
||||
return
|
||||
|
||||
run_cmd(["git", "commit", "-m", f"fix: resolve nix evaluation warnings ({warning_hash})"])
|
||||
run_cmd(["git", "push", "origin", branch, "--force"])
|
||||
|
||||
body = f"""## Nix Evaluation Warnings
|
||||
|
||||
Detected in [build_systems run]({run_url}):
|
||||
|
||||
```
|
||||
{warning_text}
|
||||
```
|
||||
|
||||
## LLM Analysis (qwen3-coder:30b)
|
||||
|
||||
{reasoning}
|
||||
|
||||
---
|
||||
*Auto-generated by fix_eval_warnings. Review carefully before merging.*"""
|
||||
|
||||
run_cmd(
|
||||
[
|
||||
"gh",
|
||||
"pr",
|
||||
"create",
|
||||
"--title",
|
||||
f"fix: resolve nix eval warnings ({warning_hash})",
|
||||
"--label",
|
||||
"automated",
|
||||
"--label",
|
||||
"eval-warning-fix",
|
||||
"--body",
|
||||
body,
|
||||
]
|
||||
)
|
||||
logger.info("PR created on branch %s", branch)
|
||||
|
||||
|
||||
def main(
|
||||
run_id: Annotated[str, typer.Option("--run-id", help="GitHub Actions run ID")],
|
||||
repo: Annotated[str, typer.Option("--repo", help="GitHub repository (owner/repo)")],
|
||||
ollama_url: Annotated[str, typer.Option("--ollama-url", help="Ollama API base URL")],
|
||||
run_url: Annotated[str, typer.Option("--run-url", help="URL to the triggering build run")],
|
||||
log_level: Annotated[str, typer.Option("--log-level", "-l", help="Log level")] = "INFO",
|
||||
) -> None:
|
||||
"""Detect Nix evaluation warnings and create PRs with LLM-suggested fixes."""
|
||||
configure_logger(log_level)
|
||||
|
||||
logs = download_logs(run_id, repo)
|
||||
warnings = parse_warnings(logs)
|
||||
if not warnings:
|
||||
return
|
||||
|
||||
warning_hash = compute_warning_hash(warnings)
|
||||
if check_duplicate_pr(warning_hash):
|
||||
return
|
||||
|
||||
files = extract_referenced_files(warnings)
|
||||
llm_response = query_ollama(warnings, files, ollama_url)
|
||||
if not llm_response:
|
||||
return
|
||||
|
||||
changes = parse_changes(llm_response)
|
||||
applied = apply_changes(changes)
|
||||
if applied == 0:
|
||||
logger.info("No changes could be applied")
|
||||
return
|
||||
|
||||
create_pr(warning_hash, warnings, llm_response, run_url)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
typer.run(main)
|
||||
@@ -1 +0,0 @@
|
||||
"""Tuya heater control service."""
|
||||
@@ -1,69 +0,0 @@
|
||||
"""TinyTuya device controller for heater."""
|
||||
|
||||
import logging
|
||||
|
||||
import tinytuya
|
||||
|
||||
from python.heater.models import ActionResult, DeviceConfig, HeaterStatus
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# DPS mapping for heater
|
||||
DPS_POWER = "1" # bool: on/off
|
||||
DPS_SETPOINT = "101" # int: target temp (read-only)
|
||||
DPS_STATE = "102" # str: "Stop", "Heat", etc.
|
||||
DPS_UNKNOWN = "104" # int: unknown
|
||||
DPS_ERROR = "108" # int: last error code
|
||||
|
||||
|
||||
class HeaterController:
|
||||
"""Controls a Tuya heater device via local network."""
|
||||
|
||||
def __init__(self, config: DeviceConfig) -> None:
|
||||
"""Initialize the controller."""
|
||||
self.device = tinytuya.Device(config.device_id, config.ip, config.local_key)
|
||||
self.device.set_version(config.version)
|
||||
self.device.set_socketTimeout(0.5)
|
||||
self.device.set_socketRetryLimit(1)
|
||||
|
||||
def status(self) -> HeaterStatus:
|
||||
"""Get current heater status."""
|
||||
data = self.device.status()
|
||||
|
||||
if "Error" in data:
|
||||
logger.error("Device error: %s", data)
|
||||
return HeaterStatus(power=False, raw_dps={"error": data["Error"]})
|
||||
|
||||
dps = data.get("dps", {})
|
||||
return HeaterStatus(
|
||||
power=bool(dps.get(DPS_POWER, False)),
|
||||
setpoint=dps.get(DPS_SETPOINT),
|
||||
state=dps.get(DPS_STATE),
|
||||
error_code=dps.get(DPS_ERROR),
|
||||
raw_dps=dps,
|
||||
)
|
||||
|
||||
def turn_on(self) -> ActionResult:
|
||||
"""Turn heater on."""
|
||||
try:
|
||||
self.device.set_value(index=DPS_POWER, value=True)
|
||||
return ActionResult(success=True, action="on", power=True)
|
||||
except Exception as error:
|
||||
logger.exception("Failed to turn on")
|
||||
return ActionResult(success=False, action="on", error=str(error))
|
||||
|
||||
def turn_off(self) -> ActionResult:
|
||||
"""Turn heater off."""
|
||||
try:
|
||||
self.device.set_value(index=DPS_POWER, value=False)
|
||||
return ActionResult(success=True, action="off", power=False)
|
||||
except Exception as error:
|
||||
logger.exception("Failed to turn off")
|
||||
return ActionResult(success=False, action="off", error=str(error))
|
||||
|
||||
def toggle(self) -> ActionResult:
|
||||
"""Toggle heater power state."""
|
||||
status = self.status()
|
||||
if status.power:
|
||||
return self.turn_off()
|
||||
return self.turn_on()
|
||||
@@ -1,85 +0,0 @@
|
||||
"""FastAPI heater control service."""
|
||||
|
||||
import logging
|
||||
from collections.abc import AsyncIterator
|
||||
from contextlib import asynccontextmanager
|
||||
from typing import Annotated
|
||||
|
||||
import typer
|
||||
import uvicorn
|
||||
from fastapi import FastAPI, HTTPException
|
||||
|
||||
from python.common import configure_logger
|
||||
from python.heater.controller import HeaterController
|
||||
from python.heater.models import ActionResult, DeviceConfig, HeaterStatus
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def create_app(config: DeviceConfig) -> FastAPI:
|
||||
"""Create FastAPI application."""
|
||||
|
||||
@asynccontextmanager
|
||||
async def lifespan(app: FastAPI) -> AsyncIterator[None]:
|
||||
app.state.controller = HeaterController(config)
|
||||
yield
|
||||
|
||||
app = FastAPI(
|
||||
title="Heater Control API",
|
||||
description="Fast local control for Tuya heater",
|
||||
lifespan=lifespan,
|
||||
)
|
||||
|
||||
@app.get("/status")
|
||||
def get_status() -> HeaterStatus:
|
||||
return app.state.controller.status()
|
||||
|
||||
@app.post("/on")
|
||||
def heater_on() -> ActionResult:
|
||||
result = app.state.controller.turn_on()
|
||||
if not result.success:
|
||||
raise HTTPException(status_code=500, detail=result.error)
|
||||
return result
|
||||
|
||||
@app.post("/off")
|
||||
def heater_off() -> ActionResult:
|
||||
result = app.state.controller.turn_off()
|
||||
if not result.success:
|
||||
raise HTTPException(status_code=500, detail=result.error)
|
||||
return result
|
||||
|
||||
@app.post("/toggle")
|
||||
def heater_toggle() -> ActionResult:
|
||||
result = app.state.controller.toggle()
|
||||
if not result.success:
|
||||
raise HTTPException(status_code=500, detail=result.error)
|
||||
return result
|
||||
|
||||
return app
|
||||
|
||||
|
||||
def serve(
|
||||
host: Annotated[str, typer.Option("--host", "-h", help="Host to bind to")],
|
||||
port: Annotated[int, typer.Option("--port", "-p", help="Port to bind to")] = 8124,
|
||||
log_level: Annotated[str, typer.Option("--log-level", "-l", help="Log level")] = "INFO",
|
||||
device_id: Annotated[str | None, typer.Option("--device-id", envvar="TUYA_DEVICE_ID")] = None,
|
||||
device_ip: Annotated[str | None, typer.Option("--device-ip", envvar="TUYA_DEVICE_IP")] = None,
|
||||
local_key: Annotated[str | None, typer.Option("--local-key", envvar="TUYA_LOCAL_KEY")] = None,
|
||||
) -> None:
|
||||
"""Start the heater control API server."""
|
||||
configure_logger(log_level)
|
||||
|
||||
logger.info("Starting heater control API server")
|
||||
|
||||
if not device_id or not device_ip or not local_key:
|
||||
error = "Must provide device ID, IP, and local key"
|
||||
raise typer.Exit(error)
|
||||
|
||||
config = DeviceConfig(device_id=device_id, ip=device_ip, local_key=local_key)
|
||||
|
||||
app = create_app(config)
|
||||
uvicorn.run(app, host=host, port=port)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
typer.run(serve)
|
||||
@@ -1,31 +0,0 @@
|
||||
"""Pydantic models for heater API."""
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
|
||||
class DeviceConfig(BaseModel):
|
||||
"""Tuya device configuration."""
|
||||
|
||||
device_id: str
|
||||
ip: str
|
||||
local_key: str
|
||||
version: float = 3.5
|
||||
|
||||
|
||||
class HeaterStatus(BaseModel):
|
||||
"""Current heater status."""
|
||||
|
||||
power: bool
|
||||
setpoint: int | None = None
|
||||
state: str | None = None # "Stop", "Heat", etc.
|
||||
error_code: int | None = None
|
||||
raw_dps: dict[str, object] = Field(default_factory=dict)
|
||||
|
||||
|
||||
class ActionResult(BaseModel):
|
||||
"""Result of a heater action."""
|
||||
|
||||
success: bool
|
||||
action: str
|
||||
power: bool | None = None
|
||||
error: str | None = None
|
||||
@@ -1,13 +0,0 @@
|
||||
"""ORM package exports."""
|
||||
|
||||
from python.orm.data_science_dev.base import DataScienceDevBase
|
||||
from python.orm.richie.base import RichieBase
|
||||
from python.orm.signal_bot.base import SignalBotBase
|
||||
from python.orm.van_inventory.base import VanInventoryBase
|
||||
|
||||
__all__ = [
|
||||
"DataScienceDevBase",
|
||||
"RichieBase",
|
||||
"SignalBotBase",
|
||||
"VanInventoryBase",
|
||||
]
|
||||
@@ -1,51 +0,0 @@
|
||||
"""Shared ORM definitions."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from os import getenv
|
||||
from typing import cast
|
||||
|
||||
from sqlalchemy import create_engine
|
||||
from sqlalchemy.engine import URL, Engine
|
||||
|
||||
NAMING_CONVENTION = {
|
||||
"ix": "ix_%(table_name)s_%(column_0_name)s",
|
||||
"uq": "uq_%(table_name)s_%(column_0_name)s",
|
||||
"ck": "ck_%(table_name)s_%(constraint_name)s",
|
||||
"fk": "fk_%(table_name)s_%(column_0_name)s_%(referred_table_name)s",
|
||||
"pk": "pk_%(table_name)s",
|
||||
}
|
||||
|
||||
|
||||
def get_connection_info(name: str) -> tuple[str, str, str, str, str | None]:
|
||||
"""Get connection info from environment variables."""
|
||||
database = getenv(f"{name}_DB")
|
||||
host = getenv(f"{name}_HOST")
|
||||
port = getenv(f"{name}_PORT")
|
||||
username = getenv(f"{name}_USER")
|
||||
password = getenv(f"{name}_PASSWORD")
|
||||
|
||||
if None in (database, host, port, username):
|
||||
error = f"Missing environment variables for Postgres connection.\n{database=}\n{host=}\n{port=}\n{username=}\n"
|
||||
raise ValueError(error)
|
||||
return cast("tuple[str, str, str, str, str | None]", (database, host, port, username, password))
|
||||
|
||||
|
||||
def get_postgres_engine(*, name: str = "POSTGRES", pool_pre_ping: bool = True) -> Engine:
|
||||
"""Create a SQLAlchemy engine from environment variables."""
|
||||
database, host, port, username, password = get_connection_info(name)
|
||||
|
||||
url = URL.create(
|
||||
drivername="postgresql+psycopg",
|
||||
username=username,
|
||||
password=password,
|
||||
host=host,
|
||||
port=int(port),
|
||||
database=database,
|
||||
)
|
||||
|
||||
return create_engine(
|
||||
url=url,
|
||||
pool_pre_ping=pool_pre_ping,
|
||||
pool_recycle=1800,
|
||||
)
|
||||
@@ -1,11 +0,0 @@
|
||||
"""Data science dev database ORM exports."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from python.orm.data_science_dev.base import DataScienceDevBase, DataScienceDevTableBase, DataScienceDevTableBaseBig
|
||||
|
||||
__all__ = [
|
||||
"DataScienceDevBase",
|
||||
"DataScienceDevTableBase",
|
||||
"DataScienceDevTableBaseBig",
|
||||
]
|
||||
@@ -1,52 +0,0 @@
|
||||
"""Data science dev database ORM base."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import datetime
|
||||
|
||||
from sqlalchemy import BigInteger, DateTime, MetaData, func
|
||||
from sqlalchemy.ext.declarative import AbstractConcreteBase
|
||||
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column
|
||||
|
||||
from python.orm.common import NAMING_CONVENTION
|
||||
|
||||
|
||||
class DataScienceDevBase(DeclarativeBase):
|
||||
"""Base class for data_science_dev database ORM models."""
|
||||
|
||||
schema_name = "main"
|
||||
|
||||
metadata = MetaData(
|
||||
schema=schema_name,
|
||||
naming_convention=NAMING_CONVENTION,
|
||||
)
|
||||
|
||||
|
||||
class _TableMixin:
|
||||
"""Shared timestamp columns for all table bases."""
|
||||
|
||||
created: Mapped[datetime] = mapped_column(
|
||||
DateTime(timezone=True),
|
||||
server_default=func.now(),
|
||||
)
|
||||
updated: Mapped[datetime] = mapped_column(
|
||||
DateTime(timezone=True),
|
||||
server_default=func.now(),
|
||||
onupdate=func.now(),
|
||||
)
|
||||
|
||||
|
||||
class DataScienceDevTableBase(_TableMixin, AbstractConcreteBase, DataScienceDevBase):
|
||||
"""Table with Integer primary key."""
|
||||
|
||||
__abstract__ = True
|
||||
|
||||
id: Mapped[int] = mapped_column(primary_key=True)
|
||||
|
||||
|
||||
class DataScienceDevTableBaseBig(_TableMixin, AbstractConcreteBase, DataScienceDevBase):
|
||||
"""Table with BigInteger primary key."""
|
||||
|
||||
__abstract__ = True
|
||||
|
||||
id: Mapped[int] = mapped_column(BigInteger, primary_key=True)
|
||||
@@ -1,14 +0,0 @@
|
||||
"""init."""
|
||||
|
||||
from python.orm.data_science_dev.congress.bill import Bill, BillText
|
||||
from python.orm.data_science_dev.congress.legislator import Legislator, LegislatorSocialMedia
|
||||
from python.orm.data_science_dev.congress.vote import Vote, VoteRecord
|
||||
|
||||
__all__ = [
|
||||
"Bill",
|
||||
"BillText",
|
||||
"Legislator",
|
||||
"LegislatorSocialMedia",
|
||||
"Vote",
|
||||
"VoteRecord",
|
||||
]
|
||||
@@ -1,66 +0,0 @@
|
||||
"""Bill model - legislation introduced in Congress."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import date
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from sqlalchemy import ForeignKey, Index, UniqueConstraint
|
||||
from sqlalchemy.orm import Mapped, mapped_column, relationship
|
||||
|
||||
from python.orm.data_science_dev.base import DataScienceDevTableBase
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from python.orm.data_science_dev.congress.vote import Vote
|
||||
|
||||
|
||||
class Bill(DataScienceDevTableBase):
|
||||
"""Legislation with congress number, type, titles, status, and sponsor."""
|
||||
|
||||
__tablename__ = "bill"
|
||||
|
||||
congress: Mapped[int]
|
||||
bill_type: Mapped[str]
|
||||
number: Mapped[int]
|
||||
|
||||
title: Mapped[str | None]
|
||||
title_short: Mapped[str | None]
|
||||
official_title: Mapped[str | None]
|
||||
|
||||
status: Mapped[str | None]
|
||||
status_at: Mapped[date | None]
|
||||
|
||||
sponsor_bioguide_id: Mapped[str | None]
|
||||
|
||||
subjects_top_term: Mapped[str | None]
|
||||
|
||||
votes: Mapped[list[Vote]] = relationship(
|
||||
"Vote",
|
||||
back_populates="bill",
|
||||
)
|
||||
bill_texts: Mapped[list[BillText]] = relationship(
|
||||
"BillText",
|
||||
back_populates="bill",
|
||||
cascade="all, delete-orphan",
|
||||
)
|
||||
|
||||
__table_args__ = (
|
||||
UniqueConstraint("congress", "bill_type", "number", name="uq_bill_congress_type_number"),
|
||||
Index("ix_bill_congress", "congress"),
|
||||
)
|
||||
|
||||
|
||||
class BillText(DataScienceDevTableBase):
|
||||
"""Stores different text versions of a bill (introduced, enrolled, etc.)."""
|
||||
|
||||
__tablename__ = "bill_text"
|
||||
|
||||
bill_id: Mapped[int] = mapped_column(ForeignKey("main.bill.id", ondelete="CASCADE"))
|
||||
version_code: Mapped[str]
|
||||
version_name: Mapped[str | None]
|
||||
text_content: Mapped[str | None]
|
||||
date: Mapped[date | None]
|
||||
|
||||
bill: Mapped[Bill] = relationship("Bill", back_populates="bill_texts")
|
||||
|
||||
__table_args__ = (UniqueConstraint("bill_id", "version_code", name="uq_bill_text_bill_id_version_code"),)
|
||||
@@ -1,66 +0,0 @@
|
||||
"""Legislator model - members of Congress."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import date
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from sqlalchemy import ForeignKey, Text
|
||||
from sqlalchemy.orm import Mapped, mapped_column, relationship
|
||||
|
||||
from python.orm.data_science_dev.base import DataScienceDevTableBase
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from python.orm.data_science_dev.congress.vote import VoteRecord
|
||||
|
||||
|
||||
class Legislator(DataScienceDevTableBase):
|
||||
"""Members of Congress with identification and current term info."""
|
||||
|
||||
__tablename__ = "legislator"
|
||||
|
||||
bioguide_id: Mapped[str] = mapped_column(Text, unique=True, index=True)
|
||||
|
||||
thomas_id: Mapped[str | None]
|
||||
lis_id: Mapped[str | None]
|
||||
govtrack_id: Mapped[int | None]
|
||||
opensecrets_id: Mapped[str | None]
|
||||
fec_ids: Mapped[str | None]
|
||||
|
||||
first_name: Mapped[str]
|
||||
last_name: Mapped[str]
|
||||
official_full_name: Mapped[str | None]
|
||||
nickname: Mapped[str | None]
|
||||
|
||||
birthday: Mapped[date | None]
|
||||
gender: Mapped[str | None]
|
||||
|
||||
current_party: Mapped[str | None]
|
||||
current_state: Mapped[str | None]
|
||||
current_district: Mapped[int | None]
|
||||
current_chamber: Mapped[str | None]
|
||||
|
||||
social_media_accounts: Mapped[list[LegislatorSocialMedia]] = relationship(
|
||||
"LegislatorSocialMedia",
|
||||
back_populates="legislator",
|
||||
cascade="all, delete-orphan",
|
||||
)
|
||||
vote_records: Mapped[list[VoteRecord]] = relationship(
|
||||
"VoteRecord",
|
||||
back_populates="legislator",
|
||||
cascade="all, delete-orphan",
|
||||
)
|
||||
|
||||
|
||||
class LegislatorSocialMedia(DataScienceDevTableBase):
|
||||
"""Social media account linked to a legislator."""
|
||||
|
||||
__tablename__ = "legislator_social_media"
|
||||
|
||||
legislator_id: Mapped[int] = mapped_column(ForeignKey("main.legislator.id"))
|
||||
platform: Mapped[str]
|
||||
account_name: Mapped[str]
|
||||
url: Mapped[str | None]
|
||||
source: Mapped[str]
|
||||
|
||||
legislator: Mapped[Legislator] = relationship(back_populates="social_media_accounts")
|
||||
@@ -1,79 +0,0 @@
|
||||
"""Vote model - roll call votes in Congress."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import date
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from sqlalchemy import ForeignKey, Index, UniqueConstraint
|
||||
from sqlalchemy.orm import Mapped, mapped_column, relationship
|
||||
|
||||
from python.orm.data_science_dev.base import DataScienceDevBase, DataScienceDevTableBase
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from python.orm.data_science_dev.congress.bill import Bill
|
||||
from python.orm.data_science_dev.congress.legislator import Legislator
|
||||
from python.orm.data_science_dev.congress.vote import Vote
|
||||
|
||||
|
||||
class VoteRecord(DataScienceDevBase):
|
||||
"""Links a vote to a legislator with their position (Yea, Nay, etc.)."""
|
||||
|
||||
__tablename__ = "vote_record"
|
||||
|
||||
vote_id: Mapped[int] = mapped_column(
|
||||
ForeignKey("main.vote.id", ondelete="CASCADE"),
|
||||
primary_key=True,
|
||||
)
|
||||
legislator_id: Mapped[int] = mapped_column(
|
||||
ForeignKey("main.legislator.id", ondelete="CASCADE"),
|
||||
primary_key=True,
|
||||
)
|
||||
position: Mapped[str]
|
||||
|
||||
vote: Mapped[Vote] = relationship("Vote", back_populates="vote_records")
|
||||
legislator: Mapped[Legislator] = relationship("Legislator", back_populates="vote_records")
|
||||
|
||||
|
||||
class Vote(DataScienceDevTableBase):
|
||||
"""Roll call votes with counts and optional bill linkage."""
|
||||
|
||||
__tablename__ = "vote"
|
||||
|
||||
congress: Mapped[int]
|
||||
chamber: Mapped[str]
|
||||
session: Mapped[int]
|
||||
number: Mapped[int]
|
||||
|
||||
vote_type: Mapped[str | None]
|
||||
question: Mapped[str | None]
|
||||
result: Mapped[str | None]
|
||||
result_text: Mapped[str | None]
|
||||
|
||||
vote_date: Mapped[date]
|
||||
|
||||
yea_count: Mapped[int | None]
|
||||
nay_count: Mapped[int | None]
|
||||
not_voting_count: Mapped[int | None]
|
||||
present_count: Mapped[int | None]
|
||||
|
||||
bill_id: Mapped[int | None] = mapped_column(ForeignKey("main.bill.id"))
|
||||
|
||||
bill: Mapped[Bill | None] = relationship("Bill", back_populates="votes")
|
||||
vote_records: Mapped[list[VoteRecord]] = relationship(
|
||||
"VoteRecord",
|
||||
back_populates="vote",
|
||||
cascade="all, delete-orphan",
|
||||
)
|
||||
|
||||
__table_args__ = (
|
||||
UniqueConstraint(
|
||||
"congress",
|
||||
"chamber",
|
||||
"session",
|
||||
"number",
|
||||
name="uq_vote_congress_chamber_session_number",
|
||||
),
|
||||
Index("ix_vote_date", "vote_date"),
|
||||
Index("ix_vote_congress_chamber", "congress", "chamber"),
|
||||
)
|
||||
@@ -1,16 +0,0 @@
|
||||
"""Data science dev database ORM models."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from python.orm.data_science_dev.congress import Bill, BillText, Legislator, Vote, VoteRecord
|
||||
from python.orm.data_science_dev.posts import partitions # noqa: F401 — registers partition classes in metadata
|
||||
from python.orm.data_science_dev.posts.tables import Posts
|
||||
|
||||
__all__ = [
|
||||
"Bill",
|
||||
"BillText",
|
||||
"Legislator",
|
||||
"Posts",
|
||||
"Vote",
|
||||
"VoteRecord",
|
||||
]
|
||||
@@ -1,11 +0,0 @@
|
||||
"""Posts module — weekly-partitioned posts table and partition ORM models."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from python.orm.data_science_dev.posts.failed_ingestion import FailedIngestion
|
||||
from python.orm.data_science_dev.posts.tables import Posts
|
||||
|
||||
__all__ = [
|
||||
"FailedIngestion",
|
||||
"Posts",
|
||||
]
|
||||
@@ -1,33 +0,0 @@
|
||||
"""Shared column definitions for the posts partitioned table family."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import datetime
|
||||
|
||||
from sqlalchemy import BigInteger, SmallInteger, Text
|
||||
from sqlalchemy.orm import Mapped, mapped_column
|
||||
|
||||
|
||||
class PostsColumns:
|
||||
"""Mixin providing all posts columns. Used by both the parent table and partitions."""
|
||||
|
||||
post_id: Mapped[int] = mapped_column(BigInteger, primary_key=True)
|
||||
user_id: Mapped[int] = mapped_column(BigInteger)
|
||||
instance: Mapped[str]
|
||||
date: Mapped[datetime] = mapped_column(primary_key=True)
|
||||
text: Mapped[str] = mapped_column(Text)
|
||||
langs: Mapped[str | None]
|
||||
like_count: Mapped[int]
|
||||
reply_count: Mapped[int]
|
||||
repost_count: Mapped[int]
|
||||
reply_to: Mapped[int | None] = mapped_column(BigInteger)
|
||||
replied_author: Mapped[int | None] = mapped_column(BigInteger)
|
||||
thread_root: Mapped[int | None] = mapped_column(BigInteger)
|
||||
thread_root_author: Mapped[int | None] = mapped_column(BigInteger)
|
||||
repost_from: Mapped[int | None] = mapped_column(BigInteger)
|
||||
reposted_author: Mapped[int | None] = mapped_column(BigInteger)
|
||||
quotes: Mapped[int | None] = mapped_column(BigInteger)
|
||||
quoted_author: Mapped[int | None] = mapped_column(BigInteger)
|
||||
labels: Mapped[str | None]
|
||||
sent_label: Mapped[int | None] = mapped_column(SmallInteger)
|
||||
sent_score: Mapped[float | None]
|
||||
@@ -1,17 +0,0 @@
|
||||
"""Table for storing JSONL lines that failed during post ingestion."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from sqlalchemy import Text
|
||||
from sqlalchemy.orm import Mapped, mapped_column
|
||||
|
||||
from python.orm.data_science_dev.base import DataScienceDevTableBase
|
||||
|
||||
|
||||
class FailedIngestion(DataScienceDevTableBase):
|
||||
"""Stores raw JSONL lines and their error messages when ingestion fails."""
|
||||
|
||||
__tablename__ = "failed_ingestion"
|
||||
|
||||
raw_line: Mapped[str] = mapped_column(Text)
|
||||
error: Mapped[str] = mapped_column(Text)
|
||||
@@ -1,71 +0,0 @@
|
||||
"""Dynamically generated ORM classes for each weekly partition of the posts table.
|
||||
|
||||
Each class maps to a PostgreSQL partition table (e.g. posts_2024_01).
|
||||
These are real ORM models tracked by Alembic autogenerate.
|
||||
|
||||
Uses ISO week numbering (datetime.isocalendar().week). ISO years can have
|
||||
52 or 53 weeks, and week boundaries are always Monday to Monday.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
from datetime import UTC, datetime
|
||||
|
||||
from python.orm.data_science_dev.base import DataScienceDevBase
|
||||
from python.orm.data_science_dev.posts.columns import PostsColumns
|
||||
|
||||
PARTITION_START_YEAR = 2023
|
||||
PARTITION_END_YEAR = 2026
|
||||
|
||||
_current_module = sys.modules[__name__]
|
||||
|
||||
|
||||
def iso_weeks_in_year(year: int) -> int:
|
||||
"""Return the number of ISO weeks in a given year (52 or 53)."""
|
||||
dec_28 = datetime(year, 12, 28, tzinfo=UTC)
|
||||
return dec_28.isocalendar().week
|
||||
|
||||
|
||||
def week_bounds(year: int, week: int) -> tuple[datetime, datetime]:
|
||||
"""Return (start, end) datetimes for an ISO week.
|
||||
|
||||
Start = Monday 00:00:00 UTC of the given ISO week.
|
||||
End = Monday 00:00:00 UTC of the following ISO week.
|
||||
"""
|
||||
start = datetime.fromisocalendar(year, week, 1).replace(tzinfo=UTC)
|
||||
if week < iso_weeks_in_year(year):
|
||||
end = datetime.fromisocalendar(year, week + 1, 1).replace(tzinfo=UTC)
|
||||
else:
|
||||
end = datetime.fromisocalendar(year + 1, 1, 1).replace(tzinfo=UTC)
|
||||
return start, end
|
||||
|
||||
|
||||
def _build_partition_classes() -> dict[str, type]:
|
||||
"""Generate one ORM class per ISO week partition."""
|
||||
classes: dict[str, type] = {}
|
||||
|
||||
for year in range(PARTITION_START_YEAR, PARTITION_END_YEAR + 1):
|
||||
for week in range(1, iso_weeks_in_year(year) + 1):
|
||||
class_name = f"PostsWeek{year}W{week:02d}"
|
||||
table_name = f"posts_{year}_{week:02d}"
|
||||
|
||||
partition_class = type(
|
||||
class_name,
|
||||
(PostsColumns, DataScienceDevBase),
|
||||
{
|
||||
"__tablename__": table_name,
|
||||
"__table_args__": ({"implicit_returning": False},),
|
||||
},
|
||||
)
|
||||
|
||||
classes[class_name] = partition_class
|
||||
|
||||
return classes
|
||||
|
||||
|
||||
# Generate all partition classes and register them on this module
|
||||
_partition_classes = _build_partition_classes()
|
||||
for _name, _cls in _partition_classes.items():
|
||||
setattr(_current_module, _name, _cls)
|
||||
__all__ = list(_partition_classes.keys())
|
||||
@@ -1,13 +0,0 @@
|
||||
"""Posts parent table with PostgreSQL weekly range partitioning on date column."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from python.orm.data_science_dev.base import DataScienceDevBase
|
||||
from python.orm.data_science_dev.posts.columns import PostsColumns
|
||||
|
||||
|
||||
class Posts(PostsColumns, DataScienceDevBase):
|
||||
"""Parent partitioned table for posts, partitioned by week on `date`."""
|
||||
|
||||
__tablename__ = "posts"
|
||||
__table_args__ = ({"postgresql_partition_by": "RANGE (date)"},)
|
||||
@@ -1,24 +0,0 @@
|
||||
"""Richie database ORM exports."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from python.orm.richie.base import RichieBase, TableBase, TableBaseBig, TableBaseSmall
|
||||
from python.orm.richie.contact import (
|
||||
Contact,
|
||||
ContactNeed,
|
||||
ContactRelationship,
|
||||
Need,
|
||||
RelationshipType,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"Contact",
|
||||
"ContactNeed",
|
||||
"ContactRelationship",
|
||||
"Need",
|
||||
"RelationshipType",
|
||||
"RichieBase",
|
||||
"TableBase",
|
||||
"TableBaseBig",
|
||||
"TableBaseSmall",
|
||||
]
|
||||
@@ -1,60 +0,0 @@
|
||||
"""Richie database ORM base."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import datetime
|
||||
|
||||
from sqlalchemy import BigInteger, DateTime, MetaData, SmallInteger, func
|
||||
from sqlalchemy.ext.declarative import AbstractConcreteBase
|
||||
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column
|
||||
|
||||
from python.orm.common import NAMING_CONVENTION
|
||||
|
||||
|
||||
class RichieBase(DeclarativeBase):
|
||||
"""Base class for richie database ORM models."""
|
||||
|
||||
schema_name = "main"
|
||||
|
||||
metadata = MetaData(
|
||||
schema=schema_name,
|
||||
naming_convention=NAMING_CONVENTION,
|
||||
)
|
||||
|
||||
|
||||
class _TableMixin:
|
||||
"""Shared timestamp columns for all table bases."""
|
||||
|
||||
created: Mapped[datetime] = mapped_column(
|
||||
DateTime(timezone=True),
|
||||
server_default=func.now(),
|
||||
)
|
||||
updated: Mapped[datetime] = mapped_column(
|
||||
DateTime(timezone=True),
|
||||
server_default=func.now(),
|
||||
onupdate=func.now(),
|
||||
)
|
||||
|
||||
|
||||
class TableBaseSmall(_TableMixin, AbstractConcreteBase, RichieBase):
|
||||
"""Table with SmallInteger primary key."""
|
||||
|
||||
__abstract__ = True
|
||||
|
||||
id: Mapped[int] = mapped_column(SmallInteger, primary_key=True)
|
||||
|
||||
|
||||
class TableBase(_TableMixin, AbstractConcreteBase, RichieBase):
|
||||
"""Table with Integer primary key."""
|
||||
|
||||
__abstract__ = True
|
||||
|
||||
id: Mapped[int] = mapped_column(primary_key=True)
|
||||
|
||||
|
||||
class TableBaseBig(_TableMixin, AbstractConcreteBase, RichieBase):
|
||||
"""Table with BigInteger primary key."""
|
||||
|
||||
__abstract__ = True
|
||||
|
||||
id: Mapped[int] = mapped_column(BigInteger, primary_key=True)
|
||||
@@ -1,168 +0,0 @@
|
||||
"""Contact database models."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from enum import StrEnum
|
||||
|
||||
from sqlalchemy import ForeignKey, String
|
||||
from sqlalchemy.orm import Mapped, mapped_column, relationship
|
||||
|
||||
from python.orm.richie.base import RichieBase, TableBase
|
||||
|
||||
|
||||
class RelationshipType(StrEnum):
|
||||
"""Relationship types with default closeness weights.
|
||||
|
||||
Default weight is an integer 1-10 where 10 = closest relationship.
|
||||
Users can override this per-relationship in the UI.
|
||||
"""
|
||||
|
||||
SPOUSE = "spouse"
|
||||
PARTNER = "partner"
|
||||
PARENT = "parent"
|
||||
CHILD = "child"
|
||||
SIBLING = "sibling"
|
||||
BEST_FRIEND = "best_friend"
|
||||
GRANDPARENT = "grandparent"
|
||||
GRANDCHILD = "grandchild"
|
||||
AUNT_UNCLE = "aunt_uncle"
|
||||
NIECE_NEPHEW = "niece_nephew"
|
||||
COUSIN = "cousin"
|
||||
IN_LAW = "in_law"
|
||||
CLOSE_FRIEND = "close_friend"
|
||||
FRIEND = "friend"
|
||||
MENTOR = "mentor"
|
||||
MENTEE = "mentee"
|
||||
BUSINESS_PARTNER = "business_partner"
|
||||
COLLEAGUE = "colleague"
|
||||
MANAGER = "manager"
|
||||
DIRECT_REPORT = "direct_report"
|
||||
CLIENT = "client"
|
||||
ACQUAINTANCE = "acquaintance"
|
||||
NEIGHBOR = "neighbor"
|
||||
EX = "ex"
|
||||
OTHER = "other"
|
||||
|
||||
@property
|
||||
def default_weight(self) -> int:
|
||||
"""Return the default closeness weight (1-10) for this relationship type."""
|
||||
weights = {
|
||||
RelationshipType.SPOUSE: 10,
|
||||
RelationshipType.PARTNER: 10,
|
||||
RelationshipType.PARENT: 9,
|
||||
RelationshipType.CHILD: 9,
|
||||
RelationshipType.SIBLING: 9,
|
||||
RelationshipType.BEST_FRIEND: 8,
|
||||
RelationshipType.GRANDPARENT: 7,
|
||||
RelationshipType.GRANDCHILD: 7,
|
||||
RelationshipType.AUNT_UNCLE: 7,
|
||||
RelationshipType.NIECE_NEPHEW: 7,
|
||||
RelationshipType.COUSIN: 7,
|
||||
RelationshipType.IN_LAW: 7,
|
||||
RelationshipType.CLOSE_FRIEND: 6,
|
||||
RelationshipType.FRIEND: 6,
|
||||
RelationshipType.MENTOR: 5,
|
||||
RelationshipType.MENTEE: 5,
|
||||
RelationshipType.BUSINESS_PARTNER: 5,
|
||||
RelationshipType.COLLEAGUE: 4,
|
||||
RelationshipType.MANAGER: 4,
|
||||
RelationshipType.DIRECT_REPORT: 4,
|
||||
RelationshipType.CLIENT: 4,
|
||||
RelationshipType.ACQUAINTANCE: 3,
|
||||
RelationshipType.NEIGHBOR: 3,
|
||||
RelationshipType.EX: 2,
|
||||
RelationshipType.OTHER: 2,
|
||||
}
|
||||
return weights.get(self, 5)
|
||||
|
||||
@property
|
||||
def display_name(self) -> str:
|
||||
"""Return a human-readable display name."""
|
||||
return self.value.replace("_", " ").title()
|
||||
|
||||
|
||||
class ContactNeed(RichieBase):
|
||||
"""Association table: Contact <-> Need."""
|
||||
|
||||
__tablename__ = "contact_need"
|
||||
|
||||
contact_id: Mapped[int] = mapped_column(
|
||||
ForeignKey("main.contact.id", ondelete="CASCADE"),
|
||||
primary_key=True,
|
||||
)
|
||||
need_id: Mapped[int] = mapped_column(
|
||||
ForeignKey("main.need.id", ondelete="CASCADE"),
|
||||
primary_key=True,
|
||||
)
|
||||
|
||||
|
||||
class ContactRelationship(RichieBase):
|
||||
"""Association table: Contact <-> Contact with relationship type and weight."""
|
||||
|
||||
__tablename__ = "contact_relationship"
|
||||
|
||||
contact_id: Mapped[int] = mapped_column(
|
||||
ForeignKey("main.contact.id", ondelete="CASCADE"),
|
||||
primary_key=True,
|
||||
)
|
||||
related_contact_id: Mapped[int] = mapped_column(
|
||||
ForeignKey("main.contact.id", ondelete="CASCADE"),
|
||||
primary_key=True,
|
||||
)
|
||||
relationship_type: Mapped[str] = mapped_column(String(100))
|
||||
closeness_weight: Mapped[int] = mapped_column(default=5)
|
||||
|
||||
|
||||
class Contact(TableBase):
|
||||
"""Contact model."""
|
||||
|
||||
__tablename__ = "contact"
|
||||
|
||||
name: Mapped[str]
|
||||
|
||||
age: Mapped[int | None]
|
||||
bio: Mapped[str | None]
|
||||
current_job: Mapped[str | None]
|
||||
gender: Mapped[str | None]
|
||||
goals: Mapped[str | None]
|
||||
legal_name: Mapped[str | None]
|
||||
profile_pic: Mapped[str | None]
|
||||
safe_conversation_starters: Mapped[str | None]
|
||||
self_sufficiency_score: Mapped[int | None]
|
||||
social_structure_style: Mapped[str | None]
|
||||
ssn: Mapped[str | None]
|
||||
suffix: Mapped[str | None]
|
||||
timezone: Mapped[str | None]
|
||||
topics_to_avoid: Mapped[str | None]
|
||||
|
||||
needs: Mapped[list[Need]] = relationship(
|
||||
"Need",
|
||||
secondary=ContactNeed.__table__,
|
||||
back_populates="contacts",
|
||||
)
|
||||
|
||||
related_to: Mapped[list[ContactRelationship]] = relationship(
|
||||
"ContactRelationship",
|
||||
foreign_keys=[ContactRelationship.contact_id],
|
||||
cascade="all, delete-orphan",
|
||||
)
|
||||
related_from: Mapped[list[ContactRelationship]] = relationship(
|
||||
"ContactRelationship",
|
||||
foreign_keys=[ContactRelationship.related_contact_id],
|
||||
cascade="all, delete-orphan",
|
||||
)
|
||||
|
||||
|
||||
class Need(TableBase):
|
||||
"""Need/accommodation model (e.g., light sensitive, ADHD)."""
|
||||
|
||||
__tablename__ = "need"
|
||||
|
||||
name: Mapped[str]
|
||||
description: Mapped[str | None]
|
||||
|
||||
contacts: Mapped[list[Contact]] = relationship(
|
||||
"Contact",
|
||||
secondary=ContactNeed.__table__,
|
||||
back_populates="needs",
|
||||
)
|
||||
@@ -1,16 +0,0 @@
|
||||
"""Signal bot database ORM exports."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from python.orm.signal_bot.base import SignalBotBase, SignalBotTableBase, SignalBotTableBaseSmall
|
||||
from python.orm.signal_bot.models import DeadLetterMessage, DeviceRole, RoleRecord, SignalDevice
|
||||
|
||||
__all__ = [
|
||||
"DeadLetterMessage",
|
||||
"DeviceRole",
|
||||
"RoleRecord",
|
||||
"SignalBotBase",
|
||||
"SignalBotTableBase",
|
||||
"SignalBotTableBaseSmall",
|
||||
"SignalDevice",
|
||||
]
|
||||
@@ -1,52 +0,0 @@
|
||||
"""Signal bot database ORM base."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import datetime
|
||||
|
||||
from sqlalchemy import DateTime, MetaData, SmallInteger, func
|
||||
from sqlalchemy.ext.declarative import AbstractConcreteBase
|
||||
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column
|
||||
|
||||
from python.orm.common import NAMING_CONVENTION
|
||||
|
||||
|
||||
class SignalBotBase(DeclarativeBase):
|
||||
"""Base class for signal_bot database ORM models."""
|
||||
|
||||
schema_name = "main"
|
||||
|
||||
metadata = MetaData(
|
||||
schema=schema_name,
|
||||
naming_convention=NAMING_CONVENTION,
|
||||
)
|
||||
|
||||
|
||||
class _TableMixin:
|
||||
"""Shared timestamp columns for all table bases."""
|
||||
|
||||
created: Mapped[datetime] = mapped_column(
|
||||
DateTime(timezone=True),
|
||||
server_default=func.now(),
|
||||
)
|
||||
updated: Mapped[datetime] = mapped_column(
|
||||
DateTime(timezone=True),
|
||||
server_default=func.now(),
|
||||
onupdate=func.now(),
|
||||
)
|
||||
|
||||
|
||||
class SignalBotTableBaseSmall(_TableMixin, AbstractConcreteBase, SignalBotBase):
|
||||
"""Table with SmallInteger primary key."""
|
||||
|
||||
__abstract__ = True
|
||||
|
||||
id: Mapped[int] = mapped_column(SmallInteger, primary_key=True)
|
||||
|
||||
|
||||
class SignalBotTableBase(_TableMixin, AbstractConcreteBase, SignalBotBase):
|
||||
"""Table with Integer primary key."""
|
||||
|
||||
__abstract__ = True
|
||||
|
||||
id: Mapped[int] = mapped_column(primary_key=True)
|
||||
@@ -1,62 +0,0 @@
|
||||
"""Signal bot device, role, and dead letter ORM models."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import datetime
|
||||
|
||||
from sqlalchemy import DateTime, Enum, ForeignKey, SmallInteger, String, Text, UniqueConstraint
|
||||
from sqlalchemy.orm import Mapped, mapped_column, relationship
|
||||
|
||||
from python.orm.signal_bot.base import SignalBotTableBase, SignalBotTableBaseSmall
|
||||
from python.signal_bot.models import MessageStatus, TrustLevel
|
||||
|
||||
|
||||
class RoleRecord(SignalBotTableBaseSmall):
|
||||
"""Lookup table for RBAC roles, keyed by smallint."""
|
||||
|
||||
__tablename__ = "role"
|
||||
|
||||
name: Mapped[str] = mapped_column(String(50), unique=True)
|
||||
|
||||
|
||||
class DeviceRole(SignalBotTableBase):
|
||||
"""Association between a device and a role."""
|
||||
|
||||
__tablename__ = "device_role"
|
||||
__table_args__ = (
|
||||
UniqueConstraint("device_id", "role_id", name="uq_device_role_device_role"),
|
||||
{"schema": "main"},
|
||||
)
|
||||
|
||||
device_id: Mapped[int] = mapped_column(ForeignKey("main.signal_device.id"))
|
||||
role_id: Mapped[int] = mapped_column(SmallInteger, ForeignKey("main.role.id"))
|
||||
|
||||
|
||||
class SignalDevice(SignalBotTableBase):
|
||||
"""A Signal device tracked by phone number and safety number."""
|
||||
|
||||
__tablename__ = "signal_device"
|
||||
|
||||
phone_number: Mapped[str] = mapped_column(String(50), unique=True)
|
||||
safety_number: Mapped[str | None]
|
||||
trust_level: Mapped[TrustLevel] = mapped_column(
|
||||
Enum(TrustLevel, name="trust_level", create_constraint=False, native_enum=False),
|
||||
default=TrustLevel.UNVERIFIED,
|
||||
)
|
||||
last_seen: Mapped[datetime] = mapped_column(DateTime(timezone=True))
|
||||
|
||||
roles: Mapped[list[RoleRecord]] = relationship(secondary=DeviceRole.__table__)
|
||||
|
||||
|
||||
class DeadLetterMessage(SignalBotTableBase):
|
||||
"""A Signal message that failed processing and was sent to the dead letter queue."""
|
||||
|
||||
__tablename__ = "dead_letter_message"
|
||||
|
||||
source: Mapped[str]
|
||||
message: Mapped[str] = mapped_column(Text)
|
||||
received_at: Mapped[datetime] = mapped_column(DateTime(timezone=True))
|
||||
status: Mapped[MessageStatus] = mapped_column(
|
||||
Enum(MessageStatus, name="message_status", create_constraint=False, native_enum=False),
|
||||
default=MessageStatus.UNPROCESSED,
|
||||
)
|
||||
@@ -1 +0,0 @@
|
||||
"""Van inventory database ORM exports."""
|
||||
@@ -1,39 +0,0 @@
|
||||
"""Van inventory database ORM base."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import datetime
|
||||
|
||||
from sqlalchemy import DateTime, MetaData, func
|
||||
from sqlalchemy.ext.declarative import AbstractConcreteBase
|
||||
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column
|
||||
|
||||
from python.orm.common import NAMING_CONVENTION
|
||||
|
||||
|
||||
class VanInventoryBase(DeclarativeBase):
|
||||
"""Base class for van_inventory database ORM models."""
|
||||
|
||||
schema_name = "main"
|
||||
|
||||
metadata = MetaData(
|
||||
schema=schema_name,
|
||||
naming_convention=NAMING_CONVENTION,
|
||||
)
|
||||
|
||||
|
||||
class VanTableBase(AbstractConcreteBase, VanInventoryBase):
|
||||
"""Abstract concrete base for van_inventory tables with IDs and timestamps."""
|
||||
|
||||
__abstract__ = True
|
||||
|
||||
id: Mapped[int] = mapped_column(primary_key=True)
|
||||
created: Mapped[datetime] = mapped_column(
|
||||
DateTime(timezone=True),
|
||||
server_default=func.now(),
|
||||
)
|
||||
updated: Mapped[datetime] = mapped_column(
|
||||
DateTime(timezone=True),
|
||||
server_default=func.now(),
|
||||
onupdate=func.now(),
|
||||
)
|
||||
@@ -1,46 +0,0 @@
|
||||
"""Van inventory ORM models."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from sqlalchemy import ForeignKey, UniqueConstraint
|
||||
from sqlalchemy.orm import Mapped, mapped_column, relationship
|
||||
|
||||
from python.orm.van_inventory.base import VanTableBase
|
||||
|
||||
|
||||
class Item(VanTableBase):
|
||||
"""A food item in the van."""
|
||||
|
||||
__tablename__ = "items"
|
||||
|
||||
name: Mapped[str] = mapped_column(unique=True)
|
||||
quantity: Mapped[float] = mapped_column(default=0)
|
||||
unit: Mapped[str]
|
||||
category: Mapped[str | None]
|
||||
|
||||
meal_ingredients: Mapped[list[MealIngredient]] = relationship(back_populates="item")
|
||||
|
||||
|
||||
class Meal(VanTableBase):
|
||||
"""A meal that can be made from items in the van."""
|
||||
|
||||
__tablename__ = "meals"
|
||||
|
||||
name: Mapped[str] = mapped_column(unique=True)
|
||||
instructions: Mapped[str | None]
|
||||
|
||||
ingredients: Mapped[list[MealIngredient]] = relationship(back_populates="meal")
|
||||
|
||||
|
||||
class MealIngredient(VanTableBase):
|
||||
"""Links a meal to the items it requires, with quantities."""
|
||||
|
||||
__tablename__ = "meal_ingredients"
|
||||
__table_args__ = (UniqueConstraint("meal_id", "item_id"),)
|
||||
|
||||
meal_id: Mapped[int] = mapped_column(ForeignKey("meals.id"))
|
||||
item_id: Mapped[int] = mapped_column(ForeignKey("items.id"))
|
||||
quantity_needed: Mapped[float]
|
||||
|
||||
meal: Mapped[Meal] = relationship(back_populates="ingredients")
|
||||
item: Mapped[Item] = relationship(back_populates="meal_ingredients")
|
||||
@@ -1,25 +0,0 @@
|
||||
# Unsloth fine-tuning container for Qwen 3.5 4B on RTX 3090.
|
||||
#
|
||||
# Build:
|
||||
# docker build -f python/prompt_bench/Dockerfile.finetune -t bill-finetune .
|
||||
#
|
||||
# Run:
|
||||
# docker run --rm --device=nvidia.com/gpu=all --ipc=host \
|
||||
# -v $(pwd)/output:/workspace/output \
|
||||
# -v $(pwd)/output/finetune_dataset.jsonl:/workspace/dataset.jsonl:ro \
|
||||
# -v /zfs/models/hf:/models \
|
||||
# bill-finetune \
|
||||
# --dataset /workspace/dataset.jsonl \
|
||||
# --output-dir /workspace/output/qwen-bill-summarizer
|
||||
|
||||
FROM ghcr.io/unslothai/unsloth:latest
|
||||
|
||||
RUN pip install --no-cache-dir typer
|
||||
|
||||
WORKDIR /workspace
|
||||
COPY python/prompt_bench/finetune.py python/prompt_bench/finetune.py
|
||||
COPY python/prompt_bench/summarization_prompts.py python/prompt_bench/summarization_prompts.py
|
||||
COPY python/prompt_bench/__init__.py python/prompt_bench/__init__.py
|
||||
COPY python/__init__.py python/__init__.py
|
||||
|
||||
ENTRYPOINT ["python", "-m", "python.prompt_bench.finetune"]
|
||||
@@ -1 +0,0 @@
|
||||
"""Prompt benchmarking system for evaluating LLMs via vLLM."""
|
||||
@@ -1,233 +0,0 @@
|
||||
"""Submit an OpenAI Batch API bill-summarization job over compressed text.
|
||||
|
||||
Reads the first N bills from a CSV with a `text_content` column, compresses
|
||||
each via `bill_token_compression.compress_bill_text`, builds a JSONL file of
|
||||
summarization requests, and submits it as an asynchronous Batch API job
|
||||
against `/v1/chat/completions`. Also writes a CSV of per-bill pre/post-
|
||||
compression token counts.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import csv
|
||||
import json
|
||||
import logging
|
||||
import re
|
||||
import sys
|
||||
from os import getenv
|
||||
from pathlib import Path
|
||||
from typing import Annotated
|
||||
|
||||
import httpx
|
||||
import typer
|
||||
from tiktoken import Encoding, get_encoding
|
||||
|
||||
from python.prompt_bench.bill_token_compression import compress_bill_text
|
||||
from python.prompt_bench.summarization_prompts import SUMMARIZATION_SYSTEM_PROMPT, SUMMARIZATION_USER_TEMPLATE
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
OPENAI_API_BASE = "https://api.openai.com/v1"
|
||||
|
||||
|
||||
def load_bills(csv_path: Path, count: int = 0) -> list[tuple[str, str]]:
|
||||
"""Return (bill_id, text_content) tuples with non-empty text.
|
||||
|
||||
If `count` is 0 or negative, all rows are returned.
|
||||
"""
|
||||
csv.field_size_limit(sys.maxsize)
|
||||
bills: list[tuple[str, str]] = []
|
||||
with csv_path.open(newline="", encoding="utf-8") as handle:
|
||||
reader = csv.DictReader(handle)
|
||||
for row in reader:
|
||||
text_content = (row.get("text_content") or "").strip()
|
||||
if not text_content:
|
||||
continue
|
||||
bill_id = row.get("bill_id") or row.get("id") or f"row-{len(bills)}"
|
||||
version_code = row.get("version_code") or ""
|
||||
unique_id = f"{bill_id}-{version_code}" if version_code else bill_id
|
||||
bills.append((unique_id, text_content))
|
||||
if count > 0 and len(bills) >= count:
|
||||
break
|
||||
return bills
|
||||
|
||||
|
||||
def safe_filename(value: str) -> str:
|
||||
"""Make a string safe for use as a filename or batch custom_id."""
|
||||
return re.sub(r"[^A-Za-z0-9._-]+", "_", value).strip("_") or "unnamed"
|
||||
|
||||
|
||||
def build_request(custom_id: str, model: str, bill_text: str) -> dict:
|
||||
"""Build one OpenAI batch request line."""
|
||||
return {
|
||||
"custom_id": custom_id,
|
||||
"method": "POST",
|
||||
"url": "/v1/chat/completions",
|
||||
"body": {
|
||||
"model": model,
|
||||
"messages": [
|
||||
{"role": "system", "content": SUMMARIZATION_SYSTEM_PROMPT},
|
||||
{"role": "user", "content": SUMMARIZATION_USER_TEMPLATE.format(text_content=bill_text)},
|
||||
],
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def write_jsonl(path: Path, lines: list[dict]) -> None:
|
||||
"""Write a list of dicts as JSONL."""
|
||||
with path.open("w", encoding="utf-8") as handle:
|
||||
for line in lines:
|
||||
handle.write(json.dumps(line, ensure_ascii=False))
|
||||
handle.write("\n")
|
||||
|
||||
|
||||
def upload_file(client: httpx.Client, path: Path) -> str:
|
||||
"""Upload a JSONL file to the OpenAI Files API and return its file id."""
|
||||
with path.open("rb") as handle:
|
||||
response = client.post(
|
||||
f"{OPENAI_API_BASE}/files",
|
||||
files={"file": (path.name, handle, "application/jsonl")},
|
||||
data={"purpose": "batch"},
|
||||
)
|
||||
response.raise_for_status()
|
||||
return response.json()["id"]
|
||||
|
||||
|
||||
def prepare_requests(
|
||||
bills: list[tuple[str, str]],
|
||||
*,
|
||||
model: str,
|
||||
encoder: Encoding,
|
||||
) -> tuple[list[dict], list[dict]]:
|
||||
"""Build (request_lines, token_rows) from bills.
|
||||
|
||||
Each bill is compressed before being turned into a request line.
|
||||
Each `token_rows` entry has chars + token counts for one bill so the caller
|
||||
can write a per-bill CSV.
|
||||
"""
|
||||
request_lines: list[dict] = []
|
||||
token_rows: list[dict] = []
|
||||
for bill_id, text_content in bills:
|
||||
raw_token_count = len(encoder.encode(text_content))
|
||||
compressed_text = compress_bill_text(text_content)
|
||||
compressed_token_count = len(encoder.encode(compressed_text))
|
||||
token_rows.append(
|
||||
{
|
||||
"bill_id": bill_id,
|
||||
"raw_chars": len(text_content),
|
||||
"compressed_chars": len(compressed_text),
|
||||
"raw_tokens": raw_token_count,
|
||||
"compressed_tokens": compressed_token_count,
|
||||
"token_ratio": (compressed_token_count / raw_token_count) if raw_token_count else None,
|
||||
},
|
||||
)
|
||||
safe_id = safe_filename(bill_id)
|
||||
request_lines.append(build_request(safe_id, model, compressed_text))
|
||||
return request_lines, token_rows
|
||||
|
||||
|
||||
def write_token_csv(path: Path, token_rows: list[dict]) -> tuple[int, int]:
|
||||
"""Write per-bill token counts to CSV. Returns (raw_total, compressed_total)."""
|
||||
with path.open("w", newline="", encoding="utf-8") as handle:
|
||||
writer = csv.DictWriter(
|
||||
handle,
|
||||
fieldnames=["bill_id", "raw_chars", "compressed_chars", "raw_tokens", "compressed_tokens", "token_ratio"],
|
||||
)
|
||||
writer.writeheader()
|
||||
writer.writerows(token_rows)
|
||||
raw_total = sum(row["raw_tokens"] for row in token_rows)
|
||||
compressed_total = sum(row["compressed_tokens"] for row in token_rows)
|
||||
return raw_total, compressed_total
|
||||
|
||||
|
||||
def create_batch(client: httpx.Client, input_file_id: str, description: str) -> dict:
|
||||
"""Create a batch job and return its full response payload."""
|
||||
response = client.post(
|
||||
f"{OPENAI_API_BASE}/batches",
|
||||
json={
|
||||
"input_file_id": input_file_id,
|
||||
"endpoint": "/v1/chat/completions",
|
||||
"completion_window": "24h",
|
||||
"metadata": {"description": description},
|
||||
},
|
||||
)
|
||||
response.raise_for_status()
|
||||
return response.json()
|
||||
|
||||
|
||||
def main(
|
||||
csv_path: Annotated[Path, typer.Option("--csv", help="Bills CSV path")] = Path("bills.csv"),
|
||||
output_dir: Annotated[Path, typer.Option("--output-dir", help="Where to write JSONL + metadata")] = Path(
|
||||
"output/openai_batch",
|
||||
),
|
||||
model: Annotated[str, typer.Option(help="OpenAI model id")] = "gpt-5-mini",
|
||||
count: Annotated[int, typer.Option(help="Max bills to process, 0 = all")] = 0,
|
||||
log_level: Annotated[str, typer.Option(help="Log level")] = "INFO",
|
||||
) -> None:
|
||||
"""Submit an OpenAI Batch job of compressed bill summaries."""
|
||||
logging.basicConfig(level=log_level, format="%(asctime)s %(levelname)s %(name)s: %(message)s")
|
||||
|
||||
api_key = getenv("CLOSEDAI_TOKEN") or getenv("OPENAI_API_KEY")
|
||||
if not api_key:
|
||||
message = "Neither CLOSEDAI_TOKEN nor OPENAI_API_KEY is set"
|
||||
raise typer.BadParameter(message)
|
||||
if not csv_path.is_file():
|
||||
message = f"CSV not found: {csv_path}"
|
||||
raise typer.BadParameter(message)
|
||||
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
logger.info("Loading %d bills from %s", count, csv_path)
|
||||
bills = load_bills(csv_path, count)
|
||||
if len(bills) < count:
|
||||
logger.warning("Only %d bills available (requested %d)", len(bills), count)
|
||||
|
||||
encoder = get_encoding("o200k_base")
|
||||
request_lines, token_rows = prepare_requests(bills, model=model, encoder=encoder)
|
||||
|
||||
token_csv_path = output_dir / "token_counts.csv"
|
||||
raw_tokens_total, compressed_tokens_total = write_token_csv(token_csv_path, token_rows)
|
||||
logger.info(
|
||||
"Token counts: raw=%d compressed=%d ratio=%.3f -> %s",
|
||||
raw_tokens_total,
|
||||
compressed_tokens_total,
|
||||
(compressed_tokens_total / raw_tokens_total) if raw_tokens_total else 0.0,
|
||||
token_csv_path,
|
||||
)
|
||||
|
||||
jsonl_path = output_dir / "requests.jsonl"
|
||||
write_jsonl(jsonl_path, request_lines)
|
||||
logger.info("Wrote %s (%d bills)", jsonl_path, len(request_lines))
|
||||
|
||||
headers = {"Authorization": f"Bearer {api_key}"}
|
||||
with httpx.Client(headers=headers, timeout=httpx.Timeout(300.0)) as client:
|
||||
logger.info("Uploading JSONL")
|
||||
file_id = upload_file(client, jsonl_path)
|
||||
logger.info("Uploaded: %s", file_id)
|
||||
|
||||
logger.info("Creating batch")
|
||||
batch = create_batch(client, file_id, f"compressed bill summaries x{len(request_lines)} ({model})")
|
||||
logger.info("Batch created: %s", batch["id"])
|
||||
|
||||
metadata = {
|
||||
"model": model,
|
||||
"count": len(bills),
|
||||
"jsonl": str(jsonl_path),
|
||||
"input_file_id": file_id,
|
||||
"batch_id": batch["id"],
|
||||
"raw_tokens_total": raw_tokens_total,
|
||||
"compressed_tokens_total": compressed_tokens_total,
|
||||
"batch": batch,
|
||||
}
|
||||
metadata_path = output_dir / "batch.json"
|
||||
metadata_path.write_text(json.dumps(metadata, indent=2))
|
||||
logger.info("Wrote metadata to %s", metadata_path)
|
||||
|
||||
|
||||
def cli() -> None:
|
||||
"""Typer entry point."""
|
||||
typer.run(main)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
cli()
|
||||
@@ -1,162 +0,0 @@
|
||||
"""Lossless-ish text compression for Congressional bill text."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
|
||||
STATES = (
|
||||
"Alabama",
|
||||
"Alaska",
|
||||
"Arizona",
|
||||
"Arkansas",
|
||||
"California",
|
||||
"Colorado",
|
||||
"Connecticut",
|
||||
"Delaware",
|
||||
"Florida",
|
||||
"Georgia",
|
||||
"Hawaii",
|
||||
"Idaho",
|
||||
"Illinois",
|
||||
"Indiana",
|
||||
"Iowa",
|
||||
"Kansas",
|
||||
"Kentucky",
|
||||
"Louisiana",
|
||||
"Maine",
|
||||
"Maryland",
|
||||
"Massachusetts",
|
||||
"Michigan",
|
||||
"Minnesota",
|
||||
"Mississippi",
|
||||
"Missouri",
|
||||
"Montana",
|
||||
"Nebraska",
|
||||
"Nevada",
|
||||
"New Hampshire",
|
||||
"New Jersey",
|
||||
"New Mexico",
|
||||
"New York",
|
||||
"North Carolina",
|
||||
"North Dakota",
|
||||
"Ohio",
|
||||
"Oklahoma",
|
||||
"Oregon",
|
||||
"Pennsylvania",
|
||||
"Rhode Island",
|
||||
"South Carolina",
|
||||
"South Dakota",
|
||||
"Tennessee",
|
||||
"Texas",
|
||||
"Utah",
|
||||
"Vermont",
|
||||
"Virginia",
|
||||
"Washington",
|
||||
"West Virginia",
|
||||
"Wisconsin",
|
||||
"Wyoming",
|
||||
"Puerto Rico",
|
||||
"Guam",
|
||||
"American Samoa",
|
||||
"District of Columbia",
|
||||
"US Virgin Islands",
|
||||
)
|
||||
STATE_PATTERNS = [(re.compile(re.escape(state), re.IGNORECASE), state) for state in STATES]
|
||||
|
||||
|
||||
def normalize_state_names(text: str) -> str:
|
||||
"""Replace any casing of state names with title case."""
|
||||
for pattern, replacement in STATE_PATTERNS:
|
||||
text = pattern.sub(replacement, text)
|
||||
return text
|
||||
|
||||
|
||||
def strip_number_commas(text: str) -> str:
|
||||
"""Remove commas from numeric thousands separators."""
|
||||
return re.sub(r"(\d{1,3}(?:,\d{3})+)", lambda match: match.group().replace(",", ""), text)
|
||||
|
||||
|
||||
def strip_horizontal_rules(text: str) -> str:
|
||||
"""Remove ASCII horizontal-rule lines built from underscores, dashes, equals, or asterisks."""
|
||||
return re.sub(r"^\s*[_\-=\*]{3,}\s*$", "", text, flags=re.MULTILINE)
|
||||
|
||||
|
||||
def collapse_double_dashes(text: str) -> str:
|
||||
"""Replace ``--`` em-dash stand-ins with a single space so they don't tokenize oddly."""
|
||||
return text.replace("--", " ")
|
||||
|
||||
|
||||
def collapse_inline_whitespace(text: str) -> str:
|
||||
"""Collapse runs of horizontal whitespace (spaces, tabs) into a single space, leaving newlines intact."""
|
||||
return re.sub(r"[^\S\n]+", " ", text)
|
||||
|
||||
|
||||
def collapse_blank_lines(text: str) -> str:
|
||||
"""Collapse three-or-more consecutive newlines down to a blank-line separator."""
|
||||
return re.sub(r"\n{3,}", "\n\n", text)
|
||||
|
||||
|
||||
def trim_line_edges(text: str) -> str:
|
||||
"""Strip spaces immediately before and after newline characters on every line."""
|
||||
text = re.sub(r" +\n", "\n", text)
|
||||
return re.sub(r"\n +", "\n", text)
|
||||
|
||||
|
||||
def shorten_section_markers(text: str) -> str:
|
||||
"""Rewrite ``Sec. 12.`` style section headings as the more compact ``SEC 12``."""
|
||||
return re.sub(r"(?i)sec\.\s*(\d+[a-zA-Z]?)\.", r"SEC \1", text)
|
||||
|
||||
|
||||
def unwrap_parens(text: str) -> str:
|
||||
"""Strip parentheses around short alphanumeric labels like ``(a)`` or ``(12)``."""
|
||||
return re.sub(r"\(([a-zA-Z0-9]+)\)", r"\1", text)
|
||||
|
||||
|
||||
def strip_typeset_quotes(text: str) -> str:
|
||||
"""Remove the `` and '' typeset quote markers used in the GPO bill format."""
|
||||
return text.replace("``", "").replace("''", "")
|
||||
|
||||
|
||||
def normalize_usc_acronym(text: str) -> str:
|
||||
"""Collapse ``U.S.C.`` to ``USC`` to save tokens on the common citation."""
|
||||
return text.replace("U.S.C.", "USC")
|
||||
|
||||
|
||||
def normalize_us_acronym(text: str) -> str:
|
||||
"""Normalize the various ``U.S.``/``U. S.`` spellings to the bare ``US`` form."""
|
||||
for acronym in ("U. S.", "u. s.", "U.S. ", "u.s. "):
|
||||
text = text.replace(acronym, "US ")
|
||||
return text
|
||||
|
||||
|
||||
def collapse_ellipses(text: str) -> str:
|
||||
"""Collapse runs of two-or-more periods (``...``, ``....``) down to a single period."""
|
||||
return re.sub(r"\.{2,}", ".", text)
|
||||
|
||||
|
||||
COMPRESSION_STEPS = (
|
||||
strip_horizontal_rules,
|
||||
collapse_double_dashes,
|
||||
collapse_inline_whitespace,
|
||||
collapse_blank_lines,
|
||||
trim_line_edges,
|
||||
shorten_section_markers,
|
||||
unwrap_parens,
|
||||
strip_typeset_quotes,
|
||||
normalize_usc_acronym,
|
||||
normalize_us_acronym,
|
||||
strip_number_commas,
|
||||
collapse_ellipses,
|
||||
normalize_state_names,
|
||||
)
|
||||
|
||||
|
||||
def compress_bill_text(text: str) -> str:
|
||||
"""Apply lossless-ish whitespace and boilerplate compression to bill text.
|
||||
|
||||
Runs every transform in :data:`COMPRESSION_STEPS` in order, then strips
|
||||
leading/trailing whitespace from the final result.
|
||||
"""
|
||||
for step in COMPRESSION_STEPS:
|
||||
text = step(text)
|
||||
return text.strip()
|
||||
@@ -1,236 +0,0 @@
|
||||
"""Run two interactive OpenAI chat-completion sweeps over bill text.
|
||||
|
||||
Reads the first N bills from a CSV with a `text_content` column and sends two
|
||||
sweeps through `/v1/chat/completions` concurrently — one with the raw bill
|
||||
text, one with the compressed bill text. Each request's prompt is saved to
|
||||
disk alongside the OpenAI response id so the prompts and responses can be
|
||||
correlated later.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import csv
|
||||
import json
|
||||
import logging
|
||||
import re
|
||||
import sys
|
||||
import time
|
||||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||
from os import getenv
|
||||
from pathlib import Path
|
||||
from typing import Annotated
|
||||
|
||||
import httpx
|
||||
import typer
|
||||
|
||||
from python.prompt_bench.bill_token_compression import compress_bill_text
|
||||
from python.prompt_bench.summarization_prompts import SUMMARIZATION_SYSTEM_PROMPT, SUMMARIZATION_USER_TEMPLATE
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
OPENAI_API_BASE = "https://api.openai.com/v1"
|
||||
DEFAULT_MODEL = "gpt-5.4-mini"
|
||||
DEFAULT_COUNT = 100
|
||||
SEED = 42
|
||||
|
||||
|
||||
def load_bills(csv_path: Path, count: int) -> list[tuple[str, str]]:
|
||||
"""Return up to `count` (bill_id, text_content) tuples with non-empty text."""
|
||||
csv.field_size_limit(sys.maxsize)
|
||||
bills: list[tuple[str, str]] = []
|
||||
with csv_path.open(newline="", encoding="utf-8") as handle:
|
||||
reader = csv.DictReader(handle)
|
||||
for row in reader:
|
||||
text_content = (row.get("text_content") or "").strip()
|
||||
if not text_content:
|
||||
continue
|
||||
bill_id = row.get("bill_id") or row.get("id") or f"row-{len(bills)}"
|
||||
version_code = row.get("version_code") or ""
|
||||
unique_id = f"{bill_id}-{version_code}" if version_code else bill_id
|
||||
bills.append((unique_id, text_content))
|
||||
if len(bills) >= count:
|
||||
break
|
||||
return bills
|
||||
|
||||
|
||||
def build_messages(bill_text: str) -> list[dict]:
|
||||
"""Return the system + user message pair for a bill."""
|
||||
return [
|
||||
{"role": "system", "content": SUMMARIZATION_SYSTEM_PROMPT},
|
||||
{"role": "user", "content": SUMMARIZATION_USER_TEMPLATE.format(text_content=bill_text)},
|
||||
]
|
||||
|
||||
|
||||
def safe_filename(value: str) -> str:
|
||||
"""Make a string safe for use as a filename."""
|
||||
return re.sub(r"[^A-Za-z0-9._-]+", "_", value).strip("_") or "unnamed"
|
||||
|
||||
|
||||
def run_one_request(
|
||||
client: httpx.Client,
|
||||
*,
|
||||
bill_id: str,
|
||||
label: str,
|
||||
bill_text: str,
|
||||
model: str,
|
||||
output_path: Path,
|
||||
) -> tuple[bool, float, str | None]:
|
||||
"""Send one chat-completion request and persist prompt + response.
|
||||
|
||||
Returns (success, elapsed_seconds, response_id).
|
||||
"""
|
||||
messages = build_messages(bill_text)
|
||||
payload = {
|
||||
"model": model,
|
||||
"messages": messages,
|
||||
"seed": SEED,
|
||||
}
|
||||
start = time.monotonic()
|
||||
record: dict = {
|
||||
"bill_id": bill_id,
|
||||
"label": label,
|
||||
"model": model,
|
||||
"seed": SEED,
|
||||
"input_chars": len(bill_text),
|
||||
"messages": messages,
|
||||
}
|
||||
try:
|
||||
response = client.post(f"{OPENAI_API_BASE}/chat/completions", json=payload)
|
||||
response.raise_for_status()
|
||||
body = response.json()
|
||||
except httpx.HTTPStatusError as error:
|
||||
elapsed = time.monotonic() - start
|
||||
record["error"] = {
|
||||
"status_code": error.response.status_code,
|
||||
"body": error.response.text,
|
||||
"elapsed_seconds": elapsed,
|
||||
}
|
||||
output_path.write_text(json.dumps(record, ensure_ascii=False, indent=2))
|
||||
logger.exception("HTTP error for %s/%s after %.2fs", label, bill_id, elapsed)
|
||||
return False, elapsed, None
|
||||
except Exception as error:
|
||||
elapsed = time.monotonic() - start
|
||||
record["error"] = {"message": str(error), "elapsed_seconds": elapsed}
|
||||
output_path.write_text(json.dumps(record, ensure_ascii=False, indent=2))
|
||||
logger.exception("Failed: %s/%s after %.2fs", label, bill_id, elapsed)
|
||||
return False, elapsed, None
|
||||
|
||||
elapsed = time.monotonic() - start
|
||||
response_id = body.get("id")
|
||||
record["response_id"] = response_id
|
||||
record["elapsed_seconds"] = elapsed
|
||||
record["usage"] = body.get("usage")
|
||||
record["response"] = body
|
||||
output_path.write_text(json.dumps(record, ensure_ascii=False, indent=2))
|
||||
logger.info("Done: %s/%s id=%s in %.2fs", label, bill_id, response_id, elapsed)
|
||||
return True, elapsed, response_id
|
||||
|
||||
|
||||
def main(
|
||||
csv_path: Annotated[Path, typer.Option("--csv", help="Bills CSV path")] = Path("bills.csv"),
|
||||
output_dir: Annotated[Path, typer.Option("--output-dir", help="Where to write per-request JSON")] = Path(
|
||||
"output/openai_runs",
|
||||
),
|
||||
model: Annotated[str, typer.Option(help="OpenAI model id")] = DEFAULT_MODEL,
|
||||
count: Annotated[int, typer.Option(help="Number of bills per set")] = DEFAULT_COUNT,
|
||||
concurrency: Annotated[int, typer.Option(help="Concurrent in-flight requests")] = 16,
|
||||
log_level: Annotated[str, typer.Option(help="Log level")] = "INFO",
|
||||
) -> None:
|
||||
"""Run two interactive OpenAI sweeps (compressed + uncompressed) over bill text."""
|
||||
logging.basicConfig(level=log_level, format="%(asctime)s %(levelname)s %(name)s: %(message)s")
|
||||
|
||||
api_key = getenv("CLOSEDAI_TOKEN") or getenv("OPENAI_API_KEY")
|
||||
if not api_key:
|
||||
message = "Neither CLOSEDAI_TOKEN nor OPENAI_API_KEY is set"
|
||||
raise typer.BadParameter(message)
|
||||
if not csv_path.is_file():
|
||||
message = f"CSV not found: {csv_path}"
|
||||
raise typer.BadParameter(message)
|
||||
|
||||
compressed_dir = output_dir / "compressed"
|
||||
uncompressed_dir = output_dir / "uncompressed"
|
||||
compressed_dir.mkdir(parents=True, exist_ok=True)
|
||||
uncompressed_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
logger.info("Loading %d bills from %s", count, csv_path)
|
||||
bills = load_bills(csv_path, count)
|
||||
if len(bills) < count:
|
||||
logger.warning("Only %d bills available (requested %d)", len(bills), count)
|
||||
|
||||
tasks: list[tuple[str, str, str, Path]] = []
|
||||
for bill_id, text_content in bills:
|
||||
filename = f"{safe_filename(bill_id)}.json"
|
||||
tasks.append((bill_id, "compressed", compress_bill_text(text_content), compressed_dir / filename))
|
||||
tasks.append((bill_id, "uncompressed", text_content, uncompressed_dir / filename))
|
||||
|
||||
logger.info("Submitting %d requests at concurrency=%d", len(tasks), concurrency)
|
||||
|
||||
headers = {"Authorization": f"Bearer {api_key}"}
|
||||
completed = 0
|
||||
failed = 0
|
||||
index: list[dict] = []
|
||||
wall_start = time.monotonic()
|
||||
with (
|
||||
httpx.Client(headers=headers, timeout=httpx.Timeout(300.0)) as client,
|
||||
ThreadPoolExecutor(
|
||||
max_workers=concurrency,
|
||||
) as executor,
|
||||
):
|
||||
future_to_task = {
|
||||
executor.submit(
|
||||
run_one_request,
|
||||
client,
|
||||
bill_id=bill_id,
|
||||
label=label,
|
||||
bill_text=bill_text,
|
||||
model=model,
|
||||
output_path=output_path,
|
||||
): (bill_id, label, output_path)
|
||||
for bill_id, label, bill_text, output_path in tasks
|
||||
}
|
||||
for future in as_completed(future_to_task):
|
||||
bill_id, label, output_path = future_to_task[future]
|
||||
success, elapsed, response_id = future.result()
|
||||
if success:
|
||||
completed += 1
|
||||
else:
|
||||
failed += 1
|
||||
index.append(
|
||||
{
|
||||
"bill_id": bill_id,
|
||||
"label": label,
|
||||
"response_id": response_id,
|
||||
"elapsed_seconds": elapsed,
|
||||
"success": success,
|
||||
"path": str(output_path),
|
||||
},
|
||||
)
|
||||
wall_elapsed = time.monotonic() - wall_start
|
||||
|
||||
summary = {
|
||||
"model": model,
|
||||
"count": len(bills),
|
||||
"completed": completed,
|
||||
"failed": failed,
|
||||
"wall_seconds": wall_elapsed,
|
||||
"concurrency": concurrency,
|
||||
"results": index,
|
||||
}
|
||||
summary_path = output_dir / "summary.json"
|
||||
summary_path.write_text(json.dumps(summary, indent=2))
|
||||
logger.info(
|
||||
"Done: completed=%d failed=%d wall=%.1fs summary=%s",
|
||||
completed,
|
||||
failed,
|
||||
wall_elapsed,
|
||||
summary_path,
|
||||
)
|
||||
|
||||
|
||||
def cli() -> None:
|
||||
"""Typer entry point."""
|
||||
typer.run(main)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
cli()
|
||||
@@ -1 +0,0 @@
|
||||
"""Prompt benchmarking system for evaluating LLMs via vLLM."""
|
||||
@@ -1,165 +0,0 @@
|
||||
"""Docker container lifecycle management for Unsloth fine-tuning."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import subprocess
|
||||
from pathlib import Path
|
||||
from typing import Annotated
|
||||
|
||||
import typer
|
||||
|
||||
from python.prompt_bench.containers.lib import check_gpu_free
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
CONTAINER_NAME = "bill-finetune"
|
||||
FINETUNE_IMAGE = "bill-finetune:latest"
|
||||
DOCKERFILE_PATH = "/home/richie/dotfiles/python/prompt_bench/Dockerfile.finetune"
|
||||
DEFAULT_HF_CACHE = Path("/zfs/models/hf")
|
||||
|
||||
|
||||
def build_image() -> None:
|
||||
"""Build the fine-tuning Docker image."""
|
||||
logger.info("Building fine-tuning image: %s", FINETUNE_IMAGE)
|
||||
result = subprocess.run(
|
||||
["docker", "build", "-f", DOCKERFILE_PATH, "-t", FINETUNE_IMAGE, "."],
|
||||
text=True,
|
||||
check=False,
|
||||
)
|
||||
if result.returncode != 0:
|
||||
message = "Failed to build fine-tuning image"
|
||||
raise RuntimeError(message)
|
||||
logger.info("Image built: %s", FINETUNE_IMAGE)
|
||||
|
||||
|
||||
def start_finetune(
|
||||
*,
|
||||
dataset_path: Path,
|
||||
output_dir: Path,
|
||||
hf_cache: Path = DEFAULT_HF_CACHE,
|
||||
) -> None:
|
||||
"""Run the fine-tuning container.
|
||||
|
||||
Args:
|
||||
dataset_path: Host path to the fine-tuning JSONL dataset.
|
||||
output_dir: Host path where the trained model will be saved.
|
||||
hf_cache: Host path to HuggingFace model cache (bind-mounted to avoid re-downloading).
|
||||
validation_split: Fraction of data held out for validation.
|
||||
"""
|
||||
dataset_path = dataset_path.resolve()
|
||||
output_dir = output_dir.resolve()
|
||||
|
||||
if not dataset_path.is_file():
|
||||
message = f"Dataset not found: {dataset_path}"
|
||||
raise FileNotFoundError(message)
|
||||
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
stop_finetune()
|
||||
|
||||
hf_cache = hf_cache.resolve()
|
||||
hf_cache.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
command = [
|
||||
"docker",
|
||||
"run",
|
||||
"--name",
|
||||
CONTAINER_NAME,
|
||||
"--device=nvidia.com/gpu=all",
|
||||
"--ipc=host",
|
||||
"-v",
|
||||
f"{hf_cache}:/root/.cache/huggingface",
|
||||
"-v",
|
||||
f"{output_dir}:/workspace/output/qwen-bill-summarizer",
|
||||
"-v",
|
||||
f"{dataset_path}:/workspace/dataset.jsonl:ro",
|
||||
FINETUNE_IMAGE,
|
||||
"--dataset",
|
||||
"/workspace/dataset.jsonl",
|
||||
"--output-dir",
|
||||
"/workspace/output/qwen-bill-summarizer",
|
||||
]
|
||||
|
||||
logger.info("Starting fine-tuning container")
|
||||
logger.info(" Dataset: %s", dataset_path)
|
||||
logger.info(" Output: %s", output_dir)
|
||||
|
||||
result = subprocess.run(command, text=True, check=False)
|
||||
if result.returncode != 0:
|
||||
message = f"Fine-tuning container exited with code {result.returncode}"
|
||||
raise RuntimeError(message)
|
||||
logger.info("Fine-tuning complete. Model saved to %s", output_dir)
|
||||
|
||||
|
||||
def stop_finetune() -> None:
|
||||
"""Stop and remove the fine-tuning container."""
|
||||
logger.info("Stopping fine-tuning container")
|
||||
subprocess.run(["docker", "stop", CONTAINER_NAME], capture_output=True, check=False)
|
||||
subprocess.run(["docker", "rm", "-f", CONTAINER_NAME], capture_output=True, check=False)
|
||||
|
||||
|
||||
def logs_finetune() -> str | None:
|
||||
"""Return recent logs from the fine-tuning container, or None if not running."""
|
||||
result = subprocess.run(
|
||||
["docker", "logs", "--tail", "50", CONTAINER_NAME],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
check=False,
|
||||
)
|
||||
if result.returncode != 0:
|
||||
return None
|
||||
return result.stdout + result.stderr
|
||||
|
||||
|
||||
app = typer.Typer(help="Fine-tuning container management.")
|
||||
|
||||
|
||||
@app.command()
|
||||
def build() -> None:
|
||||
"""Build the fine-tuning Docker image."""
|
||||
build_image()
|
||||
|
||||
|
||||
@app.command()
|
||||
def run(
|
||||
dataset: Annotated[Path, typer.Option(help="Fine-tuning JSONL")] = Path(
|
||||
"/home/richie/dotfiles/data/finetune_dataset.jsonl"
|
||||
),
|
||||
output_dir: Annotated[Path, typer.Option(help="Where to save the trained model")] = Path(
|
||||
"/home/richie/dotfiles/data/output/qwen-bill-summarizer",
|
||||
),
|
||||
hf_cache: Annotated[Path, typer.Option(help="Host path to HuggingFace model cache")] = DEFAULT_HF_CACHE,
|
||||
log_level: Annotated[str, typer.Option(help="Log level")] = "INFO",
|
||||
) -> None:
|
||||
"""Run fine-tuning inside a Docker container."""
|
||||
logging.basicConfig(level=log_level, format="%(asctime)s %(levelname)s %(name)s: %(message)s")
|
||||
check_gpu_free()
|
||||
start_finetune(
|
||||
dataset_path=dataset,
|
||||
output_dir=output_dir,
|
||||
hf_cache=hf_cache,
|
||||
)
|
||||
|
||||
@app.command()
|
||||
def stop() -> None:
|
||||
"""Stop and remove the fine-tuning container."""
|
||||
stop_finetune()
|
||||
|
||||
|
||||
@app.command()
|
||||
def logs() -> None:
|
||||
"""Show recent logs from the fine-tuning container."""
|
||||
output = logs_finetune()
|
||||
if output is None:
|
||||
typer.echo("No running fine-tuning container found.")
|
||||
raise typer.Exit(code=1)
|
||||
typer.echo(output)
|
||||
|
||||
|
||||
def cli() -> None:
|
||||
"""Typer entry point."""
|
||||
app()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
cli()
|
||||
@@ -1,23 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import subprocess
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def check_gpu_free() -> None:
|
||||
"""Warn if GPU-heavy processes (e.g. Ollama) are running."""
|
||||
result = subprocess.run(
|
||||
["nvidia-smi", "--query-compute-apps=pid,process_name", "--format=csv,noheader"],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
check=False,
|
||||
)
|
||||
if result.returncode != 0:
|
||||
logger.warning("Could not query GPU processes: %s", result.stderr.strip())
|
||||
return
|
||||
processes = result.stdout.strip()
|
||||
if processes:
|
||||
logger.warning("GPU processes detected:\n%s", processes)
|
||||
logger.warning("Consider stopping Ollama (sudo systemctl stop ollama) before benchmarking")
|
||||
@@ -1,70 +0,0 @@
|
||||
"""Docker container lifecycle management for vLLM."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import subprocess
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
CONTAINER_NAME = "vllm-bench"
|
||||
VLLM_IMAGE = "vllm/vllm-openai:v0.19.0"
|
||||
|
||||
|
||||
def start_vllm(
|
||||
*,
|
||||
model: str,
|
||||
port: int,
|
||||
model_dir: str,
|
||||
gpu_memory_utilization: float,
|
||||
) -> None:
|
||||
"""Start a vLLM container serving the given model.
|
||||
|
||||
Args:
|
||||
model: HuggingFace model directory name (relative to model_dir).
|
||||
port: Host port to bind.
|
||||
model_dir: Host path containing HuggingFace model directories.
|
||||
gpu_memory_utilization: Fraction of GPU memory to use (0-1).
|
||||
"""
|
||||
command = [
|
||||
"docker",
|
||||
"run",
|
||||
"-d",
|
||||
"--name",
|
||||
CONTAINER_NAME,
|
||||
"--device=nvidia.com/gpu=all",
|
||||
"--ipc=host",
|
||||
"-v",
|
||||
f"{model_dir}:/models",
|
||||
"-p",
|
||||
f"{port}:8000",
|
||||
VLLM_IMAGE,
|
||||
"--model",
|
||||
f"/models/{model}",
|
||||
"--served-model-name",
|
||||
model,
|
||||
"--gpu-memory-utilization",
|
||||
str(gpu_memory_utilization),
|
||||
"--max-model-len",
|
||||
"4096",
|
||||
]
|
||||
logger.info("Starting vLLM container with model: %s", model)
|
||||
stop_vllm()
|
||||
result = subprocess.run(command, capture_output=True, text=True, check=False)
|
||||
if result.returncode != 0:
|
||||
msg = f"Failed to start vLLM container: {result.stderr.strip()}"
|
||||
raise RuntimeError(msg)
|
||||
logger.info("vLLM container started: %s", result.stdout.strip()[:12])
|
||||
|
||||
|
||||
def stop_vllm() -> None:
|
||||
"""Stop and remove the vLLM benchmark container."""
|
||||
logger.info("Stopping vLLM container")
|
||||
subprocess.run(["docker", "stop", CONTAINER_NAME], capture_output=True, check=False)
|
||||
subprocess.run(["docker", "rm", "-f", CONTAINER_NAME], capture_output=True, check=False)
|
||||
subprocess.run(
|
||||
["docker", "network", "disconnect", "-f", "bridge", CONTAINER_NAME],
|
||||
capture_output=True,
|
||||
check=False,
|
||||
)
|
||||
logger.info("vLLM container stopped and removed")
|
||||
@@ -1,75 +0,0 @@
|
||||
"""HuggingFace model downloader."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from pathlib import Path
|
||||
from typing import Annotated
|
||||
|
||||
import typer
|
||||
from huggingface_hub import snapshot_download
|
||||
|
||||
from python.prompt_bench.models import BenchmarkConfig
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def local_model_path(repo: str, model_dir: str) -> Path:
|
||||
"""Return the local directory path for a HuggingFace repo."""
|
||||
return Path(model_dir) / repo
|
||||
|
||||
|
||||
def is_model_present(repo: str, model_dir: str) -> bool:
|
||||
"""Check if a model has already been downloaded."""
|
||||
path = local_model_path(repo, model_dir)
|
||||
return path.exists() and any(path.iterdir())
|
||||
|
||||
|
||||
def download_model(repo: str, model_dir: str) -> Path:
|
||||
"""Download a HuggingFace model to the local model directory.
|
||||
|
||||
Skips the download if the model directory already exists and contains files.
|
||||
"""
|
||||
local_path = local_model_path(repo, model_dir)
|
||||
|
||||
if is_model_present(repo, model_dir):
|
||||
logger.info("Model already exists: %s", local_path)
|
||||
return local_path
|
||||
|
||||
logger.info("Downloading model: %s -> %s", repo, local_path)
|
||||
snapshot_download(
|
||||
repo_id=repo,
|
||||
local_dir=str(local_path),
|
||||
)
|
||||
logger.info("Download complete: %s", repo)
|
||||
return local_path
|
||||
|
||||
|
||||
def download_all(config: BenchmarkConfig) -> None:
|
||||
"""Download every model listed in the config, top to bottom."""
|
||||
for repo in config.models:
|
||||
download_model(repo, config.model_dir)
|
||||
|
||||
|
||||
def main(
|
||||
config: Annotated[Path, typer.Option(help="Path to TOML config file")] = Path("bench.toml"),
|
||||
log_level: Annotated[str, typer.Option(help="Log level")] = "INFO",
|
||||
) -> None:
|
||||
"""Download all models listed in the benchmark config."""
|
||||
logging.basicConfig(level=log_level, format="%(asctime)s %(levelname)s %(name)s: %(message)s")
|
||||
|
||||
if not config.is_file():
|
||||
message = f"Config file does not exist: {config}"
|
||||
raise typer.BadParameter(message)
|
||||
|
||||
benchmark_config = BenchmarkConfig.from_toml(config)
|
||||
download_all(benchmark_config)
|
||||
|
||||
|
||||
def cli() -> None:
|
||||
"""Typer entry point."""
|
||||
typer.run(main)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
cli()
|
||||
@@ -1,214 +0,0 @@
|
||||
"""Fine-tune Qwen 3.5 4B on bill summarization data using Unsloth.
|
||||
|
||||
Loads a ChatML-style JSONL dataset (system/user/assistant messages),
|
||||
applies QLoRA with 4-bit quantization, and saves the merged model
|
||||
in HuggingFace format. Designed for a single RTX 3090 (24GB).
|
||||
|
||||
Usage:
|
||||
python -m python.prompt_bench.finetune \
|
||||
--dataset output/finetune_dataset.jsonl \
|
||||
--output-dir output/qwen-bill-summarizer
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Annotated
|
||||
|
||||
import tomllib
|
||||
import typer
|
||||
from unsloth import FastLanguageModel
|
||||
from datasets import Dataset
|
||||
from transformers import TrainingArguments
|
||||
from trl import SFTTrainer
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class LoraConfig:
|
||||
"""LoRA adapter hyperparameters."""
|
||||
|
||||
rank: int
|
||||
alpha: int
|
||||
dropout: float
|
||||
targets: list[str]
|
||||
|
||||
|
||||
@dataclass
|
||||
class TrainingConfig:
|
||||
"""Training loop hyperparameters."""
|
||||
|
||||
learning_rate: float
|
||||
epochs: int
|
||||
batch_size: int
|
||||
gradient_accumulation: int
|
||||
max_seq_length: int
|
||||
warmup_ratio: float
|
||||
weight_decay: float
|
||||
logging_steps: int
|
||||
save_steps: int
|
||||
|
||||
|
||||
@dataclass
|
||||
class FinetuneConfig:
|
||||
"""Top-level finetune configuration."""
|
||||
|
||||
base_model: str
|
||||
lora: LoraConfig
|
||||
training: TrainingConfig
|
||||
|
||||
@classmethod
|
||||
def from_toml(cls, config_path: Path) -> FinetuneConfig:
|
||||
"""Load finetune config from a TOML file."""
|
||||
raw = tomllib.loads(config_path.read_text())["finetune"]
|
||||
return cls(
|
||||
base_model=raw["base_model"],
|
||||
lora=LoraConfig(**raw["lora"]),
|
||||
training=TrainingConfig(**raw["training"]),
|
||||
)
|
||||
|
||||
|
||||
def _messages_to_chatml(messages: list[dict]) -> str:
|
||||
r"""Convert a message list to Qwen ChatML format.
|
||||
|
||||
Produces:
|
||||
<|im_start|>system\n...\n<|im_end|>
|
||||
<|im_start|>user\n...\n<|im_end|>
|
||||
<|im_start|>assistant\n...\n<|im_end|>
|
||||
"""
|
||||
parts = []
|
||||
for message in messages:
|
||||
role = message["role"]
|
||||
content = message["content"]
|
||||
parts.append(f"<|im_start|>{role}\n{content}<|im_end|>")
|
||||
return "\n".join(parts)
|
||||
|
||||
|
||||
def load_dataset_from_jsonl(path: Path) -> Dataset:
|
||||
"""Load a ChatML JSONL file into a HuggingFace Dataset.
|
||||
|
||||
Each line must have {"messages": [{"role": ..., "content": ...}, ...]}.
|
||||
Pre-formats into a `text` column with the Qwen ChatML template applied,
|
||||
which SFTTrainer consumes directly.
|
||||
"""
|
||||
records = []
|
||||
with path.open(encoding="utf-8") as handle:
|
||||
for raw_line in handle:
|
||||
stripped = raw_line.strip()
|
||||
if stripped:
|
||||
entry = json.loads(stripped)
|
||||
records.append({"text": _messages_to_chatml(entry["messages"])})
|
||||
logger.info("Loaded %d examples from %s", len(records), path)
|
||||
return Dataset.from_list(records)
|
||||
|
||||
|
||||
def main(
|
||||
dataset_path: Annotated[Path, typer.Option("--dataset", help="Fine-tuning JSONL")] = Path(
|
||||
"output/finetune_dataset.jsonl",
|
||||
),
|
||||
validation_split: Annotated[float, typer.Option("--val-split", help="Fraction held out for validation")] = 0.1,
|
||||
output_dir: Annotated[Path, typer.Option("--output-dir", help="Where to save the merged model")] = Path(
|
||||
"output/qwen-bill-summarizer",
|
||||
),
|
||||
config_path: Annotated[
|
||||
Path,
|
||||
typer.Option("--config", help="TOML config file"),
|
||||
] = Path(__file__).parent / "config.toml",
|
||||
save_gguf: Annotated[bool, typer.Option("--save-gguf/--no-save-gguf", help="Also save GGUF")] = False,
|
||||
) -> None:
|
||||
"""Fine-tune Qwen 3.5 4B on bill summarization with Unsloth + QLoRA."""
|
||||
logging.basicConfig(level="INFO", format="%(asctime)s %(levelname)s %(name)s: %(message)s")
|
||||
|
||||
if not dataset_path.is_file():
|
||||
message = f"Dataset not found: {dataset_path}"
|
||||
raise typer.BadParameter(message)
|
||||
|
||||
config = FinetuneConfig.from_toml(config_path)
|
||||
|
||||
logger.info("Loading base model: %s", config.base_model)
|
||||
model, tokenizer = FastLanguageModel.from_pretrained(
|
||||
model_name=config.base_model,
|
||||
max_seq_length=config.training.max_seq_length,
|
||||
load_in_4bit=True,
|
||||
dtype=None,
|
||||
)
|
||||
|
||||
logger.info("Applying LoRA (rank=%d, alpha=%d)", config.lora.rank, config.lora.alpha)
|
||||
model = FastLanguageModel.get_peft_model(
|
||||
model,
|
||||
r=config.lora.rank,
|
||||
lora_alpha=config.lora.alpha,
|
||||
lora_dropout=config.lora.dropout,
|
||||
target_modules=config.lora.targets,
|
||||
bias="none",
|
||||
use_gradient_checkpointing="unsloth",
|
||||
random_state=42,
|
||||
)
|
||||
|
||||
full_dataset = load_dataset_from_jsonl(dataset_path)
|
||||
split = full_dataset.train_test_split(test_size=validation_split, seed=42)
|
||||
train_dataset = split["train"]
|
||||
validation_dataset = split["test"]
|
||||
logger.info("Split: %d train, %d validation", len(train_dataset), len(validation_dataset))
|
||||
training_args = TrainingArguments(
|
||||
output_dir=str(output_dir / "checkpoints"),
|
||||
num_train_epochs=config.training.epochs,
|
||||
per_device_train_batch_size=config.training.batch_size,
|
||||
gradient_accumulation_steps=config.training.gradient_accumulation,
|
||||
learning_rate=config.training.learning_rate,
|
||||
warmup_ratio=config.training.warmup_ratio,
|
||||
weight_decay=config.training.weight_decay,
|
||||
lr_scheduler_type="cosine",
|
||||
logging_steps=config.training.logging_steps,
|
||||
save_steps=config.training.save_steps,
|
||||
save_total_limit=3,
|
||||
eval_strategy="steps",
|
||||
eval_steps=config.training.save_steps,
|
||||
load_best_model_at_end=True,
|
||||
bf16=True,
|
||||
optim="adamw_8bit",
|
||||
seed=42,
|
||||
report_to="none",
|
||||
)
|
||||
|
||||
trainer = SFTTrainer(
|
||||
model=model,
|
||||
tokenizer=tokenizer,
|
||||
train_dataset=train_dataset,
|
||||
eval_dataset=validation_dataset,
|
||||
args=training_args,
|
||||
max_seq_length=config.training.max_seq_length,
|
||||
packing=True,
|
||||
)
|
||||
|
||||
logger.info(
|
||||
"Starting training: %d train, %d val, %d epochs",
|
||||
len(train_dataset),
|
||||
len(validation_dataset),
|
||||
config.training.epochs,
|
||||
)
|
||||
trainer.train()
|
||||
|
||||
merged_path = str(output_dir / "merged")
|
||||
logger.info("Saving merged model to %s", merged_path)
|
||||
model.save_pretrained_merged(merged_path, tokenizer, save_method="merged_16bit")
|
||||
|
||||
if save_gguf:
|
||||
gguf_path = str(output_dir / "gguf")
|
||||
logger.info("Saving GGUF to %s", gguf_path)
|
||||
model.save_pretrained_gguf(gguf_path, tokenizer, quantization_method="q4_k_m")
|
||||
|
||||
logger.info("Done! Model saved to %s", output_dir)
|
||||
|
||||
|
||||
def cli() -> None:
|
||||
"""Typer entry point."""
|
||||
typer.run(main)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
cli()
|
||||
@@ -1,215 +0,0 @@
|
||||
"""CLI entry point for the prompt benchmarking system."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import time
|
||||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||
from pathlib import Path
|
||||
from typing import Annotated
|
||||
|
||||
import typer
|
||||
|
||||
from python.prompt_bench.containers.lib import check_gpu_free
|
||||
from python.prompt_bench.containers.vllm import start_vllm, stop_vllm
|
||||
from python.prompt_bench.downloader import is_model_present
|
||||
from python.prompt_bench.models import BenchmarkConfig
|
||||
from python.prompt_bench.vllm_client import VLLMClient
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def discover_prompts(input_dir: Path) -> list[Path]:
|
||||
"""Find all .txt files in the input directory."""
|
||||
prompts = list(input_dir.glob("*.txt"))
|
||||
if not prompts:
|
||||
message = f"No .txt files found in {input_dir}"
|
||||
raise FileNotFoundError(message)
|
||||
return prompts
|
||||
|
||||
|
||||
def _run_prompt(
|
||||
client: VLLMClient,
|
||||
prompt_path: Path,
|
||||
*,
|
||||
repo: str,
|
||||
model_dir_name: str,
|
||||
model_output: Path,
|
||||
temperature: float,
|
||||
) -> tuple[bool, float]:
|
||||
"""Run a single prompt. Returns (success, elapsed_seconds)."""
|
||||
filename = prompt_path.name
|
||||
output_path = model_output / filename
|
||||
start = time.monotonic()
|
||||
try:
|
||||
prompt_text = prompt_path.read_text()
|
||||
response = client.complete(prompt_text, model_dir_name, temperature=temperature)
|
||||
output_path.write_text(response)
|
||||
elapsed = time.monotonic() - start
|
||||
logger.info("Completed: %s / %s in %.2fs", repo, filename, elapsed)
|
||||
except Exception:
|
||||
elapsed = time.monotonic() - start
|
||||
error_path = model_output / f"{filename}.error"
|
||||
logger.exception("Failed: %s / %s after %.2fs", repo, filename, elapsed)
|
||||
error_path.write_text(f"Error processing {filename}")
|
||||
return False, elapsed
|
||||
return True, elapsed
|
||||
|
||||
|
||||
def benchmark_model(
|
||||
client: VLLMClient,
|
||||
prompts: list[Path],
|
||||
*,
|
||||
repo: str,
|
||||
model_dir_name: str,
|
||||
model_output: Path,
|
||||
temperature: float,
|
||||
concurrency: int,
|
||||
) -> tuple[int, int]:
|
||||
"""Run all prompts against a single model in parallel.
|
||||
|
||||
vLLM batches concurrent requests internally, so submitting many at once is
|
||||
significantly faster than running them serially.
|
||||
"""
|
||||
pending = [prompt for prompt in prompts if not (model_output / prompt.name).exists()]
|
||||
skipped = len(prompts) - len(pending)
|
||||
if skipped:
|
||||
logger.info("Skipping %d prompts with existing output for %s", skipped, repo)
|
||||
|
||||
if not pending:
|
||||
logger.info("Nothing to do for %s", repo)
|
||||
return 0, 0
|
||||
|
||||
completed = 0
|
||||
failed = 0
|
||||
latencies: list[float] = []
|
||||
|
||||
wall_start = time.monotonic()
|
||||
with ThreadPoolExecutor(max_workers=concurrency) as executor:
|
||||
futures = [
|
||||
executor.submit(
|
||||
_run_prompt,
|
||||
client,
|
||||
prompt_path,
|
||||
repo=repo,
|
||||
model_dir_name=model_dir_name,
|
||||
model_output=model_output,
|
||||
temperature=temperature,
|
||||
)
|
||||
for prompt_path in pending
|
||||
]
|
||||
for future in as_completed(futures):
|
||||
success, elapsed = future.result()
|
||||
latencies.append(elapsed)
|
||||
if success:
|
||||
completed += 1
|
||||
else:
|
||||
failed += 1
|
||||
wall_elapsed = time.monotonic() - wall_start
|
||||
|
||||
attempted = completed + failed
|
||||
avg_latency = sum(latencies) / attempted
|
||||
throughput = attempted / wall_elapsed if wall_elapsed > 0 else 0.0
|
||||
timing = {
|
||||
"repo": repo,
|
||||
"wall_seconds": wall_elapsed,
|
||||
"attempted": attempted,
|
||||
"completed": completed,
|
||||
"failed": failed,
|
||||
"avg_latency_seconds": avg_latency,
|
||||
"throughput_prompts_per_second": throughput,
|
||||
"concurrency": concurrency,
|
||||
}
|
||||
timing_path = model_output / "_timing.json"
|
||||
timing_path.write_text(json.dumps(timing, indent=2))
|
||||
|
||||
return completed, failed
|
||||
|
||||
|
||||
def run_benchmark(
|
||||
config: BenchmarkConfig,
|
||||
input_dir: Path,
|
||||
output_dir: Path,
|
||||
) -> None:
|
||||
"""Execute the benchmark across all models and prompts."""
|
||||
prompts = discover_prompts(input_dir)
|
||||
logger.info("Found %d prompts in %s", len(prompts), input_dir)
|
||||
|
||||
check_gpu_free()
|
||||
|
||||
total_completed = 0
|
||||
total_failed = 0
|
||||
|
||||
for repo in config.models:
|
||||
if not is_model_present(repo, config.model_dir):
|
||||
logger.warning("Skipping (not downloaded): %s", repo)
|
||||
continue
|
||||
|
||||
model_output = output_dir / repo
|
||||
model_output.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
logger.info("=== Benchmarking model: %s ===", repo)
|
||||
|
||||
stop_vllm()
|
||||
try:
|
||||
start_vllm(
|
||||
model=repo,
|
||||
port=config.port,
|
||||
model_dir=config.model_dir,
|
||||
gpu_memory_utilization=config.gpu_memory_utilization,
|
||||
)
|
||||
except RuntimeError:
|
||||
logger.exception("Failed to start vLLM for %s, skipping", repo)
|
||||
continue
|
||||
logger.info("vLLM started for %s", repo)
|
||||
try:
|
||||
with VLLMClient(port=config.port, timeout=config.timeout) as client:
|
||||
client.wait_ready(max_wait=config.vllm_startup_timeout)
|
||||
completed, failed = benchmark_model(
|
||||
client,
|
||||
prompts,
|
||||
repo=repo,
|
||||
model_dir_name=repo,
|
||||
model_output=model_output,
|
||||
temperature=config.temperature,
|
||||
concurrency=config.concurrency,
|
||||
)
|
||||
total_completed += completed
|
||||
total_failed += failed
|
||||
finally:
|
||||
stop_vllm()
|
||||
|
||||
logger.info("=== Benchmark complete ===")
|
||||
logger.info("Completed: %d | Failed: %d", total_completed, total_failed)
|
||||
|
||||
|
||||
def main(
|
||||
input_dir: Annotated[Path, typer.Argument(help="Directory containing input .txt prompt files")],
|
||||
config: Annotated[Path, typer.Option(help="Path to TOML config file")] = Path("bench.toml"),
|
||||
output_dir: Annotated[Path, typer.Option(help="Output directory for results")] = Path("output"),
|
||||
log_level: Annotated[str, typer.Option(help="Log level")] = "INFO",
|
||||
) -> None:
|
||||
"""Run prompts through multiple LLMs via vLLM and save results."""
|
||||
logging.basicConfig(level=log_level, format="%(asctime)s %(levelname)s %(name)s: %(message)s")
|
||||
|
||||
if not input_dir.is_dir():
|
||||
message = f"Input directory does not exist: {input_dir}"
|
||||
raise typer.BadParameter(message)
|
||||
if not config.is_file():
|
||||
message = f"Config file does not exist: {config}"
|
||||
raise typer.BadParameter(message)
|
||||
|
||||
benchmark_config = BenchmarkConfig.from_toml(config)
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
run_benchmark(benchmark_config, input_dir, output_dir)
|
||||
|
||||
|
||||
def cli() -> None:
|
||||
"""Typer entry point."""
|
||||
typer.run(main)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
cli()
|
||||
@@ -1,30 +0,0 @@
|
||||
"""Pydantic models for benchmark configuration."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import tomllib
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
class BenchmarkConfig(BaseModel):
|
||||
"""Top-level benchmark configuration loaded from TOML."""
|
||||
|
||||
models: list[str]
|
||||
model_dir: str = "/zfs/models/hf"
|
||||
port: int = 8000
|
||||
gpu_memory_utilization: float = 0.90
|
||||
temperature: float = 0.0
|
||||
timeout: int = 300
|
||||
concurrency: int = 4
|
||||
vllm_startup_timeout: int = 900
|
||||
|
||||
@classmethod
|
||||
def from_toml(cls, config_path: Path) -> BenchmarkConfig:
|
||||
"""Load benchmark config from a TOML file."""
|
||||
raw = tomllib.loads(config_path.read_text())["bench"]
|
||||
return cls(**raw)
|
||||
@@ -1,34 +0,0 @@
|
||||
SUMMARIZATION_SYSTEM_PROMPT = """You are a legislative analyst extracting policy substance from Congressional bill text.
|
||||
|
||||
Your job is to compress a bill into a dense, neutral structured summary that captures every distinct policy action — including secondary effects that might be buried in subsections.
|
||||
|
||||
EXTRACTION RULES:
|
||||
- IGNORE: whereas clauses, congressional findings that are purely political statements, recitals, preambles, citations of existing law by number alone, and procedural boilerplate.
|
||||
- FOCUS ON: operative verbs — what the bill SHALL do, PROHIBIT, REQUIRE, AUTHORIZE, AMEND, APPROPRIATE, or ESTABLISH.
|
||||
- SURFACE ALL THREADS: If the bill touches multiple policy areas, list each thread separately. Do not collapse them.
|
||||
- BE CONCRETE: Name the affected population, the mechanism, and the direction (expands/restricts/maintains).
|
||||
- STAY NEUTRAL: No political framing. Describe what the text does, not what its sponsors claim it does.
|
||||
|
||||
OUTPUT FORMAT — plain structured text, not JSON:
|
||||
|
||||
OPERATIVE ACTIONS:
|
||||
[Numbered list of what the bill actually does, one action per line, max 20 words each]
|
||||
|
||||
AFFECTED POPULATIONS:
|
||||
[Who gains something, who loses something, or whose behavior is regulated]
|
||||
|
||||
MECHANISMS:
|
||||
[How it works: new funding, mandate, prohibition, amendment to existing statute, grant program, study commission, etc.]
|
||||
|
||||
POLICY THREADS:
|
||||
[List each distinct policy domain this bill touches, even minor ones. Use plain language, not domain codes.]
|
||||
|
||||
SYMBOLIC/PROCEDURAL ONLY:
|
||||
[Yes or No — is this bill primarily a resolution, designation, or awareness declaration with no operative effect?]
|
||||
|
||||
LENGTH TARGET: 150-250 words total. Be ruthless about cutting. Density over completeness."""
|
||||
|
||||
SUMMARIZATION_USER_TEMPLATE = """Summarize the following Congressional bill according to your instructions.
|
||||
|
||||
BILL TEXT:
|
||||
{text_content}"""
|
||||
@@ -1,114 +0,0 @@
|
||||
"""Build a fine-tuning JSONL dataset from batch request + output files.
|
||||
|
||||
Joins the original request JSONL (system + user messages) with the batch
|
||||
output JSONL (assistant completions) by custom_id to produce a ChatML-style
|
||||
messages JSONL suitable for fine-tuning.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
from pathlib import Path
|
||||
from typing import Annotated
|
||||
|
||||
import typer
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
HTTP_OK = 200
|
||||
|
||||
|
||||
def load_requests(path: Path) -> dict[str, list[dict]]:
|
||||
"""Parse request JSONL into {custom_id: messages}."""
|
||||
results: dict[str, list[dict]] = {}
|
||||
with path.open(encoding="utf-8") as handle:
|
||||
for raw_line in handle:
|
||||
stripped = raw_line.strip()
|
||||
if not stripped:
|
||||
continue
|
||||
record = json.loads(stripped)
|
||||
custom_id = record["custom_id"]
|
||||
messages = record["body"]["messages"]
|
||||
results[custom_id] = messages
|
||||
return results
|
||||
|
||||
|
||||
def load_completions(path: Path) -> dict[str, str]:
|
||||
"""Parse batch output JSONL into {custom_id: assistant_content}."""
|
||||
results: dict[str, str] = {}
|
||||
with path.open(encoding="utf-8") as handle:
|
||||
for line_number, raw_line in enumerate(handle, 1):
|
||||
stripped = raw_line.strip()
|
||||
if not stripped:
|
||||
continue
|
||||
record = json.loads(stripped)
|
||||
custom_id = record["custom_id"]
|
||||
response = record.get("response", {})
|
||||
if response.get("status_code") != HTTP_OK:
|
||||
logger.warning("Skipping %s (line %d): status %s", custom_id, line_number, response.get("status_code"))
|
||||
continue
|
||||
body = response.get("body", {})
|
||||
choices = body.get("choices", [])
|
||||
if not choices:
|
||||
logger.warning("Skipping %s (line %d): no choices", custom_id, line_number)
|
||||
continue
|
||||
content = choices[0].get("message", {}).get("content", "")
|
||||
if not content:
|
||||
logger.warning("Skipping %s (line %d): empty content", custom_id, line_number)
|
||||
continue
|
||||
results[custom_id] = content
|
||||
return results
|
||||
|
||||
|
||||
def main(
|
||||
requests_path: Annotated[Path, typer.Option("--requests", help="Batch request JSONL")] = Path(
|
||||
"output/openai_batch/requests.jsonl",
|
||||
),
|
||||
batch_output: Annotated[Path, typer.Option("--batch-output", help="Batch output JSONL")] = Path(
|
||||
"batch_69d84558d91c819091d53f08d78f9fd6_output.jsonl",
|
||||
),
|
||||
output_path: Annotated[Path, typer.Option("--output", help="Fine-tuning JSONL output")] = Path(
|
||||
"output/finetune_dataset.jsonl",
|
||||
),
|
||||
log_level: Annotated[str, typer.Option(help="Log level")] = "INFO",
|
||||
) -> None:
|
||||
"""Build fine-tuning dataset by joining request and output JSONL files."""
|
||||
logging.basicConfig(level=log_level, format="%(asctime)s %(levelname)s %(name)s: %(message)s")
|
||||
|
||||
logger.info("Loading requests from %s", requests_path)
|
||||
requests = load_requests(requests_path)
|
||||
logger.info("Loaded %d requests", len(requests))
|
||||
|
||||
logger.info("Loading completions from %s", batch_output)
|
||||
completions = load_completions(batch_output)
|
||||
logger.info("Loaded %d completions", len(completions))
|
||||
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
matched = 0
|
||||
skipped = 0
|
||||
|
||||
with output_path.open("w", encoding="utf-8") as handle:
|
||||
for custom_id, messages in requests.items():
|
||||
assistant_content = completions.get(custom_id)
|
||||
if assistant_content is None:
|
||||
skipped += 1
|
||||
continue
|
||||
|
||||
example = {
|
||||
"messages": [*messages, {"role": "assistant", "content": assistant_content}],
|
||||
}
|
||||
handle.write(json.dumps(example, ensure_ascii=False))
|
||||
handle.write("\n")
|
||||
matched += 1
|
||||
|
||||
logger.info("Wrote %d examples to %s (skipped %d unmatched)", matched, output_path, skipped)
|
||||
|
||||
|
||||
def cli() -> None:
|
||||
"""Typer entry point."""
|
||||
typer.run(main)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
cli()
|
||||
@@ -1,97 +0,0 @@
|
||||
"""Sum token usage across compressed and uncompressed run directories."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import Annotated
|
||||
|
||||
import typer
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class UsageTotals:
|
||||
"""Aggregate usage counters for a directory of run records."""
|
||||
|
||||
files: int = 0
|
||||
errors: int = 0
|
||||
prompt_tokens: int = 0
|
||||
cached_tokens: int = 0
|
||||
completion_tokens: int = 0
|
||||
reasoning_tokens: int = 0
|
||||
total_tokens: int = 0
|
||||
per_file: list[tuple[str, int, int, int]] = field(default_factory=list)
|
||||
|
||||
|
||||
def tally_directory(directory: Path) -> UsageTotals:
|
||||
"""Return aggregated usage stats for every JSON record in a directory."""
|
||||
totals = UsageTotals()
|
||||
decoder = json.JSONDecoder()
|
||||
for path in sorted(directory.glob("*.json")):
|
||||
text = path.read_text().lstrip()
|
||||
record, _ = decoder.raw_decode(text)
|
||||
totals.files += 1
|
||||
usage = record.get("usage")
|
||||
if not usage:
|
||||
totals.errors += 1
|
||||
continue
|
||||
prompt_tokens = usage.get("prompt_tokens", 0)
|
||||
completion_tokens = usage.get("completion_tokens", 0)
|
||||
total_tokens = usage.get("total_tokens", 0)
|
||||
cached_tokens = (usage.get("prompt_tokens_details") or {}).get("cached_tokens", 0)
|
||||
reasoning_tokens = (usage.get("completion_tokens_details") or {}).get("reasoning_tokens", 0)
|
||||
totals.prompt_tokens += prompt_tokens
|
||||
totals.completion_tokens += completion_tokens
|
||||
totals.total_tokens += total_tokens
|
||||
totals.cached_tokens += cached_tokens
|
||||
totals.reasoning_tokens += reasoning_tokens
|
||||
totals.per_file.append((path.name, prompt_tokens, completion_tokens, total_tokens))
|
||||
return totals
|
||||
|
||||
|
||||
def log_totals(label: str, totals: UsageTotals) -> None:
|
||||
"""Log a one-block summary for a directory."""
|
||||
counted = totals.files - totals.errors
|
||||
average_total = totals.total_tokens / counted if counted else 0
|
||||
logger.info("[%s]", label)
|
||||
logger.info(" files : %d (with usage: %d, errors: %d)", totals.files, counted, totals.errors)
|
||||
logger.info(" prompt tokens : %d", totals.prompt_tokens)
|
||||
logger.info(" cached tokens : %d", totals.cached_tokens)
|
||||
logger.info(" completion tok : %d", totals.completion_tokens)
|
||||
logger.info(" reasoning tok : %d", totals.reasoning_tokens)
|
||||
logger.info(" total tokens : %d", totals.total_tokens)
|
||||
logger.info(" avg total/file : %.1f", average_total)
|
||||
|
||||
|
||||
def main(
|
||||
runs_dir: Annotated[Path, typer.Option("--runs-dir")] = Path("output/openai_runs_temp_1"),
|
||||
log_level: Annotated[str, typer.Option("--log-level")] = "INFO",
|
||||
) -> None:
|
||||
"""Print token usage totals for the compressed and uncompressed run directories."""
|
||||
logging.basicConfig(level=log_level, format="%(message)s")
|
||||
|
||||
grand = UsageTotals()
|
||||
for label in ("compressed", "uncompressed"):
|
||||
directory = runs_dir / label
|
||||
if not directory.is_dir():
|
||||
logger.warning("%s: directory not found at %s", label, directory)
|
||||
continue
|
||||
totals = tally_directory(directory)
|
||||
log_totals(label, totals)
|
||||
grand.files += totals.files
|
||||
grand.errors += totals.errors
|
||||
grand.prompt_tokens += totals.prompt_tokens
|
||||
grand.cached_tokens += totals.cached_tokens
|
||||
grand.completion_tokens += totals.completion_tokens
|
||||
grand.reasoning_tokens += totals.reasoning_tokens
|
||||
grand.total_tokens += totals.total_tokens
|
||||
|
||||
log_totals("grand total", grand)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
typer.run(main)
|
||||
@@ -1,68 +0,0 @@
|
||||
"""OpenAI-compatible client for vLLM's API."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import time
|
||||
from typing import Self
|
||||
|
||||
import httpx
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
READY_POLL_INTERVAL = 2.0
|
||||
|
||||
|
||||
class VLLMClient:
|
||||
"""Talk to a vLLM server via its OpenAI-compatible API.
|
||||
|
||||
Args:
|
||||
host: vLLM host.
|
||||
port: vLLM port.
|
||||
timeout: Per-request timeout in seconds.
|
||||
"""
|
||||
|
||||
def __init__(self, *, host: str = "localhost", port: int = 8000, timeout: int = 300) -> None:
|
||||
"""Create a client connected to a vLLM server."""
|
||||
self._client = httpx.Client(base_url=f"http://{host}:{port}", timeout=timeout)
|
||||
|
||||
def wait_ready(self, max_wait: int) -> None:
|
||||
"""Poll /v1/models until the server is ready or timeout."""
|
||||
deadline = time.monotonic() + max_wait
|
||||
while time.monotonic() < deadline:
|
||||
try:
|
||||
response = self._client.get("/v1/models")
|
||||
if response.is_success:
|
||||
logger.info("vLLM server is ready")
|
||||
return
|
||||
except httpx.TransportError:
|
||||
pass
|
||||
time.sleep(READY_POLL_INTERVAL)
|
||||
msg = f"vLLM server not ready after {max_wait}s"
|
||||
raise TimeoutError(msg)
|
||||
|
||||
def complete(self, prompt: str, model: str, *, temperature: float = 0.0, max_tokens: int = 4096) -> str:
|
||||
"""Send a prompt to /v1/completions and return the response text."""
|
||||
payload = {
|
||||
"model": model,
|
||||
"prompt": prompt,
|
||||
"temperature": temperature,
|
||||
"max_tokens": max_tokens,
|
||||
}
|
||||
logger.info("Sending prompt to %s (%d chars)", model, len(prompt))
|
||||
response = self._client.post("/v1/completions", json=payload)
|
||||
response.raise_for_status()
|
||||
data = response.json()
|
||||
return data["choices"][0]["text"]
|
||||
|
||||
def close(self) -> None:
|
||||
"""Close the HTTP client."""
|
||||
self._client.close()
|
||||
|
||||
def __enter__(self) -> Self:
|
||||
"""Enter the context manager."""
|
||||
return self
|
||||
|
||||
def __exit__(self, *args: object) -> None:
|
||||
"""Close the HTTP client on exit."""
|
||||
self.close()
|
||||
File renamed without changes.
File renamed without changes.
File renamed without changes.
Loaded 100 of 186 files, more files were not shown because too many files have changed in this diff.
Show more
Reference in new issue
Block a user