Update intermediate edit proposal model

src/bpmn_assistant/core/schemas.py

@@ -1,6 +1,6 @@
 from typing import List, Optional, Union
 
-from pydantic import BaseModel
+from pydantic import BaseModel, RootModel
 from typing_extensions import Literal
 
 TaskType = Literal["task", "userTask", "serviceTask"]
@@ -88,3 +88,22 @@ class ProcessModel(BaseModel):
     """
 
     process: List[BPMNElement]
+
+
+class EditProposal(BaseModel):
+    """
+    Represents an edit proposal for a BPMN process.
+    """
+
+    function: str
+    arguments: dict
+
+
+class StopSignal(BaseModel):
+    """
+    Represents a stop signal for the BPMN editing process.
+    """
+
+    stop: Literal[True]
+
+IntermediateEditProposal = RootModel[Union[EditProposal, StopSignal]]

src/bpmn_assistant/services/process_editing/bpmn_editing_service.py

@@ -1,9 +1,5 @@
-from typing import Literal, Optional
-
-from pydantic import BaseModel
-
 from bpmn_assistant.config import logger
-from bpmn_assistant.core import LLMFacade
+from bpmn_assistant.core import EditProposal, IntermediateEditProposal, LLMFacade
 from bpmn_assistant.core.exceptions import ProcessException
 from bpmn_assistant.prompts import PromptTemplateProcessor
 from bpmn_assistant.services.process_editing import (
@@ -16,17 +12,6 @@
 from bpmn_assistant.services.validate_bpmn import validate_element
 
 
-class EditProposal(BaseModel):
-    function: str
-    arguments: dict
-
-
-class IntermediateEditProposal(BaseModel):
-    function: str
-    arguments: dict
-    stop: Optional[Literal[True]] = None
-
-
 class BpmnEditingService:
     def __init__(self, llm_facade: LLMFacade, process: list, change_request: str):
         self.llm_facade = llm_facade