fixing functionality of pydantic model

Author: jonahjung22

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

@@ -1,9 +1,9 @@
 from __future__ import annotations
-from typing import TYPE_CHECKING, cast, Literal, Optional
+from typing import TYPE_CHECKING, cast, Literal, Optional, Any
 from pydantic import BaseModel
 
 if TYPE_CHECKING:
-    from typing import Any
+    pass
 
 
 PARAMETER_SCHEMAS: dict[str, dict[str, Any]] = {
@@ -116,15 +116,19 @@ class ParameterSchema(BaseModel):
     """Pydantic model for parameter schema definition."""
     type: Literal['boolean', 'integer', 'float', 'string', 'array', 'object']
     description: str
-    min: Optional[float] = None
-    max: Optional[float] = None
-
 
 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."""
+    status: str
+    message: str
+    model_id: str
+    parameters: dict[str, Any]
+
 
 def get_parameter_schema(param_name: str) -> ParameterSchema:
     """

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, GetModelParametersResponse
+from .parameter_schemas import get_parameters_with_schemas, coerce_parameter_value, GetModelParametersResponse, UpdateModelParametersResponse
 from ..config_manager import ConfigManager
 from ..config import UpdateConfigRequest
 
@@ -118,18 +118,16 @@ def put(self):
             )
             config_manager.update_config(update_request)
 
-            # TODO: Define the response type as a Pydantic model to prevent
-            # breaking API changes. Response body is currently unused by frontend
-            # but kept for API compatibility.
-            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(
+                status="success",
+                message=f"Parameters saved for model {model_id}",
+                model_id=model_id,
+                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")