Add logger import and update edge type in Graph class
This commit is contained in:
parent
952f8d39e5
commit
4e5181307f
1 changed files with 5 additions and 4 deletions
|
|
@ -2,6 +2,8 @@ from collections import defaultdict, deque
|
||||||
from typing import Dict, Generator, List, Type, Union
|
from typing import Dict, Generator, List, Type, Union
|
||||||
|
|
||||||
from langchain.chains.base import Chain
|
from langchain.chains.base import Chain
|
||||||
|
from loguru import logger
|
||||||
|
|
||||||
from langflow.graph.edge.base import ContractEdge, Edge
|
from langflow.graph.edge.base import ContractEdge, Edge
|
||||||
from langflow.graph.graph.constants import lazy_load_vertex_dict
|
from langflow.graph.graph.constants import lazy_load_vertex_dict
|
||||||
from langflow.graph.graph.utils import process_flow
|
from langflow.graph.graph.utils import process_flow
|
||||||
|
|
@ -9,7 +11,6 @@ from langflow.graph.vertex.base import Vertex
|
||||||
from langflow.graph.vertex.types import FileToolVertex, LLMVertex, ToolkitVertex
|
from langflow.graph.vertex.types import FileToolVertex, LLMVertex, ToolkitVertex
|
||||||
from langflow.interface.tools.constants import FILE_TOOLS
|
from langflow.interface.tools.constants import FILE_TOOLS
|
||||||
from langflow.utils import payload
|
from langflow.utils import payload
|
||||||
from loguru import logger
|
|
||||||
|
|
||||||
|
|
||||||
class Graph:
|
class Graph:
|
||||||
|
|
@ -192,13 +193,13 @@ class Graph:
|
||||||
neighbors[neighbor] += 1
|
neighbors[neighbor] += 1
|
||||||
return neighbors
|
return neighbors
|
||||||
|
|
||||||
def _build_edges(self) -> List[Edge]:
|
def _build_edges(self) -> List[ContractEdge]:
|
||||||
"""Builds the edges of the graph."""
|
"""Builds the edges of the graph."""
|
||||||
# Edge takes two vertices as arguments, so we need to build the vertices first
|
# Edge takes two vertices as arguments, so we need to build the vertices first
|
||||||
# and then build the edges
|
# and then build the edges
|
||||||
# if we can't find a vertex, we raise an error
|
# if we can't find a vertex, we raise an error
|
||||||
|
|
||||||
edges: List[Edge] = []
|
edges: List[ContractEdge] = []
|
||||||
for edge in self._edges:
|
for edge in self._edges:
|
||||||
source = self.get_vertex(edge["source"])
|
source = self.get_vertex(edge["source"])
|
||||||
target = self.get_vertex(edge["target"])
|
target = self.get_vertex(edge["target"])
|
||||||
|
|
@ -206,7 +207,7 @@ class Graph:
|
||||||
raise ValueError(f"Source vertex {edge['source']} not found")
|
raise ValueError(f"Source vertex {edge['source']} not found")
|
||||||
if target is None:
|
if target is None:
|
||||||
raise ValueError(f"Target vertex {edge['target']} not found")
|
raise ValueError(f"Target vertex {edge['target']} not found")
|
||||||
edges.append(Edge(source, target, edge))
|
edges.append(ContractEdge(source, target, edge))
|
||||||
return edges
|
return edges
|
||||||
|
|
||||||
def _get_vertex_class(self, vertex_type: str, vertex_base_type: str) -> Type[Vertex]:
|
def _get_vertex_class(self, vertex_type: str, vertex_base_type: str) -> Type[Vertex]:
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue