cronos3k-comfyui-agentazall/tests/test_integration_nodes.py

"""Tests for AgentAZAll ComfyUI integration nodes (SystemPrompt, ToolParser, FileBrowser)."""

import pytest

from agentazall.helpers import agent_day, ensure_dirs, today_str
from agentazall.config import INBOX, REMEMBER
from agentazall.messages import format_message

from nodes.config_nodes import AgentAZAll_WhoAmI
from nodes.memory_nodes import AgentAZAll_Remember
from nodes.integration_nodes import (
    AgentAZAll_SystemPrompt,
    AgentAZAll_ToolParser,
    AgentAZAll_FileBrowser,
)


class TestSystemPrompt:

    def test_minimal(self, agent_config):
        """SystemPrompt returns prompt with tool instructions."""
        node = AgentAZAll_SystemPrompt()
        (result,) = node.build_prompt(agent_config)
        assert "Available Tools" in result
        assert "[TOOL:remember" in result

    def test_with_base_prompt(self, agent_config):
        """SystemPrompt includes base prompt at the top."""
        node = AgentAZAll_SystemPrompt()
        (result,) = node.build_prompt(
            agent_config, base_prompt="You are an art director."
        )
        assert result.startswith("You are an art director.")

    def test_with_memories(self, agent_config):
        """SystemPrompt includes stored memories."""
        remember = AgentAZAll_Remember()
        remember.remember(agent_config, "The dragon has red scales.", "dragon")

        node = AgentAZAll_SystemPrompt()
        (result,) = node.build_prompt(agent_config, include_memories=True)
        assert "red scales" in result
        assert "Your Memories" in result

    def test_with_inbox(self, agent_config):
        """SystemPrompt includes inbox summary."""
        d = today_str()
        ensure_dirs(agent_config, d)
        inbox_dir = agent_day(agent_config, d) / INBOX
        content, msg_id = format_message(
            "sender@localhost", agent_config["agent_name"],
            "Review this", "Please check.",
        )
        (inbox_dir / f"{msg_id}.txt").write_text(content, encoding="utf-8")

        node = AgentAZAll_SystemPrompt()
        (result,) = node.build_prompt(agent_config, include_inbox=True)
        assert "Inbox" in result
        assert "Review this" in result

    def test_with_identity(self, agent_config):
        """SystemPrompt includes agent identity."""
        whoami = AgentAZAll_WhoAmI()
        whoami.whoami(agent_config, "I am the story editor.")

        node = AgentAZAll_SystemPrompt()
        (result,) = node.build_prompt(agent_config)
        assert "story editor" in result
        assert "Your Identity" in result


class TestToolParser:

    def test_no_tools(self, agent_config):
        """ToolParser with no tool calls returns clean text."""
        node = AgentAZAll_ToolParser()
        (results, clean) = node.parse_tools(
            agent_config, "Just a normal response with no tools."
        )
        assert "No tool calls" in results
        assert clean == "Just a normal response with no tools."

    def test_remember_tool(self, agent_config):
        """ToolParser executes remember tool and stores memory."""
        node = AgentAZAll_ToolParser()
        llm_output = (
            'I think this is important. '
            '[TOOL:remember title="key-insight"]The castle has three towers.[/TOOL] '
            'Let me continue.'
        )
        (results, clean) = node.parse_tools(agent_config, llm_output)

        assert "[remember:key-insight] Stored" in results
        assert "three towers" not in clean
        assert "I think this is important" in clean
        assert "Let me continue" in clean

        # Verify memory was actually stored
        from nodes.memory_nodes import AgentAZAll_Recall
        recall = AgentAZAll_Recall()
        (recall_result,) = recall.recall(agent_config, "castle")
        assert "three towers" in recall_result

    def test_note_tool(self, agent_config):
        """ToolParser executes note tool."""
        node = AgentAZAll_ToolParser()
        llm_output = '[TOOL:note name="outline"]Chapter 1: Dawn\nChapter 2: Dusk[/TOOL]'
        (results, clean) = node.parse_tools(agent_config, llm_output)

        assert "[note:outline] Written" in results

        from nodes.memory_nodes import AgentAZAll_Note
        note = AgentAZAll_Note()
        (content,) = note.note(agent_config, "outline", "")
        assert "Chapter 1" in content

    def test_send_tool(self, agent_config):
        """ToolParser executes send tool and queues message."""
        node = AgentAZAll_ToolParser()
        llm_output = '[TOOL:send to="editor@localhost" subject="Draft ready"]Here is chapter 1.[/TOOL]'
        (results, clean) = node.parse_tools(agent_config, llm_output)

        assert "[send:editor@localhost]" in results
        assert "Draft ready" in results

        # Verify outbox has the message
        outbox = agent_day(agent_config, today_str()) / "outbox"
        files = list(outbox.glob("*.txt"))
        assert len(files) == 1


class TestFileBrowser:

    def test_empty_tree(self, agent_config):
        """FileBrowser shows directory structure."""
        node = AgentAZAll_FileBrowser()
        (result,) = node.browse(agent_config, False)
        assert "Mailbox:" in result

    def test_tree_with_content(self, agent_config):
        """FileBrowser shows files with annotations."""
        # Store a memory and a message
        remember = AgentAZAll_Remember()
        remember.remember(agent_config, "Test memory content.", "test-mem")

        d = today_str()
        inbox_dir = agent_day(agent_config, d) / INBOX
        content, msg_id = format_message(
            "alice@localhost", agent_config["agent_name"],
            "Hello", "Hi!",
        )
        (inbox_dir / f"{msg_id}.txt").write_text(content, encoding="utf-8")

        node = AgentAZAll_FileBrowser()
        (result,) = node.browse(agent_config, False)

        assert "remember" in result
        assert "inbox" in result
        assert "test-mem.txt" in result

    def test_tree_annotations(self, agent_config):
        """FileBrowser annotates message files with sender/subject."""
        d = today_str()
        ensure_dirs(agent_config, d)
        inbox_dir = agent_day(agent_config, d) / INBOX
        content, msg_id = format_message(
            "bob@localhost", agent_config["agent_name"],
            "Review Request", "Please review.",
        )
        (inbox_dir / f"{msg_id}.txt").write_text(content, encoding="utf-8")

        node = AgentAZAll_FileBrowser()
        (result,) = node.browse(agent_config, False)

        assert "bob@localhost" in result
        assert "Review Request" in result