fix: openai tool tokens (#2026)

This commit is contained in:
Yeuoly 2024-01-14 15:51:05 +08:00 committed by GitHub
parent d8c0d722d2
commit 149102927b
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 26 additions and 2 deletions

View File

@ -765,7 +765,6 @@ class OpenAILargeLanguageModel(_CommonOpenAI, LargeLanguageModel):
num_tokens = 0 num_tokens = 0
for tool in tools: for tool in tools:
num_tokens += len(encoding.encode('type')) num_tokens += len(encoding.encode('type'))
num_tokens += len(encoding.encode(tool.get("type")))
num_tokens += len(encoding.encode('function')) num_tokens += len(encoding.encode('function'))
# calculate num tokens for function object # calculate num tokens for function object

View File

@ -327,10 +327,35 @@ def test_get_num_tokens():
UserPromptMessage( UserPromptMessage(
content='Hello World!' content='Hello World!'
) )
],
tools=[
PromptMessageTool(
name='get_weather',
description='Determine weather in my location',
parameters={
"type": "object",
"properties": {
"location": {
"type": "string",
"description": "The city and state e.g. San Francisco, CA"
},
"unit": {
"type": "string",
"enum": [
"c",
"f"
]
}
},
"required": [
"location"
]
}
),
] ]
) )
assert num_tokens == 21 assert num_tokens == 72
@pytest.mark.parametrize('setup_openai_mock', [['chat', 'remote']], indirect=True) @pytest.mark.parametrize('setup_openai_mock', [['chat', 'remote']], indirect=True)
def test_fine_tuned_models(setup_openai_mock): def test_fine_tuned_models(setup_openai_mock):