From 234f0a405197d291e9c7e6ce55dcb9b93b42746e Mon Sep 17 00:00:00 2001
From: Sean Shi <69175566+seanshi-scale@users.noreply.github.com>
Date: Mon, 11 Mar 2024 18:12:38 -0700
Subject: [PATCH] [Client] Add num_prompt_tokens to the client's
 CompletionOutputs (#467)

* add prompt token, untested

* comment

* remove stop_str stuff, it doesn't do anything with the public api, and it breaks on certain frameworks when hosted locally
---
 clients/python/llmengine/__init__.py   | 2 +-
 clients/python/llmengine/data_types.py | 9 +++++++++
 clients/python/pyproject.toml          | 2 +-
 clients/python/setup.py                | 2 +-
 4 files changed, 12 insertions(+), 3 deletions(-)

diff --git a/clients/python/llmengine/__init__.py b/clients/python/llmengine/__init__.py
index 998388ac..17dacfa9 100644
--- a/clients/python/llmengine/__init__.py
+++ b/clients/python/llmengine/__init__.py
@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-__version__ = "0.0.0b26"
+__version__ = "0.0.0b27"
 
 import os
 from typing import Sequence
diff --git a/clients/python/llmengine/data_types.py b/clients/python/llmengine/data_types.py
index 06c0b805..70abd6cb 100644
--- a/clients/python/llmengine/data_types.py
+++ b/clients/python/llmengine/data_types.py
@@ -305,6 +305,11 @@ class CompletionOutput(BaseModel):
     text: str
     """The text of the completion."""
 
+    # We're not guaranteed to have `num_prompt_tokens` in the response in all cases, so to be safe, set a default.
+    # If we send request to api.spellbook.scale.com, we don't get this back.
+    num_prompt_tokens: Optional[int] = None
+    """Number of tokens in the prompt."""
+
     num_completion_tokens: int
     """Number of tokens in the completion."""
 
@@ -353,6 +358,10 @@ class CompletionStreamOutput(BaseModel):
     finished: bool
     """Whether the completion is finished."""
 
+    # We're not guaranteed to have `num_prompt_tokens` in the response in all cases, so to be safe, set a default.
+    num_prompt_tokens: Optional[int] = None
+    """Number of tokens in the prompt."""
+
     num_completion_tokens: Optional[int] = None
     """Number of tokens in the completion."""
 
diff --git a/clients/python/pyproject.toml b/clients/python/pyproject.toml
index a2fdc9ce..2563b814 100644
--- a/clients/python/pyproject.toml
+++ b/clients/python/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "scale-llm-engine"
-version = "0.0.0.beta26"
+version = "0.0.0.beta27"
 description = "Scale LLM Engine Python client"
 license = "Apache-2.0"
 authors = ["Phil Chen <phil.chen@scale.com>"]
diff --git a/clients/python/setup.py b/clients/python/setup.py
index 9afe6136..257516fc 100644
--- a/clients/python/setup.py
+++ b/clients/python/setup.py
@@ -3,6 +3,6 @@
 setup(
     name="scale-llm-engine",
     python_requires=">=3.7",
-    version="0.0.0.beta26",
+    version="0.0.0.beta27",
     packages=find_packages(),
 )