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
@@ -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)