scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-10 22:41:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 129

use llm_router and accept model parameter in ai_policy_suggester

Browse source
This commit is contained in:
Ishaan Jaffer 2026-02-19 11:55:33 -08:00
parent ca5f5b3825
commit 72f58ef316
1 changed files with 10 additions and 4 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -4,9 +4,10 @@ based on user-provided attack examples and descriptions.
"""
import json
from typing import List
from typing import List, Optional
from litellm._logging import verbose_proxy_logger
from litellm.constants import DEFAULT_COMPETITOR_DISCOVERY_MODEL
SUGGEST_TOOL = {
"type": "function",
@ -53,15 +54,20 @@ class AiPolicySuggester:
templates: list,
attack_examples: List[str],
description: str,
model: Optional[str] = None,
) -> dict:
import litellm
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 litellm.acompletion(
model="gpt-4o-mini",
response = await llm_router.acompletion(
model=model,
messages=[
{"role": "system", "content": system_prompt},
{"role": "user", "content": user_prompt},