From c399658afc7abb77aac169b8a6926fbf2ddec5df Mon Sep 17 00:00:00 2001 From: David Strand Date: Thu, 16 Jan 2025 05:54:07 -0600 Subject: [PATCH] Initial commit for ComfyUI-Autogen nodes --- __init__.py | 24 ++++++++ autogen_assistant_agent.py | 114 +++++++++++++++++++++++++++++++++++++ autogen_code_executor.py | 40 +++++++++++++ autogen_groupchat.py | 85 +++++++++++++++++++++++++++ autogen_model.py | 52 +++++++++++++++++ requirements.txt | 4 ++ tools/duck_duck_go.py | 33 +++++++++++ 7 files changed, 352 insertions(+) create mode 100644 __init__.py create mode 100644 autogen_assistant_agent.py create mode 100644 autogen_code_executor.py create mode 100644 autogen_groupchat.py create mode 100644 autogen_model.py create mode 100644 requirements.txt create mode 100644 tools/duck_duck_go.py diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..704cec8 --- /dev/null +++ b/__init__.py @@ -0,0 +1,24 @@ +from .autogen_model import AutogenModel +from .autogen_assistant_agent import AutogenAssistantAgent +from .autogen_groupchat import AutogenGroupChat +from .autogen_code_executor import AutogenCodeExecutor + +NODE_CLASS_MAPPINGS = { + "AutogenModel": AutogenModel, + "AutogenAssistantAgent": AutogenAssistantAgent, + "AutogenGroupChat": AutogenGroupChat, + "AutogenCodeExecutor": AutogenCodeExecutor, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "AutogenModel": "Autogen Model", + "AutogenAssistantAgent": "Autogen Assistant Agent", + "AutogenGroupChat": "Autogen Group Chat", + "AutogenCodeExecutor": "Autogen Code Executor", +} + +# Optional: Specify a web directory for static assets (if applicable) +WEB_DIRECTORY = "./web" + +# Export all variables for ComfyUI to recognize +__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS', 'WEB_DIRECTORY'] diff --git a/autogen_assistant_agent.py b/autogen_assistant_agent.py new file mode 100644 index 0000000..4a2022d --- /dev/null +++ b/autogen_assistant_agent.py @@ -0,0 +1,114 @@ +import os +import asyncio +from autogen_agentchat.agents import AssistantAgent + +class AutogenAssistantAgent: + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "autogen_model": ("MODEL", {"default": None}), + "agent_name": ("STRING", {"default": "assistant_agent", "multiline": False}), + }, + "optional": { + "system_message": ("STRING", {"default": "You are a helpful assistant.", "multiline": True}), + "tool_1": ("STRING", {"default": "None", "options": cls._get_available_tools()}), + "tool_2": ("STRING", {"default": "None", "options": cls._get_available_tools()}), + "temperature": ("FLOAT", {"default": 0.1, "min": 0.0, "max": 1.0, "step": 0.1}), + "max_tokens": ("INT", {"default": 1200, "min": 200, "max": 16200, "step": 200}), + }, + } + + RETURN_TYPES = ("AGENT",) + RETURN_NAMES = ("assistant_agent",) + FUNCTION = "initialize_assistant_agent" + CATEGORY = "Autogen" + + # Class-level cache for tools + _tools_cache = None + _tools_last_checked = None + _tools_dir_path = tools_dir = os.path.join(os.path.dirname(__file__), "tools") + + @classmethod + def _get_available_tools(cls): + """Scan the tools directory and list available tools only if it has changed.""" + tools_dir = cls._tools_dir_path + + if not os.path.exists(tools_dir): + print(f"[AutogenAssistantAgent] Tools directory not found: {tools_dir}") + return ["None"] + + # Get the last modification time of the directory + last_modified = os.stat(tools_dir).st_mtime + + # Only rescan if directory has been modified or cache is empty + if cls._tools_cache is None or cls._tools_last_checked != last_modified: + print(f"[AutogenAssistantAgent] Scanning tools directory: {tools_dir}") + cls._tools_cache = ["None"] + [ + f.split(".")[0] for f in os.listdir(tools_dir) if f.endswith(".py") + ] + cls._tools_last_checked = last_modified + print(f"[AutogenAssistantAgent] Available tools: {cls._tools_cache}") + + return cls._tools_cache + + @staticmethod + def _load_tool(tool_name): + """Dynamically load a tool from the tools directory.""" + tools_dir = os.path.join(os.path.dirname(__file__), "tools") + tool_path = os.path.join(tools_dir, f"{tool_name}.py") + + if not os.path.exists(tool_path): + print(f"[AutogenAssistantAgent] Tool {tool_name} not found in: {os.path.abspath(tools_dir)}") + return None + + import importlib.util + spec = importlib.util.spec_from_file_location(tool_name, tool_path) + tool_module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(tool_module) + + if hasattr(tool_module, "tool_name") and hasattr(tool_module, "execute"): + return tool_module + else: + print(f"[AutogenAssistantAgent] Invalid tool interface in {tool_name}.") + return None + + def initialize_assistant_agent( + self, + autogen_model, + agent_name, + system_message=None, + tool_1=None, + tool_2=None, + temperature=0.1, + max_tokens=1200, + ): + if autogen_model is None: + raise ValueError("A model client is required for Autogen AssistantAgent.") + + # Load tools + tools = [] + for tool_name in [tool_1, tool_2]: + if tool_name != "None": + tool = self._load_tool(tool_name) + if tool: + tools.append(tool) + + # Initialize AssistantAgent + assistant_agent = AssistantAgent( + name=agent_name, + system_message=system_message, + model_client=autogen_model, + tools=[tool.execute for tool in tools] # Directly reference tool.execute + ) + + # Apply optional settings + assistant_agent.temperature = temperature + assistant_agent.max_tokens = max_tokens + + print(f"[AutogenAssistantAgent] Initialized assistant agent: {assistant_agent}") + return (assistant_agent,) + diff --git a/autogen_code_executor.py b/autogen_code_executor.py new file mode 100644 index 0000000..4691bd1 --- /dev/null +++ b/autogen_code_executor.py @@ -0,0 +1,40 @@ +from autogen_agentchat.agents import CodeExecutorAgent +from autogen_ext.code_executors.local import LocalCommandLineCodeExecutor +import os + +class AutogenCodeExecutor: + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "agent_name": ("STRING", {"default": "code_executor", "multiline": False}), + }, + "optional": { + "working_directory": ("STRING", {"default": os.getcwd(), "multiline": False}), + "timeout": ("INT", {"default": 10, "min": 1, "max": 60, "step": 1}), + }, + } + + RETURN_TYPES = ("AGENT",) + RETURN_NAMES = ("code_executor_agent",) + FUNCTION = "initialize_code_executor" + CATEGORY = "Autogen" + + def initialize_code_executor(self, agent_name, working_directory=None, allow_network=False, timeout=10): + # Initialize the LocalCommandLineCodeExecutor + local_executor = LocalCommandLineCodeExecutor( + work_dir=working_directory or os.getcwd(), + timeout=timeout, + ) + + # Create the CodeExecutorAgent + code_executor_agent = CodeExecutorAgent( + name=agent_name, + code_executor=local_executor, + ) + + print(f"[AutogenCodeExecutor] Initialized code executor agent: {code_executor_agent}") + return (code_executor_agent,) diff --git a/autogen_groupchat.py b/autogen_groupchat.py new file mode 100644 index 0000000..ee6c3d4 --- /dev/null +++ b/autogen_groupchat.py @@ -0,0 +1,85 @@ +from autogen_agentchat.teams import RoundRobinGroupChat +from autogen_agentchat.conditions import MaxMessageTermination, TextMentionTermination +from PIL import Image +import asyncio +import sys +import io + +class AutogenGroupChat: + def __init__(self): + self.chat_log = [] # Store interactive chat log + self.last_message = None # Store the last message + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "agent_1": ("AGENT", {"default": None}), + "agent_2": ("AGENT", {"default": None}), + "task_input": ("STRING", {"default": "Provide a task description here.", "multiline": True}), + }, + "optional": { + "agent_3": ("AGENT", {"default": None}), + "agent_4": ("AGENT", {"default": None}), + "max_messages": ("INT", {"default": 10, "min": 1, "max": 100, "step": 1}), + "terminate_string": ("STRING", {"default": "TERMINATE", "multiline": False}), + }, + } + + RETURN_TYPES = ("STRING", "STRING") + RETURN_NAMES = ("console_output", "last_chat") + FUNCTION = "run_group_chat" + CATEGORY = "Autogen" + + def run_group_chat(self, agent_1, agent_2, task_input, agent_3=None, agent_4=None, + max_messages=10, terminate_string="TERMINATE"): + # Redirect console output + old_stdout = sys.stdout + sys.stdout = io.StringIO() + + try: + # Initialize participants + participants = [agent_1, agent_2] + if agent_3: + participants.append(agent_3) + if agent_4: + participants.append(agent_4) + + # Create termination conditions + max_msg_termination = MaxMessageTermination(max_messages=max_messages) + text_termination = TextMentionTermination(terminate_string) + combined_termination = max_msg_termination | text_termination + + # Initialize PlannerGroupChat + team = RoundRobinGroupChat(participants=participants, termination_condition=combined_termination) + + print(f"[AutogenGroupChat] Initialized with {len(participants)} participants.") + + chat_output = [] + + async def capture_chat(): + async for task_result in team.run_stream(task=task_input): + if hasattr(task_result, "messages"): + for message in task_result.messages: + self.chat_log.append(f"{message.source}: {message.content}") + self.last_message = f"{message.source}: {message.content}" + print(f"[Interactive Chat] {self.last_message}") + chat_output.append(f"{message.source}: {message.content}") + + # Clean and process the last message + if message.content.strip().endswith(terminate_string): + self.last_message = message.content.strip()[:-len(terminate_string)].strip() + else: + self.last_message = message.content.strip() + + asyncio.run(capture_chat()) + + except RuntimeError as e: + print(f"[AutogenGroupChat] Runtime error: {e}") + return "Error during execution", self.last_message, sys.stdout.getvalue() + + # Capture console output and reset stdout + console_output = sys.stdout.getvalue() + sys.stdout = old_stdout + + return console_output, self.last_message \ No newline at end of file diff --git a/autogen_model.py b/autogen_model.py new file mode 100644 index 0000000..549c7af --- /dev/null +++ b/autogen_model.py @@ -0,0 +1,52 @@ +from autogen_ext.models import OpenAIChatCompletionClient + +class AutogenModel: + @classmethod + def INPUT_TYPES(cls): + """ + Define input parameters for configuring the OpenAI model. + """ + return { + "required": { + "model_name": ([ + "gpt-4o", + "gpt-4o-mini", + "chatgpt-4o-latest", + "gpt-3.5-turbo", + "deepseek-chat", + "deepseek-coder", + ], {"default": "gpt-4o-mini"}), + "base_url": ("STRING", {"multiline": False, "default": "https://api.openai.com/v1"}), + "api_key": ("STRING", {"multiline": False, "default": "openai_api_key"}), + "vision": (["True", "False"], {"default": "False"}), + "function_calling": (["True", "False"], {"default": "False"}), + "json_output": (["True", "False"], {"default": "False"}), + }, + } + + RETURN_TYPES = ("MODEL",) + RETURN_NAMES = ("autogen_model",) + FUNCTION = "initialize_model" + CATEGORY = "Autogen" + + def initialize_model(self, model_name, base_url, api_key, vision, function_calling, json_output): + """ + Initialize the OpenAI model with the provided configuration. + """ + try: + # Initialize the OpenAI client + autogen_model = OpenAIChatCompletionClient( + model=model_name, + base_url=base_url, + api_key=api_key, + model_capabilities={ + "vision": vision == "True", + "function_calling": function_calling == "True", + "json_output": json_output == "True", + }, + ) + print(f"[AutogenModel] Successfully initialized model: {model_name}") + return (autogen_model,) + except Exception as e: + print(f"[AutogenModel] Failed to initialize model: {e}") + return (None,) diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..e522320 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,4 @@ +autogen-core +autogen-ext +autogen-agentchat +requests diff --git a/tools/duck_duck_go.py b/tools/duck_duck_go.py new file mode 100644 index 0000000..3b9fdbf --- /dev/null +++ b/tools/duck_duck_go.py @@ -0,0 +1,33 @@ +from duckduckgo_search import DDGS + +# Define the tool name +tool_name = "duck_duck_go" + +# Define the execute function +def execute(query: str, region: str = 'en-us', safesearch: str = 'off', max_results: int = 5): + """ + Performs a DuckDuckGo search using the duckduckgo-search library. + Args: + query (str): The search query. + result_limit (int): Maximum number of results to return (default is 5). + Returns: + dict: A dictionary containing search results (title, URL, and snippet). + """ + try: + ddg = DDGS() + # Perform the search using duckduckgo-search + results = ddg.text(keywords=query, region=region, safesearch=safesearch, max_results=max_results) + + if not results: + return {"results": [], "message": "No results found."} + + # Format the results + formatted_results = [ + {"title": result.get("title"), "url": result.get("href"), "snippet": result.get("body")} + for result in results + ] + + return {"results": formatted_results} + + except Exception as e: + return {"error": f"Error occurred while searching: {e}"}