graphai-lib 0.0.8__py3-none-any.whl → 0.0.9rc1__py3-none-any.whl

graphai/graph.py

@@ -1,8 +1,23 @@
+from __future__ import annotations
 from typing import Any, Protocol, Type
+from graphlib import TopologicalSorter, CycleError
 from graphai.callback import Callback
 from graphai.utils import logger
 
 
+# to fix mypy error
+class _HasName(Protocol):
+    name: str
+
+
+class GraphError(Exception):
+    pass
+
+
+class GraphCompileError(GraphError):
+    pass
+
+
 class NodeProtocol(Protocol):
     """Protocol defining the interface of a decorated node."""
 
@@ -20,6 +35,26 @@ class NodeProtocol(Protocol):
     ) -> dict[str, Any]: ...
 
 
+def _name_of(x: Any) -> str | None:
+    """Return the node name if x is a str or has .name, else None."""
+    if x is None:
+        return None
+    if isinstance(x, str):
+        return x
+    name = getattr(x, "name", None)
+    return name if isinstance(name, str) else None
+
+
+def _require_name(x: Any, kind: str) -> str:
+    """Like _name_of, but raises a helpful compile error when missing."""
+    s = _name_of(x)
+    if s is None:
+        raise GraphCompileError(
+            f"Edge {kind} must be a node name (str) or object with .name"
+        )
+    return s
+
+
 class Graph:
     def __init__(
         self, max_steps: int = 10, initial_state: dict[str, Any] | None = None
@@ -37,22 +72,22 @@ class Graph:
         """Get the current graph state."""
         return self.state
 
-    def set_state(self, state: dict[str, Any]) -> "Graph":
+    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]) -> "Graph":
+    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) -> "Graph":
+    def reset_state(self) -> Graph:
         """Reset the graph state to an empty dict."""
         self.state = {}
         return self
 
-    def add_node(self, node: NodeProtocol) -> "Graph":
+    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
@@ -68,7 +103,9 @@ class Graph:
             self.end_nodes.append(node)
         return self
 
-    def add_edge(self, source: NodeProtocol | str, destination: NodeProtocol | str) -> "Graph":
+    def add_edge(
+        self, source: NodeProtocol | str, destination: NodeProtocol | str
+    ) -> Graph:
         """Adds an edge between two nodes that already exist in the graph.
 
         Args:
@@ -89,9 +126,7 @@ class Graph:
             else:
                 source_name = str(source)
         if source_node is None:
-            raise ValueError(
-                f"Node with name '{source_name}' not found."
-            )
+            raise ValueError(f"Node with name '{source_name}' not found.")
         # get destination node from graph
         destination_name: str
         if isinstance(destination, str):
@@ -105,9 +140,7 @@ class Graph:
             else:
                 destination_name = str(destination)
         if destination_node is None:
-            raise ValueError(
-                f"Node with name '{destination_name}' not found."
-            )
+            raise ValueError(f"Node with name '{destination_name}' not found.")
         edge = Edge(source_node, destination_node)
         self.edges.append(edge)
         return self
@@ -117,7 +150,7 @@ class Graph:
         sources: list[NodeProtocol],
         router: NodeProtocol,
         destinations: list[NodeProtocol],
-    ) -> "Graph":
+    ) -> 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]
@@ -125,26 +158,162 @@ class Graph:
             self.add_edge(router, destination)
         return self
 
-    def set_start_node(self, node: NodeProtocol) -> "Graph":
+    def set_start_node(self, node: NodeProtocol) -> Graph:
         self.start_node = node
         return self
 
-    def set_end_node(self, node: NodeProtocol) -> "Graph":
+    def set_end_node(self, node: NodeProtocol) -> Graph:
         self.end_node = node
         return 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
+        """
+        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
+        Returns self on success; raises GraphCompileError otherwise.
+        """
+        # nodes map
+        nodes = getattr(self, "nodes", None)
+        if not isinstance(nodes, dict) or not nodes:
+            raise GraphCompileError("No nodes have been added to the graph")
+
+        start_name: str | None = None
+        # Bind and narrow the attribute for mypy
+        start_node: _HasName | None = getattr(self, "start_node", None)
+        if start_node is not None:
+            start_name = start_node.name
+        else:
+            starts = [
+                name
+                for name, n in nodes.items()
+                if getattr(n, "is_start", False) or getattr(n, "start", False)
+            ]
+            if len(starts) > 1:
+                raise GraphCompileError(f"Multiple start nodes defined: {starts}")
+            if len(starts) == 1:
+                start_name = starts[0]
+
+        if not start_name:
+            raise GraphCompileError("No start node defined")
+
+        # at least one end node
+        if not any(
+            getattr(n, "is_end", False) or getattr(n, "end", False)
+            for n in nodes.values()
+        ):
+            raise GraphCompileError("No end node defined")
+
+        # normalize edges into adjacency {src: set(dst)}
+        raw_edges = getattr(self, "edges", None)
+        adj: dict[str, set[str]] = {name: set() for name in nodes.keys()}
+
+        def _add_edge(src: str, dst: str) -> None:
+            if src not in nodes:
+                raise GraphCompileError(f"Edge references unknown source node: {src}")
+            if dst not in nodes:
+                raise GraphCompileError(
+                    f"Edge from {src} references unknown node(s): ['{dst}']"
+                )
+            adj[src].add(dst)
+
+        if raw_edges is None:
+            pass
+        elif isinstance(raw_edges, dict):
+            for raw_src, dsts in raw_edges.items():
+                src = _require_name(raw_src, "source")
+                dst_iter = (
+                    [dsts]
+                    if isinstance(dsts, (str,)) or getattr(dsts, "name", None)
+                    else list(dsts)
+                )
+                for d in dst_iter:
+                    dst = _require_name(d, "destination")
+                    _add_edge(src, dst)
+        else:
+            # generic iterable of “edge records”
+            try:
+                iterator = iter(raw_edges)
+            except TypeError:
+                raise GraphCompileError("Internal edge map has unsupported type")
+
+            for item in iterator:
+                # (src, dst) OR (src, Iterable[dst])
+                if isinstance(item, (tuple, list)) and len(item) == 2:
+                    raw_src, rhs = item
+                    src = _require_name(raw_src, "source")
+
+                    if isinstance(rhs, str) or getattr(rhs, "name", None):
+                        dst = _require_name(rhs, "destination")
+                        _add_edge(src, rhs)
+                    else:
+                        # assume iterable of dsts (strings or node-like)
+                        try:
+                            for d in rhs:
+                                dst = _require_name(d, "destination")
+                                _add_edge(src, d)
+                        except TypeError:
+                            raise GraphCompileError(
+                                "Edge tuple second item must be a destination or an iterable of destinations"
+                            )
+                    continue
+
+                # Mapping-style: {"source": "...", "destination": "..."} or {"src": "...", "dst": "..."}
+                if isinstance(item, dict):
+                    src = _require_name(item.get("source", item.get("src")), "source")
+                    dst = _require_name(
+                        item.get("destination", item.get("dst")), "destination"
+                    )
+                    _add_edge(src, dst)
+                    continue
+
+                # Object with attributes .source/.destination (or .src/.dst)
+                if hasattr(item, "source") or hasattr(item, "src"):
+                    src = _require_name(
+                        getattr(item, "source", getattr(item, "src", None)), "source"
+                    )
+                    dst = _require_name(
+                        getattr(item, "destination", getattr(item, "dst", None)),
+                        "destination",
+                    )
+                    _add_edge(src, dst)
+                    continue
+
+                # If none matched, this is an unsupported edge record
+                raise GraphCompileError(
+                    "Edges must be dict[str, Iterable[str]] or an iterable of (src, dst), "
+                    "(src, Iterable[dst]), mapping{'source'/'destination'}, or objects with .source/.destination"
+                )
+
+        # cycle detection
+        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") from e
+
+        # reachability from start
+        seen: set[str] = set()
+        stack = [start_name]
+        while stack:
+            cur = stack.pop()
+            if cur in seen:
+                continue
+            seen.add(cur)
+            stack.extend(adj.get(cur, ()))
+
+        unreachable = sorted(set(nodes.keys()) - seen)
+        if unreachable:
+            raise GraphCompileError(f"Unreachable nodes: {unreachable}")
 
-    def _is_valid(self):
-        # Implement validation logic, e.g., checking for cycles, disconnected components, etc.
-        return True
+        return self
 
     def _validate_output(self, output: dict[str, Any], node_name: str):
         if not isinstance(output, dict):
@@ -249,20 +418,24 @@ class Graph:
             f"No outgoing edge found for current node '{current_node.name}'."
         )
 
-    def visualize(self):
+    def visualize(self, *, save_path: str | None = None):
+        """Render the current graph. If matplotlib is not installed,
+        raise a helpful error telling users to install the viz extra.
+        Optionally save to a file via `save_path`.
+        """
         try:
-            import networkx as nx
-        except ImportError:
+            import matplotlib.pyplot as plt
+        except ImportError as e:
             raise ImportError(
-                "NetworkX is required for visualization. Please install it with 'pip install networkx'."
-            )
+                "Graph visualization requires matplotlib. Install it with: `pip install matplotlib`"
+            ) from e
 
         try:
-            import matplotlib.pyplot as plt
-        except ImportError:
+            import networkx as nx
+        except ImportError as e:
             raise ImportError(
-                "Matplotlib is required for visualization. Please install it with 'pip install matplotlib'."
-            )
+                "NetworkX is required for visualization. Please install it with `pip install networkx`."
+            ) from e
 
         G: Any = nx.DiGraph()
 
@@ -328,8 +501,12 @@ class Graph:
             arrowsize=20,
         )
 
-        plt.axis("off")
-        plt.show()
+        if save_path:
+            plt.savefig(save_path, bbox_inches="tight")
+        else:
+            plt.axis("off")
+            plt.show()
+        plt.close()
 
 
 class Edge:

graphai/utils.py

@@ -1,9 +1,18 @@
+from enum import Enum
 import inspect
 import os
-from typing import Any, Callable
+import sys
+from typing import Any, Callable, Union, get_args, get_origin
 from pydantic import BaseModel, Field
+from pydantic_core import PydanticUndefined
 import logging
-import sys
+
+
+# we support python 3.10 so we define our own StrEnum (introduced in 3.11)
+class StrEnum(str, Enum):
+    """Backport of StrEnum for Python < 3.11"""
+    def __str__(self):
+        return self.value
 
 
 class ColoredFormatter(logging.Formatter):
@@ -119,6 +128,9 @@ class Parameter(BaseModel):
             }
         }
 
+class OpenAIAPI(StrEnum):
+    COMPLETIONS = "completions"
+    RESPONSES = "responses"
 
 class FunctionSchema(BaseModel):
     """Class that consumes a function and can return a schema required by
@@ -167,29 +179,95 @@ class FunctionSchema(BaseModel):
         )
 
     @classmethod
-    def from_pydantic(cls, model: BaseModel) -> "FunctionSchema":
+    def from_pydantic(cls, model: type[BaseModel]) -> "FunctionSchema":
+        """Create a FunctionSchema from a Pydantic model class.
+        
+        :param model: The Pydantic model class to convert
+        :type model: type[BaseModel]
+        :return: FunctionSchema instance
+        :rtype: FunctionSchema
+        """
+        # Extract model metadata
+        name = model.__name__
+        description = model.__doc__ or ""
+        
+        # Build parameters list
+        parameters = []
         signature_parts = []
-        for field_name, field_model in model.__annotations__.items():
-            field_info = model.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}"
+        
+        for field_name, field_info in model.model_fields.items():
+            # Get the field type
+            field_type = model.__annotations__.get(field_name)
+            
+            # Determine the type name - handle Optional and other generic types
+            type_name = str(field_type)
+            
+            # Try to extract the actual type from Optional[T] -> T
+            origin = get_origin(field_type)
+            args = get_args(field_type)
+            
+            if origin is Union:
+                # This is likely Optional[T] which is Union[T, None]
+                non_none_types = [arg for arg in args if arg is not type(None)]
+                if non_none_types:
+                    actual_type = non_none_types[0]
+                    if hasattr(actual_type, '__name__'):
+                        type_name = actual_type.__name__
+                    else:
+                        type_name = str(actual_type)
+            elif field_type and hasattr(field_type, '__name__'):
+                type_name = field_type.__name__
+            
+            # Check if field is required (no default value)
+            # In Pydantic v2, PydanticUndefined means no default
+            is_required = (
+                field_info.default is PydanticUndefined 
+                and field_info.default_factory is None
+            )
+            
+            # Get the actual default value
+            if field_info.default is not PydanticUndefined and field_info.default is not None:
+                default_value = field_info.default
+            elif field_info.default_factory is not None:
+                # For default_factory, we can't always call it without arguments
+                # Just use a placeholder to indicate there's a factory
+                try:
+                    # Try calling with no arguments (common case)
+                    default_value = field_info.default_factory()  # type: ignore[call-arg]
+                except TypeError:
+                    # If it needs arguments, just indicate it has a factory default
+                    default_value = "<factory>"
+            else:
+                default_value = inspect.Parameter.empty
+            
+            # Add parameter
+            parameters.append(
+                Parameter(
+                    name=field_name,
+                    description=field_info.description,
+                    type=type_name,
+                    default=default_value,
+                    required=is_required,
                 )
+            )
+            
+            # Build signature part
+            if default_value != inspect.Parameter.empty:
+                signature_parts.append(f"{field_name}: {type_name} = {repr(default_value)}")
             else:
-                signature_part = f"{field_name}: {field_model.__name__}"
-            signature_parts.append(signature_part)
-        signature = f"({', '.join(signature_parts)}) -> str"
+                signature_parts.append(f"{field_name}: {type_name}")
+        
+        signature = f"({', '.join(signature_parts)}) -> dict"
+        
         return cls.model_construct(
-            name=model.__class__.__name__,
-            description=model.__doc__ or "",
+            name=name,
+            description=description,
             signature=signature,
-            output="",  # TODO: Implement output
-            parameters=[],
+            output="dict",
+            parameters=parameters,
         )
 
-    def to_dict(self) -> dict:
+    def to_dict(self) -> dict[str, Any]:
         schema_dict = {
             "type": "function",
             "function": {
@@ -210,14 +288,42 @@ class FunctionSchema(BaseModel):
         }
         return schema_dict
 
-    def to_openai(self) -> dict:
-        return self.to_dict()
+    def to_openai(self, api: OpenAIAPI=OpenAIAPI.COMPLETIONS) -> dict[str, Any]:
+        """Convert the function schema into OpenAI-compatible formats. Supports
+        both completions and responses APIs.
+
+        :param api: The API to convert to.
+        :type api: OpenAIAPI
+        :return: The function schema in OpenAI-compatible format.
+        :rtype: dict
+        """
+        if api == "completions":
+            return self.to_dict()
+        elif api == "responses":
+            return {
+                "type": "function",
+                "name": self.name,
+                "description": self.description,
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        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
+                    ],
+                },
+            }
+        else:
+            raise ValueError(f"Unrecognized OpenAI API: {api}")
 
 
 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[str, Any]]:
     if format in DEFAULT:
         return [
             FunctionSchema.from_callable(callable).to_dict() for callable in callables

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

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

graphai_lib-0.0.9rc1.dist-info/RECORD

@@ -0,0 +1,11 @@
+graphai/__init__.py,sha256=UbqXq7iGIYe1GyTPcpgLSXbgWovggFsAbTMtr4JQm3M,160
+graphai/callback.py,sha256=Wl0JCmE8NcFifKmP9-a5bFa0WKVHTdrSClHVRmIEPpc,7323
+graphai/graph.py,sha256=Bm_Jwa5EMUACqXTZSoibVQfHSwS1rV-ExfxqYRIkPDY,18944
+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.9rc1.dist-info/METADATA,sha256=T8nS6wOKwd7ZfzVvD_h4m38bowb9zzr7TutdIz5GWJQ,913
+graphai_lib-0.0.9rc1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+graphai_lib-0.0.9rc1.dist-info/top_level.txt,sha256=TXlqmhLViX-3xGH2g5w6cavRd-QMf229Hl88jdMOGt8,8
+graphai_lib-0.0.9rc1.dist-info/RECORD,,

graphai_lib-0.0.8.dist-info/RECORD

@@ -1,10 +0,0 @@
-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,,