test_vercel_ai_translator.py
1 import json 2 from unittest import mock 3 4 import pytest 5 6 from mlflow.entities.span import Span 7 from mlflow.tracing.constant import SpanAttributeKey 8 from mlflow.tracing.otel.translation import translate_span_when_storing 9 10 11 @pytest.mark.parametrize( 12 ("attributes", "expected_inputs", "expected_outputs"), 13 [ 14 # 1. generateText 15 ( 16 { 17 "ai.operationId": "ai.generateText", 18 "ai.prompt": '{"prompt":"Why is the sky blue?"}', 19 "ai.response.text": "Because of the scattering of light by the atmosphere.", 20 "ai.response.finishReason": "length", 21 }, 22 { 23 "prompt": "Why is the sky blue?", 24 }, 25 "Because of the scattering of light by the atmosphere.", 26 ), 27 # 2. generateText.doGenerate 28 ( 29 { 30 "ai.operationId": "ai.generateText.doGenerate", 31 "ai.prompt.messages": ( 32 '[{"role":"user","content":[{"type":"text","text":"Why is the sky blue?"}]}]' 33 ), 34 "ai.response.text": "Because of the scattering of light by the atmosphere.", 35 "ai.response.finishReason": "length", 36 "ai.response.id": "resp_0c4162a99c227acc00691324c9eaac81a3a3191fef81ca2987", 37 "ai.response.model": "gpt-4-turbo-2024-04-09", 38 }, 39 { 40 "messages": [ 41 {"role": "user", "content": [{"type": "text", "text": "Why is the sky blue?"}]} 42 ] 43 }, 44 { 45 "text": "Because of the scattering of light by the atmosphere.", 46 "finishReason": "length", 47 "id": "resp_0c4162a99c227acc00691324c9eaac81a3a3191fef81ca2987", 48 "model": "gpt-4-turbo-2024-04-09", 49 }, 50 ), 51 # 3. generateText with tool calls 52 ( 53 { 54 "ai.operationId": "ai.generateText.doGenerate", 55 "ai.prompt.messages": ( 56 '[{"role":"user","content":[{"type":"text","text":' 57 '"What is the weather in SF?"}]}]' 58 ), 59 "ai.prompt.tools": [ 60 ( 61 '{"type":"function","name":"weather","description":"Get the weather in ' 62 'a location","inputSchema":{"type":"object","properties":{"location":' 63 '{"type":"string","description":"The location to get the weather for"}},' 64 '"required":["location"],"additionalProperties":false,"$schema":' 65 '"http://json-schema.org/draft-07/schema#"}}' 66 ) 67 ], 68 "ai.prompt.toolChoice": '{"type":"auto"}', 69 "ai.response.toolCalls": ( 70 '[{"toolCallId":"call_PHKlxvzLK8w4PHH8CuvHXUzE","toolName":"weather",' 71 '"input":"{\\"location\\":\\"San Francisco\\"}"}]' 72 ), 73 "ai.response.finishReason": "tool-calls", 74 }, 75 { 76 "messages": [ 77 { 78 "role": "user", 79 "content": [{"type": "text", "text": "What is the weather in SF?"}], 80 } 81 ], 82 "tools": [ 83 { 84 "type": "function", 85 "name": "weather", 86 "description": "Get the weather in a location", 87 "inputSchema": { 88 "type": "object", 89 "properties": { 90 "location": { 91 "type": "string", 92 "description": "The location to get the weather for", 93 } 94 }, 95 "required": ["location"], 96 "additionalProperties": False, 97 "$schema": "http://json-schema.org/draft-07/schema#", 98 }, 99 } 100 ], 101 "toolChoice": {"type": "auto"}, 102 }, 103 { 104 "toolCalls": [ 105 { 106 "input": '{"location":"San Francisco"}', 107 "toolName": "weather", 108 "toolCallId": "call_PHKlxvzLK8w4PHH8CuvHXUzE", 109 } 110 ], 111 "finishReason": "tool-calls", 112 }, 113 ), 114 # 4. generateText with tool call results 115 ( 116 { 117 "ai.operationId": "ai.generateText.doGenerate", 118 "ai.prompt.messages": ( 119 '[{"role":"user","content":[{"type":"text",' 120 '"text":"What is the weather in San Francisco?"}]},' 121 '{"role":"assistant","content":[{"type":"tool-call","toolCallId":"call_123",' 122 '"toolName":"weather","input":{"location":"San Francisco"}}]},' 123 '{"role":"tool","content":[{"type":"tool-result","toolCallId":"call_123",' 124 '"toolName":"weather","output":{"type":"json",' 125 '"value":{"location":"San Francisco","temperature":76}}}]}]' 126 ), 127 "ai.prompt.toolChoice": '{"type":"auto"}', 128 "ai.response.text": "The current temperature in San Francisco is 76°F.", 129 "ai.response.finishReason": "stop", 130 }, 131 { 132 "messages": [ 133 { 134 "role": "user", 135 "content": [ 136 {"type": "text", "text": "What is the weather in San Francisco?"} 137 ], 138 }, 139 { 140 "role": "assistant", 141 "content": [ 142 { 143 "type": "tool-call", 144 "toolCallId": "call_123", 145 "toolName": "weather", 146 "input": {"location": "San Francisco"}, 147 } 148 ], 149 }, 150 { 151 "role": "tool", 152 "content": [ 153 { 154 "type": "tool-result", 155 "toolCallId": "call_123", 156 "toolName": "weather", 157 "output": { 158 "type": "json", 159 "value": {"location": "San Francisco", "temperature": 76}, 160 }, 161 } 162 ], 163 }, 164 ], 165 "toolChoice": {"type": "auto"}, 166 }, 167 { 168 "text": "The current temperature in San Francisco is 76°F.", 169 "finishReason": "stop", 170 }, 171 ), 172 # 5. Tool execution span 173 ( 174 { 175 "ai.operationId": "ai.toolCall", 176 "ai.toolCall.args": '{"location":"San Francisco"}', 177 "ai.toolCall.result": '{"location":"San Francisco","temperature":76}', 178 }, 179 { 180 "location": "San Francisco", 181 }, 182 { 183 "location": "San Francisco", 184 "temperature": 76, 185 }, 186 ), 187 ], 188 ) 189 def test_parse_vercel_ai_generate_text(attributes, expected_inputs, expected_outputs): 190 span = mock.Mock(spec=Span) 191 span.parent_id = "parent_123" 192 span_dict = {"attributes": {k: json.dumps(v) for k, v in attributes.items()}} 193 span.to_dict.return_value = span_dict 194 195 result = translate_span_when_storing(span) 196 inputs = json.loads(result["attributes"][SpanAttributeKey.INPUTS]) 197 assert inputs == expected_inputs 198 outputs = json.loads(result["attributes"][SpanAttributeKey.OUTPUTS]) 199 assert outputs == expected_outputs