src/v2/output_schema.py

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

class LLMCorrectionOutput(BaseModel):
    text: List[str] = Field(
        ...,
        description="A list of exact text fragments from the Jupyter notebook cells where corrections need to be applied. Each fragment must be minimal and only include the part with issues."
    )
    corrected_text: List[str] = Field(
        ...,
        description="A list of corrected text fragments, aligned by index with `text`. Each entry must contain only the corrected version."
    )
    is_grammar_error: List[bool] = Field(
        ...,
        description="A list of booleans aligned by index with `text`. True if the issue is a grammatical/punctuation/capitalization/spelling error, False if it is a stylistic enhancement."
    )

class LLMFactualCheckOutput(BaseModel):
      text: List[str] = Field(
            ...,
            description="Exact text fragments from the notebook that contain factual errors."
        )
      corrected_text: List[str] = Field(
            ...,
            description="Corrected factual statements aligned with `text`."
        )

class EvaluationSuggestions(BaseModel):
    key_suggestions: List[str] = Field(
        description="A list of actionable suggestions that highlight the most critical improvements across business context, objectives, conclusions & recommendations, and alignment."
    )

class NotebookFlowEvaluation(BaseModel):
    suggestions: List[str] = Field(
        ..., 
        description="Actionable suggestions to improve the notebook's logical flow, instructional design, or overall quality."
    )