Coverage for src / harness_utils / config.py: 78%

1"""Configuration schema for harness-utils.""" 

2 

3from dataclasses import dataclass, field 

4from pathlib import Path 

5from typing import Any 

6 

7 

8@dataclass 

9class TruncationConfig: 

10 """Configuration for Tier 1: Output truncation.""" 

11 

12 max_lines: int = 2000 

13 max_bytes: int = 50 * 1024 # 50KB 

14 direction: str = "head" # "head" or "tail" 

15 

16 

17@dataclass 

18class PruningConfig: 

19 """Configuration for Tier 2: Selective pruning.""" 

20 

21 prune_protect: int = 40_000 # Keep recent 40K tokens 

22 prune_minimum: int = 20_000 # Only prune if saves 20K+ tokens 

23 protect_turns: int = 2 # Protect last 2 turns 

24 protected_tools: list[str] = field(default_factory=lambda: [ 

25 "skill_execution", 

26 "subtask_invocation" 

27 ]) 

28 

29 

30@dataclass 

31class TokenConfig: 

32 """Configuration for token estimation.""" 

33 

34 chars_per_token: int = 4 

35 

36 

37@dataclass 

38class ModelLimitsConfig: 

39 """Configuration for model limits.""" 

40 

41 default_context_limit: int = 200_000 

42 default_output_limit: int = 8_192 

43 

44 

45@dataclass 

46class StorageConfig: 

47 """Configuration for storage layer.""" 

48 

49 base_path: Path = field(default_factory=lambda: Path("data")) 

50 retention_days: int = 7 # For truncated outputs 

51 

52 

53@dataclass 

54class CompactionConfig: 

55 """Configuration for context compaction.""" 

56 

57 auto: bool = True # Enable auto-summarization 

58 prune: bool = True # Enable pruning 

59 

60 

61@dataclass 

62class HarnessConfig: 

63 """Main configuration for harness-utils. 

64 

65 Provides all configuration parameters for context window management 

66 with sensible defaults from the CTXWINARCH.md specification. 

67 """ 

68 

69 truncation: TruncationConfig = field(default_factory=TruncationConfig) 

70 pruning: PruningConfig = field(default_factory=PruningConfig) 

71 tokens: TokenConfig = field(default_factory=TokenConfig) 

72 model_limits: ModelLimitsConfig = field(default_factory=ModelLimitsConfig) 

73 storage: StorageConfig = field(default_factory=StorageConfig) 

74 compaction: CompactionConfig = field(default_factory=CompactionConfig) 

75 

76 @classmethod 

77 def from_dict(cls, data: dict[str, Any]) -> "HarnessConfig": 

78 """Create configuration from dictionary. 

79 

80 Args: 

81 data: Configuration dictionary 

82 

83 Returns: 

84 HarnessConfig instance 

85 """ 

86 config = cls() 

87 

88 if "truncation" in data: 

89 config.truncation = TruncationConfig(**data["truncation"]) 

90 if "pruning" in data: 

91 protected = data["pruning"].get("protected_tools") 

92 pruning_data = {k: v for k, v in data["pruning"].items() 

93 if k != "protected_tools"} 

94 if protected: 

95 pruning_data["protected_tools"] = protected 

96 config.pruning = PruningConfig(**pruning_data) 

97 if "tokens" in data: 

98 config.tokens = TokenConfig(**data["tokens"]) 

99 if "model_limits" in data: 

100 config.model_limits = ModelLimitsConfig(**data["model_limits"]) 

101 if "storage" in data: 

102 storage_data = data["storage"].copy() 

103 if "base_path" in storage_data: 

104 storage_data["base_path"] = Path(storage_data["base_path"]) 

105 config.storage = StorageConfig(**storage_data) 

106 if "compaction" in data: 

107 config.compaction = CompactionConfig(**data["compaction"]) 

108 

109 return config 

110 

111 @classmethod 

112 def from_toml(cls, path: Path) -> "HarnessConfig": 

113 """Load configuration from TOML file. 

114 

115 Args: 

116 path: Path to TOML configuration file 

117 

118 Returns: 

119 HarnessConfig instance 

120 """ 

121 import tomllib 

122 

123 with open(path, "rb") as f: 

124 data = tomllib.load(f) 

125 

126 return cls.from_dict(data) 

127 

128 @classmethod 

129 def from_json(cls, path: Path) -> "HarnessConfig": 

130 """Load configuration from JSON file. 

131 

132 Args: 

133 path: Path to JSON configuration file 

134 

135 Returns: 

136 HarnessConfig instance 

137 """ 

138 import json 

139 

140 with open(path, "r") as f: 

141 data = json.load(f) 

142 

143 return cls.from_dict(data)