soprano-sdk 0.1.94__py3-none-any.whl → 0.1.96__py3-none-any.whl

soprano_sdk/core/engine.py

@@ -0,0 +1,225 @@
+from typing import Optional, Dict, Any, Tuple
+
+import yaml
+from jinja2 import Environment
+from langgraph.checkpoint.memory import InMemorySaver
+from langgraph.constants import START
+from langgraph.graph import StateGraph
+from langgraph.graph.state import CompiledStateGraph
+
+from .constants import WorkflowKeys
+from .state import create_state_model
+from ..nodes.factory import NodeFactory
+from ..routing.router import WorkflowRouter
+from ..utils.function import FunctionRepository
+from ..utils.logger import logger
+from ..utils.tool import ToolRepository
+from ..validation import validate_workflow
+
+
+class WorkflowEngine:
+    def __init__(self, yaml_path: str, configs: dict):
+        self.yaml_path = yaml_path
+        self.configs = configs
+        logger.info(f"Loading workflow from: {yaml_path}")
+
+        try:
+            with open(yaml_path, 'r') as f:
+                self.config = yaml.safe_load(f)
+
+            logger.info("Validating workflow configuration")
+            validate_workflow(self.config)
+
+            self.workflow_name = self.config['name']
+            self.workflow_description = self.config['description']
+            self.workflow_version = self.config['version']
+            self.data_fields = self.config['data']
+            self.steps = self.config['steps']
+            self.outcomes = self.config['outcomes']
+            self.metadata = self.config.get('metadata', {})
+
+            self.StateType = create_state_model(self.data_fields)
+
+            self.step_map = {step['id']: step for step in self.steps}
+            self.outcome_map = {outcome['id']: outcome for outcome in self.outcomes}
+
+            self.function_repository = FunctionRepository()
+            self.tool_repository = None
+            if tool_config := self.config.get("tool_config"):
+                self.tool_repository = ToolRepository(tool_config)
+
+            self.context_store = {}
+            self.collect_input_fields = self._get_collect_input_fields()
+
+            logger.info(
+                f"Workflow loaded: {self.workflow_name} v{self.workflow_version} "
+                f"({len(self.steps)} steps, {len(self.outcomes)} outcomes)"
+            )
+
+        except Exception as e:
+            raise e
+
+    def get_config_value(self, key, default_value: Optional[Any]=None):
+        if value := self.configs.get(key) :
+            return value
+
+        if value := self.config.get(key) :
+            return value
+
+        return default_value
+
+    def _get_collect_input_fields(self) -> set:
+        fields = set()
+        for step in self.steps:
+            if step.get('action') == 'collect_input_with_agent' and (field := step.get('field')):
+                fields.add(field)
+        return fields
+
+    def update_context(self, context: Dict[str, Any]):
+        self.context_store.update(context)
+        logger.info(f"Context updated: {context}")
+
+    def remove_context_field(self, field_name: str):
+        if field_name in self.context_store:
+            del self.context_store[field_name]
+            logger.info(f"Removed context field: {field_name}")
+    
+    def get_context_value(self, field_name: str):
+        value = self.context_store.get(field_name, None)
+        if value is not None:
+            logger.info(f"Retrieved context value for '{field_name}': {value}")
+        return value
+
+    def build_graph(self, checkpointer=None):
+        logger.info("Building workflow graph")
+
+        try:
+            builder = StateGraph(self.StateType)
+
+            collector_nodes = []
+
+            logger.info("Adding nodes to graph")
+            for step in self.steps:
+                step_id = step['id']
+                action = step['action']
+
+                if action == 'collect_input_with_agent':
+                    collector_nodes.append(step_id)
+
+                node_fn = NodeFactory.create(step, engine_context=self)
+                builder.add_node(step_id, node_fn)
+
+                logger.info(f"Added node: {step_id} (action: {action})")
+
+            first_step_id = self.steps[0]['id']
+            builder.add_edge(START, first_step_id)
+            logger.info(f"Set entry point: {first_step_id}")
+
+            logger.info("Adding routing edges")
+            for step in self.steps:
+                step_id = step['id']
+
+                router = WorkflowRouter(step, self.step_map, self.outcome_map)
+                route_fn = router.create_route_function()
+                routing_map = router.get_routing_map(collector_nodes)
+
+                builder.add_conditional_edges(step_id, route_fn, routing_map)
+
+                logger.info(
+                    f"Added routing for {step_id}: {len(routing_map)} destinations"
+                )
+
+            if checkpointer is None:
+                checkpointer = InMemorySaver()
+                logger.info("Using InMemorySaver for state persistence")
+            else:
+                logger.info(f"Using custom checkpointer: {type(checkpointer).__name__}")
+
+            graph = builder.compile(checkpointer=checkpointer)
+
+            logger.info("Workflow graph built successfully")
+            return graph
+
+        except Exception as e:
+            raise RuntimeError(f"Failed to build workflow graph: {e}")
+
+    def get_outcome_message(self, state: Dict[str, Any]) -> str:
+        outcome_id = state.get(WorkflowKeys.OUTCOME_ID)
+        step_id = state.get(WorkflowKeys.STEP_ID)
+
+        outcome = self.outcome_map.get(outcome_id)
+        if outcome and 'message' in outcome:
+            message = outcome['message']
+            template_loader = self.get_config_value("template_loader", Environment())
+            message = template_loader.from_string(message).render(state)
+            logger.info(f"Outcome message generated in step {step_id}: {message}")
+            return message
+
+        if error := state.get("error"):
+            logger.info(f"Outcome error found in step {step_id}: {error}")
+            return f"{error}"
+
+        if message := state.get(WorkflowKeys.MESSAGES):
+            logger.info(f"Outcome message found in step {step_id}: {message}")
+            return f"{message}"
+
+        logger.error(f"No outcome message found in step {step_id}")
+        return "{'error': 'Unable to complete the request'}"
+
+    def get_step_info(self, step_id: str) -> Optional[Dict[str, Any]]:
+        return self.step_map.get(step_id)
+
+    def get_outcome_info(self, outcome_id: str) -> Optional[Dict[str, Any]]:
+        return self.outcome_map.get(outcome_id)
+
+    def list_steps(self) -> list:
+        return [step['id'] for step in self.steps]
+
+    def list_outcomes(self) -> list:
+        return [outcome['id'] for outcome in self.outcomes]
+
+    def get_workflow_info(self) -> Dict[str, Any]:
+        return {
+            'name': self.workflow_name,
+            'description': self.workflow_description,
+            'version': self.workflow_version,
+            'steps': len(self.steps),
+            'outcomes': len(self.outcomes),
+            'data_fields': [f['name'] for f in self.data_fields],
+            'metadata': self.metadata
+        }
+
+    def get_tool_policy(self) -> str:
+        tool_config = self.config.get('tool_config')
+        if not tool_config:
+            raise ValueError("Tool config is not provided in the YAML")
+        return tool_config.get('usage_policy')
+
+
+def load_workflow(yaml_path: str, checkpointer=None, config=None) -> Tuple[CompiledStateGraph, WorkflowEngine]:
+    """
+    Load a workflow from YAML configuration.
+
+    This is the main entry point for using the framework.
+
+    Args:
+        yaml_path: Path to the workflow YAML file
+        checkpointer: Optional checkpointer for state persistence.
+                     Defaults to InMemorySaver() if not provided.
+                     Example: MongoDBSaver for production persistence.
+
+    Returns:
+        Tuple of (compiled_graph, engine) where:
+        - compiled_graph: LangGraph ready for execution
+        - engine: WorkflowEngine instance for introspection
+
+    Example:
+        ```python
+        graph, engine = load_workflow("workflow.yaml")
+        result = graph.invoke({}, config={"configurable": {"thread_id": "123"}})
+        message = engine.get_outcome_message(result)
+        ```
+    """
+    engine = WorkflowEngine(yaml_path, configs=config)
+    graph = engine.build_graph(checkpointer=checkpointer)
+    return graph, engine

soprano_sdk/core/rollback_strategies.py

@@ -0,0 +1,259 @@
+import copy
+import logging
+import uuid
+from abc import ABC, abstractmethod
+from datetime import datetime
+from typing import Dict, Any, List
+
+from soprano_sdk.core.constants import WorkflowKeys, ActionType
+from ..utils.logger import logger
+
+class RollbackStrategy(ABC):
+    @abstractmethod
+    def rollback_to_node(
+        self,
+        state: Dict[str, Any],
+        target_node: str,
+        node_execution_order: List[str],
+        node_field_map: Dict[str, str],
+        workflow_steps: List[Dict[str, Any]]
+    ) -> Dict[str, Any]:
+        pass
+
+    @abstractmethod
+    def should_save_snapshot(self) -> bool:
+        pass
+
+    @abstractmethod
+    def save_snapshot(self, state: Dict[str, Any], node_id: str, execution_index: int) -> None:
+        pass
+
+    @abstractmethod
+    def get_strategy_name(self) -> str:
+        pass
+
+
+def _restore_from_snapshot(snapshot: Dict[str, Any]) -> Dict[str, Any]:
+    return copy.deepcopy(snapshot['state'])
+
+
+def _clear_future_executions(
+        state: Dict[str, Any],
+    target_node: str,
+    workflow_steps: List[Dict[str, Any]]
+) -> Dict[str, Any]:
+    target_step_index = next(
+        (i for i, step in enumerate(workflow_steps) if step['id'] == target_node),
+        None
+    )
+
+    if target_step_index is None:
+        logger.warning(f"Target node {target_node} not found in workflow steps")
+        return state
+
+    future_steps = workflow_steps[target_step_index:]
+
+    logger.info(f"Future steps to clear: {[s['id'] for s in future_steps]}")
+
+    for step in future_steps:
+        action = step.get('action')
+
+        if action == ActionType.COLLECT_INPUT_WITH_AGENT.value:
+            field_name = step.get('field')
+            if field_name:
+                state[field_name] = None
+
+                conv_key = f"{field_name}_conversation"
+                conversations = state.get(WorkflowKeys.CONVERSATIONS, {})
+                if conv_key in conversations:
+                    del conversations[conv_key]
+                    logger.info(f"Cleared conversation: {conv_key}")
+
+        elif action == ActionType.CALL_FUNCTION.value:
+            output_field = step.get('output')
+            if output_field:
+                state[output_field] = None
+                logger.info(f"Cleared computed field: {output_field}")
+
+    return state
+
+
+class HistoryBasedRollback(RollbackStrategy):
+    def get_strategy_name(self) -> str:
+        return "history_based"
+
+    def should_save_snapshot(self) -> bool:
+        return True
+
+    def save_snapshot(self, state: Dict[str, Any], node_id: str, execution_index: int) -> None:
+        state_history = state.get(WorkflowKeys.STATE_HISTORY, [])
+
+        snapshot = {
+            'snapshot_id': str(uuid.uuid4()),
+            'node_about_to_execute': node_id,
+            'execution_index': execution_index,
+            'timestamp': datetime.now().isoformat(),
+            'state': copy.deepcopy(state),
+        }
+
+        state_history.append(snapshot)
+        state[WorkflowKeys.STATE_HISTORY] = state_history
+
+        logger.info(f"Saved snapshot #{len(state_history)-1} before executing {node_id}")
+
+    def rollback_to_node(
+        self,
+        state: Dict[str, Any],
+        target_node: str,
+        node_execution_order: List[str],
+        node_field_map: Dict[str, str],
+        workflow_steps: List[Dict[str, Any]]
+    ) -> Dict[str, Any]:
+        state_history = state.get(WorkflowKeys.STATE_HISTORY, [])
+
+        if not state_history:
+            logger.warning("No state history available for rollback")
+            return {}
+
+        logger.info(f"Looking for snapshot before node '{target_node}'")
+
+        target_snapshot = None
+        target_index = None
+
+        for i, snapshot in enumerate(state_history):
+            if snapshot.get('node_about_to_execute') == target_node:
+                target_snapshot = snapshot
+                target_index = i
+                break
+
+        if target_snapshot is None:
+            logger.warning(f"No snapshot found before node '{target_node}'")
+            return {}
+
+        logger.info(f"Found snapshot at index {target_index}")
+        restored_state = _restore_from_snapshot(target_snapshot)
+
+        restored_state = _clear_future_executions(
+            restored_state,
+            target_node,
+            workflow_steps
+        )
+
+        restored_state[WorkflowKeys.STATE_HISTORY] = state_history[:target_index + 1]
+
+        logger.info(f"Successfully rolled back to {target_node}")
+
+        return restored_state
+
+
+def _build_dependency_graph(
+    workflow_steps: List[Dict[str, Any]]
+) -> Dict[str, List[str]]:
+    graph = {}
+
+    for step in workflow_steps:
+        field = step.get('field') or step.get('output')
+
+        if not field:
+            continue
+
+        depends_on = step.get('depends_on')
+
+        if depends_on:
+            if isinstance(depends_on, str):
+                depends_on_list = [depends_on]
+            elif isinstance(depends_on, list):
+                depends_on_list = depends_on
+            else:
+                logger.warning(f"Invalid depends_on type for field '{field}': {type(depends_on)}")
+                depends_on_list = []
+
+            for parent_field in depends_on_list:
+                if parent_field not in graph:
+                    graph[parent_field] = []
+                if field not in graph[parent_field]:
+                    graph[parent_field].append(field)
+
+        if field not in graph:
+            graph[field] = []
+
+    return graph
+
+
+def _find_all_dependents(
+    field: str,
+    dependency_graph: Dict[str, List[str]]
+) -> set:
+    all_dependents = set()
+    visited = set()
+
+    def _recurse(current_field: str):
+        if current_field in visited:
+            return
+        visited.add(current_field)
+
+        direct_dependents = dependency_graph.get(current_field, [])
+
+        for dependent in direct_dependents:
+            all_dependents.add(dependent)
+            _recurse(dependent)
+
+    _recurse(field)
+
+    return all_dependents
+
+
+def _clear_field_conversation(state: Dict[str, Any], field: str) -> None:
+    conv_key = f"{field}_conversation"
+    conversations = state.get(WorkflowKeys.CONVERSATIONS, {})
+
+    if conv_key in conversations:
+        del conversations[conv_key]
+        logger.info(f"Cleared conversation: {conv_key}")
+
+
+class DependencyBasedRollback(RollbackStrategy):
+    def get_strategy_name(self) -> str:
+        return "dependency_based"
+
+    def should_save_snapshot(self) -> bool:
+        return False
+
+    def save_snapshot(self, state: Dict[str, Any], node_id: str, execution_index: int) -> None:
+        return None
+
+    def rollback_to_node(
+        self,
+        state: Dict[str, Any],
+        target_node: str,
+        node_execution_order: List[str],
+        node_field_map: Dict[str, str],
+        workflow_steps: List[Dict[str, Any]]
+    ) -> Dict[str, Any]:
+        target_field = node_field_map.get(target_node)
+
+        if not target_field:
+            logger.warning(f"No field found for target node '{target_node}'")
+            return state
+
+        logger.info(f"Rolling back to node '{target_node}' (field: '{target_field}')")
+
+        dependency_graph = _build_dependency_graph(workflow_steps)
+
+        logger.info(f"Dependency graph: {dependency_graph}")
+
+        dependent_fields = _find_all_dependents(target_field, dependency_graph)
+
+        logger.info(f"Fields dependent on '{target_field}': {dependent_fields}")
+
+        state[target_field] = None
+        _clear_field_conversation(state, target_field)
+
+        for field in dependent_fields:
+            state[field] = None
+            _clear_field_conversation(state, field)
+            logger.info(f"Cleared dependent field: {field}")
+
+        logger.info(f"Successfully rolled back to {target_node} using dependency graph")
+
+        return state

soprano_sdk/core/state.py

@@ -0,0 +1,71 @@
+import types
+from typing import Annotated, Optional, Dict, List, Any, Type
+
+from typing_extensions import TypedDict
+
+from .constants import DataType
+
+
+def replace(_, right):
+    return right
+
+
+def get_state_value(state: Dict[str, Any], key: str, default: Any = None) -> Any:
+    return state.get(key, default)
+
+def set_state_value(state: Dict[str, Any], key: str, value: Any) -> None:
+    state[key] = value
+
+def create_state_model(data_fields: List[dict]):
+    type_mapping = {
+        DataType.TEXT.value: Optional[str],
+        DataType.NUMBER.value: Optional[int],
+        DataType.DOUBLE.value: Optional[float],
+        DataType.BOOLEAN.value: Optional[bool],
+        DataType.LIST.value: Optional[List[Any]],
+        DataType.DICT.value: Optional[Dict[str, Any]],
+        DataType.ANY.value: Optional[Any]
+    }
+
+    fields = {}
+    for field_def in data_fields:
+        field_name = field_def['name']
+        field_type = field_def['type']
+        python_type = type_mapping.get(field_type, Optional[str])
+        fields[field_name] = Annotated[python_type, replace]
+
+    fields['_step_id'] = Annotated[Optional[str], replace]
+    fields['_status'] = Annotated[Optional[str], replace]
+    fields['_outcome_id'] = Annotated[Optional[str], replace]
+
+    fields['_messages'] = Annotated[List[str], replace]
+    fields['_conversations'] = Annotated[Dict[str, List[Dict[str, str]]], replace]
+    fields['_state_history'] = Annotated[List[Dict[str, Any]], replace]
+    fields['_collector_nodes'] = Annotated[Dict[str, str], replace]
+    fields['_attempt_counts'] = Annotated[Dict[str, int], replace]
+    fields['_node_execution_order'] = Annotated[List[str], replace]
+    fields['_node_field_map'] = Annotated[Dict[str, str], replace]
+    fields['_computed_fields'] = Annotated[List[str], replace]
+    fields['error'] = Annotated[Optional[Dict[str, str]], replace]
+
+    return types.new_class('WorkflowState', (TypedDict,), {}, lambda ns: ns.update({'__annotations__': fields}))
+
+
+def initialize_state(state: Dict[str, Any]) -> Dict[str, Any]:
+    fields_to_initialize = {
+        '_state_history': [],
+        '_collector_nodes': {},
+        '_conversations': {},
+        '_messages': [],
+        '_attempt_counts': {},
+        '_node_execution_order': [],
+        '_node_field_map': {},
+        '_computed_fields': [],
+        'error': None
+    }
+
+    for field_name, default_value in fields_to_initialize.items():
+        if not get_state_value(state, field_name):
+            set_state_value(state, field_name, default_value.copy() if isinstance(default_value, (list, dict)) else default_value)
+
+    return state