graphai-lib 0.0.9rc2__py3-none-any.whl → 0.0.10rc1__py3-none-any.whl

graphai/callback.py

@@ -1,13 +1,70 @@
 import asyncio
+from dataclasses import dataclass
+from enum import Enum
+import json
 from pydantic import Field
 from typing import Any
 from collections.abc import AsyncIterator
+import warnings
 
 
 log_stream = True
 
 
+class StrEnum(Enum):
+    def __str__(self) -> str:
+        return str(self.value)
+
+class GraphEventType(StrEnum):
+    START = "start"
+    END = "end"
+    START_NODE = "start_node"
+    END_NODE = "end_node"
+    CALLBACK = "callback"
+
+@dataclass
+class GraphEvent:
+    """A graph event emitted for specific graph events such as start node or end node,
+    and used by the callback to emit user-defined events.
+
+    :param type: The type of event, can be start_node, end_node, or callback.
+    :type type: GraphEventType
+    :param identifier: The identifier of the event, this is set typically by a callback
+        handler and can be used to distinguish between different events. For example, a
+        conversation/session ID could be used.
+    :type identifier: str
+    :param token: The token associated with the event, such as LLM streamed output.
+    :type token: str | None
+    :param params: The parameters associated with the event, such as tool call parameters
+        or event metadata.
+    :type params: dict[str, Any] | None
+    """
+    type: GraphEventType
+    identifier: str
+    token: str | None = None
+    params: dict[str, Any] | None = None
+
+    def encode(self, charset: str = "utf-8") -> bytes:
+        """Encodes the event as a JSON string, important for compatability with FastAPI
+        and starlette.
+
+        :param charset: The character set to use for encoding the event.
+        :type charset: str
+        """
+        event_dict = {
+            "type": self.type.value if hasattr(self.type, "value") else str(self.type),
+            "identifier": self.identifier,
+            "token": self.token,
+            "params": self.params,
+        }
+        data = f"data: {json.dumps(event_dict, ensure_ascii=False, separators=(',', ':'))}\n\n"
+        return data.encode(charset)
+
+
 class Callback:
+    """The original callback handler class. Outputs a stream of structured text
+    tokens. It is recommended to use the newer `EventCallback` handler instead.
+    """
     identifier: str = Field(
         default="graphai",
         description=(
@@ -67,6 +124,11 @@ class Callback:
         special_token_format: str = "<{identifier}:{token}:{params}>",
         token_format: str = "{token}",
     ):
+        warnings.warn(
+            "The `Callback` class is deprecated and will be removed in " +
+            "v0.1.0. Use the `EventCallback` class instead.",
+            DeprecationWarning
+        )
         self.identifier = identifier
         self.special_token_format = special_token_format
         self.token_format = token_format
@@ -198,3 +260,100 @@ class Callback:
         return self.special_token_format.format(
             identifier=identifier, token=name, params=params_str
         )
+
+
+class EventCallback(Callback):
+    """The event callback handler class. Outputs a stream of structured text
+    tokens. It is recommended to use the newer `EventCallback` handler instead.
+    """
+    def __init__(
+        self,
+        identifier: str = "graphai",
+        special_token_format: str | None = None,
+        token_format: str | None = None,
+    ):
+        warnings.warn(
+            "The `special_token_format` and `token_format` parameters are " +
+            "deprecated and will be removed in v0.1.0.",
+            DeprecationWarning
+        )
+        if special_token_format is None:
+            special_token_format = "<{identifier}:{token}:{params}>"
+        if token_format is None:
+            token_format = "{token}"
+        super().__init__(identifier, special_token_format, token_format)
+        self.events: list[GraphEvent] = []
+
+    def __call__(self, token: str, node_name: str | None = None):
+        if self._done:
+            raise RuntimeError("Cannot add tokens to a closed stream")
+        self._check_node_name(node_name=node_name)
+        event = GraphEvent(type=GraphEventType.CALLBACK, identifier=self.identifier, token=token, params=None)
+        # otherwise we just assume node is correct and send token
+        self.queue.put_nowait(event)
+
+    async def acall(self, token: str, node_name: str | None = None):
+        # TODO JB: do we need to have `node_name` param?
+        if self._done:
+            raise RuntimeError("Cannot add tokens to a closed stream")
+        self._check_node_name(node_name=node_name)
+        event = GraphEvent(type=GraphEventType.CALLBACK, identifier=self.identifier, token=token, params=None)
+        # otherwise we just assume node is correct and send token
+        self.queue.put_nowait(event)
+
+    async def aiter(self) -> AsyncIterator[GraphEvent]:  # type: ignore[override]
+        """Used by receiver to get the tokens from the stream queue. Creates
+        a generator that yields tokens from the queue until the END token is
+        received.
+        """
+        while True:  # Keep going until we see the END token
+            try:
+                if self._done and self.queue.empty():
+                    break
+                token = await self.queue.get()
+                yield token
+                self.queue.task_done()
+                if token.type == GraphEventType.END:
+                    break
+            except asyncio.CancelledError:
+                break
+        self._done = True  # Mark as done after processing all tokens
+
+    async def start_node(self, node_name: str, active: bool = True):
+        """Starts a new node and emits the start token."""
+        if self._done:
+            raise RuntimeError("Cannot start node on a closed stream")
+        self.current_node_name = node_name
+        if self.first_token:
+            self.first_token = False
+        self.active = active
+        if self.active:
+            token = GraphEvent(type=GraphEventType.START_NODE, identifier=self.identifier, token=self.current_node_name, params=None)
+            self.queue.put_nowait(token)
+
+    async def end_node(self, node_name: str):
+        """Emits the end token for the current node."""
+        if self._done:
+            raise RuntimeError("Cannot end node on a closed stream")
+        # self.current_node_name = node_name
+        if self.active:
+            token = GraphEvent(type=GraphEventType.END_NODE, identifier=self.identifier, token=self.current_node_name, params=None)
+            self.queue.put_nowait(token)
+
+    async def close(self):
+        """Close the stream and prevent further tokens from being added.
+        This will send an END token and set the done flag to True.
+        """
+        if self._done:
+            return
+        end_token = GraphEvent(type=GraphEventType.END, identifier=self.identifier)
+        self._done = True  # Set done before putting the end token
+        self.queue.put_nowait(end_token)
+        # Don't wait for queue.join() as it can cause deadlock
+        # The stream will close when aiter processes the END token
+
+    async def _build_special_token(
+        self, name: str, params: dict[str, Any] | None = None
+    ):
+        raise NotImplementedError("This method is not implemented for the `EventCallback` class.")
+

graphai/graph.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
-from typing import Any, Protocol 
+import asyncio
+from typing import Any, Iterable, Protocol
+from graphlib import TopologicalSorter, CycleError
 from graphai.callback import Callback
 from graphai.utils import logger
 
@@ -165,14 +167,14 @@ class Graph:
         self.end_node = node
         return self
 
-    def compile(self) -> "Graph":
+    def compile(self, *, strict: bool = False) -> Graph:
         """
         Validate the graph:
         - exactly one start node present (or Graph.start_node set)
         - at least one end node present
         - all edges reference known nodes
-        - no cycles
         - all nodes reachable from the start
+          (optional) **no cycles** when strict=True
         Returns self on success; raises GraphCompileError otherwise.
         """
         # nodes map
@@ -301,6 +303,17 @@ class Graph:
         if unreachable:
             raise GraphCompileError(f"Unreachable nodes: {unreachable}")
 
+        # optional cycle detection (strict mode)
+        if strict:
+            preds: dict[str, set[str]] = {n: set() for n in nodes.keys()}
+            for s, ds in adj.items():
+                for d in ds:
+                    preds[d].add(s)
+            try:
+                list(TopologicalSorter(preds).static_order())
+            except CycleError as e:
+                raise GraphCompileError("cycle detected in graph (strict mode)") from e
+
         return self
 
     def _validate_output(self, output: dict[str, Any], node_name: str):
@@ -362,6 +375,32 @@ class Graph:
             del state["callback"]
         return state
 
+    async def execute_many(
+        self, inputs: Iterable[dict[str, Any]], *, concurrency: int = 5
+    ) -> list[Any]:
+        """Execute the graph on many inputs concurrently.
+
+        :param inputs: An iterable of input dicts to feed into the graph.
+        :type inputs: Iterable[dict]
+        :param concurrency: Maximum number of graph executions to run at once.
+        :type concurrency: int
+        :param state: Optional shared state to pass to each execution.
+            If you want isolated state per execution, pass None
+            and the graph's normal semantics will apply.
+        :type state: Optional[Any]
+        :returns: The list of results in the same order as ``inputs``.
+        :rtype: list[Any]
+        """
+
+        sem = asyncio.Semaphore(concurrency)
+
+        async def _run_one(inp: dict[str, Any]) -> Any:
+            async with sem:
+                return await self.execute(input=inp)
+
+        tasks = [asyncio.create_task(_run_one(i)) for i in inputs]
+        return await asyncio.gather(*tasks)
+
     def get_callback(self):
         """Get a new instance of the callback class.
 

{graphai_lib-0.0.9rc2.dist-info → graphai_lib-0.0.10rc1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: graphai-lib
-Version: 0.0.9rc2
+Version: 0.0.10rc1
 Summary: Not an AI framework
 Requires-Python: <3.14,>=3.10
 Description-Content-Type: text/markdown

graphai_lib-0.0.10rc1.dist-info/RECORD

@@ -0,0 +1,11 @@
+graphai/__init__.py,sha256=UbqXq7iGIYe1GyTPcpgLSXbgWovggFsAbTMtr4JQm3M,160
+graphai/callback.py,sha256=xQhK1xbjkouemaIB5hjPZpo3thxiA6VVu9TpHslOsww,13851
+graphai/graph.py,sha256=CSdRYlQnfVHNw4XyKMEZ1rJqqg6rOx-aT8w9xumFhBs,20138
+graphai/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+graphai/utils.py,sha256=LIFg-fQalU9sB5DCuk6is48OdpEgNX95i9h-YddFbvM,11717
+graphai/nodes/__init__.py,sha256=IaMUryAqTZlcEqh-ZS6A4NIYG18JZwzo145dzxsYjAk,74
+graphai/nodes/base.py,sha256=SoKfOdRu5EIJ_z8xIz5zbNXcxPI2l9MKTQDeaQI-2no,7494
+graphai_lib-0.0.10rc1.dist-info/METADATA,sha256=2sIQkGfW7m0UoOEioVono53cVqpDkxRlxoDVlt2Tl1o,914
+graphai_lib-0.0.10rc1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+graphai_lib-0.0.10rc1.dist-info/top_level.txt,sha256=TXlqmhLViX-3xGH2g5w6cavRd-QMf229Hl88jdMOGt8,8
+graphai_lib-0.0.10rc1.dist-info/RECORD,,

graphai_lib-0.0.9rc2.dist-info/RECORD

@@ -1,11 +0,0 @@
-graphai/__init__.py,sha256=UbqXq7iGIYe1GyTPcpgLSXbgWovggFsAbTMtr4JQm3M,160
-graphai/callback.py,sha256=Wl0JCmE8NcFifKmP9-a5bFa0WKVHTdrSClHVRmIEPpc,7323
-graphai/graph.py,sha256=O0hZ_29ln8oTYnk5EX5bdhO6-Cr9pm-6IVCq-ZPWbgY,18526
-graphai/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-graphai/utils.py,sha256=LIFg-fQalU9sB5DCuk6is48OdpEgNX95i9h-YddFbvM,11717
-graphai/nodes/__init__.py,sha256=IaMUryAqTZlcEqh-ZS6A4NIYG18JZwzo145dzxsYjAk,74
-graphai/nodes/base.py,sha256=SoKfOdRu5EIJ_z8xIz5zbNXcxPI2l9MKTQDeaQI-2no,7494
-graphai_lib-0.0.9rc2.dist-info/METADATA,sha256=_3lwoAyLzSSgFq8GQHW_FAGd6A0kKiGQB3XWCnirJRs,913
-graphai_lib-0.0.9rc2.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-graphai_lib-0.0.9rc2.dist-info/top_level.txt,sha256=TXlqmhLViX-3xGH2g5w6cavRd-QMf229Hl88jdMOGt8,8
-graphai_lib-0.0.9rc2.dist-info/RECORD,,