From 4ad2636f52a29e5a948d6bee10af83da42a3b5ad Mon Sep 17 00:00:00 2001
From: Michael Anstis <manstis@redhat.com>
Date: Thu, 10 Jul 2025 11:54:33 +0100
Subject: [PATCH] Add work-around for tool_calls with granite models and /query
 endpoint.

---
 src/app/endpoints/query.py             | 2 ++
 tests/unit/app/endpoints/test_query.py | 5 +++++
 2 files changed, 7 insertions(+)

diff --git a/src/app/endpoints/query.py b/src/app/endpoints/query.py
index 3727c9cb..7020e252 100644
--- a/src/app/endpoints/query.py
+++ b/src/app/endpoints/query.py
@@ -31,6 +31,7 @@
 from utils.endpoints import check_configuration_loaded, get_system_prompt
 from utils.mcp_headers import mcp_headers_dependency
 from utils.suid import get_suid
+from utils.types import GraniteToolParser
 
 logger = logging.getLogger("app.endpoints.handlers")
 router = APIRouter(tags=["query"])
@@ -83,6 +84,7 @@ def get_agent(
         model=model_id,
         instructions=system_prompt,
         input_shields=available_shields if available_shields else [],
+        tool_parser=GraniteToolParser.get_parser(model_id),
         enable_session_persistence=True,
     )
     conversation_id = agent.create_session(get_suid())
diff --git a/tests/unit/app/endpoints/test_query.py b/tests/unit/app/endpoints/test_query.py
index c1a8e774..4599a391 100644
--- a/tests/unit/app/endpoints/test_query.py
+++ b/tests/unit/app/endpoints/test_query.py
@@ -913,6 +913,7 @@ def test_get_agent_cache_miss_with_conversation_id(
         model="test_model",
         instructions="test_prompt",
         input_shields=["shield1"],
+        tool_parser=None,
         enable_session_persistence=True,
     )
 
@@ -963,6 +964,7 @@ def test_get_agent_no_conversation_id(setup_configuration, prepare_agent_mocks,
         model="test_model",
         instructions="test_prompt",
         input_shields=["shield1"],
+        tool_parser=None,
         enable_session_persistence=True,
     )
 
@@ -1013,6 +1015,7 @@ def test_get_agent_empty_shields(setup_configuration, prepare_agent_mocks, mocke
         model="test_model",
         instructions="test_prompt",
         input_shields=[],
+        tool_parser=None,
         enable_session_persistence=True,
     )
 
@@ -1064,6 +1067,7 @@ def test_get_agent_multiple_mcp_servers(
         model="test_model",
         instructions="test_prompt",
         input_shields=["shield1", "shield2"],
+        tool_parser=None,
         enable_session_persistence=True,
     )
 
@@ -1109,5 +1113,6 @@ def test_get_agent_session_persistence_enabled(
         model="test_model",
         instructions="test_prompt",
         input_shields=["shield1"],
+        tool_parser=None,
         enable_session_persistence=True,
     )