| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | # ruff: noqa |
| 5 | import json |
| 6 | import random |
| 7 | import string |
| 8 | |
| 9 | from vllm import LLM |
| 10 | from vllm.sampling_params import SamplingParams |
| 11 | |
| 12 | # This script is an offline demo for function calling |
| 13 | # |
| 14 | # If you want to run a server/client setup, please follow this code: |
| 15 | # |
| 16 | # - Server: |
| 17 | # |
| 18 | # ```bash |
| 19 | # vllm serve mistralai/Mistral-7B-Instruct-v0.3 --tokenizer-mode mistral --load-format mistral --config-format mistral |
| 20 | # ``` |
| 21 | # |
| 22 | # - Client: |
| 23 | # |
| 24 | # ```bash |
| 25 | # curl --location 'http://<your-node-url>:8000/v1/chat/completions' \ |
| 26 | # --header 'Content-Type: application/json' \ |
| 27 | # --header 'Authorization: Bearer token' \ |
| 28 | # --data '{ |
| 29 | # "model": "mistralai/Mistral-7B-Instruct-v0.3" |
| 30 | # "messages": [ |
| 31 | # { |
| 32 | # "role": "user", |
| 33 | # "content": [ |
| 34 | # {"type" : "text", "text": "Describe this image in detail please."}, |
| 35 | # {"type": "image_url", "image_url": {"url": "https://s3.amazonaws.com/cms.ipressroom.com/338/files/201808/5b894ee1a138352221103195_A680%7Ejogging-edit/A680%7Ejogging-edit_hero.jpg"}}, |
| 36 | # {"type" : "text", "text": "and this one as well. Answer in French."}, |
| 37 | # {"type": "image_url", "image_url": {"url": "https://www.wolframcloud.com/obj/resourcesystem/images/a0e/a0ee3983-46c6-4c92-b85d-059044639928/6af8cfb971db031b.png"}} |
| 38 | # ] |
| 39 | # } |
| 40 | # ] |
| 41 | # }' |
| 42 | # ``` |
| 43 | # |
| 44 | # Usage: |
| 45 | # python demo.py simple |
| 46 | # python demo.py advanced |
| 47 | |
| 48 | model_name = "mistralai/Mistral-7B-Instruct-v0.3" |
| 49 | # or switch to "mistralai/Mistral-Nemo-Instruct-2407" |
| 50 | # or "mistralai/Mistral-Large-Instruct-2407" |
| 51 | # or any other mistral model with function calling ability |
| 52 | |
| 53 | sampling_params = SamplingParams(max_tokens=8192, temperature=0.0) |
| 54 | llm = LLM( |
| 55 | model=model_name, |
| 56 | tokenizer_mode="mistral", |
| 57 | config_format="mistral", |
| 58 | load_format="mistral", |
| 59 | ) |
| 60 | |
| 61 | |
| 62 | def generate_random_id(length=9): |
| 63 | characters = string.ascii_letters + string.digits |
| 64 | random_id = "".join(random.choice(characters) for _ in range(length)) |
| 65 | return random_id |
| 66 | |
| 67 | |
| 68 | # simulate an API that can be called |
| 69 | def get_current_weather(city: str, state: str, unit: "str"): |
| 70 | return ( |
| 71 | f"The weather in {city}, {state} is 85 degrees {unit}. It is " |
| 72 | "partly cloudly, with highs in the 90's." |
| 73 | ) |
| 74 | |
| 75 | |
| 76 | tool_functions = {"get_current_weather": get_current_weather} |
| 77 | |
| 78 | tools = [ |
| 79 | { |
| 80 | "type": "function", |
| 81 | "function": { |
| 82 | "name": "get_current_weather", |
| 83 | "description": "Get the current weather in a given location", |
| 84 | "parameters": { |
| 85 | "type": "object", |
| 86 | "properties": { |
| 87 | "city": { |
| 88 | "type": "string", |
| 89 | "description": "The city to find the weather for, e.g. 'San Francisco'", |
| 90 | }, |
| 91 | "state": { |
| 92 | "type": "string", |
| 93 | "description": "the two-letter abbreviation for the state that the city is" |
| 94 | " in, e.g. 'CA' which would mean 'California'", |
| 95 | }, |
| 96 | "unit": { |
| 97 | "type": "string", |
| 98 | "description": "The unit to fetch the temperature in", |
| 99 | "enum": ["celsius", "fahrenheit"], |
| 100 | }, |
| 101 | }, |
| 102 | "required": ["city", "state", "unit"], |
| 103 | }, |
| 104 | }, |
| 105 | } |
| 106 | ] |
| 107 | |
| 108 | messages = [ |
| 109 | { |
| 110 | "role": "user", |
| 111 | "content": "Can you tell me what the temperate will be in Dallas, in fahrenheit?", |
| 112 | } |
| 113 | ] |
| 114 | |
| 115 | outputs = llm.chat(messages, sampling_params=sampling_params, tools=tools) |
| 116 | output = outputs[0].outputs[0].text.strip() |
| 117 | |
| 118 | # append the assistant message |
| 119 | messages.append( |
| 120 | { |
| 121 | "role": "assistant", |
| 122 | "content": output, |
| 123 | } |
| 124 | ) |
| 125 | |
| 126 | # let's now actually parse and execute the model's output simulating an API call by using the |
| 127 | # above defined function |
| 128 | tool_calls = json.loads(output) |
| 129 | tool_answers = [ |
| 130 | tool_functions[call["name"]](**call["arguments"]) for call in tool_calls |
| 131 | ] |
| 132 | |
| 133 | # append the answer as a tool message and let the LLM give you an answer |
| 134 | messages.append( |
| 135 | { |
| 136 | "role": "tool", |
| 137 | "content": "\n\n".join(tool_answers), |
| 138 | "tool_call_id": generate_random_id(), |
| 139 | } |
| 140 | ) |
| 141 | |
| 142 | outputs = llm.chat(messages, sampling_params, tools=tools) |
| 143 | |
| 144 | print(outputs[0].outputs[0].text.strip()) |
| 145 | # yields |
| 146 | # 'The weather in Dallas, TX is 85 degrees Fahrenheit. ' |
| 147 | # 'It is partly cloudly, with highs in the 90's.' |