mirror of
https://github.com/MODSetter/SurfSense.git
synced 2026-05-31 19:45:15 +02:00
refactor(automations): streamline model eligibility handling in automation creation
- Removed the eligibility gate for model selection in the automation creation process, allowing users to choose models directly in the builder. - Updated the `AutomationBuilderForm` to incorporate model selection logic, ensuring that selected models are validated and preserved during automation creation and editing. - Simplified the `AutomationsContent` and `AutomationNewContent` components by eliminating unnecessary eligibility checks and alerts. - Enhanced the user experience by integrating model selection directly into the automation approval process, ensuring that only billable models are used. - Refactored related tests to cover new model selection behavior and ensure proper validation of user-selected models.
This commit is contained in:
DESKTOP-RTLN3BA\$punk
parent
fade9d1b9d
commit
9d1a01eb0c
13 changed files with 887 additions and 263 deletions
|
|
@ -32,8 +32,7 @@ from app.agents.multi_agent_chat.subagents.shared.hitl.approvals.self_gated impo
|
|||
)
|
||||
from app.automations.schemas.api import AutomationCreate
|
||||
from app.automations.services.automation import AutomationService
|
||||
from app.automations.services.model_policy import get_automation_model_eligibility
|
||||
from app.db import SearchSpace, User, async_session_maker
|
||||
from app.db import User, async_session_maker
|
||||
from app.utils.content_utils import extract_text_content
|
||||
|
||||
from .prompt import build_draft_prompt
|
||||
|
|
@ -99,26 +98,11 @@ def create_create_automation_tool(
|
|||
declined. Acknowledge once and stop — do NOT retry or pitch
|
||||
variants without a fresh user request.
|
||||
"""
|
||||
# --- 0. Eligibility gate (fail fast, before drafting + HITL) ---
|
||||
# Automations may only use premium or BYOK models. Check up front so we
|
||||
# don't make the user draft + approve a card that can't be saved.
|
||||
async with async_session_maker() as session:
|
||||
search_space = await session.get(SearchSpace, search_space_id)
|
||||
if search_space is None:
|
||||
return {
|
||||
"status": "error",
|
||||
"message": "search space not found in this session",
|
||||
}
|
||||
eligibility = get_automation_model_eligibility(search_space)
|
||||
if not eligibility["allowed"]:
|
||||
reasons = " ".join(v["reason"] for v in eligibility["violations"])
|
||||
return {
|
||||
"status": "error",
|
||||
"message": (
|
||||
f"{reasons} Update the search space's model settings to a "
|
||||
"premium or your own (BYOK) model, then try again."
|
||||
),
|
||||
}
|
||||
# Models are chosen per-automation on the approval card (premium/BYOK
|
||||
# selectors) and validated when persisted by ``AutomationService.create``
|
||||
# — so there's no fail-fast search-space eligibility gate here. The
|
||||
# search space's current chat/role model selection no longer constrains
|
||||
# whether an automation can be drafted or saved.
|
||||
|
||||
# --- 1. Draft via sub-LLM ---
|
||||
prompt = build_draft_prompt(search_space_id=search_space_id, intent=intent)
|
||||
|
|
|
|||
|
|
@ -22,6 +22,7 @@ from app.automations.schemas.definition.envelope import AutomationModels
|
|||
from app.automations.services.model_policy import (
|
||||
AutomationModelPolicyError,
|
||||
assert_automation_models_billable,
|
||||
assert_models_billable,
|
||||
get_automation_model_eligibility,
|
||||
)
|
||||
from app.automations.triggers import get_trigger
|
||||
|
|
@ -43,16 +44,23 @@ class AutomationService:
|
|||
await self._authorize(
|
||||
payload.search_space_id, Permission.AUTOMATIONS_CREATE.value
|
||||
)
|
||||
search_space = await self._assert_models_billable(payload.search_space_id)
|
||||
|
||||
# Snapshot the search space's current (already-validated) model prefs onto
|
||||
# the definition so runs are insulated from later chat/search-space model
|
||||
# changes. Captured ids are guaranteed billable by the check above.
|
||||
payload.definition.models = AutomationModels(
|
||||
agent_llm_id=search_space.agent_llm_id or 0,
|
||||
image_generation_config_id=search_space.image_generation_config_id or 0,
|
||||
vision_llm_config_id=search_space.vision_llm_config_id or 0,
|
||||
)
|
||||
# Capture the model profile onto the definition so runs are insulated
|
||||
# from later chat/search-space model changes. Two sources:
|
||||
# 1. Explicit per-automation selection in ``payload.definition.models``
|
||||
# (manual builder + chat approval card). Validate the chosen ids.
|
||||
# 2. Fallback (no selection): snapshot the search space's current prefs.
|
||||
# Either way the captured ids are guaranteed billable (premium/BYOK).
|
||||
selected_models = payload.definition.models
|
||||
if selected_models is not None:
|
||||
self._assert_selected_models_billable(selected_models)
|
||||
else:
|
||||
search_space = await self._assert_models_billable(payload.search_space_id)
|
||||
payload.definition.models = AutomationModels(
|
||||
agent_llm_id=search_space.agent_llm_id or 0,
|
||||
image_generation_config_id=search_space.image_generation_config_id or 0,
|
||||
vision_llm_config_id=search_space.vision_llm_config_id or 0,
|
||||
)
|
||||
|
||||
automation = Automation(
|
||||
search_space_id=payload.search_space_id,
|
||||
|
|
@ -122,13 +130,22 @@ class AutomationService:
|
|||
automation.status = data["status"]
|
||||
if "definition" in data:
|
||||
new_def = patch.definition.model_dump(mode="json", by_alias=True)
|
||||
# Preserve the captured model snapshot across edits so a definition
|
||||
# change never silently re-binds the automation to the current chat
|
||||
# model selection. Backend-managed; survives whether or not the
|
||||
# client round-trips ``models``.
|
||||
# Model snapshot handling on edit:
|
||||
# * absent in the patch -> preserve the captured snapshot
|
||||
# (a non-model definition change never silently re-binds the
|
||||
# automation to the current chat/search-space selection).
|
||||
# * unchanged from the snapshot -> keep as-is, no re-validation
|
||||
# (so editing an automation whose captured model later drifted
|
||||
# out of premium isn't blocked by an unrelated name/schedule edit).
|
||||
# * genuinely changed -> validate the new selection (422 on a
|
||||
# non-billable pick), then accept it.
|
||||
existing_models = (automation.definition or {}).get("models")
|
||||
if existing_models is not None:
|
||||
new_def["models"] = existing_models
|
||||
provided_models = new_def.get("models")
|
||||
if provided_models is None:
|
||||
if existing_models is not None:
|
||||
new_def["models"] = existing_models
|
||||
elif provided_models != existing_models:
|
||||
self._assert_selected_models_billable(patch.definition.models)
|
||||
automation.definition = new_def
|
||||
automation.version += 1
|
||||
|
||||
|
|
@ -199,6 +216,22 @@ class AutomationService:
|
|||
raise HTTPException(status_code=422, detail=str(exc)) from exc
|
||||
return search_space
|
||||
|
||||
def _assert_selected_models_billable(self, models: AutomationModels) -> None:
|
||||
"""Reject creation when an explicitly selected model isn't billable.
|
||||
|
||||
Used when the client supplies ``definition.models`` (per-automation
|
||||
selection from the builder or chat approval card). Same policy as the
|
||||
search-space path: premium global or BYOK only, no free/Auto.
|
||||
"""
|
||||
try:
|
||||
assert_models_billable(
|
||||
agent_llm_id=models.agent_llm_id,
|
||||
image_generation_config_id=models.image_generation_config_id,
|
||||
vision_llm_config_id=models.vision_llm_config_id,
|
||||
)
|
||||
except AutomationModelPolicyError as exc:
|
||||
raise HTTPException(status_code=422, detail=str(exc)) from exc
|
||||
|
||||
async def _authorize(self, search_space_id: int, permission: str) -> None:
|
||||
await check_permission(
|
||||
self.session,
|
||||
|
|
|
|||
|
|
@ -16,7 +16,10 @@ from fastapi import HTTPException
|
|||
|
||||
import app.automations.services.automation as automation_mod
|
||||
from app.automations.schemas.api import AutomationCreate, AutomationUpdate
|
||||
from app.automations.schemas.definition.envelope import AutomationDefinition
|
||||
from app.automations.schemas.definition.envelope import (
|
||||
AutomationDefinition,
|
||||
AutomationModels,
|
||||
)
|
||||
from app.automations.schemas.definition.plan_step import PlanStep
|
||||
from app.automations.services.automation import AutomationService
|
||||
from app.automations.services.model_policy import AutomationModelPolicyError
|
||||
|
|
@ -175,6 +178,100 @@ async def test_create_treats_unset_prefs_as_auto_zero(
|
|||
}
|
||||
|
||||
|
||||
async def test_create_honors_selected_models_when_provided(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""When the payload carries ``definition.models`` they are validated + kept.
|
||||
|
||||
The search-space snapshot path is bypassed entirely (no
|
||||
``assert_automation_models_billable`` call).
|
||||
"""
|
||||
|
||||
def _fail_snapshot(_ss):
|
||||
raise AssertionError("snapshot path should not run when models are provided")
|
||||
|
||||
monkeypatch.setattr(
|
||||
automation_mod, "assert_automation_models_billable", _fail_snapshot
|
||||
)
|
||||
validated: dict[str, Any] = {}
|
||||
|
||||
def _assert_ok(*, agent_llm_id, image_generation_config_id, vision_llm_config_id):
|
||||
validated["ids"] = (
|
||||
agent_llm_id,
|
||||
image_generation_config_id,
|
||||
vision_llm_config_id,
|
||||
)
|
||||
|
||||
monkeypatch.setattr(automation_mod, "assert_models_billable", _assert_ok)
|
||||
|
||||
async def _noop_authorize(self, *_a, **_k):
|
||||
return None
|
||||
|
||||
async def _return_added(self, _aid):
|
||||
return self.session.added[-1]
|
||||
|
||||
monkeypatch.setattr(AutomationService, "_authorize", _noop_authorize)
|
||||
monkeypatch.setattr(AutomationService, "_get_with_triggers_or_raise", _return_added)
|
||||
|
||||
service = _service(SimpleNamespace(agent_llm_id=-99))
|
||||
payload = AutomationCreate(
|
||||
search_space_id=1,
|
||||
name="A",
|
||||
definition=_definition(
|
||||
models=AutomationModels(
|
||||
agent_llm_id=-1,
|
||||
image_generation_config_id=7,
|
||||
vision_llm_config_id=-2,
|
||||
)
|
||||
),
|
||||
)
|
||||
|
||||
automation = await service.create(payload)
|
||||
|
||||
assert validated["ids"] == (-1, 7, -2)
|
||||
assert automation.definition["models"] == {
|
||||
"agent_llm_id": -1,
|
||||
"image_generation_config_id": 7,
|
||||
"vision_llm_config_id": -2,
|
||||
}
|
||||
|
||||
|
||||
async def test_create_rejects_unbillable_selected_models(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""A non-billable explicit selection maps the policy error to HTTP 422."""
|
||||
|
||||
def _raise(*, agent_llm_id, image_generation_config_id, vision_llm_config_id):
|
||||
raise AutomationModelPolicyError(
|
||||
[{"kind": "llm", "config_id": -3, "reason": "free model"}]
|
||||
)
|
||||
|
||||
monkeypatch.setattr(automation_mod, "assert_models_billable", _raise)
|
||||
|
||||
async def _noop_authorize(self, *_a, **_k):
|
||||
return None
|
||||
|
||||
monkeypatch.setattr(AutomationService, "_authorize", _noop_authorize)
|
||||
|
||||
service = _service(SimpleNamespace(agent_llm_id=-3))
|
||||
payload = AutomationCreate(
|
||||
search_space_id=1,
|
||||
name="A",
|
||||
definition=_definition(
|
||||
models=AutomationModels(
|
||||
agent_llm_id=-3,
|
||||
image_generation_config_id=7,
|
||||
vision_llm_config_id=-2,
|
||||
)
|
||||
),
|
||||
)
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await service.create(payload)
|
||||
|
||||
assert exc_info.value.status_code == 422
|
||||
|
||||
|
||||
async def test_update_preserves_captured_models(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
|
|
@ -211,6 +308,166 @@ async def test_update_preserves_captured_models(
|
|||
assert result.version == 4
|
||||
|
||||
|
||||
async def test_update_honors_changed_models_when_valid(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""A definition edit with a *changed* models block validates + keeps it."""
|
||||
existing = SimpleNamespace(
|
||||
search_space_id=1,
|
||||
definition={
|
||||
"name": "A",
|
||||
"plan": [],
|
||||
"models": {
|
||||
"agent_llm_id": -1,
|
||||
"image_generation_config_id": 5,
|
||||
"vision_llm_config_id": -1,
|
||||
},
|
||||
},
|
||||
version=3,
|
||||
)
|
||||
validated: dict[str, Any] = {}
|
||||
|
||||
def _assert_ok(*, agent_llm_id, image_generation_config_id, vision_llm_config_id):
|
||||
validated["ids"] = (
|
||||
agent_llm_id,
|
||||
image_generation_config_id,
|
||||
vision_llm_config_id,
|
||||
)
|
||||
|
||||
monkeypatch.setattr(automation_mod, "assert_models_billable", _assert_ok)
|
||||
|
||||
async def _noop_authorize(self, *_a, **_k):
|
||||
return None
|
||||
|
||||
async def _return_existing(self, _aid):
|
||||
return existing
|
||||
|
||||
monkeypatch.setattr(AutomationService, "_authorize", _noop_authorize)
|
||||
monkeypatch.setattr(
|
||||
AutomationService, "_get_with_triggers_or_raise", _return_existing
|
||||
)
|
||||
|
||||
service = _service(SimpleNamespace())
|
||||
patch = AutomationUpdate(
|
||||
definition=_definition(
|
||||
models=AutomationModels(
|
||||
agent_llm_id=-2,
|
||||
image_generation_config_id=9,
|
||||
vision_llm_config_id=-2,
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
result = await service.update(7, patch)
|
||||
|
||||
assert validated["ids"] == (-2, 9, -2)
|
||||
assert result.definition["models"] == {
|
||||
"agent_llm_id": -2,
|
||||
"image_generation_config_id": 9,
|
||||
"vision_llm_config_id": -2,
|
||||
}
|
||||
assert result.version == 4
|
||||
|
||||
|
||||
async def test_update_rejects_changed_unbillable_models(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""A *changed* non-billable models block is rejected with HTTP 422."""
|
||||
existing = SimpleNamespace(
|
||||
search_space_id=1,
|
||||
definition={
|
||||
"name": "A",
|
||||
"plan": [],
|
||||
"models": {
|
||||
"agent_llm_id": -1,
|
||||
"image_generation_config_id": 5,
|
||||
"vision_llm_config_id": -1,
|
||||
},
|
||||
},
|
||||
version=3,
|
||||
)
|
||||
|
||||
def _raise(*, agent_llm_id, image_generation_config_id, vision_llm_config_id):
|
||||
raise AutomationModelPolicyError(
|
||||
[{"kind": "llm", "config_id": -7, "reason": "free model"}]
|
||||
)
|
||||
|
||||
monkeypatch.setattr(automation_mod, "assert_models_billable", _raise)
|
||||
|
||||
async def _noop_authorize(self, *_a, **_k):
|
||||
return None
|
||||
|
||||
async def _return_existing(self, _aid):
|
||||
return existing
|
||||
|
||||
monkeypatch.setattr(AutomationService, "_authorize", _noop_authorize)
|
||||
monkeypatch.setattr(
|
||||
AutomationService, "_get_with_triggers_or_raise", _return_existing
|
||||
)
|
||||
|
||||
service = _service(SimpleNamespace())
|
||||
patch = AutomationUpdate(
|
||||
definition=_definition(
|
||||
models=AutomationModels(
|
||||
agent_llm_id=-7,
|
||||
image_generation_config_id=5,
|
||||
vision_llm_config_id=-1,
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await service.update(7, patch)
|
||||
|
||||
assert exc_info.value.status_code == 422
|
||||
|
||||
|
||||
async def test_update_keeps_unchanged_models_without_revalidation(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""An unchanged models block is kept as-is and is NOT re-validated.
|
||||
|
||||
Lets users edit an automation whose captured model later drifted out of
|
||||
premium without an unrelated edit tripping the policy check.
|
||||
"""
|
||||
captured = {
|
||||
"agent_llm_id": -1,
|
||||
"image_generation_config_id": 5,
|
||||
"vision_llm_config_id": -1,
|
||||
}
|
||||
existing = SimpleNamespace(
|
||||
search_space_id=1,
|
||||
definition={"name": "A", "plan": [], "models": captured},
|
||||
version=3,
|
||||
)
|
||||
|
||||
def _fail(*_a, **_k):
|
||||
raise AssertionError("unchanged models must not be re-validated")
|
||||
|
||||
monkeypatch.setattr(automation_mod, "assert_models_billable", _fail)
|
||||
|
||||
async def _noop_authorize(self, *_a, **_k):
|
||||
return None
|
||||
|
||||
async def _return_existing(self, _aid):
|
||||
return existing
|
||||
|
||||
monkeypatch.setattr(AutomationService, "_authorize", _noop_authorize)
|
||||
monkeypatch.setattr(
|
||||
AutomationService, "_get_with_triggers_or_raise", _return_existing
|
||||
)
|
||||
|
||||
service = _service(SimpleNamespace())
|
||||
patch = AutomationUpdate(
|
||||
definition=_definition(models=AutomationModels(**captured))
|
||||
)
|
||||
|
||||
result = await service.update(7, patch)
|
||||
|
||||
assert result.definition["models"] == captured
|
||||
assert result.version == 4
|
||||
|
||||
|
||||
async def test_model_eligibility_authorizes_and_returns_payload(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue