From 2b646429c07f60e84018035aed5db8a8839abbce Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Fri, 9 Feb 2024 12:00:24 -0300 Subject: [PATCH] Refactor langflow processing and langfuse callback initialization --- src/backend/langflow/processing/base.py | 8 +++++--- src/backend/langflow/services/plugins/langfuse_plugin.py | 3 +-- 2 files changed, 6 insertions(+), 5 deletions(-) diff --git a/src/backend/langflow/processing/base.py b/src/backend/langflow/processing/base.py index 06dc0a063..fad53c254 100644 --- a/src/backend/langflow/processing/base.py +++ b/src/backend/langflow/processing/base.py @@ -2,10 +2,13 @@ from typing import TYPE_CHECKING, List, Union from langchain.agents.agent import AgentExecutor from langchain.callbacks.base import BaseCallbackHandler +from loguru import logger + 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 +from langflow.processing.process import fix_memory_inputs, format_actions +from langflow.services.deps import get_plugins_service if TYPE_CHECKING: from langfuse.callback import CallbackHandler # type: ignore @@ -28,13 +31,12 @@ def setup_callbacks(sync, trace_id, **kwargs): def get_langfuse_callback(trace_id): from langflow.services.deps import get_plugins_service - from langfuse.callback import CreateTrace logger.debug("Initializing langfuse callback") if langfuse := get_plugins_service().get("langfuse"): logger.debug("Langfuse credentials found") try: - trace = langfuse.trace(CreateTrace(name="langflow-" + trace_id, id=trace_id)) + trace = langfuse.trace(name="langflow-" + trace_id, id=trace_id) return trace.getNewHandler() except Exception as exc: logger.error(f"Error initializing langfuse callback: {exc}") diff --git a/src/backend/langflow/services/plugins/langfuse_plugin.py b/src/backend/langflow/services/plugins/langfuse_plugin.py index 424279342..e5d3afdd6 100644 --- a/src/backend/langflow/services/plugins/langfuse_plugin.py +++ b/src/backend/langflow/services/plugins/langfuse_plugin.py @@ -64,14 +64,13 @@ class LangfusePlugin(CallbackPlugin): def get_callback(self, _id: Optional[str] = None): if _id is None: _id = "default" - from langfuse.callback import CreateTrace # type: ignore logger.debug("Initializing langfuse callback") try: langfuse_instance = self.get() if langfuse_instance is not None and hasattr(langfuse_instance, "trace"): - trace = langfuse_instance.trace(CreateTrace(name="langflow-" + _id, id=_id)) + trace = langfuse_instance.trace(name="langflow-" + _id, id=_id) if trace: return trace.getNewHandler()