src/models/analysis_models.py

from typing import Iterator, List, Optional
from enum import Enum
from pydantic import BaseModel, Field


class InputModel(BaseModel):
    problem_statement: str = Field(
        default=None,
        description="Contains the description of the problem statement or task"
    )

class MLTaskType(str, Enum):
    CLASSIFICATION = "classification"
    REGRESSION = "regression"
    CLUSTERING = "clustering"
    NLP = "natural_language_processing"
    COMPUTER_VISION = "computer_vision"
    TIME_SERIES = "time_series"
    ANOMALY_DETECTION = "anomaly_detection"
    RECOMMENDATION = "recommendation"
    OTHER = "other"


class ModelResponseStatus(BaseModel):
    """Technical specification for ML implementation"""
    data_source: str = Field(
        # default="...",
        description="Required data sources and their characteristics"
    )
    data_format: str = Field(
        # default="...",
        description="Expected format of input data"
    )
    additional_data_requirement: bool = Field(
        # default=False,
        description="Whether additional data is needed"
    )
    constraints: str = Field(
        # default="...",
        description="Business and technical constraints"
    )
    task: MLTaskType = Field(
        # default=MLTaskType.OTHER,
        description="Type of ML task"
    )
    models: List[str] = Field(
        # default=["..."],
        description="Suggested ML models"
    )
    hyperparameters: List[str] = Field(
        # default=["..."],
        description="Key hyperparameters to consider"
    )
    eval_metrics: List[str] = Field(
        # default=["..."],
        description="Evaluation metrics for the solution"
    )
    technical_requirements: str = Field(
        # default="...",
        description="Technical implementation requirements"
    )


class RequirementsAnalysis(BaseModel):
    """Initial analysis of business requirements"""
    model_response: ModelResponseStatus
    unclear_points: List[str] = Field(
        default_factory=list,
        description="Points needing clarification"
    )
    search_queries: List[str] = Field(
        default_factory=list,
        description="Topics to research"
    )
    business_understanding: str = Field(
        description="Summary of business problem understanding"
    )


class TechnicalResearch(BaseModel):
    """Results from technical research"""
    model_response: ModelResponseStatus
    research_findings: str = Field(
        description="Key findings from research"
    )
    reference_implementations: List[str] = Field(
        default_factory=list,
        description="Similar implementation examples found"
    )
    sources: List[str] = Field(
        default_factory=list,
        description="Sources of information"
    )


# Implementation Planning Models
class ComponentType(str, Enum):
    DATA_PIPELINE = "data_pipeline"
    PREPROCESSOR = "preprocessor"
    MODEL = "model"
    EVALUATOR = "evaluator"
    INFERENCE = "inference"
    MONITORING = "monitoring"
    UTILITY = "utility"


class ParameterSpec(BaseModel):
    """Specification for a single parameter"""
    name: str = Field(description="Name of the parameter")
    param_type: str = Field(description="Type of the parameter")
    description: str = Field(description="Description of the parameter")
    default_value: str = Field(description="Default value if any")
    required: bool = Field(description="Whether the parameter is required")


class ConfigParam(BaseModel):
    """Specification for a configuration parameter"""
    name: str = Field(description="Name of the configuration parameter")
    value_type: str = Field(description="Type of value expected")
    description: str = Field(description="Description of the configuration parameter")
    default: str = Field(description="Default value if any")


class FunctionSpec(BaseModel):
    """Detailed specification for a single function"""
    name: str = Field(description="Name of the function")
    description: str = Field(description="Detailed description of function's purpose")
    input_params: List[ParameterSpec] = Field(
        description="List of input parameters and their specifications"
    )
    return_type: str = Field(description="Return type and description")
    dependencies: List[str] = Field(
        description="Required dependencies/imports"
    )
    error_handling: List[str] = Field(
        description="Expected errors and handling strategies"
    )


class ComponentSpec(BaseModel):
    """Specification for a component (module) of the system"""
    name: str = Field(description="Name of the component")
    type: ComponentType = Field(description="Type of component")
    description: str = Field(description="Detailed description of component's purpose")
    functions: List[FunctionSpec] = Field(description="Functions within this component")
    dependencies: List[str] = Field(
        description="External package dependencies"
    )
    config_params: List[ConfigParam] = Field(
        description="Configuration parameters needed"
    )


class ImplementationPlan(BaseModel):
    """Complete implementation plan for the ML system"""
    components: List[ComponentSpec] = Field(description="System components")
    system_requirements: List[str] = Field(
        description="System-level requirements and dependencies"
    )
    deployment_notes: str = Field(
        description="Notes on deployment and infrastructure"
    )
    testing_strategy: str = Field(
        description="Strategy for testing components"
    )
    implementation_order: List[str] = Field(
        description="Suggested order of implementation"
    )