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:
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user