From f963e6d70502ae36bd4c6e7830a83cb46ff51759 Mon Sep 17 00:00:00 2001
From: Siraj R Aizlewood <siraj@aurelio.ai>
Date: Tue, 30 Apr 2024 04:40:32 +0400
Subject: [PATCH] Another pytest.
---
tests/unit/llms/test_llm_openai.py | 20 +++++++++++++++++++-
1 file changed, 19 insertions(+), 1 deletion(-)
diff --git a/tests/unit/llms/test_llm_openai.py b/tests/unit/llms/test_llm_openai.py
index 25573bd9..6c01b2bd 100644
--- a/tests/unit/llms/test_llm_openai.py
+++ b/tests/unit/llms/test_llm_openai.py
@@ -178,4 +178,22 @@ class TestOpenAILLM:
assert convert_param_type_to_json_type("NoneType") == "null"
assert convert_param_type_to_json_type("list") == "array"
# Test conversion of a type not explicitly handled
- assert convert_param_type_to_json_type("dict") == "object"
\ No newline at end of file
+ assert convert_param_type_to_json_type("dict") == "object"
+
+ def test_extract_function_inputs(self, openai_llm, mocker):
+ query = "fetch user data"
+ function_schema = {"function": "get_user_data", "args": ["user_id"]}
+
+ # Mock the __call__ method to return a JSON string as expected
+ mocker.patch.object(OpenAILLM, '__call__', return_value='{"user_id": "123"}')
+ result = openai_llm.extract_function_inputs(query, function_schema)
+
+ # Ensure the __call__ method is called with the correct parameters
+ expected_messages = [
+ Message(role="system", content="You are an intelligent AI. Given a command or request from the user, call the function to complete the request."),
+ Message(role="user", content=query)
+ ]
+ openai_llm.__call__.assert_called_once_with(messages=expected_messages, function_schema=function_schema)
+
+ # Check if the result is as expected
+ assert result == {"user_id": "123"}, "The function inputs should match the expected dictionary."
\ No newline at end of file
--
GitLab