schemas.py

from pydantic import BaseModel, Field
from typing import List, Optional

# Shared model schemas


class RequirementInfo(BaseModel):
    """Represents an extracted requirement info"""
    context: str = Field(..., description="Context for the requirement.")
    requirement: str = Field(..., description="The requirement itself.")
    document: str = Field(...,
                          description="The document the requirement is extracted from.")


class ReqGroupingCategory(BaseModel):
    """Represents the category of requirements grouped together"""
    id: int = Field(..., description="ID of the grouping category")
    title: str = Field(..., description="Title given to the grouping category")
    requirements: List[RequirementInfo] = Field(
        ..., description="List of grouped requirements")


class SolutionModel(BaseModel):
    Context: str = Field(...,
                         description="Full context provided for this category.")
    Requirements: List[str] = Field(...,
                                    description="List of each requirement as string.")
    Problem_Description: str = Field(..., alias="Problem Description",
                                     description="Description of the problem being solved.")
    Solution_Description: str = Field(..., alias="Solution Description",
                                      description="Detailed description of the solution.")
    References: list[str] = Field(
        ..., description="References to documents used for the solution.")
    Category_Id: int = Field(
        ..., description="ID of the requirements category the solution is based on")

    class Config:
        validate_by_name = True  # Enables alias handling on input/output

# ============================================================= Categorize requirements endpoint


class ReqGroupingRequest(BaseModel):
    """Request schema of a requirement grouping call."""
    requirements: list[RequirementInfo]
    max_n_categories: Optional[int] = Field(
        default=None, description="Max number of categories to construct. Defaults to None")


class ReqGroupingResponse(BaseModel):
    """Response of a requirement grouping call."""
    categories: List[ReqGroupingCategory]


# INFO: keep in sync with prompt
class _ReqGroupingCategory(BaseModel):
    title: str = Field(..., description="Title given to the grouping category")
    items: list[int] = Field(
        ..., description="List of the IDs of the requirements belonging to the category.")


class _ReqGroupingOutput(BaseModel):
    categories: list[_ReqGroupingCategory] = Field(
        ..., description="List of grouping categories")


# =========================================================== Criticize solution endpoint

class CriticizeSolutionsRequest(BaseModel):
    solutions: list[SolutionModel]


class _SolutionCriticism(BaseModel):
    technical_challenges: List[str] = Field(
        ..., description="Technical challenges encountered by the solution")
    weaknesses: List[str] = Field(...,
                                  description="Identified weaknesses of the solution")
    limitations: List[str] = Field(...,
                                   description="Identified limitations of the solution")


class _SolutionCriticismOutput(BaseModel):
    criticisms: List[_SolutionCriticism]

# response format


class SolutionCriticism(BaseModel):
    solution: SolutionModel
    criticism: _SolutionCriticism


class CritiqueResponse(BaseModel):
    critiques: List[SolutionCriticism]


# =================================================================== search solution response endpoint

class _SolutionSearchOutput(BaseModel):
    solution: SolutionModel


class _SearchedSolutionModel(BaseModel):
    """"Internal model used for solutions searched using gemini"""
    requirement_ids: List[int] = Field(...,
                                       description="List of each requirement ID addressed by the solution")
    problem_description: str = Field(...,
                                     description="Description of the problem being solved.")
    solution_description: str = Field(...,
                                      description="Detailed description of the solution.")


class SolutionSearchResponse(BaseModel):
    """Response model for solution search"""
    solutions: list[SolutionModel]


# ================================================================= refine solution endpoints


class _RefinedSolutionModel(BaseModel):
    """Internal model used for solution refining"""

    problem_description: str = Field(...,
                                     description="New description of the problem being solved.")
    solution_description: str = Field(...,
                                      description="New detailed description of the solution.")