475 lines
21 KiB
Python
475 lines
21 KiB
Python
import os
|
|
import json
|
|
import time
|
|
import asyncio
|
|
import shutil
|
|
import base64
|
|
import httpx
|
|
import logging
|
|
from typing import Dict, Any, Optional, List, Callable, Awaitable
|
|
from core.metrics import AI_REQUESTS_TOTAL, AI_LATENCY_SECONDS
|
|
from db.models import AIProviderProfile
|
|
|
|
try:
|
|
from google.antigravity import Agent as AgyAgent, LocalAgentConfig as AgyLocalAgentConfig
|
|
HAS_AGY_SDK = True
|
|
except ImportError:
|
|
HAS_AGY_SDK = False
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
class LLMClient:
|
|
def __init__(
|
|
self,
|
|
provider: Optional[str] = None,
|
|
api_key: Optional[str] = None,
|
|
model: Optional[str] = None,
|
|
base_url: Optional[str] = None,
|
|
reasoning_effort: Optional[str] = None,
|
|
repo: Optional[Any] = None,
|
|
on_fallback_alert: Optional[Callable[[AIProviderProfile, AIProviderProfile, str, int], Awaitable[None]]] = None,
|
|
on_chain_failure_alert: Optional[Callable[[List[AIProviderProfile], str], Awaitable[None]]] = None,
|
|
):
|
|
self.repo = repo
|
|
self.provider = provider or os.getenv("AI_PROVIDER", "openai").lower()
|
|
self.api_key = api_key or os.getenv("AI_API_KEY", "")
|
|
default_model = "antigravity" if self.provider == "agy" else ("orcarouter/auto" if self.provider == "openai" else "gemini-1.5-flash")
|
|
self.model = model or os.getenv("AI_MODEL", default_model)
|
|
default_base_url = "http://localhost:8000/v1" if self.provider == "agy" else "https://api.orcarouter.ai/v1"
|
|
self.base_url = base_url or os.getenv("AI_BASE_URL", default_base_url)
|
|
self.reasoning_effort = (reasoning_effort or os.getenv("AI_REASONING_EFFORT", "")).strip().lower()
|
|
fallbacks = [m.strip() for m in os.getenv("AI_FALLBACK_MODELS", "").split(",") if m.strip()]
|
|
self.models = [self.model] + [m for m in fallbacks if m != self.model]
|
|
self.max_retries_per_model = int(os.getenv("AI_MAX_RETRIES", "2"))
|
|
self.retry_backoff_seconds = float(os.getenv("AI_RETRY_BACKOFF_SECONDS", "2"))
|
|
self.last_used_model = self.model
|
|
self.last_used_provider = self.provider
|
|
self.on_fallback_alert = on_fallback_alert
|
|
self.on_chain_failure_alert = on_chain_failure_alert
|
|
|
|
async def sync_config_from_repo(self):
|
|
if not self.repo:
|
|
return
|
|
try:
|
|
active_profile = await self.repo.get_active_provider_profile()
|
|
if active_profile:
|
|
self.provider = active_profile.provider_type.strip().lower()
|
|
self.model = active_profile.model.strip()
|
|
self.base_url = active_profile.base_url.strip()
|
|
self.api_key = active_profile.api_key.strip()
|
|
self.reasoning_effort = active_profile.reasoning_effort.strip().lower()
|
|
else:
|
|
db_provider = await self.repo.get_setting("ai_provider")
|
|
if db_provider:
|
|
self.provider = db_provider.strip().lower()
|
|
|
|
db_model = await self.repo.get_setting("ai_model")
|
|
if db_model:
|
|
self.model = db_model.strip()
|
|
|
|
db_base_url = await self.repo.get_setting("ai_base_url")
|
|
if db_base_url is not None:
|
|
self.base_url = db_base_url.strip()
|
|
|
|
db_api_key = await self.repo.get_setting("ai_api_key")
|
|
if db_api_key is not None:
|
|
self.api_key = db_api_key.strip()
|
|
|
|
db_reasoning = await self.repo.get_setting("ai_reasoning_effort")
|
|
if db_reasoning is not None:
|
|
self.reasoning_effort = db_reasoning.strip().lower()
|
|
|
|
fallbacks = [m.strip() for m in os.getenv("AI_FALLBACK_MODELS", "").split(",") if m.strip()]
|
|
self.models = [self.model] + [m for m in fallbacks if m != self.model]
|
|
except Exception as e:
|
|
logger.debug(f"Could not sync AI config from DB: {e}")
|
|
|
|
async def get_fallback_chain(self) -> List[AIProviderProfile]:
|
|
"""Return the ordered list of AI provider profiles manually chained via fallback_provider_id."""
|
|
if self.repo:
|
|
try:
|
|
profiles = await self.repo.get_provider_profiles()
|
|
if profiles:
|
|
profile_map = {p.id: p for p in profiles if p.id is not None}
|
|
active = next((p for p in profiles if p.is_active), profiles[0])
|
|
chain = [active]
|
|
visited = {active.id}
|
|
curr = active
|
|
while curr.fallback_provider_id and curr.fallback_provider_id in profile_map:
|
|
next_p = profile_map[curr.fallback_provider_id]
|
|
if next_p.id in visited:
|
|
break # prevent infinite cycle
|
|
chain.append(next_p)
|
|
visited.add(next_p.id)
|
|
curr = next_p
|
|
return chain
|
|
except Exception as e:
|
|
logger.debug(f"Error fetching provider profiles from repo: {e}")
|
|
|
|
# Fallback to current memory/env configuration
|
|
return [
|
|
AIProviderProfile(
|
|
id=0,
|
|
name="Default Provider",
|
|
provider_type=self.provider,
|
|
model=self.model,
|
|
base_url=self.base_url,
|
|
api_key=self.api_key,
|
|
reasoning_effort=self.reasoning_effort,
|
|
is_active=True
|
|
)
|
|
]
|
|
|
|
async def generate_json(
|
|
self,
|
|
prompt: str,
|
|
system_prompt: Optional[str] = None,
|
|
action_name: str = "general",
|
|
image_path: Optional[str] = None
|
|
) -> Dict[str, Any]:
|
|
"""Send prompt (and optional image) to LLM, falling back across provider chain and notifying admins on failure."""
|
|
await self.sync_config_from_repo()
|
|
chain = await self.get_fallback_chain()
|
|
start_time = time.time()
|
|
status = "error"
|
|
last_error: Optional[Exception] = None
|
|
result_text: Optional[str] = None
|
|
successful_profile: Optional[AIProviderProfile] = None
|
|
|
|
try:
|
|
for i, profile in enumerate(chain):
|
|
p_type = profile.provider_type.strip().lower()
|
|
p_model = profile.model.strip()
|
|
p_base_url = profile.base_url.strip()
|
|
p_key = profile.api_key.strip()
|
|
p_effort = profile.reasoning_effort.strip().lower()
|
|
|
|
for attempt in range(self.max_retries_per_model + 1):
|
|
try:
|
|
if p_type == "gemini" and "orcarouter" not in p_base_url:
|
|
result = await self._call_gemini(
|
|
prompt=prompt,
|
|
system_prompt=system_prompt,
|
|
model=p_model,
|
|
api_key=p_key or self.api_key,
|
|
reasoning_effort=p_effort or self.reasoning_effort,
|
|
image_path=image_path
|
|
)
|
|
elif p_type == "agy":
|
|
if not os.path.exists("/.dockerenv") and (shutil.which("agy") or os.path.exists("/home/mamad/.local/bin/agy")):
|
|
result = await self._call_agy_cli(
|
|
prompt=prompt,
|
|
system_prompt=system_prompt,
|
|
effort=p_effort or self.reasoning_effort,
|
|
image_path=image_path
|
|
)
|
|
else:
|
|
result = await self._call_openai(
|
|
prompt=prompt,
|
|
system_prompt=system_prompt,
|
|
model=p_model,
|
|
base_url=p_base_url,
|
|
api_key=p_key,
|
|
reasoning_effort=p_effort or self.reasoning_effort,
|
|
is_agy=True,
|
|
image_path=image_path
|
|
)
|
|
else:
|
|
result = await self._call_openai(
|
|
prompt=prompt,
|
|
system_prompt=system_prompt,
|
|
model=p_model,
|
|
base_url=p_base_url,
|
|
api_key=p_key or self.api_key,
|
|
reasoning_effort=p_effort or self.reasoning_effort,
|
|
is_agy=False,
|
|
image_path=image_path
|
|
)
|
|
|
|
status = "success"
|
|
self.last_used_model = p_model
|
|
self.last_used_provider = p_type
|
|
successful_profile = profile
|
|
result_text = json.dumps(result, ensure_ascii=False) if isinstance(result, dict) else str(result)
|
|
return result
|
|
|
|
except Exception as e:
|
|
last_error = e
|
|
retryable = self._is_retryable(e)
|
|
logger.warning(
|
|
f"LLM call failed (provider={p_type}, model={p_model}, attempt={attempt + 1}/{self.max_retries_per_model + 1}, retryable={retryable}): {e}"
|
|
)
|
|
if not retryable or attempt == self.max_retries_per_model:
|
|
break
|
|
await asyncio.sleep(self.retry_backoff_seconds * (attempt + 1))
|
|
|
|
# Profile failed across all retries; trigger fallback alert if more providers exist
|
|
if i + 1 < len(chain):
|
|
next_profile = chain[i + 1]
|
|
logger.warning(
|
|
f"AI Provider '{profile.name}' ({p_type}/{p_model}) failed: {last_error}. Falling back to '{next_profile.name}' ({next_profile.provider_type}/{next_profile.model}) (step {i + 1})."
|
|
)
|
|
if self.on_fallback_alert:
|
|
try:
|
|
await self.on_fallback_alert(
|
|
profile,
|
|
next_profile,
|
|
str(last_error),
|
|
i + 1
|
|
)
|
|
except Exception as alert_err:
|
|
logger.error(f"Error executing AI fallback alert callback: {alert_err}")
|
|
|
|
# All providers exhausted
|
|
if self.on_chain_failure_alert and len(chain) > 1:
|
|
try:
|
|
await self.on_chain_failure_alert(
|
|
chain,
|
|
str(last_error)
|
|
)
|
|
except Exception as alert_err:
|
|
logger.error(f"Error executing AI chain failure alert callback: {alert_err}")
|
|
|
|
logger.error(f"LLM generation failed across all providers in chain: {last_error}")
|
|
raise last_error if last_error else RuntimeError("LLM generation failed across all providers in chain")
|
|
|
|
finally:
|
|
duration = time.time() - start_time
|
|
AI_LATENCY_SECONDS.labels(action=action_name).observe(duration)
|
|
AI_REQUESTS_TOTAL.labels(action=action_name, status=status).inc()
|
|
if self.repo:
|
|
try:
|
|
used_prov = successful_profile.provider_type if successful_profile else self.last_used_provider
|
|
used_mod = successful_profile.model if successful_profile else (self.last_used_model or self.model)
|
|
await self.repo.record_ai_log(
|
|
action_name=action_name,
|
|
provider=used_prov,
|
|
model=used_mod,
|
|
prompt=prompt,
|
|
system_prompt=system_prompt,
|
|
response_text=result_text,
|
|
duration_sec=duration,
|
|
status=status,
|
|
error_message=str(last_error) if last_error and status != "success" else None,
|
|
)
|
|
except Exception as log_err:
|
|
logger.error(f"Failed to record AI log in database: {log_err}", exc_info=True)
|
|
|
|
@staticmethod
|
|
def _is_retryable(error: Exception) -> bool:
|
|
"""Rate limits, upstream outages and transport errors are worth another shot."""
|
|
if isinstance(error, httpx.HTTPStatusError):
|
|
return error.response.status_code in (408, 409, 425, 429, 500, 502, 503, 504)
|
|
return isinstance(error, (httpx.TransportError, json.JSONDecodeError, KeyError, ValueError))
|
|
|
|
@staticmethod
|
|
def _parse_json_content(content: Optional[str]) -> Dict[str, Any]:
|
|
"""Parse a model reply that may be empty or wrapped in a markdown code fence."""
|
|
text = (content or "").strip()
|
|
if not text:
|
|
raise ValueError("model returned an empty response")
|
|
|
|
if text.startswith("```"):
|
|
text = text.split("\n", 1)[-1] if "\n" in text else text
|
|
text = text.rsplit("```", 1)[0].strip()
|
|
|
|
try:
|
|
return json.loads(text)
|
|
except json.JSONDecodeError:
|
|
# Some models prepend prose before the JSON object.
|
|
start, end = text.find("{"), text.rfind("}")
|
|
if start != -1 and end > start:
|
|
return json.loads(text[start:end + 1])
|
|
raise
|
|
|
|
async def _call_gemini(
|
|
self,
|
|
prompt: str,
|
|
system_prompt: Optional[str] = None,
|
|
model: Optional[str] = None,
|
|
api_key: Optional[str] = None,
|
|
reasoning_effort: Optional[str] = None,
|
|
image_path: Optional[str] = None,
|
|
) -> Dict[str, Any]:
|
|
model = model or self.model
|
|
key = (api_key or self.api_key or "").strip()
|
|
eff = (reasoning_effort or self.reasoning_effort or "").strip().lower()
|
|
url = f"https://generativelanguage.googleapis.com/v1beta/models/{model}:generateContent?key={key}"
|
|
gen_config: Dict[str, Any] = {
|
|
"responseMimeType": "application/json",
|
|
"temperature": 0.2,
|
|
}
|
|
if eff:
|
|
budget_map = {"low": 1024, "medium": 2048, "high": 4096}
|
|
budget = budget_map.get(eff, 2048)
|
|
gen_config["thinkingConfig"] = {"thinkingBudget": budget}
|
|
|
|
parts: List[Dict[str, Any]] = [{"text": prompt}]
|
|
if image_path and os.path.isfile(image_path):
|
|
try:
|
|
mime = "image/jpeg"
|
|
lower_p = image_path.lower()
|
|
if lower_p.endswith(".png"):
|
|
mime = "image/png"
|
|
elif lower_p.endswith(".webp"):
|
|
mime = "image/webp"
|
|
with open(image_path, "rb") as f:
|
|
b64 = base64.b64encode(f.read()).decode("utf-8")
|
|
parts.append({"inlineData": {"mimeType": mime, "data": b64}})
|
|
except Exception as img_err:
|
|
logger.warning(f"Failed to read image {image_path} for Gemini inlineData: {img_err}")
|
|
|
|
payload: Dict[str, Any] = {
|
|
"contents": [
|
|
{
|
|
"parts": parts
|
|
}
|
|
],
|
|
"generationConfig": gen_config
|
|
}
|
|
if system_prompt:
|
|
payload["systemInstruction"] = {
|
|
"parts": [{"text": system_prompt}]
|
|
}
|
|
|
|
async with httpx.AsyncClient(timeout=60.0) as client:
|
|
resp = await client.post(url, json=payload)
|
|
resp.raise_for_status()
|
|
data = resp.json()
|
|
raw_text = data["candidates"][0]["content"]["parts"][0]["text"]
|
|
return self._parse_json_content(raw_text)
|
|
|
|
def _resolve_openai_url(self, base_url: Optional[str], is_agy: bool = False) -> str:
|
|
base = (base_url or "").strip()
|
|
if not base:
|
|
if is_agy or self.provider == "agy":
|
|
base = "http://host.docker.internal:8088/v1" if os.path.exists("/.dockerenv") else "http://localhost:8088/v1"
|
|
else:
|
|
base = "https://api.orcarouter.ai/v1"
|
|
if os.path.exists("/.dockerenv"):
|
|
base = base.replace("localhost", "host.docker.internal").replace("127.0.0.1", "host.docker.internal")
|
|
base = base.rstrip("/")
|
|
return base if base.endswith("/chat/completions") else f"{base}/chat/completions"
|
|
|
|
async def _call_openai(
|
|
self,
|
|
prompt: str,
|
|
system_prompt: Optional[str] = None,
|
|
model: Optional[str] = None,
|
|
base_url: Optional[str] = None,
|
|
api_key: Optional[str] = None,
|
|
reasoning_effort: Optional[str] = None,
|
|
is_agy: bool = False,
|
|
image_path: Optional[str] = None,
|
|
) -> Dict[str, Any]:
|
|
model = model or self.model
|
|
url = self._resolve_openai_url(base_url or self.base_url, is_agy=is_agy)
|
|
eff = (reasoning_effort or self.reasoning_effort or "").strip().lower()
|
|
key = (api_key or self.api_key or "").strip()
|
|
|
|
headers = {
|
|
"Content-Type": "application/json"
|
|
}
|
|
if key:
|
|
headers["Authorization"] = f"Bearer {key}"
|
|
|
|
messages = []
|
|
if system_prompt:
|
|
messages.append({"role": "system", "content": system_prompt})
|
|
|
|
# Format user content: text + optional image
|
|
if image_path and os.path.isfile(image_path):
|
|
try:
|
|
mime = "image/jpeg"
|
|
lower_p = image_path.lower()
|
|
if lower_p.endswith(".png"):
|
|
mime = "image/png"
|
|
elif lower_p.endswith(".webp"):
|
|
mime = "image/webp"
|
|
with open(image_path, "rb") as f:
|
|
b64 = base64.b64encode(f.read()).decode("utf-8")
|
|
user_content: List[Dict[str, Any]] = [
|
|
{"type": "text", "text": prompt},
|
|
{"type": "image_url", "image_url": {"url": f"data:{mime};base64,{b64}"}}
|
|
]
|
|
messages.append({"role": "user", "content": user_content})
|
|
except Exception as img_err:
|
|
logger.warning(f"Failed to read image {image_path} for vision payload: {img_err}")
|
|
messages.append({"role": "user", "content": prompt})
|
|
else:
|
|
messages.append({"role": "user", "content": prompt})
|
|
|
|
payload: Dict[str, Any] = {
|
|
"model": model,
|
|
"messages": messages,
|
|
"response_format": {"type": "json_object"},
|
|
"temperature": 0.2,
|
|
}
|
|
if eff in ("low", "medium", "high"):
|
|
payload["reasoning_effort"] = eff
|
|
|
|
async with httpx.AsyncClient(timeout=60.0) as client:
|
|
try:
|
|
resp = await client.post(url, headers=headers, json=payload)
|
|
resp.raise_for_status()
|
|
except httpx.HTTPStatusError as e:
|
|
# Some local servers don't accept response_format: json_object
|
|
if e.response.status_code == 400 and "response_format" in payload:
|
|
payload.pop("response_format", None)
|
|
resp = await client.post(url, headers=headers, json=payload)
|
|
resp.raise_for_status()
|
|
else:
|
|
raise
|
|
|
|
data = resp.json()
|
|
content = data["choices"][0]["message"]["content"]
|
|
return self._parse_json_content(content)
|
|
|
|
async def _call_agy_cli(
|
|
self,
|
|
prompt: str,
|
|
system_prompt: Optional[str] = None,
|
|
effort: Optional[str] = None,
|
|
image_path: Optional[str] = None,
|
|
) -> Dict[str, Any]:
|
|
agy_bin = shutil.which("agy") or ("/usr/local/bin/agy" if os.path.exists("/usr/local/bin/agy") else ("/home/mamad/.local/bin/agy" if os.path.exists("/home/mamad/.local/bin/agy") else None))
|
|
if not agy_bin:
|
|
raise RuntimeError("agy binary not found in PATH or mounted paths")
|
|
full_prompt = f"System Instruction: {system_prompt}\n\nUser Prompt: {prompt}" if system_prompt else prompt
|
|
cmd = [agy_bin, "-p", full_prompt, "--output-format", "json"]
|
|
if image_path and os.path.isfile(image_path):
|
|
cmd.extend(["-f", image_path])
|
|
eff = (effort or self.reasoning_effort or "").strip().lower()
|
|
if eff in ("low", "medium", "high"):
|
|
cmd.extend(["--effort", eff])
|
|
env = os.environ.copy()
|
|
if not env.get("HOME"):
|
|
env["HOME"] = "/root"
|
|
proc = await asyncio.create_subprocess_exec(
|
|
*cmd,
|
|
stdout=asyncio.subprocess.PIPE,
|
|
stderr=asyncio.subprocess.PIPE,
|
|
env=env
|
|
)
|
|
stdout, stderr = await proc.communicate()
|
|
if proc.returncode != 0:
|
|
raise RuntimeError(f"agy execution failed (code {proc.returncode}): {stderr.decode('utf-8')}")
|
|
data = json.loads(stdout.decode("utf-8"))
|
|
raw_text = data.get("response", "")
|
|
return self._parse_json_content(raw_text)
|
|
|
|
async def _call_agy_sdk(self, prompt: str, system_prompt: Optional[str] = None) -> Dict[str, Any]:
|
|
if not HAS_AGY_SDK:
|
|
raise RuntimeError("google-antigravity SDK is not installed")
|
|
config = AgyLocalAgentConfig(system_instructions=system_prompt or "")
|
|
async with AgyAgent(config) as agent:
|
|
response = await agent.chat(prompt)
|
|
chunks = []
|
|
async for token in response:
|
|
chunks.append(token)
|
|
raw_text = "".join(chunks)
|
|
return self._parse_json_content(raw_text)
|
|
|
|
|
|
|