Removing json_schema_type from Enums to flatten them out in the schema

Author: ashwinb

models/datatypes.py

@@ -11,7 +11,7 @@
 from pydantic import BaseModel, ConfigDict, Field
 from typing_extensions import Annotated
 
-from .schema_utils import json_schema_type
+from .schema_utils import json_schema_type, register_schema
 
 
 @json_schema_type
@@ -32,10 +32,13 @@ class TopKSamplingStrategy(BaseModel):
     top_k: int = Field(..., ge=1)
 
 
-SamplingStrategy = Annotated[
-    Union[GreedySamplingStrategy, TopPSamplingStrategy, TopKSamplingStrategy],
-    Field(discriminator="type"),
-]
+SamplingStrategy = register_schema(
+    Annotated[
+        Union[GreedySamplingStrategy, TopPSamplingStrategy, TopKSamplingStrategy],
+        Field(discriminator="type"),
+    ],
+    name="SamplingStrategy",
+)
 
 
 @json_schema_type
@@ -46,15 +49,6 @@ class SamplingParams(BaseModel):
     repetition_penalty: Optional[float] = 1.0
 
 
-@json_schema_type(
-    schema={
-        "description": """
-The format in which weights are specified. This does not necessarily
-always equal what quantization is desired at runtime since there
-can be on-the-fly conversions done.
-""",
-    }
-)
 class CheckpointQuantizationFormat(Enum):
     # default format
     bf16 = "bf16"
@@ -67,7 +61,6 @@ class CheckpointQuantizationFormat(Enum):
     int4 = "int4"
 
 
-@json_schema_type
 class ModelFamily(Enum):
     llama2 = "llama2"
     llama3 = "llama3"
@@ -77,7 +70,6 @@ class ModelFamily(Enum):
     safety = "safety"
 
 
-@json_schema_type
 class CoreModelId(Enum):
     """Each of these models is a unique "SKU". These root models can be served in various garbs (especially by quantizing them)"""
 
@@ -187,11 +179,6 @@ def model_family(model_id) -> ModelFamily:
         raise ValueError(f"Unknown model family for {model_id}")
 
 
-@json_schema_type(
-    schema={
-        "description": "The model family and SKU of the model along with other parameters corresponding to the model."
-    }
-)
 class Model(BaseModel):
     core_model_id: CoreModelId
     description: str

models/llama3/api/datatypes.py

@@ -20,15 +20,13 @@
 from ...schema_utils import json_schema_type
 
 
-@json_schema_type
 class Role(Enum):
     system = "system"
     user = "user"
     assistant = "assistant"
     tool = "tool"
 
 
-@json_schema_type
 class BuiltinTool(Enum):
     brave_search = "brave_search"
     wolfram_alpha = "wolfram_alpha"
@@ -82,13 +80,10 @@ def validate_field(cls, v):
         return v
 
 
-@json_schema_type
 class ToolPromptFormat(Enum):
-    """This Enum refers to the prompt format for calling custom / zero shot tools
+    """Prompt format for calling custom / zero shot tools.
 
-    `json` --
-        Refers to the json format for calling tools.
-        The json format takes the form like
+    :cvar json: JSON format for calling tools. It takes the form:
         {
             "type": "function",
             "function" : {
@@ -97,22 +92,19 @@ class ToolPromptFormat(Enum):
                 "parameters": {...}
             }
         }
-
-    `function_tag` --
-        This is an example of how you could define
-        your own user defined format for making tool calls.
-        The function_tag format looks like this,
+    :cvar function_tag: Function tag format, pseudo-XML. This looks like:
         <function=function_name>(parameters)</function>
 
-    The detailed prompts for each of these formats are added to llama cli
+    :cvar python_list: Python list. The output is a valid Python expression that can be
+        evaluated to a list. Each element in the list is a function call. Example:
+        ["function_name(param1, param2)", "function_name(param1, param2)"]
     """
 
     json = "json"
     function_tag = "function_tag"
     python_list = "python_list"
 
 
-@json_schema_type
 class StopReason(Enum):
     end_of_turn = "end_of_turn"
     end_of_message = "end_of_message"