From 72ca18c40628a411ea7522761ef5e96fccbbffa2 Mon Sep 17 00:00:00 2001
From: Xi Yan <xiyan@meta.com>
Date: Thu, 30 Jan 2025 15:27:45 -0800
Subject: [PATCH] openapi schema

---
 .../openapi_generator/strong_typing/schema.py | 41 ++++++++++++++++---
 docs/resources/llama-stack-spec.html          |  1 +
 docs/resources/llama-stack-spec.yaml          |  1 +
 llama_stack/apis/common/content_types.py      |  4 +-
 4 files changed, 39 insertions(+), 8 deletions(-)

diff --git a/docs/openapi_generator/strong_typing/schema.py b/docs/openapi_generator/strong_typing/schema.py
index f4393041f..795e160e3 100644
--- a/docs/openapi_generator/strong_typing/schema.py
+++ b/docs/openapi_generator/strong_typing/schema.py
@@ -248,7 +248,9 @@ class JsonSchemaGenerator:
                 type_schema.update(self._metadata_to_schema(m))
         return type_schema
 
-    def _simple_type_to_schema(self, typ: TypeLike) -> Optional[Schema]:
+    def _simple_type_to_schema(
+        self, typ: TypeLike, json_schema_extra: Optional[dict] = None
+    ) -> Optional[Schema]:
         """
         Returns the JSON schema associated with a simple, unrestricted type.
 
@@ -264,6 +266,11 @@ class JsonSchemaGenerator:
         elif typ is float:
             return {"type": "number"}
         elif typ is str:
+            if json_schema_extra and "contentEncoding" in json_schema_extra:
+                return {
+                    "type": "string",
+                    "contentEncoding": json_schema_extra["contentEncoding"],
+                }
             return {"type": "string"}
         elif typ is bytes:
             return {"type": "string", "contentEncoding": "base64"}
@@ -303,7 +310,12 @@ class JsonSchemaGenerator:
             # not a simple type
             return None
 
-    def type_to_schema(self, data_type: TypeLike, force_expand: bool = False) -> Schema:
+    def type_to_schema(
+        self,
+        data_type: TypeLike,
+        force_expand: bool = False,
+        json_schema_extra: Optional[dict] = None,
+    ) -> Schema:
         """
         Returns the JSON schema associated with a type.
 
@@ -313,7 +325,7 @@ class JsonSchemaGenerator:
         """
 
         # short-circuit for common simple types
-        schema = self._simple_type_to_schema(data_type)
+        schema = self._simple_type_to_schema(data_type, json_schema_extra)
         if schema is not None:
             return schema
 
@@ -486,7 +498,6 @@ class JsonSchemaGenerator:
         property_docstrings = get_class_property_docstrings(
             typ, self.options.property_description_fun
         )
-
         properties: Dict[str, Schema] = {}
         required: List[str] = []
         for property_name, property_type in get_class_properties(typ):
@@ -502,11 +513,27 @@ class JsonSchemaGenerator:
             else:
                 output_name = property_name
 
+            # check if there's additional field annotation
+            json_schema_extra = None
+            if "model_fields" in members:
+                f = members["model_fields"]
+                if output_name in f:
+                    json_schema_extra = f[output_name].json_schema_extra
+
+            if json_schema_extra:
+                from rich.pretty import pprint
+
+                pprint(json_schema_extra)
+
             if is_type_optional(property_type):
                 optional_type: type = unwrap_optional_type(property_type)
-                property_def = self.type_to_schema(optional_type)
+                property_def = self.type_to_schema(
+                    optional_type, json_schema_extra=json_schema_extra
+                )
             else:
-                property_def = self.type_to_schema(property_type)
+                property_def = self.type_to_schema(
+                    property_type, json_schema_extra=json_schema_extra
+                )
                 required.append(output_name)
 
             # check if attribute has a default value initializer
@@ -537,6 +564,8 @@ class JsonSchemaGenerator:
 
             properties[output_name] = property_def
 
+        # print("properties", properties)
+
         schema = {"type": "object"}
         if len(properties) > 0:
             schema["properties"] = typing.cast(JsonType, properties)
diff --git a/docs/resources/llama-stack-spec.html b/docs/resources/llama-stack-spec.html
index 8f69d152e..5b3771340 100644
--- a/docs/resources/llama-stack-spec.html
+++ b/docs/resources/llama-stack-spec.html
@@ -2451,6 +2451,7 @@
                             },
                             "data": {
                                 "type": "string",
+                                "contentEncoding": "base64",
                                 "description": "base64 encoded image data as string"
                             }
                         },
diff --git a/docs/resources/llama-stack-spec.yaml b/docs/resources/llama-stack-spec.yaml
index 937729e2b..01232c001 100644
--- a/docs/resources/llama-stack-spec.yaml
+++ b/docs/resources/llama-stack-spec.yaml
@@ -1478,6 +1478,7 @@ components:
                 Note that URL could have length limits.
             data:
               type: string
+              contentEncoding: base64
               description: base64 encoded image data as string
           additionalProperties: false
           description: >-
diff --git a/llama_stack/apis/common/content_types.py b/llama_stack/apis/common/content_types.py
index 04a5913d6..d005c2781 100644
--- a/llama_stack/apis/common/content_types.py
+++ b/llama_stack/apis/common/content_types.py
@@ -28,8 +28,8 @@ class _URLOrData(BaseModel):
 
     url: Optional[URL] = None
     # data is a base64 encoded string
-    # TODO: annotate with contentEncoding="base64" in OpenAPI schema
-    data: Optional[str] = None
+    # TODO: generate contentEncoding="base64" in OpenAPI schema
+    data: Optional[str] = Field(contentEncoding="base64", default=None)
 
     @model_validator(mode="before")
     @classmethod