RekomenciBackend/src/template_project/adapters/data_gateways/vacancy.py

from collections.abc import Sequence
from typing import override

from sqlalchemy import label, select
from sqlalchemy.ext.asyncio import AsyncSession

from template_project.adapters.data_gateways.tables import vacancy_embedding_table, vacancy_table
from template_project.application.vacancy.data_gateway import VacancyDataGateway
from template_project.application.vacancy.data_structure import SuitableVacancy
from template_project.application.vacancy.entity import Vacancy, VacancyEmbedding


class DefaultVacancyDataGateway(VacancyDataGateway):
    def __init__(self, session: AsyncSession) -> None:
        self._session = session

    @override
    async def get_suitable(self, vector: list[float]) -> Sequence[SuitableVacancy]:
        distance_expr = vacancy_embedding_table.c.vector.cosine_distance(vector)
        similarity_expr = 1 - distance_expr
        statement = (
            select(Vacancy, label("resume_similarity", similarity_expr))
            .join(VacancyEmbedding, vacancy_embedding_table.c.vacancy_id == vacancy_table.c.id)
            .where(similarity_expr >= 0.5)
            .order_by(distance_expr.asc())
            .limit(100)
        )
        result = await self._session.execute(statement)
        return [
            SuitableVacancy(
                vacancy=row[0],
                resume_similarity=row[1],
            )
            for row in result.all()
        ]