From cd43aa729a3fc7286a44a68e46c6192ceadd8071 Mon Sep 17 00:00:00 2001
From: Gabriel Luiz Freitas Almeida <gabriel@logspace.ai>
Date: Thu, 1 Feb 2024 12:13:05 -0300
Subject: [PATCH] Map ChatVertex connections and update external vertex
 attributes

---
 src/backend/langflow/graph/graph/base.py | 34 +++++++++++++++++++++---
 1 file changed, 31 insertions(+), 3 deletions(-)

diff --git a/src/backend/langflow/graph/graph/base.py b/src/backend/langflow/graph/graph/base.py
index 288d77d0f..4fb30503d 100644
--- a/src/backend/langflow/graph/graph/base.py
+++ b/src/backend/langflow/graph/graph/base.py
@@ -2,15 +2,14 @@ from collections import defaultdict, deque
 from typing import Dict, Generator, List, Type, Union
 
 from langchain.chains.base import Chain
-from loguru import logger
-
 from langflow.graph.edge.base import ContractEdge
 from langflow.graph.graph.constants import lazy_load_vertex_dict
 from langflow.graph.graph.utils import process_flow
 from langflow.graph.vertex.base import Vertex
-from langflow.graph.vertex.types import FileToolVertex, LLMVertex, ToolkitVertex
+from langflow.graph.vertex.types import ChatVertex, FileToolVertex, LLMVertex, ToolkitVertex
 from langflow.interface.tools.constants import FILE_TOOLS
 from langflow.utils import payload
+from loguru import logger
 
 
 class Graph:
@@ -21,6 +20,8 @@ class Graph:
         nodes: List[Dict],
         edges: List[Dict[str, str]],
     ) -> None:
+        self.inputs = []
+        self.outputs = []
         self._vertices = nodes
         self._edges = edges
         self.raw_graph_data = {"nodes": nodes, "edges": edges}
@@ -81,6 +82,31 @@ class Graph:
         self._build_vertex_params()
         # remove invalid vertices
         self._validate_vertices()
+        # Now that we have the vertices and edges
+        # We need to map the vertices that are connected to
+        # to ChatVertex instances
+        self._map_chat_vertices()
+
+    def _map_chat_vertices(self) -> None:
+        """Maps the vertices that are connected to ChatVertex instances."""
+        # For each edge, we need to check if the source or target vertex is a ChatVertex
+        # If it is, we need to update the other vertex `is_external` attribute
+        # and store the id of the ChatVertex in the attributes self.inputs and self.outputs
+        for edge in self.edges:
+            source_vertex = self.get_vertex(edge.source_id)
+            target_vertex = self.get_vertex(edge.target_id)
+            if isinstance(source_vertex, ChatVertex):
+                # The source vertex is a ChatVertex
+                # thus the target vertex is an external vertex
+                # and the source vertex is an input
+                target_vertex.has_external_input = True
+                self.inputs.append(source_vertex.id)
+            if isinstance(target_vertex, ChatVertex):
+                # The target vertex is a ChatVertex
+                # thus the source vertex is an external vertex
+                # and the target vertex is an output
+                source_vertex.has_external_output = True
+                self.outputs.append(target_vertex.id)
 
     def _build_vertex_params(self) -> None:
         """Identifies and handles the LLM vertex within the graph."""
@@ -284,3 +310,5 @@ class Graph:
             current_layer += 1  # Next layer
 
         return layers
+        return layers
+        return layers