schemas.py

from pydantic import BaseModel, Field
from typing import Any, List, Dict, Literal, Optional


# --------------------------------------- Document related endpoints ---------------------------------------

class GetMeetingsRequest(BaseModel):
    working_group: Literal["SA1", "SA2", "SA3", "SA4", "SA5", "SA6",
                           "CT1", "CT2", "CT3", "CT4", "CT5", "CT6", "RAN1", "RAN2"]


class GetMeetingsResponse(BaseModel):
    meetings: Dict[str, str] = Field(
        ..., description="Mapping of meetings with as key their display name and value the FTP meeting name.")


class GetMeetingDocsRequest(BaseModel):
    working_group: str
    meeting: str


class GetMeetingDocsResponse(BaseModel):
    data: List[Dict[Any, Any]]


class DocInfo(BaseModel):
    """
        Schema for describing a document to download.
    """
    # Document name
    document: str
    # Document URL
    url: str
    # Document type
    type: str
    # Agenda item this document belongs to.
    agenda_item: str


class DownloadDocsRequest(BaseModel):
    documents: List[DocInfo] = Field(
        description="List of documents to download")
    sort_by_agenda_item: bool = Field(default=False, description="Whether to sort the files by their agenda item.")

# --------------------------------------


class ExtractRequirementsRequest(BaseModel):
    documents: List[DocInfo]


class DocRequirements(BaseModel):
    document: str
    context: str
    requirements: List[str]


class ExtractRequirementsResponse(BaseModel):
    requirements: List[DocRequirements]

# --------------------------------------


class RequirementInfo(BaseModel):
    req_id: int = Field(..., description="The ID of this requirement")
    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 ReqSearchLLMResponse(BaseModel):
    selected: List[int]


class ReqSearchRequest(BaseModel):
    query: str
    requirements: List[RequirementInfo]


class ReqSearchResponse(BaseModel):
    requirements: List[RequirementInfo]

# --------------------------------------


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 covered by the solution as a string.")
    problem_description: str = Field(...,
                                     description="Description of the problem being solved.")
    solution_description: str = Field(...,
                                      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")
    disable_sort_checks: bool = Field(default=False, description="Disable sort checks when grouping requirements")


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")


# =================================================================== bootstrap solution response
# Helpers for integration with insights finder.

class InsightFinderConstraintItem(BaseModel):
    title: str
    description: str


class InsightFinderConstraintsList(BaseModel):
    constraints: list[InsightFinderConstraintItem]

# TODO: aller voir la doc, breakage API


class Technology(BaseModel):
    """Represents a single technology entry with its details."""
    title: str = Field(..., alias="name")
    purpose: str
    advantages: str
    limitations: str


class TechnologyData(BaseModel):
    """Represents the top-level object containing a list of technologies."""
    technologies: List[Technology]


class _SolutionBootstrapOutput(BaseModel):
    solution: SolutionModel


class _BootstrappedSolutionModel(BaseModel):
    """"Internal model used for solutions bootstrapped using """
    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 SolutionBootstrapResponse(BaseModel):
    """Response model for solution bootstrapping"""
    solutions: list[SolutionModel]


class SolutionBootstrapRequest(BaseModel):
    categories: List[ReqGroupingCategory]
    user_constraints: Optional[str] = Field(
        default=None, description="Additional user constraints to respect when generating the solutions.")

# =========================================================== 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]


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


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


# ==========================================================================

class _RefinedSolutionModel(BaseModel):
    """Internal model used for bootstrapped 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.")

# ===========================================================================


class PriorArtSearchRequest(BaseModel):
    topics: list[str] = Field(
        ..., description="The list of topics to search for to create an exhaustive prior art search for a problem draft.")
    mode: Literal['prior_art', 'fto'] = Field(default="fto", description="")


class PriorArtSearchResponse(BaseModel):
    # Final consolidation report contents
    content: str
    # Individual search topic contents
    topic_contents: list[dict]
    pass