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

graphai/__init__.py

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

graphai/callback.py

@@ -1,12 +1,12 @@
 import asyncio
 from pydantic import Field
-from typing import Optional
+from typing import Optional, Any
 from collections.abc import AsyncIterator
-from semantic_router.utils.logger import logger
 
 
 log_stream = True
 
+
 class Callback:
     identifier: str = Field(
         default="graphai",
@@ -14,7 +14,7 @@ class Callback:
             "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}>",
@@ -32,8 +32,8 @@ class Callback:
         examples=[
             "<{identifier}:{token}:{params}>",
             "<[{identifier} | {token} | {params}]>",
-            "<{token}:{params}>"
-        ]
+            "<{token}:{params}>",
+        ],
     )
     token_format: str = Field(
         default="{token}",
@@ -41,27 +41,23 @@ class Callback:
             "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
+        exclude=True,
     )
     _current_node_name: Optional[str] = Field(
-        default=None,
-        description="The name of the current node.",
-        exclude=True
+        default=None, description="The name of the current node.", exclude=True
     )
     _active: bool = Field(
-        default=True,
-        description="Whether the callback is active.",
-        exclude=True
+        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
+        exclude=True,
     )
     queue: asyncio.Queue
 
@@ -83,7 +79,7 @@ class Callback:
     @property
     def first_token(self) -> bool:
         return self._first_token
-    
+
     @first_token.setter
     def first_token(self, value: bool):
         self._first_token = value
@@ -91,7 +87,7 @@ class Callback:
     @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
@@ -99,7 +95,7 @@ class Callback:
     @property
     def active(self) -> bool:
         return self._active
-    
+
     @active.setter
     def active(self, value: bool):
         self._active = value
@@ -110,7 +106,7 @@ class Callback:
         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:
@@ -118,16 +114,13 @@ class Callback:
         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.
         """
-        end_token = await self._build_special_token(
-            name="END",
-            params=None
-        )
+        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():
@@ -142,8 +135,7 @@ class Callback:
         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.
-        """
+        """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
@@ -152,27 +144,23 @@ class Callback:
         self.active = active
         if self.active:
             token = await self._build_special_token(
-                name=f"{self.current_node_name}:start",
-                params=None
+                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
+                name=self.current_node_name, params=None
             )
             self.queue.put_nowait(node_token)
-    
+
     async def end_node(self, node_name: str):
-        """Emits the end token for the current node.
-        """
+        """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
+        # self.current_node_name = node_name
         if self.active:
             node_token = await self._build_special_token(
-                name=f"{self.current_node_name}:end",
-                params=None
+                name=f"{self.current_node_name}:end", params=None
             )
             self.queue.put_nowait(node_token)
 
@@ -182,10 +170,7 @@ class Callback:
         """
         if self._done:
             return
-        end_token = await self._build_special_token(
-            name="END",
-            params=None
-        )
+        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
@@ -198,8 +183,10 @@ class Callback:
                 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):
+
+    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:
@@ -209,7 +196,5 @@ class Callback:
         else:
             identifier = ""
         return self.special_token_format.format(
-            identifier=identifier,
-            token=name,
-            params=params_str
+            identifier=identifier, token=name, params=params_str
         )

graphai/graph.py

@@ -1,17 +1,33 @@
-from typing import List, Dict, Any, Optional
-from graphai.nodes.base import _Node
+from typing import List, Dict, Any, Optional, Protocol, Type
 from graphai.callback import Callback
-from semantic_router.utils.logger import logger
+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]: ...
 
 
 class Graph:
-    def __init__(self, max_steps: int = 10, initial_state: Optional[Dict[str, Any]] = None):
-        self.nodes: Dict[str, _Node] = {}
+    def __init__(
+        self, max_steps: int = 10, initial_state: Optional[Dict[str, Any]] = None
+    ):
+        self.nodes: Dict[str, NodeProtocol] = {}
         self.edges: List[Any] = []
-        self.start_node: Optional[_Node] = None
-        self.end_nodes: List[_Node] = []
-        self.Callback = Callback
-        self.callback = None
+        self.start_node: Optional[NodeProtocol] = None
+        self.end_nodes: List[NodeProtocol] = []
+        self.Callback: Type[Callback] = Callback
         self.max_steps = max_steps
         self.state = initial_state or {}
 
@@ -32,7 +48,7 @@ class Graph:
         """Reset the graph state to an empty dict."""
         self.state = {}
 
-    def add_node(self, node):
+    def add_node(self, node: NodeProtocol):
         if node.name in self.nodes:
             raise Exception(f"Node with name '{node.name}' already exists.")
         self.nodes[node.name] = node
@@ -47,9 +63,9 @@ class Graph:
         if node.is_end:
             self.end_nodes.append(node)
 
-    def add_edge(self, source: _Node | str, destination: _Node | str):
+    def add_edge(self, source: NodeProtocol | str, destination: NodeProtocol | 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.
@@ -60,7 +76,7 @@ class Graph:
             source_node = self.nodes.get(source)
         else:
             # Check if it's a node-like object by looking for required attributes
-            if hasattr(source, 'name'):
+            if hasattr(source, "name"):
                 source_node = self.nodes.get(source.name)
         if source_node is None:
             raise ValueError(
@@ -71,7 +87,7 @@ class Graph:
             destination_node = self.nodes.get(destination)
         else:
             # Check if it's a node-like object by looking for required attributes
-            if hasattr(destination, 'name'):
+            if hasattr(destination, "name"):
                 destination_node = self.nodes.get(destination.name)
         if destination_node is None:
             raise ValueError(
@@ -80,17 +96,19 @@ class Graph:
         edge = Edge(source_node, destination_node)
         self.edges.append(edge)
 
-    def add_router(self, sources: list[_Node], router: _Node, destinations: List[_Node]):
+    def add_router(
+        self, sources: list[NodeProtocol], router: NodeProtocol, destinations: List[NodeProtocol]
+    ):
         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):
+    def set_start_node(self, node: NodeProtocol):
         self.start_node = node
 
-    def set_end_node(self, node: _Node):
+    def set_end_node(self, node: NodeProtocol):
         self.end_node = node
 
     def compile(self):
@@ -112,11 +130,15 @@ class Graph:
                 f"Instead, got {type(output)} from '{output}'."
             )
 
-    async def execute(self, input):
+    async def execute(self, input, callback: Callback | None = None):
         # 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()
+        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
@@ -124,11 +146,13 @@ class Graph:
             # 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)
+                await callback.start_node(node_name=current_node.name)
                 # Include graph's internal state in the node execution context
-                output = await current_node.invoke(input=state, callback=self.callback, state=self.state)
+                output = await current_node.invoke(
+                    input=state, callback=callback, state=self.state
+                )
                 self._validate_output(output=output, node_name=current_node.name)
-                await self.callback.end_node(node_name=current_node.name)
+                await callback.end_node(node_name=current_node.name)
             else:
                 # Include graph's internal state in the node execution context
                 output = await current_node.invoke(input=state, state=self.state)
@@ -153,24 +177,38 @@ class Graph:
                     "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()
+        if callback and "callback" in state:
+            await callback.close()
             del state["callback"]
         return state
 
     def get_callback(self):
-        self.callback = self.Callback()
-        return self.callback
+        """Get a new instance of the callback class.
+
+        :return: A new instance of the callback class.
+        :rtype: Callback
+        """
+        callback = self.Callback()
+        return callback
+
+    def set_callback(self, callback_class: Type[Callback]):
+        """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.
+
+        :param callback_class: The callback class to use as the default callback.
+        :type callback_class: Type[Callback]
+        """
+        self.Callback = callback_class
 
-    def _get_node_by_name(self, node_name: str) -> _Node:
+    def _get_node_by_name(self, node_name: str) -> NodeProtocol:
         """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.
         """
@@ -191,12 +229,16 @@ class Graph:
         try:
             import networkx as nx
         except ImportError:
-            raise ImportError("NetworkX is required for visualization. Please install it with 'pip install networkx'.")
+            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'.")
+            raise ImportError(
+                "Matplotlib is required for visualization. Please install it with 'pip install matplotlib'."
+            )
 
         G = nx.DiGraph()
 
@@ -207,7 +249,9 @@ class Graph:
             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.")
+            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))
@@ -241,20 +285,30 @@ class Graph:
                 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.")
+            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)
+        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.axis("off")
         plt.show()
 
 
-
 class Edge:
     def __init__(self, source, destination):
         self.source = source
-        self.destination = destination
+        self.destination = destination

graphai/nodes/__init__.py

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

graphai/nodes/base.py

@@ -1,5 +1,6 @@
 import inspect
 from typing import Any, Callable, Dict, Optional
+from pydantic import Field
 
 from graphai.callback import Callback
 from graphai.utils import FunctionSchema
@@ -9,7 +10,11 @@ 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"}
+        init_params = {
+            name: arg
+            for name, arg in init_signature.parameters.items()
+            if name != "self"
+        }
         return init_params
 
     def __call__(cls, *args, **kwargs):
@@ -33,18 +38,32 @@ class _Node:
         stream: bool = False,
         name: str | None = None,
     ) -> Callable:
-        """Decorator validating node structure.
-        """
+        """Decorator validating node structure."""
         if not callable(func):
             raise ValueError("Node must be a callable function.")
-        
+
         func_signature = inspect.signature(func)
-        schema = FunctionSchema(func)
+        schema: FunctionSchema = FunctionSchema.from_callable(func)
 
         class NodeClass:
             _func_signature = func_signature
-            is_router = None
-            _stream = stream
+            is_router: bool = Field(
+                default=False, description="Whether the node is a router."
+            )
+            # following attributes will be overridden by the decorator
+            name: str | None = Field(default=None, description="The name of the node.")
+            is_start: bool = Field(
+                default=False, description="Whether the node is the start of the graph."
+            )
+            is_end: bool = Field(
+                default=False, description="Whether the node is the end of the graph."
+            )
+            schema: FunctionSchema | None = Field(
+                default=None, description="The schema of the node."
+            )
+            stream: bool = Field(
+                default=False, description="Whether the node includes streaming object."
+            )
 
             def __init__(self):
                 self._expected_params = set(self._func_signature.parameters.keys())
@@ -56,9 +75,13 @@ class _Node:
 
             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}
+                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'
+                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}
@@ -87,7 +110,6 @@ class _Node:
                     )
                 return filtered_params
 
-
             @classmethod
             def get_signature(cls):
                 """Returns the signature of the decorated function as LLM readable
@@ -97,15 +119,24 @@ class _Node:
                 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}")
+                            signature_components.append(
+                                f"{param.name}: {param.annotation}"
+                            )
                         else:
-                            signature_components.append(f"{param.name}: {param.annotation} = {param.default}")
+                            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, state: Optional[Dict[str, Any]] = None):
+            async def invoke(
+                cls,
+                input: Dict[str, Any],
+                callback: Optional[Callback] = None,
+                state: Optional[Dict[str, Any]] = None,
+            ):
                 if callback:
                     if stream:
                         input["callback"] = callback
@@ -116,13 +147,16 @@ class _Node:
                 # Add state to the input if present and the parameter exists in the function signature
                 if state is not None and "state" in cls._func_signature.parameters:
                     input["state"] = state
-                
+
                 instance = cls()
                 out = await instance.execute(**input)
                 return out
 
         NodeClass.__name__ = func.__name__
-        NodeClass.name = name or 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
         NodeClass.is_end = end
@@ -141,8 +175,11 @@ class _Node:
     ):
         # 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, name=name) -> Callable:
+        def wrap(
+            func: Callable, start=start, end=end, stream=stream, name=name
+        ) -> Callable:
             return self._node(func=func, start=start, end=end, stream=stream, name=name)
+
         if func:
             # Decorator is called without parenthesis
             return wrap(func=func, start=start, end=end, stream=stream, name=name)

graphai/utils.py

@@ -1,11 +1,81 @@
 import inspect
-from typing import Any, Callable, Dict, List, Union, Optional
+from typing import Any, Callable, List, Optional
 from pydantic import BaseModel, Field
+import logging
+
+import colorlog
+
+
+class CustomFormatter(colorlog.ColoredFormatter):
+    """Custom formatter for the logger."""
+
+    def __init__(self):
+        super().__init__(
+            "%(log_color)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 add_coloured_handler(logger):
+    """Add a coloured handler to the logger."""
+    formatter = CustomFormatter()
+    console_handler = logging.StreamHandler()
+    console_handler.setFormatter(formatter)
+    logger.addHandler(console_handler)
+    return logger
+
+
+def setup_custom_logger(name):
+    """Setup a custom logger."""
+    logger = logging.getLogger(name)
+
+    if not logger.hasHandlers():
+        add_coloured_handler(logger)
+        logger.setLevel(logging.INFO)
+        logger.propagate = False
+
+    return logger
+
+
+logger: logging.Logger = setup_custom_logger(__name__)
+
+
+def openai_type_mapping(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"
 
 
 class Parameter(BaseModel):
-    class Config:
-        arbitrary_types_allowed = True
+    """Parameter for a function.
+
+    :param name: The name of the parameter.
+    :type name: str
+    :param description: The description of the parameter.
+    :type description: Optional[str]
+    :param type: The type of the parameter.
+    :type type: str
+    :param default: The default value of the parameter.
+    :type default: Any
+    :param required: Whether the parameter is required.
+    :type required: bool
+    """
 
     name: str = Field(description="The name of the parameter")
     description: Optional[str] = Field(
@@ -15,15 +85,22 @@ class Parameter(BaseModel):
     default: Any = Field(description="The default value of the parameter")
     required: bool = Field(description="Whether the parameter is required")
 
-    def to_openai(self):
+    def to_dict(self) -> dict[str, Any]:
+        """Convert the parameter to a dictionary for an standard dictionary-based function schema.
+        This is the most common format used by LLM providers, including OpenAI, Ollama, and others.
+
+        :return: The parameter in dictionary format.
+        :rtype: dict[str, Any]
+        """
         return {
             self.name: {
                 "description": self.description,
-                "type": self.type,
+                "type": openai_type_mapping(self.type),
             }
         }
 
-class FunctionSchema:
+
+class FunctionSchema(BaseModel):
     """Class that consumes a function and can return a schema required by
     different LLMs for function calling.
     """
@@ -32,35 +109,68 @@ class FunctionSchema:
     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")
+    parameters: list[Parameter] = Field(description="The parameters of the function")
+
+    @classmethod
+    def from_callable(cls, function: Callable) -> "FunctionSchema":
+        """Initialize the FunctionSchema.
 
-    def __init__(self, function: Union[Callable, BaseModel]):
-        self.function = function
+        :param function: The function to consume.
+        :type function: Callable
+        """
         if callable(function):
-            self._process_function(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,
+                    )
+                )
+            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")
 
-    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,
+    @classmethod
+    def from_pydantic(cls, model: BaseModel) -> "FunctionSchema":
+        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}"
                 )
-            )
-        self.parameters = parameters
+            else:
+                signature_part = f"{field_name}: {field_model.__name__}"
+            signature_parts.append(signature_part)
+        signature = f"({', '.join(signature_parts)}) -> str"
+        return cls.model_construct(
+            name=model.__class__.__name__,
+            description=model.__doc__ or "",
+            signature=signature,
+            output="",  # TODO: Implement output
+            parameters=[],
+        )
 
-    def to_openai(self):
+    def to_dict(self) -> dict:
         schema_dict = {
             "type": "function",
             "function": {
@@ -69,15 +179,7 @@ class FunctionSchema:
                 "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
+                        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
@@ -87,39 +189,17 @@ class FunctionSchema:
         }
         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 to_openai(self) -> dict:
+        return self.to_dict()
 
 
-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
+DEFAULT = set(["default", "openai", "ollama", "litellm"])
 
-        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
+
+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
+        ]
+    else:
+        raise ValueError(f"Format {format} not supported")

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

@@ -1,12 +1,13 @@
 Metadata-Version: 2.4
 Name: graphai-lib
-Version: 0.0.4
+Version: 0.0.6
 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
+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"
@@ -16,7 +17,7 @@ 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"
+Requires-Dist: types-networkx>=3.4.2.20250319; extra == "dev"
 Provides-Extra: docs
 Requires-Dist: pydoc-markdown>=4.8.2; python_version < "3.12" and extra == "docs"
 

graphai_lib-0.0.6.dist-info/RECORD

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

{graphai_lib-0.0.4.dist-info → graphai_lib-0.0.6.dist-info}/WHEEL

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

graphai_lib-0.0.4.dist-info/RECORD

@@ -1,10 +0,0 @@
-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/nodes/__init__.py,sha256=4826Ubk5yUfbVH7F8DmoTKQyax624Q2QJHsGxqgQ_ng,73
-graphai/nodes/base.py,sha256=SZdYhFfXdtFmabFbMRcEGd8_h8w-g6s4I7hMEo6JCk8,6331
-graphai_lib-0.0.4.dist-info/METADATA,sha256=wfxK82ZO-slBl-xVNmdbYGVvPo7YKm8-SLfdrkNWtKs,982
-graphai_lib-0.0.4.dist-info/WHEEL,sha256=1tXe9gY0PYatrMPMDd6jXqjfpz_B-Wqm32CPfRC58XU,91
-graphai_lib-0.0.4.dist-info/top_level.txt,sha256=TXlqmhLViX-3xGH2g5w6cavRd-QMf229Hl88jdMOGt8,8
-graphai_lib-0.0.4.dist-info/RECORD,,