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

1""" 

2Console tracing utilities for agents_sdk_models. 

3 

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 

12 

13 

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) 

20 

21 

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 

47 

48 

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 

58 

59 def on_trace_start(self, trace): 

60 # No-op for trace start 

61 pass 

62 

63 def on_trace_end(self, trace): 

64 # No-op for trace end 

65 pass 

66 

67 def on_span_start(self, span): 

68 # No-op for span start 

69 pass 

70 

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() 

110 

111 def shutdown(self): 

112 # No-op for shutdown 

113 pass 

114 

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() 

119 

120 

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()]) 

130 

131 

132def disable_tracing(): 

133 """ 

134 English: Disable all tracing. 

135 日本語: トレーシング機能をすべて無効化します。 

136 """ 

137 set_tracing_disabled(True) 

138 

139 

140# Automatically enable console tracing when this module is imported 

141enable_console_tracing()