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
@@ -1,7 +1,7 @@
"""AI Participant Handler — bridges the kommunikation plugin with the AI Assistant.
When a message is received in a conversation that includes the 'ai' participant,
this handler generates an LLM response using litellm.acompletion (non-streaming)
this handler generates an LLM response using llm_complete (non-streaming)
and returns it as a new message in the conversation.
"""
@@ -11,8 +11,7 @@ import logging
import uuid
from typing import Any
import litellm
from app.ai.llm_client import llm_complete
from app.core.db import create_db_session
from app.plugins.builtins.kommunikation.contracts import ParticipantHandler
@@ -136,11 +135,18 @@ class AIParticipantHandler(ParticipantHandler):
if provider.base_url:
params["api_base"] = provider.base_url
# Ensure non-streaming for acompletion
params["stream"] = False
# Ensure non-streaming
params.pop("stream", None)
response = await litellm.acompletion(**params)
response_text = response.choices[0].message.content or ""
result = await llm_complete(
model=params.get("model", "gpt-4o-mini"),
messages=params.get("messages", []),
temperature=params.get("temperature", 0.7),
max_tokens=params.get("max_tokens", 2048),
api_key=params.get("api_key"),
api_base=params.get("api_base"),
)
response_text = result["content"] or ""
if not response_text.strip():
response_text = "*(keine Antwort generiert)*"
+27 -20
View File
@@ -18,6 +18,8 @@ import litellm
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession
from app.ai.llm_client import llm_complete
from app.core.permissions import check_permission
from app.plugins.builtins.ai_assistant.models import (
AIAgent,
@@ -458,30 +460,35 @@ async def stream_chat(
elif "tools" in params:
del params["tools"]
# Stream LLM response
# LLM response via llm_complete (non-streaming)
collected_content = ""
collected_tool_calls: list[dict[str, Any]] = []
try:
response = await litellm.acompletion(**params)
async for chunk in response:
delta = chunk.choices[0].delta
if delta.content:
collected_content += delta.content
yield f"data: {json.dumps({'type': 'token', 'content': delta.content})}\n\n"
if delta.tool_calls:
for tc in delta.tool_calls:
idx = tc.index
while len(collected_tool_calls) <= idx:
collected_tool_calls.append({"id": "", "function": {"name": "", "arguments": ""}})
if tc.id:
collected_tool_calls[idx]["id"] = tc.id
if tc.function:
if tc.function.name:
collected_tool_calls[idx]["function"]["name"] += tc.function.name
if tc.function.arguments:
collected_tool_calls[idx]["function"]["arguments"] += tc.function.arguments
result = await llm_complete(
model=params.get("model", "gpt-4o-mini"),
messages=params.get("messages", []),
temperature=params.get("temperature", 0.7),
max_tokens=params.get("max_tokens", 2048),
api_key=params.get("api_key"),
api_base=params.get("api_base"),
tools=params.get("tools"),
)
collected_content = result["content"]
if collected_content:
yield f"data: {json.dumps({'type': 'token', 'content': collected_content})}\n\n"
# Extract tool calls from raw response
raw_response = result["raw_response"]
if hasattr(raw_response.choices[0].message, "tool_calls") and raw_response.choices[0].message.tool_calls:
for tc in raw_response.choices[0].message.tool_calls:
collected_tool_calls.append({
"id": tc.id or "",
"function": {
"name": tc.function.name if tc.function else "",
"arguments": tc.function.arguments if tc.function and tc.function.arguments else "",
},
})
except Exception as exc:
logger.error("LLM streaming error: %s", exc)
logger.error("LLM error: %s", exc)
yield f"data: {json.dumps({'type': 'error', 'content': str(exc)})}\n\n"
await save_message(db, session.id, "assistant", f"Error: {exc}", tenant_id, model_used=model_id)
await db.commit()
@@ -117,8 +117,6 @@ async def search_related_handler(arguments: dict[str, Any], context: dict[str, A
async def summarize_mail_thread_handler(arguments: dict[str, Any], context: dict[str, Any]) -> str:
"""Summarize a mail thread using LLM."""
try:
import litellm
db, tenant_id, _ = await _get_db_and_tenant(context)
thread_id = arguments["thread_id"]
limit = arguments.get("limit", 20)
@@ -139,7 +137,9 @@ async def summarize_mail_thread_handler(arguments: dict[str, Any], context: dict
for m in mails
)
response = await litellm.acompletion(
from app.ai.llm_client import llm_complete
result = await llm_complete(
model="gpt-4o-mini",
messages=[
{
@@ -151,7 +151,7 @@ async def summarize_mail_thread_handler(arguments: dict[str, Any], context: dict
temperature=0.3,
max_tokens=300,
)
summary = response.choices[0].message.content or ""
summary = result["content"] or ""
return json.dumps({"summary": summary, "count": len(mails)}, default=str)
except Exception as e:
logger.exception("summarize_mail_thread_handler failed")
+19 -23
View File
@@ -16,7 +16,7 @@ import logging
import uuid
from typing import Any
import litellm
from app.ai.llm_client import llm_complete
from sqlalchemy import select
from app.core.db import create_db_session
@@ -201,29 +201,25 @@ async def deep_analysis(
model_parts = model.split("/", 1)
model = f"{provider_type}/{model_parts[-1]}"
litellm_kwargs: dict[str, Any] = dict(
model=model,
messages=[
{"role": "system", "content": DEEP_SYSTEM_PROMPT},
{
"role": "user",
"content": json.dumps(
extended_context, default=str, ensure_ascii=False
),
},
],
temperature=0.3,
max_tokens=1500,
response_format={"type": "json_object"},
)
if api_key:
litellm_kwargs["api_key"] = api_key
if api_base:
litellm_kwargs["api_base"] = api_base
try:
response = await litellm.acompletion(**litellm_kwargs)
content = response.choices[0].message.content
result = await llm_complete(
model=model,
messages=[
{"role": "system", "content": DEEP_SYSTEM_PROMPT},
{
"role": "user",
"content": json.dumps(
extended_context, default=str, ensure_ascii=False
),
},
],
temperature=0.3,
max_tokens=1500,
response_format={"type": "json_object"},
api_key=api_key,
api_base=api_base,
)
content = result["content"]
if not content:
return
# Strip markdown code fences if present
+17 -20
View File
@@ -15,6 +15,7 @@ from datetime import UTC, datetime, timedelta
from typing import Any
import litellm
from app.ai.llm_client import llm_complete
from sqlalchemy import func, select, text
from sqlalchemy.ext.asyncio import AsyncSession
@@ -421,27 +422,23 @@ async def generate_suggestion(
model_parts = model.split("/", 1)
model = f"{provider_type}/{model_parts[-1]}"
litellm_kwargs: dict[str, Any] = dict(
model=model,
messages=[
{"role": "system", "content": SYSTEM_PROMPT},
{
"role": "user",
"content": json.dumps(context_data, default=str, ensure_ascii=False),
},
],
temperature=0.3,
max_tokens=1500,
response_format={"type": "json_object"},
)
if api_key:
litellm_kwargs["api_key"] = api_key
if api_base:
litellm_kwargs["api_base"] = api_base
try:
response = await litellm.acompletion(**litellm_kwargs)
content = response.choices[0].message.content
result = await llm_complete(
model=model,
messages=[
{"role": "system", "content": SYSTEM_PROMPT},
{
"role": "user",
"content": json.dumps(context_data, default=str, ensure_ascii=False),
},
],
temperature=0.3,
max_tokens=1500,
response_format={"type": "json_object"},
api_key=api_key,
api_base=api_base,
)
content = result["content"]
if not content:
return None
# Strip markdown code fences if present (e.g. ```json ... ```)
+14 -25
View File
@@ -17,6 +17,7 @@ from typing import Any
from sqlalchemy import func, select
from app.core.db import get_session_factory
from app.ai.llm_client import llm_complete
logger = logging.getLogger(__name__)
@@ -154,48 +155,36 @@ async def run_agent(
system_prompt = agent.system_prompt or "You are a helpful AI assistant."
user_prompt = f"Context: {context_data}"
# Use litellm directly for more control
import litellm
litellm_model = agent.model or "gpt-4o"
if agent.provider and agent.provider != "openai":
litellm_model = f"{agent.provider}/{litellm_model}"
kwargs: dict[str, Any] = {
"model": litellm_model,
"messages": [
result = await llm_complete(
model=litellm_model,
messages=[
{"role": "system", "content": system_prompt},
{"role": "user", "content": user_prompt},
],
"temperature": 0.3,
"max_tokens": agent.max_tokens or 1000,
}
if agent.api_key:
kwargs["api_key"] = agent.api_key
if agent.api_base:
kwargs["api_base"] = agent.api_base
response = await litellm.acompletion(**kwargs)
content = response.choices[0].message.content
temperature=0.3,
max_tokens=agent.max_tokens or 1000,
api_key=agent.api_key or None,
api_base=agent.api_base or None,
)
content = result["content"]
# Track cost
if hasattr(response, "usage") and response.usage:
usage = response.usage
# Estimate cost (simplified)
input_cost = (usage.prompt_tokens or 0) * 0.00001 / 1000
output_cost = (usage.completion_tokens or 0) * 0.00003 / 1000
result_data["cost_usd"] = round(input_cost + output_cost, 6)
result_data["cost_usd"] = result["cost_usd"]
result_data["llm_response"] = content
# Execute tool calls if LLM returned function calls
if hasattr(response.choices[0].message, "tool_calls") and response.choices[0].message.tool_calls:
raw_response = result["raw_response"]
if hasattr(raw_response.choices[0].message, "tool_calls") and raw_response.choices[0].message.tool_calls:
from app.plugins.builtins.ai_assistant.contracts import get_tool_registry
registry = get_tool_registry()
tool_call_count: dict[str, int] = {}
for tc in response.choices[0].message.tool_calls:
for tc in raw_response.choices[0].message.tool_calls:
tool_name = tc.function.name
# ── Safety Check 4: Infinite Loop Detection ──
tool_call_count[tool_name] = tool_call_count.get(tool_name, 0) + 1
@@ -1,26 +1,48 @@
"""Embedding pipeline using LiteLLM with OpenRouter for embeddings."""
"""Embedding pipeline using LiteLLM with OpenRouter for embeddings.
Delegates credential lookup, model building, and embedding generation to
the centralised ``app.ai.llm_client`` module. The wrapper functions here
preserve backward compatibility for existing call sites.
"""
from __future__ import annotations
import os
import logging
import os
import uuid
from typing import Any, TYPE_CHECKING
import litellm
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from sqlalchemy.ext.asyncio import AsyncSession
from app.ai.llm_client import (
EMBEDDING_DIMENSIONS,
MAX_INPUT_CHARS,
OPENROUTER_EMBEDDING_MODEL,
build_model as _central_build_model,
get_api_credentials as _central_get_api_credentials,
llm_embed,
)
logger = logging.getLogger(__name__)
MAX_INPUT_CHARS = 8000
# Re-export constants for backward compatibility
__all__ = [
"MAX_INPUT_CHARS",
"OPENROUTER_API_KEY",
"OPENROUTER_BASE_URL",
"OPENROUTER_EMBEDDING_MODEL",
"EMBEDDING_DIMENSIONS",
"_get_api_credentials",
"_build_model",
"generate_embedding",
"generate_embeddings_batch",
"index_entity",
]
# 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)
# Re-export for backward compatibility (consumers may import these directly)
OPENROUTER_API_KEY = os.environ.get("API_KEY_OPENROUTER", "")
OPENROUTER_BASE_URL = "https://openrouter.ai/api/v1"
async def _get_api_credentials(
@@ -28,34 +50,19 @@ async def _get_api_credentials(
) -> tuple[str | None, str | None, str | None]:
"""Get API key, base_url and provider_type for embeddings.
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)
Thin wrapper delegating to ``app.ai.llm_client.get_api_credentials``.
Kept for backward compatibility with existing call sites.
"""
# 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
return await _central_get_api_credentials(db, tenant_id)
def _build_model(model: str, provider_type: str | None) -> str:
"""Build litellm model string with provider prefix."""
if provider_type:
model_parts = model.split("/", 1)
return f"{provider_type}/{model_parts[-1]}"
return model
"""Build litellm model string with provider prefix.
Thin wrapper delegating to ``app.ai.llm_client.build_model``.
Kept for backward compatibility with existing call sites.
"""
return _central_build_model(model, provider_type)
async def generate_embedding(
@@ -77,30 +84,15 @@ async def generate_embedding(
Returns:
Embedding vector as list of floats.
"""
model = model or OPENROUTER_EMBEDDING_MODEL
truncated = text[:MAX_INPUT_CHARS]
try:
api_key, api_base, provider_type = await _get_api_credentials(db, tenant_id)
litellm_model = _build_model(model, provider_type)
litellm_kwargs: dict[str, Any] = dict(
model=litellm_model,
input=truncated,
)
if api_key:
litellm_kwargs["api_key"] = api_key
if api_base:
litellm_kwargs["api_base"] = api_base
# Request 768 dimensions to match DB vector(768) column
if 'text-embedding-3' in litellm_model:
litellm_kwargs['dimensions'] = EMBEDDING_DIMENSIONS
response = await litellm.aembedding(**litellm_kwargs)
return response.data[0]["embedding"]
except Exception:
logger.warning("Failed to generate embedding", exc_info=True)
return []
embeddings = await llm_embed(
texts=text,
model=model,
db=db,
tenant_id=tenant_id,
)
if embeddings and embeddings[0]:
return embeddings[0]
return []
async def generate_embeddings_batch(
@@ -120,29 +112,12 @@ async def generate_embeddings_batch(
Returns:
List of embedding vectors.
"""
model = model or OPENROUTER_EMBEDDING_MODEL
truncated = [t[:MAX_INPUT_CHARS] for t in texts]
try:
api_key, api_base, provider_type = await _get_api_credentials(db, tenant_id)
litellm_model = _build_model(model, provider_type)
litellm_kwargs: dict[str, Any] = dict(
model=litellm_model,
input=truncated,
)
if api_key:
litellm_kwargs["api_key"] = api_key
if api_base:
litellm_kwargs["api_base"] = api_base
if 'text-embedding-3' in litellm_model:
litellm_kwargs['dimensions'] = EMBEDDING_DIMENSIONS
response = await litellm.aembedding(**litellm_kwargs)
return [d["embedding"] for d in response.data]
except Exception:
logger.warning("Failed to generate batch embeddings", exc_info=True)
return [[] for _ in texts]
return await llm_embed(
texts=texts,
model=model,
db=db,
tenant_id=tenant_id,
)
async def index_entity(
@@ -201,5 +176,5 @@ async def index_entity(
await db.commit()
return True
except Exception:
logger.exception("Failed to index entity %s/%s", entity_type, entity_id)
logger.warning("Failed to index entity %s/%s", entity_type, entity_id, exc_info=True)
return False
@@ -8,7 +8,7 @@ import logging
import uuid
from typing import Any
import litellm
from app.ai.llm_client import llm_complete
from sqlalchemy.ext.asyncio import AsyncSession
logger = logging.getLogger(__name__)
@@ -87,7 +87,7 @@ async def llm_analyze_query(
api_key, api_base, provider_type = await _get_api_credentials(db, tenant_id)
model = _build_model(DEFAULT_LLM_MODEL, provider_type)
litellm_kwargs: dict[str, Any] = dict(
result = await llm_complete(
model=model,
messages=[
{"role": "system", "content": QUERY_ANALYZE_SYSTEM},
@@ -96,14 +96,10 @@ async def llm_analyze_query(
temperature=0.1,
max_tokens=500,
response_format={"type": "json_object"},
api_key=api_key,
api_base=api_base,
)
if api_key:
litellm_kwargs["api_key"] = api_key
if api_base:
litellm_kwargs["api_base"] = api_base
response = await litellm.acompletion(**litellm_kwargs)
content = response.choices[0].message.content
content = result["content"]
# Strip markdown code fences if present
content = content.strip()
if content.startswith("```"):
@@ -139,7 +135,7 @@ async def llm_aggregate_results(
]
user_msg = json.dumps({"query": query, "results": compact})
litellm_kwargs: dict[str, Any] = dict(
result = await llm_complete(
model=model,
messages=[
{"role": "system", "content": RESULT_AGGREGATE_SYSTEM},
@@ -148,14 +144,10 @@ async def llm_aggregate_results(
temperature=0.1,
max_tokens=1000,
response_format={"type": "json_object"},
api_key=api_key,
api_base=api_base,
)
if api_key:
litellm_kwargs["api_key"] = api_key
if api_base:
litellm_kwargs["api_base"] = api_base
response = await litellm.acompletion(**litellm_kwargs)
content = response.choices[0].message.content
content = result["content"]
# Strip markdown code fences if present
content = content.strip()
if content.startswith("```"):