feat(B-LLM): Zentraler LLM Client — llm_complete() + llm_embed() + Migration + Tests + Doku
Check Cross-Plugin Imports / check (push) Has been cancelled

B-LLM: llm_client.py um generische llm_complete() und llm_embed() erweitert
- Provider-Auswahl, API-Key-Auflösung, Error-Handling, Cost-Tracking
- Retry mit Exponential-Backoff für transient errors
- Timeout konfigurierbar
- Helper: get_api_credentials(), build_model(), _classify_error()

B-LLM-MIG: Alle 8 direkten litellm.acompletion() Calls auf llm_complete() umgestellt
- agent_runner.py, query_understanding.py (2x), ai_proactive (3x), ai_assistant (2x)
- 0 verbleibende direkte litellm.acompletion() Calls außerhalb llm_client.py

B-LLM-TEST: 39 Tests in test_llm_client.py — alle grün
- Mock mode, error handling, embed, helpers, backward compat

B-LLM-DOC: Plugin-Dev-Guide Kapitel 7 (LLM Integration) hinzugefügt
This commit is contained in:
Agent Zero
2026-08-13 16:22:05 +02:00
parent 3d8210637e
commit e3ca3b3d28
12 changed files with 1230 additions and 234 deletions
+427 -30
View File
@@ -6,19 +6,429 @@ tests to run without external API dependencies.
LiteLLM provides a unified interface to OpenAI, Anthropic, Google, Azure,
AWS Bedrock, Ollama, and many more providers.
Generic functions:
- ``llm_complete()`` — generic chat completion with retry, cost tracking
- ``llm_embed()`` — generic text embedding
- ``get_api_credentials()`` — centralised credential/provider lookup
- ``build_model()`` — centralised LiteLLM model string builder
"""
from __future__ import annotations
import asyncio
import json
import logging
import os
from typing import Any
import uuid
from typing import Any, TYPE_CHECKING
import litellm
if TYPE_CHECKING:
from sqlalchemy.ext.asyncio import AsyncSession
logger = logging.getLogger(__name__)
# ──────────────────────────────────────────────────────────────────────────
# Constants
# ──────────────────────────────────────────────────────────────────────────
MAX_INPUT_CHARS = 8000
# OpenRouter for embeddings (Ollama Cloud has no embedding endpoint)
OPENROUTER_API_KEY = os.environ.get("API_KEY_OPENROUTER", "")
OPENROUTER_BASE_URL = "https://openrouter.ai/api/v1"
OPENROUTER_EMBEDDING_MODEL = os.environ.get(
"SEARCH_EMBEDDING_MODEL", "openai/text-embedding-3-small"
)
EMBEDDING_DIMENSIONS = 768 # Must match DB column vector(768)
# Default retry settings
DEFAULT_TIMEOUT = 30
DEFAULT_MAX_RETRIES = 2
BASE_BACKOFF_SECONDS = 1.0
# Transient error keywords for retry classification
_TRANSIENT_KEYWORDS = frozenset(
{
"timeout",
"timed out",
"rate limit",
"rate_limit",
"429",
"503",
"502",
"504",
"service unavailable",
"overloaded",
"connection reset",
"connection aborted",
"temporary",
}
)
# Permanent error keywords — fail immediately, no retry
_PERMANENT_KEYWORDS = frozenset(
{
"authentication",
"auth",
"401",
"403",
"unauthorized",
"forbidden",
"invalid api key",
"invalid_api_key",
"validation",
"invalid_request",
"400",
"bad request",
"model_not_found",
"not found",
}
)
# ──────────────────────────────────────────────────────────────────────────
# Centralised helper functions
# ──────────────────────────────────────────────────────────────────────────
async def get_api_credentials(
db: AsyncSession | None,
tenant_id: uuid.UUID | None,
) -> tuple[str | None, str | None, str | None]:
"""Get API key, base_url and provider_type for LLM/embedding calls.
Priority:
1. OpenRouter env var (API_KEY_OPENROUTER) dedicated embedding provider
2. Default AI provider from DB (fallback)
3. API_KEY_OLLAMA_CLOUD env var (last resort)
Args:
db: Optional async DB session for provider lookup.
tenant_id: Optional tenant ID for provider lookup.
Returns:
Tuple of (api_key, api_base, provider_type) — any may be ``None``.
"""
# OpenRouter is the primary embedding provider
if OPENROUTER_API_KEY:
return OPENROUTER_API_KEY, OPENROUTER_BASE_URL, "openai"
# Fallback to DB provider
if db and tenant_id:
try:
from app.plugins.builtins.ai_assistant.contracts import get_default_provider
provider = await get_default_provider(db, tenant_id)
if provider and provider.api_key:
return provider.api_key, provider.base_url, provider.provider_type
except Exception:
logger.debug("Failed to get provider from DB, falling back to env")
env_key = os.environ.get("API_KEY_OLLAMA_CLOUD", "")
return (env_key if env_key else None), None, None
def build_model(model: str, provider_type: str | None) -> str:
"""Build LiteLLM model string with provider prefix.
If ``provider_type`` is given, strips any existing prefix from ``model``
and prepends ``provider_type``.
Args:
model: Model name, optionally already prefixed (e.g. ``openai/gpt-4o``).
provider_type: Provider prefix to apply (e.g. ``openai``, ``anthropic``).
Returns:
LiteLLM-compatible model string (e.g. ``openai/gpt-4o``).
"""
if provider_type:
model_parts = model.split("/", 1)
return f"{provider_type}/{model_parts[-1]}"
return model
def _classify_error(exc: Exception) -> str:
"""Classify an exception as ``transient`` or ``permanent``.
Uses string matching on the exception message/type name against known
patterns. Falls back to ``transient`` for unknown errors (safer to retry).
Args:
exc: The exception to classify.
Returns:
``"transient"`` or ``"permanent"``.
"""
msg = str(exc).lower()
exc_type_name = type(exc).__name__.lower()
# Check permanent first — auth errors should never be retried
if any(kw in msg or kw in exc_type_name for kw in _PERMANENT_KEYWORDS):
return "permanent"
if any(kw in msg or kw in exc_type_name for kw in _TRANSIENT_KEYWORDS):
return "transient"
# asyncio.TimeoutError is always transient
if isinstance(exc, (asyncio.TimeoutError, TimeoutError)):
return "transient"
# Default: treat as transient (safe to retry)
return "transient"
def _extract_cost_usd(response: Any, model: str) -> float:
"""Extract cost in USD from a LiteLLM response.
Uses ``litellm.completion_cost`` when available, otherwise returns 0.0.
Args:
response: LiteLLM response object.
model: Model string used for the call.
Returns:
Estimated cost in USD, or 0.0 if unavailable.
"""
try:
cost = litellm.completion_cost(response)
if cost is not None:
return float(cost)
except Exception:
logger.debug("litellm.completion_cost failed, using fallback")
return 0.0
def _extract_usage(response: Any) -> dict[str, int]:
"""Extract token usage from a LiteLLM response.
Args:
response: LiteLLM response object.
Returns:
Dict with ``prompt_tokens``, ``completion_tokens``, ``total_tokens``.
"""
usage = getattr(response, "usage", None)
if usage is None:
return {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}
prompt_tokens = getattr(usage, "prompt_tokens", 0) or 0
completion_tokens = getattr(usage, "completion_tokens", 0) or 0
total_tokens = getattr(usage, "total_tokens", 0) or (prompt_tokens + completion_tokens)
return {
"prompt_tokens": prompt_tokens,
"completion_tokens": completion_tokens,
"total_tokens": total_tokens,
}
# ──────────────────────────────────────────────────────────────────────────
# Generic LLM functions
# ──────────────────────────────────────────────────────────────────────────
async def llm_complete(
model: str,
messages: list[dict[str, Any]],
tools: list[dict[str, Any]] | None = None,
temperature: float = 0.3,
max_tokens: int = 1000,
api_key: str | None = None,
api_base: str | None = None,
provider: str | None = None,
response_format: dict[str, Any] | None = None,
timeout: int = DEFAULT_TIMEOUT,
max_retries: int = DEFAULT_MAX_RETRIES,
) -> dict[str, Any]:
"""Generic LLM chat completion via LiteLLM with retry and cost tracking.
Supports 100+ providers through LiteLLM's unified interface.
Transient errors (timeout, rate-limit) are retried with exponential
backoff. Permanent errors (auth, validation) fail immediately.
Args:
model: Model name (e.g. ``gpt-4o``, ``openai/gpt-4o``).
messages: Chat messages list (``[{"role": ..., "content": ...}]``).
tools: Optional list of tool/function definitions.
temperature: Sampling temperature (default 0.3).
max_tokens: Maximum tokens to generate (default 1000).
api_key: Override API key. If ``None``, uses env/DB lookup.
api_base: Override API base URL.
provider: Provider prefix (e.g. ``openai``, ``anthropic``).
response_format: Optional response format spec (e.g. JSON mode).
timeout: Request timeout in seconds (default 30).
max_retries: Max retry attempts for transient errors (default 2).
Returns:
Dict with keys: ``content``, ``usage``, ``cost_usd``, ``model``,
``raw_response`` (the LiteLLM response object for advanced use).
Raises:
Exception: Permanent errors or after exhausting retries.
"""
# Build LiteLLM model string
litellm_model = build_model(model, provider)
# Build kwargs
kwargs: dict[str, Any] = {
"model": litellm_model,
"messages": messages,
"temperature": temperature,
"max_tokens": max_tokens,
"timeout": timeout,
}
if api_key:
kwargs["api_key"] = api_key
if api_base:
kwargs["api_base"] = api_base
if tools:
kwargs["tools"] = tools
if response_format:
kwargs["response_format"] = response_format
last_exc: Exception | None = None
for attempt in range(max_retries + 1):
try:
response = await litellm.acompletion(**kwargs)
content = response.choices[0].message.content or ""
usage = _extract_usage(response)
cost_usd = _extract_cost_usd(response, litellm_model)
logger.debug(
"llm_complete success: model=%s tokens=%d cost=$%.6f attempt=%d",
litellm_model,
usage["total_tokens"],
cost_usd,
attempt + 1,
)
return {
"content": content,
"usage": usage,
"cost_usd": cost_usd,
"model": litellm_model,
"raw_response": response,
}
except Exception as exc:
last_exc = exc
error_class = _classify_error(exc)
if error_class == "permanent" or attempt >= max_retries:
logger.error(
"llm_complete failed (permanent/exhausted): model=%s attempt=%d error=%s",
litellm_model,
attempt + 1,
exc,
)
raise
# Transient error — retry with exponential backoff
backoff = BASE_BACKOFF_SECONDS * (2**attempt)
logger.warning(
"llm_complete transient error (attempt %d/%d), retrying in %.1fs: %s",
attempt + 1,
max_retries + 1,
backoff,
exc,
)
await asyncio.sleep(backoff)
# Should not reach here, but satisfy type checker
assert last_exc is not None
raise last_exc
async def llm_embed(
texts: str | list[str],
model: str | None = None,
db: AsyncSession | None = None,
tenant_id: uuid.UUID | None = None,
api_key: str | None = None,
api_base: str | None = None,
provider: str | None = None,
dimensions: int | None = None,
timeout: int = DEFAULT_TIMEOUT,
) -> list[list[float]]:
"""Generic text embedding via LiteLLM.
Handles both single-text and batch embedding. Uses centralised
credential lookup when ``api_key`` is not provided.
Args:
texts: Single text string or list of texts to embed.
model: Embedding model name (default: ``openai/text-embedding-3-small``).
db: Optional DB session for API key lookup.
tenant_id: Optional tenant ID for API key lookup.
api_key: Override API key. If ``None``, uses env/DB lookup.
api_base: Override API base URL.
provider: Provider prefix override.
dimensions: Override embedding dimensions.
timeout: Request timeout in seconds (default 30).
Returns:
List of embedding vectors (each a list of floats). For a single
text input, returns a one-element list.
"""
# Normalise to list input
single_input = isinstance(texts, str)
text_list = [texts] if single_input else texts
if not text_list:
return []
# Truncate inputs
truncated = [t[:MAX_INPUT_CHARS] for t in text_list]
# Resolve model
embedding_model = model or OPENROUTER_EMBEDDING_MODEL
# Resolve credentials
if not api_key:
resolved_key, resolved_base, resolved_provider = await get_api_credentials(
db, tenant_id
)
api_key = resolved_key
if not api_base:
api_base = resolved_base
if not provider:
provider = resolved_provider
# Build LiteLLM model string
litellm_model = build_model(embedding_model, provider)
# Build kwargs
litellm_kwargs: dict[str, Any] = {
"model": litellm_model,
"input": truncated[0] if single_input else truncated,
"timeout": timeout,
}
if api_key:
litellm_kwargs["api_key"] = api_key
if api_base:
litellm_kwargs["api_base"] = api_base
# Request specific dimensions for text-embedding-3 models
effective_dims = dimensions or EMBEDDING_DIMENSIONS
if "text-embedding-3" in litellm_model:
litellm_kwargs["dimensions"] = effective_dims
try:
response = await litellm.aembedding(**litellm_kwargs)
embeddings = [d["embedding"] for d in response.data]
logger.debug(
"llm_embed success: model=%s count=%d dims=%d",
litellm_model,
len(embeddings),
len(embeddings[0]) if embeddings else 0,
)
return embeddings
except Exception:
logger.warning("llm_embed failed: model=%s", litellm_model, exc_info=True)
return [[] for _ in text_list]
# ──────────────────────────────────────────────────────────────────────────
# LLMClient class (AI Copilot — backward compatible)
# ──────────────────────────────────────────────────────────────────────────
class LLMResponse:
"""Structured LLM response containing proposed actions."""
@@ -90,7 +500,7 @@ class LLMClient:
)
async def _api_generate(self, query: str, context: dict[str, Any]) -> LLMResponse:
"""Call LLM via LiteLLM unified interface.
"""Call LLM via ``llm_complete()`` (delegates to LiteLLM).
Supports 100+ providers through a single API:
- OpenAI: "openai/gpt-4o"
@@ -103,37 +513,24 @@ class LLMClient:
system_prompt = self._build_system_prompt(context)
user_prompt = f"User request: {query}\n\nRespond with proposed actions as JSON."
# Build LiteLLM model string: "provider/model" or just "model" for OpenAI compat
if self.provider and self.provider != "openai":
litellm_model = f"{self.provider}/{self.model}"
else:
litellm_model = self.model
# Build kwargs for litellm.acompletion
kwargs: dict[str, Any] = {
"model": litellm_model,
"messages": [
{"role": "system", "content": system_prompt},
{"role": "user", "content": user_prompt},
],
"temperature": 0.3,
"max_tokens": 1000,
}
# Add API key if set
if self.api_key:
kwargs["api_key"] = self.api_key
# Add API base if set (for self-hosted or custom endpoints)
if self.api_base:
kwargs["api_base"] = self.api_base
messages = [
{"role": "system", "content": system_prompt},
{"role": "user", "content": user_prompt},
]
try:
response = await litellm.acompletion(**kwargs)
content = response.choices[0].message.content
return self._parse_llm_response(content)
result = await llm_complete(
model=self.model,
messages=messages,
temperature=0.3,
max_tokens=1000,
api_key=self.api_key or None,
api_base=self.api_base or None,
provider=self.provider,
)
return self._parse_llm_response(result["content"])
except Exception as e:
logger.error("LiteLLM API call failed: %s", e)
logger.error("LLM API call failed: %s", e)
# Fall back to mock mode on API error
return LLMResponse(
message=f"LLM API call failed: {e}. Falling back to keyword matching.",