Coverage for src\agents_sdk_models\tracing.py: 72%

71 statements  

« prev     ^ index     » next       coverage.py v7.8.0, created at 2025-06-04 17:38 +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 agents_sdk_models.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 data = span.span_data 

74 instr = "" 

75 prompt = "" 

76 output = "" 

77 if isinstance(data, GenerationSpanData): 

78 instr = _merge_msgs(data.input, "system") 

79 prompt = _merge_msgs(data.input, "user") 

80 output = "\n".join(extract_output_texts(data.output)) 

81 elif isinstance(data, ResponseSpanData) and getattr(data, 'response', None): 

82 instr = data.response.instructions or "" 

83 prompt = _merge_msgs(data.input, "user") 

84 output = "\n".join(extract_output_texts(data.response.output)) 

85 else: 

86 # Irrelevant span type 

87 return 

88 # Color-coded output with localized labels 

89 if instr: 

90 instr_label = get_message("trace_instruction", DEFAULT_LANGUAGE) 

91 self.output_stream.write(f"\033[93m{instr_label} {instr}\033[0m\n") 

92 if prompt: 

93 prompt_label = get_message("trace_prompt", DEFAULT_LANGUAGE) 

94 self.output_stream.write(f"\033[94m{prompt_label} {prompt}\033[0m\n") 

95 if output: 

96 output_label = get_message("trace_output", DEFAULT_LANGUAGE) 

97 self.output_stream.write(f"\033[92m{output_label} {output}\033[0m\n") 

98 self.output_stream.flush() 

99 # # Log span end marker with error info 

100 # info = span.export() or {} 

101 # span_data = info.get('span_data') or span.span_data.export() 

102 # name = span_data.get('name') if isinstance(span_data, dict) else None 

103 # error = info.get('error', None) 

104 # self.output_stream.write(f"[SPAN END] name={name}, error={error}\n") 

105 # self.output_stream.flush() 

106 

107 def shutdown(self): 

108 # No-op for shutdown 

109 pass 

110 

111 def force_flush(self): 

112 # Forces flush of the output stream 

113 if hasattr(self, 'output_stream'): 

114 self.output_stream.flush() 

115 

116 

117def enable_console_tracing(): 

118 """ 

119 English: Enable console tracing by registering ConsoleTracingProcessor and enabling tracing. 

120 日本語: ConsoleTracingProcessorを登録してトレーシングを有効化します。 

121 """ 

122 # Enable tracing in Agents SDK 

123 set_tracing_disabled(False) 

124 # Register console tracing processor 

125 set_trace_processors([ConsoleTracingProcessor()]) 

126 

127 

128def disable_tracing(): 

129 """ 

130 English: Disable all tracing. 

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

132 """ 

133 set_tracing_disabled(True) 

134 

135 

136# Automatically enable console tracing when this module is imported 

137enable_console_tracing()