From 6ca57809086e80f4e2966aeb97fe937d2bd806d1 Mon Sep 17 00:00:00 2001
From: Gabriel Luiz Freitas Almeida <gabriel@langflow.org>
Date: Thu, 18 Jul 2024 16:22:41 -0300
Subject: [PATCH] refactor(callback): simplify import statements and type
 annotations for better readability (#2793)

* refactor: add LogFunctionType protocol for type hint

The LogFunctionType protocol is introduced to define the structure of a logging function that takes a message and an optional name parameter. This protocol will be used to ensure consistency and compatibility with other parts of the codebase when logging messages.

* refactor(callback.py): simplify import statements and type annotations for better readability
---
 src/backend/base/langflow/base/agents/callback.py | 6 +++---
 src/backend/base/langflow/schema/log.py           | 7 ++++++-
 2 files changed, 9 insertions(+), 4 deletions(-)

diff --git a/src/backend/base/langflow/base/agents/callback.py b/src/backend/base/langflow/base/agents/callback.py
index 7f75e7745..5d7260b6c 100644
--- a/src/backend/base/langflow/base/agents/callback.py
+++ b/src/backend/base/langflow/base/agents/callback.py
@@ -1,16 +1,16 @@
-from typing import Any, Callable, Concatenate, Dict, List
+from typing import Any, Dict, List
 from uuid import UUID
 
 from langchain.callbacks.base import AsyncCallbackHandler
 from langchain_core.agents import AgentAction, AgentFinish
 
-from langflow.schema.log import LoggableType
+from langflow.schema.log import LogFunctionType
 
 
 class AgentAsyncHandler(AsyncCallbackHandler):
     """Async callback handler that can be used to handle callbacks from langchain."""
 
-    def __init__(self, log_function: Callable[Concatenate[LoggableType | list[LoggableType], ...], None] | None = None):
+    def __init__(self, log_function: LogFunctionType | None = None):
         self.log_function = log_function
 
     async def on_tool_start(
diff --git a/src/backend/base/langflow/schema/log.py b/src/backend/base/langflow/schema/log.py
index 771ff0670..4cc17544e 100644
--- a/src/backend/base/langflow/schema/log.py
+++ b/src/backend/base/langflow/schema/log.py
@@ -1,5 +1,10 @@
-from typing import Union
+from typing import Optional, Union
 
 from pydantic import BaseModel
+from typing_extensions import Protocol
 
 LoggableType = Union[str, dict, list, int, float, bool, None, BaseModel]
+
+
+class LogFunctionType(Protocol):
+    def __call__(self, message: Union[LoggableType, list[LoggableType]], *, name: Optional[str] = None) -> None: ...