tests/meta_prompt_graph_test.py

import unittest
import pprint
import logging
from unittest.mock import MagicMock
from unittest.mock import patch

# Assuming the necessary imports are made for the classes and functions used in meta_prompt_graph.py
from meta_prompt import AgentState, MetaPromptGraph

from langchain_openai import ChatOpenAI

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

    def test_prompt_node(self):
        llms = {
            MetaPromptGraph.NODE_PROMPT_INITIAL_DEVELOPER: MagicMock(
                invoke=MagicMock(return_value=MagicMock(content="Mocked response content"))
            )
        }

        # Create an instance of MetaPromptGraph with the mocked language model and template
        graph = MetaPromptGraph(llms=llms)

        # Create a mock AgentState
        state = AgentState(user_message="Test message", expected_output="Expected output")

        # Invoke the _prompt_node method with the mock node, target attribute, and state
        updated_state = graph._prompt_node(
            MetaPromptGraph.NODE_PROMPT_INITIAL_DEVELOPER, "output", state
        )

        # Assertions
        assert updated_state.output == "Mocked response content", \
            "The output attribute should be updated with the mocked response content"

    def test_output_history_analyzer(self):
        # Setup
        llms = {
            "output_history_analyzer": MagicMock(invoke=lambda prompt: MagicMock(content="""# Analysis

    This analysis compares two outputs to the expected output based on specific criteria.

    # Preferred Output ID: B"""))
        }
        prompts = {}
        meta_prompt_graph = MetaPromptGraph(llms=llms, 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."
        )

        # Invoke the output history analyzer node
        updated_state = meta_prompt_graph._output_history_analyzer(state)

        # Assertions
        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):
        llms = {
            MetaPromptGraph.NODE_PROMPT_ANALYZER: MagicMock(
                invoke=lambda prompt: MagicMock(content="Accept: Yes"))
        }
        meta_prompt_graph = MetaPromptGraph(llms)
        state = AgentState(output="Test output", expected_output="Expected output")
        updated_state = meta_prompt_graph._prompt_analyzer(state)
        assert updated_state.accepted == True

    def test_workflow_execution(self):
        MODEL_NAME = "anthropic/claude-3.5-sonnet:haiku"
        # MODEL_NAME = "meta-llama/llama-3-70b-instruct"
        # MODEL_NAME = "deepseek/deepseek-chat"
        # MODEL_NAME = "google/gemma-2-9b-it"
        # MODEL_NAME = "recursal/eagle-7b"
        # MODEL_NAME = "meta-llama/llama-3-8b-instruct"
        llm = ChatOpenAI(model_name=MODEL_NAME)

        meta_prompt_graph = MetaPromptGraph(llms=llm)
        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."
            )
        output_state = meta_prompt_graph(input_state, recursion_limit=25)

        pprint.pp(output_state)
        # if output_state has key 'best_system_message', print it
        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'])

        # try another similar user message with the generated 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 = llm.invoke(messages)

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

    def test_workflow_execution_with_llms(self):
        optimizer_llm = ChatOpenAI(model_name="anthropic/claude-3.5-sonnet:haiku", temperature=0.5)
        executor_llm = ChatOpenAI(model_name="meta-llama/llama-3-8b-instruct", temperature=0.01)

        llms = {
            MetaPromptGraph.NODE_PROMPT_INITIAL_DEVELOPER: optimizer_llm,
            MetaPromptGraph.NODE_PROMPT_DEVELOPER: optimizer_llm,
            MetaPromptGraph.NODE_PROMPT_EXECUTOR: executor_llm,
            MetaPromptGraph.NODE_OUTPUT_HISTORY_ANALYZER: optimizer_llm,
            MetaPromptGraph.NODE_PROMPT_ANALYZER: optimizer_llm,
            MetaPromptGraph.NODE_PROMPT_SUGGESTER: optimizer_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."
            )
        output_state = meta_prompt_graph(input_state, recursion_limit=25)

        pprint.pp(output_state)
        # if output_state has key 'best_system_message', print it
        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'])

        # try another similar user message with the generated 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 attr 'content' in result
        assert hasattr(result, 'content'), \
            "The result should have the attribute 'content'"
        print(result.content)

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