@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 194 lines 5.78 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Set up this example by starting a vLLM OpenAI-compatible server with tool call
4 options enabled. For example:
5
6 IMPORTANT: for mistral, you must use one of the provided mistral tool call
7 templates, or your own - the model default doesn't work for tool calls with vLLM
8 See the vLLM docs on OpenAI server & tool calling for more details.
9
10 vllm serve mistralai/Mistral-7B-Instruct-v0.3 \
11 --chat-template examples/tool_chat_template_mistral.jinja \
12 --enable-auto-tool-choice --tool-call-parser mistral
13
14 OR
15 vllm serve NousResearch/Hermes-2-Pro-Llama-3-8B \
16 --chat-template examples/tool_chat_template_hermes.jinja \
17 --enable-auto-tool-choice --tool-call-parser hermes
18 """
19
20 import json
21 from typing import Any
22
23 from openai import OpenAI
24
25 # Modify OpenAI's API key and API base to use vLLM's API server.
26 openai_api_key = "EMPTY"
27 openai_api_base = "http://localhost:8000/v1"
28
29 properties = {
30 "city": {
31 "type": "string",
32 "description": "The city to find the weather for, e.g. 'San Francisco'",
33 },
34 "state": {
35 "type": "string",
36 "description": "the two-letter abbreviation for the state that the city is"
37 " in, e.g. 'CA' which would mean 'California'",
38 },
39 "unit": {
40 "type": "string",
41 "description": "The unit to fetch the temperature in",
42 "enum": ["celsius", "fahrenheit"],
43 },
44 }
45
46 tools = [
47 {
48 "type": "function",
49 "function": {
50 "name": "get_current_weather",
51 "description": "Get the current weather in a given location",
52 "parameters": {
53 "type": "object",
54 "properties": properties,
55 "required": ["city", "state", "unit"],
56 },
57 },
58 }
59 ]
60
61 messages = [
62 {"role": "user", "content": "Hi! How are you doing today?"},
63 {"role": "assistant", "content": "I'm doing well! How can I help you?"},
64 {
65 "role": "user",
66 "content": (
67 "Can you tell me what the temperate will be in Dallas, in fahrenheit?"
68 ),
69 },
70 ]
71
72
73 def get_current_weather(city: str, state: str, unit: "str"):
74 return (
75 "The weather in Dallas, Texas is 85 degrees fahrenheit. It is "
76 "partly cloudly, with highs in the 90's."
77 )
78
79
80 def handle_tool_calls_stream(
81 client: OpenAI,
82 messages: list[dict[str, str]],
83 model: str,
84 tools: list[dict[str, Any]],
85 ) -> list[Any]:
86 tool_calls_stream = client.chat.completions.create(
87 messages=messages, model=model, tools=tools, stream=True
88 )
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: {arguments[tool_call_idx]}")
110 tool_call_idx = chunk.choices[0].delta.tool_calls[0].index
111 arguments.append("")
112 if tool_call.id:
113 print(f"streamed tool call id: {tool_call.id} ")
114
115 if tool_call.function:
116 if tool_call.function.name:
117 print(f"streamed tool call name: {tool_call.function.name}")
118
119 if tool_call.function.arguments:
120 arguments[tool_call_idx] += tool_call.function.arguments
121
122 return arguments
123
124
125 def main():
126 # Initialize OpenAI client
127 client = OpenAI(
128 # defaults to os.environ.get("OPENAI_API_KEY")
129 api_key=openai_api_key,
130 base_url=openai_api_base,
131 )
132
133 # Get available models and select one
134 models = client.models.list()
135 model = models.data[0].id
136
137 chat_completion = client.chat.completions.create(
138 messages=messages, model=model, tools=tools
139 )
140
141 print("-" * 70)
142 print("Chat completion results:")
143 print(chat_completion)
144 print("-" * 70)
145
146 # Stream tool calls
147 chunks = handle_tool_calls_stream(client, messages, model, tools)
148 print("-" * 70)
149
150 # Handle arguments from streamed tool calls
151 arguments = handle_tool_calls_arguments(chunks)
152
153 if len(arguments):
154 print(f"streamed tool call arguments: {arguments[-1]}\n")
155
156 print("-" * 70)
157
158 # Add tool call results to the conversation
159 messages.append(
160 {
161 "role": "assistant",
162 "tool_calls": chat_completion.choices[0].message.tool_calls,
163 "reasoning": chat_completion.choices[0].message.reasoning,
164 }
165 )
166
167 # Now, simulate a tool call
168 available_tools = {"get_current_weather": get_current_weather}
169
170 completion_tool_calls = chat_completion.choices[0].message.tool_calls
171 for call in completion_tool_calls:
172 tool_to_call = available_tools[call.function.name]
173 args = json.loads(call.function.arguments)
174 result = tool_to_call(**args)
175 print("tool_to_call result: ", result)
176 messages.append(
177 {
178 "role": "tool",
179 "content": result,
180 "tool_call_id": call.id,
181 "name": call.function.name,
182 }
183 )
184
185 chat_completion_2 = client.chat.completions.create(
186 messages=messages, model=model, tools=tools, stream=False
187 )
188 print("Chat completion2 results:")
189 print(chat_completion_2)
190 print("-" * 70)
191
192
193 if __name__ == "__main__":
194 main()