diff --git a/src/backend/langflow/interface/agents/custom.py b/src/backend/langflow/interface/agents/custom.py index 5fab9bf77..d85ba8a56 100644 --- a/src/backend/langflow/interface/agents/custom.py +++ b/src/backend/langflow/interface/agents/custom.py @@ -1,3 +1,4 @@ +from abc import ABC from typing import Any, List, Optional from langchain import LLMChain @@ -33,7 +34,25 @@ from langchain.tools.python.tool import PythonAstREPLTool from langchain.tools.sql_database.prompt import QUERY_CHECKER -class JsonAgent(AgentExecutor): +class CustomAgentExecutor(AgentExecutor, ABC): + """Custom agent executor""" + + @staticmethod + def function_name(): + return "CustomAgentExecutor" + + @classmethod + def initialize(cls, *args, **kwargs): + pass + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + def run(self, *args, **kwargs): + return super().run(*args, **kwargs) + + +class JsonAgent(CustomAgentExecutor): """Json agent""" @staticmethod @@ -69,7 +88,7 @@ class JsonAgent(AgentExecutor): return super().run(*args, **kwargs) -class CSVAgent(AgentExecutor): +class CSVAgent(CustomAgentExecutor): """CSV agent""" @staticmethod @@ -117,7 +136,7 @@ class CSVAgent(AgentExecutor): return super().run(*args, **kwargs) -class VectorStoreAgent(AgentExecutor): +class VectorStoreAgent(CustomAgentExecutor): """Vector Store agent""" @staticmethod @@ -155,7 +174,7 @@ class VectorStoreAgent(AgentExecutor): return super().run(*args, **kwargs) -class SQLAgent(AgentExecutor): +class SQLAgent(CustomAgentExecutor): """SQL agent""" @staticmethod @@ -227,7 +246,7 @@ class SQLAgent(AgentExecutor): return super().run(*args, **kwargs) -class VectorStoreRouterAgent(AgentExecutor): +class VectorStoreRouterAgent(CustomAgentExecutor): """Vector Store Router Agent""" @staticmethod @@ -266,7 +285,7 @@ class VectorStoreRouterAgent(AgentExecutor): return super().run(*args, **kwargs) -class InitializeAgent(AgentExecutor): +class InitializeAgent(CustomAgentExecutor): """Implementation of initialize_agent function""" @staticmethod diff --git a/src/backend/langflow/interface/run.py b/src/backend/langflow/interface/run.py index 05e4b49d6..af80b7acd 100644 --- a/src/backend/langflow/interface/run.py +++ b/src/backend/langflow/interface/run.py @@ -2,12 +2,11 @@ import contextlib import io from typing import Any, Dict -from chromadb.errors import NotEnoughElementsException +from chromadb.errors import NotEnoughElementsException # type: ignore from langflow.api.callback import AsyncStreamingLLMCallbackHandler, StreamingLLMCallbackHandler # type: ignore from langflow.cache.base import compute_dict_hash, load_cache, memoize_dict from langflow.graph.graph import Graph -from langflow.interface import loading from langflow.utils.logger import logger