use annotate and field discriminator

This commit is contained in:
Dinesh Yeduguru 2024-11-01 13:39:38 -07:00 committed by Dinesh Yeduguru
parent ef64863616
commit 7696f31284
7 changed files with 56 additions and 22 deletions

View file

@ -4,7 +4,7 @@
# This source code is licensed under the terms described in the LICENSE file in
# the root directory of this source tree.
from typing import Dict, List, Optional, Protocol
from typing import Dict, List, Literal, Optional, Protocol
from llama_models.llama3.api.datatypes import URL
@ -32,6 +32,7 @@ class DatasetDef(BaseModel):
@json_schema_type
class DatasetDefWithProvider(DatasetDef):
type: Literal["dataset"] = "dataset"
provider_id: str = Field(
description="ID of the provider which serves this dataset",
)

View file

@ -4,7 +4,7 @@
# This source code is licensed under the terms described in the LICENSE file in
# the root directory of this source tree.
from typing import Dict, List, Optional, Protocol, runtime_checkable
from typing import Dict, List, Literal, Optional, Protocol, runtime_checkable
from llama_models.schema_utils import json_schema_type, webmethod
from pydantic import BaseModel, Field
@ -25,6 +25,7 @@ class ModelDef(BaseModel):
@json_schema_type
class ModelDefWithProvider(ModelDef):
type: Literal["model"] = "model"
provider_id: str = Field(
description="The provider ID for this model",
)

View file

@ -4,12 +4,12 @@
# This source code is licensed under the terms described in the LICENSE file in
# the root directory of this source tree.
from typing import Dict, List, Optional, Protocol, runtime_checkable
from typing import Dict, List, Literal, Optional, Protocol, runtime_checkable
from llama_models.schema_utils import json_schema_type, webmethod
from pydantic import BaseModel, Field
from llama_stack.apis.common.type_system import JsonType, ParamType
from pydantic import BaseModel, Field
@json_schema_type
@ -53,6 +53,7 @@ class ScoringFnDef(BaseModel):
@json_schema_type
class ScoringFnDefWithProvider(ScoringFnDef):
type: Literal["scoring_fn"] = "scoring_fn"
provider_id: str = Field(
description="ID of the provider which serves this dataset",
)

View file

@ -5,7 +5,7 @@
# the root directory of this source tree.
from enum import Enum
from typing import Dict, List, Optional, Protocol, runtime_checkable
from typing import Dict, List, Literal, Optional, Protocol, runtime_checkable
from llama_models.schema_utils import json_schema_type, webmethod
from pydantic import BaseModel, Field
@ -34,6 +34,7 @@ class ShieldDef(BaseModel):
@json_schema_type
class ShieldDefWithProvider(ShieldDef):
type: Literal["shield"] = "shield"
provider_id: str = Field(
description="The provider ID for this shield type",
)