Coverage for integrations / agent_engine / upgrade_tools.py: 40.6%
64 statements
« prev ^ index » next coverage.py v7.14.0, created at 2026-05-12 04:49 +0000
1"""
2Unified Agent Goal Engine - Upgrade AutoGen Tools
410 tools for the upgrade goal type. Follows ip_protection_tools.py pattern.
5"""
8def check_upgrade_status() -> dict:
9 """Check current upgrade pipeline status and detect new versions."""
10 try:
11 from .upgrade_orchestrator import get_upgrade_orchestrator
12 orch = get_upgrade_orchestrator()
13 status = orch.get_status()
14 new_version = orch.check_for_new_version()
15 return {
16 'success': True,
17 'pipeline': status,
18 'new_version': new_version,
19 }
20 except Exception as e:
21 return {'success': False, 'error': str(e)}
24def capture_benchmark(version: str = '', git_sha: str = '',
25 tier: str = 'fast') -> dict:
26 """Capture a benchmark snapshot for the given version.
28 Args:
29 version: Version tag (auto-detected if empty)
30 git_sha: Git commit SHA
31 tier: 'fast' (upgrade gate) or 'heavy' (milestone) or 'all'
32 """
33 try:
34 from .benchmark_registry import get_benchmark_registry
35 registry = get_benchmark_registry()
36 if not version:
37 from .upgrade_orchestrator import get_upgrade_orchestrator
38 version = get_upgrade_orchestrator()._detect_version()
39 snapshot = registry.capture_snapshot(version, git_sha, tier=tier)
40 return {'success': True, 'snapshot': snapshot}
41 except Exception as e:
42 return {'success': False, 'error': str(e)}
45def compare_benchmarks(old_version: str, new_version: str) -> dict:
46 """Compare benchmarks between two versions. Reports regressions."""
47 try:
48 from .benchmark_registry import get_benchmark_registry
49 registry = get_benchmark_registry()
50 safe, reason = registry.is_upgrade_safe(old_version, new_version)
51 return {'success': True, 'safe': safe, 'reason': reason}
52 except Exception as e:
53 return {'success': False, 'error': str(e)}
56def start_upgrade(new_version: str, git_sha: str = '') -> dict:
57 """Start the 7-stage upgrade pipeline."""
58 try:
59 from .upgrade_orchestrator import get_upgrade_orchestrator
60 return get_upgrade_orchestrator().start_upgrade(new_version, git_sha)
61 except Exception as e:
62 return {'success': False, 'error': str(e)}
65def advance_upgrade_pipeline() -> dict:
66 """Execute the next stage of the upgrade pipeline."""
67 try:
68 from .upgrade_orchestrator import get_upgrade_orchestrator
69 return get_upgrade_orchestrator().advance_pipeline()
70 except Exception as e:
71 return {'success': False, 'error': str(e)}
74def check_canary_health() -> dict:
75 """Check health of canary nodes during deployment."""
76 try:
77 from .upgrade_orchestrator import get_upgrade_orchestrator
78 return {
79 'success': True,
80 **get_upgrade_orchestrator().check_canary_health_status(),
81 }
82 except Exception as e:
83 return {'success': False, 'error': str(e)}
86def rollback_upgrade(reason: str = '') -> dict:
87 """Safely rollback the upgrade at any stage."""
88 try:
89 from .upgrade_orchestrator import get_upgrade_orchestrator
90 return get_upgrade_orchestrator().rollback(reason)
91 except Exception as e:
92 return {'success': False, 'error': str(e)}
95def get_benchmark_history() -> dict:
96 """Get all stored benchmark snapshots."""
97 try:
98 import os
99 import json
100 from .benchmark_registry import BENCHMARK_DIR
101 snapshots = []
102 if os.path.isdir(BENCHMARK_DIR):
103 for fname in sorted(os.listdir(BENCHMARK_DIR)):
104 if fname.endswith('.json'):
105 fpath = os.path.join(BENCHMARK_DIR, fname)
106 try:
107 with open(fpath) as f:
108 data = json.load(f)
109 snapshots.append({
110 'version': data.get('version', fname.replace('.json', '')),
111 'timestamp': data.get('timestamp', 0),
112 'tier': data.get('tier', 'unknown'),
113 'benchmark_count': len(data.get('benchmarks', {})),
114 })
115 except Exception:
116 pass
117 return {'success': True, 'snapshots': snapshots}
118 except Exception as e:
119 return {'success': False, 'error': str(e)}
122def register_benchmark(name: str, repo_url: str,
123 requires_gpu: bool = False,
124 min_vram_gb: float = 0.0,
125 run_command: str = '',
126 metrics_file: str = '') -> dict:
127 """Coding agent dynamically installs and registers a benchmark from a git repo."""
128 try:
129 from .benchmark_registry import get_benchmark_registry
130 registry = get_benchmark_registry()
131 installed = registry.discover_and_install(
132 repo_url=repo_url, name=name,
133 requires_gpu=requires_gpu, min_vram_gb=min_vram_gb,
134 run_command=run_command, metrics_file=metrics_file)
135 return {
136 'success': installed,
137 'name': name,
138 'installed': installed,
139 }
140 except Exception as e:
141 return {'success': False, 'error': str(e)}
144def list_benchmarks() -> dict:
145 """List all registered benchmarks with availability status."""
146 try:
147 from .benchmark_registry import get_benchmark_registry
148 return {
149 'success': True,
150 'benchmarks': get_benchmark_registry().list_benchmarks(),
151 }
152 except Exception as e:
153 return {'success': False, 'error': str(e)}
156# Tool descriptors for AutoGen registration
157UPGRADE_TOOLS = [
158 {'name': 'check_upgrade_status',
159 'description': 'Check upgrade pipeline status and detect new versions.',
160 'function': check_upgrade_status},
161 {'name': 'capture_benchmark',
162 'description': 'Capture benchmark snapshot for a version.',
163 'function': capture_benchmark},
164 {'name': 'compare_benchmarks',
165 'description': 'Compare benchmarks between two versions.',
166 'function': compare_benchmarks},
167 {'name': 'start_upgrade',
168 'description': 'Start the 7-stage upgrade pipeline.',
169 'function': start_upgrade},
170 {'name': 'advance_upgrade_pipeline',
171 'description': 'Execute the next pipeline stage.',
172 'function': advance_upgrade_pipeline},
173 {'name': 'check_canary_health',
174 'description': 'Check canary node health during deployment.',
175 'function': check_canary_health},
176 {'name': 'rollback_upgrade',
177 'description': 'Safely rollback upgrade at any stage.',
178 'function': rollback_upgrade},
179 {'name': 'get_benchmark_history',
180 'description': 'Get all stored benchmark snapshots.',
181 'function': get_benchmark_history},
182 {'name': 'register_benchmark',
183 'description': 'Dynamically install a benchmark from a git repo.',
184 'function': register_benchmark},
185 {'name': 'list_benchmarks',
186 'description': 'List all registered benchmarks.',
187 'function': list_benchmarks},
188]