|
1 | 1 | import logging |
2 | 2 | from functools import wraps |
3 | | -from typing import Any, Callable, Mapping, Optional, Sequence, TypeVar |
| 3 | +from typing import Any, Callable, Optional, TypeVar |
4 | 4 | from typing_extensions import ParamSpec |
5 | 5 |
|
6 | 6 | from opentelemetry.trace import Span, Tracer |
@@ -67,56 +67,54 @@ def wrapper(*args: P.args, **kwargs: P.kwargs) -> Optional[R]: |
67 | 67 | # **log_inputs, |
68 | 68 | # log_status="incomplete", |
69 | 69 | # ) |
70 | | - token = set_trace_id(init_log["id"]) |
71 | | - |
72 | | - span.set_attribute(HUMANLOOP_PATH_KEY, decorator_path) |
73 | | - span.set_attribute(HUMANLOOP_FILE_TYPE_KEY, file_type) |
74 | | - |
75 | | - func_output: Optional[R] |
76 | | - log_output: str |
77 | | - log_error: Optional[str] |
78 | | - log_output_message: ChatMessage |
79 | | - try: |
80 | | - func_output = func(*args, **kwargs) |
81 | | - if ( |
82 | | - isinstance(func_output, dict) |
83 | | - and len(func_output.keys()) == 2 |
84 | | - and "role" in func_output |
85 | | - and "content" in func_output |
86 | | - ): |
87 | | - log_output_message = ChatMessage(**func_output) |
88 | | - log_output = None |
89 | | - else: |
90 | | - log_output = process_output(func=func, output=func_output) |
| 70 | + with set_trace_id(init_log["id"]): |
| 71 | + span.set_attribute(HUMANLOOP_PATH_KEY, decorator_path) |
| 72 | + span.set_attribute(HUMANLOOP_FILE_TYPE_KEY, file_type) |
| 73 | + |
| 74 | + func_output: Optional[R] |
| 75 | + log_output: str |
| 76 | + log_error: Optional[str] |
| 77 | + log_output_message: ChatMessage |
| 78 | + try: |
| 79 | + func_output = func(*args, **kwargs) |
| 80 | + if ( |
| 81 | + isinstance(func_output, dict) |
| 82 | + and len(func_output.keys()) == 2 |
| 83 | + and "role" in func_output |
| 84 | + and "content" in func_output |
| 85 | + ): |
| 86 | + log_output_message = ChatMessage(**func_output) |
| 87 | + log_output = None |
| 88 | + else: |
| 89 | + log_output = process_output(func=func, output=func_output) |
| 90 | + log_output_message = None |
| 91 | + log_error = None |
| 92 | + except Exception as e: |
| 93 | + logger.error(f"Error calling {func.__name__}: {e}") |
| 94 | + output = None |
91 | 95 | log_output_message = None |
92 | | - log_error = None |
93 | | - except Exception as e: |
94 | | - logger.error(f"Error calling {func.__name__}: {e}") |
95 | | - output = None |
96 | | - log_output_message = None |
97 | | - log_error = str(e) |
98 | | - |
99 | | - flow_log = { |
100 | | - "inputs": {k: v for k, v in args_to_func.items() if k != "messages"}, |
101 | | - "messages": args_to_func.get("messages"), |
102 | | - "log_status": "complete", |
103 | | - "output": log_output, |
104 | | - "error": log_error, |
105 | | - "output_message": log_output_message, |
106 | | - "id": init_log["id"], |
107 | | - } |
108 | | - |
109 | | - # Write the Flow Log to the Span on HL_LOG_OT_KEY |
110 | | - if flow_log: |
111 | | - write_to_opentelemetry_span( |
112 | | - span=span, |
113 | | - key=HUMANLOOP_LOG_KEY, |
114 | | - value=flow_log, # type: ignore |
115 | | - ) |
116 | | - |
117 | | - context_api.detach(token=token) |
118 | | - # Return the output of the decorated function |
119 | | - return output |
| 96 | + log_error = str(e) |
| 97 | + |
| 98 | + flow_log = { |
| 99 | + "inputs": {k: v for k, v in args_to_func.items() if k != "messages"}, |
| 100 | + "messages": args_to_func.get("messages"), |
| 101 | + "log_status": "complete", |
| 102 | + "output": log_output, |
| 103 | + "error": log_error, |
| 104 | + "output_message": log_output_message, |
| 105 | + "id": init_log["id"], |
| 106 | + } |
| 107 | + |
| 108 | + # Write the Flow Log to the Span on HL_LOG_OT_KEY |
| 109 | + if flow_log: |
| 110 | + write_to_opentelemetry_span( |
| 111 | + span=span, |
| 112 | + key=HUMANLOOP_LOG_KEY, |
| 113 | + value=flow_log, # type: ignore |
| 114 | + ) |
| 115 | + |
| 116 | + # Return the output of the decorated function |
| 117 | + return func_output |
120 | 118 |
|
121 | 119 | wrapper.file = File( # type: ignore |
122 | 120 | path=decorator_path, |
|
0 commit comments