""" Data models for issue context estimation. Defines enums and data classes used by the context estimator. """ from dataclasses import dataclass from enum import Enum from typing import Optional class ComplexityLevel(Enum): """Implementation complexity levels with token allocations.""" LOW = 10000 # Simple, straightforward implementation MEDIUM = 20000 # Moderate complexity, some edge cases HIGH = 30000 # Complex logic, many edge cases, algorithms class TestLevel(Enum): """Test requirement levels with token allocations.""" LOW = 5000 # Basic unit tests MEDIUM = 10000 # Unit + integration tests HIGH = 15000 # Unit + integration + E2E tests class DocLevel(Enum): """Documentation requirement levels with token allocations.""" NONE = 0 # No documentation needed LIGHT = 2000 # Inline comments, basic docstrings MEDIUM = 3000 # API docs, usage examples HEAVY = 5000 # Comprehensive docs, guides, tutorials @dataclass class EstimationInput: """ Input parameters for context estimation. Attributes: files_to_modify: Number of files expected to be modified implementation_complexity: Complexity level of implementation test_requirements: Level of testing required documentation: Level of documentation required """ files_to_modify: int implementation_complexity: ComplexityLevel test_requirements: TestLevel documentation: DocLevel def __post_init__(self): """Validate input parameters.""" if self.files_to_modify < 0: raise ValueError("files_to_modify must be >= 0") if not isinstance(self.implementation_complexity, ComplexityLevel): raise TypeError( f"implementation_complexity must be ComplexityLevel, " f"got {type(self.implementation_complexity)}" ) if not isinstance(self.test_requirements, TestLevel): raise TypeError( f"test_requirements must be TestLevel, " f"got {type(self.test_requirements)}" ) if not isinstance(self.documentation, DocLevel): raise TypeError( f"documentation must be DocLevel, " f"got {type(self.documentation)}" ) @dataclass class EstimationResult: """ Result of context estimation. Contains breakdown of token allocation by category and total estimate with safety buffer applied. Attributes: files_context: Tokens allocated for file context (files × 7000) implementation_tokens: Tokens allocated for implementation test_tokens: Tokens allocated for tests doc_tokens: Tokens allocated for documentation base_estimate: Sum of all components before buffer buffer_tokens: Safety buffer (30% of base) total_estimate: Final estimate with buffer applied recommended_agent: Suggested agent based on total (haiku/sonnet/opus) """ files_context: int implementation_tokens: int test_tokens: int doc_tokens: int base_estimate: int buffer_tokens: int total_estimate: int recommended_agent: str def to_dict(self) -> dict: """Convert result to dictionary for serialization.""" return { "files_context": self.files_context, "implementation_tokens": self.implementation_tokens, "test_tokens": self.test_tokens, "doc_tokens": self.doc_tokens, "base_estimate": self.base_estimate, "buffer_tokens": self.buffer_tokens, "total_estimate": self.total_estimate, "recommended_agent": self.recommended_agent, } @dataclass class ValidationResult: """ Result of validating estimate against actual usage. Used for historical validation and accuracy tracking. Attributes: issue_number: Issue number validated estimated_tokens: Estimated context tokens actual_tokens: Actual tokens used (if known) percentage_error: Percentage error (estimated vs actual) within_tolerance: Whether error is within ±20% tolerance notes: Optional notes about validation """ issue_number: int estimated_tokens: int actual_tokens: Optional[int] = None percentage_error: Optional[float] = None within_tolerance: Optional[bool] = None notes: Optional[str] = None def __post_init__(self): """Calculate derived fields if actual_tokens provided.""" if self.actual_tokens is not None: self.percentage_error = abs( self.estimated_tokens - self.actual_tokens ) / self.actual_tokens self.within_tolerance = self.percentage_error <= 0.20 def to_dict(self) -> dict: """Convert result to dictionary for serialization.""" return { "issue_number": self.issue_number, "estimated_tokens": self.estimated_tokens, "actual_tokens": self.actual_tokens, "percentage_error": ( f"{self.percentage_error:.2%}" if self.percentage_error else None ), "within_tolerance": self.within_tolerance, "notes": self.notes, }