Api/app/modules/documentHandeler/dependencies/dependencies.py

from fastapi import Depends, Request
from transformers import AutoModel, AutoModelForMaskedLM, AutoTokenizer

from app.infrastructure.repository.document_handeler_repository import (
    DocumentHandelerRepository,
)
from app.modules.denseEmbeddings.denseEmbeddings import DenseEmbeddings
from app.modules.documentHandeler.controllers.document_handeler_controller import (
    DocumentHandelerController,
)
from app.modules.documentHandeler.features.createEmbeddings_feature import (
    CreateEmbeddingsFeature,
)
from app.modules.documentHandeler.features.deleteDocument_feature import (
    DeleteDocumentFeature,
)
from app.modules.documentHandeler.features.getAllChunkedText_feature import (
    GetAllChunkedTextFeature,
)
from app.qdrant import QdrantConnectionDb


def get_qdrant_connection_db() -> QdrantConnectionDb:
    return QdrantConnectionDb()


def get_document_handeler_repository(
    qdrant_connection_db: QdrantConnectionDb = Depends(get_qdrant_connection_db),
):
    return DocumentHandelerRepository(qdrant_connection_db)


def get_dense_model(request: Request) -> AutoModel:
    return request.scope["state"]["dense_model"]


def get_sparse_model(request: Request) -> AutoModelForMaskedLM:
    return request.scope["state"]["sparse_model"]


def get_dense_tokenizer(request: Request) -> AutoTokenizer:
    return request.scope["state"]["dense_tokenizer"]


def get_sparse_tokenizer(request: Request) -> AutoTokenizer:
    return request.scope["state"]["sparse_tokenizer"]


def get_dense_embeddings(
    dense_model: AutoModel = Depends(get_dense_model),
    dense_tokenizer: AutoTokenizer = Depends(get_dense_tokenizer),
    sparse_model: AutoModelForMaskedLM = Depends(get_sparse_model),
    sparse_tokenizer: AutoTokenizer = Depends(get_sparse_tokenizer),
):
    return DenseEmbeddings(
        dense_model=dense_model,
        dense_tokenizer=dense_tokenizer,
        sparse_model=sparse_model,
        sparse_tokenizer=sparse_tokenizer,
    )


def get_all_chunked_text_feature(
    document_handeler_repository: DocumentHandelerRepository = Depends(
        get_document_handeler_repository
    ),
):
    return GetAllChunkedTextFeature(document_handeler_repository)


def get_create_embeddings_feature(
    dense_embeddings: DenseEmbeddings = Depends(get_dense_embeddings),
    document_handeler_repository: DocumentHandelerRepository = Depends(
        get_document_handeler_repository
    ),
):
    return CreateEmbeddingsFeature(dense_embeddings, document_handeler_repository)


def get_delete_document_feature(
    document_handeler_repository: DocumentHandelerRepository = Depends(
        get_document_handeler_repository
    ),
):
    return DeleteDocumentFeature(document_handeler_repository)


def get_document_handeler_controller(
    delete_document_feature: DeleteDocumentFeature = Depends(
        get_delete_document_feature
    ),
    create_embeddings_feature: CreateEmbeddingsFeature = Depends(
        get_create_embeddings_feature
    ),
    get_all_chunked_text_feature: GetAllChunkedTextFeature = Depends(
        get_all_chunked_text_feature
    ),
):
    return DocumentHandelerController(
        delete_document_feature=delete_document_feature,
        create_embeddings_feature=create_embeddings_feature,
        get_all_chunked_text_feature=get_all_chunked_text_feature,
    )


def get_create_embeddings_feature(
    dense_embeddings: DenseEmbeddings = Depends(get_dense_embeddings),
    document_handeler_repository: DocumentHandelerRepository = Depends(
        get_document_handeler_repository
    ),
):
    return CreateEmbeddingsFeature(dense_embeddings, document_handeler_repository)