mirror of
https://github.com/Rarebuffalo/securelens-backend.git
synced 2026-06-19 07:00:30 +00:00
add support for custom AI api base url and Agent Router integrations
This commit is contained in:
@@ -21,6 +21,7 @@ async def call_ai(
|
||||
temperature: float = 0.3,
|
||||
json_mode: bool = False,
|
||||
conversation_history: Optional[list] = None,
|
||||
api_base: Optional[str] = None,
|
||||
) -> str:
|
||||
"""
|
||||
Single entry-point for all AI calls in the CLI.
|
||||
@@ -34,6 +35,7 @@ async def call_ai(
|
||||
json_mode : Ask the model to respond with valid JSON only
|
||||
conversation_history : Optional list of {"role": ..., "content": ...} dicts
|
||||
for multi-turn chat sessions
|
||||
api_base : Optional custom API base URL (e.g. for Agent Router)
|
||||
"""
|
||||
import litellm
|
||||
|
||||
@@ -49,6 +51,9 @@ async def call_ai(
|
||||
"api_key": api_key if api_key else None,
|
||||
}
|
||||
|
||||
if api_base:
|
||||
kwargs["api_base"] = api_base
|
||||
|
||||
if json_mode:
|
||||
kwargs["response_format"] = {"type": "json_object"}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user