railtracks

 42class Session:
 43    """
 44    The main class for managing an execution session.
 45
 46    This class is responsible for setting up all the necessary components for running a Railtracks execution, including the coordinator, publisher, and state management.
 47
 48    For the configuration parameters of the setting. It will follow this precedence:
 49    1. The parameters in the `Session` constructor.
 50    2. The parameters in global context variables.
 51    3. The default values.
 52
 53    Default Values:
 54    - `name`: None
 55    - `timeout`: 150.0 seconds
 56    - `end_on_error`: False
 57    - `logging_setting`: "INFO"
 58    - `log_file`: None (logs will not be written to a file)
 59    - `broadcast_callback`: None (no callback for broadcast messages)
 60    - `prompt_injection`: True (the prompt will be automatically injected from context variables)
 61    - `save_state`: True (the state of the execution will be saved to a file at the end of the run in the `.railtracks/data/sessions/` directory)
 62
 63
 64    Args:
 65        name (str | None, optional): Optional name for the session. This name will be included in the saved state file if `save_state` is True.
 66        context (Dict[str, Any], optional): A dictionary of global context variables to be used during the execution.
 67        timeout (float, optional): The maximum number of seconds to wait for a response to your top-level request.
 68        end_on_error (bool, optional): If True, the execution will stop when an exception is encountered.
 69        logging_setting (AllowableLogLevels, optional): The setting for the level of logging you would like to have. This will override the module-level logging settings for the duration of this session.
 70        log_file (str | os.PathLike | None, optional): The file to which the logs will be written.
 71        broadcast_callback (Callable[[str], None] | Callable[[str], Coroutine[None, None, None]] | None, optional): A callback function that will be called with the broadcast messages.
 72        prompt_injection (bool, optional): If True, the prompt will be automatically injected from context variables.
 73        save_state (bool, optional): If True, the state of the execution will be saved to a file at the end of the run in the `.railtracks/data/sessions/` directory.
 74    """
 75
 76    def __init__(
 77        self,
 78        context: Dict[str, Any] | None = None,
 79        *,
 80        name: str | None = None,
 81        timeout: float | None = None,
 82        end_on_error: bool | None = None,
 83        logging_setting: AllowableLogLevels | None = None,
 84        log_file: str | os.PathLike | None = None,
 85        broadcast_callback: (
 86            Callable[[str], None] | Callable[[str], Coroutine[None, None, None]] | None
 87        ) = None,
 88        prompt_injection: bool | None = None,
 89        save_state: bool | None = None,
 90    ):
 91        # first lets read from defaults if nessecary for the provided input config
 92
 93        self.executor_config = self.global_config_precedence(
 94            timeout=timeout,
 95            end_on_error=end_on_error,
 96            logging_setting=logging_setting,
 97            log_file=log_file,
 98            broadcast_callback=broadcast_callback,
 99            prompt_injection=prompt_injection,
100            save_state=save_state,
101        )
102
103        if context is None:
104            context = {}
105
106        self.name = name
107
108        self._has_custom_logging = logging_setting is not None or log_file is not None
109
110        if self._has_custom_logging:
111            mark_session_logging_override(
112                session_level=self.executor_config.logging_setting,
113                session_log_file=self.executor_config.log_file,
114            )
115
116        self.publisher: RTPublisher = RTPublisher()
117
118        self._identifier = str(uuid.uuid4())
119
120        executor_info = ExecutionInfo.create_new()
121        self.coordinator = Coordinator(
122            execution_modes={"async": AsyncioExecutionStrategy()}
123        )
124        self.rt_state = RTState(
125            executor_info, self.executor_config, self.coordinator, self.publisher
126        )
127
128        self.coordinator.start(self.publisher)
129        self._setup_subscriber()
130        register_globals(
131            session_id=self._identifier,
132            rt_publisher=self.publisher,
133            parent_id=None,
134            executor_config=self.executor_config,
135            global_context_vars=context,
136        )
137
138        self._start_time = time.time()
139
140        logger.debug("Session %s is initialized" % self._identifier)
141
142    @classmethod
143    def global_config_precedence(
144        cls,
145        timeout: float | None,
146        end_on_error: bool | None,
147        logging_setting: AllowableLogLevels | None,
148        log_file: str | os.PathLike | None,
149        broadcast_callback: (
150            Callable[[str], None] | Callable[[str], Coroutine[None, None, None]] | None
151        ),
152        prompt_injection: bool | None,
153        save_state: bool | None,
154    ) -> ExecutorConfig:
155        """
156        Uses the following precedence order to determine the configuration parameters:
157        1. The parameters in the method parameters.
158        2. The parameters in global context variables.
159        3. The default values.
160        """
161        global_executor_config = get_global_config()
162
163        return global_executor_config.precedence_overwritten(
164            timeout=timeout,
165            end_on_error=end_on_error,
166            logging_setting=logging_setting,
167            log_file=log_file,
168            subscriber=broadcast_callback,
169            prompt_injection=prompt_injection,
170            save_state=save_state,
171        )
172
173    def __enter__(self):
174        return self
175
176    def __exit__(self, exc_type, exc_val, exc_tb):
177        if self.executor_config.save_state:
178            try:
179                railtracks_dir = Path(".railtracks")
180                sessions_dir = railtracks_dir / "data" / "sessions"
181                sessions_dir.mkdir(
182                    parents=True, exist_ok=True
183                )  # Creates directory structure if doesn't exist, skips otherwise.
184
185                # Try to create file path with name, fallback to identifier only if there's an issue
186                try:
187                    file_path = (
188                        sessions_dir / f"{self.name}_{self._identifier}.json"
189                        if self.name
190                        else sessions_dir / f"{self._identifier}.json"
191                    )
192                    file_path.touch()
193                except FileNotFoundError:
194                    logger.warning(
195                        get_message(
196                            ExceptionMessageKey.INVALID_SESSION_FILE_NAME_WARN
197                        ).format(name=self.name, identifier=self._identifier)
198                    )
199                    file_path = sessions_dir / f"{self._identifier}.json"
200
201                logger.info("Saving execution info to %s" % file_path)
202
203                file_path.write_text(json.dumps(self.payload()))
204            except Exception as e:
205                logger.error(
206                    "Error while saving to execution info to file",
207                    exc_info=e,
208                )
209
210        self._close()
211
212    def _setup_subscriber(self):
213        """
214        Prepares and attaches the saved broadcast_callback to the publisher attached to this runner.
215        """
216
217        if self.executor_config.subscriber is not None:
218            self.publisher.subscribe(
219                stream_subscriber(self.executor_config.subscriber),
220                name="Streaming Subscriber",
221            )
222
223    def _close(self):
224        """
225        Closes the runner and cleans up all resources.
226
227        - Shuts down the state object
228        - Detaches logging handlers so they aren't duplicated
229        - Deletes all the global variables that were registered in the context
230        """
231        # the publisher should have already been closed in `_run_base`
232        self.rt_state.shutdown()
233
234        if self._has_custom_logging:
235            restore_module_logging()
236
237        delete_globals()
238        # by deleting all of the state variables we are ensuring that the next time we create a runner it is fresh
239
240    @property
241    def info(self) -> ExecutionInfo:
242        """
243        Returns the current state of the runner.
244
245        This is useful for debugging and viewing the current state of the run.
246        """
247        return self.rt_state.info
248
249    def payload(self) -> Dict[str, Any]:
250        """
251        Gets the complete json payload tied to this session.
252
253        The outputted json schema is maintained in (link here)
254        """
255        info = self.info
256
257        run_list = info.graph_serialization()
258
259        full_dict = {
260            "session_id": self._identifier,
261            "session_name": self.name,
262            "start_time": self._start_time,
263            "end_time": time.time(),
264            "runs": run_list,
265        }
266
267        return json.loads(json.dumps(full_dict))

328def session(
329    func: Callable[_P, Coroutine[Any, Any, _TOutput]] | None = None,
330    *,
331    name: str | None = None,
332    context: Dict[str, Any] | None = None,
333    timeout: float | None = None,
334    end_on_error: bool | None = None,
335    logging_setting: AllowableLogLevels | None = None,
336    log_file: str | os.PathLike | None = None,
337    broadcast_callback: (
338        Callable[[str], None] | Callable[[str], Coroutine[None, None, None]] | None
339    ) = None,
340    prompt_injection: bool | None = None,
341    save_state: bool | None = None,
342) -> (
343    Callable[_P, Coroutine[Any, Any, Tuple[_TOutput, Session]]]
344    | Callable[
345        [Callable[_P, Coroutine[Any, Any, _TOutput]]],
346        Callable[_P, Coroutine[Any, Any, Tuple[_TOutput, Session]]],
347    ]
348):
349    """
350    This decorator automatically creates and manages a Session context for the decorated function,
351    allowing async functions to use Railtracks operations without manually managing the session lifecycle.
352
353    Can be used as:
354    - @session (without parentheses) - uses default settings
355    - @session() (with empty parentheses) - uses default settings
356    - @session(name="my_task", timeout=30) (with configuration parameters)
357
358    When using this decorator, the function returns a tuple containing:
359    1. The original function's return value
360    2. The Session object used during execution
361
362    This allows access to session information (like execution state, logs, etc.) after the function completes,
363    while maintaining the simplicity of decorator usage.
364
365    Args:
366        name (str | None, optional): Optional name for the session. This name will be included in the saved state file if `save_state` is True.
367        context (Dict[str, Any], optional): A dictionary of global context variables to be used during the execution.
368        timeout (float, optional): The maximum number of seconds to wait for a response to your top-level request.
369        end_on_error (bool, optional): If True, the execution will stop when an exception is encountered.
370        logging_setting (AllowableLogLevels, optional): The setting for the level of logging you would like to have. This will override the module-level logging settings for the duration of this session.
371        log_file (str | os.PathLike | None, optional): The file to which the logs will be written.
372        broadcast_callback (Callable[[str], None] | Callable[[str], Coroutine[None, None, None]] | None, optional): A callback function that will be called with the broadcast messages.
373        prompt_injection (bool, optional): If True, the prompt will be automatically injected from context variables.
374        save_state (bool, optional): If True, the state of the execution will be saved to a file at the end of the run in the `.railtracks/data/sessions/` directory.
375
376    Returns:
377        When used as @session (without parentheses): Returns the decorated function that returns (result, session).
378        When used as @session(...) (with parameters): Returns a decorator function that takes an async function
379        and returns a new async function that returns (result, session).
380    """
381
382    def decorator(
383        target_func: Callable[_P, Coroutine[Any, Any, _TOutput]],
384    ) -> Callable[_P, Coroutine[Any, Any, Tuple[_TOutput, Session]]]:
385        # Validate that the decorated function is async
386        if not inspect.iscoroutinefunction(target_func):
387            raise TypeError(
388                f"@session decorator can only be applied to async functions. "
389                f"Function '{target_func.__name__}' is not async. "
390                f"Add 'async' keyword to your function definition."
391            )
392
393        @wraps(target_func)
394        async def wrapper(
395            *args: _P.args, **kwargs: _P.kwargs
396        ) -> Tuple[_TOutput, Session]:
397            session_obj = Session(
398                context=context,
399                timeout=timeout,
400                end_on_error=end_on_error,
401                logging_setting=logging_setting,
402                log_file=log_file,
403                broadcast_callback=broadcast_callback,
404                name=name,
405                prompt_injection=prompt_injection,
406                save_state=save_state,
407            )
408
409            with session_obj:
410                result = await target_func(*args, **kwargs)
411                return result, session_obj
412
413        return wrapper
414
415    # If used as @session without parentheses
416    if func is not None:
417        return decorator(func)
418
419    # If used as @session(...)
420    return decorator

 60async def call(
 61    node_: Callable[_P, Node[_TOutput]] | RTFunction[_P, _TOutput],
 62    *args: _P.args,
 63    **kwargs: _P.kwargs,
 64) -> _TOutput:
 65    """
 66    Call a node from within a node inside the framework. This will return a coroutine that you can interact with
 67    in whatever way using async/await logic.
 68
 69    Usage:
 70    ```python
 71    # for sequential operation
 72    result = await call(NodeA, "hello world", 42)
 73
 74    # for parallel operation
 75    tasks = [call(NodeA, "hello world", i) for i in range(10)]
 76    results = await asyncio.gather(*tasks)
 77    ```
 78
 79    Args:
 80        node: The node type you would like to create. This could be a function decorated with `@function_node`, a function, or a Node instance.
 81        *args: The arguments to pass to the node
 82        **kwargs: The keyword arguments to pass to the node
 83    """
 84    node: Callable[_P, Node[_TOutput]]
 85    # this entire section is a bit of a typing nightmare becuase all overloads we provide.
 86    if isinstance(node_, FunctionType):
 87        node = extract_node_from_function(node_)
 88    else:
 89        node = node_
 90
 91    # if the context is none then we will need to create a wrapper for the state object to work with.
 92    if not is_context_present():
 93        # we have to use lazy import here to prevent a circular import issue. This is a must have unfortunately.
 94        from railtracks import Session
 95
 96        with Session():
 97            result = await _start(node, args=args, kwargs=kwargs)
 98            return result
 99
100    # if the context is not active then we know this is the top level request
101    if not is_context_active():
102        result = await _start(node, args=args, kwargs=kwargs)
103        return result
104
105    # if the context is active then we can just run the node
106    result = await _run(node, args=args, kwargs=kwargs)
107    return result

 6async def broadcast(item: str):
 7    """
 8    Streams the given message
 9
10    This will trigger the broadcast_callback callback you have already provided.
11
12    Args:
13        item (str): The item you want to stream.
14    """
15    publisher = get_publisher()
16
17    await publisher.publish(Streaming(node_id=get_parent_id(), streamed_object=item))

27async def call_batch(
28    node: Callable[..., Node[_TOutput]]
29    | Callable[..., _TOutput]
30    | _AsyncNodeAttachedFunc[_P, _TOutput]
31    | _SyncNodeAttachedFunc[_P, _TOutput],
32    *iterables: Iterable[Any],
33    return_exceptions: bool = True,
34):
35    """
36    Complete a node over multiple iterables, allowing for parallel execution.
37
38    Note the results will be returned in the order of the iterables, not the order of completion.
39
40    If one of the nodes returns an exception, the thrown exception will be included as a response.
41
42    Args:
43        node: The node type to create.
44        *iterables: The iterables to map the node over.
45        return_exceptions: If True, exceptions will be returned as part of the results.
46            If False, exceptions will be raised immediately, and you will lose access to the results.
47            Defaults to true.
48
49    Returns:
50        An iterable of results from the node.
51
52    Usage:
53        ```python
54        results = await batch(NodeA, ["hello world"] * 10)
55        for result in results:
56            handle(result)
57        ```
58    """
59    # this is big typing disaster but there is no way around it. Try if if you want to.
60    contracts = [call(node, *args) for args in zip(*iterables)]
61
62    results = await asyncio.gather(*contracts, return_exceptions=return_exceptions)
63    return results

 19class ExecutionInfo:
 20    """
 21    A class that contains the full details of the state of a run at any given point in time.
 22
 23    The class is designed to be used as a snapshot of state that can be used to display the state of the run, or to
 24    create a graphical representation of the system.
 25    """
 26
 27    def __init__(
 28        self,
 29        request_forest: RequestForest,
 30        node_forest: NodeForest,
 31        stamper: StampManager,
 32    ):
 33        self.request_forest = request_forest
 34        self.node_forest = node_forest
 35        self.stamper = stamper
 36
 37    @classmethod
 38    def default(cls) -> ExecutionInfo:
 39        """Creates a new "empty" instance of the ExecutionInfo class with the default values."""
 40        return cls.create_new()
 41
 42    @classmethod
 43    def create_new(
 44        cls,
 45    ) -> ExecutionInfo:
 46        """
 47        Creates a new empty instance of state variables with the provided executor configuration.
 48
 49        """
 50        request_heap = RequestForest()
 51        node_heap = NodeForest()
 52        stamper = StampManager()
 53
 54        return ExecutionInfo(
 55            request_forest=request_heap,
 56            node_forest=node_heap,
 57            stamper=stamper,
 58        )
 59
 60    @property
 61    def answer(self):
 62        """Convenience method to access the answer of the run."""
 63        return self.request_forest.answer
 64
 65    @property
 66    def all_stamps(self) -> List[Stamp]:
 67        """Convenience method to access all the stamps of the run."""
 68        return self.stamper.all_stamps
 69
 70    @property
 71    def name(self):
 72        """
 73        Gets the name of the graph by pulling the name of the insertion request. It will raise a ValueError if the insertion
 74        request is not present or there are multiple insertion requests.
 75        """
 76        insertion_requests = self.insertion_requests
 77
 78        if len(insertion_requests) >= 2:
 79            raise ValueError(
 80                "You cannot get the name of a graph with multiple insertion requests"
 81            )
 82
 83        if len(insertion_requests) == 0:
 84            raise ValueError(
 85                "You cannot get the name of a graph with no insertion requests"
 86            )
 87
 88        i_r = insertion_requests[0]
 89
 90        return self.node_forest.get_node_type(i_r.sink_id).name()
 91
 92    @property
 93    def insertion_requests(self):
 94        """A convenience method to access all the insertion requests of the run."""
 95        return self.request_forest.insertion_request
 96
 97    def _get_info(self, ids: List[str] | str | None = None) -> ExecutionInfo:
 98        """
 99        Gets a subset of the current state based on the provided node ids. It will contain all the children of the provided node ids
100
101        Note: If no ids are provided, the full state is returned.
102
103        Args:
104            ids (List[str] | str | None): A list of node ids to filter the state by. If None, the full state is returned.
105
106        Returns:
107            ExecutionInfo: A new instance of ExecutionInfo containing only the children of the provided ids.
108
109        """
110        if ids is None:
111            return self
112        else:
113            # firstly lets
114            if isinstance(ids, str):
115                ids = [ids]
116
117            # we need to quickly check to make sure these ids are valid
118            for identifier in ids:
119                if identifier not in self.request_forest:
120                    raise ValueError(
121                        f"Identifier '{identifier}' not found in the current state."
122                    )
123
124            new_node_forest, new_request_forest = create_sub_state_info(
125                self.node_forest.heap(),
126                self.request_forest.heap(),
127                ids,
128            )
129            return ExecutionInfo(
130                node_forest=new_node_forest,
131                request_forest=new_request_forest,
132                stamper=self.stamper,
133            )
134
135    def _to_graph(self) -> Tuple[List[Vertex], List[Edge]]:
136        """
137        Converts the current state into its graph representation.
138
139        Returns:
140            List[Node]: An iterable of nodes in the graph.
141            List[Edge]: An iterable of edges in the graph.
142        """
143        return self.node_forest.to_vertices(), self.request_forest.to_edges()
144
145    def graph_serialization(self) -> dict[str, Any]:
146        """
147                Creates a string (JSON) representation of this info object designed to be used to construct a graph for this
148                info object.
149
150                Some important notes about its structure are outlined below:
151                - The `nodes` key contains a list of all the nodes in the graph, represented as `Vertex` objects.
152                - The `edges` key contains a list of all the edges in the graph, represented as `Edge` objects.
153                - The `stamps` key contains an ease of use list of all the stamps associated with the run, represented as `Stamp` objects.
154
155                - The "nodes" and "requests" key will be outlined with normal graph details like connections and identifiers in addition to a loose details object.
156                - However, both will carry an addition param called "stamp" which is a timestamp style object.
157                - They also will carry a "parent" param which is a recursive structure that allows you to traverse the graph in time.
158
159
160        ```
161        """
162        parent_nodes = [x.identifier for x in self.insertion_requests]
163
164        infos = [self._get_info(parent_node) for parent_node in parent_nodes]
165
166        runs = []
167
168        for info, parent_node_id in zip(infos, parent_nodes):
169            insertion_requests = info.request_forest.insertion_request
170
171            assert len(insertion_requests) == 1
172            parent_request = insertion_requests[0]
173
174            all_parents = parent_request.get_all_parents()
175
176            start_time = all_parents[-1].stamp.time
177
178            assert len([x for x in all_parents if x.status == "Completed"]) <= 1
179            end_time = None
180            for req in all_parents:
181                if req.status in ["Completed", "Failed"]:
182                    end_time = req.stamp.time
183                    break
184
185            entry = {
186                "name": info.name,
187                "run_id": parent_node_id,
188                "nodes": info.node_forest.to_vertices(),
189                "status": parent_request.status,
190                "edges": info.request_forest.to_edges(),
191                "steps": _get_stamps_from_forests(
192                    info.node_forest, info.request_forest
193                ),
194                "start_time": start_time,
195                "end_time": end_time,
196            }
197            runs.append(entry)
198
199        return json.loads(
200            json.dumps(
201                runs,
202                cls=RTJSONEncoder,
203            )
204        )

10class ExecutorConfig:
11    def __init__(
12        self,
13        *,
14        timeout: float = 150.0,
15        end_on_error: bool = False,
16        logging_setting: AllowableLogLevels = "INFO",
17        log_file: str | os.PathLike | None = None,
18        broadcast_callback: (
19            Callable[[str], None] | Callable[[str], Coroutine[None, None, None]] | None
20        ) = None,
21        prompt_injection: bool = True,
22        save_state: bool = True,
23    ):
24        """
25        ExecutorConfig is special configuration object designed to allow customization of the executor in the RT system.
26
27        Args:
28            timeout (float): The maximum number of seconds to wait for a response to your top level request
29            end_on_error (bool): If true, the executor will stop execution when an exception is encountered.
30            logging_setting (AllowableLogLevels): The setting for the level of logging you would like to have.
31            log_file (str | os.PathLike | None): The file to which the logs will be written. If None, no file will be created.
32            broadcast_callback (Callable or Coroutine): A function or coroutine that will handle streaming messages.
33            prompt_injection (bool): If true, prompts can be injected with global context
34            save_state (bool): If true, the state of the executor will be saved to disk.
35        """
36        self.timeout = timeout
37        self.end_on_error = end_on_error
38        self.logging_setting = logging_setting
39        self.subscriber = broadcast_callback
40        self.log_file = log_file
41        self.prompt_injection = prompt_injection
42        self.save_state = save_state
43
44    @property
45    def logging_setting(self) -> AllowableLogLevels:
46        return self._logging_setting
47
48    @logging_setting.setter
49    def logging_setting(self, value: AllowableLogLevels):
50        if value not in str_to_log_level:
51            raise ValueError(
52                f"logging_setting must be one of {str_to_log_level}, got {value}"
53            )
54        self._logging_setting: AllowableLogLevels = value
55
56    def precedence_overwritten(
57        self,
58        *,
59        timeout: float | None = None,
60        end_on_error: bool | None = None,
61        logging_setting: AllowableLogLevels | None = None,
62        log_file: str | os.PathLike | None = None,
63        subscriber: (
64            Callable[[str], None] | Callable[[str], Coroutine[None, None, None]] | None
65        ) = None,
66        prompt_injection: bool | None = None,
67        save_state: bool | None = None,
68    ):
69        """
70        If any of the parameters are provided (not None), it will create a new update the current instance with the new values and return a deep copied reference to it.
71        """
72        return ExecutorConfig(
73            timeout=timeout if timeout is not None else self.timeout,
74            end_on_error=end_on_error
75            if end_on_error is not None
76            else self.end_on_error,
77            logging_setting=logging_setting
78            if logging_setting is not None
79            else self.logging_setting,
80            log_file=log_file if log_file is not None else self.log_file,
81            broadcast_callback=subscriber
82            if subscriber is not None
83            else self.subscriber,
84            prompt_injection=prompt_injection
85            if prompt_injection is not None
86            else self.prompt_injection,
87            save_state=save_state if save_state is not None else self.save_state,
88        )
89
90    def __repr__(self):
91        return (
92            f"ExecutorConfig(timeout={self.timeout}, end_on_error={self.end_on_error}, "
93            f"logging_setting={self.logging_setting}, log_file={self.log_file}, "
94            f"prompt_injection={self.prompt_injection}, "
95            f"save_state={self.save_state})"
96        )