graphai-lib 0.0.1__py3-none-any.whl → 0.0.3__py3-none-any.whl

graphai/callback.py

@@ -1,4 +1,5 @@
 import asyncio
+from pydantic import Field
 from typing import Optional
 from collections.abc import AsyncIterator
 from semantic_router.utils.logger import logger
@@ -7,20 +8,113 @@ from semantic_router.utils.logger import logger
 log_stream = True
 
 class Callback:
-    first_token = True
-    current_node_name: Optional[str] = None
-    active: bool = True
+    identifier: str = Field(
+        default="graphai",
+        description=(
+            "The identifier for special tokens. This allows us to easily "
+            "identify special tokens in the stream so we can handle them "
+            "correctly in any downstream process."
+        )
+    )
+    special_token_format: str = Field(
+        default="<{identifier}:{token}:{params}>",
+        description=(
+            "The format for special tokens. This is used to format special "
+            "tokens so they can be easily identified in the stream. "
+            "The format is a string with three possible components:\n"
+            "- {identifier}: An identifier shared by all special tokens, "
+            "by default this is 'graphai'.\n"
+            "- {token}: The special token type to be streamed. This may "
+            "be a tool name, identifier for start/end nodes, etc.\n"
+            "- {params}: Any additional parameters to be streamed. The parameters "
+            "are formatted as a comma-separated list of key-value pairs."
+        ),
+        examples=[
+            "<{identifier}:{token}:{params}>",
+            "<[{identifier} | {token} | {params}]>",
+            "<{token}:{params}>"
+        ]
+    )
+    token_format: str = Field(
+        default="{token}",
+        description=(
+            "The format for streamed tokens. This is used to format the "
+            "tokens typically returned from LLMs. By default, no special "
+            "formatting is applied."
+        )
+    )
+    _first_token: bool = Field(
+        default=True,
+        description="Whether this is the first token in the stream.",
+        exclude=True
+    )
+    _current_node_name: Optional[str] = Field(
+        default=None,
+        description="The name of the current node.",
+        exclude=True
+    )
+    _active: bool = Field(
+        default=True,
+        description="Whether the callback is active.",
+        exclude=True
+    )
+    _done: bool = Field(
+        default=False,
+        description="Whether the stream is done and should be closed.",
+        exclude=True
+    )
     queue: asyncio.Queue
 
-    def __init__(self):
+    def __init__(
+        self,
+        identifier: str = "graphai",
+        special_token_format: str = "<{identifier}:{token}:{params}>",
+        token_format: str = "{token}",
+    ):
+        self.identifier = identifier
+        self.special_token_format = special_token_format
+        self.token_format = token_format
         self.queue = asyncio.Queue()
+        self._done = False
+        self._first_token = True
+        self._current_node_name = None
+        self._active = True
+
+    @property
+    def first_token(self) -> bool:
+        return self._first_token
+    
+    @first_token.setter
+    def first_token(self, value: bool):
+        self._first_token = value
+
+    @property
+    def current_node_name(self) -> Optional[str]:
+        return self._current_node_name
+    
+    @current_node_name.setter
+    def current_node_name(self, value: Optional[str]):
+        self._current_node_name = value
+
+    @property
+    def active(self) -> bool:
+        return self._active
+    
+    @active.setter
+    def active(self, value: bool):
+        self._active = value
 
     def __call__(self, token: str, node_name: Optional[str] = None):
+        if self._done:
+            raise RuntimeError("Cannot add tokens to a closed stream")
         self._check_node_name(node_name=node_name)
         # otherwise we just assume node is correct and send token
         self.queue.put_nowait(token)
     
     async def acall(self, token: str, node_name: Optional[str] = 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)
         # otherwise we just assume node is correct and send token
         self.queue.put_nowait(token)
@@ -30,29 +124,72 @@ class Callback:
         a generator that yields tokens from the queue until the END token is
         received.
         """
-        while True:
-            token = await self.queue.get()
-            yield token
-            self.queue.task_done()
-            if token == "<graphai:END>":
+        end_token = await self._build_special_token(
+            name="END",
+            params=None
+        )
+        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 == end_token:
+                    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:
-            # TODO JB: not sure if we need self.first_token
             self.first_token = False
         self.active = active
         if self.active:
-            self.queue.put_nowait(f"<graphai:start:{node_name}>")
+            token = await self._build_special_token(
+                name=f"{self.current_node_name}:start",
+                params=None
+            )
+            self.queue.put_nowait(token)
+            # TODO JB: should we use two tokens here?
+            node_token = await self._build_special_token(
+                name=self.current_node_name,
+                params=None
+            )
+            self.queue.put_nowait(node_token)
     
     async def end_node(self, node_name: str):
-        self.current_node_name = None
+        """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:
-            self.queue.put_nowait(f"<graphai:end:{node_name}>")
+            node_token = await self._build_special_token(
+                name=f"{self.current_node_name}:end",
+                params=None
+            )
+            self.queue.put_nowait(node_token)
 
     async def close(self):
-        self.queue.put_nowait("<graphai:END>")
+        """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 = await self._build_special_token(
+            name="END",
+            params=None
+        )
+        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
 
     def _check_node_name(self, node_name: Optional[str] = None):
         if node_name:
@@ -60,4 +197,19 @@ class Callback:
             if self.current_node_name != node_name:
                 raise ValueError(
                     f"Node name mismatch: {self.current_node_name} != {node_name}"
-                )
+                )
+            
+    async def _build_special_token(self, name: str, params: dict[str, any] | None = None):
+        if params:
+            params_str = ",".join([f"{k}={v}" for k, v in params.items()])
+        else:
+            params_str = ""
+        if self.identifier:
+            identifier = self.identifier
+        else:
+            identifier = ""
+        return self.special_token_format.format(
+            identifier=identifier,
+            token=name,
+            params=params_str
+        )

graphai/graph.py

@@ -1,21 +1,41 @@
-from typing import List, Dict, Any
+from typing import List, Dict, Any, Optional
 from graphai.nodes.base import _Node
 from graphai.callback import Callback
 from semantic_router.utils.logger import logger
 
 
 class Graph:
-    def __init__(self, max_steps: int = 10):
-        self.nodes = []
-        self.edges = []
-        self.start_node = None
-        self.end_nodes = []
+    def __init__(self, max_steps: int = 10, initial_state: Optional[Dict[str, Any]] = None):
+        self.nodes: Dict[str, _Node] = {}
+        self.edges: List[Any] = []
+        self.start_node: Optional[_Node] = None
+        self.end_nodes: List[_Node] = []
         self.Callback = Callback
         self.callback = None
         self.max_steps = max_steps
+        self.state = initial_state or {}
+
+    # Allow getting and setting the graph's internal state
+    def get_state(self) -> Dict[str, Any]:
+        """Get the current graph state."""
+        return self.state
+
+    def set_state(self, state: Dict[str, Any]):
+        """Set the graph state."""
+        self.state = state
+
+    def update_state(self, values: Dict[str, Any]):
+        """Update the graph state with new values."""
+        self.state.update(values)
+
+    def reset_state(self):
+        """Reset the graph state to an empty dict."""
+        self.state = {}
 
     def add_node(self, node):
-        self.nodes.append(node)
+        if node.name in self.nodes:
+            raise Exception(f"Node with name '{node.name}' already exists.")
+        self.nodes[node.name] = node
         if node.is_start:
             if self.start_node is not None:
                 raise Exception(
@@ -27,10 +47,37 @@ class Graph:
         if node.is_end:
             self.end_nodes.append(node)
 
-    def add_edge(self, source: _Node, destination: _Node):
-        # TODO add logic to check that source and destination are nodes
-        # and they exist in the graph object already
-        edge = Edge(source, destination)
+    def add_edge(self, source: _Node | str, destination: _Node | str):
+        """Adds an edge between two nodes that already exist in the graph.
+        
+        Args:
+            source: The source node or its name.
+            destination: The destination node or its name.
+        """
+        source_node, destination_node = None, None
+        # get source node from graph
+        if isinstance(source, str):
+            source_node = self.nodes.get(source)
+        else:
+            # Check if it's a node-like object by looking for required attributes
+            if hasattr(source, 'name'):
+                source_node = self.nodes.get(source.name)
+        if source_node is None:
+            raise ValueError(
+                f"Node with name '{source.name if hasattr(source, 'name') else source}' not found."
+            )
+        # get destination node from graph
+        if isinstance(destination, str):
+            destination_node = self.nodes.get(destination)
+        else:
+            # Check if it's a node-like object by looking for required attributes
+            if hasattr(destination, 'name'):
+                destination_node = self.nodes.get(destination.name)
+        if destination_node is None:
+            raise ValueError(
+                f"Node with name '{destination.name if hasattr(destination, 'name') else destination}' not found."
+            )
+        edge = Edge(source_node, destination_node)
         self.edges.append(edge)
 
     def add_router(self, sources: list[_Node], router: _Node, destinations: List[_Node]):
@@ -71,17 +118,20 @@ class Graph:
             self.callback = self.get_callback()
         current_node = self.start_node
         state = input
+        # Don't reset the graph state if it was initialized with initial_state
         steps = 0
         while True:
             # we invoke the node here
             if current_node.stream:
                 # add callback tokens and param here if we are streaming
                 await self.callback.start_node(node_name=current_node.name)
-                output = await current_node.invoke(input=state, callback=self.callback)
+                # Include graph's internal state in the node execution context
+                output = await current_node.invoke(input=state, callback=self.callback, state=self.state)
                 self._validate_output(output=output, node_name=current_node.name)
                 await self.callback.end_node(node_name=current_node.name)
             else:
-                output = await current_node.invoke(input=state)
+                # Include graph's internal state in the node execution context
+                output = await current_node.invoke(input=state, state=self.state)
                 self._validate_output(output=output, node_name=current_node.name)
             # add output to state
             state = {**state, **output}
@@ -113,10 +163,21 @@ class Graph:
         return self.callback
 
     def _get_node_by_name(self, node_name: str) -> _Node:
-        for node in self.nodes:
-            if node.name == node_name:
-                return node
-        raise Exception(f"Node with name {node_name} not found.")
+        """Get a node by its name.
+        
+        Args:
+            node_name: The name of the node to find.
+            
+        Returns:
+            The node with the given name.
+            
+        Raises:
+            Exception: If no node with the given name is found.
+        """
+        node = self.nodes.get(node_name)
+        if node is None:
+            raise Exception(f"Node with name {node_name} not found.")
+        return node
 
     def _get_next_node(self, current_node):
         for edge in self.edges:
@@ -139,7 +200,7 @@ class Graph:
 
         G = nx.DiGraph()
 
-        for node in self.nodes:
+        for node in self.nodes.values():
             G.add_node(node.name)
 
         for edge in self.edges:
@@ -173,10 +234,11 @@ class Graph:
                     pos[node] = (pos[node][0] - x_center, pos[node][1])
 
             # Scale the layout
-            max_x = max(abs(p[0]) for p in pos.values())
-            max_y = max(abs(p[1]) for p in pos.values())
-            scale = min(0.8 / max_x, 0.8 / max_y)
-            pos = {node: (x * scale, y * scale) for node, (x, y) in pos.items()}
+            max_x = max(abs(p[0]) for p in pos.values()) if pos else 1
+            max_y = max(abs(p[1]) for p in pos.values()) if pos else 1
+            if max_x > 0 and max_y > 0:
+                scale = min(0.8 / max_x, 0.8 / max_y)
+                pos = {node: (x * scale, y * scale) for node, (x, y) in pos.items()}
 
         else:
             print("Warning: The graph contains cycles. Visualization will use a spring layout.")

graphai_lib-0.0.3.dist-info/METADATA

@@ -0,0 +1,32 @@
+Metadata-Version: 2.4
+Name: graphai-lib
+Version: 0.0.3
+Summary: Not an AI framework
+Requires-Python: <3.14,>=3.10
+Description-Content-Type: text/markdown
+Requires-Dist: semantic-router>=0.1.5
+Requires-Dist: networkx>=3.4.2
+Requires-Dist: matplotlib>=3.10.0
+Provides-Extra: dev
+Requires-Dist: ipykernel>=6.25.0; extra == "dev"
+Requires-Dist: ruff>=0.1.5; extra == "dev"
+Requires-Dist: pytest>=8.2.0; extra == "dev"
+Requires-Dist: pytest-mock>=3.12.0; extra == "dev"
+Requires-Dist: pytest-cov>=4.1.0; extra == "dev"
+Requires-Dist: pytest-xdist>=3.5.0; extra == "dev"
+Requires-Dist: pytest-asyncio>=0.24.0; extra == "dev"
+Requires-Dist: mypy>=1.7.1; extra == "dev"
+Requires-Dist: black[jupyter]<24.5.0,>=23.12.1; extra == "dev"
+
+# Philosophy
+
+1. Async-first
+2. Minimize abstractions
+3. One way to do one thing
+4. Graph-based AI
+
+## Installation
+
+```
+pip install -qU graphai-lib
+```

graphai_lib-0.0.3.dist-info/RECORD

@@ -0,0 +1,8 @@
+graphai/__init__.py,sha256=EHigFOWewDXLZXbdfjZH9kdLPhw6NT0ChS77lNAVAA8,109
+graphai/callback.py,sha256=K-h44pyL2VLXwJzIB_bcVYp5R6xv8zNca5FmN6994Uk,7598
+graphai/graph.py,sha256=EALHEhbXAaJmTvm7cL3Tdh0moRIw7lIyJDCNnCts2QA,10335
+graphai/utils.py,sha256=zrgpk82rIn7lwh631KhN-OgMAJMdbm0k5GPL1eMf2sQ,4522
+graphai_lib-0.0.3.dist-info/METADATA,sha256=fOjGeptDi6YD-fYZD9NtFv4_lzLWjMtvBnfwoQ9zLfo,879
+graphai_lib-0.0.3.dist-info/WHEEL,sha256=1tXe9gY0PYatrMPMDd6jXqjfpz_B-Wqm32CPfRC58XU,91
+graphai_lib-0.0.3.dist-info/top_level.txt,sha256=TXlqmhLViX-3xGH2g5w6cavRd-QMf229Hl88jdMOGt8,8
+graphai_lib-0.0.3.dist-info/RECORD,,

{graphai_lib-0.0.1.dist-info → graphai_lib-0.0.3.dist-info}/WHEEL

@@ -1,4 +1,5 @@
 Wheel-Version: 1.0
-Generator: poetry-core 1.9.1
+Generator: setuptools (77.0.3)
 Root-Is-Purelib: true
 Tag: py3-none-any
+

graphai_lib-0.0.3.dist-info/top_level.txt

@@ -0,0 +1 @@
+graphai

graphai/nodes/__init__.py

@@ -1,3 +0,0 @@
-from graphai.nodes.base import node, router
-
-__all__ = ["node", "router"]

graphai/nodes/base.py

@@ -1,148 +0,0 @@
-import inspect
-from typing import Any, Callable, Dict, Optional
-
-from graphai.callback import Callback
-from graphai.utils import FunctionSchema
-
-
-class NodeMeta(type):
-    @staticmethod
-    def positional_to_kwargs(cls_type, args) -> Dict[str, Any]:
-        init_signature = inspect.signature(cls_type.__init__)
-        init_params = {name: arg for name, arg in init_signature.parameters.items() if name != "self"}
-        return init_params
-
-    def __call__(cls, *args, **kwargs):
-        named_positional_args = NodeMeta.positional_to_kwargs(cls, args)
-        kwargs.update(named_positional_args)
-        return super().__call__(**kwargs)
-
-
-class _Node:
-    def __init__(
-        self,
-        is_router: bool = False,
-    ):
-        self.is_router = is_router
-
-    def _node(
-        self,
-        func: Callable,
-        start: bool = False,
-        end: bool = False,
-        stream: bool = False,
-    ) -> Callable:
-        """Decorator validating node structure.
-        """
-        if not callable(func):
-            raise ValueError("Node must be a callable function.")
-        
-        func_signature = inspect.signature(func)
-        schema = FunctionSchema(func)
-
-        class NodeClass:
-            _func_signature = func_signature
-            is_router = None
-            _stream = stream
-
-            def __init__(self):
-                self._expected_params = set(self._func_signature.parameters.keys())
-
-            async def execute(self, *args, **kwargs):
-                # Prepare arguments, including callback if stream is True
-                params_dict = await self._parse_params(*args, **kwargs)
-                return await func(**params_dict)  # Pass only the necessary arguments
-
-            async def _parse_params(self, *args, **kwargs) -> Dict[str, Any]:
-                # filter out unexpected keyword args
-                expected_kwargs = {k: v for k, v in kwargs.items() if k in self._expected_params}
-                # Convert args to kwargs based on the function signature
-                args_names = list(self._func_signature.parameters.keys())[1:len(args)+1]  # skip 'self'
-                expected_args_kwargs = dict(zip(args_names, args))
-                # Combine filtered args and kwargs
-                combined_params = {**expected_args_kwargs, **expected_kwargs}
-
-                # Bind the current instance attributes to the function signature
-                if "callback" in self._expected_params and not stream:
-                    raise ValueError(
-                        f"Node {func.__name__}: requires stream=True when callback is defined."
-                    )
-                bound_params = self._func_signature.bind_partial(**combined_params)
-                # get the default parameters (if any)
-                bound_params.apply_defaults()
-                params_dict = bound_params.arguments.copy()
-                # Filter arguments to match the next node's parameters
-                filtered_params = {
-                    k: v for k, v in params_dict.items() if k in self._expected_params
-                }
-                # confirm all required parameters are present
-                missing_params = [
-                    p for p in self._expected_params if p not in filtered_params
-                ]
-                # if anything is missing we raise an error
-                if missing_params:
-                    raise ValueError(
-                        f"Missing required parameters for the {func.__name__} node: {', '.join(missing_params)}"
-                    )
-                return filtered_params
-
-
-            @classmethod
-            def get_signature(cls):
-                """Returns the signature of the decorated function as LLM readable
-                string.
-                """
-                signature_components = []
-                if NodeClass._func_signature:
-                    for param in NodeClass._func_signature.parameters.values():
-                        if param.default is param.empty:
-                            signature_components.append(f"{param.name}: {param.annotation}")
-                        else:
-                            signature_components.append(f"{param.name}: {param.annotation} = {param.default}")
-                else:
-                    return "No signature"
-                return "\n".join(signature_components)
-
-            @classmethod
-            async def invoke(cls, input: Dict[str, Any], callback: Optional[Callback] = None):
-                if callback:
-                    if stream:
-                        input["callback"] = callback
-                    else:
-                        raise ValueError(
-                            f"Error in node {func.__name__}. When callback provided, stream must be True."
-                        )
-                instance = cls()
-                out = await instance.execute(**input)
-                return out
-
-        NodeClass.__name__ = func.__name__
-        NodeClass.name = func.__name__
-        NodeClass.__doc__ = func.__doc__
-        NodeClass.is_start = start
-        NodeClass.is_end = end
-        NodeClass.is_router = self.is_router
-        NodeClass.stream = stream
-        NodeClass.schema = schema
-        return NodeClass
-
-    def __call__(
-        self,
-        func: Optional[Callable] = None,
-        start: bool = False,
-        end: bool = False,
-        stream: bool = False,
-    ):
-        # We must wrap the call to the decorator in a function for it to work
-        # correctly with or without parenthesis
-        def wrap(func: Callable, start=start, end=end, stream=stream) -> Callable:
-            return self._node(func=func, start=start, end=end, stream=stream)
-        if func:
-            # Decorator is called without parenthesis
-            return wrap(func=func, start=start, end=end, stream=stream)
-        # Decorator is called with parenthesis
-        return wrap
-
-
-node = _Node()
-router = _Node(is_router=True)

graphai_lib-0.0.1.dist-info/METADATA

@@ -1,25 +0,0 @@
-Metadata-Version: 2.1
-Name: graphai-lib
-Version: 0.0.1
-Summary: 
-License: MIT
-Author: Aurelio AI
-Author-email: hello@aurelio.ai
-Requires-Python: >=3.10,<3.14
-Classifier: License :: OSI Approved :: MIT License
-Classifier: Programming Language :: Python :: 3
-Classifier: Programming Language :: Python :: 3.10
-Classifier: Programming Language :: Python :: 3.11
-Classifier: Programming Language :: Python :: 3.12
-Classifier: Programming Language :: Python :: 3.13
-Requires-Dist: matplotlib (>=3.10.0,<4.0.0)
-Requires-Dist: networkx (>=3.4.2,<4.0.0)
-Requires-Dist: semantic-router (>=0.1.0.dev4)
-Description-Content-Type: text/markdown
-
-# Philosophy
-
-1. Async-first
-2. Minimize abstractions
-3. One way to do one thing
-4. Graph-based AI

graphai_lib-0.0.1.dist-info/RECORD

@@ -1,9 +0,0 @@
-graphai/__init__.py,sha256=EHigFOWewDXLZXbdfjZH9kdLPhw6NT0ChS77lNAVAA8,109
-graphai/callback.py,sha256=0pFe2ocxHzz5rdv2w2iYuiiSDSTiVwJO81Pky3qktmw,2160
-graphai/graph.py,sha256=Gn_bl0iAtcH9QT54G2WHJDbxv0TPkVwuPOgAfu4bHT4,7687
-graphai/nodes/__init__.py,sha256=4826Ubk5yUfbVH7F8DmoTKQyax624Q2QJHsGxqgQ_ng,73
-graphai/nodes/base.py,sha256=VHceReX-FcRTrq8eWF6vfDyRfc4F8Zdcya5UXSPFbmw,5937
-graphai/utils.py,sha256=zrgpk82rIn7lwh631KhN-OgMAJMdbm0k5GPL1eMf2sQ,4522
-graphai_lib-0.0.1.dist-info/METADATA,sha256=AszW28j6lxyA6iyl_KdfCEYSAKsiyMUSbY_5_9G6zbE,732
-graphai_lib-0.0.1.dist-info/WHEEL,sha256=Nq82e9rUAnEjt98J6MlVmMCZb-t9cYE2Ir1kpBmnWfs,88
-graphai_lib-0.0.1.dist-info/RECORD,,