Compare commits

...

124 Commits

Author SHA1 Message Date
Richie 8d43fcbc09 add .ebook_search_bm25 to gitignore
treefmt / nix fmt (pull_request) Failing after 6s
pytest / pytest (pull_request) Successful in 41s
build_systems / build-rhapsody-in-green (pull_request) Successful in 1m40s
build_systems / build-brain (pull_request) Successful in 1m1s
build_systems / build-bob (pull_request) Successful in 1m5s
build_systems / build-leviathan (pull_request) Successful in 1m35s
build_systems / build-jeeves (pull_request) Successful in 3m33s
2026-06-12 13:15:55 -04:00
Richie e936b5850c updated python 2026-06-12 13:15:55 -04:00
Richie 03c85c5ebd setup tests 2026-06-12 13:15:55 -04:00
Richie 5d8a758b89 build api and frountend 2026-06-12 13:15:55 -04:00
Richie 051b5c7dc0 added answer.py and config 2026-06-12 13:15:55 -04:00
Richie c250ff387c added __init__ 2026-06-12 13:15:55 -04:00
Richie 116f101fd8 made llm_interface.py 2026-06-12 13:15:55 -04:00
Richie ed77efa158 added rerank 2026-06-12 13:15:55 -04:00
Richie 9d1abb1f26 built ingest 2026-06-12 13:15:55 -04:00
Richie 9a284db314 built rag search setup 2026-06-12 13:15:55 -04:00
Richie 1f701cf458 set up embedding system 2026-06-12 13:15:55 -04:00
Richie d3f72379c7 built BM25 search foundation 2026-06-12 13:15:55 -04:00
Richie b47392dd52 clean up 2026-06-12 13:15:55 -04:00
Richie eee1b8f8e0 added ebook embedding to orm 2026-06-12 13:15:55 -04:00
Richie 57e98dc667 removed hedgedoc 2026-06-12 13:15:55 -04:00
Richie fde5826dc7 adding embedding Models to jeeves 2026-06-12 13:15:55 -04:00
Richie 5f08932007 updated series_index to float and added UniqueConstraint to audiobook and audiobook_author 2026-06-12 13:15:55 -04:00
Richie c25f973d4a fixed omnibus for audio books 2026-06-12 13:15:55 -04:00
Richie 62483e7894 moved installer to python dir 2026-06-12 13:15:55 -04:00
Richie ebe5b21fc4 deleted frontend dir 2026-06-12 13:15:55 -04:00
Richie eb06174152 added llm_tool_calling.py 2026-06-12 13:15:55 -04:00
Richie bdc8ec31fc built workflow 2026-06-12 13:15:55 -04:00
Richie 25b9576c5e Add catalog.py for manually adding authors and series to the database. 2026-06-12 13:15:55 -04:00
Richie d7cc46253b adding audiobook data to DB 2026-06-12 13:15:55 -04:00
Richie f77c9657a3 chore: update flake.lock
treefmt / nix fmt (pull_request) Successful in 6s
pytest / pytest (pull_request) Successful in 27s
build_systems / build-brain (pull_request) Successful in 41s
build_systems / build-bob (pull_request) Successful in 46s
build_systems / build-leviathan (pull_request) Successful in 46s
build_systems / build-rhapsody-in-green (pull_request) Successful in 52s
build_systems / build-jeeves (pull_request) Successful in 2m34s
pytest / pytest (push) Successful in 26s
build_systems / build-leviathan (push) Successful in 41s
treefmt / nix fmt (push) Successful in 6s
build_systems / build-brain (push) Successful in 31s
build_systems / build-bob (push) Successful in 32s
build_systems / build-rhapsody-in-green (push) Successful in 48s
build_systems / build-jeeves (push) Successful in 2m18s
2026-06-12 13:06:18 -04:00
Richie f908f969d3 opening port for vllm
treefmt / nix fmt (pull_request) Successful in 7s
pytest / pytest (pull_request) Successful in 35s
build_systems / build-brain (pull_request) Successful in 47s
build_systems / build-bob (pull_request) Successful in 50s
build_systems / build-leviathan (pull_request) Successful in 56s
build_systems / build-rhapsody-in-green (pull_request) Successful in 1m1s
build_systems / build-jeeves (pull_request) Successful in 2m31s
treefmt / nix fmt (push) Successful in 6s
pytest / pytest (push) Successful in 25s
build_systems / build-brain (push) Successful in 31s
build_systems / build-bob (push) Successful in 34s
build_systems / build-leviathan (push) Successful in 42s
build_systems / build-rhapsody-in-green (push) Successful in 45s
build_systems / build-jeeves (push) Successful in 2m13s
2026-06-08 19:43:07 -04:00
Richie 3cf49c5479 fixing update-flake-lock.yaml permissions
pytest / pytest (pull_request) Successful in 25s
treefmt / nix fmt (pull_request) Successful in 6s
build_systems / build-bob (pull_request) Successful in 45s
build_systems / build-brain (pull_request) Successful in 46s
build_systems / build-leviathan (pull_request) Successful in 54s
build_systems / build-rhapsody-in-green (pull_request) Successful in 59s
build_systems / build-jeeves (pull_request) Successful in 2m31s
treefmt / nix fmt (push) Successful in 6s
pytest / pytest (push) Successful in 29s
build_systems / build-brain (push) Successful in 31s
build_systems / build-bob (push) Successful in 33s
build_systems / build-leviathan (push) Successful in 42s
build_systems / build-rhapsody-in-green (push) Successful in 48s
build_systems / build-jeeves (push) Successful in 2m34s
2026-06-07 11:21:10 -04:00
Richie b34354f5e5 adding storage to bob
treefmt / nix fmt (pull_request) Successful in 6s
pytest / pytest (pull_request) Successful in 26s
build_systems / build-brain (pull_request) Successful in 47s
build_systems / build-bob (pull_request) Successful in 47s
build_systems / build-leviathan (pull_request) Successful in 53s
build_systems / build-rhapsody-in-green (pull_request) Successful in 1m0s
build_systems / build-jeeves (pull_request) Successful in 2m32s
treefmt / nix fmt (push) Successful in 7s
build_systems / build-brain (push) Successful in 8s
pytest / pytest (push) Successful in 25s
build_systems / build-bob (push) Successful in 33s
build_systems / build-leviathan (push) Successful in 41s
build_systems / build-rhapsody-in-green (push) Successful in 46s
build_systems / build-jeeves (push) Successful in 2m19s
2026-06-07 10:48:47 -04:00
Richie 44826464de flake update fro claud code
treefmt / nix fmt (pull_request) Successful in 5s
pytest / pytest (pull_request) Successful in 26s
build_systems / build-brain (pull_request) Successful in 46s
build_systems / build-bob (pull_request) Successful in 48s
build_systems / build-leviathan (pull_request) Successful in 53s
build_systems / build-rhapsody-in-green (pull_request) Successful in 1m0s
build_systems / build-jeeves (pull_request) Successful in 2m33s
treefmt / nix fmt (push) Successful in 6s
pytest / pytest (push) Successful in 25s
build_systems / build-brain (push) Successful in 31s
build_systems / build-bob (push) Successful in 34s
build_systems / build-leviathan (push) Successful in 41s
build_systems / build-rhapsody-in-green (push) Successful in 46s
build_systems / build-jeeves (push) Successful in 2m19s
2026-06-07 10:01:51 -04:00
Richie 3de0ffccb0 adding workflow dispatch for gitea_flake_lock.py
treefmt / nix fmt (pull_request) Successful in 6s
pytest / pytest (pull_request) Successful in 25s
build_systems / build-brain (pull_request) Successful in 47s
build_systems / build-bob (pull_request) Successful in 49s
build_systems / build-leviathan (pull_request) Successful in 54s
build_systems / build-rhapsody-in-green (pull_request) Successful in 1m0s
pytest / pytest (push) Successful in 29s
build_systems / build-bob (push) Successful in 32s
build_systems / build-jeeves (pull_request) Successful in 2m35s
build_systems / build-leviathan (push) Successful in 44s
treefmt / nix fmt (push) Successful in 6s
build_systems / build-rhapsody-in-green (push) Successful in 18s
build_systems / build-brain (push) Successful in 32s
build_systems / build-jeeves (push) Successful in 2m19s
2026-06-06 22:56:34 -04:00
Richie c6c98b3e26 updated Primary nic
pytest / pytest (pull_request) Successful in 26s
build_systems / build-bob (pull_request) Successful in 49s
build_systems / build-rhapsody-in-green (pull_request) Successful in 1m1s
treefmt / nix fmt (pull_request) Successful in 6s
build_systems / build-brain (pull_request) Successful in 48s
build_systems / build-leviathan (pull_request) Successful in 55s
build_systems / build-jeeves (pull_request) Successful in 2m39s
treefmt / nix fmt (push) Successful in 6s
build_systems / build-rhapsody-in-green (push) Successful in 14s
pytest / pytest (push) Successful in 25s
build_systems / build-brain (push) Successful in 29s
build_systems / build-bob (push) Successful in 33s
build_systems / build-leviathan (push) Successful in 41s
build_systems / build-jeeves (push) Successful in 2m19s
2026-06-04 18:10:41 -04:00
Richie d459f3d675 adding brave 2026-06-04 18:10:41 -04:00
Richie 33e4b37cce fixing jeeves dns 2026-06-04 18:10:41 -04:00
Richie 2a8e7e7f2b updated my ssh_config.nix
treefmt / nix fmt (pull_request) Successful in 7s
build_systems / build-brain (pull_request) Successful in 50s
build_systems / build-bob (pull_request) Successful in 50s
build_systems / build-leviathan (pull_request) Successful in 1m1s
build_systems / build-jeeves (pull_request) Successful in 2m49s
pytest / pytest (pull_request) Successful in 31s
build_systems / build-rhapsody-in-green (pull_request) Successful in 1m8s
treefmt / nix fmt (push) Successful in 5s
build_systems / build-leviathan (push) Successful in 12s
pytest / pytest (push) Successful in 24s
build_systems / build-bob (push) Successful in 32s
build_systems / build-rhapsody-in-green (push) Successful in 51s
build_systems / build-jeeves (push) Successful in 2m23s
build_systems / build-brain (push) Successful in 30s
2026-06-03 22:11:19 -04:00
Richie 07759353be flake update
treefmt / nix fmt (pull_request) Successful in 6s
pytest / pytest (pull_request) Successful in 28s
build_systems / build-leviathan (pull_request) Successful in 18m28s
build_systems / build-rhapsody-in-green (pull_request) Successful in 19m7s
build_systems / build-brain (pull_request) Successful in 19m51s
build_systems / build-jeeves (pull_request) Successful in 19m54s
build_systems / build-bob (pull_request) Successful in 30s
build_systems / build-bob (push) Successful in 30s
treefmt / nix fmt (push) Successful in 5s
build_systems / build-brain (push) Successful in 29s
pytest / pytest (push) Successful in 25s
build_systems / build-leviathan (push) Successful in 39s
build_systems / build-rhapsody-in-green (push) Successful in 46s
build_systems / build-jeeves (push) Successful in 2m19s
2026-05-29 22:30:33 -04:00
Richie 38fb14520e removed --reload
treefmt / nix fmt (pull_request) Successful in 6s
pytest / pytest (pull_request) Successful in 27s
build_systems / build-brain (pull_request) Successful in 52s
build_systems / build-bob (pull_request) Successful in 54s
build_systems / build-leviathan (pull_request) Successful in 1m4s
build_systems / build-rhapsody-in-green (pull_request) Successful in 1m5s
build_systems / build-jeeves (pull_request) Successful in 2m45s
build_systems / build-bob (push) Successful in 34s
build_systems / build-brain (push) Successful in 32s
treefmt / nix fmt (push) Successful in 6s
pytest / pytest (push) Successful in 26s
build_systems / build-leviathan (push) Successful in 43s
build_systems / build-rhapsody-in-green (push) Successful in 47s
build_systems / build-jeeves (push) Successful in 2m26s
2026-05-29 20:26:32 -04:00
Richie 006ae6079a moved nornsight off my_python 2026-05-29 20:15:51 -04:00
Richie 7d507fb7e1 adding nornsight.nix
treefmt / nix fmt (pull_request) Successful in 6s
build_systems / build-brain (pull_request) Successful in 51s
build_systems / build-bob (pull_request) Successful in 56s
pytest / pytest (pull_request) Successful in 28s
build_systems / build-leviathan (pull_request) Successful in 1m24s
build_systems / build-rhapsody-in-green (pull_request) Successful in 1m30s
build_systems / build-jeeves (pull_request) Successful in 2m45s
2026-05-29 18:39:27 -04:00
Richie 0f69022e51 disabled terminal bell
treefmt / nix fmt (pull_request) Successful in 7s
pytest / pytest (pull_request) Successful in 29s
build_systems / build-brain (pull_request) Successful in 48s
build_systems / build-bob (pull_request) Successful in 48s
build_systems / build-jeeves (pull_request) Successful in 2m42s
build_systems / build-brain (push) Successful in 30s
build_systems / build-leviathan (pull_request) Successful in 1m0s
build_systems / build-rhapsody-in-green (pull_request) Successful in 1m4s
treefmt / nix fmt (push) Successful in 6s
build_systems / build-bob (push) Successful in 33s
pytest / pytest (push) Successful in 25s
build_systems / build-leviathan (push) Successful in 41s
build_systems / build-rhapsody-in-green (push) Successful in 46s
build_systems / build-jeeves (push) Successful in 2m23s
2026-05-29 13:52:46 -04:00
Richie a260ae2470 adding ffmpeg to jeeves and rhapsody-in-green
treefmt / nix fmt (pull_request) Successful in 7s
build_systems / build-bob (pull_request) Successful in 32s
pytest / pytest (pull_request) Successful in 26s
build_systems / build-brain (pull_request) Successful in 44s
build_systems / build-leviathan (pull_request) Successful in 55s
build_systems / build-rhapsody-in-green (pull_request) Successful in 1m30s
build_systems / build-jeeves (pull_request) Successful in 2m40s
treefmt / nix fmt (push) Successful in 6s
build_systems / build-bob (push) Successful in 33s
build_systems / build-brain (push) Successful in 34s
pytest / pytest (push) Successful in 26s
build_systems / build-leviathan (push) Successful in 44s
build_systems / build-rhapsody-in-green (push) Successful in 45s
build_systems / build-jeeves (push) Successful in 2m21s
2026-05-28 22:14:59 -04:00
Richie 820b4a53d2 adding photos to syncthing
treefmt / nix fmt (pull_request) Successful in 6s
pytest / pytest (pull_request) Successful in 1m16s
build_systems / build-jeeves (pull_request) Successful in 5m29s
build_systems / build-brain (pull_request) Successful in 6m4s
build_systems / build-rhapsody-in-green (pull_request) Successful in 16m47s
build_systems / build-leviathan (pull_request) Successful in 16m49s
build_systems / build-bob (pull_request) Successful in 31s
treefmt / nix fmt (push) Successful in 6s
build_systems / build-bob (push) Successful in 31s
build_systems / build-brain (push) Successful in 32s
pytest / pytest (push) Successful in 26s
build_systems / build-leviathan (push) Successful in 40s
build_systems / build-rhapsody-in-green (push) Successful in 14s
build_systems / build-jeeves (push) Successful in 2m33s
2026-05-28 22:08:46 -04:00
Richie ea77e83f06 setting forceImportRoot to false
pytest / pytest (pull_request) Successful in 53s
treefmt / nix fmt (pull_request) Successful in 9s
build_systems / build-brain (pull_request) Successful in 2m33s
build_systems / build-bob (pull_request) Successful in 2m41s
build_systems / build-leviathan (pull_request) Successful in 3m22s
build_systems / build-rhapsody-in-green (pull_request) Successful in 3m32s
build_systems / build-jeeves (pull_request) Successful in 8m52s
build_systems / build-bob (push) Successful in 33s
treefmt / nix fmt (push) Successful in 6s
build_systems / build-brain (push) Successful in 31s
pytest / pytest (push) Successful in 26s
build_systems / build-leviathan (push) Successful in 41s
build_systems / build-rhapsody-in-green (push) Successful in 47s
build_systems / build-jeeves (push) Successful in 2m28s
2026-05-14 15:12:53 -04:00
Richie a9da208bc3 added --accept-flake-config to nixos-rebuild step
treefmt / nix fmt (pull_request) Successful in 9s
pytest / pytest (pull_request) Successful in 1m17s
build_systems / build-brain (pull_request) Successful in 2m14s
build_systems / build-bob (pull_request) Successful in 2m25s
build_systems / build-leviathan (pull_request) Successful in 4m32s
build_systems / build-rhapsody-in-green (pull_request) Successful in 4m35s
build_systems / build-jeeves (pull_request) Successful in 8m45s
pytest / pytest (push) Successful in 1m1s
treefmt / nix fmt (push) Successful in 8s
build_systems / build-bob (push) Successful in 44s
build_systems / build-leviathan (push) Successful in 38s
build_systems / build-brain (push) Successful in 1m39s
build_systems / build-rhapsody-in-green (push) Successful in 3m0s
build_systems / build-jeeves (push) Successful in 7m3s
2026-05-14 13:39:13 -04:00
Richie 739d7dd28c droped whisper from my_python 2026-05-14 13:38:41 -04:00
Richie 651599796e moved ./llm_tools.nix to gui only
treefmt / nix fmt (pull_request) Successful in 9s
pytest / pytest (pull_request) Successful in 1m24s
build_systems / build-brain (pull_request) Successful in 4m7s
build_systems / build-leviathan (pull_request) Successful in 4m11s
build_systems / build-rhapsody-in-green (pull_request) Successful in 4m41s
build_systems / build-jeeves (pull_request) Successful in 8m38s
build_systems / build-bob (pull_request) Failing after 14m11s
2026-05-14 12:58:15 -04:00
Richie b9d440597c removed llm tools from gui
treefmt / nix fmt (pull_request) Successful in 9s
pytest / pytest (pull_request) Successful in 1m4s
build_systems / build-brain (pull_request) Successful in 2m31s
build_systems / build-leviathan (pull_request) Successful in 3m21s
build_systems / build-rhapsody-in-green (pull_request) Successful in 3m21s
build_systems / build-jeeves (pull_request) Successful in 6m55s
build_systems / build-bob (pull_request) Failing after 16m4s
2026-05-13 10:03:15 -04:00
Richie 311cc5d7a7 adding pi-coding-agenta
treefmt / nix fmt (pull_request) Successful in 6s
pytest / pytest (pull_request) Successful in 1m24s
build_systems / build-brain (pull_request) Successful in 6m28s
build_systems / build-leviathan (pull_request) Failing after 7m21s
build_systems / build-rhapsody-in-green (pull_request) Failing after 7m22s
build_systems / build-jeeves (pull_request) Successful in 11m47s
build_systems / build-bob (pull_request) Failing after 19m3s
2026-05-13 08:57:45 -04:00
Richie fb2519046d moved codex and opencode to master pkgs 2026-05-13 08:56:18 -04:00
Richie bc6b1585ec flake update 2026-05-10 13:49:53 -04:00
Richie d71330a85a updated firefox configPath
treefmt / nix fmt (pull_request) Successful in 6s
pytest / pytest (pull_request) Successful in 29s
build_systems / build-brain (pull_request) Successful in 5m41s
build_systems / build-leviathan (pull_request) Successful in 5m43s
build_systems / build-jeeves (pull_request) Successful in 6m58s
build_systems / build-rhapsody-in-green (pull_request) Successful in 27m16s
build_systems / build-bob (pull_request) Failing after 12m14s
2026-05-10 12:36:54 -04:00
Richie df51aa5200 removing sunshine
sunshine is a cool idea but has been causing annoying ui glitches and started preventing the display manning for starting
Its a cool idea in theory but not useful enough for me to want to debug
2026-05-10 12:31:06 -04:00
Richie e93cc816db flake update 2026-05-09 17:38:13 -04:00
Richie 19050b4cf4 removing llms from rhapsody-in-green 2026-05-07 18:06:21 -04:00
Richie 6676c15f75 adding qwen3.6:27b 2026-05-07 18:05:00 -04:00
Richie 27e487e322 removing signal_bot
treefmt / nix fmt (pull_request) Successful in 5s
pytest / pytest (pull_request) Successful in 27s
build_systems / build-bob (pull_request) Successful in 48s
build_systems / build-brain (pull_request) Successful in 46s
build_systems / build-leviathan (pull_request) Successful in 54s
build_systems / build-rhapsody-in-green (pull_request) Successful in 1m0s
build_systems / build-jeeves (pull_request) Successful in 2m34s
treefmt / nix fmt (push) Successful in 5s
build_systems / build-bob (push) Successful in 34s
build_systems / build-brain (push) Successful in 31s
pytest / pytest (push) Successful in 27s
build_systems / build-leviathan (push) Successful in 40s
build_systems / build-rhapsody-in-green (push) Successful in 43s
build_systems / build-jeeves (push) Successful in 2m31s
2026-05-03 21:23:20 -04:00
Richie 4f28050eff added nixfmt and nix
build_systems / build-bob (pull_request) Failing after 52s
build_systems / build-brain (pull_request) Failing after 50s
pytest / pytest (pull_request) Failing after 4s
treefmt / nix fmt (pull_request) Failing after 4s
build_systems / build-leviathan (pull_request) Failing after 57s
build_systems / build-rhapsody-in-green (pull_request) Failing after 52s
build_systems / build-jeeves (pull_request) Failing after 3m17s
2026-05-03 20:47:03 -04:00
Richie b58ea60557 adding hostPackages
pytest / pytest (pull_request) Failing after 10s
treefmt / nix fmt (pull_request) Failing after 13s
build_systems / build-brain (pull_request) Failing after 29s
build_systems / build-bob (pull_request) Failing after 29s
build_systems / build-rhapsody-in-green (pull_request) Failing after 46s
build_systems / build-jeeves (pull_request) Failing after 2m29s
build_systems / build-leviathan (pull_request) Failing after 35s
2026-05-03 19:16:37 -04:00
Richie e95eedffe4 updated br-nix-builder
build_systems / build-bob (pull_request) Failing after 2s
build_systems / build-brain (pull_request) Failing after 1s
build_systems / build-jeeves (pull_request) Failing after 1s
build_systems / build-leviathan (pull_request) Failing after 1s
build_systems / build-rhapsody-in-green (pull_request) Failing after 1s
treefmt / nix fmt (pull_request) Failing after 2s
pytest / pytest (pull_request) Failing after 9s
2026-05-03 16:30:51 -04:00
Richie 1abd53987c made nix_builders not ephemeral and depended on gitea 2026-05-03 16:29:56 -04:00
Richie d1a3e7338a added permittedInsecurePackages for discord-canary 2026-05-03 00:39:23 -04:00
Richie 687ef0c167 moved acme_challenge backend 2026-05-03 00:39:19 -04:00
Richie 3a86148352 working nix builder 2026-05-02 17:10:02 -04:00
Richie fe9a2912e1 added words to spell check 2026-04-30 12:46:55 -04:00
Richie 29a99fc210 flake lock update 2026-04-30 12:46:55 -04:00
Richie d7651bf588 set update.nix to gitea 2026-04-30 12:46:55 -04:00
Richie 2865dcbe9c set dbus.implementation = "dbus"; 2026-04-30 12:46:55 -04:00
Richie d920b77bab removed verilux 2026-04-30 12:46:55 -04:00
Richie 1b53167b53 updated nix builders 2026-04-30 12:46:55 -04:00
Richie 9dabb9dc07 updated actions 2026-04-30 12:46:55 -04:00
Richie 95630fe151 made Prometheus require zfs-media-database-prometheus.mount 2026-04-30 10:16:37 -04:00
Richie d3a889f100 fixed typo 2026-04-30 10:16:37 -04:00
Richie 6ce0671f51 ran treefmt 2026-04-30 10:16:37 -04:00
Richie 25ab6b2ab6 added gitlens.pushRepositories key shourtcut 2026-04-30 10:16:37 -04:00
Richie 374d7e8d38 setting up resource monitoring for bob and jeeves 2026-04-30 10:16:37 -04:00
Richie 957110b7e9 increasing kitty scrollback_lines 2026-04-28 12:07:03 -04:00
Richie e7dc60f2c3 adding tiktoken 2026-04-28 12:07:03 -04:00
Richie 353a9d6787 adding pgvector 2026-04-27 13:02:04 -04:00
Richie 9f2d3a3c89 updated .gitignore 2026-04-25 16:33:45 -04:00
Richie 73e221716f adding nornsight 2026-04-25 16:33:45 -04:00
Richie 0d0ed5445a moved models 2026-04-19 21:05:56 -04:00
Richie 9e4c6f6f56 adding qwen3.6 2026-04-19 21:05:56 -04:00
Richie 1cf4b99d18 updating signing.format for programs.git 2026-04-19 10:35:48 -04:00
Richie b536fb9f09 removed fallbackToPassword = true; 2026-04-19 08:08:33 -04:00
github-actions[bot] c41a2ce3bd flake.lock: Update
Flake lock file updates:

• Updated input 'firefox-addons':
    'gitlab:rycee/nur-expressions/81e28f4?dir=pkgs/firefox-addons' (2026-03-20)
  → 'gitlab:rycee/nur-expressions/0581568?dir=pkgs/firefox-addons' (2026-04-17)
• Updated input 'home-manager':
    'github:nix-community/home-manager/9670de2' (2026-03-20)
  → 'github:nix-community/home-manager/565e534' (2026-04-17)
• Updated input 'nixos-hardware':
    'github:nixos/nixos-hardware/2d4b471' (2026-03-20)
  → 'github:nixos/nixos-hardware/c775c27' (2026-04-06)
• Updated input 'nixpkgs':
    'github:nixos/nixpkgs/b40629e' (2026-03-18)
  → 'github:nixos/nixpkgs/4bd9165' (2026-04-14)
• Updated input 'nixpkgs-master':
    'github:nixos/nixpkgs/8620c0b' (2026-03-21)
  → 'github:nixos/nixpkgs/025c852' (2026-04-17)
• Updated input 'sops-nix':
    'github:Mic92/sops-nix/29b6519' (2026-03-19)
  → 'github:Mic92/sops-nix/d4971dd' (2026-04-13)
2026-04-19 08:08:33 -04:00
Richie 8ef776f859 updating download-buffer-size 2026-04-18 22:12:47 -04:00
Richie d350c2d074 adding codex 2026-04-18 20:14:11 -04:00
Richie 93d6914e9d enabling appimages 2026-04-18 20:14:11 -04:00
Richie 7db063a240 setting up whisper transcriber 2026-04-18 19:09:02 -04:00
Richie dfe5997e0b removing brain_substituter.nix from bob 2026-04-18 12:00:59 -04:00
Richie 68671a1e84 adding steve 2026-04-18 11:56:56 -04:00
Richie bcc2227cfd updating syncthing phone id 2026-04-18 11:53:20 -04:00
Richie d6eec926e7 made web_services dir 2026-04-13 19:12:32 -04:00
Richie 5ddf1c4cab ran tree fmt 2026-04-13 19:12:32 -04:00
Richie 5a2171b9c7 updated gitea ssh settings 2026-04-13 19:12:32 -04:00
Richie 95c6ade154 moving off cloudflare tunnel 2026-04-13 19:12:32 -04:00
Richie a0bbc2896a added math the bob 2026-04-13 19:08:42 -04:00
Richie 736596c387 made bob a server 2026-04-13 19:08:42 -04:00
Richie 67622c0e51 setting up hedgedoc 2026-04-11 11:42:08 -04:00
Richie d2f447a1af disabling kafka 2026-04-11 11:11:21 -04:00
Richie af365fce9a setup sunshine.nix 2026-04-03 17:12:24 -04:00
Richie 6430049e92 updated postgres snapshot settings 2026-03-30 14:07:08 -04:00
Richie 26e4620f8f fixed systemd sandboxing 2026-03-30 14:07:08 -04:00
Richie 93fc700fa2 removed preStart step 2026-03-30 14:07:08 -04:00
Richie 8d1c1fc628 added mountpoint= to postgres zfs create 2026-03-30 14:07:08 -04:00
Richie dda318753b improving postgres wal 2026-03-30 14:07:08 -04:00
Richie 261ff139f7 removed ds table from richie DB 2026-03-29 15:54:54 -04:00
Richie ba8ff35109 updated ingest_congress to use congress-legislators for legislator info 2026-03-29 15:54:54 -04:00
Richie e368402eea adding LegislatorSocialMedia 2026-03-29 15:54:54 -04:00
Richie dd9329d218 fixed tests 2026-03-29 15:54:54 -04:00
Richie 89f6627bed converted session.execute(select to session.scalars(select 2026-03-29 15:54:54 -04:00
Richie c5babf8bad ran treefmt 2026-03-29 15:54:54 -04:00
Richie dae38ffd9b added ingest_congress.py 2026-03-29 15:54:54 -04:00
Richie ca62cc36a7 adding congress data to new DS DB 2026-03-29 15:54:54 -04:00
Richie 035410f39e adding nemotron-3-nano 2026-03-29 15:54:54 -04:00
Richie e40ab757ca making more generic exception handling 2026-03-29 15:54:54 -04:00
Richie 345ba94a59 ran ingest_posts 2026-03-29 15:54:54 -04:00
Richie f2084206b6 adding tables for 2023 2026-03-29 15:54:54 -04:00
Richie 50e764146a added ingest_posts.py 2026-03-29 15:54:54 -04:00
Richie ea97b5eb19 adding 2026 partitions 2026-03-29 15:54:54 -04:00
Richie 1ef2512daa adding post table 2026-03-29 15:54:54 -04:00
Richie f9a9e5395c added media/temp for fast dir when working with data 2026-03-29 15:54:54 -04:00
Richie d8e166a340 adding data_science_dev 2026-03-29 15:54:54 -04:00
Richie c266ba79f4 updated snapshot_config.toml 2026-03-29 14:12:06 -04:00
Richie f627a5ac6e enabling kafka 2026-03-26 09:59:31 -04:00
165 changed files with 18072 additions and 579 deletions
+1 -1
View File
@@ -23,6 +23,6 @@ jobs:
steps:
- uses: actions/checkout@v4
- name: Build default package
run: "nixos-rebuild build --flake ./#${{ matrix.system }}"
run: "nixos-rebuild build --accept-flake-config --flake ./#${{ matrix.system }}"
- name: copy to nix-cache
run: nix copy --accept-flake-config --to unix:///host-nix/var/nix/daemon-socket/socket .#nixosConfigurations.${{ matrix.system }}.config.system.build.toplevel
-30
View File
@@ -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 }}"
+7 -13
View File
@@ -6,24 +6,18 @@ on:
jobs:
merge:
runs-on: ubuntu-latest
runs-on: self-hosted
permissions:
contents: write
pull-requests: write
steps:
- name: Checkout repository
uses: actions/checkout@v4
- name: merge_flake_lock_update
run: |
pr_number=$(gh pr list --state open --author RichieCahill --label flake_lock_update --json number --jq '.[0].number')
echo "pr_number=$pr_number" >> $GITHUB_ENV
if [ -n "$pr_number" ]; then
gh pr merge "$pr_number" --rebase
else
echo "No open PR found with label flake_lock_update"
fi
run: >-
nix develop .#devShells.x86_64-linux.default -c
python -m python.gitea_flake_lock merge
--repo "${{ github.repository }}"
env:
GITHUB_TOKEN: ${{ secrets.GH_TOKEN_FOR_UPDATES }}
GITEA_TOKEN: ${{ secrets.GITEA_TOKEN }}
GITEA_URL: https://gitea.tmmworkshop.com
+1 -1
View File
@@ -1,13 +1,13 @@
name: pytest
on:
workflow_dispatch:
push:
branches:
- main
pull_request:
branches:
- main
merge_group:
jobs:
pytest:
+14 -11
View File
@@ -6,18 +6,21 @@ on:
jobs:
lockfile:
runs-on: ubuntu-latest
runs-on: self-hosted
permissions:
actions: write
contents: write
pull-requests: write
steps:
- name: Checkout repository
uses: actions/checkout@v4
- name: Install Nix
uses: DeterminateSystems/nix-installer-action@main
- name: Update flake.lock
uses: DeterminateSystems/update-flake-lock@main
with:
token: ${{ secrets.GH_TOKEN_FOR_UPDATES }}
pr-title: "Update flake.lock"
pr-labels: |
dependencies
automated
flake_lock_update
run: nix flake update
- name: Create or update flake.lock PR
env:
GITEA_TOKEN: ${{ secrets.GITEA_TOKEN }}
GITEA_URL: https://gitea.tmmworkshop.com
run: >-
nix develop .#devShells.x86_64-linux.default -c
python -m python.gitea_flake_lock update
--repo "${{ github.repository }}"
+4
View File
@@ -169,3 +169,7 @@ test.*
# Frontend build output
frontend/dist/
frontend/node_modules/
# data from testing llms
data/*
.ebook_search_bm25
+2 -1
View File
@@ -40,7 +40,6 @@
"cgroupdriver",
"charliermarsh",
"Checkpointing",
"cloudflared",
"codellama",
"codezombiech",
"compactmode",
@@ -204,6 +203,7 @@
"peerconnection",
"PESKYFOX",
"PGID",
"pgvector",
"pipewire",
"pkgs",
"plugdev",
@@ -308,6 +308,7 @@
"usernamehw",
"userprefs",
"vaninventory",
"vdev",
"vfat",
"victron",
"virt",
+12 -2
View File
@@ -23,7 +23,10 @@
boot = {
tmp.useTmpfs = true;
kernelPackages = lib.mkDefault pkgs.linuxPackages_6_12;
zfs.package = lib.mkDefault pkgs.zfs_2_4;
zfs = {
package = lib.mkDefault pkgs.zfs_2_4;
forceImportRoot = lib.mkDefault false;
};
};
hardware.enableRedistributableFirmware = true;
@@ -37,10 +40,17 @@
nixpkgs = {
overlays = builtins.attrValues outputs.overlays;
config.allowUnfree = true;
config = {
allowUnfree = true;
permittedInsecurePackages = [
"openssl-1.1.1w" # This is for discord-canary
];
};
};
services = {
dbus.implementation = "dbus";
# firmware update
fwupd.enable = true;
+1
View File
@@ -34,6 +34,7 @@ in
warn-dirty = false;
flake-registry = ""; # disable global flake registries
connect-timeout = 10;
download-buffer-size = 536870912;
fallback = true;
};
+256
View File
@@ -0,0 +1,256 @@
{
config,
lib,
pkgs,
...
}:
let
monitoringInterface = "ztwfunumly";
nodeTextfileDir = "/var/lib/prometheus-node-exporter-textfile";
mkProcessNameTemplate =
perPid: template: if perPid then "${template}:{{.PID}}:{{.StartTime}}" else template;
mkProcessMatchers = perPid: [
{
name = mkProcessNameTemplate perPid "{{.Username}}:{{.Matches.Module}}";
cmdline = [ "^/nix/store[^ ]*/bin/python[^ ]* -m (?P<Module>[^ ]+)" ];
}
{
name = mkProcessNameTemplate perPid "{{.Username}}:{{.Matches.Wrapped}}";
cmdline = [
"^/nix/store[^ ]*/bin/python[^ ]* /nix/store[^ ]*/bin/\\.?(?P<Wrapped>[^ /]+?)(?:-wrapped)?(?:\\s|$)"
];
}
{
name = mkProcessNameTemplate perPid "{{.Username}}:{{.Matches.Wrapped}}";
cmdline = [
"^/nix/store[^ ]*/bin/node /nix/store[^ ]*-(?P<Wrapped>[A-Za-z0-9._+-]+)-[0-9][^ /]*/"
];
}
{
name = mkProcessNameTemplate perPid "{{.Username}}:{{.Matches.Wrapped}}";
cmdline = [ "^/nix/store[^ ]*/(?:bin/|lib/[^ ]*/)?\\.?(?P<Wrapped>[^ /]+?)(?:-wrapped)?(?:\\s|$)" ];
}
{
name = mkProcessNameTemplate perPid "{{.Username}}:{{.ExeBase}}";
cmdline = [ ".+" ];
}
];
perPidConfig = pkgs.writeText "process-exporter-per-pid.yaml" (
builtins.toJSON {
process_names = mkProcessMatchers true;
}
);
zpoolLatencyScript = pkgs.writeShellScript "zpool-latency-exporter" ''
set -euo pipefail
out_dir=${lib.escapeShellArg nodeTextfileDir}
host=${lib.escapeShellArg config.networking.hostName}
tmp_file="$(mktemp "$out_dir/zpool.prom.XXXXXX")"
trap 'rm -f "$tmp_file"' EXIT
pools="$(zpool list -H -o name | paste -sd, -)"
cat >"$tmp_file" <<'EOF'
# HELP zpool_iostat_total_wait_read_ns Average total read wait time reported by zpool iostat.
# TYPE zpool_iostat_total_wait_read_ns gauge
# HELP zpool_iostat_total_wait_write_ns Average total write wait time reported by zpool iostat.
# TYPE zpool_iostat_total_wait_write_ns gauge
# HELP zpool_iostat_disk_wait_read_ns Average disk read wait time reported by zpool iostat.
# TYPE zpool_iostat_disk_wait_read_ns gauge
# HELP zpool_iostat_disk_wait_write_ns Average disk write wait time reported by zpool iostat.
# TYPE zpool_iostat_disk_wait_write_ns gauge
# HELP zpool_iostat_syncq_wait_read_ns Average synchronous queue read wait time reported by zpool iostat.
# TYPE zpool_iostat_syncq_wait_read_ns gauge
# HELP zpool_iostat_syncq_wait_write_ns Average synchronous queue write wait time reported by zpool iostat.
# TYPE zpool_iostat_syncq_wait_write_ns gauge
# HELP zpool_iostat_asyncq_wait_read_ns Average asynchronous queue read wait time reported by zpool iostat.
# TYPE zpool_iostat_asyncq_wait_read_ns gauge
# HELP zpool_iostat_asyncq_wait_write_ns Average asynchronous queue write wait time reported by zpool iostat.
# TYPE zpool_iostat_asyncq_wait_write_ns gauge
EOF
zpool iostat -Hplvy -y 1 1 | awk -F '\t' -v host="$host" -v pools="$pools" '
function esc(str, out) {
out = str
gsub(/\\/, "\\\\", out)
gsub(/"/, "\\\"", out)
return out
}
function emit(metric, pool, vdev, value) {
if (value == "" || value == "-") {
return
}
printf "%s{host=\"%s\",pool=\"%s\",vdev=\"%s\"} %s\n",
metric,
esc(host),
esc(pool),
esc(vdev),
value
}
BEGIN {
split(pools, pool_names, ",")
for (idx in pool_names) {
if (pool_names[idx] != "") {
known_pools[pool_names[idx]] = 1
}
}
}
NF == 0 {
next
}
{
row_name = $1
if (row_name in known_pools) {
current_pool = row_name
current_vdev = "_pool"
} else if (current_pool == "") {
next
} else {
current_vdev = row_name
}
emit("zpool_iostat_total_wait_read_ns", current_pool, current_vdev, $8)
emit("zpool_iostat_total_wait_write_ns", current_pool, current_vdev, $9)
emit("zpool_iostat_disk_wait_read_ns", current_pool, current_vdev, $10)
emit("zpool_iostat_disk_wait_write_ns", current_pool, current_vdev, $11)
emit("zpool_iostat_syncq_wait_read_ns", current_pool, current_vdev, $12)
emit("zpool_iostat_syncq_wait_write_ns", current_pool, current_vdev, $13)
emit("zpool_iostat_asyncq_wait_read_ns", current_pool, current_vdev, $14)
emit("zpool_iostat_asyncq_wait_write_ns", current_pool, current_vdev, $15)
}
' >>"$tmp_file"
mv "$tmp_file" "$out_dir/zpool.prom"
trap - EXIT
'';
in
{
networking.firewall.interfaces.${monitoringInterface}.allowedTCPPorts = [
9100
9134
9256
9257
9633
];
services.prometheus.exporters = {
node = {
enable = true;
enabledCollectors = [
"pressure"
"processes"
"systemd"
];
extraFlags = [ "--collector.textfile.directory=${nodeTextfileDir}" ];
};
process = {
enable = true;
user = "root";
group = "root";
settings.process_names = mkProcessMatchers false;
extraFlags = [
"-gather-smaps=false"
"-remove-empty-groups=true"
"-threads=false"
];
};
smartctl.enable = true;
zfs.enable = true;
};
programs.atop = {
enable = true;
atopService.enable = true;
atopRotateTimer.enable = true;
atopacctService.enable = true;
settings.interval = 30;
};
systemd = {
services = {
prometheus-process-pid-exporter = {
description = "Prometheus process exporter with per-PID naming";
wantedBy = [ "multi-user.target" ];
after = [ "network.target" ];
serviceConfig = {
ExecStart = ''
${pkgs.prometheus-process-exporter}/bin/process-exporter \
--web.listen-address 0.0.0.0:9257 \
--config.path ${perPidConfig} \
-children=false \
-gather-smaps=false \
-remove-empty-groups=true \
-threads=false
'';
User = "root";
Group = "root";
Restart = "always";
WorkingDirectory = "/tmp";
CapabilityBoundingSet = [ "" ];
DeviceAllow = [ "" ];
LockPersonality = true;
MemoryDenyWriteExecute = true;
NoNewPrivileges = true;
PrivateDevices = true;
PrivateTmp = true;
ProtectClock = true;
ProtectControlGroups = true;
ProtectHome = true;
ProtectHostname = true;
ProtectKernelLogs = true;
ProtectKernelModules = true;
ProtectKernelTunables = true;
ProtectSystem = "strict";
RemoveIPC = true;
RestrictAddressFamilies = [
"AF_INET"
"AF_INET6"
];
RestrictNamespaces = true;
RestrictRealtime = true;
RestrictSUIDSGID = true;
SystemCallArchitectures = "native";
UMask = "0077";
};
};
zpool-latency-exporter = {
description = "Exports ZFS latency metrics for node_exporter textfile collection";
after = [ "zfs-import.target" ];
requires = [ "zfs-import.target" ];
path = [
config.boot.zfs.package
pkgs.coreutils
pkgs.gawk
];
serviceConfig = {
Type = "oneshot";
ExecStart = zpoolLatencyScript;
};
};
};
timers.zpool-latency-exporter = {
wantedBy = [ "timers.target" ];
timerConfig = {
OnBootSec = "2m";
OnUnitActiveSec = "60s";
Unit = "zpool-latency-exporter.service";
};
};
tmpfiles.rules = [ "d ${nodeTextfileDir} 0755 root root - -" ];
};
}
+1 -1
View File
@@ -12,7 +12,7 @@
brain.id = "SSCGIPI-IV3VYKB-TRNIJE3-COV4T2H-CDBER7F-I2CGHYA-NWOEUDU-3T5QAAN"; # cspell:disable-line
ipad.id = "KI76T3X-SFUGV2L-VSNYTKR-TSIUV5L-SHWD3HE-GQRGRCN-GY4UFMD-CW6Z6AX"; # cspell:disable-line
jeeves.id = "ICRHXZW-ECYJCUZ-I4CZ64R-3XRK7CG-LL2HAAK-FGOHD22-BQA4AI6-5OAL6AG"; # cspell:disable-line
phone.id = "TBRULKD-7DZPGGZ-F6LLB7J-MSO54AY-7KLPBIN-QOFK6PX-W2HBEWI-PHM2CQI"; # cspell:disable-line
phone.id = "JPVQKQW-CFXOJXT-Q5G5F3H-QIDHDRE-GKHPTQB-GXZUQSP-U7FR7F7-INP3AAH"; # cspell:disable-line
rhapsody-in-green.id = "ASL3KC4-3XEN6PA-7BQBRKE-A7JXLI6-DJT43BY-Q4WPOER-7UALUAZ-VTPQ6Q4"; # cspell:disable-line
};
};
+1 -1
View File
@@ -4,7 +4,7 @@
flags = [ "--accept-flake-config" ];
randomizedDelaySec = "1h";
persistent = true;
flake = "github:RichieCahill/dotfiles";
flake = "git+https://gitea.tmmworkshop.com/richie/dotfiles?ref=main";
allowReboot = true;
dates = "Sat *-*-* 06:00:00";
};
+76
View File
@@ -0,0 +1,76 @@
# ZFS failed root import recovery
## Fast path
If the machine fails to boot because ZFS refuses to import `root_pool`:
### GRUB
1. At the bootloader menu, select the normal NixOS entry.
2. Press `e`.
3. Find the line that starts with `linux`.
4. Append this to the end of that line:
```text
zfs_force=1
```
5. Boot once with `Ctrl+x` or `F10`.
### systemd-boot
1. At the bootloader menu, highlight the normal NixOS entry.
2. Press `e`.
3. Append this to the end of the options line:
```text
zfs_force=1
```
4. Press `Enter` to boot once.
## After boot
Run:
```bash
sudo zpool status
sudo zpool import
journalctl -b | rg "ZFS|zfs|import|root_pool"
```
## Expected result
`sudo zpool status` should show `root_pool` as `ONLINE`.
## Reboot test
Run:
```bash
sudo reboot
```
Do not add `zfs_force=1` the second time.
## If it still fails
Boot once more with:
```text
zfs_force=1
```
Then run:
```bash
sudo zpool status -v
sudo zpool history | tail -n 50
journalctl -b | rg "ZFS|zfs|import|root_pool"
```
## Notes
- Root pool name is `root_pool`.
- This is a one-time recovery path after disk moves, controller changes, dirty exports, or interrupted imports.
- Some hosts also need the LUKS unlock USB key inserted before boot.
Generated
+42 -26
View File
@@ -8,11 +8,11 @@
},
"locked": {
"dir": "pkgs/firefox-addons",
"lastModified": 1773979456,
"narHash": "sha256-9kBMJ5IvxqNlkkj/swmE8uK1Sc7TL/LIRUI958m7uBM=",
"lastModified": 1781150628,
"narHash": "sha256-b4mp8l3qWuSCyYYo9HSngDtcB3PpecYiOXjULrjwwlw=",
"owner": "rycee",
"repo": "nur-expressions",
"rev": "81e28f47ac18d9e89513929c77e711e657b64851",
"rev": "753319310f4673a2dabbfab87482187b40bf9bac",
"type": "gitlab"
},
"original": {
@@ -29,11 +29,11 @@
]
},
"locked": {
"lastModified": 1774007980,
"narHash": "sha256-FOnZjElEI8pqqCvB6K/1JRHTE8o4rer8driivTpq2uo=",
"lastModified": 1781189114,
"narHash": "sha256-5inaamLgUMWy+MOBE9ChF9QAF1o/74LFuHkI0W/9rqc=",
"owner": "nix-community",
"repo": "home-manager",
"rev": "9670de2921812bc4e0452f6e3efd8c859696c183",
"rev": "486595d2cf49cfcd649b58a284fa11ac0e34da22",
"type": "github"
},
"original": {
@@ -43,12 +43,15 @@
}
},
"nixos-hardware": {
"inputs": {
"nixpkgs": "nixpkgs"
},
"locked": {
"lastModified": 1774018263,
"narHash": "sha256-HHYEwK1A22aSaxv2ibhMMkKvrDGKGlA/qObG4smrSqc=",
"lastModified": 1781168557,
"narHash": "sha256-LOnLQ2tpYF9gqIDDr3+j3DbpJJr/QCH6zPRT2GzEUOE=",
"owner": "nixos",
"repo": "nixos-hardware",
"rev": "2d4b4717b2534fad5c715968c1cece04a172b365",
"rev": "6358ff76821101c178e3ab4919a62799bfe3652e",
"type": "github"
},
"original": {
@@ -60,27 +63,24 @@
},
"nixpkgs": {
"locked": {
"lastModified": 1773821835,
"narHash": "sha256-TJ3lSQtW0E2JrznGVm8hOQGVpXjJyXY2guAxku2O9A4=",
"owner": "nixos",
"repo": "nixpkgs",
"rev": "b40629efe5d6ec48dd1efba650c797ddbd39ace0",
"type": "github"
"lastModified": 1767892417,
"narHash": "sha256-8bW3q88CEg2u4hSP66Vf4lpbLonHz7hqDNBMcCY7E9U=",
"rev": "3497aa5c9457a9d88d71fa93a4a8368816fbeeba",
"type": "tarball",
"url": "https://releases.nixos.org/nixos/unstable/nixos-26.05pre924538.3497aa5c9457/nixexprs.tar.xz"
},
"original": {
"owner": "nixos",
"ref": "nixos-unstable",
"repo": "nixpkgs",
"type": "github"
"type": "tarball",
"url": "https://channels.nixos.org/nixos-unstable/nixexprs.tar.xz"
}
},
"nixpkgs-master": {
"locked": {
"lastModified": 1774051532,
"narHash": "sha256-d3CGMweyYIcPuTj5BKq+1Lx4zwlgL31nVtN647tOZKo=",
"lastModified": 1781229721,
"narHash": "sha256-ORvqDbb/LYxiJljGIejapjkc/kJbVote2N1WSb9W45I=",
"owner": "nixos",
"repo": "nixpkgs",
"rev": "8620c0b5cc8fbe76502442181be1d0514bc3a1b7",
"rev": "173d0ad7a974f8543a9ab01d2271b2e290341b33",
"type": "github"
},
"original": {
@@ -106,12 +106,28 @@
"type": "github"
}
},
"nixpkgs_2": {
"locked": {
"lastModified": 1781074563,
"narHash": "sha256-md8WlXOlfnIeHeOScMTTHFyf2d6iaTwPl2apR5EQ3P4=",
"owner": "nixos",
"repo": "nixpkgs",
"rev": "9ae611a455b90cf061d8f332b977e387bda8e1ca",
"type": "github"
},
"original": {
"owner": "nixos",
"ref": "nixos-unstable",
"repo": "nixpkgs",
"type": "github"
}
},
"root": {
"inputs": {
"firefox-addons": "firefox-addons",
"home-manager": "home-manager",
"nixos-hardware": "nixos-hardware",
"nixpkgs": "nixpkgs",
"nixpkgs": "nixpkgs_2",
"nixpkgs-master": "nixpkgs-master",
"nixpkgs-stable": "nixpkgs-stable",
"sops-nix": "sops-nix",
@@ -125,11 +141,11 @@
]
},
"locked": {
"lastModified": 1773889674,
"narHash": "sha256-+ycaiVAk3MEshJTg35cBTUa0MizGiS+bgpYw/f8ohkg=",
"lastModified": 1780547341,
"narHash": "sha256-Gq8KNx5A7hBB3uGJaj6eQfLDIz5YdLu92gqBcvHvoUo=",
"owner": "Mic92",
"repo": "sops-nix",
"rev": "29b6519f3e0780452bca0ac0be4584f04ac16cc5",
"rev": "9ed65852b6257fbeae4355bc24ecfea307ca759a",
"type": "github"
},
"original": {
-24
View File
@@ -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?
+30 -1
View File
@@ -17,14 +17,41 @@
python-env = final: _prev: {
my_python = final.python314.withPackages (
ps: with ps; [
ps:
let
bm25s = ps.buildPythonPackage rec {
pname = "bm25s";
version = "0.3.9";
pyproject = true;
src = final.fetchPypi {
inherit pname version;
hash = "sha256-iVxnnZUrfeg1XttfPhpiCh4vKU0dQrkZvwghzOLi9Zc=";
};
build-system = [ ps.setuptools ];
dependencies = with ps; [
numpy
scipy
];
pythonImportsCheck = [ "bm25s" ];
};
in
with ps;
[
alembic
apprise
apscheduler
beautifulsoup4
ebooklib
fastapi
fastapi-cli
httpx
mypy
numpy
orjson
pgvector
polars
psycopg
pydantic
@@ -38,8 +65,10 @@
scalene
sqlalchemy
sqlalchemy
bm25s
tenacity
textual
tiktoken
tinytuya
typer
websockets
+3 -3
View File
@@ -26,6 +26,7 @@ dependencies = [
[project.scripts]
database = "python.database_cli:app"
van-inventory = "python.van_inventory.main:serve"
whisper-transcribe = "python.tools.whisper.transcribe:main"
[dependency-groups]
dev = [
@@ -50,6 +51,7 @@ lint.ignore = [
"COM812", # (TEMP) conflicts when used with the formatter
"ISC001", # (TEMP) conflicts when used with the formatter
"S603", # (PERM) This is known to cause a false positive
"S607", # (PERM) This is becoming a consistent annoyance
]
[tool.ruff.lint.per-file-ignores]
@@ -78,9 +80,7 @@ lint.ignore = [
"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/alembic/**" = [
"INP001", # (perm) this creates LSP issues for alembic
]
@@ -0,0 +1,50 @@
"""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 ###
@@ -0,0 +1,72 @@
"""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}")
@@ -0,0 +1,153 @@
"""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 ###
@@ -0,0 +1,58 @@
"""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
View File
@@ -81,6 +81,7 @@ def include_name(
"""
if type_ == "schema":
# allows a database with multiple schemas to have separate alembic revisions
return name == target_metadata.schema
return True
@@ -0,0 +1,187 @@
"""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 ###
@@ -0,0 +1,93 @@
"""adding audiobook libreary metadata.
Revision ID: d7864d1ffc17
Revises: c8a794340928
Create Date: 2026-06-03 20:24:09.200837
"""
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 = "d7864d1ffc17"
down_revision: str | None = "c8a794340928"
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(
"audiobook_author",
sa.Column("name", 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.PrimaryKeyConstraint("id", name=op.f("pk_audiobook_author")),
sa.UniqueConstraint("name", name=op.f("uq_audiobook_author_name")),
schema=schema,
)
op.create_table(
"audiobook_series",
sa.Column("name", sa.String(), nullable=False),
sa.Column("author_id", sa.Integer(), 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(
["author_id"],
[f"{schema}.audiobook_author.id"],
name=op.f("fk_audiobook_series_author_id_audiobook_author"),
ondelete="CASCADE",
),
sa.PrimaryKeyConstraint("id", name=op.f("pk_audiobook_series")),
sa.UniqueConstraint("author_id", "name", name=op.f("uq_audiobook_series_author_id")),
schema=schema,
)
op.create_table(
"audiobook",
sa.Column("title", sa.String(), nullable=False),
sa.Column("author_id", sa.Integer(), nullable=False),
sa.Column("series_id", sa.Integer(), nullable=True),
sa.Column("series_index", sa.Integer(), 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(
["author_id"],
[f"{schema}.audiobook_author.id"],
name=op.f("fk_audiobook_author_id_audiobook_author"),
ondelete="CASCADE",
),
sa.ForeignKeyConstraint(
["series_id"],
[f"{schema}.audiobook_series.id"],
name=op.f("fk_audiobook_series_id_audiobook_series"),
ondelete="SET NULL",
),
sa.PrimaryKeyConstraint("id", name=op.f("pk_audiobook")),
schema=schema,
)
# ### end Alembic commands ###
def downgrade() -> None:
"""Downgrade."""
# ### commands auto generated by Alembic - please adjust! ###
op.drop_table("audiobook", schema=schema)
op.drop_table("audiobook_series", schema=schema)
op.drop_table("audiobook_author", schema=schema)
# ### end Alembic commands ###
@@ -0,0 +1,200 @@
"""add ebook search tables.
Revision ID: 2db132cace1a
Revises: b3c60cc5beb5
Create Date: 2026-06-10 22:10:54.379159
"""
from __future__ import annotations
from typing import TYPE_CHECKING
import pgvector
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 = "2db132cace1a"
down_revision: str | None = "b3c60cc5beb5"
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(
"ebook_embedding_model",
sa.Column("name", sa.String(), nullable=False),
sa.Column("dimension", sa.Integer(), nullable=False),
sa.Column("is_default", sa.Boolean(), 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_ebook_embedding_model")),
sa.UniqueConstraint("name", name=op.f("uq_ebook_embedding_model_name")),
schema=schema,
)
op.create_table(
"ebook_source",
sa.Column("title", sa.String(), nullable=False),
sa.Column("author", sa.String(), nullable=True),
sa.Column("language", sa.String(), nullable=True),
sa.Column("publisher", sa.String(), nullable=True),
sa.Column("identifier", sa.String(), nullable=True),
sa.Column("file_path", sa.String(), nullable=False),
sa.Column("file_sha256", sa.String(length=64), nullable=False),
sa.Column("file_mtime", sa.DateTime(timezone=True), nullable=False),
sa.Column("file_size", sa.BigInteger(), 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_ebook_source")),
sa.UniqueConstraint("file_path", name=op.f("uq_ebook_source_file_path")),
sa.UniqueConstraint("file_sha256", name=op.f("uq_ebook_source_file_sha256")),
schema=schema,
)
op.create_table(
"ebook_chapter",
sa.Column("source_id", sa.Integer(), nullable=False),
sa.Column("spine_index", sa.Integer(), nullable=False),
sa.Column("title", sa.String(), nullable=True),
sa.Column("href", 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.ForeignKeyConstraint(
["source_id"],
[f"{schema}.ebook_source.id"],
name=op.f("fk_ebook_chapter_source_id_ebook_source"),
ondelete="CASCADE",
),
sa.PrimaryKeyConstraint("id", name=op.f("pk_ebook_chapter")),
sa.UniqueConstraint("source_id", "spine_index", name=op.f("uq_ebook_chapter_source_id")),
schema=schema,
)
op.create_table(
"ebook_chunk",
sa.Column("source_id", sa.Integer(), nullable=False),
sa.Column("chapter_id", sa.Integer(), nullable=True),
sa.Column("chunk_index", sa.Integer(), nullable=False),
sa.Column("text", sa.String(), nullable=False),
sa.Column("token_start", sa.Integer(), nullable=False),
sa.Column("token_count", sa.Integer(), nullable=False),
sa.Column("page_label", sa.String(), nullable=True),
sa.Column("content_sha256", sa.String(length=64), nullable=False),
sa.Column("search_text", sa.String(), nullable=False),
sa.Column("id", sa.BigInteger(), 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(
["chapter_id"],
[f"{schema}.ebook_chapter.id"],
name=op.f("fk_ebook_chunk_chapter_id_ebook_chapter"),
ondelete="SET NULL",
),
sa.ForeignKeyConstraint(
["source_id"],
[f"{schema}.ebook_source.id"],
name=op.f("fk_ebook_chunk_source_id_ebook_source"),
ondelete="CASCADE",
),
sa.PrimaryKeyConstraint("id", name=op.f("pk_ebook_chunk")),
sa.UniqueConstraint("source_id", "chunk_index", name="uq_ebook_chunk_source_id_chunk_index"),
sa.UniqueConstraint("source_id", "content_sha256", name="uq_ebook_chunk_source_id_content_sha256"),
schema=schema,
)
op.create_table(
"ebook_chunk_embedding_1024",
sa.Column("chunk_id", sa.BigInteger(), nullable=False),
sa.Column("model_id", sa.Integer(), nullable=False),
sa.Column("embedding", pgvector.sqlalchemy.vector.VECTOR(dim=1024), nullable=False),
sa.Column("id", sa.BigInteger(), 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(
["chunk_id"],
[f"{schema}.ebook_chunk.id"],
name=op.f("fk_ebook_chunk_embedding_1024_chunk_id_ebook_chunk"),
ondelete="CASCADE",
),
sa.ForeignKeyConstraint(
["model_id"],
[f"{schema}.ebook_embedding_model.id"],
name=op.f("fk_ebook_chunk_embedding_1024_model_id_ebook_embedding_model"),
ondelete="CASCADE",
),
sa.PrimaryKeyConstraint("id", name=op.f("pk_ebook_chunk_embedding_1024")),
sa.UniqueConstraint("chunk_id", "model_id", name=op.f("uq_ebook_chunk_embedding_1024_chunk_id")),
schema=schema,
)
op.create_table(
"ebook_chunk_embedding_2560",
sa.Column("chunk_id", sa.BigInteger(), nullable=False),
sa.Column("model_id", sa.Integer(), nullable=False),
sa.Column("embedding", pgvector.sqlalchemy.vector.VECTOR(dim=2560), nullable=False),
sa.Column("id", sa.BigInteger(), 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(
["chunk_id"],
[f"{schema}.ebook_chunk.id"],
name=op.f("fk_ebook_chunk_embedding_2560_chunk_id_ebook_chunk"),
ondelete="CASCADE",
),
sa.ForeignKeyConstraint(
["model_id"],
[f"{schema}.ebook_embedding_model.id"],
name=op.f("fk_ebook_chunk_embedding_2560_model_id_ebook_embedding_model"),
ondelete="CASCADE",
),
sa.PrimaryKeyConstraint("id", name=op.f("pk_ebook_chunk_embedding_2560")),
sa.UniqueConstraint("chunk_id", "model_id", name=op.f("uq_ebook_chunk_embedding_2560_chunk_id")),
schema=schema,
)
op.create_table(
"ebook_chunk_embedding_4096",
sa.Column("chunk_id", sa.BigInteger(), nullable=False),
sa.Column("model_id", sa.Integer(), nullable=False),
sa.Column("embedding", pgvector.sqlalchemy.vector.VECTOR(dim=4096), nullable=False),
sa.Column("id", sa.BigInteger(), 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(
["chunk_id"],
[f"{schema}.ebook_chunk.id"],
name=op.f("fk_ebook_chunk_embedding_4096_chunk_id_ebook_chunk"),
ondelete="CASCADE",
),
sa.ForeignKeyConstraint(
["model_id"],
[f"{schema}.ebook_embedding_model.id"],
name=op.f("fk_ebook_chunk_embedding_4096_model_id_ebook_embedding_model"),
ondelete="CASCADE",
),
sa.PrimaryKeyConstraint("id", name=op.f("pk_ebook_chunk_embedding_4096")),
sa.UniqueConstraint("chunk_id", "model_id", name=op.f("uq_ebook_chunk_embedding_4096_chunk_id")),
schema=schema,
)
# ### end Alembic commands ###
def downgrade() -> None:
"""Downgrade."""
# ### commands auto generated by Alembic - please adjust! ###
op.drop_table("ebook_chunk_embedding_4096", schema=schema)
op.drop_table("ebook_chunk_embedding_2560", schema=schema)
op.drop_table("ebook_chunk_embedding_1024", schema=schema)
op.drop_table("ebook_chunk", schema=schema)
op.drop_table("ebook_chapter", schema=schema)
op.drop_table("ebook_source", schema=schema)
op.drop_table("ebook_embedding_model", schema=schema)
# ### end Alembic commands ###
@@ -0,0 +1,63 @@
"""updated series_index to float and added UniqueConstraint to audiobook and audiobook_author.
Revision ID: b3c60cc5beb5
Revises: d7864d1ffc17
Create Date: 2026-06-10 20:02:43.073725
"""
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 = "b3c60cc5beb5"
down_revision: str | None = "d7864d1ffc17"
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(
"audiobook",
"series_index",
existing_type=sa.INTEGER(),
type_=sa.Float(),
existing_nullable=False,
schema=schema,
)
op.create_unique_constraint(
op.f("uq_audiobook_author_id"),
"audiobook",
["author_id", "series_id", "title"],
schema=schema,
postgresql_nulls_not_distinct=True,
)
# ### end Alembic commands ###
def downgrade() -> None:
"""Downgrade."""
# ### commands auto generated by Alembic - please adjust! ###
op.drop_constraint(op.f("uq_audiobook_author_id"), "audiobook", schema=schema, type_="unique")
op.alter_column(
"audiobook",
"series_index",
existing_type=sa.Float(),
type_=sa.INTEGER(),
existing_nullable=False,
schema=schema,
)
# ### end Alembic commands ###
+3
View File
@@ -0,0 +1,3 @@
"""Data science CLI tools."""
from __future__ import annotations
+613
View File
@@ -0,0 +1,613 @@
"""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()
+247
View File
@@ -0,0 +1,247 @@
"""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()
+7
View File
@@ -90,6 +90,13 @@ DATABASES: dict[str, DatabaseConfig] = {
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",
),
}
+1
View File
@@ -0,0 +1 @@
"""EPUB search package."""
+57
View File
@@ -0,0 +1,57 @@
"""Grounded answer generation."""
from __future__ import annotations
import logging
from typing import TYPE_CHECKING
from python.ebook_search.llm_interface import request_chat_completion
if TYPE_CHECKING:
from python.ebook_search.config import EbookSearchConfig
from python.ebook_search.search import SearchResult
logger = logging.getLogger(__name__)
def answer_query(query: str, results: list[SearchResult], config: EbookSearchConfig) -> str:
"""Answer a question using only retrieved chunks."""
if not config.answer_enabled:
logger.info("ebook_answer_skipped_disabled")
return "Answer generation is disabled. Source chunks are shown below."
if not results:
logger.info("ebook_answer_skipped_no_results")
return "No relevant sources were found."
logger.info(
"ebook_answer_request_start base_url=%s model=%s sources=%s query_length=%s",
config.vllm_base_url,
config.chat_model,
len(results),
len(query),
)
context = "\n\n".join(
f"[{index}] {result.source_title}{' - ' + result.chapter_title if result.chapter_title else ''}\n{result.text}"
for index, result in enumerate(results, start=1)
)
content = request_chat_completion(
config,
[
{
"role": "system",
"content": (
"Answer only from the provided context. Cite sources with bracketed numbers like [1]. "
"If the context is insufficient, say so."
),
},
{"role": "user", "content": f"Question:\n{query}\n\nContext:\n{context}"},
],
)
logger.info(
"ebook_answer_request_complete model=%s answer_length=%s",
config.chat_model,
len(content),
)
return content or "The model returned an empty answer."
+1
View File
@@ -0,0 +1 @@
"""Web and external API adapters for EPUB search."""
+58
View File
@@ -0,0 +1,58 @@
"""Background BM25 refresh tasks for the web app."""
from __future__ import annotations
import logging
from threading import Timer
from typing import TYPE_CHECKING
from sqlalchemy.orm import Session
from python.ebook_search.bm25_corpus import refresh_bm25_corpus
if TYPE_CHECKING:
from fastapi import FastAPI
from sqlalchemy.engine import Engine
from python.ebook_search.config import EbookSearchConfig
logger = logging.getLogger(__name__)
def schedule_bm25_refresh(app: FastAPI) -> None:
"""Schedule a delayed BM25 corpus refresh, replacing any pending refresh."""
existing_timer = getattr(app.state, "bm25_refresh_timer", None)
if existing_timer is not None:
existing_timer.cancel()
timer = Timer(app.state.config.bm25_refresh_delay_seconds, refresh_bm25_for_app, args=(app,))
timer.daemon = True
timer.start()
app.state.bm25_refresh_timer = timer
logger.info(
"ebook_bm25_refresh_scheduled delay_seconds=%s",
app.state.config.bm25_refresh_delay_seconds,
)
def cancel_bm25_refresh(app: FastAPI) -> None:
"""Cancel any pending BM25 corpus refresh."""
existing_timer = getattr(app.state, "bm25_refresh_timer", None)
if existing_timer is not None:
existing_timer.cancel()
app.state.bm25_refresh_timer = None
logger.info("ebook_bm25_refresh_cancelled")
def refresh_bm25_for_app(app: FastAPI) -> None:
"""Refresh the BM25 corpus using the app engine and config."""
try:
refresh_bm25_for_engine(app.state.engine, app.state.config)
except Exception:
logger.exception("ebook_bm25_refresh_failed")
def refresh_bm25_for_engine(engine: Engine, config: EbookSearchConfig) -> None:
"""Refresh the BM25 corpus using a SQLAlchemy engine."""
with Session(engine) as session:
refresh_bm25_corpus(session, config)
+75
View File
@@ -0,0 +1,75 @@
"""FastAPI HTMX app for EPUB search."""
from __future__ import annotations
import logging
from contextlib import asynccontextmanager
from typing import TYPE_CHECKING, Annotated
import typer
import uvicorn
from fastapi import FastAPI
from fastapi.staticfiles import StaticFiles
from sqlalchemy.orm import Session
from python.common import configure_logger
from python.ebook_search.api.bm25_tasks import cancel_bm25_refresh
from python.ebook_search.api.routes import register_admin_routes, register_page_routes, register_search_routes
from python.ebook_search.api.web import STATIC_DIR
from python.ebook_search.bm25_corpus import ensure_bm25_corpus
from python.ebook_search.config import load_config
from python.orm.common import get_postgres_engine
if TYPE_CHECKING:
from collections.abc import AsyncIterator
logger = logging.getLogger(__name__)
@asynccontextmanager
async def lifespan(app: FastAPI) -> AsyncIterator[None]:
"""Manage application startup and shutdown resources."""
logger.info("ebook_search_startup")
app.state.engine = get_postgres_engine(name="RICHIE")
with Session(app.state.engine) as session:
ensure_bm25_corpus(session, app.state.config)
try:
yield
finally:
logger.info("ebook_search_shutdown")
cancel_bm25_refresh(app)
app.state.engine.dispose()
def create_app() -> FastAPI:
"""Create the EPUB search web app."""
app = FastAPI(title="EPUB Search", lifespan=lifespan)
app.mount("/static", StaticFiles(directory=STATIC_DIR), name="static")
app.state.config = load_config()
logger.info(
"ebook_search_config_loaded top_k=%s embedding_model=%s rerank_enabled=%s answer_enabled=%s library_paths=%s",
app.state.config.top_k,
app.state.config.embedding_model,
app.state.config.rerank.enabled,
app.state.config.answer_enabled,
len(app.state.config.library_paths),
)
register_page_routes(app)
register_search_routes(app)
register_admin_routes(app)
return app
def serve(
host: Annotated[str, typer.Option("--host", "-h", help="Host to bind to")] = "127.0.0.1",
port: Annotated[int, typer.Option("--port", "-p", help="Port to bind to")] = 8070,
log_level: Annotated[str, typer.Option("--log-level", "-l", help="Log level")] = "INFO",
) -> None:
"""Start the EPUB search server."""
configure_logger(log_level)
uvicorn.run(create_app(), host=host, port=port)
if __name__ == "__main__":
typer.run(serve)
@@ -0,0 +1,16 @@
"""EPUB search web route modules."""
from python.ebook_search.api.routes import admin, page, search
register_admin_routes = admin.register_admin_routes
register_page_routes = page.register_page_routes
register_search_routes = search.register_search_routes
__all__ = [
"admin",
"page",
"register_admin_routes",
"register_page_routes",
"register_search_routes",
"search",
]
+116
View File
@@ -0,0 +1,116 @@
"""Admin routes for the EPUB search web UI."""
from __future__ import annotations
import logging
from dataclasses import replace
from typing import TYPE_CHECKING
from fastapi import APIRouter, Request
from fastapi.responses import HTMLResponse
from sqlalchemy.orm import Session
from python.ebook_search.api.bm25_tasks import schedule_bm25_refresh
from python.ebook_search.api.web import templates
from python.ebook_search.embeddings import embed_missing_chunks, embedding_model_stats
from python.ebook_search.ingest import ingest_configured_paths
if TYPE_CHECKING:
from fastapi import FastAPI
logger = logging.getLogger(__name__)
router = APIRouter(prefix="/admin")
EMBED_ALL_BATCH_SIZE = 32
def register_admin_routes(app: FastAPI) -> None:
"""Register admin routes on the app."""
app.include_router(router)
@router.get("", response_class=HTMLResponse)
def admin(request: Request) -> HTMLResponse:
"""Render the admin page."""
with Session(request.app.state.engine) as session:
stats = embedding_model_stats(session)
logger.info("ebook_admin_page_loaded models=%s", len(stats))
return templates.TemplateResponse(request, "admin.html", {"config": request.app.state.config, "stats": stats})
@router.post("/scan", response_class=HTMLResponse)
def scan_library(request: Request) -> HTMLResponse:
"""Scan configured library paths for EPUB changes."""
try:
with Session(request.app.state.engine) as session:
count = ingest_configured_paths(session, request.app.state.config)
session.commit()
except Exception as error:
logger.exception("ebook_admin_scan_failed")
return templates.TemplateResponse(request, "partials/error.html", {"message": str(error)}, status_code=500)
logger.info("ebook_admin_scan_complete changed_files=%s", count)
if count > 0:
schedule_bm25_refresh(request.app)
return templates.TemplateResponse(request, "partials/admin_status.html", {"message": f"Indexed {count} EPUBs"})
@router.post("/embed-missing", response_class=HTMLResponse)
def embed_missing(request: Request) -> HTMLResponse:
"""Embed chunks missing vectors for the configured model."""
try:
with Session(request.app.state.engine) as session:
count = embed_missing_chunks(session, request.app.state.config)
session.commit()
except Exception as error:
logger.exception("ebook_admin_embed_missing_failed")
return templates.TemplateResponse(request, "partials/error.html", {"message": str(error)}, status_code=500)
logger.info("ebook_admin_embed_missing_complete chunks=%s", count)
return templates.TemplateResponse(
request,
"partials/admin_status.html",
{"message": f"Embedded {count} chunks"},
)
@router.post("/embed-all", response_class=HTMLResponse)
def embed_all(request: Request) -> HTMLResponse:
"""Embed all chunks missing vectors in fixed-size batches."""
total = 0
batches = 0
config = replace(request.app.state.config, embedding_batch_size=EMBED_ALL_BATCH_SIZE)
try:
with Session(request.app.state.engine) as session:
while True:
count = embed_missing_chunks(session, config)
if count == 0:
break
session.commit()
total += count
batches += 1
logger.info(
"ebook_admin_embed_all_batch_complete batch=%s chunks=%s total_chunks=%s",
batches,
count,
total,
)
except Exception as error:
logger.exception(
"ebook_admin_embed_all_failed batches=%s chunks=%s",
batches,
total,
)
return templates.TemplateResponse(
request,
"partials/error.html",
{"message": f"Embed all failed after {total} chunks in {batches} batches: {error}"},
status_code=500,
)
logger.info("ebook_admin_embed_all_complete batches=%s chunks=%s", batches, total)
return templates.TemplateResponse(
request,
"partials/admin_status.html",
{"message": f"Embedded {total} chunks in {batches} batches of {EMBED_ALL_BATCH_SIZE}"},
)
+66
View File
@@ -0,0 +1,66 @@
"""Page routes for the EPUB search web UI."""
from __future__ import annotations
import logging
from typing import TYPE_CHECKING
from fastapi import APIRouter, Request
from fastapi.responses import HTMLResponse
from sqlalchemy import select
from sqlalchemy.orm import Session
from python.ebook_search.api.web import templates
from python.orm.richie import EbookSource
if TYPE_CHECKING:
from fastapi import FastAPI
logger = logging.getLogger(__name__)
router = APIRouter()
def register_page_routes(app: FastAPI) -> None:
"""Register page routes on the app."""
app.include_router(router)
@router.get("/", response_class=HTMLResponse)
def index(request: Request) -> HTMLResponse:
"""Render the search page."""
return templates.TemplateResponse(request, "search.html", {"config": request.app.state.config})
@router.get("/books", response_class=HTMLResponse)
def books(request: Request) -> HTMLResponse:
"""Render the indexed books page."""
with Session(request.app.state.engine) as session:
sources = list(session.scalars(select(EbookSource).order_by(EbookSource.title)).all())
logger.info("ebook_books_page_loaded count=%s", len(sources))
return templates.TemplateResponse(request, "books.html", {"sources": sources})
@router.get("/books/{source_id}", response_class=HTMLResponse)
def book_detail(source_id: int, request: Request) -> HTMLResponse:
"""Render details for one indexed book."""
with Session(request.app.state.engine) as session:
source = session.get(EbookSource, source_id)
if source is not None:
chapter_count = len(source.chapters)
chunk_count = len(source.chunks)
else:
chapter_count = 0
chunk_count = 0
logger.info(
"ebook_book_detail_loaded source_id=%s found=%s chapters=%s chunks=%s",
source_id,
source is not None,
chapter_count,
chunk_count,
)
return templates.TemplateResponse(
request,
"book_detail.html",
{"chapter_count": chapter_count, "chunk_count": chunk_count, "source": source},
)
+66
View File
@@ -0,0 +1,66 @@
"""Search routes for the EPUB search web UI."""
from __future__ import annotations
import logging
from dataclasses import replace
from time import perf_counter
from typing import TYPE_CHECKING, Annotated
from fastapi import APIRouter, Form, Request
from fastapi.responses import HTMLResponse
from python.ebook_search.answer import answer_query
from python.ebook_search.api.web import templates
from python.ebook_search.search import search_ebooks
from python.ebook_search.timing import runtime_step_from_start
if TYPE_CHECKING:
from fastapi import FastAPI
logger = logging.getLogger(__name__)
router = APIRouter()
def register_search_routes(app: FastAPI) -> None:
"""Register search routes on the app."""
app.include_router(router)
@router.post("/search", response_class=HTMLResponse)
def search(
request: Request,
query: Annotated[str, Form()],
rerank: Annotated[str | None, Form()] = None,
) -> HTMLResponse:
"""Run a search and render HTMX results."""
try:
response = search_ebooks(request.app.state.engine, query, request.app.state.config, rerank=rerank == "true")
except Exception as error:
logger.exception("ebook_search_request_failed")
return templates.TemplateResponse(request, "partials/error.html", {"message": str(error)}, status_code=500)
answer_start = perf_counter()
if request.app.state.config.answer_enabled:
try:
answer = answer_query(query, response.results, request.app.state.config)
except RuntimeError as error:
logger.warning("ebook_answer_request_failed_falling_back error=%s", error)
answer = "Answer generation failed. Source chunks are still shown below."
else:
logger.info("ebook_answer_skipped_disabled")
answer = "Answer generation is disabled. Source chunks are shown below."
answer_step_name = "Answer generation" if request.app.state.config.answer_enabled else "Answer skipped"
response = replace(
response,
timings=(*response.timings, runtime_step_from_start(answer_step_name, answer_start)),
)
logger.info(
"ebook_search_request_complete results=%s rank_label=%s runtime_ms=%.1f",
len(response.results),
response.rank_label,
response.total_runtime_ms,
)
return templates.TemplateResponse(request, "partials/results.html", {"answer": answer, "response": response})
+140
View File
@@ -0,0 +1,140 @@
body {
margin: 0;
background: #f7f7f4;
color: #202124;
font-family: system-ui, -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif;
}
main {
max-width: 960px;
margin: 0 auto;
padding: 24px;
}
nav {
display: flex;
gap: 12px;
align-items: center;
margin-bottom: 20px;
}
nav form {
margin: 0;
}
.actions {
display: flex;
flex-wrap: wrap;
gap: 12px;
margin-bottom: 24px;
}
textarea {
display: block;
width: 100%;
margin: 8px 0 12px;
}
button {
padding: 8px 14px;
}
.check {
display: inline-flex;
gap: 8px;
align-items: center;
margin-right: 12px;
}
.rank-label {
margin-top: 24px;
font-weight: 700;
}
.results {
padding-left: 24px;
}
.meta,
.scores,
.status {
color: #626a73;
}
.scores {
display: flex;
flex-wrap: wrap;
gap: 8px;
margin: 12px 0;
}
.scores div {
display: inline-flex;
gap: 4px;
align-items: baseline;
}
.scores dt {
font-weight: 700;
}
.scores dd {
margin: 0;
}
.runtime {
margin-top: 16px;
}
.timing-chart {
display: grid;
gap: 8px;
padding: 0;
list-style: none;
}
.timing-chart li {
display: grid;
grid-template-columns: minmax(150px, 1fr) minmax(160px, 2fr) auto auto;
gap: 8px;
align-items: center;
}
.timing-bar {
height: 10px;
overflow: hidden;
background: #e5e5df;
}
.timing-bar span {
display: block;
height: 100%;
background: #3767c8;
}
.timing-value,
.timing-remaining {
color: #626a73;
font-variant-numeric: tabular-nums;
}
table {
width: 100%;
border-collapse: collapse;
}
th,
td {
padding: 8px;
border-bottom: 1px solid #d8d8d2;
text-align: left;
}
th {
font-weight: 700;
}
.error {
color: #9f1d20;
font-weight: 700;
}
@@ -0,0 +1,57 @@
<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="UTF-8">
<meta name="viewport" content="width=device-width, initial-scale=1.0">
<title>EPUB Admin</title>
<script src="https://unpkg.com/htmx.org@2.0.4"></script>
<link rel="stylesheet" href="/static/style.css">
</head>
<body>
<main>
<nav>
<a href="/">Search</a>
<a href="/books">Books</a>
<a href="/admin">Admin</a>
</nav>
<h1>Admin</h1>
<section id="admin-status"></section>
<section class="actions">
<form hx-post="/admin/scan" hx-target="#admin-status" hx-swap="innerHTML">
<button type="submit">Scan</button>
</form>
<form hx-post="/admin/embed-missing" hx-target="#admin-status" hx-swap="innerHTML">
<button type="submit">Embed</button>
</form>
<form hx-post="/admin/embed-all" hx-target="#admin-status" hx-swap="innerHTML">
<button type="submit">Embed all</button>
</form>
</section>
<section>
<h2>Embeddings</h2>
<table>
<thead>
<tr>
<th>Model</th>
<th>Dimensions</th>
<th>Embedded</th>
<th>Missing</th>
<th>Total chunks</th>
</tr>
</thead>
<tbody>
{% for item in stats %}
<tr>
<td>{{ item.model_name }}</td>
<td>{{ item.dimension }}</td>
<td>{{ item.embedded_chunks }}</td>
<td>{{ item.missing_chunks }}</td>
<td>{{ item.total_chunks }}</td>
</tr>
{% endfor %}
</tbody>
</table>
</section>
</main>
</body>
</html>
@@ -0,0 +1,32 @@
<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="UTF-8">
<meta name="viewport" content="width=device-width, initial-scale=1.0">
<title>{% if source %}{{ source.title }}{% else %}Book not found{% endif %}</title>
<link rel="stylesheet" href="/static/style.css">
</head>
<body>
<main>
<nav>
<a href="/">Search</a>
<a href="/books">Books</a>
<a href="/admin">Admin</a>
</nav>
{% if source %}
<h1>{{ source.title }}</h1>
<p class="meta">{{ source.author or "Unknown author" }}</p>
<dl>
<dt>File</dt>
<dd>{{ source.file_path }}</dd>
<dt>Chapters</dt>
<dd>{{ chapter_count }}</dd>
<dt>Chunks</dt>
<dd>{{ chunk_count }}</dd>
</dl>
{% else %}
<h1>Book not found</h1>
{% endif %}
</main>
</body>
</html>
@@ -0,0 +1,31 @@
<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="UTF-8">
<meta name="viewport" content="width=device-width, initial-scale=1.0">
<title>EPUB Books</title>
<link rel="stylesheet" href="/static/style.css">
</head>
<body>
<main>
<nav>
<a href="/">Search</a>
<a href="/books">Books</a>
<a href="/admin">Admin</a>
</nav>
<h1>Books</h1>
{% if sources %}
<ol class="results">
{% for source in sources %}
<li>
<h2><a href="/books/{{ source.id }}">{{ source.title }}</a></h2>
<p class="meta">{{ source.author or "Unknown author" }}</p>
</li>
{% endfor %}
</ol>
{% else %}
<p>No EPUBs indexed.</p>
{% endif %}
</main>
</body>
</html>
@@ -0,0 +1 @@
<p class="status">{{ message }}</p>
@@ -0,0 +1 @@
<p class="error">{{ message }}</p>
@@ -0,0 +1,74 @@
<div class="rank-label">{{ response.rank_label }}</div>
{% if response.timings %}
<section class="runtime">
<h2>Runtime</h2>
<p class="meta">Total {{ "%.1f"|format(response.total_runtime_ms) }} ms</p>
<ol class="timing-chart">
{% set total = response.total_runtime_ms %}
{% set ns = namespace(remaining=total) %}
{% for step in response.timings %}
{% set width = (step.duration_ms / total * 100) if total else 0 %}
{% if step.counts_toward_total %}
{% set ns.remaining = ns.remaining - step.duration_ms %}
{% endif %}
<li>
<span class="timing-label">{{ step.name }}</span>
<span class="timing-bar"><span style="width: {{ "%.2f"|format(width) }}%"></span></span>
<span class="timing-value">{{ "%.1f"|format(step.duration_ms) }} ms</span>
<span class="timing-remaining">{{ "%.1f"|format([ns.remaining, 0]|max) }} ms left</span>
</li>
{% endfor %}
</ol>
</section>
{% endif %}
<section class="answer">
<h2>Answer</h2>
<p>{{ answer }}</p>
</section>
{% if response.results %}
<ol class="results">
{% for result in response.results %}
<li>
<h2>{{ result.source_title }}</h2>
<p class="meta">
{% if result.source_author %}{{ result.source_author }}{% endif %}
{% if result.chapter_title %} · {{ result.chapter_title }}{% endif %}
{% if result.page_label %} · page {{ result.page_label }}{% endif %}
</p>
<p>{{ result.text }}</p>
<dl class="scores">
<div>
<dt>final</dt>
<dd>{{ "%.3f"|format(result.score) }}</dd>
</div>
{% if result.rerank_score is not none %}
<div>
<dt>rerank</dt>
<dd>{{ "%.3f"|format(result.rerank_score) }}</dd>
</div>
{% endif %}
{% if result.vector_score is not none %}
<div>
<dt>vector cosine</dt>
<dd>{{ "%.3f"|format(result.vector_score) }}</dd>
</div>
{% endif %}
{% if result.bm25_score is not none %}
<div>
<dt>BM25</dt>
<dd>{{ "%.6f"|format(result.bm25_score) }}</dd>
</div>
{% endif %}
{% if result.fused_score is not none %}
<div>
<dt>RRF</dt>
<dd>{{ "%.3f"|format(result.fused_score) }}</dd>
</div>
{% endif %}
</dl>
</li>
{% endfor %}
</ol>
{% else %}
<p>No results.</p>
{% endif %}
@@ -0,0 +1,30 @@
<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="UTF-8">
<meta name="viewport" content="width=device-width, initial-scale=1.0">
<title>EPUB Search</title>
<script src="https://unpkg.com/htmx.org@2.0.4"></script>
<link rel="stylesheet" href="/static/style.css">
</head>
<body>
<main>
<nav>
<a href="/">Search</a>
<a href="/books">Books</a>
<a href="/admin">Admin</a>
</nav>
<h1>EPUB Search</h1>
<form hx-post="/search" hx-target="#results" hx-swap="innerHTML">
<label for="query">Search</label>
<textarea id="query" name="query" rows="4" required></textarea>
<label class="check">
<input type="checkbox" name="rerank" value="true" {% if config.rerank.enabled %}checked{% endif %}>
Rerank
</label>
<button type="submit">Search</button>
</form>
<section id="results"></section>
</main>
</body>
</html>
+13
View File
@@ -0,0 +1,13 @@
"""Shared web UI resources for EPUB search."""
from __future__ import annotations
from pathlib import Path
from fastapi.templating import Jinja2Templates
PACKAGE_DIR = Path(__file__).resolve().parent
TEMPLATE_DIR = PACKAGE_DIR / "templates"
STATIC_DIR = PACKAGE_DIR / "static"
templates = Jinja2Templates(directory=TEMPLATE_DIR)
+249
View File
@@ -0,0 +1,249 @@
"""Persisted BM25 corpus management."""
from __future__ import annotations
import json
import logging
import shutil
import tempfile
from dataclasses import dataclass
from datetime import UTC, datetime
from functools import cache
from pathlib import Path
from typing import TYPE_CHECKING
import bm25s
from sqlalchemy import func, select, union_all
from python.orm.richie import EbookChapter, EbookChunk, EbookSource
if TYPE_CHECKING:
from sqlalchemy.orm import Session
from python.ebook_search.config import EbookSearchConfig
logger = logging.getLogger(__name__)
MANIFEST_NAME = "manifest.json"
REQUIRED_INDEX_FILES = frozenset(
{
"data.csc.index.npy",
"indices.csc.index.npy",
"indptr.csc.index.npy",
"params.index.json",
"vocab.index.json",
"corpus.jsonl",
}
)
@dataclass(frozen=True)
class BM25Manifest:
"""Metadata describing a persisted BM25 corpus."""
created_at: datetime
db_updated_at: datetime | None
chunk_count: int
@dataclass(frozen=True)
class BM25Corpus:
"""Loaded persisted BM25 corpus and retriever."""
retriever: object | None
records: tuple[dict[str, object], ...]
manifest: BM25Manifest
class BM25CorpusUnavailableError(RuntimeError):
"""Raised when the persisted BM25 corpus cannot be loaded."""
def bm25_index_path(config: EbookSearchConfig) -> Path:
"""Return the configured BM25 index path relative to the current working directory."""
path = Path(config.bm25_index_dir).expanduser()
if path.is_absolute():
return path
return Path.cwd() / path
def ensure_bm25_corpus(session: Session, config: EbookSearchConfig) -> None:
"""Create or refresh the persisted BM25 corpus when it is missing or stale."""
index_path = bm25_index_path(config)
manifest = read_bm25_manifest(index_path)
db_updated_at = corpus_last_updated_at(session)
if not bm25_index_exists(index_path, manifest):
logger.info("ebook_bm25_index_missing path=%s", index_path)
refresh_bm25_corpus(session, config, db_updated_at=db_updated_at)
return
if db_updated_at is not None and manifest is not None and manifest.created_at < db_updated_at:
logger.info(
"ebook_bm25_index_stale path=%s created_at=%s db_updated_at=%s",
index_path,
manifest.created_at.isoformat(),
db_updated_at.isoformat(),
)
refresh_bm25_corpus(session, config, db_updated_at=db_updated_at)
return
logger.info(
"ebook_bm25_index_current path=%s chunks=%s created_at=%s",
index_path,
manifest.chunk_count if manifest else 0,
manifest.created_at.isoformat() if manifest else None,
)
def refresh_bm25_corpus(
session: Session,
config: EbookSearchConfig,
*,
db_updated_at: datetime | None = None,
) -> BM25Manifest:
"""Rebuild and persist the BM25 corpus from the current database chunks."""
index_path = bm25_index_path(config)
records = fetch_bm25_corpus_records(session)
manifest = BM25Manifest(
created_at=datetime.now(tz=UTC),
db_updated_at=db_updated_at if db_updated_at is not None else corpus_last_updated_at(session),
chunk_count=len(records),
)
write_bm25_corpus(index_path, records, manifest)
logger.info(
"ebook_bm25_index_refreshed path=%s chunks=%s created_at=%s note=%s",
index_path,
manifest.chunk_count,
manifest.created_at.isoformat(),
"restart_service_to_use_refreshed_bm25_cache",
)
return manifest
@cache
def load_bm25_corpus(config: EbookSearchConfig) -> BM25Corpus:
"""Load the BM25 corpus into memory once per process.
This cache intentionally does not notice later on-disk corpus refreshes. Restart the service after rebuilding the
BM25 corpus for searches to use the new index.
"""
index_path = bm25_index_path(config)
logger.info(
"ebook_bm25_corpus_cache_load path=%s note=%s",
index_path,
"restart_service_after_bm25_refresh",
)
manifest = read_bm25_manifest(index_path)
if manifest is None or not bm25_index_exists(index_path, manifest):
msg = f"BM25 corpus is not available: {index_path}"
raise BM25CorpusUnavailableError(msg)
if manifest.chunk_count == 0:
return BM25Corpus(retriever=None, records=(), manifest=manifest)
retriever = bm25s.BM25.load(index_path, load_corpus=True, mmap=True)
records = tuple(dict(record) for record in retriever.corpus)
return BM25Corpus(retriever=retriever, records=records, manifest=manifest)
def score_bm25_corpus(query: str, corpus: BM25Corpus, *, limit: int) -> list[tuple[dict[str, object], float]]:
"""Score a query against a loaded BM25 corpus."""
if corpus.retriever is None or not corpus.records:
return []
k = min(limit, len(corpus.records))
documents, scores = corpus.retriever.retrieve(
bm25s.tokenize(query, show_progress=False),
corpus=list(corpus.records),
k=k,
show_progress=False,
)
results: list[tuple[dict[str, object], float]] = []
for document, score in zip(documents[0], scores[0], strict=True):
score_value = float(score)
if score_value <= 0:
continue
results.append((dict(document), score_value))
return results
def fetch_bm25_corpus_records(session: Session) -> list[dict[str, object]]:
"""Fetch BM25 corpus records from the database."""
statement = (
select(
EbookChunk.id.label("chunk_id"),
EbookChunk.text.label("text"),
EbookSource.title.label("source_title"),
EbookSource.author.label("source_author"),
EbookChapter.title.label("chapter_title"),
EbookChunk.page_label.label("page_label"),
func.concat_ws(
" ",
EbookSource.title,
EbookSource.author,
EbookChapter.title,
EbookChunk.search_text,
).label("bm25_text"),
)
.select_from(EbookChunk)
.join(EbookSource, EbookSource.id == EbookChunk.source_id)
.outerjoin(EbookChapter, EbookChapter.id == EbookChunk.chapter_id)
.order_by(EbookChunk.id)
)
return [dict(row) for row in session.execute(statement).mappings()]
def corpus_last_updated_at(session: Session) -> datetime | None:
"""Return the latest source/chapter/chunk update timestamp relevant to BM25 text."""
update_times = union_all(
select(func.max(EbookSource.updated).label("updated")),
select(func.max(EbookChapter.updated).label("updated")),
select(func.max(EbookChunk.updated).label("updated")),
).subquery()
return session.scalar(select(func.max(update_times.c.updated)))
def write_bm25_corpus(index_path: Path, records: list[dict[str, object]], manifest: BM25Manifest) -> None:
"""Write a BM25 corpus and manifest atomically."""
index_path.parent.mkdir(parents=True, exist_ok=True)
temp_path = Path(tempfile.mkdtemp(prefix=f"{index_path.name}.", dir=index_path.parent))
try:
if records:
retriever = bm25s.BM25()
texts = [str(record["bm25_text"]) for record in records]
retriever.index(bm25s.tokenize(texts, show_progress=False), show_progress=False)
retriever.save(temp_path, corpus=records, show_progress=False)
write_bm25_manifest(temp_path, manifest)
if index_path.exists():
shutil.rmtree(index_path)
temp_path.rename(index_path)
except Exception:
shutil.rmtree(temp_path, ignore_errors=True)
raise
def read_bm25_manifest(index_path: Path) -> BM25Manifest | None:
"""Read the BM25 manifest if it exists and is valid."""
manifest_path = index_path / MANIFEST_NAME
if not manifest_path.exists():
return None
body = json.loads(manifest_path.read_text(encoding="utf-8"))
return BM25Manifest(
created_at=datetime.fromisoformat(str(body["created_at"])),
db_updated_at=datetime.fromisoformat(str(body["db_updated_at"])) if body.get("db_updated_at") else None,
chunk_count=int(body["chunk_count"]),
)
def write_bm25_manifest(index_path: Path, manifest: BM25Manifest) -> None:
"""Write the BM25 manifest to an index directory."""
body = {
"created_at": manifest.created_at.isoformat(),
"db_updated_at": manifest.db_updated_at.isoformat() if manifest.db_updated_at else None,
"chunk_count": manifest.chunk_count,
}
(index_path / MANIFEST_NAME).write_text(json.dumps(body, indent=2, sort_keys=True), encoding="utf-8")
def bm25_index_exists(index_path: Path, manifest: BM25Manifest | None) -> bool:
"""Return whether a usable persisted BM25 index exists."""
if manifest is None or not index_path.is_dir():
return False
if manifest.chunk_count == 0:
return True
return all((index_path / file_name).exists() for file_name in REQUIRED_INDEX_FILES)
+117
View File
@@ -0,0 +1,117 @@
"""Configuration for the EPUB search app."""
from __future__ import annotations
from dataclasses import dataclass
from os import getenv
def getenv_bool(name: str, *, default: bool) -> bool:
"""Read a boolean environment variable with a default fallback."""
value = getenv(name)
if value is None:
return default
return value.strip().lower() in {"1", "true", "yes", "on"}
def getenv_int(name: str, *, default: int) -> int:
"""Read an integer environment variable with a default fallback."""
value = getenv(name)
if value is None or not value.strip():
return default
return int(value)
@dataclass(frozen=True)
class RerankConfig:
"""vLLM reranker settings."""
enabled: bool = False
base_url: str = "http://192.168.90.25:8001"
model: str = "qwen3-reranker-06b"
candidates: int = 24
timeout_seconds: float = 30.0
@dataclass(frozen=True)
class EbookSearchConfig:
"""Runtime settings for EPUB search."""
rerank: RerankConfig
top_k: int = 12
library_paths: tuple[str, ...] = ()
vllm_base_url: str = "https://ollama.com/v1"
vllm_api_key: str = "not-needed"
chat_model: str = "deepseek-v4-flash"
answer_enabled: bool = True
embedding_base_url: str = "http://192.168.90.25:8000/v1"
embedding_api_key: str = "not-needed"
embedding_model: str = "qwen3-embedding-0.6b"
embedding_batch_size: int = 32
bm25_index_dir: str = ".ebook_search_bm25"
bm25_refresh_delay_seconds: int = 60
def load_rerank_config() -> RerankConfig:
"""Load reranker config from environment variables."""
return RerankConfig(
enabled=getenv_bool("EBOOK_SEARCH_RERANK_ENABLED", default=False),
base_url=getenv("EBOOK_SEARCH_RERANK_BASE_URL", "http://192.168.90.25:8001"),
model=getenv("EBOOK_SEARCH_RERANK_MODEL", "qwen3-reranker-06b"),
candidates=getenv_int("EBOOK_SEARCH_RERANK_CANDIDATES", default=24),
timeout_seconds=float(getenv_int("EBOOK_SEARCH_RERANK_TIMEOUT_SECONDS", default=30)),
)
def load_config() -> EbookSearchConfig:
"""Load EPUB search config from environment variables."""
return EbookSearchConfig(
rerank=load_rerank_config(),
top_k=getenv_int("EBOOK_SEARCH_TOP_K", default=12),
library_paths=library_paths_from_env(),
vllm_base_url=getenv("EBOOK_SEARCH_VLLM_BASE_URL", "https://ollama.com/v1"),
vllm_api_key=getenv("EBOOK_SEARCH_VLLM_API_KEY") or getenv("OLLAMA_API_KEY") or "not-needed",
chat_model=getenv("EBOOK_SEARCH_CHAT_MODEL", "deepseek-v4-flash"),
answer_enabled=getenv_bool("EBOOK_SEARCH_ANSWER_ENABLED", default=True),
embedding_base_url=getenv("EBOOK_SEARCH_EMBEDDING_BASE_URL", "http://192.168.90.25:8000/v1"),
embedding_api_key=getenv("EBOOK_SEARCH_EMBEDDING_API_KEY", "not-needed"),
embedding_model=normalize_embedding_model(),
embedding_batch_size=getenv_int("EBOOK_SEARCH_EMBEDDING_BATCH_SIZE", default=32),
bm25_index_dir=getenv("EBOOK_SEARCH_BM25_INDEX_DIR", ".ebook_search_bm25"),
bm25_refresh_delay_seconds=getenv_int("EBOOK_SEARCH_BM25_REFRESH_DELAY_SECONDS", default=60),
)
def normalize_embedding_model(default: str = "qwen3-embedding-0.6b") -> str:
"""Normalize supported embedding aliases to provider model names."""
aliases = {
"Qwen3-Embedding-0.6B": "qwen3-embedding-0.6b",
"Qwen3-Embedding-4B": "qwen3-embedding-4b",
"Qwen3-Embedding-8B": "qwen3-embedding-8b",
"Qwen/Qwen3-Embedding-0.6B": "qwen3-embedding-0.6b",
"Qwen/Qwen3-Embedding-4B": "qwen3-embedding-4b",
"Qwen/Qwen3-Embedding-8B": "qwen3-embedding-8b",
"qwen3-embedding:0.6b": "qwen3-embedding-0.6b",
"qwen3-embedding:4b": "qwen3-embedding-4b",
"qwen3-embedding:8b": "qwen3-embedding-8b",
"qwen3-embedding-0.6b": "qwen3-embedding-0.6b",
"qwen3-embedding-4b": "qwen3-embedding-4b",
"qwen3-embedding-8b": "qwen3-embedding-8b",
}
model = getenv("EBOOK_SEARCH_EMBEDDING_MODEL", default)
standard_model = aliases.get(model)
if standard_model is None:
error = f"Embedding model {model} is not supported. Supported models are {aliases.keys()}"
raise ValueError(error)
return standard_model
def library_paths_from_env() -> tuple[str, ...]:
"""Read configured EPUB library paths from the environment."""
value = getenv("EBOOK_SEARCH_LIBRARY_PATHS")
if value is None:
return ()
return tuple(path for path in value.split(":") if path)
+170
View File
@@ -0,0 +1,170 @@
"""Embedding model helpers."""
from __future__ import annotations
import logging
from dataclasses import dataclass
from typing import TYPE_CHECKING
from sqlalchemy import func, select
from sqlalchemy.dialects.postgresql import insert
from python.ebook_search.llm_interface import request_embeddings
from python.orm.richie import (
EbookChunk,
EbookChunkEmbedding1024,
EbookChunkEmbedding2560,
EbookChunkEmbedding4096,
EbookEmbeddingModel,
)
logger = logging.getLogger(__name__)
if TYPE_CHECKING:
from collections.abc import Sequence
from sqlalchemy.orm import Session
from python.ebook_search.config import EbookSearchConfig
MODEL_DIMENSIONS = {
"qwen3-embedding-0.6b": 1024,
"qwen3-embedding-4b": 2560,
"qwen3-embedding-8b": 4096,
}
def get_embedding_table(
dimension: int,
) -> type[EbookChunkEmbedding1024 | EbookChunkEmbedding2560 | EbookChunkEmbedding4096]:
"""Return the embedding table mapped to an embedding dimension."""
embedding_tables = {
1024: EbookChunkEmbedding1024,
2560: EbookChunkEmbedding2560,
4096: EbookChunkEmbedding4096,
}
table = embedding_tables.get(dimension)
if not table:
msg = f"Embedding dimension {dimension} is not supported"
raise ValueError(msg)
return table
@dataclass(frozen=True)
class EmbeddingModelStats:
"""Embedding coverage for one model."""
model_name: str
dimension: int
embedded_chunks: int
total_chunks: int
@property
def missing_chunks(self) -> int:
"""Return chunks missing this embedding model."""
return max(self.total_chunks - self.embedded_chunks, 0)
def embed_texts(texts: Sequence[str], config: EbookSearchConfig) -> list[list[float]]:
"""Embed text with the configured vLLM embedding model."""
logger.info(
"ebook_embed_request_start base_url=%s model=%s count=%s",
config.embedding_base_url,
config.embedding_model,
len(texts),
)
vectors = request_embeddings(texts, config)
expected_dimension = MODEL_DIMENSIONS[config.embedding_model]
for vector in vectors:
if len(vector) != expected_dimension:
msg = f"Expected {expected_dimension} dimensions, got {len(vector)}"
raise ValueError(msg)
logger.info(
"ebook_embed_request_complete model=%s count=%s dimension=%s",
config.embedding_model,
len(vectors),
expected_dimension,
)
return vectors
def embed_query(query: str, config: EbookSearchConfig) -> list[float]:
"""Embed a search query with the Qwen retrieval instruction."""
instructed_query = f"Instruct: Retrieve relevant passages for the query.\nQuery: {query}"
return embed_texts([instructed_query], config)[0]
def ensure_embedding_models(session: Session) -> None:
"""Ensure supported embedding model rows exist."""
for name, dimension in MODEL_DIMENSIONS.items():
existing = session.scalar(select(EbookEmbeddingModel).where(EbookEmbeddingModel.name == name))
if existing is None:
session.add(EbookEmbeddingModel(name=name, dimension=dimension, is_default=name == "qwen3-embedding-0.6b"))
logger.info("ebook_embedding_model_created model=%s dimension=%s", name, dimension)
session.flush()
def embedding_model_stats(session: Session) -> list[EmbeddingModelStats]:
"""Return embedding coverage counts for every supported model."""
total_chunks = session.scalar(select(func.count(EbookChunk.id))) or 0
models = {
model.name: model
for model in session.scalars(
select(EbookEmbeddingModel)
.where(EbookEmbeddingModel.name.in_(MODEL_DIMENSIONS))
.order_by(EbookEmbeddingModel.name)
)
}
stats: list[EmbeddingModelStats] = []
for model_name, dimension in MODEL_DIMENSIONS.items():
model = models.get(model_name)
embedded_chunks = 0
if model is not None:
table = get_embedding_table(dimension)
embedded_chunks = session.scalar(select(func.count(table.id)).where(table.model_id == model.id)) or 0
stats.append(
EmbeddingModelStats(
model_name=model_name,
dimension=dimension,
embedded_chunks=embedded_chunks,
total_chunks=total_chunks,
)
)
return stats
def embed_missing_chunks(session: Session, config: EbookSearchConfig) -> int:
"""Embed chunks missing embeddings for the configured model."""
ensure_embedding_models(session)
model = session.scalar(select(EbookEmbeddingModel).where(EbookEmbeddingModel.name == config.embedding_model))
if model is None:
supported_models = ", ".join(MODEL_DIMENSIONS)
msg = f"Unknown embedding model: {config.embedding_model}. Supported models: {supported_models}"
raise ValueError(msg)
table = get_embedding_table(model.dimension)
chunks = list(
session.scalars(
select(EbookChunk)
.outerjoin(table, (table.chunk_id == EbookChunk.id) & (table.model_id == model.id))
.where(table.id.is_(None))
.order_by(EbookChunk.id)
.limit(config.embedding_batch_size)
)
)
if not chunks:
logger.info("ebook_embed_missing_none model=%s", config.embedding_model)
return 0
logger.info("ebook_embed_missing_batch_start model=%s count=%s", config.embedding_model, len(chunks))
vectors = embed_texts([chunk.text for chunk in chunks], config)
rows = [
{"chunk_id": chunk.id, "model_id": model.id, "embedding": vector}
for chunk, vector in zip(chunks, vectors, strict=True)
]
statement = insert(table).values(rows).on_conflict_do_nothing(index_elements=["chunk_id", "model_id"])
session.execute(statement)
session.flush()
logger.info("ebook_embed_missing_batch_complete model=%s count=%s", config.embedding_model, len(rows))
return len(rows)
+95
View File
@@ -0,0 +1,95 @@
"""EPUB parsing helpers."""
from __future__ import annotations
import re
from dataclasses import dataclass
from typing import TYPE_CHECKING
from bs4 import BeautifulSoup
from ebooklib import ITEM_DOCUMENT, epub
if TYPE_CHECKING:
from pathlib import Path
WHITESPACE_RE = re.compile(r"\s+")
@dataclass(frozen=True)
class ParsedChapter:
"""Text extracted from one EPUB spine document."""
title: str | None
href: str | None
text: str
page_labels: tuple[str, ...]
@dataclass(frozen=True)
class ParsedEpub:
"""Parsed EPUB metadata and text."""
title: str
author: str | None
language: str | None
publisher: str | None
identifier: str | None
chapters: tuple[ParsedChapter, ...]
def parse_epub(path: Path) -> ParsedEpub:
"""Parse EPUB metadata and spine text."""
book = epub.read_epub(path)
chapters = []
for item in book.get_items_of_type(ITEM_DOCUMENT):
soup = BeautifulSoup(item.get_content(), "html.parser")
title = chapter_title(soup)
page_labels = tuple(extract_page_labels(soup))
text = clean_text(soup.get_text(" "))
if text:
chapters.append(ParsedChapter(title=title, href=item.get_name(), text=text, page_labels=page_labels))
return ParsedEpub(
title=metadata_value(book, "title") or path.stem,
author=metadata_value(book, "creator"),
language=metadata_value(book, "language"),
publisher=metadata_value(book, "publisher"),
identifier=metadata_value(book, "identifier"),
chapters=tuple(chapters),
)
def metadata_value(book: epub.EpubBook, name: str) -> str | None:
"""Return the first non-empty Dublin Core metadata value for a name."""
values = book.get_metadata("DC", name)
if not values:
return None
value = values[0][0]
return str(value).strip() or None
def chapter_title(soup: BeautifulSoup) -> str | None:
"""Extract the best available title from an EPUB document soup."""
heading = soup.find(["h1", "h2", "h3"])
if heading is None:
title = soup.find("title")
if title is None:
return None
return clean_text(title.get_text(" ")) or None
return clean_text(heading.get_text(" ")) or None
def extract_page_labels(soup: BeautifulSoup) -> list[str]:
"""Extract EPUB page-break labels from a document soup."""
labels: list[str] = []
for tag in soup.find_all(attrs={"epub:type": "pagebreak"}):
label = tag.get("title") or tag.get("aria-label") or tag.get_text(" ")
clean = clean_text(str(label))
if clean:
labels.append(clean)
return labels
def clean_text(text: str) -> str:
"""Normalize whitespace in extracted EPUB text."""
return WHITESPACE_RE.sub(" ", text).strip()
+190
View File
@@ -0,0 +1,190 @@
"""EPUB ingestion into Richie DB."""
from __future__ import annotations
import hashlib
import logging
from dataclasses import dataclass
from datetime import UTC, datetime
from pathlib import Path
from typing import TYPE_CHECKING
import tiktoken
from sqlalchemy import or_, select
from python.ebook_search.epub_parse import parse_epub
from python.orm.richie import EbookChapter, EbookChunk, EbookSource
logger = logging.getLogger(__name__)
DEFAULT_CHUNK_TOKENS = 700
DEFAULT_CHUNK_OVERLAP = 100
if TYPE_CHECKING:
from sqlalchemy.orm import Session
from python.ebook_search.config import EbookSearchConfig
from python.ebook_search.epub_parse import ParsedChapter
@dataclass(frozen=True)
class TextChunk:
"""A token-bounded chunk of text."""
text: str
token_start: int
token_count: int
def chunk_text(
text: str,
*,
chunk_tokens: int = DEFAULT_CHUNK_TOKENS,
overlap_tokens: int = DEFAULT_CHUNK_OVERLAP,
) -> list[TextChunk]:
"""Split text into overlapping token chunks."""
if chunk_tokens <= 0:
msg = "chunk_tokens must be positive"
raise ValueError(msg)
if overlap_tokens < 0 or overlap_tokens >= chunk_tokens:
msg = "overlap_tokens must be non-negative and smaller than chunk_tokens"
raise ValueError(msg)
encoding = tiktoken.get_encoding("cl100k_base")
tokens = encoding.encode(text)
if not tokens:
return []
chunks: list[TextChunk] = []
step = chunk_tokens - overlap_tokens
for start in range(0, len(tokens), step):
chunk = tokens[start : start + chunk_tokens]
if not chunk:
continue
chunks.append(
TextChunk(
text=encoding.decode(chunk).strip(),
token_start=start,
token_count=len(chunk),
)
)
if start + chunk_tokens >= len(tokens):
break
return [chunk for chunk in chunks if chunk.text]
def ingest_configured_paths(session: Session, config: EbookSearchConfig) -> int:
"""Ingest every EPUB found under configured library paths."""
count = 0
for library_path in config.library_paths:
path = Path(library_path).expanduser()
logger.info("ebook_ingest_path_start path=%s", path)
if path.is_file() and path.suffix.lower() == ".epub":
count += int(ingest_file(session, path))
elif path.is_dir():
for epub_path in sorted(path.rglob("*.epub")):
count += int(ingest_file(session, epub_path))
else:
logger.warning("ebook_ingest_path_missing path=%s", path)
logger.info("ebook_ingest_paths_complete changed_files=%s configured_paths=%s", count, len(config.library_paths))
return count
def ingest_file(session: Session, path: Path) -> bool:
"""Ingest one EPUB file. Return True when the database changed."""
resolved_path = path.expanduser().resolve()
logger.info("ebook_ingest_file_start path=%s", resolved_path)
file_hash = sha256_file(resolved_path)
existing = find_existing_source(session, resolved_path, file_hash)
if existing is not None and existing.file_sha256 == file_hash:
stat = resolved_path.stat()
existing.file_path = str(resolved_path)
existing.file_mtime = datetime.fromtimestamp(stat.st_mtime, tz=UTC)
existing.file_size = stat.st_size
session.flush()
logger.info("ebook_ingest_file_unchanged source_id=%s path=%s", existing.id, resolved_path)
return False
if existing is not None:
logger.info("ebook_ingest_file_replacing source_id=%s path=%s", existing.id, resolved_path)
session.delete(existing)
session.flush()
stat = resolved_path.stat()
parsed = parse_epub(resolved_path)
source = EbookSource(
title=parsed.title,
author=parsed.author,
language=parsed.language,
publisher=parsed.publisher,
identifier=parsed.identifier,
file_path=str(resolved_path),
file_sha256=file_hash,
file_mtime=datetime.fromtimestamp(stat.st_mtime, tz=UTC),
file_size=stat.st_size,
)
session.add(source)
session.flush()
chunk_index = 0
for spine_index, parsed_chapter in enumerate(parsed.chapters):
chapter = EbookChapter(
source_id=source.id,
spine_index=spine_index,
title=parsed_chapter.title,
href=parsed_chapter.href,
)
session.add(chapter)
session.flush()
chunk_index = add_chapter_chunks(session, source, chapter, parsed_chapter, chunk_index)
session.flush()
logger.info(
"ebook_ingest_file_complete source_id=%s path=%s chapters=%s chunks=%s",
source.id,
resolved_path,
len(parsed.chapters),
chunk_index,
)
return True
def find_existing_source(session: Session, path: Path, file_hash: str) -> EbookSource | None:
"""Find an existing source by canonical path or file hash."""
return session.scalar(
select(EbookSource).where(or_(EbookSource.file_path == str(path), EbookSource.file_sha256 == file_hash))
)
def add_chapter_chunks(
session: Session,
source: EbookSource,
chapter: EbookChapter,
parsed_chapter: ParsedChapter,
chunk_index: int,
) -> int:
"""Add chunk rows for one parsed chapter and return the next chunk index."""
page_label = parsed_chapter.page_labels[0] if parsed_chapter.page_labels else None
for text_chunk in chunk_text(parsed_chapter.text):
session.add(
EbookChunk(
source_id=source.id,
chapter_id=chapter.id,
chunk_index=chunk_index,
text=text_chunk.text,
token_start=text_chunk.token_start,
token_count=text_chunk.token_count,
page_label=page_label,
content_sha256=hashlib.sha256(text_chunk.text.encode()).hexdigest(),
search_text=f"{source.title} {source.author or ''} {chapter.title or ''} {text_chunk.text}",
)
)
chunk_index += 1
return chunk_index
def sha256_file(path: Path) -> str:
"""Calculate the SHA-256 digest for a file."""
digest = hashlib.sha256()
with path.open("rb") as file:
for block in iter(lambda: file.read(1024 * 1024), b""):
digest.update(block)
return digest.hexdigest()
+143
View File
@@ -0,0 +1,143 @@
"""LLM provider HTTP adapters."""
from __future__ import annotations
import logging
from typing import TYPE_CHECKING
import httpx
if TYPE_CHECKING:
from collections.abc import Sequence
from python.ebook_search.config import EbookSearchConfig, RerankConfig
logger = logging.getLogger(__name__)
def auth_headers(api_key: str) -> dict[str, str]:
"""Build authorization headers when an API key is configured."""
if api_key == "not-needed":
return {}
return {"Authorization": f"Bearer {api_key}"}
def request_embeddings(texts: Sequence[str], config: EbookSearchConfig) -> list[list[float]]:
"""Request embeddings from the configured OpenAI-compatible endpoint."""
try:
response = httpx.post(
f"{config.embedding_base_url.rstrip('/')}/embeddings",
headers=auth_headers(config.embedding_api_key),
json={"model": config.embedding_model, "input": list(texts)},
timeout=60,
)
response.raise_for_status()
return embedding_vectors_from_response(response.json())
except (httpx.HTTPError, ValueError, KeyError, TypeError) as error:
logger.exception(
"ebook_embed_request_failed base_url=%s model=%s count=%s",
config.embedding_base_url,
config.embedding_model,
len(texts),
)
msg = f"Embedding request failed. base_url={config.embedding_base_url} model={config.embedding_model}"
raise RuntimeError(msg) from error
def embedding_vectors_from_response(body: object) -> list[list[float]]:
"""Extract embedding vectors from an OpenAI-compatible embedding response."""
if not isinstance(body, dict):
msg = "Embedding response is not an object"
raise TypeError(msg)
data = body["data"]
if not isinstance(data, list):
msg = "Embedding response data is not a list"
raise TypeError(msg)
vectors: list[list[float]] = []
for item in data:
if not isinstance(item, dict):
msg = "Embedding item is not an object"
raise TypeError(msg)
embedding = item["embedding"]
if not isinstance(embedding, list):
msg = "Embedding value is not a list"
raise TypeError(msg)
vectors.append([float(value) for value in embedding])
return vectors
def request_rerank(
query: str,
documents: Sequence[str],
config: RerankConfig,
) -> object | None:
"""Request rerank scores from the configured vLLM endpoint."""
payload = {
"model": config.model,
"query": query,
"documents": list(documents),
}
response = httpx.post(
f"{config.base_url.rstrip('/')}/rerank",
json=payload,
timeout=config.timeout_seconds,
)
response.raise_for_status()
try:
return response.json()
except ValueError:
logger.debug("ebook_rerank_response_invalid_json", extra={"response": response.text})
return None
def request_chat_completion(
config: EbookSearchConfig,
messages: Sequence[dict[str, str]],
) -> str:
"""Request a chat completion from the configured OpenAI-compatible endpoint."""
try:
response = httpx.post(
f"{config.vllm_base_url.rstrip('/')}/chat/completions",
headers=auth_headers(config.vllm_api_key),
json={
"model": config.chat_model,
"messages": list(messages),
"temperature": 0,
},
timeout=60,
)
response.raise_for_status()
return chat_content_from_response(response.json())
except (httpx.HTTPError, ValueError, KeyError, TypeError) as error:
msg = f"Chat request failed. base_url={config.vllm_base_url} model={config.chat_model}"
raise RuntimeError(msg) from error
def chat_content_from_response(body: object) -> str:
"""Extract text content from an OpenAI-compatible chat response."""
if not isinstance(body, dict):
msg = "Chat response is not an object"
raise TypeError(msg)
choices = body["choices"]
if not isinstance(choices, list) or not choices:
msg = "Chat response has no choices"
raise ValueError(msg)
first = choices[0]
if not isinstance(first, dict):
msg = "Chat choice is not an object"
raise TypeError(msg)
message = first["message"]
if not isinstance(message, dict):
msg = "Chat message is not an object"
raise TypeError(msg)
content = message.get("content") or ""
if not isinstance(content, str):
msg = "Chat content is not text"
raise TypeError(msg)
return content
+127
View File
@@ -0,0 +1,127 @@
"""vLLM-backed optional reranking."""
from __future__ import annotations
import logging
from dataclasses import dataclass, replace
from typing import TYPE_CHECKING
from python.ebook_search.llm_interface import request_rerank
if TYPE_CHECKING:
from python.ebook_search.config import RerankConfig
from python.ebook_search.search import SearchResult
logger = logging.getLogger(__name__)
@dataclass(frozen=True)
class RerankResult:
"""A relevance score for one candidate chunk."""
chunk_id: int
score: float
def rerank_chunks(query: str, candidates: list[SearchResult], config: RerankConfig) -> list[SearchResult]:
"""Rerank candidates with a vLLM rerank endpoint."""
if not candidates:
return []
logger.info(
"ebook_rerank_request_start base_url=%s model=%s candidates=%s",
config.base_url,
config.model,
len(candidates),
)
scores = score_candidates(query, candidates, config)
results = sorted(
(
replace(
result,
score=final_rerank_score(result, scores[result.chunk_id].score, candidates),
rerank_score=scores[result.chunk_id].score,
)
for result in candidates
),
key=lambda result: result.score,
reverse=True,
)
logger.info(
"ebook_rerank_request_complete base_url=%s model=%s candidates=%s",
config.base_url,
config.model,
len(results),
)
return results
def score_candidates(
query: str,
candidates: list[SearchResult],
config: RerankConfig,
) -> dict[int, RerankResult]:
"""Score candidate chunks with the configured rerank API."""
body = request_rerank(query, [candidate.text for candidate in candidates], config)
if body is None:
return zero_rerank_scores(candidates)
scores = parse_vllm_scores(body, candidates)
for result in scores.values():
logger.debug("ebook_rerank_candidate_scored chunk_id=%s score=%s", result.chunk_id, result.score)
return scores
def parse_vllm_scores(body: object, candidates: list[SearchResult]) -> dict[int, RerankResult]:
"""Parse vLLM rerank scores into chunk-id keyed results."""
if not isinstance(body, dict):
logger.debug("ebook_rerank_response_not_object", extra={"response": body})
return zero_rerank_scores(candidates)
results = body.get("results") or body.get("data")
if not isinstance(results, list):
logger.debug("ebook_rerank_response_missing_results", extra={"response": body})
return zero_rerank_scores(candidates)
scores = zero_rerank_scores(candidates)
for item in results:
if not isinstance(item, dict):
continue
index = item.get("index")
score = item.get("relevance_score", item.get("score"))
if not isinstance(index, int) or index < 0 or index >= len(candidates):
continue
if not isinstance(score, int | float):
continue
chunk_id = candidates[index].chunk_id
scores[chunk_id] = RerankResult(chunk_id=chunk_id, score=clamp_score(float(score)))
return scores
def zero_rerank_scores(candidates: list[SearchResult]) -> dict[int, RerankResult]:
"""Return zero relevance scores for all candidate chunks."""
return {candidate.chunk_id: RerankResult(chunk_id=candidate.chunk_id, score=0.0) for candidate in candidates}
def clamp_score(score: float) -> float:
"""Clamp a rerank score into the supported 0.0 to 1.0 range."""
return min(max(score, 0.0), 1.0)
def final_rerank_score(result: SearchResult, rerank_score: float, candidates: list[SearchResult]) -> float:
"""Combine rerank relevance with normalized hybrid retrieval evidence."""
return rerank_score * normalized_hybrid_score(result, candidates)
def normalized_hybrid_score(result: SearchResult, candidates: list[SearchResult]) -> float:
"""Normalize a candidate hybrid score against the rerank candidate set."""
hybrid_scores = [
candidate.fused_score if candidate.fused_score is not None else candidate.score for candidate in candidates
]
low = min(hybrid_scores)
high = max(hybrid_scores)
if high == low:
return 1.0
score = result.fused_score if result.fused_score is not None else result.score
return (score - low) / (high - low)
+371
View File
@@ -0,0 +1,371 @@
"""Hybrid search orchestration."""
from __future__ import annotations
import logging
import re
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass, replace
from typing import TYPE_CHECKING
from pgvector.sqlalchemy import Vector
from sqlalchemy import literal, select
from sqlalchemy.orm import Session
from python.ebook_search.bm25_corpus import (
load_bm25_corpus,
score_bm25_corpus,
)
from python.ebook_search.embeddings import MODEL_DIMENSIONS, embed_query, get_embedding_table
from python.ebook_search.rerank import rerank_chunks
from python.ebook_search.timing import RuntimeStep, timed_result
from python.orm.richie import (
EbookChapter,
EbookChunk,
EbookEmbeddingModel,
EbookSource,
)
if TYPE_CHECKING:
from collections.abc import Mapping
from sqlalchemy.engine import Engine
from python.ebook_search.config import EbookSearchConfig
logger = logging.getLogger(__name__)
BM25_CANDIDATE_LIMIT = 120
@dataclass(frozen=True)
class SearchResult:
"""One source chunk returned by search."""
chunk_id: int
text: str
source_title: str
score: float = 0.0
vector_score: float | None = None
bm25_score: float | None = None
fused_score: float | None = None
rerank_score: float | None = None
source_author: str | None = None
chapter_title: str | None = None
page_label: str | None = None
rank_source: str = "Hybrid"
@dataclass(frozen=True)
class SearchResponse:
"""Search output for the UI."""
query: str
results: list[SearchResult]
rank_label: str
timings: tuple[RuntimeStep, ...] = ()
@property
def total_runtime_ms(self) -> float:
"""Return total measured runtime for the response."""
return sum(step.duration_ms for step in self.timings if step.counts_toward_total)
@dataclass(frozen=True)
class RetrievalResponse:
"""Parallel retrieval output for vector and BM25 candidates."""
vector_results: list[SearchResult]
lexical_results: list[SearchResult]
timings: tuple[RuntimeStep, ...]
def search_ebooks(
engine: Engine,
query: str,
config: EbookSearchConfig,
*,
rerank: bool = False,
) -> SearchResponse:
"""Run hybrid vector/BM25 search and optional reranking."""
if not query.strip():
logger.info("ebook_search_empty_query")
return SearchResponse(query=query, results=[], rank_label="Hybrid")
logger.info("ebook_search_start query_length=%s rerank=%s", len(query), rerank)
timings: list[RuntimeStep] = []
retrieval_query, timing = timed_result("Query preparation", retrieval_query_from_text, query)
timings.append(timing)
retrieval, timing = timed_result(
"Hybrid retrieval",
parallel_retrieval,
engine,
retrieval_query,
config,
)
timings.extend(retrieval.timings)
timings.append(timing)
fused, timing = timed_result(
"Reciprocal rank fusion",
reciprocal_rank_fusion,
retrieval.vector_results,
retrieval.lexical_results,
)
timings.append(timing)
if config.rerank.enabled and rerank:
response, timing = timed_result("Rerank", apply_rerank, query, fused, config)
else:
response, timing = timed_result("Rerank skipped", skip_rerank, query, fused, config)
timings.append(timing)
response = replace(response, timings=tuple(timings))
logger.info(
"ebook_search_complete vector_candidates=%s lexical_candidates=%s "
"fused_candidates=%s returned=%s rank_label=%s runtime_ms=%.1f",
len(retrieval.vector_results),
len(retrieval.lexical_results),
len(fused),
len(response.results),
response.rank_label,
response.total_runtime_ms,
)
return response
def parallel_retrieval(engine: Engine, query: str, config: EbookSearchConfig) -> RetrievalResponse:
"""Run vector and BM25 candidate retrieval concurrently with separate database sessions."""
with ThreadPoolExecutor(max_workers=2, thread_name_prefix="ebook-search") as executor:
vector_future = executor.submit(
timed_result,
"Embedding + vector search",
vector_candidates,
engine,
query,
config,
)
bm25_future = executor.submit(
timed_result,
"BM25 search",
bm25_candidates,
query,
config,
)
vector_results, vector_timing = vector_future.result()
lexical_results, lexical_timing = bm25_future.result()
logger.info(
"ebook_parallel_retrieval_complete vector_candidates=%s lexical_candidates=%s",
len(vector_results),
len(lexical_results),
)
return RetrievalResponse(
vector_results=vector_results,
lexical_results=lexical_results,
timings=(
replace(vector_timing, counts_toward_total=False),
replace(lexical_timing, counts_toward_total=False),
),
)
def skip_rerank(
query: str,
candidates: list[SearchResult],
config: EbookSearchConfig,
) -> SearchResponse:
"""Return fused hybrid results without reranking."""
logger.info("ebook_rerank_skipped candidates=%s", len(candidates))
return SearchResponse(query=query, results=candidates[: config.top_k], rank_label="Hybrid")
def apply_rerank(
query: str,
candidates: list[SearchResult],
config: EbookSearchConfig,
) -> SearchResponse:
"""Rerank already-fused hybrid candidates."""
reranked = rerank_chunks(query, candidates[: config.rerank.candidates], config.rerank)
logger.info(
"ebook_rerank_complete input_candidates=%s returned=%s",
min(len(candidates), config.rerank.candidates),
len(reranked),
)
return SearchResponse(
query=query,
results=[replace(result, rank_source="Hybrid + rerank") for result in reranked[: config.top_k]],
rank_label="Hybrid + rerank",
)
def vector_candidates(engine: Engine, query: str, config: EbookSearchConfig) -> list[SearchResult]:
"""Return pgvector cosine candidates for a normalized query."""
with Session(engine) as session:
model = session.scalar(select(EbookEmbeddingModel).where(EbookEmbeddingModel.name == config.embedding_model))
if model is None:
msg = f"Embedding model is not registered: {config.embedding_model}"
raise ValueError(msg)
expected_dimension = MODEL_DIMENSIONS[config.embedding_model]
if model.dimension != expected_dimension:
msg = f"Model row dimension {model.dimension} does not match configured dimension {expected_dimension}"
raise ValueError(msg)
embedding = embed_query(query, config)
limit = max(config.rerank.candidates, config.top_k) * 4
embedding_table = get_embedding_table(model.dimension)
embedding_param = literal(embedding, type_=Vector(model.dimension))
distance = embedding_table.embedding.op("<=>")(embedding_param)
score = (literal(1.0) - distance).label("score")
statement = (
select(
EbookChunk.id.label("chunk_id"),
EbookChunk.text.label("text"),
EbookSource.title.label("source_title"),
EbookSource.author.label("source_author"),
EbookChapter.title.label("chapter_title"),
EbookChunk.page_label.label("page_label"),
score,
)
.select_from(embedding_table)
.join(EbookChunk, EbookChunk.id == embedding_table.chunk_id)
.join(EbookSource, EbookSource.id == EbookChunk.source_id)
.outerjoin(EbookChapter, EbookChapter.id == EbookChunk.chapter_id)
.where(embedding_table.model_id == model.id)
.order_by(distance)
.limit(limit)
)
rows = session.execute(statement).mappings()
results = [search_result_from_row(row) for row in rows]
logger.info(
"ebook_vector_search_complete model=%s dimension=%s candidates=%s",
config.embedding_model,
model.dimension,
len(results),
)
return results
def bm25_candidates(query: str, config: EbookSearchConfig) -> list[SearchResult]:
"""Return BM25-ranked lexical candidates using the persisted corpus."""
corpus = load_bm25_corpus(config)
if not corpus.records:
logger.info("ebook_bm25_search_complete corpus=0 candidates=0")
return []
scored_records = score_bm25_corpus(query, corpus, limit=BM25_CANDIDATE_LIMIT)
results = [
replace(search_result_from_row(record), score=score, vector_score=None, bm25_score=score)
for record, score in scored_records
]
max_score = results[0].bm25_score if results else 0.0
logger.info(
"ebook_bm25_search_complete corpus=%s candidates=%s max_score=%.6f",
len(corpus.records),
len(results),
max_score,
)
return results
def reciprocal_rank_fusion(
vector_results: list[SearchResult],
lexical_results: list[SearchResult],
*,
rank_constant: int = 60,
) -> list[SearchResult]:
"""Fuse vector and lexical rankings with Reciprocal Rank Fusion."""
by_chunk: dict[int, SearchResult] = {}
scores: dict[int, float] = {}
vector_scores: dict[int, float] = {}
bm25_scores: dict[int, float] = {}
for rank, result in enumerate(vector_results, start=1):
by_chunk.setdefault(result.chunk_id, result)
vector_scores[result.chunk_id] = result.vector_score if result.vector_score is not None else result.score
scores[result.chunk_id] = scores.get(result.chunk_id, 0.0) + (1 / (rank_constant + rank))
for rank, result in enumerate(lexical_results, start=1):
by_chunk.setdefault(result.chunk_id, result)
bm25_scores[result.chunk_id] = result.bm25_score if result.bm25_score is not None else result.score
scores[result.chunk_id] = scores.get(result.chunk_id, 0.0) + (1 / (rank_constant + rank))
return sorted(
(
replace(
result,
score=scores[result.chunk_id],
vector_score=vector_scores.get(result.chunk_id),
bm25_score=bm25_scores.get(result.chunk_id),
fused_score=scores[result.chunk_id],
rank_source="Hybrid",
)
for result in by_chunk.values()
),
key=lambda result: result.score,
reverse=True,
)
def search_result_from_row(row: Mapping[str, object]) -> SearchResult:
"""Convert a database row mapping into a search result."""
return SearchResult(
chunk_id=int(row["chunk_id"]),
text=str(row["text"]),
source_title=str(row["source_title"]),
source_author=optional_str(row["source_author"]),
chapter_title=optional_str(row["chapter_title"]),
page_label=optional_str(row["page_label"]),
score=float(row["score"]) if "score" in row else 0.0,
vector_score=float(row["score"]) if "score" in row else None,
)
def optional_str(value: object) -> str | None:
"""Convert nullable database values to optional strings."""
if value is None:
return None
return str(value)
TOKEN_RE = re.compile(r"[A-Za-z0-9_]+")
def tokens(text_value: str) -> list[str]:
"""Extract tokens from a text value.
This is a simple approximation of the tokenization used by PostgreSQL's full-text search,
which is sufficient for BM25 candidate retrieval. It lowercases tokens and includes alphanumeric characters and
underscores.
"""
return [match.group(0).lower() for match in TOKEN_RE.finditer(text_value)]
QUERY_STOP_WORDS = {
"a",
"an",
"and",
"are",
"as",
"at",
"does",
"for",
"in",
"is",
"of",
"the",
"to",
"what",
"when",
"where",
"which",
"who",
"why",
}
def retrieval_query_from_text(query: str) -> str:
"""Remove generic question words while preserving entity and series terms."""
keywords = [token for token in tokens(query) if token not in QUERY_STOP_WORDS]
if not keywords:
return query
return " ".join(keywords)
+36
View File
@@ -0,0 +1,36 @@
"""Runtime timing helpers for EPUB search."""
from __future__ import annotations
from dataclasses import dataclass
from time import perf_counter
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from collections.abc import Callable
@dataclass(frozen=True)
class RuntimeStep:
"""Elapsed runtime for one named search step."""
name: str
duration_ms: float
counts_toward_total: bool = True
def runtime_step_from_start(name: str, start_seconds: float) -> RuntimeStep:
"""Create a runtime step from a prior perf_counter timestamp."""
return RuntimeStep(name=name, duration_ms=(perf_counter() - start_seconds) * 1000)
def timed_result[T, **P](
name: str,
operation: Callable[P, T],
*args: P.args,
**kwargs: P.kwargs,
) -> tuple[T, RuntimeStep]:
"""Run an operation and return its result plus elapsed runtime."""
start_seconds = perf_counter()
result = operation(*args, **kwargs)
return result, runtime_step_from_start(name, start_seconds)
+347
View File
@@ -0,0 +1,347 @@
"""Small Gitea API client for repository automation."""
from __future__ import annotations
from dataclasses import dataclass
from typing import Self
from urllib.parse import quote
import httpx
DEFAULT_PAGE_SIZE = 100
EXPECTED_NO_CONTENT = 204
EXPECTED_CREATED = 201
EXPECTED_OK = 200
@dataclass(frozen=True)
class CreatedIssue:
"""Issue data returned by Gitea."""
number: int | None
html_url: str | None
title: str
@dataclass(frozen=True)
class PullRequest:
"""Pull request data returned by Gitea."""
number: int
title: str
html_url: str | None
labels: tuple[str, ...]
head_branch: str | None
base_branch: str | None
@dataclass(frozen=True)
class WorkflowJob:
"""Workflow job data returned by Gitea Actions."""
id: int
name: str
run_id: int | None
status: str | None
conclusion: str | None
class GiteaError(RuntimeError):
"""Raised when Gitea rejects an API request."""
def split_repo_name(repo: str) -> tuple[str, str]:
"""Split an owner/repo string into its parts."""
owner, separator, repo_name = repo.partition("/")
if not separator or not owner or not repo_name:
msg = f"Invalid repository name: {repo}"
raise ValueError(msg)
return owner, repo_name
class GiteaClient:
"""HTTP client for the subset of Gitea APIs used in this repository."""
def __init__(
self,
*,
base_url: str,
token: str,
timeout: int = 30,
transport: httpx.BaseTransport | None = None,
) -> None:
"""Initialize the Gitea client."""
self._client = httpx.Client(
base_url=base_url.rstrip("/"),
timeout=timeout,
headers={"Authorization": f"token {token}"},
transport=transport,
)
def create_issue(
self,
*,
owner: str,
repo: str,
title: str,
body: str,
labels: list[int] | None = None,
) -> CreatedIssue:
"""Create a Gitea issue."""
payload: dict[str, object] = {"title": title, "body": body, "labels": labels or []}
response = self._request(
"POST",
f"/api/v1/repos/{owner}/{repo}/issues",
expected_statuses={EXPECTED_CREATED},
json=payload,
)
data = response.json()
return CreatedIssue(
number=_optional_int(data.get("number")),
html_url=_optional_str(data.get("html_url")),
title=str(data.get("title", title)),
)
def resolve_label_ids(self, *, owner: str, repo: str, labels: list[str]) -> list[int]:
"""Resolve label names to Gitea label IDs."""
if not labels:
return []
available_labels: dict[str, int] = {}
page = 1
while True:
response = self._request(
"GET",
f"/api/v1/repos/{owner}/{repo}/labels",
params={"page": page, "limit": DEFAULT_PAGE_SIZE},
)
batch = response.json()
if not batch:
break
for label in batch:
label_name = str(label.get("name", ""))
label_id = _optional_int(label.get("id"))
if label_name and label_id is not None:
available_labels[label_name] = label_id
if len(batch) < DEFAULT_PAGE_SIZE:
break
page += 1
missing = [label for label in labels if label not in available_labels]
if missing:
missing_names = ", ".join(sorted(missing))
msg = f"Missing Gitea labels: {missing_names}"
raise GiteaError(msg)
return [available_labels[label] for label in labels]
def list_open_pull_requests(
self,
*,
owner: str,
repo: str,
labels: list[str] | None = None,
head: str | None = None,
) -> list[PullRequest]:
"""List open pull requests for a repository."""
expected_labels = set(labels or [])
pull_requests: list[PullRequest] = []
page = 1
while True:
response = self._request(
"GET",
f"/api/v1/repos/{owner}/{repo}/pulls",
params={"state": "open", "page": page, "limit": DEFAULT_PAGE_SIZE},
)
batch = response.json()
if not batch:
break
for item in batch:
pull_request = _pull_request_from_api(item)
if head and pull_request.head_branch != head:
continue
if expected_labels and not expected_labels.issubset(set(pull_request.labels)):
continue
pull_requests.append(pull_request)
if len(batch) < DEFAULT_PAGE_SIZE:
break
page += 1
return pull_requests
def create_pull_request(
self,
*,
owner: str,
repo: str,
title: str,
body: str,
head: str,
base: str,
labels: list[str] | None = None,
) -> PullRequest:
"""Create a pull request."""
payload: dict[str, object] = {
"title": title,
"body": body,
"head": head,
"base": base,
}
if labels:
payload["labels"] = self.resolve_label_ids(owner=owner, repo=repo, labels=labels)
response = self._request(
"POST",
f"/api/v1/repos/{owner}/{repo}/pulls",
expected_statuses={EXPECTED_CREATED},
json=payload,
)
return _pull_request_from_api(response.json())
def merge_pull_request(
self,
*,
owner: str,
repo: str,
number: int,
merge_method: str = "rebase",
head_commit_id: str | None = None,
delete_branch_after_merge: bool = False,
) -> None:
"""Merge a pull request."""
payload: dict[str, object] = {
"Do": merge_method,
"delete_branch_after_merge": delete_branch_after_merge,
}
if head_commit_id:
payload["head_commit_id"] = head_commit_id
self._request(
"POST",
f"/api/v1/repos/{owner}/{repo}/pulls/{number}/merge",
json=payload,
)
def dispatch_workflow(self, *, owner: str, repo: str, workflow_id: str, ref: str) -> None:
"""Trigger a workflow_dispatch run."""
workflow_path = quote(workflow_id, safe="")
self._request(
"POST",
f"/api/v1/repos/{owner}/{repo}/actions/workflows/{workflow_path}/dispatches",
expected_statuses={EXPECTED_OK, EXPECTED_NO_CONTENT},
json={"ref": ref},
)
def list_run_jobs(self, *, owner: str, repo: str, run_id: str | int) -> list[WorkflowJob]:
"""List workflow jobs for a specific run."""
jobs: list[WorkflowJob] = []
page = 1
while True:
response = self._request(
"GET",
f"/api/v1/repos/{owner}/{repo}/actions/jobs",
params={"page": page, "limit": DEFAULT_PAGE_SIZE},
)
payload = response.json()
batch = payload.get("jobs", [])
if not batch:
break
for item in batch:
if str(item.get("run_id")) != str(run_id):
continue
jobs.append(_workflow_job_from_api(item))
if len(batch) < DEFAULT_PAGE_SIZE:
break
page += 1
return jobs
def download_job_logs(self, *, owner: str, repo: str, job_id: int) -> str:
"""Download logs for a workflow job."""
response = self._request(
"GET",
f"/api/v1/repos/{owner}/{repo}/actions/jobs/{job_id}/logs",
)
return response.text
def close(self) -> None:
"""Close the underlying 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."""
self.close()
def _request(
self,
method: str,
path: str,
*,
expected_statuses: set[int] | None = None,
**kwargs: object,
) -> httpx.Response:
"""Send an HTTP request and validate the response status."""
response = self._client.request(method, path, **kwargs)
statuses = expected_statuses or {EXPECTED_OK}
if response.status_code not in statuses:
msg = f"Gitea request failed ({response.status_code}): {response.text}"
raise GiteaError(msg)
return response
def _pull_request_from_api(data: dict[str, object]) -> PullRequest:
"""Convert Gitea API pull-request data into a dataclass."""
number = _optional_int(data.get("number")) or _optional_int(data.get("index"))
if number is None:
msg = "Gitea pull request payload is missing a number"
raise GiteaError(msg)
labels = tuple(str(label.get("name", "")) for label in data.get("labels", []))
head = data.get("head", {})
base = data.get("base", {})
return PullRequest(
number=number,
title=str(data.get("title", "")),
html_url=_optional_str(data.get("html_url")),
labels=tuple(label for label in labels if label),
head_branch=_optional_str(head.get("ref")) or _optional_str(data.get("head_branch")),
base_branch=_optional_str(base.get("ref")) or _optional_str(data.get("base_branch")),
)
def _workflow_job_from_api(data: dict[str, object]) -> WorkflowJob:
"""Convert Gitea API workflow-job data into a dataclass."""
job_id = _optional_int(data.get("id"))
if job_id is None:
msg = "Gitea workflow job payload is missing an ID"
raise GiteaError(msg)
return WorkflowJob(
id=job_id,
name=str(data.get("name", "")),
run_id=_optional_int(data.get("run_id")),
status=_optional_str(data.get("status")),
conclusion=_optional_str(data.get("conclusion")),
)
def _optional_int(value: object) -> int | None:
"""Convert an API value to an integer when present."""
if value is None:
return None
return int(value)
def _optional_str(value: object) -> str | None:
"""Convert an API value to a string when present."""
if value is None:
return None
return str(value)
+148
View File
@@ -0,0 +1,148 @@
"""Automation helpers for flake.lock pull requests on Gitea."""
from __future__ import annotations
import subprocess
from os import getenv
from typing import Annotated
import typer
from python.gitea import GiteaClient, PullRequest, split_repo_name
DEFAULT_BASE_BRANCH = "main"
DEFAULT_BRANCH = "automation/update-flake-lock"
DEFAULT_GITEA_URL = "https://gitea.tmmworkshop.com"
PR_LABELS = ["dependencies", "automated", "flake_lock_update"]
PR_CHECK_WORKFLOWS = ["build_systems.yml", "treefmt.yml", "pytest.yml"]
PR_TITLE = "Update flake.lock"
PR_BODY = "Automated flake.lock update."
app = typer.Typer(add_completion=False)
def run_cmd(cmd: list[str], *, check: bool = True) -> subprocess.CompletedProcess[str]:
"""Run a subprocess command."""
return subprocess.run(cmd, capture_output=True, text=True, check=check)
def ensure_flake_lock_pull_request(
client: GiteaClient,
*,
owner: str,
repo: str,
branch: str,
base: str,
) -> PullRequest:
"""Return an existing flake.lock PR for the branch or create one."""
pull_requests = client.list_open_pull_requests(owner=owner, repo=repo, head=branch)
if pull_requests:
return pull_requests[0]
return client.create_pull_request(
owner=owner,
repo=repo,
title=PR_TITLE,
body=PR_BODY,
head=branch,
base=base,
labels=PR_LABELS,
)
def find_flake_lock_pull_request(client: GiteaClient, *, owner: str, repo: str) -> PullRequest | None:
"""Find the first open flake.lock pull request."""
pull_requests = client.list_open_pull_requests(owner=owner, repo=repo, labels=["flake_lock_update"])
if not pull_requests:
return None
return pull_requests[0]
def dispatch_pull_request_checks(client: GiteaClient, *, owner: str, repo: str, branch: str) -> None:
"""Dispatch the workflows that normally run for pull requests."""
for workflow in PR_CHECK_WORKFLOWS:
client.dispatch_workflow(owner=owner, repo=repo, workflow_id=workflow, ref=branch)
def has_worktree_changes() -> bool:
"""Return whether `flake.lock` has worktree changes."""
result = run_cmd(["git", "diff", "--quiet", "--", "flake.lock"], check=False)
return result.returncode != 0
def commit_flake_lock_update(*, branch: str) -> None:
"""Commit the updated lock file to the automation branch."""
run_cmd(["git", "config", "user.name", "gitea-actions[bot]"])
run_cmd(["git", "config", "user.email", "gitea-actions@tmmworkshop.com"])
run_cmd(["git", "checkout", "-B", branch])
run_cmd(["git", "add", "flake.lock"])
run_cmd(["git", "commit", "-m", "chore: update flake.lock"])
def push_branch(*, branch: str) -> None:
"""Push the automation branch to origin."""
run_cmd(["git", "push", "origin", f"HEAD:{branch}", "--force"])
def _required_gitea_token() -> str:
"""Read the required Gitea token from the environment."""
token = getenv("GITEA_TOKEN")
if token:
return token
msg = "GITEA_TOKEN environment variable is required"
raise RuntimeError(msg)
@app.command()
def update(
repo: Annotated[str, typer.Option("--repo", help="Gitea repository in owner/repo form")],
base: Annotated[str, typer.Option("--base", help="Base branch")] = DEFAULT_BASE_BRANCH,
branch: Annotated[str, typer.Option("--branch", help="Automation branch")] = DEFAULT_BRANCH,
) -> None:
"""Commit flake.lock changes and ensure a pull request exists."""
if not has_worktree_changes():
typer.echo("No flake.lock changes detected")
return
commit_flake_lock_update(branch=branch)
push_branch(branch=branch)
owner, repo_name = split_repo_name(repo)
with GiteaClient(
base_url=getenv("GITEA_URL", DEFAULT_GITEA_URL),
token=_required_gitea_token(),
) as client:
pull_request = ensure_flake_lock_pull_request(
client,
owner=owner,
repo=repo_name,
branch=branch,
base=base,
)
# We can remove this if Gitea fixes the following issue:
# https://github.com/go-gitea/gitea/issues/33963
dispatch_pull_request_checks(client, owner=owner, repo=repo_name, branch=branch)
typer.echo(pull_request.html_url or f"Pull request #{pull_request.number}")
@app.command()
def merge(
repo: Annotated[str, typer.Option("--repo", help="Gitea repository in owner/repo form")],
) -> None:
"""Merge the first open flake.lock pull request."""
owner, repo_name = split_repo_name(repo)
with GiteaClient(
base_url=getenv("GITEA_URL", DEFAULT_GITEA_URL),
token=_required_gitea_token(),
) as client:
pull_request = find_flake_lock_pull_request(client, owner=owner, repo=repo_name)
if not pull_request:
typer.echo("No open PR found with label flake_lock_update")
return
client.merge_pull_request(owner=owner, repo=repo_name, number=pull_request.number, merge_method="rebase")
typer.echo(f"Merged PR #{pull_request.number}")
if __name__ == "__main__":
app()
+2
View File
@@ -1,10 +1,12 @@
"""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",
+11
View File
@@ -0,0 +1,11 @@
"""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",
]
+52
View File
@@ -0,0 +1,52 @@
"""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)
@@ -0,0 +1,14 @@
"""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",
]
@@ -0,0 +1,66 @@
"""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"),)
@@ -0,0 +1,66 @@
"""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")
@@ -0,0 +1,79 @@
"""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"),
)
+16
View File
@@ -0,0 +1,16 @@
"""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",
]
@@ -0,0 +1,11 @@
"""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",
]
@@ -0,0 +1,33 @@
"""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]
@@ -0,0 +1,17 @@
"""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)
@@ -0,0 +1,71 @@
"""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())
@@ -0,0 +1,13 @@
"""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)"},)
+20 -5
View File
@@ -2,8 +2,8 @@
from __future__ import annotations
from python.orm.richie.audiobook import Audiobook, AudiobookAuthor, AudiobookSeries
from python.orm.richie.base import RichieBase, TableBase, TableBaseBig, TableBaseSmall
from python.orm.richie.congress import Bill, Legislator, Vote, VoteRecord
from python.orm.richie.contact import (
Contact,
ContactNeed,
@@ -11,19 +11,34 @@ from python.orm.richie.contact import (
Need,
RelationshipType,
)
from python.orm.richie.ebook import (
EbookChapter,
EbookChunk,
EbookChunkEmbedding1024,
EbookChunkEmbedding2560,
EbookChunkEmbedding4096,
EbookEmbeddingModel,
EbookSource,
)
__all__ = [
"Bill",
"Audiobook",
"AudiobookAuthor",
"AudiobookSeries",
"Contact",
"ContactNeed",
"ContactRelationship",
"Legislator",
"EbookChapter",
"EbookChunk",
"EbookChunkEmbedding1024",
"EbookChunkEmbedding2560",
"EbookChunkEmbedding4096",
"EbookEmbeddingModel",
"EbookSource",
"Need",
"RelationshipType",
"RichieBase",
"TableBase",
"TableBaseBig",
"TableBaseSmall",
"Vote",
"VoteRecord",
]
+55
View File
@@ -0,0 +1,55 @@
"""Audiobook catalog models."""
from __future__ import annotations
from sqlalchemy import ForeignKey, String, UniqueConstraint
from sqlalchemy.orm import Mapped, mapped_column, relationship
from python.orm.richie.base import TableBase
class AudiobookAuthor(TableBase):
"""Canonical audiobook author."""
__tablename__ = "audiobook_author"
__table_args__ = (UniqueConstraint("name"),)
name: Mapped[str] = mapped_column(String, unique=True)
books: Mapped[list[Audiobook]] = relationship("Audiobook", back_populates="author")
series: Mapped[list[AudiobookSeries]] = relationship("AudiobookSeries", back_populates="author")
class AudiobookSeries(TableBase):
"""Canonical audiobook series."""
__tablename__ = "audiobook_series"
__table_args__ = (UniqueConstraint("author_id", "name"),)
name: Mapped[str] = mapped_column(String)
author_id: Mapped[int] = mapped_column(ForeignKey("main.audiobook_author.id", ondelete="CASCADE"))
author: Mapped[AudiobookAuthor] = relationship("AudiobookAuthor", back_populates="series")
books: Mapped[list[Audiobook]] = relationship("Audiobook", back_populates="series")
class Audiobook(TableBase):
"""Canonical audiobook title."""
__tablename__ = "audiobook"
__table_args__ = (
UniqueConstraint(
"author_id",
"series_id",
"title",
postgresql_nulls_not_distinct=True,
),
)
title: Mapped[str] = mapped_column(String)
author_id: Mapped[int] = mapped_column(ForeignKey("main.audiobook_author.id", ondelete="CASCADE"))
series_id: Mapped[int | None] = mapped_column(ForeignKey("main.audiobook_series.id", ondelete="SET NULL"))
series_index: Mapped[float] = mapped_column(default=0.0)
author: Mapped[AudiobookAuthor] = relationship("AudiobookAuthor", back_populates="books")
series: Mapped[AudiobookSeries | None] = relationship("AudiobookSeries", back_populates="books")
-150
View File
@@ -1,150 +0,0 @@
"""Congress Tracker database models."""
from __future__ import annotations
from datetime import date
from sqlalchemy import ForeignKey, Index, Text, UniqueConstraint
from sqlalchemy.orm import Mapped, mapped_column, relationship
from python.orm.richie.base import RichieBase, TableBase
class Legislator(TableBase):
"""Legislator model - members of Congress."""
__tablename__ = "legislator"
# Natural key - bioguide ID is the authoritative identifier
bioguide_id: Mapped[str] = mapped_column(Text, unique=True, index=True)
# Other IDs for cross-referencing
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] # JSON array stored as string
# Name info
first_name: Mapped[str]
last_name: Mapped[str]
official_full_name: Mapped[str | None]
nickname: Mapped[str | None]
# Bio
birthday: Mapped[date | None]
gender: Mapped[str | None] # M/F
# Current term info (denormalized for query efficiency)
current_party: Mapped[str | None]
current_state: Mapped[str | None]
current_district: Mapped[int | None] # House only
current_chamber: Mapped[str | None] # rep/sen
# Relationships
vote_records: Mapped[list[VoteRecord]] = relationship(
"VoteRecord",
back_populates="legislator",
cascade="all, delete-orphan",
)
class Bill(TableBase):
"""Bill model - legislation introduced in Congress."""
__tablename__ = "bill"
# Composite natural key: congress + bill_type + number
congress: Mapped[int]
bill_type: Mapped[str] # hr, s, hres, sres, hjres, sjres
number: Mapped[int]
# Bill info
title: Mapped[str | None]
title_short: Mapped[str | None]
official_title: Mapped[str | None]
# Status
status: Mapped[str | None]
status_at: Mapped[date | None]
# Sponsor
sponsor_bioguide_id: Mapped[str | None]
# Subjects
subjects_top_term: Mapped[str | None]
# Relationships
votes: Mapped[list[Vote]] = relationship(
"Vote",
back_populates="bill",
)
__table_args__ = (
UniqueConstraint("congress", "bill_type", "number", name="uq_bill_congress_type_number"),
Index("ix_bill_congress", "congress"),
)
class Vote(TableBase):
"""Vote model - roll call votes in Congress."""
__tablename__ = "vote"
# Composite natural key: congress + chamber + session + number
congress: Mapped[int]
chamber: Mapped[str] # house/senate
session: Mapped[int]
number: Mapped[int]
# Vote details
vote_type: Mapped[str | None]
question: Mapped[str | None]
result: Mapped[str | None]
result_text: Mapped[str | None]
# Timing
vote_date: Mapped[date]
# Vote counts (denormalized for efficiency)
yea_count: Mapped[int | None]
nay_count: Mapped[int | None]
not_voting_count: Mapped[int | None]
present_count: Mapped[int | None]
# Related bill (optional - not all votes are on bills)
bill_id: Mapped[int | None] = mapped_column(ForeignKey("main.bill.id"))
# Relationships
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"),
)
class VoteRecord(RichieBase):
"""Association table: Vote <-> Legislator with position."""
__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] # Yea, Nay, Not Voting, Present
# Relationships
vote: Mapped[Vote] = relationship("Vote", back_populates="vote_records")
legislator: Mapped[Legislator] = relationship("Legislator", back_populates="vote_records")
+130
View File
@@ -0,0 +1,130 @@
"""EPUB search models."""
from __future__ import annotations
from datetime import datetime
from pgvector.sqlalchemy import Vector
from sqlalchemy import BigInteger, Boolean, DateTime, ForeignKey, String, UniqueConstraint
from sqlalchemy.orm import Mapped, mapped_column, relationship
from python.orm.richie.base import TableBase, TableBaseBig
class EbookSource(TableBase):
"""One indexed EPUB file."""
__tablename__ = "ebook_source"
__table_args__ = (
UniqueConstraint("file_path"),
UniqueConstraint("file_sha256"),
)
title: Mapped[str]
author: Mapped[str | None]
language: Mapped[str | None]
publisher: Mapped[str | None]
identifier: Mapped[str | None]
file_path: Mapped[str]
file_sha256: Mapped[str] = mapped_column(String(64))
file_mtime: Mapped[datetime] = mapped_column(DateTime(timezone=True))
file_size: Mapped[int] = mapped_column(BigInteger)
chapters: Mapped[list[EbookChapter]] = relationship(
"EbookChapter",
back_populates="source",
cascade="all, delete-orphan",
passive_deletes=True,
)
chunks: Mapped[list[EbookChunk]] = relationship(
"EbookChunk",
back_populates="source",
cascade="all, delete-orphan",
passive_deletes=True,
)
class EbookChapter(TableBase):
"""A chapter or spine document inside an EPUB."""
__tablename__ = "ebook_chapter"
__table_args__ = (UniqueConstraint("source_id", "spine_index"),)
source_id: Mapped[int] = mapped_column(ForeignKey("main.ebook_source.id", ondelete="CASCADE"))
spine_index: Mapped[int]
title: Mapped[str | None]
href: Mapped[str | None]
source: Mapped[EbookSource] = relationship("EbookSource", back_populates="chapters")
chunks: Mapped[list[EbookChunk]] = relationship(
"EbookChunk",
back_populates="chapter",
cascade="all, delete-orphan",
passive_deletes=True,
)
class EbookChunk(TableBaseBig):
"""A searchable text chunk."""
__tablename__ = "ebook_chunk"
__table_args__ = (
UniqueConstraint("source_id", "chunk_index", name="uq_ebook_chunk_source_id_chunk_index"),
UniqueConstraint("source_id", "content_sha256", name="uq_ebook_chunk_source_id_content_sha256"),
)
source_id: Mapped[int] = mapped_column(ForeignKey("main.ebook_source.id", ondelete="CASCADE"))
chapter_id: Mapped[int | None] = mapped_column(ForeignKey("main.ebook_chapter.id", ondelete="SET NULL"))
chunk_index: Mapped[int]
text: Mapped[str]
token_start: Mapped[int]
token_count: Mapped[int]
page_label: Mapped[str | None]
content_sha256: Mapped[str] = mapped_column(String(64))
search_text: Mapped[str]
source: Mapped[EbookSource] = relationship("EbookSource", back_populates="chunks")
chapter: Mapped[EbookChapter | None] = relationship("EbookChapter", back_populates="chunks")
class EbookEmbeddingModel(TableBase):
"""A supported embedding model."""
__tablename__ = "ebook_embedding_model"
name: Mapped[str] = mapped_column(String, unique=True)
dimension: Mapped[int]
is_default: Mapped[bool] = mapped_column(Boolean, default=False)
class EbookChunkEmbedding1024(TableBaseBig):
"""1024-dimensional chunk embedding."""
__tablename__ = "ebook_chunk_embedding_1024"
__table_args__ = (UniqueConstraint("chunk_id", "model_id"),)
chunk_id: Mapped[int] = mapped_column(ForeignKey("main.ebook_chunk.id", ondelete="CASCADE"))
model_id: Mapped[int] = mapped_column(ForeignKey("main.ebook_embedding_model.id", ondelete="CASCADE"))
embedding: Mapped[list[float]] = mapped_column(Vector(1024))
class EbookChunkEmbedding2560(TableBaseBig):
"""2560-dimensional chunk embedding."""
__tablename__ = "ebook_chunk_embedding_2560"
__table_args__ = (UniqueConstraint("chunk_id", "model_id"),)
chunk_id: Mapped[int] = mapped_column(ForeignKey("main.ebook_chunk.id", ondelete="CASCADE"))
model_id: Mapped[int] = mapped_column(ForeignKey("main.ebook_embedding_model.id", ondelete="CASCADE"))
embedding: Mapped[list[float]] = mapped_column(Vector(2560))
class EbookChunkEmbedding4096(TableBaseBig):
"""4096-dimensional chunk embedding."""
__tablename__ = "ebook_chunk_embedding_4096"
__table_args__ = (UniqueConstraint("chunk_id", "model_id"),)
chunk_id: Mapped[int] = mapped_column(ForeignKey("main.ebook_chunk.id", ondelete="CASCADE"))
model_id: Mapped[int] = mapped_column(ForeignKey("main.ebook_embedding_model.id", ondelete="CASCADE"))
embedding: Mapped[list[float]] = mapped_column(Vector(4096))
+17 -19
View File
@@ -63,9 +63,9 @@ class DeviceRegistry:
return
with Session(self.engine) as session:
device = session.execute(
device = session.scalars(
select(SignalDevice).where(SignalDevice.phone_number == phone_number)
).scalar_one_or_none()
).one_or_none()
if device:
if device.safety_number != safety_number and device.trust_level != TrustLevel.BLOCKED:
@@ -99,9 +99,9 @@ class DeviceRegistry:
Returns True if the device was found and verified.
"""
with Session(self.engine) as session:
device = session.execute(
device = session.scalars(
select(SignalDevice).where(SignalDevice.phone_number == phone_number)
).scalar_one_or_none()
).one_or_none()
if not device:
logger.warning(f"Cannot verify unknown device: {phone_number}")
@@ -139,9 +139,9 @@ class DeviceRegistry:
def grant_role(self, phone_number: str, role: Role) -> bool:
"""Add a role to a device. Called by admin over SSH."""
with Session(self.engine) as session:
device = session.execute(
device = session.scalars(
select(SignalDevice).where(SignalDevice.phone_number == phone_number)
).scalar_one_or_none()
).one_or_none()
if not device:
logger.warning(f"Cannot grant role for unknown device: {phone_number}")
@@ -150,7 +150,7 @@ class DeviceRegistry:
if any(record.name == role for record in device.roles):
return True
role_record = session.execute(select(RoleRecord).where(RoleRecord.name == role)).scalar_one_or_none()
role_record = session.scalars(select(RoleRecord).where(RoleRecord.name == role)).one_or_none()
if not role_record:
logger.warning(f"Unknown role: {role}")
@@ -165,9 +165,9 @@ class DeviceRegistry:
def revoke_role(self, phone_number: str, role: Role) -> bool:
"""Remove a role from a device. Called by admin over SSH."""
with Session(self.engine) as session:
device = session.execute(
device = session.scalars(
select(SignalDevice).where(SignalDevice.phone_number == phone_number)
).scalar_one_or_none()
).one_or_none()
if not device:
logger.warning(f"Cannot revoke role for unknown device: {phone_number}")
@@ -182,16 +182,16 @@ class DeviceRegistry:
def set_roles(self, phone_number: str, roles: list[Role]) -> bool:
"""Replace all roles for a device. Called by admin over SSH."""
with Session(self.engine) as session:
device = session.execute(
device = session.scalars(
select(SignalDevice).where(SignalDevice.phone_number == phone_number)
).scalar_one_or_none()
).one_or_none()
if not device:
logger.warning(f"Cannot set roles for unknown device: {phone_number}")
return False
role_names = [str(role) for role in roles]
records = list(session.execute(select(RoleRecord).where(RoleRecord.name.in_(role_names))).scalars().all())
records = session.scalars(select(RoleRecord).where(RoleRecord.name.in_(role_names))).all()
device.roles = records
session.commit()
self._update_cache(phone_number, device)
@@ -203,7 +203,7 @@ class DeviceRegistry:
def list_devices(self) -> list[SignalDevice]:
"""Return all known devices."""
with Session(self.engine) as session:
return list(session.execute(select(SignalDevice)).scalars().all())
return list(session.scalars(select(SignalDevice)).all())
def sync_identities(self) -> None:
"""Pull identity list from signal-cli and record any new ones."""
@@ -226,9 +226,7 @@ class DeviceRegistry:
def _load_device(self, phone_number: str) -> SignalDevice | None:
"""Fetch a device by phone number (with joined roles)."""
with Session(self.engine) as session:
return session.execute(
select(SignalDevice).where(SignalDevice.phone_number == phone_number)
).scalar_one_or_none()
return session.scalars(select(SignalDevice).where(SignalDevice.phone_number == phone_number)).one_or_none()
def _update_cache(self, phone_number: str, device: SignalDevice) -> None:
"""Refresh the cache entry for a device."""
@@ -244,9 +242,9 @@ class DeviceRegistry:
def _set_trust(self, phone_number: str, level: str, log_msg: str | None = None) -> bool:
"""Update the trust level for a device."""
with Session(self.engine) as session:
device = session.execute(
device = session.scalars(
select(SignalDevice).where(SignalDevice.phone_number == phone_number)
).scalar_one_or_none()
).one_or_none()
if not device:
return False
@@ -269,7 +267,7 @@ def sync_roles(engine: Engine) -> None:
expected = {role.value for role in Role}
with Session(engine) as session:
existing = {record.name for record in session.execute(select(RoleRecord)).scalars().all()}
existing = set(session.scalars(select(RoleRecord.name)).all())
to_add = expected - existing
to_remove = existing - expected
+1
View File
@@ -0,0 +1 @@
"""Audiobook tools."""
+471
View File
@@ -0,0 +1,471 @@
"""Convert Audible AAX downloads into Audiobookshelf-friendly M4B files."""
from __future__ import annotations
import json
import logging
import re
import shutil
import subprocess
from concurrent.futures import ThreadPoolExecutor
from dataclasses import asdict, dataclass
from os import getenv
from pathlib import Path # noqa: TC003 This is required for the typer CLI
from typing import TYPE_CHECKING, Annotated, Any
from uuid import uuid7
import typer
from python.common import configure_logger
from python.orm.common import get_postgres_engine
from python.tools.audiobook.metadata_agent import (
AgentConfig,
StandardBookMetadata,
standard_book_metadata,
write_agent_log,
)
if TYPE_CHECKING:
from sqlalchemy.engine import Engine
logger = logging.getLogger(__name__)
SENSITIVE_COMMAND_ARGUMENTS = {"-activation_bytes"}
BOOK_RANGE_PATTERN = re.compile(r"(?:^|-)books?-(?P<start>[1-9]\d*)-(?P<end>[1-9]\d*)(?:-|$)")
@dataclass(frozen=True)
class ConversionConfig:
"""Runtime settings for one conversion command."""
resolved_output: Path
ollama_api_key: str
agent_config: AgentConfig
engine: Engine
activation_bytes: str | None
dry_run: bool
overwrite: bool
work_directory_name: str = ".audible_convert"
dry_run_directory_name: str = "dry-run"
temp_directory_name: str = "tmp"
log_directory_name: str = "logs"
review_directory_name: str = "review"
@dataclass(frozen=True)
class ConcurrentConversionResult:
"""Result from running ffmpeg and metadata resolution together."""
metadata: StandardBookMetadata | None
conversion_error: Exception | None
metadata_error: Exception | None
class CommandExecutionError(RuntimeError):
"""Command failed without exposing sensitive arguments."""
def __init__(self, arguments: list[str], returncode: int) -> None:
"""Create a redacted command failure."""
self.arguments = tuple(arguments)
self.returncode = returncode
command = " ".join(redact_command_arguments(arguments))
super().__init__(f"Command failed with exit code {returncode}: {command}")
def main(
input_directory: Annotated[Path, typer.Argument(help="Directory audible-cli downloads AAX files into.")],
output_directory: Annotated[Path, typer.Argument(help="Audiobook output directory.")],
*,
dry_run: Annotated[
bool,
typer.Option("--dry-run", help="Print planned output files and write marker files without converting."),
] = False,
overwrite: Annotated[bool, typer.Option("--overwrite", help="Overwrite existing M4B files.")] = False,
) -> None:
"""Convert AAX files from a download directory into M4B files."""
configure_logger()
resolved_input = input_directory.resolve(strict=True)
resolved_output = output_directory.resolve()
if not dry_run:
resolved_output.mkdir(parents=True, exist_ok=True)
ollama_api_key = getenv("OLLAMA_API_KEY")
if not ollama_api_key:
msg = "OLLAMA_API_KEY is required for audiobook metadata resolution"
raise RuntimeError(msg)
config = ConversionConfig(
resolved_output=resolved_output,
ollama_api_key=ollama_api_key,
agent_config=AgentConfig(),
engine=get_postgres_engine(name="RICHIE"),
activation_bytes=getenv("AUDIBLE_ACTIVATION_BYTES"),
dry_run=dry_run,
overwrite=overwrite,
)
aax_files = sorted(resolved_input.glob("*.aax"))
if not aax_files:
logger.info("No AAX files found in %s", resolved_input)
return
for aax_file in aax_files:
logger.info("Converting %s", aax_file)
convert_aax_file_with_agent(aax_file, config)
def run_command(arguments: list[str], *, capture: bool = False) -> subprocess.CompletedProcess[str]:
"""Run a command and return the completed process.
Args:
arguments: Command and arguments to run.
capture: Whether to capture stdout and stderr.
Returns:
The completed process.
"""
logger.debug("%s", " ".join(redact_command_arguments(arguments)))
try:
return subprocess.run(arguments, check=True, capture_output=capture, text=True)
except subprocess.CalledProcessError as error:
raise CommandExecutionError(arguments, error.returncode) from error
def redact_command_arguments(arguments: list[str]) -> list[str]:
"""Return command arguments with sensitive values redacted."""
redacted = []
redact_next = False
for argument in arguments:
if redact_next:
redacted.append("<redacted>")
redact_next = False
continue
redacted.append(argument)
redact_next = argument in SENSITIVE_COMMAND_ARGUMENTS
return redacted
def read_metadata(aax_file: Path) -> dict[str, str]:
"""Read ffprobe format tags from an AAX file.
Args:
aax_file: AAX file to inspect.
Returns:
Lower-cased metadata tag names mapped to their values.
"""
completed = run_command(
[
"ffprobe",
"-v",
"quiet",
"-print_format",
"json",
"-show_format",
str(aax_file),
],
capture=True,
)
ffprobe_data: dict[str, Any] = json.loads(completed.stdout)
tags = ffprobe_data.get("format", {}).get("tags", {})
return {str(key).lower(): str(value) for key, value in tags.items()}
def output_stem(metadata: StandardBookMetadata) -> str:
"""Build the output stem for a book.
Args:
metadata: Book metadata.
Returns:
Output stem in author-series_01-title form.
"""
index_slug = series_index_slug(metadata.series_index, metadata.title)
return f"{metadata.author}-{metadata.series}_{index_slug}-{metadata.title}"
def series_index_slug(series_index: float, title: str = "") -> str:
"""Return a filename-safe series index."""
if title_range := title_series_range_slug(series_index, title):
return title_range
index = float(series_index)
if index.is_integer():
return f"{int(index):02}"
return f"{int(index):02}.5"
def title_series_range_slug(series_index: float, title: str) -> str | None:
"""Return a series range slug found in an omnibus title."""
index = float(series_index)
if not index.is_integer():
return None
first_index = int(index)
for match in BOOK_RANGE_PATTERN.finditer(title):
start = int(match.group("start"))
end = int(match.group("end"))
if start == first_index and end > start:
return f"{start:02}-{end:02}"
return None
def metadata_output_path(output_directory: Path, metadata: StandardBookMetadata) -> Path:
"""Build the final M4B path from resolved metadata."""
stem = output_stem(metadata)
return output_directory / stem / f"{stem}.m4b"
def convert_aax_file(
aax_file: Path,
destination: Path,
activation_bytes: str | None,
*,
overwrite: bool,
) -> None:
"""Convert an AAX file into an M4B file.
Args:
aax_file: Source AAX file.
destination: Destination M4B file.
activation_bytes: Optional Audible activation bytes for ffmpeg.
overwrite: Whether to overwrite an existing M4B.
"""
if destination.exists() and not overwrite:
logger.info("Skipping existing file %s", destination)
return
destination.parent.mkdir(parents=True, exist_ok=True)
arguments = ["ffmpeg", "-hide_banner", "-y" if overwrite else "-n"]
if activation_bytes:
arguments.extend(["-activation_bytes", activation_bytes])
arguments.extend(["-i", str(aax_file), "-map_metadata", "0", "-c", "copy", str(destination)])
run_command(arguments)
def write_review_file(
*,
destination: Path | None,
ffprobe_metadata: dict[str, str],
log_file: Path,
metadata: StandardBookMetadata | None,
reason: str,
review_file: Path,
source: Path,
temp_file: Path | None,
) -> None:
"""Write a manual review file for an unresolved conversion."""
review_file.parent.mkdir(parents=True, exist_ok=True)
payload = {
"destination": str(destination) if destination else None,
"ffprobe_metadata": ffprobe_metadata,
"metadata": asdict(metadata) if metadata else None,
"reason": reason,
"source": str(source),
"temp_file": str(temp_file) if temp_file else None,
}
review_file.write_text(json.dumps(payload, indent=2, sort_keys=True), encoding="utf-8")
write_agent_log(log_file, "review_written", path=str(review_file), reason=reason)
def cleanup_temp_output(temp_file: Path) -> None:
"""Remove a run's temporary output directory."""
shutil.rmtree(temp_file.parent, ignore_errors=True)
def dry_run_aax_file_with_agent(
aax_file: Path,
ffprobe_metadata: dict[str, str],
engine: Engine,
config: ConversionConfig,
log_file: Path,
review_file: Path,
) -> None:
"""Resolve and print the planned output path without converting."""
metadata = standard_book_metadata(
aax_file.name,
ffprobe_metadata,
engine,
log_file,
config.ollama_api_key,
config.agent_config,
)
destination = None if metadata.needs_review else metadata_output_path(config.resolved_output, metadata)
if metadata.needs_review:
write_review_file(
destination=destination,
ffprobe_metadata=ffprobe_metadata,
log_file=log_file,
metadata=metadata,
reason="metadata_needs_review",
review_file=review_file,
source=aax_file,
temp_file=None,
)
typer.echo(f"{aax_file} -> REVIEW {review_file}")
else:
stem = output_stem(metadata)
dry_run_file = (
config.resolved_output / config.work_directory_name / config.dry_run_directory_name / stem / f"{stem}.m4b"
)
dry_run_file.parent.mkdir(parents=True, exist_ok=True)
dry_run_file.write_text(f"{destination}\n", encoding="utf-8")
write_agent_log(
log_file,
"dry_run_file_written",
destination=str(destination),
path=str(dry_run_file),
)
typer.echo(f"{aax_file} -> {destination}")
def convert_temp_file_and_resolve_metadata(
aax_file: Path,
temp_file: Path,
ffprobe_metadata: dict[str, str],
config: ConversionConfig,
log_file: Path,
) -> ConcurrentConversionResult:
"""Run ffmpeg and metadata resolution in parallel."""
conversion_error: Exception | None = None
metadata_error: Exception | None = None
metadata: StandardBookMetadata | None = None
with ThreadPoolExecutor(max_workers=2) as executor:
conversion_future = executor.submit(
convert_aax_file,
aax_file,
temp_file,
config.activation_bytes,
overwrite=True,
)
metadata_future = executor.submit(
standard_book_metadata,
aax_file.name,
ffprobe_metadata,
config.engine,
log_file,
config.ollama_api_key,
config.agent_config,
)
conversion_error = conversion_future.exception()
if conversion_error is None:
conversion_future.result()
metadata_error = metadata_future.exception()
if metadata_error is None:
metadata = metadata_future.result()
return ConcurrentConversionResult(
metadata=metadata,
conversion_error=conversion_error,
metadata_error=metadata_error,
)
def convert_aax_file_with_agent(aax_file: Path, config: ConversionConfig) -> None:
"""Convert one AAX file using the metadata agent for the final path."""
run_id = uuid7().hex
log_file = config.resolved_output / config.work_directory_name / config.log_directory_name / f"{run_id}.jsonl"
review_file = config.resolved_output / config.work_directory_name / config.review_directory_name / f"{run_id}.json"
write_agent_log(log_file, "conversion_start", source=str(aax_file), dry_run=config.dry_run)
try:
ffprobe_metadata = read_metadata(aax_file)
except Exception as error:
logger.exception("ffprobe failed")
write_review_file(
destination=None,
ffprobe_metadata={},
log_file=log_file,
metadata=None,
reason=f"ffprobe_failed: {error}",
review_file=review_file,
source=aax_file,
temp_file=None,
)
return
if config.dry_run:
dry_run_aax_file_with_agent(
aax_file,
ffprobe_metadata,
config.engine,
config,
log_file,
review_file,
)
return
temp_file = (
config.resolved_output / config.work_directory_name / config.temp_directory_name / run_id / "converted.m4b"
)
temp_file.parent.mkdir(parents=True, exist_ok=True)
result = convert_temp_file_and_resolve_metadata(aax_file, temp_file, ffprobe_metadata, config, log_file)
if result.conversion_error:
reason = f"ffmpeg_failed: {result.conversion_error}"
write_review_file(
destination=None,
ffprobe_metadata=ffprobe_metadata,
log_file=log_file,
metadata=result.metadata,
reason=reason,
review_file=review_file,
source=aax_file,
temp_file=temp_file if temp_file.exists() else None,
)
return
if result.metadata_error:
write_review_file(
destination=None,
ffprobe_metadata=ffprobe_metadata,
log_file=log_file,
metadata=None,
reason=f"metadata_failed: {result.metadata_error}",
review_file=review_file,
source=aax_file,
temp_file=temp_file,
)
return
if result.metadata is None or result.metadata.needs_review:
write_review_file(
destination=None,
ffprobe_metadata=ffprobe_metadata,
log_file=log_file,
metadata=result.metadata,
reason="metadata_needs_review",
review_file=review_file,
source=aax_file,
temp_file=temp_file,
)
return
destination = metadata_output_path(config.resolved_output, result.metadata)
if destination.exists() and not config.overwrite:
write_agent_log(log_file, "destination_exists", destination=str(destination))
cleanup_temp_output(temp_file)
return
destination.parent.mkdir(parents=True, exist_ok=True)
try:
temp_file.replace(destination)
except Exception as error: # noqa: BLE001
write_review_file(
destination=destination,
ffprobe_metadata=ffprobe_metadata,
log_file=log_file,
metadata=result.metadata,
reason=f"rename_failed: {error}",
review_file=review_file,
source=aax_file,
temp_file=temp_file if temp_file.exists() else None,
)
else:
cleanup_temp_output(temp_file)
write_agent_log(log_file, "conversion_complete", destination=str(destination))
if __name__ == "__main__":
typer.run(main)
+176
View File
@@ -0,0 +1,176 @@
"""Import audiobook catalog authors and series from CSV files."""
from __future__ import annotations
import csv
import logging
from pathlib import Path # noqa: TC003 This is required for the typer CLI
from typing import Annotated
import typer
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.richie import AudiobookAuthor, AudiobookSeries
logger = logging.getLogger(__name__)
AUTHOR_NAME_COLUMN = "author_name"
ID_COLUMN = "id"
NAME_COLUMN = "name"
class CatalogImportError(ValueError):
"""CSV catalog import failed validation."""
def main(
authors_csv: Annotated[Path, typer.Argument(help="CSV with name and optional id.")],
series_csv: Annotated[Path, typer.Argument(help="CSV with name, author_name, and optional id.")],
) -> None:
"""Upsert audiobook authors and series from CSV files."""
configure_logger()
try:
engine = get_postgres_engine(name="RICHIE")
with Session(engine) as session:
author_count = upsert_authors_from_csv(session, authors_csv)
series_count = upsert_series_from_csv(session, series_csv)
session.commit()
except CatalogImportError as error:
typer.echo(str(error), err=True)
raise typer.Exit(code=1) from error
logger.info("Upserted %s authors and %s series", author_count, series_count)
def upsert_authors_from_csv(session: Session, authors_csv: Path) -> int:
"""Upsert authors from a CSV file."""
count = 0
for row_number, row in csv_rows(authors_csv):
name = required_csv_value(row, authors_csv, row_number, NAME_COLUMN)
upsert_author(session, name, csv_id(row, authors_csv, row_number))
count += 1
return count
def upsert_series_from_csv(session: Session, series_csv: Path) -> int:
"""Upsert series from a CSV file."""
count = 0
for row_number, row in csv_rows(series_csv):
series_name = required_csv_value(row, series_csv, row_number, NAME_COLUMN)
author_name = required_csv_value(row, series_csv, row_number, AUTHOR_NAME_COLUMN)
author = find_author_by_name(session, author_name)
if author is None:
msg = f"{series_csv}:{row_number}: author not found: {author_name}"
raise CatalogImportError(msg)
upsert_series(session, series_name, author, csv_id(row, series_csv, row_number))
count += 1
return count
def upsert_author(session: Session, name: str, author_id: int | None) -> AudiobookAuthor:
"""Upsert one author by id or exact name."""
if author_id is not None:
author = session.get(AudiobookAuthor, author_id)
if author is None:
author = AudiobookAuthor(id=author_id, name=name)
session.add(author)
else:
author.name = name
session.flush()
return author
author = find_author_by_name(session, name)
if author is None:
author = AudiobookAuthor(name=name)
session.add(author)
session.flush()
return author
def upsert_series(
session: Session,
name: str,
author: AudiobookAuthor,
series_id: int | None,
) -> AudiobookSeries:
"""Upsert one series by id or exact author/name match."""
if series_id is not None:
series = session.get(AudiobookSeries, series_id)
if series is None:
series = AudiobookSeries(id=series_id, name=name, author=author)
session.add(series)
else:
series.name = name
series.author = author
session.flush()
return series
series = find_series_by_name_and_author(session, name, author.id)
if series is None:
series = AudiobookSeries(name=name, author=author)
session.add(series)
session.flush()
return series
def find_author_by_name(session: Session, name: str) -> AudiobookAuthor | None:
"""Find one author by exact name."""
return session.scalar(select(AudiobookAuthor).where(AudiobookAuthor.name == name))
def find_series_by_name_and_author(
session: Session,
name: str,
author_id: int,
) -> AudiobookSeries | None:
"""Find one series by exact name and author."""
return session.scalar(
select(AudiobookSeries).where(
AudiobookSeries.name == name,
AudiobookSeries.author_id == author_id,
),
)
def csv_rows(csv_path: Path) -> list[tuple[int, dict[str, str | None]]]:
"""Read a CSV file as numbered rows."""
with csv_path.open(newline="", encoding="utf-8") as file:
reader = csv.DictReader(file)
if reader.fieldnames is None:
msg = f"{csv_path}: missing CSV header"
raise CatalogImportError(msg)
return [(row_number, row) for row_number, row in enumerate(reader, start=2)]
def required_csv_value(
row: dict[str, str | None],
csv_path: Path,
row_number: int,
column: str,
) -> str:
"""Read a required CSV value."""
value = row.get(column)
if value and value.strip():
return value.strip()
msg = f"{csv_path}:{row_number}: missing required column value: {column}"
raise CatalogImportError(msg)
def csv_id(row: dict[str, str | None], csv_path: Path, row_number: int) -> int | None:
"""Read an optional id field from a CSV row."""
value = row.get(ID_COLUMN)
if value is None or not value.strip():
return None
try:
return int(value)
except ValueError as error:
msg = f"{csv_path}:{row_number}: id must be an integer: {value}"
raise CatalogImportError(msg) from error
return None
if __name__ == "__main__":
typer.run(main)
+599
View File
@@ -0,0 +1,599 @@
"""LLM tool calling support for audiobook metadata resolution."""
from __future__ import annotations
import json
import re
import time
from collections.abc import Callable
from dataclasses import dataclass
from typing import TYPE_CHECKING
from sqlalchemy import or_, select
from python.orm.richie import Audiobook, AudiobookAuthor, AudiobookSeries
if TYPE_CHECKING:
from pathlib import Path
from sqlalchemy.orm import Session
from python.tools.audiobook.metadata_agent import AgentConfig
CATALOG_SLUG_PATTERN = re.compile(r"^[a-z0-9]+(?:_[a-z0-9]+)*$")
TITLE_SLUG_PATTERN = re.compile(r"^[a-z0-9]+(?:-[a-z0-9]+)*$")
LogWriter = Callable[..., None]
class MetadataResolutionError(ValueError):
"""Metadata resolution failed validation."""
@dataclass(frozen=True)
class EnsuredBook:
"""Book row plus whether it was created."""
book: Audiobook
action: str
class CatalogToolRegistry:
"""Controlled catalog tools exposed to the metadata model."""
def __init__(
self,
session: Session,
log_path: Path,
config: AgentConfig,
write_log: LogWriter,
) -> None:
"""Create a registry bound to one database session and audit log."""
self.session = session
self.log_path = log_path
self.config = config
self.write_log = write_log
self.seen_author_ids: set[int] = set()
self.seen_series_ids: set[int] = set()
self.seen_book_ids: set[int] = set()
self.created_author_ids: set[int] = set()
self.created_series_ids: set[int] = set()
self.created_book_ids: set[int] = set()
def tool_schemas(self) -> list[dict[str, object]]:
"""Return Ollama tool schemas."""
schemas = [
{
"type": "function",
"function": {
"name": "search_authors",
"description": "Search canonical audiobook authors by slug or noisy source text.",
"parameters": {
"type": "object",
"properties": {"query": {"type": "string"}},
"required": ["query"],
},
},
},
{
"type": "function",
"function": {
"name": "search_series",
"description": "Search canonical audiobook series by slug or noisy source text.",
"parameters": {
"type": "object",
"properties": {
"query": {"type": "string"},
"author_id": {"type": ["integer", "null"]},
},
"required": ["query"],
},
},
},
{
"type": "function",
"function": {
"name": "search_books",
"description": "Search canonical audiobook titles with optional author and series filters.",
"parameters": {
"type": "object",
"properties": {
"query": {"type": "string"},
"author_id": {"type": ["integer", "null"]},
"series_id": {"type": ["integer", "null"]},
},
"required": ["query"],
},
},
},
{
"type": "function",
"function": {
"name": "ensure_author",
"description": "Normalize an author name to a catalog slug, then return or create that author.",
"parameters": {
"type": "object",
"properties": {"name": {"type": "string"}},
"required": ["name"],
},
},
},
{
"type": "function",
"function": {
"name": "ensure_series",
"description": "Normalize a series name to a catalog slug, then return or create it for an author.",
"parameters": {
"type": "object",
"properties": {
"name": {"type": "string"},
"author_id": {"type": "integer"},
},
"required": ["name", "author_id"],
},
},
},
{
"type": "function",
"function": {
"name": "ensure_book",
"description": "Normalize a title to a book slug, then return or create it for an author/series.",
"parameters": {
"type": "object",
"properties": {
"title": {"type": "string"},
"author_id": {"type": "integer"},
"series_id": {"type": ["integer", "null"]},
"series_index": {"type": "number", "multipleOf": 0.5},
},
"required": ["title", "author_id", "series_id", "series_index"],
},
},
},
]
enabled_tool_names = set(self.config.tool_names)
return [schema for schema in schemas if schema["function"]["name"] in enabled_tool_names]
def run(self, name: str, arguments: dict[str, object]) -> list[dict[str, object]]:
"""Run one catalog tool and audit the call."""
handlers = {
"search_authors": self.run_search_authors,
"search_series": self.run_search_series,
"search_books": self.run_search_books,
"ensure_author": self.run_ensure_author,
"ensure_series": self.run_ensure_series,
"ensure_book": self.run_ensure_book,
}
handler = handlers.get(name)
if handler is None:
self.write_log(self.log_path, "tool_error", tool=name, arguments=arguments, error="unknown_tool")
msg = f"Unknown audiobook metadata tool: {name}"
raise MetadataResolutionError(msg)
if name not in self.config.tool_names:
self.write_log(self.log_path, "tool_error", tool=name, arguments=arguments, error="tool_not_enabled")
msg = f"Audiobook metadata tool is not enabled: {name}"
raise MetadataResolutionError(msg)
started = time.perf_counter()
self.write_log(self.log_path, "tool_call", tool=name, arguments=arguments)
result = handler(arguments)
duration_ms = round((time.perf_counter() - started) * 1000, 3)
self.write_log(
self.log_path,
"tool_result",
tool=name,
duration_ms=duration_ms,
result_count=len(result),
preview=result[:3],
)
return result
def get_author(self, author_id: int) -> AudiobookAuthor | None:
"""Return an author by id."""
return self.session.get(AudiobookAuthor, author_id)
def get_book(self, book_id: int) -> Audiobook | None:
"""Return a book by id."""
return self.session.get(Audiobook, book_id)
def get_series(self, series_id: int) -> AudiobookSeries | None:
"""Return a series by id."""
return self.session.get(AudiobookSeries, series_id)
def prune_unused_created_rows(self, *, author_id: int, book_id: int | None, series_id: int | None) -> None:
"""Remove catalog rows created during this run but not used by final metadata."""
used_book_ids = {book_id} if book_id is not None else set()
for created_book_id in self.created_book_ids - used_book_ids:
if book := self.get_book(created_book_id):
self.session.delete(book)
self.session.flush()
used_series_ids = {series_id} if series_id is not None else set()
for created_series_id in self.created_series_ids - used_series_ids:
series = self.get_series(created_series_id)
if series and not series.books:
self.session.delete(series)
self.session.flush()
for created_author_id in self.created_author_ids - {author_id}:
author = self.get_author(created_author_id)
if author and not author.books and not author.series:
self.session.delete(author)
def run_search_authors(self, arguments: dict[str, object]) -> list[dict[str, object]]:
"""Search authors from tool arguments and remember returned ids."""
query = required_string(arguments, "query")
statement = select(AudiobookAuthor).order_by(AudiobookAuthor.name).limit(self.config.max_tool_results)
if terms := query_terms(query):
statement = statement.where(or_(*(AudiobookAuthor.name.ilike(f"%{term}%") for term in terms)))
authors = self.session.scalars(statement).all()
self.seen_author_ids.update(author.id for author in authors)
return [{"id": author.id, "name": author.name} for author in authors]
def run_search_series(self, arguments: dict[str, object]) -> list[dict[str, object]]:
"""Search series from tool arguments and remember returned ids."""
query = required_string(arguments, "query")
author_id = optional_int(arguments.get("author_id"), "author_id")
statement = select(AudiobookSeries).order_by(AudiobookSeries.name).limit(self.config.max_tool_results)
if terms := query_terms(query):
statement = statement.where(or_(*(AudiobookSeries.name.ilike(f"%{term}%") for term in terms)))
if author_id is not None:
statement = statement.where(AudiobookSeries.author_id == author_id)
series_rows = self.session.scalars(statement).all()
self.seen_series_ids.update(series.id for series in series_rows)
self.seen_author_ids.update(series.author_id for series in series_rows)
return [
{
"id": series.id,
"name": series.name,
"author_id": series.author_id,
"author": series.author.name,
}
for series in series_rows
]
def run_search_books(self, arguments: dict[str, object]) -> list[dict[str, object]]:
"""Search books from tool arguments and remember returned ids."""
query = required_string(arguments, "query")
author_id = optional_int(arguments.get("author_id"), "author_id")
series_id = optional_int(arguments.get("series_id"), "series_id")
statement = select(Audiobook).order_by(Audiobook.title).limit(self.config.max_tool_results)
if terms := query_terms(query):
statement = statement.where(or_(*(Audiobook.title.ilike(f"%{term}%") for term in terms)))
if author_id is not None:
statement = statement.where(Audiobook.author_id == author_id)
if series_id is not None:
statement = statement.where(Audiobook.series_id == series_id)
books = self.session.scalars(statement).all()
self.seen_book_ids.update(book.id for book in books)
self.seen_author_ids.update(book.author_id for book in books)
self.seen_series_ids.update(book.series_id for book in books if book.series_id is not None)
return [
{
"id": book.id,
"title": book.title,
"author_id": book.author_id,
"author": book.author.name,
"series_id": book.series_id,
"series": book.series.name if book.series else self.config.standalone_series,
"series_index": book.series_index,
}
for book in books
]
def run_ensure_author(self, arguments: dict[str, object]) -> list[dict[str, object]]:
"""Ensure an author from tool arguments and return a tool result."""
name = normalize_catalog_slug(required_string(arguments, "name"))
validate_catalog_slug(name, "author")
author = self.session.scalar(select(AudiobookAuthor).where(AudiobookAuthor.name == name))
action = "existing"
if author is None:
author = AudiobookAuthor(name=name)
self.session.add(author)
self.session.flush()
self.created_author_ids.add(author.id)
action = "created"
self.seen_author_ids.add(author.id)
return [{"id": author.id, "name": author.name, "action": action}]
def run_ensure_series(self, arguments: dict[str, object]) -> list[dict[str, object]]:
"""Ensure a series from tool arguments and return a tool result."""
name = normalize_catalog_slug(required_string(arguments, "name"))
author_id = required_int(arguments, "author_id")
validate_catalog_slug(name, "series")
author = self.required_author(author_id)
series = self.find_series_by_catalog_slug(name, author.id)
action = "existing"
if series is None:
series = AudiobookSeries(name=name, author=author)
self.session.add(series)
self.session.flush()
self.created_series_ids.add(series.id)
action = "created"
self.seen_author_ids.add(author.id)
self.seen_series_ids.add(series.id)
return [self.series_result(series, action)]
def run_ensure_book(self, arguments: dict[str, object]) -> list[dict[str, object]]:
"""Ensure a book from tool arguments and return a tool result."""
title = required_string(arguments, "title")
author_id = required_int(arguments, "author_id")
series_id = optional_int(arguments.get("series_id"), "series_id")
series_index = required_series_index(arguments, "series_index")
ensured = self.ensure_book(title, author_id, series_id, series_index)
return [self.book_result(ensured.book, ensured.action)]
def ensure_book(
self,
title: str,
author_id: int,
series_id: int | None,
series_index: float,
) -> EnsuredBook:
"""Return an existing book row, or create it after validating ownership."""
title = normalize_title_slug(title)
validate_title_slug(title)
author = self.required_author(author_id)
series = None
if series_id is None:
if series_index != 0:
msg = "standalone books must use series_index 0"
raise MetadataResolutionError(msg)
else:
series = self.required_series(series_id)
if series.author_id != author.id:
msg = f"series_id {series_id} does not belong to author_id {author_id}"
raise MetadataResolutionError(msg)
if series_index <= 0:
msg = "series books must use a positive series_index"
raise MetadataResolutionError(msg)
statement = select(Audiobook).where(
Audiobook.title == title,
Audiobook.author_id == author.id,
)
if series is None:
statement = statement.where(Audiobook.series_id.is_(None))
else:
statement = statement.where(Audiobook.series_id == series.id)
book = self.session.scalar(statement)
if book is None:
book = Audiobook(title=title, author=author, series=series, series_index=series_index)
self.session.add(book)
self.session.flush()
self.created_book_ids.add(book.id)
action = "created"
else:
action = "existing"
self.seen_book_ids.add(book.id)
self.seen_author_ids.add(author.id)
if book.series_id is not None:
self.seen_series_ids.add(book.series_id)
return EnsuredBook(book=book, action=action)
def required_author(self, author_id: int) -> AudiobookAuthor:
"""Return an author or fail metadata resolution."""
author = self.get_author(author_id)
if author is None:
msg = f"author_id {author_id} does not exist"
raise MetadataResolutionError(msg)
return author
def required_series(self, series_id: int) -> AudiobookSeries:
"""Return a series or fail metadata resolution."""
series = self.get_series(series_id)
if series is None:
msg = f"series_id {series_id} does not exist"
raise MetadataResolutionError(msg)
return series
def find_series_by_catalog_slug(self, name: str, author_id: int) -> AudiobookSeries | None:
"""Return a series by exact slug or underscore-insensitive slug."""
exact = self.session.scalar(
select(AudiobookSeries).where(
AudiobookSeries.name == name,
AudiobookSeries.author_id == author_id,
),
)
if exact is not None:
return exact
compact_name = compact_catalog_slug(name)
series_rows = self.session.scalars(
select(AudiobookSeries).where(AudiobookSeries.author_id == author_id).order_by(AudiobookSeries.name),
).all()
for series in series_rows:
if compact_catalog_slug(series.name) == compact_name:
return series
return None
def series_result(self, series: AudiobookSeries, action: str) -> dict[str, object]:
"""Build a normalized series tool result."""
return {
"id": series.id,
"name": series.name,
"author_id": series.author_id,
"author": series.author.name,
"action": action,
}
def book_result(self, book: Audiobook, action: str) -> dict[str, object]:
"""Build a normalized book tool result."""
return {
"id": book.id,
"title": book.title,
"author_id": book.author_id,
"author": book.author.name,
"series_id": book.series_id,
"series": book.series.name if book.series else self.config.standalone_series,
"series_index": book.series_index,
"action": action,
}
def run_tool_calls(
messages: list[dict[str, object]],
message: dict[str, object],
tool_calls: list[tuple[str, dict[str, object]]],
registry: CatalogToolRegistry,
log_path: Path,
write_log: LogWriter,
) -> str | None:
"""Run tool calls, append tool messages, and return fatal error text when stopped."""
messages.append(message)
for tool_name, arguments in tool_calls:
try:
tool_result = registry.run(tool_name, arguments)
except MetadataResolutionError as error:
if is_fatal_tool_error(error):
return str(error)
write_log(log_path, "tool_error", tool=tool_name, arguments=arguments, error=str(error))
messages.append(
{
"role": "tool",
"tool_name": tool_name,
"content": json.dumps({"error": str(error)}, sort_keys=True),
},
)
continue
messages.append(
{
"role": "tool",
"tool_name": tool_name,
"content": json.dumps(tool_result, sort_keys=True),
},
)
return None
def parse_tool_calls(message: dict[str, object]) -> list[tuple[str, dict[str, object]]]:
"""Parse Ollama tool calls from a response message."""
raw_tool_calls = message.get("tool_calls") or []
if not isinstance(raw_tool_calls, list):
msg = "tool_calls must be a list"
raise MetadataResolutionError(msg)
tool_calls = []
for raw_call in raw_tool_calls:
if not isinstance(raw_call, dict):
msg = "tool call must be an object"
raise MetadataResolutionError(msg)
function = raw_call.get("function")
if not isinstance(function, dict):
msg = "tool call is missing function"
raise MetadataResolutionError(msg)
name = function.get("name")
if not isinstance(name, str) or not name:
msg = "tool call is missing function name"
raise MetadataResolutionError(msg)
arguments = parse_tool_arguments(function.get("arguments", {}))
tool_calls.append((name, arguments))
return tool_calls
def parse_tool_arguments(raw_arguments: object) -> dict[str, object]:
"""Parse tool call arguments returned by Ollama."""
if isinstance(raw_arguments, dict):
return {str(key): value for key, value in raw_arguments.items()}
if isinstance(raw_arguments, str):
parsed = json.loads(raw_arguments) if raw_arguments else {}
if isinstance(parsed, dict):
return {str(key): value for key, value in parsed.items()}
msg = "tool arguments must be an object"
raise MetadataResolutionError(msg)
def validate_title_slug(title: str) -> None:
"""Validate a canonical book title slug."""
if not TITLE_SLUG_PATTERN.fullmatch(title):
msg = f"title slug is invalid: {title}"
raise MetadataResolutionError(msg)
def validate_catalog_slug(value: str, label: str) -> None:
"""Validate a canonical catalog slug."""
if not CATALOG_SLUG_PATTERN.fullmatch(value):
msg = f"{label} slug is invalid: {value}"
raise MetadataResolutionError(msg)
def normalize_catalog_slug(value: str) -> str:
"""Normalize noisy catalog names into lower snake-case slugs."""
return re.sub(r"[^a-z0-9]+", "_", value.strip().casefold()).strip("_")
def compact_catalog_slug(value: str) -> str:
"""Return a catalog slug comparison key that ignores underscores."""
return normalize_catalog_slug(value).replace("_", "")
def normalize_title_slug(value: str) -> str:
"""Normalize noisy book titles into lower kebab-case slugs."""
return re.sub(r"[^a-z0-9]+", "-", value.strip().casefold()).strip("-")
def is_fatal_tool_error(error: MetadataResolutionError) -> bool:
"""Return whether a tool error should stop the agent immediately."""
message = str(error)
return message.startswith(
(
"Unknown audiobook metadata tool",
"Audiobook metadata tool is not enabled",
),
)
def query_terms(query: str) -> tuple[str, ...]:
"""Return text variants useful for matching noisy audiobook metadata."""
normalized = query.strip().casefold()
underscore_slug = normalize_catalog_slug(normalized)
compact_slug = compact_catalog_slug(normalized)
hyphen_slug = normalize_title_slug(normalized)
return tuple(dict.fromkeys(term for term in (normalized, underscore_slug, compact_slug, hyphen_slug) if term))
def required_string(data: dict[str, object], key: str) -> str:
"""Read a required string field."""
value = data.get(key)
if not isinstance(value, str) or not value.strip():
msg = f"{key} must be a non-empty string"
raise MetadataResolutionError(msg)
return value.strip()
def required_int(data: dict[str, object], key: str) -> int:
"""Read a required integer field."""
value = data.get(key)
if isinstance(value, bool) or not isinstance(value, int):
msg = f"{key} must be an integer"
raise MetadataResolutionError(msg)
return value
def required_series_index(data: dict[str, object], key: str) -> float:
"""Read a required whole-number or half-number series index."""
value = data.get(key)
if isinstance(value, bool) or not isinstance(value, int | float):
msg = f"{key} must be a number"
raise MetadataResolutionError(msg)
series_index = float(value)
if not (series_index * 2).is_integer():
msg = f"{key} must be a whole number or .5 increment"
raise MetadataResolutionError(msg)
return series_index
def optional_int(value: object, key: str) -> int | None:
"""Read an optional integer field."""
if value is None:
return None
if isinstance(value, bool) or not isinstance(value, int):
msg = f"{key} must be an integer or null"
raise MetadataResolutionError(msg)
return value
+575
View File
@@ -0,0 +1,575 @@
"""Resolve audiobook metadata with a controlled Ollama tool loop."""
from __future__ import annotations
import json
import re
from dataclasses import asdict, dataclass, is_dataclass, replace
from os import PathLike
from typing import TYPE_CHECKING
import httpx
from sqlalchemy.orm import Session
from python.common import utcnow
from python.tools.audiobook.llm_tool_calling import (
CatalogToolRegistry,
MetadataResolutionError,
normalize_title_slug,
optional_int,
parse_tool_calls,
required_int,
required_series_index,
required_string,
run_tool_calls,
validate_catalog_slug,
validate_title_slug,
)
if TYPE_CHECKING:
from pathlib import Path
from sqlalchemy.engine import Engine
from python.orm.richie import AudiobookAuthor
FENCED_JSON_PATTERN = re.compile(r"^```(?:json)?\s*(?P<json>.*?)\s*```$", re.IGNORECASE | re.DOTALL)
@dataclass(frozen=True)
class AgentConfig:
"""Runtime settings for the audiobook metadata agent."""
model: str = "deepseek-v4-flash:cloud"
ollama_chat_url: str = "https://ollama.com/api/chat"
http_timeout_seconds: int = 300
max_agent_turns: int = 8
max_tool_results: int = 10
min_confidence: float = 0.85
invalid_final_retries: int = 1
standalone_series: str = "standalone"
tool_names: tuple[str, ...] = (
"search_authors",
"search_series",
"search_books",
"ensure_author",
"ensure_series",
"ensure_book",
)
@dataclass(frozen=True)
class StandardBookMetadata:
"""Canonical metadata for the final audiobook path."""
author_id: int
author: str
book_id: int | None
title: str
series_id: int | None
series: str
series_index: float
confidence: float
needs_review: bool
evidence: list[str]
@dataclass(frozen=True)
class FinalMetadataFields:
"""Raw model fields after schema validation."""
author_id: int
book_id: int | None
title: str
series_id: int | None
series_index: float
confidence: float
evidence: list[str]
@dataclass(frozen=True)
class ResolvedBookFields:
"""Book fields after optional catalog book resolution."""
book_id: int | None
title: str
series_id: int | None
series_index: float
@dataclass(frozen=True)
class AgentStepResult:
"""Outcome from one model response."""
metadata: StandardBookMetadata | None
invalid_final_count: int
should_continue: bool
def standard_book_metadata(
aax_file_name: str,
aax_metadata_from_ffprobe: dict[str, str],
engine: Engine,
log_path: Path,
ollama_api_key: str,
config: AgentConfig,
) -> StandardBookMetadata:
"""Resolve canonical audiobook metadata with the configured Ollama Cloud model."""
with Session(engine) as session:
registry = CatalogToolRegistry(session, log_path, config, write_agent_log)
agent = AudiobookMetadataAgent(
registry=registry, log_path=log_path, ollama_api_key=ollama_api_key, config=config
)
metadata = agent.run(aax_file_name, aax_metadata_from_ffprobe)
if metadata.needs_review:
session.rollback()
else:
registry.prune_unused_created_rows(
author_id=metadata.author_id,
book_id=metadata.book_id,
series_id=metadata.series_id,
)
session.commit()
return metadata
class AudiobookMetadataAgent:
"""Ollama-backed metadata resolver with a fixed local tool registry."""
def __init__(
self,
*,
registry: CatalogToolRegistry,
log_path: Path,
ollama_api_key: str,
config: AgentConfig,
) -> None:
"""Create an Ollama metadata agent."""
self._registry = registry
self._log_path = log_path
self._ollama_api_key = ollama_api_key
self._config = config
def run(self, aax_file_name: str, aax_metadata_from_ffprobe: dict[str, str]) -> StandardBookMetadata:
"""Resolve metadata for one AAX file."""
messages = [
{"role": "system", "content": system_prompt()},
{"role": "user", "content": user_prompt(aax_file_name, aax_metadata_from_ffprobe)},
]
invalid_final_count = 0
result: StandardBookMetadata | None = None
for turn in range(1, self._config.max_agent_turns + 1):
step = self.run_step(messages, turn, invalid_final_count)
invalid_final_count = step.invalid_final_count
if step.should_continue:
continue
result = step.metadata
break
if result is None:
return self.force_final_response(messages)
return result
def run_step(
self,
messages: list[dict[str, object]],
turn: int,
invalid_final_count: int,
) -> AgentStepResult:
"""Run one model turn and return the next agent-loop action."""
data = self.chat(messages, turn)
message = data.get("message")
if not isinstance(message, dict):
return AgentStepResult(
metadata=review_metadata("Ollama response did not include a message", self._config),
invalid_final_count=invalid_final_count,
should_continue=False,
)
try:
tool_calls = parse_tool_calls(message)
except (json.JSONDecodeError, MetadataResolutionError) as error:
return AgentStepResult(
metadata=review_metadata(str(error), self._config),
invalid_final_count=invalid_final_count,
should_continue=False,
)
if tool_calls:
fatal_error = run_tool_calls(messages, message, tool_calls, self._registry, self._log_path, write_agent_log)
if fatal_error is not None:
return AgentStepResult(
metadata=review_metadata(fatal_error, self._config),
invalid_final_count=invalid_final_count,
should_continue=False,
)
return AgentStepResult(metadata=None, invalid_final_count=invalid_final_count, should_continue=True)
return self.handle_final_message(messages, message, invalid_final_count)
def handle_final_message(
self,
messages: list[dict[str, object]],
message: dict[str, object],
invalid_final_count: int,
) -> AgentStepResult:
"""Validate a final model message or request one retry."""
content = message.get("content")
if not isinstance(content, str):
return AgentStepResult(
metadata=review_metadata("Ollama final response did not include string content", self._config),
invalid_final_count=invalid_final_count,
should_continue=False,
)
try:
resolved = self.validate_final(parse_final_json_content(content))
except (json.JSONDecodeError, MetadataResolutionError) as error:
return self.handle_invalid_final(messages, error, invalid_final_count)
write_agent_log(self._log_path, "final_metadata", metadata=resolved)
return AgentStepResult(metadata=resolved, invalid_final_count=invalid_final_count, should_continue=False)
def handle_invalid_final(
self,
messages: list[dict[str, object]],
error: json.JSONDecodeError | MetadataResolutionError,
invalid_final_count: int,
) -> AgentStepResult:
"""Log invalid final JSON and either retry or return review metadata."""
invalid_final_count += 1
write_agent_log(
self._log_path,
"final_validation_error",
error=str(error),
invalid_final_count=invalid_final_count,
)
if invalid_final_count > self._config.invalid_final_retries:
return AgentStepResult(
metadata=review_metadata(str(error), self._config),
invalid_final_count=invalid_final_count,
should_continue=False,
)
messages.append(
{
"role": "user",
"content": (
"Your previous final answer was invalid. Return only valid JSON matching the required "
f"schema. Validation error: {error}"
),
},
)
return AgentStepResult(metadata=None, invalid_final_count=invalid_final_count, should_continue=True)
def force_final_response(self, messages: list[dict[str, object]]) -> StandardBookMetadata:
"""Request a no-tool final answer after the normal turn limit."""
messages.append({"role": "user", "content": forced_final_prompt()})
write_agent_log(self._log_path, "forced_final_request", reason="max_turns")
data = self.chat(messages, self._config.max_agent_turns + 1, tools_enabled=False)
message = data.get("message")
if not isinstance(message, dict):
return review_metadata("Ollama forced final response did not include a message", self._config)
content = message.get("content")
if not isinstance(content, str):
return review_metadata("Ollama forced final response did not include string content", self._config)
try:
resolved = self.validate_final(parse_final_json_content(content))
except (json.JSONDecodeError, MetadataResolutionError) as error:
return review_metadata(f"Ollama forced final response was invalid: {error}", self._config)
write_agent_log(self._log_path, "final_metadata", metadata=resolved)
return resolved
def chat(self, messages: list[dict[str, object]], turn: int, *, tools_enabled: bool = True) -> dict[str, object]:
"""Send one chat request to Ollama and log the request and response."""
payload = {
"model": self._config.model,
"messages": messages,
"stream": False,
"options": {"temperature": 0.1},
}
tool_names = []
if tools_enabled:
payload["tools"] = self._registry.tool_schemas()
tool_names = self._config.tool_names
write_agent_log(
self._log_path,
"model_request",
model=self._config.model,
turn=turn,
message_count=len(messages),
tool_names=tool_names,
tools_enabled=tools_enabled,
)
write_agent_log(
self._log_path,
"llm_messages_sent",
model=self._config.model,
turn=turn,
messages=messages,
tools_enabled=tools_enabled,
)
response = httpx.post(
self._config.ollama_chat_url,
headers={"Authorization": f"Bearer {self._ollama_api_key}"},
json=payload,
timeout=self._config.http_timeout_seconds,
)
response.raise_for_status()
raw_data = response.json()
if not isinstance(raw_data, dict):
return {}
data = {str(key): value for key, value in raw_data.items()}
message = data.get("message", {})
content = message.get("content") if isinstance(message, dict) else ""
write_agent_log(
self._log_path,
"llm_message_received",
model=self._config.model,
turn=turn,
message=message,
)
write_agent_log(
self._log_path,
"model_response",
model=self._config.model,
turn=turn,
has_tool_calls=bool(isinstance(message, dict) and message.get("tool_calls")),
content_chars=len(content) if isinstance(content, str) else 0,
)
return data
def validate_final(self, raw_metadata: object) -> StandardBookMetadata:
"""Validate final model metadata against catalog rows."""
fields = parse_final_metadata_fields(raw_metadata)
fields = replace(fields, title=normalize_title_slug(fields.title))
author = self.validate_author(fields.author_id)
validate_title_slug(fields.title)
book_fields = self.resolve_book_fields(fields)
series = self.validate_series(fields.author_id, book_fields.series_id, book_fields.series_index)
return StandardBookMetadata(
author_id=fields.author_id,
author=author.name,
book_id=book_fields.book_id,
title=book_fields.title,
series_id=book_fields.series_id,
series=series,
series_index=book_fields.series_index,
confidence=fields.confidence,
needs_review=fields.confidence < self._config.min_confidence,
evidence=fields.evidence,
)
def validate_author(self, author_id: int) -> AudiobookAuthor:
"""Validate that an author id was seen and exists."""
if author_id not in self._registry.seen_author_ids:
msg = f"author_id {author_id} was not returned by search_authors"
raise MetadataResolutionError(msg)
author = self._registry.get_author(author_id)
if author is None:
msg = f"author_id {author_id} does not exist"
raise MetadataResolutionError(msg)
validate_catalog_slug(author.name, "author")
return author
def resolve_book_fields(self, fields: FinalMetadataFields) -> ResolvedBookFields:
"""Resolve final book fields from a seen book id or created book."""
if fields.book_id is None:
ensured = self._registry.ensure_book(
fields.title,
fields.author_id,
fields.series_id,
fields.series_index,
)
return ResolvedBookFields(
book_id=ensured.book.id,
title=ensured.book.title,
series_id=ensured.book.series_id,
series_index=ensured.book.series_index,
)
if fields.book_id not in self._registry.seen_book_ids:
msg = f"book_id {fields.book_id} was not returned by search_books"
raise MetadataResolutionError(msg)
book = self._registry.get_book(fields.book_id)
if book is None:
msg = f"book_id {fields.book_id} does not exist"
raise MetadataResolutionError(msg)
if book.author_id != fields.author_id:
msg = f"book_id {fields.book_id} does not belong to author_id {fields.author_id}"
raise MetadataResolutionError(msg)
return ResolvedBookFields(
book_id=fields.book_id,
title=book.title,
series_id=book.series_id,
series_index=book.series_index,
)
def validate_series(self, author_id: int, series_id: int | None, series_index: float) -> str:
"""Validate final series fields and return the canonical series slug."""
if series_id is None:
if series_index != 0:
msg = "standalone books must use series_index 0"
raise MetadataResolutionError(msg)
return self._config.standalone_series
if series_id not in self._registry.seen_series_ids:
msg = f"series_id {series_id} was not returned by search_series"
raise MetadataResolutionError(msg)
series = self._registry.get_series(series_id)
if series is None:
msg = f"series_id {series_id} does not exist"
raise MetadataResolutionError(msg)
if series.author_id != author_id:
msg = f"series_id {series_id} does not belong to author_id {author_id}"
raise MetadataResolutionError(msg)
if series_index <= 0:
msg = "series books must use a positive series_index"
raise MetadataResolutionError(msg)
validate_catalog_slug(series.name, "series")
return series.name
def write_agent_log(log_path: Path, event: str, **fields: object) -> None:
"""Append one JSONL audit event."""
log_path.parent.mkdir(parents=True, exist_ok=True)
record = {
"created": utcnow().isoformat(),
"event": event,
**{key: json_log_value(value) for key, value in fields.items()},
}
with log_path.open("a", encoding="utf-8") as file:
file.write(json.dumps(record, sort_keys=True))
file.write("\n")
def json_log_value(value: object) -> object:
"""Return a JSON-serializable value for audit logs."""
if is_dataclass(value) and not isinstance(value, type):
return json_log_value(asdict(value))
if isinstance(value, dict):
return {str(key): json_log_value(item) for key, item in value.items()}
if isinstance(value, list | tuple):
return [json_log_value(item) for item in value]
if isinstance(value, set):
return [json_log_value(item) for item in sorted(value, key=str)]
if isinstance(value, PathLike):
return str(value)
return value
def system_prompt() -> str:
"""Return the stable system prompt."""
return """You standardize Audible audiobook metadata against a private catalog.
Rules:
- You must use the provided tools before returning final metadata.
- Only use author_id, series_id, or book_id values returned by tools.
- Return final metadata as JSON only. Do not wrap it in Markdown.
- The final JSON object must contain author_id, book_id, title, series_id, series_index, confidence, and evidence.
- title must be a canonical title slug using lower-case words separated by hyphens.
- Use series_id null and series_index 0 for standalone books.
- If you use a series_id, series_index must be a whole number or .5 value greater than 0.
- Treat series slugs that differ only by underscores as the same series. Prefer the existing catalog row instead of
creating a new series.
- Detect omnibus or box-set editions that contain multiple numbered novels, books, or novellas.
- For an omnibus, make a best-effort range from the filename, tags, and catalog rows. Keep series_index as the
first covered book number and include the range in the title when the source title includes it, for example
books-1-3.
- Be careful with omnibuses of novels or novellas later published as one book: keep the omnibus as the audiobook's
book record unless catalog rows clearly identify a better match.
- Do not create publisher collections or author collections as series unless the book metadata clearly gives a
numbered series.
- Series belong to authors. Use a series_id only when it belongs to the selected author_id.
- Always search for the author before creating one. If no exact author slug exists, call ensure_author.
- Always search for a series with author_id before creating one. If no exact series slug exists, call ensure_series.
- Always search for a book before creating one. If no exact title slug exists, call ensure_book.
- If a tool returns an error, correct your tool arguments or final metadata before continuing.
- confidence must be a number from 0 to 1.
- evidence must be a short list of strings explaining which filename, tags, and catalog rows support the answer."""
def forced_final_prompt() -> str:
"""Return the no-tools finalization prompt."""
return (
"Stop calling tools. Return final metadata as JSON only using the tool results already provided. "
"If search_books returned no matching rows but author and series are known, use book_id null and resolve "
"the title slug from the AAX filename and ffprobe tags. The validator will create the missing book. "
"Use only author_id and series_id values returned by earlier tool results."
)
def user_prompt(aax_file_name: str, metadata: dict[str, str]) -> str:
"""Build the user prompt from source metadata."""
return (
"Resolve this Audible audiobook.\n\n"
f"AAX file name: {aax_file_name}\n\n"
"ffprobe format tags:\n"
f"{json.dumps(metadata, indent=2, sort_keys=True)}"
)
def parse_final_json_content(content: str) -> object:
"""Parse final model content, accepting bare or fenced JSON."""
stripped = content.strip()
if match := FENCED_JSON_PATTERN.fullmatch(stripped):
stripped = match.group("json").strip()
return json.loads(stripped)
def parse_final_metadata_fields(raw_metadata: object) -> FinalMetadataFields:
"""Parse the model's final JSON object into typed fields."""
if not isinstance(raw_metadata, dict):
msg = "Final metadata must be a JSON object"
raise MetadataResolutionError(msg)
data = {str(key): value for key, value in raw_metadata.items()}
return FinalMetadataFields(
author_id=required_int(data, "author_id"),
book_id=optional_int(data.get("book_id"), "book_id"),
title=required_string(data, "title"),
series_id=optional_int(data.get("series_id"), "series_id"),
series_index=required_series_index(data, "series_index"),
confidence=required_float(data, "confidence"),
evidence=required_string_list(data, "evidence"),
)
def review_metadata(reason: str, config: AgentConfig) -> StandardBookMetadata:
"""Return a metadata result that must be reviewed manually."""
return StandardBookMetadata(
author_id=0,
author="unknown_author",
book_id=None,
title="unknown-title",
series_id=None,
series=config.standalone_series,
series_index=0,
confidence=0,
needs_review=True,
evidence=[reason],
)
def required_float(data: dict[str, object], key: str) -> float:
"""Read a required float field."""
value = data.get(key)
if isinstance(value, bool) or not isinstance(value, int | float):
msg = f"{key} must be a number"
raise MetadataResolutionError(msg)
confidence = float(value)
if confidence < 0 or confidence > 1:
msg = f"{key} must be between 0 and 1"
raise MetadataResolutionError(msg)
return confidence
def required_string_list(data: dict[str, object], key: str) -> list[str]:
"""Read a required list of strings."""
value = data.get(key)
if not isinstance(value, list) or not value or not all(isinstance(item, str) for item in value):
msg = f"{key} must be a non-empty list of strings"
raise MetadataResolutionError(msg)
strings = [item.strip() for item in value if item.strip()]
if not strings:
msg = f"{key} must include at least one non-empty string"
raise MetadataResolutionError(msg)
return strings
+4 -2
View File
@@ -34,8 +34,9 @@ def main(config_file: Path) -> None:
logger.error(msg)
signal_alert(msg)
continue
get_snapshots_to_delete(dataset, get_count_lookup(config_file, dataset.name))
count_lookup = get_count_lookup(config_file, dataset.name)
logger.info(f"using {count_lookup} for {dataset.name}")
get_snapshots_to_delete(dataset, count_lookup)
except Exception:
logger.exception("snapshot_manager failed")
signal_alert("snapshot_manager failed")
@@ -99,6 +100,7 @@ def get_snapshots_to_delete(
"""
snapshots = dataset.get_snapshots()
logger.info(f"calculating snapshots for {dataset.name} to be deleted")
if not snapshots:
logger.info(f"{dataset.name} has no snapshots")
return
+17
View File
@@ -0,0 +1,17 @@
FROM nvidia/cuda:12.4.1-cudnn-runtime-ubuntu22.04
ENV DEBIAN_FRONTEND=noninteractive \
PYTHONDONTWRITEBYTECODE=1 \
PYTHONUNBUFFERED=1
RUN apt-get update \
&& apt-get install -y --no-install-recommends python3 python3-pip ffmpeg \
&& rm -rf /var/lib/apt/lists/*
RUN pip3 install --no-cache-dir --upgrade pip \
&& pip3 install --no-cache-dir faster-whisper requests
WORKDIR /app
COPY python/tools/whisper/inference.py /app/inference.py
ENTRYPOINT ["python3", "/app/inference.py"]
@@ -0,0 +1,2 @@
*
!python/tools/whisper/inference.py
+1
View File
@@ -0,0 +1 @@
"""Whisper transcription tools (host orchestrator and container entrypoint)."""
+136
View File
@@ -0,0 +1,136 @@
"""Container entrypoint that transcribes a directory of audio files with faster-whisper.
Run inside the whisper-transcribe docker image; segment timestamps are grouped
into one-minute buckets so the output reads as ``[HH:MM:00] text``.
"""
from __future__ import annotations
import argparse
import logging
from pathlib import Path
from faster_whisper import WhisperModel
logger = logging.getLogger(__name__)
AUDIO_EXTENSIONS = {".mp3", ".wav", ".m4a", ".flac", ".ogg", ".opus", ".mp4", ".mkv", ".webm", ".aac"}
BUCKET_SECONDS = 60
BEAM_SIZE = 5
SECONDS_PER_HOUR = 3600
SECONDS_PER_MINUTE = 60
def format_timestamp(total_seconds: float) -> str:
"""Render a whole-minute timestamp as ``HH:MM:00``.
Args:
total_seconds: Offset in seconds from the start of the audio.
Returns:
A zero-padded ``HH:MM:00`` string.
"""
hours = int(total_seconds // SECONDS_PER_HOUR)
minutes = int((total_seconds % SECONDS_PER_HOUR) // SECONDS_PER_MINUTE)
return f"{hours:02d}:{minutes:02d}:00"
def transcribe_file(model: WhisperModel, audio_path: Path, output_path: Path) -> None:
"""Transcribe one audio file and write the bucketed transcript to disk.
Args:
model: Loaded faster-whisper model.
audio_path: Source audio file.
output_path: Destination ``.txt`` path.
"""
logger.info("Transcribing %s", audio_path)
segments, info = model.transcribe(
str(audio_path),
language="en",
beam_size=BEAM_SIZE,
vad_filter=True,
)
logger.info("Duration %.1fs", info.duration)
buckets: dict[int, list[str]] = {}
for segment in segments:
bucket = int(segment.start // BUCKET_SECONDS)
buckets.setdefault(bucket, []).append(segment.text.strip())
lines = [f"[{format_timestamp(bucket * BUCKET_SECONDS)}] {' '.join(buckets[bucket])}" for bucket in sorted(buckets)]
output_path.write_text("\n\n".join(lines) + "\n", encoding="utf-8")
logger.info("Wrote %s", output_path)
def find_audio_files(input_directory: Path) -> list[Path]:
"""Collect every audio file under ``input_directory``.
Args:
input_directory: Directory to walk recursively.
Returns:
Sorted list of audio file paths.
"""
return sorted(
path for path in input_directory.rglob("*") if path.is_file() and path.suffix.lower() in AUDIO_EXTENSIONS
)
def configure_container_logger() -> None:
"""Configure logging for the container (stdout, INFO)."""
logging.basicConfig(
level=logging.INFO,
format="%(asctime)s %(levelname)s %(message)s",
)
def parse_arguments() -> argparse.Namespace:
"""Parse CLI arguments for the container entrypoint.
Returns:
Parsed argparse namespace.
"""
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--input", type=Path, default=Path("/audio"))
parser.add_argument("--output", type=Path, default=Path("/output"))
parser.add_argument("--model", default="large-v3")
parser.add_argument(
"--download-only",
action="store_true",
help="Download the model into the cache volume and exit without transcribing.",
)
return parser.parse_args()
def main() -> None:
"""Load the model, then either exit (download-only) or transcribe the directory."""
configure_container_logger()
arguments = parse_arguments()
logger.info("Loading model %s on CUDA", arguments.model)
model = WhisperModel(arguments.model, device="cuda", compute_type="float16")
if arguments.download_only:
logger.info("Model ready; exiting (download-only mode)")
return
arguments.output.mkdir(parents=True, exist_ok=True)
audio_files = find_audio_files(arguments.input)
if not audio_files:
logger.warning("No audio files found in %s", arguments.input)
return
logger.info("Found %d audio file(s)", len(audio_files))
for audio_path in audio_files:
relative = audio_path.relative_to(arguments.input)
output_path = arguments.output / relative.with_suffix(".txt")
output_path.parent.mkdir(parents=True, exist_ok=True)
if output_path.exists():
logger.info("Skip %s (already transcribed)", relative)
continue
transcribe_file(model, audio_path, output_path)
if __name__ == "__main__":
main()
+167
View File
@@ -0,0 +1,167 @@
"""Build and run the whisper transcription docker container on demand.
The container is started fresh for each invocation and removed on exit
(``docker run --rm``). The model is cached in a named docker volume so
only the first run pays the download cost.
"""
from __future__ import annotations
import logging
import subprocess
from pathlib import Path
from typing import Annotated
import typer
from python.common import configure_logger
logger = logging.getLogger(__name__)
class Config:
"""Paths and names for the whisper-transcribe Docker workflow."""
image_tag = "whisper-transcribe:latest"
model_volume = "whisper-models"
repo_root = Path(__file__).resolve().parents[3]
dockerfile = Path(__file__).resolve().parent / "Dockerfile"
huggingface_cache = "/root/.cache/huggingface"
def run_docker(arguments: list[str]) -> None:
"""Run a docker subcommand, streaming output and raising on failure.
Args:
arguments: Arguments to pass to the ``docker`` binary.
Raises:
subprocess.CalledProcessError: If docker exits non-zero.
"""
logger.info("docker %s", " ".join(arguments))
subprocess.run(["docker", *arguments], check=True)
def build_image() -> None:
"""Build the whisper-transcribe image using the repo root as build context."""
logger.info("Building image %s", Config.image_tag)
run_docker(
[
"build",
"--tag",
Config.image_tag,
"--file",
str(Config.dockerfile),
str(Config.repo_root),
],
)
def model_cache_present(model: str) -> bool:
"""Check whether the given model is already downloaded in the cache volume.
Args:
model: faster-whisper model name (e.g. ``large-v3``).
Returns:
True if the HuggingFace cache directory for the model exists in the volume.
"""
cache_directory = f"hub/models--Systran--faster-whisper-{model}"
completed = subprocess.run(
[
"docker",
"run",
"--rm",
"--volume",
f"{Config.model_volume}:/cache",
"alpine",
"test",
"-d",
f"/cache/{cache_directory}",
],
check=False,
)
return completed.returncode == 0
def download_model(model: str) -> None:
"""Download the model into the cache volume and exit.
Args:
model: faster-whisper model name.
"""
logger.info("Downloading model %s into volume %s", model, Config.model_volume)
run_docker(
[
"run",
"--rm",
"--device=nvidia.com/gpu=all",
"--ipc=host",
"--volume",
f"{Config.model_volume}:{Config.huggingface_cache}",
Config.image_tag,
"--model",
model,
"--download-only",
],
)
def transcribe(input_directory: Path, output_directory: Path, model: str) -> None:
"""Run transcription on every audio file under ``input_directory``.
Args:
input_directory: Host path containing audio files (mounted read-only).
output_directory: Host path for ``.txt`` transcripts.
model: faster-whisper model name.
"""
logger.info("Transcribing %s -> %s (model=%s)", input_directory, output_directory, model)
run_docker(
[
"run",
"--rm",
"--device=nvidia.com/gpu=all",
"--ipc=host",
"--volume",
f"{input_directory}:/audio:ro",
"--volume",
f"{output_directory}:/output",
"--volume",
f"{Config.model_volume}:{Config.huggingface_cache}",
Config.image_tag,
"--model",
model,
],
)
def main(
input_directory: Annotated[Path, typer.Argument(help="Directory of audio files to transcribe.")],
output_directory: Annotated[Path, typer.Argument(help="Directory to write .txt transcripts to.")],
model: Annotated[str, typer.Option(help="faster-whisper model name.")] = "large-v3",
*,
force_download: Annotated[
bool,
typer.Option("--force-download", help="Re-download the model even if already cached."),
] = False,
) -> None:
"""Build the image, ensure the model is cached, then transcribe and stop."""
configure_logger()
resolved_input = input_directory.resolve(strict=True)
output_directory.mkdir(parents=True, exist_ok=True)
resolved_output = output_directory.resolve()
build_image()
if force_download or not model_cache_present(model):
download_model(model)
else:
logger.info("Model %s already cached in volume %s", model, Config.model_volume)
transcribe(resolved_input, resolved_output, model)
logger.info("Done. Container stopped.")
if __name__ == "__main__":
typer.run(main)
+15 -4
View File
@@ -1,28 +1,39 @@
{ inputs, ... }:
{ inputs, pkgs, ... }:
{
imports = [
"${inputs.self}/users/math"
"${inputs.self}/users/richie"
"${inputs.self}/users/steve"
"${inputs.self}/common/global"
"${inputs.self}/common/optional/desktop.nix"
"${inputs.self}/common/optional/docker.nix"
"${inputs.self}/common/optional/scanner.nix"
"${inputs.self}/common/optional/monitoring-agent.nix"
"${inputs.self}/common/optional/steam.nix"
"${inputs.self}/common/optional/syncthing_base.nix"
"${inputs.self}/common/optional/systemd-boot.nix"
"${inputs.self}/common/optional/update.nix"
"${inputs.self}/common/optional/yubikey.nix"
"${inputs.self}/common/optional/zerotier.nix"
"${inputs.self}/common/optional/brain_substituter.nix"
"${inputs.self}/common/optional/nvidia.nix"
./hardware.nix
./syncthing.nix
./llms.nix
];
boot = {
kernelPackages = pkgs.linuxPackages_6_18;
zfs.package = pkgs.zfs_2_4;
};
networking = {
hostName = "bob";
hostId = "7c678a41";
firewall.enable = true;
firewall = {
enable = true;
allowedTCPPorts = [
8000
];
};
networkmanager.enable = true;
};
+5 -1
View File
@@ -28,9 +28,13 @@
allowDiscards = true;
keyFileSize = 4096;
keyFile = "/dev/disk/by-id/usb-Samsung_Flash_Drive_FIT_0374620080067131-0:0";
fallbackToPassword = true;
};
};
zfs.extraPools = [
"storage"
];
kernelModules = [ "kvm-amd" ];
extraModulePackages = [ ];
};
+6 -2
View File
@@ -4,7 +4,7 @@
host = "0.0.0.0";
enable = true;
syncModels = true;
syncModels = false;
loadModels = [
"codellama:7b"
"deepscaler:1.5b"
@@ -23,6 +23,7 @@
"magistral:24b"
"ministral-3:14b"
"nemotron-3-nano:30b"
"nemotron-3-nano:4b"
"nemotron-cascade-2:30b"
"qwen3-coder:30b"
"qwen3-embedding:0.6b"
@@ -41,11 +42,14 @@
"qwen3:8b"
"qwen3.5:27b"
"qwen3.5:35b"
"qwen3.6:27b"
"qwen3.6:35b"
"rinex20/translategemma3:12b"
"translategemma:12b"
"translategemma:27b"
"translategemma:4b"
];
models = "/zfs/models";
models = "/zfs/storage/models";
openFirewall = true;
};
}
+10
View File
@@ -31,5 +31,15 @@
];
fsWatcherEnabled = true;
};
"recordings" = {
path = "/home/richie/recordings";
devices = [
"jeeves"
"phone"
"rhapsody-in-green"
];
fsWatcherEnabled = true;
};
};
}
-1
View File
@@ -26,7 +26,6 @@
allowDiscards = true;
keyFileSize = 4096;
keyFile = "/dev/disk/by-id/usb-USB_SanDisk_3.2Gen1_03021630090925173333-0:0";
fallbackToPassword = true;
};
};
kernelModules = [ "kvm-intel" ];
+11 -2
View File
@@ -4,17 +4,21 @@ let
in
{
imports = [
"${inputs.self}/users/richie"
"${inputs.self}/users/math"
"${inputs.self}/users/dov"
"${inputs.self}/users/math"
"${inputs.self}/users/richie"
"${inputs.self}/users/steve"
"${inputs.self}/common/global"
"${inputs.self}/common/optional/docker.nix"
"${inputs.self}/common/optional/monitoring-agent.nix"
"${inputs.self}/common/optional/ssh_decrypt.nix"
"${inputs.self}/common/optional/syncthing_base.nix"
"${inputs.self}/common/optional/update.nix"
"${inputs.self}/common/optional/zerotier.nix"
./monitoring
./docker
./services
./web_services
./hardware.nix
./networking.nix
./programs.nix
@@ -35,5 +39,10 @@ in
zerotierone.joinNetworks = [ "a09acf02330d37b9" ];
};
users.groups = {
nornsight = { };
nornsight-admin = { };
};
system.stateVersion = "24.05";
}
-1
View File
@@ -9,7 +9,6 @@ let
inherit device;
keyFileSize = 4096;
keyFile = "/dev/disk/by-id/usb-XIAO_USB_Drive_24587CE29074-0:0";
fallbackToPassword = true;
};
makeLuksSSD =
device:
@@ -0,0 +1,426 @@
{
"annotations": {
"list": [
{
"builtIn": 1,
"datasource": {
"type": "grafana",
"uid": "-- Grafana --"
},
"enable": true,
"hide": true,
"iconColor": "rgba(0, 211, 255, 1)",
"name": "Annotations & Alerts",
"type": "dashboard"
}
]
},
"editable": true,
"fiscalYearStartMonth": 0,
"graphTooltip": 0,
"links": [],
"panels": [
{
"datasource": {
"type": "prometheus",
"uid": "prom-main"
},
"fieldConfig": {
"defaults": {
"unit": "percent"
},
"overrides": []
},
"gridPos": {
"h": 8,
"w": 6,
"x": 0,
"y": 0
},
"id": 1,
"options": {
"legend": {
"displayMode": "list",
"placement": "bottom"
},
"tooltip": {
"mode": "multi"
}
},
"targets": [
{
"datasource": {
"type": "prometheus",
"uid": "prom-main"
},
"editorMode": "code",
"expr": "100 * (1 - avg by (instance) (rate(node_cpu_seconds_total{mode=\"idle\"}[5m])))",
"legendFormat": "{{instance}}",
"range": true,
"refId": "A"
}
],
"title": "CPU Used",
"type": "timeseries"
},
{
"datasource": {
"type": "prometheus",
"uid": "prom-main"
},
"fieldConfig": {
"defaults": {
"unit": "percent"
},
"overrides": []
},
"gridPos": {
"h": 8,
"w": 6,
"x": 6,
"y": 0
},
"id": 2,
"options": {
"legend": {
"displayMode": "list",
"placement": "bottom"
},
"tooltip": {
"mode": "multi"
}
},
"targets": [
{
"datasource": {
"type": "prometheus",
"uid": "prom-main"
},
"editorMode": "code",
"expr": "100 * (1 - (node_memory_MemAvailable_bytes / node_memory_MemTotal_bytes))",
"legendFormat": "{{instance}}",
"range": true,
"refId": "A"
}
],
"title": "RAM Used",
"type": "timeseries"
},
{
"datasource": {
"type": "prometheus",
"uid": "prom-main"
},
"fieldConfig": {
"defaults": {
"unit": "percent"
},
"overrides": []
},
"gridPos": {
"h": 8,
"w": 6,
"x": 12,
"y": 0
},
"id": 3,
"options": {
"legend": {
"displayMode": "list",
"placement": "bottom"
},
"tooltip": {
"mode": "multi"
}
},
"targets": [
{
"datasource": {
"type": "prometheus",
"uid": "prom-main"
},
"editorMode": "code",
"expr": "100 * (1 - (node_memory_SwapFree_bytes / node_memory_SwapTotal_bytes))",
"legendFormat": "{{instance}}",
"range": true,
"refId": "A"
}
],
"title": "Swap Used",
"type": "timeseries"
},
{
"datasource": {
"type": "prometheus",
"uid": "prom-main"
},
"fieldConfig": {
"defaults": {
"unit": "short"
},
"overrides": []
},
"gridPos": {
"h": 8,
"w": 6,
"x": 18,
"y": 0
},
"id": 4,
"options": {
"legend": {
"displayMode": "list",
"placement": "bottom"
},
"tooltip": {
"mode": "multi"
}
},
"targets": [
{
"datasource": {
"type": "prometheus",
"uid": "prom-main"
},
"editorMode": "code",
"expr": "node_load1",
"legendFormat": "{{instance}} load1",
"range": true,
"refId": "A"
},
{
"datasource": {
"type": "prometheus",
"uid": "prom-main"
},
"editorMode": "code",
"expr": "node_load5",
"legendFormat": "{{instance}} load5",
"range": true,
"refId": "B"
},
{
"datasource": {
"type": "prometheus",
"uid": "prom-main"
},
"editorMode": "code",
"expr": "node_load15",
"legendFormat": "{{instance}} load15",
"range": true,
"refId": "C"
}
],
"title": "Load",
"type": "timeseries"
},
{
"datasource": {
"type": "prometheus",
"uid": "prom-main"
},
"fieldConfig": {
"defaults": {
"unit": "Bps"
},
"overrides": []
},
"gridPos": {
"h": 9,
"w": 12,
"x": 0,
"y": 8
},
"id": 5,
"options": {
"legend": {
"displayMode": "list",
"placement": "bottom"
},
"tooltip": {
"mode": "multi"
}
},
"targets": [
{
"datasource": {
"type": "prometheus",
"uid": "prom-main"
},
"editorMode": "code",
"expr": "sum by (instance) (rate(node_disk_read_bytes_total[5m]))",
"legendFormat": "{{instance}} read",
"range": true,
"refId": "A"
},
{
"datasource": {
"type": "prometheus",
"uid": "prom-main"
},
"editorMode": "code",
"expr": "sum by (instance) (rate(node_disk_written_bytes_total[5m]))",
"legendFormat": "{{instance}} write",
"range": true,
"refId": "B"
}
],
"title": "Disk Throughput",
"type": "timeseries"
},
{
"datasource": {
"type": "prometheus",
"uid": "prom-main"
},
"fieldConfig": {
"defaults": {
"unit": "percent"
},
"overrides": []
},
"gridPos": {
"h": 9,
"w": 12,
"x": 12,
"y": 8
},
"id": 6,
"options": {
"cellHeight": "sm",
"showHeader": true,
"sortBy": [
{
"desc": true,
"displayName": "Value"
}
]
},
"targets": [
{
"datasource": {
"type": "prometheus",
"uid": "prom-main"
},
"editorMode": "code",
"expr": "100 * (1 - (node_filesystem_avail_bytes{mountpoint=~\"(/|/home|/var|/zfs.*)\",fstype!=\"\"} / node_filesystem_size_bytes{mountpoint=~\"(/|/home|/var|/zfs.*)\",fstype!=\"\"}))",
"format": "table",
"instant": true,
"legendFormat": "{{instance}} {{mountpoint}}",
"refId": "A"
}
],
"title": "Filesystem Usage",
"type": "table"
},
{
"datasource": {
"type": "prometheus",
"uid": "prom-main"
},
"fieldConfig": {
"defaults": {
"unit": "percentunit"
},
"overrides": []
},
"gridPos": {
"h": 10,
"w": 12,
"x": 0,
"y": 17
},
"id": 7,
"options": {
"cellHeight": "sm",
"showHeader": true,
"sortBy": [
{
"desc": true,
"displayName": "Value"
}
]
},
"targets": [
{
"datasource": {
"type": "prometheus",
"uid": "prom-main"
},
"editorMode": "code",
"expr": "topk(10, rate(namedprocess_namegroup_cpu_seconds_total[5m]))",
"format": "table",
"instant": true,
"legendFormat": "{{instance}} {{groupname}}",
"refId": "A"
}
],
"title": "Top Grouped CPU",
"type": "table"
},
{
"datasource": {
"type": "prometheus",
"uid": "prom-main"
},
"fieldConfig": {
"defaults": {
"unit": "bytes"
},
"overrides": []
},
"gridPos": {
"h": 10,
"w": 12,
"x": 12,
"y": 17
},
"id": 8,
"options": {
"cellHeight": "sm",
"showHeader": true,
"sortBy": [
{
"desc": true,
"displayName": "Value"
}
]
},
"targets": [
{
"datasource": {
"type": "prometheus",
"uid": "prom-main"
},
"editorMode": "code",
"expr": "topk(10, namedprocess_namegroup_memory_bytes{memtype=\"resident\"})",
"format": "table",
"instant": true,
"legendFormat": "{{instance}} {{groupname}}",
"refId": "A"
}
],
"title": "Top Grouped Memory",
"type": "table"
}
],
"refresh": "30s",
"schemaVersion": 39,
"style": "dark",
"tags": [
"monitoring"
],
"templating": {
"list": []
},
"time": {
"from": "now-24h",
"to": "now"
},
"timepicker": {},
"timezone": "",
"title": "Overview",
"uid": "monitor-overview",
"version": 1,
"weekStart": ""
}

Some files were not shown because too many files have changed in this diff Show More