From 5b598d2708fcefbd58aec81f970b0798e481306b Mon Sep 17 00:00:00 2001 From: Carmen Iniesta Date: Fri, 6 Nov 2026 11:59:45 +0100 Subject: [PATCH 1/2] migrate o3 mini migrate o3 mini --- ..._migrate_chatgpt_o3_mini_to_gpt_5_6_sol.py | 27 +++ tests/db/test_o3_mini_migration.py | 226 ++++++++++++++++++ 2 files changed, 253 insertions(+) create mode 100644 alembic/versions/50540732f342_migrate_chatgpt_o3_mini_to_gpt_5_6_sol.py create mode 100644 tests/db/test_o3_mini_migration.py diff --git a/alembic/versions/50540732f342_migrate_chatgpt_o3_mini_to_gpt_5_6_sol.py b/alembic/versions/50540732f342_migrate_chatgpt_o3_mini_to_gpt_5_6_sol.py new file mode 100644 index 0000000..5af480e --- /dev/null +++ b/alembic/versions/50540732f342_migrate_chatgpt_o3_mini_to_gpt_5_6_sol.py @@ -0,0 +1,27 @@ +"""migrate chatgpt o3 mini to gpt 5 6 sol + +Revision ID: 50540732f342 +Revises: b3781ec6a8ad +Create Date: 2026-11-05 17:12:25.297473 + +""" + +from typing import Sequence, Union + +from alembic import op +from nucliadb_agentic_api.db.llm_migration_utils import migrate_llm_models + +# revision identifiers, used by Alembic. +revision: str = "50540732f342" +down_revision: Union[str, Sequence[str], None] = "b3781ec6a8ad" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + migrate_llm_models(op.get_bind(), {"chatgpt-o3-mini": "chatgpt-5.6-sol"}) + + +def downgrade() -> None: + """Downgrade schema.""" + pass diff --git a/tests/db/test_o3_mini_migration.py b/tests/db/test_o3_mini_migration.py new file mode 100644 index 0000000..d5f2465 --- /dev/null +++ b/tests/db/test_o3_mini_migration.py @@ -0,0 +1,226 @@ +import copy +import json +from pathlib import Path +from unittest.mock import Mock + +import pytest +import sqlalchemy as sa +from alembic.config import Config +from alembic.migration import MigrationContext +from alembic.operations import Operations +from alembic.script import ScriptDirectory +from sqlalchemy.engine import Connection + +from nucliadb_agentic_api.db.agentic_configs import agentic_config_table +from nucliadb_agentic_api.db.llm_migration_utils import ( + migrate_llm_models, + walk_and_replace, +) +from nucliadb_agentic_api.db.sources import sources_table + +REPLACEMENTS = {"chatgpt-o3-mini": "chatgpt-5.6-sol"} + + +@pytest.fixture +def migration(): + root = Path(__file__).resolve().parents[2] + scripts = ScriptDirectory.from_config(Config(str(root / "alembic.ini"))) + revision = scripts.get_revision("50540732f342") + assert revision is not None + return revision + + +def test_revision_follows_previous_head(migration): + assert migration.down_revision == "b3781ec6a8ad" + assert migration.branch_labels == set() + assert migration.dependencies is None + + +def test_upgrade_calls_helper_with_only_o3_mini_mapping(migration, monkeypatch): + connection = Mock(spec=Connection) + get_bind = Mock(return_value=connection) + migrate = Mock() + monkeypatch.setattr(migration.module.op, "get_bind", get_bind) + monkeypatch.setattr(migration.module, "migrate_llm_models", migrate) + + migration.module.upgrade() + + get_bind.assert_called_once_with() + migrate.assert_called_once_with(connection, REPLACEMENTS) + + +def test_downgrade_does_not_reverse_model_replacements(migration, monkeypatch): + migrate = Mock() + get_bind = Mock() + monkeypatch.setattr(migration.module, "migrate_llm_models", migrate) + monkeypatch.setattr(migration.module.op, "get_bind", get_bind) + + migration.module.downgrade() + + migrate.assert_not_called() + get_bind.assert_not_called() + + +def test_o3_mini_replacement_recurses_through_objects_and_nested_lists(): + model = {"_type": "llm_config", "model_id": "chatgpt-o3-mini", "temperature": 0.2} + config = { + "smart_agent": { + "models": {"planner": copy.deepcopy(model)}, + "sources": [[{"tool_choice_model": copy.deepcopy(model)}]], + }, + "summarize": {"model": copy.deepcopy(model)}, + "metadata": [None, 42, False, "chatgpt-o3-mini"], + } + expected = copy.deepcopy(config) + expected["smart_agent"]["models"]["planner"]["model_id"] = "chatgpt-5.6-sol" + expected["smart_agent"]["sources"][0][0]["tool_choice_model"]["model_id"] = ( + "chatgpt-5.6-sol" + ) + expected["summarize"]["model"]["model_id"] = "chatgpt-5.6-sol" + + assert walk_and_replace(config, REPLACEMENTS) + assert config == expected + assert not walk_and_replace(config, REPLACEMENTS) + + +@pytest.mark.parametrize( + "config", + [ + {"model_id": "chatgpt-o3-mini"}, + {"_type": "custom", "model_id": "chatgpt-o3-mini"}, + {"_type": "LLMConfig", "model_id": "chatgpt-o3-mini"}, + {"_type": "llm_config", "model_id": "chatgpt-o3-mini/custom-id"}, + {"_type": "llm_config", "model_id": "custom-id/chatgpt-o3-mini"}, + {"_type": "llm_config", "model_id": "chatgpt-azure-o3-mini"}, + {"_type": "llm_config", "model_id": "chatgpt-o3-mini-extra"}, + {"_type": "llm_config", "model_id": "chatgpt-5.6-sol"}, + {"_type": "llm_config", "model_id": None}, + {"_type": "llm_config", "model_id": 42}, + {"_type": "llm_config"}, + ], +) +def test_o3_mini_replacement_preserves_other_models_and_discriminators(config): + original = copy.deepcopy(config) + + assert not walk_and_replace(config, REPLACEMENTS) + assert config == original + + +@pytest.mark.parametrize("serialize", [False, True]) +def test_o3_mini_migration_updates_only_matching_agent_and_source_rows(serialize): + config = { + "models": [{"_type": "llm_config", "model_id": "chatgpt-o3-mini"}], + "custom": {"_type": "llm_config", "model_id": "chatgpt-o3-mini/custom-id"}, + } + expected = copy.deepcopy(config) + expected["models"][0]["model_id"] = "chatgpt-5.6-sol" + connection = Mock(spec=Connection) + executions = [] + + def execute(statement, params=None): + sql = str(statement) + executions.append((sql, params)) + if sql.startswith("SELECT"): + stored = json.dumps(config) if serialize else copy.deepcopy(config) + return Mock( + fetchall=Mock( + return_value=[ + ("account", "kbid", "changed-id", stored), + ("other-account", "kbid", "changed-id", expected), + ("account", "kbid", "null-id", None), + ] + ) + ) + return Mock() + + connection.execute.side_effect = execute + + assert migrate_llm_models(connection, REPLACEMENTS) == 2 + assert len(executions) == 4 + for (select_sql, _), (update_sql, params), table, id_column in ( + (*executions[:2], "agentic_config_table", "agentic_id"), + (*executions[2:], "sources_table", "source_id"), + ): + assert select_sql == ( + f"SELECT account, kbid, {id_column}, config FROM {table} " + 'WHERE config::text LIKE \'%"_type": "llm_config"%\'' + ) + assert update_sql == ( + f"UPDATE {table} SET config = CAST(:config AS jsonb) " + f"WHERE account = :account AND kbid = :kbid AND {id_column} = :{id_column}" + ) + assert params == { + "account": "account", + "kbid": "kbid", + id_column: "changed-id", + "config": json.dumps(expected), + } + connection.commit.assert_not_called() + + +@pytest.mark.parametrize( + ("table", "id_column", "extra_values"), + [ + (agentic_config_table, "agentic_id", {}), + (sources_table, "source_id", {"type": "mcp"}), + ], +) +def test_upgrade_persists_agent_and_source_configs( + agentic_pg_dsn, migration, table, id_column, extra_values +): + preserved = { + "custom": {"_type": "llm_config", "model_id": "chatgpt-o3-mini/custom-id"}, + "untyped": {"model_id": "chatgpt-o3-mini"}, + "other_type": {"_type": "custom", "model_id": "chatgpt-o3-mini"}, + "azure": {"_type": "llm_config", "model_id": "chatgpt-azure-o3-mini"}, + } + config = { + **preserved, + "nested": [[{"model": {"_type": "llm_config", "model_id": "chatgpt-o3-mini"}}]], + } + expected = copy.deepcopy(config) + expected["nested"][0][0]["model"]["model_id"] = "chatgpt-5.6-sol" + row_id = "o3-mini-migration" + rows = [ + ("o3-account", "o3-kbid", config), + ("other-account", "o3-kbid", preserved), + ("o3-account", "other-kbid", preserved), + ] + engine = sa.create_engine(agentic_pg_dsn) + try: + with engine.begin() as conn: + conn.execute( + table.insert(), + [ + { + "account": account, + "kbid": kbid, + id_column: row_id, + "config": stored, + **extra_values, + } + for account, kbid, stored in rows + ], + ) + + with ( + engine.begin() as conn, + Operations.context(MigrationContext.configure(conn)), + ): + migration.module.upgrade() + + with engine.connect() as conn: + stored_rows = conn.execute( + sa.select(table.c.account, table.c.kbid, table.c.config).where( + table.c[id_column] == row_id + ) + ).all() + assert {(account, kbid): stored for account, kbid, stored in stored_rows} == { + ("o3-account", "o3-kbid"): expected, + ("other-account", "o3-kbid"): preserved, + ("o3-account", "other-kbid"): preserved, + } + finally: + with engine.begin() as conn: + conn.execute(table.delete().where(table.c[id_column] == row_id)) + engine.dispose() From 8ea4491d9d6651ab5e3a7571f6d1bd0187cf9a47 Mon Sep 17 00:00:00 2001 From: Carmen Iniesta Date: Fri, 6 Nov 2026 15:42:11 +0100 Subject: [PATCH 2/2] Migrate OpenAI o1 agent configurations to GPT-5.6 Sol Add a separate migration after the o3-mini retirement and reuse the exact-model recursive helper for agent and source configurations. Extend existing regressions to both retired models, including persisted updates, custom and Azure preservation, and repeated execution. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- ...61535_migrate_chatgpt_o1_to_gpt_5_6_sol.py | 27 ++++++ ...ion.py => test_openai_model_migrations.py} | 86 +++++++++++++------ 2 files changed, 88 insertions(+), 25 deletions(-) create mode 100644 alembic/versions/eff180d61535_migrate_chatgpt_o1_to_gpt_5_6_sol.py rename tests/db/{test_o3_mini_migration.py => test_openai_model_migrations.py} (71%) diff --git a/alembic/versions/eff180d61535_migrate_chatgpt_o1_to_gpt_5_6_sol.py b/alembic/versions/eff180d61535_migrate_chatgpt_o1_to_gpt_5_6_sol.py new file mode 100644 index 0000000..8d48f5e --- /dev/null +++ b/alembic/versions/eff180d61535_migrate_chatgpt_o1_to_gpt_5_6_sol.py @@ -0,0 +1,27 @@ +"""migrate chatgpt o1 to gpt 5 6 sol + +Revision ID: eff180d61535 +Revises: 50540732f342 +Create Date: 2026-11-06 15:29:58.883376 + +""" + +from typing import Sequence, Union + +from alembic import op +from nucliadb_agentic_api.db.llm_migration_utils import migrate_llm_models + +# revision identifiers, used by Alembic. +revision: str = "eff180d61535" +down_revision: Union[str, Sequence[str], None] = "50540732f342" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + migrate_llm_models(op.get_bind(), {"chatgpt-o1": "chatgpt-5.6-sol"}) + + +def downgrade() -> None: + """Downgrade schema.""" + pass diff --git a/tests/db/test_o3_mini_migration.py b/tests/db/test_openai_model_migrations.py similarity index 71% rename from tests/db/test_o3_mini_migration.py rename to tests/db/test_openai_model_migrations.py index d5f2465..ce9423d 100644 --- a/tests/db/test_o3_mini_migration.py +++ b/tests/db/test_openai_model_migrations.py @@ -18,25 +18,46 @@ ) from nucliadb_agentic_api.db.sources import sources_table -REPLACEMENTS = {"chatgpt-o3-mini": "chatgpt-5.6-sol"} + +@pytest.fixture(params=["chatgpt-o3-mini", "chatgpt-o1"]) +def old_model(request): + return request.param @pytest.fixture -def migration(): +def replacements(old_model): + return {old_model: "chatgpt-5.6-sol"} + + +@pytest.fixture +def migration(old_model): root = Path(__file__).resolve().parents[2] scripts = ScriptDirectory.from_config(Config(str(root / "alembic.ini"))) - revision = scripts.get_revision("50540732f342") + revision = scripts.get_revision( + { + "chatgpt-o3-mini": "50540732f342", + "chatgpt-o1": "eff180d61535", + }[old_model] + ) assert revision is not None return revision -def test_revision_follows_previous_head(migration): - assert migration.down_revision == "b3781ec6a8ad" +def test_revision_follows_previous_head(migration, old_model): + assert ( + migration.down_revision + == { + "chatgpt-o3-mini": "b3781ec6a8ad", + "chatgpt-o1": "50540732f342", + }[old_model] + ) assert migration.branch_labels == set() assert migration.dependencies is None -def test_upgrade_calls_helper_with_only_o3_mini_mapping(migration, monkeypatch): +def test_upgrade_calls_helper_with_only_retired_model_mapping( + migration, monkeypatch, replacements +): connection = Mock(spec=Connection) get_bind = Mock(return_value=connection) migrate = Mock() @@ -46,7 +67,7 @@ def test_upgrade_calls_helper_with_only_o3_mini_mapping(migration, monkeypatch): migration.module.upgrade() get_bind.assert_called_once_with() - migrate.assert_called_once_with(connection, REPLACEMENTS) + migrate.assert_called_once_with(connection, replacements) def test_downgrade_does_not_reverse_model_replacements(migration, monkeypatch): @@ -61,15 +82,15 @@ def test_downgrade_does_not_reverse_model_replacements(migration, monkeypatch): get_bind.assert_not_called() -def test_o3_mini_replacement_recurses_through_objects_and_nested_lists(): - model = {"_type": "llm_config", "model_id": "chatgpt-o3-mini", "temperature": 0.2} +def test_replacement_recurses_through_objects_and_nested_lists(old_model, replacements): + model = {"_type": "llm_config", "model_id": old_model, "temperature": 0.2} config = { "smart_agent": { "models": {"planner": copy.deepcopy(model)}, "sources": [[{"tool_choice_model": copy.deepcopy(model)}]], }, "summarize": {"model": copy.deepcopy(model)}, - "metadata": [None, 42, False, "chatgpt-o3-mini"], + "metadata": [None, 42, False, old_model], } expected = copy.deepcopy(config) expected["smart_agent"]["models"]["planner"]["model_id"] = "chatgpt-5.6-sol" @@ -78,9 +99,9 @@ def test_o3_mini_replacement_recurses_through_objects_and_nested_lists(): ) expected["summarize"]["model"]["model_id"] = "chatgpt-5.6-sol" - assert walk_and_replace(config, REPLACEMENTS) + assert walk_and_replace(config, replacements) assert config == expected - assert not walk_and_replace(config, REPLACEMENTS) + assert not walk_and_replace(config, replacements) @pytest.mark.parametrize( @@ -93,24 +114,35 @@ def test_o3_mini_replacement_recurses_through_objects_and_nested_lists(): {"_type": "llm_config", "model_id": "custom-id/chatgpt-o3-mini"}, {"_type": "llm_config", "model_id": "chatgpt-azure-o3-mini"}, {"_type": "llm_config", "model_id": "chatgpt-o3-mini-extra"}, + {"model_id": "chatgpt-o1"}, + {"_type": "custom", "model_id": "chatgpt-o1"}, + {"_type": "LLMConfig", "model_id": "chatgpt-o1"}, + {"_type": "llm_config", "model_id": "chatgpt-o1/custom-id"}, + {"_type": "llm_config", "model_id": "custom-id/chatgpt-o1"}, + {"_type": "llm_config", "model_id": "chatgpt-azure-o1"}, + {"_type": "llm_config", "model_id": "chatgpt-o1-extra"}, + {"_type": "llm_config", "model_id": "chatgpt4o"}, + {"_type": "llm_config", "model_id": "chatgpt4o-mini"}, {"_type": "llm_config", "model_id": "chatgpt-5.6-sol"}, {"_type": "llm_config", "model_id": None}, {"_type": "llm_config", "model_id": 42}, {"_type": "llm_config"}, ], ) -def test_o3_mini_replacement_preserves_other_models_and_discriminators(config): +def test_replacement_preserves_other_models_and_discriminators(config, replacements): original = copy.deepcopy(config) - assert not walk_and_replace(config, REPLACEMENTS) + assert not walk_and_replace(config, replacements) assert config == original @pytest.mark.parametrize("serialize", [False, True]) -def test_o3_mini_migration_updates_only_matching_agent_and_source_rows(serialize): +def test_migration_updates_only_matching_agent_and_source_rows( + serialize, old_model, replacements +): config = { - "models": [{"_type": "llm_config", "model_id": "chatgpt-o3-mini"}], - "custom": {"_type": "llm_config", "model_id": "chatgpt-o3-mini/custom-id"}, + "models": [{"_type": "llm_config", "model_id": old_model}], + "custom": {"_type": "llm_config", "model_id": f"{old_model}/custom-id"}, } expected = copy.deepcopy(config) expected["models"][0]["model_id"] = "chatgpt-5.6-sol" @@ -135,7 +167,7 @@ def execute(statement, params=None): connection.execute.side_effect = execute - assert migrate_llm_models(connection, REPLACEMENTS) == 2 + assert migrate_llm_models(connection, replacements) == 2 assert len(executions) == 4 for (select_sql, _), (update_sql, params), table, id_column in ( (*executions[:2], "agentic_config_table", "agentic_id"), @@ -166,21 +198,24 @@ def execute(statement, params=None): ], ) def test_upgrade_persists_agent_and_source_configs( - agentic_pg_dsn, migration, table, id_column, extra_values + agentic_pg_dsn, migration, table, id_column, extra_values, old_model ): preserved = { - "custom": {"_type": "llm_config", "model_id": "chatgpt-o3-mini/custom-id"}, - "untyped": {"model_id": "chatgpt-o3-mini"}, - "other_type": {"_type": "custom", "model_id": "chatgpt-o3-mini"}, - "azure": {"_type": "llm_config", "model_id": "chatgpt-azure-o3-mini"}, + "custom": {"_type": "llm_config", "model_id": f"{old_model}/custom-id"}, + "untyped": {"model_id": old_model}, + "other_type": {"_type": "custom", "model_id": old_model}, + "azure": { + "_type": "llm_config", + "model_id": old_model.replace("chatgpt-", "chatgpt-azure-"), + }, } config = { **preserved, - "nested": [[{"model": {"_type": "llm_config", "model_id": "chatgpt-o3-mini"}}]], + "nested": [[{"model": {"_type": "llm_config", "model_id": old_model}}]], } expected = copy.deepcopy(config) expected["nested"][0][0]["model"]["model_id"] = "chatgpt-5.6-sol" - row_id = "o3-mini-migration" + row_id = f"{old_model}-migration" rows = [ ("o3-account", "o3-kbid", config), ("other-account", "o3-kbid", preserved), @@ -208,6 +243,7 @@ def test_upgrade_persists_agent_and_source_configs( Operations.context(MigrationContext.configure(conn)), ): migration.module.upgrade() + migration.module.upgrade() with engine.connect() as conn: stored_rows = conn.execute(