Update types.py

Author: Smartappli

llama_cpp/server/types.py

@@ -1,18 +1,18 @@
 from __future__ import annotations
 
-from typing import Dict, List, Literal, Optional, Union
+from typing import Dict, List, Optional, Union
 
 from pydantic import BaseModel, Field
-from typing_extensions import TypedDict
+from typing_extensions import Literal, TypedDict
 
 import llama_cpp
 
 model_field = Field(
-    description="The model to use for generating completions.", default=None,
+    description="The model to use for generating completions.", default=None
 )
 
 max_tokens_field = Field(
-    default=16, ge=1, description="The maximum number of tokens to generate.",
+    default=16, ge=1, description="The maximum number of tokens to generate."
 )
 
 min_tokens_field = Field(
@@ -96,7 +96,7 @@
 )
 
 mirostat_eta_field = Field(
-    default=0.1, ge=0.001, le=1.0, description="Mirostat learning rate",
+    default=0.1, ge=0.001, le=1.0, description="Mirostat learning rate"
 )
 
 grammar = Field(
@@ -106,15 +106,15 @@
 
 
 class CreateCompletionRequest(BaseModel):
-    prompt: str | list[str] = Field(
-        default="", description="The prompt to generate completions for.",
+    prompt: Union[str, List[str]] = Field(
+        default="", description="The prompt to generate completions for."
     )
-    suffix: str | None = Field(
+    suffix: Optional[str] = Field(
         default=None,
         description="A suffix to append to the generated text. If None, no suffix is appended. Useful for chatbots.",
     )
-    max_tokens: int | None = Field(
-        default=16, ge=0, description="The maximum number of tokens to generate.",
+    max_tokens: Optional[int] = Field(
+        default=16, ge=0, description="The maximum number of tokens to generate."
     )
     min_tokens: int = min_tokens_field
     temperature: float = temperature_field
@@ -124,172 +124,172 @@ class CreateCompletionRequest(BaseModel):
         default=False,
         description="Whether to echo the prompt in the generated text. Useful for chatbots.",
     )
-    stop: str | list[str] | None = stop_field
+    stop: Optional[Union[str, List[str]]] = stop_field
     stream: bool = stream_field
-    logprobs: int | None = Field(
+    logprobs: Optional[int] = Field(
         default=None,
         ge=0,
         description="The number of logprobs to generate. If None, no logprobs are generated.",
     )
-    presence_penalty: float | None = presence_penalty_field
-    frequency_penalty: float | None = frequency_penalty_field
-    logit_bias: dict[str, float] | None = Field(None)
-    seed: int | None = Field(None)
+    presence_penalty: Optional[float] = presence_penalty_field
+    frequency_penalty: Optional[float] = frequency_penalty_field
+    logit_bias: Optional[Dict[str, float]] = Field(None)
+    seed: Optional[int] = Field(None)
 
     # ignored or currently unsupported
-    model: str | None = model_field
-    n: int | None = 1
-    best_of: int | None = 1
-    user: str | None = Field(default=None)
+    model: Optional[str] = model_field
+    n: Optional[int] = 1
+    best_of: Optional[int] = 1
+    user: Optional[str] = Field(default=None)
 
     # llama.cpp specific parameters
     top_k: int = top_k_field
     repeat_penalty: float = repeat_penalty_field
-    logit_bias_type: Literal["input_ids", "tokens"] | None = Field(None)
+    logit_bias_type: Optional[Literal["input_ids", "tokens"]] = Field(None)
     mirostat_mode: int = mirostat_mode_field
     mirostat_tau: float = mirostat_tau_field
     mirostat_eta: float = mirostat_eta_field
-    grammar: str | None = None
+    grammar: Optional[str] = None
 
     model_config = {
         "json_schema_extra": {
             "examples": [
                 {
                     "prompt": "\n\n### Instructions:\nWhat is the capital of France?\n\n### Response:\n",
                     "stop": ["\n", "###"],
-                },
-            ],
-        },
+                }
+            ]
+        }
     }
 
 
 class CreateEmbeddingRequest(BaseModel):
-    model: str | None = model_field
-    input: str | list[str] = Field(description="The input to embed.")
-    user: str | None = Field(default=None)
+    model: Optional[str] = model_field
+    input: Union[str, List[str]] = Field(description="The input to embed.")
+    user: Optional[str] = Field(default=None)
 
     model_config = {
         "json_schema_extra": {
             "examples": [
                 {
                     "input": "The food was delicious and the waiter...",
-                },
-            ],
-        },
+                }
+            ]
+        }
     }
 
 
 class ChatCompletionRequestMessage(BaseModel):
     role: Literal["system", "user", "assistant", "function"] = Field(
-        default="user", description="The role of the message.",
+        default="user", description="The role of the message."
     )
-    content: str | None = Field(
-        default="", description="The content of the message.",
+    content: Optional[str] = Field(
+        default="", description="The content of the message."
     )
 
 
 class CreateChatCompletionRequest(BaseModel):
-    messages: list[llama_cpp.ChatCompletionRequestMessage] = Field(
-        default=[], description="A list of messages to generate completions for.",
+    messages: List[llama_cpp.ChatCompletionRequestMessage] = Field(
+        default=[], description="A list of messages to generate completions for."
     )
-    functions: list[llama_cpp.ChatCompletionFunction] | None = Field(
+    functions: Optional[List[llama_cpp.ChatCompletionFunction]] = Field(
         default=None,
         description="A list of functions to apply to the generated completions.",
     )
-    function_call: llama_cpp.ChatCompletionRequestFunctionCall | None = Field(
+    function_call: Optional[llama_cpp.ChatCompletionRequestFunctionCall] = Field(
         default=None,
         description="A function to apply to the generated completions.",
     )
-    tools: list[llama_cpp.ChatCompletionTool] | None = Field(
+    tools: Optional[List[llama_cpp.ChatCompletionTool]] = Field(
         default=None,
         description="A list of tools to apply to the generated completions.",
     )
-    tool_choice: llama_cpp.ChatCompletionToolChoiceOption | None = Field(
+    tool_choice: Optional[llama_cpp.ChatCompletionToolChoiceOption] = Field(
         default=None,
         description="A tool to apply to the generated completions.",
     )  # TODO: verify
-    max_tokens: int | None = Field(
+    max_tokens: Optional[int] = Field(
         default=None,
         description="The maximum number of tokens to generate. Defaults to inf",
     )
     min_tokens: int = min_tokens_field
-    logprobs: bool | None = Field(
+    logprobs: Optional[bool] = Field(
         default=False,
         description="Whether to output the logprobs or not. Default is True",
     )
-    top_logprobs: int | None = Field(
+    top_logprobs: Optional[int] = Field(
         default=None,
         ge=0,
         description="The number of logprobs to generate. If None, no logprobs are generated. logprobs need to set to True.",
     )
     temperature: float = temperature_field
     top_p: float = top_p_field
     min_p: float = min_p_field
-    stop: str | list[str] | None = stop_field
+    stop: Optional[Union[str, List[str]]] = stop_field
     stream: bool = stream_field
-    presence_penalty: float | None = presence_penalty_field
-    frequency_penalty: float | None = frequency_penalty_field
-    logit_bias: dict[str, float] | None = Field(None)
-    seed: int | None = Field(None)
-    response_format: llama_cpp.ChatCompletionRequestResponseFormat | None = Field(
+    presence_penalty: Optional[float] = presence_penalty_field
+    frequency_penalty: Optional[float] = frequency_penalty_field
+    logit_bias: Optional[Dict[str, float]] = Field(None)
+    seed: Optional[int] = Field(None)
+    response_format: Optional[llama_cpp.ChatCompletionRequestResponseFormat] = Field(
         default=None,
     )
 
     # ignored or currently unsupported
-    model: str | None = model_field
-    n: int | None = 1
-    user: str | None = Field(None)
+    model: Optional[str] = model_field
+    n: Optional[int] = 1
+    user: Optional[str] = Field(None)
 
     # llama.cpp specific parameters
     top_k: int = top_k_field
     repeat_penalty: float = repeat_penalty_field
-    logit_bias_type: Literal["input_ids", "tokens"] | None = Field(None)
+    logit_bias_type: Optional[Literal["input_ids", "tokens"]] = Field(None)
     mirostat_mode: int = mirostat_mode_field
     mirostat_tau: float = mirostat_tau_field
     mirostat_eta: float = mirostat_eta_field
-    grammar: str | None = None
+    grammar: Optional[str] = None
 
     model_config = {
         "json_schema_extra": {
             "examples": [
                 {
                     "messages": [
                         ChatCompletionRequestMessage(
-                            role="system", content="You are a helpful assistant.",
+                            role="system", content="You are a helpful assistant."
                         ).model_dump(),
                         ChatCompletionRequestMessage(
-                            role="user", content="What is the capital of France?",
+                            role="user", content="What is the capital of France?"
                         ).model_dump(),
-                    ],
-                },
-            ],
-        },
+                    ]
+                }
+            ]
+        }
     }
 
 
 class ModelData(TypedDict):
     id: str
     object: Literal["model"]
     owned_by: str
-    permissions: list[str]
+    permissions: List[str]
 
 
 class ModelList(TypedDict):
     object: Literal["list"]
-    data: list[ModelData]
+    data: List[ModelData]
 
 
 class TokenizeInputRequest(BaseModel):
-    model: str | None = model_field
+    model: Optional[str] = model_field
     input: str = Field(description="The input to tokenize.")
 
     model_config = {
-        "json_schema_extra": {"examples": [{"input": "How many tokens in this query?"}]},
+        "json_schema_extra": {"examples": [{"input": "How many tokens in this query?"}]}
     }
 
 
 class TokenizeInputResponse(BaseModel):
-    tokens: list[int] = Field(description="A list of tokens.")
+    tokens: List[int] = Field(description="A list of tokens.")
 
     model_config = {"json_schema_extra": {"example": {"tokens": [123, 321, 222]}}}
 
@@ -301,8 +301,8 @@ class TokenizeInputCountResponse(BaseModel):
 
 
 class DetokenizeInputRequest(BaseModel):
-    model: str | None = model_field
-    tokens: list[int] = Field(description="A list of toekns to detokenize.")
+    model: Optional[str] = model_field
+    tokens: List[int] = Field(description="A list of toekns to detokenize.")
 
     model_config = {"json_schema_extra": {"example": [{"tokens": [123, 321, 222]}]}}
 
@@ -311,5 +311,5 @@ class DetokenizeInputResponse(BaseModel):
     text: str = Field(description="The detokenized text.")
 
     model_config = {
-        "json_schema_extra": {"example": {"text": "How many tokens in this query?"}},
+        "json_schema_extra": {"example": {"text": "How many tokens in this query?"}}
     }