@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 169 lines 5.27 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """An example demonstrates how to use tool calling with reasoning models
4 like QwQ-32B. The reasoning will not be parsed by the tool
5 calling process; only the final output will be parsed.
6
7 To run this example, you need to start the vLLM server with both
8 the reasoning parser and tool calling enabled.
9
10 ```bash
11 vllm serve Qwen/QwQ-32B \
12 --reasoning-parser deepseek_r1 \
13 --enable-auto-tool-choice --tool-call-parser hermes
14
15 ```
16
17 """
18
19 from openai import OpenAI
20
21
22 # Now, simulate a tool call
23 def get_current_weather(city: str, state: str, unit: "str"):
24 return (
25 "The weather in Dallas, Texas is 85 degrees fahrenheit. It is "
26 "partly cloudly, with highs in the 90's."
27 )
28
29
30 available_tools = {"get_current_weather": get_current_weather}
31
32 # Modify OpenAI's API key and API base to use vLLM's API server.
33 openai_api_key = "EMPTY"
34 openai_api_base = "http://localhost:8000/v1"
35
36 properties = {
37 "city": {
38 "type": "string",
39 "description": "The city to find the weather for, e.g. 'San Francisco'",
40 },
41 "state": {
42 "type": "string",
43 "description": "the two-letter abbreviation for the state that the city is"
44 " in, e.g. 'CA' which would mean 'California'",
45 },
46 "unit": {
47 "type": "string",
48 "description": "The unit to fetch the temperature in",
49 "enum": ["celsius", "fahrenheit"],
50 },
51 }
52
53 tools = [
54 {
55 "type": "function",
56 "function": {
57 "name": "get_current_weather",
58 "description": "Get the current weather in a given location",
59 "parameters": {
60 "type": "object",
61 "properties": properties,
62 "required": ["city", "state", "unit"],
63 },
64 },
65 }
66 ]
67 messages = [
68 {"role": "user", "content": "Hi! How are you doing today?"},
69 {"role": "assistant", "content": "I'm doing well! How can I help you?"},
70 {
71 "role": "user",
72 "content": (
73 "Can you tell me what the temperate will be in Dallas, in fahrenheit?"
74 ),
75 },
76 ]
77
78
79 def extract_reasoning_and_calls(chunks: list):
80 reasoning = ""
81 tool_call_idx = -1
82 arguments = []
83 function_names = []
84 for chunk in chunks:
85 if chunk.choices[0].delta.tool_calls:
86 tool_call = chunk.choices[0].delta.tool_calls[0]
87 if tool_call.index != tool_call_idx:
88 tool_call_idx = chunk.choices[0].delta.tool_calls[0].index
89 arguments.append("")
90 function_names.append("")
91
92 if tool_call.function:
93 if tool_call.function.name:
94 function_names[tool_call_idx] = tool_call.function.name
95
96 if tool_call.function.arguments:
97 arguments[tool_call_idx] += tool_call.function.arguments
98 else:
99 if hasattr(chunk.choices[0].delta, "reasoning"):
100 reasoning += chunk.choices[0].delta.reasoning
101 return reasoning, arguments, function_names
102
103
104 def main():
105 client = OpenAI(
106 api_key=openai_api_key,
107 base_url=openai_api_base,
108 )
109
110 models = client.models.list()
111 model = models.data[0].id
112
113 print("---------Full Generate With Automatic Function Calling-------------")
114 tool_calls = client.chat.completions.create(
115 messages=messages, model=model, tools=tools
116 )
117 print(f"reasoning: {tool_calls.choices[0].message.reasoning}")
118 print(f"function name: {tool_calls.choices[0].message.tool_calls[0].function.name}")
119 print(
120 f"function arguments: "
121 f"{tool_calls.choices[0].message.tool_calls[0].function.arguments}"
122 )
123
124 print("----------Stream Generate With Automatic Function Calling-----------")
125 tool_calls_stream = client.chat.completions.create(
126 messages=messages, model=model, tools=tools, stream=True
127 )
128
129 chunks = list(tool_calls_stream)
130
131 reasoning, arguments, function_names = extract_reasoning_and_calls(chunks)
132
133 print(f"reasoning: {reasoning}")
134 print(f"function name: {function_names[0]}")
135 print(f"function arguments: {arguments[0]}")
136
137 print("----------Full Generate With Named Function Calling-----------------")
138 tool_calls = client.chat.completions.create(
139 messages=messages,
140 model=model,
141 tools=tools,
142 tool_choice={"type": "function", "function": {"name": "get_current_weather"}},
143 )
144
145 tool_call = tool_calls.choices[0].message.tool_calls[0].function
146 print(f"reasoning: {tool_calls.choices[0].message.reasoning}")
147 print(f"function name: {tool_call.name}")
148 print(f"function arguments: {tool_call.arguments}")
149 print("----------Stream Generate With Named Function Calling--------------")
150
151 tool_calls_stream = client.chat.completions.create(
152 messages=messages,
153 model=model,
154 tools=tools,
155 tool_choice={"type": "function", "function": {"name": "get_current_weather"}},
156 stream=True,
157 )
158
159 chunks = list(tool_calls_stream)
160
161 reasoning, arguments, function_names = extract_reasoning_and_calls(chunks)
162 print(f"reasoning: {reasoning}")
163 print(f"function name: {function_names[0]}")
164 print(f"function arguments: {arguments[0]}")
165 print("\n\n")
166
167
168 if __name__ == "__main__":
169 main()