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

1"""

2 Console tracing utilities for agents_sdk_models.

3

4 English: Provides ConsoleTracingProcessor for color-coded output of span data and utility functions to enable/disable tracing.

5日本語: Spanデータの色分け出力を行う ConsoleTracingProcessor とトレーシングの有効化/無効化ユーティリティを提供します。

6"""

7 from agents.tracing import TracingProcessor, set_trace_processors

8 from agents.tracing.span_data import GenerationSpanData, ResponseSpanData

9 from agents import set_tracing_disabled

10 from agents_sdk_models.message import get_message, DEFAULT_LANGUAGE # Import for localized trace labels

11 import sys

12

13

14 def _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

22 def 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

49 class 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

117 def 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

128 def 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

137 enable_console_tracing()