|
30 | 30 | from importlib_metadata import version as v |
31 | 31 |
|
32 | 32 | from langtrace_python_sdk.constants import LANGTRACE_SDK_NAME |
| 33 | +from langtrace.trace_attributes import SpanAttributes |
33 | 34 |
|
34 | 35 |
|
35 | 36 | def generic_patch( |
@@ -78,8 +79,22 @@ def traced_method(wrapped, instance, args, kwargs): |
78 | 79 | try: |
79 | 80 | # Attempt to call the original method |
80 | 81 | result = wrapped(*args, **kwargs) |
| 82 | + |
81 | 83 | if trace_output: |
82 | 84 | span.set_attribute("langchain.outputs", to_json_string(result)) |
| 85 | + if hasattr(result, 'usage'): |
| 86 | + prompt_tokens = result.usage.prompt_tokens |
| 87 | + completion_tokens = result.usage.completion_tokens |
| 88 | + span.set_attribute(SpanAttributes.LLM_USAGE_PROMPT_TOKENS, prompt_tokens) |
| 89 | + span.set_attribute(SpanAttributes.LLM_USAGE_COMPLETION_TOKENS, completion_tokens) |
| 90 | + |
| 91 | + elif result.generations[0][0].text: |
| 92 | + span.set_attribute(SpanAttributes.LLM_USAGE_COMPLETION_TOKENS, instance.get_num_tokens(result.generations[0][0].text)) |
| 93 | + elif isinstance(args[0][0], str): |
| 94 | + span.set_attribute(SpanAttributes.LLM_USAGE_PROMPT_TOKENS, instance.get_num_tokens(args[0][0])) |
| 95 | + |
| 96 | + else: |
| 97 | + span.set_attribute(SpanAttributes.LLM_USAGE_PROMPT_TOKENS, instance.get_num_tokens(args[0][0].text)) |
83 | 98 |
|
84 | 99 | span.set_status(StatusCode.OK) |
85 | 100 | return result |
@@ -156,6 +171,7 @@ def traced_method(wrapped, instance, args, kwargs): |
156 | 171 | try: |
157 | 172 | # Attempt to call the original method |
158 | 173 | result = wrapped(*args, **kwargs) |
| 174 | + |
159 | 175 | if trace_output: |
160 | 176 | outputs = {} |
161 | 177 | if isinstance(result, dict): |
|
0 commit comments