llm-review/models.py
2025-04-01 14:57:26 -04:00

93 lines
2.7 KiB
Python

import enum
from typing import Literal, Optional, Union
from pydantic import BaseModel, Field
class ResponseWithThinking[T](BaseModel):
thinking: str
output: T
class Importance(enum.Enum):
CRITICAL = "Critical"
MAJOR = "Major"
MINOR = "Minor"
TRIVIAL = "Trivial"
class Comment(BaseModel):
"""
A comment on a Design Review for a design.
"""
description: str = Field(
...,
description=(
"A detailed description of the issue/improvement in GitHub Flavored "
"Markdown. This should include a description of the issue, why it is "
"important, the reasoning behind why this is an issue or why this "
"improvement is necessary, and any other relevant information. "
),
)
summary: str = Field(
...,
description=(
"A 1-2 line summary of the issue/improvement in GitHub Flavored Markdown."
),
)
page: int = Field(
...,
description=(
"The page number that this comment is about. When critiquing a "
"specific page, you MUST provide the page number. You can set this "
"field to 0 ONLY if the comment is about the overall design."
),
)
elements: Optional[str] = Field(
None,
description=(
"The component, pin or net that should be reviewed. If this comment"
"is about a specific component, refer to it by its designator. If "
"it is about a specific pin on a component, refer to it in the form "
"`U1A.1` where `U1A` is the designator of the component and `1` is "
"the pin number. If it is about a net, refer to it as the name of "
"the net. If this comment is not specific to an element, you can "
"leave this blank. When referencing multiple elements, they MUST be "
"separated by exactly a comma and a space. For example: `U1A.1, "
"U2A.2, U3A.3`."
),
)
suggestion: Optional[str] = Field(
None,
description=(
"If you have a concrete, actionable suggestion, provide it here in "
"GitHub Flavored Markdown. This MUST be one paragraph of 1-2 sentences. "
"Otherwise, you may leave this blank."
),
)
importance: Importance = Field(..., description="The importance of the comment.")
class Sample[T](BaseModel):
"""
An example of a T to be provided to the LLM for training.
"""
sample: T
rating: Union[Literal["Good"], Literal["Bad"]] = Field(
..., description="The rating of the sample."
)
explanation: str = Field(
...,
description="An explanation for why this sample is rated as it is.",
)
COMMENT_EXAMPLES = []