From 149102927b8430735fb52747355304e82a8b73e4 Mon Sep 17 00:00:00 2001
From: Yeuoly <45712896+Yeuoly@users.noreply.github.com>
Date: Sun, 14 Jan 2024 15:51:05 +0800
Subject: [PATCH] fix: openai tool tokens (#2026)

---
 .../model_providers/openai/llm/llm.py         |  1 -
 .../model_runtime/openai/test_llm.py          | 27 ++++++++++++++++++-
 2 files changed, 26 insertions(+), 2 deletions(-)

diff --git a/api/core/model_runtime/model_providers/openai/llm/llm.py b/api/core/model_runtime/model_providers/openai/llm/llm.py
index 97397c2274..4b6ba4be9c 100644
--- a/api/core/model_runtime/model_providers/openai/llm/llm.py
+++ b/api/core/model_runtime/model_providers/openai/llm/llm.py
@@ -765,7 +765,6 @@ class OpenAILargeLanguageModel(_CommonOpenAI, LargeLanguageModel):
         num_tokens = 0
         for tool in tools:
             num_tokens += len(encoding.encode('type'))
-            num_tokens += len(encoding.encode(tool.get("type")))
             num_tokens += len(encoding.encode('function'))
 
             # calculate num tokens for function object
diff --git a/api/tests/integration_tests/model_runtime/openai/test_llm.py b/api/tests/integration_tests/model_runtime/openai/test_llm.py
index c7a9b48776..55afd69167 100644
--- a/api/tests/integration_tests/model_runtime/openai/test_llm.py
+++ b/api/tests/integration_tests/model_runtime/openai/test_llm.py
@@ -327,10 +327,35 @@ def test_get_num_tokens():
             UserPromptMessage(
                 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)
 def test_fine_tuned_models(setup_openai_mock):