refa: resolve tenant model refs consistently (#16744)

This commit is contained in:
buua436
2026-07-09 14:02:08 +08:00
committed by GitHub
parent 794fcc2517
commit 6a77523bf0
51 changed files with 300 additions and 225 deletions
@@ -40,7 +40,7 @@ from rag.svr.task_executor_refactor.task_context import TaskContext
from api.db.services.doc_metadata_service import DocMetadataService
from api.db.services.llm_service import LLMBundle
from api.db.joint_services.tenant_model_service import get_model_config_from_provider_instance
from api.db.joint_services.tenant_model_service import resolve_model_config
from rag.prompts.generator import gen_metadata, keyword_extraction, question_proposal, content_tagging
from rag.graphrag.utils import get_llm_cache, set_llm_cache, get_tags_from_cache, set_tags_to_cache
@@ -56,7 +56,7 @@ async def extract_keywords(docs: List[Dict], ctx: TaskContext) -> None:
st = timer()
ctx.progress_cb(msg="Start to generate keywords for every chunk ...")
chat_model_config = get_model_config_from_provider_instance(ctx.tenant_id, LLMType.CHAT, ctx.llm_id)
chat_model_config = resolve_model_config(ctx.tenant_id, LLMType.CHAT, ctx.llm_id)
with LLMBundle(ctx.tenant_id, chat_model_config, lang=ctx.language) as chat_model:
async def doc_keyword_extraction(chat_mdl, d, topn):
@@ -98,7 +98,7 @@ async def generate_questions(docs: List[Dict], ctx: TaskContext) -> None:
st = timer()
ctx.progress_cb(msg="Start to generate questions for every chunk ...")
chat_model_config = get_model_config_from_provider_instance(ctx.tenant_id, LLMType.CHAT, ctx.llm_id)
chat_model_config = resolve_model_config(ctx.tenant_id, LLMType.CHAT, ctx.llm_id)
with LLMBundle(ctx.tenant_id, chat_model_config, lang=ctx.language) as chat_model:
async def doc_question_proposal(chat_mdl, d, topn):
@@ -178,7 +178,7 @@ async def generate_metadata(docs: List[Dict], ctx: TaskContext) -> None:
st = timer()
ctx.progress_cb(msg="Start to generate meta-data for every chunk ...")
chat_model_config = get_model_config_from_provider_instance(ctx.tenant_id, LLMType.CHAT, ctx.llm_id)
chat_model_config = resolve_model_config(ctx.tenant_id, LLMType.CHAT, ctx.llm_id)
with LLMBundle(ctx.tenant_id, chat_model_config, lang=ctx.language) as chat_model:
metadata_conf = build_metadata_config(ctx.parser_config)
@@ -266,7 +266,7 @@ async def apply_tags(docs: List[Dict], ctx: TaskContext) -> None:
set_tags_to_cache(kb_ids, all_tags)
else:
all_tags = json.loads(all_tags)
chat_model_config = get_model_config_from_provider_instance(tenant_id, LLMType.CHAT, ctx.llm_id)
chat_model_config = resolve_model_config(tenant_id, LLMType.CHAT, ctx.llm_id)
with LLMBundle(ctx.tenant_id, chat_model_config, lang=ctx.language) as chat_model:
docs_to_tag = []
for doc in docs:
@@ -940,7 +940,7 @@ async def run_document_structure_compile(handler, embedding_model: LLMBundle) ->
chat_llm_id = _resolve_template_chat_llm_id(parser_cfg, ctx)
if chat_llm_id not in llm_bundle_cache:
try:
cfg = get_model_config_from_provider_instance(
cfg = resolve_model_config(
ctx.tenant_id,
LLMType.CHAT,
chat_llm_id,
@@ -38,7 +38,7 @@ from api.db.services.canvas_service import UserCanvasService
from api.db.services.document_service import DocumentService
from api.db.services.doc_metadata_service import DocMetadataService
from api.db.services.pipeline_operation_log_service import PipelineOperationLogService
from api.db.joint_services.tenant_model_service import get_model_config_from_provider_instance, get_model_config_by_id
from api.db.joint_services.tenant_model_service import resolve_model_config, get_model_config_by_id
from common.connection_utils import timeout
from common.constants import LLMType, PipelineTaskType
from common.metadata_utils import update_metadata_to
@@ -244,13 +244,13 @@ class DataflowService:
embedding_id = kb.embd_id
if kb.tenant_embd_id:
try:
embd_model_config = get_model_config_by_id(ctx.tenant_id, kb.tenant_embd_id)
embd_model_config = get_model_config_by_id(ctx.tenant_id, LLMType.EMBEDDING, kb.tenant_embd_id)
except LookupError:
embd_model_config = get_model_config_from_provider_instance(
embd_model_config = resolve_model_config(
ctx.tenant_id, LLMType.EMBEDDING, embedding_id
)
else:
embd_model_config = get_model_config_from_provider_instance(
embd_model_config = resolve_model_config(
ctx.tenant_id, LLMType.EMBEDDING, embedding_id
)
from api.db.services.llm_service import LLMBundle
@@ -645,7 +645,7 @@ async def run_wiki(
from api.db.services.llm_service import LLMBundle
from api.db.joint_services.tenant_model_service import (
get_tenant_default_model_by_type,
get_model_config_from_provider_instance,
resolve_model_config,
)
from api.apps.restful_apis.chunk_api import _compilation_template_kind
@@ -705,7 +705,7 @@ async def run_wiki(
if key == "__tenant_default__":
cfg = get_tenant_default_model_by_type(ctx.tenant_id, LLMType.CHAT)
else:
cfg = get_model_config_from_provider_instance(
cfg = resolve_model_config(
ctx.tenant_id,
LLMType.CHAT,
key,
+13 -7
View File
@@ -42,7 +42,7 @@ from api.db.services.compilation_template_group_service import CompilationTempla
from api.db.joint_services.memory_message_service import handle_save_to_memory_task
from api.db.joint_services.tenant_model_service import (
get_tenant_default_model_by_type,
get_model_config_from_provider_instance,
resolve_model_config,
get_model_config_by_id,
)
from api.db.services.llm_service import LLMBundle
@@ -321,11 +321,17 @@ class TaskHandler:
try:
if ctx.tenant_embd_id:
try:
embd_model_config = get_model_config_by_id(task_tenant_id, ctx.tenant_embd_id)
embd_model_config = get_model_config_by_id(
task_tenant_id, LLMType.EMBEDDING, ctx.tenant_embd_id
)
except LookupError:
embd_model_config = get_model_config_from_provider_instance(task_tenant_id, LLMType.EMBEDDING, task_embedding_id)
embd_model_config = resolve_model_config(
task_tenant_id, LLMType.EMBEDDING, task_embedding_id
)
elif task_embedding_id:
embd_model_config = get_model_config_from_provider_instance(task_tenant_id, LLMType.EMBEDDING, task_embedding_id)
embd_model_config = resolve_model_config(
task_tenant_id, LLMType.EMBEDDING, task_embedding_id
)
else:
embd_model_config = get_tenant_default_model_by_type(task_tenant_id, LLMType.EMBEDDING)
embedding_model = LLMBundle(task_tenant_id, embd_model_config, lang=task_language)
@@ -381,7 +387,7 @@ class TaskHandler:
return
# Bind LLM for raptor
chat_model_config = get_model_config_from_provider_instance(task_tenant_id, LLMType.CHAT, kb_task_llm_id)
chat_model_config = resolve_model_config(task_tenant_id, LLMType.CHAT, kb_task_llm_id)
with LLMBundle(task_tenant_id, chat_model_config, lang=ctx.language) as chat_model:
# Run RAPTOR
raptor_service = RaptorService(ctx=ctx)
@@ -497,7 +503,7 @@ class TaskHandler:
graphrag_conf = kb_parser_config.get("graphrag", {})
start_ts = timer()
chat_model_config = get_model_config_from_provider_instance(task_tenant_id, LLMType.CHAT, kb_task_llm_id)
chat_model_config = resolve_model_config(task_tenant_id, LLMType.CHAT, kb_task_llm_id)
with LLMBundle(task_tenant_id, chat_model_config, lang=task_language) as chat_model:
with_resolution = graphrag_conf.get("resolution", False)
with_community = graphrag_conf.get("community", False)
@@ -780,7 +786,7 @@ class TaskHandler:
def _build_toc(cls, ctx: TaskContext, docs: List[Dict], progress_cb: Callable) -> Optional[Dict]:
"""Build table of contents."""
progress_cb(msg="Start to generate table of content ...")
chat_model_config = get_model_config_from_provider_instance(ctx.tenant_id, LLMType.CHAT, ctx.llm_id)
chat_model_config = resolve_model_config(ctx.tenant_id, LLMType.CHAT, ctx.llm_id)
with LLMBundle(ctx.tenant_id, chat_model_config, lang=ctx.language) as chat_mdl:
docs = sorted(
docs,