diff --git a/pilot/utils/function_calling.py b/pilot/utils/function_calling.py
index 3c7c3d2..4ef7ae6 100644
--- a/pilot/utils/function_calling.py
+++ b/pilot/utils/function_calling.py
@@ -1,6 +1,6 @@
 import json
 import re
-from typing import Literal, NotRequired, TypedDict, Callable
+from typing import Literal, Optional, TypedDict, Callable
 
 JsonType = str | int | float | bool | None | list["JsonType"] | dict[str, "JsonType"]
 
@@ -10,14 +10,14 @@ class FunctionParameters(TypedDict):
 
     type: Literal["object"]
     properties: dict[str, JsonType]
-    required: NotRequired[list[str]]
+    required: Optional[list[str]]
 
 
 class FunctionType(TypedDict):
     """Function type"""
 
     name: str
-    description: NotRequired[str]
+    description: Optional[str]
     parameters: FunctionParameters
 
 
diff --git a/pilot/utils/test_function_calling.py b/pilot/utils/test_function_calling.py
index dfd0382..36b5491 100644
--- a/pilot/utils/test_function_calling.py
+++ b/pilot/utils/test_function_calling.py
@@ -1,5 +1,3 @@
-from local_llm_function_calling.prompter import CompletionModelPrompter, InstructModelPrompter
-
 from const.function_calls import ARCHITECTURE, DEV_STEPS
 from .function_calling import parse_agent_response, JsonPrompter
 
@@ -61,81 +59,40 @@ class TestFunctionCalling:
         assert name == 'John'
 
 
-def test_completion_function_prompt():
+def test_json_prompter():
     # Given
-    prompter = CompletionModelPrompter()
+    prompter = JsonPrompter()
 
     # When
     prompt = prompter.prompt('Create a web-based chat app', ARCHITECTURE['definitions'])  # , 'process_technologies')
 
     # Then
-    assert prompt == '''Create a web-based chat app
+    assert prompt == '''Help choose the appropriate function to call to answer the user's question.
+The response should contain only the JSON object, with no additional text or explanation.
 
 Available functions:
-process_technologies - Print the list of technologies that are created.
-```jsonschema
-{
-    "technologies": {
-        "type": "array",
-        "description": "List of technologies.",
-        "items": {
-            "type": "string",
-            "description": "technology"
-        }
-    }
-}
-```
+- process_technologies - Print the list of technologies that are created.
 
-Function call: 
-
-Function call: '''
+Create a web-based chat app'''
 
 
-def test_instruct_function_prompter():
+def test_llama_json_prompter():
     # Given
-    prompter = InstructModelPrompter()
+    prompter = JsonPrompter(is_llama=True)
 
     # When
     prompt = prompter.prompt('Create a web-based chat app', ARCHITECTURE['definitions'])  # , 'process_technologies')
 
     # Then
-    assert prompt == '''Your task is to call a function when needed. You will be provided with a list of functions. Available functions:
-process_technologies - Print the list of technologies that are created.
-```jsonschema
-{
-    "technologies": {
-        "type": "array",
-        "description": "List of technologies.",
-        "items": {
-            "type": "string",
-            "description": "technology"
-        }
-    }
-}
-```
+    assert prompt == '''[INST] <<SYS>>
+Help choose the appropriate function to call to answer the user's question.
+The response should contain only the JSON object, with no additional text or explanation.
 
-Create a web-based chat app
+Available functions:
+- process_technologies - Print the list of technologies that are created.
+<</SYS>>
 
-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]'''
+Create a web-based chat app [/INST]'''
 
 
 def test_json_prompter_named():