File size: 4,822 Bytes
41c1aed 594f2fe 41c1aed 594f2fe f6a7399 72683de 594f2fe 72683de 594f2fe 41c1aed 594f2fe 72683de 594f2fe f6a7399 72683de 03acc5b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 |
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.") |