Agent Zero
abbe7a18fc
- P0: hooks.py 3-tuple fix, trigger_dispatcher Contract, contacts/plugin unregister_actions_by_owner - P0: 5 test files — check_permission mocks removed, hardcoded DB credential → env var - P1: attachment_service DmsFile via Contract helper, restore_registry/history_hooks dedup - P1: mail/plugin restore unregister, mcp_client datetime.now(UTC), saved_views/filters patterns - P1: ProtectedRoute fail-closed, 13 test assertion fixes (bcrypt, DB-URLs, SECRET_KEYs) - P2: deprecated notifications → post_system_message (3 files), forgejo Base, report_generator lazy import - P2: webhooks permissions, deps.py/roles.py plugin perms removed, import_export default - P2: address/tags/entity_links patterns removed, worker.py Contract-Umgehungen fixed - P2: 28 frontend TODOs (hardcoded constants, deprecated notification API) - P3: dead code, duplicates, deprecated imports, private attr, __import__ inline - P3: 8 frontend TODOs (LucideIcons, inline styles, XSS, i18n) - ruff: 838 → 0 (612 auto-fix + 246 manual + 27 F821 regression fix) - F821: 30 → 0 (AutomationDefinition, DmsFile, user_id, Path, Any, String) - Contract-Umgehungen: 2 neue gefunden (worker.py:169, worker.py:280) und gefixt
874 lines
32 KiB
Python
874 lines
32 KiB
Python
"""Configurable LLM client — supports LiteLLM (100+ providers) or mock/stub mode.
|
||
|
||
Reads AI_MODEL, AI_API_KEY, AI_PROVIDER from environment. If not set, uses mock mode
|
||
which returns predefined actions based on keyword matching. This allows
|
||
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
|
||
import uuid
|
||
from datetime import UTC, datetime
|
||
from typing import TYPE_CHECKING, Any
|
||
|
||
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
|
||
# ──────────────────────────────────────────────────────────────────────────
|
||
|
||
|
||
# ── Cost Overrun Protection (B.17) ───────────────────────────────────────────
|
||
|
||
|
||
def _get_cost_key(tenant_id: uuid.UUID | str) -> str:
|
||
"""Build the Redis cost-tracking key for the current month."""
|
||
now = datetime.now(UTC)
|
||
month_str = now.strftime("%Y-%m")
|
||
return f"cost:tenant:{tenant_id}:month:{month_str}"
|
||
|
||
|
||
async def get_tenant_monthly_cost(tenant_id: uuid.UUID | str) -> float:
|
||
"""Get the accumulated LLM cost for a tenant in the current month.
|
||
|
||
Reads from Redis key ``cost:tenant:{tenant_id}:month:{YYYY-MM}``.
|
||
Returns 0.0 if Redis is unavailable or the key doesn't exist.
|
||
"""
|
||
try:
|
||
from app.core.auth import get_redis
|
||
r = get_redis()
|
||
key = _get_cost_key(tenant_id)
|
||
val = await r.get(key)
|
||
return float(val) if val else 0.0
|
||
except Exception:
|
||
logger.debug("Failed to get tenant monthly cost from Redis")
|
||
return 0.0
|
||
|
||
|
||
async def _check_tenant_budget(
|
||
tenant_id: uuid.UUID | str | None,
|
||
db: AsyncSession | None = None,
|
||
) -> None:
|
||
"""Check if tenant has exceeded their LLM budget.
|
||
|
||
Raises ``ValueError("Tenant LLM budget exceeded")`` if the tenant's
|
||
accumulated monthly cost exceeds the configured budget and
|
||
``llm_hard_cutoff`` is enabled.
|
||
"""
|
||
if tenant_id is None:
|
||
return
|
||
|
||
from app.config import get_settings
|
||
settings = get_settings()
|
||
|
||
if settings.llm_monthly_budget_usd <= 0:
|
||
return # No budget limit configured
|
||
|
||
current_cost = await get_tenant_monthly_cost(tenant_id)
|
||
if current_cost >= settings.llm_monthly_budget_usd:
|
||
if settings.llm_hard_cutoff:
|
||
logger.warning(
|
||
"Tenant %s LLM budget exceeded: $%.2f >= $%.2f (hard cutoff)",
|
||
tenant_id, current_cost, settings.llm_monthly_budget_usd,
|
||
)
|
||
raise ValueError("Tenant LLM budget exceeded")
|
||
else:
|
||
logger.warning(
|
||
"Tenant %s LLM budget exceeded: $%.2f >= $%.2f (soft limit, no cutoff)",
|
||
tenant_id, current_cost, settings.llm_monthly_budget_usd,
|
||
)
|
||
|
||
|
||
async def _track_tenant_cost(
|
||
tenant_id: uuid.UUID | str | None,
|
||
cost_usd: float,
|
||
db: AsyncSession | None = None,
|
||
) -> float:
|
||
"""Track LLM cost in Redis and check for alert thresholds.
|
||
|
||
Increments ``cost:tenant:{tenant_id}:month:{YYYY-MM}`` by ``cost_usd``.
|
||
Returns the new total cost for the month.
|
||
"""
|
||
if tenant_id is None or cost_usd <= 0:
|
||
return 0.0
|
||
|
||
try:
|
||
from app.core.auth import get_redis
|
||
r = get_redis()
|
||
key = _get_cost_key(tenant_id)
|
||
new_total = await r.incrbyfloat(key, cost_usd)
|
||
# Set TTL to 35 days so old months auto-expire
|
||
await r.expire(key, 35 * 24 * 3600)
|
||
|
||
# Check cost alert thresholds
|
||
await _check_cost_alerts(tenant_id, new_total, db)
|
||
|
||
return float(new_total)
|
||
except Exception:
|
||
logger.debug("Failed to track tenant cost in Redis")
|
||
return 0.0
|
||
|
||
|
||
async def _check_cost_alerts(
|
||
tenant_id: uuid.UUID | str,
|
||
current_cost: float,
|
||
db: AsyncSession | None = None,
|
||
) -> None:
|
||
"""Send cost alerts at 50%, 80%, and 100% of the tenant budget.
|
||
|
||
Each threshold alert is sent only once per month (tracked via Redis flag
|
||
``cost:alerted:{tenant_id}:{threshold}``).
|
||
"""
|
||
from app.config import get_settings
|
||
settings = get_settings()
|
||
|
||
budget = settings.llm_monthly_budget_usd
|
||
if budget <= 0:
|
||
return
|
||
|
||
thresholds = [0.50, 0.80, 1.00]
|
||
now = datetime.now(UTC)
|
||
month_str = now.strftime("%Y-%m")
|
||
|
||
try:
|
||
from app.core.auth import get_redis
|
||
r = get_redis()
|
||
|
||
for threshold in thresholds:
|
||
threshold_cost = budget * threshold
|
||
if current_cost >= threshold_cost:
|
||
alert_key = f"cost:alerted:{tenant_id}:{threshold}:{month_str}"
|
||
already_alerted = await r.set(alert_key, "1", nx=True, ex=35 * 24 * 3600)
|
||
if already_alerted:
|
||
# This is a new alert — send notification
|
||
pct = int(threshold * 100)
|
||
logger.info(
|
||
"Cost alert: tenant %s reached %d%% of budget ($%.2f / $%.2f)",
|
||
tenant_id, pct, current_cost, budget,
|
||
)
|
||
# Best-effort system notification
|
||
if db is not None:
|
||
try:
|
||
# Need a user_id — try to find an admin for this tenant
|
||
from sqlalchemy import select as sa_select
|
||
|
||
from app.core.notifications import post_system_message
|
||
from app.models.user import User, UserTenant
|
||
|
||
async with db.begin_nested() if db.in_transaction() else _NoopCtx():
|
||
result = await db.execute(
|
||
sa_select(User.id).join(UserTenant, UserTenant.user_id == User.id)
|
||
.where(UserTenant.tenant_id == tenant_id)
|
||
.where(User.is_active == True) # noqa: E712
|
||
.limit(1)
|
||
)
|
||
admin_row = result.first()
|
||
if admin_row:
|
||
await post_system_message(
|
||
db=db,
|
||
tenant_id=tenant_id if isinstance(tenant_id, uuid.UUID) else uuid.UUID(str(tenant_id)),
|
||
user_id=admin_row[0],
|
||
message_type="cost_alert",
|
||
title=f"LLM Cost Alert: {pct}% of budget reached",
|
||
body=f"Current monthly LLM cost: ${current_cost:.2f} / ${budget:.2f} ({pct}%).",
|
||
severity="warning" if pct < 100 else "error",
|
||
)
|
||
except Exception:
|
||
logger.debug("Failed to send cost alert notification")
|
||
except Exception:
|
||
logger.debug("Failed to check cost alerts")
|
||
|
||
|
||
class _NoopCtx:
|
||
"""No-op async context manager for optional transaction nesting."""
|
||
async def __aenter__(self):
|
||
return self
|
||
async def __aexit__(self, *args):
|
||
return False
|
||
|
||
|
||
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.contracts import get_contract
|
||
ai_contract = get_contract("ai_assistant")
|
||
if ai_contract is not None:
|
||
provider = await ai_contract.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
|
||
|
||
|
||
async def get_provider_compliance(
|
||
db: AsyncSession | None,
|
||
tenant_id: uuid.UUID | None,
|
||
) -> dict[str, Any] | None:
|
||
"""Get compliance metadata for the active AI provider.
|
||
|
||
Returns a dict with keys: ``region``, ``hosting_type``, ``dpa_status``,
|
||
``retention_policy``, ``training_on_customer_data``, ``transfer_notice``,
|
||
``allowed_data_classes``.
|
||
|
||
Returns ``None`` if no DB provider is configured (env-based fallback).
|
||
"""
|
||
if not (db and tenant_id):
|
||
return None
|
||
try:
|
||
from app.plugins.builtins.contracts import get_contract
|
||
ai_contract = get_contract("ai_assistant")
|
||
if ai_contract is None:
|
||
return None
|
||
provider = await ai_contract.get_default_provider(db, tenant_id)
|
||
if provider is None:
|
||
return None
|
||
return {
|
||
"region": getattr(provider, "region", "unknown"),
|
||
"hosting_type": getattr(provider, "hosting_type", "cloud"),
|
||
"dpa_status": getattr(provider, "dpa_status", "none"),
|
||
"retention_policy": getattr(provider, "retention_policy", ""),
|
||
"training_on_customer_data": getattr(provider, "training_on_customer_data", False),
|
||
"transfer_notice": getattr(provider, "transfer_notice", ""),
|
||
"allowed_data_classes": getattr(provider, "allowed_data_classes", []),
|
||
}
|
||
except Exception:
|
||
logger.debug("Failed to get provider compliance metadata")
|
||
return None
|
||
|
||
|
||
def check_data_class_allowed(
|
||
compliance: dict[str, Any] | None,
|
||
data_class: str,
|
||
) -> bool:
|
||
"""Check whether the configured provider may process *data_class*.
|
||
|
||
Uses :func:`app.core.sensitive_data.check_provider_compliance`.
|
||
Returns ``True`` if compliance metadata is unavailable (fail-open for
|
||
backward compatibility and mock mode).
|
||
"""
|
||
from app.core.sensitive_data import check_provider_compliance
|
||
|
||
if compliance is None:
|
||
return True
|
||
return check_provider_compliance(
|
||
compliance.get("allowed_data_classes"),
|
||
data_class,
|
||
)
|
||
|
||
|
||
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( # noqa: ASYNC109
|
||
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, # noqa: ASYNC109
|
||
max_retries: int = DEFAULT_MAX_RETRIES,
|
||
trace_id: str | None = None,
|
||
tenant_id: uuid.UUID | str | None = None,
|
||
db: AsyncSession | None = None,
|
||
) -> 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).
|
||
trace_id: Optional trace ID for request correlation/logging.
|
||
tenant_id: Optional tenant ID for cost tracking and budget enforcement.
|
||
db: Optional DB session for cost alert notifications.
|
||
|
||
Returns:
|
||
Dict with keys: ``content``, ``usage``, ``cost_usd``, ``model``,
|
||
``raw_response`` (the LiteLLM response object for advanced use).
|
||
|
||
Raises:
|
||
ValueError: If tenant LLM budget is exceeded (hard cutoff).
|
||
Exception: Permanent errors or after exhausting retries.
|
||
"""
|
||
# Check tenant budget before making the call
|
||
await _check_tenant_budget(tenant_id, db)
|
||
|
||
# Build LiteLLM model string
|
||
litellm_model = build_model(model, provider)
|
||
|
||
# Log trace_id correlation if provided
|
||
if trace_id:
|
||
logger.debug("llm_complete trace_id=%s model=%s", trace_id, litellm_model)
|
||
|
||
# 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 trace_id=%s",
|
||
litellm_model,
|
||
usage["total_tokens"],
|
||
cost_usd,
|
||
attempt + 1,
|
||
trace_id or "-",
|
||
)
|
||
# Track cost in Redis for tenant budget enforcement
|
||
if tenant_id is not None and cost_usd > 0:
|
||
await _track_tenant_cost(tenant_id, cost_usd, db)
|
||
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( # noqa: ASYNC109
|
||
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, # noqa: ASYNC109
|
||
trace_id: str | None = None,
|
||
) -> 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 and cost alert notifications.
|
||
tenant_id: Optional tenant ID for cost tracking and budget enforcement.
|
||
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).
|
||
trace_id: Optional trace ID for request correlation/logging.
|
||
|
||
Returns:
|
||
List of embedding vectors (each a list of floats). For a single
|
||
text input, returns a one-element list.
|
||
"""
|
||
# Check tenant budget before making the call
|
||
await _check_tenant_budget(tenant_id, db)
|
||
|
||
if trace_id:
|
||
logger.debug("llm_embed trace_id=%s", trace_id)
|
||
|
||
# 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]
|
||
# Track embedding cost for tenant budget enforcement
|
||
if tenant_id is not None:
|
||
try:
|
||
emb_cost = _extract_cost_usd(response, litellm_model)
|
||
if emb_cost > 0:
|
||
await _track_tenant_cost(tenant_id, emb_cost, db)
|
||
except Exception:
|
||
pass
|
||
logger.debug(
|
||
"llm_embed success: model=%s count=%d dims=%d trace_id=%s",
|
||
litellm_model,
|
||
len(embeddings),
|
||
len(embeddings[0]) if embeddings else 0,
|
||
trace_id or "-",
|
||
)
|
||
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."""
|
||
|
||
def __init__(
|
||
self, message: str, proposed_actions: list[dict[str, Any]], confidence: float = 0.8
|
||
):
|
||
self.message = message
|
||
self.proposed_actions = proposed_actions
|
||
self.confidence = confidence
|
||
|
||
def to_dict(self) -> dict[str, Any]:
|
||
return {
|
||
"message": self.message,
|
||
"proposed_actions": self.proposed_actions,
|
||
"confidence": self.confidence,
|
||
}
|
||
|
||
|
||
class LLMClient:
|
||
"""LLM client that translates natural language to proposed API actions.
|
||
|
||
Modes:
|
||
- If AI_MODEL and AI_API_KEY are set: calls LiteLLM chat completions API
|
||
- Otherwise: mock/stub mode with keyword-based action mapping
|
||
|
||
LiteLLM model format: "provider/model_name" (e.g. "openai/gpt-4o", "anthropic/claude-3-sonnet", "ollama/llama3")
|
||
"""
|
||
|
||
def __init__(
|
||
self, model: str | None = None, api_key: str | None = None, api_base: str | None = None,
|
||
provider: str | None = None,
|
||
) -> None:
|
||
self.model = model or os.environ.get("AI_MODEL", "")
|
||
self.api_key = api_key or os.environ.get("AI_API_KEY", "")
|
||
self.api_base = api_base or os.environ.get("AI_API_BASE", "")
|
||
self.provider = provider or os.environ.get("AI_PROVIDER", "openai")
|
||
self.is_mock = not bool(self.model and self.api_key)
|
||
|
||
async def generate(self, user_query: str, context: dict[str, Any] | None = None) -> LLMResponse:
|
||
"""Generate proposed actions from natural language query.
|
||
|
||
Args:
|
||
user_query: Natural language input from user
|
||
context: Optional context (e.g. current page, selected entity)
|
||
|
||
Returns:
|
||
LLMResponse with message and proposed_actions list
|
||
"""
|
||
if self.is_mock:
|
||
return await self._mock_generate(user_query, context or {})
|
||
return await self._api_generate(user_query, context or {})
|
||
|
||
async def _mock_generate(self, query: str, context: dict[str, Any]) -> LLMResponse:
|
||
"""Mock/stub mode — keyword-based action mapping for tests."""
|
||
from app.ai.action_mapper import map_query_to_actions
|
||
|
||
actions = map_query_to_actions(query, context)
|
||
if actions:
|
||
return LLMResponse(
|
||
message=f"I found {len(actions)} possible action(s) based on your request.",
|
||
proposed_actions=actions,
|
||
confidence=0.85,
|
||
)
|
||
return LLMResponse(
|
||
message="I couldn't determine a specific action from your request. Could you be more specific?",
|
||
proposed_actions=[],
|
||
confidence=0.3,
|
||
)
|
||
|
||
async def _api_generate(self, query: str, context: dict[str, Any]) -> LLMResponse:
|
||
"""Call LLM via ``llm_complete()`` (delegates to LiteLLM).
|
||
|
||
Supports 100+ providers through a single API:
|
||
- OpenAI: "openai/gpt-4o"
|
||
- Anthropic: "anthropic/claude-3-sonnet"
|
||
- Google: "gemini/gemini-pro"
|
||
- Azure: "azure/<deployment-name>"
|
||
- Ollama: "ollama/llama3"
|
||
- And many more.
|
||
"""
|
||
system_prompt = self._build_system_prompt(context)
|
||
user_prompt = f"User request: {query}\n\nRespond with proposed actions as JSON."
|
||
|
||
messages = [
|
||
{"role": "system", "content": system_prompt},
|
||
{"role": "user", "content": user_prompt},
|
||
]
|
||
|
||
try:
|
||
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("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.",
|
||
proposed_actions=[],
|
||
confidence=0.1,
|
||
)
|
||
|
||
def _build_system_prompt(self, context: dict[str, Any]) -> str:
|
||
"""Build system prompt describing available API actions."""
|
||
available_apis = [
|
||
{"method": "GET", "path": "/api/v1/contacts", "description": "List contacts (persons and companies)"},
|
||
{"method": "POST", "path": "/api/v1/contacts", "description": "Create a contact (person or company)"},
|
||
{
|
||
"method": "GET",
|
||
"path": "/api/v1/contacts/{id}",
|
||
"description": "Get contact details",
|
||
},
|
||
{
|
||
"method": "PATCH",
|
||
"path": "/api/v1/contacts/{id}",
|
||
"description": "Update a contact",
|
||
},
|
||
{
|
||
"method": "DELETE",
|
||
"path": "/api/v1/contacts/{id}",
|
||
"description": "Delete a contact",
|
||
},
|
||
{"method": "GET", "path": "/api/v1/workflows", "description": "List workflows"},
|
||
{"method": "POST", "path": "/api/v1/workflows", "description": "Create a workflow"},
|
||
{"method": "GET", "path": "/api/v1/calendar/entries", "description": "List calendar entries"},
|
||
{"method": "POST", "path": "/api/v1/calendar/entries", "description": "Create a calendar entry"},
|
||
{"method": "GET", "path": "/api/v1/dms/files", "description": "List DMS files"},
|
||
]
|
||
context_str = json.dumps(context) if context else "{}"
|
||
return (
|
||
"You are an AI copilot for LeoCRM. Based on the user's natural language request, "
|
||
"propose one or more API actions. Always respond with a JSON object containing: "
|
||
'"message": a human-readable summary, '
|
||
'"proposed_actions": an array of {method, path, body, description, confidence}. '
|
||
f"Available API endpoints: {json.dumps(available_apis)}. "
|
||
f"Current context: {context_str}. "
|
||
"Never execute actions directly — only propose them for user confirmation."
|
||
)
|
||
|
||
def _parse_llm_response(self, content: str) -> LLMResponse:
|
||
"""Parse LLM JSON response into LLMResponse."""
|
||
try:
|
||
parsed = json.loads(content)
|
||
return LLMResponse(
|
||
message=parsed.get("message", "Here are the proposed actions."),
|
||
proposed_actions=parsed.get("proposed_actions", []),
|
||
confidence=parsed.get("confidence", 0.8),
|
||
)
|
||
except (json.JSONDecodeError, KeyError):
|
||
logger.warning("Failed to parse LLM response as JSON: %s", content[:200])
|
||
return LLMResponse(
|
||
message=content,
|
||
proposed_actions=[],
|
||
confidence=0.3,
|
||
)
|
||
|
||
|
||
# Global client instance
|
||
_client: LLMClient | None = None
|
||
|
||
|
||
def get_llm_client() -> LLMClient:
|
||
"""Get or create the global LLM client instance."""
|
||
global _client
|
||
if _client is None:
|
||
_client = LLMClient()
|
||
return _client
|
||
|
||
|
||
def reset_llm_client() -> None:
|
||
"""Reset the global client (for testing)."""
|
||
global _client
|
||
_client = None
|