[Fix] Fix lambda node failures (#67)

Fix lambda node failures

Author: psriramsnc

sygra/core/graph/backend_factory.py

@@ -10,13 +10,12 @@ class BackendFactory(ABC):
     """
 
     @abstractmethod
-    def create_lambda_runnable(self, function_to_execute, node_config):
+    def create_lambda_runnable(self, exec_wrapper):
         """
         Abstract method to create a Lambda runnable.
 
         Args:
-            function_to_execute: Python function to execute, if it is a class it should be callable(__call__)
-            node_config:node config dictionary
+            exec_wrapper: Async function to execute
 
         Returns:
             Any: backend specific runnable object like Runnable for backend=Langgraph

sygra/core/graph/langgraph/langgraph_factory.py

@@ -1,4 +1,3 @@
-from functools import partial
 from typing import Any
 
 from langchain_core.messages import BaseMessage
@@ -18,18 +17,17 @@ class LangGraphFactory(BackendFactory):
     A factory class to convert Nodes into Runnable objects which LangGraph framework can execute.
     """
 
-    def create_lambda_runnable(self, function_to_execute, node_config):
+    def create_lambda_runnable(self, exec_wrapper):
         """
         Abstract method to create a Lambda runnable.
 
         Args:
-            function_to_execute: Python function to execute, if it is a class it should be callable(__call__)
-            node_config:node config dictionary
+            exec_wrapper: Async function to execute
 
         Returns:
             Any: backend specific runnable object like Runnable for backend=Langgraph
         """
-        return RunnableLambda(partial(function_to_execute, node_config))
+        return RunnableLambda(lambda x: x, afunc=exec_wrapper)
 
     def create_llm_runnable(self, exec_wrapper):
         """

sygra/core/graph/nodes/lambda_node.py

@@ -38,7 +38,7 @@ async def _exec_wrapper(self, state: dict[str, Any]) -> dict[str, Any]:
         success = True
 
         try:
-            result: dict[str, Any] = self.func_to_execute(state)
+            result: dict[str, Any] = self.func_to_execute(self.node_config, state)
             return result
         except Exception:
             success = False
@@ -53,7 +53,7 @@ def to_backend(self) -> Any:
         Returns:
              Any: platform specific runnable object like Runnable in LangGraph.
         """
-        return utils.backend_factory.create_lambda_runnable(self._exec_wrapper, self.node_config)
+        return utils.backend_factory.create_lambda_runnable(self._exec_wrapper)
 
     def validate_node(self):
         """

sygra/core/models/lite_llm/vllm_model.py

@@ -22,7 +22,7 @@ def __init__(self, model_config: dict[str, Any]) -> None:
         utils.validate_required_keys(["url", "auth_token"], model_config, "model")
         self.model_config = model_config
         self.auth_token = str(model_config.get("auth_token")).replace("Bearer ", "")
-        self.model_serving_name = model_config.get("model_serving_name", self.name())
+        self.model_name = model_config.get("model_serving_name", self.name())
 
     def _validate_completions_api_model_support(self) -> None:
         logger.info(f"Model {self.name()} supports completion API.")

tests/core/models/lite_llm/test_litellm_vllm.py

@@ -56,11 +56,11 @@ def test_init(self):
         self.assertEqual(model.generation_params, self.base_config["parameters"])
         self.assertEqual(model.name(), "vllm_model")
         self.assertEqual(model.auth_token, "test_token_123")
-        self.assertEqual(model.model_serving_name, "vllm_model")
+        self.assertEqual(model.model_name, "vllm_model")
 
     def test_init_with_custom_serving_name(self):
         model = LiteLLMVLLM(self.serving_name_config)
-        self.assertEqual(model.model_serving_name, "custom_serving_name")
+        self.assertEqual(model.model_name, "custom_serving_name")
 
     @patch("sygra.core.models.lite_llm.vllm_model.logger")
     def test_init_with_completions_api(self, mock_logger):