6
6
7
7
from deprecated import deprecated
8
8
from unstract .sdk .constants import Command , LogLevel , LogStage , ToolEnv
9
+ from unstract .sdk .exceptions import SdkError
9
10
from unstract .sdk .utils import Utils
10
11
from unstract .sdk .utils .common_utils import UNSTRACT_TO_PY_LOG_LEVEL
11
12
@@ -19,12 +20,14 @@ class StreamMixin:
19
20
to stdout.
20
21
"""
21
22
22
- def __init__ (self , log_level : LogLevel = LogLevel .INFO , ** kwargs ) -> None :
23
- """Args:
24
- log_level (LogLevel): The log level for filtering of log messages.
25
- The default is INFO.
26
- Allowed values are DEBUG, INFO, WARN, ERROR, and FATAL.
23
+ def __init__ (
24
+ self , log_level : LogLevel = LogLevel .INFO , ** kwargs : dict [str , Any ]
25
+ ) -> None :
26
+ """Constructor for StreamMixin.
27
27
28
+ Args:
29
+ log_level (LogLevel): The log level for filtering of log messages.
30
+ The default is INFO. Allowed values are DEBUG, INFO, WARN, ERROR, and FATAL.
28
31
"""
29
32
self .log_level = log_level
30
33
self ._exec_by_tool = Utils .str_to_bool (
@@ -35,7 +38,7 @@ def __init__(self, log_level: LogLevel = LogLevel.INFO, **kwargs) -> None:
35
38
super ().__init__ (** kwargs )
36
39
37
40
@property
38
- def is_exec_by_tool (self ):
41
+ def is_exec_by_tool (self ) -> bool :
39
42
"""Flag to determine if SDK library is used in a tool's context.
40
43
41
44
Returns:
@@ -51,20 +54,43 @@ def _configure_logger(self) -> None:
51
54
return
52
55
handler = logging .StreamHandler ()
53
56
handler .setLevel (level = UNSTRACT_TO_PY_LOG_LEVEL [self .log_level ])
57
+
58
+ # Determine if OpenTelemetry trace context should be included in logs
59
+ otel_trace_context = (
60
+ " trace_id:%(otelTraceID)s span_id:%(otelSpanID)s"
61
+ if os .environ .get ("OTEL_TRACES_EXPORTER" , "none" ).lower () != "none"
62
+ else ""
63
+ )
64
+
54
65
handler .setFormatter (
55
66
logging .Formatter (
56
- "[%(asctime)s] %(levelname)s in %(module)s: %(message)s" ,
67
+ "%(levelname)s : [%(asctime)s]"
68
+ "[pid:%(process)d tid:%(thread)d]" + otel_trace_context + " "
69
+ "%(name)s:- %(message)s"
57
70
)
58
71
)
59
72
rootlogger .addHandler (handler )
60
73
rootlogger .setLevel (level = UNSTRACT_TO_PY_LOG_LEVEL [self .log_level ])
61
74
75
+ noisy_lib_list = [
76
+ "asyncio" ,
77
+ "aiobotocore" ,
78
+ "boto3" ,
79
+ "botocore" ,
80
+ "fsspec" ,
81
+ "requests" ,
82
+ "s3fs" ,
83
+ "urllib3" ,
84
+ ]
85
+ for noisy_lib in noisy_lib_list :
86
+ logging .getLogger (noisy_lib ).setLevel (logging .WARNING )
87
+
62
88
def stream_log (
63
89
self ,
64
90
log : str ,
65
91
level : LogLevel = LogLevel .INFO ,
66
92
stage : str = LogStage .TOOL_RUN ,
67
- ** kwargs : Any ,
93
+ ** kwargs : dict [ str , Any ] ,
68
94
) -> None :
69
95
"""Streams a log message using the Unstract protocol LOG to stdout.
70
96
@@ -106,7 +132,7 @@ def stream_error_and_exit(self, message: str) -> None:
106
132
if self ._exec_by_tool :
107
133
exit (1 )
108
134
else :
109
- raise RuntimeError ( "RuntimeError from SDK, check the above log for details " )
135
+ raise SdkError ( f" SDK Error: { message } " )
110
136
111
137
def get_env_or_die (self , env_key : str ) -> str :
112
138
"""Returns the value of an env variable.
@@ -126,8 +152,7 @@ def get_env_or_die(self, env_key: str) -> str:
126
152
127
153
@staticmethod
128
154
def stream_spec (spec : str ) -> None :
129
- """Streams JSON schema of the tool using the Unstract protocol SPEC to
130
- stdout.
155
+ """Streams JSON schema of tool using Unstract protocol SPEC to stdout.
131
156
132
157
Args:
133
158
spec (str): The JSON schema of the tool.
@@ -145,8 +170,7 @@ def stream_spec(spec: str) -> None:
145
170
146
171
@staticmethod
147
172
def stream_properties (properties : str ) -> None :
148
- """Streams the properties of the tool using the Unstract protocol
149
- PROPERTIES to stdout.
173
+ """Streams tool properties JSON.
150
174
151
175
Args:
152
176
properties (str): The properties of the tool.
@@ -164,8 +188,10 @@ def stream_properties(properties: str) -> None:
164
188
165
189
@staticmethod
166
190
def stream_variables (variables : str ) -> None :
167
- """Streams JSON schema of the tool's variables using the Unstract
168
- protocol VARIABLES to stdout.
191
+ """Stream JSON variables.
192
+
193
+ Streams JSON schema of the tool's variables using the
194
+ Unstract protocol VARIABLES to stdout.
169
195
170
196
Args:
171
197
variables (str): The tool's runtime variables.
@@ -183,8 +209,7 @@ def stream_variables(variables: str) -> None:
183
209
184
210
@staticmethod
185
211
def stream_icon (icon : str ) -> None :
186
- """Streams the icon of the tool using the Unstract protocol ICON to
187
- stdout.
212
+ """Streams tool's icon JSON.
188
213
189
214
Args:
190
215
icon (str): The icon of the tool.
@@ -201,7 +226,7 @@ def stream_icon(icon: str) -> None:
201
226
print (json .dumps (record ))
202
227
203
228
@staticmethod
204
- def stream_update (message : str , state : str , ** kwargs : Any ) -> None :
229
+ def stream_update (message : str , state : str , ** kwargs : dict [ str , Any ] ) -> None :
205
230
"""Streams a log message using the Unstract protocol UPDATE to stdout.
206
231
207
232
Args:
@@ -219,9 +244,8 @@ def stream_update(message: str, state: str, **kwargs: Any) -> None:
219
244
220
245
@staticmethod
221
246
@deprecated (version = "0.4.4" , reason = "Unused in workflow execution" )
222
- def stream_cost (cost : float , cost_units : str , ** kwargs : Any ) -> None :
223
- """Streams the cost of the tool using the Unstract protocol COST to
224
- stdout.
247
+ def stream_cost (cost : float , cost_units : str , ** kwargs : dict [str , Any ]) -> None :
248
+ """Streams tool cost (deprecated).
225
249
226
250
Args:
227
251
cost (float): The cost of the tool.
@@ -242,9 +266,10 @@ def stream_cost(cost: float, cost_units: str, **kwargs: Any) -> None:
242
266
243
267
@staticmethod
244
268
@deprecated (version = "0.4.4" , reason = "Unused in workflow execution" )
245
- def stream_single_step_message (message : str , ** kwargs : Any ) -> None :
246
- """Streams a single step message using the Unstract protocol
247
- SINGLE_STEP_MESSAGE to stdout.
269
+ def stream_single_step_message (message : str , ** kwargs : dict [str , Any ]) -> None :
270
+ """Stream single step message.
271
+
272
+ Streams a single step message to stdout.
248
273
249
274
Args:
250
275
message (str): The single step message.
@@ -263,9 +288,8 @@ def stream_single_step_message(message: str, **kwargs: Any) -> None:
263
288
264
289
@staticmethod
265
290
@deprecated (version = "0.4.4" , reason = "Use `BaseTool.write_to_result()` instead" )
266
- def stream_result (result : dict [Any , Any ], ** kwargs : Any ) -> None :
267
- """Streams the result of the tool using the Unstract protocol RESULT to
268
- stdout.
291
+ def stream_result (result : dict [Any , Any ], ** kwargs : dict [str , Any ]) -> None :
292
+ """Streams tool result (review if required).
269
293
270
294
Args:
271
295
result (dict): The result of the tool. Refer to the
0 commit comments