fix(server): minor type fixes

abetlen · abetlen · commit d11ccc30361a · 2024-03-23T17:14:15.000-04:00
diff --git a/llama_cpp/server/app.py b/llama_cpp/server/app.py
@@ -493,7 +493,7 @@ async def tokenize(
 ) -> TokenizeInputResponse:
     tokens = llama_proxy(body.model).tokenize(body.input.encode("utf-8"), special=True)
 
-    return {"tokens": tokens}
+    return TokenizeInputResponse(tokens=tokens)
 
 
 @router.post(
@@ -508,7 +508,7 @@ async def count_query_tokens(
 ) -> TokenizeInputCountResponse:
     tokens = llama_proxy(body.model).tokenize(body.input.encode("utf-8"), special=True)
 
-    return {"count": len(tokens)}
+    return TokenizeInputCountResponse(count=len(tokens))
 
 
 @router.post(
@@ -523,4 +523,4 @@ async def detokenize(
 ) -> DetokenizeInputResponse:
     text = llama_proxy(body.model).detokenize(body.tokens).decode("utf-8")
 
-    return {"text": text}
+    return DetokenizeInputResponse(text=text)
diff --git a/llama_cpp/server/types.py b/llama_cpp/server/types.py
@@ -268,7 +268,7 @@ class ModelList(TypedDict):
 
 class TokenizeInputRequest(BaseModel):
     model: Optional[str] = model_field
-    input: Optional[str] = Field(description="The input to tokenize.")
+    input: str = Field(description="The input to tokenize.")
 
     model_config = {
         "json_schema_extra": {"examples": [{"input": "How many tokens in this query?"}]}
