From 667e7c42ff1d48036766ccd63061ddcf0a59af60 Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Wed, 6 Dec 2023 11:05:11 -0300 Subject: [PATCH] Refactor callback handling in base.py --- src/backend/langflow/processing/base.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/src/backend/langflow/processing/base.py b/src/backend/langflow/processing/base.py index c04bca3d6..4bcbb3e7e 100644 --- a/src/backend/langflow/processing/base.py +++ b/src/backend/langflow/processing/base.py @@ -2,8 +2,10 @@ from typing import TYPE_CHECKING, List, Union from langchain.agents.agent import AgentExecutor from langchain.callbacks.base import BaseCallbackHandler -from langflow.api.v1.callback import AsyncStreamingLLMCallbackHandler, StreamingLLMCallbackHandler +from langflow.api.v1.callback import (AsyncStreamingLLMCallbackHandler, + StreamingLLMCallbackHandler) from langflow.processing.process import fix_memory_inputs, format_actions +from langflow.services.deps import get_plugins_service from loguru import logger if TYPE_CHECKING: @@ -18,9 +20,10 @@ def setup_callbacks(sync, trace_id, **kwargs): else: callbacks.append(AsyncStreamingLLMCallbackHandler(**kwargs)) - if langfuse_callback := get_langfuse_callback(trace_id=trace_id): - logger.debug("Langfuse callback loaded") - callbacks.append(langfuse_callback) + plugin_service = get_plugins_service() + plugin_callbacks = plugin_service.get_callbacks(_id=trace_id) + if plugin_callbacks: + callbacks.extend(plugin_callbacks) return callbacks