graphai-lib 0.0.1__tar.gz

graphai_lib-0.0.1/PKG-INFO

@@ -0,0 +1,25 @@
+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/README.md

@@ -0,0 +1,6 @@
+# Philosophy
+
+1. Async-first
+2. Minimize abstractions
+3. One way to do one thing
+4. Graph-based AI

graphai_lib-0.0.1/graphai/__init__.py

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

graphai_lib-0.0.1/graphai/callback.py

@@ -0,0 +1,63 @@
+import asyncio
+from typing import Optional
+from collections.abc import AsyncIterator
+from semantic_router.utils.logger import logger
+
+
+log_stream = True
+
+class Callback:
+    first_token = True
+    current_node_name: Optional[str] = None
+    active: bool = True
+    queue: asyncio.Queue
+
+    def __init__(self):
+        self.queue = asyncio.Queue()
+
+    def __call__(self, token: str, node_name: Optional[str] = None):
+        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):
+        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 aiter(self) -> AsyncIterator[str]:
+        """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:
+            token = await self.queue.get()
+            yield token
+            self.queue.task_done()
+            if token == "<graphai:END>":
+                break
+
+    async def start_node(self, node_name: str, active: bool = True):
+        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}>")
+    
+    async def end_node(self, node_name: str):
+        self.current_node_name = None
+        if self.active:
+            self.queue.put_nowait(f"<graphai:end:{node_name}>")
+
+    async def close(self):
+        self.queue.put_nowait("<graphai:END>")
+
+    def _check_node_name(self, node_name: Optional[str] = None):
+        if node_name:
+            # we confirm this is the current node
+            if self.current_node_name != node_name:
+                raise ValueError(
+                    f"Node name mismatch: {self.current_node_name} != {node_name}"
+                )

graphai_lib-0.0.1/graphai/graph.py

@@ -0,0 +1,198 @@
+from typing import List, Dict, Any
+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 = []
+        self.Callback = Callback
+        self.callback = None
+        self.max_steps = max_steps
+
+    def add_node(self, node):
+        self.nodes.append(node)
+        if node.is_start:
+            if self.start_node is not None:
+                raise Exception(
+                    "Multiple start nodes are not allowed. Start node "
+                    f"'{self.start_node.name}' already exists, so new start "
+                    f"node '{node.name}' can not be added to the graph."
+                )
+            self.start_node = node
+        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)
+        self.edges.append(edge)
+
+    def add_router(self, sources: list[_Node], router: _Node, destinations: List[_Node]):
+        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)
+
+    def set_start_node(self, node: _Node):
+        self.start_node = node
+
+    def set_end_node(self, node: _Node):
+        self.end_node = node
+
+    def compile(self):
+        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.")
+
+    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):
+        if not isinstance(output, dict):
+            raise ValueError(
+                f"Expected dictionary output from node {node_name}. "
+                f"Instead, got {type(output)} from '{output}'."
+            )
+
+    async def execute(self, input):
+        # TODO JB: may need to add init callback here to init the queue on every new execution
+        if self.callback is None:
+            self.callback = self.get_callback()
+        current_node = self.start_node
+        state = input
+        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)
+                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)
+                self._validate_output(output=output, node_name=current_node.name)
+            # add output to state
+            state = {**state, **output}
+            if current_node.is_end:
+                # finish loop if this was an end node
+                break
+            if current_node.is_router:
+                # if we have a router node we let the router decide the next node
+                next_node_name = str(output["choice"])
+                del output["choice"]
+                current_node = self._get_node_by_name(node_name=next_node_name)
+            else:
+                # otherwise, we have linear path
+                current_node = self._get_next_node(current_node=current_node)
+            steps += 1
+            if steps >= self.max_steps:
+                raise Exception(
+                    f"Max steps reached: {self.max_steps}. You can modify this "
+                    "by setting `max_steps` when initializing the Graph object."
+                )
+        # TODO JB: may need to add end callback here to close the queue for every execution
+        if self.callback and "callback" in state:
+            await self.callback.close()
+            del state["callback"]
+        return state
+
+    def get_callback(self):
+        self.callback = self.Callback()
+        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.")
+
+    def _get_next_node(self, current_node):
+        for edge in self.edges:
+            if edge.source == current_node:
+                return edge.destination
+        raise Exception(
+            f"No outgoing edge found for current node '{current_node.name}'."
+        )
+
+    def visualize(self):
+        try:
+            import networkx as nx
+        except ImportError:
+            raise ImportError("NetworkX is required for visualization. Please install it with 'pip install networkx'.")
+
+        try:
+            import matplotlib.pyplot as plt
+        except ImportError:
+            raise ImportError("Matplotlib is required for visualization. Please install it with 'pip install matplotlib'.")
+
+        G = nx.DiGraph()
+
+        for node in self.nodes:
+            G.add_node(node.name)
+
+        for edge in self.edges:
+            G.add_edge(edge.source.name, edge.destination.name)
+
+        if nx.is_directed_acyclic_graph(G):
+            logger.info("The graph is acyclic. Visualization will use a topological layout.")
+            # Use topological layout if acyclic
+            # Compute the topological generations
+            generations = list(nx.topological_generations(G))
+            y_max = len(generations)
+
+            # Create a dictionary to store the y-coordinate for each node
+            y_coord = {}
+            for i, generation in enumerate(generations):
+                for node in generation:
+                    y_coord[node] = y_max - i - 1
+
+            # Set up the layout
+            pos = {}
+            for i, generation in enumerate(generations):
+                x = 0
+                for node in generation:
+                    pos[node] = (x, y_coord[node])
+                    x += 1
+
+            # Center each level horizontally
+            for i, generation in enumerate(generations):
+                x_center = sum(pos[node][0] for node in generation) / len(generation)
+                for node in generation:
+                    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()}
+
+        else:
+            print("Warning: The graph contains cycles. Visualization will use a spring layout.")
+            pos = nx.spring_layout(G, k=1, iterations=50)
+
+        plt.figure(figsize=(8, 6))
+        nx.draw(G, pos, with_labels=True, node_color='lightblue', 
+                node_size=3000, font_size=8, font_weight='bold', 
+                arrows=True, edge_color='gray', arrowsize=20)
+
+        plt.axis('off')
+        plt.show()
+
+
+
+class Edge:
+    def __init__(self, source, destination):
+        self.source = source
+        self.destination = destination

graphai_lib-0.0.1/graphai/nodes/__init__.py

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

graphai_lib-0.0.1/graphai/nodes/base.py

@@ -0,0 +1,148 @@
+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/graphai/utils.py

@@ -0,0 +1,125 @@
+import inspect
+from typing import Any, Callable, Dict, List, Union, Optional
+from pydantic import BaseModel, Field
+
+
+class Parameter(BaseModel):
+    class Config:
+        arbitrary_types_allowed = True
+
+    name: str = Field(description="The name of the parameter")
+    description: Optional[str] = Field(
+        default=None, description="The description of the parameter"
+    )
+    type: str = Field(description="The type of the parameter")
+    default: Any = Field(description="The default value of the parameter")
+    required: bool = Field(description="Whether the parameter is required")
+
+    def to_openai(self):
+        return {
+            self.name: {
+                "description": self.description,
+                "type": self.type,
+            }
+        }
+
+class FunctionSchema:
+    """Class that consumes a function and can return a schema required by
+    different LLMs for function calling.
+    """
+
+    name: str = Field(description="The name of the function")
+    description: str = Field(description="The description of the function")
+    signature: str = Field(description="The signature of the function")
+    output: str = Field(description="The output of the function")
+    parameters: List[Parameter] = Field(description="The parameters of the function")
+
+    def __init__(self, function: Union[Callable, BaseModel]):
+        self.function = function
+        if callable(function):
+            self._process_function(function)
+        elif isinstance(function, BaseModel):
+            raise NotImplementedError("Pydantic BaseModel not implemented yet.")
+        else:
+            raise TypeError("Function must be a Callable or BaseModel")
+
+    def _process_function(self, function: Callable):
+        self.name = function.__name__
+        self.description = str(inspect.getdoc(function))
+        self.signature = str(inspect.signature(function))
+        self.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,
+                )
+            )
+        self.parameters = parameters
+
+    def to_openai(self):
+        schema_dict = {
+            "type": "function",
+            "function": {
+                "name": self.name,
+                "description": self.description,
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        param.name: {
+                            "description": (
+                                param.description
+                                if isinstance(param.description, str)
+                                else "None provided"
+                            ),
+                            "type": self._openai_type_mapping(param.type),
+                        }
+                        for param in self.parameters
+                    },
+                    "required": [
+                        param.name for param in self.parameters if param.required
+                    ],
+                },
+            },
+        }
+        return schema_dict
+
+    def _openai_type_mapping(self, param_type: str) -> str:
+        if param_type == "int":
+            return "number"
+        elif param_type == "float":
+            return "number"
+        elif param_type == "str":
+            return "string"
+        elif param_type == "bool":
+            return "boolean"
+        else:
+            return "object"
+
+
+def get_schema_pydantic(model: BaseModel) -> Dict[str, Any]:
+    signature_parts = []
+    for field_name, field_model in model.__annotations__.items():
+        field_info = model.__fields__[field_name]
+        default_value = field_info.default
+
+        if default_value:
+            default_repr = repr(default_value)
+            signature_part = (
+                f"{field_name}: {field_model.__name__} = {default_repr}"
+            )
+        else:
+            signature_part = f"{field_name}: {field_model.__name__}"
+
+        signature_parts.append(signature_part)
+    signature = f"({', '.join(signature_parts)}) -> str"
+    schema = FunctionSchema(
+        name=model.__class__.__name__,
+        description=model.__doc__,
+        signature=signature,
+        output="",  # TODO: Implement output
+    )
+    return schema

graphai_lib-0.0.1/pyproject.toml

@@ -0,0 +1,29 @@
+[tool.poetry]
+name = "graphai-lib"
+version = "0.0.1"
+description = ""
+authors = ["Aurelio AI <hello@aurelio.ai>"]
+readme = "README.md"
+packages = [{include = "graphai"}]
+license = "MIT"
+
+[tool.poetry.dependencies]
+python = ">=3.10,<3.14"
+semantic-router = ">=0.1.0.dev4"
+networkx = "^3.4.2"
+matplotlib = "^3.10.0"
+
+[tool.poetry.group.dev.dependencies]
+ipykernel = "^6.25.0"
+ruff = "^0.1.5"
+pytest = "^8.2"
+pytest-mock = "^3.12.0"
+pytest-cov = "^4.1.0"
+pytest-xdist = "^3.5.0"
+pytest-asyncio = "^0.24.0"
+mypy = "^1.7.1"
+black = {extras = ["jupyter"], version = ">=23.12.1,<24.5.0"}
+
+[build-system]
+requires = ["poetry-core"]
+build-backend = "poetry.core.masonry.api"