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`: "REGULAR"
 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` 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.
 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` 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        prepare_logger(
109            setting=self.executor_config.logging_setting,
110            path=self.executor_config.log_file,
111        )
112        self.publisher: RTPublisher = RTPublisher()
113
114        self._identifier = str(uuid.uuid4())
115
116        executor_info = ExecutionInfo.create_new()
117        self.coordinator = Coordinator(
118            execution_modes={"async": AsyncioExecutionStrategy()}
119        )
120        self.rt_state = RTState(
121            executor_info, self.executor_config, self.coordinator, self.publisher
122        )
123
124        self.coordinator.start(self.publisher)
125        self._setup_subscriber()
126        register_globals(
127            session_id=self._identifier,
128            rt_publisher=self.publisher,
129            parent_id=None,
130            executor_config=self.executor_config,
131            global_context_vars=context,
132        )
133
134        self._start_time = time.time()
135
136        logger.debug("Session %s is initialized" % self._identifier)
137
138    @classmethod
139    def global_config_precedence(
140        cls,
141        timeout: float | None,
142        end_on_error: bool | None,
143        logging_setting: AllowableLogLevels | None,
144        log_file: str | os.PathLike | None,
145        broadcast_callback: (
146            Callable[[str], None] | Callable[[str], Coroutine[None, None, None]] | None
147        ),
148        prompt_injection: bool | None,
149        save_state: bool | None,
150    ) -> ExecutorConfig:
151        """
152        Uses the following precedence order to determine the configuration parameters:
153        1. The parameters in the method parameters.
154        2. The parameters in global context variables.
155        3. The default values.
156        """
157        global_executor_config = get_global_config()
158
159        return global_executor_config.precedence_overwritten(
160            timeout=timeout,
161            end_on_error=end_on_error,
162            logging_setting=logging_setting,
163            log_file=log_file,
164            subscriber=broadcast_callback,
165            prompt_injection=prompt_injection,
166            save_state=save_state,
167        )
168
169    def __enter__(self):
170        return self
171
172    def __exit__(self, exc_type, exc_val, exc_tb):
173        if self.executor_config.save_state:
174            try:
175                railtracks_dir = Path(".railtracks")
176                railtracks_dir.mkdir(
177                    exist_ok=True
178                )  # Creates if doesn't exist, skips otherwise.
179
180                # Try to create file path with name, fallback to identifier only if there's an issue
181                try:
182                    file_path = (
183                        railtracks_dir / f"{self.name}_{self._identifier}.json"
184                        if self.name
185                        else railtracks_dir / f"{self._identifier}.json"
186                    )
187                    file_path.touch()
188                except FileNotFoundError:
189                    logger.warning(
190                        get_message(
191                            ExceptionMessageKey.INVALID_SESSION_FILE_NAME_WARN
192                        ).format(name=self.name, identifier=self._identifier)
193                    )
194                    file_path = railtracks_dir / f"{self._identifier}.json"
195
196                logger.info("Saving execution info to %s" % file_path)
197
198                file_path.write_text(json.dumps(self.payload()))
199            except Exception as e:
200                logger.error(
201                    "Error while saving to execution info to file",
202                    exc_info=e,
203                )
204
205        self._close()
206
207    def _setup_subscriber(self):
208        """
209        Prepares and attaches the saved broadcast_callback to the publisher attached to this runner.
210        """
211
212        if self.executor_config.subscriber is not None:
213            self.publisher.subscribe(
214                stream_subscriber(self.executor_config.subscriber),
215                name="Streaming Subscriber",
216            )
217
218    def _close(self):
219        """
220        Closes the runner and cleans up all resources.
221
222        - Shuts down the state object
223        - Detaches logging handlers so they aren't duplicated
224        - Deletes all the global variables that were registered in the context
225        """
226        # the publisher should have already been closed in `_run_base`
227        self.rt_state.shutdown()
228        detach_logging_handlers()
229        delete_globals()
230        # by deleting all of the state variables we are ensuring that the next time we create a runner it is fresh
231
232    @property
233    def info(self) -> ExecutionInfo:
234        """
235        Returns the current state of the runner.
236
237        This is useful for debugging and viewing the current state of the run.
238        """
239        return self.rt_state.info
240
241    def payload(self) -> Dict[str, Any]:
242        """
243        Gets the complete json payload tied to this session.
244
245        The outputted json schema is maintained in (link here)
246        """
247        info = self.info
248
249        run_list = info.graph_serialization()
250
251        full_dict = {
252            "session_id": self._identifier,
253            "session_name": self.name,
254            "start_time": self._start_time,
255            "end_time": time.time(),
256            "runs": run_list,
257        }
258
259        return json.loads(json.dumps(full_dict))

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

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