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() ]