Coverage for src / moai_adk / core / jit_context_loader.py: 36.36%

374 statements  

« prev     ^ index     » next       coverage.py v7.12.0, created at 2025-11-20 20:52 +0900

1#!/usr/bin/env python3 

2""" 

3JIT Context Loading System 

4 

5Phase-based token optimization system that achieves 85%+ efficiency through 

6intelligent context loading, skill filtering, and budget management. 

7""" 

8 

9import hashlib 

10import json 

11import logging 

12import os 

13import sys 

14import time 

15from collections import OrderedDict 

16from dataclasses import dataclass, field 

17from datetime import datetime 

18from enum import Enum 

19from pathlib import Path 

20from typing import Any, Dict, List, Optional, Tuple 

21 

22import psutil 

23 

24logger = logging.getLogger(__name__) 

25 

26class Phase(Enum): 

27 """Development phases with different token budgets""" 

28 SPEC = "spec" # 30K budget 

29 RED = "red" # 25K budget 

30 GREEN = "green" # 25K budget 

31 REFACTOR = "refactor" # 20K budget 

32 SYNC = "sync" # 40K budget 

33 DEBUG = "debug" # 15K budget 

34 PLANNING = "planning" # 35K budget 

35 

36@dataclass 

37class PhaseConfig: 

38 """Configuration for each development phase""" 

39 max_tokens: int 

40 essential_skills: List[str] 

41 essential_documents: List[str] 

42 cache_clear_on_phase_change: bool = False 

43 context_compression: bool = True 

44 

45@dataclass 

46class ContextMetrics: 

47 """Metrics for context loading performance""" 

48 load_time: float 

49 token_count: int 

50 cache_hit: bool 

51 phase: str 

52 skills_loaded: int 

53 docs_loaded: int 

54 compression_ratio: float = 1.0 

55 memory_usage: int = 0 

56 

57@dataclass 

58class SkillInfo: 

59 """Information about a skill for filtering decisions""" 

60 name: str 

61 path: str 

62 size: int 

63 tokens: int 

64 categories: List[str] 

65 dependencies: List[str] = field(default_factory=list) 

66 priority: int = 1 # 1=high, 2=medium, 3=low 

67 last_used: Optional[datetime] = None 

68 

69@dataclass 

70class ContextEntry: 

71 """Entry in the context cache""" 

72 key: str 

73 content: Any 

74 token_count: int 

75 created_at: datetime 

76 last_accessed: datetime 

77 access_count: int = 0 

78 phase: Optional[str] = None 

79 

80class PhaseDetector: 

81 """Intelligently detects current development phase from context""" 

82 

83 def __init__(self): 

84 self.phase_patterns = { 

85 Phase.SPEC: [ 

86 r'/moai:1-plan', 

87 r'SPEC-\d+', 

88 r'spec|requirements|design', 

89 r'create.*spec|define.*requirements', 

90 r'plan.*feature|design.*system' 

91 ], 

92 Phase.RED: [ 

93 r'/moai:2-run.*RED', 

94 r'test.*fail|failing.*test', 

95 r'red.*phase|tdd.*red', 

96 r'2-run.*RED|write.*test', 

97 r'create.*test.*failure' 

98 ], 

99 Phase.GREEN: [ 

100 r'/moai:2-run.*GREEN', 

101 r'test.*pass|passing.*test', 

102 r'green.*phase|tdd.*green', 

103 r'2-run.*GREEN|minimal.*implementation', 

104 r'make.*test.*pass|implement.*minimum' 

105 ], 

106 Phase.REFACTOR: [ 

107 r'/moai:2-run.*REFACTOR', 

108 r'refactor|clean.*code', 

109 r'quality.*improvement|improve.*code', 

110 r'optimize.*code|code.*cleanup' 

111 ], 

112 Phase.SYNC: [ 

113 r'/moai:3-sync', 

114 r'documentation.*sync|sync.*docs', 

115 r'generate.*docs|create.*documentation', 

116 r'update.*documentation' 

117 ], 

118 Phase.DEBUG: [ 

119 r'debug|troubleshoot', 

120 r'error.*analysis|analyze.*error', 

121 r'fix.*bug|error.*investigation', 

122 r'problem.*solving' 

123 ], 

124 Phase.PLANNING: [ 

125 r'plan.*implementation|implementation.*plan', 

126 r'architecture.*design|design.*architecture', 

127 r'task.*decomposition|breakdown.*task', 

128 r'plan.*development|development.*planning' 

129 ] 

130 } 

131 

132 self.last_phase = Phase.SPEC 

133 self.phase_history = [] 

134 self.max_history = 10 

135 

136 def detect_phase(self, user_input: str, conversation_history: List[str] = None) -> Phase: 

137 """Detect current phase from user input and conversation context""" 

138 import re 

139 

140 # Combine user input with recent conversation history 

141 context = user_input.lower() 

142 if conversation_history: 

143 recent_history = " ".join(conversation_history[-3:]) # Last 3 messages 

144 context += " " + recent_history.lower() 

145 

146 # Score each phase based on pattern matches 

147 phase_scores = {} 

148 for phase, patterns in self.phase_patterns.items(): 

149 score = 0 

150 for pattern in patterns: 

151 matches = len(re.findall(pattern, context, re.IGNORECASE)) 

152 score += matches 

153 phase_scores[phase] = score 

154 

155 # Find phase with highest score 

156 if phase_scores: 

157 best_phase = max(phase_scores, key=phase_scores.get) 

158 else: 

159 best_phase = self.last_phase 

160 

161 # If no clear winner (all scores 0), use last known phase 

162 if phase_scores[best_phase] == 0: 

163 best_phase = self.last_phase 

164 

165 # Update history 

166 if best_phase != self.last_phase: 

167 self.phase_history.append({ 

168 'from': self.last_phase.value, 

169 'to': best_phase.value, 

170 'timestamp': datetime.now(), 

171 'context': user_input[:100] + "..." if len(user_input) > 100 else user_input 

172 }) 

173 

174 if len(self.phase_history) > self.max_history: 

175 self.phase_history.pop(0) 

176 

177 self.last_phase = best_phase 

178 return best_phase 

179 

180 def get_phase_config(self, phase: Phase) -> PhaseConfig: 

181 """Get configuration for a specific phase""" 

182 configs = { 

183 Phase.SPEC: PhaseConfig( 

184 max_tokens=30000, 

185 essential_skills=[ 

186 "moai-foundation-ears", 

187 "moai-foundation-specs", 

188 "moai-essentials-review", 

189 "moai-domain-backend", 

190 "moai-lang-python", 

191 "moai-core-spec-authoring" 

192 ], 

193 essential_documents=[ 

194 ".moai/specs/template.md", 

195 ".claude/skills/moai-foundation-ears/SKILL.md" 

196 ], 

197 cache_clear_on_phase_change=True 

198 ), 

199 Phase.RED: PhaseConfig( 

200 max_tokens=25000, 

201 essential_skills=[ 

202 "moai-domain-testing", 

203 "moai-foundation-trust", 

204 "moai-essentials-review", 

205 "moai-core-code-reviewer", 

206 "moai-essentials-debug", 

207 "moai-lang-python" 

208 ], 

209 essential_documents=[ 

210 ".moai/specs/{spec_id}/spec.md", 

211 ".claude/skills/moai-domain-testing/SKILL.md" 

212 ] 

213 ), 

214 Phase.GREEN: PhaseConfig( 

215 max_tokens=25000, 

216 essential_skills=[ 

217 "moai-lang-python", 

218 "moai-domain-backend", 

219 "moai-essentials-review" 

220 ], 

221 essential_documents=[ 

222 ".moai/specs/{spec_id}/spec.md" 

223 ] 

224 ), 

225 Phase.REFACTOR: PhaseConfig( 

226 max_tokens=20000, 

227 essential_skills=[ 

228 "moai-essentials-refactor", 

229 "moai-essentials-review", 

230 "moai-core-code-reviewer", 

231 "moai-essentials-debug" 

232 ], 

233 essential_documents=[ 

234 "src/{module}/current_implementation.py", 

235 ".claude/skills/moai-essentials-refactor/SKILL.md" 

236 ] 

237 ), 

238 Phase.SYNC: PhaseConfig( 

239 max_tokens=40000, 

240 essential_skills=[ 

241 "moai-docs-unified", 

242 "moai-nextra-architecture", 

243 "moai-core-spec-authoring", 

244 "moai-cc-configuration" 

245 ], 

246 essential_documents=[ 

247 ".moai/specs/{spec_id}/implementation.md", 

248 ".moai/specs/{spec_id}/test-cases.md" 

249 ], 

250 cache_clear_on_phase_change=True 

251 ), 

252 Phase.DEBUG: PhaseConfig( 

253 max_tokens=15000, 

254 essential_skills=[ 

255 "moai-essentials-debug", 

256 "moai-core-code-reviewer" 

257 ], 

258 essential_documents=[ 

259 ".moai/logs/latest_error.log" 

260 ] 

261 ), 

262 Phase.PLANNING: PhaseConfig( 

263 max_tokens=35000, 

264 essential_skills=[ 

265 "moai-core-practices", 

266 "moai-essentials-review", 

267 "moai-foundation-specs" 

268 ], 

269 essential_documents=[ 

270 ".moai/specs/{spec_id}/spec.md" 

271 ] 

272 ) 

273 } 

274 

275 return configs.get(phase, configs[Phase.SPEC]) 

276 

277class SkillFilterEngine: 

278 """Intelligently filters and selects skills based on phase and context""" 

279 

280 def __init__(self, skills_dir: str = ".claude/skills"): 

281 self.skills_dir = Path(skills_dir) 

282 self.skills_cache = {} 

283 self.skill_index = {} 

284 self.phase_preferences = self._load_phase_preferences() 

285 self._build_skill_index() 

286 

287 def _load_phase_preferences(self) -> Dict[str, Dict[str, int]]: 

288 """Load phase-based skill preferences""" 

289 return { 

290 "spec": { 

291 "moai-foundation-ears": 1, 

292 "moai-foundation-specs": 1, 

293 "moai-essentials-review": 2, 

294 "moai-domain-backend": 2, 

295 "moai-lang-python": 3, 

296 "moai-core-spec-authoring": 1 

297 }, 

298 "red": { 

299 "moai-domain-testing": 1, 

300 "moai-foundation-trust": 1, 

301 "moai-essentials-review": 2, 

302 "moai-core-code-reviewer": 2, 

303 "moai-essentials-debug": 2, 

304 "moai-lang-python": 3 

305 }, 

306 "green": { 

307 "moai-lang-python": 1, 

308 "moai-domain-backend": 1, 

309 "moai-essentials-review": 2 

310 }, 

311 "refactor": { 

312 "moai-essentials-refactor": 1, 

313 "moai-essentials-review": 2, 

314 "moai-core-code-reviewer": 2, 

315 "moai-essentials-debug": 3 

316 }, 

317 "sync": { 

318 "moai-docs-unified": 1, 

319 "moai-nextra-architecture": 1, 

320 "moai-core-spec-authoring": 2, 

321 "moai-cc-configuration": 2 

322 }, 

323 "debug": { 

324 "moai-essentials-debug": 1, 

325 "moai-core-code-reviewer": 2 

326 }, 

327 "planning": { 

328 "moai-core-practices": 1, 

329 "moai-essentials-review": 2, 

330 "moai-foundation-specs": 2 

331 } 

332 } 

333 

334 def _build_skill_index(self): 

335 """Build index of all available skills with metadata""" 

336 if not self.skills_dir.exists(): 

337 logger.warning(f"Skills directory not found: {self.skills_dir}") 

338 return 

339 

340 for skill_dir in self.skills_dir.iterdir(): 

341 if skill_dir.is_dir(): 

342 skill_file = skill_dir / "SKILL.md" 

343 if skill_file.exists(): 

344 skill_info = self._analyze_skill(skill_file) 

345 if skill_info: 

346 self.skill_index[skill_info.name] = skill_info 

347 

348 def _analyze_skill(self, skill_file: Path) -> Optional[SkillInfo]: 

349 """Analyze a skill file to extract metadata""" 

350 try: 

351 stat = skill_file.stat() 

352 

353 # Read skill content to extract categories 

354 with open(skill_file, 'r', encoding='utf-8') as f: 

355 content = f.read() 

356 

357 # Extract skill name from directory 

358 skill_name = skill_file.parent.name 

359 

360 # Estimate tokens (rough approximation: 1 token ≈ 4 characters) 

361 estimated_tokens = len(content) // 4 

362 

363 # Extract categories from content (look for keywords) 

364 categories = [] 

365 if any(keyword in content.lower() for keyword in ['python', 'javascript', 'typescript', 'go', 'rust']): 

366 categories.append('language') 

367 if any(keyword in content.lower() for keyword in ['backend', 'frontend', 'database', 'security']): 

368 categories.append('domain') 

369 if any(keyword in content.lower() for keyword in ['testing', 'debug', 'refactor', 'review']): 

370 categories.append('development') 

371 if any(keyword in content.lower() for keyword in ['foundation', 'essential', 'core']): 

372 categories.append('core') 

373 

374 return SkillInfo( 

375 name=skill_name, 

376 path=str(skill_file), 

377 size=stat.st_size, 

378 tokens=estimated_tokens, 

379 categories=categories, 

380 priority=1 

381 ) 

382 

383 except Exception as e: 

384 logger.error(f"Error analyzing skill {skill_file}: {e}") 

385 return None 

386 

387 def filter_skills(self, phase: Phase, token_budget: int, context: Dict[str, Any] = None) -> List[SkillInfo]: 

388 """Filter skills based on phase, token budget, and context""" 

389 phase_name = phase.value 

390 preferences = self.phase_preferences.get(phase_name, {}) 

391 

392 # Get all relevant skills for this phase 

393 relevant_skills = [] 

394 for skill_name, skill_info in self.skill_index.items(): 

395 # Check if skill is relevant for this phase 

396 if skill_name in preferences: 

397 # Apply priority from preferences 

398 skill_info.priority = preferences[skill_name] 

399 relevant_skills.append(skill_info) 

400 

401 # Sort by priority and token efficiency 

402 relevant_skills.sort(key=lambda s: (s.priority, s.tokens)) 

403 

404 # Select skills within token budget 

405 selected_skills = [] 

406 used_tokens = 0 

407 

408 for skill in relevant_skills: 

409 if used_tokens + skill.tokens <= token_budget: 

410 selected_skills.append(skill) 

411 used_tokens += skill.tokens 

412 else: 

413 break 

414 

415 return selected_skills 

416 

417 def get_skill_stats(self) -> Dict[str, Any]: 

418 """Get statistics about available skills""" 

419 total_skills = len(self.skill_index) 

420 total_tokens = sum(skill.tokens for skill in self.skill_index.values()) 

421 

422 categories = {} 

423 for skill in self.skill_index.values(): 

424 for category in skill.categories: 

425 categories[category] = categories.get(category, 0) + 1 

426 

427 return { 

428 'total_skills': total_skills, 

429 'total_tokens': total_tokens, 

430 'categories': categories, 

431 'average_tokens_per_skill': total_tokens / total_skills if total_skills > 0 else 0 

432 } 

433 

434class TokenBudgetManager: 

435 """Manages token budgets and usage across phases""" 

436 

437 def __init__(self, max_total_tokens: int = 180000): 

438 self.max_total_tokens = max_total_tokens 

439 self.phase_budgets = self._initialize_phase_budgets() 

440 self.current_usage = 0 

441 self.usage_history = [] 

442 self.budget_warnings = [] 

443 

444 def _initialize_phase_budgets(self) -> Dict[str, int]: 

445 """Initialize token budgets for each phase""" 

446 return { 

447 "spec": 30000, 

448 "red": 25000, 

449 "green": 25000, 

450 "refactor": 20000, 

451 "sync": 40000, 

452 "debug": 15000, 

453 "planning": 35000 

454 } 

455 

456 def check_budget(self, phase: Phase, requested_tokens: int) -> Tuple[bool, int]: 

457 """Check if requested tokens fit in budget""" 

458 phase_budget = self.phase_budgets.get(phase.value, 30000) 

459 

460 if requested_tokens <= phase_budget: 

461 return True, phase_budget - requested_tokens 

462 else: 

463 return False, phase_budget 

464 

465 def record_usage(self, phase: Phase, tokens_used: int, context: str = ""): 

466 """Record actual token usage""" 

467 usage_entry = { 

468 'phase': phase.value, 

469 'tokens': tokens_used, 

470 'timestamp': datetime.now(), 

471 'context': context[:100] 

472 } 

473 

474 self.usage_history.append(usage_entry) 

475 self.current_usage += tokens_used 

476 

477 # Keep only recent history (last 50 entries) 

478 if len(self.usage_history) > 50: 

479 self.usage_history.pop(0) 

480 

481 # Check for budget warnings 

482 phase_budget = self.phase_budgets.get(phase.value, 30000) 

483 if tokens_used > phase_budget: 

484 warning = f"Phase {phase.value} exceeded budget: {tokens_used} > {phase_budget}" 

485 self.budget_warnings.append({ 

486 'warning': warning, 

487 'timestamp': datetime.now() 

488 }) 

489 logger.warning(warning) 

490 

491 def get_efficiency_metrics(self) -> Dict[str, Any]: 

492 """Calculate and return efficiency metrics""" 

493 if not self.usage_history: 

494 return { 

495 'efficiency_score': 0, 

496 'average_phase_efficiency': {}, 

497 'budget_compliance': 100, 

498 'total_usage': 0 

499 } 

500 

501 # Calculate efficiency by phase 

502 phase_usage = {} 

503 phase_efficiency = {} 

504 

505 for entry in self.usage_history: 

506 phase = entry['phase'] 

507 tokens = entry['tokens'] 

508 budget = self.phase_budgets.get(phase, 30000) 

509 

510 if phase not in phase_usage: 

511 phase_usage[phase] = {'total': 0, 'count': 0, 'over_budget': 0} 

512 

513 phase_usage[phase]['total'] += tokens 

514 phase_usage[phase]['count'] += 1 

515 if tokens > budget: 

516 phase_usage[phase]['over_budget'] += 1 

517 

518 # Calculate efficiency scores 

519 for phase, usage in phase_usage.items(): 

520 budget = self.phase_budgets.get(phase, 30000) 

521 actual = usage['total'] / usage['count'] if usage['count'] > 0 else 0 

522 efficiency = min(100, (budget / actual * 100) if actual > 0 else 100) 

523 phase_efficiency[phase] = efficiency 

524 

525 # Overall efficiency 

526 overall_efficiency = sum(phase_efficiency.values()) / len(phase_efficiency) if phase_efficiency else 0 

527 

528 # Budget compliance 

529 total_entries = len(self.usage_history) 

530 over_budget_entries = sum(usage['over_budget'] for usage in phase_usage.values()) 

531 budget_compliance = ((total_entries - over_budget_entries) / total_entries * 100) if total_entries > 0 else 100 

532 

533 return { 

534 'efficiency_score': overall_efficiency, 

535 'average_phase_efficiency': phase_efficiency, 

536 'budget_compliance': budget_compliance, 

537 'total_usage': self.current_usage, 

538 'phase_usage': phase_usage 

539 } 

540 

541class ContextCache: 

542 """LRU Cache for context entries with phase-aware eviction""" 

543 

544 def __init__(self, max_size: int = 100, max_memory_mb: int = 50): 

545 self.max_size = max_size 

546 self.max_memory_bytes = max_memory_mb * 1024 * 1024 

547 self.cache = OrderedDict() 

548 self.current_memory = 0 

549 self.hits = 0 

550 self.misses = 0 

551 self.evictions = 0 

552 

553 def _calculate_memory_usage(self, entry: ContextEntry) -> int: 

554 """Calculate memory usage of a cache entry""" 

555 

556 # Rough estimation of memory usage 

557 content_size = sys.getsizeof(entry.content) 

558 entry_overhead = sys.getsizeof(entry.key) + sys.getsizeof(str(entry.token_count)) 

559 total_size = content_size + entry_overhead 

560 

561 return total_size 

562 

563 def get(self, key: str) -> Optional[ContextEntry]: 

564 """Get entry from cache""" 

565 if key in self.cache: 

566 entry = self.cache[key] 

567 entry.last_accessed = datetime.now() 

568 entry.access_count += 1 

569 self.cache.move_to_end(key) 

570 self.hits += 1 

571 return entry 

572 

573 self.misses += 1 

574 return None 

575 

576 def put(self, key: str, content: Any, token_count: int, phase: Optional[str] = None): 

577 """Put entry in cache with LRU eviction""" 

578 entry = ContextEntry( 

579 key=key, 

580 content=content, 

581 token_count=token_count, 

582 created_at=datetime.now(), 

583 last_accessed=datetime.now(), 

584 phase=phase 

585 ) 

586 

587 entry_memory = self._calculate_memory_usage(entry) 

588 

589 # Check if we need to evict entries 

590 while (len(self.cache) >= self.max_size or 

591 self.current_memory + entry_memory > self.max_memory_bytes): 

592 if not self.cache: 

593 break 

594 

595 oldest_key = next(iter(self.cache)) 

596 oldest_entry = self.cache.pop(oldest_key) 

597 self.current_memory -= self._calculate_memory_usage(oldest_entry) 

598 self.evictions += 1 

599 

600 # Add new entry 

601 self.cache[key] = entry 

602 self.current_memory += entry_memory 

603 

604 def clear_phase(self, phase: str): 

605 """Clear all entries for a specific phase""" 

606 keys_to_remove = [key for key, entry in self.cache.items() if entry.phase == phase] 

607 for key in keys_to_remove: 

608 entry = self.cache.pop(key) 

609 self.current_memory -= self._calculate_memory_usage(entry) 

610 

611 def clear(self): 

612 """Clear all cache entries""" 

613 self.cache.clear() 

614 self.current_memory = 0 

615 

616 def get_stats(self) -> Dict[str, Any]: 

617 """Get cache statistics""" 

618 total_requests = self.hits + self.misses 

619 hit_rate = (self.hits / total_requests * 100) if total_requests > 0 else 0 

620 

621 return { 

622 'entries': len(self.cache), 

623 'memory_usage_bytes': self.current_memory, 

624 'memory_usage_mb': self.current_memory / (1024 * 1024), 

625 'hit_rate': hit_rate, 

626 'hits': self.hits, 

627 'misses': self.misses, 

628 'evictions': self.evictions 

629 } 

630 

631class JITContextLoader: 

632 """Main JIT Context Loading System orchestrator""" 

633 

634 def __init__(self, cache_size: int = 100, cache_memory_mb: int = 50): 

635 self.phase_detector = PhaseDetector() 

636 self.skill_filter = SkillFilterEngine() 

637 self.token_manager = TokenBudgetManager() 

638 self.context_cache = ContextCache(cache_size, cache_memory_mb) 

639 

640 self.metrics_history = [] 

641 self.current_phase = Phase.SPEC 

642 

643 # Performance monitoring 

644 self.performance_stats = { 

645 'total_loads': 0, 

646 'average_load_time': 0, 

647 'cache_hit_rate': 0, 

648 'efficiency_score': 0 

649 } 

650 

651 async def load_context(self, user_input: str, conversation_history: List[str] = None, 

652 context: Dict[str, Any] = None) -> Tuple[Dict[str, Any], ContextMetrics]: 

653 """Load optimized context based on current phase and requirements""" 

654 start_time = time.time() 

655 

656 # Detect current phase 

657 self.current_phase = self.phase_detector.detect_phase(user_input, conversation_history or []) 

658 phase_config = self.phase_detector.get_phase_config(self.current_phase) 

659 

660 # Generate cache key 

661 cache_key = self._generate_cache_key(self.current_phase, user_input, context or {}) 

662 

663 # Check cache first 

664 cached_entry = self.context_cache.get(cache_key) 

665 if cached_entry: 

666 metrics = ContextMetrics( 

667 load_time=time.time() - start_time, 

668 token_count=cached_entry.token_count, 

669 cache_hit=True, 

670 phase=self.current_phase.value, 

671 skills_loaded=len(cached_entry.content.get('skills', [])), 

672 docs_loaded=len(cached_entry.content.get('documents', [])), 

673 memory_usage=psutil.Process().memory_info().rss 

674 ) 

675 

676 self._record_metrics(metrics) 

677 return cached_entry.content, metrics 

678 

679 # Load fresh context 

680 context_data = await self._build_context(self.current_phase, phase_config, context or {}) 

681 

682 # Calculate total tokens 

683 total_tokens = self._calculate_total_tokens(context_data) 

684 

685 # Check token budget 

686 within_budget, remaining_budget = self.token_manager.check_budget(self.current_phase, total_tokens) 

687 

688 if not within_budget: 

689 # Apply aggressive optimization 

690 context_data = await self._optimize_context_aggressively(context_data, remaining_budget) 

691 total_tokens = self._calculate_total_tokens(context_data) 

692 

693 # Cache the result 

694 self.context_cache.put(cache_key, context_data, total_tokens, self.current_phase.value) 

695 

696 # Record usage 

697 load_time = time.time() - start_time 

698 self.token_manager.record_usage(self.current_phase, total_tokens, user_input) 

699 

700 metrics = ContextMetrics( 

701 load_time=load_time, 

702 token_count=total_tokens, 

703 cache_hit=False, 

704 phase=self.current_phase.value, 

705 skills_loaded=len(context_data.get('skills', [])), 

706 docs_loaded=len(context_data.get('documents', [])), 

707 memory_usage=psutil.Process().memory_info().rss 

708 ) 

709 

710 self._record_metrics(metrics) 

711 return context_data, metrics 

712 

713 def _generate_cache_key(self, phase: Phase, user_input: str, context: Dict[str, Any]) -> str: 

714 """Generate unique cache key for context request""" 

715 key_data = { 

716 'phase': phase.value, 

717 'input_hash': hashlib.md5(user_input.encode()).hexdigest()[:16], 

718 'context_keys': sorted(context.keys()), 

719 'timestamp': datetime.now().strftime('%Y%m%d') # Daily cache 

720 } 

721 

722 return hashlib.md5(json.dumps(key_data, sort_keys=True).encode()).hexdigest() 

723 

724 async def _build_context(self, phase: Phase, phase_config: PhaseConfig, 

725 context: Dict[str, Any]) -> Dict[str, Any]: 

726 """Build optimized context for the current phase""" 

727 context_data = { 

728 'phase': phase.value, 

729 'skills': [], 

730 'documents': [], 

731 'metadata': { 

732 'created_at': datetime.now().isoformat(), 

733 'max_tokens': phase_config.max_tokens, 

734 'compression_enabled': phase_config.context_compression 

735 } 

736 } 

737 

738 # Filter and load skills 

739 skills = self.skill_filter.filter_skills(phase, phase_config.max_tokens // 2, context) 

740 for skill in skills: 

741 skill_content = await self._load_skill_content(skill) 

742 if skill_content: 

743 context_data['skills'].append({ 

744 'name': skill.name, 

745 'content': skill_content, 

746 'tokens': skill.tokens, 

747 'categories': skill.categories, 

748 'priority': skill.priority 

749 }) 

750 

751 # Load essential documents 

752 for doc_path in phase_config.essential_documents: 

753 doc_content = await self._load_document(doc_path, context) 

754 if doc_content: 

755 context_data['documents'].append({ 

756 'path': doc_path, 

757 'content': doc_content['content'], 

758 'tokens': doc_content['tokens'], 

759 'type': doc_content['type'] 

760 }) 

761 

762 return context_data 

763 

764 async def _load_skill_content(self, skill: SkillInfo) -> Optional[str]: 

765 """Load content of a skill file""" 

766 try: 

767 with open(skill.path, 'r', encoding='utf-8') as f: 

768 content = f.read() 

769 return content 

770 except Exception as e: 

771 logger.error(f"Error loading skill {skill.name}: {e}") 

772 return None 

773 

774 async def _load_document(self, doc_path: str, context: Dict[str, Any]) -> Optional[Dict[str, Any]]: 

775 """Load and process a document""" 

776 try: 

777 # Apply context variable substitution 

778 formatted_path = doc_path.format( 

779 spec_id=context.get('spec_id', 'SPEC-XXX'), 

780 module=context.get('module', 'unknown'), 

781 language=context.get('language', 'python') 

782 ) 

783 

784 if os.path.exists(formatted_path): 

785 with open(formatted_path, 'r', encoding='utf-8') as f: 

786 content = f.read() 

787 

788 return { 

789 'content': content, 

790 'tokens': len(content) // 4, # Rough token estimation 

791 'type': self._detect_document_type(formatted_path) 

792 } 

793 except Exception as e: 

794 logger.error(f"Error loading document {doc_path}: {e}") 

795 return None 

796 

797 def _detect_document_type(self, file_path: str) -> str: 

798 """Detect document type from file path""" 

799 if file_path.endswith('.md'): 

800 return 'markdown' 

801 elif file_path.endswith('.py'): 

802 return 'python' 

803 elif file_path.endswith('.json'): 

804 return 'json' 

805 elif file_path.endswith('.yaml') or file_path.endswith('.yml'): 

806 return 'yaml' 

807 else: 

808 return 'text' 

809 

810 async def _optimize_context_aggressively(self, context_data: Dict[str, Any], 

811 token_budget: int) -> Dict[str, Any]: 

812 """Apply aggressive optimization to fit token budget""" 

813 current_tokens = self._calculate_total_tokens(context_data) 

814 

815 if current_tokens <= token_budget: 

816 return context_data 

817 

818 # Strategy 1: Remove low-priority skills 

819 skills = context_data.get('skills', []) 

820 skills.sort(key=lambda s: s.get('priority', 3)) # Remove low priority first 

821 

822 optimized_skills = [] 

823 used_tokens = 0 

824 

825 for skill in skills: 

826 if used_tokens + skill.get('tokens', 0) <= token_budget * 0.7: # Reserve 30% for docs 

827 optimized_skills.append(skill) 

828 used_tokens += skill.get('tokens', 0) 

829 else: 

830 break 

831 

832 context_data['skills'] = optimized_skills 

833 

834 # Strategy 2: Compress documents if still over budget 

835 if self._calculate_total_tokens(context_data) > token_budget: 

836 documents = context_data.get('documents', []) 

837 for doc in documents: 

838 if doc.get('tokens', 0) > 1000: # Only compress large documents 

839 doc['content'] = self._compress_text(doc['content']) 

840 doc['tokens'] = len(doc['content']) // 4 

841 doc['compressed'] = True 

842 

843 return context_data 

844 

845 def _compress_text(self, text: str) -> str: 

846 """Simple text compression by removing redundancy""" 

847 lines = text.split('\n') 

848 compressed_lines = [] 

849 

850 for line in lines: 

851 stripped = line.strip() 

852 # Skip empty lines and common comment patterns 

853 if (stripped and 

854 not stripped.startswith('#') and 

855 not stripped.startswith('//') and 

856 len(stripped) > 10): 

857 compressed_lines.append(stripped) 

858 

859 return '\n'.join(compressed_lines) 

860 

861 def _calculate_total_tokens(self, context_data: Dict[str, Any]) -> int: 

862 """Calculate total tokens in context data""" 

863 total_tokens = 0 

864 

865 # Count skill tokens 

866 for skill in context_data.get('skills', []): 

867 total_tokens += skill.get('tokens', 0) 

868 

869 # Count document tokens 

870 for doc in context_data.get('documents', []): 

871 total_tokens += doc.get('tokens', 0) 

872 

873 # Add overhead (approximate) 

874 total_tokens += 1000 # Metadata and structure overhead 

875 

876 return total_tokens 

877 

878 def _record_metrics(self, metrics: ContextMetrics): 

879 """Record performance metrics""" 

880 self.metrics_history.append(metrics) 

881 

882 # Keep only recent metrics (last 100) 

883 if len(self.metrics_history) > 100: 

884 self.metrics_history.pop(0) 

885 

886 # Update performance stats 

887 self.performance_stats['total_loads'] += 1 

888 self.performance_stats['average_load_time'] = ( 

889 sum(m.load_time for m in self.metrics_history) / len(self.metrics_history) 

890 ) 

891 

892 cache_hits = sum(1 for m in self.metrics_history if m.cache_hit) 

893 self.performance_stats['cache_hit_rate'] = ( 

894 cache_hits / len(self.metrics_history) * 100 

895 ) 

896 

897 # Update efficiency score from token manager 

898 efficiency_metrics = self.token_manager.get_efficiency_metrics() 

899 self.performance_stats['efficiency_score'] = efficiency_metrics['efficiency_score'] 

900 

901 def get_comprehensive_stats(self) -> Dict[str, Any]: 

902 """Get comprehensive system statistics""" 

903 return { 

904 'performance': self.performance_stats, 

905 'cache': self.context_cache.get_stats(), 

906 'token_efficiency': self.token_manager.get_efficiency_metrics(), 

907 'skill_filter': self.skill_filter.get_skill_stats(), 

908 'current_phase': self.current_phase.value, 

909 'phase_history': self.phase_detector.phase_history[-5:], # Last 5 phase changes 

910 'metrics_count': len(self.metrics_history) 

911 } 

912 

913# Global instance for easy import 

914jit_context_loader = JITContextLoader() 

915 

916# Convenience functions 

917async def load_optimized_context(user_input: str, conversation_history: List[str] = None, 

918 context: Dict[str, Any] = None) -> Tuple[Dict[str, Any], ContextMetrics]: 

919 """Load optimized context using global JIT loader instance""" 

920 return await jit_context_loader.load_context(user_input, conversation_history, context) 

921 

922def get_jit_stats() -> Dict[str, Any]: 

923 """Get comprehensive JIT system statistics""" 

924 return jit_context_loader.get_comprehensive_stats() 

925 

926def clear_jit_cache(phase: Optional[str] = None): 

927 """Clear JIT cache (optionally for specific phase)""" 

928 if phase: 

929 jit_context_loader.context_cache.clear_phase(phase) 

930 else: 

931 jit_context_loader.context_cache.clear() 

932 

933# Initial setup 

934def initialize_jit_system(): 

935 """Initialize JIT Context Loading System""" 

936 logger.info("Initializing JIT Context Loading System...") 

937 

938 stats = get_jit_stats() 

939 logger.info(f"JIT System initialized with {stats['skill_filter']['total_skills']} skills") 

940 logger.info(f"Cache configured: {stats['cache']['entries']} entries, {stats['cache']['memory_usage_mb']:.1f}MB") 

941 

942 return True