insight_and_tasks/data_models/tasks.py

# data_models/tasks.py
from enum import Enum
from typing import List, Optional, Literal
from pydantic import BaseModel, Field, field_validator
from datetime import datetime

# Using Literal for more precise string enums if preferred over Enum class for Pydantic
# However, Enum provides better structure and can be used with Field choices.

class EffortLevel(str, Enum):
    """Estimated effort level for a task."""
    SMALL = "Small"
    MEDIUM = "Medium"
    LARGE = "Large"

class TaskType(str, Enum):
    """Type of task, indicating its nature."""
    INITIATIVE = "initiative"  # Action-oriented, new projects/changes
    TRACKING = "tracking"      # Measurement-focused, monitoring existing metrics/processes

class DataSubject(str, Enum):
    """Specifies the primary data domain a tracking task relates to."""
    FOLLOWER_STATS = "follower_stats"
    POSTS = "posts"
    MENTIONS = "mentions"
    GENERAL = "general" # For initiatives or tasks not tied to a single data type

class TimelineCategory(str, Enum):
    """Categorization of task timelines."""
    IMMEDIATE = "Immediate"       # (e.g., 1-2 weeks)
    SHORT_TERM = "Short-term"     # (e.g., rest of current quarter, up to 3 months)
    MEDIUM_TERM = "Medium-term"   # (e.g., next quarter, 3-6 months)
    LONG_TERM = "Long-term"       # (e.g., 6+ months)

class PriorityLevel(str, Enum):
    """Priority level for tasks."""
    HIGH = "High"
    MEDIUM = "Medium"
    LOW = "Low"

class Task(BaseModel):
    """
    Represents a single actionable task derived from analysis.
    """
    task_category: str = Field(
        description="The broader category or theme of the task (e.g., Content Strategy, Audience Engagement, Reputation Management, Performance Monitoring)."
    )
    task_description: str = Field( # Renamed from 'task' for clarity
        description="A concise yet clear description of the specific action to be taken."
    )
    objective_deliverable: str = Field(
        description="The clear, measurable objective this task aims to achieve and the specific deliverable(s) expected upon completion."
    )
    effort: EffortLevel = Field(
        description="Estimated effort required to complete the task (Small, Medium, Large)."
    )
    timeline: TimelineCategory = Field(
        description="Projected timeline for task completion, considering urgency and dependencies."
    )
    responsible_party: str = Field(
        description="The team, role, or individual suggested to be responsible for executing this task (e.g., Marketing Team, Content Creation Lead, Social Media Manager)."
    )
    success_criteria_metrics: str = Field(
        description="Specific, measurable criteria and metrics that will be used to determine if the task was successfully completed and achieved its objective."
    )
    dependencies_prerequisites: Optional[str] = Field(
        default=None,
        description="Any other tasks, resources, or conditions that must be met before this task can begin or be completed."
    )
    priority: PriorityLevel = Field(
        description="The priority level of the task (High, Medium, Low)."
    )
    priority_justification: str = Field(
        description="A brief explanation for the assigned priority level, linking it to impact or urgency."
    )
    why_proposed: str = Field(
        description="The rationale behind proposing this task, clearly linking it back to specific findings or insights from the data analysis."
    )
    task_type: TaskType = Field(
        description="Indicates whether this task is a new 'initiative' or ongoing 'tracking' of performance/metrics."
    )
    data_subject: Optional[DataSubject] = Field(
        default=None, 
        description="For 'tracking' tasks, specifies the primary data subject (e.g., follower_stats, posts, mentions). Can be 'general' or null for 'initiative' tasks."
    )

    @field_validator('data_subject')
    @classmethod
    def check_data_subject_for_tracking(cls, value: Optional[DataSubject], values) -> Optional[DataSubject]:
        # Pydantic v2 uses `values.data` to get other field values if needed before validation
        # For Pydantic v1, it would be `values.get('task_type')`
        # This example assumes Pydantic v2 structure for `values` if needed, but here we only need `task_type`
        # which should already be validated or available.
        # For simplicity, accessing it via `values.data.get('task_type')` in Pydantic v2 context.
        # If using Pydantic v1, it's `values.get('task_type')`.
        # Let's assume `values` is a dict-like object containing other fields.
        
        # The validator structure depends on Pydantic version.
        # For Pydantic v2, it's `info: ValidationInfo` and `info.data.get('task_type')`
        # For Pydantic v1, `values` is a dict.
        # For this example, let's assume `values` is a dict of the fields.
        task_type_value = None
        if hasattr(values, 'data'): # Pydantic v2 way
             task_type_value = values.data.get('task_type')
        elif isinstance(values, dict): # Pydantic v1 way (or if it's passed as a dict)
             task_type_value = values.get('task_type')


        if task_type_value == TaskType.TRACKING and value is None:
            raise ValueError("For 'tracking' tasks, 'data_subject' must be specified.")
        if task_type_value == TaskType.INITIATIVE and value is DataSubject.GENERAL:
            # This is acceptable, or you might want to enforce it to be None
            pass
        return value

class KeyResult(BaseModel):
    """
    A measurable outcome that contributes to an Objective.
    """
    key_result_description: str = Field( # Renamed from 'key_result'
        description="A clear, specific, and measurable description of the key result."
    )
    tasks: List[Task] = Field(
        default_factory=list,
        description="A list of specific tasks that will be undertaken to achieve this key result."
    )
    target_metric: Optional[str] = Field(
        default=None,
        description="The primary metric used to measure the achievement of this key result (e.g., 'Follower Growth Rate', 'Average Engagement Rate')."
    )
    target_value: Optional[str] = Field( # Can be numeric or descriptive (e.g., "Increase by 10%", "Achieve 5%")
        default=None,
        description="The specific target value for the metric (e.g., '5%', '1000 new followers')."
    )

class OKR(BaseModel):
    """
    Defines an Objective and its associated Key Results (OKRs).
    """
    objective_description: str = Field( # Renamed from 'objective'
        description="A high-level, qualitative goal that the team aims to achieve. Should be aspirational and motivating."
    )
    key_results: List[KeyResult] = Field(
        default_factory=list,
        description="A list of 2-5 specific, measurable, achievable, relevant, and time-bound (SMART) key results that define success for the objective."
    )
    objective_timeline: TimelineCategory = Field(
        description="The overall timeline category for achieving this objective."
    )
    objective_owner: Optional[str] = Field(
        default=None,
        description="The team or individual primarily responsible for this objective."
    )


class TaskExtractionOutput(BaseModel):
    """
    Structured output from the TaskExtractionAgent, including context and OKRs.
    """
    current_quarter_info: str = Field(
        description="Information about the current quarter and days remaining (e.g., 'Q2 2025, 45 days remaining')."
    )
    okrs: List[OKR] = Field(
        default_factory=list,
        description="A list of Objectives and Key Results (OKRs) derived from the analysis."
    )
    overall_strategic_focus: Optional[str] = Field(
        default=None,
        description="A brief summary of the main strategic focus areas identified from the tasks."
    )
    generation_timestamp: str = Field(
        default_factory=lambda: datetime.utcnow().isoformat(),
        description="Timestamp of when this task extraction output was generated."
    )