From e29cec292c922824499cfd5e7d68fae8901f7279 Mon Sep 17 00:00:00 2001 From: Zach Flanders Date: Tue, 28 Jul 2026 14:50:35 -0500 Subject: [PATCH 1/2] update db router to ensure llm models always use the same database for read/write --- usaspending_api/routers/replicas.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/usaspending_api/routers/replicas.py b/usaspending_api/routers/replicas.py index 931774088c..2b3d95d8e5 100644 --- a/usaspending_api/routers/replicas.py +++ b/usaspending_api/routers/replicas.py @@ -1,6 +1,8 @@ from django.db import DEFAULT_DB_ALIAS +from django.db.models import Model from usaspending_api.download.models.download_job import DownloadJob +from usaspending_api.llm.models.db_models import AIModel, Message, Prompts, Session, ToolUse from usaspending_api.references.models import FilterHash @@ -17,12 +19,12 @@ class ReadReplicaRouter: def __init__(self): self.usaspending_databases = [self.writable_database] + self.read_replicas - def db_for_read(self, model, **hints) -> str: + def db_for_read(self, model: Model, **hints) -> str: """ FilterHash and DownloadJob are writable tables so always read from source (default) to mitigate replication lag. Otherwise, choose a connection randomly. """ - if model in [FilterHash, DownloadJob]: + if model in [FilterHash, DownloadJob, AIModel, Message, Prompts, Session, ToolUse]: return self.writable_database if len(self.read_replicas) > 0: @@ -30,14 +32,14 @@ def db_for_read(self, model, **hints) -> str: else: return self.writable_database - def db_for_write(self, model, **hints) -> str: + def db_for_write(self, model: Model, **hints) -> str: return self.writable_database - def allow_relation(self, obj1, obj2, **hints): + def allow_relation(self, obj1: Model, obj2: Model, **hints) -> bool: """Relations are currently only allowed in USAspending. Cross database relations are not allowed.""" return obj1._state.db in self.usaspending_databases and obj2._state.db == obj1._state.db - def allow_migrate(self, db, app_label, model_name=None, **hints): + def allow_migrate(self, db: str, app_label: str, model_name: str | None = None, **hints) -> bool: """Migrations should only run in USAspending against the writable database.""" return db == self.writable_database From 14980640e480e4e246fa4f36d4844cbdc5f20232 Mon Sep 17 00:00:00 2001 From: Zach Flanders Date: Tue, 28 Jul 2026 14:58:20 -0500 Subject: [PATCH 2/2] Update comment to reference the LLM models --- usaspending_api/routers/replicas.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/usaspending_api/routers/replicas.py b/usaspending_api/routers/replicas.py index 2b3d95d8e5..865ef5544c 100644 --- a/usaspending_api/routers/replicas.py +++ b/usaspending_api/routers/replicas.py @@ -21,7 +21,7 @@ def __init__(self): def db_for_read(self, model: Model, **hints) -> str: """ - FilterHash and DownloadJob are writable tables so always read from source (default) to + FilterHash, DownloadJob, and the LLM app models are writable tables so always read from source (default) to mitigate replication lag. Otherwise, choose a connection randomly. """ if model in [FilterHash, DownloadJob, AIModel, Message, Prompts, Session, ToolUse]: