From d62f2f5154434c7cf41b070a8487b9f0a5734778 Mon Sep 17 00:00:00 2001
From: Max Chis <maxachis@gmail.com>
Date: Sun, 5 Oct 2025 17:13:01 -0400
Subject: [PATCH] Add standalone route for agency suggestions

---
 .../annotate/all/get/queries/agency/core.py   |  7 ++-
 .../all/get/queries/agency/requester.py       | 44 +++++++++++++++++--
 src/api/endpoints/annotate/routes.py          | 16 +++++++
 3 files changed, 61 insertions(+), 6 deletions(-)

diff --git a/src/api/endpoints/annotate/all/get/queries/agency/core.py b/src/api/endpoints/annotate/all/get/queries/agency/core.py
index 236aae88..28cfbd2d 100644
--- a/src/api/endpoints/annotate/all/get/queries/agency/core.py
+++ b/src/api/endpoints/annotate/all/get/queries/agency/core.py
@@ -16,15 +16,18 @@ class GetAgencySuggestionsQueryBuilder(QueryBuilderBase):
 
     def __init__(
         self,
-        url_id: int
+        url_id: int,
+        location_id: int | None = None
     ):
         super().__init__()
         self.url_id = url_id
+        self.location_id = location_id
 
     async def run(self, session: AsyncSession) -> AgencyAnnotationResponseOuterInfo:
         requester = GetAgencySuggestionsRequester(
             session,
-            url_id=self.url_id
+            url_id=self.url_id,
+            location_id=self.location_id
         )
 
         user_suggestions: list[AgencyAnnotationUserSuggestion] = \
diff --git a/src/api/endpoints/annotate/all/get/queries/agency/requester.py b/src/api/endpoints/annotate/all/get/queries/agency/requester.py
index bec13508..fc309e50 100644
--- a/src/api/endpoints/annotate/all/get/queries/agency/requester.py
+++ b/src/api/endpoints/annotate/all/get/queries/agency/requester.py
@@ -9,6 +9,7 @@
     SuggestionsWithHighestConfidenceCTE
 from src.db.helpers.session import session_helper as sh
 from src.db.models.impl.agency.sqlalchemy import Agency
+from src.db.models.impl.link.agency_location.sqlalchemy import LinkAgencyLocation
 from src.db.models.impl.link.user_suggestion_not_found.agency.sqlalchemy import LinkUserSuggestionAgencyNotFound
 from src.db.models.impl.url.suggestion.agency.user import UserUrlAgencySuggestion
 from src.db.templates.requester import RequesterBase
@@ -16,9 +17,15 @@
 
 class GetAgencySuggestionsRequester(RequesterBase):
 
-    def __init__(self, session: AsyncSession, url_id: int):
+    def __init__(
+        self,
+        session: AsyncSession,
+        url_id: int,
+        location_id: int
+    ):
         super().__init__(session)
         self.url_id = url_id
+        self.location_id = location_id
 
     async def get_user_agency_suggestions(self) -> list[AgencyAnnotationUserSuggestion]:
         query = (
@@ -31,7 +38,22 @@ async def get_user_agency_suggestions(self) -> list[AgencyAnnotationUserSuggesti
                 Agency,
                 Agency.agency_id == UserUrlAgencySuggestion.agency_id
             )
-            .where(
+
+        )
+
+        if self.location_id is not None:
+            query = (
+                query.join(
+                    LinkAgencyLocation,
+                    LinkAgencyLocation.agency_id == UserUrlAgencySuggestion.agency_id
+                )
+                .where(
+                    LinkAgencyLocation.location_id == self.location_id
+                )
+            )
+
+        query = (
+            query.where(
                 UserUrlAgencySuggestion.url_id == self.url_id
             )
             .group_by(
@@ -64,11 +86,25 @@ async def get_auto_agency_suggestions(self) -> list[AgencyAnnotationAutoSuggesti
                 cte.confidence,
                 Agency.name.label("agency_name"),
             )
-            .outerjoin(
+            .join(
                 Agency,
                 Agency.agency_id == cte.agency_id
             )
-            .where(
+        )
+
+        if self.location_id is not None:
+            query = (
+                query.join(
+                    LinkAgencyLocation,
+                    LinkAgencyLocation.agency_id == cte.agency_id
+                )
+                .where(
+                    LinkAgencyLocation.location_id == self.location_id
+                )
+            )
+
+        query = (
+            query.where(
                 cte.url_id == self.url_id
             )
             .order_by(
diff --git a/src/api/endpoints/annotate/routes.py b/src/api/endpoints/annotate/routes.py
index 50798990..6972314d 100644
--- a/src/api/endpoints/annotate/routes.py
+++ b/src/api/endpoints/annotate/routes.py
@@ -1,7 +1,9 @@
 from fastapi import APIRouter, Depends, Query
 
 from src.api.dependencies import get_async_core
+from src.api.endpoints.annotate.all.get.models.agency import AgencyAnnotationResponseOuterInfo
 from src.api.endpoints.annotate.all.get.models.response import GetNextURLForAllAnnotationResponse
+from src.api.endpoints.annotate.all.get.queries.agency.core import GetAgencySuggestionsQueryBuilder
 from src.api.endpoints.annotate.all.post.models.request import AllAnnotationPostInfo
 from src.core.core import AsyncCore
 from src.security.dtos.access_info import AccessInfo
@@ -59,4 +61,18 @@ async def annotate_url_for_all_annotations_and_get_next_url(
         batch_id=batch_id,
         user_id=access_info.user_id,
         url_id=anno_url_id
+    )
+
+@annotate_router.get("/suggestions/agencies/{url_id}")
+async def get_agency_suggestions(
+    url_id: int,
+    async_core: AsyncCore = Depends(get_async_core),
+    access_info: AccessInfo = Depends(get_access_info),
+    location_id: int | None = Query(default=None)
+) -> AgencyAnnotationResponseOuterInfo:
+    return await async_core.adb_client.run_query_builder(
+        GetAgencySuggestionsQueryBuilder(
+            url_id=url_id,
+            location_id=location_id
+        )
     )
\ No newline at end of file