From 5388d212bcc5822f07a9b8fda92f809a8509397c Mon Sep 17 00:00:00 2001
From: Yukt Mitash <yuktmitash17@meta.com>
Date: Mon, 10 Mar 2025 14:31:06 -0700
Subject: [PATCH] Updated the api spec to include COT object

---
 docs/_static/llama-stack-spec.html       | 28 +++++++++++++++++++++++-
 docs/_static/llama-stack-spec.yaml       | 20 +++++++++++++++++
 llama_stack/apis/common/content_types.py | 14 +++++++++++-
 3 files changed, 60 insertions(+), 2 deletions(-)

diff --git a/docs/_static/llama-stack-spec.html b/docs/_static/llama-stack-spec.html
index 1a8169090..5af0dac68 100644
--- a/docs/_static/llama-stack-spec.html
+++ b/docs/_static/llama-stack-spec.html
@@ -3706,6 +3706,28 @@
                 ],
                 "title": "AppendRowsRequest"
             },
+            "ChainOfThoughtContentItem": {
+                "type": "object",
+                "properties": {
+                    "type": {
+                        "type": "string",
+                        "const": "chainOfThought",
+                        "default": "chainOfThought",
+                        "description": "Discriminator type of the content item. Always \"ChainOfThought\""
+                    },
+                    "text": {
+                        "type": "string",
+                        "description": "Chain of Thought content"
+                    }
+                },
+                "additionalProperties": false,
+                "required": [
+                    "type",
+                    "text"
+                ],
+                "title": "ChainOfThoughtContentItem",
+                "description": "A Chain of Thought content item"
+            },
             "CompletionMessage": {
                 "type": "object",
                 "properties": {
@@ -3861,13 +3883,17 @@
                     },
                     {
                         "$ref": "#/components/schemas/TextContentItem"
+                    },
+                    {
+                        "$ref": "#/components/schemas/ChainOfThoughtContentItem"
                     }
                 ],
                 "discriminator": {
                     "propertyName": "type",
                     "mapping": {
                         "image": "#/components/schemas/ImageContentItem",
-                        "text": "#/components/schemas/TextContentItem"
+                        "text": "#/components/schemas/TextContentItem",
+                        "chainOfThought": "#/components/schemas/ChainOfThoughtContentItem"
                     }
                 }
             },
diff --git a/docs/_static/llama-stack-spec.yaml b/docs/_static/llama-stack-spec.yaml
index d6001c00d..c42c1fea5 100644
--- a/docs/_static/llama-stack-spec.yaml
+++ b/docs/_static/llama-stack-spec.yaml
@@ -2518,6 +2518,24 @@ components:
         - dataset_id
         - rows
       title: AppendRowsRequest
+    ChainOfThoughtContentItem:
+      type: object
+      properties:
+        type:
+          type: string
+          const: chainOfThought
+          default: chainOfThought
+          description: >-
+            Discriminator type of the content item. Always "ChainOfThought"
+        text:
+          type: string
+          description: Chain of Thought content
+      additionalProperties: false
+      required:
+        - type
+        - text
+      title: ChainOfThoughtContentItem
+      description: A Chain of Thought content item
     CompletionMessage:
       type: object
       properties:
@@ -2637,11 +2655,13 @@ components:
       oneOf:
         - $ref: '#/components/schemas/ImageContentItem'
         - $ref: '#/components/schemas/TextContentItem'
+        - $ref: '#/components/schemas/ChainOfThoughtContentItem'
       discriminator:
         propertyName: type
         mapping:
           image: '#/components/schemas/ImageContentItem'
           text: '#/components/schemas/TextContentItem'
+          chainOfThought: '#/components/schemas/ChainOfThoughtContentItem'
     JsonSchemaResponseFormat:
       type: object
       properties:
diff --git a/llama_stack/apis/common/content_types.py b/llama_stack/apis/common/content_types.py
index 0d0afa894..283359ebc 100644
--- a/llama_stack/apis/common/content_types.py
+++ b/llama_stack/apis/common/content_types.py
@@ -50,6 +50,18 @@ class ImageContentItem(BaseModel):
     image: _URLOrData
 
 
+@json_schema_type
+class ChainOfThoughtContentItem(BaseModel):
+    """A Chain of Thought content item
+
+    :param type: Discriminator type of the content item. Always "ChainOfThought"
+    :param text: Chain of Thought content
+    """
+
+    type: Literal["chainOfThought"] = "chainOfThought"
+    text: str
+
+
 @json_schema_type
 class TextContentItem(BaseModel):
     """A text content item
@@ -65,7 +77,7 @@ class TextContentItem(BaseModel):
 # other modalities can be added here
 InterleavedContentItem = register_schema(
     Annotated[
-        Union[ImageContentItem, TextContentItem],
+        Union[ImageContentItem, TextContentItem, ChainOfThoughtContentItem],
         Field(discriminator="type"),
     ],
     name="InterleavedContentItem",