apunkt/SurfSense
1
0
Fork 0
mirror of https://github.com/MODSetter/SurfSense.git synced 2026-05-31 19:45:15 +02:00
Code Issues Projects Releases Packages Wiki Activity

refactor(automations): streamline model eligibility handling in automation creation

Browse source 
- 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 2026-05-29 20:27:40 -07:00
parent fade9d1b9d
commit 9d1a01eb0c
13 changed files with 887 additions and 263 deletions
Download patch file Download diff file Expand all files Collapse all files

28
surfsense_backend/app/agents/multi_agent_chat/main_agent/tools/automation/create.py
View file

@ -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)

63
surfsense_backend/app/automations/services/automation.py
View file

@ -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,