Add batch link subtasks for location/agency id tasks

ENV.md

@@ -102,6 +102,7 @@ Agency ID Subtasks are collectively disabled by the `URL_AGENCY_IDENTIFICATION_T
 | `AGENCY_ID_NLP_LOCATION_MATCH_FLAG` | Enables the NLP location match subtask for agency identification. |
 | `AGENCY_ID_CKAN_FLAG`               | Enables the CKAN subtask for agency identification.               |
 | `AGENCY_ID_MUCKROCK_FLAG`           | Enables the MuckRock subtask for agency identification.           |
+| `AGENCY_ID_BATCH_LINK_FLAG` | Enables the Batch Link subtask for agency identification. |
 
 
 ### Location ID Subtasks
@@ -111,6 +112,7 @@ Location ID Subtasks are collectively disabled by the `URL_LOCATION_IDENTIFICATI
 | Flag                                  | Description                                                         | 
 |---------------------------------------|---------------------------------------------------------------------| 
 | `LOCATION_ID_NLP_LOCATION_MATCH_FLAG` | Enables the NLP location match subtask for location identification. |
+| `LOCATION_ID_BATCH_LINK_FLAG` | Enables the Batch Link subtask for location identification. |
 
 
 ## Foreign Data Wrapper (FDW)

alembic/versions/2025_10_11_1438-8b2adc95c5d7_add_batch_link_subtasks.py

@@ -0,0 +1,34 @@
+"""Add batch link subtasks
+
+Revision ID: 8b2adc95c5d7
+Revises: 7c4049508bfc
+Create Date: 2025-10-11 14:38:01.874040
+
+"""
+from typing import Sequence, Union
+
+from alembic import op
+import sqlalchemy as sa
+
+from src.util.alembic_helpers import add_enum_value
+
+# revision identifiers, used by Alembic.
+revision: str = '8b2adc95c5d7'
+down_revision: Union[str, None] = '7c4049508bfc'
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    add_enum_value(
+        enum_name="agency_auto_suggestion_method",
+        enum_value="batch_link"
+    )
+    add_enum_value(
+        enum_name="auto_location_id_subtask_type",
+        enum_value="batch_link"
+    )
+
+
+def downgrade() -> None:
+    pass

src/core/tasks/url/operators/agency_identification/subtasks/flags/mappings.py

@@ -4,5 +4,6 @@
     AutoAgencyIDSubtaskType.HOMEPAGE_MATCH: "AGENCY_ID_HOMEPAGE_MATCH_FLAG",
     AutoAgencyIDSubtaskType.NLP_LOCATION_MATCH: "AGENCY_ID_NLP_LOCATION_MATCH_FLAG",
     AutoAgencyIDSubtaskType.CKAN: "AGENCY_ID_CKAN_FLAG",
-    AutoAgencyIDSubtaskType.MUCKROCK: "AGENCY_ID_MUCKROCK_FLAG"
+    AutoAgencyIDSubtaskType.MUCKROCK: "AGENCY_ID_MUCKROCK_FLAG",
+    AutoAgencyIDSubtaskType.BATCH_LINK: "AGENCY_ID_BATCH_LINK_FLAG"
 }

src/core/tasks/url/operators/agency_identification/subtasks/impl/batch_link/core.py

@@ -0,0 +1,48 @@
+from src.core.tasks.url.operators.agency_identification.subtasks.impl.batch_link.params import \
+    AgencyBatchLinkSubtaskParams
+from src.core.tasks.url.operators.agency_identification.subtasks.impl.batch_link.query import \
+    GetLocationBatchLinkSubtaskParamsQueryBuilder
+from src.core.tasks.url.operators.agency_identification.subtasks.models.subtask import AutoAgencyIDSubtaskData
+from src.core.tasks.url.operators.agency_identification.subtasks.models.suggestion import AgencySuggestion
+from src.core.tasks.url.operators.agency_identification.subtasks.templates.subtask import AgencyIDSubtaskOperatorBase
+from src.db.client.async_ import AsyncDatabaseClient
+from src.db.models.impl.url.suggestion.agency.subtask.enum import AutoAgencyIDSubtaskType
+from src.db.models.impl.url.suggestion.agency.subtask.pydantic import URLAutoAgencyIDSubtaskPydantic
+
+
+class AgencyBatchLinkSubtaskOperator(AgencyIDSubtaskOperatorBase):
+
+    def __init__(
+        self,
+        adb_client: AsyncDatabaseClient,
+        task_id: int
+    ):
+        super().__init__(adb_client=adb_client, task_id=task_id)
+
+    async def inner_logic(self) -> None:
+        params: list[AgencyBatchLinkSubtaskParams] = await self._get_params()
+        self.linked_urls = [param.url_id for param in params]
+        subtask_data_list: list[AutoAgencyIDSubtaskData] = []
+        for param in params:
+            subtask_data: AutoAgencyIDSubtaskData = AutoAgencyIDSubtaskData(
+                pydantic_model=URLAutoAgencyIDSubtaskPydantic(
+                    task_id=self.task_id,
+                    url_id=param.url_id,
+                    type=AutoAgencyIDSubtaskType.BATCH_LINK,
+                    agencies_found=True,
+                ),
+                suggestions=[
+                    AgencySuggestion(
+                        agency_id=param.agency_id,
+                        confidence=80,
+                    )
+                ],
+            )
+            subtask_data_list.append(subtask_data)
+
+        await self._upload_subtask_data(subtask_data_list)
+
+    async def _get_params(self) -> list[AgencyBatchLinkSubtaskParams]:
+        return await self.adb_client.run_query_builder(
+            GetLocationBatchLinkSubtaskParamsQueryBuilder()
+        )

src/core/tasks/url/operators/agency_identification/subtasks/impl/batch_link/params.py

@@ -0,0 +1,6 @@
+from pydantic import BaseModel
+
+
+class AgencyBatchLinkSubtaskParams(BaseModel):
+    url_id: int
+    agency_id: int

src/core/tasks/url/operators/agency_identification/subtasks/impl/batch_link/query.py

@@ -0,0 +1,45 @@
+from typing import Sequence
+
+from sqlalchemy import select, RowMapping
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from src.core.tasks.url.operators.agency_identification.subtasks.impl.batch_link.params import \
+    AgencyBatchLinkSubtaskParams
+from src.core.tasks.url.operators.agency_identification.subtasks.queries.survey.queries.ctes.eligible import \
+    EligibleContainer
+from src.db.models.impl.link.agency_batch.sqlalchemy import LinkAgencyBatch
+from src.db.models.impl.link.batch_url.sqlalchemy import LinkBatchURL
+from src.db.queries.base.builder import QueryBuilderBase
+from src.db.helpers.session import session_helper as sh
+
+class GetLocationBatchLinkSubtaskParamsQueryBuilder(QueryBuilderBase):
+
+    async def run(self, session: AsyncSession) -> list[AgencyBatchLinkSubtaskParams]:
+        container = EligibleContainer()
+        query = (
+            select(
+                container.url_id,
+                LinkAgencyBatch.agency_id,
+            )
+            .select_from(container.cte)
+            .join(
+                LinkBatchURL,
+                LinkBatchURL.url_id == container.url_id,
+            )
+            .join(
+                LinkAgencyBatch,
+                LinkAgencyBatch.batch_id == LinkBatchURL.batch_id,
+            )
+            .where(
+                container.batch_link,
+            )
+            .limit(500)
+        )
+        results: Sequence[RowMapping] = await sh.mappings(session, query=query)
+        return [
+            AgencyBatchLinkSubtaskParams(
+                url_id=mapping["id"],
+                agency_id=mapping["agency_id"],
+            )
+            for mapping in results
+        ]

src/core/tasks/url/operators/agency_identification/subtasks/impl/ckan_/query.py

@@ -3,13 +3,10 @@
 from sqlalchemy import select, RowMapping
 from sqlalchemy.ext.asyncio import AsyncSession
 
-from src.collectors.enums import CollectorType
 from src.core.tasks.url.operators.agency_identification.subtasks.impl.ckan_.params import CKANAgencyIDSubtaskParams
 from src.core.tasks.url.operators.agency_identification.subtasks.queries.survey.queries.ctes.eligible import \
     EligibleContainer
 from src.db.helpers.session import session_helper as sh
-from src.db.models.impl.batch.sqlalchemy import Batch
-from src.db.models.impl.link.batch_url.sqlalchemy import LinkBatchURL
 from src.db.models.impl.url.core.sqlalchemy import URL
 from src.db.queries.base.builder import QueryBuilderBase
 

src/core/tasks/url/operators/agency_identification/subtasks/loader.py

@@ -1,4 +1,6 @@
 from src.collectors.impl.muckrock.api_interface.core import MuckrockAPIInterface
+from src.core.tasks.url.operators.agency_identification.subtasks.impl.batch_link.core import \
+    AgencyBatchLinkSubtaskOperator
 from src.core.tasks.url.operators.agency_identification.subtasks.impl.ckan_.core import CKANAgencyIDSubtaskOperator
 from src.core.tasks.url.operators.agency_identification.subtasks.impl.homepage_match_.core import \
     HomepageMatchSubtaskOperator
@@ -52,6 +54,15 @@ def _load_nlp_location_match_subtask(self, task_id: int) -> NLPLocationMatchSubt
             adb_client=self.adb_client,
         )
 
+    def _load_batch_link_subtask(
+        self,
+        task_id: int
+    ) -> AgencyBatchLinkSubtaskOperator:
+        return AgencyBatchLinkSubtaskOperator(
+            task_id=task_id,
+            adb_client=self.adb_client,
+        )
+
 
     async def load_subtask(
         self,
@@ -68,4 +79,6 @@ async def load_subtask(
                 return self._load_nlp_location_match_subtask(task_id)
             case AutoAgencyIDSubtaskType.HOMEPAGE_MATCH:
                 return self._load_homepage_match_subtask(task_id)
+            case AutoAgencyIDSubtaskType.BATCH_LINK:
+                return self._load_batch_link_subtask(task_id)
         raise ValueError(f"Unknown subtask type: {subtask_type}")

src/core/tasks/url/operators/agency_identification/subtasks/queries/survey/constants.py

@@ -5,7 +5,8 @@
     AutoAgencyIDSubtaskType.CKAN,
     AutoAgencyIDSubtaskType.MUCKROCK,
     AutoAgencyIDSubtaskType.HOMEPAGE_MATCH,
-    AutoAgencyIDSubtaskType.NLP_LOCATION_MATCH
+    AutoAgencyIDSubtaskType.NLP_LOCATION_MATCH,
+    AutoAgencyIDSubtaskType.BATCH_LINK
 ]
 
 SUBTASK_HIERARCHY_MAPPING: dict[AutoAgencyIDSubtaskType, int] = {

src/core/tasks/url/operators/agency_identification/subtasks/queries/survey/queries/ctes/eligible.py

@@ -4,6 +4,8 @@
     HIGH_CONFIDENCE_ANNOTATIONS_EXISTS_CONTAINER
 from src.core.tasks.url.operators._shared.ctes.validated import \
     VALIDATED_EXISTS_CONTAINER
+from src.core.tasks.url.operators.agency_identification.subtasks.queries.survey.queries.ctes.subtask.impl.batch_link import \
+    BATCH_LINK_SUBTASK_CONTAINER
 from src.core.tasks.url.operators.agency_identification.subtasks.queries.survey.queries.ctes.subtask.impl.ckan import \
     CKAN_SUBTASK_CONTAINER
 from src.core.tasks.url.operators.agency_identification.subtasks.queries.survey.queries.ctes.subtask.impl.homepage import \
@@ -24,6 +26,7 @@
                 MUCKROCK_SUBTASK_CONTAINER.eligible_query.label("muckrock"),
                 HOMEPAGE_SUBTASK_CONTAINER.eligible_query.label("homepage"),
                 NLP_LOCATION_CONTAINER.eligible_query.label("nlp_location"),
+                BATCH_LINK_SUBTASK_CONTAINER.eligible_query.label("batch_link"),
             )
             .where(
                 HIGH_CONFIDENCE_ANNOTATIONS_EXISTS_CONTAINER.not_exists_query,
@@ -44,6 +47,10 @@
     def ckan(self) -> Column[bool]:
         return self._cte.c['ckan']
 
+    @property
+    def batch_link(self) -> Column[bool]:
+        return self._cte.c['batch_link']
+
     @property
     def muckrock(self) -> Column[bool]:
         return self._cte.c['muckrock']

src/core/tasks/url/operators/agency_identification/subtasks/queries/survey/queries/ctes/subtask/impl/batch_link.py

@@ -0,0 +1,31 @@
+from sqlalchemy import select
+
+from src.core.tasks.url.operators._shared.container.subtask.eligible import URLsSubtaskEligibleCTEContainer
+from src.core.tasks.url.operators.agency_identification.subtasks.queries.survey.queries.ctes.subtask.helpers import \
+    get_exists_subtask_query
+from src.db.models.impl.link.agency_batch.sqlalchemy import LinkAgencyBatch
+from src.db.models.impl.link.batch_url.sqlalchemy import LinkBatchURL
+from src.db.models.impl.url.core.sqlalchemy import URL
+from src.db.models.impl.url.suggestion.agency.subtask.enum import AutoAgencyIDSubtaskType
+
+cte = (
+    select(
+        URL.id,
+        get_exists_subtask_query(
+            AutoAgencyIDSubtaskType.BATCH_LINK,
+        )
+    )
+    .join(
+        LinkBatchURL,
+        LinkBatchURL.url_id == URL.id,
+    )
+    .join(
+        LinkAgencyBatch,
+        LinkAgencyBatch.batch_id == LinkBatchURL.batch_id,
+    )
+    .cte("batch_link_eligible")
+)
+
+BATCH_LINK_SUBTASK_CONTAINER = URLsSubtaskEligibleCTEContainer(
+    cte,
+)

src/core/tasks/url/operators/agency_identification/subtasks/queries/survey/queries/eligible_counts.py

@@ -21,5 +21,6 @@ def sum_count(col: ColumnElement[bool], subtask_type: AutoAgencyIDSubtaskType) -
         sum_count(container.muckrock, AutoAgencyIDSubtaskType.MUCKROCK),
         sum_count(container.homepage, AutoAgencyIDSubtaskType.HOMEPAGE_MATCH),
         sum_count(container.nlp_location, AutoAgencyIDSubtaskType.NLP_LOCATION_MATCH),
+        sum_count(container.batch_link, AutoAgencyIDSubtaskType.BATCH_LINK)
     )
 )

src/core/tasks/url/operators/location_id/subtasks/flags/mappings.py

@@ -2,4 +2,5 @@
 
 SUBTASK_TO_ENV_FLAG: dict[LocationIDSubtaskType, str] = {
     LocationIDSubtaskType.NLP_LOCATION_FREQUENCY: "LOCATION_ID_NLP_LOCATION_MATCH_FLAG",
+    LocationIDSubtaskType.BATCH_LINK: "LOCATION_ID_BATCH_LINK_FLAG",
 }

src/core/tasks/url/operators/location_id/subtasks/impl/batch_link/core.py

@@ -0,0 +1,56 @@
+from src.core.tasks.url.operators.location_id.subtasks.impl.batch_link.inputs import LocationBatchLinkInput
+from src.core.tasks.url.operators.location_id.subtasks.impl.batch_link.query import GetLocationBatchLinkQueryBuilder
+from src.core.tasks.url.operators.location_id.subtasks.impl.nlp_location_freq.constants import ITERATIONS_PER_SUBTASK
+from src.core.tasks.url.operators.location_id.subtasks.models.subtask import AutoLocationIDSubtaskData
+from src.core.tasks.url.operators.location_id.subtasks.models.suggestion import LocationSuggestion
+from src.core.tasks.url.operators.location_id.subtasks.templates.subtask import LocationIDSubtaskOperatorBase
+from src.db.client.async_ import AsyncDatabaseClient
+from src.db.models.impl.url.suggestion.location.auto.subtask.enums import LocationIDSubtaskType
+from src.db.models.impl.url.suggestion.location.auto.subtask.pydantic import AutoLocationIDSubtaskPydantic
+
+
+class LocationBatchLinkSubtaskOperator(LocationIDSubtaskOperatorBase):
+
+    def __init__(
+        self,
+        task_id: int,
+        adb_client: AsyncDatabaseClient,
+    ):
+        super().__init__(adb_client=adb_client, task_id=task_id)
+
+    async def inner_logic(self) -> None:
+        for iteration in range(ITERATIONS_PER_SUBTASK):
+            inputs: list[LocationBatchLinkInput] = await self._get_from_db()
+            if len(inputs) == 0:
+                break
+            await self.run_subtask_iteration(inputs)
+
+    async def run_subtask_iteration(
+        self,
+        inputs: list[LocationBatchLinkInput]
+    ) -> None:
+        self.linked_urls.extend([input_.url_id for input_ in inputs])
+        subtask_data_list: list[AutoLocationIDSubtaskData] = []
+        for input_ in inputs:
+            subtask_data_list.append(
+                AutoLocationIDSubtaskData(
+                    pydantic_model=AutoLocationIDSubtaskPydantic(
+                        url_id=input_.url_id,
+                        task_id=self.task_id,
+                        locations_found=True,
+                        type=LocationIDSubtaskType.BATCH_LINK,
+                    ),
+                    suggestions=[
+                        LocationSuggestion(
+                            location_id=input_.location_id,
+                            confidence=80,
+                        )
+                    ]
+                )
+            )
+
+        await self._upload_subtask_data(subtask_data_list)
+
+    async def _get_from_db(self) -> list[LocationBatchLinkInput]:
+        query = GetLocationBatchLinkQueryBuilder()
+        return await self.adb_client.run_query_builder(query)

src/core/tasks/url/operators/location_id/subtasks/impl/batch_link/inputs.py

@@ -0,0 +1,6 @@
+from pydantic import BaseModel
+
+
+class LocationBatchLinkInput(BaseModel):
+    location_id: int
+    url_id: int