Simplify model parameter REST API (#1475)

* Simplifying the model parameter rest api

* implemented pydantic model and fixed type inputs

* fixing functionality of pydantic model

* simple code revision

* Fixing parameter response passing

* only using parameters in the response

Author: jonahjung22

packages/jupyter-ai/jupyter_ai/model_providers/parameter_schemas.py

@@ -1,19 +1,16 @@
 from __future__ import annotations
-from typing import TYPE_CHECKING, cast
-
-if TYPE_CHECKING:
-    from typing import Any
-
+from typing import Literal, Any
+from pydantic import BaseModel
 
 PARAMETER_SCHEMAS: dict[str, dict[str, Any]] = {
     "temperature": {
-        "type": "number",
+        "type": "float",
         "min": 0,
         "max": 2,
         "description": "Controls randomness in the output. Lower values make it more focused and deterministic."
     },
     "top_p": {
-        "type": "number", 
+        "type": "float", 
         "min": 0,
         "max": 1,
         "description": "Nucleus sampling parameter. Consider tokens with top_p probability mass."
@@ -69,13 +66,13 @@
     # },
 
     "presence_penalty": {
-        "type": "number",
+        "type": "float",
         "min": -2,
         "max": 2,
         "description": "Penalize new tokens based on whether they appear in the text so far."
     },
     "frequency_penalty": {
-        "type": "number",
+        "type": "float",
         "min": -2,
         "max": 2,
         "description": "Penalize new tokens based on their frequency in the text so far."
@@ -110,26 +107,35 @@
     }
 }
 
-def get_parameter_schema(param_name: str) -> dict[str, Any]:
+class ParameterSchema(BaseModel):
+    """Pydantic model for parameter schema definition."""
+    type: Literal['boolean', 'integer', 'float', 'string', 'array', 'object']
+    description: str
+
+class GetModelParametersResponse(BaseModel):
+    """Pydantic model for GET model parameters response."""
+    parameters: dict[str, ParameterSchema]
+    parameter_names: list[str]
+
+class UpdateModelParametersResponse(BaseModel):
+    """Pydantic model for PUT model parameters response."""
+    parameters: dict[str, Any]
+
+def get_parameter_schema(param_name: str) -> ParameterSchema:
     """
     Get the schema for a specific parameter.
-
-    TODO: Define a Pydantic model for the parameter schema, e.g.
-    `ParameterSchema`. Update the return type annotation to `ParameterSchema`.
     """
     schema = PARAMETER_SCHEMAS.get(param_name)
     if schema is None:
-        return {
-            "type": "string",
-            "description": f"Parameter {param_name} (schema not defined)"
-        }
-    return schema
+        return ParameterSchema(
+            type="string",
+            description=f"Parameter {param_name} (schema not defined)"
+        )
+    return ParameterSchema(**schema)
 
-def get_parameters_with_schemas(param_names: list[str]) -> dict[str, Any]:
+def get_parameters_with_schemas(param_names: list[str]) -> dict[str, ParameterSchema]:
     """
     Get schemas for a list of parameter names.
-
-    TODO: Update the return type annotation to `dict[str, ParameterSchema]`.
     """
     return {
         name: get_parameter_schema(name) 

packages/jupyter-ai/jupyter_ai/model_providers/parameters_rest_api.py

@@ -3,7 +3,7 @@
 import json
 
 from litellm.litellm_core_utils.get_supported_openai_params import get_supported_openai_params
-from .parameter_schemas import get_parameters_with_schemas, coerce_parameter_value
+from .parameter_schemas import get_parameters_with_schemas, coerce_parameter_value, GetModelParametersResponse, UpdateModelParametersResponse
 from ..config_manager import ConfigManager
 from ..config import UpdateConfigRequest
 
@@ -56,23 +56,14 @@ def get(self):
             # Get parameter schemas with types, defaults, and descriptions
             parameters_with_schemas = get_parameters_with_schemas(parameter_names)
 
-            # TODO: Define the response type as a Pydantic model to prevent
-            # breaking API changes, e.g. `GetModelParametersResponse`.
-            #
-            # TODO: replace 'number' with 'float' in parameter schemas for
-            # clarity. 'number' is ambiguous as to whether it is an integer or
-            # float. Make sure to update the frontend type.
-            #
-            # TODO: Drop 'count' from the response (and the corresponding types
-            # across the frontend & backend).
-            response = {
-                "parameters": parameters_with_schemas,
-                "parameter_names": parameter_names,
-                "count": len(parameter_names)
-            }
+            # Create Pydantic response model
+            response = GetModelParametersResponse(
+                parameters=parameters_with_schemas,
+                parameter_names=parameter_names
+            )
 
             self.set_header("Content-Type", "application/json")
-            self.finish(json.dumps(response))
+            self.finish(response.model_dump_json())
 
         except Exception as e:
             self.log.exception("Failed to get model parameters")
@@ -127,21 +118,13 @@ def put(self):
             )
             config_manager.update_config(update_request)
 
-            # TODO: Determine what a response could be used for, and remove the
-            # response if it is unnecessary. Right now the response body is
-            # ignored by the frontend.
-            #
-            # TODO: Define the response type as a Pydantic model to prevent
-            # breaking API changes, e.g. `UpdateModelParametersResponse`.
-            response = {
-                "status": "success",
-                "message": f"Parameters saved for model {model_id}",
-                "model_id": model_id,
-                "parameters": coerced_parameters,
-            }
+            # Create Pydantic response model for API compatibility
+            response = UpdateModelParametersResponse(
+                parameters=coerced_parameters
+            )
 
             self.set_header("Content-Type", "application/json")
-            self.finish(json.dumps(response))
+            self.finish(response.model_dump_json())
 
         except json.JSONDecodeError:
             raise HTTPError(400, "Invalid JSON in request body")

packages/jupyter-ai/src/handler.ts

@@ -200,20 +200,16 @@ export namespace AiService {
   export type GetModelParametersResponse = {
     parameters: Record<string, ParameterSchema>;
     parameter_names: string[];
-    count: number;
   };
 
   export type ParameterSchema = {
-    type: 'boolean' | 'integer' | 'number' | 'string' | 'array' | 'object';
+    type: 'boolean' | 'integer' | 'float' | 'string' | 'array' | 'object';
     description: string;
     min?: number;
     max?: number;
   };
 
   export type UpdateModelParametersResponse = {
-    status: string;
-    message: string;
-    model_id: string;
     parameters: Record<string, any>;
   };