lfx-nightly 0.2.0.dev25__py3-none-any.whl

lfx/graph/edge/schema.py

@@ -0,0 +1,119 @@
+from typing import Any
+
+from pydantic import ConfigDict, Field, field_validator
+from typing_extensions import TypedDict
+
+from lfx.helpers.base_model import BaseModel
+
+
+class SourceHandleDict(TypedDict, total=False):
+    baseClasses: list[str]
+    dataType: str
+    id: str
+    name: str | None
+    output_types: list[str]
+
+
+class TargetHandleDict(TypedDict):
+    fieldName: str
+    id: str
+    inputTypes: list[str] | None
+    type: str
+
+
+class LoopTargetHandleDict(TypedDict):
+    dataType: str
+    id: str
+    name: str
+    output_types: list[str]
+
+
+class EdgeDataDetails(TypedDict):
+    sourceHandle: SourceHandleDict
+    targetHandle: TargetHandleDict | LoopTargetHandleDict
+
+
+class EdgeData(TypedDict, total=False):
+    source: str
+    target: str
+    data: EdgeDataDetails
+
+
+class ResultPair(BaseModel):
+    result: Any
+    extra: Any
+
+
+class Payload(BaseModel):
+    result_pairs: list[ResultPair] = []
+
+    def __iter__(self):
+        return iter(self.result_pairs)
+
+    def add_result_pair(self, result: Any, extra: Any | None = None) -> None:
+        self.result_pairs.append(ResultPair(result=result, extra=extra))
+
+    def get_last_result_pair(self) -> ResultPair:
+        return self.result_pairs[-1]
+
+    # format all but the last result pair
+    # into a string
+    def format(self, sep: str = "\n") -> str:
+        # Result: the result
+        # Extra: the extra if it exists don't show if it doesn't
+        return sep.join(
+            [
+                f"Result: {result_pair.result}\nExtra: {result_pair.extra}"
+                if result_pair.extra is not None
+                else f"Result: {result_pair.result}"
+                for result_pair in self.result_pairs[:-1]
+            ]
+        )
+
+
+class TargetHandle(BaseModel):
+    model_config = ConfigDict(populate_by_name=True)
+    field_name: str = Field(..., alias="fieldName", description="Field name for the target handle.")
+    id: str = Field(..., description="Unique identifier for the target handle.")
+    input_types: list[str] = Field(
+        default_factory=list, alias="inputTypes", description="List of input types for the target handle."
+    )
+    type: str = Field(None, description="Type of the target handle.")
+
+    @classmethod
+    def from_loop_target_handle(cls, target_handle: LoopTargetHandleDict) -> "TargetHandle":
+        # The target handle is a loop edge
+        # The target handle is a dict with the following keys:
+        # - name: str
+        # - id: str
+        # - inputTypes: list[str]
+        # - type: str
+        # It is built from an Output, which is why it has a different structure
+        return cls(
+            field_name=target_handle.get("name"),
+            id=target_handle.get("id"),
+            input_types=target_handle.get("output_types"),
+        )
+
+
+class SourceHandle(BaseModel):
+    model_config = ConfigDict(populate_by_name=True)
+    base_classes: list[str] = Field(
+        default_factory=list, alias="baseClasses", description="List of base classes for the source handle."
+    )
+    data_type: str = Field(..., alias="dataType", description="Data type for the source handle.")
+    id: str = Field(..., description="Unique identifier for the source handle.")
+    name: str | None = Field(None, description="Name of the source handle.")
+    output_types: list[str] = Field(default_factory=list, description="List of output types for the source handle.")
+
+    @field_validator("name", mode="before")
+    @classmethod
+    def validate_name(cls, v, info):
+        if info.data["data_type"] == "GroupNode":
+            # 'OpenAIModel-u4iGV_text_output'
+            splits = v.split("_", 1)
+            if len(splits) != 2:  # noqa: PLR2004
+                msg = f"Invalid source handle name {v}"
+                raise ValueError(msg)
+            v = splits[1]
+        return v

lfx/graph/graph/ascii.py

@@ -0,0 +1,202 @@
+"""This code is adapted from the DVC project.
+
+Original source:
+https://github.com/iterative/dvc/blob/c5bac1c8cfdb2c0f54d52ac61ff754e6f583822a/dvc/dagascii.py
+
+The original code has been modified to focus on drawing a Directed Acyclic Graph (DAG) in ASCII
+using the `grandalf` library. Non-essential parts have been removed, and the code has been
+refactored to suit this specific use case.
+
+
+"""
+
+import math
+
+from grandalf.graphs import Edge as GrandalfEdge
+from grandalf.graphs import Graph as GrandalfGraph
+from grandalf.graphs import Vertex as GrandalfVertex
+from grandalf.layouts import SugiyamaLayout
+from grandalf.routing import EdgeViewer, route_with_lines
+
+MINIMUM_EDGE_VIEW_POINTS = 2
+
+
+class VertexViewer:
+    """Class to define vertex box boundaries that will be accounted for during graph building by grandalf."""
+
+    HEIGHT = 3  # top and bottom box edges + text
+
+    def __init__(self, name) -> None:
+        self._h = self.HEIGHT  # top and bottom box edges + text
+        self._w = len(name) + 2  # right and left bottom edges + text
+
+    @property
+    def h(self):
+        return self._h
+
+    @property
+    def w(self):
+        return self._w
+
+
+class AsciiCanvas:
+    """Class for drawing in ASCII."""
+
+    def __init__(self, cols, lines) -> None:
+        if cols <= 1:
+            msg = "cols must be greater than 1"
+            raise ValueError(msg)
+        if lines <= 1:
+            msg = "lines must be greater than 1"
+            raise ValueError(msg)
+        self.cols = cols
+        self.lines = lines
+        self.canvas = [[" "] * cols for _ in range(lines)]
+
+    def get_lines(self):
+        return map("".join, self.canvas)
+
+    def draws(self):
+        return "\n".join(self.get_lines())
+
+    def draw(self) -> None:
+        """Draws ASCII canvas on the screen."""
+        lines = self.get_lines()
+        print("\n".join(lines))  # noqa: T201
+
+    def point(self, x, y, char) -> None:
+        """Create a point on ASCII canvas."""
+        if len(char) != 1:
+            msg = "char must be a single character"
+            raise ValueError(msg)
+        if x < 0 or x >= self.cols:
+            msg = "x is out of bounds"
+            raise ValueError(msg)
+        if y < 0 or y >= self.lines:
+            msg = "y is out of bounds"
+            raise ValueError(msg)
+        self.canvas[y][x] = char
+
+    def line(self, x0, y0, x1, y1, char) -> None:
+        """Create a line on ASCII canvas."""
+        if x0 > x1:
+            x1, x0 = x0, x1
+            y1, y0 = y0, y1
+
+        dx = x1 - x0
+        dy = y1 - y0
+
+        if dx == 0 and dy == 0:
+            self.point(x0, y0, char)
+        elif abs(dx) >= abs(dy):
+            for x in range(x0, x1 + 1):
+                y = y0 + round((x - x0) * dy / float(dx)) if dx else y0
+                self.point(x, y, char)
+        else:
+            for y in range(min(y0, y1), max(y0, y1) + 1):
+                x = x0 + round((y - y0) * dx / float(dy)) if dy else x0
+                self.point(x, y, char)
+
+    def text(self, x, y, text) -> None:
+        """Print a text on ASCII canvas."""
+        for i, char in enumerate(text):
+            self.point(x + i, y, char)
+
+    def box(self, x0, y0, width, height) -> None:
+        """Create a box on ASCII canvas."""
+        if width <= 1:
+            msg = "width must be greater than 1"
+            raise ValueError(msg)
+        if height <= 1:
+            msg = "height must be greater than 1"
+            raise ValueError(msg)
+        width -= 1
+        height -= 1
+
+        for x in range(x0, x0 + width):
+            self.point(x, y0, "-")
+            self.point(x, y0 + height, "-")
+        for y in range(y0, y0 + height):
+            self.point(x0, y, "|")
+            self.point(x0 + width, y, "|")
+        self.point(x0, y0, "+")
+        self.point(x0 + width, y0, "+")
+        self.point(x0, y0 + height, "+")
+        self.point(x0 + width, y0 + height, "+")
+
+
+def build_sugiyama_layout(vertexes, edges):
+    vertexes = {v: GrandalfVertex(v) for v in vertexes}
+    edges = [GrandalfEdge(vertexes[s], vertexes[e]) for s, e in edges]
+    graph = GrandalfGraph(vertexes.values(), edges)
+
+    for vertex in vertexes.values():
+        vertex.view = VertexViewer(vertex.data)
+
+    minw = min(v.view.w for v in vertexes.values())
+
+    for edge in edges:
+        edge.view = EdgeViewer()
+
+    sug = SugiyamaLayout(graph.C[0])
+    roots = [v for v in sug.g.sV if len(v.e_in()) == 0]
+    sug.init_all(roots=roots, optimize=True)
+
+    sug.yspace = VertexViewer.HEIGHT
+    sug.xspace = minw
+    sug.route_edge = route_with_lines
+
+    sug.draw()
+    return sug
+
+
+def draw_graph(vertexes, edges, *, return_ascii=True):
+    """Build a DAG and draw it in ASCII."""
+    sug = build_sugiyama_layout(vertexes, edges)
+
+    xlist = []
+    ylist = []
+
+    for vertex in sug.g.sV:
+        xlist.extend([vertex.view.xy[0] - vertex.view.w / 2.0, vertex.view.xy[0] + vertex.view.w / 2.0])
+        ylist.extend([vertex.view.xy[1], vertex.view.xy[1] + vertex.view.h])
+
+    for edge in sug.g.sE:
+        for x, y in edge.view._pts:
+            xlist.append(x)
+            ylist.append(y)
+
+    minx = min(xlist)
+    miny = min(ylist)
+    maxx = max(xlist)
+    maxy = max(ylist)
+
+    canvas_cols = math.ceil(maxx - minx) + 1
+    canvas_lines = round(maxy - miny)
+
+    canvas = AsciiCanvas(canvas_cols, canvas_lines)
+
+    for edge in sug.g.sE:
+        if len(edge.view._pts) < MINIMUM_EDGE_VIEW_POINTS:
+            msg = "edge.view._pts must have at least 2 points"
+            raise ValueError(msg)
+        for index in range(1, len(edge.view._pts)):
+            start = edge.view._pts[index - 1]
+            end = edge.view._pts[index]
+            canvas.line(
+                round(start[0] - minx),
+                round(start[1] - miny),
+                round(end[0] - minx),
+                round(end[1] - miny),
+                "*",
+            )
+
+    for vertex in sug.g.sV:
+        x = vertex.view.xy[0] - vertex.view.w / 2.0
+        y = vertex.view.xy[1]
+        canvas.box(round(x - minx), round(y - miny), vertex.view.w, vertex.view.h)
+        canvas.text(round(x - minx) + 1, round(y - miny) + 1, vertex.data)
+    if return_ascii:
+        return canvas.draws()
+    canvas.draw()
+    return None