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()