Coverage for langsmith/integrations/openai_agents_sdk/_openai_agent_utils.py: 18%

1import json 

2import logging 

3from typing import Any, 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, list): 

30 if len(data) == 0: 

31 return {} 

32 # Check if this is a list of output blocks (reasoning, message, etc.) 

33 if len(data) > 0 and isinstance(data[0], dict): 

34 if "type" in data[0]: 

35 return {default_key: data} 

36 elif len(data) == 1: 

37 return data[0] 

38 return {default_key: data} 

39 elif isinstance(data, dict): 

40 data_ = data 

41 elif isinstance(data, str): 

42 try: 

43 parsed_json = json.loads(data) 

44 if isinstance(parsed_json, dict): 

45 data_ = parsed_json 

46 else: 

47 data_ = {default_key: data} 

48 except json.JSONDecodeError: 

49 data_ = {default_key: data} 

50 elif ( 

51 data is not None 

52 and hasattr(data, "model_dump") 

53 and callable(data.model_dump) 

54 and not isinstance(data, type) 

55 ): 

56 try: 

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

58 except Exception as e: 

59 logger.debug( 

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

61 ) 

62 data_ = {default_key: data} 

63 else: 

64 data_ = {default_key: data} 

65 

66 return data_ 

67 

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

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

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

71 return "chain" 

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

73 return "tool" 

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

75 return "llm" 

76 else: 

77 return "chain" 

78 

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

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

81 return span.span_data.name 

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

83 if span_type == "generation": 

84 return "Generation" 

85 elif span_type == "response": 

86 return "Response" 

87 elif span_type == "handoff": 

88 return "Handoff" 

89 else: 

90 return "Span" 

91 

92 def _extract_function_span_data( 

93 span_data: tracing.FunctionSpanData, 

94 ) -> dict[str, Any]: 

95 return { 

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

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

98 } 

99 

100 def _extract_generation_span_data( 

101 span_data: tracing.GenerationSpanData, 

102 ) -> dict[str, Any]: 

103 data = { 

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

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

106 "invocation_params": { 

107 "model": span_data.model, 

108 "model_config": span_data.model_config, 

109 }, 

110 } 

111 if span_data.usage: 

112 from langsmith.wrappers._openai import _create_usage_metadata 

113 

114 if "metadata" not in data: 

115 data["metadata"] = {} 

116 data["metadata"]["usage_metadata"] = _create_usage_metadata(span_data.usage) 

117 return data 

118 

119 def _extract_response_span_data( 

120 span_data: tracing.ResponseSpanData, 

121 ) -> dict[str, Any]: 

122 data: dict[str, Any] = {} 

123 if span_data.input is not None: 

124 data["inputs"] = { 

125 "input": span_data.input, 

126 "instructions": ( 

127 span_data.response.instructions 

128 if span_data.response is not None 

129 and span_data.response.instructions 

130 else "" 

131 ), 

132 } 

133 if span_data.response is not None: 

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

135 output_data = response.pop("output", []) 

136 data["outputs"] = parse_io(output_data, "output") 

137 data["invocation_params"] = { 

138 k: v 

139 for k, v in response.items() 

140 if k 

141 in ( 

142 "max_output_tokens", 

143 "model", 

144 "parallel_tool_calls", 

145 "reasoning", 

146 "temperature", 

147 "text", 

148 "tool_choice", 

149 "tools", 

150 "top_p", 

151 "truncation", 

152 ) 

153 } 

154 metadata = { 

155 k: v 

156 for k, v in response.items() 

157 if k 

158 not in ( 

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

160 ) 

161 } 

162 metadata.update( 

163 { 

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

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

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

167 "ls_model_type": "chat", 

168 "ls_provider": "openai", 

169 } 

170 ) 

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

172 from langsmith.wrappers._openai import _create_usage_metadata 

173 

174 metadata["usage_metadata"] = _create_usage_metadata(usage) 

175 data["metadata"] = metadata 

176 

177 return data 

178 

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

180 return { 

181 "invocation_params": { 

182 "tools": span_data.tools, 

183 "handoffs": span_data.handoffs, 

184 }, 

185 "metadata": { 

186 "output_type": span_data.output_type, 

187 }, 

188 } 

189 

190 def _extract_handoff_span_data( 

191 span_data: tracing.HandoffSpanData, 

192 ) -> dict[str, Any]: 

193 return { 

194 "inputs": { 

195 "from_agent": span_data.from_agent, 

196 "to_agent": span_data.to_agent, 

197 } 

198 } 

199 

200 def _extract_guardrail_span_data( 

201 span_data: tracing.GuardrailSpanData, 

202 ) -> dict[str, Any]: 

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

204 

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

206 return {"metadata": span_data.data} 

207 

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

209 data: dict[str, Any] = {} 

210 

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

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

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

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

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

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

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

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

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

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

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

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

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

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

225 else: 

226 return {} 

227 

228 return data