graphai-lib 0.0.6__py3-none-any.whl → 0.0.8__py3-none-any.whl

graphai/__init__.py

@@ -1,4 +1,5 @@
+from graphai.callback import Callback
 from graphai.graph import Graph
 from graphai.nodes import node, router
 
-__all__ = ["node", "router", "Graph"]
+__all__ = ["node", "router", "Callback", "Graph"]

graphai/callback.py

@@ -1,6 +1,6 @@
 import asyncio
 from pydantic import Field
-from typing import Optional, Any
+from typing import Any
 from collections.abc import AsyncIterator
 
 
@@ -48,7 +48,7 @@ class Callback:
         description="Whether this is the first token in the stream.",
         exclude=True,
     )
-    _current_node_name: Optional[str] = Field(
+    _current_node_name: str | None = Field(
         default=None, description="The name of the current node.", exclude=True
     )
     _active: bool = Field(
@@ -85,11 +85,11 @@ class Callback:
         self._first_token = value
 
     @property
-    def current_node_name(self) -> Optional[str]:
+    def current_node_name(self) -> str | None:
         return self._current_node_name
 
     @current_node_name.setter
-    def current_node_name(self, value: Optional[str]):
+    def current_node_name(self, value: str | None):
         self._current_node_name = value
 
     @property
@@ -100,14 +100,14 @@ class Callback:
     def active(self, value: bool):
         self._active = value
 
-    def __call__(self, token: str, node_name: Optional[str] = None):
+    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)
         # 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):
+    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")
@@ -176,7 +176,7 @@ class Callback:
         # 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):
+    def _check_node_name(self, node_name: str | None = None):
         if node_name:
             # we confirm this is the current node
             if self.current_node_name != node_name:

graphai/graph.py

@@ -1,54 +1,58 @@
-from typing import List, Dict, Any, Optional, Protocol, Type
+from typing import Any, Protocol, Type
 from graphai.callback import Callback
 from graphai.utils import logger
 
 
 class NodeProtocol(Protocol):
     """Protocol defining the interface of a decorated node."""
+
     name: str
     is_start: bool
     is_end: bool
     is_router: bool
     stream: bool
-    
+
     async def invoke(
-        self, 
-        input: Dict[str, Any], 
-        callback: Optional[Callback] = None, 
-        state: Optional[Dict[str, Any]] = None
-    ) -> Dict[str, Any]: ...
+        self,
+        input: dict[str, Any],
+        callback: Callback | None = None,
+        state: dict[str, Any] | None = None,
+    ) -> dict[str, Any]: ...
 
 
 class Graph:
     def __init__(
-        self, max_steps: int = 10, initial_state: Optional[Dict[str, Any]] = None
+        self, max_steps: int = 10, initial_state: dict[str, Any] | None = None
     ):
-        self.nodes: Dict[str, NodeProtocol] = {}
-        self.edges: List[Any] = []
-        self.start_node: Optional[NodeProtocol] = None
-        self.end_nodes: List[NodeProtocol] = []
+        self.nodes: dict[str, NodeProtocol] = {}
+        self.edges: list[Any] = []
+        self.start_node: NodeProtocol | None = None
+        self.end_nodes: list[NodeProtocol] = []
         self.Callback: Type[Callback] = Callback
         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]:
+    def get_state(self) -> dict[str, Any]:
         """Get the current graph state."""
         return self.state
 
-    def set_state(self, state: Dict[str, Any]):
+    def set_state(self, state: dict[str, Any]) -> "Graph":
         """Set the graph state."""
         self.state = state
+        return self
 
-    def update_state(self, values: Dict[str, Any]):
+    def update_state(self, values: dict[str, Any]) -> "Graph":
         """Update the graph state with new values."""
         self.state.update(values)
+        return self
 
-    def reset_state(self):
+    def reset_state(self) -> "Graph":
         """Reset the graph state to an empty dict."""
         self.state = {}
+        return self
 
-    def add_node(self, node: NodeProtocol):
+    def add_node(self, node: NodeProtocol) -> "Graph":
         if node.name in self.nodes:
             raise Exception(f"Node with name '{node.name}' already exists.")
         self.nodes[node.name] = node
@@ -62,8 +66,9 @@ class Graph:
             self.start_node = node
         if node.is_end:
             self.end_nodes.append(node)
+        return self
 
-    def add_edge(self, source: NodeProtocol | str, destination: NodeProtocol | str):
+    def add_edge(self, source: NodeProtocol | str, destination: NodeProtocol | str) -> "Graph":
         """Adds an edge between two nodes that already exist in the graph.
 
         Args:
@@ -72,58 +77,76 @@ class Graph:
         """
         source_node, destination_node = None, None
         # get source node from graph
+        source_name: str
         if isinstance(source, str):
             source_node = self.nodes.get(source)
+            source_name = 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)
+                source_name = source.name
+            else:
+                source_name = str(source)
         if source_node is None:
             raise ValueError(
-                f"Node with name '{source.name if hasattr(source, 'name') else source}' not found."
+                f"Node with name '{source_name}' not found."
             )
         # get destination node from graph
+        destination_name: str
         if isinstance(destination, str):
             destination_node = self.nodes.get(destination)
+            destination_name = 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)
+                destination_name = destination.name
+            else:
+                destination_name = str(destination)
         if destination_node is None:
             raise ValueError(
-                f"Node with name '{destination.name if hasattr(destination, 'name') else destination}' not found."
+                f"Node with name '{destination_name}' not found."
             )
         edge = Edge(source_node, destination_node)
         self.edges.append(edge)
+        return self
 
     def add_router(
-        self, sources: list[NodeProtocol], router: NodeProtocol, destinations: List[NodeProtocol]
-    ):
+        self,
+        sources: list[NodeProtocol],
+        router: NodeProtocol,
+        destinations: list[NodeProtocol],
+    ) -> "Graph":
         if not router.is_router:
             raise TypeError("A router object must be passed to the router parameter.")
         [self.add_edge(source, router) for source in sources]
         for destination in destinations:
             self.add_edge(router, destination)
+        return self
 
-    def set_start_node(self, node: NodeProtocol):
+    def set_start_node(self, node: NodeProtocol) -> "Graph":
         self.start_node = node
+        return self
 
-    def set_end_node(self, node: NodeProtocol):
+    def set_end_node(self, node: NodeProtocol) -> "Graph":
         self.end_node = node
+        return self
 
-    def compile(self):
+    def compile(self) -> "Graph":
         if not self.start_node:
             raise Exception("Start node not defined.")
         if not self.end_nodes:
             raise Exception("No end nodes defined.")
         if not self._is_valid():
             raise Exception("Graph is not valid.")
+        return self
 
     def _is_valid(self):
         # Implement validation logic, e.g., checking for cycles, disconnected components, etc.
         return True
 
-    def _validate_output(self, output: Dict[str, Any], node_name: str):
+    def _validate_output(self, output: dict[str, Any], node_name: str):
         if not isinstance(output, dict):
             raise ValueError(
                 f"Expected dictionary output from node {node_name}. "
@@ -134,11 +157,11 @@ class Graph:
         # TODO JB: may need to add init callback here to init the queue on every new execution
         if callback is None:
             callback = self.get_callback()
-        
+
         # Type assertion to tell the type checker that start_node is not None after compile()
         assert self.start_node is not None, "Graph must be compiled before execution"
         current_node = self.start_node
-        
+
         state = input
         # Don't reset the graph state if it was initialized with initial_state
         steps = 0
@@ -191,7 +214,7 @@ class Graph:
         callback = self.Callback()
         return callback
 
-    def set_callback(self, callback_class: Type[Callback]):
+    def set_callback(self, callback_class: type[Callback]) -> "Graph":
         """Set the callback class that is returned by the `get_callback` method and used
         as the default callback when no callback is passed to the `execute` method.
 
@@ -199,6 +222,7 @@ class Graph:
         :type callback_class: Type[Callback]
         """
         self.Callback = callback_class
+        return self
 
     def _get_node_by_name(self, node_name: str) -> NodeProtocol:
         """Get a node by its name.
@@ -240,7 +264,7 @@ class Graph:
                 "Matplotlib is required for visualization. Please install it with 'pip install matplotlib'."
             )
 
-        G = nx.DiGraph()
+        G: Any = nx.DiGraph()
 
         for node in self.nodes.values():
             G.add_node(node.name)
@@ -264,11 +288,11 @@ class Graph:
                     y_coord[node] = y_max - i - 1
 
             # Set up the layout
-            pos = {}
+            pos: dict[Any, tuple[float, float]] = {}
             for i, generation in enumerate(generations):
                 x = 0
                 for node in generation:
-                    pos[node] = (x, y_coord[node])
+                    pos[node] = (float(x), float(y_coord[node]))
                     x += 1
 
             # Center each level horizontally

graphai/nodes/base.py

@@ -1,5 +1,5 @@
 import inspect
-from typing import Any, Callable, Dict, Optional
+from typing import Any, Callable
 from pydantic import Field
 
 from graphai.callback import Callback
@@ -8,7 +8,7 @@ from graphai.utils import FunctionSchema
 
 class NodeMeta(type):
     @staticmethod
-    def positional_to_kwargs(cls_type, args) -> Dict[str, Any]:
+    def positional_to_kwargs(cls_type, args) -> dict[str, Any]:
         init_signature = inspect.signature(cls_type.__init__)
         init_params = {
             name: arg
@@ -73,7 +73,7 @@ class _Node:
                 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]:
+            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
@@ -133,9 +133,9 @@ class _Node:
             @classmethod
             async def invoke(
                 cls,
-                input: Dict[str, Any],
-                callback: Optional[Callback] = None,
-                state: Optional[Dict[str, Any]] = None,
+                input: dict[str, Any],
+                callback: Callback | None = None,
+                state: dict[str, Any] | None = None,
             ):
                 if callback:
                     if stream:
@@ -154,8 +154,6 @@ class _Node:
 
         NodeClass.__name__ = func.__name__
         node_class_name = name or func.__name__
-        if node_class_name is None:
-            raise ValueError("Unexpected error: node name not set.")
         NodeClass.name = node_class_name
         NodeClass.__doc__ = func.__doc__
         NodeClass.is_start = start
@@ -167,7 +165,7 @@ class _Node:
 
     def __call__(
         self,
-        func: Optional[Callable] = None,
+        func: Callable | None = None,
         start: bool = False,
         end: bool = False,
         stream: bool = False,

graphai/utils.py

@@ -1,33 +1,43 @@
 import inspect
-from typing import Any, Callable, List, Optional
+import os
+from typing import Any, Callable
 from pydantic import BaseModel, Field
 import logging
+import sys
 
-import colorlog
 
+class ColoredFormatter(logging.Formatter):
+    """Custom colored formatter for the logger using ANSI escape codes."""
 
-class CustomFormatter(colorlog.ColoredFormatter):
-    """Custom formatter for the logger."""
+    # ANSI escape codes for colors
+    COLORS = {
+        "DEBUG": "\033[36m",  # Cyan
+        "INFO": "\033[32m",  # Green
+        "WARNING": "\033[33m",  # Yellow
+        "ERROR": "\033[31m",  # Red
+        "CRITICAL": "\033[1;31m",  # Bold Red
+    }
+    RESET = "\033[0m"
 
     def __init__(self):
         super().__init__(
-            "%(log_color)s%(asctime)s %(levelname)s %(name)s %(message)s",
+            "%(asctime)s %(levelname)s %(name)s %(message)s",
             datefmt="%Y-%m-%d %H:%M:%S",
-            log_colors={
-                "DEBUG": "cyan",
-                "INFO": "green",
-                "WARNING": "yellow",
-                "ERROR": "red",
-                "CRITICAL": "bold_red",
-            },
-            reset=True,
-            style="%",
         )
 
+    def format(self, record):
+        # Check if the output supports color (TTY)
+        if hasattr(sys.stderr, "isatty") and sys.stderr.isatty():
+            levelname = record.levelname
+            if levelname in self.COLORS:
+                record.levelname = f"{self.COLORS[levelname]}{levelname}{self.RESET}"
+                record.msg = f"{self.COLORS[levelname]}{record.msg}{self.RESET}"
+        return super().format(record)
+
 
 def add_coloured_handler(logger):
     """Add a coloured handler to the logger."""
-    formatter = CustomFormatter()
+    formatter = ColoredFormatter()
     console_handler = logging.StreamHandler()
     console_handler.setFormatter(formatter)
     logger.addHandler(console_handler)
@@ -40,7 +50,17 @@ def setup_custom_logger(name):
 
     if not logger.hasHandlers():
         add_coloured_handler(logger)
-        logger.setLevel(logging.INFO)
+        
+        # Set log level from environment variable, default to INFO
+        log_level = os.getenv("GRAPHAI_LOG_LEVEL", "INFO").upper()
+        level_map = {
+            "DEBUG": logging.DEBUG,
+            "INFO": logging.INFO,
+            "WARNING": logging.WARNING,
+            "ERROR": logging.ERROR,
+            "CRITICAL": logging.CRITICAL,
+        }
+        logger.setLevel(level_map.get(log_level, logging.INFO))
         logger.propagate = False
 
     return logger
@@ -68,7 +88,7 @@ class Parameter(BaseModel):
     :param name: The name of the parameter.
     :type name: str
     :param description: The description of the parameter.
-    :type description: Optional[str]
+    :type description: str | None
     :param type: The type of the parameter.
     :type type: str
     :param default: The default value of the parameter.
@@ -78,7 +98,7 @@ class Parameter(BaseModel):
     """
 
     name: str = Field(description="The name of the parameter")
-    description: Optional[str] = Field(
+    description: str | None = Field(
         default=None, description="The description of the parameter"
     )
     type: str = Field(description="The type of the parameter")
@@ -118,34 +138,33 @@ class FunctionSchema(BaseModel):
         :param function: The function to consume.
         :type function: Callable
         """
-        if callable(function):
-            name = function.__name__
-            description = str(inspect.getdoc(function))
-            if description is None or description == "":
-                logger.warning(f"Function {name} has no docstring")
-            signature = str(inspect.signature(function))
-            output = str(inspect.signature(function).return_annotation)
-            parameters = []
-            for param in inspect.signature(function).parameters.values():
-                parameters.append(
-                    Parameter(
-                        name=param.name,
-                        type=param.annotation.__name__,
-                        default=param.default,
-                        required=param.default is inspect.Parameter.empty,
-                    )
+        if not callable(function):
+            raise TypeError("Function must be a Callable")
+        
+        name = function.__name__
+        doc = inspect.getdoc(function)
+        description = str(doc) if doc else ""
+        if not description:
+            logger.warning(f"Function {name} has no docstring")
+        signature = str(inspect.signature(function))
+        output = str(inspect.signature(function).return_annotation)
+        parameters = []
+        for param in inspect.signature(function).parameters.values():
+            parameters.append(
+                Parameter(
+                    name=param.name,
+                    type=param.annotation.__name__,
+                    default=param.default,
+                    required=param.default is inspect.Parameter.empty,
                 )
-            return cls.model_construct(
-                name=name,
-                description=description,
-                signature=signature,
-                output=output,
-                parameters=parameters,
             )
-        elif isinstance(function, BaseModel):
-            raise NotImplementedError("Pydantic BaseModel not implemented yet.")
-        else:
-            raise TypeError("Function must be a Callable or BaseModel")
+        return cls.model_construct(
+            name=name,
+            description=description,
+            signature=signature,
+            output=output,
+            parameters=parameters,
+        )
 
     @classmethod
     def from_pydantic(cls, model: BaseModel) -> "FunctionSchema":
@@ -179,7 +198,9 @@ class FunctionSchema(BaseModel):
                 "parameters": {
                     "type": "object",
                     "properties": {
-                        k: v for param in self.parameters for k, v in param.to_dict().items()
+                        k: v
+                        for param in self.parameters
+                        for k, v in param.to_dict().items()
                     },
                     "required": [
                         param.name for param in self.parameters if param.required
@@ -196,7 +217,7 @@ class FunctionSchema(BaseModel):
 DEFAULT = set(["default", "openai", "ollama", "litellm"])
 
 
-def get_schemas(callables: List[Callable], format: str = "default") -> list[dict]:
+def get_schemas(callables: list[Callable], format: str = "default") -> list[dict]:
     if format in DEFAULT:
         return [
             FunctionSchema.from_callable(callable).to_dict() for callable in callables

{graphai_lib-0.0.6.dist-info → graphai_lib-0.0.8.dist-info}/METADATA

@@ -1,13 +1,10 @@
 Metadata-Version: 2.4
 Name: graphai-lib
-Version: 0.0.6
+Version: 0.0.8
 Summary: Not an AI framework
 Requires-Python: <3.14,>=3.10
 Description-Content-Type: text/markdown
-Requires-Dist: networkx>=3.4.2
-Requires-Dist: matplotlib>=3.10.0
 Requires-Dist: pydantic>=2.11.1
-Requires-Dist: colorlog>=6.9.0
 Provides-Extra: dev
 Requires-Dist: ipykernel>=6.25.0; extra == "dev"
 Requires-Dist: ruff>=0.1.5; extra == "dev"

graphai_lib-0.0.8.dist-info/RECORD

@@ -0,0 +1,10 @@
+graphai/__init__.py,sha256=UbqXq7iGIYe1GyTPcpgLSXbgWovggFsAbTMtr4JQm3M,160
+graphai/callback.py,sha256=Wl0JCmE8NcFifKmP9-a5bFa0WKVHTdrSClHVRmIEPpc,7323
+graphai/graph.py,sha256=4pY07zSs8VakN9Pma3xgiFOgKqkrHMi5C04umo4UN8c,12342
+graphai/utils.py,sha256=_q3pE7rgeXzGLYE41LWn87rlW3t74hQQwqHI1PvVjrY,7554
+graphai/nodes/__init__.py,sha256=IaMUryAqTZlcEqh-ZS6A4NIYG18JZwzo145dzxsYjAk,74
+graphai/nodes/base.py,sha256=SoKfOdRu5EIJ_z8xIz5zbNXcxPI2l9MKTQDeaQI-2no,7494
+graphai_lib-0.0.8.dist-info/METADATA,sha256=5OmA9X6HwtZxOr2GexbLFsBv7TBL8bMzE5OpLZ5oReE,910
+graphai_lib-0.0.8.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+graphai_lib-0.0.8.dist-info/top_level.txt,sha256=TXlqmhLViX-3xGH2g5w6cavRd-QMf229Hl88jdMOGt8,8
+graphai_lib-0.0.8.dist-info/RECORD,,

graphai_lib-0.0.6.dist-info/RECORD

@@ -1,10 +0,0 @@
-graphai/__init__.py,sha256=kZJ21W6gwN-eRzvrWQf8xDTPCIWTIuyWq1IGDS9tn7Y,110
-graphai/callback.py,sha256=M2gEpj7uVvANg2dVgxKFMUSPIM362YSilShMIFfrr8s,7351
-graphai/graph.py,sha256=p5iAXQqx0266bAS2zGztv34UXW3-TiA7d2XIsKxAW54,11744
-graphai/utils.py,sha256=LlL-Wx643nIeRFAl2xcv0crNQcA_0563epRo8ZsyL40,6898
-graphai/nodes/__init__.py,sha256=IaMUryAqTZlcEqh-ZS6A4NIYG18JZwzo145dzxsYjAk,74
-graphai/nodes/base.py,sha256=-ZOfJhxews5CGutYB5lfoIVvJ6dqdYJWeeKzDMz9odg,7624
-graphai_lib-0.0.6.dist-info/METADATA,sha256=DZ-RIueKiVmtUJN2vXm02ZD-_aB5uiduWwN6HxriP4w,1006
-graphai_lib-0.0.6.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-graphai_lib-0.0.6.dist-info/top_level.txt,sha256=TXlqmhLViX-3xGH2g5w6cavRd-QMf229Hl88jdMOGt8,8
-graphai_lib-0.0.6.dist-info/RECORD,,