From dc9e296661a4cd6c0e70fe2dff302bcff5b14fad Mon Sep 17 00:00:00 2001
From: James Briggs <35938317+jamescalam@users.noreply.github.com>
Date: Sat, 13 Jan 2024 17:59:44 +0000
Subject: [PATCH] added extract function inputs test for llamacpp

---
 tests/unit/llms/test_llm_llamacpp.py | 16 ++++++++++++++++
 1 file changed, 16 insertions(+)

diff --git a/tests/unit/llms/test_llm_llamacpp.py b/tests/unit/llms/test_llm_llamacpp.py
index 507a5273..e2bad635 100644
--- a/tests/unit/llms/test_llm_llamacpp.py
+++ b/tests/unit/llms/test_llm_llamacpp.py
@@ -30,3 +30,19 @@ class TestLlamaCppLLM:
 
     def test_llamacpp_llm_grammar(self, llamacpp_llm):
         llamacpp_llm._grammar()
+
+    def test_llamacpp_extract_function_inputs(self, llamacpp_llm, mocker):
+        llamacpp_llm.llm.create_chat_completion = mocker.Mock(
+            return_value={"choices": [{"message": {"content": "{'timezone': 'America/New_York'}"}}]}
+        )
+        test_schema = {
+            "name": "get_time",
+            "description": 'Finds the current time in a specific timezone.\n\n:param timezone: The timezone to find the current time in, should\n    be a valid timezone from the IANA Time Zone Database like\n    "America/New_York" or "Europe/London". Do NOT put the place\n    name itself like "rome", or "new york", you must provide\n    the IANA format.\n:type timezone: str\n:return: The current time in the specified timezone.',
+            "signature": "(timezone: str) -> str",
+            "output": "<class 'str'>",
+        }
+        test_query = "What time is it in America/New_York?"
+
+        llamacpp_llm.extract_function_inputs(
+            query=test_query, function_schema=test_schema
+        )
-- 
GitLab