Refactored AgentConvo.postprecess_response() into function_calling.process_json_response()

pilot/utils/test_function_calling.py

@@ -1,7 +1,64 @@
 from local_llm_function_calling.prompter import CompletionModelPrompter, InstructModelPrompter
 
 from const.function_calls import ARCHITECTURE, DEV_STEPS
-from .function_calling import JsonPrompter
+from .function_calling import parse_agent_response, LlamaInstructPrompter
+
+
+class TestFunctionCalling:
+    def test_parse_agent_response_text(self):
+        # Given
+        response = {'text': 'Hello world!'}
+
+        # When
+        response = parse_agent_response(response, None)
+
+        # Then
+        assert response == 'Hello world!'
+
+    def test_parse_agent_response_json(self):
+        # Given
+        response = {'text': '{"greeting": "Hello world!"}'}
+        function_calls = {'definitions': [], 'functions': {}}
+
+        # When
+        response = parse_agent_response(response, function_calls)
+
+        # Then
+        assert response == 'Hello world!'
+
+    def test_parse_agent_response_json_markdown(self):
+        # Given
+        response = {'text': '```json\n{"greeting": "Hello world!"}\n```'}
+        function_calls = {'definitions': [], 'functions': {}}
+
+        # When
+        response = parse_agent_response(response, function_calls)
+
+        # Then
+        assert response == 'Hello world!'
+
+    def test_parse_agent_response_markdown(self):
+        # Given
+        response = {'text': '```\n{"greeting": "Hello world!"}\n```'}
+        function_calls = {'definitions': [], 'functions': {}}
+
+        # When
+        response = parse_agent_response(response, function_calls)
+
+        # Then
+        assert response == 'Hello world!'
+
+    def test_parse_agent_response_multiple_args(self):
+        # Given
+        response = {'text': '{"greeting": "Hello", "name": "John"}'}
+        function_calls = {'definitions': [], 'functions': {}}
+
+        # When
+        greeting, name = parse_agent_response(response, function_calls)
+
+        # Then
+        assert greeting == 'Hello'
+        assert name == 'John'
 
 
 def test_completion_function_prompt():
@@ -62,23 +119,23 @@ Create a web-based chat app
 Function call: '''
 
 
-def test_json_prompter():
-    # Given
-    prompter = JsonPrompter()
-
-    # When
-    prompt = prompter.prompt('Create a web-based chat app', ARCHITECTURE['definitions'])  # , 'process_technologies')
-
-    # Then
-    assert prompt == '''[INST] <<SYS>>
-Help choose the appropriate function to call to answer the user's question.
-In your response you must only use JSON output and provide no notes or commentary.
-
-Available functions:
-- process_technologies - Print the list of technologies that are created.
-<</SYS>>
-
-Create a web-based chat app [/INST]'''
+# def test_json_prompter():
+#     # Given
+#     prompter = JsonPrompter()
+#
+#     # When
+#     prompt = prompter.prompt('Create a web-based chat app', ARCHITECTURE['definitions'])  # , 'process_technologies')
+#
+#     # Then
+#     assert prompt == '''[INST] <<SYS>>
+# Help choose the appropriate function to call to answer the user's question.
+# In your response you must only use JSON output and provide no notes or commentary.
+#
+# Available functions:
+# - process_technologies - Print the list of technologies that are created.
+# <</SYS>>
+#
+# Create a web-based chat app [/INST]'''
 
 
 def test_llama_instruct_function_prompter_named():