Coverage for src/refinire/core/trace_registry.py: 49%

203 statements  

« prev     ^ index     » next       coverage.py v7.9.1, created at 2025-06-15 18:51 +0900

1#!/usr/bin/env python3 

2""" 

3Trace Registry - Storage and search functionality for traces 

4トレースレジストリ - トレースの保存・検索機能 

5 

6Provides centralized trace management with search capabilities 

7集中化されたトレース管理と検索機能を提供 

8""" 

9 

10import json 

11from typing import Dict, List, Optional, Any, Union 

12from datetime import datetime, timedelta 

13from dataclasses import dataclass, asdict 

14from pathlib import Path 

15import threading 

16 

17 

18@dataclass 

19class TraceMetadata: 

20 """ 

21 Trace metadata for search and management 

22 検索・管理用のトレースメタデータ 

23 """ 

24 trace_id: str # Unique trace identifier / ユニークなトレース識別子 

25 flow_name: Optional[str] # Flow name / フロー名 

26 flow_id: Optional[str] # Flow instance ID / フローインスタンスID 

27 agent_names: List[str] # List of agent names used / 使用されたエージェント名のリスト 

28 start_time: datetime # Trace start time / トレース開始時刻 

29 end_time: Optional[datetime] # Trace end time / トレース終了時刻 

30 status: str # Trace status (running, completed, error) / トレースステータス 

31 total_spans: int # Number of spans in trace / トレース内のスパン数 

32 error_count: int # Number of error spans / エラースパン数 

33 duration_seconds: Optional[float] # Total duration / 総実行時間 

34 tags: Dict[str, Any] # Custom tags for filtering / フィルタリング用カスタムタグ 

35 artifacts: Dict[str, Any] # Trace artifacts / トレース成果物 

36 

37 

38class TraceRegistry: 

39 """ 

40 Registry for storing and searching traces 

41 トレースの保存・検索用レジストリ 

42  

43 Provides functionality to: 

44 以下の機能を提供: 

45 - Store trace metadata / トレースメタデータの保存 

46 - Search by flow name, agent name, tags / フロー名、エージェント名、タグによる検索 

47 - Filter by time range, status / 時間範囲、ステータスによるフィルタ 

48 - Export/import trace data / トレースデータのエクスポート/インポート 

49 """ 

50 

51 def __init__(self, storage_path: Optional[str] = None): 

52 """ 

53 Initialize trace registry 

54 トレースレジストリを初期化 

55  

56 Args: 

57 storage_path: Path to store trace data / トレースデータの保存パス 

58 """ 

59 self.traces: Dict[str, TraceMetadata] = {} 

60 self.storage_path = Path(storage_path) if storage_path else None 

61 self._lock = threading.Lock() 

62 

63 # Load existing traces if storage path exists 

64 # 保存パスが存在する場合、既存のトレースを読み込み 

65 if self.storage_path and self.storage_path.exists(): 

66 self.load_traces() 

67 

68 def register_trace( 

69 self, 

70 trace_id: str, 

71 flow_name: Optional[str] = None, 

72 flow_id: Optional[str] = None, 

73 agent_names: Optional[List[str]] = None, 

74 tags: Optional[Dict[str, Any]] = None 

75 ) -> None: 

76 """ 

77 Register a new trace 

78 新しいトレースを登録 

79  

80 Args: 

81 trace_id: Unique trace identifier / ユニークなトレース識別子 

82 flow_name: Flow name / フロー名 

83 flow_id: Flow instance ID / フローインスタンスID 

84 agent_names: List of agent names / エージェント名のリスト 

85 tags: Custom tags / カスタムタグ 

86 """ 

87 with self._lock: 

88 metadata = TraceMetadata( 

89 trace_id=trace_id, 

90 flow_name=flow_name, 

91 flow_id=flow_id, 

92 agent_names=agent_names or [], 

93 start_time=datetime.now(), 

94 end_time=None, 

95 status="running", 

96 total_spans=0, 

97 error_count=0, 

98 duration_seconds=None, 

99 tags=tags or {}, 

100 artifacts={} 

101 ) 

102 self.traces[trace_id] = metadata 

103 self._save_if_configured() 

104 

105 def update_trace( 

106 self, 

107 trace_id: str, 

108 status: Optional[str] = None, 

109 total_spans: Optional[int] = None, 

110 error_count: Optional[int] = None, 

111 artifacts: Optional[Dict[str, Any]] = None, 

112 add_agent_names: Optional[List[str]] = None, 

113 add_tags: Optional[Dict[str, Any]] = None 

114 ) -> None: 

115 """ 

116 Update trace metadata 

117 トレースメタデータを更新 

118  

119 Args: 

120 trace_id: Trace identifier / トレース識別子 

121 status: New status / 新しいステータス 

122 total_spans: Total span count / 総スパン数 

123 error_count: Error span count / エラースパン数 

124 artifacts: Trace artifacts / トレース成果物 

125 add_agent_names: Additional agent names / 追加エージェント名 

126 add_tags: Additional tags / 追加タグ 

127 """ 

128 with self._lock: 

129 if trace_id not in self.traces: 

130 return 

131 

132 trace = self.traces[trace_id] 

133 

134 if status: 

135 trace.status = status 

136 if status in ["completed", "error"]: 

137 trace.end_time = datetime.now() 

138 if trace.start_time: 

139 trace.duration_seconds = (trace.end_time - trace.start_time).total_seconds() 

140 

141 if total_spans is not None: 

142 trace.total_spans = total_spans 

143 

144 if error_count is not None: 

145 trace.error_count = error_count 

146 

147 if artifacts: 

148 trace.artifacts.update(artifacts) 

149 

150 if add_agent_names: 

151 trace.agent_names.extend(add_agent_names) 

152 trace.agent_names = list(set(trace.agent_names)) # Remove duplicates 

153 

154 if add_tags: 

155 trace.tags.update(add_tags) 

156 

157 self._save_if_configured() 

158 

159 def search_by_flow_name(self, flow_name: str, exact_match: bool = False) -> List[TraceMetadata]: 

160 """ 

161 Search traces by flow name 

162 フロー名でトレースを検索 

163  

164 Args: 

165 flow_name: Flow name to search / 検索するフロー名 

166 exact_match: Whether to use exact matching / 完全一致を使用するか 

167  

168 Returns: 

169 List[TraceMetadata]: Matching traces / マッチするトレース 

170 """ 

171 with self._lock: 

172 results = [] 

173 for trace in self.traces.values(): 

174 if trace.flow_name: 

175 if exact_match: 

176 if trace.flow_name == flow_name: 

177 results.append(trace) 

178 else: 

179 if flow_name.lower() in trace.flow_name.lower(): 

180 results.append(trace) 

181 return results 

182 

183 def search_by_agent_name(self, agent_name: str, exact_match: bool = False) -> List[TraceMetadata]: 

184 """ 

185 Search traces by agent name 

186 エージェント名でトレースを検索 

187  

188 Args: 

189 agent_name: Agent name to search / 検索するエージェント名 

190 exact_match: Whether to use exact matching / 完全一致を使用するか 

191  

192 Returns: 

193 List[TraceMetadata]: Matching traces / マッチするトレース 

194 """ 

195 with self._lock: 

196 results = [] 

197 for trace in self.traces.values(): 

198 for trace_agent in trace.agent_names: 

199 if exact_match: 

200 if trace_agent == agent_name: 

201 results.append(trace) 

202 break 

203 else: 

204 if agent_name.lower() in trace_agent.lower(): 

205 results.append(trace) 

206 break 

207 return results 

208 

209 def search_by_tags(self, tags: Dict[str, Any], match_all: bool = True) -> List[TraceMetadata]: 

210 """ 

211 Search traces by tags 

212 タグでトレースを検索 

213  

214 Args: 

215 tags: Tags to search for / 検索するタグ 

216 match_all: Whether all tags must match / すべてのタグがマッチする必要があるか 

217  

218 Returns: 

219 List[TraceMetadata]: Matching traces / マッチするトレース 

220 """ 

221 with self._lock: 

222 results = [] 

223 for trace in self.traces.values(): 

224 if match_all: 

225 # All search tags must be present and match 

226 # すべての検索タグが存在し、マッチする必要がある 

227 if all( 

228 key in trace.tags and trace.tags[key] == value 

229 for key, value in tags.items() 

230 ): 

231 results.append(trace) 

232 else: 

233 # At least one search tag must match 

234 # 少なくとも1つの検索タグがマッチする必要がある 

235 if any( 

236 key in trace.tags and trace.tags[key] == value 

237 for key, value in tags.items() 

238 ): 

239 results.append(trace) 

240 return results 

241 

242 def search_by_time_range( 

243 self, 

244 start_time: Optional[datetime] = None, 

245 end_time: Optional[datetime] = None 

246 ) -> List[TraceMetadata]: 

247 """ 

248 Search traces by time range 

249 時間範囲でトレースを検索 

250  

251 Args: 

252 start_time: Search from this time / この時刻から検索 

253 end_time: Search until this time / この時刻まで検索 

254  

255 Returns: 

256 List[TraceMetadata]: Matching traces / マッチするトレース 

257 """ 

258 with self._lock: 

259 results = [] 

260 for trace in self.traces.values(): 

261 # Check if trace start time is within range 

262 # トレース開始時刻が範囲内かチェック 

263 if start_time and trace.start_time < start_time: 

264 continue 

265 if end_time and trace.start_time > end_time: 

266 continue 

267 results.append(trace) 

268 return results 

269 

270 def search_by_status(self, status: str) -> List[TraceMetadata]: 

271 """ 

272 Search traces by status 

273 ステータスでトレースを検索 

274  

275 Args: 

276 status: Status to search for / 検索するステータス 

277  

278 Returns: 

279 List[TraceMetadata]: Matching traces / マッチするトレース 

280 """ 

281 with self._lock: 

282 return [trace for trace in self.traces.values() if trace.status == status] 

283 

284 def get_trace(self, trace_id: str) -> Optional[TraceMetadata]: 

285 """ 

286 Get specific trace by ID 

287 IDで特定のトレースを取得 

288  

289 Args: 

290 trace_id: Trace identifier / トレース識別子 

291  

292 Returns: 

293 TraceMetadata | None: Trace metadata if found / 見つかった場合のトレースメタデータ 

294 """ 

295 with self._lock: 

296 return self.traces.get(trace_id) 

297 

298 def get_all_traces(self) -> List[TraceMetadata]: 

299 """ 

300 Get all traces 

301 すべてのトレースを取得 

302  

303 Returns: 

304 List[TraceMetadata]: All traces / すべてのトレース 

305 """ 

306 with self._lock: 

307 return list(self.traces.values()) 

308 

309 def get_recent_traces(self, hours: int = 24) -> List[TraceMetadata]: 

310 """ 

311 Get recent traces within specified hours 

312 指定時間内の最近のトレースを取得 

313  

314 Args: 

315 hours: Number of hours to look back / 遡る時間数 

316  

317 Returns: 

318 List[TraceMetadata]: Recent traces / 最近のトレース 

319 """ 

320 cutoff_time = datetime.now() - timedelta(hours=hours) 

321 return self.search_by_time_range(start_time=cutoff_time) 

322 

323 def complex_search( 

324 self, 

325 flow_name: Optional[str] = None, 

326 agent_name: Optional[str] = None, 

327 tags: Optional[Dict[str, Any]] = None, 

328 status: Optional[str] = None, 

329 start_time: Optional[datetime] = None, 

330 end_time: Optional[datetime] = None, 

331 max_results: Optional[int] = None 

332 ) -> List[TraceMetadata]: 

333 """ 

334 Complex search with multiple criteria 

335 複数条件による複雑な検索 

336  

337 Args: 

338 flow_name: Flow name filter / フロー名フィルタ 

339 agent_name: Agent name filter / エージェント名フィルタ 

340 tags: Tags filter / タグフィルタ 

341 status: Status filter / ステータスフィルタ 

342 start_time: Start time filter / 開始時刻フィルタ 

343 end_time: End time filter / 終了時刻フィルタ 

344 max_results: Maximum number of results / 最大結果数 

345  

346 Returns: 

347 List[TraceMetadata]: Matching traces / マッチするトレース 

348 """ 

349 with self._lock: 

350 results = list(self.traces.values()) 

351 

352 # Apply filters 

353 # フィルタを適用 

354 if flow_name: 

355 results = [t for t in results if t.flow_name and flow_name.lower() in t.flow_name.lower()] 

356 

357 if agent_name: 

358 results = [ 

359 t for t in results 

360 if any(agent_name.lower() in agent.lower() for agent in t.agent_names) 

361 ] 

362 

363 if tags: 

364 results = [ 

365 t for t in results 

366 if all(key in t.tags and t.tags[key] == value for key, value in tags.items()) 

367 ] 

368 

369 if status: 

370 results = [t for t in results if t.status == status] 

371 

372 if start_time: 

373 results = [t for t in results if t.start_time >= start_time] 

374 

375 if end_time: 

376 results = [t for t in results if t.start_time <= end_time] 

377 

378 # Sort by start time (newest first) 

379 # 開始時刻でソート(新しい順) 

380 results.sort(key=lambda t: t.start_time, reverse=True) 

381 

382 # Limit results 

383 # 結果数を制限 

384 if max_results: 

385 results = results[:max_results] 

386 

387 return results 

388 

389 def get_statistics(self) -> Dict[str, Any]: 

390 """ 

391 Get trace statistics 

392 トレース統計を取得 

393  

394 Returns: 

395 Dict[str, Any]: Statistics / 統計情報 

396 """ 

397 with self._lock: 

398 total_traces = len(self.traces) 

399 if total_traces == 0: 

400 return {"total_traces": 0} 

401 

402 statuses = {} 

403 flow_names = set() 

404 agent_names = set() 

405 total_spans = 0 

406 total_errors = 0 

407 durations = [] 

408 

409 for trace in self.traces.values(): 

410 # Status distribution 

411 # ステータス分布 

412 statuses[trace.status] = statuses.get(trace.status, 0) + 1 

413 

414 # Collect names 

415 # 名前を収集 

416 if trace.flow_name: 

417 flow_names.add(trace.flow_name) 

418 agent_names.update(trace.agent_names) 

419 

420 # Aggregate metrics 

421 # メトリクスを集計 

422 total_spans += trace.total_spans 

423 total_errors += trace.error_count 

424 

425 if trace.duration_seconds is not None: 

426 durations.append(trace.duration_seconds) 

427 

428 avg_duration = sum(durations) / len(durations) if durations else 0 

429 

430 return { 

431 "total_traces": total_traces, 

432 "status_distribution": statuses, 

433 "unique_flow_names": len(flow_names), 

434 "unique_agent_names": len(agent_names), 

435 "total_spans": total_spans, 

436 "total_errors": total_errors, 

437 "average_duration_seconds": avg_duration, 

438 "flow_names": list(flow_names), 

439 "agent_names": list(agent_names) 

440 } 

441 

442 def export_traces(self, file_path: str, format: str = "json") -> None: 

443 """ 

444 Export traces to file 

445 トレースをファイルにエクスポート 

446  

447 Args: 

448 file_path: Output file path / 出力ファイルパス 

449 format: Export format (json, csv) / エクスポート形式 

450 """ 

451 with self._lock: 

452 if format == "json": 

453 data = { 

454 "export_time": datetime.now().isoformat(), 

455 "traces": [asdict(trace) for trace in self.traces.values()] 

456 } 

457 with open(file_path, 'w', encoding='utf-8') as f: 

458 json.dump(data, f, indent=2, default=str) 

459 else: 

460 raise ValueError(f"Unsupported export format: {format}") 

461 

462 def import_traces(self, file_path: str, format: str = "json") -> int: 

463 """ 

464 Import traces from file 

465 ファイルからトレースをインポート 

466  

467 Args: 

468 file_path: Input file path / 入力ファイルパス 

469 format: Import format (json) / インポート形式 

470  

471 Returns: 

472 int: Number of imported traces / インポートされたトレース数 

473 """ 

474 with self._lock: 

475 if format == "json": 

476 with open(file_path, 'r', encoding='utf-8') as f: 

477 data = json.load(f) 

478 

479 imported_count = 0 

480 for trace_data in data.get("traces", []): 

481 # Convert datetime strings back to datetime objects 

482 # 日時文字列をdatetimeオブジェクトに戻す 

483 if isinstance(trace_data.get("start_time"), str): 

484 trace_data["start_time"] = datetime.fromisoformat(trace_data["start_time"]) 

485 if isinstance(trace_data.get("end_time"), str): 

486 trace_data["end_time"] = datetime.fromisoformat(trace_data["end_time"]) 

487 

488 trace = TraceMetadata(**trace_data) 

489 self.traces[trace.trace_id] = trace 

490 imported_count += 1 

491 

492 self._save_if_configured() 

493 return imported_count 

494 else: 

495 raise ValueError(f"Unsupported import format: {format}") 

496 

497 def cleanup_old_traces(self, days: int = 30) -> int: 

498 """ 

499 Remove traces older than specified days 

500 指定日数より古いトレースを削除 

501  

502 Args: 

503 days: Number of days to keep / 保持する日数 

504  

505 Returns: 

506 int: Number of removed traces / 削除されたトレース数 

507 """ 

508 with self._lock: 

509 cutoff_time = datetime.now() - timedelta(days=days) 

510 old_trace_ids = [ 

511 trace_id for trace_id, trace in self.traces.items() 

512 if trace.start_time < cutoff_time 

513 ] 

514 

515 for trace_id in old_trace_ids: 

516 del self.traces[trace_id] 

517 

518 self._save_if_configured() 

519 return len(old_trace_ids) 

520 

521 def _save_if_configured(self) -> None: 

522 """ 

523 Save traces to storage if configured 

524 設定されている場合、トレースを保存 

525 """ 

526 if self.storage_path: 

527 self.save_traces() 

528 

529 def save_traces(self) -> None: 

530 """ 

531 Save traces to storage 

532 トレースをストレージに保存 

533 """ 

534 if not self.storage_path: 

535 return 

536 

537 self.storage_path.parent.mkdir(parents=True, exist_ok=True) 

538 self.export_traces(str(self.storage_path), "json") 

539 

540 def load_traces(self) -> int: 

541 """ 

542 Load traces from storage 

543 ストレージからトレースを読み込み 

544  

545 Returns: 

546 int: Number of loaded traces / 読み込まれたトレース数 

547 """ 

548 if not self.storage_path or not self.storage_path.exists(): 

549 return 0 

550 

551 return self.import_traces(str(self.storage_path), "json") 

552 

553 

554# Global trace registry instance 

555# グローバルトレースレジストリインスタンス 

556_global_registry: Optional[TraceRegistry] = None 

557 

558 

559def get_global_registry() -> TraceRegistry: 

560 """ 

561 Get global trace registry instance 

562 グローバルトレースレジストリインスタンスを取得 

563  

564 Returns: 

565 TraceRegistry: Global registry instance / グローバルレジストリインスタンス 

566 """ 

567 global _global_registry 

568 if _global_registry is None: 

569 _global_registry = TraceRegistry() 

570 return _global_registry 

571 

572 

573def set_global_registry(registry: TraceRegistry) -> None: 

574 """ 

575 Set global trace registry instance 

576 グローバルトレースレジストリインスタンスを設定 

577  

578 Args: 

579 registry: Registry instance to set as global / グローバルに設定するレジストリインスタンス 

580 """ 

581 global _global_registry 

582 _global_registry = registry