mirror of
https://github.com/lfnovo/open-notebook.git
synced 2026-07-19 12:46:07 +00:00
bb48d3d578
* feat: per-transformation custom model selection Fixes #776 * fix: validate model_id exists when creating or updating a transformation The create and update endpoints persisted model_id without checking the referenced model exists, so an invalid reference was stored silently and only surfaced later as a 404 at execution time. Add the same existence check execute_transformation already performs (Model.get -> 404 if missing) to both the create and update paths, so a bad model_id is rejected up front with a clear 404. Update still allows clearing model_id to None. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * test: cover model_id existence validation on transformation create/update The create test now mocks Model.get so the new existence check passes and asserts it was awaited with the supplied model_id. The update test expects Model.get to be awaited twice (once validating the update, once at execute) rather than once. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * docs: add CHANGELOG entry for per-transformation model selection (#776) --------- Co-authored-by: Matt Van Horn <455140+mvanhorn@users.noreply.github.com> Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com> Co-authored-by: Luis Novo <lfnovo@gmail.com>
242 lines
7.7 KiB
Python
242 lines
7.7 KiB
Python
from datetime import datetime
|
|
from unittest.mock import AsyncMock, patch
|
|
|
|
from fastapi.testclient import TestClient
|
|
|
|
from open_notebook.domain.transformation import Transformation
|
|
|
|
|
|
def _client() -> TestClient:
|
|
from api.main import app
|
|
|
|
return TestClient(app)
|
|
|
|
|
|
def _transformation(model_id: str | None = None) -> Transformation:
|
|
return Transformation(
|
|
id="transformation:123",
|
|
name="summary",
|
|
title="Summary",
|
|
description="Summarize the source",
|
|
prompt="Summarize this",
|
|
apply_default=False,
|
|
model_id=model_id,
|
|
created=datetime(2026, 1, 1, 12, 0, 0),
|
|
updated=datetime(2026, 1, 1, 12, 0, 0),
|
|
)
|
|
|
|
|
|
def _create_payload(model_id: str | None = None) -> dict:
|
|
payload = {
|
|
"name": "summary",
|
|
"title": "Summary",
|
|
"description": "Summarize the source",
|
|
"prompt": "Summarize this",
|
|
"apply_default": False,
|
|
}
|
|
if model_id is not None:
|
|
payload["model_id"] = model_id
|
|
return payload
|
|
|
|
|
|
def test_create_transformation_with_model_id_persists_and_reads_back():
|
|
client = _client()
|
|
saved_transformations: list[Transformation] = []
|
|
|
|
async def capture_save(transformation: Transformation):
|
|
saved_transformations.append(transformation)
|
|
transformation.id = "transformation:created"
|
|
transformation.created = datetime(2026, 1, 1, 12, 0, 0)
|
|
transformation.updated = datetime(2026, 1, 1, 12, 0, 0)
|
|
|
|
with (
|
|
patch.object(
|
|
Transformation, "save", autospec=True, side_effect=capture_save
|
|
),
|
|
patch(
|
|
"api.routers.transformations.Model.get",
|
|
new_callable=AsyncMock,
|
|
return_value=object(),
|
|
) as mock_model_get,
|
|
):
|
|
response = client.post(
|
|
"/api/transformations", json=_create_payload("model:local")
|
|
)
|
|
|
|
assert response.status_code == 200
|
|
assert saved_transformations[0].model_id == "model:local"
|
|
assert response.json()["model_id"] == "model:local"
|
|
mock_model_get.assert_awaited_once_with("model:local")
|
|
|
|
with patch(
|
|
"api.routers.transformations.Transformation.get",
|
|
new_callable=AsyncMock,
|
|
return_value=saved_transformations[0],
|
|
):
|
|
response = client.get("/api/transformations/transformation:created")
|
|
|
|
assert response.status_code == 200
|
|
assert response.json()["model_id"] == "model:local"
|
|
|
|
|
|
def test_execute_without_request_model_uses_stored_model():
|
|
client = _client()
|
|
transformation = _transformation(model_id="model:stored")
|
|
|
|
with (
|
|
patch(
|
|
"api.routers.transformations.Transformation.get",
|
|
new_callable=AsyncMock,
|
|
return_value=transformation,
|
|
),
|
|
patch(
|
|
"api.routers.transformations.Model.get",
|
|
new_callable=AsyncMock,
|
|
return_value=object(),
|
|
) as mock_model_get,
|
|
patch(
|
|
"api.routers.transformations.transformation_graph.ainvoke",
|
|
new_callable=AsyncMock,
|
|
return_value={"output": "Stored model output"},
|
|
) as mock_ainvoke,
|
|
):
|
|
response = client.post(
|
|
"/api/transformations/execute",
|
|
json={
|
|
"transformation_id": transformation.id,
|
|
"input_text": "Input text",
|
|
},
|
|
)
|
|
|
|
assert response.status_code == 200
|
|
assert response.json()["model_id"] == "model:stored"
|
|
mock_model_get.assert_awaited_once_with("model:stored")
|
|
assert (
|
|
mock_ainvoke.call_args.kwargs["config"]["configurable"]["model_id"]
|
|
== "model:stored"
|
|
)
|
|
|
|
|
|
def test_execute_without_any_model_uses_default_transformation_model_fallback():
|
|
client = _client()
|
|
transformation = _transformation(model_id=None)
|
|
|
|
with (
|
|
patch(
|
|
"api.routers.transformations.Transformation.get",
|
|
new_callable=AsyncMock,
|
|
return_value=transformation,
|
|
),
|
|
patch(
|
|
"api.routers.transformations.Model.get",
|
|
new_callable=AsyncMock,
|
|
) as mock_model_get,
|
|
patch(
|
|
"api.routers.transformations.transformation_graph.ainvoke",
|
|
new_callable=AsyncMock,
|
|
return_value={"output": "Default model output"},
|
|
) as mock_ainvoke,
|
|
):
|
|
response = client.post(
|
|
"/api/transformations/execute",
|
|
json={
|
|
"transformation_id": transformation.id,
|
|
"input_text": "Input text",
|
|
},
|
|
)
|
|
|
|
assert response.status_code == 200
|
|
assert response.json()["model_id"] is None
|
|
mock_model_get.assert_not_awaited()
|
|
assert mock_ainvoke.call_args.kwargs["config"]["configurable"]["model_id"] is None
|
|
|
|
|
|
def test_update_transformation_model_id_is_used_by_subsequent_execution():
|
|
client = _client()
|
|
transformation = _transformation(model_id="model:old")
|
|
|
|
async def save_update(transformation: Transformation):
|
|
transformation.updated = datetime(2026, 1, 1, 12, 5, 0)
|
|
|
|
with (
|
|
patch(
|
|
"api.routers.transformations.Transformation.get",
|
|
new_callable=AsyncMock,
|
|
return_value=transformation,
|
|
),
|
|
patch.object(
|
|
Transformation, "save", autospec=True, side_effect=save_update
|
|
),
|
|
patch(
|
|
"api.routers.transformations.Model.get",
|
|
new_callable=AsyncMock,
|
|
return_value=object(),
|
|
) as mock_model_get,
|
|
patch(
|
|
"api.routers.transformations.transformation_graph.ainvoke",
|
|
new_callable=AsyncMock,
|
|
return_value={"output": "Updated model output"},
|
|
) as mock_ainvoke,
|
|
):
|
|
update_response = client.put(
|
|
f"/api/transformations/{transformation.id}",
|
|
json={"model_id": "model:new"},
|
|
)
|
|
execute_response = client.post(
|
|
"/api/transformations/execute",
|
|
json={
|
|
"transformation_id": transformation.id,
|
|
"input_text": "Input text",
|
|
},
|
|
)
|
|
|
|
assert update_response.status_code == 200
|
|
assert update_response.json()["model_id"] == "model:new"
|
|
assert execute_response.status_code == 200
|
|
assert execute_response.json()["model_id"] == "model:new"
|
|
# model:new is validated once at update time and looked up again at execute.
|
|
assert mock_model_get.await_count == 2
|
|
mock_model_get.assert_awaited_with("model:new")
|
|
assert (
|
|
mock_ainvoke.call_args.kwargs["config"]["configurable"]["model_id"]
|
|
== "model:new"
|
|
)
|
|
|
|
|
|
def test_execute_request_model_overrides_stored_model():
|
|
client = _client()
|
|
transformation = _transformation(model_id="model:stored")
|
|
|
|
with (
|
|
patch(
|
|
"api.routers.transformations.Transformation.get",
|
|
new_callable=AsyncMock,
|
|
return_value=transformation,
|
|
),
|
|
patch(
|
|
"api.routers.transformations.Model.get",
|
|
new_callable=AsyncMock,
|
|
return_value=object(),
|
|
) as mock_model_get,
|
|
patch(
|
|
"api.routers.transformations.transformation_graph.ainvoke",
|
|
new_callable=AsyncMock,
|
|
return_value={"output": "Override output"},
|
|
) as mock_ainvoke,
|
|
):
|
|
response = client.post(
|
|
"/api/transformations/execute",
|
|
json={
|
|
"transformation_id": transformation.id,
|
|
"input_text": "Input text",
|
|
"model_id": "model:override",
|
|
},
|
|
)
|
|
|
|
assert response.status_code == 200
|
|
assert response.json()["model_id"] == "model:override"
|
|
mock_model_get.assert_awaited_once_with("model:override")
|
|
assert (
|
|
mock_ainvoke.call_args.kwargs["config"]["configurable"]["model_id"]
|
|
== "model:override"
|
|
)
|