From 287623b5136638a7c11303c1ed34f3588d7d4923 Mon Sep 17 00:00:00 2001
From: Xi Yan <xiyan@meta.com>
Date: Tue, 22 Oct 2024 08:16:14 -0700
Subject: [PATCH] type system api

---
 llama_stack/apis/common/type_system.py | 65 ++++++++++++++++++++++++++
 1 file changed, 65 insertions(+)
 create mode 100644 llama_stack/apis/common/type_system.py

diff --git a/llama_stack/apis/common/type_system.py b/llama_stack/apis/common/type_system.py
new file mode 100644
index 000000000..35a26e9ef
--- /dev/null
+++ b/llama_stack/apis/common/type_system.py
@@ -0,0 +1,65 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# 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, Literal, Union
+
+from pydantic import BaseModel, Field
+from typing_extensions import Annotated
+
+
+class StringType(BaseModel):
+    type: Literal["string"] = "string"
+
+
+class NumberType(BaseModel):
+    type: Literal["number"] = "number"
+
+
+class BooleanType(BaseModel):
+    type: Literal["boolean"] = "boolean"
+
+
+class ArrayType(BaseModel):
+    type: Literal["array"] = "array"
+    items: "ParamType"
+
+
+class ObjectType(BaseModel):
+    type: Literal["object"] = "object"
+    properties: Dict[str, "ParamType"] = Field(default_factory=dict)
+
+
+class JsonType(BaseModel):
+    type: Literal["json"] = "json"
+
+
+class UnionType(BaseModel):
+    type: Literal["union"] = "union"
+    options: List["ParamType"] = Field(default_factory=list)
+
+
+class CustomType(BaseModel):
+    type: Literal["custom"] = "custom"
+    validator_class: str
+
+
+ParamType = Annotated[
+    Union[
+        StringType,
+        NumberType,
+        BooleanType,
+        ArrayType,
+        ObjectType,
+        JsonType,
+        UnionType,
+        CustomType,
+    ],
+    Field(discriminator="type"),
+]
+
+ArrayType.model_rebuild()
+ObjectType.model_rebuild()
+UnionType.model_rebuild()