Skip to content

Commit 61fcaaa

Browse files
committed
adds optional no_tools parameter on /query and /streaming_query to skip all tool calling
1 parent 423272b commit 61fcaaa

File tree

5 files changed

+688
-45
lines changed

5 files changed

+688
-45
lines changed

src/app/endpoints/query.py

Lines changed: 34 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -80,6 +80,7 @@ def get_agent( # pylint: disable=too-many-arguments,too-many-positional-argumen
8080
available_input_shields: list[str],
8181
available_output_shields: list[str],
8282
conversation_id: str | None,
83+
no_tools: bool = False,
8384
) -> tuple[Agent, str]:
8485
"""Get existing agent or create a new one with session persistence."""
8586
if conversation_id is not None:
@@ -99,7 +100,7 @@ def get_agent( # pylint: disable=too-many-arguments,too-many-positional-argumen
99100
instructions=system_prompt,
100101
input_shields=available_input_shields if available_input_shields else [],
101102
output_shields=available_output_shields if available_output_shields else [],
102-
tool_parser=GraniteToolParser.get_parser(model_id),
103+
tool_parser=None if no_tools else GraniteToolParser.get_parser(model_id),
103104
enable_session_persistence=True,
104105
)
105106
conversation_id = agent.create_session(get_suid())
@@ -288,36 +289,47 @@ def retrieve_response( # pylint: disable=too-many-locals
288289
available_input_shields,
289290
available_output_shields,
290291
query_request.conversation_id,
292+
query_request.no_tools or False,
291293
)
292294

293-
# preserve compatibility when mcp_headers is not provided
294-
if mcp_headers is None:
295+
# bypass tools and MCP servers if no_tools is True
296+
if query_request.no_tools:
295297
mcp_headers = {}
296-
mcp_headers = handle_mcp_headers_with_toolgroups(mcp_headers, configuration)
297-
if not mcp_headers and token:
298-
for mcp_server in configuration.mcp_servers:
299-
mcp_headers[mcp_server.url] = {
300-
"Authorization": f"Bearer {token}",
301-
}
302-
303-
agent.extra_headers = {
304-
"X-LlamaStack-Provider-Data": json.dumps(
305-
{
306-
"mcp_headers": mcp_headers,
307-
}
308-
),
309-
}
298+
agent.extra_headers = {}
299+
toolgroups = None
300+
else:
301+
# preserve compatibility when mcp_headers is not provided
302+
if mcp_headers is None:
303+
mcp_headers = {}
304+
mcp_headers = handle_mcp_headers_with_toolgroups(mcp_headers, configuration)
305+
if not mcp_headers and token:
306+
for mcp_server in configuration.mcp_servers:
307+
mcp_headers[mcp_server.url] = {
308+
"Authorization": f"Bearer {token}",
309+
}
310+
311+
agent.extra_headers = {
312+
"X-LlamaStack-Provider-Data": json.dumps(
313+
{
314+
"mcp_headers": mcp_headers,
315+
}
316+
),
317+
}
318+
319+
vector_db_ids = [vector_db.identifier for vector_db in client.vector_dbs.list()]
320+
toolgroups = (get_rag_toolgroups(vector_db_ids) or []) + [
321+
mcp_server.name for mcp_server in configuration.mcp_servers
322+
]
323+
# Convert empty list to None for consistency with existing behavior
324+
if not toolgroups:
325+
toolgroups = None
310326

311-
vector_db_ids = [vector_db.identifier for vector_db in client.vector_dbs.list()]
312-
toolgroups = (get_rag_toolgroups(vector_db_ids) or []) + [
313-
mcp_server.name for mcp_server in configuration.mcp_servers
314-
]
315327
response = agent.create_turn(
316328
messages=[UserMessage(role="user", content=query_request.query)],
317329
session_id=conversation_id,
318330
documents=query_request.get_documents(),
319331
stream=False,
320-
toolgroups=toolgroups or None,
332+
toolgroups=toolgroups,
321333
)
322334

323335
# Check for validation errors in the response

src/app/endpoints/streaming_query.py

Lines changed: 36 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -58,6 +58,7 @@ async def get_agent(
5858
available_input_shields: list[str],
5959
available_output_shields: list[str],
6060
conversation_id: str | None,
61+
no_tools: bool = False,
6162
) -> tuple[AsyncAgent, str]:
6263
"""Get existing agent or create a new one with session persistence."""
6364
if conversation_id is not None:
@@ -76,7 +77,7 @@ async def get_agent(
7677
instructions=system_prompt,
7778
input_shields=available_input_shields if available_input_shields else [],
7879
output_shields=available_output_shields if available_output_shields else [],
79-
tool_parser=GraniteToolParser.get_parser(model_id),
80+
tool_parser=None if no_tools else GraniteToolParser.get_parser(model_id),
8081
enable_session_persistence=True,
8182
)
8283
conversation_id = await agent.create_session(get_suid())
@@ -532,41 +533,53 @@ async def retrieve_response(
532533
available_input_shields,
533534
available_output_shields,
534535
query_request.conversation_id,
536+
query_request.no_tools or False,
535537
)
536538

537-
# preserve compatibility when mcp_headers is not provided
538-
if mcp_headers is None:
539+
# bypass tools and MCP servers if no_tools is True
540+
if query_request.no_tools:
539541
mcp_headers = {}
542+
agent.extra_headers = {}
543+
toolgroups = None
544+
else:
545+
# preserve compatibility when mcp_headers is not provided
546+
if mcp_headers is None:
547+
mcp_headers = {}
540548

541-
mcp_headers = handle_mcp_headers_with_toolgroups(mcp_headers, configuration)
549+
mcp_headers = handle_mcp_headers_with_toolgroups(mcp_headers, configuration)
542550

543-
if not mcp_headers and token:
544-
for mcp_server in configuration.mcp_servers:
545-
mcp_headers[mcp_server.url] = {
546-
"Authorization": f"Bearer {token}",
547-
}
551+
if not mcp_headers and token:
552+
for mcp_server in configuration.mcp_servers:
553+
mcp_headers[mcp_server.url] = {
554+
"Authorization": f"Bearer {token}",
555+
}
548556

549-
agent.extra_headers = {
550-
"X-LlamaStack-Provider-Data": json.dumps(
551-
{
552-
"mcp_headers": mcp_headers,
553-
}
554-
),
555-
}
557+
agent.extra_headers = {
558+
"X-LlamaStack-Provider-Data": json.dumps(
559+
{
560+
"mcp_headers": mcp_headers,
561+
}
562+
),
563+
}
564+
565+
logger.debug("Session ID: %s", conversation_id)
566+
vector_db_ids = [
567+
vector_db.identifier for vector_db in await client.vector_dbs.list()
568+
]
569+
toolgroups = (get_rag_toolgroups(vector_db_ids) or []) + [
570+
mcp_server.name for mcp_server in configuration.mcp_servers
571+
]
572+
# Convert empty list to None for consistency with existing behavior
573+
if not toolgroups:
574+
toolgroups = None
556575

557576
logger.debug("Session ID: %s", conversation_id)
558-
vector_db_ids = [
559-
vector_db.identifier for vector_db in await client.vector_dbs.list()
560-
]
561-
toolgroups = (get_rag_toolgroups(vector_db_ids) or []) + [
562-
mcp_server.name for mcp_server in configuration.mcp_servers
563-
]
564577
response = await agent.create_turn(
565578
messages=[UserMessage(role="user", content=query_request.query)],
566579
session_id=conversation_id,
567580
documents=query_request.get_documents(),
568581
stream=True,
569-
toolgroups=toolgroups or None,
582+
toolgroups=toolgroups,
570583
)
571584

572585
return response, conversation_id

src/models/requests.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -69,6 +69,7 @@ class QueryRequest(BaseModel):
6969
model: The optional model.
7070
system_prompt: The optional system prompt.
7171
attachments: The optional attachments.
72+
no_tools: Whether to bypass all tools and MCP servers (default: False).
7273
7374
Example:
7475
```python
@@ -82,6 +83,7 @@ class QueryRequest(BaseModel):
8283
model: Optional[str] = None
8384
system_prompt: Optional[str] = None
8485
attachments: Optional[list[Attachment]] = None
86+
no_tools: Optional[bool] = False
8587
# media_type is not used in 'lightspeed-stack' that only supports application/json.
8688
# the field is kept here to enable compatibility with 'road-core' clients.
8789
media_type: Optional[str] = None
@@ -97,6 +99,7 @@ class QueryRequest(BaseModel):
9799
"provider": "openai",
98100
"model": "model-name",
99101
"system_prompt": "You are a helpful assistant",
102+
"no_tools": False,
100103
"attachments": [
101104
{
102105
"attachment_type": "log",

0 commit comments

Comments
 (0)