Skip to content

Commit 2568055

Browse files
reidliu41yangw-dev
reidliu41
authored and
yangw-dev
committed
[Misc] refactor examples series - Chat Completion Client With Tools (vllm-project#16829)
Signed-off-by: reidliu41 <[email protected]> Co-authored-by: reidliu41 <[email protected]> Signed-off-by: Yang Wang <[email protected]>
1 parent ee3b53c commit 2568055

File tree

1 file changed

+112
-83
lines changed

1 file changed

+112
-83
lines changed

‎examples/online_serving/openai_chat_completion_client_with_tools.py

Lines changed: 112 additions & 83 deletions
Original file line numberDiff line numberDiff line change
@@ -17,22 +17,14 @@
1717
--enable-auto-tool-choice --tool-call-parser hermes
1818
"""
1919
import json
20+
from typing import Any
2021

2122
from openai import OpenAI
2223

2324
# Modify OpenAI's API key and API base to use vLLM's API server.
2425
openai_api_key = "EMPTY"
2526
openai_api_base = "http://localhost:8000/v1"
2627

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-
3628
tools = [{
3729
"type": "function",
3830
"function": {
@@ -78,86 +70,123 @@
7870
"Can you tell me what the temperate will be in Dallas, in fahrenheit?"
7971
}]
8072

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-
})
13573

136-
137-
# Now, simulate a tool call
13874
def get_current_weather(city: str, state: str, unit: 'str'):
13975
return ("The weather in Dallas, Texas is 85 degrees fahrenheit. It is "
14076
"partly cloudly, with highs in the 90's.")
14177

14278

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
151161
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
156164
})
157165

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()

0 commit comments

Comments
 (0)