Coverage for langsmith/wrappers/_agent_utils.py: 0%

1import json 

2import logging 

3from typing import Any, Dict, Literal 

4 

5try: 

6 from agents import tracing # type: ignore[import] 

7 

8 HAVE_AGENTS = True 

9except ImportError: 

10 HAVE_AGENTS = False 

11 

12logger = logging.getLogger(__name__) 

13 

14RunTypeT = Literal["tool", "chain", "llm", "retriever", "embedding", "prompt", "parser"] 

15 

16if HAVE_AGENTS: 

17 

18 def parse_io(data: Any, default_key: str = "output") -> Dict: 

19 """Parse inputs or outputs into a dictionary format. 

20 

21 Args: 

22 data: The data to parse (can be inputs or outputs) 

23 default_key: The default key to use if data is not a dict 

24 ("input" or "output") 

25 

26 Returns: 

27 Dict: The parsed data as a dictionary 

28 """ 

29 if isinstance(data, dict): 

30 data_ = data 

31 elif isinstance(data, str): 

32 try: 

33 parsed_json = json.loads(data) 

34 if isinstance(parsed_json, dict): 

35 data_ = parsed_json 

36 else: 

37 data_ = {default_key: data} 

38 except json.JSONDecodeError: 

39 data_ = {default_key: data} 

40 elif ( 

41 data is not None 

42 and hasattr(data, "model_dump") 

43 and callable(data.model_dump) 

44 and not isinstance(data, type) 

45 ): 

46 try: 

47 data_ = data.model_dump(exclude_none=True, mode="json") 

48 except Exception as e: 

49 logger.debug( 

50 f"Failed to use model_dump to serialize {type(data)} to JSON: {e}" 

51 ) 

52 data_ = {default_key: data} 

53 else: 

54 data_ = {default_key: data} 

55 

56 return data_ 

57 

58 def get_run_type(span: tracing.Span) -> RunTypeT: 

59 span_type = getattr(span.span_data, "type", None) 

60 if span_type in ["agent", "handoff", "custom"]: 

61 return "chain" 

62 elif span_type in ["function", "guardrail"]: 

63 return "tool" 

64 elif span_type in ["generation", "response"]: 

65 return "llm" 

66 else: 

67 return "chain" 

68 

69 def get_run_name(span: tracing.Span) -> str: 

70 if hasattr(span.span_data, "name") and span.span_data.name: 

71 return span.span_data.name 

72 span_type = getattr(span.span_data, "type", None) 

73 if span_type == "generation": 

74 return "Generation" 

75 elif span_type == "response": 

76 return "Response" 

77 elif span_type == "handoff": 

78 return "Handoff" 

79 else: 

80 return "Span" 

81 

82 def _extract_function_span_data( 

83 span_data: tracing.FunctionSpanData, 

84 ) -> Dict[str, Any]: 

85 return { 

86 "inputs": parse_io(span_data.input, "input"), 

87 "outputs": parse_io(span_data.output, "output"), 

88 } 

89 

90 def _extract_generation_span_data( 

91 span_data: tracing.GenerationSpanData, 

92 ) -> Dict[str, Any]: 

93 data = { 

94 "inputs": parse_io(span_data.input, "input"), 

95 "outputs": parse_io(span_data.output, "output"), 

96 "invocation_params": { 

97 "model": span_data.model, 

98 "model_config": span_data.model_config, 

99 }, 

100 } 

101 if span_data.usage: 

102 data["outputs"]["usage_metadata"] = { 

103 "total_tokens": span_data.usage.get("total_tokens"), 

104 "input_tokens": span_data.usage.get("prompt_tokens"), 

105 "output_tokens": span_data.usage.get("completion_tokens"), 

106 } 

107 return data 

108 

109 def _extract_response_span_data( 

110 span_data: tracing.ResponseSpanData, 

111 ) -> Dict[str, Any]: 

112 data: Dict[str, Any] = {} 

113 if span_data.input is not None: 

114 data["inputs"] = { 

115 "input": span_data.input, 

116 "instructions": span_data.response.instructions, 

117 } 

118 if span_data.response is not None: 

119 response = span_data.response.model_dump(exclude_none=True, mode="json") 

120 data["outputs"] = {"output": response.pop("output", [])} 

121 if usage := response.pop("usage", None): 

122 # tokens -> token 

123 if "output_tokens_details" in usage: 

124 usage["output_token_details"] = usage.pop("output_tokens_details") 

125 usage["output_token_details"]["reasoning"] = usage[ 

126 "output_token_details" 

127 ].pop("reasoning_tokens", 0) 

128 if "input_tokens_details" in usage: 

129 usage["input_token_details"] = usage.pop("input_tokens_details") 

130 usage["input_token_details"]["cache_read"] = usage[ 

131 "input_token_details" 

132 ].pop("cached_tokens", 0) 

133 data["outputs"]["usage_metadata"] = usage 

134 

135 data["invocation_params"] = { 

136 k: v 

137 for k, v in response.items() 

138 if k 

139 in ( 

140 "max_output_tokens", 

141 "model", 

142 "parallel_tool_calls", 

143 "reasoning", 

144 "temperature", 

145 "text", 

146 "tool_choice", 

147 "tools", 

148 "top_p", 

149 "truncation", 

150 ) 

151 } 

152 metadata = { 

153 k: v 

154 for k, v in response.items() 

155 if k 

156 not in ( 

157 {"output", "usage", "instructions"}.union(data["invocation_params"]) 

158 ) 

159 } 

160 metadata.update( 

161 { 

162 "ls_model_name": data["invocation_params"].get("model"), 

163 "ls_max_tokens": data["invocation_params"].get("max_output_tokens"), 

164 "ls_temperature": data["invocation_params"].get("temperature"), 

165 "ls_model_type": "chat", 

166 "ls_provider": "openai", 

167 } 

168 ) 

169 data["metadata"] = metadata 

170 

171 return data 

172 

173 def _extract_agent_span_data(span_data: tracing.AgentSpanData) -> Dict[str, Any]: 

174 return { 

175 "invocation_params": { 

176 "tools": span_data.tools, 

177 "handoffs": span_data.handoffs, 

178 }, 

179 "metadata": { 

180 "output_type": span_data.output_type, 

181 }, 

182 } 

183 

184 def _extract_handoff_span_data( 

185 span_data: tracing.HandoffSpanData, 

186 ) -> Dict[str, Any]: 

187 return { 

188 "inputs": { 

189 "from_agent": span_data.from_agent, 

190 "to_agent": span_data.to_agent, 

191 } 

192 } 

193 

194 def _extract_guardrail_span_data( 

195 span_data: tracing.GuardrailSpanData, 

196 ) -> Dict[str, Any]: 

197 return {"metadata": {"triggered": span_data.triggered}} 

198 

199 def _extract_custom_span_data(span_data: tracing.CustomSpanData) -> Dict[str, Any]: 

200 return {"metadata": span_data.data} 

201 

202 def extract_span_data(span: tracing.Span) -> Dict[str, Any]: 

203 data: Dict[str, Any] = {} 

204 

205 if isinstance(span.span_data, tracing.FunctionSpanData): 

206 data.update(_extract_function_span_data(span.span_data)) 

207 elif isinstance(span.span_data, tracing.GenerationSpanData): 

208 data.update(_extract_generation_span_data(span.span_data)) 

209 elif isinstance(span.span_data, tracing.ResponseSpanData): 

210 data.update(_extract_response_span_data(span.span_data)) 

211 elif isinstance(span.span_data, tracing.AgentSpanData): 

212 data.update(_extract_agent_span_data(span.span_data)) 

213 elif isinstance(span.span_data, tracing.HandoffSpanData): 

214 data.update(_extract_handoff_span_data(span.span_data)) 

215 elif isinstance(span.span_data, tracing.GuardrailSpanData): 

216 data.update(_extract_guardrail_span_data(span.span_data)) 

217 elif isinstance(span.span_data, tracing.CustomSpanData): 

218 data.update(_extract_custom_span_data(span.span_data)) 

219 else: 

220 return {} 

221 

222 return data