Files
RekomenciBackend/src/template_project/ml/routes/predict.py
T
ivankirpichnikov effbcfbc2d не помнб
2025-11-22 18:15:29 +03:00

122 lines
4.1 KiB
Python

from decimal import Decimal
from dishka import FromDishka
from dishka.integrations.fastapi import DishkaRoute
from fastapi import APIRouter
from pydantic import BaseModel, Field
from template_project.application.resume.entity import ResumeId
from template_project.ml.interactors.predict_salary import (
PredictSalaryInteractor,
PredictSalaryRequest,
VacancyInput,
)
router = APIRouter(route_class=DishkaRoute, tags=["Prediction"])
class VacancyInputModel(BaseModel):
vacancy_id: str = Field(description="Vacancy ID", examples=["vacancy_123"])
from_salary: Decimal = Field(description="Minimum salary", examples=[Decimal(100000)])
to_salary: Decimal = Field(description="Maximum salary", examples=[Decimal(150000)])
key_skills: list[str] = Field(description="List of key skills", examples=[["Python", "FastAPI", "PostgreSQL"]])
resume_similarity: float = Field(
ge=0.0, le=1.0, description="Resume similarity score (0.0 to 1.0)", examples=[0.85]
)
model_config = {
"json_schema_extra": {
"example": {
"vacancy_id": "vacancy_123",
"from_salary": "100000",
"to_salary": "150000",
"key_skills": ["Python", "FastAPI", "PostgreSQL"],
"resume_similarity": 0.85,
}
}
}
class PredictSalaryRequestModel(BaseModel):
resume_id: ResumeId = Field(description="Resume ID", examples=["01234567-89ab-cdef-0123-456789abcdef"])
key_skills: list[str] = Field(
min_length=1, description="List of key skills from resume", examples=[["Python", "FastAPI", "PostgreSQL"]]
)
vacancies: list[VacancyInputModel] = Field(
min_length=1, description="List of relevant vacancies", examples=[[]]
)
model_config = {
"json_schema_extra": {
"example": {
"resume_id": "01234567-89ab-cdef-0123-456789abcdef",
"key_skills": ["Python", "FastAPI", "PostgreSQL"],
"vacancies": [
{
"vacancy_id": "vacancy_123",
"from_salary": "100000",
"to_salary": "150000",
"key_skills": ["Python", "FastAPI", "PostgreSQL", "Docker"],
"resume_similarity": 0.85,
}
],
}
}
}
class PredictSalaryResponseModel(BaseModel):
salary_from: Decimal = Field(description="Minimum predicted salary", examples=[Decimal(100000)])
salary_to: Decimal = Field(description="Maximum predicted salary", examples=[Decimal(150000)])
recommended_skills: list[str] = Field(
description="Top 3 recommended skills", examples=[["Kubernetes", "Redis", "Docker"]]
)
model_config = {
"json_schema_extra": {
"example": {
"salary_from": "100000",
"salary_to": "150000",
"recommended_skills": ["Kubernetes", "Redis", "Docker"],
}
}
}
@router.post(
"/predict_salary",
summary="Predict salary",
description="Predict salary range and recommend skills based on resume and relevant vacancies",
responses={
200: {"description": "Salary prediction generated successfully", "model": PredictSalaryResponseModel},
},
)
async def predict_salary(
request: PredictSalaryRequestModel,
interactor: FromDishka[PredictSalaryInteractor],
) -> PredictSalaryResponseModel:
vacancy_inputs = [
VacancyInput(
vacancy_id=vacancy.vacancy_id,
from_salary=vacancy.from_salary,
to_salary=vacancy.to_salary,
key_skills=vacancy.key_skills,
resume_similarity=vacancy.resume_similarity,
)
for vacancy in request.vacancies
]
predict_request = PredictSalaryRequest(
resume_id=request.resume_id,
key_skills=request.key_skills,
vacancies=vacancy_inputs,
)
response = await interactor.execute(predict_request)
return PredictSalaryResponseModel(
salary_from=response.salary_from,
salary_to=response.salary_to,
recommended_skills=response.recommended_skills,
)