add metrics to all response types

docs/_static/llama-stack-spec.html

@@ -3096,6 +3096,12 @@
                         "items": {
                             "$ref": "#/components/schemas/ChatCompletionResponse"
                         }
+                    },
+                    "metrics": {
+                        "type": "array",
+                        "items": {
+                            "$ref": "#/components/schemas/MetricEvent"
+                        }
                     }
                 },
                 "additionalProperties": false,
@@ -3116,6 +3122,12 @@
                             "$ref": "#/components/schemas/TokenLogProbs"
                         },
                         "description": "Optional log probabilities for generated tokens"
+                    },
+                    "metrics": {
+                        "type": "array",
+                        "items": {
+                            "$ref": "#/components/schemas/MetricEvent"
+                        }
                     }
                 },
                 "additionalProperties": false,
@@ -3185,6 +3197,12 @@
                         "items": {
                             "$ref": "#/components/schemas/CompletionResponse"
                         }
+                    },
+                    "metrics": {
+                        "type": "array",
+                        "items": {
+                            "$ref": "#/components/schemas/MetricEvent"
+                        }
                     }
                 },
                 "additionalProperties": false,
@@ -3214,6 +3232,12 @@
                             "$ref": "#/components/schemas/TokenLogProbs"
                         },
                         "description": "Optional log probabilities for generated tokens"
+                    },
+                    "metrics": {
+                        "type": "array",
+                        "items": {
+                            "$ref": "#/components/schemas/MetricEvent"
+                        }
                     }
                 },
                 "additionalProperties": false,
@@ -3391,6 +3415,12 @@
                     "event": {
                         "$ref": "#/components/schemas/ChatCompletionResponseEvent",
                         "description": "The event containing the new content"
+                    },
+                    "metrics": {
+                        "type": "array",
+                        "items": {
+                            "$ref": "#/components/schemas/MetricEvent"
+                        }
                     }
                 },
                 "additionalProperties": false,
@@ -3556,6 +3586,12 @@
                             "$ref": "#/components/schemas/TokenLogProbs"
                         },
                         "description": "Optional log probabilities for generated tokens"
+                    },
+                    "metrics": {
+                        "type": "array",
+                        "items": {
+                            "$ref": "#/components/schemas/MetricEvent"
+                        }
                     }
                 },
                 "additionalProperties": false,
@@ -3794,6 +3830,12 @@
                 "properties": {
                     "agent_id": {
                         "type": "string"
+                    },
+                    "metrics": {
+                        "type": "array",
+                        "items": {
+                            "$ref": "#/components/schemas/MetricEvent"
+                        }
                     }
                 },
                 "additionalProperties": false,
@@ -3818,6 +3860,12 @@
                 "properties": {
                     "session_id": {
                         "type": "string"
+                    },
+                    "metrics": {
+                        "type": "array",
+                        "items": {
+                            "$ref": "#/components/schemas/MetricEvent"
+                        }
                     }
                 },
                 "additionalProperties": false,
@@ -4108,6 +4156,12 @@
                     },
                     "content": {
                         "$ref": "#/components/schemas/InterleavedContent"
+                    },
+                    "metrics": {
+                        "type": "array",
+                        "items": {
+                            "$ref": "#/components/schemas/MetricEvent"
+                        }
                     }
                 },
                 "additionalProperties": false,
@@ -4418,6 +4472,12 @@
                 "properties": {
                     "event": {
                         "$ref": "#/components/schemas/AgentTurnResponseEvent"
+                    },
+                    "metrics": {
+                        "type": "array",
+                        "items": {
+                            "$ref": "#/components/schemas/MetricEvent"
+                        }
                     }
                 },
                 "additionalProperties": false,
@@ -4495,6 +4555,12 @@
                             }
                         },
                         "description": "List of embedding vectors, one per input content. Each embedding is a list of floats. The dimensionality of the embedding is model-specific; you can check model metadata using /models/{model_id}"
+                    },
+                    "metrics": {
+                        "type": "array",
+                        "items": {
+                            "$ref": "#/components/schemas/MetricEvent"
+                        }
                     }
                 },
                 "additionalProperties": false,
@@ -4822,6 +4888,12 @@
                         "additionalProperties": {
                             "$ref": "#/components/schemas/ScoringResult"
                         }
+                    },
+                    "metrics": {
+                        "type": "array",
+                        "items": {
+                            "$ref": "#/components/schemas/MetricEvent"
+                        }
                     }
                 },
                 "additionalProperties": false,
@@ -4949,6 +5021,12 @@
                                 "memory_retrieval": "#/components/schemas/MemoryRetrievalStep"
                             }
                         }
+                    },
+                    "metrics": {
+                        "type": "array",
+                        "items": {
+                            "$ref": "#/components/schemas/MetricEvent"
+                        }
                     }
                 },
                 "additionalProperties": false,
@@ -5625,6 +5703,12 @@
                         "additionalProperties": {
                             "$ref": "#/components/schemas/SpanWithStatus"
                         }
+                    },
+                    "metrics": {
+                        "type": "array",
+                        "items": {
+                            "$ref": "#/components/schemas/MetricEvent"
+                        }
                     }
                 },
                 "additionalProperties": false,
@@ -5803,6 +5887,12 @@
                         "items": {
                             "$ref": "#/components/schemas/Checkpoint"
                         }
+                    },
+                    "metrics": {
+                        "type": "array",
+                        "items": {
+                            "$ref": "#/components/schemas/MetricEvent"
+                        }
                     }
                 },
                 "additionalProperties": false,
@@ -5872,6 +5962,12 @@
                         "items": {
                             "$ref": "#/components/schemas/Checkpoint"
                         }
+                    },
+                    "metrics": {
+                        "type": "array",
+                        "items": {
+                            "$ref": "#/components/schemas/MetricEvent"
+                        }
                     }
                 },
                 "additionalProperties": false,
@@ -7081,6 +7177,12 @@
                         "items": {
                             "type": "number"
                         }
+                    },
+                    "metrics": {
+                        "type": "array",
+                        "items": {
+                            "$ref": "#/components/schemas/MetricEvent"
+                        }
                     }
                 },
                 "additionalProperties": false,
@@ -7145,6 +7247,12 @@
                         "items": {
                             "$ref": "#/components/schemas/Span"
                         }
+                    },
+                    "metrics": {
+                        "type": "array",
+                        "items": {
+                            "$ref": "#/components/schemas/MetricEvent"
+                        }
                     }
                 },
                 "additionalProperties": false,
@@ -7160,6 +7268,12 @@
                         "items": {
                             "$ref": "#/components/schemas/Trace"
                         }
+                    },
+                    "metrics": {
+                        "type": "array",
+                        "items": {
+                            "$ref": "#/components/schemas/MetricEvent"
+                        }
                     }
                 },
                 "additionalProperties": false,
@@ -7536,6 +7650,12 @@
                 "properties": {
                     "violation": {
                         "$ref": "#/components/schemas/SafetyViolation"
+                    },
+                    "metrics": {
+                        "type": "array",
+                        "items": {
+                            "$ref": "#/components/schemas/MetricEvent"
+                        }
                     }
                 },
                 "additionalProperties": false
@@ -7628,6 +7748,12 @@
                         "additionalProperties": {
                             "$ref": "#/components/schemas/ScoringResult"
                         }
+                    },
+                    "metrics": {
+                        "type": "array",
+                        "items": {
+                            "$ref": "#/components/schemas/MetricEvent"
+                        }
                     }
                 },
                 "additionalProperties": false,
@@ -7676,6 +7802,12 @@
                         "additionalProperties": {
                             "$ref": "#/components/schemas/ScoringResult"
                         }
+                    },
+                    "metrics": {
+                        "type": "array",
+                        "items": {
+                            "$ref": "#/components/schemas/MetricEvent"
+                        }
                     }
                 },
                 "additionalProperties": false,
@@ -7931,6 +8063,12 @@
                                 }
                             ]
                         }
+                    },
+                    "metrics": {
+                        "type": "array",
+                        "items": {
+                            "$ref": "#/components/schemas/MetricEvent"
+                        }
                     }
                 },
                 "additionalProperties": false,

docs/_static/llama-stack-spec.yaml

@@ -1919,6 +1919,10 @@ components:
           type: array
           items:
             $ref: '#/components/schemas/ChatCompletionResponse'
+        metrics:
+          type: array
+          items:
+            $ref: '#/components/schemas/MetricEvent'
       additionalProperties: false
       required:
         - batch
@@ -1934,6 +1938,10 @@ components:
             $ref: '#/components/schemas/TokenLogProbs'
           description: >-
             Optional log probabilities for generated tokens
+        metrics:
+          type: array
+          items:
+            $ref: '#/components/schemas/MetricEvent'
       additionalProperties: false
       required:
         - completion_message
@@ -1984,6 +1992,10 @@ components:
           type: array
           items:
             $ref: '#/components/schemas/CompletionResponse'
+        metrics:
+          type: array
+          items:
+            $ref: '#/components/schemas/MetricEvent'
       additionalProperties: false
       required:
         - batch
@@ -2006,6 +2018,10 @@ components:
             $ref: '#/components/schemas/TokenLogProbs'
           description: >-
             Optional log probabilities for generated tokens
+        metrics:
+          type: array
+          items:
+            $ref: '#/components/schemas/MetricEvent'
       additionalProperties: false
       required:
         - content
@@ -2176,6 +2192,10 @@ components:
         event:
           $ref: '#/components/schemas/ChatCompletionResponseEvent'
           description: The event containing the new content
+        metrics:
+          type: array
+          items:
+            $ref: '#/components/schemas/MetricEvent'
       additionalProperties: false
       required:
         - event
@@ -2303,6 +2323,10 @@ components:
             $ref: '#/components/schemas/TokenLogProbs'
           description: >-
             Optional log probabilities for generated tokens
+        metrics:
+          type: array
+          items:
+            $ref: '#/components/schemas/MetricEvent'
       additionalProperties: false
       required:
         - delta
@@ -2449,6 +2473,10 @@ components:
       properties:
         agent_id:
           type: string
+        metrics:
+          type: array
+          items:
+            $ref: '#/components/schemas/MetricEvent'
       additionalProperties: false
       required:
         - agent_id
@@ -2465,6 +2493,10 @@ components:
       properties:
         session_id:
           type: string
+        metrics:
+          type: array
+          items:
+            $ref: '#/components/schemas/MetricEvent'
       additionalProperties: false
       required:
         - session_id
@@ -2653,6 +2685,10 @@ components:
             - type: string
         content:
           $ref: '#/components/schemas/InterleavedContent'
+        metrics:
+          type: array
+          items:
+            $ref: '#/components/schemas/MetricEvent'
       additionalProperties: false
       required:
         - call_id
@@ -2847,6 +2883,10 @@ components:
       properties:
         event:
           $ref: '#/components/schemas/AgentTurnResponseEvent'
+        metrics:
+          type: array
+          items:
+            $ref: '#/components/schemas/MetricEvent'
       additionalProperties: false
       required:
         - event
@@ -2910,6 +2950,10 @@ components:
             List of embedding vectors, one per input content. Each embedding is a
             list of floats. The dimensionality of the embedding is model-specific;
             you can check model metadata using /models/{model_id}
+        metrics:
+          type: array
+          items:
+            $ref: '#/components/schemas/MetricEvent'
       additionalProperties: false
       required:
         - embeddings
@@ -3117,6 +3161,10 @@ components:
           type: object
           additionalProperties:
             $ref: '#/components/schemas/ScoringResult'
+        metrics:
+          type: array
+          items:
+            $ref: '#/components/schemas/MetricEvent'
       additionalProperties: false
       required:
         - generations
@@ -3188,6 +3236,10 @@ components:
               tool_execution: '#/components/schemas/ToolExecutionStep'
               shield_call: '#/components/schemas/ShieldCallStep'
               memory_retrieval: '#/components/schemas/MemoryRetrievalStep'
+        metrics:
+          type: array
+          items:
+            $ref: '#/components/schemas/MetricEvent'
       additionalProperties: false
       required:
         - step
@@ -3602,6 +3654,10 @@ components:
           type: object
           additionalProperties:
             $ref: '#/components/schemas/SpanWithStatus'
+        metrics:
+          type: array
+          items:
+            $ref: '#/components/schemas/MetricEvent'
       additionalProperties: false
       required:
         - data
@@ -3714,6 +3770,10 @@ components:
           type: array
           items:
             $ref: '#/components/schemas/Checkpoint'
+        metrics:
+          type: array
+          items:
+            $ref: '#/components/schemas/MetricEvent'
       additionalProperties: false
       required:
         - job_uuid
@@ -3756,6 +3816,10 @@ components:
           type: array
           items:
             $ref: '#/components/schemas/Checkpoint'
+        metrics:
+          type: array
+          items:
+            $ref: '#/components/schemas/MetricEvent'
       additionalProperties: false
       required:
         - job_uuid
@@ -4514,6 +4578,10 @@ components:
           type: array
           items:
             type: number
+        metrics:
+          type: array
+          items:
+            $ref: '#/components/schemas/MetricEvent'
       additionalProperties: false
       required:
         - chunks
@@ -4552,6 +4620,10 @@ components:
           type: array
           items:
             $ref: '#/components/schemas/Span'
+        metrics:
+          type: array
+          items:
+            $ref: '#/components/schemas/MetricEvent'
       additionalProperties: false
       required:
         - data
@@ -4562,6 +4634,10 @@ components:
           type: array
           items:
             $ref: '#/components/schemas/Trace'
+        metrics:
+          type: array
+          items:
+            $ref: '#/components/schemas/MetricEvent'
       additionalProperties: false
       required:
         - data
@@ -4776,6 +4852,10 @@ components:
       properties:
         violation:
           $ref: '#/components/schemas/SafetyViolation'
+        metrics:
+          type: array
+          items:
+            $ref: '#/components/schemas/MetricEvent'
       additionalProperties: false
     SaveSpansToDatasetRequest:
       type: object
@@ -4829,6 +4909,10 @@ components:
           type: object
           additionalProperties:
             $ref: '#/components/schemas/ScoringResult'
+        metrics:
+          type: array
+          items:
+            $ref: '#/components/schemas/MetricEvent'
       additionalProperties: false
       required:
         - results
@@ -4859,6 +4943,10 @@ components:
           type: object
           additionalProperties:
             $ref: '#/components/schemas/ScoringResult'
+        metrics:
+          type: array
+          items:
+            $ref: '#/components/schemas/MetricEvent'
       additionalProperties: false
       required:
         - results
@@ -5008,6 +5096,10 @@ components:
               - type: string
               - type: array
               - type: object
+        metrics:
+          type: array
+          items:
+            $ref: '#/components/schemas/MetricEvent'
       additionalProperties: false
       required:
         - synthetic_data

docs/openapi_generator/pyopenapi/generator.py

@@ -55,6 +55,7 @@ from .specification import (
     Tag,
     TagGroup,
 )
+from llama_stack.apis.inference import ChatCompletionResponse
 
 register_schema(
     ipaddress.IPv4Address,
@@ -106,6 +107,27 @@ class SchemaBuilder:
         self.schema_generator = schema_generator
         self.schemas = {}
 
+    def _is_response_type(self, type_or_name: Union[str, type]) -> bool:
+        """Helper method to check if a type or type name is a response type.
+        
+        Args:
+            type_or_name: Either a string type name or a type object
+            
+        Returns:
+            bool: True if the type represents a response, False otherwise
+        """
+        name = type_or_name if isinstance(type_or_name, str) else type_or_name.__name__
+        return not name.startswith('List') and any(name.endswith(suffix) for suffix in ('Response', 'ResponseStreamChunk'))
+
+    def _add_metrics_to_schema(self, schema: Dict[str, Any]) -> None:
+        """Helper method to add metrics field to a schema"""
+        if 'properties' not in schema:
+            schema['properties'] = {}
+        schema['properties']['metrics'] = {
+            'type': 'array',
+            'items': {'$ref': '#/components/schemas/MetricEvent'}
+        }
+
     def classdef_to_schema(self, typ: type) -> Schema:
         """
         Converts a type to a JSON schema.
@@ -114,7 +136,18 @@ class SchemaBuilder:
 
         type_schema, type_definitions = self.schema_generator.classdef_to_schema(typ)
 
-        # append schema to list of known schemas, to be used in OpenAPI's Components Object section
+        # Add metrics field to all response schemas
+        if self._is_response_type(typ):
+            if isinstance(type_schema, dict) and '$ref' in type_schema:
+                # If it's a reference, modify the schema in type_definitions
+                ref_name = type_schema['$ref'].split('/')[-1]
+                if ref_name in type_definitions:
+                    self._add_metrics_to_schema(type_definitions[ref_name])
+            else:
+                # Direct schema case
+                self._add_metrics_to_schema(type_schema)
+
+        # Register all schemas, including modified ones
         for ref, schema in type_definitions.items():
             self._add_ref(ref, schema)
 
@@ -122,6 +155,14 @@ class SchemaBuilder:
 
     def classdef_to_named_schema(self, name: str, typ: type) -> Schema:
         schema = self.classdef_to_schema(typ)
+        # If this is a Response type, ensure the registered schema has metrics
+        if self._is_response_type(typ):
+            if isinstance(schema, dict) and '$ref' in schema:
+                ref_name = schema['$ref'].split('/')[-1]
+                if ref_name in self.schemas:
+                    self._add_metrics_to_schema(self.schemas[ref_name])
+            else:
+                self._add_metrics_to_schema(schema)
         self._add_ref(name, schema)
         return schema
 
@@ -154,6 +195,10 @@ class SchemaBuilder:
 
     def _add_ref(self, type_name: str, type_schema: Schema) -> None:
         if type_name not in self.schemas:
+            # If adding a Response type, ensure it has metrics field
+            if self._is_response_type(type_name):
+                if isinstance(type_schema, dict):
+                    self._add_metrics_to_schema(type_schema)
             self.schemas[type_name] = type_schema