fix: unified_search + ai_proactive get API key from DB, fix model names for Ollama Cloud

- query_understanding.py: get API key/base_url/provider_type from ai_providers DB
- embedding.py: get API key from DB, pass db+tenant_id through call chain
- routes.py: pass db+tenant_id to llm_analyze_query and llm_aggregate_results
- search_engine.py: pass db+tenant_id to generate_embedding
- unified_search/jobs.py: pass db+tenant_id to generate_embedding
- Fix all default model names: ollama/deepseek-v4 -> ollama/deepseek-v4-flash
- Ollama Cloud has no embedding endpoint; embedding calls fail gracefully
This commit is contained in:
Agent Zero
2026-07-19 02:22:25 +02:00
parent ef4f0cc494
commit 4a43745b50
9 changed files with 147 additions and 37 deletions
+2 -2
View File
@@ -186,7 +186,7 @@ async def deep_analysis(
extended_context["similar"] = {}
# Generate extended suggestion with deep analysis prompt
model = settings.model or "ollama/deepseek-v4"
model = settings.model or "ollama/deepseek-v4-flash"
# Get API key from DB (like ai_assistant does) or fall back to env
from app.plugins.builtins.ai_proactive.services import _get_llm_api_key
@@ -195,7 +195,7 @@ async def deep_analysis(
api_key = os.environ.get('API_KEY_OLLAMA_CLOUD', '') or None
# Build model string with provider prefix
model = settings.model or "ollama/deepseek-v4"
model = settings.model or "ollama/deepseek-v4-flash"
if provider_type:
model_parts = model.split("/", 1)
model = f"{provider_type}/{model_parts[-1]}"
+1 -1
View File
@@ -108,4 +108,4 @@ class ProactiveSettings(Base, TenantMixin):
rate_limit_seconds: Mapped[int] = mapped_column(
Integer, nullable=False, default=10
)
model: Mapped[str] = mapped_column(String(100), nullable=False, default="ollama/deepseek-v4")
model: Mapped[str] = mapped_column(String(100), nullable=False, default="ollama/deepseek-v4-flash")
+1 -1
View File
@@ -74,7 +74,7 @@ class SettingsResponse(BaseModel):
rate_limit_seconds: int
model: str
available_models: list[str] = Field(default_factory=lambda: [
'ollama/deepseek-v4',
'ollama/deepseek-v4-flash',
'ollama/deepseek-v4-pro',
'ollama/llama3.2',
'ollama/gpt-4o-mini',
@@ -113,7 +113,7 @@ async def get_user_settings(
suggestion_categories=["mail", "tasks", "contacts", "companies", "insights"],
confidence_threshold=0.5,
rate_limit_seconds=10,
model="ollama/deepseek-v4",
model="ollama/deepseek-v4-flash",
)
db.add(settings)
await db.flush()
@@ -395,7 +395,7 @@ async def generate_suggestion(
Returns dict with suggestion_type, title, content, confidence, actions
or None on failure.
"""
model = settings.model or "ollama/deepseek-v4"
model = settings.model or "ollama/deepseek-v4-flash"
# Get API key from DB (like ai_assistant does) or fall back to env
api_key = None
@@ -407,7 +407,7 @@ async def generate_suggestion(
api_key = os.environ.get('API_KEY_OLLAMA_CLOUD', '') or None
# Build model string with provider prefix (like ai_assistant build_litellm_params)
model = settings.model or "ollama/deepseek-v4"
model = settings.model or "ollama/deepseek-v4-flash"
if provider_type:
model_parts = model.split("/", 1)
model = f"{provider_type}/{model_parts[-1]}"
@@ -5,7 +5,7 @@ from __future__ import annotations
import os
import logging
import uuid
from typing import TYPE_CHECKING
from typing import Any, TYPE_CHECKING
import litellm
@@ -17,15 +17,49 @@ logger = logging.getLogger(__name__)
MAX_INPUT_CHARS = 8000
DEFAULT_EMBEDDING_MODEL = os.environ.get('SEARCH_EMBEDDING_MODEL', 'ollama/nomic-embed-text')
OLLAMA_API_KEY = os.environ.get('API_KEY_OLLAMA_CLOUD', '')
async def generate_embedding(text: str, model: str | None = None) -> list[float]:
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 from the default AI provider in DB.
Falls back to API_KEY_OLLAMA_CLOUD env var.
Returns (api_key, base_url, provider_type).
"""
if db and tenant_id:
try:
from app.plugins.builtins.ai_assistant.services 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:
model_parts = model.split("/", 1)
return f"{provider_type}/{model_parts[-1]}"
return model
async def generate_embedding(
text: str,
model: str | None = None,
db: "AsyncSession | None" = None,
tenant_id: "uuid.UUID | None" = None,
) -> list[float]:
"""Generate a single embedding via LiteLLM.
Args:
text: Input text (truncated to 8000 chars).
model: Embedding model name (default: ollama/nomic-embed-text).
db: Optional DB session for API key lookup.
tenant_id: Optional tenant ID for API key lookup.
Returns:
Embedding vector as list of floats.
@@ -33,25 +67,38 @@ async def generate_embedding(text: str, model: str | None = None) -> list[float]
model = model or DEFAULT_EMBEDDING_MODEL
truncated = text[:MAX_INPUT_CHARS]
try:
response = await litellm.aembedding(
model=model,
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,
api_key=OLLAMA_API_KEY,
)
if api_key:
litellm_kwargs["api_key"] = api_key
if api_base:
litellm_kwargs["api_base"] = api_base
response = await litellm.aembedding(**litellm_kwargs)
return response.data[0]["embedding"]
except Exception:
logger.exception("Failed to generate embedding")
logger.warning("Failed to generate embedding (Ollama Cloud may not support embeddings)", exc_info=True)
return []
async def generate_embeddings_batch(
texts: list[str], model: str | None = None
texts: list[str],
model: str | None = None,
db: "AsyncSession | None" = None,
tenant_id: "uuid.UUID | None" = None,
) -> list[list[float]]:
"""Generate embeddings for multiple texts in a single API call.
Args:
texts: List of input texts.
model: Embedding model name (default: ollama/nomic-embed-text).
db: Optional DB session for API key lookup.
tenant_id: Optional tenant ID for API key lookup.
Returns:
List of embedding vectors.
@@ -59,14 +106,22 @@ async def generate_embeddings_batch(
model = model or DEFAULT_EMBEDDING_MODEL
truncated = [t[:MAX_INPUT_CHARS] for t in texts]
try:
response = await litellm.aembedding(
model=model,
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,
api_key=OLLAMA_API_KEY,
)
if api_key:
litellm_kwargs["api_key"] = api_key
if api_base:
litellm_kwargs["api_base"] = api_base
response = await litellm.aembedding(**litellm_kwargs)
return [d["embedding"] for d in response.data]
except Exception:
logger.exception("Failed to generate batch embeddings")
logger.warning("Failed to generate batch embeddings", exc_info=True)
return [[] for _ in texts]
@@ -74,7 +129,7 @@ async def index_entity(
entity_type: str,
entity_id: uuid.UUID,
tenant_id: uuid.UUID,
db: AsyncSession,
db: "AsyncSession",
) -> bool:
"""Generate and store embedding for a single entity.
@@ -97,7 +152,7 @@ async def index_entity(
logger.debug("Empty embedding text for %s/%s", entity_type, entity_id)
return False
embedding = await generate_embedding(text)
embedding = await generate_embedding(text, db=db, tenant_id=tenant_id)
if not embedding:
return False
+1 -1
View File
@@ -87,7 +87,7 @@ async def index_file(ctx: dict[str, Any], file_id: str) -> None:
embedding_text = f"{name} {content_text[:5000]}"
if embedding_text.strip():
embedding = await generate_embedding(embedding_text)
embedding = await generate_embedding(embedding_text, db=db, tenant_id=tenant_id)
if embedding:
await db.execute(
text("UPDATE files SET embedding = cast(:emb AS vector) WHERE id = :fid"),
@@ -5,14 +5,15 @@ from __future__ import annotations
import os
import json
import logging
import uuid
from typing import Any
import litellm
from sqlalchemy.ext.asyncio import AsyncSession
logger = logging.getLogger(__name__)
DEFAULT_LLM_MODEL = os.environ.get('SEARCH_LLM_MODEL', 'ollama/deepseek-v4')
OLLAMA_API_KEY = os.environ.get('API_KEY_OLLAMA_CLOUD', '')
DEFAULT_LLM_MODEL = os.environ.get('SEARCH_LLM_MODEL', 'ollama/deepseek-v4-flash')
QUERY_ANALYZE_SYSTEM = (
"Du bist ein Query-Analyzer fuer ein CRM. "
@@ -27,6 +28,34 @@ RESULT_AGGREGATE_SYSTEM = (
)
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 from the default AI provider in DB.
Falls back to API_KEY_OLLAMA_CLOUD env var.
Returns (api_key, base_url, provider_type).
"""
if db and tenant_id:
try:
from app.plugins.builtins.ai_assistant.services 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:
model_parts = model.split("/", 1)
return f"{provider_type}/{model_parts[-1]}"
return model
def _fallback_query_analysis(query: str) -> dict[str, Any]:
return {
"normalized_query": query,
@@ -45,14 +74,21 @@ def _fallback_aggregate(results: list[dict], query: str) -> dict[str, Any]:
}
async def llm_analyze_query(query: str) -> dict[str, Any]:
async def llm_analyze_query(
query: str,
db: AsyncSession | None = None,
tenant_id: uuid.UUID | None = None,
) -> dict[str, Any]:
"""Analyze a search query using LLM for intent, entities, and semantic terms.
Falls back to a simple dict if LLM fails.
"""
try:
response = await litellm.acompletion(
model=DEFAULT_LLM_MODEL,
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(
model=model,
messages=[
{"role": "system", "content": QUERY_ANALYZE_SYSTEM},
{"role": "user", "content": query},
@@ -60,16 +96,26 @@ async def llm_analyze_query(query: str) -> dict[str, Any]:
temperature=0.1,
max_tokens=500,
response_format={"type": "json_object"},
api_key=OLLAMA_API_KEY,
)
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
return json.loads(content)
except Exception:
logger.warning("LLM query analysis failed, using fallback")
logger.warning("LLM query analysis failed, using fallback", exc_info=True)
return _fallback_query_analysis(query)
async def llm_aggregate_results(results: list[dict], query: str) -> dict[str, Any]:
async def llm_aggregate_results(
results: list[dict],
query: str,
db: AsyncSession | None = None,
tenant_id: uuid.UUID | None = None,
) -> dict[str, Any]:
"""Aggregate search results using LLM for summary, facets, and suggestions.
Falls back to a simple dict if LLM fails.
@@ -77,14 +123,18 @@ async def llm_aggregate_results(results: list[dict], query: str) -> dict[str, An
if not results:
return _fallback_aggregate(results, query)
try:
api_key, api_base, provider_type = await _get_api_credentials(db, tenant_id)
model = _build_model(DEFAULT_LLM_MODEL, provider_type)
# Truncate results to avoid token overflow
compact = [
{"entity_type": r.get("entity_type"), "title": r.get("title", "")[:100]}
for r in results[:50]
]
user_msg = json.dumps({"query": query, "results": compact})
response = await litellm.acompletion(
model=DEFAULT_LLM_MODEL,
litellm_kwargs: dict[str, Any] = dict(
model=model,
messages=[
{"role": "system", "content": RESULT_AGGREGATE_SYSTEM},
{"role": "user", "content": user_msg},
@@ -92,10 +142,15 @@ async def llm_aggregate_results(results: list[dict], query: str) -> dict[str, An
temperature=0.1,
max_tokens=1000,
response_format={"type": "json_object"},
api_key=OLLAMA_API_KEY,
)
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
return json.loads(content)
except Exception:
logger.warning("LLM result aggregation failed, using fallback")
logger.warning("LLM result aggregation failed, using fallback", exc_info=True)
return _fallback_aggregate(results, query)
@@ -51,7 +51,7 @@ async def search(
tenant_id = uuid.UUID(current_user["tenant_id"])
# KI query understanding
query_analysis = await llm_analyze_query(req.query)
query_analysis = await llm_analyze_query(req.query, db=db, tenant_id=tenant_id)
# Hybrid search
results = await hybrid_search(
@@ -63,7 +63,7 @@ async def search(
)
# KI result aggregation
aggregation = await llm_aggregate_results(results, req.query)
aggregation = await llm_aggregate_results(results, req.query, db=db, tenant_id=tenant_id)
search_results = [
SearchResult(
@@ -95,7 +95,7 @@ async def hybrid_search(
query_text = normalized_query
if semantic_terms:
query_text = f"{normalized_query} {' '.join(semantic_terms)}"
query_embedding = await generate_embedding(query_text)
query_embedding = await generate_embedding(query_text, db=db, tenant_id=tenant_id)
all_results: list[dict[str, Any]] = []
fetch_limit = limit * 2