A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 13:30 UTC
[Misc] refactor examples series - Chat Completion Client With Tools (#16829)
Signed-off-by: reidliu41 <reid201711@gmail.com> Co-authored-by: reidliu41 <reid201711@gmail.com>
Reid committed
Apr 19, 2025 at 07:24 UTC
5a5e29de88826a1d3d7aa4f1d621067401999e98
1 file changed
+112
-83
examples/online_serving/openai_chat_completion_client_with_tools.py
+112
-83
@@ -17,6 +17,7 @@ vllm serve --model NousResearch/Hermes-2-Pro-Llama-3-8B \
17
--enable-auto-tool-choice --tool-call-parser hermes
18
"""
19
import json
20
+from typing import Any
21
22
from openai import OpenAI
23
@@ -24,15 +25,6 @@ from openai import OpenAI
25
openai_api_key = "EMPTY"
26
openai_api_base = "http://localhost:8000/v1"
27
27
-client = OpenAI(
28
- # defaults to os.environ.get("OPENAI_API_KEY")
29
- api_key=openai_api_key,
30
- base_url=openai_api_base,
31
-)
32
-
33
-models = client.models.list()
34
-model = models.data[0].id
35
-
28
tools = [{
29
"type": "function",
30
"function": {
@@ -78,86 +70,123 @@ messages = [{
70
"Can you tell me what the temperate will be in Dallas, in fahrenheit?"
71
}]
72
81
-chat_completion = client.chat.completions.create(messages=messages,
82
- model=model,
83
- tools=tools)
84
-
85
-print("Chat completion results:")
86
-print(chat_completion)
87
-print("\n\n")
88
-
89
-tool_calls_stream = client.chat.completions.create(messages=messages,
90
- model=model,
91
- tools=tools,
92
- stream=True)
93
-
94
-chunks = []
95
-for chunk in tool_calls_stream:
96
- chunks.append(chunk)
97
- if chunk.choices[0].delta.tool_calls:
98
- print(chunk.choices[0].delta.tool_calls[0])
99
- else:
100
- print(chunk.choices[0].delta)
101
-
102
-arguments = []
103
-tool_call_idx = -1
104
-for chunk in chunks:
105
-
106
- if chunk.choices[0].delta.tool_calls:
107
- tool_call = chunk.choices[0].delta.tool_calls[0]
108
-
109
- if tool_call.index != tool_call_idx:
110
- if tool_call_idx >= 0:
111
- print(
112
- f"streamed tool call arguments: {arguments[tool_call_idx]}"
113
- )
114
- tool_call_idx = chunk.choices[0].delta.tool_calls[0].index
115
- arguments.append("")
116
- if tool_call.id:
117
- print(f"streamed tool call id: {tool_call.id} ")
118
-
119
- if tool_call.function:
120
- if tool_call.function.name:
121
- print(f"streamed tool call name: {tool_call.function.name}")
122
-
123
- if tool_call.function.arguments:
124
- arguments[tool_call_idx] += tool_call.function.arguments
125
-
126
-if len(arguments):
127
- print(f"streamed tool call arguments: {arguments[-1]}")
128
-
129
-print("\n\n")
130
-
131
-messages.append({
132
- "role": "assistant",
133
- "tool_calls": chat_completion.choices[0].message.tool_calls
134
-})
73
136
-
137
-# Now, simulate a tool call
74
def get_current_weather(city: str, state: str, unit: 'str'):
75
return ("The weather in Dallas, Texas is 85 degrees fahrenheit. It is "
76
"partly cloudly, with highs in the 90's.")
77
78
143
-available_tools = {"get_current_weather": get_current_weather}
144
-
145
-completion_tool_calls = chat_completion.choices[0].message.tool_calls
146
-for call in completion_tool_calls:
147
- tool_to_call = available_tools[call.function.name]
148
- args = json.loads(call.function.arguments)
149
- result = tool_to_call(**args)
150
- print(result)
79
+def handle_tool_calls_stream(
80
+ client: OpenAI,
81
+ messages: list[dict[str, str]],
82
+ model: str,
83
+ tools: list[dict[str, Any]],
84
+) -> list[Any]:
85
+ tool_calls_stream = client.chat.completions.create(messages=messages,
86
+ model=model,
87
+ tools=tools,
88
+ stream=True)
89
+ chunks = []
90
+ print("chunks: ")
91
+ for chunk in tool_calls_stream:
92
+ chunks.append(chunk)
93
+ if chunk.choices[0].delta.tool_calls:
94
+ print(chunk.choices[0].delta.tool_calls[0])
95
+ else:
96
+ print(chunk.choices[0].delta)
97
+ return chunks
98
+
99
+
100
+def handle_tool_calls_arguments(chunks: list[Any]) -> list[str]:
101
+ arguments = []
102
+ tool_call_idx = -1
103
+ print("arguments: ")
104
+ for chunk in chunks:
105
+ if chunk.choices[0].delta.tool_calls:
106
+ tool_call = chunk.choices[0].delta.tool_calls[0]
107
+ if tool_call.index != tool_call_idx:
108
+ if tool_call_idx >= 0:
109
+ print(f"streamed tool call arguments: "
110
+ f"{arguments[tool_call_idx]}")
111
+ tool_call_idx = chunk.choices[0].delta.tool_calls[0].index
112
+ arguments.append("")
113
+ if tool_call.id:
114
+ print(f"streamed tool call id: {tool_call.id} ")
115
+
116
+ if tool_call.function:
117
+ if tool_call.function.name:
118
+ print(
119
+ f"streamed tool call name: {tool_call.function.name}")
120
+
121
+ if tool_call.function.arguments:
122
+ arguments[tool_call_idx] += tool_call.function.arguments
123
+
124
+ return arguments
125
+
126
+
127
+def main():
128
+ # Initialize OpenAI client
129
+ client = OpenAI(
130
+ # defaults to os.environ.get("OPENAI_API_KEY")
131
+ api_key=openai_api_key,
132
+ base_url=openai_api_base,
133
+ )
134
+
135
+ # Get available models and select one
136
+ models = client.models.list()
137
+ model = models.data[0].id
138
+
139
+ chat_completion = client.chat.completions.create(messages=messages,
140
+ model=model,
141
+ tools=tools)
142
+
143
+ print("-" * 70)
144
+ print("Chat completion results:")
145
+ print(chat_completion)
146
+ print("-" * 70)
147
+
148
+ # Stream tool calls
149
+ chunks = handle_tool_calls_stream(client, messages, model, tools)
150
+ print("-" * 70)
151
+
152
+ # Handle arguments from streamed tool calls
153
+ arguments = handle_tool_calls_arguments(chunks)
154
+
155
+ if len(arguments):
156
+ print(f"streamed tool call arguments: {arguments[-1]}\n")
157
+
158
+ print("-" * 70)
159
+
160
+ # Add tool call results to the conversation
161
messages.append({
152
- "role": "tool",
153
- "content": result,
154
- "tool_call_id": call.id,
155
- "name": call.function.name
162
+ "role": "assistant",
163
+ "tool_calls": chat_completion.choices[0].message.tool_calls
164
})
165
158
-chat_completion_2 = client.chat.completions.create(messages=messages,
159
- model=model,
160
- tools=tools,
161
- stream=False)
162
-print("\n\n")
163
-print(chat_completion_2)
166
+ # Now, simulate a tool call
167
+ available_tools = {"get_current_weather": get_current_weather}
168
+
169
+ completion_tool_calls = chat_completion.choices[0].message.tool_calls
170
+ for call in completion_tool_calls:
171
+ tool_to_call = available_tools[call.function.name]
172
+ args = json.loads(call.function.arguments)
173
+ result = tool_to_call(**args)
174
+ print("tool_to_call result: ", result)
175
+ messages.append({
176
+ "role": "tool",
177
+ "content": result,
178
+ "tool_call_id": call.id,
179
+ "name": call.function.name
180
+ })
181
+
182
+ chat_completion_2 = client.chat.completions.create(messages=messages,
183
+ model=model,
184
+ tools=tools,
185
+ stream=False)
186
+ print("Chat completion2 results:")
187
+ print(chat_completion_2)
188
+ print("-" * 70)
189
+
190
+
191
+if __name__ == "__main__":
192
+ main()