tests/meta_prompt_graph_test.py

import unittest
from unittest.mock import MagicMock, Mock
import functools
import pprint
from langchain_core.language_models import BaseLanguageModel
from langchain_openai import ChatOpenAI
from meta_prompt import *
from meta_prompt.consts import NODE_ACCEPTANCE_CRITERIA_DEVELOPER
from langgraph.graph import END
import os

class TestMetaPromptGraph(unittest.TestCase):
    def setUp(self):
        # logging.basicConfig(level=logging.DEBUG)
        pass


    def test_prompt_node(self):
        """
        Test the _prompt_node method of MetaPromptGraph.

        This test case sets up a mock language model that returns a response content
        and verifies that the updated state has the output attribute updated with
        the mocked response content.
        """
        llm = Mock(spec=BaseLanguageModel)
        llm.config_specs = []
        llm.invoke = lambda x, y=None: "Mocked response content"

        llms = {
            NODE_PROMPT_INITIAL_DEVELOPER: llm
        }

        graph = MetaPromptGraph(llms=llms)
        state = AgentState(
            user_message="Test message", expected_output="Expected output"
        )
        updated_state = graph._prompt_node(
            NODE_PROMPT_INITIAL_DEVELOPER, "output", state
        )

        assert (
            updated_state['output'] == "Mocked response content"
        ), "The output attribute should be updated with the mocked response content"


    def test_output_history_analyzer(self):
        """
        Test the _output_history_analyzer method of MetaPromptGraph.

        This test case sets up a mock language model that returns an analysis
        response and verifies that the updated state has the best output, best
        system message, and best output age updated correctly.
        """
        llm = Mock(spec=BaseLanguageModel)
        llm.config_specs = []
        llm.invoke = lambda x, y: "{\"closerOutputID\": 2, \"analysis\": \"The output should use the `reverse()` method.\"}"
        prompts = {}
        meta_prompt_graph = MetaPromptGraph(llms=llm, prompts=prompts)
        state = AgentState(
            user_message="How do I reverse a list in Python?",
            expected_output="Use the `[::-1]` slicing technique or the `list.reverse()` method.",
            output="To reverse a list in Python, you can use the `[::-1]` slicing.",
            system_message="To reverse a list, use slicing or the reverse method.",
            best_output="To reverse a list in Python, use the `reverse()` method.",
            best_system_message="To reverse a list, use the `reverse()` method.",
            acceptance_criteria="The output should correctly describe how to reverse a list in Python.",
        )

        updated_state = meta_prompt_graph._output_history_analyzer(state)

        assert (
            updated_state['best_output'] == state['output']
        ), "Best output should be updated to the current output."
        assert (
            updated_state['best_system_message'] == state['system_message']
        ), "Best system message should be updated to the current system message."
        assert (
            updated_state['best_output_age'] == 0
        ), "Best output age should be reset to 0."


    def test_prompt_analyzer_accept(self):
        """
        Test the _prompt_analyzer method of MetaPromptGraph when the prompt analyzer
        accepts the output.

        This test case sets up a mock language model that returns an acceptance
        response and verifies that the updated state has the accepted attribute
        set to True.
        """
        # llms = {
        #     NODE_PROMPT_ANALYZER: lambda prompt: "{\"Accept\": \"Yes\"}"
        # }
        llm = Mock(spec=BaseLanguageModel)
        llm.config_specs = []
        llm.invoke = lambda x, y: "{\"Accept\": \"Yes\"}"
        meta_prompt_graph = MetaPromptGraph(llms=llm)
        state = AgentState(
            output="Test output", expected_output="Expected output",
            acceptance_criteria="Acceptance criteria: ...",
            system_message="System message: ...",
            max_output_age=2
        )
        updated_state = meta_prompt_graph._prompt_analyzer(state)
        assert updated_state['accepted'] is True


    def test_get_node_names(self):
        """
        Test the get_node_names method of MetaPromptGraph.

        This test case verifies that the get_node_names method returns the
        correct list of node names.
        """
        graph = MetaPromptGraph()
        node_names = graph.get_node_names()
        self.assertEqual(node_names, META_PROMPT_NODES)


    def test_workflow_execution(self):
        """
        Test the workflow execution of the MetaPromptGraph.

        This test case sets up a MetaPromptGraph with a single language model and
        executes it with a given input state. It then verifies that the output
        state contains the expected keys and values.
        """
        model_name = os.getenv("TEST_MODEL_NAME_EXECUTOR")
        raw_llm = ChatOpenAI(model_name=model_name)

        llms = {
            NODE_PROMPT_INITIAL_DEVELOPER: raw_llm,
            NODE_ACCEPTANCE_CRITERIA_DEVELOPER: raw_llm,
            NODE_PROMPT_DEVELOPER: raw_llm,
            NODE_PROMPT_EXECUTOR: raw_llm,
            NODE_OUTPUT_HISTORY_ANALYZER: raw_llm,
            NODE_PROMPT_ANALYZER: raw_llm,
            NODE_PROMPT_SUGGESTER: raw_llm,
        }

        meta_prompt_graph = MetaPromptGraph(llms=llms)
        input_state = AgentState(
            user_message="How do I reverse a list in Python?",
            expected_output="Use the `[::-1]` slicing technique or the "
                            "`list.reverse()` method.",
            acceptance_criteria="Similar in meaning, text length and style.",
            max_output_age=2
        )
        output_state = meta_prompt_graph(input_state, recursion_limit=25)

        pprint.pp(output_state)
        assert (
            "best_system_message" in output_state
        ), "The output state should contain the key 'best_system_message'"
        assert (
            output_state["best_system_message"] is not None
        ), "The best system message should not be None"
        if (
            "best_system_message" in output_state
            and output_state["best_system_message"] is not None
        ):
            print(output_state["best_system_message"])

        user_message = "How can I create a list of numbers in Python?"
        messages = [("system", output_state["best_system_message"]), ("human", user_message)]
        result = raw_llm.invoke(messages)

        assert hasattr(result, "content"), "The result should have the attribute 'content'"
        print(result.content)


    def test_workflow_execution_with_llms(self):
        """
        Test the workflow execution of the MetaPromptGraph with multiple LLMs.

        This test case sets up a MetaPromptGraph with multiple language models and
        executes it with a given input state. It then verifies that the output
        state contains the expected keys and values.
        """
        optimizer_llm = ChatOpenAI(
            model_name=os.getenv("TEST_MODEL_NAME_OPTIMIZER"), temperature=0.5
        )
        executor_llm = ChatOpenAI(
            model_name=os.getenv("TEST_MODEL_NAME_EXECUTOR"), temperature=0.01
        )

        llms = {
            NODE_PROMPT_INITIAL_DEVELOPER: optimizer_llm,
            NODE_ACCEPTANCE_CRITERIA_DEVELOPER: optimizer_llm,
            NODE_PROMPT_DEVELOPER: optimizer_llm,
            NODE_PROMPT_EXECUTOR: executor_llm,
            NODE_OUTPUT_HISTORY_ANALYZER: optimizer_llm,
            NODE_PROMPT_ANALYZER: optimizer_llm.bind(response_format={"type": "json_object"}),
            NODE_PROMPT_SUGGESTER: optimizer_llm,
        }

        meta_prompt_graph = MetaPromptGraph(llms=llms)
        input_state = AgentState(
            max_output_age=2,
            user_message="How do I reverse a list in Python?",
            expected_output="Use the `[::-1]` slicing technique or the "
                            "`list.reverse()` method.",
            # acceptance_criteria="Similar in meaning, text length and style."
        )
        output_state = meta_prompt_graph(input_state, recursion_limit=25)

        pprint.pp(output_state)
        assert (
            "best_system_message" in output_state
        ), "The output state should contain the key 'best_system_message'"
        assert (
            output_state["best_system_message"] is not None
        ), "The best system message should not be None"
        if (
            "best_system_message" in output_state
            and output_state["best_system_message"] is not None
        ):
            print(output_state["best_system_message"])

        user_message = "How can I create a list of numbers in Python?"
        messages = [("system", output_state["best_system_message"]), ("human", user_message)]
        result = executor_llm.invoke(messages)

        assert hasattr(result, "content"), "The result should have the attribute 'content'"
        print(result.content)
        

    def test_simple_workflow_execution(self):
        """
        Test the simple workflow execution of the MetaPromptGraph.

        This test case sets up a MetaPromptGraph with a mock LLM and executes it
        with a given input state. It then verifies that the output state contains
        the expected keys and values.
        """
        # Create a mock LLM that returns predefined responses based on the input messages
        llm = Mock(spec=BaseLanguageModel)
        llm.config_specs = []
        responses = [
            "Explain how to reverse a list in Python.",  # NODE_PROMPT_INITIAL_DEVELOPER
            "Here's one way: `my_list[::-1]`",  # NODE_PROMPT_EXECUTOR
            "{\"Accept\": \"Yes\"}",  # NODE_PPROMPT_ANALYZER
        ]
        # everytime llm.invoke was called, it returns a item in responses
        llm.invoke = lambda x, y=None: responses.pop(0)

        meta_prompt_graph = MetaPromptGraph(llms=llm)
        input_state = AgentState(
            user_message="How do I reverse a list in Python?",
            expected_output="The output should use the `reverse()` method.",
            acceptance_criteria="The output should be correct and efficient.",
            max_output_age=2
        )

        output_state = meta_prompt_graph(input_state)

        self.assertIsNotNone(output_state['best_system_message'])
        self.assertIsNotNone(output_state['best_output'])

        pprint.pp(output_state["best_output"])
        

    def test_iterated_workflow_execution(self):
        """
        Test the iterated workflow execution of the MetaPromptGraph.

        This test case sets up a MetaPromptGraph with a mock LLM and executes it
        with a given input state. It then verifies that the output state contains
        the expected keys and values. The test case simulates an iterated workflow
        where the LLM provides multiple responses based on the input messages.
        """
        # Create a mock LLM that returns predefined responses based on the input messages
        llm = Mock(spec=BaseLanguageModel)
        llm.config_specs = []
        responses = [
            "Explain how to reverse a list in Python.",  # NODE_PROMPT_INITIAL_DEVELOPER
            "Here's one way: `my_list[::-1]`",  # NODE_PROMPT_EXECUTOR
            "{\"Accept\": \"No\"}",  # NODE_PPROMPT_ANALYZER
            "Try using the `reverse()` method instead.",  # NODE_PROMPT_SUGGESTER
            "Explain how to reverse a list in Python. Output in a Markdown List.",  # NODE_PROMPT_DEVELOPER
            "Here's one way: `my_list.reverse()`",  # NODE_PROMPT_EXECUTOR
            "{\"closerOutputID\": 2, \"analysis\": \"The output should use the `reverse()` method.\"}", # NODE_OUTPUT_HISTORY_ANALYZER
            "{\"Accept\": \"Yes\"}",  # NODE_PPROMPT_ANALYZER
        ]
        llm.invoke = lambda x, y = None: responses.pop(0)

        meta_prompt_graph = MetaPromptGraph(llms=llm)
        input_state = AgentState(
            user_message="How do I reverse a list in Python?",
            expected_output="The output should use the `reverse()` method.",
            acceptance_criteria="The output should be correct and efficient.",
            max_output_age=2
        )

        output_state = meta_prompt_graph(input_state)

        self.assertIsNotNone(output_state['best_system_message'])
        self.assertIsNotNone(output_state['best_output'])

        pprint.pp(output_state["best_output"])

    def test_create_acceptance_criteria_workflow(self):
        """
        Test the _create_acceptance_criteria_workflow method of MetaPromptGraph.

        This test case verifies that the workflow created by the _create_acceptance_criteria_workflow method
        contains the correct node and edge.
        """

        llms = {
            NODE_ACCEPTANCE_CRITERIA_DEVELOPER: ChatOpenAI(model_name=os.getenv("TEST_MODEL_NAME_ACCEPTANCE_CRITERIA_DEVELOPER"))
        }
        meta_prompt_graph = MetaPromptGraph(llms=llms)
        workflow = meta_prompt_graph._create_workflow_for_node(NODE_ACCEPTANCE_CRITERIA_DEVELOPER)

        # Check if the workflow contains the correct node
        self.assertIn(NODE_ACCEPTANCE_CRITERIA_DEVELOPER, workflow.nodes)

        # Check if the workflow contains the correct edge
        self.assertIn((NODE_ACCEPTANCE_CRITERIA_DEVELOPER, END), workflow.edges)

        # compile the workflow
        graph = workflow.compile()
        print(graph)

        # invoke the workflow
        state = AgentState(
            user_message="How do I reverse a list in Python?",
            expected_output="The output should use the `reverse()` method.",
            # system_message="Create acceptance criteria for the task of reversing a list in Python."
        )
        output_state = graph.invoke(state)

        # check if the output state contains the acceptance criteria
        self.assertIsNotNone(output_state['acceptance_criteria'])

        # check if the acceptance criteria includes string '`reverse()`'
        self.assertIn('`reverse()`', output_state['acceptance_criteria'])

        pprint.pp(output_state["acceptance_criteria"])


    def test_run_acceptance_criteria_graph(self):
        """Test the run_acceptance_criteria_graph method of MetaPromptGraph.

        This test case verifies that the run_acceptance_criteria_graph method
        returns a state with acceptance criteria.
        """
        llm = Mock(spec=BaseLanguageModel)
        llm.config_specs = []
        llm.invoke = lambda x, y: "{\"Acceptance criteria\": \"Acceptance criteria: ...\"}"
        meta_prompt_graph = MetaPromptGraph(llms=llm)
        state = AgentState(
            user_message="How do I reverse a list in Python?",
            expected_output="The output should use the `reverse()` method.",
        )
        output_state = meta_prompt_graph.run_node_graph(NODE_ACCEPTANCE_CRITERIA_DEVELOPER, state)

        # Check if the output state contains the acceptance criteria
        self.assertIsNotNone(output_state["acceptance_criteria"])

        # Check if the acceptance criteria includes the expected content
        self.assertIn("Acceptance criteria: ...", output_state["acceptance_criteria"])


    def test_run_prompt_initial_developer_graph(self):
        """Test the run_prompt_initial_developer_graph method of MetaPromptGraph.

        This test case verifies that the run_prompt_initial_developer_graph method
        returns a state with an initial developer prompt.
        """
        llm = Mock(spec=BaseLanguageModel)
        llm.config_specs = []
        llm.invoke = lambda x, y: "{\"Initial developer prompt\": \"Initial developer prompt: ...\"}"
        meta_prompt_graph = MetaPromptGraph(llms=llm)
        state = AgentState(user_message="How do I reverse a list in Python?")
        output_state = meta_prompt_graph.run_node_graph(NODE_PROMPT_INITIAL_DEVELOPER, state)

        # Check if the output state contains the initial developer prompt
        self.assertIsNotNone(output_state['system_message'])

        # Check if the initial developer prompt includes the expected content
        self.assertIn("Initial developer prompt: ...", output_state['system_message'])


if __name__ == '__main__':
    unittest.main()