File size: 5,968 Bytes
41c1aed 594f2fe 41c1aed 594f2fe f6a7399 72683de 8cc6fe4 72683de 594f2fe 72683de 594f2fe 41c1aed 594f2fe 72683de 594f2fe f6a7399 72683de 03acc5b 2334311 03acc5b 2334311 36dc4ec d4b51a7 36dc4ec d4b51a7 36dc4ec |
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 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 |
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[dict] = 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]
# ================================================================ search solution endpoint v2
class SolutionSearchV2Request(BaseModel):
"""Response of a requirement grouping call."""
categories: List[ReqGroupingCategory]
user_constraints: Optional[str] = Field(
default=None, description="Additional user constraints to respect when generating the solutions.")
# ================================================================= 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.")
# ================================================================= search solutions using insight finder endpoints
# Helpers to extract constraints for Insights Finder
class InsightFinderConstraintItem(BaseModel):
title: str
description: str
class InsightFinderConstraintsList(BaseModel):
constraints: list[InsightFinderConstraintItem]
# =================================================
class Technology(BaseModel):
"""Represents a single technology entry with its details."""
title: str
purpose: str
key_components: str
advantages: str
limitations: str
id: int
class TechnologyData(BaseModel):
"""Represents the top-level object containing a list of technologies."""
technologies: List[Technology]
|