Compare commits
27
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
6d6226ed14 | ||
|
|
39ab7358bb | ||
|
|
4488441a84 | ||
|
|
681a2d8d12 | ||
|
|
f3e36f7ec3 | ||
|
|
d8693736d9 | ||
|
|
c7cd63f8e4 | ||
|
|
5c73fc9e9b | ||
|
|
e596652582 | ||
|
|
507c3f5406 | ||
|
|
7c1618cfd6 | ||
|
|
086224a4c9 | ||
|
|
c11895f144 | ||
|
|
dab18c1385 | ||
|
|
e34ed6c597 | ||
|
|
a19ace7959 | ||
|
|
0f84c7ce41 | ||
|
|
cc31483973 | ||
|
|
790fb6680d | ||
|
|
927c5cbc60 | ||
|
|
8db44916c8 | ||
|
|
f20c45bea9 | ||
|
|
6a0e3ebcdb | ||
|
|
d9115f7c91 | ||
|
|
1eecf7181d | ||
|
|
3e164831b5 | ||
|
|
872e55da1d |
@@ -3,19 +3,6 @@
|
||||
.mypy_cache
|
||||
.pytest_cache
|
||||
.ruff_cache
|
||||
.venv
|
||||
**/.venv
|
||||
.env
|
||||
.cache
|
||||
.claude
|
||||
.coverage
|
||||
.vscode
|
||||
.stfolder
|
||||
.literotica_data
|
||||
esphome
|
||||
htmlcov
|
||||
data
|
||||
ebooks
|
||||
__pycache__
|
||||
**/__pycache__
|
||||
*.pyc
|
||||
|
||||
@@ -18,7 +18,6 @@ jobs:
|
||||
- "brain"
|
||||
- "jeeves"
|
||||
- "rhapsody-in-green"
|
||||
- "portal-1"
|
||||
continue-on-error: true
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
|
||||
@@ -19,5 +19,5 @@ jobs:
|
||||
python -m python.gitea_flake_lock merge
|
||||
--repo "${{ github.repository }}"
|
||||
env:
|
||||
JEEVES_BOT_TOKEN: ${{ secrets.JEEVES_BOT_TOKEN }}
|
||||
GITEA_TOKEN: ${{ secrets.GITEA_TOKEN }}
|
||||
GITEA_URL: https://gitea.tmmworkshop.com
|
||||
|
||||
@@ -1,26 +0,0 @@
|
||||
name: test ebook search
|
||||
|
||||
on:
|
||||
workflow_dispatch:
|
||||
push:
|
||||
branches:
|
||||
- main
|
||||
pull_request:
|
||||
branches:
|
||||
- main
|
||||
|
||||
env:
|
||||
UV_PYTHON_DOWNLOADS: never
|
||||
UV_CACHE_DIR: /var/cache/uv
|
||||
UV_LINK_MODE: copy
|
||||
|
||||
jobs:
|
||||
test-ebook-search:
|
||||
runs-on: self-hosted
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- name: Install dependencies
|
||||
run: nix develop .#devShells.x86_64-linux.ebook-search -c uv sync --locked --project python/ebook_search/docker
|
||||
- name: Run ebook search tests
|
||||
run: nix develop .#devShells.x86_64-linux.ebook-search -c uv run --project python/ebook_search/docker --no-sync pytest tests/ebook_search --override-ini addopts="-n auto -ra"
|
||||
@@ -18,7 +18,7 @@ jobs:
|
||||
run: nix flake update
|
||||
- name: Create or update flake.lock PR
|
||||
env:
|
||||
JEEVES_BOT_TOKEN: ${{ secrets.JEEVES_BOT_TOKEN }}
|
||||
GITEA_TOKEN: ${{ secrets.GITEA_TOKEN }}
|
||||
GITEA_URL: https://gitea.tmmworkshop.com
|
||||
run: >-
|
||||
nix develop .#devShells.x86_64-linux.default -c
|
||||
|
||||
@@ -173,6 +173,3 @@ frontend/node_modules/
|
||||
# data from testing llms
|
||||
data/*
|
||||
.ebook_search_bm25
|
||||
|
||||
# gems data
|
||||
.gems
|
||||
|
||||
@@ -1,9 +1,6 @@
|
||||
# Generate AGE keys from SSH keys with:
|
||||
# ssh-keygen -A
|
||||
# nix-shell -p ssh-to-age --run 'cat /etc/ssh/ssh_host_ed25519_key.pub | ssh-to-age'
|
||||
# update keys after addin/removing a key
|
||||
# nix-shell -p sops --run "sops updatekeys users/secrets.yaml" users/secrets.yaml
|
||||
|
||||
keys:
|
||||
- &admin_richie age1u8zj599elqqvcmhxn8zuwrufsz8w8w366d3ayrljjejljt2q45kq8mxw9c # cspell:disable-line
|
||||
|
||||
@@ -11,7 +8,6 @@ keys:
|
||||
- &system_brain age1jhf7vm0005j60mjq63696frrmjhpy8kpc2d66mw044lqap5mjv4snmwvwm # cspell:disable-line
|
||||
- &system_jeeves age13lmqgc3jvkyah5e3vcwmj4s5wsc2akctcga0lpc0x8v8du3fxprqp4ldkv # cspell:disable-line
|
||||
- &system_rhapsody age1ufnewppysaq2wwcl4ugngjz8pfzc5a35yg7luq0qmuqvctajcycs5lf6k4 # cspell:disable-line
|
||||
- &system_portal_1 age1vyav6kxtvt3z4vtnkkjj38eu8hlts5m7ygyckhskvalg2gpjk52su53d0a # cspell:disable-line
|
||||
|
||||
creation_rules:
|
||||
- path_regex: users/secrets\.yaml$
|
||||
@@ -22,4 +18,3 @@ creation_rules:
|
||||
- *system_brain
|
||||
- *system_jeeves
|
||||
- *system_rhapsody
|
||||
- *system_portal_1
|
||||
|
||||
Vendored
+3
@@ -10,6 +10,7 @@
|
||||
"aiounifi",
|
||||
"alsa",
|
||||
"apiclient",
|
||||
"apscheduler",
|
||||
"archlinux",
|
||||
"ashift",
|
||||
"asrouter",
|
||||
@@ -336,6 +337,8 @@
|
||||
"yubioath",
|
||||
"yzhang",
|
||||
"zeroconf",
|
||||
"zerotier",
|
||||
"zerotierone",
|
||||
"zoxide",
|
||||
"zram",
|
||||
"zstd"
|
||||
|
||||
Generated
-1686
File diff suppressed because it is too large
Load Diff
@@ -1,4 +0,0 @@
|
||||
[workspace]
|
||||
resolver = "2"
|
||||
|
||||
members = ["rust/*"]
|
||||
@@ -17,11 +17,16 @@
|
||||
./nix.nix
|
||||
./programs.nix
|
||||
./ssh.nix
|
||||
./snapshot_manager.nix
|
||||
];
|
||||
|
||||
boot = {
|
||||
tmp.useTmpfs = lib.mkDefault true;
|
||||
tmp.useTmpfs = true;
|
||||
kernelPackages = lib.mkDefault pkgs.linuxPackages_6_12;
|
||||
zfs = {
|
||||
package = lib.mkDefault pkgs.zfs_2_4;
|
||||
forceImportRoot = lib.mkDefault false;
|
||||
};
|
||||
};
|
||||
|
||||
hardware.enableRedistributableFirmware = true;
|
||||
@@ -37,6 +42,9 @@
|
||||
overlays = builtins.attrValues outputs.overlays;
|
||||
config = {
|
||||
allowUnfree = true;
|
||||
permittedInsecurePackages = [
|
||||
"openssl-1.1.1w" # This is for discord-canary
|
||||
];
|
||||
};
|
||||
};
|
||||
|
||||
@@ -45,6 +53,16 @@
|
||||
|
||||
# firmware update
|
||||
fwupd.enable = true;
|
||||
|
||||
snapshot_manager = {
|
||||
enable = lib.mkDefault true;
|
||||
PYTHONPATH = "${inputs.self}/";
|
||||
};
|
||||
|
||||
zfs = {
|
||||
trim.enable = lib.mkDefault true;
|
||||
autoScrub.enable = lib.mkDefault true;
|
||||
};
|
||||
};
|
||||
|
||||
powerManagement.powertop.enable = lib.mkDefault true;
|
||||
|
||||
@@ -31,10 +31,6 @@ in
|
||||
"flakes"
|
||||
"ca-derivations"
|
||||
];
|
||||
system-features = lib.mkAfter [
|
||||
"gccarch-x86-64-v2"
|
||||
"gccarch-x86-64-v3"
|
||||
];
|
||||
warn-dirty = false;
|
||||
flake-registry = ""; # disable global flake registries
|
||||
connect-timeout = 10;
|
||||
|
||||
@@ -22,12 +22,6 @@ hourly = 0
|
||||
daily = 0
|
||||
monthly = 0
|
||||
|
||||
["root_pool/nix_build"]
|
||||
15_min = 1
|
||||
hourly = 0
|
||||
daily = 0
|
||||
monthly = 0
|
||||
|
||||
["root_pool/var"]
|
||||
15_min = 8
|
||||
hourly = 24
|
||||
@@ -17,7 +17,6 @@
|
||||
logDriver = "local";
|
||||
storageDriver = "overlay2";
|
||||
daemon.settings = {
|
||||
live-restore = false;
|
||||
experimental = true;
|
||||
exec-opts = [ "native.cgroupdriver=systemd" ];
|
||||
log-opts = {
|
||||
|
||||
@@ -5,7 +5,7 @@
|
||||
...
|
||||
}:
|
||||
let
|
||||
monitoringInterface = "tailscale0";
|
||||
monitoringInterface = "ztwfunumly";
|
||||
nodeTextfileDir = "/var/lib/prometheus-node-exporter-textfile";
|
||||
|
||||
mkProcessNameTemplate =
|
||||
|
||||
@@ -1,37 +0,0 @@
|
||||
{
|
||||
config,
|
||||
inputs,
|
||||
...
|
||||
}:
|
||||
{
|
||||
nix.settings = {
|
||||
trusted-substituters = [ "http://jeeves:5000" ];
|
||||
substituters = [ "http://jeeves:5000/?priority=1&want-mass-query=true" ];
|
||||
trusted-public-keys = [ "cache.tmmworkshop.com:jHffkpgbmEdstQPoihJPYW9TQe6jnQbWR2LqkNGV3iA=" ];
|
||||
};
|
||||
|
||||
services.tailscale = {
|
||||
enable = true;
|
||||
openFirewall = true;
|
||||
authKeyFile = config.sops.secrets.tailscale_auth_key.path;
|
||||
|
||||
# OAuth client secrets create ephemeral nodes by default. NixOS machines
|
||||
# are persistent and should enroll without interactive device approval.
|
||||
authKeyParameters = {
|
||||
ephemeral = false;
|
||||
preauthorized = true;
|
||||
};
|
||||
|
||||
extraUpFlags = [ "--advertise-tags=tag:nixos" ];
|
||||
};
|
||||
|
||||
sops = {
|
||||
age.sshKeyPaths = [ "/etc/ssh/ssh_host_ed25519_key" ];
|
||||
|
||||
secrets.tailscale_auth_key = {
|
||||
sopsFile = "${inputs.self}/users/secrets.yaml";
|
||||
owner = "root";
|
||||
mode = "0400";
|
||||
};
|
||||
};
|
||||
}
|
||||
@@ -1,9 +0,0 @@
|
||||
{
|
||||
nixpkgs.hostPlatform = {
|
||||
system = "x86_64-linux";
|
||||
gcc = {
|
||||
arch = "x86-64-v3";
|
||||
tune = "generic";
|
||||
};
|
||||
};
|
||||
}
|
||||
@@ -0,0 +1,11 @@
|
||||
{
|
||||
services.zerotierone = {
|
||||
enable = true;
|
||||
joinNetworks = [ "e4da7455b2ae64ca" ];
|
||||
};
|
||||
nix.settings = {
|
||||
trusted-substituters = [ "http://192.168.90.40:5000" ];
|
||||
substituters = [ "http://192.168.90.40:5000/?priority=1&want-mass-query=true" ];
|
||||
trusted-public-keys = [ "cache.tmmworkshop.com:jHffkpgbmEdstQPoihJPYW9TQe6jnQbWR2LqkNGV3iA=" ];
|
||||
};
|
||||
}
|
||||
@@ -1,26 +0,0 @@
|
||||
{
|
||||
inputs,
|
||||
lib,
|
||||
pkgs,
|
||||
...
|
||||
}:
|
||||
{
|
||||
imports = [ ./snapshot.nix ];
|
||||
|
||||
boot.zfs = {
|
||||
package = lib.mkDefault pkgs.zfs_2_4;
|
||||
forceImportRoot = lib.mkDefault false;
|
||||
};
|
||||
|
||||
services = {
|
||||
snapshot_manager = {
|
||||
enable = lib.mkDefault true;
|
||||
PYTHONPATH = "${inputs.self}/";
|
||||
};
|
||||
|
||||
zfs = {
|
||||
trim.enable = lib.mkDefault true;
|
||||
autoScrub.enable = lib.mkDefault true;
|
||||
};
|
||||
};
|
||||
}
|
||||
Generated
+28
-26
@@ -1,23 +1,25 @@
|
||||
{
|
||||
"nodes": {
|
||||
"disko": {
|
||||
"firefox-addons": {
|
||||
"inputs": {
|
||||
"nixpkgs": [
|
||||
"nixpkgs"
|
||||
]
|
||||
},
|
||||
"locked": {
|
||||
"lastModified": 1781152676,
|
||||
"narHash": "sha256-RxWs5ND31KzTG7wvMM+PMfUjyNpmIEr999lqNARaM5o=",
|
||||
"owner": "nix-community",
|
||||
"repo": "disko",
|
||||
"rev": "ff8702b4de27f72b4c78573dfb89ec74e36abdf1",
|
||||
"type": "github"
|
||||
"dir": "pkgs/firefox-addons",
|
||||
"lastModified": 1782964936,
|
||||
"narHash": "sha256-wXEBDr7/dFQYhVpDwCKc9fkrYQQE4x0bdirX1bsLBGA=",
|
||||
"owner": "rycee",
|
||||
"repo": "nur-expressions",
|
||||
"rev": "64feee871e0373dd6121e412c3fb12e372d1bfb5",
|
||||
"type": "gitlab"
|
||||
},
|
||||
"original": {
|
||||
"owner": "nix-community",
|
||||
"repo": "disko",
|
||||
"type": "github"
|
||||
"dir": "pkgs/firefox-addons",
|
||||
"owner": "rycee",
|
||||
"repo": "nur-expressions",
|
||||
"type": "gitlab"
|
||||
}
|
||||
},
|
||||
"home-manager": {
|
||||
@@ -27,11 +29,11 @@
|
||||
]
|
||||
},
|
||||
"locked": {
|
||||
"lastModified": 1788651960,
|
||||
"narHash": "sha256-v9wJd32eZ2bvhBzVOd7TIjLQd011P7nwOhjKtWlci5I=",
|
||||
"lastModified": 1783005591,
|
||||
"narHash": "sha256-NcLHV5uBAeggDUE2wPbKszjfyaSLsoqaYt7izOphkZw=",
|
||||
"owner": "nix-community",
|
||||
"repo": "home-manager",
|
||||
"rev": "2c0350c759688177331b8f5242311fae8877bdb3",
|
||||
"rev": "f469c79b955609d6a8fdd9e689be76a93b1621d7",
|
||||
"type": "github"
|
||||
},
|
||||
"original": {
|
||||
@@ -45,11 +47,11 @@
|
||||
"nixpkgs": "nixpkgs"
|
||||
},
|
||||
"locked": {
|
||||
"lastModified": 1788860136,
|
||||
"narHash": "sha256-MhPMOFV4pVkygWEbQ8t1De/uQ9cWF1u++tRe2L5tG48=",
|
||||
"lastModified": 1782562157,
|
||||
"narHash": "sha256-a7+T6QSeowynwZ1ZJJbP8T8ntAytvrui8kFGJmIZt2c=",
|
||||
"owner": "nixos",
|
||||
"repo": "nixos-hardware",
|
||||
"rev": "62173785b9a18c78b4a15aca2623d02bceb9d077",
|
||||
"rev": "a9cf7546a938c737b079e738de73934a13de9784",
|
||||
"type": "github"
|
||||
},
|
||||
"original": {
|
||||
@@ -74,11 +76,11 @@
|
||||
},
|
||||
"nixpkgs-master": {
|
||||
"locked": {
|
||||
"lastModified": 1788892992,
|
||||
"narHash": "sha256-cIMFh9gyU4/aLeB3JCcsWM3tTAvD9pAq9Smr1Wa8aIU=",
|
||||
"lastModified": 1783021952,
|
||||
"narHash": "sha256-8PghAtSGGZ0umfVI8Qbd7ZbFrfZPiH1UwtVbgLeikDA=",
|
||||
"owner": "nixos",
|
||||
"repo": "nixpkgs",
|
||||
"rev": "dff6994123e257ec9901c271bc2b52e64d7c8f05",
|
||||
"rev": "f136374c679c54171a3ace589d15e9e79a8bd086",
|
||||
"type": "github"
|
||||
},
|
||||
"original": {
|
||||
@@ -106,11 +108,11 @@
|
||||
},
|
||||
"nixpkgs_2": {
|
||||
"locked": {
|
||||
"lastModified": 1788752844,
|
||||
"narHash": "sha256-VaWGJ6+cIYN2erfSecbRV+4ljI185Ty2wUrXyvQbgOw=",
|
||||
"lastModified": 1782723713,
|
||||
"narHash": "sha256-oPXCU/SSUokcGaJREHibG1CBX3+s/W7orDWQOZDsEeQ=",
|
||||
"owner": "nixos",
|
||||
"repo": "nixpkgs",
|
||||
"rev": "dc5d91f840324650bac8c379428c7037a416959a",
|
||||
"rev": "b5aa0fbd538984f6e3d201be0005b4463d8b09f8",
|
||||
"type": "github"
|
||||
},
|
||||
"original": {
|
||||
@@ -122,7 +124,7 @@
|
||||
},
|
||||
"root": {
|
||||
"inputs": {
|
||||
"disko": "disko",
|
||||
"firefox-addons": "firefox-addons",
|
||||
"home-manager": "home-manager",
|
||||
"nixos-hardware": "nixos-hardware",
|
||||
"nixpkgs": "nixpkgs_2",
|
||||
@@ -139,11 +141,11 @@
|
||||
]
|
||||
},
|
||||
"locked": {
|
||||
"lastModified": 1788337237,
|
||||
"narHash": "sha256-gkSH8VUtCo6hnysNmb9DbTuDepH2t5pv+QWjP75xKAk=",
|
||||
"lastModified": 1782165805,
|
||||
"narHash": "sha256-478kKQBvK6SYTOdN2h9jhKJv94nbXRbFMfuL1WshErg=",
|
||||
"owner": "Mic92",
|
||||
"repo": "sops-nix",
|
||||
"rev": "fbf759290e0cb0a98dfc813a4eb7d53ad1dacb57",
|
||||
"rev": "56b24064fdcaedca53553b1a6d607fd23b613a24",
|
||||
"type": "github"
|
||||
},
|
||||
"original": {
|
||||
|
||||
@@ -26,13 +26,13 @@
|
||||
inputs.nixpkgs.follows = "nixpkgs";
|
||||
};
|
||||
|
||||
sops-nix = {
|
||||
url = "github:Mic92/sops-nix";
|
||||
firefox-addons = {
|
||||
url = "gitlab:rycee/nur-expressions?dir=pkgs/firefox-addons";
|
||||
inputs.nixpkgs.follows = "nixpkgs";
|
||||
};
|
||||
|
||||
disko = {
|
||||
url = "github:nix-community/disko";
|
||||
sops-nix = {
|
||||
url = "github:Mic92/sops-nix";
|
||||
inputs.nixpkgs.follows = "nixpkgs";
|
||||
};
|
||||
};
|
||||
|
||||
@@ -15,19 +15,17 @@
|
||||
};
|
||||
};
|
||||
|
||||
patches = import ./patches;
|
||||
test-exclusions = import ./test-exclusions.nix;
|
||||
x86-64-v3-workarounds = import ./x86-64-v3-workarounds.nix;
|
||||
|
||||
python-env = final: _prev: {
|
||||
my_python = final.python314.withPackages (
|
||||
ps: with ps; [
|
||||
ps:
|
||||
with ps;
|
||||
[
|
||||
alembic
|
||||
apprise
|
||||
apscheduler
|
||||
fastapi
|
||||
fastapi-cli
|
||||
httpx
|
||||
jinja2
|
||||
mypy
|
||||
pgvector
|
||||
psycopg
|
||||
@@ -38,13 +36,11 @@
|
||||
pytest-mock
|
||||
pytest-xdist
|
||||
python-multipart
|
||||
pydantic-settings
|
||||
ruff
|
||||
sqlalchemy
|
||||
tenacity
|
||||
tinytuya
|
||||
typer
|
||||
uvicorn
|
||||
websockets
|
||||
]
|
||||
);
|
||||
|
||||
@@ -1,20 +0,0 @@
|
||||
Subject: [PATCH] abseil: include BMI2 intrinsics through the public header
|
||||
|
||||
GCC and Clang reject direct inclusion of bmi2intrin.h. Include immintrin.h
|
||||
instead so that the compiler supplies the required intrinsic setup when
|
||||
BMI2 is enabled, including builds targeting x86-64-v3.
|
||||
|
||||
This patch is shared by the vendored Abseil copies in Electron, rusty_v8
|
||||
(Deno), and Signal's WebRTC build.
|
||||
|
||||
--- a/third_party/abseil-cpp/absl/container/internal/raw_hash_set.h
|
||||
+++ b/third_party/abseil-cpp/absl/container/internal/raw_hash_set.h
|
||||
@@ -226,7 +226,7 @@
|
||||
#endif
|
||||
|
||||
#ifdef __BMI2__
|
||||
-#include <bmi2intrin.h>
|
||||
+#include <immintrin.h>
|
||||
#endif // __BMI2__
|
||||
|
||||
namespace absl {
|
||||
@@ -1,38 +0,0 @@
|
||||
# Abseil accepted the upstream fix: https://github.com/abseil/abseil-cpp/pull/2071
|
||||
# Keep this workaround until Electron, Deno's rusty_v8, and Signal's WebRTC
|
||||
# update their bundled Abseil copies to include it.
|
||||
{ prev }:
|
||||
let
|
||||
patchAbseilBmi2Include =
|
||||
package:
|
||||
package.overrideAttrs (old: {
|
||||
# GCC and Clang require the public umbrella header for BMI2 intrinsics.
|
||||
patches = (old.patches or [ ]) ++ [ ./bmi2-public-header.patch ];
|
||||
});
|
||||
|
||||
electron43Unwrapped = patchAbseilBmi2Include prev.electron_43.unwrapped;
|
||||
electron43 = prev.electron_43.override {
|
||||
electron-unwrapped = electron43Unwrapped;
|
||||
};
|
||||
|
||||
signalCallPackage =
|
||||
path: args:
|
||||
let
|
||||
package = prev.callPackage path args;
|
||||
in
|
||||
if builtins.baseNameOf path == "webrtc.nix" then patchAbseilBmi2Include package else package;
|
||||
in
|
||||
prev.lib.optionalAttrs ((prev.stdenv.hostPlatform.gcc.arch or null) == "x86-64-v3") {
|
||||
deno =
|
||||
let
|
||||
librusty_v8 = patchAbseilBmi2Include prev.deno.passthru.librusty_v8;
|
||||
in
|
||||
prev.deno.override { inherit librusty_v8; };
|
||||
|
||||
electron_43 = electron43;
|
||||
|
||||
signal-desktop = prev.signal-desktop.override {
|
||||
electron_43 = electron43;
|
||||
callPackage = signalCallPackage;
|
||||
};
|
||||
}
|
||||
@@ -1,12 +0,0 @@
|
||||
_final: prev:
|
||||
(import ./abseil { inherit prev; })
|
||||
// {
|
||||
gnutls = import ./gnutls { inherit (prev) gnutls; };
|
||||
|
||||
pythonPackagesExtensions = prev.pythonPackagesExtensions ++ [
|
||||
(_pythonFinal: pythonPrev: {
|
||||
scipy = import ./scipy { inherit (pythonPrev) scipy; };
|
||||
torchcodec = import ./torchcodec { inherit (pythonPrev) torchcodec; };
|
||||
})
|
||||
];
|
||||
}
|
||||
@@ -1,160 +0,0 @@
|
||||
# GnuTLS UDP server readiness
|
||||
|
||||
Under load, the test client can start before `gnutls-serv` binds its UDP
|
||||
socket, and the first handshake fails with `Connection refused`.
|
||||
`serv-udp.sh` currently waits a fixed four seconds; elapsed time does not
|
||||
establish server readiness. `udp-server-readiness.patch` replaces that wait
|
||||
with polling for the local IPv4 UDP endpoint.
|
||||
|
||||
## Scope and waiting behavior
|
||||
|
||||
The patch changes the existing `wait_udp_server()` and adds a new
|
||||
`check_if_udp_port_bound()` beside it in `tests/scripts/common.sh`.
|
||||
`serv-udp.sh` is its only caller in 3.8.13. The TCP helpers `wait_server()`
|
||||
and `wait_for_port()`, including their existing sleeps, are unchanged.
|
||||
Both original DTLS handshake checks remain unchanged.
|
||||
|
||||
Each iteration checks process liveness and the socket **before sleeping**.
|
||||
A ready socket returns immediately. An unsuccessful check sleeps two
|
||||
seconds only if another attempt remains: at most 90 attempts, consistent
|
||||
with the existing `wait_server()` budget implemented by `wait_for_port()`,
|
||||
with no sleep after the final check. Server exit fails early; exhausting the
|
||||
budget fails and terminates the server. No handshake is retried, and no
|
||||
protocol timeout is changed. Once bound, the kernel can queue datagrams
|
||||
while the server is scheduled; the probe itself sends no packets.
|
||||
|
||||
The existing `have_port_finder()` prefers `ss`, then `netstat`. If neither
|
||||
exists, it prints `neither ss nor netstat found` and exits **77 (skip)**.
|
||||
In the normal test flow, port selection calls it before launching a server.
|
||||
The probe runs in a subshell so that, even if this skip occurs after launch,
|
||||
the waiting helper can terminate and reap the server before exiting 77.
|
||||
|
||||
## Why an IPv4 socket is expected
|
||||
|
||||
This is specific to the server used by this test, not a general rule that
|
||||
IPv6 sockets cannot serve IPv4 clients. The client explicitly uses
|
||||
`127.0.0.1`. The server's `--udp` path calls `udp_server()`, which calls
|
||||
`listen_socket(..., SOCK_DGRAM)`. That function iterates the wildcard
|
||||
addresses returned by `getaddrinfo(NULL, ..., AI_PASSIVE)`:
|
||||
|
||||
| Server build / Linux setting | Binding behavior |
|
||||
| --- | --- |
|
||||
| IPv6 enabled, `net.ipv6.bindv6only=0` | Requests `IPV6_V6ONLY=1` on the IPv6 socket, binds `[::]:PORT`, and separately binds `0.0.0.0:PORT`. It overrides the system's dual-stack default. |
|
||||
| IPv6 enabled, `net.ipv6.bindv6only=1` | The same explicit socket option and separate IPv4/IPv6 binds. |
|
||||
| `HAVE_IPV6` undefined | Skips every address family except `AF_INET`; only the IPv4 wildcard is attempted. |
|
||||
|
||||
`udp_server()` uses `wait_for_connection()`, which puts **every listener**
|
||||
from that list into `select()` and returns a readable socket for `recvfrom()`;
|
||||
it does not permanently choose one socket based on `getaddrinfo()` order.
|
||||
|
||||
The first two cases were traced with the actual GnuTLS 3.8.13 binary in
|
||||
separate Linux network namespaces: `setsockopt(IPV6_V6ONLY, [1])` and both
|
||||
UDP binds returned success under each setting. The no-IPv6 case was checked
|
||||
in source, not by building a second binary. The same bind implementation
|
||||
was checked directly on GitLab master.
|
||||
|
||||
Thus, successful normal startup for this invocation provides an explicit
|
||||
IPv4 socket; a lone IPv6 wildcard is not the expected success path.
|
||||
There is one portability caveat: upstream discards the return value of
|
||||
`setsockopt(IPV6_V6ONLY)`. On a platform where that call fails and the server
|
||||
ends up with only a dual-stack socket, this helper would time out despite
|
||||
IPv4 reachability. Such a platform needs additional handling before this
|
||||
patch can claim support. Blindly accepting every IPv6 wildcard would also
|
||||
accept IPv6-only sockets before the separate IPv4 bind finishes.
|
||||
|
||||
Source: [`src/serv.c`, `listen_socket()`](https://gitlab.com/gnutls/gnutls/-/blob/master/src/serv.c#L937),
|
||||
[`src/udp-serv.c`](https://gitlab.com/gnutls/gnutls/-/blob/master/src/udp-serv.c),
|
||||
and [`tests/serv-udp.sh`](https://gitlab.com/gnutls/gnutls/-/blob/master/tests/serv-udp.sh).
|
||||
|
||||
## Port matching and ownership limit
|
||||
|
||||
Only `-an` is passed to the socket-listing tool: BSD `netstat -u` selects
|
||||
Unix-domain sockets, whereas Linux `netstat -u` selects UDP. The parser
|
||||
handles the extra state column in `ss`, Linux colon-separated endpoints,
|
||||
and BSD dot-separated endpoints, including `*.PORT`. It matches the full
|
||||
local port and rejects TCP, IPv6 entries, peer ports, and longer numbers.
|
||||
|
||||
A live PID plus a bound port does **not** prove that PID owns the socket.
|
||||
Existing `GETPORT` selection checks for an unused port and uses a test
|
||||
port-lock directory; `launch_bare_server()` also calls
|
||||
`wait_for_free_port()` before starting the process. These are advisory:
|
||||
the launcher does not enforce the latter's result, and another process
|
||||
can bind between the check and launch. The patch does not close that race
|
||||
or add nonportable PID parsing. An unrelated process can satisfy the
|
||||
socket check; the real handshakes remain the functional check and may
|
||||
fail (or reach the wrong server). This is a startup-order fix, not a
|
||||
socket-ownership guarantee.
|
||||
|
||||
## Reproduction and focused checks
|
||||
|
||||
Apply the patch to an unpacked source tree, then run the companion checks
|
||||
with Python's standard library and a shell:
|
||||
|
||||
```sh
|
||||
patch --fuzz=0 -d /path/to/gnutls -p1 < udp-server-readiness.patch
|
||||
SHELL=/bin/sh python3 verify-readiness.py /path/to/gnutls/tests/scripts/common.sh -v
|
||||
```
|
||||
|
||||
Set `NETSTAT=/path/to/netstat` to exercise one outside `PATH`. The checks
|
||||
cover Linux/BSD output samples, false matches, immediate readiness,
|
||||
missing tools, process exit, timeout cleanup, and real IPv4 UDP sockets
|
||||
whose bind is delayed six seconds. The missing-tools fixture is skipped
|
||||
if an absolute fallback `ss` path cannot be hidden with `PATH`. Native
|
||||
BSD execution remains untested.
|
||||
|
||||
To reproduce with GnuTLS itself, run `tests/serv-udp.sh` with `SERV` pointing
|
||||
to a wrapper that sleeps six seconds, then `exec`s `gnutls-serv` with all
|
||||
arguments. Set `CLI` to the matching `gnutls-cli`, `srcdir` to the source
|
||||
`tests` directory, and `abs_top_builddir` to a writable build directory.
|
||||
With GnuTLS 3.8.13, the original helper failed the first handshake with
|
||||
`Connection refused`; the patched helper passed both with the same binaries.
|
||||
|
||||
## GnuTLS submission
|
||||
|
||||
Development and merge requests are on [GitLab](https://gitlab.com/gnutls/gnutls).
|
||||
[`CONTRIBUTING.md` on master](https://gitlab.com/gnutls/gnutls/-/blob/master/CONTRIBUTING.md)
|
||||
was read directly for this review. It requires the contributor's DCO
|
||||
`Signed-off-by`, successful and failure test coverage, consistent coding
|
||||
style, and adequate documentation; GitLab CI runs for merge requests.
|
||||
Its commenting guidance asks for comments explaining non-obvious behavior
|
||||
or protocol expectations. It does not prescribe an additional special
|
||||
test-suite comment. The patch now explains its IPv4 binding assumption
|
||||
next to the probe.
|
||||
|
||||
The submission will contain the shell patch, without the Python verifier
|
||||
or a new Python test dependency. The existing `serv-udp.sh` supplies the
|
||||
functional success check. Running it through the six-second startup
|
||||
wrapper supplies a reproducible regression case: it fails before the fix
|
||||
and passes after it. The local verifier was used to validate socket-output
|
||||
parsing and the helper's success, process-exit, skip-cleanup, and timeout
|
||||
branches. Those branch checks are local evidence, not new automated
|
||||
coverage in the upstream suite; the MR must state that distinction.
|
||||
|
||||
No dedicated unit-test harness for these shell helpers was found in the
|
||||
3.8.13 tests inspected. That does not establish that Python cannot be used
|
||||
upstream; keeping this submission dependency-free is a scope choice. Use
|
||||
the existing test and before/after reproduction as the submission's
|
||||
coverage argument, retaining the platform limitations above. Apply the
|
||||
patch in an upstream checkout and include those results with the
|
||||
contributor's own sign-off. No MR or sign-off has been created.
|
||||
|
||||
## Local NixOS integration and build results
|
||||
|
||||
`overlays/default.nix` imports the `overlays/patches` overlay, which loads
|
||||
`gnutls/default.nix` to apply the patch and keep `serv-udp.sh` enabled.
|
||||
The patch itself has no Nix dependencies and applies to 3.8.13 and GitLab
|
||||
master without fuzz.
|
||||
|
||||
The final patch was rebuilt with:
|
||||
|
||||
```sh
|
||||
nix build --no-link -L .#nixosConfigurations.jeeves.pkgs.gnutls
|
||||
```
|
||||
|
||||
That x86-64-v3 build passed: 927 tests, 796 passes, 131 existing skips,
|
||||
zero failures/errors, and `PASS: serv-udp.sh`. The patch bytes in the built
|
||||
derivation were compared with the repository artifact; both have SHA-256
|
||||
`59013d47fd446f2dd065012a2259ccc1898fedc8a053a630e13efa0076368760`.
|
||||
All seven local checks passed, including skip cleanup and exactly 90
|
||||
probes with 89 sleeps on timeout. The six-second before/after reproduction
|
||||
was also repeated successfully with the final helper.
|
||||
@@ -1,6 +0,0 @@
|
||||
{ gnutls }:
|
||||
gnutls.overrideAttrs (old: {
|
||||
# Keep the UDP handshake test enabled on loaded builders by waiting for
|
||||
# the server to bind its socket. Kept as a standalone patch for upstream.
|
||||
patches = (old.patches or [ ]) ++ [ ./udp-server-readiness.patch ];
|
||||
})
|
||||
@@ -1,70 +0,0 @@
|
||||
Subject: [PATCH] tests: wait for the UDP server socket before connecting
|
||||
|
||||
A fixed four-second sleep does not guarantee that gnutls-serv has bound
|
||||
its UDP socket on a busy builder. Poll the local IPv4 UDP endpoint using
|
||||
the existing ss/netstat discovery, with the same retry budget as the TCP
|
||||
helper. Fail early if the server exits, and retain the original handshake
|
||||
checks in serv-udp.sh.
|
||||
|
||||
Use flags common to ss and BSD/Linux netstat. Match the local endpoint
|
||||
and complete port number, excluding TCP, IPv6-only and peer endpoints.
|
||||
|
||||
--- a/tests/scripts/common.sh
|
||||
+++ b/tests/scripts/common.sh
|
||||
@@ -185,10 +185,55 @@
|
||||
fi
|
||||
}
|
||||
|
||||
+check_if_udp_port_bound() {
|
||||
+ local PORT=$1
|
||||
+ have_port_finder
|
||||
+ # Use only -an, which is shared by ss and BSD/Linux netstat. UDP has
|
||||
+ # no LISTEN state. Match the local IPv4 endpoint, not a peer port or
|
||||
+ # a longer port number. serv-udp.sh connects to 127.0.0.1;
|
||||
+ # listen_socket() in serv.c binds IPv4 separately and requests
|
||||
+ # IPV6_V6ONLY=1 for its IPv6 socket.
|
||||
+ $PFCMD -an | awk -v port="$PORT" '
|
||||
+ $1 == "udp" || $1 == "udp4" {
|
||||
+ # ss includes a state column; netstat does not.
|
||||
+ address = ($2 == "UNCONN" || $2 == "ESTAB") ? $5 : $4
|
||||
+ if (address ~ ("^[0-9.]+[.:]" port "$") ||
|
||||
+ address == "*." port)
|
||||
+ found = 1
|
||||
+ }
|
||||
+ END { exit !found }
|
||||
+ '
|
||||
+}
|
||||
+
|
||||
wait_udp_server() {
|
||||
local PID=$1
|
||||
+ local ret
|
||||
trap "test -n \"${PID}\" && kill ${PID};exit 1" 1 15 2
|
||||
- sleep 4
|
||||
+ local i=0
|
||||
+ # Use the same retry budget as wait_for_port(), but also stop if the
|
||||
+ # server exits before binding its socket.
|
||||
+ while test $i -lt 90; do
|
||||
+ if ! kill -0 "$PID" 2>/dev/null; then
|
||||
+ fail "" "UDP server $PID exited before binding port $PORT"
|
||||
+ fi
|
||||
+ # Contain have_port_finder's exit so a skip also stops the server.
|
||||
+ if (check_if_udp_port_bound "$PORT"); then
|
||||
+ return 0
|
||||
+ else
|
||||
+ ret=$?
|
||||
+ if test "$ret" = 77; then
|
||||
+ kill "$PID" 2>/dev/null || :
|
||||
+ wait "$PID" 2>/dev/null || :
|
||||
+ exit 77
|
||||
+ fi
|
||||
+ fi
|
||||
+ i=$((i + 1))
|
||||
+ if test $i -lt 90; then
|
||||
+ echo "try $i: waiting for UDP port $PORT"
|
||||
+ sleep 2
|
||||
+ fi
|
||||
+ done
|
||||
+ fail "$PID" "UDP server $PORT did not come up"
|
||||
}
|
||||
|
||||
create_testdir() {
|
||||
@@ -1,180 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Exercise patched common.sh without building GnuTLS (Python standard library only).
|
||||
|
||||
Usage: python3 verify-readiness.py /path/to/patched/tests/scripts/common.sh
|
||||
Set SHELL to test another shell, and NETSTAT to test a netstat outside PATH.
|
||||
"""
|
||||
|
||||
# Use unittest so this upstream companion tool needs no pytest installation.
|
||||
# ruff: noqa: PT009
|
||||
|
||||
import os
|
||||
import shutil
|
||||
import socket
|
||||
import subprocess
|
||||
import sys
|
||||
import tempfile
|
||||
import time
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
COMMON = str(Path(sys.argv.pop(1)).resolve())
|
||||
SHELL = os.environ.get("SHELL", "/bin/sh")
|
||||
|
||||
|
||||
class ReadinessTests(unittest.TestCase):
|
||||
"""Check endpoint parsing and the server startup lifecycle."""
|
||||
|
||||
def setUp(self) -> None:
|
||||
"""Create a socket-listing fixture for each check."""
|
||||
self.tmp = tempfile.TemporaryDirectory()
|
||||
self.addCleanup(self.tmp.cleanup)
|
||||
self.root = Path(self.tmp.name)
|
||||
self.fixture = self.root / "sockets"
|
||||
self.fixture.write_text("")
|
||||
self.finder = self.root / "port-finder"
|
||||
self.finder.write_text('#!/bin/sh\ncat "$SOCKET_FIXTURE"\n')
|
||||
self.finder.chmod(0o755)
|
||||
|
||||
def run_shell(self, body: str, **env: str) -> subprocess.CompletedProcess[str]:
|
||||
"""Source the actual helper and run a shell scenario."""
|
||||
return subprocess.run(
|
||||
[SHELL, "-c", '. "$COMMON"\n' + body],
|
||||
env={
|
||||
**os.environ,
|
||||
"COMMON": COMMON,
|
||||
"SOCKET_FIXTURE": str(self.fixture),
|
||||
"PFCMD": str(self.finder),
|
||||
"PORT": "12345",
|
||||
**env,
|
||||
},
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=20,
|
||||
check=False,
|
||||
)
|
||||
|
||||
def test_socket_formats_and_false_matches(self) -> None:
|
||||
"""Accept IPv4 UDP local endpoints and reject unrelated sockets."""
|
||||
cases = [
|
||||
("udp UNCONN 0 0 0.0.0.0:12345 0.0.0.0:*", True),
|
||||
("udp UNCONN 0 0 127.0.0.1:12345 0.0.0.0:*", True),
|
||||
("udp 0 0 0.0.0.0:12345 0.0.0.0:*", True),
|
||||
("udp4 0 0 *.12345 *.*", True),
|
||||
("udp 0 0 127.0.0.1.12345 *.*", True),
|
||||
("udp 0 0 *.12345 *.*", True),
|
||||
("udp UNCONN 0 0 0.0.0.0:123456 0.0.0.0:*", False),
|
||||
("udp 0 0 0.0.0.0:123456 0.0.0.0:*", False),
|
||||
("udp ESTAB 0 0 127.0.0.1:54321 127.0.0.1:12345", False),
|
||||
("udp 0 0 127.0.0.1:54321 127.0.0.1:12345", False),
|
||||
("tcp LISTEN 0 128 0.0.0.0:12345 0.0.0.0:*", False),
|
||||
("tcp 0 0 0.0.0.0:12345 0.0.0.0:* LISTEN", False),
|
||||
("udp UNCONN 0 0 [::]:12345 [::]:*", False),
|
||||
("udp UNCONN 0 0 *:12345 *:*", False),
|
||||
("udp6 0 0 :::12345 :::*", False),
|
||||
("udp6 0 0 *.12345 *.*", False),
|
||||
("", False),
|
||||
]
|
||||
for row, ready in cases:
|
||||
with self.subTest(row=row):
|
||||
self.fixture.write_text(row + "\n")
|
||||
result = self.run_shell('check_if_udp_port_bound "$PORT"')
|
||||
self.assertEqual(result.returncode, 0 if ready else 1, result.stderr)
|
||||
|
||||
def test_exited_server_fails_immediately(self) -> None:
|
||||
"""Fail without sleeping when the server has already exited."""
|
||||
result = self.run_shell(
|
||||
'true &\npid=$!\nwait "$pid"\nsleep() { echo "unexpected sleep" >&2; }\nwait_udp_server "$pid"'
|
||||
)
|
||||
self.assertEqual(result.returncode, 1)
|
||||
self.assertIn("exited before binding", result.stderr)
|
||||
self.assertNotIn("unexpected sleep", result.stderr)
|
||||
|
||||
def test_ready_socket_does_not_sleep(self) -> None:
|
||||
"""Check readiness before the first sleep."""
|
||||
self.fixture.write_text("udp UNCONN 0 0 0.0.0.0:12345 0.0.0.0:*\n")
|
||||
result = self.run_shell('sleep() { echo "unexpected sleep" >&2; }\nwait_udp_server "$$"')
|
||||
self.assertEqual(result.returncode, 0, result.stderr)
|
||||
self.assertNotIn("unexpected sleep", result.stderr)
|
||||
|
||||
def test_missing_port_finders_skip(self) -> None:
|
||||
"""Skip and stop the live server when no finder is available."""
|
||||
# have_port_finder also tries these paths independently of PATH.
|
||||
if any(os.access(f"{directory}/ss", os.X_OK) for directory in ("/sbin", "/usr/sbin", "/usr/local/sbin")):
|
||||
self.skipTest("an absolute ss path cannot be hidden by this PATH-only fixture")
|
||||
with subprocess.Popen(["sleep", "60"]) as server:
|
||||
try:
|
||||
result = self.run_shell(
|
||||
'unset PFCMD\nPATH=/nonexistent\nwait_udp_server "$SERVER_PID"',
|
||||
SERVER_PID=str(server.pid),
|
||||
)
|
||||
self.assertEqual(result.returncode, 77)
|
||||
self.assertIn("neither ss nor netstat found", result.stderr)
|
||||
server.wait(timeout=3)
|
||||
self.assertLess(server.returncode, 0)
|
||||
finally:
|
||||
if server.poll() is None:
|
||||
server.kill()
|
||||
|
||||
def test_timeout_is_bounded_and_cleans_up(self) -> None:
|
||||
"""Stop polling after the retry budget and terminate the server."""
|
||||
# Only accelerate the polling delay; keep a real live server process.
|
||||
self.finder.write_text('#!/bin/sh\necho probe >&2\ncat "$SOCKET_FIXTURE"\n')
|
||||
with subprocess.Popen(["sleep", "60"]) as server:
|
||||
try:
|
||||
result = self.run_shell(
|
||||
'sleep() { echo polling-sleep; }\nwait_udp_server "$SERVER_PID"',
|
||||
SERVER_PID=str(server.pid),
|
||||
)
|
||||
self.assertEqual(result.returncode, 1)
|
||||
self.assertIn("did not come up", result.stderr)
|
||||
self.assertEqual(result.stderr.count("probe\n"), 90)
|
||||
self.assertEqual(result.stdout.count("polling-sleep"), 89)
|
||||
server.wait(timeout=3)
|
||||
self.assertLess(server.returncode, 0)
|
||||
finally:
|
||||
if server.poll() is None:
|
||||
server.kill()
|
||||
|
||||
def test_server_exits_while_waiting(self) -> None:
|
||||
"""Detect a startup failure that happens after polling begins."""
|
||||
result = self.run_shell('sleep 1 &\npid=$!\nwait_udp_server "$pid"')
|
||||
self.assertEqual(result.returncode, 1)
|
||||
self.assertIn("exited before binding", result.stderr)
|
||||
self.assertIn("waiting for UDP port", result.stdout)
|
||||
|
||||
def test_real_socket_delayed_beyond_four_seconds(self) -> None:
|
||||
"""Wait for a real delayed bind with each installed port finder."""
|
||||
finders = [shutil.which("ss"), os.environ.get("NETSTAT") or shutil.which("netstat")]
|
||||
finders = [finder for finder in finders if finder]
|
||||
if not finders:
|
||||
self.skipTest("neither ss nor netstat available")
|
||||
for finder in finders:
|
||||
with self.subTest(finder=finder):
|
||||
with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as sock:
|
||||
sock.bind(("127.0.0.1", 0))
|
||||
port = sock.getsockname()[1]
|
||||
code = (
|
||||
"import socket,time,sys; time.sleep(6); "
|
||||
"s=socket.socket(socket.AF_INET,socket.SOCK_DGRAM); "
|
||||
"s.bind(('127.0.0.1',int(sys.argv[1]))); time.sleep(30)"
|
||||
)
|
||||
with subprocess.Popen([sys.executable, "-c", code, str(port)]) as server:
|
||||
try:
|
||||
started = time.monotonic()
|
||||
result = self.run_shell(
|
||||
'wait_udp_server "$SERVER_PID"',
|
||||
SERVER_PID=str(server.pid),
|
||||
PORT=str(port),
|
||||
PFCMD=finder,
|
||||
)
|
||||
self.assertEqual(result.returncode, 0, result.stderr)
|
||||
self.assertGreaterEqual(time.monotonic() - started, 6)
|
||||
self.assertIsNone(server.poll())
|
||||
finally:
|
||||
server.terminate()
|
||||
server.wait(timeout=3)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,5 +0,0 @@
|
||||
{ scipy }:
|
||||
scipy.overridePythonAttrs (old: {
|
||||
# Keep the STFT tests enabled with tolerances for x86-64-v3 rounding.
|
||||
patches = (old.patches or [ ]) ++ [ ./stft-test-tolerances.patch ];
|
||||
})
|
||||
@@ -1,51 +0,0 @@
|
||||
Subject: [PATCH] signal: allow floating-point rounding in STFT tests
|
||||
|
||||
Keep the STFT tests enabled for x86-64-v3 builds. Allow two float32
|
||||
epsilons of relative error when comparing inverse-STFT implementations;
|
||||
float64 and the existing i686 override remain unchanged. Allow one
|
||||
float64 epsilon of absolute error in all three scaling round trips,
|
||||
which otherwise require exact zeros (observed residual: 4e-17 for a
|
||||
signal with amplitude 2).
|
||||
|
||||
Upstream issue: https://github.com/scipy/scipy/issues/25488
|
||||
|
||||
--- a/scipy/signal/tests/_scipy_spectral_test_shim.py
|
||||
+++ b/scipy/signal/tests/_scipy_spectral_test_shim.py
|
||||
@@ -294,7 +294,7 @@
|
||||
|
||||
# Adapted tolerances to account for resolution loss:
|
||||
atol = np.finfo(x.dtype).resolution*2 # instead of default atol = 0
|
||||
- rtol = 1e-7 # default for np.allclose()
|
||||
+ rtol = max(1e-7, 2 * np.finfo(x.dtype).eps)
|
||||
|
||||
# Relax atol on 32-Bit platforms a bit to pass CI tests.
|
||||
# - Not clear why there are discrepancies (in the FFT maybe?)
|
||||
--- a/scipy/signal/tests/test_spectral.py
|
||||
+++ b/scipy/signal/tests/test_spectral.py
|
||||
@@ -2044,7 +2044,7 @@
|
||||
|
||||
# Test round trip:
|
||||
x1 = istft(Zs, boundary=True, scaling='spectrum')[1]
|
||||
- assert_allclose(x1, x)
|
||||
+ assert_allclose(x1, x, atol=np.finfo(x.dtype).eps)
|
||||
|
||||
# For a Hann-windowed 256 sample length FFT, we expect a peak at
|
||||
# frequency 64 (since it is 1/4 the length of X) with a height of 1
|
||||
@@ -2074,7 +2074,7 @@
|
||||
|
||||
# Test round trip:
|
||||
x1 = istft(Zp, input_onesided=False, boundary=True, scaling='psd')[1]
|
||||
- assert_allclose(x1, x)
|
||||
+ assert_allclose(x1, x, atol=np.finfo(x.dtype).eps)
|
||||
|
||||
# The power of the one-sided psd-scaled STFT can be determined
|
||||
# analogously (note that the two sides are not of equal shape):
|
||||
@@ -2094,7 +2094,7 @@
|
||||
|
||||
# Test round trip:
|
||||
x1 = istft(Zp0, input_onesided=True, boundary=True, scaling='psd')[1]
|
||||
- assert_allclose(x1, x)
|
||||
+ assert_allclose(x1, x, atol=np.finfo(x.dtype).eps)
|
||||
|
||||
|
||||
class TestSampledSpectralRepresentations:
|
||||
@@ -1,115 +0,0 @@
|
||||
# Test exclusions for the locally rebuilt x86-64-v3 package set.
|
||||
#
|
||||
# Selecting x86-64-v3 changes every affected derivation, so the normal
|
||||
# nixpkgs binary cache cannot be used and upstream test suites run locally.
|
||||
# The jeeves builder uses /tmp/nix-builds so filesystem tests run on tmpfs
|
||||
# instead of ZFS with normalization=formD and utf8only=on; those tests remain
|
||||
# enabled. The remaining workarounds cover UDP readiness, resource-sensitive
|
||||
# parser and nested-worker races, and mismatched timeout clocks, plus
|
||||
# architecture-dependent floating-point differences whose risk we accept for
|
||||
# our workloads. Keep these exceptions visible until their causes are fixed.
|
||||
_final: prev: {
|
||||
prometheus = prev.prometheus.overrideAttrs (
|
||||
old:
|
||||
let
|
||||
assets = old.passthru.assets.overrideAttrs (assetsOld: {
|
||||
# CodeMirror's bounded synchronous parser can return an incomplete tree
|
||||
# when these cases run on a heavily loaded builder.
|
||||
postPatch = (assetsOld.postPatch or "") + ''
|
||||
substituteInPlace module/codemirror-promql/src/complete/hybrid.test.ts \
|
||||
--replace-fail "it(value.title, () => {" \
|
||||
"(value.title === 'autocomplete topk params 2' ? it.skip : it)(value.title, () => {"
|
||||
substituteInPlace module/codemirror-promql/src/parser/vector.test.ts \
|
||||
--replace-fail "it(value.binaryExpr, () => {" \
|
||||
"(value.binaryExpr === 'foo * on(test,blub) bar' ? it.skip : it)(value.binaryExpr, () => {"
|
||||
'';
|
||||
});
|
||||
in
|
||||
{
|
||||
postPatch = builtins.replaceStrings [ "${old.passthru.assets}" ] [ "${assets}" ] (
|
||||
builtins.unsafeDiscardStringContext old.postPatch
|
||||
);
|
||||
passthru = old.passthru // {
|
||||
inherit assets;
|
||||
};
|
||||
}
|
||||
);
|
||||
|
||||
pythonPackagesExtensions = prev.pythonPackagesExtensions ++ [
|
||||
(_pythonFinal: pythonPrev: {
|
||||
backrefs = pythonPrev.backrefs.overridePythonAttrs (old: {
|
||||
# regex measures its timeout in process CPU time, while this test used
|
||||
# wall time and could miss the timeout when a busy builder descheduled it.
|
||||
postPatch = (old.postPatch or "") + ''
|
||||
substituteInPlace tests/test_bregex.py \
|
||||
--replace-fail "time.time()" "time.process_time()"
|
||||
'';
|
||||
});
|
||||
|
||||
pytest-xdist = pythonPrev.pytest-xdist.overridePythonAttrs (old: {
|
||||
# The suite exercises its own worker pools. Run the outer suite with one
|
||||
# worker and allow inner workers more time on heavily loaded builders.
|
||||
postPatch = (old.postPatch or "") + ''
|
||||
substituteInPlace testing/test_remote.py \
|
||||
--replace-fail "WAIT_TIMEOUT = 10.0" "WAIT_TIMEOUT = 60.0"
|
||||
'';
|
||||
preCheck = builtins.replaceStrings [ "--numprocesses=$NIX_BUILD_CORES" ] [ "--numprocesses=1" ] (
|
||||
old.preCheck or ""
|
||||
);
|
||||
# This test deliberately crashes workers past the restart limit and
|
||||
# races while checking which replacement message was emitted.
|
||||
disabledTests = (old.disabledTests or [ ]) ++ [
|
||||
"test_max_worker_restart_tests_queued"
|
||||
];
|
||||
});
|
||||
|
||||
jupyter-server = pythonPrev.jupyter-server.overridePythonAttrs (old: {
|
||||
# The kernel reply arrived after the one-second outer deadline when the
|
||||
# builder was heavily loaded. Keep the regression test but allow it the
|
||||
# same margin as the other resource-sensitive tests.
|
||||
postPatch = (old.postPatch or "") + ''
|
||||
substituteInPlace tests/services/kernels/test_connection.py \
|
||||
--replace-fail \
|
||||
"await asyncio.wait_for(asyncio.wrap_future(conn2.request_kernel_info()), timeout=1.0)" \
|
||||
"await asyncio.wait_for(asyncio.wrap_future(conn2.request_kernel_info()), timeout=10.0)"
|
||||
'';
|
||||
});
|
||||
|
||||
sentry-sdk = pythonPrev.sentry-sdk.overridePythonAttrs (old: {
|
||||
# This test globally mocks threading.current_thread while another
|
||||
# thread is running. On Python 3.14, Thread.join can race with that
|
||||
# mock and exhaust its single side effect before the worker removes it.
|
||||
disabledTests = (old.disabledTests or [ ]) ++ [
|
||||
"test_get_current_thread_meta_main_thread"
|
||||
];
|
||||
});
|
||||
|
||||
torchaudio = pythonPrev.torchaudio.overridePythonAttrs (old: {
|
||||
# x86-64-v3 pitch shifting produces batch-versus-single-item numerical
|
||||
# differences up to 2.9e-6. We accept that audio-precision risk for our
|
||||
# workloads.
|
||||
disabledTests = (old.disabledTests or [ ]) ++ [
|
||||
"test_batch_pitch_shift"
|
||||
];
|
||||
});
|
||||
|
||||
torchcodec = pythonPrev.torchcodec.overridePythonAttrs (old: {
|
||||
# For these 8 kHz MP3 cases, the x86-64-v3 API and CLI codec paths
|
||||
# differ in 0.8% of decoded samples. Retain the original tolerance for
|
||||
# 99% of samples and accept the localized audio-precision risk.
|
||||
postPatch = (old.postPatch or "") + ''
|
||||
substituteInPlace test/test_encoders.py \
|
||||
--replace-fail \
|
||||
'if sys.platform == "darwin":' \
|
||||
'if sys.platform == "darwin" or (
|
||||
format == "mp3"
|
||||
and sample_rate == 8_000
|
||||
and asset is SINE_MONO_S32
|
||||
and bit_rate in (None, 0)
|
||||
and num_channels in (None, 1)
|
||||
):'
|
||||
'';
|
||||
});
|
||||
})
|
||||
];
|
||||
}
|
||||
@@ -1,20 +0,0 @@
|
||||
# Output-validation workarounds for packages rebuilt with x86-64-v3.
|
||||
_final: prev:
|
||||
let
|
||||
removeSiblingOutputChecks =
|
||||
package:
|
||||
package.overrideAttrs (old: {
|
||||
# Nix 2.34 can validate a partial multi-output rebuild against only the
|
||||
# outputs still being realised. PostgreSQL's checks then reject valid
|
||||
# sibling names such as "out" and "lib". Keep the test suite and
|
||||
# disallowed-requisite checks; accept the loss of cross-output checks.
|
||||
outputChecks = builtins.mapAttrs (
|
||||
_output: checks: builtins.removeAttrs checks [ "disallowedReferences" ]
|
||||
) (old.outputChecks or { });
|
||||
});
|
||||
|
||||
in
|
||||
prev.lib.optionalAttrs ((prev.stdenv.hostPlatform.gcc.arch or null) == "x86-64-v3") {
|
||||
postgresql = removeSiblingOutputChecks prev.postgresql;
|
||||
postgresql_18 = removeSiblingOutputChecks prev.postgresql_18;
|
||||
}
|
||||
+3
-2
@@ -11,6 +11,7 @@ license = "MIT"
|
||||
dependencies = [
|
||||
"alembic",
|
||||
"apprise",
|
||||
"apscheduler",
|
||||
"beautifulsoup4",
|
||||
"bm25s",
|
||||
"ebooklib",
|
||||
@@ -19,6 +20,7 @@ dependencies = [
|
||||
"httpx",
|
||||
"jinja2",
|
||||
"pgvector",
|
||||
"polars",
|
||||
"psycopg[binary]",
|
||||
"pydantic",
|
||||
"pydantic-settings",
|
||||
@@ -63,7 +65,6 @@ lint.ignore = [
|
||||
"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
|
||||
"CPY001", # (PERM) I don't include the license in every file
|
||||
]
|
||||
|
||||
[tool.ruff.lint.per-file-ignores]
|
||||
@@ -118,7 +119,7 @@ exclude_lines = [
|
||||
]
|
||||
|
||||
[tool.pytest.ini_options]
|
||||
addopts = "-n auto -ra --ignore=tests/ebook_search"
|
||||
addopts = "-n auto -ra"
|
||||
asyncio_mode = "auto"
|
||||
testpaths = ["tests"]
|
||||
# --cov=system_tools --cov-report=term-missing --cov-report=xml --cov-report=html --cov-branch
|
||||
|
||||
@@ -1,55 +0,0 @@
|
||||
"""remove spaCy-ner.
|
||||
|
||||
Revision ID: 751260fc3228
|
||||
Revises: dddee09eddcc
|
||||
Create Date: 2026-07-09 23:03:39.554083
|
||||
|
||||
"""
|
||||
|
||||
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 = "751260fc3228"
|
||||
down_revision: str | None = "dddee09eddcc"
|
||||
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_column("candidate_phrases", "source_spacy_noun_chunk", schema=schema)
|
||||
op.drop_column("candidate_phrases", "source_spacy_ner", schema=schema)
|
||||
op.drop_column("candidate_phrases", "spacy_label", schema=schema)
|
||||
# ### end Alembic commands ###
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
"""Downgrade."""
|
||||
# ### commands auto generated by Alembic - please adjust! ###
|
||||
op.add_column(
|
||||
"candidate_phrases", sa.Column("spacy_label", sa.VARCHAR(), autoincrement=False, nullable=True), schema=schema
|
||||
)
|
||||
op.add_column(
|
||||
"candidate_phrases",
|
||||
sa.Column("source_spacy_ner", sa.BOOLEAN(), autoincrement=False, nullable=False),
|
||||
schema=schema,
|
||||
)
|
||||
op.add_column(
|
||||
"candidate_phrases",
|
||||
sa.Column("source_spacy_noun_chunk", sa.BOOLEAN(), autoincrement=False, nullable=False),
|
||||
schema=schema,
|
||||
)
|
||||
# ### end Alembic commands ###
|
||||
@@ -0,0 +1 @@
|
||||
"""FastAPI applications."""
|
||||
@@ -0,0 +1,56 @@
|
||||
"""FastAPI interface for Contact database."""
|
||||
|
||||
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 python.api.routers import contact_router, views_router
|
||||
from python.common import configure_logger
|
||||
from python.fastapi_tools import ZstdMiddleware
|
||||
from python.orm.common import get_postgres_engine
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import AsyncIterator
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def create_app() -> FastAPI:
|
||||
"""Create and configure the FastAPI application."""
|
||||
|
||||
@asynccontextmanager
|
||||
async def lifespan(app: FastAPI) -> AsyncIterator[None]:
|
||||
"""Manage application lifespan."""
|
||||
app.state.engine = get_postgres_engine()
|
||||
yield
|
||||
app.state.engine.dispose()
|
||||
|
||||
app = FastAPI(title="Contact Database API", lifespan=lifespan)
|
||||
app.add_middleware(ZstdMiddleware)
|
||||
|
||||
app.include_router(contact_router)
|
||||
app.include_router(views_router)
|
||||
|
||||
return app
|
||||
|
||||
|
||||
def serve(
|
||||
host: Annotated[str, typer.Option("--host", "-h", help="Host to bind to")],
|
||||
port: Annotated[int, typer.Option("--port", "-p", help="Port to bind to")] = 8000,
|
||||
log_level: Annotated[str, typer.Option("--log-level", "-l", help="Log level")] = "INFO",
|
||||
) -> None:
|
||||
"""Start the Contact API server."""
|
||||
configure_logger(log_level)
|
||||
|
||||
app = create_app()
|
||||
uvicorn.run(app, host=host, port=port)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
typer.run(serve)
|
||||
@@ -0,0 +1,6 @@
|
||||
"""API routers."""
|
||||
|
||||
from python.api.routers.contact import router as contact_router
|
||||
from python.api.routers.views import router as views_router
|
||||
|
||||
__all__ = ["contact_router", "views_router"]
|
||||
@@ -0,0 +1,481 @@
|
||||
"""Contact API router."""
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
from fastapi import APIRouter, HTTPException, Request
|
||||
from fastapi.responses import HTMLResponse
|
||||
from fastapi.templating import Jinja2Templates
|
||||
from pydantic import BaseModel
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy.orm import selectinload
|
||||
|
||||
from python.fastapi_tools.db import DbSession # noqa: TC001 this is a FastAPI needed at runtime
|
||||
from python.orm.richie.contact import Contact, ContactRelationship, Need, RelationshipType
|
||||
|
||||
TEMPLATES_DIR = Path(__file__).parent.parent / "templates"
|
||||
templates = Jinja2Templates(directory=TEMPLATES_DIR)
|
||||
|
||||
|
||||
def _is_htmx(request: Request) -> bool:
|
||||
"""Check if the request is from HTMX."""
|
||||
return request.headers.get("HX-Request") == "true"
|
||||
|
||||
|
||||
class NeedBase(BaseModel):
|
||||
"""Base schema for Need."""
|
||||
|
||||
name: str
|
||||
description: str | None = None
|
||||
|
||||
|
||||
class NeedCreate(NeedBase):
|
||||
"""Schema for creating a Need."""
|
||||
|
||||
|
||||
class NeedResponse(NeedBase):
|
||||
"""Schema for Need response."""
|
||||
|
||||
id: int
|
||||
|
||||
model_config = {"from_attributes": True}
|
||||
|
||||
|
||||
class ContactRelationshipCreate(BaseModel):
|
||||
"""Schema for creating a contact relationship."""
|
||||
|
||||
related_contact_id: int
|
||||
relationship_type: RelationshipType
|
||||
closeness_weight: int | None = None
|
||||
|
||||
|
||||
class ContactRelationshipUpdate(BaseModel):
|
||||
"""Schema for updating a contact relationship."""
|
||||
|
||||
relationship_type: RelationshipType | None = None
|
||||
closeness_weight: int | None = None
|
||||
|
||||
|
||||
class ContactRelationshipResponse(BaseModel):
|
||||
"""Schema for contact relationship response."""
|
||||
|
||||
contact_id: int
|
||||
related_contact_id: int
|
||||
relationship_type: str
|
||||
closeness_weight: int
|
||||
|
||||
model_config = {"from_attributes": True}
|
||||
|
||||
|
||||
class RelationshipTypeInfo(BaseModel):
|
||||
"""Information about a relationship type."""
|
||||
|
||||
value: str
|
||||
display_name: str
|
||||
default_weight: int
|
||||
|
||||
|
||||
class GraphNode(BaseModel):
|
||||
"""Node in the relationship graph."""
|
||||
|
||||
id: int
|
||||
name: str
|
||||
current_job: str | None = None
|
||||
|
||||
|
||||
class GraphEdge(BaseModel):
|
||||
"""Edge in the relationship graph."""
|
||||
|
||||
source: int
|
||||
target: int
|
||||
relationship_type: str
|
||||
closeness_weight: int
|
||||
|
||||
|
||||
class GraphData(BaseModel):
|
||||
"""Complete graph data for visualization."""
|
||||
|
||||
nodes: list[GraphNode]
|
||||
edges: list[GraphEdge]
|
||||
|
||||
|
||||
class ContactBase(BaseModel):
|
||||
"""Base schema for Contact."""
|
||||
|
||||
name: str
|
||||
age: int | None = None
|
||||
bio: str | None = None
|
||||
current_job: str | None = None
|
||||
gender: str | None = None
|
||||
goals: str | None = None
|
||||
legal_name: str | None = None
|
||||
profile_pic: str | None = None
|
||||
safe_conversation_starters: str | None = None
|
||||
self_sufficiency_score: int | None = None
|
||||
social_structure_style: str | None = None
|
||||
ssn: str | None = None
|
||||
suffix: str | None = None
|
||||
timezone: str | None = None
|
||||
topics_to_avoid: str | None = None
|
||||
|
||||
|
||||
class ContactCreate(ContactBase):
|
||||
"""Schema for creating a Contact."""
|
||||
|
||||
need_ids: list[int] = []
|
||||
|
||||
|
||||
class ContactUpdate(BaseModel):
|
||||
"""Schema for updating a Contact."""
|
||||
|
||||
name: str | None = None
|
||||
age: int | None = None
|
||||
bio: str | None = None
|
||||
current_job: str | None = None
|
||||
gender: str | None = None
|
||||
goals: str | None = None
|
||||
legal_name: str | None = None
|
||||
profile_pic: str | None = None
|
||||
safe_conversation_starters: str | None = None
|
||||
self_sufficiency_score: int | None = None
|
||||
social_structure_style: str | None = None
|
||||
ssn: str | None = None
|
||||
suffix: str | None = None
|
||||
timezone: str | None = None
|
||||
topics_to_avoid: str | None = None
|
||||
need_ids: list[int] | None = None
|
||||
|
||||
|
||||
class ContactResponse(ContactBase):
|
||||
"""Schema for Contact response with relationships."""
|
||||
|
||||
id: int
|
||||
needs: list[NeedResponse] = []
|
||||
related_to: list[ContactRelationshipResponse] = []
|
||||
related_from: list[ContactRelationshipResponse] = []
|
||||
|
||||
model_config = {"from_attributes": True}
|
||||
|
||||
|
||||
class ContactListResponse(ContactBase):
|
||||
"""Schema for Contact list response."""
|
||||
|
||||
id: int
|
||||
|
||||
model_config = {"from_attributes": True}
|
||||
|
||||
|
||||
router = APIRouter(prefix="/api", tags=["contacts"])
|
||||
|
||||
|
||||
@router.post("/needs", response_model=NeedResponse)
|
||||
def create_need(need: NeedCreate, db: DbSession) -> Need:
|
||||
"""Create a new need."""
|
||||
db_need = Need(name=need.name, description=need.description)
|
||||
db.add(db_need)
|
||||
db.commit()
|
||||
db.refresh(db_need)
|
||||
return db_need
|
||||
|
||||
|
||||
@router.get("/needs", response_model=list[NeedResponse])
|
||||
def list_needs(db: DbSession) -> list[Need]:
|
||||
"""List all needs."""
|
||||
return list(db.scalars(select(Need)).all())
|
||||
|
||||
|
||||
@router.get("/needs/{need_id}", response_model=NeedResponse)
|
||||
def get_need(need_id: int, db: DbSession) -> Need:
|
||||
"""Get a need by ID."""
|
||||
need = db.get(Need, need_id)
|
||||
if not need:
|
||||
raise HTTPException(status_code=404, detail="Need not found")
|
||||
return need
|
||||
|
||||
|
||||
@router.delete("/needs/{need_id}", response_model=None)
|
||||
def delete_need(need_id: int, request: Request, db: DbSession) -> dict[str, bool] | HTMLResponse:
|
||||
"""Delete a need by ID."""
|
||||
need = db.get(Need, need_id)
|
||||
if not need:
|
||||
raise HTTPException(status_code=404, detail="Need not found")
|
||||
db.delete(need)
|
||||
db.commit()
|
||||
if _is_htmx(request):
|
||||
return HTMLResponse("")
|
||||
return {"deleted": True}
|
||||
|
||||
|
||||
@router.post("/contacts", response_model=ContactResponse)
|
||||
def create_contact(contact: ContactCreate, db: DbSession) -> Contact:
|
||||
"""Create a new contact."""
|
||||
need_ids = contact.need_ids
|
||||
contact_data = contact.model_dump(exclude={"need_ids"})
|
||||
db_contact = Contact(**contact_data)
|
||||
|
||||
if need_ids:
|
||||
needs = list(db.scalars(select(Need).where(Need.id.in_(need_ids))).all())
|
||||
db_contact.needs = needs
|
||||
|
||||
db.add(db_contact)
|
||||
db.commit()
|
||||
db.refresh(db_contact)
|
||||
return db_contact
|
||||
|
||||
|
||||
@router.get("/contacts", response_model=list[ContactListResponse])
|
||||
def list_contacts(
|
||||
db: DbSession,
|
||||
skip: int = 0,
|
||||
limit: int = 100,
|
||||
) -> list[Contact]:
|
||||
"""List all contacts with pagination."""
|
||||
return list(db.scalars(select(Contact).offset(skip).limit(limit)).all())
|
||||
|
||||
|
||||
@router.get("/contacts/{contact_id}", response_model=ContactResponse)
|
||||
def get_contact(contact_id: int, db: DbSession) -> Contact:
|
||||
"""Get a contact by ID with all relationships."""
|
||||
contact = db.scalar(
|
||||
select(Contact)
|
||||
.where(Contact.id == contact_id)
|
||||
.options(
|
||||
selectinload(Contact.needs),
|
||||
selectinload(Contact.related_to),
|
||||
selectinload(Contact.related_from),
|
||||
)
|
||||
)
|
||||
if not contact:
|
||||
raise HTTPException(status_code=404, detail="Contact not found")
|
||||
return contact
|
||||
|
||||
|
||||
@router.patch("/contacts/{contact_id}", response_model=ContactResponse)
|
||||
def update_contact(
|
||||
contact_id: int,
|
||||
contact: ContactUpdate,
|
||||
db: DbSession,
|
||||
) -> Contact:
|
||||
"""Update a contact by ID."""
|
||||
db_contact = db.get(Contact, contact_id)
|
||||
if not db_contact:
|
||||
raise HTTPException(status_code=404, detail="Contact not found")
|
||||
|
||||
update_data = contact.model_dump(exclude_unset=True)
|
||||
need_ids = update_data.pop("need_ids", None)
|
||||
|
||||
for key, value in update_data.items():
|
||||
setattr(db_contact, key, value)
|
||||
|
||||
if need_ids is not None:
|
||||
needs = list(db.scalars(select(Need).where(Need.id.in_(need_ids))).all())
|
||||
db_contact.needs = needs
|
||||
|
||||
db.commit()
|
||||
db.refresh(db_contact)
|
||||
return db_contact
|
||||
|
||||
|
||||
@router.delete("/contacts/{contact_id}", response_model=None)
|
||||
def delete_contact(contact_id: int, request: Request, db: DbSession) -> dict[str, bool] | HTMLResponse:
|
||||
"""Delete a contact by ID."""
|
||||
contact = db.get(Contact, contact_id)
|
||||
if not contact:
|
||||
raise HTTPException(status_code=404, detail="Contact not found")
|
||||
db.delete(contact)
|
||||
db.commit()
|
||||
if _is_htmx(request):
|
||||
return HTMLResponse("")
|
||||
return {"deleted": True}
|
||||
|
||||
|
||||
@router.post("/contacts/{contact_id}/needs/{need_id}")
|
||||
def add_need_to_contact(
|
||||
contact_id: int,
|
||||
need_id: int,
|
||||
db: DbSession,
|
||||
) -> dict[str, bool]:
|
||||
"""Add a need to a contact."""
|
||||
contact = db.get(Contact, contact_id)
|
||||
if not contact:
|
||||
raise HTTPException(status_code=404, detail="Contact not found")
|
||||
|
||||
need = db.get(Need, need_id)
|
||||
if not need:
|
||||
raise HTTPException(status_code=404, detail="Need not found")
|
||||
|
||||
if need not in contact.needs:
|
||||
contact.needs.append(need)
|
||||
db.commit()
|
||||
|
||||
return {"added": True}
|
||||
|
||||
|
||||
@router.delete("/contacts/{contact_id}/needs/{need_id}", response_model=None)
|
||||
def remove_need_from_contact(
|
||||
contact_id: int,
|
||||
need_id: int,
|
||||
request: Request,
|
||||
db: DbSession,
|
||||
) -> dict[str, bool] | HTMLResponse:
|
||||
"""Remove a need from a contact."""
|
||||
contact = db.get(Contact, contact_id)
|
||||
if not contact:
|
||||
raise HTTPException(status_code=404, detail="Contact not found")
|
||||
|
||||
need = db.get(Need, need_id)
|
||||
if not need:
|
||||
raise HTTPException(status_code=404, detail="Need not found")
|
||||
|
||||
if need in contact.needs:
|
||||
contact.needs.remove(need)
|
||||
db.commit()
|
||||
|
||||
if _is_htmx(request):
|
||||
return HTMLResponse("")
|
||||
return {"removed": True}
|
||||
|
||||
|
||||
@router.post(
|
||||
"/contacts/{contact_id}/relationships",
|
||||
response_model=ContactRelationshipResponse,
|
||||
)
|
||||
def add_contact_relationship(
|
||||
contact_id: int,
|
||||
relationship: ContactRelationshipCreate,
|
||||
db: DbSession,
|
||||
) -> ContactRelationship:
|
||||
"""Add a relationship between two contacts."""
|
||||
contact = db.get(Contact, contact_id)
|
||||
if not contact:
|
||||
raise HTTPException(status_code=404, detail="Contact not found")
|
||||
|
||||
related_contact = db.get(Contact, relationship.related_contact_id)
|
||||
if not related_contact:
|
||||
raise HTTPException(status_code=404, detail="Related contact not found")
|
||||
|
||||
if contact_id == relationship.related_contact_id:
|
||||
raise HTTPException(status_code=400, detail="Cannot relate contact to itself")
|
||||
|
||||
# Use provided weight or default from relationship type
|
||||
weight = relationship.closeness_weight
|
||||
if weight is None:
|
||||
weight = relationship.relationship_type.default_weight
|
||||
|
||||
db_relationship = ContactRelationship(
|
||||
contact_id=contact_id,
|
||||
related_contact_id=relationship.related_contact_id,
|
||||
relationship_type=relationship.relationship_type.value,
|
||||
closeness_weight=weight,
|
||||
)
|
||||
db.add(db_relationship)
|
||||
db.commit()
|
||||
db.refresh(db_relationship)
|
||||
return db_relationship
|
||||
|
||||
|
||||
@router.get(
|
||||
"/contacts/{contact_id}/relationships",
|
||||
response_model=list[ContactRelationshipResponse],
|
||||
)
|
||||
def get_contact_relationships(
|
||||
contact_id: int,
|
||||
db: DbSession,
|
||||
) -> list[ContactRelationship]:
|
||||
"""Get all relationships for a contact."""
|
||||
contact = db.get(Contact, contact_id)
|
||||
if not contact:
|
||||
raise HTTPException(status_code=404, detail="Contact not found")
|
||||
|
||||
outgoing = list(db.scalars(select(ContactRelationship).where(ContactRelationship.contact_id == contact_id)).all())
|
||||
incoming = list(
|
||||
db.scalars(select(ContactRelationship).where(ContactRelationship.related_contact_id == contact_id)).all()
|
||||
)
|
||||
return outgoing + incoming
|
||||
|
||||
|
||||
@router.patch(
|
||||
"/contacts/{contact_id}/relationships/{related_contact_id}",
|
||||
response_model=ContactRelationshipResponse,
|
||||
)
|
||||
def update_contact_relationship(
|
||||
contact_id: int,
|
||||
related_contact_id: int,
|
||||
update: ContactRelationshipUpdate,
|
||||
db: DbSession,
|
||||
) -> ContactRelationship:
|
||||
"""Update a relationship between two contacts."""
|
||||
relationship = db.scalar(
|
||||
select(ContactRelationship).where(
|
||||
ContactRelationship.contact_id == contact_id,
|
||||
ContactRelationship.related_contact_id == related_contact_id,
|
||||
)
|
||||
)
|
||||
if not relationship:
|
||||
raise HTTPException(status_code=404, detail="Relationship not found")
|
||||
|
||||
if update.relationship_type is not None:
|
||||
relationship.relationship_type = update.relationship_type.value
|
||||
if update.closeness_weight is not None:
|
||||
relationship.closeness_weight = update.closeness_weight
|
||||
|
||||
db.commit()
|
||||
db.refresh(relationship)
|
||||
return relationship
|
||||
|
||||
|
||||
@router.delete("/contacts/{contact_id}/relationships/{related_contact_id}", response_model=None)
|
||||
def remove_contact_relationship(
|
||||
contact_id: int,
|
||||
related_contact_id: int,
|
||||
request: Request,
|
||||
db: DbSession,
|
||||
) -> dict[str, bool] | HTMLResponse:
|
||||
"""Remove a relationship between two contacts."""
|
||||
relationship = db.scalar(
|
||||
select(ContactRelationship).where(
|
||||
ContactRelationship.contact_id == contact_id,
|
||||
ContactRelationship.related_contact_id == related_contact_id,
|
||||
)
|
||||
)
|
||||
if not relationship:
|
||||
raise HTTPException(status_code=404, detail="Relationship not found")
|
||||
|
||||
db.delete(relationship)
|
||||
db.commit()
|
||||
if _is_htmx(request):
|
||||
return HTMLResponse("")
|
||||
return {"deleted": True}
|
||||
|
||||
|
||||
@router.get("/relationship-types")
|
||||
def list_relationship_types() -> list[RelationshipTypeInfo]:
|
||||
"""List all available relationship types with their default weights."""
|
||||
return [
|
||||
RelationshipTypeInfo(
|
||||
value=rt.value,
|
||||
display_name=rt.display_name,
|
||||
default_weight=rt.default_weight,
|
||||
)
|
||||
for rt in RelationshipType
|
||||
]
|
||||
|
||||
|
||||
@router.get("/graph")
|
||||
def get_relationship_graph(db: DbSession) -> GraphData:
|
||||
"""Get all contacts and relationships as graph data for visualization."""
|
||||
contacts = list(db.scalars(select(Contact)).all())
|
||||
relationships = list(db.scalars(select(ContactRelationship)).all())
|
||||
|
||||
nodes = [GraphNode(id=c.id, name=c.name, current_job=c.current_job) for c in contacts]
|
||||
|
||||
edges = [
|
||||
GraphEdge(
|
||||
source=rel.contact_id,
|
||||
target=rel.related_contact_id,
|
||||
relationship_type=rel.relationship_type,
|
||||
closeness_weight=rel.closeness_weight,
|
||||
)
|
||||
for rel in relationships
|
||||
]
|
||||
|
||||
return GraphData(nodes=nodes, edges=edges)
|
||||
@@ -0,0 +1,345 @@
|
||||
"""HTMX server-rendered view router."""
|
||||
|
||||
from pathlib import Path
|
||||
from typing import Annotated, Any
|
||||
|
||||
from fastapi import APIRouter, Form, HTTPException, Request
|
||||
from fastapi.responses import HTMLResponse, RedirectResponse
|
||||
from fastapi.templating import Jinja2Templates
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy.orm import Session, selectinload
|
||||
|
||||
from python.fastapi_tools.db import DbSession # noqa: TC001 this is a FastAPI needed at runtime
|
||||
from python.orm.richie.contact import Contact, ContactRelationship, Need, RelationshipType
|
||||
|
||||
TEMPLATES_DIR = Path(__file__).parent.parent / "templates"
|
||||
templates = Jinja2Templates(directory=TEMPLATES_DIR)
|
||||
|
||||
router = APIRouter(tags=["views"])
|
||||
|
||||
FAMILIAL_TYPES = {
|
||||
"parent",
|
||||
"child",
|
||||
"sibling",
|
||||
"grandparent",
|
||||
"grandchild",
|
||||
"aunt_uncle",
|
||||
"niece_nephew",
|
||||
"cousin",
|
||||
"in_law",
|
||||
}
|
||||
FRIEND_TYPES = {"best_friend", "close_friend", "friend", "acquaintance", "neighbor"}
|
||||
PARTNER_TYPES = {"spouse", "partner"}
|
||||
PROFESSIONAL_TYPES = {"mentor", "mentee", "business_partner", "colleague", "manager", "direct_report", "client"}
|
||||
|
||||
CONTACT_STRING_FIELDS = (
|
||||
"name",
|
||||
"legal_name",
|
||||
"suffix",
|
||||
"gender",
|
||||
"current_job",
|
||||
"timezone",
|
||||
"profile_pic",
|
||||
"bio",
|
||||
"goals",
|
||||
"social_structure_style",
|
||||
"safe_conversation_starters",
|
||||
"topics_to_avoid",
|
||||
"ssn",
|
||||
)
|
||||
|
||||
CONTACT_INT_FIELDS = ("age", "self_sufficiency_score")
|
||||
|
||||
|
||||
def _group_relationships(relationships: list[ContactRelationship]) -> dict[str, list[ContactRelationship]]:
|
||||
"""Group relationships by category."""
|
||||
groups: dict[str, list[ContactRelationship]] = {
|
||||
"familial": [],
|
||||
"partners": [],
|
||||
"friends": [],
|
||||
"professional": [],
|
||||
"other": [],
|
||||
}
|
||||
for rel in relationships:
|
||||
if rel.relationship_type in FAMILIAL_TYPES:
|
||||
groups["familial"].append(rel)
|
||||
elif rel.relationship_type in PARTNER_TYPES:
|
||||
groups["partners"].append(rel)
|
||||
elif rel.relationship_type in FRIEND_TYPES:
|
||||
groups["friends"].append(rel)
|
||||
elif rel.relationship_type in PROFESSIONAL_TYPES:
|
||||
groups["professional"].append(rel)
|
||||
else:
|
||||
groups["other"].append(rel)
|
||||
return groups
|
||||
|
||||
|
||||
def _build_contact_name_map(database: Session, contact: Contact) -> dict[int, str]:
|
||||
"""Build a mapping of contact IDs to names for relationship display."""
|
||||
related_ids = {rel.related_contact_id for rel in contact.related_to}
|
||||
related_ids |= {rel.contact_id for rel in contact.related_from}
|
||||
related_ids.discard(contact.id)
|
||||
|
||||
if not related_ids:
|
||||
return {}
|
||||
|
||||
related_contacts = list(database.scalars(select(Contact).where(Contact.id.in_(related_ids))).all())
|
||||
return {related.id: related.name for related in related_contacts}
|
||||
|
||||
|
||||
def _get_relationship_type_display() -> dict[str, str]:
|
||||
"""Build a mapping of relationship type values to display names."""
|
||||
return {rel_type.value: rel_type.display_name for rel_type in RelationshipType}
|
||||
|
||||
|
||||
async def _parse_contact_form(request: Request) -> dict[str, Any]:
|
||||
"""Parse contact form data from a multipart/form request."""
|
||||
form_data = await request.form()
|
||||
result: dict[str, Any] = {}
|
||||
|
||||
for field in CONTACT_STRING_FIELDS:
|
||||
value = form_data.get(field, "")
|
||||
result[field] = str(value) if value else None
|
||||
|
||||
for field in CONTACT_INT_FIELDS:
|
||||
value = form_data.get(field, "")
|
||||
result[field] = int(value) if value else None
|
||||
|
||||
result["need_ids"] = [int(value) for value in form_data.getlist("need_ids")]
|
||||
return result
|
||||
|
||||
|
||||
def _save_contact_from_form(database: Session, contact: Contact, form_result: dict[str, Any]) -> None:
|
||||
"""Apply parsed form data to a Contact and save associated needs."""
|
||||
need_ids = form_result.pop("need_ids")
|
||||
|
||||
for key, value in form_result.items():
|
||||
setattr(contact, key, value)
|
||||
|
||||
if need_ids:
|
||||
contact.needs = list(database.scalars(select(Need).where(Need.id.in_(need_ids))).all())
|
||||
else:
|
||||
contact.needs = []
|
||||
|
||||
|
||||
@router.get("/", response_class=HTMLResponse)
|
||||
@router.get("/contacts", response_class=HTMLResponse)
|
||||
def contact_list_page(request: Request, database: DbSession) -> HTMLResponse:
|
||||
"""Render the contacts list page."""
|
||||
contacts = list(database.scalars(select(Contact)).all())
|
||||
return templates.TemplateResponse(request, "contact_list.html", {"contacts": contacts})
|
||||
|
||||
|
||||
@router.get("/contacts/new", response_class=HTMLResponse)
|
||||
def new_contact_page(request: Request, database: DbSession) -> HTMLResponse:
|
||||
"""Render the new contact form page."""
|
||||
all_needs = list(database.scalars(select(Need)).all())
|
||||
return templates.TemplateResponse(request, "contact_form.html", {"contact": None, "all_needs": all_needs})
|
||||
|
||||
|
||||
@router.post("/htmx/contacts/new")
|
||||
async def create_contact_form(request: Request, database: DbSession) -> RedirectResponse:
|
||||
"""Handle the create contact form submission."""
|
||||
form_result = await _parse_contact_form(request)
|
||||
contact = Contact()
|
||||
_save_contact_from_form(database, contact, form_result)
|
||||
|
||||
database.add(contact)
|
||||
database.commit()
|
||||
database.refresh(contact)
|
||||
return RedirectResponse(url=f"/contacts/{contact.id}", status_code=303)
|
||||
|
||||
|
||||
@router.get("/contacts/{contact_id}", response_class=HTMLResponse)
|
||||
def contact_detail_page(contact_id: int, request: Request, database: DbSession) -> HTMLResponse:
|
||||
"""Render the contact detail page."""
|
||||
contact = database.scalar(
|
||||
select(Contact)
|
||||
.where(Contact.id == contact_id)
|
||||
.options(
|
||||
selectinload(Contact.needs),
|
||||
selectinload(Contact.related_to),
|
||||
selectinload(Contact.related_from),
|
||||
)
|
||||
)
|
||||
if not contact:
|
||||
raise HTTPException(status_code=404, detail="Contact not found")
|
||||
|
||||
contact_names = _build_contact_name_map(database, contact)
|
||||
grouped_relationships = _group_relationships(contact.related_to)
|
||||
all_contacts = list(database.scalars(select(Contact)).all())
|
||||
all_needs = list(database.scalars(select(Need)).all())
|
||||
available_needs = [need for need in all_needs if need not in contact.needs]
|
||||
|
||||
return templates.TemplateResponse(
|
||||
request,
|
||||
"contact_detail.html",
|
||||
{
|
||||
"contact": contact,
|
||||
"contact_names": contact_names,
|
||||
"grouped_relationships": grouped_relationships,
|
||||
"all_contacts": all_contacts,
|
||||
"available_needs": available_needs,
|
||||
"relationship_types": list(RelationshipType),
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
@router.get("/contacts/{contact_id}/edit", response_class=HTMLResponse)
|
||||
def edit_contact_page(contact_id: int, request: Request, database: DbSession) -> HTMLResponse:
|
||||
"""Render the edit contact form page."""
|
||||
contact = database.scalar(select(Contact).where(Contact.id == contact_id).options(selectinload(Contact.needs)))
|
||||
if not contact:
|
||||
raise HTTPException(status_code=404, detail="Contact not found")
|
||||
|
||||
all_needs = list(database.scalars(select(Need)).all())
|
||||
return templates.TemplateResponse(request, "contact_form.html", {"contact": contact, "all_needs": all_needs})
|
||||
|
||||
|
||||
@router.post("/htmx/contacts/{contact_id}/edit")
|
||||
async def update_contact_form(contact_id: int, request: Request, database: DbSession) -> RedirectResponse:
|
||||
"""Handle the edit contact form submission."""
|
||||
contact = database.get(Contact, contact_id)
|
||||
if not contact:
|
||||
raise HTTPException(status_code=404, detail="Contact not found")
|
||||
|
||||
form_result = await _parse_contact_form(request)
|
||||
_save_contact_from_form(database, contact, form_result)
|
||||
|
||||
database.commit()
|
||||
return RedirectResponse(url=f"/contacts/{contact_id}", status_code=303)
|
||||
|
||||
|
||||
@router.post("/htmx/contacts/{contact_id}/add-need", response_class=HTMLResponse)
|
||||
def add_need_to_contact_htmx(
|
||||
contact_id: int,
|
||||
request: Request,
|
||||
database: DbSession,
|
||||
need_id: Annotated[int, Form()],
|
||||
) -> HTMLResponse:
|
||||
"""Add a need to a contact and return updated manage-needs partial."""
|
||||
contact = database.scalar(select(Contact).where(Contact.id == contact_id).options(selectinload(Contact.needs)))
|
||||
if not contact:
|
||||
raise HTTPException(status_code=404, detail="Contact not found")
|
||||
|
||||
need = database.get(Need, need_id)
|
||||
if not need:
|
||||
raise HTTPException(status_code=404, detail="Need not found")
|
||||
|
||||
if need not in contact.needs:
|
||||
contact.needs.append(need)
|
||||
database.commit()
|
||||
database.refresh(contact)
|
||||
|
||||
return templates.TemplateResponse(request, "partials/manage_needs.html", {"contact": contact})
|
||||
|
||||
|
||||
@router.post("/htmx/contacts/{contact_id}/add-relationship", response_class=HTMLResponse)
|
||||
def add_relationship_htmx(
|
||||
contact_id: int,
|
||||
request: Request,
|
||||
database: DbSession,
|
||||
related_contact_id: Annotated[int, Form()],
|
||||
relationship_type: Annotated[str, Form()],
|
||||
) -> HTMLResponse:
|
||||
"""Add a relationship and return updated manage-relationships partial."""
|
||||
contact = database.scalar(select(Contact).where(Contact.id == contact_id).options(selectinload(Contact.related_to)))
|
||||
if not contact:
|
||||
raise HTTPException(status_code=404, detail="Contact not found")
|
||||
|
||||
related_contact = database.get(Contact, related_contact_id)
|
||||
if not related_contact:
|
||||
raise HTTPException(status_code=404, detail="Related contact not found")
|
||||
|
||||
rel_type = RelationshipType(relationship_type)
|
||||
weight = rel_type.default_weight
|
||||
|
||||
relationship = ContactRelationship(
|
||||
contact_id=contact_id,
|
||||
related_contact_id=related_contact_id,
|
||||
relationship_type=relationship_type,
|
||||
closeness_weight=weight,
|
||||
)
|
||||
database.add(relationship)
|
||||
database.commit()
|
||||
database.refresh(contact)
|
||||
|
||||
contact_names = _build_contact_name_map(database, contact)
|
||||
return templates.TemplateResponse(
|
||||
request,
|
||||
"partials/manage_relationships.html",
|
||||
{"contact": contact, "contact_names": contact_names},
|
||||
)
|
||||
|
||||
|
||||
@router.post("/htmx/contacts/{contact_id}/relationships/{related_contact_id}/weight")
|
||||
def update_relationship_weight_htmx(
|
||||
contact_id: int,
|
||||
related_contact_id: int,
|
||||
database: DbSession,
|
||||
closeness_weight: Annotated[int, Form()],
|
||||
) -> HTMLResponse:
|
||||
"""Update a relationship's closeness weight from HTMX range input."""
|
||||
relationship = database.scalar(
|
||||
select(ContactRelationship).where(
|
||||
ContactRelationship.contact_id == contact_id,
|
||||
ContactRelationship.related_contact_id == related_contact_id,
|
||||
)
|
||||
)
|
||||
if not relationship:
|
||||
raise HTTPException(status_code=404, detail="Relationship not found")
|
||||
|
||||
relationship.closeness_weight = closeness_weight
|
||||
database.commit()
|
||||
return HTMLResponse("")
|
||||
|
||||
|
||||
@router.post("/htmx/needs", response_class=HTMLResponse)
|
||||
def create_need_htmx(
|
||||
request: Request,
|
||||
database: DbSession,
|
||||
name: Annotated[str, Form()],
|
||||
description: Annotated[str, Form()] = "",
|
||||
) -> HTMLResponse:
|
||||
"""Create a need via form data and return updated needs list."""
|
||||
need = Need(name=name, description=description or None)
|
||||
database.add(need)
|
||||
database.commit()
|
||||
needs = list(database.scalars(select(Need)).all())
|
||||
return templates.TemplateResponse(request, "partials/need_items.html", {"needs": needs})
|
||||
|
||||
|
||||
@router.get("/needs", response_class=HTMLResponse)
|
||||
def needs_page(request: Request, database: DbSession) -> HTMLResponse:
|
||||
"""Render the needs list page."""
|
||||
needs = list(database.scalars(select(Need)).all())
|
||||
return templates.TemplateResponse(request, "need_list.html", {"needs": needs})
|
||||
|
||||
|
||||
@router.get("/graph", response_class=HTMLResponse)
|
||||
def graph_page(request: Request, database: DbSession) -> HTMLResponse:
|
||||
"""Render the relationship graph page."""
|
||||
contacts = list(database.scalars(select(Contact)).all())
|
||||
relationships = list(database.scalars(select(ContactRelationship)).all())
|
||||
|
||||
graph_data = {
|
||||
"nodes": [{"id": contact.id, "name": contact.name, "current_job": contact.current_job} for contact in contacts],
|
||||
"edges": [
|
||||
{
|
||||
"source": rel.contact_id,
|
||||
"target": rel.related_contact_id,
|
||||
"relationship_type": rel.relationship_type,
|
||||
"closeness_weight": rel.closeness_weight,
|
||||
}
|
||||
for rel in relationships
|
||||
],
|
||||
}
|
||||
|
||||
return templates.TemplateResponse(
|
||||
request,
|
||||
"graph.html",
|
||||
{
|
||||
"graph_data": graph_data,
|
||||
"relationship_type_display": _get_relationship_type_display(),
|
||||
},
|
||||
)
|
||||
@@ -0,0 +1,198 @@
|
||||
<!DOCTYPE html>
|
||||
<html lang="en" data-theme="light">
|
||||
<head>
|
||||
<meta charset="UTF-8">
|
||||
<meta name="viewport" content="width=device-width, initial-scale=1.0">
|
||||
<title>{% block title %}Contact Database{% endblock %}</title>
|
||||
<script src="https://unpkg.com/htmx.org@2.0.4"></script>
|
||||
<style>
|
||||
:root {
|
||||
--color-bg: #f5f5f5;
|
||||
--color-bg-card: #ffffff;
|
||||
--color-bg-hover: #f0f0f0;
|
||||
--color-bg-muted: #f9f9f9;
|
||||
--color-bg-error: #ffe0e0;
|
||||
--color-text: #333333;
|
||||
--color-text-muted: #666666;
|
||||
--color-text-error: #cc0000;
|
||||
--color-border: #dddddd;
|
||||
--color-border-light: #eeeeee;
|
||||
--color-border-lighter: #f0f0f0;
|
||||
--color-primary: #0066cc;
|
||||
--color-primary-hover: #0055aa;
|
||||
--color-danger: #cc3333;
|
||||
--color-danger-hover: #aa2222;
|
||||
--color-tag-bg: #e0e0e0;
|
||||
--shadow: 0 1px 3px rgba(0, 0, 0, 0.1);
|
||||
font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, sans-serif;
|
||||
line-height: 1.5;
|
||||
color: var(--color-text);
|
||||
background-color: var(--color-bg);
|
||||
}
|
||||
[data-theme="dark"] {
|
||||
--color-bg: #1a1a1a;
|
||||
--color-bg-card: #2d2d2d;
|
||||
--color-bg-hover: #3d3d3d;
|
||||
--color-bg-muted: #252525;
|
||||
--color-bg-error: #4a2020;
|
||||
--color-text: #e0e0e0;
|
||||
--color-text-muted: #a0a0a0;
|
||||
--color-text-error: #ff6b6b;
|
||||
--color-border: #404040;
|
||||
--color-border-light: #353535;
|
||||
--color-border-lighter: #303030;
|
||||
--color-primary: #4da6ff;
|
||||
--color-primary-hover: #7dbfff;
|
||||
--color-danger: #ff6b6b;
|
||||
--color-danger-hover: #ff8a8a;
|
||||
--color-tag-bg: #404040;
|
||||
--shadow: 0 1px 3px rgba(0, 0, 0, 0.3);
|
||||
}
|
||||
* { box-sizing: border-box; }
|
||||
body { margin: 0; background: var(--color-bg); color: var(--color-text); }
|
||||
.app { max-width: 1000px; margin: 0 auto; padding: 20px; }
|
||||
nav { display: flex; align-items: center; gap: 20px; padding: 15px 0; border-bottom: 1px solid var(--color-border); margin-bottom: 20px; }
|
||||
nav a { color: var(--color-primary); text-decoration: none; font-weight: 500; }
|
||||
nav a:hover { text-decoration: underline; }
|
||||
.theme-toggle { margin-left: auto; }
|
||||
main { background: var(--color-bg-card); padding: 20px; border-radius: 8px; box-shadow: var(--shadow); }
|
||||
.header { display: flex; justify-content: space-between; align-items: center; margin-bottom: 20px; }
|
||||
.header h1 { margin: 0; }
|
||||
a { color: var(--color-primary); }
|
||||
a:hover { text-decoration: underline; }
|
||||
|
||||
.btn { display: inline-block; padding: 8px 16px; border: 1px solid var(--color-border); border-radius: 4px; background: var(--color-bg-card); color: var(--color-text); text-decoration: none; cursor: pointer; font-size: 14px; margin-left: 8px; }
|
||||
.btn:hover { background: var(--color-bg-hover); }
|
||||
.btn-primary { background: var(--color-primary); border-color: var(--color-primary); color: white; }
|
||||
.btn-primary:hover { background: var(--color-primary-hover); }
|
||||
.btn-danger { background: var(--color-danger); border-color: var(--color-danger); color: white; }
|
||||
.btn-danger:hover { background: var(--color-danger-hover); }
|
||||
.btn-small { padding: 4px 8px; font-size: 12px; }
|
||||
.btn:disabled { opacity: 0.6; cursor: not-allowed; }
|
||||
|
||||
table { width: 100%; border-collapse: collapse; }
|
||||
th, td { padding: 12px; text-align: left; border-bottom: 1px solid var(--color-border-light); }
|
||||
th { font-weight: 600; background: var(--color-bg-muted); }
|
||||
tr:hover { background: var(--color-bg-muted); }
|
||||
|
||||
.error { background: var(--color-bg-error); color: var(--color-text-error); padding: 10px; border-radius: 4px; margin-bottom: 20px; }
|
||||
.tag { display: inline-block; background: var(--color-tag-bg); padding: 2px 8px; border-radius: 12px; font-size: 12px; color: var(--color-text-muted); }
|
||||
|
||||
.add-form { display: flex; gap: 10px; margin-top: 15px; flex-wrap: wrap; }
|
||||
.add-form select, .add-form input { padding: 8px; border: 1px solid var(--color-border); border-radius: 4px; min-width: 200px; background: var(--color-bg-card); color: var(--color-text); }
|
||||
|
||||
.form-group { margin-bottom: 20px; }
|
||||
.form-group label { display: block; font-weight: 500; margin-bottom: 5px; }
|
||||
.form-group input, .form-group textarea, .form-group select { width: 100%; padding: 10px; border: 1px solid var(--color-border); border-radius: 4px; font-size: 14px; background: var(--color-bg-card); color: var(--color-text); }
|
||||
.form-group textarea { resize: vertical; }
|
||||
.form-row { display: grid; grid-template-columns: 1fr 1fr; gap: 20px; }
|
||||
.checkbox-group { display: flex; flex-wrap: wrap; gap: 15px; }
|
||||
.checkbox-label { display: flex; align-items: center; gap: 5px; cursor: pointer; }
|
||||
.form-actions { display: flex; gap: 10px; margin-top: 30px; padding-top: 20px; border-top: 1px solid var(--color-border-light); }
|
||||
|
||||
.need-form { background: var(--color-bg-muted); padding: 20px; border-radius: 4px; margin-bottom: 20px; }
|
||||
.need-items { list-style: none; padding: 0; }
|
||||
.need-items li { display: flex; justify-content: space-between; align-items: flex-start; padding: 15px; border: 1px solid var(--color-border-light); border-radius: 4px; margin-bottom: 10px; }
|
||||
.need-info p { margin: 5px 0 0; color: var(--color-text-muted); font-size: 14px; }
|
||||
|
||||
.graph-container { width: 100%; }
|
||||
.graph-hint { color: var(--color-text-muted); font-size: 14px; margin-bottom: 15px; }
|
||||
.selected-info { margin-top: 15px; padding: 15px; background: var(--color-bg-muted); border-radius: 8px; }
|
||||
.selected-info h3 { margin: 0 0 10px; }
|
||||
.selected-info p { margin: 5px 0; color: var(--color-text-muted); }
|
||||
.legend { margin-top: 20px; padding: 15px; background: var(--color-bg-muted); border-radius: 8px; }
|
||||
.legend h4 { margin: 0 0 10px; font-size: 14px; }
|
||||
.legend-items { display: flex; flex-wrap: wrap; gap: 15px; }
|
||||
.legend-item { display: flex; align-items: center; gap: 8px; font-size: 12px; color: var(--color-text-muted); }
|
||||
.legend-line { width: 30px; border-radius: 2px; }
|
||||
|
||||
.id-card { width: 100%; }
|
||||
.id-card-inner { background: linear-gradient(135deg, #0a0a0f 0%, #1a1a2e 50%, #0a0a0f 100%); background-image: radial-gradient(white 1px, transparent 1px), linear-gradient(135deg, #0a0a0f 0%, #1a1a2e 50%, #0a0a0f 100%); background-size: 50px 50px, 100% 100%; color: #fff; border-radius: 12px; padding: 25px; min-height: 500px; position: relative; overflow: hidden; }
|
||||
.id-card-header { display: flex; justify-content: space-between; align-items: flex-start; margin-bottom: 15px; }
|
||||
.id-card-header-left { flex: 1; }
|
||||
.id-card-header-right { display: flex; flex-direction: column; align-items: flex-end; gap: 10px; }
|
||||
.id-card-title { font-size: 2.5rem; font-weight: 700; margin: 0; color: #fff; text-shadow: 2px 2px 4px rgba(0,0,0,0.5); }
|
||||
.id-profile-pic { width: 80px; height: 80px; border-radius: 8px; object-fit: cover; border: 2px solid rgba(255,255,255,0.3); }
|
||||
.id-profile-placeholder { width: 80px; height: 80px; border-radius: 8px; background: linear-gradient(135deg, #4ecdc4 0%, #44a8a0 100%); display: flex; align-items: center; justify-content: center; border: 2px solid rgba(255,255,255,0.3); }
|
||||
.id-profile-placeholder span { font-size: 2rem; font-weight: 700; color: #fff; text-shadow: 1px 1px 2px rgba(0,0,0,0.3); }
|
||||
.id-card-actions { display: flex; gap: 8px; }
|
||||
.id-card-actions .btn { background: rgba(255,255,255,0.1); border-color: rgba(255,255,255,0.3); color: #fff; }
|
||||
.id-card-actions .btn:hover { background: rgba(255,255,255,0.2); }
|
||||
.id-card-body { display: grid; grid-template-columns: 1fr 1.5fr; gap: 30px; }
|
||||
.id-card-left { display: flex; flex-direction: column; gap: 8px; }
|
||||
.id-field { font-size: 1rem; line-height: 1.4; }
|
||||
.id-field-block { margin-top: 15px; font-size: 0.95rem; line-height: 1.5; }
|
||||
.id-label { color: #4ecdc4; font-weight: 500; }
|
||||
.id-card-right { display: flex; flex-direction: column; gap: 20px; }
|
||||
.id-bio { font-size: 0.9rem; line-height: 1.6; color: #e0e0e0; }
|
||||
.id-relationships { margin-top: 10px; }
|
||||
.id-section-title { font-size: 1.5rem; margin: 0 0 15px; color: #fff; border-bottom: 1px solid rgba(255,255,255,0.2); padding-bottom: 8px; }
|
||||
.id-rel-group { margin-bottom: 12px; font-size: 0.9rem; line-height: 1.6; }
|
||||
.id-rel-label { color: #a0a0a0; }
|
||||
.id-rel-group a { color: #4ecdc4; text-decoration: none; }
|
||||
.id-rel-group a:hover { text-decoration: underline; }
|
||||
.id-rel-type { color: #888; font-size: 0.85em; }
|
||||
.id-card-warnings { margin-top: 30px; padding-top: 20px; border-top: 1px solid rgba(255,255,255,0.2); display: flex; flex-wrap: wrap; gap: 20px; }
|
||||
.id-warning { display: flex; align-items: center; gap: 8px; font-size: 0.9rem; color: #ff6b6b; }
|
||||
.warning-dot { width: 8px; height: 8px; background: #ff6b6b; border-radius: 50%; flex-shrink: 0; }
|
||||
.warning-desc { color: #ccc; }
|
||||
|
||||
.id-card-manage { margin-top: 20px; background: var(--color-bg-muted); border-radius: 8px; padding: 15px; }
|
||||
.id-card-manage summary { cursor: pointer; font-weight: 600; font-size: 1.1rem; padding: 5px 0; }
|
||||
.id-card-manage[open] summary { margin-bottom: 15px; border-bottom: 1px solid var(--color-border-light); padding-bottom: 10px; }
|
||||
.manage-section { margin-bottom: 25px; }
|
||||
.manage-section h3 { margin: 0 0 15px; font-size: 1rem; }
|
||||
.manage-relationships { display: flex; flex-direction: column; gap: 10px; margin-bottom: 15px; }
|
||||
.manage-rel-item { display: flex; align-items: center; gap: 12px; padding: 10px; background: var(--color-bg-card); border-radius: 6px; flex-wrap: wrap; }
|
||||
.manage-rel-item a { font-weight: 500; min-width: 120px; }
|
||||
.weight-control { display: flex; align-items: center; gap: 8px; font-size: 12px; color: var(--color-text-muted); }
|
||||
.weight-control input[type="range"] { width: 80px; cursor: pointer; }
|
||||
.weight-value { min-width: 20px; text-align: center; font-weight: 600; }
|
||||
.manage-needs-list { list-style: none; padding: 0; margin: 0 0 15px; }
|
||||
.manage-needs-list li { display: flex; align-items: center; gap: 12px; padding: 10px; background: var(--color-bg-card); border-radius: 6px; margin-bottom: 8px; }
|
||||
.manage-needs-list li .btn { margin-left: auto; }
|
||||
|
||||
.htmx-indicator { display: none; }
|
||||
.htmx-request .htmx-indicator { display: inline; }
|
||||
.htmx-request.htmx-indicator { display: inline; }
|
||||
|
||||
@media (max-width: 768px) {
|
||||
.id-card-body { grid-template-columns: 1fr; }
|
||||
.id-card-title { font-size: 1.8rem; }
|
||||
.id-card-header { flex-direction: column; gap: 15px; }
|
||||
}
|
||||
</style>
|
||||
</head>
|
||||
<body>
|
||||
<div class="app">
|
||||
<nav>
|
||||
<a href="/contacts">Contacts</a>
|
||||
<a href="/graph">Graph</a>
|
||||
<a href="/needs">Needs</a>
|
||||
<button class="btn btn-small theme-toggle" onclick="toggleTheme()">
|
||||
<span id="theme-label">Dark</span>
|
||||
</button>
|
||||
</nav>
|
||||
|
||||
<main id="main-content">
|
||||
{% block content %}{% endblock %}
|
||||
</main>
|
||||
</div>
|
||||
|
||||
<script>
|
||||
function toggleTheme() {
|
||||
const html = document.documentElement;
|
||||
const current = html.getAttribute('data-theme');
|
||||
const next = current === 'light' ? 'dark' : 'light';
|
||||
html.setAttribute('data-theme', next);
|
||||
localStorage.setItem('theme', next);
|
||||
document.getElementById('theme-label').textContent = next === 'light' ? 'Dark' : 'Light';
|
||||
}
|
||||
(function() {
|
||||
const saved = localStorage.getItem('theme') || 'light';
|
||||
document.documentElement.setAttribute('data-theme', saved);
|
||||
document.getElementById('theme-label').textContent = saved === 'light' ? 'Dark' : 'Light';
|
||||
})();
|
||||
</script>
|
||||
</body>
|
||||
</html>
|
||||
@@ -0,0 +1,204 @@
|
||||
{% extends "base.html" %}
|
||||
{% block title %}{{ contact.name }}{% endblock %}
|
||||
{% block content %}
|
||||
<div class="id-card">
|
||||
<div class="id-card-inner">
|
||||
<div class="id-card-header">
|
||||
<div class="id-card-header-left">
|
||||
<h1 class="id-card-title">I.D.: {{ contact.name }}</h1>
|
||||
</div>
|
||||
<div class="id-card-header-right">
|
||||
{% if contact.profile_pic %}
|
||||
<img src="{{ contact.profile_pic }}" alt="{{ contact.name }}'s profile" class="id-profile-pic">
|
||||
{% else %}
|
||||
<div class="id-profile-placeholder">
|
||||
<span>{{ contact.name[0]|upper }}</span>
|
||||
</div>
|
||||
{% endif %}
|
||||
<div class="id-card-actions">
|
||||
<a href="/contacts/{{ contact.id }}/edit" class="btn btn-small">Edit</a>
|
||||
<a href="/contacts" class="btn btn-small">Back</a>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<div class="id-card-body">
|
||||
<div class="id-card-left">
|
||||
{% if contact.legal_name %}
|
||||
<div class="id-field">Legal name: {{ contact.legal_name }}</div>
|
||||
{% endif %}
|
||||
{% if contact.suffix %}
|
||||
<div class="id-field">Suffix: {{ contact.suffix }}</div>
|
||||
{% endif %}
|
||||
{% if contact.gender %}
|
||||
<div class="id-field">Gender: {{ contact.gender }}</div>
|
||||
{% endif %}
|
||||
{% if contact.age %}
|
||||
<div class="id-field">Age: {{ contact.age }}</div>
|
||||
{% endif %}
|
||||
{% if contact.current_job %}
|
||||
<div class="id-field">Job: {{ contact.current_job }}</div>
|
||||
{% endif %}
|
||||
{% if contact.social_structure_style %}
|
||||
<div class="id-field">Social style: {{ contact.social_structure_style }}</div>
|
||||
{% endif %}
|
||||
{% if contact.self_sufficiency_score is not none %}
|
||||
<div class="id-field">Self-Sufficiency: {{ contact.self_sufficiency_score }}</div>
|
||||
{% endif %}
|
||||
{% if contact.timezone %}
|
||||
<div class="id-field">Timezone: {{ contact.timezone }}</div>
|
||||
{% endif %}
|
||||
{% if contact.safe_conversation_starters %}
|
||||
<div class="id-field-block">
|
||||
<span class="id-label">Safe con starters:</span> {{ contact.safe_conversation_starters }}
|
||||
</div>
|
||||
{% endif %}
|
||||
{% if contact.topics_to_avoid %}
|
||||
<div class="id-field-block">
|
||||
<span class="id-label">Topics to avoid:</span> {{ contact.topics_to_avoid }}
|
||||
</div>
|
||||
{% endif %}
|
||||
{% if contact.goals %}
|
||||
<div class="id-field-block">
|
||||
<span class="id-label">Goals:</span> {{ contact.goals }}
|
||||
</div>
|
||||
{% endif %}
|
||||
</div>
|
||||
|
||||
<div class="id-card-right">
|
||||
{% if contact.bio %}
|
||||
<div class="id-bio">
|
||||
<span class="id-label">Bio:</span> {{ contact.bio }}
|
||||
</div>
|
||||
{% endif %}
|
||||
|
||||
<div class="id-relationships">
|
||||
<h2 class="id-section-title">Relationships</h2>
|
||||
|
||||
{% if grouped_relationships.familial %}
|
||||
<div class="id-rel-group">
|
||||
<span class="id-rel-label">Familial:</span>
|
||||
{% for rel in grouped_relationships.familial %}
|
||||
<a href="/contacts/{{ rel.related_contact_id }}">{{ contact_names[rel.related_contact_id] }}</a><span class="id-rel-type">({{ rel.relationship_type|replace("_", " ")|title }})</span>{% if not loop.last %}, {% endif %}
|
||||
{% endfor %}
|
||||
</div>
|
||||
{% endif %}
|
||||
|
||||
{% if grouped_relationships.partners %}
|
||||
<div class="id-rel-group">
|
||||
<span class="id-rel-label">Partners:</span>
|
||||
{% for rel in grouped_relationships.partners %}
|
||||
<a href="/contacts/{{ rel.related_contact_id }}">{{ contact_names[rel.related_contact_id] }}</a>{% if not loop.last %}, {% endif %}
|
||||
{% endfor %}
|
||||
</div>
|
||||
{% endif %}
|
||||
|
||||
{% if grouped_relationships.friends %}
|
||||
<div class="id-rel-group">
|
||||
<span class="id-rel-label">Friends:</span>
|
||||
{% for rel in grouped_relationships.friends %}
|
||||
<a href="/contacts/{{ rel.related_contact_id }}">{{ contact_names[rel.related_contact_id] }}</a>{% if not loop.last %}, {% endif %}
|
||||
{% endfor %}
|
||||
</div>
|
||||
{% endif %}
|
||||
|
||||
{% if grouped_relationships.professional %}
|
||||
<div class="id-rel-group">
|
||||
<span class="id-rel-label">Professional:</span>
|
||||
{% for rel in grouped_relationships.professional %}
|
||||
<a href="/contacts/{{ rel.related_contact_id }}">{{ contact_names[rel.related_contact_id] }}</a><span class="id-rel-type">({{ rel.relationship_type|replace("_", " ")|title }})</span>{% if not loop.last %}, {% endif %}
|
||||
{% endfor %}
|
||||
</div>
|
||||
{% endif %}
|
||||
|
||||
{% if grouped_relationships.other %}
|
||||
<div class="id-rel-group">
|
||||
<span class="id-rel-label">Other:</span>
|
||||
{% for rel in grouped_relationships.other %}
|
||||
<a href="/contacts/{{ rel.related_contact_id }}">{{ contact_names[rel.related_contact_id] }}</a><span class="id-rel-type">({{ rel.relationship_type|replace("_", " ")|title }})</span>{% if not loop.last %}, {% endif %}
|
||||
{% endfor %}
|
||||
</div>
|
||||
{% endif %}
|
||||
|
||||
{% if contact.related_from %}
|
||||
<div class="id-rel-group">
|
||||
<span class="id-rel-label">Known by:</span>
|
||||
{% for rel in contact.related_from %}
|
||||
<a href="/contacts/{{ rel.contact_id }}">{{ contact_names[rel.contact_id] }}</a>{% if not loop.last %}, {% endif %}
|
||||
{% endfor %}
|
||||
</div>
|
||||
{% endif %}
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
{% if contact.needs %}
|
||||
<div class="id-card-warnings">
|
||||
{% for need in contact.needs %}
|
||||
<div class="id-warning">
|
||||
<span class="warning-dot"></span>
|
||||
Warning: {{ need.name }}
|
||||
{% if need.description %}<span class="warning-desc"> - {{ need.description }}</span>{% endif %}
|
||||
</div>
|
||||
{% endfor %}
|
||||
</div>
|
||||
{% endif %}
|
||||
</div>
|
||||
|
||||
<details class="id-card-manage">
|
||||
<summary>Manage Contact</summary>
|
||||
|
||||
<div class="manage-section">
|
||||
<h3>Manage Relationships</h3>
|
||||
<div id="manage-relationships" class="manage-relationships">
|
||||
{% include "partials/manage_relationships.html" %}
|
||||
</div>
|
||||
|
||||
{% if all_contacts %}
|
||||
<form hx-post="/htmx/contacts/{{ contact.id }}/add-relationship"
|
||||
hx-target="#manage-relationships"
|
||||
hx-swap="innerHTML"
|
||||
class="add-form">
|
||||
<select name="related_contact_id" required>
|
||||
<option value="">Select contact...</option>
|
||||
{% for other in all_contacts %}
|
||||
{% if other.id != contact.id %}
|
||||
<option value="{{ other.id }}">{{ other.name }}</option>
|
||||
{% endif %}
|
||||
{% endfor %}
|
||||
</select>
|
||||
<select name="relationship_type" required>
|
||||
<option value="">Select relationship type...</option>
|
||||
{% for rel_type in relationship_types %}
|
||||
<option value="{{ rel_type.value }}">{{ rel_type.display_name }}</option>
|
||||
{% endfor %}
|
||||
</select>
|
||||
<button type="submit" class="btn btn-primary">Add Relationship</button>
|
||||
</form>
|
||||
{% endif %}
|
||||
</div>
|
||||
|
||||
<div class="manage-section">
|
||||
<h3>Manage Needs/Warnings</h3>
|
||||
<div id="manage-needs">
|
||||
{% include "partials/manage_needs.html" %}
|
||||
</div>
|
||||
|
||||
{% if available_needs %}
|
||||
<form hx-post="/htmx/contacts/{{ contact.id }}/add-need"
|
||||
hx-target="#manage-needs"
|
||||
hx-swap="innerHTML"
|
||||
class="add-form">
|
||||
<select name="need_id" required>
|
||||
<option value="">Select a need...</option>
|
||||
{% for need in available_needs %}
|
||||
<option value="{{ need.id }}">{{ need.name }}</option>
|
||||
{% endfor %}
|
||||
</select>
|
||||
<button type="submit" class="btn btn-primary">Add Need</button>
|
||||
</form>
|
||||
{% endif %}
|
||||
</div>
|
||||
</details>
|
||||
</div>
|
||||
{% endblock %}
|
||||
@@ -0,0 +1,115 @@
|
||||
{% extends "base.html" %}
|
||||
{% block title %}{{ "Edit " + contact.name if contact else "New Contact" }}{% endblock %}
|
||||
{% block content %}
|
||||
<div class="contact-form">
|
||||
<h1>{{ "Edit Contact" if contact else "New Contact" }}</h1>
|
||||
|
||||
{% if contact %}
|
||||
<form method="post" action="/htmx/contacts/{{ contact.id }}/edit">
|
||||
{% else %}
|
||||
<form method="post" action="/htmx/contacts/new">
|
||||
{% endif %}
|
||||
|
||||
<div class="form-group">
|
||||
<label for="name">Name *</label>
|
||||
<input id="name" name="name" type="text" value="{{ contact.name if contact else '' }}" required>
|
||||
</div>
|
||||
|
||||
<div class="form-row">
|
||||
<div class="form-group">
|
||||
<label for="legal_name">Legal Name</label>
|
||||
<input id="legal_name" name="legal_name" type="text" value="{{ contact.legal_name or '' }}">
|
||||
</div>
|
||||
<div class="form-group">
|
||||
<label for="suffix">Suffix</label>
|
||||
<input id="suffix" name="suffix" type="text" value="{{ contact.suffix or '' }}">
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<div class="form-row">
|
||||
<div class="form-group">
|
||||
<label for="age">Age</label>
|
||||
<input id="age" name="age" type="number" value="{{ contact.age if contact and contact.age is not none else '' }}">
|
||||
</div>
|
||||
<div class="form-group">
|
||||
<label for="gender">Gender</label>
|
||||
<input id="gender" name="gender" type="text" value="{{ contact.gender or '' }}">
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<div class="form-group">
|
||||
<label for="current_job">Current Job</label>
|
||||
<input id="current_job" name="current_job" type="text" value="{{ contact.current_job or '' }}">
|
||||
</div>
|
||||
|
||||
<div class="form-group">
|
||||
<label for="timezone">Timezone</label>
|
||||
<input id="timezone" name="timezone" type="text" value="{{ contact.timezone or '' }}">
|
||||
</div>
|
||||
|
||||
<div class="form-group">
|
||||
<label for="profile_pic">Profile Picture URL</label>
|
||||
<input id="profile_pic" name="profile_pic" type="url" placeholder="https://example.com/photo.jpg" value="{{ contact.profile_pic or '' }}">
|
||||
</div>
|
||||
|
||||
<div class="form-group">
|
||||
<label for="bio">Bio</label>
|
||||
<textarea id="bio" name="bio" rows="3">{{ contact.bio or '' }}</textarea>
|
||||
</div>
|
||||
|
||||
<div class="form-group">
|
||||
<label for="goals">Goals</label>
|
||||
<textarea id="goals" name="goals" rows="3">{{ contact.goals or '' }}</textarea>
|
||||
</div>
|
||||
|
||||
<div class="form-group">
|
||||
<label for="social_structure_style">Social Structure Style</label>
|
||||
<input id="social_structure_style" name="social_structure_style" type="text" value="{{ contact.social_structure_style or '' }}">
|
||||
</div>
|
||||
|
||||
<div class="form-group">
|
||||
<label for="self_sufficiency_score">Self-Sufficiency Score (1-10)</label>
|
||||
<input id="self_sufficiency_score" name="self_sufficiency_score" type="number" min="1" max="10" value="{{ contact.self_sufficiency_score if contact and contact.self_sufficiency_score is not none else '' }}">
|
||||
</div>
|
||||
|
||||
<div class="form-group">
|
||||
<label for="safe_conversation_starters">Safe Conversation Starters</label>
|
||||
<textarea id="safe_conversation_starters" name="safe_conversation_starters" rows="2">{{ contact.safe_conversation_starters or '' }}</textarea>
|
||||
</div>
|
||||
|
||||
<div class="form-group">
|
||||
<label for="topics_to_avoid">Topics to Avoid</label>
|
||||
<textarea id="topics_to_avoid" name="topics_to_avoid" rows="2">{{ contact.topics_to_avoid or '' }}</textarea>
|
||||
</div>
|
||||
|
||||
<div class="form-group">
|
||||
<label for="ssn">SSN</label>
|
||||
<input id="ssn" name="ssn" type="text" value="{{ contact.ssn or '' }}">
|
||||
</div>
|
||||
|
||||
{% if all_needs %}
|
||||
<div class="form-group">
|
||||
<label>Needs/Accommodations</label>
|
||||
<div class="checkbox-group">
|
||||
{% for need in all_needs %}
|
||||
<label class="checkbox-label">
|
||||
<input type="checkbox" name="need_ids" value="{{ need.id }}"
|
||||
{% if contact and need in contact.needs %}checked{% endif %}>
|
||||
{{ need.name }}
|
||||
</label>
|
||||
{% endfor %}
|
||||
</div>
|
||||
</div>
|
||||
{% endif %}
|
||||
|
||||
<div class="form-actions">
|
||||
<button type="submit" class="btn btn-primary">Save</button>
|
||||
{% if contact %}
|
||||
<a href="/contacts/{{ contact.id }}" class="btn">Cancel</a>
|
||||
{% else %}
|
||||
<a href="/contacts" class="btn">Cancel</a>
|
||||
{% endif %}
|
||||
</div>
|
||||
</form>
|
||||
</div>
|
||||
{% endblock %}
|
||||
@@ -0,0 +1,14 @@
|
||||
{% extends "base.html" %}
|
||||
{% block title %}Contacts{% endblock %}
|
||||
{% block content %}
|
||||
<div class="contact-list">
|
||||
<div class="header">
|
||||
<h1>Contacts</h1>
|
||||
<a href="/contacts/new" class="btn btn-primary">Add Contact</a>
|
||||
</div>
|
||||
|
||||
<div id="contact-table">
|
||||
{% include "partials/contact_table.html" %}
|
||||
</div>
|
||||
</div>
|
||||
{% endblock %}
|
||||
@@ -0,0 +1,198 @@
|
||||
{% extends "base.html" %}
|
||||
{% block title %}Relationship Graph{% endblock %}
|
||||
{% block content %}
|
||||
<div class="graph-container">
|
||||
<div class="header">
|
||||
<h1>Relationship Graph</h1>
|
||||
</div>
|
||||
<p class="graph-hint">Drag nodes to reposition. Closer relationships have shorter, darker edges.</p>
|
||||
<canvas id="graph-canvas" width="900" height="600"
|
||||
style="border: 1px solid var(--color-border); border-radius: 8px; background: var(--color-bg); cursor: grab;">
|
||||
</canvas>
|
||||
<div id="selected-info"></div>
|
||||
<div class="legend">
|
||||
<h4>Relationship Closeness (1-10)</h4>
|
||||
<div class="legend-items">
|
||||
<div class="legend-item">
|
||||
<span class="legend-line" style="background: hsl(220, 70%, 40%); height: 4px; display: inline-block;"></span>
|
||||
<span>10 - Very Close (Spouse, Partner)</span>
|
||||
</div>
|
||||
<div class="legend-item">
|
||||
<span class="legend-line" style="background: hsl(220, 70%, 52%); height: 3px; display: inline-block;"></span>
|
||||
<span>7 - Close (Family, Best Friend)</span>
|
||||
</div>
|
||||
<div class="legend-item">
|
||||
<span class="legend-line" style="background: hsl(220, 70%, 64%); height: 2px; display: inline-block;"></span>
|
||||
<span>4 - Moderate (Friend, Colleague)</span>
|
||||
</div>
|
||||
<div class="legend-item">
|
||||
<span class="legend-line" style="background: hsl(220, 70%, 72%); height: 1px; display: inline-block;"></span>
|
||||
<span>2 - Distant (Acquaintance)</span>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<script>
|
||||
(function() {
|
||||
const RELATIONSHIP_DISPLAY = {{ relationship_type_display|tojson }};
|
||||
const graphData = {{ graph_data|tojson }};
|
||||
|
||||
const canvas = document.getElementById('graph-canvas');
|
||||
const ctx = canvas.getContext('2d');
|
||||
const width = canvas.width;
|
||||
const height = canvas.height;
|
||||
const centerX = width / 2;
|
||||
const centerY = height / 2;
|
||||
|
||||
const nodes = graphData.nodes.map(function(node) {
|
||||
return Object.assign({}, node, {
|
||||
x: centerX + (Math.random() - 0.5) * 300,
|
||||
y: centerY + (Math.random() - 0.5) * 300,
|
||||
vx: 0,
|
||||
vy: 0
|
||||
});
|
||||
});
|
||||
|
||||
const nodeMap = new Map(nodes.map(function(node) { return [node.id, node]; }));
|
||||
|
||||
const edges = graphData.edges.map(function(edge) {
|
||||
const sourceNode = nodeMap.get(edge.source);
|
||||
const targetNode = nodeMap.get(edge.target);
|
||||
if (!sourceNode || !targetNode) return null;
|
||||
return Object.assign({}, edge, { sourceNode: sourceNode, targetNode: targetNode });
|
||||
}).filter(function(edge) { return edge !== null; });
|
||||
|
||||
let dragNode = null;
|
||||
let selectedNode = null;
|
||||
|
||||
const repulsion = 5000;
|
||||
const springStrength = 0.05;
|
||||
const baseSpringLength = 150;
|
||||
const damping = 0.9;
|
||||
const centerPull = 0.01;
|
||||
|
||||
function simulate() {
|
||||
for (const node of nodes) { node.vx = 0; node.vy = 0; }
|
||||
for (let i = 0; i < nodes.length; i++) {
|
||||
for (let j = i + 1; j < nodes.length; j++) {
|
||||
const dx = nodes[j].x - nodes[i].x;
|
||||
const dy = nodes[j].y - nodes[i].y;
|
||||
const dist = Math.sqrt(dx * dx + dy * dy) || 1;
|
||||
const force = repulsion / (dist * dist);
|
||||
const fx = (dx / dist) * force;
|
||||
const fy = (dy / dist) * force;
|
||||
nodes[i].vx -= fx; nodes[i].vy -= fy;
|
||||
nodes[j].vx += fx; nodes[j].vy += fy;
|
||||
}
|
||||
}
|
||||
for (const edge of edges) {
|
||||
const dx = edge.targetNode.x - edge.sourceNode.x;
|
||||
const dy = edge.targetNode.y - edge.sourceNode.y;
|
||||
const dist = Math.sqrt(dx * dx + dy * dy) || 1;
|
||||
const normalizedWeight = edge.closeness_weight / 10;
|
||||
const idealLength = baseSpringLength * (1.5 - normalizedWeight);
|
||||
const displacement = dist - idealLength;
|
||||
const force = springStrength * displacement;
|
||||
const fx = (dx / dist) * force;
|
||||
const fy = (dy / dist) * force;
|
||||
edge.sourceNode.vx += fx; edge.sourceNode.vy += fy;
|
||||
edge.targetNode.vx -= fx; edge.targetNode.vy -= fy;
|
||||
}
|
||||
for (const node of nodes) {
|
||||
node.vx += (centerX - node.x) * centerPull;
|
||||
node.vy += (centerY - node.y) * centerPull;
|
||||
}
|
||||
for (const node of nodes) {
|
||||
if (node === dragNode) continue;
|
||||
node.x += node.vx * damping;
|
||||
node.y += node.vy * damping;
|
||||
node.x = Math.max(30, Math.min(width - 30, node.x));
|
||||
node.y = Math.max(30, Math.min(height - 30, node.y));
|
||||
}
|
||||
}
|
||||
|
||||
function getEdgeColor(weight) {
|
||||
const normalized = weight / 10;
|
||||
return 'hsl(220, 70%, ' + (80 - normalized * 40) + '%)';
|
||||
}
|
||||
|
||||
function draw() {
|
||||
ctx.clearRect(0, 0, width, height);
|
||||
for (const edge of edges) {
|
||||
const lineWidth = 1 + (edge.closeness_weight / 10) * 3;
|
||||
ctx.strokeStyle = getEdgeColor(edge.closeness_weight);
|
||||
ctx.lineWidth = lineWidth;
|
||||
ctx.beginPath();
|
||||
ctx.moveTo(edge.sourceNode.x, edge.sourceNode.y);
|
||||
ctx.lineTo(edge.targetNode.x, edge.targetNode.y);
|
||||
ctx.stroke();
|
||||
const midX = (edge.sourceNode.x + edge.targetNode.x) / 2;
|
||||
const midY = (edge.sourceNode.y + edge.targetNode.y) / 2;
|
||||
ctx.fillStyle = '#666';
|
||||
ctx.font = '10px sans-serif';
|
||||
ctx.textAlign = 'center';
|
||||
const label = RELATIONSHIP_DISPLAY[edge.relationship_type] || edge.relationship_type;
|
||||
ctx.fillText(label, midX, midY - 5);
|
||||
}
|
||||
for (const node of nodes) {
|
||||
const isSelected = node === selectedNode;
|
||||
const radius = isSelected ? 25 : 20;
|
||||
ctx.beginPath();
|
||||
ctx.arc(node.x, node.y, radius, 0, Math.PI * 2);
|
||||
ctx.fillStyle = isSelected ? '#0066cc' : '#fff';
|
||||
ctx.fill();
|
||||
ctx.strokeStyle = '#0066cc';
|
||||
ctx.lineWidth = 2;
|
||||
ctx.stroke();
|
||||
ctx.fillStyle = isSelected ? '#fff' : '#333';
|
||||
ctx.font = '12px sans-serif';
|
||||
ctx.textAlign = 'center';
|
||||
ctx.textBaseline = 'middle';
|
||||
const name = node.name.length > 10 ? node.name.slice(0, 9) + '\u2026' : node.name;
|
||||
ctx.fillText(name, node.x, node.y);
|
||||
}
|
||||
}
|
||||
|
||||
function animate() {
|
||||
simulate();
|
||||
draw();
|
||||
requestAnimationFrame(animate);
|
||||
}
|
||||
animate();
|
||||
|
||||
function getNodeAt(x, y) {
|
||||
for (const node of nodes) {
|
||||
const dx = x - node.x;
|
||||
const dy = y - node.y;
|
||||
if (dx * dx + dy * dy < 400) return node;
|
||||
}
|
||||
return null;
|
||||
}
|
||||
|
||||
canvas.addEventListener('mousedown', function(event) {
|
||||
const rect = canvas.getBoundingClientRect();
|
||||
const node = getNodeAt(event.clientX - rect.left, event.clientY - rect.top);
|
||||
if (node) {
|
||||
dragNode = node;
|
||||
selectedNode = node;
|
||||
const infoDiv = document.getElementById('selected-info');
|
||||
let html = '<div class="selected-info"><h3>' + node.name + '</h3>';
|
||||
if (node.current_job) html += '<p>Job: ' + node.current_job + '</p>';
|
||||
html += '<a href="/contacts/' + node.id + '">View details</a></div>';
|
||||
infoDiv.innerHTML = html;
|
||||
}
|
||||
});
|
||||
|
||||
canvas.addEventListener('mousemove', function(event) {
|
||||
if (!dragNode) return;
|
||||
const rect = canvas.getBoundingClientRect();
|
||||
dragNode.x = event.clientX - rect.left;
|
||||
dragNode.y = event.clientY - rect.top;
|
||||
});
|
||||
|
||||
canvas.addEventListener('mouseup', function() { dragNode = null; });
|
||||
canvas.addEventListener('mouseleave', function() { dragNode = null; });
|
||||
})();
|
||||
</script>
|
||||
{% endblock %}
|
||||
@@ -0,0 +1,31 @@
|
||||
{% extends "base.html" %}
|
||||
{% block title %}Needs{% endblock %}
|
||||
{% block content %}
|
||||
<div class="need-list">
|
||||
<div class="header">
|
||||
<h1>Needs / Accommodations</h1>
|
||||
<button class="btn btn-primary" onclick="document.getElementById('need-form').toggleAttribute('hidden')">Add Need</button>
|
||||
</div>
|
||||
|
||||
<form id="need-form" hidden
|
||||
hx-post="/htmx/needs"
|
||||
hx-target="#need-items"
|
||||
hx-swap="innerHTML"
|
||||
hx-on::after-request="if(event.detail.successful) this.reset()"
|
||||
class="need-form">
|
||||
<div class="form-group">
|
||||
<label for="name">Name *</label>
|
||||
<input id="name" name="name" type="text" placeholder="e.g., Light Sensitive, ADHD" required>
|
||||
</div>
|
||||
<div class="form-group">
|
||||
<label for="description">Description</label>
|
||||
<textarea id="description" name="description" placeholder="Optional description..." rows="2"></textarea>
|
||||
</div>
|
||||
<button type="submit" class="btn btn-primary">Create</button>
|
||||
</form>
|
||||
|
||||
<div id="need-items">
|
||||
{% include "partials/need_items.html" %}
|
||||
</div>
|
||||
</div>
|
||||
{% endblock %}
|
||||
@@ -0,0 +1,33 @@
|
||||
{% if contacts %}
|
||||
<table>
|
||||
<thead>
|
||||
<tr>
|
||||
<th>Name</th>
|
||||
<th>Job</th>
|
||||
<th>Timezone</th>
|
||||
<th>Actions</th>
|
||||
</tr>
|
||||
</thead>
|
||||
<tbody>
|
||||
{% for contact in contacts %}
|
||||
<tr id="contact-row-{{ contact.id }}">
|
||||
<td><a href="/contacts/{{ contact.id }}">{{ contact.name }}</a></td>
|
||||
<td>{{ contact.current_job or "-" }}</td>
|
||||
<td>{{ contact.timezone or "-" }}</td>
|
||||
<td>
|
||||
<a href="/contacts/{{ contact.id }}/edit" class="btn">Edit</a>
|
||||
<button class="btn btn-danger"
|
||||
hx-delete="/api/contacts/{{ contact.id }}"
|
||||
hx-target="#contact-row-{{ contact.id }}"
|
||||
hx-swap="outerHTML"
|
||||
hx-confirm="Delete this contact?">
|
||||
Delete
|
||||
</button>
|
||||
</td>
|
||||
</tr>
|
||||
{% endfor %}
|
||||
</tbody>
|
||||
</table>
|
||||
{% else %}
|
||||
<p>No contacts yet.</p>
|
||||
{% endif %}
|
||||
@@ -0,0 +1,14 @@
|
||||
<ul class="manage-needs-list">
|
||||
{% for need in contact.needs %}
|
||||
<li id="contact-need-{{ need.id }}">
|
||||
<strong>{{ need.name }}</strong>
|
||||
{% if need.description %}<span> - {{ need.description }}</span>{% endif %}
|
||||
<button class="btn btn-small btn-danger"
|
||||
hx-delete="/api/contacts/{{ contact.id }}/needs/{{ need.id }}"
|
||||
hx-target="#contact-need-{{ need.id }}"
|
||||
hx-swap="outerHTML">
|
||||
Remove
|
||||
</button>
|
||||
</li>
|
||||
{% endfor %}
|
||||
</ul>
|
||||
@@ -0,0 +1,23 @@
|
||||
{% for rel in contact.related_to %}
|
||||
<div class="manage-rel-item" id="rel-{{ contact.id }}-{{ rel.related_contact_id }}">
|
||||
<a href="/contacts/{{ rel.related_contact_id }}">{{ contact_names[rel.related_contact_id] }}</a>
|
||||
<span class="tag">{{ rel.relationship_type|replace("_", " ")|title }}</span>
|
||||
<label class="weight-control">
|
||||
<span>Closeness:</span>
|
||||
<input type="range" min="1" max="10" value="{{ rel.closeness_weight }}"
|
||||
hx-post="/htmx/contacts/{{ contact.id }}/relationships/{{ rel.related_contact_id }}/weight"
|
||||
hx-trigger="change"
|
||||
hx-include="this"
|
||||
name="closeness_weight"
|
||||
hx-swap="none"
|
||||
oninput="this.nextElementSibling.textContent = this.value">
|
||||
<span class="weight-value">{{ rel.closeness_weight }}</span>
|
||||
</label>
|
||||
<button class="btn btn-small btn-danger"
|
||||
hx-delete="/api/contacts/{{ contact.id }}/relationships/{{ rel.related_contact_id }}"
|
||||
hx-target="#rel-{{ contact.id }}-{{ rel.related_contact_id }}"
|
||||
hx-swap="outerHTML">
|
||||
Remove
|
||||
</button>
|
||||
</div>
|
||||
{% endfor %}
|
||||
@@ -0,0 +1,21 @@
|
||||
{% if needs %}
|
||||
<ul class="need-items">
|
||||
{% for need in needs %}
|
||||
<li id="need-item-{{ need.id }}">
|
||||
<div class="need-info">
|
||||
<strong>{{ need.name }}</strong>
|
||||
{% if need.description %}<p>{{ need.description }}</p>{% endif %}
|
||||
</div>
|
||||
<button class="btn btn-danger"
|
||||
hx-delete="/api/needs/{{ need.id }}"
|
||||
hx-target="#need-item-{{ need.id }}"
|
||||
hx-swap="outerHTML"
|
||||
hx-confirm="Delete this need?">
|
||||
Delete
|
||||
</button>
|
||||
</li>
|
||||
{% endfor %}
|
||||
</ul>
|
||||
{% else %}
|
||||
<p>No needs defined yet.</p>
|
||||
{% endif %}
|
||||
@@ -6,7 +6,6 @@ import logging
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from python.ebook_search.llm_interface import request_chat_completion
|
||||
from python.ebook_search.prompts import load_prompt
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import httpx
|
||||
@@ -33,8 +32,11 @@ async def answer_query(
|
||||
return "No relevant sources were found."
|
||||
|
||||
logger.info(
|
||||
f"ebook_answer_request_start {config.vllm_base_url=} {config.chat_model=} sources={len(results)} "
|
||||
f"query_length={len(query)}"
|
||||
"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}"
|
||||
@@ -43,8 +45,21 @@ async def answer_query(
|
||||
content = await request_chat_completion(
|
||||
client,
|
||||
config,
|
||||
load_prompt("answer").messages(query=query, context=context),
|
||||
[
|
||||
{
|
||||
"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(f"ebook_answer_request_complete {config.chat_model=} answer_length={len(content)}")
|
||||
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."
|
||||
|
||||
@@ -36,7 +36,10 @@ def schedule_bm25_refresh(app: FastAPI) -> None:
|
||||
app.state.bm25_refresh_task = loop.create_task(refresh_bm25_for_app(app))
|
||||
|
||||
app.state.bm25_refresh_timer = loop.call_later(app.state.config.bm25_refresh_delay_seconds, start_refresh)
|
||||
logger.info(f"ebook_bm25_refresh_scheduled {app.state.config.bm25_refresh_delay_seconds=}")
|
||||
logger.info(
|
||||
"ebook_bm25_refresh_scheduled delay_seconds=%s",
|
||||
app.state.config.bm25_refresh_delay_seconds,
|
||||
)
|
||||
|
||||
|
||||
def cancel_bm25_refresh(app: FastAPI) -> None:
|
||||
|
||||
@@ -6,6 +6,7 @@ from typing import Annotated
|
||||
|
||||
import httpx
|
||||
from fastapi import Depends, Request
|
||||
from sqlalchemy.ext.asyncio import AsyncEngine
|
||||
|
||||
from python.ebook_search.config import EbookSearchConfig
|
||||
|
||||
@@ -15,10 +16,16 @@ def get_config(request: Request) -> EbookSearchConfig:
|
||||
return request.app.state.config
|
||||
|
||||
|
||||
def get_engine(request: Request) -> AsyncEngine:
|
||||
"""Get the database engine from app state."""
|
||||
return request.app.state.engine
|
||||
|
||||
|
||||
def get_http_client(request: Request) -> httpx.AsyncClient:
|
||||
"""Get the shared LLM HTTP client from app state."""
|
||||
return request.app.state.http_client
|
||||
|
||||
|
||||
AppConfig = Annotated[EbookSearchConfig, Depends(get_config)]
|
||||
AppEngine = Annotated[AsyncEngine, Depends(get_engine)]
|
||||
AppHttpClient = Annotated[httpx.AsyncClient, Depends(get_http_client)]
|
||||
|
||||
@@ -65,12 +65,12 @@ def start_book_phrase_judgment(app: FastAPI, background_tasks: BackgroundTasks,
|
||||
"""
|
||||
state = get_judge_task_state(app)
|
||||
if source_id in state.running_book_ids:
|
||||
logger.info(f"ebook_book_phrase_judgment_already_running {source_id=}")
|
||||
logger.info("ebook_book_phrase_judgment_already_running source_id=%s", source_id)
|
||||
return False
|
||||
state.running_book_ids.add(source_id)
|
||||
state.outcome_messages.pop(source_id, None)
|
||||
background_tasks.add_task(judge_book_phrases_for_app, app, source_id)
|
||||
logger.info(f"ebook_book_phrase_judgment_queued {source_id=}")
|
||||
logger.info("ebook_book_phrase_judgment_queued source_id=%s", source_id)
|
||||
return True
|
||||
|
||||
|
||||
@@ -85,8 +85,12 @@ async def judge_book_phrases_for_app(app: FastAPI, source_id: int) -> None:
|
||||
try:
|
||||
result = await judge_candidate_phrases_for_books(app.state.engine, app.state.config, source_ids=[source_id])
|
||||
logger.info(
|
||||
f"ebook_book_phrase_judgment_complete {source_id=} {result.candidates_judged=} {result.protected_phrases=} "
|
||||
f"{result.phrase_mentions=} {result.books_failed=}"
|
||||
"ebook_book_phrase_judgment_complete source_id=%s judged=%s protected=%s mentions=%s failed=%s",
|
||||
source_id,
|
||||
result.candidates_judged,
|
||||
result.protected_phrases,
|
||||
result.phrase_mentions,
|
||||
result.books_failed,
|
||||
)
|
||||
if result.books_failed:
|
||||
message = "Judging failed; see server logs for details"
|
||||
@@ -95,7 +99,7 @@ async def judge_book_phrases_for_app(app: FastAPI, source_id: int) -> None:
|
||||
f"Judged {result.candidates_judged} candidates; {result.protected_phrases} protected phrases promoted"
|
||||
)
|
||||
except Exception:
|
||||
logger.exception(f"ebook_book_phrase_judgment_task_failed {source_id=}")
|
||||
logger.exception("ebook_book_phrase_judgment_task_failed source_id=%s", source_id)
|
||||
message = "Judging failed; see server logs for details"
|
||||
state.running_book_ids.discard(source_id)
|
||||
state.outcome_messages[source_id] = message
|
||||
|
||||
@@ -37,9 +37,16 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]:
|
||||
config = load_config()
|
||||
app.state.config = config
|
||||
logger.info(
|
||||
f"ebook_search_config_loaded {config.top_k=} {config.embedding_model=} {config.embedding_base_url=} "
|
||||
f"{config.vllm_base_url=} {config.rerank.enabled=} {config.phrase_matching_enabled=} {config.answer_enabled=} "
|
||||
f"library_paths={len(config.library_paths)}"
|
||||
"ebook_search_config_loaded top_k=%s embedding_model=%s embedding_base_url=%s vllm_base_url=%s "
|
||||
"rerank_enabled=%s phrase_matching_enabled=%s answer_enabled=%s library_paths=%s",
|
||||
config.top_k,
|
||||
config.embedding_model,
|
||||
config.embedding_base_url,
|
||||
config.vllm_base_url,
|
||||
config.rerank.enabled,
|
||||
config.phrase_matching_enabled,
|
||||
config.answer_enabled,
|
||||
len(config.library_paths),
|
||||
)
|
||||
if not config.library_paths:
|
||||
logger.warning("ebook_search_no_library_paths_configured")
|
||||
|
||||
@@ -10,18 +10,16 @@ from fastapi.responses import HTMLResponse
|
||||
from python.ebook_search.api.bm25_tasks import schedule_bm25_refresh
|
||||
from python.ebook_search.api.dependencies import ( # noqa: TC001 FastAPI resolves these annotated dependencies at runtime
|
||||
AppConfig,
|
||||
AppEngine,
|
||||
AppHttpClient,
|
||||
)
|
||||
from python.ebook_search.api.web import error_response, templates
|
||||
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
|
||||
from python.ebook_search.protected_phrases.generate_ngrams import generate_candidate_phrases_for_books
|
||||
from python.ebook_search.protected_phrases.judge_ngrams import judge_candidate_phrases_for_books
|
||||
from python.ebook_search.protected_phrases.store import book_ids_pending_first_judgment, corpus_phrase_stats
|
||||
from python.fastapi_tools import ( # noqa: TC001 FastAPI resolves these annotated dependencies at runtime
|
||||
AppAsyncEngine,
|
||||
AsyncDbSession,
|
||||
)
|
||||
from python.fastapi_tools import AsyncDbSession # noqa: TC001 FastAPI resolves this annotated dependency at runtime
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -34,8 +32,10 @@ async def admin(request: Request, config: AppConfig, session: AsyncDbSession) ->
|
||||
stats = await embedding_model_stats(session)
|
||||
phrase_stats = await corpus_phrase_stats(session)
|
||||
logger.info(
|
||||
f"ebook_admin_page_loaded models={len(stats)} {phrase_stats.candidate_phrases=} "
|
||||
f"{phrase_stats.protected_phrases=}"
|
||||
"ebook_admin_page_loaded models=%s candidate_phrases=%s protected_phrases=%s",
|
||||
len(stats),
|
||||
phrase_stats.candidate_phrases,
|
||||
phrase_stats.protected_phrases,
|
||||
)
|
||||
return templates.TemplateResponse(
|
||||
request,
|
||||
@@ -52,26 +52,65 @@ async def scan_library(request: Request, config: AppConfig, session: AsyncDbSess
|
||||
await session.commit()
|
||||
except Exception as error:
|
||||
logger.exception("ebook_admin_scan_failed")
|
||||
return error_response(request, error)
|
||||
return templates.TemplateResponse(request, "partials/error.html", {"message": str(error)}, status_code=500)
|
||||
|
||||
logger.info(f"ebook_admin_scan_complete {count=}")
|
||||
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("/phrases/generate-all", response_class=HTMLResponse)
|
||||
async def generate_all_phrases(request: Request, config: AppConfig, engine: AppAsyncEngine) -> HTMLResponse:
|
||||
async def generate_all_phrases(request: Request, config: AppConfig, session: AsyncDbSession) -> HTMLResponse:
|
||||
"""Regenerate candidate phrases for every indexed book without LLM judging."""
|
||||
return await run_phrase_generation(request, config, session, only_missing=False)
|
||||
|
||||
|
||||
@router.post("/phrases/generate-missing", response_class=HTMLResponse)
|
||||
async def generate_missing_phrases(request: Request, config: AppConfig, session: AsyncDbSession) -> HTMLResponse:
|
||||
"""Generate candidate phrases only for books that have none yet."""
|
||||
return await run_phrase_generation(request, config, session, only_missing=True)
|
||||
|
||||
|
||||
async def run_phrase_generation(
|
||||
request: Request,
|
||||
config: AppConfig,
|
||||
session: AsyncDbSession,
|
||||
*,
|
||||
only_missing: bool,
|
||||
) -> HTMLResponse:
|
||||
"""Run candidate phrase generation and render the outcome as an admin status partial.
|
||||
|
||||
Args:
|
||||
request (Request): Current request, for template rendering.
|
||||
config (AppConfig): Runtime phrase-tuning settings.
|
||||
session (AsyncDbSession): Active database session.
|
||||
only_missing (bool): Only generate for books without candidates instead of every book.
|
||||
|
||||
Returns:
|
||||
HTMLResponse: Status partial describing the generation outcome.
|
||||
"""
|
||||
try:
|
||||
result = await generate_candidate_phrases_for_books(engine, config)
|
||||
result = await generate_candidate_phrases_for_books(session, config, only_missing=only_missing)
|
||||
await session.commit()
|
||||
except Exception as error:
|
||||
logger.exception("ebook_admin_generate_phrases_failed")
|
||||
return error_response(request, error)
|
||||
await session.rollback()
|
||||
logger.exception("ebook_admin_generate_phrases_failed only_missing=%s", only_missing)
|
||||
return templates.TemplateResponse(request, "partials/error.html", {"message": str(error)}, status_code=500)
|
||||
|
||||
logger.info(
|
||||
f"ebook_admin_generate_phrases_complete {result.books_seen=} {result.books_built=} {result.candidate_phrases=}"
|
||||
"ebook_admin_generate_phrases_complete only_missing=%s books_seen=%s books_built=%s candidates=%s",
|
||||
only_missing,
|
||||
result.books_seen,
|
||||
result.books_built,
|
||||
result.candidate_phrases,
|
||||
)
|
||||
if only_missing and result.books_seen == 0:
|
||||
return templates.TemplateResponse(
|
||||
request,
|
||||
"partials/admin_status.html",
|
||||
{"message": "All books already have candidate phrases"},
|
||||
)
|
||||
return templates.TemplateResponse(
|
||||
request,
|
||||
"partials/admin_status.html",
|
||||
@@ -85,7 +124,7 @@ async def generate_all_phrases(request: Request, config: AppConfig, engine: AppA
|
||||
|
||||
|
||||
@router.post("/phrases/judge-all", response_class=HTMLResponse)
|
||||
async def judge_all_phrases(request: Request, engine: AppAsyncEngine, config: AppConfig) -> HTMLResponse:
|
||||
async def judge_all_phrases(request: Request, engine: AppEngine, config: AppConfig) -> HTMLResponse:
|
||||
"""Judge unjudged candidate phrases across every indexed book."""
|
||||
return await run_phrase_judgment(request, engine, config, source_ids=None)
|
||||
|
||||
@@ -93,7 +132,7 @@ async def judge_all_phrases(request: Request, engine: AppAsyncEngine, config: Ap
|
||||
@router.post("/phrases/judge-missing", response_class=HTMLResponse)
|
||||
async def judge_missing_phrases(
|
||||
request: Request,
|
||||
engine: AppAsyncEngine,
|
||||
engine: AppEngine,
|
||||
config: AppConfig,
|
||||
session: AsyncDbSession,
|
||||
) -> HTMLResponse:
|
||||
@@ -110,7 +149,7 @@ async def judge_missing_phrases(
|
||||
|
||||
async def run_phrase_judgment(
|
||||
request: Request,
|
||||
engine: AppAsyncEngine,
|
||||
engine: AppEngine,
|
||||
config: AppConfig,
|
||||
*,
|
||||
source_ids: list[int] | None,
|
||||
@@ -119,7 +158,7 @@ async def run_phrase_judgment(
|
||||
|
||||
Args:
|
||||
request (Request): Current request, for template rendering.
|
||||
engine (AppAsyncEngine): Engine used to open per-book judging sessions.
|
||||
engine (AppEngine): Engine used to open per-book judging sessions.
|
||||
config (AppConfig): Runtime phrase-tuning settings.
|
||||
source_ids (list[int] | None): Books to judge; ``None`` judges every indexed book.
|
||||
|
||||
@@ -130,11 +169,17 @@ async def run_phrase_judgment(
|
||||
result = await judge_candidate_phrases_for_books(engine, config, source_ids=source_ids)
|
||||
except Exception as error:
|
||||
logger.exception("ebook_admin_judge_phrases_failed")
|
||||
return error_response(request, error)
|
||||
return templates.TemplateResponse(request, "partials/error.html", {"message": str(error)}, status_code=500)
|
||||
|
||||
logger.info(
|
||||
f"ebook_admin_judge_phrases_complete {result.books_seen=} {result.books_judged=} {result.books_failed=} "
|
||||
f"{result.candidates_judged=} {result.protected_phrases=} {result.phrase_mentions=}"
|
||||
"ebook_admin_judge_phrases_complete books_seen=%s books_judged=%s books_failed=%s candidates_judged=%s "
|
||||
"protected=%s mentions=%s",
|
||||
result.books_seen,
|
||||
result.books_judged,
|
||||
result.books_failed,
|
||||
result.candidates_judged,
|
||||
result.protected_phrases,
|
||||
result.phrase_mentions,
|
||||
)
|
||||
return templates.TemplateResponse(
|
||||
request,
|
||||
@@ -163,9 +208,9 @@ async def embed_missing(
|
||||
await session.commit()
|
||||
except Exception as error:
|
||||
logger.exception("ebook_admin_embed_missing_failed")
|
||||
return error_response(request, error)
|
||||
return templates.TemplateResponse(request, "partials/error.html", {"message": str(error)}, status_code=500)
|
||||
|
||||
logger.info(f"ebook_admin_embed_missing_complete {count=}")
|
||||
logger.info("ebook_admin_embed_missing_complete chunks=%s", count)
|
||||
return templates.TemplateResponse(
|
||||
request,
|
||||
"partials/admin_status.html",
|
||||
@@ -191,12 +236,26 @@ async def embed_all(
|
||||
await session.commit()
|
||||
total += count
|
||||
batches += 1
|
||||
logger.info(f"ebook_admin_embed_all_batch_complete {batches=} {count=} {total=}")
|
||||
logger.info(
|
||||
"ebook_admin_embed_all_batch_complete batch=%s chunks=%s total_chunks=%s",
|
||||
batches,
|
||||
count,
|
||||
total,
|
||||
)
|
||||
except Exception as error:
|
||||
logger.exception(f"ebook_admin_embed_all_failed {batches=} {total=}")
|
||||
return error_response(request, f"Embed all failed after {total} chunks in {batches} batches: {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(f"ebook_admin_embed_all_complete {batches=} {total=}")
|
||||
logger.info("ebook_admin_embed_all_complete batches=%s chunks=%s", batches, total)
|
||||
return templates.TemplateResponse(
|
||||
request,
|
||||
"partials/admin_status.html",
|
||||
|
||||
@@ -60,7 +60,14 @@ async def ready(config: AppConfig, session: AsyncDbSession, client: AppHttpClien
|
||||
status = "ready"
|
||||
status_code = HTTPStatus.OK
|
||||
|
||||
logger.info(f"ebook_ready_check {status=} {database_ok=} {embedding_ok=} {chat_status=} {bm25_status=}")
|
||||
logger.info(
|
||||
"ebook_ready_check status=%s database=%s embedding=%s chat=%s bm25=%s",
|
||||
status,
|
||||
database_ok,
|
||||
embedding_ok,
|
||||
chat_status,
|
||||
bm25_status,
|
||||
)
|
||||
return JSONResponse(content={"status": status, "checks": checks}, status_code=status_code)
|
||||
|
||||
|
||||
@@ -76,7 +83,7 @@ async def check_database(session: AsyncSession) -> bool:
|
||||
try:
|
||||
await session.execute(select(literal(1)))
|
||||
except SQLAlchemyError as error:
|
||||
logger.warning(f"ebook_ready_database_unavailable {error=}")
|
||||
logger.warning("ebook_ready_database_unavailable error=%s", error)
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
@@ -15,7 +15,6 @@ from python.ebook_search.api.dependencies import (
|
||||
from python.ebook_search.api.judge_tasks import is_judging_book, pop_book_judgment_outcome, start_book_phrase_judgment
|
||||
from python.ebook_search.api.web import templates
|
||||
from python.ebook_search.protected_phrases.generate_ngrams import recalculate_candidate_phrases_for_book
|
||||
from python.ebook_search.protected_phrases.store import count_protected_phrases
|
||||
from python.fastapi_tools import AsyncDbSession # noqa: TC001 FastAPI resolves this annotated dependency at runtime
|
||||
from python.orm.richie import EbookCandidatePhrase, EbookChapter, EbookChunk, EbookProtectedPhrase, EbookSource
|
||||
|
||||
@@ -37,7 +36,7 @@ async def index(request: Request, config: AppConfig) -> HTMLResponse:
|
||||
async def books(request: Request, session: AsyncDbSession) -> HTMLResponse:
|
||||
"""Render the indexed books page."""
|
||||
sources = list((await session.scalars(select(EbookSource).order_by(EbookSource.title))).all())
|
||||
logger.info(f"ebook_books_page_loaded count={len(sources)}")
|
||||
logger.info("ebook_books_page_loaded count=%s", len(sources))
|
||||
return templates.TemplateResponse(request, "books.html", {"sources": sources})
|
||||
|
||||
|
||||
@@ -72,6 +71,14 @@ async def get_judged_candidate_count(session: AsyncSession, book_id: int) -> int
|
||||
)
|
||||
|
||||
|
||||
async def get_protected_count(session: AsyncSession, book_id: int) -> int:
|
||||
"""Return the number of protected phrases for one book."""
|
||||
return (
|
||||
await session.scalar(select(func.count(EbookProtectedPhrase.id)).where(EbookProtectedPhrase.book_id == book_id))
|
||||
or 0
|
||||
)
|
||||
|
||||
|
||||
async def get_candidates(session: AsyncSession, book_id: int) -> list[EbookCandidatePhrase]:
|
||||
"""Return the indexed candidates for one book."""
|
||||
return list(
|
||||
@@ -115,7 +122,7 @@ async def book_detail(source_id: int, request: Request, session: AsyncDbSession)
|
||||
chunk_count = await get_chunk_count(session, source.id)
|
||||
candidate_count = await get_candidate_count(session, source.id)
|
||||
judged_candidate_count = await get_judged_candidate_count(session, source.id)
|
||||
protected_count = await count_protected_phrases(session, source.id)
|
||||
protected_count = await get_protected_count(session, source.id)
|
||||
candidates = await get_candidates(session, source.id)
|
||||
protected_phrases = await get_protected_phrases(session, source.id)
|
||||
else:
|
||||
@@ -127,8 +134,14 @@ async def book_detail(source_id: int, request: Request, session: AsyncDbSession)
|
||||
candidates = []
|
||||
protected_phrases = []
|
||||
logger.info(
|
||||
f"ebook_book_detail_loaded {source_id=} found={source is not None} {chapter_count=} {chunk_count=} "
|
||||
f"{candidate_count=} {judged_candidate_count=} {protected_count=}"
|
||||
"ebook_book_detail_loaded source_id=%s found=%s chapters=%s chunks=%s candidates=%s judged=%s protected=%s",
|
||||
source_id,
|
||||
source is not None,
|
||||
chapter_count,
|
||||
chunk_count,
|
||||
candidate_count,
|
||||
judged_candidate_count,
|
||||
protected_count,
|
||||
)
|
||||
return templates.TemplateResponse(
|
||||
request,
|
||||
@@ -155,14 +168,16 @@ async def recalculate_book_phrases(source_id: int, config: AppConfig, session: A
|
||||
if source is None:
|
||||
raise HTTPException(status_code=404, detail="Book not found")
|
||||
|
||||
try:
|
||||
result = await recalculate_candidate_phrases_for_book(session, source, config)
|
||||
except ValueError as error:
|
||||
raise HTTPException(status_code=409, detail=str(error)) from error
|
||||
result = await recalculate_candidate_phrases_for_book(session, source, config, use_process_pool=True)
|
||||
logger.info(
|
||||
f"ebook_book_phrase_recalculation_complete {source_id=} {result.candidate_phrases=} "
|
||||
f"{result.deleted_candidates=} {result.deleted_protected_phrases=} {result.deleted_aliases=} "
|
||||
f"{result.deleted_mentions=}"
|
||||
"ebook_book_phrase_recalculation_complete source_id=%s candidates=%s deleted_candidates=%s "
|
||||
"deleted_protected=%s deleted_aliases=%s deleted_mentions=%s",
|
||||
source_id,
|
||||
result.candidate_phrases,
|
||||
result.deleted_candidates,
|
||||
result.deleted_protected_phrases,
|
||||
result.deleted_aliases,
|
||||
result.deleted_mentions,
|
||||
)
|
||||
return RedirectResponse(
|
||||
url=f"/books/{source_id}?phrases_recalculated={result.candidate_phrases}",
|
||||
@@ -183,5 +198,5 @@ async def judge_book_phrases(
|
||||
raise HTTPException(status_code=404, detail="Book not found")
|
||||
|
||||
started = start_book_phrase_judgment(request.app, background_tasks, source.id)
|
||||
logger.info(f"ebook_book_phrase_judgment_requested {source_id=} {started=}")
|
||||
logger.info("ebook_book_phrase_judgment_requested source_id=%s started=%s", source_id, started)
|
||||
return RedirectResponse(url=f"/books/{source_id}", status_code=303)
|
||||
|
||||
@@ -13,9 +13,10 @@ from fastapi.responses import HTMLResponse
|
||||
from python.ebook_search.answer import answer_query
|
||||
from python.ebook_search.api.dependencies import ( # noqa: TC001 FastAPI resolves these annotated dependencies at runtime
|
||||
AppConfig,
|
||||
AppEngine,
|
||||
AppHttpClient,
|
||||
)
|
||||
from python.ebook_search.api.web import error_response, templates
|
||||
from python.ebook_search.api.web import templates
|
||||
from python.ebook_search.guardrails import (
|
||||
CitationReport,
|
||||
is_confident,
|
||||
@@ -24,7 +25,6 @@ from python.ebook_search.guardrails import (
|
||||
)
|
||||
from python.ebook_search.search import SearchResponse, search_ebooks
|
||||
from python.ebook_search.timing import runtime_step_from_start
|
||||
from python.fastapi_tools import AppAsyncEngine # noqa: TC001 FastAPI resolves this annotated dependency at runtime
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import httpx
|
||||
@@ -49,8 +49,9 @@ async def build_answer(
|
||||
|
||||
if not is_confident(response.results, config):
|
||||
logger.info(
|
||||
f"ebook_answer_low_confidence confidence={retrieval_confidence(response.results):.4f} "
|
||||
f"{config.min_retrieval_confidence=:.4f}"
|
||||
"ebook_answer_low_confidence confidence=%.4f threshold=%.4f",
|
||||
retrieval_confidence(response.results),
|
||||
config.min_retrieval_confidence,
|
||||
)
|
||||
answer = (
|
||||
"Retrieval confidence is low for this query, so answer generation was skipped. "
|
||||
@@ -61,14 +62,18 @@ async def build_answer(
|
||||
try:
|
||||
answer = await answer_query(client, query, response.results, config)
|
||||
except RuntimeError as error:
|
||||
logger.warning(f"ebook_answer_request_failed_falling_back {error=}")
|
||||
logger.warning("ebook_answer_request_failed_falling_back error=%s", error)
|
||||
return "Answer generation failed. Source chunks are still shown below.", False, None
|
||||
|
||||
citation_report = None
|
||||
if config.validate_citations_enabled and response.results:
|
||||
citation_report = validate_citations(answer, len(response.results))
|
||||
if citation_report.invalid or not citation_report.grounded:
|
||||
logger.warning(f"ebook_answer_citation_issue {citation_report.invalid=} {citation_report.grounded=}")
|
||||
logger.warning(
|
||||
"ebook_answer_citation_issue invalid=%s grounded=%s",
|
||||
citation_report.invalid,
|
||||
citation_report.grounded,
|
||||
)
|
||||
return answer, False, citation_report
|
||||
|
||||
|
||||
@@ -76,12 +81,11 @@ async def build_answer(
|
||||
async def search(
|
||||
request: Request,
|
||||
config: AppConfig,
|
||||
engine: AppAsyncEngine,
|
||||
engine: AppEngine,
|
||||
client: AppHttpClient,
|
||||
query: Annotated[str, Form()],
|
||||
*,
|
||||
rerank: Annotated[bool, Form()] = False,
|
||||
phrase_matching: Annotated[bool, Form()] = False,
|
||||
rerank: Annotated[str | None, Form()] = None,
|
||||
phrase_matching: Annotated[str | None, Form()] = None,
|
||||
) -> HTMLResponse:
|
||||
"""Run a search and render HTMX results."""
|
||||
try:
|
||||
@@ -90,12 +94,12 @@ async def search(
|
||||
client,
|
||||
query,
|
||||
config,
|
||||
rerank=rerank,
|
||||
phrase_matching=phrase_matching,
|
||||
rerank=rerank == "true",
|
||||
phrase_matching=phrase_matching == "true",
|
||||
)
|
||||
except Exception as error:
|
||||
logger.exception("ebook_search_request_failed")
|
||||
return error_response(request, error)
|
||||
return templates.TemplateResponse(request, "partials/error.html", {"message": str(error)}, status_code=500)
|
||||
|
||||
answer_start = perf_counter()
|
||||
answer, low_confidence, citation_report = await build_answer(client, query, response, config)
|
||||
@@ -106,10 +110,12 @@ async def search(
|
||||
)
|
||||
|
||||
for step in response.timings:
|
||||
logger.info(f"ebook_search_timing {step.name=} {step.duration_ms=:.1f}")
|
||||
logger.info("ebook_search_timing step=%r runtime_ms=%.1f", step.name, step.duration_ms)
|
||||
logger.info(
|
||||
f"ebook_search_request_complete results={len(response.results)} {response.rank_label=} "
|
||||
f"{response.total_runtime_ms=:.1f}"
|
||||
"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,
|
||||
|
||||
@@ -60,6 +60,13 @@ head %}
|
||||
>
|
||||
<button type="submit">Regenerate all phrases</button>
|
||||
</form>
|
||||
<form
|
||||
hx-post="/admin/phrases/generate-missing"
|
||||
hx-target="#admin-status"
|
||||
hx-swap="innerHTML"
|
||||
>
|
||||
<button type="submit">Add missing phrases</button>
|
||||
</form>
|
||||
<form
|
||||
hx-post="/admin/phrases/judge-all"
|
||||
hx-target="#admin-status"
|
||||
|
||||
@@ -3,14 +3,9 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from fastapi.templating import Jinja2Templates
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from fastapi import Request
|
||||
from fastapi.responses import HTMLResponse
|
||||
|
||||
PACKAGE_DIR = Path(__file__).resolve().parent
|
||||
TEMPLATE_DIR = PACKAGE_DIR / "templates"
|
||||
STATIC_DIR = PACKAGE_DIR / "static"
|
||||
@@ -26,8 +21,3 @@ def static_version(filename: str) -> int:
|
||||
|
||||
templates = Jinja2Templates(directory=TEMPLATE_DIR)
|
||||
templates.env.globals["static_version"] = static_version
|
||||
|
||||
|
||||
def error_response(request: Request, message: object) -> HTMLResponse:
|
||||
"""Render the shared error partial for a failed UI request."""
|
||||
return templates.TemplateResponse(request, "partials/error.html", {"message": str(message)}, status_code=500)
|
||||
|
||||
@@ -15,7 +15,6 @@ from typing import TYPE_CHECKING
|
||||
import bm25s
|
||||
from sqlalchemy import func, select, union_all
|
||||
|
||||
from python.ebook_search.chunk_records import CHUNK_RECORD_COLUMNS
|
||||
from python.orm.richie import EbookChapter, EbookChunk, EbookSource
|
||||
|
||||
if TYPE_CHECKING:
|
||||
@@ -81,19 +80,23 @@ async def ensure_bm25_corpus(session: AsyncSession, config: EbookSearchConfig) -
|
||||
manifest = read_bm25_manifest(index_path)
|
||||
db_updated_at = await corpus_last_updated_at(session)
|
||||
if not bm25_index_exists(index_path, manifest):
|
||||
logger.info(f"ebook_bm25_index_missing {index_path=}")
|
||||
logger.info("ebook_bm25_index_missing path=%s", index_path)
|
||||
await 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(
|
||||
f"ebook_bm25_index_stale {index_path=} created_at={manifest.created_at.isoformat()} "
|
||||
f"db_updated_at={db_updated_at.isoformat()}"
|
||||
"ebook_bm25_index_stale path=%s created_at=%s db_updated_at=%s",
|
||||
index_path,
|
||||
manifest.created_at.isoformat(),
|
||||
db_updated_at.isoformat(),
|
||||
)
|
||||
await refresh_bm25_corpus(session, config, db_updated_at=db_updated_at)
|
||||
return
|
||||
logger.info(
|
||||
f"ebook_bm25_index_current {index_path=} chunks={manifest.chunk_count if manifest else 0} "
|
||||
f"created_at={manifest.created_at.isoformat() if manifest else None}"
|
||||
"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,
|
||||
)
|
||||
|
||||
|
||||
@@ -116,7 +119,10 @@ async def refresh_bm25_corpus(
|
||||
)
|
||||
await asyncio.to_thread(write_bm25_corpus, index_path, records, texts, manifest)
|
||||
logger.info(
|
||||
f"ebook_bm25_index_refreshed {index_path=} {manifest.chunk_count=} created_at={manifest.created_at.isoformat()}"
|
||||
"ebook_bm25_index_refreshed path=%s chunks=%s created_at=%s",
|
||||
index_path,
|
||||
manifest.chunk_count,
|
||||
manifest.created_at.isoformat(),
|
||||
)
|
||||
return manifest
|
||||
|
||||
@@ -129,7 +135,7 @@ def load_bm25_corpus(config: EbookSearchConfig) -> BM25Corpus:
|
||||
"""
|
||||
index_path = bm25_index_path(config)
|
||||
active_index_path = get_current_bm25_index(index_path)
|
||||
logger.info(f"ebook_bm25_corpus_cache_load {index_path=} {active_index_path=}")
|
||||
logger.info("ebook_bm25_corpus_cache_load path=%s active_path=%s", index_path, active_index_path)
|
||||
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}"
|
||||
@@ -170,7 +176,13 @@ async def fetch_bm25_corpus_records(session: AsyncSession) -> tuple[list[dict[st
|
||||
"""
|
||||
statement = (
|
||||
select(
|
||||
*CHUNK_RECORD_COLUMNS,
|
||||
EbookChunk.id.label("chunk_id"),
|
||||
EbookChunk.text.label("text"),
|
||||
EbookSource.id.label("source_id"),
|
||||
EbookSource.title.label("source_title"),
|
||||
EbookSource.author.label("source_author"),
|
||||
EbookChapter.title.label("chapter_title"),
|
||||
EbookChunk.page_label.label("page_label"),
|
||||
EbookChunk.search_text.label("bm25_text"),
|
||||
)
|
||||
.select_from(EbookChunk)
|
||||
|
||||
@@ -1,13 +0,0 @@
|
||||
"""Shared database columns used to build search-result records."""
|
||||
|
||||
from python.orm.richie import EbookChapter, EbookChunk, EbookSource
|
||||
|
||||
CHUNK_RECORD_COLUMNS = (
|
||||
EbookChunk.id.label("chunk_id"),
|
||||
EbookChunk.text.label("text"),
|
||||
EbookSource.id.label("source_id"),
|
||||
EbookSource.title.label("source_title"),
|
||||
EbookSource.author.label("source_author"),
|
||||
EbookChapter.title.label("chapter_title"),
|
||||
EbookChunk.page_label.label("page_label"),
|
||||
)
|
||||
@@ -2,6 +2,7 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from os import getenv
|
||||
from typing import Annotated, Self
|
||||
|
||||
from pydantic import AliasChoices, Field, field_validator, model_validator
|
||||
@@ -31,13 +32,18 @@ def normalize_embedding_alias(model: str) -> str:
|
||||
return standard_model
|
||||
|
||||
|
||||
def normalize_embedding_model(default: str = "qwen3-embedding-0.6b") -> str:
|
||||
"""Normalize the configured embedding alias to its provider model name."""
|
||||
return normalize_embedding_alias(getenv("EBOOK_SEARCH_EMBEDDING_MODEL", default))
|
||||
|
||||
|
||||
class RerankConfig(BaseSettings):
|
||||
"""vLLM reranker settings."""
|
||||
|
||||
model_config = SettingsConfigDict(env_prefix="EBOOK_SEARCH_RERANK_", frozen=True, protected_namespaces=())
|
||||
|
||||
enabled: bool = True
|
||||
base_url: str = "http://bob:8001"
|
||||
base_url: str = "http://192.168.90.25:8001"
|
||||
model: str = "qwen3-reranker-06b"
|
||||
candidates: int = 24
|
||||
timeout_seconds: float = 30.0
|
||||
@@ -67,7 +73,7 @@ class EbookSearchConfig(BaseSettings):
|
||||
)
|
||||
chat_model: str = "deepseek-v4-flash"
|
||||
answer_enabled: bool = True
|
||||
embedding_base_url: str = "http://bob:8000/v1"
|
||||
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
|
||||
@@ -91,8 +97,6 @@ class EbookSearchConfig(BaseSettings):
|
||||
phrase_min_tokens: int = 2
|
||||
phrase_max_tokens: int = 5
|
||||
phrase_max_entity_tokens: int = 8
|
||||
phrase_yake_top_k: int = 1000
|
||||
phrase_yake_dedup_limit: float = 0.85
|
||||
phrase_raw_ngram_min_count: int = 2
|
||||
phrase_raw_count_score_threshold: int = 3
|
||||
phrase_raw_count_high_score_threshold: int = 10
|
||||
@@ -101,10 +105,6 @@ class EbookSearchConfig(BaseSettings):
|
||||
phrase_target_protected_per_book: int = 100
|
||||
phrase_default_allow_nested: bool = False
|
||||
phrase_default_suppress_children: bool = True
|
||||
phrase_bad_start_score_penalty: float = 10.0
|
||||
phrase_bad_end_score_penalty: float = 10.0
|
||||
phrase_multi_source_score_bonus: float = 2.0
|
||||
phrase_multi_source_min_sources: int = 2
|
||||
|
||||
@field_validator("library_paths", mode="before")
|
||||
@classmethod
|
||||
|
||||
@@ -1,17 +1,12 @@
|
||||
FROM python:3.14-slim AS base
|
||||
|
||||
COPY --from=ghcr.io/astral-sh/uv:0.11.26 /uv /uvx /bin/
|
||||
FROM python:3.14-slim
|
||||
|
||||
ENV PYTHONDONTWRITEBYTECODE=1 \
|
||||
PYTHONUNBUFFERED=1 \
|
||||
PIP_NO_CACHE_DIR=1 \
|
||||
APP_DIR=/home/richie/dotfiles \
|
||||
UV_PROJECT_ENVIRONMENT=/opt/venv \
|
||||
UV_PYTHON_DOWNLOADS=never \
|
||||
UV_NO_CACHE=1
|
||||
|
||||
# Separate ENV instruction so ${APP_DIR} and ${PATH} from above resolve.
|
||||
ENV PYTHONPATH=${APP_DIR} \
|
||||
PATH=/opt/venv/bin:${PATH}
|
||||
EBOOK_SEARCH_HOST=0.0.0.0 \
|
||||
EBOOK_SEARCH_PORT=8070 \
|
||||
EBOOK_SEARCH_BM25_INDEX_DIR=/data/bm25
|
||||
|
||||
WORKDIR ${APP_DIR}
|
||||
|
||||
@@ -19,29 +14,29 @@ RUN apt-get update \
|
||||
&& apt-get install -y --no-install-recommends build-essential curl \
|
||||
&& rm -rf /var/lib/apt/lists/*
|
||||
|
||||
COPY python/ebook_search/docker/pyproject.toml python/ebook_search/docker/uv.lock ./
|
||||
|
||||
RUN uv sync --locked --no-dev
|
||||
|
||||
|
||||
FROM base AS test
|
||||
|
||||
RUN uv sync --locked
|
||||
|
||||
COPY pyproject.toml README.md LICENSE ./
|
||||
COPY python ./python
|
||||
COPY tests/__init__.py ./tests/__init__.py
|
||||
COPY tests/ebook_search ./tests/ebook_search
|
||||
|
||||
CMD ["pytest"]
|
||||
|
||||
|
||||
FROM base AS runtime
|
||||
|
||||
ENV EBOOK_SEARCH_HOST=0.0.0.0 \
|
||||
EBOOK_SEARCH_PORT=8070 \
|
||||
EBOOK_SEARCH_BM25_INDEX_DIR=/data/bm25
|
||||
|
||||
COPY python ./python
|
||||
RUN python -m pip install --upgrade pip \
|
||||
&& python -m pip install \
|
||||
"alembic" \
|
||||
"beautifulsoup4" \
|
||||
"bm25s" \
|
||||
"ebooklib" \
|
||||
"fastapi" \
|
||||
"httpx" \
|
||||
"jinja2" \
|
||||
"pgvector" \
|
||||
"psycopg[binary]" \
|
||||
"pydantic" \
|
||||
"pydantic-settings" \
|
||||
"python-multipart" \
|
||||
"sqlalchemy[asyncio]" \
|
||||
"tiktoken" \
|
||||
"typer" \
|
||||
"uvicorn[standard]" \
|
||||
"yake" \
|
||||
&& python -m pip install --no-deps --editable "${APP_DIR}"
|
||||
|
||||
RUN useradd --create-home --uid 10001 app \
|
||||
&& mkdir -p /data \
|
||||
|
||||
@@ -3,27 +3,26 @@
|
||||
Run the EPUB search app against the existing Postgres database on `jeeves`:
|
||||
|
||||
```sh
|
||||
python -m python.ebook_search.docker.containers start --library-path /path/to/epubs --build
|
||||
ebook-search-containers start --library-path /path/to/epubs --build
|
||||
```
|
||||
|
||||
All ebook-search Docker files live in this directory:
|
||||
|
||||
- `Dockerfile` — multi-stage: `test` (runs pytest) and `runtime` (default target, the app image)
|
||||
- `Dockerfile`
|
||||
- `docker-compose.yml`
|
||||
- `containers.py` — Typer lifecycle CLI
|
||||
- `pyproject.toml` / `uv.lock` — the container's uv-locked dependencies
|
||||
- `containers.py`
|
||||
- `container.py`
|
||||
|
||||
The app listens on `http://localhost:8070`.
|
||||
|
||||
Useful lifecycle commands:
|
||||
|
||||
```sh
|
||||
python -m python.ebook_search.docker.containers build
|
||||
python -m python.ebook_search.docker.containers start --library-path /path/to/epubs
|
||||
python -m python.ebook_search.docker.containers test
|
||||
python -m python.ebook_search.docker.containers logs
|
||||
python -m python.ebook_search.docker.containers ps
|
||||
python -m python.ebook_search.docker.containers stop
|
||||
ebook-search-containers build
|
||||
ebook-search-containers start --library-path /path/to/epubs
|
||||
ebook-search-containers logs
|
||||
ebook-search-containers ps
|
||||
ebook-search-containers stop
|
||||
```
|
||||
|
||||
Direct compose usage from the repo root:
|
||||
@@ -32,46 +31,10 @@ Direct compose usage from the repo root:
|
||||
docker compose -f python/ebook_search/docker/docker-compose.yml ps
|
||||
```
|
||||
|
||||
## Dependencies
|
||||
|
||||
The image builds its environment with uv from `pyproject.toml` + `uv.lock` in this
|
||||
directory — this is the source of truth for the container's dependencies. To add or
|
||||
update a dependency, edit `pyproject.toml` here and regenerate the lock (uv is
|
||||
available in the `ebook-search` dev shell):
|
||||
|
||||
```sh
|
||||
nix develop .#ebook-search -c uv lock --project python/ebook_search/docker
|
||||
```
|
||||
|
||||
## Tests
|
||||
|
||||
The main pytest suite excludes `tests/ebook_search` (its dependencies are no longer
|
||||
in the nix dev shell). The `test ebook search` CI workflow runs them in a uv env
|
||||
built from the lockfile in this directory — same commands work locally from the
|
||||
repo root (the `--override-ini` drops the main suite's ignore):
|
||||
|
||||
```sh
|
||||
uv sync --locked --project python/ebook_search/docker
|
||||
uv run --project python/ebook_search/docker --no-sync pytest tests/ebook_search --override-ini addopts="-n auto -ra"
|
||||
```
|
||||
|
||||
They can also run inside the Docker `test` image, which validates the image itself:
|
||||
|
||||
```sh
|
||||
python -m python.ebook_search.docker.containers test
|
||||
```
|
||||
|
||||
or the raw docker equivalent:
|
||||
|
||||
```sh
|
||||
docker build --file python/ebook_search/docker/Dockerfile --target test --tag ebook-search:test .
|
||||
docker run --rm ebook-search:test
|
||||
```
|
||||
|
||||
## Configuration
|
||||
|
||||
The compose service loads the repo root `.env` into the container via `env_file`.
|
||||
The compose service also loads the repo root `.env` into the container via `env_file`.
|
||||
|
||||
Mount your EPUB directory by setting `EBOOK_LIBRARY_HOST_PATH` in an env file or on the command line. The container sees it as `/library`, and `EBOOK_SEARCH_LIBRARY_PATHS` is set to `/library` inside the container.
|
||||
|
||||
Database connection settings are controlled by `RICHIE_DB`, `RICHIE_HOST`, `RICHIE_PORT`, `RICHIE_USER`, and `RICHIE_PASSWORD`. The default host is `jeeves`.
|
||||
|
||||
Startup runs the Richie Alembic migrations automatically after creating the `main` schema and `vector` extension.
|
||||
|
||||
@@ -32,7 +32,7 @@ def docker_run(
|
||||
capture_output: bool = False,
|
||||
) -> subprocess.CompletedProcess[str]:
|
||||
"""Run docker with repo-root cwd and consistent error handling."""
|
||||
logger.info(f"docker {' '.join(arguments)}")
|
||||
logger.info("docker %s", " ".join(arguments))
|
||||
return subprocess.run(
|
||||
["docker", *arguments],
|
||||
cwd=get_repo_dir(),
|
||||
@@ -73,23 +73,6 @@ def build_image() -> None:
|
||||
raise RuntimeError(msg)
|
||||
|
||||
|
||||
def build_test_image() -> None:
|
||||
"""Build the ebook search test Docker image."""
|
||||
dockerfile = Path(__file__).resolve().with_name("Dockerfile")
|
||||
result = docker_run(["build", "--file", str(dockerfile), "--target", "test", "--tag", "ebook-search:test", "."])
|
||||
if result.returncode != 0:
|
||||
msg = "Failed to build ebook search test image"
|
||||
raise RuntimeError(msg)
|
||||
|
||||
|
||||
def run_test_image() -> None:
|
||||
"""Run the ebook search test suite inside Docker."""
|
||||
result = docker_run(["run", "--rm", "ebook-search:test"])
|
||||
if result.returncode != 0:
|
||||
msg = f"Ebook search tests failed with code {result.returncode}"
|
||||
raise RuntimeError(msg)
|
||||
|
||||
|
||||
def start_stack(
|
||||
*,
|
||||
library_path: Path | None = None,
|
||||
@@ -227,19 +210,6 @@ def logs(
|
||||
typer.echo(output)
|
||||
|
||||
|
||||
@app.command("test")
|
||||
def run_tests(
|
||||
*,
|
||||
build: Annotated[bool, typer.Option("--build/--no-build", help="Build the test image before running.")] = True,
|
||||
log_level: Annotated[str, typer.Option(help="Log level.")] = "INFO",
|
||||
) -> None:
|
||||
"""Run ebook search tests inside the Docker test image."""
|
||||
configure_logger(log_level)
|
||||
if build:
|
||||
build_test_image()
|
||||
run_test_image()
|
||||
|
||||
|
||||
@app.command("ps")
|
||||
def ps() -> None:
|
||||
"""Show ebook search container status."""
|
||||
|
||||
@@ -9,6 +9,8 @@ services:
|
||||
restart: unless-stopped
|
||||
ports:
|
||||
- "${EBOOK_SEARCH_PORT:-8070}:8070"
|
||||
extra_hosts:
|
||||
- "jeeves:192.168.90.40"
|
||||
env_file:
|
||||
- ../../../.env
|
||||
environment:
|
||||
|
||||
@@ -1,41 +0,0 @@
|
||||
[project]
|
||||
name = "ebook-search"
|
||||
version = "0.1.0"
|
||||
description = "Locked runtime environment for the ebook search container."
|
||||
requires-python = "~=3.14.0"
|
||||
dependencies = [
|
||||
"alembic",
|
||||
"beautifulsoup4",
|
||||
"bm25s",
|
||||
"ebooklib",
|
||||
"fastapi",
|
||||
"httpx",
|
||||
"jinja2",
|
||||
"pgvector",
|
||||
"psycopg[binary]",
|
||||
"pydantic",
|
||||
"pydantic-settings",
|
||||
"python-multipart",
|
||||
"sqlalchemy[asyncio]",
|
||||
"tiktoken",
|
||||
"typer",
|
||||
"uvicorn[standard]",
|
||||
"yake",
|
||||
]
|
||||
|
||||
[dependency-groups]
|
||||
dev = [
|
||||
"aiosqlite",
|
||||
"pytest",
|
||||
"pytest-asyncio",
|
||||
"pytest-mock",
|
||||
"pytest-xdist",
|
||||
]
|
||||
|
||||
[tool.uv]
|
||||
package = false
|
||||
|
||||
[tool.pytest.ini_options]
|
||||
addopts = "-n auto -ra"
|
||||
asyncio_mode = "auto"
|
||||
testpaths = ["tests/ebook_search"]
|
||||
Generated
-1143
File diff suppressed because it is too large
Load Diff
@@ -72,14 +72,24 @@ async def embed_texts(
|
||||
config: EbookSearchConfig,
|
||||
) -> list[list[float]]:
|
||||
"""Embed text with the configured vLLM embedding model."""
|
||||
logger.info(f"ebook_embed_request_start {config.embedding_base_url=} {config.embedding_model=} count={len(texts)}")
|
||||
logger.info(
|
||||
"ebook_embed_request_start base_url=%s model=%s count=%s",
|
||||
config.embedding_base_url,
|
||||
config.embedding_model,
|
||||
len(texts),
|
||||
)
|
||||
vectors = await request_embeddings(client, 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(f"ebook_embed_request_complete {config.embedding_model=} count={len(vectors)} {expected_dimension=}")
|
||||
logger.info(
|
||||
"ebook_embed_request_complete model=%s count=%s dimension=%s",
|
||||
config.embedding_model,
|
||||
len(vectors),
|
||||
expected_dimension,
|
||||
)
|
||||
return vectors
|
||||
|
||||
|
||||
@@ -95,7 +105,7 @@ async def ensure_embedding_models(session: AsyncSession) -> None:
|
||||
existing = await 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(f"ebook_embedding_model_created {name=} {dimension=}")
|
||||
logger.info("ebook_embedding_model_created model=%s dimension=%s", name, dimension)
|
||||
await session.flush()
|
||||
|
||||
|
||||
@@ -149,10 +159,10 @@ async def embed_missing_chunks(session: AsyncSession, client: httpx.AsyncClient,
|
||||
)
|
||||
)
|
||||
if not chunks:
|
||||
logger.info(f"ebook_embed_missing_none {config.embedding_model=}")
|
||||
logger.info("ebook_embed_missing_none model=%s", config.embedding_model)
|
||||
return 0
|
||||
|
||||
logger.info(f"ebook_embed_missing_batch_start {config.embedding_model=} count={len(chunks)}")
|
||||
logger.info("ebook_embed_missing_batch_start model=%s count=%s", config.embedding_model, len(chunks))
|
||||
vectors = await embed_texts(client, [chunk.text for chunk in chunks], config)
|
||||
rows = [
|
||||
{"chunk_id": chunk.id, "model_id": model.id, "embedding": vector}
|
||||
@@ -161,5 +171,5 @@ async def embed_missing_chunks(session: AsyncSession, client: httpx.AsyncClient,
|
||||
statement = insert(table).values(rows).on_conflict_do_nothing(index_elements=["chunk_id", "model_id"])
|
||||
await session.execute(statement)
|
||||
await session.flush()
|
||||
logger.info(f"ebook_embed_missing_batch_complete {config.embedding_model=} count={len(rows)}")
|
||||
logger.info("ebook_embed_missing_batch_complete model=%s count=%s", config.embedding_model, len(rows))
|
||||
return len(rows)
|
||||
|
||||
@@ -11,7 +11,6 @@ from pathlib import Path
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import tiktoken
|
||||
from anyio import Path as AsyncPath
|
||||
from sqlalchemy import or_, select
|
||||
|
||||
from python.ebook_search.epub_parse import parse_epub
|
||||
@@ -75,18 +74,18 @@ def chunk_text(
|
||||
return [chunk for chunk in chunks if chunk.text]
|
||||
|
||||
|
||||
async def find_library_epubs(library_path: str) -> tuple[AsyncPath, list[AsyncPath] | None]:
|
||||
"""Resolve one configured library path and collect its EPUB files asynchronously.
|
||||
def find_library_epubs(library_path: str) -> tuple[Path, list[Path] | None]:
|
||||
"""Resolve one configured library path and collect its EPUB files (blocking filesystem walk).
|
||||
|
||||
Returns:
|
||||
tuple[Path, list[Path] | None]: The expanded path and its EPUB files, or ``None`` when
|
||||
the path is neither an EPUB file nor a directory.
|
||||
"""
|
||||
path = await AsyncPath(library_path).expanduser()
|
||||
if await path.is_file() and path.suffix.lower() == ".epub":
|
||||
path = Path(library_path).expanduser()
|
||||
if path.is_file() and path.suffix.lower() == ".epub":
|
||||
return path, [path]
|
||||
if await path.is_dir():
|
||||
return path, sorted([epub_path async for epub_path in path.rglob("*.epub")])
|
||||
if path.is_dir():
|
||||
return path, sorted(path.rglob("*.epub"))
|
||||
return path, None
|
||||
|
||||
|
||||
@@ -94,45 +93,44 @@ async def ingest_configured_paths(session: AsyncSession, config: EbookSearchConf
|
||||
"""Ingest every EPUB found under configured library paths."""
|
||||
count = 0
|
||||
for library_path in config.library_paths:
|
||||
path, epub_paths = await find_library_epubs(library_path)
|
||||
logger.info(f"ebook_ingest_path_start {path=}")
|
||||
path, epub_paths = await asyncio.to_thread(find_library_epubs, library_path)
|
||||
logger.info("ebook_ingest_path_start path=%s", path)
|
||||
if epub_paths is None:
|
||||
logger.warning(f"ebook_ingest_path_missing {path=}")
|
||||
logger.warning("ebook_ingest_path_missing path=%s", path)
|
||||
continue
|
||||
for epub_path in epub_paths:
|
||||
count += int(await ingest_file(session, epub_path, config))
|
||||
logger.info(f"ebook_ingest_paths_complete {count=} configured_paths={len(config.library_paths)}")
|
||||
logger.info("ebook_ingest_paths_complete changed_files=%s configured_paths=%s", count, len(config.library_paths))
|
||||
return count
|
||||
|
||||
|
||||
async def resolve_ingest_path(path: Path | AsyncPath) -> AsyncPath:
|
||||
"""Expand and resolve an ingest path without blocking the event loop."""
|
||||
expanded_path = await AsyncPath(path).expanduser()
|
||||
return await expanded_path.resolve()
|
||||
def resolve_ingest_path(path: Path) -> Path:
|
||||
"""Expand and resolve an ingest path (blocking filesystem call)."""
|
||||
return path.expanduser().resolve()
|
||||
|
||||
|
||||
async def ingest_file(session: AsyncSession, path: Path | AsyncPath, config: EbookSearchConfig) -> bool:
|
||||
async def ingest_file(session: AsyncSession, path: Path, config: EbookSearchConfig) -> bool:
|
||||
"""Ingest one EPUB file. Return True when the database changed."""
|
||||
try:
|
||||
resolved_path = await resolve_ingest_path(path)
|
||||
logger.info(f"ebook_ingest_file_start {resolved_path=}")
|
||||
file_hash = await sha256_file(resolved_path)
|
||||
resolved_path = await asyncio.to_thread(resolve_ingest_path, path)
|
||||
logger.info("ebook_ingest_file_start path=%s", resolved_path)
|
||||
file_hash = await asyncio.to_thread(sha256_file, resolved_path)
|
||||
existing = await find_existing_source(session, resolved_path, file_hash)
|
||||
if existing is not None and existing.file_sha256 == file_hash:
|
||||
stat = await resolved_path.stat()
|
||||
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
|
||||
await session.flush()
|
||||
logger.info(f"ebook_ingest_file_unchanged {existing.id=} {resolved_path=}")
|
||||
logger.info("ebook_ingest_file_unchanged source_id=%s path=%s", existing.id, resolved_path)
|
||||
return False
|
||||
if existing is not None:
|
||||
logger.info(f"ebook_ingest_file_replacing {existing.id=} {resolved_path=}")
|
||||
logger.info("ebook_ingest_file_replacing source_id=%s path=%s", existing.id, resolved_path)
|
||||
await session.delete(existing)
|
||||
await session.flush()
|
||||
|
||||
stat = await resolved_path.stat()
|
||||
parsed = await asyncio.to_thread(parse_epub, Path(resolved_path))
|
||||
stat = resolved_path.stat()
|
||||
parsed = await asyncio.to_thread(parse_epub, resolved_path)
|
||||
source = EbookSource(
|
||||
title=parsed.title,
|
||||
author=parsed.author,
|
||||
@@ -159,21 +157,24 @@ async def ingest_file(session: AsyncSession, path: Path | AsyncPath, config: Ebo
|
||||
await session.flush()
|
||||
chunk_index = add_chapter_chunks(session, source, chapter, parsed_chapter, chunk_index, config)
|
||||
|
||||
mention_count = await index_chunk_phrase_mentions_for_book(session, source.id, config)
|
||||
await session.commit()
|
||||
mention_count = await index_chunk_phrase_mentions_for_book(session, source.id, config)
|
||||
logger.info(
|
||||
f"ebook_ingest_file_complete {source.id=} {resolved_path=} chapters={len(parsed.chapters)} {chunk_index=} "
|
||||
f"{mention_count=}"
|
||||
"ebook_ingest_file_complete source_id=%s path=%s chapters=%s chunks=%s phrase_mentions=%s",
|
||||
source.id,
|
||||
resolved_path,
|
||||
len(parsed.chapters),
|
||||
chunk_index,
|
||||
mention_count,
|
||||
)
|
||||
except Exception:
|
||||
await session.rollback()
|
||||
logger.exception(f"ebook_ingest_file_error {path=}")
|
||||
logger.exception(f"ebook_ingest_file_error path={path}")
|
||||
return False
|
||||
else:
|
||||
return True
|
||||
|
||||
|
||||
async def find_existing_source(session: AsyncSession, path: Path | AsyncPath, file_hash: str) -> EbookSource | None:
|
||||
async def find_existing_source(session: AsyncSession, path: Path, file_hash: str) -> EbookSource | None:
|
||||
"""Find an existing source by canonical path or file hash."""
|
||||
return await session.scalar(
|
||||
select(EbookSource).where(or_(EbookSource.file_path == str(path), EbookSource.file_sha256 == file_hash))
|
||||
@@ -212,10 +213,10 @@ def add_chapter_chunks(
|
||||
return chunk_index
|
||||
|
||||
|
||||
async def sha256_file(path: AsyncPath) -> str:
|
||||
"""Calculate the SHA-256 digest for a file without blocking the event loop."""
|
||||
def sha256_file(path: Path) -> str:
|
||||
"""Calculate the SHA-256 digest for a file."""
|
||||
digest = hashlib.sha256()
|
||||
async with await path.open("rb") as file:
|
||||
while block := await file.read(1024 * 1024):
|
||||
with path.open("rb") as file:
|
||||
for block in iter(lambda: file.read(1024 * 1024), b""):
|
||||
digest.update(block)
|
||||
return digest.hexdigest()
|
||||
|
||||
@@ -51,7 +51,10 @@ async def request_embeddings(
|
||||
return embedding_vectors_from_response(response.json())
|
||||
except (httpx.HTTPError, ValueError, KeyError, TypeError) as error:
|
||||
logger.exception(
|
||||
f"ebook_embed_request_failed {config.embedding_base_url=} {config.embedding_model=} count={len(texts)}"
|
||||
"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
|
||||
@@ -64,13 +67,17 @@ async def check_embedding_endpoint(
|
||||
timeout_seconds: float = 5.0,
|
||||
) -> bool:
|
||||
"""Return whether the configured embedding endpoint answers a model listing."""
|
||||
return await _check_endpoint(
|
||||
client,
|
||||
base_url=config.embedding_base_url,
|
||||
api_key=config.embedding_api_key,
|
||||
timeout_seconds=timeout_seconds,
|
||||
unavailable_log=f"ebook_embedding_endpoint_unreachable {config.embedding_base_url=}",
|
||||
)
|
||||
try:
|
||||
response = await client.get(
|
||||
f"{config.embedding_base_url.rstrip('/')}/models",
|
||||
headers=auth_headers(config.embedding_api_key),
|
||||
timeout=timeout_seconds,
|
||||
)
|
||||
response.raise_for_status()
|
||||
except httpx.HTTPError as error:
|
||||
logger.warning("ebook_embedding_endpoint_unreachable base_url=%s error=%s", config.embedding_base_url, error)
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
async def check_chat_endpoint(
|
||||
@@ -80,33 +87,15 @@ async def check_chat_endpoint(
|
||||
timeout_seconds: float = 5.0,
|
||||
) -> bool:
|
||||
"""Return whether the configured chat (answering) endpoint answers a model listing."""
|
||||
return await _check_endpoint(
|
||||
client,
|
||||
base_url=config.vllm_base_url,
|
||||
api_key=config.vllm_api_key,
|
||||
timeout_seconds=timeout_seconds,
|
||||
unavailable_log=f"ebook_chat_endpoint_unreachable {config.vllm_base_url=}",
|
||||
)
|
||||
|
||||
|
||||
async def _check_endpoint(
|
||||
client: httpx.AsyncClient,
|
||||
*,
|
||||
base_url: str,
|
||||
api_key: str,
|
||||
timeout_seconds: float,
|
||||
unavailable_log: str,
|
||||
) -> bool:
|
||||
"""Return whether an OpenAI-compatible endpoint answers a model listing."""
|
||||
try:
|
||||
response = await client.get(
|
||||
f"{base_url.rstrip('/')}/models",
|
||||
headers=auth_headers(api_key),
|
||||
f"{config.vllm_base_url.rstrip('/')}/models",
|
||||
headers=auth_headers(config.vllm_api_key),
|
||||
timeout=timeout_seconds,
|
||||
)
|
||||
response.raise_for_status()
|
||||
except httpx.HTTPError as error:
|
||||
logger.warning(f"{unavailable_log} {error=}")
|
||||
logger.warning("ebook_chat_endpoint_unreachable base_url=%s error=%s", config.vllm_base_url, error)
|
||||
return False
|
||||
return True
|
||||
|
||||
@@ -174,8 +163,6 @@ async def request_chat_completion(
|
||||
client: httpx.AsyncClient,
|
||||
config: EbookSearchConfig,
|
||||
messages: Sequence[dict[str, str]],
|
||||
*,
|
||||
response_format: dict[str, object] | None = None,
|
||||
) -> str:
|
||||
"""Request a chat completion over a shared async client.
|
||||
|
||||
@@ -183,7 +170,6 @@ async def request_chat_completion(
|
||||
client (httpx.AsyncClient): Shared async client whose connection pool bounds concurrency.
|
||||
config (EbookSearchConfig): Runtime settings supplying the endpoint, model, and auth.
|
||||
messages (Sequence[dict[str, str]]): OpenAI-style chat messages.
|
||||
response_format (dict[str, object] | None): Optional OpenAI-compatible structured output constraint.
|
||||
|
||||
Returns:
|
||||
str: The assistant message text.
|
||||
@@ -195,8 +181,11 @@ async def request_chat_completion(
|
||||
response = await client.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}
|
||||
| ({"response_format": response_format} if response_format is not None else {}),
|
||||
json={
|
||||
"model": config.chat_model,
|
||||
"messages": list(messages),
|
||||
"temperature": 0,
|
||||
},
|
||||
timeout=config.chat_timeout_seconds,
|
||||
)
|
||||
response.raise_for_status()
|
||||
|
||||
@@ -112,7 +112,7 @@ async def send_search(client: httpx.AsyncClient, query: str, *, rerank: bool) ->
|
||||
try:
|
||||
response = await client.post("/search", data=data)
|
||||
except httpx.HTTPError as error:
|
||||
logger.warning(f"ebook_loadtest_request_failed {error=}")
|
||||
logger.warning("ebook_loadtest_request_failed error=%s", error)
|
||||
return RequestResult(status_code=0, latency_ms=(time.perf_counter() - start) * 1000, ok=False)
|
||||
return RequestResult(
|
||||
status_code=response.status_code,
|
||||
@@ -192,7 +192,14 @@ def main(
|
||||
"""Load test the search endpoint and report latency and throughput."""
|
||||
configure_logger(log_level)
|
||||
queries = load_queries(queries_file)
|
||||
logger.info(f"ebook_loadtest_start {base_url=} {request_count=} {concurrency=} {rerank=} queries={len(queries)}")
|
||||
logger.info(
|
||||
"ebook_loadtest_start base_url=%s requests=%s concurrency=%s rerank=%s queries=%s",
|
||||
base_url,
|
||||
request_count,
|
||||
concurrency,
|
||||
rerank,
|
||||
len(queries),
|
||||
)
|
||||
summary = asyncio.run(
|
||||
run_load(
|
||||
base_url=base_url,
|
||||
|
||||
@@ -1,8 +0,0 @@
|
||||
"""LLM prompt templates for EPUB search."""
|
||||
|
||||
from python.ebook_search.prompts.lib import (
|
||||
Prompt,
|
||||
load_prompt,
|
||||
)
|
||||
|
||||
__all__ = ["Prompt", "load_prompt"]
|
||||
@@ -1,9 +0,0 @@
|
||||
system = """\
|
||||
Answer only from the provided context. Cite sources with bracketed numbers like [1]. \
|
||||
If the context is insufficient, say so."""
|
||||
user = """\
|
||||
Question:
|
||||
{query}
|
||||
|
||||
Context:
|
||||
{context}"""
|
||||
@@ -1,44 +0,0 @@
|
||||
"""Load and render TOML-backed LLM prompt templates."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import tomllib
|
||||
from dataclasses import dataclass
|
||||
from functools import cache
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Prompt:
|
||||
"""A system and user prompt pair loaded from TOML."""
|
||||
|
||||
system: str
|
||||
user: str
|
||||
|
||||
def messages(self, **values: str) -> list[dict[str, str]]:
|
||||
"""Render this prompt as OpenAI-style chat messages."""
|
||||
return [
|
||||
{"role": "system", "content": self.system.format(**values)},
|
||||
{"role": "user", "content": self.user.format(**values)},
|
||||
]
|
||||
|
||||
|
||||
@cache
|
||||
def _get_prompt_dir() -> Path:
|
||||
"""Return the directory containing prompt template files."""
|
||||
return Path(__file__).resolve().parent
|
||||
|
||||
|
||||
@cache
|
||||
def load_prompt(name: str) -> Prompt:
|
||||
"""Load and validate a named system and user prompt pair from TOML."""
|
||||
path = _get_prompt_dir() / f"{name}.toml"
|
||||
with path.open("rb") as file:
|
||||
body = tomllib.load(file)
|
||||
|
||||
system = body.get("system")
|
||||
user = body.get("user")
|
||||
if not isinstance(system, str) or not isinstance(user, str):
|
||||
msg = f"{path} must define string system and user prompts"
|
||||
raise TypeError(msg)
|
||||
return Prompt(system=system, user=user)
|
||||
@@ -1,10 +0,0 @@
|
||||
system = """\
|
||||
Judge whether a candidate phrase from a book should be protected for RAG retrieval. \
|
||||
Do not extract new phrases. Reject common grammar fragments, ordinary nonspecific \
|
||||
phrases, unstable fragments, and phrases kept only because they are frequent. Keep \
|
||||
people, places, organizations, factions, events, technologies, fictional conditions, \
|
||||
magic systems, formal titles, named concepts, and recurring world-specific terms. \
|
||||
Return only a JSON object with keys: keep, canonical, category, aliases, confidence, \
|
||||
importance, allow_nested, suppress_children, reason."""
|
||||
|
||||
user = "{candidate_json}"
|
||||
@@ -1 +1 @@
|
||||
"""Protected phrase extraction and matching for ebook search."""
|
||||
"""Init."""
|
||||
|
||||
@@ -29,9 +29,39 @@ if TYPE_CHECKING:
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
BAD_START_SCORE_PENALTY = 10.0
|
||||
BAD_END_SCORE_PENALTY = 10.0
|
||||
MULTI_SOURCE_SCORE_BONUS = 2.0
|
||||
MULTI_SOURCE_MIN_SOURCES = 2
|
||||
CAPITALIZED_PHRASE_RE = re.compile(r"\b(?:[A-Z][a-zA-Z']+)(?:\s+(?:of|the|and|in|on|for|[A-Z][a-zA-Z']+)){0,6}")
|
||||
|
||||
|
||||
class SpacySpan(Protocol):
|
||||
"""Small protocol for the spaCy span attributes used by this module."""
|
||||
|
||||
text: str
|
||||
|
||||
|
||||
class SpacyEntity(SpacySpan, Protocol):
|
||||
"""Small protocol for the spaCy entity attributes used by this module."""
|
||||
|
||||
label_: str
|
||||
|
||||
|
||||
class SpacyDoc(Protocol):
|
||||
"""Small protocol for the spaCy doc attributes used by this module."""
|
||||
|
||||
ents: Iterable[SpacyEntity]
|
||||
noun_chunks: Iterable[SpacySpan]
|
||||
|
||||
|
||||
class SpacyLanguage(Protocol):
|
||||
"""Small protocol for a callable spaCy language pipeline."""
|
||||
|
||||
def __call__(self, text: str) -> SpacyDoc:
|
||||
"""Parse text into a spaCy-like doc."""
|
||||
|
||||
|
||||
class YakeExtractor(Protocol):
|
||||
"""Small protocol for the YAKE extractor used by this module."""
|
||||
|
||||
@@ -60,30 +90,40 @@ def normalize_candidate_phrase(
|
||||
phrase_text: str,
|
||||
config: EbookSearchConfig,
|
||||
*,
|
||||
min_tokens: int | None = None,
|
||||
max_tokens: int | None = None,
|
||||
strip_leading_article: bool = False,
|
||||
) -> tuple[str, str, int] | None:
|
||||
"""Normalize a candidate phrase and validate token bounds.
|
||||
|
||||
Args:
|
||||
phrase_text (str): Raw phrase text to normalize.
|
||||
config (EbookSearchConfig): Runtime phrase-tuning settings.
|
||||
min_tokens (int | None): Minimum token count override; defaults to ``config.phrase_min_tokens``.
|
||||
max_tokens (int | None): Maximum token count override; defaults to ``config.phrase_max_tokens``.
|
||||
strip_leading_article (bool): Whether to drop a single leading English article.
|
||||
|
||||
Returns:
|
||||
tuple[str, str, int] | None: Display text, normalized phrase, and token count, or ``None``
|
||||
when the phrase falls outside the token bounds or is ignored.
|
||||
"""
|
||||
normalized_tokens = tokenize_with_offsets(phrase_text)
|
||||
start = 0
|
||||
if strip_leading_article and normalized_tokens and normalized_tokens[0].text in {"the", "a", "an"}:
|
||||
start = 1
|
||||
|
||||
selected_tokens = normalized_tokens[start:]
|
||||
min_count = config.phrase_min_tokens if min_tokens is None else min_tokens
|
||||
max_count = config.phrase_max_tokens if max_tokens is None else max_tokens
|
||||
if len(normalized_tokens) < config.phrase_min_tokens or len(normalized_tokens) > max_count:
|
||||
if len(selected_tokens) < min_count or len(selected_tokens) > max_count:
|
||||
return None
|
||||
|
||||
phrase_norm = " ".join(token.text for token in normalized_tokens)
|
||||
phrase_norm = " ".join(token.text for token in selected_tokens)
|
||||
if phrase_norm in get_ignored_phrases():
|
||||
return None
|
||||
|
||||
display_text = phrase_text[normalized_tokens[0].start_char : normalized_tokens[-1].end_char].strip()
|
||||
return display_text or phrase_norm, phrase_norm, len(normalized_tokens)
|
||||
display_text = phrase_text[selected_tokens[0].start_char : selected_tokens[-1].end_char].strip()
|
||||
return display_text or phrase_norm, phrase_norm, len(selected_tokens)
|
||||
|
||||
|
||||
def count_raw_ngrams(tokens: Sequence[str], config: EbookSearchConfig) -> Counter[str]:
|
||||
@@ -154,7 +194,7 @@ def extract_raw_ngrams_by_chapter(
|
||||
|
||||
|
||||
@lru_cache(maxsize=2)
|
||||
def get_yake_extractor(max_ngram: int, top_k: int, dedup_limit: float) -> KeywordExtractor:
|
||||
def get_yake_extractor(max_ngram: int, top_k: int) -> KeywordExtractor:
|
||||
"""Return a cached YAKE extractor for the given settings.
|
||||
|
||||
Constructing a ``KeywordExtractor`` loads the language's stopword list from disk, so it is
|
||||
@@ -163,32 +203,29 @@ def get_yake_extractor(max_ngram: int, top_k: int, dedup_limit: float) -> Keywor
|
||||
Args:
|
||||
max_ngram (int): Maximum n-gram size to extract.
|
||||
top_k (int): Maximum number of keyphrases to request.
|
||||
dedup_limit (float): Deduplication similarity threshold.
|
||||
|
||||
Returns:
|
||||
KeywordExtractor: A shared extractor instance for the given settings.
|
||||
"""
|
||||
return KeywordExtractor(lan="en", n=max_ngram, dedupLim=dedup_limit, top=top_k)
|
||||
return KeywordExtractor(lan="en", n=max_ngram, dedupLim=0.85, top=top_k)
|
||||
|
||||
|
||||
def extract_yake_candidates(
|
||||
book_text: str,
|
||||
config: EbookSearchConfig,
|
||||
top_k: int = 1000,
|
||||
) -> dict[str, PhraseCandidate]:
|
||||
"""Extract YAKE keyphrases when the optional YAKE package is installed.
|
||||
|
||||
Args:
|
||||
book_text (str): Full book text to extract keyphrases from.
|
||||
config (EbookSearchConfig): Runtime phrase-tuning settings.
|
||||
top_k (int): Maximum number of YAKE keyphrases to request.
|
||||
|
||||
Returns:
|
||||
dict[str, PhraseCandidate]: Candidates keyed by normalized phrase, with YAKE scores.
|
||||
"""
|
||||
extractor = get_yake_extractor(
|
||||
config.phrase_max_tokens,
|
||||
config.phrase_yake_top_k,
|
||||
config.phrase_yake_dedup_limit,
|
||||
)
|
||||
extractor = get_yake_extractor(config.phrase_max_tokens, top_k)
|
||||
out: dict[str, PhraseCandidate] = {}
|
||||
for phrase_text, yake_score in extractor.extract_keywords(book_text):
|
||||
normalized = normalize_candidate_phrase(phrase_text, config)
|
||||
@@ -205,6 +242,55 @@ def extract_yake_candidates(
|
||||
return out
|
||||
|
||||
|
||||
def extract_spacy_candidates(
|
||||
book_text: str,
|
||||
nlp: SpacyLanguage,
|
||||
config: EbookSearchConfig,
|
||||
) -> dict[str, PhraseCandidate]:
|
||||
"""Extract spaCy named entities and noun chunks from one text block.
|
||||
|
||||
Args:
|
||||
book_text (str): Text block to parse with spaCy.
|
||||
nlp (SpacyLanguage): Callable spaCy language pipeline.
|
||||
config (EbookSearchConfig): Runtime phrase-tuning settings.
|
||||
|
||||
Returns:
|
||||
dict[str, PhraseCandidate]: Candidates keyed by normalized phrase from entities and noun chunks.
|
||||
"""
|
||||
out: dict[str, PhraseCandidate] = {}
|
||||
doc = nlp(book_text)
|
||||
|
||||
for ent in doc.ents:
|
||||
normalized = normalize_candidate_phrase(
|
||||
ent.text,
|
||||
config,
|
||||
max_tokens=config.phrase_max_entity_tokens,
|
||||
)
|
||||
if normalized is None:
|
||||
continue
|
||||
phrase_text, phrase_norm, token_count = normalized
|
||||
out[phrase_norm] = PhraseCandidate(
|
||||
phrase_text=phrase_text,
|
||||
phrase_norm=phrase_norm,
|
||||
token_count=token_count,
|
||||
source_spacy_ner=True,
|
||||
spacy_label=ent.label_,
|
||||
)
|
||||
|
||||
for chunk in doc.noun_chunks:
|
||||
normalized = normalize_candidate_phrase(chunk.text, config, strip_leading_article=True)
|
||||
if normalized is None:
|
||||
continue
|
||||
phrase_text, phrase_norm, token_count = normalized
|
||||
out[phrase_norm] = PhraseCandidate(
|
||||
phrase_text=phrase_text,
|
||||
phrase_norm=phrase_norm,
|
||||
token_count=token_count,
|
||||
source_spacy_noun_chunk=True,
|
||||
)
|
||||
return out
|
||||
|
||||
|
||||
def extract_capitalized_phrases(original_text: str, config: EbookSearchConfig) -> dict[str, PhraseCandidate]:
|
||||
"""Extract capitalized phrase runs that often carry fictional terms.
|
||||
|
||||
@@ -306,12 +392,16 @@ def merge_candidate(existing: PhraseCandidate, item: PhraseCandidate) -> None:
|
||||
"""
|
||||
existing.source_raw_ngram = existing.source_raw_ngram or item.source_raw_ngram
|
||||
existing.source_yake = existing.source_yake or item.source_yake
|
||||
existing.source_spacy_ner = existing.source_spacy_ner or item.source_spacy_ner
|
||||
existing.source_spacy_noun_chunk = existing.source_spacy_noun_chunk or item.source_spacy_noun_chunk
|
||||
existing.source_capitalized = existing.source_capitalized or item.source_capitalized
|
||||
existing.source_metadata = existing.source_metadata or item.source_metadata
|
||||
existing.raw_count += item.raw_count
|
||||
existing.chapter_count = max(existing.chapter_count, item.chapter_count)
|
||||
if item.yake_score is not None:
|
||||
existing.yake_score = item.yake_score
|
||||
if item.spacy_label:
|
||||
existing.spacy_label = item.spacy_label
|
||||
|
||||
|
||||
def enrich_with_frequency_and_chapter_counts(
|
||||
@@ -486,14 +576,14 @@ def score_candidate(candidate: PhraseCandidate, config: EbookSearchConfig) -> fl
|
||||
float: Combined score from sources, frequency, and length, less any penalties.
|
||||
"""
|
||||
score = source_score(candidate) + frequency_score(candidate, config) + token_count_score(candidate, config)
|
||||
if non_raw_source_count(candidate) >= config.phrase_multi_source_min_sources:
|
||||
score += config.phrase_multi_source_score_bonus
|
||||
if non_raw_source_count(candidate) >= MULTI_SOURCE_MIN_SOURCES:
|
||||
score += MULTI_SOURCE_SCORE_BONUS
|
||||
if candidate.phrase_norm in get_ignored_phrases():
|
||||
score -= 100.0
|
||||
if has_bad_start(candidate.phrase_norm):
|
||||
score -= config.phrase_bad_start_score_penalty
|
||||
score -= BAD_START_SCORE_PENALTY
|
||||
if has_bad_end(candidate.phrase_norm):
|
||||
score -= config.phrase_bad_end_score_penalty
|
||||
score -= BAD_END_SCORE_PENALTY
|
||||
return score
|
||||
|
||||
|
||||
@@ -509,6 +599,8 @@ def non_raw_source_count(candidate: PhraseCandidate) -> int:
|
||||
return sum(
|
||||
(
|
||||
candidate.source_yake,
|
||||
candidate.source_spacy_ner,
|
||||
candidate.source_spacy_noun_chunk,
|
||||
candidate.source_capitalized,
|
||||
candidate.source_metadata,
|
||||
)
|
||||
@@ -554,6 +646,8 @@ def source_score(candidate: PhraseCandidate) -> float:
|
||||
weight
|
||||
for enabled, weight in (
|
||||
(candidate.source_yake, 2.0),
|
||||
(candidate.source_spacy_ner, 2.5),
|
||||
(candidate.source_spacy_noun_chunk, 1.5),
|
||||
(candidate.source_capitalized, 2.0),
|
||||
(candidate.source_metadata, 2.0),
|
||||
(candidate.source_raw_ngram, 0.5),
|
||||
@@ -644,6 +738,10 @@ def candidate_source_names(candidate: PhraseCandidate) -> list[str]:
|
||||
names.append("raw_ngram")
|
||||
if candidate.source_yake:
|
||||
names.append("yake")
|
||||
if candidate.source_spacy_ner:
|
||||
names.append("spacy_ner")
|
||||
if candidate.source_spacy_noun_chunk:
|
||||
names.append("spacy_noun_chunk")
|
||||
if candidate.source_capitalized:
|
||||
names.append("capitalized")
|
||||
if candidate.source_metadata:
|
||||
@@ -656,14 +754,16 @@ def extract_phrase_candidates_for_book(
|
||||
chapters: Sequence[str],
|
||||
config: EbookSearchConfig,
|
||||
*,
|
||||
nlp: SpacyLanguage | None = None,
|
||||
metadata: Mapping[str, object] | None = None,
|
||||
) -> list[PhraseCandidate]:
|
||||
"""Extract, score, and limit phrase candidates for one book.
|
||||
|
||||
Args:
|
||||
book_text (str): Full book text used for most extraction sources.
|
||||
chapters (Sequence[str]): Chapter-like text blocks used for frequency counts.
|
||||
chapters (Sequence[str]): Chapter-like text blocks used for spaCy and frequency counts.
|
||||
config (EbookSearchConfig): Runtime phrase-tuning settings.
|
||||
nlp (SpacyLanguage | None): Optional spaCy pipeline for entity and noun-chunk sources.
|
||||
metadata (Mapping[str, object] | None): Optional book metadata used as a candidate source.
|
||||
|
||||
Returns:
|
||||
@@ -671,30 +771,47 @@ def extract_phrase_candidates_for_book(
|
||||
"""
|
||||
started_at = perf_counter()
|
||||
logger.info(
|
||||
f"ebook_phrase_candidate_extract_start chapters={len(chapters)} chars={len(book_text)} "
|
||||
f"{config.phrase_min_tokens=} {config.phrase_max_tokens=} {config.protected_phrase_max_candidates_per_book=}"
|
||||
"ebook_phrase_candidate_extract_start chapters=%s chars=%s min_tokens=%s max_tokens=%s max_candidates=%s",
|
||||
len(chapters),
|
||||
len(book_text),
|
||||
config.phrase_min_tokens,
|
||||
config.phrase_max_tokens,
|
||||
config.protected_phrase_max_candidates_per_book,
|
||||
)
|
||||
raw_started_at = perf_counter()
|
||||
raw = extract_raw_ngrams_by_chapter(chapters, config)
|
||||
logger.info(
|
||||
f"ebook_phrase_candidate_extract_raw_complete candidates={len(raw)} "
|
||||
f"duration_ms={(perf_counter() - raw_started_at) * 1000:.1f}"
|
||||
"ebook_phrase_candidate_extract_raw_complete candidates=%s duration_ms=%.1f",
|
||||
len(raw),
|
||||
(perf_counter() - raw_started_at) * 1000,
|
||||
)
|
||||
yake_started_at = perf_counter()
|
||||
yake_candidates = extract_yake_candidates(book_text, config)
|
||||
logger.info(
|
||||
f"ebook_phrase_candidate_extract_yake_complete candidates={len(yake_candidates)} "
|
||||
f"duration_ms={(perf_counter() - yake_started_at) * 1000:.1f}"
|
||||
"ebook_phrase_candidate_extract_yake_complete candidates=%s duration_ms=%.1f",
|
||||
len(yake_candidates),
|
||||
(perf_counter() - yake_started_at) * 1000,
|
||||
)
|
||||
spacy_candidates: dict[str, PhraseCandidate] = {}
|
||||
if nlp is not None:
|
||||
spacy_started_at = perf_counter()
|
||||
for chapter in chapters:
|
||||
spacy_candidates = merge_candidate_sources(spacy_candidates, extract_spacy_candidates(chapter, nlp, config))
|
||||
logger.info(
|
||||
"ebook_phrase_candidate_extract_spacy_complete candidates=%s duration_ms=%.1f",
|
||||
len(spacy_candidates),
|
||||
(perf_counter() - spacy_started_at) * 1000,
|
||||
)
|
||||
capitalized_started_at = perf_counter()
|
||||
capitalized = extract_capitalized_phrases(book_text, config)
|
||||
logger.info(
|
||||
f"ebook_phrase_candidate_extract_capitalized_complete candidates={len(capitalized)} "
|
||||
f"duration_ms={(perf_counter() - capitalized_started_at) * 1000:.1f}"
|
||||
"ebook_phrase_candidate_extract_capitalized_complete candidates=%s duration_ms=%.1f",
|
||||
len(capitalized),
|
||||
(perf_counter() - capitalized_started_at) * 1000,
|
||||
)
|
||||
metadata_candidates = extract_metadata_candidates(metadata, config)
|
||||
|
||||
candidates = merge_candidate_sources(raw, yake_candidates, capitalized, metadata_candidates)
|
||||
candidates = merge_candidate_sources(raw, yake_candidates, spacy_candidates, capitalized, metadata_candidates)
|
||||
enriched_started_at = perf_counter()
|
||||
# Raw n-gram sizes were already counted per chapter above, so only enrich the remaining
|
||||
# (entity-length) sizes here instead of re-sliding every size over the whole book.
|
||||
@@ -714,11 +831,23 @@ def extract_phrase_candidates_for_book(
|
||||
: config.protected_phrase_max_candidates_per_book
|
||||
]
|
||||
logger.info(
|
||||
f"ebook_phrase_candidate_extract_complete raw={len(raw)} yake={len(yake_candidates)} "
|
||||
f"capitalized={len(capitalized)} metadata={len(metadata_candidates)} {pre_filter_count=} {filtered_too_short=} "
|
||||
f"{filtered_too_rare=} {filtered_too_common=} {filtered_junk=} min_uses={minimum_candidate_raw_count(config)} "
|
||||
f"storable={len(candidates)} limited={len(limited)} "
|
||||
f"enrich_score_ms={(perf_counter() - enriched_started_at) * 1000:.1f} "
|
||||
f"duration_ms={(perf_counter() - started_at) * 1000:.1f}"
|
||||
"ebook_phrase_candidate_extract_complete raw=%s yake=%s spacy=%s capitalized=%s metadata=%s "
|
||||
"merged=%s filtered_too_short=%s filtered_too_rare=%s filtered_too_common=%s filtered_junk=%s "
|
||||
"min_uses=%s storable=%s limited=%s enrich_score_ms=%.1f duration_ms=%.1f",
|
||||
len(raw),
|
||||
len(yake_candidates),
|
||||
len(spacy_candidates),
|
||||
len(capitalized),
|
||||
len(metadata_candidates),
|
||||
pre_filter_count,
|
||||
filtered_too_short,
|
||||
filtered_too_rare,
|
||||
filtered_too_common,
|
||||
filtered_junk,
|
||||
minimum_candidate_raw_count(config),
|
||||
len(candidates),
|
||||
len(limited),
|
||||
(perf_counter() - enriched_started_at) * 1000,
|
||||
(perf_counter() - started_at) * 1000,
|
||||
)
|
||||
return limited
|
||||
|
||||
@@ -4,11 +4,11 @@ from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
from collections import deque
|
||||
from time import perf_counter
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from python.ebook_search.protected_phrases.extraction import extract_phrase_candidates_for_book
|
||||
from python.ebook_search.protected_phrases.models import (
|
||||
@@ -16,81 +16,65 @@ from python.ebook_search.protected_phrases.models import (
|
||||
PhraseCandidateGenerationResult,
|
||||
PhraseRecalculationResult,
|
||||
)
|
||||
from python.ebook_search.protected_phrases.pool import get_extraction_pool
|
||||
from python.ebook_search.protected_phrases.pool import extract_phrase_candidates_in_pool, get_extraction_pool
|
||||
from python.ebook_search.protected_phrases.store import (
|
||||
bulk_upsert_unjudged_candidates,
|
||||
delete_phrase_data_for_book,
|
||||
load_book_chapter_texts,
|
||||
metadata_for_source_id,
|
||||
metadata_for_source,
|
||||
new_candidate_row,
|
||||
prune_unstorable_unjudged_candidate_phrases,
|
||||
)
|
||||
from python.orm.common import get_async_postgres_engine
|
||||
from python.orm.richie import EbookSource
|
||||
from python.orm.richie import EbookCandidatePhrase, EbookSource
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from sqlalchemy.ext.asyncio import AsyncEngine
|
||||
from collections.abc import Mapping, Sequence
|
||||
from concurrent.futures import Future
|
||||
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from python.ebook_search.config import EbookSearchConfig
|
||||
from python.ebook_search.protected_phrases.extraction import SpacyLanguage
|
||||
from python.ebook_search.protected_phrases.models import PhraseCandidate
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class BookHasNoChaptersError(ValueError):
|
||||
"""Raised when a book has no indexed chapter text to generate phrases from."""
|
||||
|
||||
|
||||
async def generate_candidate_phrases_for_books(
|
||||
engine: AsyncEngine,
|
||||
session: AsyncSession,
|
||||
config: EbookSearchConfig,
|
||||
*,
|
||||
only_missing: bool = False,
|
||||
) -> PhraseCandidateGenerationResult:
|
||||
"""Create or refresh candidate phrases for indexed books without calling the LLM judge.
|
||||
|
||||
Every book is submitted to the shared process pool up front and runs in parallel across the
|
||||
pool's workers; the call blocks until all books have finished. Each worker opens its own
|
||||
database engine from environment variables, loads the book's chapters, and commits the
|
||||
book's candidates independently.
|
||||
Extraction always runs concurrently in the shared process pool so a full backfill uses
|
||||
multiple cores.
|
||||
|
||||
Args:
|
||||
engine (AsyncEngine): Engine used to read the book list in this process.
|
||||
session (Session): Active database session.
|
||||
config (EbookSearchConfig): Runtime phrase-tuning settings.
|
||||
only_missing (bool): When True, only generate for books that have no candidate phrases
|
||||
yet instead of refreshing every book.
|
||||
|
||||
Returns:
|
||||
PhraseCandidateGenerationResult: Per-corpus counts of books seen, built, and candidates stored.
|
||||
|
||||
Results are collected in book order while the pool keeps working. A book failure (including
|
||||
a book with no indexed chapters) is logged and counted as not built; the remaining books
|
||||
are unaffected.
|
||||
"""
|
||||
async with AsyncSession(engine, expire_on_commit=False) as session:
|
||||
source_query = select(EbookSource.id).order_by(EbookSource.id)
|
||||
source_ids = (await session.scalars(source_query)).all()
|
||||
books_seen = len(source_ids)
|
||||
source_query = select(EbookSource).order_by(EbookSource.id)
|
||||
if only_missing:
|
||||
has_candidates = select(EbookCandidatePhrase.id).where(EbookCandidatePhrase.book_id == EbookSource.id)
|
||||
source_query = source_query.where(~has_candidates.exists())
|
||||
sources = (await session.scalars(source_query)).all()
|
||||
books_seen = len(sources)
|
||||
logger.info(
|
||||
f"ebook_candidate_phrase_generation_start {books_seen=} {config.phrase_min_tokens=} "
|
||||
f"{config.phrase_max_tokens=} {config.protected_phrase_max_candidates_per_book=}"
|
||||
"ebook_candidate_phrase_generation_start books_seen=%s min_tokens=%s max_tokens=%s max_candidates_per_book=%s",
|
||||
books_seen,
|
||||
config.phrase_min_tokens,
|
||||
config.phrase_max_tokens,
|
||||
config.protected_phrase_max_candidates_per_book,
|
||||
)
|
||||
|
||||
pool = get_extraction_pool(config.protected_phrase_extraction_workers)
|
||||
wrapped_futures = [
|
||||
(
|
||||
source_id,
|
||||
asyncio.wrap_future(pool.submit(generate_candidate_phrases_for_book_in_worker, source_id, None, config)),
|
||||
)
|
||||
for source_id in source_ids
|
||||
]
|
||||
outcomes: list[BookCandidateResult] = []
|
||||
for source_id, wrapped_future in wrapped_futures:
|
||||
await asyncio.wait([wrapped_future])
|
||||
exception = wrapped_future.exception()
|
||||
if exception is not None:
|
||||
logger.error(f"ebook_candidate_phrase_generation_book_failed {source_id=}")
|
||||
outcomes.append(BookCandidateResult())
|
||||
continue
|
||||
saved_count = wrapped_future.result()
|
||||
logger.info(f"ebook_candidate_phrase_generation_book_committed {source_id=} {saved_count=}")
|
||||
outcomes.append(BookCandidateResult(candidates=saved_count, built=True))
|
||||
outcomes = await generate_candidates_for_sources_pooled(session, sources, config)
|
||||
|
||||
result = PhraseCandidateGenerationResult(
|
||||
books_seen=books_seen,
|
||||
@@ -98,44 +82,159 @@ async def generate_candidate_phrases_for_books(
|
||||
candidate_phrases=sum(outcome.candidates for outcome in outcomes),
|
||||
)
|
||||
logger.info(
|
||||
f"ebook_candidate_phrase_generation_complete {result.books_seen=} {result.books_built=} "
|
||||
f"{result.candidate_phrases=}"
|
||||
"ebook_candidate_phrase_generation_complete books_seen=%s books_built=%s candidate_total=%s",
|
||||
result.books_seen,
|
||||
result.books_built,
|
||||
result.candidate_phrases,
|
||||
)
|
||||
return result
|
||||
|
||||
|
||||
async def generate_candidates_for_sources_pooled(
|
||||
session: AsyncSession,
|
||||
sources: Sequence[EbookSource],
|
||||
config: EbookSearchConfig,
|
||||
) -> list[BookCandidateResult]:
|
||||
"""Generate candidate phrases for many books, extracting them concurrently in worker processes.
|
||||
|
||||
Chapter loading and row persistence stay on the caller's session (serial), while the CPU-bound
|
||||
extraction runs in the shared process pool. A bounded window of in-flight books overlaps
|
||||
extraction across cores without loading every book's candidates into memory at once.
|
||||
|
||||
Args:
|
||||
session (Session): Active database session.
|
||||
sources (Sequence[EbookSource]): Indexed books to generate candidates for.
|
||||
config (EbookSearchConfig): Runtime phrase-tuning settings.
|
||||
|
||||
Returns:
|
||||
list[BookCandidateResult]: One result per book.
|
||||
"""
|
||||
pool = get_extraction_pool(config.protected_phrase_extraction_workers)
|
||||
max_in_flight = max(1, config.protected_phrase_extraction_workers) * 2
|
||||
pending: deque[tuple[EbookSource, Future[list[PhraseCandidate]]]] = deque()
|
||||
outcomes: list[BookCandidateResult] = []
|
||||
|
||||
async def drain_one() -> None:
|
||||
source, future = pending.popleft()
|
||||
extracted = await asyncio.wrap_future(future)
|
||||
outcomes.append(await store_source_candidates(session, source, extracted, config))
|
||||
|
||||
try:
|
||||
for source in sources:
|
||||
chapters = await load_book_chapter_texts(session, source.id)
|
||||
if not chapters:
|
||||
logger.warning("ebook_candidate_phrase_generation_book_empty source_id=%s", source.id)
|
||||
outcomes.append(BookCandidateResult())
|
||||
continue
|
||||
future = pool.submit(
|
||||
extract_phrase_candidates_for_book,
|
||||
"\n\n".join(chapters),
|
||||
chapters,
|
||||
config,
|
||||
metadata=metadata_for_source(source),
|
||||
)
|
||||
pending.append((source, future))
|
||||
if len(pending) >= max_in_flight:
|
||||
await drain_one()
|
||||
while pending:
|
||||
await drain_one()
|
||||
except Exception:
|
||||
for _, future in pending:
|
||||
future.cancel()
|
||||
await session.rollback()
|
||||
logger.exception("ebook_candidate_phrase_generation_pooled_failed")
|
||||
raise
|
||||
return outcomes
|
||||
|
||||
|
||||
async def store_source_candidates(
|
||||
session: AsyncSession,
|
||||
source: EbookSource,
|
||||
limited_candidates: list[PhraseCandidate],
|
||||
config: EbookSearchConfig,
|
||||
) -> BookCandidateResult:
|
||||
"""Persist and commit one book's already-extracted candidates.
|
||||
|
||||
Args:
|
||||
session (AsyncSession): Active database session.
|
||||
source (EbookSource): Book the candidates belong to.
|
||||
limited_candidates (list[PhraseCandidate]): Scored candidates to persist.
|
||||
config (EbookSearchConfig): Runtime phrase-tuning settings.
|
||||
|
||||
Returns:
|
||||
BookCandidateResult: Candidate count and that the book was committed.
|
||||
"""
|
||||
book_started_at = perf_counter()
|
||||
saved_count = await store_candidate_phrases_for_book(session, source.id, None, limited_candidates, config)
|
||||
await session.commit()
|
||||
logger.info(
|
||||
"ebook_candidate_phrase_generation_book_committed source_id=%s candidates=%s duration_ms=%.1f",
|
||||
source.id,
|
||||
saved_count,
|
||||
(perf_counter() - book_started_at) * 1000,
|
||||
)
|
||||
return BookCandidateResult(candidates=saved_count, built=True)
|
||||
|
||||
|
||||
async def recalculate_candidate_phrases_for_book(
|
||||
session: AsyncSession,
|
||||
source: EbookSource,
|
||||
config: EbookSearchConfig,
|
||||
*,
|
||||
nlp: SpacyLanguage | None = None,
|
||||
use_process_pool: bool = False,
|
||||
) -> PhraseRecalculationResult:
|
||||
"""Remove all book phrase data, regenerate candidates, and commit the completed book.
|
||||
|
||||
Args:
|
||||
session (AsyncSession): Active database session; deletion and regeneration commit on it.
|
||||
session (Session): Active database session.
|
||||
source (EbookSource): Indexed book to recalculate.
|
||||
config (EbookSearchConfig): Runtime phrase-tuning settings.
|
||||
nlp (SpacyLanguage | None): Optional spaCy pipeline for entity and noun-chunk sources.
|
||||
use_process_pool (bool): Run the CPU-bound extraction in a worker process so concurrent
|
||||
recalculations do not serialize behind the GIL. Defaults to in-process for callers
|
||||
(tests, backfills) that do not need it.
|
||||
|
||||
Returns:
|
||||
PhraseRecalculationResult: Deleted-row counts and the number of candidates regenerated.
|
||||
|
||||
Raises:
|
||||
BookHasNoChaptersError: If the book has no indexed chapters. The deletion is rolled
|
||||
back, so the book's existing phrases stay intact.
|
||||
|
||||
The deletion and regeneration share the caller's session, so they commit together; a
|
||||
regeneration failure rolls the deletion back.
|
||||
"""
|
||||
started_at = perf_counter()
|
||||
logger.info(f"ebook_candidate_phrase_recalculation_start {source.id=} {source.title=}")
|
||||
deleted = await delete_phrase_data_for_book(session, source.id)
|
||||
candidate_count = await generate_candidate_phrases_for_book(
|
||||
session,
|
||||
logger.info(
|
||||
"ebook_candidate_phrase_recalculation_start source_id=%s title=%r",
|
||||
source.id,
|
||||
series_id=None,
|
||||
config=config,
|
||||
replace_all=True,
|
||||
source.title,
|
||||
)
|
||||
try:
|
||||
deleted = await delete_phrase_data_for_book(session, source.id)
|
||||
chapters = await load_book_chapter_texts(session, source.id)
|
||||
if not chapters:
|
||||
logger.warning("ebook_candidate_phrase_recalculation_book_empty source_id=%s", source.id)
|
||||
await session.commit()
|
||||
return PhraseRecalculationResult(
|
||||
book_id=source.id,
|
||||
deleted_candidates=deleted.deleted_candidates,
|
||||
deleted_protected_phrases=deleted.deleted_protected_phrases,
|
||||
deleted_aliases=deleted.deleted_aliases,
|
||||
deleted_mentions=deleted.deleted_mentions,
|
||||
candidate_phrases=0,
|
||||
)
|
||||
|
||||
candidate_count = await generate_candidate_phrases_for_book(
|
||||
session,
|
||||
source.id,
|
||||
series_id=None,
|
||||
chapters=chapters,
|
||||
config=config,
|
||||
nlp=nlp,
|
||||
metadata=metadata_for_source(source),
|
||||
replace_all=True,
|
||||
use_process_pool=use_process_pool,
|
||||
)
|
||||
await session.commit()
|
||||
except Exception:
|
||||
await session.rollback()
|
||||
logger.exception("ebook_candidate_phrase_recalculation_failed source_id=%s", source.id)
|
||||
raise
|
||||
|
||||
result = PhraseRecalculationResult(
|
||||
book_id=source.id,
|
||||
@@ -146,106 +245,75 @@ async def recalculate_candidate_phrases_for_book(
|
||||
candidate_phrases=candidate_count,
|
||||
)
|
||||
logger.info(
|
||||
f"ebook_candidate_phrase_recalculation_complete {source.id=} {result.deleted_candidates=} "
|
||||
f"{result.deleted_protected_phrases=} {result.deleted_aliases=} {result.deleted_mentions=} "
|
||||
f"{result.candidate_phrases=} duration_ms={(perf_counter() - started_at) * 1000:.1f}"
|
||||
"ebook_candidate_phrase_recalculation_complete source_id=%s deleted_candidates=%s "
|
||||
"deleted_protected=%s deleted_aliases=%s deleted_mentions=%s candidates=%s duration_ms=%.1f",
|
||||
source.id,
|
||||
result.deleted_candidates,
|
||||
result.deleted_protected_phrases,
|
||||
result.deleted_aliases,
|
||||
result.deleted_mentions,
|
||||
result.candidate_phrases,
|
||||
(perf_counter() - started_at) * 1000,
|
||||
)
|
||||
return result
|
||||
|
||||
|
||||
def generate_candidate_phrases_for_book_in_worker(
|
||||
book_id: int,
|
||||
series_id: int | None,
|
||||
config: EbookSearchConfig,
|
||||
) -> int:
|
||||
"""Run one book's candidate generation in a pooled worker process.
|
||||
|
||||
The worker has no engine or session to inherit (neither can cross process boundaries), so
|
||||
it creates its own engine from environment variables, opens the book's session on it, and
|
||||
disposes the engine once the book is stored.
|
||||
|
||||
Args:
|
||||
book_id (int): Book the candidates belong to.
|
||||
series_id (int | None): Series scope for the stored candidates.
|
||||
config (EbookSearchConfig): Runtime phrase-tuning settings.
|
||||
|
||||
Returns:
|
||||
int: Number of candidate phrase rows stored.
|
||||
"""
|
||||
|
||||
async def generate_with_worker_engine() -> int:
|
||||
engine = get_async_postgres_engine(name="RICHIE", vector_engine=True, pool_size=1)
|
||||
try:
|
||||
async with AsyncSession(engine, expire_on_commit=False) as session:
|
||||
return await generate_candidate_phrases_for_book(
|
||||
session,
|
||||
book_id,
|
||||
series_id,
|
||||
config,
|
||||
)
|
||||
finally:
|
||||
await engine.dispose()
|
||||
|
||||
return asyncio.run(generate_with_worker_engine())
|
||||
|
||||
|
||||
async def generate_candidate_phrases_for_book(
|
||||
session: AsyncSession,
|
||||
book_id: int,
|
||||
series_id: int | None,
|
||||
chapters: Sequence[str],
|
||||
config: EbookSearchConfig,
|
||||
*,
|
||||
nlp: SpacyLanguage | None = None,
|
||||
metadata: Mapping[str, object] | None = None,
|
||||
replace_all: bool = False,
|
||||
use_process_pool: bool = False,
|
||||
) -> int:
|
||||
"""Load a book's chapters and metadata, extract candidate phrases, and store them without LLM judging.
|
||||
|
||||
The session commits only when the whole book succeeds; any failure rolls the session back,
|
||||
which also restores rows the caller deleted in the same transaction (e.g. a recalculation).
|
||||
"""Extract and store candidate phrases for one book without LLM judging.
|
||||
|
||||
Args:
|
||||
session (AsyncSession): Active database session; committed on success, rolled back on failure.
|
||||
session (Session): Active database session.
|
||||
book_id (int): Book the candidates belong to.
|
||||
series_id (int | None): Series scope for the stored candidates.
|
||||
chapters (Sequence[str]): Chapter-like text blocks used for extraction and frequency counts.
|
||||
config (EbookSearchConfig): Runtime phrase-tuning settings.
|
||||
nlp (SpacyLanguage | None): Optional spaCy pipeline for entity and noun-chunk sources.
|
||||
metadata (Mapping[str, object] | None): Optional book metadata used as a candidate source.
|
||||
replace_all (bool): When the caller has already cleared this book's candidates (e.g. a
|
||||
recalculation), skip the per-candidate existence lookup and bulk-insert new rows.
|
||||
use_process_pool (bool): Run the CPU-bound extraction in a worker process to avoid
|
||||
serializing concurrent requests behind the GIL. Ignored when ``nlp`` is set, since
|
||||
the spaCy pipeline cannot be sent to a worker process.
|
||||
|
||||
Returns:
|
||||
int: Number of candidate phrase rows stored.
|
||||
|
||||
Raises:
|
||||
BookHasNoChaptersError: If the book has no indexed chapter text.
|
||||
"""
|
||||
started_at = perf_counter()
|
||||
chapters = await load_book_chapter_texts(session, book_id)
|
||||
if not chapters:
|
||||
await session.rollback()
|
||||
message = f"book {book_id} has no indexed chapters"
|
||||
raise BookHasNoChaptersError(message)
|
||||
metadata = await metadata_for_source_id(session, book_id)
|
||||
try:
|
||||
book_text = "\n\n".join(chapters)
|
||||
candidates = extract_phrase_candidates_for_book(
|
||||
book_text = "\n\n".join(chapters)
|
||||
if use_process_pool and nlp is None:
|
||||
limited_candidates = await extract_phrase_candidates_in_pool(book_text, chapters, config, metadata=metadata)
|
||||
else:
|
||||
limited_candidates = extract_phrase_candidates_for_book(
|
||||
book_text,
|
||||
chapters,
|
||||
config,
|
||||
nlp=nlp,
|
||||
metadata=metadata,
|
||||
)
|
||||
saved_count = await store_candidate_phrases_for_book(
|
||||
session,
|
||||
book_id,
|
||||
series_id,
|
||||
candidates,
|
||||
config,
|
||||
replace_all=replace_all,
|
||||
)
|
||||
await session.commit()
|
||||
except Exception:
|
||||
await session.rollback()
|
||||
raise
|
||||
saved_count = await store_candidate_phrases_for_book(
|
||||
session,
|
||||
book_id,
|
||||
series_id,
|
||||
limited_candidates,
|
||||
config,
|
||||
replace_all=replace_all,
|
||||
)
|
||||
logger.info(
|
||||
f"ebook_candidate_phrase_generation_book_duration {book_id=} {saved_count=} "
|
||||
f"duration_ms={(perf_counter() - started_at) * 1000:.1f}"
|
||||
"ebook_candidate_phrase_generation_book_duration book_id=%s candidates=%s duration_ms=%.1f",
|
||||
book_id,
|
||||
saved_count,
|
||||
(perf_counter() - started_at) * 1000,
|
||||
)
|
||||
return saved_count
|
||||
|
||||
@@ -280,16 +348,23 @@ async def store_candidate_phrases_for_book(
|
||||
await session.flush()
|
||||
saved_count = len(rows)
|
||||
logger.info(
|
||||
f"ebook_candidate_phrase_save_start {book_id=} candidates={len(limited_candidates)} mode=bulk_insert"
|
||||
"ebook_candidate_phrase_save_start book_id=%s candidates=%s mode=bulk_insert",
|
||||
book_id,
|
||||
len(limited_candidates),
|
||||
)
|
||||
else:
|
||||
pruned_count = await prune_unstorable_unjudged_candidate_phrases(session, book_id, config)
|
||||
logger.info(
|
||||
f"ebook_candidate_phrase_save_start {book_id=} candidates={len(limited_candidates)} {pruned_count=}"
|
||||
"ebook_candidate_phrase_save_start book_id=%s candidates=%s pruned_unstorable=%s",
|
||||
book_id,
|
||||
len(limited_candidates),
|
||||
pruned_count,
|
||||
)
|
||||
saved_count = await bulk_upsert_unjudged_candidates(session, book_id, series_id, limited_candidates)
|
||||
logger.info(
|
||||
f"ebook_candidate_phrase_save_complete {book_id=} {saved_count=} "
|
||||
f"save_ms={(perf_counter() - save_started_at) * 1000:.1f}"
|
||||
"ebook_candidate_phrase_save_complete book_id=%s candidates=%s save_ms=%.1f",
|
||||
book_id,
|
||||
saved_count,
|
||||
(perf_counter() - save_started_at) * 1000,
|
||||
)
|
||||
return saved_count
|
||||
|
||||
@@ -5,7 +5,7 @@ from __future__ import annotations
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
from dataclasses import replace
|
||||
import re
|
||||
from time import perf_counter
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
@@ -14,7 +14,6 @@ from sqlalchemy import select
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from python.ebook_search.llm_interface import request_chat_completion
|
||||
from python.ebook_search.prompts import load_prompt
|
||||
from python.ebook_search.protected_phrases.extraction import (
|
||||
candidate_source_names,
|
||||
get_sample_contexts,
|
||||
@@ -45,6 +44,9 @@ if TYPE_CHECKING:
|
||||
from python.ebook_search.protected_phrases.models import PhraseCandidate
|
||||
from python.orm.richie import EbookProtectedPhrase
|
||||
|
||||
JSON_OBJECT_RE = re.compile(r"\{.*\}", re.DOTALL)
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@@ -78,8 +80,12 @@ async def judge_candidate_phrases_for_books(
|
||||
book_workers = max(1, config.phrase_judge_book_workers)
|
||||
phrase_workers = max(1, config.phrase_judge_phrase_workers)
|
||||
logger.info(
|
||||
f"ebook_candidate_phrase_judgment_start {books_seen=} {book_workers=} {phrase_workers=} "
|
||||
f"{config.protected_phrase_confidence_threshold=:.2f}"
|
||||
"ebook_candidate_phrase_judgment_start books_seen=%s book_workers=%s phrase_workers=%s "
|
||||
"confidence_threshold=%.2f",
|
||||
books_seen,
|
||||
book_workers,
|
||||
phrase_workers,
|
||||
config.protected_phrase_confidence_threshold,
|
||||
)
|
||||
|
||||
book_semaphore = asyncio.Semaphore(book_workers)
|
||||
@@ -99,8 +105,14 @@ async def judge_candidate_phrases_for_books(
|
||||
phrase_mentions=sum(outcome.mentions for outcome in outcomes),
|
||||
)
|
||||
logger.info(
|
||||
f"ebook_candidate_phrase_judgment_complete {result.books_seen=} {result.books_judged=} {result.books_failed=} "
|
||||
f"{result.candidates_judged=} {result.protected_phrases=} {result.phrase_mentions=}"
|
||||
"ebook_candidate_phrase_judgment_complete books_seen=%s books_judged=%s books_failed=%s "
|
||||
"candidates_judged=%s protected=%s mentions=%s",
|
||||
result.books_seen,
|
||||
result.books_judged,
|
||||
result.books_failed,
|
||||
result.candidates_judged,
|
||||
result.protected_phrases,
|
||||
result.phrase_mentions,
|
||||
)
|
||||
return result
|
||||
|
||||
@@ -135,7 +147,7 @@ async def judge_one_book_async(
|
||||
return BookJudgmentResult()
|
||||
return await persist_book_judgments(engine, source_id, config, judged)
|
||||
except Exception:
|
||||
logger.exception(f"ebook_candidate_phrase_judgment_book_failed {source_id=}")
|
||||
logger.exception("ebook_candidate_phrase_judgment_book_failed source_id=%s", source_id)
|
||||
return BookJudgmentResult(failed=True)
|
||||
|
||||
|
||||
@@ -161,7 +173,7 @@ async def prepare_book_judgment(
|
||||
return None
|
||||
async with AsyncSession(engine) as session:
|
||||
if not await count_unjudged_candidates(session, source_id, config):
|
||||
logger.info(f"ebook_candidate_phrase_judgment_book_skip_no_unjudged {source_id=}")
|
||||
logger.info("ebook_candidate_phrase_judgment_book_skip_no_unjudged source_id=%s", source_id)
|
||||
return None
|
||||
existing_protected = await count_protected_phrases(session, source_id)
|
||||
target_remaining: int | None = None
|
||||
@@ -169,13 +181,15 @@ async def prepare_book_judgment(
|
||||
target_remaining = max(config.phrase_target_protected_per_book - existing_protected, 0)
|
||||
if target_remaining == 0:
|
||||
logger.info(
|
||||
f"ebook_candidate_phrase_judgment_skipped_target_met {source_id=} {existing_protected=} "
|
||||
f"{config.phrase_target_protected_per_book=}"
|
||||
"ebook_candidate_phrase_judgment_skipped_target_met source_id=%s existing_protected=%s target=%s",
|
||||
source_id,
|
||||
existing_protected,
|
||||
config.phrase_target_protected_per_book,
|
||||
)
|
||||
return None
|
||||
book_text = await load_book_text(session, source_id)
|
||||
if not book_text:
|
||||
logger.warning(f"ebook_candidate_phrase_judgment_book_empty {source_id=}")
|
||||
logger.warning("ebook_candidate_phrase_judgment_book_empty source_id=%s", source_id)
|
||||
return None
|
||||
normalized_book_text = normalize_text(book_text)
|
||||
# Stored rows may predate the current junk filters and score weights, so re-filter and
|
||||
@@ -197,8 +211,15 @@ async def prepare_book_judgment(
|
||||
normalized_book_text, candidate.phrase_norm
|
||||
)
|
||||
logger.info(
|
||||
f"ebook_candidate_phrase_judgment_candidates_loaded {source_id=} candidates={len(work_items)} {skipped_junk=} "
|
||||
f"unjudged_rows={len(rows)} {existing_protected=} {target_remaining=} {judgment_limit=}"
|
||||
"ebook_candidate_phrase_judgment_candidates_loaded source_id=%s candidates=%s skipped_junk=%s "
|
||||
"unjudged_rows=%s existing_protected=%s target_remaining=%s judgment_limit=%s",
|
||||
source_id,
|
||||
len(work_items),
|
||||
skipped_junk,
|
||||
len(rows),
|
||||
existing_protected,
|
||||
target_remaining,
|
||||
judgment_limit,
|
||||
)
|
||||
return work_items, target_remaining
|
||||
|
||||
@@ -259,9 +280,7 @@ async def judge_candidate_async(
|
||||
Returns:
|
||||
LLMJudgment: The parsed judgment.
|
||||
"""
|
||||
content = await request_chat_completion(
|
||||
client, config, build_judge_messages(candidate), response_format={"type": "json_object"}
|
||||
)
|
||||
content = await request_chat_completion(client, config, build_judge_messages(candidate))
|
||||
return parse_llm_judgment(content, config)
|
||||
|
||||
|
||||
@@ -286,39 +305,39 @@ async def persist_book_judgments(
|
||||
book_started_at = perf_counter()
|
||||
async with AsyncSession(engine, expire_on_commit=False) as session:
|
||||
try:
|
||||
normalized_book_text = normalize_text(await load_book_text(session, source_id))
|
||||
protected: list[EbookProtectedPhrase] = []
|
||||
for candidate_id, candidate, judgment, promote in judged:
|
||||
filtered_judgment = replace(
|
||||
judgment,
|
||||
aliases=tuple(
|
||||
alias for alias in judgment.aliases if alias_occurs_in_book(alias, normalized_book_text)
|
||||
),
|
||||
)
|
||||
candidate_row = await save_candidate_to_db(
|
||||
session, source_id, None, candidate, judgment=filtered_judgment
|
||||
)
|
||||
candidate_row = await save_candidate_to_db(session, source_id, None, candidate, judgment=judgment)
|
||||
if promote:
|
||||
protected.append(
|
||||
await upsert_protected_phrase(
|
||||
session, source_id, None, candidate, filtered_judgment, candidate_row
|
||||
)
|
||||
await upsert_protected_phrase(session, source_id, None, candidate, judgment, candidate_row)
|
||||
)
|
||||
logger.info(
|
||||
f"ebook_candidate_phrase_judgment_candidate_complete {source_id=} {candidate_id=} "
|
||||
f"{candidate.phrase_norm=} {judgment.keep=} {judgment.confidence=:.3f} {judgment.category=} "
|
||||
f"{promote=}"
|
||||
"ebook_candidate_phrase_judgment_candidate_complete source_id=%s candidate_id=%s phrase=%r "
|
||||
"keep=%s confidence=%.3f category=%r promoted=%s",
|
||||
source_id,
|
||||
candidate_id,
|
||||
candidate.phrase_norm,
|
||||
judgment.keep,
|
||||
judgment.confidence,
|
||||
judgment.category,
|
||||
promote,
|
||||
)
|
||||
await session.flush()
|
||||
mentions = await index_chunk_phrase_mentions_for_book(session, source_id, config) if protected else 0
|
||||
await session.commit()
|
||||
except Exception:
|
||||
await session.rollback()
|
||||
logger.exception(f"ebook_candidate_phrase_judgment_book_persist_failed {source_id=}")
|
||||
logger.exception("ebook_candidate_phrase_judgment_book_persist_failed source_id=%s", source_id)
|
||||
return BookJudgmentResult(failed=True)
|
||||
logger.info(
|
||||
f"ebook_candidate_phrase_judgment_book_committed {source_id=} judged={len(judged)} protected={len(protected)} "
|
||||
f"{mentions=} duration_ms={(perf_counter() - book_started_at) * 1000:.1f}"
|
||||
"ebook_candidate_phrase_judgment_book_committed source_id=%s judged=%s protected=%s mentions=%s "
|
||||
"duration_ms=%.1f",
|
||||
source_id,
|
||||
len(judged),
|
||||
len(protected),
|
||||
mentions,
|
||||
(perf_counter() - book_started_at) * 1000,
|
||||
)
|
||||
return BookJudgmentResult(judged=len(judged), protected=len(protected), mentions=mentions, committed=True)
|
||||
|
||||
@@ -350,14 +369,24 @@ def should_protect_judged_candidate(
|
||||
accepted_token_count = len(accepted_tokens)
|
||||
if accepted_token_count < config.phrase_min_tokens:
|
||||
logger.info(
|
||||
f"ebook_candidate_phrase_judgment_candidate_skip_short_canonical {book_id=} {candidate_id=} "
|
||||
f"{candidate.phrase_norm=} {accepted_norm=} {accepted_token_count=} {config.phrase_min_tokens=}"
|
||||
"ebook_candidate_phrase_judgment_candidate_skip_short_canonical book_id=%s candidate_id=%s "
|
||||
"phrase=%r canonical=%r token_count=%s min_tokens=%s",
|
||||
book_id,
|
||||
candidate_id,
|
||||
candidate.phrase_norm,
|
||||
accepted_norm,
|
||||
accepted_token_count,
|
||||
config.phrase_min_tokens,
|
||||
)
|
||||
return False
|
||||
if is_most_common_word_phrase(accepted_tokens):
|
||||
logger.info(
|
||||
f"ebook_candidate_phrase_judgment_candidate_skip_common_canonical {book_id=} {candidate_id=} "
|
||||
f"{candidate.phrase_norm=} {accepted_norm=}"
|
||||
"ebook_candidate_phrase_judgment_candidate_skip_common_canonical book_id=%s candidate_id=%s "
|
||||
"phrase=%r canonical=%r",
|
||||
book_id,
|
||||
candidate_id,
|
||||
candidate.phrase_norm,
|
||||
accepted_norm,
|
||||
)
|
||||
return False
|
||||
return True
|
||||
@@ -380,7 +409,20 @@ def build_judge_messages(candidate: PhraseCandidate) -> list[dict[str, str]]:
|
||||
"chapter_count": candidate.chapter_count,
|
||||
"contexts": candidate.sample_contexts,
|
||||
}
|
||||
return load_prompt("phrase_judge").messages(candidate_json=json.dumps(payload, ensure_ascii=True))
|
||||
return [
|
||||
{
|
||||
"role": "system",
|
||||
"content": (
|
||||
"Judge whether a candidate phrase from a book should be protected for RAG retrieval. "
|
||||
"Do not extract new phrases. Reject common grammar fragments, ordinary nonspecific phrases, "
|
||||
"unstable fragments, and phrases kept only because they are frequent. Keep people, places, "
|
||||
"organizations, factions, events, technologies, fictional conditions, magic systems, formal titles, "
|
||||
"named concepts, and recurring world-specific terms. Return only a JSON object with keys: keep, "
|
||||
"canonical, category, aliases, confidence, importance, allow_nested, suppress_children, reason."
|
||||
),
|
||||
},
|
||||
{"role": "user", "content": json.dumps(payload, ensure_ascii=True)},
|
||||
]
|
||||
|
||||
|
||||
def parse_llm_judgment(content: str, config: EbookSearchConfig) -> LLMJudgment:
|
||||
@@ -405,14 +447,14 @@ def parse_llm_judgment(content: str, config: EbookSearchConfig) -> LLMJudgment:
|
||||
if not isinstance(aliases, list | tuple):
|
||||
aliases = ()
|
||||
return LLMJudgment(
|
||||
keep=strict_bool(body.get("keep"), default=False),
|
||||
keep=bool(body.get("keep", False)),
|
||||
canonical=optional_text(body.get("canonical")),
|
||||
category=optional_text(body.get("category")),
|
||||
aliases=tuple(str(alias) for alias in aliases if isinstance(alias, str) and alias.strip()),
|
||||
confidence=clamped_float(body.get("confidence"), default=0.0),
|
||||
importance=clamped_float(body.get("importance"), default=0.5),
|
||||
allow_nested=strict_bool(body.get("allow_nested"), default=config.phrase_default_allow_nested),
|
||||
suppress_children=strict_bool(body.get("suppress_children"), default=config.phrase_default_suppress_children),
|
||||
allow_nested=bool(body.get("allow_nested", config.phrase_default_allow_nested)),
|
||||
suppress_children=bool(body.get("suppress_children", config.phrase_default_suppress_children)),
|
||||
reason=optional_text(body.get("reason")),
|
||||
)
|
||||
|
||||
@@ -429,13 +471,14 @@ def extract_json_object(content: str) -> str:
|
||||
Raises:
|
||||
ValueError: If no JSON object is found in the response.
|
||||
"""
|
||||
return content.strip()
|
||||
|
||||
|
||||
def alias_occurs_in_book(alias: str, normalized_book_text: str) -> bool:
|
||||
"""Return whether a normalized alias occurs as a complete phrase in the source book."""
|
||||
alias_norm = normalize_text(alias)
|
||||
return bool(alias_norm) and f" {alias_norm} " in f" {normalized_book_text} "
|
||||
stripped = content.strip()
|
||||
if stripped.startswith("{") and stripped.endswith("}"):
|
||||
return stripped
|
||||
match = JSON_OBJECT_RE.search(stripped)
|
||||
if match is None:
|
||||
msg = "LLM phrase judge response did not contain a JSON object"
|
||||
raise ValueError(msg)
|
||||
return match.group(0)
|
||||
|
||||
|
||||
def optional_text(value: object) -> str | None:
|
||||
@@ -453,11 +496,6 @@ def optional_text(value: object) -> str | None:
|
||||
return stripped or None
|
||||
|
||||
|
||||
def strict_bool(value: object, *, default: bool) -> bool:
|
||||
"""Return a JSON boolean, falling back when the value has another type."""
|
||||
return value if isinstance(value, bool) else default
|
||||
|
||||
|
||||
def clamped_float(value: object, *, default: float) -> float:
|
||||
"""Coerce a JSON number into the 0.0 to 1.0 range.
|
||||
|
||||
|
||||
@@ -6,11 +6,12 @@ import logging
|
||||
from collections import defaultdict
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from sqlalchemy import and_, delete, func, or_, select, union
|
||||
from sqlalchemy import and_, delete, func, or_, select
|
||||
|
||||
from python.ebook_search.protected_phrases.config import get_ignored_phrases
|
||||
from python.ebook_search.protected_phrases.models import (
|
||||
ChunkPhraseHit,
|
||||
HydratedPhraseMatch,
|
||||
PhraseLookup,
|
||||
PhraseMatch,
|
||||
)
|
||||
@@ -28,107 +29,11 @@ if TYPE_CHECKING:
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from python.ebook_search.config import EbookSearchConfig
|
||||
from python.ebook_search.protected_phrases.text_normalization import NormalizedToken
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
async def detect_protected_phrases_for_query(
|
||||
session: AsyncSession,
|
||||
query_text: str,
|
||||
config: EbookSearchConfig,
|
||||
) -> list[PhraseMatch]:
|
||||
"""Find query phrases with indexed exact matches on canonical and alias norms.
|
||||
|
||||
Args:
|
||||
session (AsyncSession): Active database session.
|
||||
query_text (str): User query text to detect phrases in.
|
||||
config (EbookSearchConfig): Runtime phrase-tuning settings.
|
||||
|
||||
Returns:
|
||||
list[PhraseMatch]: Metadata-backed, overlap-resolved phrase matches for the query.
|
||||
"""
|
||||
tokens_ = tokenize_with_offsets(query_text)
|
||||
windows_by_norm: defaultdict[str, list[tuple[int, int]]] = defaultdict(list)
|
||||
token_texts = [token.text for token in tokens_]
|
||||
max_tokens = max(config.phrase_max_tokens, config.phrase_max_entity_tokens)
|
||||
for phrase_norm, start, end in generate_query_ngrams(
|
||||
token_texts,
|
||||
min_n=config.phrase_min_tokens,
|
||||
max_n=max_tokens,
|
||||
):
|
||||
windows_by_norm[phrase_norm].append((start, end))
|
||||
if not windows_by_norm:
|
||||
return []
|
||||
|
||||
query_norms = tuple(windows_by_norm)
|
||||
matched_norms = union(
|
||||
select(
|
||||
EbookProtectedPhrase.id.label("phrase_id"),
|
||||
EbookProtectedPhrase.phrase_norm.label("matched_norm"),
|
||||
).where(EbookProtectedPhrase.phrase_norm.in_(query_norms)),
|
||||
select(
|
||||
EbookPhraseAlias.phrase_id.label("phrase_id"),
|
||||
EbookPhraseAlias.alias_norm.label("matched_norm"),
|
||||
).where(EbookPhraseAlias.alias_norm.in_(query_norms)),
|
||||
).subquery()
|
||||
statement = select(EbookProtectedPhrase, matched_norms.c.matched_norm).join(
|
||||
matched_norms,
|
||||
matched_norms.c.phrase_id == EbookProtectedPhrase.id,
|
||||
)
|
||||
|
||||
matches: list[PhraseMatch] = []
|
||||
for phrase, matched_norm in await session.execute(statement):
|
||||
for start, end in windows_by_norm[matched_norm]:
|
||||
matches.append(
|
||||
PhraseMatch(
|
||||
phrase_id=phrase.id,
|
||||
matched_norm=matched_norm,
|
||||
phrase_text=phrase.phrase_text,
|
||||
phrase_norm=phrase.phrase_norm,
|
||||
canonical_id=phrase.canonical_id,
|
||||
phrase_type=phrase.phrase_type,
|
||||
token_count=end - start,
|
||||
confidence=phrase.confidence,
|
||||
importance=phrase.importance,
|
||||
allow_nested=phrase.allow_nested,
|
||||
suppress_children=phrase.suppress_children,
|
||||
start_token=start,
|
||||
end_token=end,
|
||||
start_char=tokens_[start].start_char,
|
||||
end_char=tokens_[end - 1].end_char,
|
||||
book_id=phrase.book_id,
|
||||
series_id=phrase.series_id,
|
||||
)
|
||||
)
|
||||
return resolve_overlaps(matches)
|
||||
|
||||
|
||||
async def index_chunk_phrase_mentions_for_book(
|
||||
session: AsyncSession,
|
||||
book_id: int,
|
||||
config: EbookSearchConfig,
|
||||
) -> int:
|
||||
"""Rebuild chunk phrase mentions for all chunks in one book.
|
||||
|
||||
Args:
|
||||
session (AsyncSession): Active database session.
|
||||
book_id (int): Book whose chunk mentions are rebuilt.
|
||||
config (EbookSearchConfig): Runtime phrase-tuning settings.
|
||||
|
||||
Returns:
|
||||
int: Total number of chunk phrase mentions indexed for the book.
|
||||
"""
|
||||
lookup = await load_phrase_lookup(session, config, book_id=book_id)
|
||||
await session.execute(delete(EbookChunkPhraseMention).where(EbookChunkPhraseMention.book_id == book_id))
|
||||
chunks = await session.scalars(select(EbookChunk).where(EbookChunk.source_id == book_id).order_by(EbookChunk.id))
|
||||
count = 0
|
||||
for chunk in chunks:
|
||||
count += index_chunk_phrase_mentions(session, chunk, lookup=lookup)
|
||||
await session.flush()
|
||||
logger.info(f"ebook_chunk_phrase_mentions_indexed {book_id=} {count=}")
|
||||
return count
|
||||
|
||||
|
||||
async def load_phrase_lookup(
|
||||
session: AsyncSession,
|
||||
config: EbookSearchConfig,
|
||||
@@ -147,29 +52,40 @@ async def load_phrase_lookup(
|
||||
Returns:
|
||||
PhraseLookup: Normalized phrase and alias maps with the token-window bounds to test.
|
||||
"""
|
||||
phrase_ids_by_norm: defaultdict[str, set[int]] = defaultdict(set)
|
||||
phrases_by_id: dict[int, EbookProtectedPhrase] = {}
|
||||
norm_to_ids: defaultdict[str, list[int]] = defaultdict(list)
|
||||
alias_to_ids: defaultdict[str, list[int]] = defaultdict(list)
|
||||
max_tokens = config.phrase_max_tokens
|
||||
|
||||
statement = select(
|
||||
EbookProtectedPhrase,
|
||||
EbookPhraseAlias.alias_norm,
|
||||
).outerjoin(EbookPhraseAlias, EbookPhraseAlias.phrase_id == EbookProtectedPhrase.id)
|
||||
phrase_statement = select(
|
||||
EbookProtectedPhrase.id,
|
||||
EbookProtectedPhrase.phrase_norm,
|
||||
EbookProtectedPhrase.token_count,
|
||||
)
|
||||
scope_filter = protected_phrase_scope_filter(book_id=book_id, series_id=series_id)
|
||||
if scope_filter is not None:
|
||||
statement = statement.where(scope_filter)
|
||||
phrase_statement = phrase_statement.where(scope_filter)
|
||||
|
||||
for phrase, alias_norm in await session.execute(statement):
|
||||
phrases_by_id[phrase.id] = phrase
|
||||
phrase_ids_by_norm[phrase.phrase_norm].add(phrase.id)
|
||||
max_tokens = max(max_tokens, phrase.token_count)
|
||||
if alias_norm is not None:
|
||||
phrase_ids_by_norm[alias_norm].add(phrase.id)
|
||||
max_tokens = max(max_tokens, len(alias_norm.split()))
|
||||
for row in await session.execute(phrase_statement):
|
||||
phrase_id = int(row.id)
|
||||
phrase_norm = str(row.phrase_norm)
|
||||
norm_to_ids[phrase_norm].append(phrase_id)
|
||||
max_tokens = max(max_tokens, int(row.token_count))
|
||||
|
||||
alias_statement = select(
|
||||
EbookPhraseAlias.alias_norm,
|
||||
EbookPhraseAlias.phrase_id,
|
||||
).join(EbookProtectedPhrase, EbookProtectedPhrase.id == EbookPhraseAlias.phrase_id)
|
||||
if scope_filter is not None:
|
||||
alias_statement = alias_statement.where(scope_filter)
|
||||
|
||||
for row in await session.execute(alias_statement):
|
||||
alias_norm = str(row.alias_norm)
|
||||
alias_to_ids[alias_norm].append(int(row.phrase_id))
|
||||
max_tokens = max(max_tokens, len(alias_norm.split()))
|
||||
|
||||
return PhraseLookup(
|
||||
phrase_ids_by_norm={key: tuple(sorted(values)) for key, values in phrase_ids_by_norm.items()},
|
||||
phrases_by_id=phrases_by_id,
|
||||
norm_to_phrase_ids={key: tuple(values) for key, values in norm_to_ids.items()},
|
||||
alias_to_phrase_ids={key: tuple(values) for key, values in alias_to_ids.items()},
|
||||
min_tokens=config.phrase_min_tokens,
|
||||
max_tokens=max_tokens,
|
||||
)
|
||||
@@ -195,144 +111,6 @@ def protected_phrase_scope_filter(*, book_id: int | None, series_id: int | None)
|
||||
return and_(*conditions)
|
||||
|
||||
|
||||
def is_inside(child: PhraseMatch, parent: PhraseMatch) -> bool:
|
||||
"""Return whether one token span is strictly inside another.
|
||||
|
||||
Args:
|
||||
child (PhraseMatch): Candidate nested match.
|
||||
parent (PhraseMatch): Candidate enclosing match.
|
||||
|
||||
Returns:
|
||||
bool: True when ``child`` lies within ``parent`` and is not the same span.
|
||||
"""
|
||||
return (
|
||||
child.start_token >= parent.start_token
|
||||
and child.end_token <= parent.end_token
|
||||
and (child.start_token, child.end_token, child.phrase_id)
|
||||
!= (parent.start_token, parent.end_token, parent.phrase_id)
|
||||
)
|
||||
|
||||
|
||||
def index_chunk_phrase_mentions(session: AsyncSession, chunk: EbookChunk, *, lookup: PhraseLookup) -> int:
|
||||
"""Store protected phrase mentions for one chunk.
|
||||
|
||||
Args:
|
||||
session (AsyncSession): Active database session.
|
||||
chunk (EbookChunk): Chunk whose text is scanned for phrase mentions.
|
||||
lookup (PhraseLookup): In-memory phrase and alias lookup maps.
|
||||
|
||||
Returns:
|
||||
int: Number of phrase mentions stored for the chunk.
|
||||
"""
|
||||
tokens_ = tokenize_with_offsets(chunk.text)
|
||||
token_texts = [token.text for token in tokens_]
|
||||
raw_matches: list[PhraseMatch] = []
|
||||
phrase_windows = generate_query_ngrams(token_texts, min_n=lookup.min_tokens, max_n=lookup.max_tokens)
|
||||
for matched_norm, start, end in phrase_windows:
|
||||
for phrase_id in lookup.phrase_ids_by_norm.get(matched_norm, ()):
|
||||
phrase = lookup.phrases_by_id[phrase_id]
|
||||
raw_matches.append(
|
||||
PhraseMatch(
|
||||
phrase_id=phrase_id,
|
||||
matched_norm=matched_norm,
|
||||
phrase_text=phrase.phrase_text,
|
||||
phrase_norm=phrase.phrase_norm,
|
||||
canonical_id=phrase.canonical_id,
|
||||
phrase_type=phrase.phrase_type,
|
||||
confidence=phrase.confidence,
|
||||
importance=phrase.importance,
|
||||
allow_nested=phrase.allow_nested,
|
||||
suppress_children=phrase.suppress_children,
|
||||
start_token=start,
|
||||
end_token=end,
|
||||
token_count=end - start,
|
||||
start_char=tokens_[start].start_char,
|
||||
end_char=tokens_[end - 1].end_char,
|
||||
book_id=phrase.book_id,
|
||||
series_id=phrase.series_id,
|
||||
)
|
||||
)
|
||||
|
||||
matches = resolve_overlaps(raw_matches)
|
||||
for match in matches:
|
||||
session.add(
|
||||
EbookChunkPhraseMention(
|
||||
chunk_id=chunk.id,
|
||||
phrase_id=match.phrase_id,
|
||||
book_id=match.book_id if match.book_id is not None else chunk.source_id,
|
||||
series_id=match.series_id,
|
||||
start_char=match.start_char if match.start_char is not None else 0,
|
||||
end_char=match.end_char,
|
||||
)
|
||||
)
|
||||
return len(matches)
|
||||
|
||||
|
||||
def resolve_overlaps(matches: Sequence[PhraseMatch]) -> list[PhraseMatch]:
|
||||
"""Resolve overlapping phrase matches without relying only on longest match.
|
||||
|
||||
Args:
|
||||
matches (Sequence[PhraseMatch]): Metadata-backed matches that may overlap.
|
||||
|
||||
Returns:
|
||||
list[PhraseMatch]: The kept, non-suppressed matches.
|
||||
"""
|
||||
sorted_matches = sorted(
|
||||
matches,
|
||||
key=lambda match: (match.start_token, -match.token_count, -match.importance, -match.confidence),
|
||||
)
|
||||
kept: list[PhraseMatch] = []
|
||||
for candidate in sorted_matches:
|
||||
if any(should_suppress(candidate, existing) for existing in kept):
|
||||
continue
|
||||
kept.append(candidate)
|
||||
return kept
|
||||
|
||||
|
||||
def should_suppress(candidate: PhraseMatch, kept: PhraseMatch) -> bool:
|
||||
"""Return whether an already-kept match should suppress a candidate.
|
||||
|
||||
Args:
|
||||
candidate (PhraseMatch): Match being considered for keeping.
|
||||
kept (PhraseMatch): Match already kept that may suppress the candidate.
|
||||
|
||||
Returns:
|
||||
bool: True when the candidate should be dropped in favor of the kept match.
|
||||
"""
|
||||
if not overlaps(candidate, kept):
|
||||
return False
|
||||
if candidate.canonical_id == kept.canonical_id:
|
||||
return rank_match(kept) >= rank_match(candidate)
|
||||
if is_inside(candidate, kept) and kept.suppress_children and not candidate.allow_nested:
|
||||
return True
|
||||
return not candidate.allow_nested and rank_match(kept) > rank_match(candidate)
|
||||
|
||||
|
||||
def overlaps(first: PhraseMatch, second: PhraseMatch) -> bool:
|
||||
"""Return whether two token spans overlap.
|
||||
|
||||
Args:
|
||||
first (PhraseMatch): First match to compare.
|
||||
second (PhraseMatch): Second match to compare.
|
||||
|
||||
Returns:
|
||||
bool: True when the two token spans share at least one token position.
|
||||
"""
|
||||
return not (first.end_token <= second.start_token or first.start_token >= second.end_token)
|
||||
|
||||
|
||||
def rank_match(match: PhraseMatch) -> tuple[float, float, int]:
|
||||
"""Rank phrase matches by importance, confidence, then token count.
|
||||
|
||||
Args:
|
||||
match (PhraseMatch): Match to build a sort key for.
|
||||
|
||||
Returns:
|
||||
tuple[float, float, int]: A comparable key of importance, confidence, and token count.
|
||||
"""
|
||||
return (match.importance, match.confidence, match.token_count)
|
||||
|
||||
|
||||
def generate_query_ngrams(
|
||||
tokens_: Sequence[str],
|
||||
min_n: int,
|
||||
@@ -358,13 +136,301 @@ def generate_query_ngrams(
|
||||
yield phrase_norm, start, end
|
||||
|
||||
|
||||
def detect_phrase_candidates(query_text: str, lookup: PhraseLookup) -> list[PhraseMatch]:
|
||||
"""Detect protected phrase windows in a user query using RAM hash lookups.
|
||||
|
||||
Args:
|
||||
query_text (str): User query text to scan.
|
||||
lookup (PhraseLookup): In-memory phrase and alias lookup maps.
|
||||
|
||||
Returns:
|
||||
list[PhraseMatch]: Unhydrated phrase matches found in the query.
|
||||
"""
|
||||
return detect_phrase_candidates_from_tokens(tokenize_with_offsets(query_text), lookup)
|
||||
|
||||
|
||||
def detect_phrase_candidates_in_text(text: str, lookup: PhraseLookup) -> list[PhraseMatch]:
|
||||
"""Detect protected phrase windows in arbitrary text with character offsets.
|
||||
|
||||
Args:
|
||||
text (str): Arbitrary text, such as a chunk, to scan.
|
||||
lookup (PhraseLookup): In-memory phrase and alias lookup maps.
|
||||
|
||||
Returns:
|
||||
list[PhraseMatch]: Unhydrated phrase matches found in the text.
|
||||
"""
|
||||
return detect_phrase_candidates_from_tokens(tokenize_with_offsets(text), lookup)
|
||||
|
||||
|
||||
def detect_phrase_candidates_from_tokens(tokens_: Sequence[NormalizedToken], lookup: PhraseLookup) -> list[PhraseMatch]:
|
||||
"""Detect protected phrase windows from already-normalized tokens.
|
||||
|
||||
Args:
|
||||
tokens_ (Sequence[NormalizedToken]): Normalized tokens with character offsets.
|
||||
lookup (PhraseLookup): In-memory phrase and alias lookup maps.
|
||||
|
||||
Returns:
|
||||
list[PhraseMatch]: Deduplicated unhydrated phrase matches with token and character spans.
|
||||
"""
|
||||
matches: list[PhraseMatch] = []
|
||||
seen: set[tuple[int | None, str, int, int]] = set()
|
||||
token_texts = [token.text for token in tokens_]
|
||||
for phrase_norm, start, end in generate_query_ngrams(token_texts, min_n=lookup.min_tokens, max_n=lookup.max_tokens):
|
||||
phrase_ids = lookup.norm_to_phrase_ids.get(phrase_norm, ())
|
||||
alias_ids = lookup.alias_to_phrase_ids.get(phrase_norm, ())
|
||||
for phrase_id in (*phrase_ids, *alias_ids):
|
||||
key = (phrase_id, phrase_norm, start, end)
|
||||
if key in seen:
|
||||
continue
|
||||
seen.add(key)
|
||||
matches.append(
|
||||
PhraseMatch(
|
||||
phrase_norm=phrase_norm,
|
||||
phrase_id=phrase_id,
|
||||
start_token=start,
|
||||
end_token=end,
|
||||
token_count=end - start,
|
||||
start_char=tokens_[start].start_char,
|
||||
end_char=tokens_[end - 1].end_char,
|
||||
)
|
||||
)
|
||||
return matches
|
||||
|
||||
|
||||
async def hydrate_matches(session: AsyncSession, matches: Sequence[PhraseMatch]) -> list[HydratedPhraseMatch]:
|
||||
"""Fetch protected phrase metadata for raw phrase matches.
|
||||
|
||||
Args:
|
||||
session (AsyncSession): Active database session.
|
||||
matches (Sequence[PhraseMatch]): Unhydrated matches to enrich.
|
||||
|
||||
Returns:
|
||||
list[HydratedPhraseMatch]: Matches with protected-phrase metadata attached.
|
||||
"""
|
||||
if not matches:
|
||||
return []
|
||||
|
||||
phrase_ids = sorted({match.phrase_id for match in matches if match.phrase_id is not None})
|
||||
if not phrase_ids:
|
||||
return []
|
||||
|
||||
rows = {
|
||||
row.id: row
|
||||
for row in await session.scalars(select(EbookProtectedPhrase).where(EbookProtectedPhrase.id.in_(phrase_ids)))
|
||||
}
|
||||
hydrated: list[HydratedPhraseMatch] = []
|
||||
for match in matches:
|
||||
if match.phrase_id is None:
|
||||
continue
|
||||
phrase = rows.get(match.phrase_id)
|
||||
if phrase is None:
|
||||
continue
|
||||
hydrated.append(
|
||||
HydratedPhraseMatch(
|
||||
phrase_id=phrase.id,
|
||||
matched_norm=match.phrase_norm,
|
||||
phrase_text=phrase.phrase_text,
|
||||
phrase_norm=phrase.phrase_norm,
|
||||
canonical_id=phrase.canonical_id,
|
||||
phrase_type=phrase.phrase_type,
|
||||
token_count=match.token_count,
|
||||
confidence=phrase.confidence,
|
||||
importance=phrase.importance,
|
||||
allow_nested=phrase.allow_nested,
|
||||
suppress_children=phrase.suppress_children,
|
||||
start_token=match.start_token,
|
||||
end_token=match.end_token,
|
||||
start_char=match.start_char,
|
||||
end_char=match.end_char,
|
||||
book_id=phrase.book_id,
|
||||
series_id=phrase.series_id,
|
||||
)
|
||||
)
|
||||
return hydrated
|
||||
|
||||
|
||||
def overlaps(first: HydratedPhraseMatch, second: HydratedPhraseMatch) -> bool:
|
||||
"""Return whether two token spans overlap.
|
||||
|
||||
Args:
|
||||
first (HydratedPhraseMatch): First match to compare.
|
||||
second (HydratedPhraseMatch): Second match to compare.
|
||||
|
||||
Returns:
|
||||
bool: True when the two token spans share at least one token position.
|
||||
"""
|
||||
return not (first.end_token <= second.start_token or first.start_token >= second.end_token)
|
||||
|
||||
|
||||
def is_inside(child: HydratedPhraseMatch, parent: HydratedPhraseMatch) -> bool:
|
||||
"""Return whether one token span is strictly inside another.
|
||||
|
||||
Args:
|
||||
child (HydratedPhraseMatch): Candidate nested match.
|
||||
parent (HydratedPhraseMatch): Candidate enclosing match.
|
||||
|
||||
Returns:
|
||||
bool: True when ``child`` lies within ``parent`` and is not the same span.
|
||||
"""
|
||||
return (
|
||||
child.start_token >= parent.start_token
|
||||
and child.end_token <= parent.end_token
|
||||
and (child.start_token, child.end_token, child.phrase_id)
|
||||
!= (parent.start_token, parent.end_token, parent.phrase_id)
|
||||
)
|
||||
|
||||
|
||||
def rank_match(match: HydratedPhraseMatch) -> tuple[float, float, int]:
|
||||
"""Rank phrase matches by importance, confidence, then token count.
|
||||
|
||||
Args:
|
||||
match (HydratedPhraseMatch): Match to build a sort key for.
|
||||
|
||||
Returns:
|
||||
tuple[float, float, int]: A comparable key of importance, confidence, and token count.
|
||||
"""
|
||||
return (match.importance, match.confidence, match.token_count)
|
||||
|
||||
|
||||
def should_suppress(candidate: HydratedPhraseMatch, kept: HydratedPhraseMatch) -> bool:
|
||||
"""Return whether an already-kept match should suppress a candidate.
|
||||
|
||||
Args:
|
||||
candidate (HydratedPhraseMatch): Match being considered for keeping.
|
||||
kept (HydratedPhraseMatch): Match already kept that may suppress the candidate.
|
||||
|
||||
Returns:
|
||||
bool: True when the candidate should be dropped in favor of the kept match.
|
||||
"""
|
||||
if not overlaps(candidate, kept):
|
||||
return False
|
||||
if candidate.canonical_id == kept.canonical_id:
|
||||
return rank_match(kept) >= rank_match(candidate)
|
||||
if is_inside(candidate, kept) and kept.suppress_children and not candidate.allow_nested:
|
||||
return True
|
||||
return not candidate.allow_nested and rank_match(kept) > rank_match(candidate)
|
||||
|
||||
|
||||
def resolve_overlaps(matches: Sequence[HydratedPhraseMatch]) -> list[HydratedPhraseMatch]:
|
||||
"""Resolve overlapping phrase matches without relying only on longest match.
|
||||
|
||||
Args:
|
||||
matches (Sequence[HydratedPhraseMatch]): Hydrated matches that may overlap.
|
||||
|
||||
Returns:
|
||||
list[HydratedPhraseMatch]: The kept, non-suppressed matches.
|
||||
"""
|
||||
sorted_matches = sorted(
|
||||
matches,
|
||||
key=lambda match: (match.start_token, -match.token_count, -match.importance, -match.confidence),
|
||||
)
|
||||
kept: list[HydratedPhraseMatch] = []
|
||||
for candidate in sorted_matches:
|
||||
if any(should_suppress(candidate, existing) for existing in kept):
|
||||
continue
|
||||
kept.append(candidate)
|
||||
return kept
|
||||
|
||||
|
||||
async def detect_protected_phrases_for_query(
|
||||
session: AsyncSession,
|
||||
query_text: str,
|
||||
config: EbookSearchConfig,
|
||||
*,
|
||||
lookup: PhraseLookup | None = None,
|
||||
book_id: int | None = None,
|
||||
series_id: int | None = None,
|
||||
) -> list[HydratedPhraseMatch]:
|
||||
"""Run the full online protected-phrase query-detection pipeline.
|
||||
|
||||
Args:
|
||||
session (AsyncSession): Active database session.
|
||||
query_text (str): User query text to detect phrases in.
|
||||
config (EbookSearchConfig): Runtime phrase-tuning settings.
|
||||
lookup (PhraseLookup | None): Optional preloaded lookup; loaded on demand when ``None``.
|
||||
book_id (int | None): Optional book scope for lookup loading.
|
||||
series_id (int | None): Optional series scope for lookup loading.
|
||||
|
||||
Returns:
|
||||
list[HydratedPhraseMatch]: Hydrated, overlap-resolved phrase matches for the query.
|
||||
"""
|
||||
active_lookup = (
|
||||
lookup
|
||||
if lookup is not None
|
||||
else await load_phrase_lookup(session, config, book_id=book_id, series_id=series_id)
|
||||
)
|
||||
return resolve_overlaps(await hydrate_matches(session, detect_phrase_candidates(query_text, active_lookup)))
|
||||
|
||||
|
||||
async def index_chunk_phrase_mentions_for_book(
|
||||
session: AsyncSession,
|
||||
book_id: int,
|
||||
config: EbookSearchConfig,
|
||||
*,
|
||||
series_id: int | None = None,
|
||||
lookup: PhraseLookup | None = None,
|
||||
) -> int:
|
||||
"""Rebuild chunk phrase mentions for all chunks in one book.
|
||||
|
||||
Args:
|
||||
session (AsyncSession): Active database session.
|
||||
book_id (int): Book whose chunk mentions are rebuilt.
|
||||
config (EbookSearchConfig): Runtime phrase-tuning settings.
|
||||
series_id (int | None): Optional series scope for lookup loading.
|
||||
lookup (PhraseLookup | None): Optional preloaded lookup; loaded on demand when ``None``.
|
||||
|
||||
Returns:
|
||||
int: Total number of chunk phrase mentions indexed for the book.
|
||||
"""
|
||||
active_lookup = (
|
||||
lookup
|
||||
if lookup is not None
|
||||
else await load_phrase_lookup(session, config, book_id=book_id, series_id=series_id)
|
||||
)
|
||||
await session.execute(delete(EbookChunkPhraseMention).where(EbookChunkPhraseMention.book_id == book_id))
|
||||
chunks = await session.scalars(select(EbookChunk).where(EbookChunk.source_id == book_id).order_by(EbookChunk.id))
|
||||
count = 0
|
||||
for chunk in chunks:
|
||||
count += await index_chunk_phrase_mentions(session, chunk, lookup=active_lookup)
|
||||
await session.flush()
|
||||
logger.info("ebook_chunk_phrase_mentions_indexed book_id=%s mentions=%s", book_id, count)
|
||||
return count
|
||||
|
||||
|
||||
async def index_chunk_phrase_mentions(session: AsyncSession, chunk: EbookChunk, *, lookup: PhraseLookup) -> int:
|
||||
"""Store protected phrase mentions for one chunk.
|
||||
|
||||
Args:
|
||||
session (AsyncSession): Active database session.
|
||||
chunk (EbookChunk): Chunk whose text is scanned for phrase mentions.
|
||||
lookup (PhraseLookup): In-memory phrase and alias lookup maps.
|
||||
|
||||
Returns:
|
||||
int: Number of phrase mentions stored for the chunk.
|
||||
"""
|
||||
raw_matches = detect_phrase_candidates_in_text(chunk.text, lookup)
|
||||
hydrated = resolve_overlaps(await hydrate_matches(session, raw_matches))
|
||||
for match in hydrated:
|
||||
session.add(
|
||||
EbookChunkPhraseMention(
|
||||
chunk_id=chunk.id,
|
||||
phrase_id=match.phrase_id,
|
||||
book_id=match.book_id if match.book_id is not None else chunk.source_id,
|
||||
series_id=match.series_id,
|
||||
start_char=match.start_char if match.start_char is not None else 0,
|
||||
end_char=match.end_char,
|
||||
)
|
||||
)
|
||||
return len(hydrated)
|
||||
|
||||
|
||||
async def phrase_hits_for_chunks(
|
||||
session: AsyncSession,
|
||||
*,
|
||||
chunk_ids: Sequence[int],
|
||||
phrase_ids: Sequence[int],
|
||||
) -> dict[int, tuple[ChunkPhraseHit, ...]]:
|
||||
"""Return matched protected phrases with mention counts by chunk id.
|
||||
"""Return matched protected phrases with mention counts by chunk id using indexed chunk mentions.
|
||||
|
||||
Args:
|
||||
session (AsyncSession): Active database session.
|
||||
@@ -395,11 +461,31 @@ async def phrase_hits_for_chunks(
|
||||
)
|
||||
hits: defaultdict[int, list[ChunkPhraseHit]] = defaultdict(list)
|
||||
for row in await session.execute(statement):
|
||||
hits[row.chunk_id].append(
|
||||
hits[int(row.chunk_id)].append(
|
||||
ChunkPhraseHit(
|
||||
phrase_id=row.phrase_id,
|
||||
phrase_text=row.phrase_text,
|
||||
mention_count=row.mention_count,
|
||||
phrase_id=int(row.phrase_id),
|
||||
phrase_text=str(row.phrase_text),
|
||||
mention_count=int(row.mention_count),
|
||||
)
|
||||
)
|
||||
return {chunk_id: tuple(chunk_hits) for chunk_id, chunk_hits in hits.items()}
|
||||
|
||||
|
||||
async def phrase_hit_counts_for_chunks(
|
||||
session: AsyncSession,
|
||||
*,
|
||||
chunk_ids: Sequence[int],
|
||||
phrase_ids: Sequence[int],
|
||||
) -> dict[int, int]:
|
||||
"""Return phrase-hit counts by chunk id using indexed chunk mentions.
|
||||
|
||||
Args:
|
||||
session (AsyncSession): Active database session.
|
||||
chunk_ids (Sequence[int]): Chunk ids to count mentions for.
|
||||
phrase_ids (Sequence[int]): Protected phrase ids to restrict the counts to.
|
||||
|
||||
Returns:
|
||||
dict[int, int]: Total mention count per chunk id.
|
||||
"""
|
||||
hits = await phrase_hits_for_chunks(session, chunk_ids=chunk_ids, phrase_ids=phrase_ids)
|
||||
return {chunk_id: sum(hit.mention_count for hit in chunk_hits) for chunk_id, chunk_hits in hits.items()}
|
||||
|
||||
@@ -8,8 +8,6 @@ from typing import TYPE_CHECKING
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Mapping
|
||||
|
||||
from python.orm.richie import EbookProtectedPhrase
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class PhraseCandidate:
|
||||
@@ -21,8 +19,11 @@ class PhraseCandidate:
|
||||
token_count (int): Number of normalized tokens in the phrase.
|
||||
source_raw_ngram (bool): Whether the raw n-gram extractor produced the phrase.
|
||||
source_yake (bool): Whether YAKE keyword extraction produced the phrase.
|
||||
source_spacy_ner (bool): Whether spaCy named-entity recognition produced the phrase.
|
||||
source_spacy_noun_chunk (bool): Whether spaCy noun chunking produced the phrase.
|
||||
source_capitalized (bool): Whether the capitalized-run extractor produced the phrase.
|
||||
source_metadata (bool): Whether book metadata produced the phrase.
|
||||
spacy_label (str | None): spaCy entity label when NER produced the phrase.
|
||||
raw_count (int): Occurrences counted across the book text.
|
||||
chapter_count (int): Number of chapters containing the phrase.
|
||||
yake_score (float | None): Raw YAKE score when available; lower is better.
|
||||
@@ -35,8 +36,11 @@ class PhraseCandidate:
|
||||
token_count: int
|
||||
source_raw_ngram: bool = False
|
||||
source_yake: bool = False
|
||||
source_spacy_ner: bool = False
|
||||
source_spacy_noun_chunk: bool = False
|
||||
source_capitalized: bool = False
|
||||
source_metadata: bool = False
|
||||
spacy_label: str | None = None
|
||||
raw_count: int = 0
|
||||
chapter_count: int = 0
|
||||
yake_score: float | None = None
|
||||
@@ -73,24 +77,47 @@ class LLMJudgment:
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class PhraseLookup:
|
||||
"""In-memory phrase metadata used for constant-time text-window checks.
|
||||
"""In-memory lookup maps used for constant-time phrase-window checks.
|
||||
|
||||
Attributes:
|
||||
phrase_ids_by_norm (Mapping[str, tuple[int, ...]]): Canonical and alias norms to phrase ids.
|
||||
phrases_by_id (Mapping[int, EbookProtectedPhrase]): Protected phrase metadata by id.
|
||||
norm_to_phrase_ids (Mapping[str, tuple[int, ...]]): Normalized phrase to protected phrase ids.
|
||||
alias_to_phrase_ids (Mapping[str, tuple[int, ...]]): Normalized alias to protected phrase ids.
|
||||
min_tokens (int): Smallest token-window size to test.
|
||||
max_tokens (int): Largest token-window size to test.
|
||||
"""
|
||||
|
||||
phrase_ids_by_norm: Mapping[str, tuple[int, ...]]
|
||||
phrases_by_id: Mapping[int, EbookProtectedPhrase]
|
||||
norm_to_phrase_ids: Mapping[str, tuple[int, ...]]
|
||||
alias_to_phrase_ids: Mapping[str, tuple[int, ...]]
|
||||
min_tokens: int
|
||||
max_tokens: int
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class PhraseMatch:
|
||||
"""A detected phrase match with protected-phrase metadata attached.
|
||||
"""An unhydrated query or chunk phrase match.
|
||||
|
||||
Attributes:
|
||||
phrase_norm (str): Normalized text of the matched window.
|
||||
start_token (int): Index of the first matched token.
|
||||
end_token (int): Index one past the last matched token.
|
||||
token_count (int): Number of tokens in the match.
|
||||
phrase_id (int | None): Matched protected phrase id when known.
|
||||
start_char (int | None): Start character offset in the source text.
|
||||
end_char (int | None): End character offset in the source text.
|
||||
"""
|
||||
|
||||
phrase_norm: str
|
||||
start_token: int
|
||||
end_token: int
|
||||
token_count: int
|
||||
phrase_id: int | None = None
|
||||
start_char: int | None = None
|
||||
end_char: int | None = None
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class HydratedPhraseMatch:
|
||||
"""A phrase match with protected-phrase metadata attached.
|
||||
|
||||
Attributes:
|
||||
phrase_id (int): Protected phrase id.
|
||||
@@ -137,8 +164,8 @@ class ChunkPhraseHit:
|
||||
|
||||
Attributes:
|
||||
phrase_id (int): Protected phrase id.
|
||||
phrase_text (str): Display text for the phrase.
|
||||
mention_count (int): Indexed mentions inside the chunk.
|
||||
phrase_text (str): Display text of the protected phrase.
|
||||
mention_count (int): Indexed mentions of the phrase in the chunk.
|
||||
"""
|
||||
|
||||
phrase_id: int
|
||||
|
||||
@@ -9,11 +9,21 @@ or server threads.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import multiprocessing
|
||||
import os
|
||||
from concurrent.futures import ProcessPoolExecutor
|
||||
from threading import Lock
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from python.ebook_search.protected_phrases.extraction import extract_phrase_candidates_for_book
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Mapping, Sequence
|
||||
|
||||
from python.ebook_search.config import EbookSearchConfig
|
||||
from python.ebook_search.protected_phrases.models import PhraseCandidate
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -45,7 +55,7 @@ def get_extraction_pool(max_workers: int) -> ProcessPoolExecutor:
|
||||
max_workers=workers,
|
||||
mp_context=multiprocessing.get_context("spawn"),
|
||||
)
|
||||
logger.info(f"ebook_phrase_extraction_pool_started {workers=}")
|
||||
logger.info("ebook_phrase_extraction_pool_started workers=%s", workers)
|
||||
return _extraction_pool.pool
|
||||
|
||||
|
||||
@@ -56,3 +66,36 @@ def shutdown_extraction_pool() -> None:
|
||||
_extraction_pool.pool.shutdown(wait=False, cancel_futures=True)
|
||||
_extraction_pool.pool = None
|
||||
logger.info("ebook_phrase_extraction_pool_shutdown")
|
||||
|
||||
|
||||
async def extract_phrase_candidates_in_pool(
|
||||
book_text: str,
|
||||
chapters: Sequence[str],
|
||||
config: EbookSearchConfig,
|
||||
*,
|
||||
metadata: Mapping[str, object] | None,
|
||||
) -> list[PhraseCandidate]:
|
||||
"""Run book phrase extraction in a worker process and await the result.
|
||||
|
||||
Only the CPU-bound extraction runs in the worker; the caller keeps all database work in the
|
||||
request process. The spaCy pipeline is not supported here because it is not picklable, so
|
||||
this always runs the non-spaCy extraction path.
|
||||
|
||||
Args:
|
||||
book_text (str): Full book text used for extraction.
|
||||
chapters (Sequence[str]): Chapter-like text blocks used for frequency counts.
|
||||
config (EbookSearchConfig): Runtime phrase-tuning settings.
|
||||
metadata (Mapping[str, object] | None): Optional book metadata used as a candidate source.
|
||||
|
||||
Returns:
|
||||
list[PhraseCandidate]: Scored candidates sorted best-first and capped per book.
|
||||
"""
|
||||
pool = get_extraction_pool(config.protected_phrase_extraction_workers)
|
||||
future = pool.submit(
|
||||
extract_phrase_candidates_for_book,
|
||||
book_text,
|
||||
list(chapters),
|
||||
config,
|
||||
metadata=dict(metadata) if metadata is not None else None,
|
||||
)
|
||||
return await asyncio.wrap_future(future)
|
||||
|
||||
@@ -7,7 +7,8 @@ import re
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from sqlalchemy import delete, func, or_, select
|
||||
from sqlalchemy.dialects.postgresql import insert
|
||||
from sqlalchemy.dialects.postgresql import insert as pg_insert
|
||||
from sqlalchemy.dialects.sqlite import insert as sqlite_insert
|
||||
|
||||
from python.ebook_search.protected_phrases.extraction import minimum_candidate_raw_count
|
||||
from python.ebook_search.protected_phrases.models import (
|
||||
@@ -28,14 +29,35 @@ from python.orm.richie import (
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Sequence
|
||||
|
||||
from sqlalchemy.dialects.postgresql.dml import Insert as PostgresInsert
|
||||
from sqlalchemy.dialects.sqlite.dml import Insert as SqliteInsert
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from python.ebook_search.config import EbookSearchConfig
|
||||
from python.ebook_search.protected_phrases.models import LLMJudgment
|
||||
from python.orm.richie.base import TableBase
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def dialect_insert(session: AsyncSession, table: type[TableBase]) -> PostgresInsert | SqliteInsert:
|
||||
"""Return a dialect-specific INSERT construct that supports ``ON CONFLICT DO UPDATE``.
|
||||
|
||||
Production runs on PostgreSQL while tests run on SQLite; both support upserts with
|
||||
compatible SQLAlchemy constructs, so the correct one is chosen from the bound dialect.
|
||||
|
||||
Args:
|
||||
session (AsyncSession): Active database session whose bind selects the dialect.
|
||||
table (type[TableBase]): Mapped table to insert into.
|
||||
|
||||
Returns:
|
||||
PostgresInsert | SqliteInsert: A dialect insert exposing ``on_conflict_do_update``.
|
||||
"""
|
||||
if session.get_bind().dialect.name == "sqlite":
|
||||
return sqlite_insert(table)
|
||||
return pg_insert(table)
|
||||
|
||||
|
||||
async def load_book_text(session: AsyncSession, book_id: int) -> str:
|
||||
"""Load a book's indexed chunk text as one string for phrase extraction.
|
||||
|
||||
@@ -273,8 +295,11 @@ def phrase_candidate_from_row(row: EbookCandidatePhrase) -> PhraseCandidate:
|
||||
token_count=row.token_count,
|
||||
source_raw_ngram=row.source_raw_ngram,
|
||||
source_yake=row.source_yake,
|
||||
source_spacy_ner=row.source_spacy_ner,
|
||||
source_spacy_noun_chunk=row.source_spacy_noun_chunk,
|
||||
source_capitalized=row.source_capitalized,
|
||||
source_metadata=row.source_metadata,
|
||||
spacy_label=row.spacy_label,
|
||||
raw_count=row.raw_count,
|
||||
chapter_count=row.chapter_count,
|
||||
yake_score=row.yake_score,
|
||||
@@ -309,8 +334,11 @@ def candidate_row_values(
|
||||
"token_count": candidate.token_count,
|
||||
"source_raw_ngram": candidate.source_raw_ngram,
|
||||
"source_yake": candidate.source_yake,
|
||||
"source_spacy_ner": candidate.source_spacy_ner,
|
||||
"source_spacy_noun_chunk": candidate.source_spacy_noun_chunk,
|
||||
"source_capitalized": candidate.source_capitalized,
|
||||
"source_metadata": candidate.source_metadata,
|
||||
"spacy_label": candidate.spacy_label,
|
||||
"raw_count": candidate.raw_count,
|
||||
"chapter_count": candidate.chapter_count,
|
||||
"yake_score": candidate.yake_score,
|
||||
@@ -355,7 +383,7 @@ async def save_candidate_to_db(
|
||||
skip_update = {"book_id", "phrase_norm"}
|
||||
if judgment is None:
|
||||
skip_update.add("llm_judged")
|
||||
insert_statement = insert(EbookCandidatePhrase).values(**values)
|
||||
insert_statement = dialect_insert(session, EbookCandidatePhrase).values(**values)
|
||||
statement = insert_statement.on_conflict_do_update(
|
||||
index_elements=["book_id", "phrase_norm"],
|
||||
set_={column: insert_statement.excluded[column] for column in values if column not in skip_update},
|
||||
@@ -401,7 +429,7 @@ async def bulk_upsert_unjudged_candidates(
|
||||
skip_update = {"book_id", "phrase_norm", "llm_judged"}
|
||||
for chunk_start in range(0, len(values), BULK_CANDIDATE_UPSERT_CHUNK):
|
||||
chunk = values[chunk_start : chunk_start + BULK_CANDIDATE_UPSERT_CHUNK]
|
||||
insert_statement = insert(EbookCandidatePhrase).values(chunk)
|
||||
insert_statement = dialect_insert(session, EbookCandidatePhrase).values(chunk)
|
||||
statement = insert_statement.on_conflict_do_update(
|
||||
index_elements=["book_id", "phrase_norm"],
|
||||
set_={column: insert_statement.excluded[column] for column in chunk[0] if column not in skip_update},
|
||||
@@ -432,8 +460,11 @@ def new_candidate_row(book_id: int, series_id: int | None, candidate: PhraseCand
|
||||
row.token_count = candidate.token_count
|
||||
row.source_raw_ngram = candidate.source_raw_ngram
|
||||
row.source_yake = candidate.source_yake
|
||||
row.source_spacy_ner = candidate.source_spacy_ner
|
||||
row.source_spacy_noun_chunk = candidate.source_spacy_noun_chunk
|
||||
row.source_capitalized = candidate.source_capitalized
|
||||
row.source_metadata = candidate.source_metadata
|
||||
row.spacy_label = candidate.spacy_label
|
||||
row.raw_count = candidate.raw_count
|
||||
row.chapter_count = candidate.chapter_count
|
||||
row.yake_score = candidate.yake_score
|
||||
@@ -487,7 +518,7 @@ async def upsert_protected_phrase(
|
||||
"suppress_children": judgment.suppress_children,
|
||||
"source_candidate_id": source_candidate.id,
|
||||
}
|
||||
insert_statement = insert(EbookProtectedPhrase).values(**values)
|
||||
insert_statement = dialect_insert(session, EbookProtectedPhrase).values(**values)
|
||||
statement = insert_statement.on_conflict_do_update(
|
||||
index_elements=["book_id", "phrase_norm"],
|
||||
set_={
|
||||
@@ -520,7 +551,7 @@ async def upsert_phrase_alias(
|
||||
if not alias_norm or alias_norm == phrase.phrase_norm:
|
||||
return None
|
||||
|
||||
insert_statement = insert(EbookPhraseAlias).values(
|
||||
insert_statement = dialect_insert(session, EbookPhraseAlias).values(
|
||||
phrase_id=phrase.id,
|
||||
alias_norm=alias_norm,
|
||||
alias_text=alias_text,
|
||||
@@ -593,8 +624,11 @@ async def prune_unstorable_unjudged_candidate_phrases(
|
||||
)
|
||||
if deleted:
|
||||
logger.info(
|
||||
f"ebook_candidate_phrase_unstorable_pruned {book_id=} {deleted=} {config.phrase_min_tokens=} "
|
||||
f"min_uses={minimum_candidate_raw_count(config)}"
|
||||
"ebook_candidate_phrase_unstorable_pruned book_id=%s deleted=%s min_tokens=%s min_uses=%s",
|
||||
book_id,
|
||||
deleted,
|
||||
config.phrase_min_tokens,
|
||||
minimum_candidate_raw_count(config),
|
||||
)
|
||||
return deleted
|
||||
|
||||
@@ -636,8 +670,12 @@ async def delete_phrase_data_for_book(session: AsyncSession, book_id: int) -> Ph
|
||||
)
|
||||
await session.flush()
|
||||
logger.info(
|
||||
f"ebook_candidate_phrase_data_deleted {book_id=} {deleted_candidates=} {deleted_protected=} {deleted_aliases=} "
|
||||
f"{deleted_mentions=}"
|
||||
"ebook_candidate_phrase_data_deleted book_id=%s candidates=%s protected=%s aliases=%s mentions=%s",
|
||||
book_id,
|
||||
deleted_candidates,
|
||||
deleted_protected,
|
||||
deleted_aliases,
|
||||
deleted_mentions,
|
||||
)
|
||||
return PhraseRecalculationResult(
|
||||
book_id=book_id,
|
||||
|
||||
@@ -35,7 +35,12 @@ async def rerank_chunks(
|
||||
if not candidates:
|
||||
return []
|
||||
|
||||
logger.info(f"ebook_rerank_request_start {config.base_url=} {config.model=} candidates={len(candidates)}")
|
||||
logger.info(
|
||||
"ebook_rerank_request_start base_url=%s model=%s candidates=%s",
|
||||
config.base_url,
|
||||
config.model,
|
||||
len(candidates),
|
||||
)
|
||||
scores = await score_candidates(client, query, candidates, config)
|
||||
results = sorted(
|
||||
(
|
||||
@@ -49,7 +54,12 @@ async def rerank_chunks(
|
||||
key=lambda result: result.score,
|
||||
reverse=True,
|
||||
)
|
||||
logger.info(f"ebook_rerank_request_complete {config.base_url=} {config.model=} candidates={len(results)}")
|
||||
logger.info(
|
||||
"ebook_rerank_request_complete base_url=%s model=%s candidates=%s",
|
||||
config.base_url,
|
||||
config.model,
|
||||
len(results),
|
||||
)
|
||||
return results
|
||||
|
||||
|
||||
@@ -66,7 +76,7 @@ async def score_candidates(
|
||||
|
||||
scores = parse_vllm_scores(body, candidates)
|
||||
for result in scores.values():
|
||||
logger.debug(f"ebook_rerank_candidate_scored {result.chunk_id=} {result.score=}")
|
||||
logger.debug("ebook_rerank_candidate_scored chunk_id=%s score=%s", result.chunk_id, result.score)
|
||||
return scores
|
||||
|
||||
|
||||
|
||||
@@ -19,7 +19,6 @@ from python.ebook_search.bm25_corpus import (
|
||||
load_bm25_corpus,
|
||||
score_bm25_corpus,
|
||||
)
|
||||
from python.ebook_search.chunk_records import CHUNK_RECORD_COLUMNS
|
||||
from python.ebook_search.embeddings import MODEL_DIMENSIONS, embed_query, get_embedding_table
|
||||
from python.ebook_search.protected_phrases.matching import (
|
||||
detect_protected_phrases_for_query,
|
||||
@@ -41,7 +40,7 @@ if TYPE_CHECKING:
|
||||
from sqlalchemy.ext.asyncio import AsyncEngine
|
||||
|
||||
from python.ebook_search.config import EbookSearchConfig
|
||||
from python.ebook_search.protected_phrases.models import PhraseMatch
|
||||
from python.ebook_search.protected_phrases.models import HydratedPhraseMatch
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -75,7 +74,7 @@ class SearchResponse:
|
||||
results: list[SearchResult]
|
||||
rank_label: str
|
||||
timings: tuple[RuntimeStep, ...] = ()
|
||||
phrase_matches: tuple[PhraseMatch, ...] = ()
|
||||
phrase_matches: tuple[HydratedPhraseMatch, ...] = ()
|
||||
|
||||
@property
|
||||
def total_runtime_ms(self) -> float:
|
||||
@@ -85,11 +84,10 @@ class SearchResponse:
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class RetrievalResponse:
|
||||
"""Parallel retrieval output for vector, BM25, and protected phrase candidates."""
|
||||
"""Parallel retrieval output for vector and BM25 candidates."""
|
||||
|
||||
vector_results: list[SearchResult]
|
||||
lexical_results: list[SearchResult]
|
||||
phrase_matches: list[PhraseMatch]
|
||||
timings: tuple[RuntimeStep, ...]
|
||||
|
||||
|
||||
@@ -99,26 +97,33 @@ async def search_ebooks(
|
||||
query: str,
|
||||
config: EbookSearchConfig,
|
||||
*,
|
||||
rerank: bool,
|
||||
phrase_matching: bool,
|
||||
rerank: bool = False,
|
||||
phrase_matching: bool | None = None,
|
||||
) -> SearchResponse:
|
||||
"""Run hybrid vector/BM25 search and optional reranking.
|
||||
|
||||
Phrase matching only runs when both the request asks for it and
|
||||
``config.phrase_matching_enabled`` allows it.
|
||||
"""
|
||||
"""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")
|
||||
|
||||
phrase_matching = phrase_matching and config.phrase_matching_enabled
|
||||
logger.info(f"ebook_search_start query_length={len(query)} {rerank=} {phrase_matching=}")
|
||||
phrase_matching_enabled = config.phrase_matching_enabled if phrase_matching is None else phrase_matching
|
||||
logger.info(
|
||||
"ebook_search_start query_length=%s rerank=%s phrase_matching=%s",
|
||||
len(query),
|
||||
rerank,
|
||||
phrase_matching_enabled,
|
||||
)
|
||||
timings: list[RuntimeStep] = []
|
||||
if phrase_matching_enabled:
|
||||
phrase_matches, timing = await async_timed_result(
|
||||
"Protected phrase detection", query_phrase_matches(engine, query, config)
|
||||
)
|
||||
else:
|
||||
phrase_matches, timing = timed_result("Protected phrase detection skipped", skip_phrase_matches)
|
||||
timings.append(timing)
|
||||
retrieval, timing = await async_timed_result(
|
||||
"Hybrid retrieval",
|
||||
parallel_retrieval(engine, client, query, config, phrase_matching=phrase_matching),
|
||||
parallel_retrieval(engine, client, query, config),
|
||||
)
|
||||
phrase_matches = retrieval.phrase_matches
|
||||
timings.extend(retrieval.timings)
|
||||
timings.append(timing)
|
||||
fused, timing = timed_result(
|
||||
@@ -129,67 +134,68 @@ async def search_ebooks(
|
||||
rank_constant=config.rrf_rank_constant,
|
||||
)
|
||||
timings.append(timing)
|
||||
phrase_boost_timing_name = "Phrase mention boost" if phrase_matching else "Phrase mention boost skipped"
|
||||
fused, timing = await async_timed_result(
|
||||
phrase_boost_timing_name,
|
||||
apply_phrase_mention_boosts(
|
||||
engine,
|
||||
fused,
|
||||
phrase_matches,
|
||||
config.phrase_hit_boost,
|
||||
phrase_matching=phrase_matching,
|
||||
),
|
||||
)
|
||||
if phrase_matching_enabled:
|
||||
fused, timing = await async_timed_result(
|
||||
"Phrase mention boost",
|
||||
apply_phrase_mention_boosts(engine, fused, phrase_matches, config.phrase_hit_boost),
|
||||
)
|
||||
else:
|
||||
fused, timing = timed_result("Phrase mention boost skipped", skip_phrase_mention_boosts, fused)
|
||||
timings.append(timing)
|
||||
rerank_enabled = config.rerank.enabled and rerank
|
||||
rerank_timing_name = "Rerank" if rerank_enabled else "Rerank skipped"
|
||||
response, timing = await async_timed_result(
|
||||
rerank_timing_name,
|
||||
apply_rerank(client, query, fused, config, rerank=rerank_enabled),
|
||||
)
|
||||
if config.rerank.enabled and rerank:
|
||||
response, timing = await async_timed_result("Rerank", apply_rerank(client, query, fused, config))
|
||||
else:
|
||||
response, timing = timed_result("Rerank skipped", skip_rerank, query, fused, config)
|
||||
timings.append(timing)
|
||||
response = replace(response, timings=tuple(timings), phrase_matches=tuple(phrase_matches))
|
||||
logger.info(
|
||||
f"ebook_search_complete vector_candidates={len(retrieval.vector_results)} "
|
||||
f"lexical_candidates={len(retrieval.lexical_results)} fused_candidates={len(fused)} {phrase_matching=} "
|
||||
f"phrase_matches={len(phrase_matches)} returned={len(response.results)} {response.rank_label=} "
|
||||
f"{response.total_runtime_ms=:.1f}"
|
||||
"ebook_search_complete vector_candidates=%s lexical_candidates=%s "
|
||||
"fused_candidates=%s phrase_matching=%s phrase_matches=%s returned=%s rank_label=%s runtime_ms=%.1f",
|
||||
len(retrieval.vector_results),
|
||||
len(retrieval.lexical_results),
|
||||
len(fused),
|
||||
phrase_matching_enabled,
|
||||
len(phrase_matches),
|
||||
len(response.results),
|
||||
response.rank_label,
|
||||
response.total_runtime_ms,
|
||||
)
|
||||
return response
|
||||
|
||||
|
||||
def skip_phrase_matches() -> list[HydratedPhraseMatch]:
|
||||
"""Return no protected phrase matches when phrase matching is disabled."""
|
||||
logger.info("ebook_protected_phrase_detection_skipped")
|
||||
return []
|
||||
|
||||
|
||||
async def query_phrase_matches(
|
||||
engine: AsyncEngine,
|
||||
query: str,
|
||||
config: EbookSearchConfig,
|
||||
*,
|
||||
phrase_matching: bool,
|
||||
) -> list[PhraseMatch]:
|
||||
) -> list[HydratedPhraseMatch]:
|
||||
"""Detect protected phrases in a query without making search fail when phrase tables are unavailable."""
|
||||
if not phrase_matching:
|
||||
logger.info("ebook_protected_phrase_detection_skipped")
|
||||
return []
|
||||
try:
|
||||
async with AsyncSession(engine) as session:
|
||||
return await detect_protected_phrases_for_query(session, query, config)
|
||||
except SQLAlchemyError as error:
|
||||
logger.warning(f"ebook_protected_phrase_detection_unavailable {error=}")
|
||||
logger.warning("ebook_protected_phrase_detection_unavailable error=%s", error)
|
||||
return []
|
||||
|
||||
|
||||
def skip_phrase_mention_boosts(candidates: list[SearchResult]) -> list[SearchResult]:
|
||||
"""Return candidates unchanged when phrase matching is disabled."""
|
||||
logger.info("ebook_phrase_boost_skipped candidates=%s", len(candidates))
|
||||
return candidates
|
||||
|
||||
|
||||
async def apply_phrase_mention_boosts(
|
||||
engine: AsyncEngine,
|
||||
candidates: list[SearchResult],
|
||||
phrase_matches: Sequence[PhraseMatch],
|
||||
phrase_matches: Sequence[HydratedPhraseMatch],
|
||||
phrase_hit_boost: float,
|
||||
*,
|
||||
phrase_matching: bool,
|
||||
) -> list[SearchResult]:
|
||||
"""Boost retrieved chunks that have indexed mentions for detected protected phrases when enabled."""
|
||||
if not phrase_matching:
|
||||
logger.info(f"ebook_phrase_boost_skipped candidates={len(candidates)}")
|
||||
return candidates
|
||||
|
||||
"""Boost retrieved chunks that have indexed mentions for detected protected phrases."""
|
||||
phrase_ids = sorted({match.phrase_id for match in phrase_matches})
|
||||
if not candidates or not phrase_ids or phrase_hit_boost <= 0:
|
||||
return candidates
|
||||
@@ -199,7 +205,7 @@ async def apply_phrase_mention_boosts(
|
||||
async with AsyncSession(engine) as session:
|
||||
phrase_hits = await phrase_hits_for_chunks(session, chunk_ids=chunk_ids, phrase_ids=phrase_ids)
|
||||
except SQLAlchemyError as error:
|
||||
logger.warning(f"ebook_phrase_boost_unavailable {error=}")
|
||||
logger.warning("ebook_phrase_boost_unavailable error=%s", error)
|
||||
return candidates
|
||||
|
||||
if not phrase_hits:
|
||||
@@ -241,62 +247,54 @@ async def parallel_retrieval(
|
||||
client: httpx.AsyncClient,
|
||||
query: str,
|
||||
config: EbookSearchConfig,
|
||||
*,
|
||||
phrase_matching: bool,
|
||||
) -> RetrievalResponse:
|
||||
"""Run vector, BM25, and protected phrase retrieval concurrently with separate database sessions.
|
||||
"""Run vector and BM25 candidate retrieval concurrently with separate database sessions.
|
||||
|
||||
BM25 scoring is pure CPU work over the cached corpus, so it runs in a worker thread
|
||||
instead of on the event loop. Protected phrase detection only depends on the query, so
|
||||
it joins the gather as a third task and returns immediately when phrase matching is disabled.
|
||||
instead of on the event loop.
|
||||
"""
|
||||
phrase_timing_name = "Protected phrase detection" if phrase_matching else "Protected phrase detection skipped"
|
||||
(
|
||||
(vector_results, vector_timing),
|
||||
(lexical_results, lexical_timing),
|
||||
(phrase_matches, phrase_timing),
|
||||
) = await asyncio.gather(
|
||||
(vector_results, vector_timing), (lexical_results, lexical_timing) = await asyncio.gather(
|
||||
async_timed_result("Embedding + vector search", vector_candidates(engine, client, query, config)),
|
||||
async_timed_result("BM25 search", asyncio.to_thread(bm25_candidates, query, config)),
|
||||
async_timed_result(
|
||||
phrase_timing_name,
|
||||
query_phrase_matches(engine, query, config, phrase_matching=phrase_matching),
|
||||
),
|
||||
)
|
||||
|
||||
logger.info(
|
||||
f"ebook_parallel_retrieval_complete vector_candidates={len(vector_results)} "
|
||||
f"lexical_candidates={len(lexical_results)} phrase_matches={len(phrase_matches)}"
|
||||
"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,
|
||||
phrase_matches=phrase_matches,
|
||||
timings=(
|
||||
replace(vector_timing, counts_toward_total=False),
|
||||
replace(lexical_timing, counts_toward_total=False),
|
||||
replace(phrase_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")
|
||||
|
||||
|
||||
async def apply_rerank(
|
||||
client: httpx.AsyncClient,
|
||||
query: str,
|
||||
candidates: list[SearchResult],
|
||||
config: EbookSearchConfig,
|
||||
*,
|
||||
rerank: bool,
|
||||
) -> SearchResponse:
|
||||
"""Rerank already-fused hybrid candidates when enabled for this request."""
|
||||
if not rerank:
|
||||
logger.info(f"ebook_rerank_skipped candidates={len(candidates)}")
|
||||
return SearchResponse(query=query, results=candidates[: config.top_k], rank_label="Hybrid")
|
||||
|
||||
"""Rerank already-fused hybrid candidates."""
|
||||
reranked = await rerank_chunks(client, query, candidates[: config.rerank.candidates], config.rerank)
|
||||
logger.info(
|
||||
f"ebook_rerank_complete input_candidates={min(len(candidates), config.rerank.candidates)} "
|
||||
f"returned={len(reranked)}"
|
||||
"ebook_rerank_complete input_candidates=%s returned=%s",
|
||||
min(len(candidates), config.rerank.candidates),
|
||||
len(reranked),
|
||||
)
|
||||
return SearchResponse(
|
||||
query=query,
|
||||
@@ -334,7 +332,13 @@ async def vector_candidates(
|
||||
score = (literal(1.0) - distance).label("score")
|
||||
statement = (
|
||||
select(
|
||||
*CHUNK_RECORD_COLUMNS,
|
||||
EbookChunk.id.label("chunk_id"),
|
||||
EbookChunk.text.label("text"),
|
||||
EbookSource.id.label("source_id"),
|
||||
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)
|
||||
@@ -348,7 +352,10 @@ async def vector_candidates(
|
||||
rows = (await session.execute(statement)).mappings()
|
||||
results = [search_result_from_row(row) for row in rows]
|
||||
logger.info(
|
||||
f"ebook_vector_search_complete {config.embedding_model=} {model.dimension=} candidates={len(results)}"
|
||||
"ebook_vector_search_complete model=%s dimension=%s candidates=%s",
|
||||
config.embedding_model,
|
||||
model.dimension,
|
||||
len(results),
|
||||
)
|
||||
return results
|
||||
|
||||
@@ -358,7 +365,7 @@ def bm25_candidates(query: str, config: EbookSearchConfig) -> list[SearchResult]
|
||||
try:
|
||||
corpus = load_bm25_corpus(config)
|
||||
except BM25CorpusUnavailableError as error:
|
||||
logger.warning(f"ebook_bm25_index_unavailable_skipping {error=}")
|
||||
logger.warning("ebook_bm25_index_unavailable_skipping error=%s", error)
|
||||
return []
|
||||
|
||||
if not corpus.records:
|
||||
@@ -373,7 +380,12 @@ def bm25_candidates(query: str, config: EbookSearchConfig) -> list[SearchResult]
|
||||
]
|
||||
|
||||
max_score = results[0].bm25_score if results else 0.0
|
||||
logger.info(f"ebook_bm25_search_complete corpus={len(corpus.records)} candidates={len(results)} {max_score=:.6f}")
|
||||
logger.info(
|
||||
"ebook_bm25_search_complete corpus=%s candidates=%s max_score=%.6f",
|
||||
len(corpus.records),
|
||||
len(results),
|
||||
max_score,
|
||||
)
|
||||
return results
|
||||
|
||||
|
||||
|
||||
@@ -1,25 +1,6 @@
|
||||
"""Reusable FastAPI tools."""
|
||||
|
||||
from python.fastapi_tools.db import (
|
||||
AppAsyncEngine,
|
||||
AppEngine,
|
||||
AsyncDbSession,
|
||||
DbSession,
|
||||
get_async_db,
|
||||
get_async_engine,
|
||||
get_db,
|
||||
get_engine,
|
||||
)
|
||||
from python.fastapi_tools.db import AsyncDbSession, DbSession, get_async_db, get_db
|
||||
from python.fastapi_tools.zstd_middleware import ZstdMiddleware
|
||||
|
||||
__all__ = [
|
||||
"AppAsyncEngine",
|
||||
"AppEngine",
|
||||
"AsyncDbSession",
|
||||
"DbSession",
|
||||
"ZstdMiddleware",
|
||||
"get_async_db",
|
||||
"get_async_engine",
|
||||
"get_db",
|
||||
"get_engine",
|
||||
]
|
||||
__all__ = ["AsyncDbSession", "DbSession", "ZstdMiddleware", "get_async_db", "get_db"]
|
||||
|
||||
@@ -5,24 +5,13 @@ from __future__ import annotations
|
||||
from typing import TYPE_CHECKING, Annotated
|
||||
|
||||
from fastapi import Depends, Request
|
||||
from sqlalchemy.engine import Engine
|
||||
from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import AsyncIterator, Iterator
|
||||
|
||||
|
||||
def get_engine(request: Request) -> Engine:
|
||||
"""Get a synchronous database engine from app state."""
|
||||
return request.app.state.engine
|
||||
|
||||
|
||||
def get_async_engine(request: Request) -> AsyncEngine:
|
||||
"""Get an asynchronous database engine from app state."""
|
||||
return request.app.state.engine
|
||||
|
||||
|
||||
def get_db(request: Request) -> Iterator[Session]:
|
||||
"""Get database session from app state."""
|
||||
with Session(request.app.state.engine) as session:
|
||||
@@ -39,7 +28,5 @@ async def get_async_db(request: Request) -> AsyncIterator[AsyncSession]:
|
||||
yield session
|
||||
|
||||
|
||||
AppEngine = Annotated[Engine, Depends(get_engine)]
|
||||
AppAsyncEngine = Annotated[AsyncEngine, Depends(get_async_engine)]
|
||||
DbSession = Annotated[Session, Depends(get_db)]
|
||||
AsyncDbSession = Annotated[AsyncSession, Depends(get_async_db)]
|
||||
|
||||
@@ -1,54 +0,0 @@
|
||||
# Gems
|
||||
|
||||
Gems is a server-rendered, turn-based resource-engine game for one to four human or AI players. It uses FastAPI,
|
||||
Jinja, HTMX, server-sent events, and SQLite.
|
||||
|
||||
The application deliberately contains no playable card deck, patron/governor set, objective set, official artwork,
|
||||
or copied rulebook text. A room host must upload a content pack they are entitled to use before starting a game.
|
||||
|
||||
## Run locally
|
||||
|
||||
```shell
|
||||
uv run gems --host 127.0.0.1 --port 8082
|
||||
```
|
||||
|
||||
The default database and installation key are created under `.gems/`. The following environment variables override
|
||||
runtime behavior:
|
||||
|
||||
- `GEMS_DATABASE_PATH`
|
||||
- `GEMS_KEY_PATH`
|
||||
- `GEMS_PUBLIC_ORIGIN`
|
||||
- `GEMS_SECURE_COOKIES`
|
||||
- `GEMS_HOST`
|
||||
- `GEMS_PORT`
|
||||
|
||||
## Content packs
|
||||
|
||||
The current schema is available from a running server at `/schemas/content-pack-v1.json`. A pack defines exactly
|
||||
five normal resources, one wild resource, cards, and optional patrons, objectives, and outposts. `patrons` is the
|
||||
canonical field name; `governors` is accepted as an input alias.
|
||||
|
||||
Cards may use only the built-in, bounded effect vocabulary:
|
||||
|
||||
- `none`
|
||||
- `virtual_wild`
|
||||
- `copy_bonus`
|
||||
- `copy_and_claim`
|
||||
- `multi_bonus`
|
||||
- `claim_free`
|
||||
- an optional discard-cards alternate cost
|
||||
|
||||
Unknown fields, resource references, executable expressions, HTML, artwork URLs, and files larger than 512 KiB are
|
||||
rejected. The normalized pack is private to its room and becomes immutable when play starts.
|
||||
|
||||
## Neutral module mapping
|
||||
|
||||
Gems calls the four optional mechanics Objectives, Outposts, Eastern Decks, and Fortifications. Lobby presets combine
|
||||
these mechanics into the familiar base, objective-race, objective-plus-outpost, eastern-plus-fortification, and
|
||||
all-module configurations. Component identities and values always come from the uploaded pack.
|
||||
|
||||
## Jeeves
|
||||
|
||||
The NixOS module runs one Uvicorn worker on `127.0.0.1:8002`, stores state in
|
||||
`/zfs/media/services/gems`, and publishes it through HAProxy at `https://gems.tmmworkshop.com`. The DNS record must
|
||||
point to Jeeves before ACME can issue the certificate.
|
||||
@@ -1 +0,0 @@
|
||||
"""Gems multiplayer card-engine game."""
|
||||
@@ -1,5 +0,0 @@
|
||||
"""AI controllers for Gems."""
|
||||
|
||||
from .runner import choose_ai_command
|
||||
|
||||
__all__ = ["choose_ai_command"]
|
||||
@@ -1,68 +0,0 @@
|
||||
"""Deterministic AI command selection at three difficulty levels."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import random
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from python.gems.domain.engine import RuleError, apply_command, score
|
||||
from python.gems.domain.legal_actions import legal_commands
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from python.gems.domain.models import ContentPack, GameCommand, GameSettings, GameState
|
||||
|
||||
|
||||
def choose_ai_command(
|
||||
state: GameState,
|
||||
pack: ContentPack,
|
||||
settings: GameSettings,
|
||||
seat: int,
|
||||
difficulty: str,
|
||||
) -> GameCommand:
|
||||
"""Choose from legal commands using only public/current-player information."""
|
||||
actions = legal_commands(state, pack, settings, seat)
|
||||
if not actions:
|
||||
error = "AI has no legal command"
|
||||
raise RuntimeError(error)
|
||||
rng = random.Random(f"{state.seed}:{state.revision}:{seat}:{difficulty}") # noqa: S311 - deterministic AI
|
||||
if difficulty == "easy":
|
||||
return rng.choice(actions)
|
||||
|
||||
ranked = sorted(actions, key=lambda item: _heuristic(item, state, pack), reverse=True)
|
||||
if difficulty == "medium":
|
||||
return ranked[0]
|
||||
|
||||
# Bounded deterministic rollout: examine at most the twelve strongest actions.
|
||||
best = ranked[0]
|
||||
best_value = float("-inf")
|
||||
for candidate in ranked[:12]:
|
||||
try:
|
||||
future = apply_command(state, candidate, pack, settings, actor_seat=seat)
|
||||
except RuleError:
|
||||
continue
|
||||
value = score(future.players[seat], pack) * 100 + _heuristic(candidate, state, pack)
|
||||
value += rng.random() * 0.001
|
||||
if value > best_value:
|
||||
best, best_value = candidate, value
|
||||
return best
|
||||
|
||||
|
||||
def _heuristic(command: GameCommand, state: GameState, pack: ContentPack) -> float:
|
||||
player = state.players[state.current_seat]
|
||||
value = 0.0
|
||||
if command.type == "purchase":
|
||||
card = pack.card(str(command.payload.get("card_id")))
|
||||
value += 40 + card.points * 25 + (8 if card.bonus_resource else 0)
|
||||
value -= sum(command.payload.get("payment", {}).values())
|
||||
elif command.type == "take_distinct":
|
||||
value += 10 + len(command.payload.get("resources", []))
|
||||
elif command.type == "reserve":
|
||||
value += 6
|
||||
elif command.type == "take_double":
|
||||
value += 4
|
||||
elif command.type == "choose":
|
||||
value += 20
|
||||
elif command.type == "decline":
|
||||
value -= 2
|
||||
value += score(player, pack) * 0.01
|
||||
return value
|
||||
@@ -1,25 +0,0 @@
|
||||
"""Runtime configuration loaded from environment variables."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
from pydantic_settings import BaseSettings, SettingsConfigDict
|
||||
|
||||
|
||||
class GemsConfig(BaseSettings):
|
||||
"""Configure persistence, public URLs, cookies, and the HTTP server."""
|
||||
|
||||
model_config = SettingsConfigDict(env_prefix="GEMS_")
|
||||
|
||||
database_path: Path = Path(".gems/gems.sqlite3")
|
||||
key_path: Path = Path(".gems/instance.key")
|
||||
public_origin: str = "http://127.0.0.1:8082"
|
||||
secure_cookies: bool = False
|
||||
host: str = "127.0.0.1"
|
||||
port: int = 8082
|
||||
|
||||
|
||||
def load_config() -> GemsConfig:
|
||||
"""Load Gems configuration from defaults and environment variables."""
|
||||
return GemsConfig()
|
||||
@@ -1,55 +0,0 @@
|
||||
"""Content-pack parsing, normalization, and JSON Schema publication."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
from pydantic import ValidationError
|
||||
|
||||
from .domain.models import ContentPack
|
||||
|
||||
MAX_PACK_BYTES = 512 * 1024
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ParsedPack:
|
||||
"""Validated content plus its canonical representation."""
|
||||
|
||||
pack: ContentPack
|
||||
canonical_json: str
|
||||
digest: str
|
||||
|
||||
|
||||
class ContentPackError(ValueError):
|
||||
"""A safe, user-visible content validation error."""
|
||||
|
||||
|
||||
def parse_content_pack(raw: bytes | str) -> ParsedPack:
|
||||
"""Validate a content pack and return stable canonical JSON."""
|
||||
data = raw.encode() if isinstance(raw, str) else raw
|
||||
if len(data) > MAX_PACK_BYTES:
|
||||
error = "Content pack exceeds the 512 KiB limit"
|
||||
raise ContentPackError(error)
|
||||
try:
|
||||
decoded: Any = json.loads(data)
|
||||
except (UnicodeDecodeError, json.JSONDecodeError) as exc:
|
||||
error = f"Invalid JSON: {exc}"
|
||||
raise ContentPackError(error) from exc
|
||||
try:
|
||||
pack = ContentPack.model_validate(decoded)
|
||||
except ValidationError as exc:
|
||||
messages = []
|
||||
for issue in exc.errors(include_url=False):
|
||||
path = ".".join(str(part) for part in issue["loc"])
|
||||
messages.append(f"{path or '$'}: {issue['msg']}")
|
||||
raise ContentPackError("\n".join(messages)) from exc
|
||||
canonical = json.dumps(pack.model_dump(mode="json", by_alias=False), sort_keys=True, separators=(",", ":"))
|
||||
return ParsedPack(pack=pack, canonical_json=canonical, digest=hashlib.sha256(canonical.encode()).hexdigest())
|
||||
|
||||
|
||||
def content_pack_schema() -> dict[str, Any]:
|
||||
"""Return the authoritative version-one JSON Schema."""
|
||||
return ContentPack.model_json_schema(by_alias=False)
|
||||
@@ -1 +0,0 @@
|
||||
"""Domain models and rules for Gems."""
|
||||
@@ -1,903 +0,0 @@
|
||||
"""Deterministic, server-authoritative Gems rule engine."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import random
|
||||
from typing import Any
|
||||
|
||||
from .models import (
|
||||
CardDefinition,
|
||||
CardEffectKind,
|
||||
ContentPack,
|
||||
GameCommand,
|
||||
GameSettings,
|
||||
GameState,
|
||||
OutpostDefinition,
|
||||
OutpostPower,
|
||||
OwnedCard,
|
||||
PendingChoice,
|
||||
PlayerState,
|
||||
)
|
||||
from .requirements import requirements_met
|
||||
|
||||
MAX_PLAYERS = 4
|
||||
TWO_PLAYER_COUNT = 2
|
||||
THREE_PLAYER_COUNT = 3
|
||||
DISTINCT_TAKE_COUNT = 3
|
||||
DOUBLE_TAKE_MINIMUM = 4
|
||||
BLIND_RESERVE_DRAW_COUNT = 2
|
||||
|
||||
|
||||
class RuleError(ValueError):
|
||||
"""A command was not legal for the current state."""
|
||||
|
||||
|
||||
def deck_key(deck: str, tier: int) -> str:
|
||||
"""Build the state key for a named deck and tier."""
|
||||
return f"{deck}:{tier}"
|
||||
|
||||
|
||||
def new_game(
|
||||
room_code: str,
|
||||
names: list[str],
|
||||
pack: ContentPack,
|
||||
settings: GameSettings,
|
||||
*,
|
||||
seed: int,
|
||||
) -> GameState:
|
||||
"""Create and deal a deterministic game."""
|
||||
if not 1 <= len(names) <= MAX_PLAYERS:
|
||||
error = "Games require one to four seats"
|
||||
raise RuleError(error)
|
||||
if settings.first_player_mode == "selected" and settings.first_player_seat >= len(names):
|
||||
error = "The selected first player is no longer in the room"
|
||||
raise RuleError(error)
|
||||
_validate_startable(pack, settings)
|
||||
rng = random.Random(seed) # noqa: S311 - deterministic seeded shuffle
|
||||
first_seat = settings.first_player_seat if settings.first_player_mode == "selected" else rng.randrange(len(names))
|
||||
normal_count = 4 if len(names) <= TWO_PLAYER_COUNT else 5 if len(names) == THREE_PLAYER_COUNT else 7
|
||||
supply = dict.fromkeys(pack.resource_ids, normal_count)
|
||||
supply[pack.wild_resource.id] = 5
|
||||
players = [
|
||||
PlayerState(
|
||||
seat=index,
|
||||
name=name,
|
||||
tokens=dict.fromkeys((*pack.resource_ids, pack.wild_resource.id), 0),
|
||||
fortifications_available=settings.fortifications_per_player if settings.modules.fortifications else 0,
|
||||
)
|
||||
for index, name in enumerate(names)
|
||||
]
|
||||
decks: dict[str, list[str]] = {}
|
||||
markets: dict[str, list[str]] = {}
|
||||
for source in ("base", "eastern"):
|
||||
if source == "eastern" and not settings.modules.eastern_decks:
|
||||
continue
|
||||
market_size = settings.base_market_size if source == "base" else settings.eastern_market_size
|
||||
for tier in (1, 2, 3):
|
||||
key = deck_key(source, tier)
|
||||
cards = [card.id for card in pack.cards if card.deck == source and card.tier == tier]
|
||||
rng.shuffle(cards)
|
||||
markets[key] = [cards.pop() for _ in range(min(market_size, len(cards)))]
|
||||
decks[key] = cards
|
||||
patrons = [patron.id for patron in pack.patrons]
|
||||
rng.shuffle(patrons)
|
||||
patrons = patrons[: min(len(patrons), len(players) + 1)]
|
||||
objectives = [objective.id for objective in pack.objectives]
|
||||
rng.shuffle(objectives)
|
||||
objectives = objectives[: settings.objective_count] if settings.modules.objectives else []
|
||||
return GameState(
|
||||
room_code=room_code,
|
||||
seed=seed,
|
||||
players=players,
|
||||
supply=supply,
|
||||
decks=decks,
|
||||
markets=markets,
|
||||
available_patrons=[] if settings.modules.objectives else patrons,
|
||||
available_objectives=objectives,
|
||||
current_seat=first_seat,
|
||||
first_seat=first_seat,
|
||||
)
|
||||
|
||||
|
||||
def _validate_startable(pack: ContentPack, settings: GameSettings) -> None:
|
||||
for tier in (1, 2, 3):
|
||||
if not any(card.deck == "base" and card.tier == tier for card in pack.cards):
|
||||
error = f"The base deck has no tier-{tier} cards"
|
||||
raise RuleError(error)
|
||||
if settings.modules.eastern_decks and not any(
|
||||
card.deck == "eastern" and card.tier == tier for card in pack.cards
|
||||
):
|
||||
error = f"The eastern deck has no tier-{tier} cards"
|
||||
raise RuleError(error)
|
||||
if settings.modules.objectives and not pack.objectives:
|
||||
error = "The Objectives module needs objective definitions"
|
||||
raise RuleError(error)
|
||||
if settings.modules.outposts and not pack.outposts:
|
||||
error = "The Outposts module needs outpost definitions"
|
||||
raise RuleError(error)
|
||||
|
||||
|
||||
def card_map(pack: ContentPack) -> dict[str, CardDefinition]:
|
||||
"""Index all cards in a content pack by identifier."""
|
||||
return {card.id: card for card in pack.cards}
|
||||
|
||||
|
||||
def bonuses(player: PlayerState, pack: ContentPack) -> dict[str, int]:
|
||||
"""Calculate effective bonuses after copies, multiples, and discards."""
|
||||
cards = card_map(pack)
|
||||
result = dict.fromkeys(pack.resource_ids, 0)
|
||||
for owned in player.cards:
|
||||
definition = cards[owned.card_id]
|
||||
resource = owned.copied_resource or definition.bonus_resource
|
||||
if resource is not None:
|
||||
result[resource] += definition.effect.amount if definition.effect.kind == CardEffectKind.MULTI_BONUS else 1
|
||||
return result
|
||||
|
||||
|
||||
def score(player: PlayerState, pack: ContentPack) -> int:
|
||||
"""Calculate a player's score from cards, patrons, and outposts."""
|
||||
cards = card_map(pack)
|
||||
patrons = {patron.id: patron for patron in pack.patrons}
|
||||
outposts = {outpost.id: outpost for outpost in pack.outposts}
|
||||
value = sum(cards[item.card_id].points for item in player.cards)
|
||||
value += sum(patrons[item].points for item in player.patrons)
|
||||
for item in player.outposts:
|
||||
definition = outposts[item]
|
||||
if definition.power == OutpostPower.POINTS_PER_OUTPOST:
|
||||
value += definition.value * len(player.outposts)
|
||||
return value
|
||||
|
||||
|
||||
def affordable_payments(player: PlayerState, card: CardDefinition, pack: ContentPack, *, double_wild: bool) -> bool:
|
||||
"""Report whether a player can cover a card's standard cost."""
|
||||
discount = bonuses(player, pack)
|
||||
shortage = sum(
|
||||
max(0, amount - discount.get(resource, 0) - player.tokens.get(resource, 0))
|
||||
for resource, amount in card.cost.items()
|
||||
)
|
||||
wild_value = 2 if double_wild else 1
|
||||
return shortage <= player.tokens.get(pack.wild_resource.id, 0) * wild_value
|
||||
|
||||
|
||||
def visible_cards(state: GameState) -> set[str]:
|
||||
"""Return the identifiers of cards currently visible in markets."""
|
||||
return {card for market in state.markets.values() for card in market}
|
||||
|
||||
|
||||
def apply_command(
|
||||
state: GameState,
|
||||
command: GameCommand,
|
||||
pack: ContentPack,
|
||||
settings: GameSettings,
|
||||
*,
|
||||
actor_seat: int,
|
||||
) -> GameState:
|
||||
"""Validate and apply one command, returning a new state snapshot."""
|
||||
if state.finished:
|
||||
error = "The game is finished"
|
||||
raise RuleError(error)
|
||||
if command.expected_revision != state.revision:
|
||||
error = "The board changed; refresh and try again"
|
||||
raise RuleError(error)
|
||||
if actor_seat != state.current_seat:
|
||||
error = "It is not your turn"
|
||||
raise RuleError(error)
|
||||
if state.pending and state.pending.seat != actor_seat:
|
||||
error = "Another player must resolve the pending choice"
|
||||
raise RuleError(error)
|
||||
|
||||
result = state.model_copy(deep=True)
|
||||
if result.pending:
|
||||
_apply_pending(result, command, pack, settings)
|
||||
elif command.type == "take_distinct":
|
||||
_take_distinct(result, command.payload, pack, settings)
|
||||
elif command.type == "take_double":
|
||||
_take_double(result, command.payload, pack, settings)
|
||||
elif command.type == "reserve":
|
||||
_reserve(result, command.payload, pack, settings)
|
||||
elif command.type == "purchase":
|
||||
_purchase(result, command.payload, pack, settings, is_conquest=False)
|
||||
else:
|
||||
error = "Choose a normal turn action"
|
||||
raise RuleError(error)
|
||||
result.revision += 1
|
||||
return result
|
||||
|
||||
|
||||
def _take_distinct(state: GameState, payload: dict[str, Any], pack: ContentPack, settings: GameSettings) -> None:
|
||||
colors = payload.get("resources")
|
||||
if not isinstance(colors, list) or len(colors) != len(set(colors)):
|
||||
error = "Choose distinct resources"
|
||||
raise RuleError(error)
|
||||
available = [resource for resource in pack.resource_ids if state.supply.get(resource, 0) > 0]
|
||||
required = DISTINCT_TAKE_COUNT if len(available) >= DISTINCT_TAKE_COUNT else None
|
||||
if required is not None and len(colors) != required:
|
||||
error = "Take exactly three different resources when possible"
|
||||
raise RuleError(error)
|
||||
if required is None and not 1 <= len(colors) <= len(available):
|
||||
error = "Choose one or more available resources"
|
||||
raise RuleError(error)
|
||||
if not set(colors) <= set(available):
|
||||
error = "A selected resource is unavailable"
|
||||
raise RuleError(error)
|
||||
player = state.players[state.current_seat]
|
||||
for resource in colors:
|
||||
state.supply[resource] -= 1
|
||||
player.tokens[resource] += 1
|
||||
state.log.append(f"{player.name} took {len(colors)} different resources")
|
||||
_after_standard_action(state, pack, settings, action="take_distinct")
|
||||
|
||||
|
||||
def _take_double(state: GameState, payload: dict[str, Any], pack: ContentPack, settings: GameSettings) -> None:
|
||||
resource = payload.get("resource")
|
||||
if resource not in pack.resource_ids or state.supply.get(resource, 0) < DOUBLE_TAKE_MINIMUM:
|
||||
error = "That resource cannot be taken twice"
|
||||
raise RuleError(error)
|
||||
player = state.players[state.current_seat]
|
||||
state.supply[resource] -= 2
|
||||
player.tokens[resource] += 2
|
||||
state.log.append(f"{player.name} took two {resource} resources")
|
||||
extra = _owned_outpost(player, pack, OutpostPower.RESOURCE_AFTER_DOUBLE)
|
||||
if settings.modules.outposts and extra:
|
||||
options = [item for item in pack.resource_ids if item != resource and state.supply.get(item, 0) > 0]
|
||||
if options:
|
||||
state.pending = PendingChoice(
|
||||
kind="resource", seat=player.seat, options=options, context={"next": "after_action"}
|
||||
)
|
||||
return
|
||||
_after_standard_action(state, pack, settings, action="take_double")
|
||||
|
||||
|
||||
def _reserve(state: GameState, payload: dict[str, Any], pack: ContentPack, settings: GameSettings) -> None:
|
||||
player = state.players[state.current_seat]
|
||||
if len(player.reserved) >= settings.reserve_limit:
|
||||
error = "Your reserve is full"
|
||||
raise RuleError(error)
|
||||
card_id = payload.get("card_id")
|
||||
source = payload.get("deck")
|
||||
drawn: list[str]
|
||||
if card_id:
|
||||
if card_id not in visible_cards(state):
|
||||
error = "That card is not visible"
|
||||
raise RuleError(error)
|
||||
_assert_not_blocked(state, card_id, player.seat)
|
||||
_remove_visible(state, card_id)
|
||||
drawn = [card_id]
|
||||
elif isinstance(source, str) and source in state.decks:
|
||||
if not state.decks[source]:
|
||||
error = "That deck is empty"
|
||||
raise RuleError(error)
|
||||
draw_count = 1
|
||||
post = _owned_outpost(player, pack, OutpostPower.BLIND_RESERVE_TWO)
|
||||
if post and (not source.startswith("eastern:") or settings.interactions.outpost_reserve_applies_eastern):
|
||||
draw_count = min(BLIND_RESERVE_DRAW_COUNT, len(state.decks[source]))
|
||||
drawn = [state.decks[source].pop() for _ in range(draw_count)]
|
||||
else:
|
||||
error = "Choose a visible card or deck"
|
||||
raise RuleError(error)
|
||||
if len(drawn) == BLIND_RESERVE_DRAW_COUNT:
|
||||
state.pending = PendingChoice(kind="reserve_keep", seat=player.seat, options=drawn, context={"deck": source})
|
||||
return
|
||||
_finish_reserve(state, drawn[0], pack, settings)
|
||||
|
||||
|
||||
def _finish_reserve(state: GameState, card_id: str, pack: ContentPack, settings: GameSettings) -> None:
|
||||
player = state.players[state.current_seat]
|
||||
player.reserved.append(card_id)
|
||||
_refill_markets(state, settings)
|
||||
wild = pack.wild_resource.id
|
||||
if state.supply[wild] > 0:
|
||||
state.supply[wild] -= 1
|
||||
player.tokens[wild] += 1
|
||||
state.log.append(f"{player.name} reserved a card")
|
||||
_after_standard_action(state, pack, settings, action="reserve")
|
||||
|
||||
|
||||
def _purchase(
|
||||
state: GameState,
|
||||
payload: dict[str, Any],
|
||||
pack: ContentPack,
|
||||
settings: GameSettings,
|
||||
*,
|
||||
is_conquest: bool,
|
||||
) -> None:
|
||||
player = state.players[state.current_seat]
|
||||
card_id = payload.get("card_id")
|
||||
if not isinstance(card_id, str):
|
||||
error = "Choose a card"
|
||||
raise RuleError(error)
|
||||
from_reserve = card_id in player.reserved
|
||||
if not from_reserve:
|
||||
if card_id not in visible_cards(state):
|
||||
error = "That card is not available"
|
||||
raise RuleError(error)
|
||||
_assert_not_blocked(state, card_id, player.seat)
|
||||
if is_conquest and state.fortifications.get(card_id, {}).get(player.seat, 0) < settings.fortifications_per_player:
|
||||
error = "All your fortifications must occupy the conquest card"
|
||||
raise RuleError(error)
|
||||
card = pack.card(card_id)
|
||||
_pay_for_purchase(state, player, card, payload, pack, settings)
|
||||
if from_reserve:
|
||||
player.reserved.remove(card_id)
|
||||
else:
|
||||
_remove_visible(state, card_id)
|
||||
_return_fortifications(state, card_id)
|
||||
owned = OwnedCard(card_id=card_id)
|
||||
player.cards.append(owned)
|
||||
player.purchased_card_count += 1
|
||||
state.turn_purchase_count += 1
|
||||
state.log.append(f"{player.name} purchased {card.label}")
|
||||
if card.effect.kind in {CardEffectKind.COPY_BONUS, CardEffectKind.COPY_AND_CLAIM}:
|
||||
# Do not allow the new copy card to satisfy its own target requirement.
|
||||
prior = player.cards.pop()
|
||||
options = sorted(resource for resource, amount in bonuses(player, pack).items() if amount)
|
||||
player.cards.append(prior)
|
||||
if not options:
|
||||
error = "This card requires an existing bonus to copy"
|
||||
raise RuleError(error)
|
||||
state.pending = PendingChoice(
|
||||
kind="copy_bonus",
|
||||
seat=player.seat,
|
||||
options=options,
|
||||
context={"card_id": card_id, "is_conquest": is_conquest},
|
||||
)
|
||||
return
|
||||
if card.effect.kind == CardEffectKind.CLAIM_FREE:
|
||||
_queue_free_card(state, card, pack, settings, is_conquest=is_conquest)
|
||||
return
|
||||
_after_purchase(state, pack, settings, is_conquest=is_conquest)
|
||||
|
||||
|
||||
def _pay_for_purchase(
|
||||
state: GameState,
|
||||
player: PlayerState,
|
||||
card: CardDefinition,
|
||||
payload: dict[str, Any],
|
||||
pack: ContentPack,
|
||||
settings: GameSettings,
|
||||
) -> None:
|
||||
if card.alternate_cost:
|
||||
discarded = payload.get("discard_cards", [])
|
||||
_pay_alternate(player, card, discarded, pack)
|
||||
if not settings.interactions.retain_outposts_after_discard:
|
||||
_reconcile_outposts(player, pack)
|
||||
else:
|
||||
payment = payload.get("payment", {})
|
||||
virtual_cards = payload.get("virtual_wild_cards", [])
|
||||
_pay_tokens(state, player, card, payment, virtual_cards, pack, settings)
|
||||
|
||||
|
||||
def _pay_tokens( # noqa: C901 - payment validation mirrors the rule sequence
|
||||
state: GameState,
|
||||
player: PlayerState,
|
||||
card: CardDefinition,
|
||||
payment: object,
|
||||
virtual_cards: object,
|
||||
pack: ContentPack,
|
||||
settings: GameSettings,
|
||||
) -> None:
|
||||
if not isinstance(payment, dict) or any(not isinstance(value, int) or value < 0 for value in payment.values()):
|
||||
error = "Provide a valid payment"
|
||||
raise RuleError(error)
|
||||
allowed = {*pack.resource_ids, pack.wild_resource.id}
|
||||
if not set(payment) <= allowed:
|
||||
error = "Payment contains an unknown resource"
|
||||
raise RuleError(error)
|
||||
if any(payment.get(resource, 0) > player.tokens.get(resource, 0) for resource in allowed):
|
||||
error = "Payment uses resources you do not have"
|
||||
raise RuleError(error)
|
||||
discount = bonuses(player, pack)
|
||||
if not isinstance(virtual_cards, list) or len(virtual_cards) != len(set(virtual_cards)):
|
||||
error = "Virtual-wild selections must be unique"
|
||||
raise RuleError(error)
|
||||
owned_by_id = {item.card_id: item for item in player.cards}
|
||||
virtual_value = 0
|
||||
for card_id in virtual_cards:
|
||||
owned = owned_by_id.get(card_id)
|
||||
if owned is None:
|
||||
error = "You do not own a selected virtual-wild card"
|
||||
raise RuleError(error)
|
||||
definition = pack.card(card_id)
|
||||
if definition.effect.kind != CardEffectKind.VIRTUAL_WILD:
|
||||
error = "A selected card does not grant virtual wild resources"
|
||||
raise RuleError(error)
|
||||
virtual_value += definition.effect.amount
|
||||
wild = payment.get(pack.wild_resource.id, 0)
|
||||
double = bool(_owned_outpost(player, pack, OutpostPower.DOUBLE_WILD))
|
||||
wild_value = 2 if double else 1
|
||||
virtual_multiplier = 2 if double and settings.interactions.virtual_wild_can_double else 1
|
||||
remaining_wild_value = wild * wild_value + virtual_value * virtual_multiplier
|
||||
for resource in pack.resource_ids:
|
||||
due = max(0, card.cost.get(resource, 0) - discount.get(resource, 0))
|
||||
normal = payment.get(resource, 0)
|
||||
if normal > due:
|
||||
error = "Payment overpays a normal resource"
|
||||
raise RuleError(error)
|
||||
remaining_wild_value -= due - normal
|
||||
if remaining_wild_value < 0:
|
||||
error = "Payment does not exactly cover the card cost"
|
||||
raise RuleError(error)
|
||||
for resource, amount in payment.items():
|
||||
player.tokens[resource] -= amount
|
||||
state.supply[resource] += amount
|
||||
if virtual_cards:
|
||||
player.cards = [item for item in player.cards if item.card_id not in virtual_cards]
|
||||
|
||||
|
||||
def _pay_alternate(player: PlayerState, card: CardDefinition, selected: object, pack: ContentPack) -> None:
|
||||
if not isinstance(selected, list) or not card.alternate_cost or len(selected) != card.alternate_cost.count:
|
||||
error = "Choose the required owned cards to discard"
|
||||
raise RuleError(error)
|
||||
if len(selected) != len(set(selected)):
|
||||
error = "A card can only be discarded once"
|
||||
raise RuleError(error)
|
||||
owned_by_id = {item.card_id: item for item in player.cards}
|
||||
for card_id in selected:
|
||||
owned = owned_by_id.get(card_id)
|
||||
if owned is None:
|
||||
error = "You do not own a selected discard"
|
||||
raise RuleError(error)
|
||||
definition = pack.card(card_id)
|
||||
effective = owned.copied_resource or definition.bonus_resource
|
||||
if effective != card.alternate_cost.discard_resource:
|
||||
error = "A selected discard has the wrong color"
|
||||
raise RuleError(error)
|
||||
# Copy cards of the effective color must be discarded first.
|
||||
matching_copies = [
|
||||
item.card_id for item in player.cards if item.copied_resource == card.alternate_cost.discard_resource
|
||||
]
|
||||
if any(item not in selected for item in matching_copies[: min(len(matching_copies), len(selected))]):
|
||||
error = "Copied cards of this color must be discarded first"
|
||||
raise RuleError(error)
|
||||
player.cards = [item for item in player.cards if item.card_id not in selected]
|
||||
|
||||
|
||||
def _apply_pending( # noqa: C901, PLR0911, PLR0912, PLR0915 - explicit pending-choice state machine
|
||||
state: GameState, command: GameCommand, pack: ContentPack, settings: GameSettings
|
||||
) -> None:
|
||||
pending = state.pending
|
||||
if pending is None or command.type not in {"choose", "decline", "purchase"}:
|
||||
error = "Resolve the pending choice"
|
||||
raise RuleError(error)
|
||||
player = state.players[state.current_seat]
|
||||
choice = command.payload.get("choice")
|
||||
if pending.kind == "reserve_keep":
|
||||
if choice not in pending.options:
|
||||
error = "Choose one of the drawn cards"
|
||||
raise RuleError(error)
|
||||
other = next(item for item in pending.options if item != choice)
|
||||
state.decks[pending.context["deck"]].insert(0, other)
|
||||
state.pending = None
|
||||
_finish_reserve(state, choice, pack, settings)
|
||||
return
|
||||
if pending.kind == "copy_bonus":
|
||||
if choice not in pending.options:
|
||||
error = "Choose an existing bonus"
|
||||
raise RuleError(error)
|
||||
card_id = pending.context["card_id"]
|
||||
owned = next(item for item in reversed(player.cards) if item.card_id == card_id)
|
||||
owned.copied_resource = choice
|
||||
definition = pack.card(card_id)
|
||||
state.pending = None
|
||||
if definition.effect.kind == CardEffectKind.COPY_AND_CLAIM:
|
||||
_queue_free_card(state, definition, pack, settings, is_conquest=bool(pending.context.get("is_conquest")))
|
||||
else:
|
||||
_after_purchase(
|
||||
state,
|
||||
pack,
|
||||
settings,
|
||||
is_conquest=bool(pending.context.get("is_conquest")),
|
||||
chained=bool(pending.context.get("chained")),
|
||||
)
|
||||
return
|
||||
if pending.kind == "free_card":
|
||||
if choice not in pending.options:
|
||||
error = "Choose an eligible free card"
|
||||
raise RuleError(error)
|
||||
if settings.interactions.fortifications_block_free_claim:
|
||||
_assert_not_blocked(state, choice, player.seat)
|
||||
_remove_visible(state, choice)
|
||||
_return_fortifications(state, choice)
|
||||
claimed = pack.card(choice)
|
||||
player.cards.append(OwnedCard(card_id=choice))
|
||||
state.turn_chain_free_count += 1
|
||||
state.log.append(f"{player.name} claimed {claimed.label} without payment")
|
||||
state.pending = None
|
||||
if claimed.effect.kind in {CardEffectKind.COPY_BONUS, CardEffectKind.COPY_AND_CLAIM}:
|
||||
prior = player.cards.pop()
|
||||
options = sorted(resource for resource, amount in bonuses(player, pack).items() if amount)
|
||||
player.cards.append(prior)
|
||||
if options:
|
||||
state.pending = PendingChoice(
|
||||
kind="copy_bonus",
|
||||
seat=player.seat,
|
||||
options=options,
|
||||
context={
|
||||
"card_id": choice,
|
||||
"is_conquest": pending.context.get("is_conquest", False),
|
||||
"chained": True,
|
||||
},
|
||||
)
|
||||
return
|
||||
_after_purchase(state, pack, settings, is_conquest=bool(pending.context.get("is_conquest")), chained=True)
|
||||
return
|
||||
if pending.kind == "resource":
|
||||
if choice not in pending.options or state.supply.get(str(choice), 0) <= 0:
|
||||
error = "Choose an available resource"
|
||||
raise RuleError(error)
|
||||
state.supply[str(choice)] -= 1
|
||||
player.tokens[str(choice)] += 1
|
||||
context = pending.context
|
||||
state.pending = None
|
||||
remaining = int(context.get("remaining", 1))
|
||||
if context.get("next") == "after_purchase" and remaining > 1:
|
||||
options = [item for item in pack.resource_ids if state.supply.get(item, 0) > 0]
|
||||
if options:
|
||||
state.pending = PendingChoice(
|
||||
kind="resource",
|
||||
seat=player.seat,
|
||||
options=options,
|
||||
context={**context, "remaining": remaining - 1},
|
||||
)
|
||||
return
|
||||
if context.get("next") == "after_purchase":
|
||||
_after_purchase(state, pack, settings, is_conquest=bool(context.get("is_conquest")), skip_resource=True)
|
||||
else:
|
||||
_after_standard_action(state, pack, settings, action="take_double")
|
||||
return
|
||||
if pending.kind == "discard_tokens":
|
||||
discard = command.payload.get("tokens", {})
|
||||
allowed = {*pack.resource_ids, pack.wild_resource.id}
|
||||
if not isinstance(discard, dict) or not set(discard) <= allowed or sum(discard.values()) != pending.amount:
|
||||
error = "Discard exactly the required number of resources"
|
||||
raise RuleError(error)
|
||||
for resource, amount in discard.items():
|
||||
if not isinstance(amount, int) or amount <= 0 or amount > player.tokens.get(resource, 0):
|
||||
error = "Invalid token discard"
|
||||
raise RuleError(error)
|
||||
for resource, amount in discard.items():
|
||||
player.tokens[resource] -= amount
|
||||
state.supply[resource] += amount
|
||||
state.pending = None
|
||||
_end_checks(state, pack, settings)
|
||||
return
|
||||
if pending.kind == "patron":
|
||||
if choice not in pending.options:
|
||||
error = "Choose an eligible patron"
|
||||
raise RuleError(error)
|
||||
player.patrons.append(choice)
|
||||
state.available_patrons.remove(choice)
|
||||
state.pending = None
|
||||
_check_outposts_or_objectives(state, pack, settings)
|
||||
return
|
||||
if pending.kind == "outpost":
|
||||
if choice not in pending.options:
|
||||
error = "Choose an eligible outpost"
|
||||
raise RuleError(error)
|
||||
player.outposts.append(choice)
|
||||
state.pending = None
|
||||
_check_objectives(state, pack, settings)
|
||||
return
|
||||
if pending.kind == "fortification":
|
||||
target = command.payload.get("card_id")
|
||||
mode = command.payload.get("mode")
|
||||
if target not in visible_cards(state):
|
||||
error = "Choose a visible card"
|
||||
raise RuleError(error)
|
||||
if (
|
||||
target in state.markets.get(deck_key("eastern", pack.card(target).tier), [])
|
||||
and not settings.interactions.fortifications_on_eastern
|
||||
):
|
||||
error = "Fortifications cannot occupy eastern cards with this setting"
|
||||
raise RuleError(error)
|
||||
if mode == "place":
|
||||
occupants = state.fortifications.get(target, {})
|
||||
if any(seat != player.seat and count for seat, count in occupants.items()):
|
||||
error = "An opponent occupies that card"
|
||||
raise RuleError(error)
|
||||
source = command.payload.get("from_card")
|
||||
if source:
|
||||
if state.fortifications.get(source, {}).get(player.seat, 0) <= 0:
|
||||
error = "You have no fortification there"
|
||||
raise RuleError(error)
|
||||
state.fortifications[source][player.seat] -= 1
|
||||
elif player.fortifications_available > 0:
|
||||
player.fortifications_available -= 1
|
||||
else:
|
||||
error = "Move one of your placed fortifications"
|
||||
raise RuleError(error)
|
||||
state.fortifications.setdefault(target, {})[player.seat] = occupants.get(player.seat, 0) + 1
|
||||
elif mode == "remove":
|
||||
occupants = state.fortifications.get(target, {})
|
||||
opponents = [(seat, count) for seat, count in occupants.items() if seat != player.seat and count == 1]
|
||||
if len(opponents) != 1:
|
||||
error = "Choose a card with exactly one opposing fortification"
|
||||
raise RuleError(error)
|
||||
seat, _ = opponents[0]
|
||||
occupants[seat] = 0
|
||||
state.players[seat].fortifications_available += 1
|
||||
else:
|
||||
error = "Choose place or remove"
|
||||
raise RuleError(error)
|
||||
remaining = int(pending.context.get("remaining", 1))
|
||||
state.pending = None
|
||||
if remaining > 1 and _queue_fortification(state, pack, settings, remaining=remaining - 1):
|
||||
return
|
||||
_after_fortification(state, pack, settings)
|
||||
return
|
||||
if pending.kind == "conquest":
|
||||
if command.type == "decline":
|
||||
state.pending = None
|
||||
_enforce_token_limit_or_checks(state, pack, settings)
|
||||
return
|
||||
if command.type != "purchase":
|
||||
error = "Purchase the conquest card or decline"
|
||||
raise RuleError(error)
|
||||
state.pending = None
|
||||
_purchase(state, command.payload, pack, settings, is_conquest=True)
|
||||
return
|
||||
error = "Unsupported pending choice"
|
||||
raise RuleError(error)
|
||||
|
||||
|
||||
def _queue_free_card(
|
||||
state: GameState, card: CardDefinition, pack: ContentPack, settings: GameSettings, *, is_conquest: bool
|
||||
) -> None:
|
||||
options = []
|
||||
for card_id in visible_cards(state):
|
||||
candidate = pack.card(card_id)
|
||||
if candidate.tier != card.effect.target_tier:
|
||||
continue
|
||||
if settings.interactions.fortifications_block_free_claim:
|
||||
occupants = state.fortifications.get(card_id, {})
|
||||
if any(seat != state.current_seat and count for seat, count in occupants.items()):
|
||||
continue
|
||||
options.append(card_id)
|
||||
if options:
|
||||
state.pending = PendingChoice(
|
||||
kind="free_card", seat=state.current_seat, options=sorted(options), context={"is_conquest": is_conquest}
|
||||
)
|
||||
else:
|
||||
_after_purchase(state, pack, settings, is_conquest=is_conquest)
|
||||
|
||||
|
||||
def _after_purchase(
|
||||
state: GameState,
|
||||
pack: ContentPack,
|
||||
settings: GameSettings,
|
||||
*,
|
||||
is_conquest: bool,
|
||||
chained: bool = False,
|
||||
skip_resource: bool = False,
|
||||
) -> None:
|
||||
player = state.players[state.current_seat]
|
||||
_refill_markets(state, settings)
|
||||
resource_post = _owned_outpost(player, pack, OutpostPower.RESOURCE_AFTER_PURCHASE)
|
||||
triggers = not is_conquest or settings.interactions.purchase_resource_on_conquest
|
||||
if settings.modules.outposts and resource_post and triggers and not skip_resource:
|
||||
options = [item for item in pack.resource_ids if state.supply.get(item, 0) > 0]
|
||||
if options:
|
||||
remaining = 1
|
||||
if chained and not settings.interactions.chained_claim_is_not_purchase:
|
||||
remaining += state.turn_chain_free_count
|
||||
state.pending = PendingChoice(
|
||||
kind="resource",
|
||||
seat=player.seat,
|
||||
options=options,
|
||||
context={"next": "after_purchase", "is_conquest": is_conquest, "remaining": remaining},
|
||||
)
|
||||
return
|
||||
if settings.modules.fortifications:
|
||||
remaining = 1
|
||||
if chained and not settings.interactions.one_fortification_decision_per_purchase_chain:
|
||||
remaining += state.turn_chain_free_count
|
||||
if _queue_fortification(state, pack, settings, remaining=remaining):
|
||||
return
|
||||
_after_fortification(state, pack, settings)
|
||||
|
||||
|
||||
def _after_fortification(state: GameState, pack: ContentPack, settings: GameSettings) -> None:
|
||||
player = state.players[state.current_seat]
|
||||
conquest = [
|
||||
card_id
|
||||
for card_id, occupants in state.fortifications.items()
|
||||
if occupants.get(player.seat, 0) >= settings.fortifications_per_player and card_id in visible_cards(state)
|
||||
]
|
||||
if settings.modules.fortifications and conquest:
|
||||
state.turn_chain_free_count = 0
|
||||
state.pending = PendingChoice(kind="conquest", seat=player.seat, options=conquest)
|
||||
return
|
||||
state.turn_chain_free_count = 0
|
||||
_enforce_token_limit_or_checks(state, pack, settings)
|
||||
|
||||
|
||||
def _queue_fortification(
|
||||
state: GameState,
|
||||
pack: ContentPack,
|
||||
settings: GameSettings,
|
||||
*,
|
||||
remaining: int,
|
||||
) -> bool:
|
||||
player = state.players[state.current_seat]
|
||||
options = [
|
||||
card_id
|
||||
for card_id in visible_cards(state)
|
||||
if settings.interactions.fortifications_on_eastern or pack.card(card_id).deck != "eastern"
|
||||
]
|
||||
can_act = False
|
||||
for card_id in options:
|
||||
occupants = state.fortifications.get(card_id, {})
|
||||
can_place = not any(seat != player.seat and count for seat, count in occupants.items())
|
||||
can_remove = any(seat != player.seat and count == 1 for seat, count in occupants.items())
|
||||
if can_place or can_remove:
|
||||
can_act = True
|
||||
break
|
||||
if not can_act:
|
||||
return False
|
||||
state.pending = PendingChoice(
|
||||
kind="fortification",
|
||||
seat=player.seat,
|
||||
options=sorted(options),
|
||||
context={"remaining": remaining},
|
||||
)
|
||||
return True
|
||||
|
||||
|
||||
def _after_standard_action(state: GameState, pack: ContentPack, settings: GameSettings, *, action: str) -> None:
|
||||
del action
|
||||
if settings.modules.fortifications and state.turn_purchase_count:
|
||||
_after_purchase(state, pack, settings, is_conquest=False)
|
||||
else:
|
||||
_enforce_token_limit_or_checks(state, pack, settings)
|
||||
|
||||
|
||||
def _enforce_token_limit_or_checks(state: GameState, pack: ContentPack, settings: GameSettings) -> None:
|
||||
player = state.players[state.current_seat]
|
||||
excess = sum(player.tokens.values()) - settings.token_limit
|
||||
if excess > 0:
|
||||
state.pending = PendingChoice(kind="discard_tokens", seat=player.seat, amount=excess)
|
||||
return
|
||||
_end_checks(state, pack, settings)
|
||||
|
||||
|
||||
def _end_checks(state: GameState, pack: ContentPack, settings: GameSettings) -> None:
|
||||
player = state.players[state.current_seat]
|
||||
bonus = bonuses(player, pack)
|
||||
eligible_patrons = [
|
||||
patron.id
|
||||
for patron in pack.patrons
|
||||
if patron.id in state.available_patrons and requirements_met(patron.requirements, bonus)
|
||||
]
|
||||
if eligible_patrons:
|
||||
state.pending = PendingChoice(kind="patron", seat=player.seat, options=eligible_patrons)
|
||||
return
|
||||
_check_outposts_or_objectives(state, pack, settings)
|
||||
|
||||
|
||||
def _check_outposts_or_objectives(state: GameState, pack: ContentPack, settings: GameSettings) -> None:
|
||||
if settings.modules.outposts and settings.interactions.outposts_before_objectives and _queue_outpost(state, pack):
|
||||
return
|
||||
_check_objectives(state, pack, settings)
|
||||
|
||||
|
||||
def _queue_outpost(state: GameState, pack: ContentPack) -> bool:
|
||||
player = state.players[state.current_seat]
|
||||
bonus = bonuses(player, pack)
|
||||
choices = [
|
||||
outpost.id
|
||||
for outpost in pack.outposts
|
||||
if outpost.id not in player.outposts and requirements_met(outpost.requirements, bonus)
|
||||
]
|
||||
if choices:
|
||||
state.pending = PendingChoice(kind="outpost", seat=player.seat, options=choices)
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def _check_objectives(state: GameState, pack: ContentPack, settings: GameSettings) -> None:
|
||||
player = state.players[state.current_seat]
|
||||
if settings.modules.objectives:
|
||||
bonus = bonuses(player, pack)
|
||||
choices = [
|
||||
objective.id
|
||||
for objective in pack.objectives
|
||||
if objective.id in state.available_objectives
|
||||
and score(player, pack) >= objective.minimum_score
|
||||
and requirements_met(objective.requirements, bonus)
|
||||
]
|
||||
if choices:
|
||||
player.objective_met = choices[0]
|
||||
if player.seat not in state.objective_qualifiers:
|
||||
state.objective_qualifiers.append(player.seat)
|
||||
if (
|
||||
settings.modules.outposts
|
||||
and not settings.interactions.outposts_before_objectives
|
||||
and _queue_outpost(state, pack)
|
||||
):
|
||||
return
|
||||
_finish_turn(state, pack, settings)
|
||||
|
||||
|
||||
def _finish_turn(state: GameState, pack: ContentPack, settings: GameSettings) -> None:
|
||||
player = state.players[state.current_seat]
|
||||
has_score = score(player, pack) >= settings.target_score
|
||||
has_objective = player.objective_met is not None
|
||||
triggered = {
|
||||
"score": has_score,
|
||||
"objective": has_objective,
|
||||
"either": has_score or has_objective,
|
||||
"both": has_score and has_objective,
|
||||
}[settings.victory_condition]
|
||||
if triggered and state.finish_at_seat is None:
|
||||
state.finish_at_seat = (state.first_seat - 1) % len(state.players)
|
||||
if state.finish_at_seat == state.current_seat:
|
||||
candidates = (
|
||||
state.objective_qualifiers if settings.victory_condition == "objective" else list(range(len(state.players)))
|
||||
)
|
||||
if not candidates:
|
||||
candidates = list(range(len(state.players)))
|
||||
best_score = max(score(state.players[seat], pack) for seat in candidates)
|
||||
candidates = [seat for seat in candidates if score(state.players[seat], pack) == best_score]
|
||||
fewest = min(state.players[seat].purchased_card_count for seat in candidates)
|
||||
state.winners = [seat for seat in candidates if state.players[seat].purchased_card_count == fewest]
|
||||
state.finished = True
|
||||
state.pending = None
|
||||
return
|
||||
state.current_seat = (state.current_seat + 1) % len(state.players)
|
||||
if state.current_seat == state.first_seat:
|
||||
state.round_number += 1
|
||||
state.turn_purchase_count = 0
|
||||
state.turn_chain_free_count = 0
|
||||
|
||||
|
||||
def _remove_visible(state: GameState, card_id: str) -> None:
|
||||
for market in state.markets.values():
|
||||
if card_id in market:
|
||||
market.remove(card_id)
|
||||
return
|
||||
error = "Card is not visible"
|
||||
raise RuleError(error)
|
||||
|
||||
|
||||
def _refill_markets(state: GameState, settings: GameSettings) -> None:
|
||||
for key, market in state.markets.items():
|
||||
desired = settings.eastern_market_size if key.startswith("eastern:") else settings.base_market_size
|
||||
deck = state.decks[key]
|
||||
while len(market) < desired and deck:
|
||||
market.append(deck.pop())
|
||||
|
||||
|
||||
def _assert_not_blocked(state: GameState, card_id: str, seat: int) -> None:
|
||||
occupants = state.fortifications.get(card_id, {})
|
||||
if any(owner != seat and count > 0 for owner, count in occupants.items()):
|
||||
error = "An opponent's fortification protects that card"
|
||||
raise RuleError(error)
|
||||
|
||||
|
||||
def _return_fortifications(state: GameState, card_id: str) -> None:
|
||||
for seat, count in state.fortifications.pop(card_id, {}).items():
|
||||
state.players[seat].fortifications_available += count
|
||||
|
||||
|
||||
def _owned_outpost(player: PlayerState, pack: ContentPack, power: OutpostPower) -> OutpostDefinition | None:
|
||||
definitions = {item.id: item for item in pack.outposts}
|
||||
return next((definitions[item] for item in player.outposts if definitions[item].power == power), None)
|
||||
|
||||
|
||||
def _reconcile_outposts(player: PlayerState, pack: ContentPack) -> None:
|
||||
definitions = {item.id: item for item in pack.outposts}
|
||||
bonus = bonuses(player, pack)
|
||||
player.outposts = [item for item in player.outposts if requirements_met(definitions[item].requirements, bonus)]
|
||||
|
||||
|
||||
def public_state(state: GameState, viewer_seat: int | None) -> dict[str, Any]:
|
||||
"""Serialize state while redacting other players' reserved cards and deck order."""
|
||||
data = state.model_dump(mode="json")
|
||||
for player in data["players"]:
|
||||
if player["seat"] != viewer_seat:
|
||||
player["reserved"] = [None] * len(player["reserved"])
|
||||
data["decks"] = {key: len(value) for key, value in state.decks.items()}
|
||||
if state.pending and state.pending.seat != viewer_seat:
|
||||
data["pending"]["options"] = []
|
||||
data["pending"]["context"] = {}
|
||||
return data
|
||||
@@ -1,159 +0,0 @@
|
||||
"""Legal command generation used by bots and server-rendered controls."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import itertools
|
||||
import math
|
||||
import uuid
|
||||
|
||||
from .engine import bonuses, visible_cards
|
||||
from .models import ContentPack, GameCommand, GameSettings, GameState, OutpostPower
|
||||
|
||||
DISTINCT_TAKE_COUNT = 3
|
||||
DOUBLE_TAKE_MINIMUM = 4
|
||||
|
||||
|
||||
def command(state: GameState, type_: str, payload: dict | None = None) -> GameCommand:
|
||||
"""Create a command targeting the state's current revision."""
|
||||
return GameCommand(
|
||||
command_id=uuid.uuid4().hex,
|
||||
expected_revision=state.revision,
|
||||
type=type_, # type: ignore[arg-type]
|
||||
payload=payload or {},
|
||||
)
|
||||
|
||||
|
||||
def legal_commands( # noqa: C901, PLR0911, PLR0912 - mirrors pending and normal action branches
|
||||
state: GameState, pack: ContentPack, settings: GameSettings, seat: int
|
||||
) -> list[GameCommand]:
|
||||
"""Enumerate meaningful legal commands without revealing hidden information."""
|
||||
if state.finished or seat != state.current_seat:
|
||||
return []
|
||||
player = state.players[seat]
|
||||
if state.pending:
|
||||
pending = state.pending
|
||||
if pending.seat != seat:
|
||||
return []
|
||||
if pending.kind == "discard_tokens":
|
||||
available = [item for item, amount in player.tokens.items() for _ in range(amount)]
|
||||
payload: dict[str, int] = {}
|
||||
for item in available[: pending.amount]:
|
||||
payload[item] = payload.get(item, 0) + 1
|
||||
return [command(state, "choose", {"tokens": payload})]
|
||||
if pending.kind == "fortification":
|
||||
pending_results = []
|
||||
for target in pending.options:
|
||||
occupants = state.fortifications.get(target, {})
|
||||
if not any(owner != seat and count for owner, count in occupants.items()):
|
||||
if player.fortifications_available:
|
||||
pending_results.append(command(state, "choose", {"card_id": target, "mode": "place"}))
|
||||
for source, source_occupants in state.fortifications.items():
|
||||
if source_occupants.get(seat, 0):
|
||||
pending_results.append(
|
||||
command(state, "choose", {"card_id": target, "from_card": source, "mode": "place"})
|
||||
)
|
||||
if any(owner != seat and count == 1 for owner, count in occupants.items()):
|
||||
pending_results.append(command(state, "choose", {"card_id": target, "mode": "remove"}))
|
||||
return pending_results
|
||||
if pending.kind == "conquest":
|
||||
conquest_results = [command(state, "decline")]
|
||||
for card_id in pending.options:
|
||||
payment = default_payment(state, pack, settings, card_id, seat)
|
||||
if payment is not None:
|
||||
conquest_results.append(command(state, "purchase", {"card_id": card_id, **payment}))
|
||||
return conquest_results
|
||||
return [command(state, "choose", {"choice": item}) for item in pending.options]
|
||||
|
||||
results: list[GameCommand] = []
|
||||
available = [resource for resource in pack.resource_ids if state.supply.get(resource, 0)]
|
||||
take_size = DISTINCT_TAKE_COUNT if len(available) >= DISTINCT_TAKE_COUNT else 1
|
||||
sizes = [take_size] if len(available) >= DISTINCT_TAKE_COUNT else list(range(1, len(available) + 1))
|
||||
for size in sizes:
|
||||
results.extend(
|
||||
command(state, "take_distinct", {"resources": list(items)})
|
||||
for items in itertools.combinations(available, size)
|
||||
)
|
||||
results.extend(
|
||||
command(state, "take_double", {"resource": resource})
|
||||
for resource in pack.resource_ids
|
||||
if state.supply.get(resource, 0) >= DOUBLE_TAKE_MINIMUM
|
||||
)
|
||||
if len(player.reserved) < settings.reserve_limit:
|
||||
results.extend(
|
||||
command(state, "reserve", {"card_id": card_id})
|
||||
for card_id in visible_cards(state)
|
||||
if not _blocked(state, card_id, seat)
|
||||
)
|
||||
results.extend(command(state, "reserve", {"deck": key}) for key, cards in state.decks.items() if cards)
|
||||
for card_id in [*visible_cards(state), *player.reserved]:
|
||||
if card_id not in player.reserved and _blocked(state, card_id, seat):
|
||||
continue
|
||||
payment = default_payment(state, pack, settings, card_id, seat)
|
||||
if payment is not None:
|
||||
results.append(command(state, "purchase", {"card_id": card_id, **payment}))
|
||||
return results
|
||||
|
||||
|
||||
def default_payment( # noqa: C901, PLR0912 - ordered payment rules are intentionally explicit
|
||||
state: GameState,
|
||||
pack: ContentPack,
|
||||
settings: GameSettings,
|
||||
card_id: str,
|
||||
seat: int,
|
||||
) -> dict | None:
|
||||
"""Return one legal colored-first payment, or None when unaffordable."""
|
||||
player = state.players[seat]
|
||||
card = pack.card(card_id)
|
||||
if card.alternate_cost:
|
||||
matching = []
|
||||
definitions = {item.id: item for item in pack.cards}
|
||||
for owned in player.cards:
|
||||
resource = owned.copied_resource or definitions[owned.card_id].bonus_resource
|
||||
if resource == card.alternate_cost.discard_resource:
|
||||
matching.append(owned.card_id)
|
||||
copies = [
|
||||
item for item in matching if next(owned for owned in player.cards if owned.card_id == item).copied_resource
|
||||
]
|
||||
ordered = [*copies, *(item for item in matching if item not in copies)]
|
||||
if len(ordered) < card.alternate_cost.count:
|
||||
return None
|
||||
return {"discard_cards": ordered[: card.alternate_cost.count]}
|
||||
discount = bonuses(player, pack)
|
||||
payment: dict[str, int] = {}
|
||||
shortage = 0
|
||||
for resource in pack.resource_ids:
|
||||
due = max(0, card.cost.get(resource, 0) - discount.get(resource, 0))
|
||||
amount = min(due, player.tokens.get(resource, 0))
|
||||
if amount:
|
||||
payment[resource] = amount
|
||||
shortage += due - amount
|
||||
double_wild = any(
|
||||
outpost.power == OutpostPower.DOUBLE_WILD and outpost.id in player.outposts for outpost in pack.outposts
|
||||
)
|
||||
wild_value = 2 if double_wild else 1
|
||||
virtual_multiplier = 2 if double_wild and settings.interactions.virtual_wild_can_double else 1
|
||||
virtual_cards: list[str] = []
|
||||
virtual_total = 0
|
||||
if shortage > player.tokens.get(pack.wild_resource.id, 0) * wild_value:
|
||||
for owned in player.cards:
|
||||
definition = pack.card(owned.card_id)
|
||||
if definition.effect.kind.value != "virtual_wild":
|
||||
continue
|
||||
virtual_cards.append(owned.card_id)
|
||||
virtual_total += definition.effect.amount * virtual_multiplier
|
||||
if virtual_total + player.tokens.get(pack.wild_resource.id, 0) * wild_value >= shortage:
|
||||
break
|
||||
remaining = max(0, shortage - virtual_total)
|
||||
wild_needed = math.ceil(remaining / wild_value)
|
||||
if wild_needed > player.tokens.get(pack.wild_resource.id, 0) or virtual_total + wild_needed * wild_value < shortage:
|
||||
return None
|
||||
if wild_needed:
|
||||
payment[pack.wild_resource.id] = wild_needed
|
||||
result: dict[str, object] = {"payment": payment}
|
||||
if virtual_cards:
|
||||
result["virtual_wild_cards"] = virtual_cards
|
||||
return result
|
||||
|
||||
|
||||
def _blocked(state: GameState, card_id: str, seat: int) -> bool:
|
||||
return any(owner != seat and count for owner, count in state.fortifications.get(card_id, {}).items())
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user