scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-07 02:59:05 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 103

refactor(panw_prisma_airs): type the tool call argument extractor

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Devin AI 2026-08-15 16:20:09 +00:00
parent dcabd457e5
commit 122a0775cd
1 changed files with 7 additions and 6 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

13
litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py
View file

@ -38,6 +38,8 @@ from litellm.types.guardrails import GuardrailEventHooks
from litellm.types.utils import ( from litellm.types.utils import (
CallTypes, CallTypes,
CallTypesLiteral, CallTypesLiteral,
ChatCompletionMessageToolCall,
ChatCompletionToolCallChunk,
Choices, Choices,
GenericGuardrailAPIInputs, GenericGuardrailAPIInputs,
ModelResponse, ModelResponse,
@ -1431,15 +1433,14 @@ class PanwPrismaAirsHandler(CustomGuardrail):
raise HTTPException(status_code=400, detail=error_detail) raise HTTPException(status_code=400, detail=error_detail)
@staticmethod @staticmethod
def _get_tool_call_arguments(tool_call) -> str | None: def _get_tool_call_arguments(
tool_call: ChatCompletionMessageToolCall | ChatCompletionToolCallChunk,
) -> str | None:
"""Read a tool call's function arguments, handling both object and dict forms.""" """Read a tool call's function arguments, handling both object and dict forms."""
if hasattr(tool_call, "function") and hasattr(tool_call.function, "arguments"):
return tool_call.function.arguments
if isinstance(tool_call, dict): if isinstance(tool_call, dict):
func: Final = tool_call.get("function") func: Final = tool_call.get("function")
if isinstance(func, dict): return func.get("arguments") if isinstance(func, dict) else None
return func.get("arguments") return tool_call.function.arguments
return None
@staticmethod @staticmethod
def _set_tool_call_arguments(tool_call, masked_text: str) -> None: def _set_tool_call_arguments(tool_call, masked_text: str) -> None: