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() # Only pass image if this provider profile explicitly enables vision/image processing active_image = image_path if (profile.supports_vision and image_path and os.path.isfile(image_path)) else None 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=active_image ) 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=active_image ) 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=active_image ) 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=active_image ) 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)