@@ -71,16 +71,26 @@ def traced_method(wrapped, instance, args, kwargs):
7171 result = wrapped (* args , ** kwargs )
7272 if trace_output :
7373 span .set_attribute ("langchain.outputs" , to_json_string (result ))
74-
75- prompt_tokens = instance .get_num_tokens (args [0 ])
76- completion_tokens = instance .get_num_tokens (result )
77- if hasattr (result , 'usage' ):
74+ prompt_tokens = (
75+ instance .get_num_tokens (args [0 ])
76+ if hasattr (instance , "get_num_tokens" )
77+ else None
78+ )
79+ completion_tokens = (
80+ instance .get_num_tokens (result )
81+ if hasattr (instance , "get_num_tokens" )
82+ else None
83+ )
84+ if hasattr (result , "usage" ):
7885 prompt_tokens = result .usage .prompt_tokens
7986 completion_tokens = result .usage .completion_tokens
8087
81- span .set_attribute (SpanAttributes .LLM_USAGE_COMPLETION_TOKENS , prompt_tokens )
82- span .set_attribute (SpanAttributes .LLM_USAGE_PROMPT_TOKENS , completion_tokens )
83-
88+ span .set_attribute (
89+ SpanAttributes .LLM_USAGE_COMPLETION_TOKENS , prompt_tokens
90+ )
91+ span .set_attribute (
92+ SpanAttributes .LLM_USAGE_PROMPT_TOKENS , completion_tokens
93+ )
8494
8595 span .set_status (StatusCode .OK )
8696 return result
@@ -102,9 +112,17 @@ def clean_empty(d):
102112 if not isinstance (d , (dict , list , tuple )):
103113 return d
104114 if isinstance (d , tuple ):
105- return tuple (val for val in (clean_empty (val ) for val in d ) if val != () and val is not None )
115+ return tuple (
116+ val
117+ for val in (clean_empty (val ) for val in d )
118+ if val != () and val is not None
119+ )
106120 if isinstance (d , list ):
107- return [val for val in (clean_empty (val ) for val in d ) if val != [] and val is not None ]
121+ return [
122+ val
123+ for val in (clean_empty (val ) for val in d )
124+ if val != [] and val is not None
125+ ]
108126 result = {}
109127 for k , val in d .items ():
110128 if isinstance (val , dict ):
@@ -120,7 +138,7 @@ def clean_empty(d):
120138 result [k ] = val .strip ()
121139 elif isinstance (val , object ):
122140 # some langchain objects have a text attribute
123- val = getattr (val , ' text' , None )
141+ val = getattr (val , " text" , None )
124142 if val is not None and val .strip () != "" :
125143 result [k ] = val .strip ()
126144 return result
0 commit comments