Add tests

src/tests/test_function_calling.py

@@ -0,0 +1,72 @@
+from llama_cpp import LlamaGrammar
+import pytest
+
+from llama_cpp_agent.function_calling import (
+    LlamaCppFunctionTool,
+    LlamaCppFunctionToolRegistry,
+    BaseModel,
+)
+
+
+class MockModel(BaseModel):
+    arg1: int
+    arg2: int
+
+    def run(self) -> int:
+        return self.arg1 * self.arg2
+
+
+@pytest.fixture
+def llama_tool_registry():
+    return LlamaCppFunctionToolRegistry(allow_parallel_function_calling=False)
+
+
+def test_llama_cpp_function_tool_init():
+    llama_tool = LlamaCppFunctionTool(MockModel)
+    assert llama_tool.model == MockModel
+    assert not llama_tool.add_params_to_result
+    assert not llama_tool.look_for_field_string
+    assert not llama_tool.has_markdown_code_block
+    assert not llama_tool.has_triple_quoted_string
+    assert llama_tool.additional_parameters == {}
+
+
+def test_llama_cpp_function_tool_call():
+    llama_tool = LlamaCppFunctionTool(MockModel)
+    result = llama_tool(arg1=10, arg2=2)
+    assert isinstance(result, MockModel)
+    assert result.arg1 == 10
+    assert result.arg2 == 2
+
+
+def test_llama_cpp_function_tool_registry_register(llama_tool_registry):
+    llama_tool = LlamaCppFunctionTool(MockModel)
+    llama_tool_registry.register_function_tool(llama_tool)
+    assert llama_tool_registry.function_tools == {"MockModel": llama_tool}
+
+
+def test_llama_cpp_function_tool_registry_get_function_tool(llama_tool_registry):
+    llama_tool = LlamaCppFunctionTool(MockModel)
+    llama_tool_registry.register_function_tool(llama_tool)
+    assert llama_tool_registry.get_function_tool("MockModel") == llama_tool
+
+
+def test_llama_cpp_function_tool_registry_finalization(llama_tool_registry):
+    llama_tool = LlamaCppFunctionTool(MockModel)
+    llama_tool_registry.register_function_tool(llama_tool)
+    llama_tool_registry.finalize()
+    assert isinstance(llama_tool_registry.grammar, LlamaGrammar)
+    print(llama_tool_registry.get_documentation())
+    assert (
+        llama_tool_registry.get_documentation() == "function: MockModel\n  params:\n    arg1 (int)\n    arg2 (int)\n\n"
+    )
+
+
+def test_llama_cpp_function_tool_registry_handle_function_call(llama_tool_registry):
+    llama_tool = LlamaCppFunctionTool(MockModel)
+    llama_tool_registry.register_function_tool(llama_tool)
+    llama_tool_registry.finalize()
+    response = '{"function": "MockModel", "params": {"arg1": 10, "arg2": 2}}'
+    output = llama_tool_registry.handle_function_call(response)
+    assert output[0]["function"] == "MockModel"
+    assert output[0]["return_value"] == 20

src/tests/test_output_parser.py

@@ -0,0 +1,63 @@
+import pytest
+from llama_cpp_agent.output_parser import (
+    sanitize_and_load_json,
+    is_empty_or_whitespace,
+    parse_json_response,
+    parse_json_response_with_markdown_code_block_or_triple_quoted_string,
+    extract_object_from_response,
+)
+
+
+def test_sanitize_and_load_json_correct_format():
+    # Test JSON in correct format
+    input_json = '{"key": "value"}'
+    expected_output = {"key": "value"}
+    assert sanitize_and_load_json(input_json) == expected_output
+
+
+def test_sanitize_and_load_json_with_newlines():
+    # Test JSON with newlines
+    input_json = '{"key": "value\nwith newline"}'
+    expected_output = {"key": "value with newline"}
+    assert sanitize_and_load_json(input_json) == expected_output
+
+
+def test_is_empty_or_whitespace():
+    # Test empty string
+    assert is_empty_or_whitespace("") is True
+    # Test string with only whitespace
+    assert is_empty_or_whitespace("   ") is True
+    # Test non-empty string
+    assert is_empty_or_whitespace("hello") is False
+
+
+def test_parse_json_response():
+    # Test JSON response parsing
+    response = '{"key": "value"}'
+    expected_output = {"key": "value"}
+    assert parse_json_response(response) == expected_output
+
+
+def test_parse_json_response_with_markdown_code_block_or_triple_quoted_string():
+    # Test parsing JSON response with Markdown code block or triple-quoted string
+    response = """{"key": "value"}
+    ```python
+    print("Hello, world!")
+    ```"""
+    expected_json = {"key": "value"}
+    expected_markdown = 'print("Hello, world!")'
+    json_obj, markdown_content = parse_json_response_with_markdown_code_block_or_triple_quoted_string(response, "```")
+    assert json_obj == expected_json
+    assert markdown_content == expected_markdown
+
+
+def test_extract_object_from_response():
+    # Test extracting object from JSON response
+    class TestObject:
+        def __init__(self, key):
+            self.key = key
+
+    response = '{"key": "value"}'
+    expected_object = TestObject("value")
+    extracted_object = extract_object_from_response(response, TestObject)
+    assert vars(extracted_object) == vars(expected_object)