Spaces:
Running
Running
| 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() |