feat: Enhance backend with context management, error handling, and refactored code

src/backend/base/langflow/base/agents/context.py

@@ -0,0 +1,109 @@
+from datetime import datetime, timezone
+from typing import Any
+
+from langchain_core.language_models import BaseLanguageModel, BaseLLM
+from langchain_core.language_models.chat_models import BaseChatModel
+from pydantic import BaseModel, Field, field_validator, model_serializer
+
+from langflow.field_typing import LanguageModel
+from langflow.schema.data import Data
+
+
+class AgentContext(BaseModel):
+    tools: dict[str, Any]
+    llm: Any
+    context: str = ""
+    iteration: int = 0
+    max_iterations: int = 5
+    thought: str = ""
+    last_action: Any = None
+    last_action_result: Any = None
+    final_answer: Any = ""
+    context_history: list[tuple[str, str, str]] = Field(default_factory=list)
+
+    @model_serializer(mode="plain")
+    def serialize_agent_context(self):
+        serliazed_llm = self.llm.to_json() if hasattr(self.llm, "to_json") else str(self.llm)
+        serliazed_tools = {k: v.to_json() if hasattr(v, "to_json") else str(v) for k, v in self.tools.items()}
+        return {
+            "tools": serliazed_tools,
+            "llm": serliazed_llm,
+            "context": self.context,
+            "iteration": self.iteration,
+            "max_iterations": self.max_iterations,
+            "thought": self.thought,
+            "last_action": self.last_action.to_json()
+            if hasattr(self.last_action, "to_json")
+            else str(self.last_action),
+            "action_result": self.last_action_result.to_json()
+            if hasattr(self.last_action_result, "to_json")
+            else str(self.last_action_result),
+            "final_answer": self.final_answer,
+            "context_history": self.context_history,
+        }
+
+    @field_validator("llm", mode="before")
+    @classmethod
+    def validate_llm(cls, v) -> LanguageModel:
+        if not isinstance(v, BaseLLM | BaseChatModel | BaseLanguageModel):
+            msg = "llm must be an instance of LanguageModel"
+            raise TypeError(msg)
+        return v
+
+    def to_data_repr(self):
+        data_objs = []
+        for name, val, time_str in self.context_history:
+            content = val.content if hasattr(val, "content") else val
+            data_objs.append(Data(name=name, value=content, timestamp=time_str))
+
+        sorted_data_objs = sorted(data_objs, key=lambda x: datetime.fromisoformat(x.timestamp), reverse=True)
+
+        sorted_data_objs.append(
+            Data(
+                name="Formatted Context",
+                value=self.get_full_context(),
+            )
+        )
+        return sorted_data_objs
+
+    def _build_tools_context(self):
+        tool_context = ""
+        for tool_name, tool_obj in self.tools.items():
+            tool_context += f"{tool_name}: {tool_obj.description}\n"
+        return tool_context
+
+    def _build_init_context(self):
+        return f"""
+{self.context}
+
+"""
+
+    def model_post_init(self, _context: Any) -> None:
+        if hasattr(self.llm, "bind_tools"):
+            self.llm = self.llm.bind_tools(self.tools.values())
+        if self.context:
+            self.update_context("Initial Context", self.context)
+
+    def update_context(self, key: str, value: str):
+        self.context_history.insert(0, (key, value, datetime.now(tz=timezone.utc).astimezone().isoformat()))
+
+    def _serialize_context_history_tuple(self, context_history_tuple: tuple[str, str, str]) -> str:
+        name, value, _ = context_history_tuple
+        if hasattr(value, "content"):
+            value = value.content
+        elif hasattr(value, "log"):
+            value = value.log
+        return f"{name}: {value}"
+
+    def get_full_context(self) -> str:
+        context_history_reversed = self.context_history[::-1]
+        context_formatted = "\n".join(
+            [
+                self._serialize_context_history_tuple(context_history_tuple)
+                for context_history_tuple in context_history_reversed
+            ]
+        )
+        return f"""
+Context:
+{context_formatted}
+"""

src/backend/base/langflow/components/embeddings/util/aiml.py → src/backend/base/langflow/base/embeddings/aiml_embeddings.py

@@ -30,14 +30,15 @@ def embed_documents(self, texts: list[str]) -> list[list[float]]:
                 try:
                     result_data = future.result()
                     if len(result_data["data"]) != 1:
-                        msg = "Expected one embedding"
+                        msg = f"Expected one embedding, got {len(result_data['data'])}"
                         raise ValueError(msg)
                     embeddings[index] = result_data["data"][0]["embedding"]
                 except (
                     httpx.HTTPStatusError,
                     httpx.RequestError,
                     json.JSONDecodeError,
                     KeyError,
+                    ValueError,
                 ):
                     logger.exception("Error occurred")
                     raise

src/backend/base/langflow/components/embeddings/aiml.py

@@ -1,6 +1,6 @@
+from langflow.base.embeddings.aiml_embeddings import AIMLEmbeddingsImpl
 from langflow.base.embeddings.model import LCEmbeddingsModel
 from langflow.base.models.aiml_constants import AIML_EMBEDDING_MODELS
-from langflow.components.embeddings.util import AIMLEmbeddingsImpl
 from langflow.field_typing import Embeddings
 from langflow.inputs.inputs import DropdownInput
 from langflow.io import SecretStrInput

src/backend/base/langflow/components/prototypes/conditional_router.py

@@ -1,5 +1,5 @@
 from langflow.custom import Component
-from langflow.io import BoolInput, DropdownInput, MessageInput, MessageTextInput, Output
+from langflow.io import BoolInput, DropdownInput, IntInput, MessageInput, MessageTextInput, Output
 from langflow.schema.message import Message
 
 
@@ -9,6 +9,10 @@ class ConditionalRouterComponent(Component):
     icon = "equal"
     name = "ConditionalRouter"
 
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.__iteration_updated = False
+
     inputs = [
         MessageTextInput(
             name="input_text",
@@ -40,13 +44,30 @@ class ConditionalRouterComponent(Component):
             display_name="Message",
             info="The message to pass through either route.",
         ),
+        IntInput(
+            name="max_iterations",
+            display_name="Max Iterations",
+            info="The maximum number of iterations for the conditional router.",
+            value=10,
+        ),
+        DropdownInput(
+            name="default_route",
+            display_name="Default Route",
+            options=["true_result", "false_result"],
+            info="The default route to take when max iterations are reached.",
+            value="false_result",
+            advanced=True,
+        ),
     ]
 
     outputs = [
         Output(display_name="True Route", name="true_result", method="true_response"),
         Output(display_name="False Route", name="false_result", method="false_response"),
     ]
 
+    def _pre_run_setup(self):
+        self.__iteration_updated = False
+
     def evaluate_condition(self, input_text: str, match_text: str, operator: str, *, case_sensitive: bool) -> bool:
         if not case_sensitive:
             input_text = input_text.lower()
@@ -64,22 +85,33 @@ def evaluate_condition(self, input_text: str, match_text: str, operator: str, *,
             return input_text.endswith(match_text)
         return False
 
+    def iterate_and_stop_once(self, route_to_stop: str):
+        if not self.__iteration_updated:
+            self.update_ctx({f"{self._id}_iteration": self.ctx.get(f"{self._id}_iteration", 0) + 1})
+            self.__iteration_updated = True
+            if self.ctx.get(f"{self._id}_iteration", 0) >= self.max_iterations and route_to_stop == self.default_route:
+                # We need to stop the other route
+                route_to_stop = "true_result" if route_to_stop == "false_result" else "false_result"
+            self.stop(route_to_stop)
+
     def true_response(self) -> Message:
         result = self.evaluate_condition(
             self.input_text, self.match_text, self.operator, case_sensitive=self.case_sensitive
         )
         if result:
             self.status = self.message
+            self.iterate_and_stop_once("false_result")
             return self.message
-        self.stop("true_result")
-        return None  # type: ignore[return-value]
+        self.iterate_and_stop_once("true_result")
+        return self.message
 
     def false_response(self) -> Message:
         result = self.evaluate_condition(
             self.input_text, self.match_text, self.operator, case_sensitive=self.case_sensitive
         )
         if not result:
             self.status = self.message
+            self.iterate_and_stop_once("true_result")
             return self.message
-        self.stop("false_result")
-        return None  # type: ignore[return-value]
+        self.iterate_and_stop_once("false_result")
+        return self.message

src/backend/base/langflow/custom/custom_component/component.py

@@ -66,6 +66,7 @@ class Component(CustomComponent):
     _output_logs: dict[str, list[Log]] = {}
     _current_output: str = ""
     _metadata: dict = {}
+    _ctx: dict = {}
 
     def __init__(self, **kwargs) -> None:
         # if key starts with _ it is a config
@@ -113,6 +114,53 @@ def __init__(self, **kwargs) -> None:
         self.set_class_code()
         self._set_output_required_inputs()
 
+    @property
+    def ctx(self):
+        if not hasattr(self, "graph") or self.graph is None:
+            msg = "Graph not found. Please build the graph first."
+            raise ValueError(msg)
+        return self.graph.context
+
+    def add_to_ctx(self, key: str, value: Any, *, overwrite: bool = False) -> None:
+        """Add a key-value pair to the context.
+
+        Args:
+            key (str): The key to add.
+            value (Any): The value to associate with the key.
+            overwrite (bool, optional): Whether to overwrite the existing value. Defaults to False.
+
+        Raises:
+            ValueError: If the graph is not built.
+        """
+        if not hasattr(self, "graph") or self.graph is None:
+            msg = "Graph not found. Please build the graph first."
+            raise ValueError(msg)
+        if key in self.graph.context and not overwrite:
+            msg = f"Key {key} already exists in context. Set overwrite=True to overwrite."
+            raise ValueError(msg)
+        self.graph.context.update({key: value})
+
+    def update_ctx(self, value_dict: dict[str, Any]) -> None:
+        """Update the context with a dictionary of values.
+
+        Args:
+            value_dict (dict[str, Any]): The dictionary of values to update.
+
+        Raises:
+            ValueError: If the graph is not built.
+        """
+        if not hasattr(self, "graph") or self.graph is None:
+            msg = "Graph not found. Please build the graph first."
+            raise ValueError(msg)
+        if not isinstance(value_dict, dict):
+            msg = "Value dict must be a dictionary"
+            raise TypeError(msg)
+
+        self.graph.context.update(value_dict)
+
+    def _pre_run_setup(self):
+        pass
+
     def set_event_manager(self, event_manager: EventManager | None = None) -> None:
         self._event_manager = event_manager
 
@@ -768,6 +816,8 @@ async def build_results(self):
     async def _build_results(self) -> tuple[dict, dict]:
         _results = {}
         _artifacts = {}
+        if hasattr(self, "_pre_run_setup"):
+            self._pre_run_setup()
         if hasattr(self, "outputs"):
             for output in self._outputs_map.values():
                 # Build the output if it's connected to some other vertex