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
« 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トレースレジストリ - トレースの保存・検索機能
6Provides centralized trace management with search capabilities
7集中化されたトレース管理と検索機能を提供
8"""
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
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 / トレース成果物
38class TraceRegistry:
39 """
40 Registry for storing and searching traces
41 トレースの保存・検索用レジストリ
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 """
51 def __init__(self, storage_path: Optional[str] = None):
52 """
53 Initialize trace registry
54 トレースレジストリを初期化
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()
63 # Load existing traces if storage path exists
64 # 保存パスが存在する場合、既存のトレースを読み込み
65 if self.storage_path and self.storage_path.exists():
66 self.load_traces()
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 新しいトレースを登録
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()
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 トレースメタデータを更新
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
132 trace = self.traces[trace_id]
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()
141 if total_spans is not None:
142 trace.total_spans = total_spans
144 if error_count is not None:
145 trace.error_count = error_count
147 if artifacts:
148 trace.artifacts.update(artifacts)
150 if add_agent_names:
151 trace.agent_names.extend(add_agent_names)
152 trace.agent_names = list(set(trace.agent_names)) # Remove duplicates
154 if add_tags:
155 trace.tags.update(add_tags)
157 self._save_if_configured()
159 def search_by_flow_name(self, flow_name: str, exact_match: bool = False) -> List[TraceMetadata]:
160 """
161 Search traces by flow name
162 フロー名でトレースを検索
164 Args:
165 flow_name: Flow name to search / 検索するフロー名
166 exact_match: Whether to use exact matching / 完全一致を使用するか
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
183 def search_by_agent_name(self, agent_name: str, exact_match: bool = False) -> List[TraceMetadata]:
184 """
185 Search traces by agent name
186 エージェント名でトレースを検索
188 Args:
189 agent_name: Agent name to search / 検索するエージェント名
190 exact_match: Whether to use exact matching / 完全一致を使用するか
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
209 def search_by_tags(self, tags: Dict[str, Any], match_all: bool = True) -> List[TraceMetadata]:
210 """
211 Search traces by tags
212 タグでトレースを検索
214 Args:
215 tags: Tags to search for / 検索するタグ
216 match_all: Whether all tags must match / すべてのタグがマッチする必要があるか
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
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 時間範囲でトレースを検索
251 Args:
252 start_time: Search from this time / この時刻から検索
253 end_time: Search until this time / この時刻まで検索
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
270 def search_by_status(self, status: str) -> List[TraceMetadata]:
271 """
272 Search traces by status
273 ステータスでトレースを検索
275 Args:
276 status: Status to search for / 検索するステータス
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]
284 def get_trace(self, trace_id: str) -> Optional[TraceMetadata]:
285 """
286 Get specific trace by ID
287 IDで特定のトレースを取得
289 Args:
290 trace_id: Trace identifier / トレース識別子
292 Returns:
293 TraceMetadata | None: Trace metadata if found / 見つかった場合のトレースメタデータ
294 """
295 with self._lock:
296 return self.traces.get(trace_id)
298 def get_all_traces(self) -> List[TraceMetadata]:
299 """
300 Get all traces
301 すべてのトレースを取得
303 Returns:
304 List[TraceMetadata]: All traces / すべてのトレース
305 """
306 with self._lock:
307 return list(self.traces.values())
309 def get_recent_traces(self, hours: int = 24) -> List[TraceMetadata]:
310 """
311 Get recent traces within specified hours
312 指定時間内の最近のトレースを取得
314 Args:
315 hours: Number of hours to look back / 遡る時間数
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)
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 複数条件による複雑な検索
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 / 最大結果数
346 Returns:
347 List[TraceMetadata]: Matching traces / マッチするトレース
348 """
349 with self._lock:
350 results = list(self.traces.values())
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()]
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 ]
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 ]
369 if status:
370 results = [t for t in results if t.status == status]
372 if start_time:
373 results = [t for t in results if t.start_time >= start_time]
375 if end_time:
376 results = [t for t in results if t.start_time <= end_time]
378 # Sort by start time (newest first)
379 # 開始時刻でソート(新しい順)
380 results.sort(key=lambda t: t.start_time, reverse=True)
382 # Limit results
383 # 結果数を制限
384 if max_results:
385 results = results[:max_results]
387 return results
389 def get_statistics(self) -> Dict[str, Any]:
390 """
391 Get trace statistics
392 トレース統計を取得
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}
402 statuses = {}
403 flow_names = set()
404 agent_names = set()
405 total_spans = 0
406 total_errors = 0
407 durations = []
409 for trace in self.traces.values():
410 # Status distribution
411 # ステータス分布
412 statuses[trace.status] = statuses.get(trace.status, 0) + 1
414 # Collect names
415 # 名前を収集
416 if trace.flow_name:
417 flow_names.add(trace.flow_name)
418 agent_names.update(trace.agent_names)
420 # Aggregate metrics
421 # メトリクスを集計
422 total_spans += trace.total_spans
423 total_errors += trace.error_count
425 if trace.duration_seconds is not None:
426 durations.append(trace.duration_seconds)
428 avg_duration = sum(durations) / len(durations) if durations else 0
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 }
442 def export_traces(self, file_path: str, format: str = "json") -> None:
443 """
444 Export traces to file
445 トレースをファイルにエクスポート
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}")
462 def import_traces(self, file_path: str, format: str = "json") -> int:
463 """
464 Import traces from file
465 ファイルからトレースをインポート
467 Args:
468 file_path: Input file path / 入力ファイルパス
469 format: Import format (json) / インポート形式
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)
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"])
488 trace = TraceMetadata(**trace_data)
489 self.traces[trace.trace_id] = trace
490 imported_count += 1
492 self._save_if_configured()
493 return imported_count
494 else:
495 raise ValueError(f"Unsupported import format: {format}")
497 def cleanup_old_traces(self, days: int = 30) -> int:
498 """
499 Remove traces older than specified days
500 指定日数より古いトレースを削除
502 Args:
503 days: Number of days to keep / 保持する日数
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 ]
515 for trace_id in old_trace_ids:
516 del self.traces[trace_id]
518 self._save_if_configured()
519 return len(old_trace_ids)
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()
529 def save_traces(self) -> None:
530 """
531 Save traces to storage
532 トレースをストレージに保存
533 """
534 if not self.storage_path:
535 return
537 self.storage_path.parent.mkdir(parents=True, exist_ok=True)
538 self.export_traces(str(self.storage_path), "json")
540 def load_traces(self) -> int:
541 """
542 Load traces from storage
543 ストレージからトレースを読み込み
545 Returns:
546 int: Number of loaded traces / 読み込まれたトレース数
547 """
548 if not self.storage_path or not self.storage_path.exists():
549 return 0
551 return self.import_traces(str(self.storage_path), "json")
554# Global trace registry instance
555# グローバルトレースレジストリインスタンス
556_global_registry: Optional[TraceRegistry] = None
559def get_global_registry() -> TraceRegistry:
560 """
561 Get global trace registry instance
562 グローバルトレースレジストリインスタンスを取得
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
573def set_global_registry(registry: TraceRegistry) -> None:
574 """
575 Set global trace registry instance
576 グローバルトレースレジストリインスタンスを設定
578 Args:
579 registry: Registry instance to set as global / グローバルに設定するレジストリインスタンス
580 """
581 global _global_registry
582 _global_registry = registry