Merge branch 'pr1573' into api_2

llama_stack/apis/agents/agents.py

@@ -370,7 +370,7 @@ class AgentTurnResumeRequest(BaseModel):
     agent_id: str
     session_id: str
     turn_id: str
-    tool_responses: Union[List[ToolResponse], List[ToolResponseMessage]]
+    tool_responses: List[ToolResponse]
     stream: Optional[bool] = False
 
 
@@ -449,7 +449,7 @@ class Agents(Protocol):
         agent_id: str,
         session_id: str,
         turn_id: str,
-        tool_responses: Union[List[ToolResponse], List[ToolResponseMessage]],
+        tool_responses: List[ToolResponse],
         stream: Optional[bool] = False,
     ) -> Union[Turn, AsyncIterator[AgentTurnResponseStreamChunk]]:
         """Resume an agent turn with executed tool call responses.
@@ -460,7 +460,6 @@ class Agents(Protocol):
         :param session_id: The ID of the session to resume.
         :param turn_id: The ID of the turn to resume.
         :param tool_responses: The tool call responses to resume the turn with.
-            NOTE: ToolResponseMessage will be deprecated. Use ToolResponse.
         :param stream: Whether to stream the response.
         :returns: A Turn object if stream is False, otherwise an AsyncIterator of AgentTurnResponseStreamChunk objects.
         """

llama_stack/apis/datasets/datasets.py

@@ -13,10 +13,10 @@ from llama_stack.apis.resource import Resource, ResourceType
 from llama_stack.schema_utils import json_schema_type, register_schema, webmethod
 
 
-class Schema(Enum):
+class DatasetPurpose(Enum):
     """
-    Schema of the dataset. Each type has a different column format.
-    :cvar messages: The dataset contains messages used for post-training. Examples:
+    Purpose of the dataset. Each type has a different column format.
+    :cvar post-training/messages: The dataset contains messages used for post-training. Examples:
         {
             "messages": [
                 {"role": "user", "content": "Hello, world!"},
@@ -25,11 +25,19 @@ class Schema(Enum):
         }
     """
 
-    messages = "messages"
+    post_training_messages = "post-training/messages"
+    eval_question_answer = "eval/question-answer"
+
     # TODO: add more schemas here
 
 
 class DatasetType(Enum):
+    """
+    Type of the dataset source.
+    :cvar huggingface: The dataset is stored in Huggingface.
+    :cvar uri: The dataset can be obtained from a URI. 
+    :cvar rows: The dataset is stored in rows. 
+    """
     huggingface = "huggingface"
     uri = "uri"
     rows = "rows"
@@ -37,19 +45,36 @@ class DatasetType(Enum):
 
 @json_schema_type
 class URIDataSource(BaseModel):
+    """A dataset that can be obtained from a URI.
+    :param uri: The dataset can be obtained from a URI. E.g.
+        - "https://mywebsite.com/mydata.jsonl"
+        - "lsfs://mydata.jsonl"
+        - "data:csv;base64,{base64_content}"
+    """
     type: Literal["uri"] = "uri"
     uri: str
 
 
 @json_schema_type
 class HuggingfaceDataSource(BaseModel):
+    """A dataset stored in Huggingface.
+    :param path: The path to the dataset in Huggingface. E.g.
+        - "llamastack/simpleqa"
+    :param params: The parameters for the dataset.
+    """
     type: Literal["huggingface"] = "huggingface"
-    dataset_path: str
+    path: str
     params: Dict[str, Any]
 
 
 @json_schema_type
 class RowsDataSource(BaseModel):
+    """A dataset stored in rows.
+    :param rows: The dataset is stored in rows. E.g.
+        - [
+            {"messages": [{"role": "user", "content": "Hello, world!"}, {"role": "assistant", "content": "Hello, world!"}]}
+        ]
+    """
     type: Literal["rows"] = "rows"
     rows: List[Dict[str, Any]]
 
@@ -64,8 +89,11 @@ DataSource = register_schema(
 
 
 class CommonDatasetFields(BaseModel):
-    schema: Schema
-    data_source: DataSource
+    """
+    Common fields for a dataset.
+    """
+    purpose: DatasetPurpose
+    source: DataSource
     metadata: Dict[str, Any] = Field(
         default_factory=dict,
         description="Any additional metadata for this dataset",
@@ -99,17 +127,18 @@ class Datasets(Protocol):
     @webmethod(route="/datasets", method="POST")
     async def register_dataset(
         self,
-        schema: Schema,
-        data_source: DataSource,
+        purpose: DatasetPurpose,
+        source: DataSource,
         metadata: Optional[Dict[str, Any]] = None,
         dataset_id: Optional[str] = None,
     ) -> Dataset:
         """
         Register a new dataset.
 
-        :param schema: The schema format of the dataset. One of
-            - messages: The dataset contains a messages column with list of messages for post-training.
-        :param data_source: The data source of the dataset. Examples:
+        :param purpose: The purpose of the dataset. One of
+            - "post-training/messages": The dataset contains a messages column with list of messages for post-training.
+            - "eval/question-answer": The dataset contains a question and answer column.
+        :param source: The data source of the dataset. Examples:
            - {
                "type": "uri",
                "uri": "https://mywebsite.com/mydata.jsonl"

llama_stack/apis/inference/inference.py

@@ -285,7 +285,7 @@ class CompletionRequest(BaseModel):
 
 
 @json_schema_type
-class CompletionResponse(BaseModel):
+class CompletionResponse(MetricResponseMixin):
     """Response from a completion request.
 
     :param content: The generated completion text
@@ -299,7 +299,7 @@ class CompletionResponse(BaseModel):
 
 
 @json_schema_type
-class CompletionResponseStreamChunk(BaseModel):
+class CompletionResponseStreamChunk(MetricResponseMixin):
     """A chunk of a streamed completion response.
 
     :param delta: New content generated since last chunk. This can be one or more tokens.
@@ -368,7 +368,7 @@ class ChatCompletionRequest(BaseModel):
 
 
 @json_schema_type
-class ChatCompletionResponseStreamChunk(MetricResponseMixin, BaseModel):
+class ChatCompletionResponseStreamChunk(MetricResponseMixin):
     """A chunk of a streamed chat completion response.
 
     :param event: The event containing the new content
@@ -378,7 +378,7 @@ class ChatCompletionResponseStreamChunk(MetricResponseMixin, BaseModel):
 
 
 @json_schema_type
-class ChatCompletionResponse(MetricResponseMixin, BaseModel):
+class ChatCompletionResponse(MetricResponseMixin):
     """Response from a chat completion request.
 
     :param completion_message: The complete response message

llama_stack/apis/telemetry/telemetry.py

@@ -96,6 +96,13 @@ class MetricEvent(EventCommon):
     unit: str
 
 
+@json_schema_type
+class MetricInResponse(BaseModel):
+    metric: str
+    value: Union[int, float]
+    unit: Optional[str] = None
+
+
 # This is a short term solution to allow inference API to return metrics
 # The ideal way to do this is to have a way for all response types to include metrics
 # and all metric events logged to the telemetry API to be inlcuded with the response
@@ -117,7 +124,7 @@ class MetricEvent(EventCommon):
 
 
 class MetricResponseMixin(BaseModel):
-    metrics: Optional[List[MetricEvent]] = None
+    metrics: Optional[List[MetricInResponse]] = None
 
 
 @json_schema_type