OpenGateLLM/api/dependencies.py at 57826faacaee5ac5f3734691336df33dc237cf7a · etalab-ia/OpenGateLLM · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
from collections.abc import AsyncGenerator
from contextvars import ContextVar

from fastapi import Depends
from sqlalchemy.ext.asyncio import AsyncSession

from api.domain.key import KeyRepository
from api.infrastructure.model import ModelProviderGateway
from api.infrastructure.postgres import PostgresKeyRepository, PostgresProviderRepository, PostgresRouterRepository, PostgresUserInfoRepository
from api.schemas.core.context import RequestContext
from api.use_cases.admin.providers import CreateProviderUseCase, DeleteProviderUseCase, GetOneProviderUseCase, GetProvidersUseCase
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


def get_request_context() -> ContextVar[RequestContext]:
    return request_context


def get_master_key() -> str:
    return configuration.settings.auth_master_key


# databases
async def get_postgres_session() -> AsyncGenerator[AsyncSession]:
    session_factory = global_context.postgres_session_factory
    async with session_factory() as postgres_session:
        try:
            yield postgres_session
            if postgres_session.in_transaction():
                await postgres_session.commit()
        except Exception:
            if postgres_session.in_transaction():
                await postgres_session.rollback()
            raise


# repositories
def _router_repository(session: AsyncSession) -> PostgresRouterRepository:
    return PostgresRouterRepository(postgres_session=session, app_title=configuration.settings.app_title)


def _user_info_repository(session: AsyncSession) -> PostgresUserInfoRepository:
    return PostgresUserInfoRepository(postgres_session=session)


def get_key_repository(postgres_session: AsyncSession = Depends(get_postgres_session)) -> KeyRepository:
    return PostgresKeyRepository(postgres_session=postgres_session)


# models use cases
def get_models_use_case(
    postgres_session: AsyncSession = Depends(get_postgres_session),
    request_context: RequestContext = Depends(get_request_context),
) -> GetModelsUseCase:
    return GetModelsUseCase(
        router_repository=_router_repository(postgres_session),
        user_id=request_context.get().user_id,
        user_info_repository=_user_info_repository(postgres_session),
    )


# routers use cases
def get_one_router_use_case_factory(postgres_session: AsyncSession = Depends(get_postgres_session)) -> GetOneRouterUseCase:
    return GetOneRouterUseCase(router_repository=_router_repository(postgres_session), user_info_repository=_user_info_repository(postgres_session))


def get_routers_use_case_factory(postgres_session: AsyncSession = Depends(get_postgres_session)) -> GetRoutersUseCase:
    return GetRoutersUseCase(router_repository=_router_repository(postgres_session), user_info_repository=_user_info_repository(postgres_session))


def create_router_use_case_factory(postgres_session: AsyncSession = Depends(get_postgres_session)) -> CreateRouterUseCase:
    return CreateRouterUseCase(router_repository=_router_repository(postgres_session), user_info_repository=_user_info_repository(postgres_session))


def delete_router_use_case_factory(postgres_session: AsyncSession = Depends(get_postgres_session)) -> DeleteRouterUseCase:
    return DeleteRouterUseCase(router_repository=_router_repository(postgres_session), user_info_repository=_user_info_repository(postgres_session))


def update_router_use_case_factory(postgres_session: AsyncSession = Depends(get_postgres_session)) -> UpdateRouterUseCase:
    return UpdateRouterUseCase(router_repository=_router_repository(postgres_session), user_info_repository=_user_info_repository(postgres_session))


# providers use cases
def create_provider_use_case_factory(postgres_session: AsyncSession = Depends(get_postgres_session)) -> CreateProviderUseCase:
    return CreateProviderUseCase(
        router_repository=_router_repository(postgres_session),
        provider_repository=PostgresProviderRepository(postgres_session=postgres_session),
        provider_gateway=ModelProviderGateway(),
        user_info_repository=_user_info_repository(postgres_session),
    )


def delete_provider_use_case_factory(postgres_session: AsyncSession = Depends(get_postgres_session)) -> DeleteProviderUseCase:
    return DeleteProviderUseCase(
        provider_repository=PostgresProviderRepository(postgres_session=postgres_session),
        user_info_repository=_user_info_repository(postgres_session),
    )


def get_one_provider_use_case_factory(postgres_session: AsyncSession = Depends(get_postgres_session)) -> GetOneProviderUseCase:
    return GetOneProviderUseCase(
        provider_repository=PostgresProviderRepository(postgres_session=postgres_session),
        user_info_repository=_user_info_repository(postgres_session),
    )


def get_providers_use_case_factory(postgres_session: AsyncSession = Depends(get_postgres_session)) -> GetProvidersUseCase:
    return GetProvidersUseCase(
        provider_repository=PostgresProviderRepository(postgres_session=postgres_session),
        user_info_repository=_user_info_repository(postgres_session),
    )