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

Merge pull request #21638 from BerriAI/fix/ai-policy-suggester-use-litellm-acompletion

Browse source 
fix(policy): use litellm.acompletion directly in AiPolicySuggester
This commit is contained in:
jquinter 2026-02-20 01:14:59 -03:00 • committed by GitHub
commit d08c3a1d68
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 2 additions and 6 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

8
litellm/proxy/management_endpoints/policy_endpoints/ai_policy_suggester.py
View file

@ -6,6 +6,7 @@ based on user-provided attack examples and descriptions.
import json
from typing import List, Optional
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.constants import DEFAULT_COMPETITOR_DISCOVERY_MODEL
@ -56,17 +57,12 @@ class AiPolicySuggester:
description: str,
model: Optional[str] = None,
) -> dict:
from litellm.proxy.proxy_server import llm_router
if llm_router is None:
raise ValueError("LLM router not initialized")
system_prompt = self._build_system_prompt(templates)
user_prompt = self._build_user_prompt(attack_examples, description)
model = model or DEFAULT_COMPETITOR_DISCOVERY_MODEL
try:
response = await llm_router.acompletion(
response = await litellm.acompletion(
model=model,
messages=[
{"role": "system", "content": system_prompt},