scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-08 03:08:45 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 114

migrate to use new aim FW API

Browse source
This commit is contained in:
drorbaron 2025-08-19 12:20:19 +03:00
parent de59691c4b
commit 6b78ade918
2 changed files with 10 additions and 6 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

9
litellm/proxy/guardrails/guardrail_hooks/aim/aim.py
View file

@ -118,7 +118,7 @@ class AimGuardrail(CustomGuardrail):
litellm_call_id=call_id,
)
response = await self.async_handler.post(
f"{self.api_base}/detect/openai/v2",
f"{self.api_base}/fw/v1/analyze",
headers=headers,
json={"messages": data.get("messages", [])},
)
@ -183,14 +183,17 @@ class AimGuardrail(CustomGuardrail):
)
call_id = request_data.get("litellm_call_id")
response = await self.async_handler.post(
f"{self.api_base}/detect/output/v2",
f"{self.api_base}/fw/v1/analyze",
headers=self._build_aim_headers(
hook=hook,
key_alias=key_alias,
user_email=user_email,
litellm_call_id=call_id,
),
json={"output": output, "messages": request_data.get("messages", [])},
json={
"messages": request_data.get("messages", [])
+ [{"role": "assistant", "content": output}]
},
)
response.raise_for_status()
res = response.json()

7
tests/local_testing/test_aim_guardrails.py
View file

@ -216,12 +216,13 @@ async def test_post_call__with_anonymized_entities__it_deanonymizes_output():
) as mock_post:
def mock_post_detect_side_effect(url, *args, **kwargs):
if url.endswith("/detect/openai/v2"):
request_body = kwargs.get("json", {})
if request_body["messages"][-1]["role"] == "user":
return response_with_detections
elif url.endswith("/detect/output/v2"):
elif request_body["messages"][-1]["role"] == "assistant":
return response_without_detections
else:
raise ValueError("Unexpected URL: {}".format(url))
raise ValueError("Unexpected request: {}".format(request_body))
mock_post.side_effect = mock_post_detect_side_effect