tool functions and pytest

.gitignore

@@ -1,2 +1,3 @@
 /model/*
 *.prof
+__pycache__

__main__.py

@@ -0,0 +1,8 @@
+
+print("running __main__.-py")
+
+from llama import main
+
+
+if __name__ == "__main__":
+    main()

llama.py

@@ -2,8 +2,9 @@ from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 import time
 import torch
 import random
-import datetime
-import json
+from tool_helper import tool_list, parse_and_execute_tool_call
+from tool_functions import register_dummy
+import utils
 
 t_start = time.time()
 
@@ -41,10 +42,9 @@ print("load took %.3fs" % (time.time() - t_start))
 max_context_length = model.config.max_position_embeddings
 
 
-# if tokenizer.chat_template is None:
-print("apply external chat template...")
-with open("chat_template.json", "r") as f:
-    tokenizer.chat_template = json.load(f)
+tokenizer.chat_template = utils.load_json_file("chat_template.json")
+
+
 
 
 print("max_context_length is %d tokens." % (max_context_length))
@@ -97,22 +97,9 @@ messages = [
 roleflip = {"role": "system", "content": "Keep the conversation going, ask for more information on the subject. Keep messages short at max 1-2 sentences. Do not thank and say goodbye."}
 
 
-def current_time():
-    """Get the current local date and time as a string."""
-    return datetime.datetime.now().strftime("%Y-%m-%d %H:%M")
-
-def random_float():
-    """Get a random float from 0..1"""
-    return str(random.random())
 
-def random_int(a: int, b: int):
-    """Return random integer in range [a, b], including both end points.
-Args:
-    a: minimum possible value
-    b: maximum possible value"""
-    return str(random.randint(a, b))
-
-tool_functions = [current_time, random_float, random_int]
+register_dummy()
+# tool_functions = [current_time, random_float, random_int]
 
 
 
@@ -139,6 +126,7 @@ def generate_incremental(inputs):
     generated_tokens = input_ids  # Initially, this is just the input tokens
 
     n = 0
+    try:
 
         # Loop to generate one token at a time
         while True:
@@ -161,10 +149,13 @@ def generate_incremental(inputs):
                 break
 
             n += 1
-        if n >= 30:
+            if n >= 15:
                 n = 0
                 torch.cuda.empty_cache()
 
+    except KeyboardInterrupt:
+        pass
+
 
     # Once done, return the full generated sequence
     input_token_count = len(inputs["input_ids"][0])
@@ -184,7 +175,7 @@ def append_generate_chat(input_text: str, role="user"):
 
     # input_text = "Hello, who are you?"
     # inputs = tokenizer(input_text, return_tensors="pt").to("cpu")  # .to("cuda") .to("cpu")
-    inputs = tokenizer.apply_chat_template(messages, return_tensors="pt", tokenize=True, return_dict=True, add_generation_prompt=True, tools=tool_functions)  #continue_final_message=True,
+    inputs = tokenizer.apply_chat_template(messages, return_tensors="pt", tokenize=True, return_dict=True, add_generation_prompt=True, tools=tool_list)  #continue_final_message=True,
     inputs = {key: value.to(model.device) for key, value in inputs.items()}
     # inputs = {key: value.to("cpu") for key, value in inputs.items()}
     # inputs["input_ids"] = inputs["input_ids"][:, 1:]
@@ -194,11 +185,23 @@ def append_generate_chat(input_text: str, role="user"):
 
     # append result to message history
     messages.append({"role": "assistant", "content": out_text})
+
     print("")
     print("generation took %.3fs (%d tokens)" % (time.time() - t_start, len(outputs[0])))
 
+    # handle tool call and check if a tool call has happened.
+    tool_result = parse_and_execute_tool_call(out_text, tool_list)
+    if tool_result != None:
+        # tool call happened
+        # tool_result = "<tool_response>%s</tool_response>" % tool_result
+        # depending on the chat template the tool response tags must or must not be passed. :(
+        append_generate_chat(tool_result, role="tool")
 
-while True:
+
+
+def main():
+    global messages
+    while True:
         # print an input prompt to receive text or commands
         input_text = input(">>> ")
         print("")
@@ -210,11 +213,12 @@ while True:
 
         elif input_text.startswith("/clear"):
             print("clearing chat history")
-        messages = [messages[0]]
+            start_msg = messages[0]
+            messages = [start_msg]
             print("")
 
         elif input_text.startswith("/history"):
-        history = tokenizer.apply_chat_template(messages, return_tensors="pt", tokenize=False, add_generation_prompt=False, tools=tool_functions)
+            history = tokenizer.apply_chat_template(messages, return_tensors="pt", tokenize=False, add_generation_prompt=False, tools=tool_list)
             print(history)
 
         elif input_text.startswith("/undo"):
@@ -240,6 +244,14 @@ while True:
         elif input_text.startswith("/more"):
             append_generate_chat(None)
 
+        elif input_text.startswith("/file"):
+            filename = input_text[len("/file "):]
+            print("read '%s' for prompt:" % filename)
+            with open(filename, "r") as f:
+                content = f.read()
+            print(content)
+            append_generate_chat(content)
+
         elif input_text.startswith("/auto"):
             messages_backup = messages
             messages = [roleflip]
@@ -264,6 +276,7 @@ while True:
             print("/undo       undo latest prompt")
             print("/regen      regenerate the last message")
             print("/more       generate more additional information")
+            print("/file       read prompt input from file")
             print("/auto       automatically advance conversation")
             print("/help       print this message")
             print("")
@@ -273,3 +286,4 @@ while True:
 
         else:
             append_generate_chat(input_text)
+

tests/__init__.py

@@ -0,0 +1 @@
+# empty

tests/helper.py

@@ -0,0 +1,7 @@
+
+
+def tool_dummy(a: int, b: str):
+    return "result_%d_%s" % (a, b)
+
+def tool_dummy2(text: str):
+    return text.upper()

tests/test_tool_function_decorator.py

@@ -0,0 +1,30 @@
+import pytest
+import tool_helper
+import tests.helper as helper
+
+
+def test_tool_function_decorator_if_clean_tool_list():
+    """ tests for the tool list to be empty. NOT strictly nessesary,
+    but I want to be warned if this is not the case anymore. Could be not the intention """
+    start_len = len(tool_helper.tool_list)
+    assert start_len == 0
+
+def test_tool_function_decorator():
+    # get length before adding tools
+    start_len = len(tool_helper.tool_list)
+
+    # add tools like it would be a decorator
+    tool_helper.tool(helper.tool_dummy)
+    tool_helper.tool(helper.tool_dummy2)
+
+    # get length after adding tools
+    end_len = len(tool_helper.tool_list)
+
+    # remove the added ones again
+    tool_helper.tool_list = tool_helper.tool_list[:-2]
+
+    assert end_len == start_len + 2
+    assert len(tool_helper.tool_list) == start_len
+
+
+

tests/test_tool_parse_exec.py

@@ -0,0 +1,89 @@
+import pytest
+import tool_helper
+from unittest import mock
+import tests.helper as helper
+
+
+
+
+
+def test_tool_dummy():
+    with mock.patch("tests.helper.tool_dummy") as mock_dummy:
+        helper.tool_dummy()
+        mock_dummy.assert_called_once()  # this will check if the mocked function on the context was called.
+
+
+def test_tool_parse_no_exec():
+    with mock.patch("tests.helper.tool_dummy") as mock_dummy:
+        tool_helper.parse_and_execute_tool_call("something else", [helper.tool_dummy, helper.tool_dummy2])
+        assert mock_dummy.call_count == 0
+
+
+def test_match_and_extract_no_match():
+    result = tool_helper._match_and_extract("something else", r"<tool_call>(.*)<\/tool_call>")
+    assert result == None
+
+
+def test_match_and_extract_matching():
+    result = tool_helper._match_and_extract("asdfsdfas   <tool_call>{json content}</tool_call>   adfafsd", r"<tool_call>(.*)<\/tool_call>")
+    assert result == "{json content}"
+
+
+def test_match_and_extract_matching2():
+    result = tool_helper._match_and_extract("<tool_call>{json content}</tool_call>", r"<tool_call>(.*)<\/tool_call>")
+    assert result == "{json content}"
+
+
+def test_match_and_extract_matching3_with_newline():
+    result = tool_helper._match_and_extract("<tool_call>\n{json content}\n</tool_call>", r"<tool_call>(.*)<\/tool_call>")
+    assert result == "\n{json content}\n"
+
+
+def test_string_malformed_faulty():
+    with mock.patch("utils.print_error") as print_error_mock:
+        result = tool_helper._execute_tool_call_str("{json_content}", [])
+        assert result == None
+        print_error_mock.assert_called_once()  # this will check if the mocked function on the context was called.
+
+
+def test_tool_call_json_1():
+    with mock.patch("utils.print_error") as print_error_mock:
+        result = tool_helper._execute_tool_call_json({"name": "tool_dummy", "arguments": {"a": 1, "b": "zwei"}}, [helper.tool_dummy, helper.tool_dummy2])
+        assert result == "result_1_zwei"
+        assert print_error_mock.call_count == 0
+
+
+def test_tool_call_json_2():
+    with mock.patch("utils.print_error") as print_error_mock:
+        result = tool_helper._execute_tool_call_json({"name": "tool_dummy2", "arguments": {"text": "some_text"}}, [helper.tool_dummy, helper.tool_dummy2])
+        assert result == "SOME_TEXT"
+        assert print_error_mock.call_count == 0
+
+
+def test_tool_call_json_non_existing_call_check():
+    with mock.patch("utils.print_error") as print_error_mock:
+        result = tool_helper._execute_tool_call_json({"name": "tool_dummy_which_is_not_existing", "arguments": {"text": "some_text"}}, [helper.tool_dummy, helper.tool_dummy2])
+        assert result == None
+        assert print_error_mock.call_count == 1  # this will check if the mocked function on the context was called.
+
+def test_tool_call_json_wrong_arguments_check():
+    with mock.patch("utils.print_error") as print_error_mock:
+        result = tool_helper._execute_tool_call_json({"name": "tool_dummy", "arguments": {"a": "must_be_an_int_but_is_string", "b": "zwei"}}, [helper.tool_dummy, helper.tool_dummy2])
+        assert result == None
+        assert print_error_mock.call_count == 1  # this will check if the mocked function on the context was called.
+
+
+
+def test_regex_multiline():
+    import re
+    pattern = r"<start>(.*)</end>"
+
+    # The text to search (spanning multiple lines)
+    text = """<start>
+    {json}
+    </end>"""
+
+    # Use re.search with re.DOTALL to match across newlines
+    match = re.search(pattern, text, re.DOTALL)
+
+    assert match.group(1).find("{json}") != -1

tool_functions.py

@@ -0,0 +1,35 @@
+import random
+import datetime
+from tool_helper import tool
+
+@tool
+def current_time():
+    """Get the current local date and time as a string."""
+    return datetime.datetime.now().strftime("%Y-%m-%d %H:%M")
+
+# @tool
+# def random_float():
+#     """Generate a random float from 0..1."""
+#     return str(random.random())
+
+@tool
+def random_float(a: float=0.0, b: float=1.0):
+    """Generate a random float in range [a, b], including both end points. Optional pass no parameter and range 0..1 will be used.
+Args:
+    a: minimum possible value
+    b: maximum possible value"""
+    return str(random.randint(a, b))
+
+@tool
+def random_int(a: int, b: int):
+    """Generate a random integer in range [a, b], including both end points.
+Args:
+    a: minimum possible value
+    b: maximum possible value"""
+    return str(random.randint(a, b))
+
+
+
+
+def register_dummy():
+    pass  # dummy function to run and be sure the decorators have run

tool_helper.py

@@ -0,0 +1,102 @@
+
+from typing import Callable, List, Optional
+import json
+import re
+import utils
+
+tool_list = []
+
+
+def tool(fn):
+    """tool function decorator"""
+    print("register tool '%s'" % fn.__name__)
+    tool_list.append(fn)
+
+# def parse_and_execute_tool_call(message: str, tools: list[function]) -> str | None:
+#     """execute tool call if needed accordint <tool_call> tag and return the content of the tool call or None if no call happened."""
+#     pass
+
+
+
+
+
+
+def parse_and_execute_tool_call(message: str, tools: List[Callable]) -> Optional[str]:
+    """
+    Execute a tool call if the <tool_call> tag is present and return the tool's response.
+    If no <tool_call> tag is found, return None.
+
+    Args:
+        message (str): The message containing the tool call.
+        tools (list[function]): A list of tool functions available for execution.
+
+    Returns:
+        Optional[str]: The content of the tool response or None if no tool call occurred.
+    """
+
+    # in case LLM responds with <tool_call></tool_call> the correct way
+    extracted = _match_and_extract(message, r"<tool_call>(.*)<\/tool_call>")
+    if extracted:
+        return _execute_tool_call_str(extracted, tools)
+
+    # in case LLM responds with <tool_call></tool_response> by accident
+    extracted = _match_and_extract(message, r"<tool_call>(.*)<\/tool_.*>")
+    if extracted:
+        return _execute_tool_call_str(extracted, tools)
+
+    # in case LLM responds with <tool_call></???> by accident
+    extracted = _match_and_extract(message, r"<tool_call>(.*)<\/.*>")
+    if extracted:
+        return _execute_tool_call_str(extracted, tools)
+
+    # in case LLM responds with <tool_call></???> by accident
+    extracted = _match_and_extract(message, r"<tool_response>(.*)<\/.*>")
+    if extracted:
+        return _execute_tool_call_str(extracted, tools)
+
+    return None
+
+
+def _match_and_extract(message: str, pattern: str) -> Optional[str]:
+    """ helper function to match regex and extract group 1 """
+    match = re.search(pattern, message, re.DOTALL)
+    if match:
+        group1 = match.group(1)
+        return group1
+    return None
+
+
+def _execute_tool_call_str(tool_call: str, tools: List[Callable]) -> Optional[str]:
+    """ execute tool call per string content. The content must be a valid json """
+    try:
+        js = json.loads(tool_call)
+        return _execute_tool_call_json(js, tools)
+    except json.JSONDecodeError:
+        utils.print_error("Json was malformed. Will be ignored.")
+        return None
+
+def _execute_tool_call_json(data: any, tools: List[Callable]) -> Optional[str]:
+    """ extract name and arguments from parsed data and call the tool, which is matched from the tools list """
+    # Extract tool name and arguments
+    tool_name = data.get("name")
+    arguments = data.get("arguments", {})
+
+    # Find the tool by name in the list of tools
+    for tool in tools:
+        if tool.__name__ == tool_name:
+            # Execute the tool
+            return _execute_tool_function(arguments, tool)
+
+    utils.print_error("Specified tool '%s' not found." % tool_name)
+    return None
+
+def _execute_tool_function(arguments: any, tool: Callable) -> Optional[str]:
+    """ Execute the tool and return the result. """
+    try:
+        result = tool(**arguments)
+        print("<tool_response>", result, "</tool_response>")
+        return result
+    except TypeError as e:
+        utils.print_error("Type error while executing function call: '%s'" % str(e))
+
+    return None

utils.py

@@ -0,0 +1,11 @@
+import json
+import sys
+
+def load_json_file(filepath: str) -> any:
+    with open(filepath, "r") as f:
+        return json.load(f)
+
+
+
+def print_error(*args, **kwargs):
+     print(*args, file=sys.stderr, **kwargs)