@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 243 lines 7.65 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 # ruff: noqa: E501
4 """Set up this example by starting a vLLM OpenAI-compatible server with tool call
5 options enabled for xLAM-2 models:
6
7 vllm serve --model Salesforce/Llama-xLAM-2-8b-fc-r --enable-auto-tool-choice --tool-call-parser xlam
8
9 OR
10
11 vllm serve --model Salesforce/xLAM-2-3b-fc-r --enable-auto-tool-choice --tool-call-parser xlam
12 """
13
14 import json
15 import time
16
17 from openai import OpenAI
18
19 # Modify OpenAI's API key and API base to use vLLM's API server.
20 openai_api_key = "empty"
21 openai_api_base = "http://localhost:8000/v1"
22
23
24 # Define tool functions
25 def get_weather(location: str, unit: str):
26 return f"Weather in {location} is 22 degrees {unit}."
27
28
29 def calculate_expression(expression: str):
30 try:
31 result = eval(expression)
32 return f"The result of {expression} is {result}"
33 except Exception as e:
34 return f"Could not calculate {expression}: {e}"
35
36
37 def translate_text(text: str, target_language: str):
38 return f"Translation of '{text}' to {target_language}: [translated content]"
39
40
41 # Define tools
42 tools = [
43 {
44 "type": "function",
45 "function": {
46 "name": "get_weather",
47 "description": "Get the current weather in a given location",
48 "parameters": {
49 "type": "object",
50 "properties": {
51 "location": {
52 "type": "string",
53 "description": "City and state, e.g., 'San Francisco, CA'",
54 },
55 "unit": {"type": "string", "enum": ["celsius", "fahrenheit"]},
56 },
57 "required": ["location", "unit"],
58 },
59 },
60 },
61 {
62 "type": "function",
63 "function": {
64 "name": "calculate_expression",
65 "description": "Calculate a mathematical expression",
66 "parameters": {
67 "type": "object",
68 "properties": {
69 "expression": {
70 "type": "string",
71 "description": "Mathematical expression to evaluate, needs to be a valid python expression",
72 }
73 },
74 "required": ["expression"],
75 },
76 },
77 },
78 {
79 "type": "function",
80 "function": {
81 "name": "translate_text",
82 "description": "Translate text to another language",
83 "parameters": {
84 "type": "object",
85 "properties": {
86 "text": {"type": "string", "description": "Text to translate"},
87 "target_language": {
88 "type": "string",
89 "description": "Target language for translation",
90 },
91 },
92 "required": ["text", "target_language"],
93 },
94 },
95 },
96 ]
97
98 # Map of function names to implementations
99 tool_functions = {
100 "get_weather": get_weather,
101 "calculate_expression": calculate_expression,
102 "translate_text": translate_text,
103 }
104
105
106 def process_response(response, tool_functions, original_query):
107 """Process a non-streaming response with possible tool calls."""
108 print("\n--- Response Output ---")
109
110 # Check if the response has content
111 if response.choices[0].message.content:
112 print(f"Content: {response.choices[0].message.content}")
113
114 # Check if the response has tool calls
115 if response.choices[0].message.tool_calls:
116 print("--------------------------------")
117 print(f"Tool calls: {response.choices[0].message.tool_calls}")
118 print("--------------------------------")
119
120 # Collect all tool calls and results before making follow-up request
121 tool_results = []
122 assistant_message = {"role": "assistant"}
123
124 if response.choices[0].message.content:
125 assistant_message["content"] = response.choices[0].message.content
126
127 assistant_tool_calls = []
128
129 # Process each tool call
130 for tool_call in response.choices[0].message.tool_calls:
131 function_name = tool_call.function.name
132 function_args = tool_call.function.arguments
133 function_id = tool_call.id
134
135 print(f"Function called: {function_name}")
136 print(f"Arguments: {function_args}")
137 print(f"Function ID: {function_id}")
138
139 # Execute the function
140 try:
141 # Parse the JSON arguments
142 args = json.loads(function_args)
143
144 # Call the function with the arguments
145 function_result = tool_functions[function_name](**args)
146 print(f"\n--- Function Result ---\n{function_result}\n")
147
148 # Add tool call to assistant message
149 assistant_tool_calls.append(
150 {
151 "id": function_id,
152 "type": "function",
153 "function": {"name": function_name, "arguments": function_args},
154 }
155 )
156
157 # Add tool result to tool_results
158 tool_results.append(
159 {
160 "role": "tool",
161 "tool_call_id": function_id,
162 "content": function_result,
163 }
164 )
165
166 except Exception as e:
167 print(f"Error executing function: {e}")
168
169 # Add tool_calls to assistant message
170 assistant_message["tool_calls"] = assistant_tool_calls
171
172 # Create a follow-up message with all function results
173 follow_up_messages = [
174 {"role": "user", "content": original_query},
175 assistant_message,
176 ]
177
178 # Add all tool results to the messages
179 follow_up_messages.extend(tool_results)
180
181 # Get completion with all tool results in a single follow-up
182 follow_up_response = client.chat.completions.create(
183 model=client.models.list().data[0].id,
184 messages=follow_up_messages,
185 stream=False,
186 )
187
188 print("\n--- Follow-up Response ---")
189 print(follow_up_response.choices[0].message.content)
190 print("--- End Follow-up ---\n")
191
192 print("--- End Response ---\n")
193
194
195 def run_test_case(query, test_name):
196 """Run a single test case with the given query."""
197 print(f"\n{'=' * 50}\nTEST CASE: {test_name}\n{'=' * 50}")
198 print(f"Query: '{query}'")
199
200 start_time = time.time()
201
202 # Create non-streaming chat completion request
203 response = client.chat.completions.create(
204 model=client.models.list().data[0].id,
205 messages=[{"role": "user", "content": query}],
206 tools=tools,
207 tool_choice="auto",
208 stream=False,
209 )
210
211 # Process the non-streaming response, passing the original query
212 process_response(response, tool_functions, query)
213
214 end_time = time.time()
215 print(f"Test completed in {end_time - start_time:.2f} seconds")
216
217
218 def main():
219 # Initialize OpenAI client
220 global client
221 client = OpenAI(
222 api_key=openai_api_key,
223 base_url=openai_api_base,
224 )
225
226 # Run test cases
227 test_cases = [
228 ("I want to know the weather in San Francisco", "Weather Information"),
229 ("Calculate 25 * 17 + 31", "Math Calculation"),
230 ("Translate 'Hello world' to Spanish", "Text Translation"),
231 ("What is the weather in Tokyo and New York in celsius", "Multiple Tool Usage"),
232 ]
233
234 # Execute all test cases
235 for query, test_name in test_cases:
236 run_test_case(query, test_name)
237 time.sleep(1) # Small delay between tests
238
239 print("\nAll tests completed.")
240
241
242 if __name__ == "__main__":
243 main()