mirror of
https://github.com/Rarebuffalo/securelens-backend.git
synced 2026-06-19 07:00:30 +00:00
90 lines
2.7 KiB
Python
90 lines
2.7 KiB
Python
"""
|
|
AI Client
|
|
=========
|
|
Thin, model-agnostic wrapper around LiteLLM.
|
|
The CLI uses this instead of directly calling litellm
|
|
so we have one place to handle retries, logging, and key injection.
|
|
"""
|
|
|
|
import json
|
|
import asyncio
|
|
import logging
|
|
from typing import Optional
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
async def call_ai(
|
|
prompt: str,
|
|
api_key: str,
|
|
model: str,
|
|
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.
|
|
|
|
Parameters
|
|
----------
|
|
prompt : The prompt to send (added as last user message)
|
|
api_key : LiteLLM-compatible API key
|
|
model : LiteLLM model string (e.g. "gemini/gemini-2.0-flash")
|
|
temperature : Creativity (0=deterministic, 1=creative)
|
|
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
|
|
|
|
litellm.suppress_debug_info = True
|
|
|
|
messages = list(conversation_history or [])
|
|
messages.append({"role": "user", "content": prompt})
|
|
|
|
kwargs: dict = {
|
|
"model": model,
|
|
"messages": messages,
|
|
"temperature": temperature,
|
|
"api_key": api_key if api_key else None,
|
|
}
|
|
|
|
if api_base:
|
|
kwargs["api_base"] = api_base
|
|
if "agentrouter.org" in api_base.lower():
|
|
kwargs["extra_headers"] = {
|
|
"Originator": "codex_cli_rs",
|
|
"User-Agent": "codex_cli_rs/0.101.0 (Mac OS 26.0.1; arm64) Apple_Terminal/464",
|
|
"Version": "0.101.0",
|
|
}
|
|
|
|
if json_mode:
|
|
kwargs["response_format"] = {"type": "json_object"}
|
|
|
|
try:
|
|
response = await litellm.acompletion(**kwargs)
|
|
return response.choices[0].message.content or ""
|
|
except Exception as e:
|
|
logger.error(f"AI call failed [{model}]: {e}")
|
|
return ""
|
|
|
|
|
|
async def call_ai_json(
|
|
prompt: str,
|
|
api_key: str,
|
|
model: str,
|
|
temperature: float = 0.2,
|
|
api_base: Optional[str] = None,
|
|
) -> Optional[dict]:
|
|
"""Convenience wrapper — calls AI in JSON mode and parses the result."""
|
|
raw = await call_ai(prompt, api_key, model, temperature=temperature, json_mode=True, api_base=api_base)
|
|
if not raw:
|
|
return None
|
|
try:
|
|
return json.loads(raw)
|
|
except json.JSONDecodeError as e:
|
|
logger.error(f"JSON parse failed: {e}\nRaw: {raw[:300]}")
|
|
return None
|