NEBIUS_API_KEY=your_api_key_goes_here

import os

if os.getenv("COLAB_RELEASE_TAG"):
   print("Running on Colab")
   RUNNING_ON_COLAB = True
else:
   print("NOT running on Colab")
   RUNNING_ON_COLAB = False

NOT running on Colab

#if RUNNING_ON_COLAB:
!pip install -q openai python-dotenv  pydantic

import os

## Recommended way of getting configuration
if RUNNING_ON_COLAB:
   from google.colab import userdata
   NEBIUS_API_KEY = userdata.get('NEBIUS_API_KEY')
else:
   from dotenv import load_dotenv
   load_dotenv()
   NEBIUS_API_KEY = os.getenv('NEBIUS_API_KEY')


## quick hack (not recommended) - you can hardcode the config key here
# NEBIUS_API_KEY = "your_key_here"

if NEBIUS_API_KEY:
  print ('✅ NEBIUS_API_KEY found')
  os.environ['NEBIUS_API_KEY'] = NEBIUS_API_KEY
else:
  raise RuntimeError ('❌ NEBIUS_API_KEY NOT found')

✅ NEBIUS_API_KEY found

from pydantic import BaseModel, Field
from typing import Literal

class GetCurrentWeatherParams(BaseModel):
    city: str = Field(..., description="The city to find the weather for, e.g. 'San Francisco'")
    unit: Literal['celsius', 'fahrenheit'] = Field(..., description="The unit to fetch the temperature in")

# Now, simulate a tool call
## in real example you will use an API to get the actual weather data
def get_current_weather(city: str, unit: str):
    return ("The weather in San Francisco is 72 degrees fahrenheit. "
            "It is sunny, with highs in the 80's.")

tools = [{
    "type": "function",
    "function": {
        "name": "get_current_weather",
        "description": "Get the current weather in a given location",
        "parameters": GetCurrentWeatherParams.model_json_schema()
    }
}]

available_tools = {"get_current_weather": get_current_weather}

%%time

import json
import os

from openai import OpenAI

client = OpenAI(
    base_url="https://api.tokenfactory.nebius.com/v1/",
    api_key=NEBIUS_API_KEY,
)

CPU times: user 490 ms, sys: 65.7 ms, total: 555 ms
Wall time: 612 ms

messages=[
    {
        "role": "user",
        "content": "Can you tell me what the temperature will be in San Francisco?",
    },
  ]

chat_completion = client.chat.completions.create(
  model = "openai/gpt-oss-20b",
  messages=messages,
  tools = tools,
  tool_choice = {
        "type": "function",
        "function": {
            "name": "get_current_weather"
        }
    }
)

messages.append({
    "role": "assistant",
    "tool_calls": chat_completion.choices[0].message.tool_calls
})


completion_tool_calls = chat_completion.choices[0].message.tool_calls
for call in completion_tool_calls:
    tool_to_call = available_tools[call.function.name]
    args = json.loads(call.function.arguments)
    result = tool_to_call(**args)
    print(result)
    messages.append({
        "role": "tool",
        "content": result,
        "tool_call_id": call.id,
        "name": call.function.name
    })

The weather in San Francisco is 72 degrees fahrenheit. It is sunny, with highs in the 80's.

messages

[{'role': 'user',
  'content': 'Can you tell me what the temperature will be in San Francisco?'},
 {'role': 'assistant',
  'tool_calls': [ChatCompletionMessageToolCall(id='chatcmpl-tool-0c160206cdea4ca6b10cc37e4b592ad7', function=Function(arguments='{"city": "San Francisco", "unit": "celsius"}', name='get_current_weather'), type='function')]},
 {'role': 'tool',
  'content': "The weather in San Francisco is 72 degrees fahrenheit. It is sunny, with highs in the 80's.",
  'tool_call_id': 'chatcmpl-tool-0c160206cdea4ca6b10cc37e4b592ad7',
  'name': 'get_current_weather'}]

Function Calling in Nebius Token Factory¶

Pre requisites¶

1 - Setup¶

1.1 - If running on Google Colab¶

1.2 - If running locally¶

2 - Install Dependencies¶

3 - Load Configuration¶

4 - Pick a Model¶

5 - Define Function Call¶

6 - Tool calling¶