Refacto router_repository et fixes

Author: Benjamin PILIA

api/dependencies.py

@@ -9,8 +9,7 @@
 from api.infrastructure.postgres import PostgresKeyRepository, PostgresProviderRepository, PostgresRouterRepository, PostgresUserInfoRepository
 from api.schemas.core.context import RequestContext
 from api.use_cases.admin.providers import CreateProviderUseCase
-from api.use_cases.admin.routers import CreateRouterUseCase, DeleteRouterUseCase, GetOneRouterUseCase, GetRoutersUseCase
-from api.use_cases.admin.routers._updaterouterusecase import UpdateRouterUseCase
+from api.use_cases.admin.routers import CreateRouterUseCase, DeleteRouterUseCase, GetOneRouterUseCase, GetRoutersUseCase, UpdateRouterUseCase
 from api.use_cases.models import GetModelsUseCase
 from api.utils.configuration import configuration
 from api.utils.context import global_context, request_context

api/domain/router/entities.py

@@ -52,8 +52,11 @@ def with_type(self, router_type: RouterType) -> "Router":
     def with_load_balancing_strategy(self, strategy: RouterLoadBalancingStrategy) -> "Router":
         return self.model_copy(update={"load_balancing_strategy": strategy})
 
-    def with_costs(self, prompt_tokens: float, completion_tokens: float) -> "Router":
-        return self.model_copy(update={"cost_prompt_tokens": prompt_tokens, "cost_completion_tokens": completion_tokens})
+    def with_cost_prompt_tokens(self, prompt_tokens: float) -> "Router":
+        return self.model_copy(update={"cost_prompt_tokens": prompt_tokens})
+
+    def with_cost_completion_tokens(self, completion_tokens: float) -> "Router":
+        return self.model_copy(update={"cost_completion_tokens": completion_tokens})
 
     def with_aliases(self, aliases: list[str]) -> "Router":
         return self.model_copy(update={"aliases": aliases})

api/infrastructure/fastapi/endpoints/admin/routers.py

@@ -39,8 +39,10 @@
     GetRoutersCommand,
     GetRoutersUseCase,
     GetRoutersUseCaseSuccess,
+    UpdateRouterCommand,
+    UpdateRouterUseCase,
+    UpdateRouterUseCaseSuccess,
 )
-from api.use_cases.admin.routers._updaterouterusecase import UpdateRouterCommand, UpdateRouterUseCase, UpdateRouterUseCaseSuccess
 from api.utils.variables import EndpointRoute
 
 logger = logging.getLogger(__name__)
@@ -216,7 +218,9 @@ async def delete_router(
 @router.patch(
     path=EndpointRoute.ADMIN_ROUTERS + "/{router_id}",
     dependencies=[Security(dependency=get_current_key)],
-    responses=get_documentation_responses([RouterNotFoundHTTPException, NotAdminUserHTTPException]),
+    responses=get_documentation_responses(
+        [RouterNotFoundHTTPException, NotAdminUserHTTPException, RouterAliasAlreadyExistsHTTPException, RouterAlreadyExistsHTTPException]
+    ),
     status_code=200,
 )
 async def update_router(
@@ -248,13 +252,13 @@ async def update_router(
         )
         raise InternalServerHTTPException()
     match result:
-        case UpdateRouterUseCaseSuccess(updated_router):
+        case UpdateRouterUseCaseSuccess(router=updated_router):
             return Router.model_validate(updated_router, from_attributes=True)
         case RouterNotFoundError(router_id=not_found_id):
             raise RouterNotFoundHTTPException(not_found_id)
         case UserIsNotAdminError():
             raise NotAdminUserHTTPException()
-        case RouterAliasAlreadyExistsError(name):
-            raise RouterAliasAlreadyExistsHTTPException(name)
+        case RouterAliasAlreadyExistsError(aliases):
+            raise RouterAliasAlreadyExistsHTTPException(aliases)
         case RouterNameAlreadyExistsError(name):
             raise RouterAlreadyExistsHTTPException(name)

api/infrastructure/postgres/_postgresrouterrepository.py

@@ -204,33 +204,60 @@ async def delete_router(self, router_id: int) -> Router | None:
         await self.postgres_session.execute(delete(RouterTable).where(RouterTable.id == router_id))
         return router
 
-    async def update_router(self, router: Router) -> Router | RouterNameAlreadyExistsError:
-        db_user_id = None if router.user_id == MASTER_USER_ID else router.user_id
+    async def update_router(self, router_to_update: Router) -> Router | RouterNameAlreadyExistsError:
+        db_user_id = None if router_to_update.user_id == MASTER_USER_ID else router_to_update.user_id
 
         try:
-            await self.postgres_session.execute(
+            update_query = (
                 update(RouterTable)
-                .where(RouterTable.id == router.id)
+                .where(RouterTable.id == router_to_update.id)
                 .values(
                     user_id=db_user_id,
-                    name=router.name,
-                    type=router.type.value,
-                    load_balancing_strategy=router.load_balancing_strategy.value,
-                    cost_prompt_tokens=router.cost_prompt_tokens,
-                    cost_completion_tokens=router.cost_completion_tokens,
+                    name=router_to_update.name,
+                    type=router_to_update.type.value,
+                    load_balancing_strategy=router_to_update.load_balancing_strategy.value,
+                    cost_prompt_tokens=router_to_update.cost_prompt_tokens,
+                    cost_completion_tokens=router_to_update.cost_completion_tokens,
+                )
+                .returning(
+                    RouterTable.id,
+                    RouterTable.name,
+                    RouterTable.user_id,
+                    RouterTable.type,
+                    RouterTable.load_balancing_strategy,
+                    RouterTable.cost_prompt_tokens,
+                    RouterTable.cost_completion_tokens,
+                    cast(func.extract("epoch", RouterTable.created), Integer).label("created"),
+                    cast(func.extract("epoch", RouterTable.updated), Integer).label("updated"),
                 )
             )
+            result = await self.postgres_session.execute(update_query)
+            row = result.one()
 
-            if router.aliases is not None:
-                await self.postgres_session.execute(delete(RouterAliasTable).where(RouterAliasTable.router_id == router.id))
-                if router.aliases:
+            if router_to_update.aliases is not None:
+                await self.postgres_session.execute(delete(RouterAliasTable).where(RouterAliasTable.router_id == router_to_update.id))
+                if router_to_update.aliases:
                     await self.postgres_session.execute(
                         insert(RouterAliasTable),
-                        [{"value": alias, "router_id": router.id} for alias in router.aliases],
+                        [{"value": alias, "router_id": router_to_update.id} for alias in router_to_update.aliases],
                     )
         except IntegrityError as e:
             if "router_name_key" in str(e.orig):
-                return RouterNameAlreadyExistsError(name=router.name)
+                return RouterNameAlreadyExistsError(name=router_to_update.name)
             raise
 
-        return router
+        return Router(
+            id=row.id,
+            name=row.name,
+            user_id=router_to_update.user_id,
+            type=RouterType(row.type),
+            aliases=router_to_update.aliases,
+            load_balancing_strategy=RouterLoadBalancingStrategy(row.load_balancing_strategy),
+            vector_size=router_to_update.vector_size,
+            max_context_length=router_to_update.max_context_length,
+            cost_prompt_tokens=row.cost_prompt_tokens or 0.0,
+            cost_completion_tokens=row.cost_completion_tokens or 0.0,
+            providers=router_to_update.providers,
+            created=row.created,
+            updated=row.updated,
+        )

api/tests/integration/postgres/test_postgresrouterrepository.py

@@ -551,7 +551,7 @@ async def test_update_router_should_return_router_with_updated_costs(self, repos
         await db_session.flush()
 
         # Act
-        result = await repository.update_router(to_router_domain(router).with_costs(0.010, 0.020))
+        result = await repository.update_router(to_router_domain(router).with_cost_prompt_tokens(0.010).with_cost_completion_tokens(0.020))
 
         # Assert
         assert isinstance(result, Router)

api/tests/unit/use_case/admin/routers/test_updaterouterusecase.py

@@ -3,7 +3,7 @@
 import pytest
 
 from api.domain.model import ModelType as RouterType
-from api.domain.router.entities import Router, RouterLoadBalancingStrategy
+from api.domain.router.entities import RouterLoadBalancingStrategy
 from api.domain.router.errors import RouterAliasAlreadyExistsError, RouterNameAlreadyExistsError, RouterNotFoundError
 from api.domain.userinfo.errors import UserIsNotAdminError
 from api.tests.unit.use_case.factories import RouterFactory, UserInfoFactory
@@ -36,7 +36,7 @@ def unauthorized_user_info():
 
 
 @pytest.fixture
-def sample_router() -> Router:
+def sample_router():
     return RouterFactory(id=42, user_id=1, name="original-name", aliases=[])
 
 
@@ -50,7 +50,8 @@ async def test_should_return_updated_router_when_user_is_admin_and_router_exists
             sample_router.with_name("new-name")
             .with_type(RouterType.TEXT_EMBEDDINGS_INFERENCE)
             .with_load_balancing_strategy(RouterLoadBalancingStrategy.LEAST_BUSY)
-            .with_costs(0.005, 0.010)
+            .with_cost_prompt_tokens(0.005)
+            .with_cost_completion_tokens(0.010)
             .with_aliases(["alias-a", "alias-b"])
         )
         use_case.user_info_repository.get_user_info.return_value = admin_user_info
@@ -169,6 +170,23 @@ async def test_should_propagate_router_name_already_exists_error_from_repository
         assert isinstance(result, RouterNameAlreadyExistsError)
         assert result.name == "taken-name"
 
+    @pytest.mark.asyncio
+    async def test_should_add_aliases_when_router_has_no_aliases_and_command_updates_aliases(self, use_case, router_repository, admin_user_info):
+        # Arrange
+        router = RouterFactory(id=42, user_id=1, aliases=None)
+        updated_router = router.with_aliases(["new-alias"])
+        use_case.user_info_repository.get_user_info.return_value = admin_user_info
+        use_case.router_repository.get_router_by_id.return_value = router
+        use_case.router_repository.get_aliases.return_value = []
+        use_case.router_repository.update_router.return_value = updated_router
+
+        # Act
+        result = await use_case.execute(command=UpdateRouterCommand(user_id=admin_user_info.id, router_id=42, aliases=["new-alias"]))
+
+        # Assert
+        assert isinstance(result, UpdateRouterUseCaseSuccess)
+        router_repository.update_router.assert_called_once_with(router=router.with_aliases(["new-alias"]))
+
     @pytest.mark.asyncio
     async def test_should_not_call_update_router_when_router_should_not_be_updated(self, use_case, router_repository, admin_user_info, sample_router):
         # Arrange

api/use_cases/admin/routers/__init__.py

@@ -2,6 +2,7 @@
 from ._deleterouterusecase import DeleteRouterCommand, DeleteRouterUseCase, DeleteRouterUseCaseSuccess
 from ._getonerouterusecase import GetOneRouterCommand, GetOneRouterUseCase, GetOneRouterUseCaseSuccess
 from ._getroutersusecase import GetRoutersCommand, GetRoutersUseCase, GetRoutersUseCaseSuccess
+from ._updaterouterusecase import UpdateRouterCommand, UpdateRouterUseCase, UpdateRouterUseCaseSuccess
 
 __all__ = [
     "CreateRouterCommand",
@@ -16,4 +17,7 @@
     "DeleteRouterCommand",
     "DeleteRouterUseCase",
     "DeleteRouterUseCaseSuccess",
+    "UpdateRouterCommand",
+    "UpdateRouterUseCase",
+    "UpdateRouterUseCaseSuccess",
 ]

api/use_cases/admin/routers/_updaterouterusecase.py

@@ -50,7 +50,7 @@ async def execute(
 
         if command.aliases:
             existing_aliases = await self.router_repository.get_aliases()
-            conflicting_aliases = set(command.aliases) & (set(existing_aliases) - set(router.aliases))
+            conflicting_aliases = set(command.aliases) & (set(existing_aliases) - set(router.aliases or []))
             if conflicting_aliases:
                 return RouterAliasAlreadyExistsError(aliases=list(conflicting_aliases))
 
@@ -62,7 +62,9 @@ async def execute(
         if command.load_balancing_strategy is not None:
             updated_router = updated_router.with_load_balancing_strategy(command.load_balancing_strategy)
         if command.cost_prompt_tokens is not None:
-            updated_router = updated_router.with_costs(command.cost_prompt_tokens, command.cost_completion_tokens)
+            updated_router = updated_router.with_cost_prompt_tokens(command.cost_prompt_tokens)
+        if command.cost_completion_tokens is not None:
+            updated_router = updated_router.with_cost_completion_tokens(command.cost_completion_tokens)
         if command.aliases is not None:
             updated_router = updated_router.with_aliases(command.aliases)