Coverage for src/refinire/core/tracing.py: 71%
75 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"""
2Console tracing utilities for agents_sdk_models.
4English: Provides ConsoleTracingProcessor for color-coded output of span data and utility functions to enable/disable tracing.
5日本語: Spanデータの色分け出力を行う ConsoleTracingProcessor とトレーシングの有効化/無効化ユーティリティを提供します。
6"""
7from agents.tracing import TracingProcessor, set_trace_processors
8from agents.tracing.span_data import GenerationSpanData, ResponseSpanData
9from agents import set_tracing_disabled
10from .message import get_message, DEFAULT_LANGUAGE # Import for localized trace labels
11import sys
14def _merge_msgs(msgs, role):
15 """
16 English: Merge message contents by role from a list of message dicts.
17 日本語: メッセージのリストから指定したroleに一致するcontentを結合します。
18 """
19 return "\n".join(m.get("content", "") for m in (msgs or []) if m.get("role") == role)
22def extract_output_texts(obj):
23 """
24 English: Recursively extract all text contents from output message objects or dicts.
25 日本語: 出力メッセージのオブジェクトや辞書からtextフィールドを再帰的に抽出します。
26 """
27 results = []
28 if isinstance(obj, list):
29 for item in obj:
30 results.extend(extract_output_texts(item))
31 elif isinstance(obj, dict):
32 if "content" in obj and isinstance(obj["content"], list):
33 results.extend(extract_output_texts(obj["content"]))
34 elif "text" in obj and isinstance(obj["text"], str):
35 results.append(obj["text"])
36 elif "content" in obj and isinstance(obj["content"], str):
37 results.append(obj["content"])
38 else:
39 content = getattr(obj, "content", None)
40 if isinstance(content, list):
41 results.extend(extract_output_texts(content))
42 elif hasattr(obj, "text") and isinstance(obj.text, str):
43 results.append(obj.text)
44 elif isinstance(content, str):
45 results.append(content)
46 return results
49class ConsoleTracingProcessor(TracingProcessor):
50 """
51 English: A tracing processor that outputs Instruction, Prompt, and Output with colors to the console.
52 日本語: Instruction、Prompt、Outputを色分けしてコンソールに出力するトレーシングプロセッサ。
53 """
54 def __init__(self, output_stream=sys.stdout):
55 # English: Initialize with an output stream for logs.
56 # 日本語: ログ出力用の出力ストリームで初期化します。
57 self.output_stream = output_stream
59 def on_trace_start(self, trace):
60 # No-op for trace start
61 pass
63 def on_trace_end(self, trace):
64 # No-op for trace end
65 pass
67 def on_span_start(self, span):
68 # No-op for span start
69 pass
71 def on_span_end(self, span):
72 # Called at the end of each span; outputs color-coded Instruction/Prompt/Output and logs span end
73 if span is None:
74 return
75 data = span.span_data
76 instr = ""
77 prompt = ""
78 output = ""
79 if isinstance(data, GenerationSpanData):
80 instr = _merge_msgs(data.input, "system")
81 prompt = _merge_msgs(data.input, "user")
82 output = "\n".join(extract_output_texts(data.output))
83 elif isinstance(data, ResponseSpanData) and getattr(data, 'response', None):
84 instr = data.response.instructions or ""
85 prompt = _merge_msgs(data.input, "user")
86 output = "\n".join(extract_output_texts(data.response.output))
87 else:
88 # Irrelevant span type
89 return
90 # Color-coded output with localized labels
91 if self.output_stream is None:
92 return
93 if instr:
94 instr_label = get_message("trace_instruction", DEFAULT_LANGUAGE)
95 self.output_stream.write(f"\033[93m{instr_label} {instr}\033[0m\n")
96 if prompt:
97 prompt_label = get_message("trace_prompt", DEFAULT_LANGUAGE)
98 self.output_stream.write(f"\033[94m{prompt_label} {prompt}\033[0m\n")
99 if output:
100 output_label = get_message("trace_output", DEFAULT_LANGUAGE)
101 self.output_stream.write(f"\033[92m{output_label} {output}\033[0m\n")
102 self.output_stream.flush()
103 # # Log span end marker with error info
104 # info = span.export() or {}
105 # span_data = info.get('span_data') or span.span_data.export()
106 # name = span_data.get('name') if isinstance(span_data, dict) else None
107 # error = info.get('error', None)
108 # self.output_stream.write(f"[SPAN END] name={name}, error={error}\n")
109 # self.output_stream.flush()
111 def shutdown(self):
112 # No-op for shutdown
113 pass
115 def force_flush(self):
116 # Forces flush of the output stream
117 if hasattr(self, 'output_stream') and self.output_stream is not None:
118 self.output_stream.flush()
121def enable_console_tracing():
122 """
123 English: Enable console tracing by registering ConsoleTracingProcessor and enabling tracing.
124 日本語: ConsoleTracingProcessorを登録してトレーシングを有効化します。
125 """
126 # Enable tracing in Agents SDK
127 set_tracing_disabled(False)
128 # Register console tracing processor
129 set_trace_processors([ConsoleTracingProcessor()])
132def disable_tracing():
133 """
134 English: Disable all tracing.
135 日本語: トレーシング機能をすべて無効化します。
136 """
137 set_tracing_disabled(True)
140# Automatically enable console tracing when this module is imported
141enable_console_tracing()