Update annotations to join user and robo suggestions for locations and agencies

Author: maxachis

src/api/endpoints/annotate/all/get/models/agency.py

@@ -10,30 +10,6 @@ class AgencyAnnotationSuggestion(BaseModel):
         le=100,
     )
 
-# TODO: Replace Usages and Delete
-class AgencyAnnotationAutoSuggestion(BaseModel):
-    agency_id: int
-    agency_name: str
-    confidence: int = Field(
-        title="The confidence of the location",
-        ge=0,
-        le=100,
-    )
-
-# TODO: Replace Usages and Delete
-class AgencyAnnotationUserSuggestion(BaseModel):
-    agency_id: int
-    agency_name: str
-    user_count: int
-
-# TODO: Replace Usages and Delete
-class AgencyAnnotationUserSuggestionOuterInfo(BaseModel):
-    suggestions: list[AgencyAnnotationUserSuggestion]
-    not_found_count: int = Field(
-        title="How many users listed the agency as not found.",
-        ge=0,
-    )
-
 class AgencyAnnotationResponseOuterInfo(BaseModel):
     suggestions: list[AgencyAnnotationSuggestion]
     not_found_count: int = Field(

src/api/endpoints/annotate/all/get/models/location.py

@@ -11,37 +11,6 @@ class LocationAnnotationSuggestion(BaseModel):
         le=100,
     )
 
-# TODO: Replace Usages and Delete
-class LocationAnnotationAutoSuggestion(BaseModel):
-    location_id: int
-    location_name: str = Field(
-        title="The full name of the location"
-    )
-    confidence: int = Field(
-        title="The confidence of the location",
-        ge=0,
-        le=100,
-    )
-
-# TODO: Replace Usages and Delete
-class LocationAnnotationUserSuggestion(BaseModel):
-    location_id: int
-    location_name: str = Field(
-        title="The full name of the location"
-    )
-    user_count: int = Field(
-        title="The number of users who suggested this location",
-        ge=1,
-    )
-
-# TODO: Replace Usages and Delete
-class LocationAnnotationUserSuggestionOuterInfo(BaseModel):
-    suggestions: list[LocationAnnotationUserSuggestion]
-    not_found_count: int = Field(
-        title="How many users listed the location as not found.",
-        ge=0,
-    )
-
 class LocationAnnotationResponseOuterInfo(BaseModel):
     suggestions: list[LocationAnnotationSuggestion]
     not_found_count: int = Field(

src/api/endpoints/annotate/all/get/queries/agency/core.py

@@ -1,13 +1,7 @@
 from sqlalchemy.ext.asyncio import AsyncSession
 
-from src.api.endpoints.annotate.all.get.models.agency import AgencyAnnotationResponseOuterInfo, \
-    AgencyAnnotationUserSuggestionOuterInfo, AgencyAnnotationUserSuggestion, AgencyAnnotationAutoSuggestion
-from src.api.endpoints.annotate.all.get.queries.agency.requester import GetAgencySuggestionsRequester
-from src.db.queries.base.builder import QueryBuilderBase
-from sqlalchemy.ext.asyncio import AsyncSession
-
-from src.api.endpoints.annotate.all.get.models.agency import AgencyAnnotationResponseOuterInfo, \
-    AgencyAnnotationUserSuggestionOuterInfo, AgencyAnnotationUserSuggestion, AgencyAnnotationAutoSuggestion
+from src.api.endpoints.annotate.all.get.models.agency import AgencyAnnotationResponseOuterInfo
+from src.api.endpoints.annotate.all.get.models.agency import AgencyAnnotationSuggestion
 from src.api.endpoints.annotate.all.get.queries.agency.requester import GetAgencySuggestionsRequester
 from src.db.queries.base.builder import QueryBuilderBase
 
@@ -30,19 +24,13 @@ async def run(self, session: AsyncSession) -> AgencyAnnotationResponseOuterInfo:
             location_id=self.location_id
         )
 
-        # TODO: Pull both in single query
-        user_suggestions: list[AgencyAnnotationUserSuggestion] = \
-            await requester.get_user_agency_suggestions()
-        auto_suggestions: list[AgencyAnnotationAutoSuggestion] = \
-            await requester.get_auto_agency_suggestions()
+        suggestions: list[AgencyAnnotationSuggestion] = \
+            await requester.get_agency_suggestions()
         not_found_count: int = \
             await requester.get_not_found_count()
         return AgencyAnnotationResponseOuterInfo(
-            user=AgencyAnnotationUserSuggestionOuterInfo(
-                suggestions=user_suggestions,
-                not_found_count=not_found_count
-            ),
-            auto=auto_suggestions,
+            suggestions=suggestions,
+            not_found_count=not_found_count
         )
 
 

src/api/endpoints/annotate/all/get/queries/agency/requester.py

@@ -1,16 +1,15 @@
 from typing import Sequence
 
-from sqlalchemy import func, select, RowMapping
+from sqlalchemy import func, select, RowMapping, or_, and_
 from sqlalchemy.ext.asyncio import AsyncSession
 
-from src.api.endpoints.annotate.all.get.models.agency import AgencyAnnotationAutoSuggestion, \
-    AgencyAnnotationUserSuggestion
-from src.api.endpoints.annotate.all.get.queries.agency.suggestions_with_highest_confidence import \
-    SuggestionsWithHighestConfidenceCTE
+from src.api.endpoints.annotate.all.get.models.agency import AgencyAnnotationSuggestion
+from src.db.helpers.query import exists_url
 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.subtask.sqlalchemy import URLAutoAgencyIDSubtask
+from src.db.models.impl.url.suggestion.agency.suggestion.sqlalchemy import AgencyIDSubtaskSuggestion
 from src.db.models.impl.url.suggestion.agency.user import UserURLAgencySuggestion
 from src.db.templates.requester import RequesterBase
 
@@ -27,102 +26,97 @@ def __init__(
         self.url_id = url_id
         self.location_id = location_id
 
-    async def get_user_agency_suggestions(self) -> list[AgencyAnnotationUserSuggestion]:
-        query = (
+    async def get_agency_suggestions(self) -> list[AgencyAnnotationSuggestion]:
+        # All agencies with either a user or robo annotation
+        valid_agencies_cte = (
             select(
-                UserURLAgencySuggestion.agency_id,
-                func.count(UserURLAgencySuggestion.user_id).label("count"),
-                Agency.name.label("agency_name"),
-            )
-            .join(
-                Agency,
-                Agency.id == UserURLAgencySuggestion.agency_id
+                Agency.id,
             )
-
-        )
-
-        if self.location_id is not None:
-            query = (
-                query.join(
-                    LinkAgencyLocation,
-                    LinkAgencyLocation.agency_id == UserURLAgencySuggestion.agency_id
-                )
-                .where(
-                    LinkAgencyLocation.location_id == self.location_id
+            .where(
+                or_(
+                    exists_url(
+                        UserURLAgencySuggestion
+                    ),
+                    exists_url(
+                        URLAutoAgencyIDSubtask
+                    )
                 )
             )
+            .cte("valid_agencies")
+        )
 
-        query = (
-            query.where(
-                UserURLAgencySuggestion.url_id == self.url_id
+        # Number of users who suggested each agency
+        user_suggestions_cte = (
+            select(
+                UserURLAgencySuggestion.url_id,
+                UserURLAgencySuggestion.agency_id,
+                func.count(UserURLAgencySuggestion.user_id).label('user_count')
             )
             .group_by(
                 UserURLAgencySuggestion.agency_id,
-                Agency.name
+                UserURLAgencySuggestion.url_id,
             )
-            .order_by(
-                func.count(UserURLAgencySuggestion.user_id).desc()
-            )
-            .limit(3)
+            .cte("user_suggestions")
         )
 
-        results: Sequence[RowMapping] = await sh.mappings(self.session, query=query)
-
-        return [
-            AgencyAnnotationUserSuggestion(
-                agency_id=autosuggestion["agency_id"],
-                user_count=autosuggestion["count"],
-                agency_name=autosuggestion["agency_name"],
+        # Maximum confidence of robo annotation, if any
+        robo_suggestions_cte = (
+            select(
+                URLAutoAgencyIDSubtask.url_id,
+                Agency.id.label("agency_id"),
+                func.max(AgencyIDSubtaskSuggestion.confidence).label('robo_confidence')
             )
-            for autosuggestion in results
-        ]
-
-
-    async def get_auto_agency_suggestions(self) -> list[AgencyAnnotationAutoSuggestion]:
-        cte = SuggestionsWithHighestConfidenceCTE()
-        query = (
+            .join(
+                AgencyIDSubtaskSuggestion,
+                AgencyIDSubtaskSuggestion.subtask_id == URLAutoAgencyIDSubtask.id
+            )
+            .join(
+                Agency,
+                Agency.id == AgencyIDSubtaskSuggestion.agency_id
+            )
+            .group_by(
+                URLAutoAgencyIDSubtask.url_id,
+                Agency.id
+            )
+            .cte("robo_suggestions")
+        )
+        # Join user and robo suggestions
+        joined_suggestions_query = (
             select(
-                cte.agency_id,
-                cte.confidence,
+                valid_agencies_cte.c.id.label("agency_id"),
                 Agency.name.label("agency_name"),
+                func.coalesce(user_suggestions_cte.c.user_count, 0).label('user_count'),
+                func.coalesce(robo_suggestions_cte.c.robo_confidence, 0).label('robo_confidence'),
             )
             .join(
                 Agency,
-                Agency.id == cte.agency_id
+                Agency.id == valid_agencies_cte.c.id
             )
-        )
-
-        if self.location_id is not None:
-            query = (
-                query.join(
-                    LinkAgencyLocation,
-                    LinkAgencyLocation.agency_id == cte.agency_id
-                )
-                .where(
-                    LinkAgencyLocation.location_id == self.location_id
+            .outerjoin(
+                user_suggestions_cte,
+                and_(
+                    user_suggestions_cte.c.url_id == self.url_id,
+                    user_suggestions_cte.c.agency_id == Agency.id
                 )
             )
-
-        query = (
-            query.where(
-                cte.url_id == self.url_id
-            )
-            .order_by(
-                cte.confidence.desc()
+            .outerjoin(
+                robo_suggestions_cte,
+                and_(
+                    robo_suggestions_cte.c.url_id == self.url_id,
+                    robo_suggestions_cte.c.agency_id == Agency.id
+                )
             )
-            .limit(3)
         )
 
-        results: Sequence[RowMapping] = await sh.mappings(self.session, query=query)
-
-        return [
-            AgencyAnnotationAutoSuggestion(
-                agency_id=autosuggestion["agency_id"],
-                confidence=autosuggestion["confidence"],
-                agency_name=autosuggestion["agency_name"],
+        # Return suggestions
+        mappings: Sequence[RowMapping] = await self.mappings(joined_suggestions_query)
+        suggestions: list[AgencyAnnotationSuggestion] = [
+            AgencyAnnotationSuggestion(
+                **mapping
             )
-            for autosuggestion in results
+            for mapping in mappings
         ]
+        return suggestions
 
     async def get_not_found_count(self) -> int:
         query = (

src/api/endpoints/annotate/all/get/queries/location_/core.py

@@ -1,13 +1,7 @@
 from sqlalchemy.ext.asyncio import AsyncSession
 
-from src.api.endpoints.annotate.all.get.models.location import LocationAnnotationResponseOuterInfo, \
-    LocationAnnotationUserSuggestion, LocationAnnotationAutoSuggestion, LocationAnnotationUserSuggestionOuterInfo
-from src.api.endpoints.annotate.all.get.queries.location_.requester import GetLocationSuggestionsRequester
-from src.db.queries.base.builder import QueryBuilderBase
-from sqlalchemy.ext.asyncio import AsyncSession
-
-from src.api.endpoints.annotate.all.get.models.location import LocationAnnotationResponseOuterInfo, \
-    LocationAnnotationUserSuggestion, LocationAnnotationAutoSuggestion
+from src.api.endpoints.annotate.all.get.models.location import LocationAnnotationResponseOuterInfo
+from src.api.endpoints.annotate.all.get.models.location import LocationAnnotationSuggestion
 from src.api.endpoints.annotate.all.get.queries.location_.requester import GetLocationSuggestionsRequester
 from src.db.queries.base.builder import QueryBuilderBase
 
@@ -21,24 +15,17 @@ def __init__(
         super().__init__()
         self.url_id = url_id
 
-
+    # TODO: Test
     async def run(self, session: AsyncSession) -> LocationAnnotationResponseOuterInfo:
         requester = GetLocationSuggestionsRequester(session)
 
-        # TODO: Pull both in single query
-        suggestions
-        user_suggestions: list[LocationAnnotationUserSuggestion] = \
-            await requester.get_user_location_suggestions(self.url_id)
-        auto_suggestions: list[LocationAnnotationAutoSuggestion] = \
-            await requester.get_auto_location_suggestions(self.url_id)
+        suggestions: list[LocationAnnotationSuggestion] = \
+            await requester.get_location_suggestions(self.url_id)
         not_found_count: int = \
             await requester.get_not_found_count(self.url_id)
 
         return LocationAnnotationResponseOuterInfo(
-            user=LocationAnnotationUserSuggestionOuterInfo(
-                suggestions=user_suggestions,
-                not_found_count=not_found_count
-            ),
-            auto=auto_suggestions
+            suggestions=suggestions,
+            not_found_count=not_found_count
         )