theroyallab/tabbyAPI
1
0
Fork 0
mirror of https://github.com/theroyallab/tabbyAPI.git synced 2026-03-15 00:07:28 +00:00
Code Issues Packages Projects Releases Wiki Activity

Adding token usage support

Browse Source
This commit is contained in:
Mehran Ziadloo
2023-11-27 20:05:05 -08:00
parent 44e7f7b0ee
commit ead503c75b
6 changed files with 34 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
OAI/types/chat_completion.py
View File

@@ -32,8 +32,6 @@ class ChatCompletionResponse(BaseModel):
created: int = Field(default_factory=lambda: int(time()))
model: str
object: str = "chat.completion"
# TODO: Add usage stats
usage: Optional[UsageStats] = None
class ChatCompletionStreamChunk(BaseModel):

2
OAI/types/common.py
View File

@@ -8,8 +8,8 @@ class LogProbs(BaseModel):
top_logprobs: List[Dict[str, float]] = Field(default_factory=list)
class UsageStats(BaseModel):
completion_tokens: int
prompt_tokens: int
completion_tokens: int
total_tokens: int
class CommonCompletionRequest(BaseModel):

2
OAI/types/completion.py
View File

@@ -22,6 +22,4 @@ class CompletionResponse(BaseModel):
created: int = Field(default_factory=lambda: int(time()))
model: str
object: str = "text_completion"
# TODO: Add usage stats
usage: Optional[UsageStats] = None

20
OAI/utils.py
View File

@@ -1,5 +1,5 @@
import os, pathlib
from OAI.types.completion import CompletionResponse, CompletionRespChoice
from OAI.types.completion import CompletionResponse, CompletionRespChoice, UsageStats
from OAI.types.chat_completion import (
ChatCompletionMessage,
ChatCompletionRespChoice,
@@ -20,9 +20,7 @@ try:
except ImportError:
_fastchat_available = False
def create_completion_response(text: str, model_name: Optional[str]):
# TODO: Add method to get token amounts in model for UsageStats
def create_completion_response(text: str, prompt_tokens: int, completion_tokens: int, model_name: Optional[str]):
choice = CompletionRespChoice(
finish_reason = "Generated",
text = text
@@ -30,14 +28,15 @@ def create_completion_response(text: str, model_name: Optional[str]):
response = CompletionResponse(
choices = [choice],
model = model_name or ""
model = model_name or "",
usage = UsageStats(prompt_tokens = prompt_tokens,
completion_tokens = completion_tokens,
total_tokens = prompt_tokens + completion_tokens)
)
return response
def create_chat_completion_response(text: str, model_name: Optional[str]):
# TODO: Add method to get token amounts in model for UsageStats
def create_chat_completion_response(text: str, prompt_tokens: int, completion_tokens: int, model_name: Optional[str]):
message = ChatCompletionMessage(
role = "assistant",
content = text
@@ -50,7 +49,10 @@ def create_chat_completion_response(text: str, model_name: Optional[str]):
response = ChatCompletionResponse(
choices = [choice],
model = model_name or ""
model = model_name or "",
usage = UsageStats(prompt_tokens = prompt_tokens,
completion_tokens = completion_tokens,
total_tokens = prompt_tokens + completion_tokens)
)
return response