zlei6/ragflow
1
0
Fork 0
mirror of https://github.com/infiniflow/ragflow.git synced 2026-07-16 12:47:19 +08:00
Code Issues Packages Projects Releases Wiki Activity

Feat: tenant llm provider (#14595)

Browse Source 
### What problem does this PR solve?

Python implementation of the Go-based model_provider API suite.

### Type of change

- [x] New Feature (non-breaking change which adds functionality)

---------

Co-authored-by: bill <yibie_jingnian@163.com>
This commit is contained in:
Lynn
2026-05-29 17:39:41 +08:00
committed by GitHub
parent b79f79d9b9
commit dc4b82523b
148 changed files with 6059 additions and 3075 deletions
Download Patch File Download Diff File Expand all files Collapse all files

109
api/db/init_data.py
View File

@@ -24,17 +24,15 @@ from copy import deepcopy
from peewee import IntegrityError
from api.db import UserTenantRole
from api.db.db_models import init_database_tables as init_web_db, LLMFactories, LLM, TenantLLM, Knowledgebase, Dialog, Memory
from api.db.db_models import init_database_tables as init_web_db, LLMFactories, LLM, TenantLLM
from api.db.services import UserService
from api.db.services.canvas_service import CanvasTemplateService
from api.db.services.document_service import DocumentService
from api.db.services.knowledgebase_service import KnowledgebaseService
from api.db.services.memory_service import MemoryService
from api.db.services.tenant_llm_service import LLMFactoriesService, TenantLLMService
from api.db.services.llm_service import LLMService, LLMBundle, get_init_tenant_llm
from api.db.services.user_service import TenantService, UserTenantService
from api.db.services.system_settings_service import SystemSettingsService
from api.db.services.dialog_service import DialogService
from api.db.template_utils import normalize_canvas_template_categories
from api.db.joint_services.memory_message_service import init_message_id_sequence, init_memory_size_cache, fix_missing_tokenized_memory
from api.db.joint_services.tenant_model_service import get_tenant_default_model_by_type
@@ -154,6 +152,8 @@ def init_llm_factory():
except Exception:
pass
break
def update_document_number_in_init():
doc_count = DocumentService.get_all_kb_doc_count()
for kb_id in KnowledgebaseService.get_all_ids():
KnowledgebaseService.update_document_number_in_init(kb_id=kb_id, doc_num=doc_count.get(kb_id, 0))
@@ -189,7 +189,8 @@ def init_web_data():
init_table()
init_llm_factory()
# init_llm_factory()
update_document_number_in_init()
# if not UserService.get_all().count():
# init_superuser()
@@ -197,7 +198,6 @@ def init_web_data():
init_message_id_sequence()
init_memory_size_cache()
fix_missing_tokenized_memory()
fix_empty_tenant_model_id()
logging.info("init web data success:{}".format(time.time() - start_time))
def init_table():
@@ -226,105 +226,6 @@ def init_table():
raise e
def fix_empty_tenant_model_id():
# knowledgebase
empty_tenant_embd_id_kbs = KnowledgebaseService.get_null_tenant_embd_id_row()
if empty_tenant_embd_id_kbs:
logging.info(f"Found {len(empty_tenant_embd_id_kbs)} empty tenant_embd_id knowledgebase.")
kb_groups: dict = {}
for obj in empty_tenant_embd_id_kbs:
if kb_groups.get((obj.tenant_id, obj.embd_id)):
kb_groups[(obj.tenant_id, obj.embd_id)].append(obj.id)
else:
kb_groups[(obj.tenant_id, obj.embd_id)] = [obj.id]
update_cnt = 0
for k, v in kb_groups.items():
tenant_llm = TenantLLMService.get_api_key(k[0], k[1])
if tenant_llm:
update_cnt += KnowledgebaseService.filter_update([Knowledgebase.id.in_(v)], {"tenant_embd_id": tenant_llm.id})
logging.info(f"Update {update_cnt} tenant_embd_id in table knowledgebase.")
# dialog
empty_tenant_llm_id_dialog = DialogService.get_null_tenant_llm_id_row()
if empty_tenant_llm_id_dialog:
logging.info(f"Found {len(empty_tenant_llm_id_dialog)} empty tenant_llm_id dialogs.")
dialog_groups: dict = {}
for obj in empty_tenant_llm_id_dialog:
if dialog_groups.get((obj.tenant_id, obj.llm_id)):
dialog_groups[(obj.tenant_id, obj.llm_id)].append(obj.id)
else:
dialog_groups[(obj.tenant_id, obj.llm_id)] = [obj.id]
update_cnt = 0
for k, v in dialog_groups.items():
tenant_llm = TenantLLMService.get_api_key(k[0], k[1])
if tenant_llm:
update_cnt += DialogService.filter_update([Dialog.id.in_(v)], {"tenant_llm_id": tenant_llm.id})
logging.info(f"Update {update_cnt} tenant_llm_id in table dialog.")
empty_tenant_rerank_id_dialog = DialogService.get_null_tenant_rerank_id_row()
if empty_tenant_rerank_id_dialog:
logging.info(f"Found {len(empty_tenant_rerank_id_dialog)} empty tenant_rerank_id dialogs.")
dialog_groups: dict = {}
for obj in empty_tenant_rerank_id_dialog:
if dialog_groups.get((obj.tenant_id, obj.rerank_id)):
dialog_groups[(obj.tenant_id, obj.rerank_id)].append(obj.id)
else:
dialog_groups[(obj.tenant_id, obj.rerank_id)] = [obj.id]
update_cnt = 0
for k, v in dialog_groups.items():
tenant_llm = TenantLLMService.get_api_key(k[0], k[1])
if tenant_llm:
update_cnt += DialogService.filter_update([Dialog.id.in_(v)], {"tenant_rerank_id": tenant_llm.id})
logging.info(f"Update {update_cnt} tenant_rerank_id in table dialog.")
# memory
empty_tenant_embd_id_memories = MemoryService.get_null_tenant_embd_id_row()
if empty_tenant_embd_id_memories:
logging.info(f"Found {len(empty_tenant_embd_id_memories)} empty tenant_embd_id memories.")
memory_groups: dict = {}
for obj in empty_tenant_embd_id_memories:
if memory_groups.get((obj.tenant_id, obj.embd_id)):
memory_groups[(obj.tenant_id, obj.embd_id)].append(obj.id)
else:
memory_groups[(obj.tenant_id, obj.embd_id)] = [obj.id]
update_cnt = 0
for k, v in memory_groups.items():
tenant_llm = TenantLLMService.get_api_key(k[0], k[1])
if tenant_llm:
update_cnt += MemoryService.filter_update([Memory.id.in_(v)], {"tenant_embd_id": tenant_llm.id})
logging.info(f"Update {update_cnt} tenant_embd_id in table memory.")
empty_tenant_llm_id_memories = MemoryService.get_null_tenant_llm_id_row()
if empty_tenant_llm_id_memories:
logging.info(f"Found {len(empty_tenant_llm_id_memories)} empty tenant_llm_id memories.")
memory_groups: dict = {}
for obj in empty_tenant_llm_id_memories:
if memory_groups.get((obj.tenant_id, obj.llm_id)):
memory_groups[(obj.tenant_id, obj.llm_id)].append(obj.id)
else:
memory_groups[(obj.tenant_id, obj.llm_id)] = [obj.id]
update_cnt = 0
for k, v in memory_groups.items():
tenant_llm = TenantLLMService.get_api_key(k[0], k[1])
if tenant_llm:
update_cnt += MemoryService.filter_update([Memory.id.in_(v)], {"tenant_llm_id": tenant_llm.id})
logging.info(f"Update {update_cnt} tenant_llm_id in table memory.")
# tenant
empty_tenant_model_id_tenants = TenantService.get_null_tenant_model_id_rows()
if empty_tenant_model_id_tenants:
logging.info(f"Found {len(empty_tenant_model_id_tenants)} empty tenant_model_id tenants.")
update_cnt = 0
for obj in empty_tenant_model_id_tenants:
tenant_dict = obj.to_dict()
update_dict = {}
for key in ["llm_id", "embd_id", "asr_id", "img2txt_id", "rerank_id", "tts_id"]:
if tenant_dict.get(key) and not tenant_dict.get(f"tenant_{key}"):
tenant_model = TenantLLMService.get_api_key(tenant_dict["id"], tenant_dict[key])
if tenant_model:
update_dict.update({f"tenant_{key}": tenant_model.id})
if update_dict:
update_cnt += TenantService.update_by_id(tenant_dict["id"], update_dict)
logging.info(f"Update {update_cnt} tenant_model_id in table tenant.")
logging.info("Fix empty tenant_model_id done.")
if __name__ == '__main__':
init_web_db()
init_web_data()