models.py

from enum import Enum
from typing import Any, Optional
from pydantic import BaseModel, Field


class TaskType(str, Enum):
    SUMMARIZATION = "summarization"
    QUESTION_ANSWERING = "question_answering"
    CODE_ANALYSIS = "code_analysis"
    INFORMATION_EXTRACTION = "information_extraction"


class TextInput(BaseModel):
    content: str = Field(description="The text content to analyze")
    source_type: Optional[str] = Field(default=None, description="Type of source")
    metadata: dict[str, Any] = Field(default_factory=dict, description="Additional metadata")


class OverviewObservation(BaseModel):
    task_id: str
    task_type: TaskType
    task_name: str
    task_description: str
    difficulty: str = "medium"
    input_text: str
    question: Optional[str] = None
    code_language: Optional[str] = None
    expected_output: Optional[str] = None
    max_steps: int = 10


class AnalysisResult(BaseModel):
    summary: Optional[str] = None
    answer: Optional[str] = None
    explanation: Optional[str] = None
    extracted_info: Optional[dict[str, Any]] = None
    confidence: float = 0.0


class OverviewAction(BaseModel):
    analysis: AnalysisResult
    submit: bool = False