mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-08-04 23:00:30 +08:00
Feat/tenant model (#13072)
### What problem does this PR solve? Add id for table tenant_llm and apply in LLMBundle. ### Type of change - [x] Refactoring --------- Co-authored-by: Yingfeng <yingfeng.zhang@gmail.com> Co-authored-by: Liu An <asiro@qq.com>
This commit is contained in:
@@ -207,6 +207,10 @@ def _load_chunk_module(monkeypatch):
|
||||
EMBEDDING = SimpleNamespace(value="embedding")
|
||||
CHAT = SimpleNamespace(value="chat")
|
||||
RERANK = SimpleNamespace(value="rerank")
|
||||
SPEECH2TEXT = SimpleNamespace(value="speech2text")
|
||||
IMAGE2TEXT = SimpleNamespace(value="image2text")
|
||||
TTS = SimpleNamespace(value="tts")
|
||||
OCR = SimpleNamespace(value="ocr")
|
||||
|
||||
constants_mod.RetCode = _DummyRetCode
|
||||
constants_mod.LLMType = _DummyLLMType
|
||||
@@ -301,6 +305,10 @@ def _load_chunk_module(monkeypatch):
|
||||
def get_embd_id(_doc_id):
|
||||
return "embed-1"
|
||||
|
||||
@staticmethod
|
||||
def get_tenant_embd_id(_doc_id):
|
||||
return 1
|
||||
|
||||
@staticmethod
|
||||
def decrement_chunk_num(*args):
|
||||
_DocumentService.decrement_calls.append(args)
|
||||
@@ -327,13 +335,24 @@ def _load_chunk_module(monkeypatch):
|
||||
|
||||
@staticmethod
|
||||
def get_by_id(_kb_id):
|
||||
return True, SimpleNamespace(pagerank=0.6)
|
||||
return True, SimpleNamespace(pagerank=0.6, tenant_embd_id=2, tenant_llm_id=1)
|
||||
|
||||
kb_service_mod.KnowledgebaseService = _KnowledgebaseService
|
||||
monkeypatch.setitem(sys.modules, "api.db.services.knowledgebase_service", kb_service_mod)
|
||||
services_pkg.knowledgebase_service = kb_service_mod
|
||||
|
||||
class _DummyLLMService:
|
||||
@staticmethod
|
||||
def query(**_kwargs):
|
||||
return [SimpleNamespace(
|
||||
llm_name="gpt-3.5-turbo",
|
||||
model_type="chat",
|
||||
max_tokens=8192,
|
||||
is_tools=True
|
||||
)]
|
||||
|
||||
llm_service_mod = ModuleType("api.db.services.llm_service")
|
||||
llm_service_mod.LLMService = _DummyLLMService
|
||||
llm_service_mod.LLMBundle = _DummyLLMBundle
|
||||
monkeypatch.setitem(sys.modules, "api.db.services.llm_service", llm_service_mod)
|
||||
services_pkg.llm_service = llm_service_mod
|
||||
@@ -343,6 +362,77 @@ def _load_chunk_module(monkeypatch):
|
||||
monkeypatch.setitem(sys.modules, "api.db.services.search_service", search_service_mod)
|
||||
services_pkg.search_service = search_service_mod
|
||||
|
||||
tenant_llm_service_mod = ModuleType("api.db.services.tenant_llm_service")
|
||||
|
||||
class _MockTableObject:
|
||||
def __init__(self, **kwargs):
|
||||
for key, value in kwargs.items():
|
||||
setattr(self, key, value)
|
||||
|
||||
def to_dict(self):
|
||||
return {k: v for k, v in self.__dict__.items()}
|
||||
|
||||
class _TenantLLMService:
|
||||
@staticmethod
|
||||
def get_by_id(tenant_model_id):
|
||||
return True, _MockTableObject(
|
||||
id=tenant_model_id,
|
||||
tenant_id="tenant-1",
|
||||
llm_factory="",
|
||||
model_type="chat",
|
||||
llm_name="gpt-3.5-turbo",
|
||||
api_key="fake-api-key",
|
||||
api_base="https://api.example.com",
|
||||
max_tokens=8192,
|
||||
used_tokens=0,
|
||||
status=1
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def get_api_key(tenant_id, model_name):
|
||||
return _MockTableObject(
|
||||
id=1,
|
||||
tenant_id=tenant_id,
|
||||
llm_factory="",
|
||||
model_type="chat",
|
||||
llm_name=model_name,
|
||||
api_key="fake-api-key",
|
||||
api_base="https://api.example.com",
|
||||
max_tokens=8192,
|
||||
used_tokens=0,
|
||||
status=1
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def split_model_name_and_factory(model_name):
|
||||
if "@" in model_name:
|
||||
parts = model_name.rsplit("@", 1)
|
||||
return parts[0], parts[1]
|
||||
return model_name, None
|
||||
|
||||
@staticmethod
|
||||
def increase_usage_by_id(model_id, used_tokens):
|
||||
return True
|
||||
|
||||
class _TenantService:
|
||||
@staticmethod
|
||||
def get_by_id(tenant_id):
|
||||
return True, SimpleNamespace(
|
||||
llm_id="gpt-3.5-turbo",
|
||||
tenant_llm_id=1,
|
||||
embd_id="text-embedding-ada-002",
|
||||
tenant_embd_id=2,
|
||||
asr_id="whisper-1",
|
||||
img2txt_id="gpt-4-vision-preview",
|
||||
rerank_id="bge-reranker",
|
||||
tts_id="tts-1"
|
||||
)
|
||||
|
||||
tenant_llm_service_mod.TenantLLMService = _TenantLLMService
|
||||
tenant_llm_service_mod.TenantService = _TenantService
|
||||
monkeypatch.setitem(sys.modules, "api.db.services.tenant_llm_service", tenant_llm_service_mod)
|
||||
services_pkg.tenant_llm_service = tenant_llm_service_mod
|
||||
|
||||
user_service_mod = ModuleType("api.db.services.user_service")
|
||||
|
||||
class _UserTenantService:
|
||||
@@ -775,7 +865,7 @@ def test_retrieval_test_branch_matrix_unit(monkeypatch):
|
||||
assert "Knowledgebase not found!" in res["message"], res
|
||||
|
||||
retriever = _Retriever(mode="ok")
|
||||
monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _kb_id: (True, SimpleNamespace(tenant_id="tenant-kb", embd_id="embd-1")), raising=False)
|
||||
monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _kb_id: (True, SimpleNamespace(tenant_id="tenant-kb", embd_id="embd-1", tenant_embd_id=2)), raising=False)
|
||||
monkeypatch.setattr(module.settings, "retriever", retriever)
|
||||
monkeypatch.setattr(module.settings, "kg_retriever", _KgRetriever(), raising=False)
|
||||
_set_request_json(
|
||||
|
||||
@@ -143,6 +143,19 @@ def _load_conversation_module(monkeypatch):
|
||||
apps_mod.login_required = lambda func: func
|
||||
monkeypatch.setitem(sys.modules, "api.apps", apps_mod)
|
||||
|
||||
# Create user_service module with TenantService stub if not already exists
|
||||
if "api.db.services.user_service" not in sys.modules:
|
||||
user_service_mod = ModuleType("api.db.services.user_service")
|
||||
user_service_mod.UserService = SimpleNamespace() # Dummy UserService class
|
||||
user_service_mod.TenantService = SimpleNamespace(
|
||||
get_info_by=lambda _uid: [],
|
||||
get_by_id=lambda _uid: (False, None)
|
||||
)
|
||||
user_service_mod.UserTenantService = SimpleNamespace(
|
||||
query=lambda **_kwargs: []
|
||||
)
|
||||
monkeypatch.setitem(sys.modules, "api.db.services.user_service", user_service_mod)
|
||||
|
||||
module_name = "test_conversation_routes_unit_module"
|
||||
module_path = repo_root / "api" / "apps" / "conversation_app.py"
|
||||
spec = importlib.util.spec_from_file_location(module_name, module_path)
|
||||
@@ -519,15 +532,15 @@ def test_sequence2txt_validation_and_transcription_paths(monkeypatch):
|
||||
|
||||
wav_file = _DummyUploadedFile("audio.wav")
|
||||
monkeypatch.setattr(module, "request", _DummyRequest(form={"stream": "false"}, files={"file": wav_file}))
|
||||
monkeypatch.setattr(module.TenantService, "get_info_by", lambda _uid: [])
|
||||
monkeypatch.setattr(sys.modules["api.db.joint_services.tenant_model_service"].TenantService, "get_by_id", lambda _uid: (False, None))
|
||||
res = _run(module.sequence2txt())
|
||||
assert res["message"] == "Tenant not found!"
|
||||
assert res["message"] == "Tenant not found"
|
||||
|
||||
wav_file = _DummyUploadedFile("audio.wav")
|
||||
monkeypatch.setattr(module, "request", _DummyRequest(form={"stream": "false"}, files={"file": wav_file}))
|
||||
monkeypatch.setattr(module.TenantService, "get_info_by", lambda _uid: [{"tenant_id": "tenant-1", "asr_id": ""}])
|
||||
monkeypatch.setattr(sys.modules["api.db.joint_services.tenant_model_service"].TenantService, "get_by_id", lambda _uid: (True, SimpleNamespace(tenant_id="tenant-1", asr_id="")))
|
||||
res = _run(module.sequence2txt())
|
||||
assert res["message"] == "No default ASR model is set"
|
||||
assert res["message"] == "No default speech2text model is set."
|
||||
|
||||
class _SyncAsr:
|
||||
def transcription(self, _path):
|
||||
@@ -538,7 +551,8 @@ def test_sequence2txt_validation_and_transcription_paths(monkeypatch):
|
||||
|
||||
wav_file = _DummyUploadedFile("audio.wav")
|
||||
monkeypatch.setattr(module, "request", _DummyRequest(form={"stream": "false"}, files={"file": wav_file}))
|
||||
monkeypatch.setattr(module.TenantService, "get_info_by", lambda _uid: [{"tenant_id": "tenant-1", "asr_id": "asr-model"}])
|
||||
monkeypatch.setattr(sys.modules["api.db.joint_services.tenant_model_service"].TenantService, "get_by_id", lambda _uid: (True, SimpleNamespace(tenant_id="tenant-1", asr_id="asr-model")))
|
||||
monkeypatch.setattr(module.TenantLLMService, "get_api_key", lambda tenant_id, model_name: SimpleNamespace(to_dict=lambda: {"llm_factory": "test", "llm_name": "asr-model"}))
|
||||
monkeypatch.setattr(module, "LLMBundle", lambda *_args, **_kwargs: _SyncAsr())
|
||||
monkeypatch.setattr(module.os, "remove", lambda _path: (_ for _ in ()).throw(RuntimeError("remove failed")))
|
||||
res = _run(module.sequence2txt())
|
||||
@@ -579,13 +593,13 @@ def test_sequence2txt_validation_and_transcription_paths(monkeypatch):
|
||||
def test_tts_request_parse_entry(monkeypatch):
|
||||
module = _load_conversation_module(monkeypatch)
|
||||
_set_request_json(monkeypatch, module, {"text": "A。B"})
|
||||
monkeypatch.setattr(module.TenantService, "get_info_by", lambda _uid: [])
|
||||
monkeypatch.setattr(sys.modules["api.db.joint_services.tenant_model_service"].TenantService, "get_by_id", lambda _uid: (False, None))
|
||||
res = _run(module.tts())
|
||||
assert res["message"] == "Tenant not found!"
|
||||
assert res["message"] == "Tenant not found"
|
||||
|
||||
monkeypatch.setattr(module.TenantService, "get_info_by", lambda _uid: [{"tenant_id": "tenant-1", "tts_id": ""}])
|
||||
monkeypatch.setattr(sys.modules["api.db.joint_services.tenant_model_service"].TenantService, "get_by_id", lambda _uid: (True, SimpleNamespace(tenant_id="tenant-1", tts_id="")))
|
||||
res = _run(module.tts())
|
||||
assert res["message"] == "No default TTS model is set"
|
||||
assert res["message"] == "No default tts model is set."
|
||||
|
||||
class _TTSOk:
|
||||
def tts(self, txt):
|
||||
@@ -593,7 +607,8 @@ def test_tts_request_parse_entry(monkeypatch):
|
||||
return []
|
||||
yield f"chunk-{txt}".encode("utf-8")
|
||||
|
||||
monkeypatch.setattr(module.TenantService, "get_info_by", lambda _uid: [{"tenant_id": "tenant-1", "tts_id": "tts-x"}])
|
||||
monkeypatch.setattr(sys.modules["api.db.joint_services.tenant_model_service"].TenantService, "get_by_id", lambda _uid: (True, SimpleNamespace(tenant_id="tenant-1", tts_id="tts-x")))
|
||||
monkeypatch.setattr(module.TenantLLMService, "get_api_key", lambda tenant_id, model_name: SimpleNamespace(to_dict=lambda: {"llm_factory": "test", "llm_name": model_name}))
|
||||
monkeypatch.setattr(module, "LLMBundle", lambda *_args, **_kwargs: _TTSOk())
|
||||
resp = _run(module.tts())
|
||||
assert resp.mimetype == "audio/mpeg"
|
||||
@@ -749,18 +764,18 @@ def test_mindmap_and_related_questions_matrix_unit(monkeypatch):
|
||||
llm_calls["options"] = options
|
||||
return "1. Alpha\n2. Beta\nignored"
|
||||
|
||||
def _fake_bundle(tenant_id, llm_type, chat_id):
|
||||
llm_calls["bundle"] = (tenant_id, llm_type, chat_id)
|
||||
def _fake_bundle(tenant_id, model_config, lang="Chinese", **kwargs):
|
||||
llm_calls["bundle"] = (tenant_id, model_config)
|
||||
return _FakeChat()
|
||||
|
||||
monkeypatch.setattr(module, "LLMBundle", _fake_bundle)
|
||||
monkeypatch.setattr(module, "load_prompt", lambda name: f"prompt-{name}")
|
||||
monkeypatch.setattr(module.TenantLLMService, "get_api_key", lambda tenant_id, model_name: SimpleNamespace(to_dict=lambda: {"llm_factory": "test", "llm_name": model_name}))
|
||||
_set_request_json(monkeypatch, module, {"question": "solar", "search_id": "search-1"})
|
||||
res = _run(module.related_questions.__wrapped__())
|
||||
assert res["code"] == 0
|
||||
assert res["data"] == ["Alpha", "Beta"]
|
||||
assert llm_calls["bundle"][0] == "user-1"
|
||||
assert llm_calls["bundle"][2] == "chat-x"
|
||||
assert llm_calls["options"] == {"temperature": 0.2}
|
||||
assert llm_calls["prompt"] == "prompt-related_question"
|
||||
assert "Keywords: solar" in llm_calls["messages"][0]["content"]
|
||||
|
||||
@@ -137,11 +137,34 @@ def _load_dialog_module(monkeypatch):
|
||||
|
||||
tenant_llm_service_mod = ModuleType("api.db.services.tenant_llm_service")
|
||||
|
||||
class _MockTableObject:
|
||||
def __init__(self, **kwargs):
|
||||
for key, value in kwargs.items():
|
||||
setattr(self, key, value)
|
||||
|
||||
def to_dict(self):
|
||||
return {k: v for k, v in self.__dict__.items()}
|
||||
|
||||
class _TenantLLMService:
|
||||
@staticmethod
|
||||
def split_model_name_and_factory(embd_id):
|
||||
return embd_id.split("@")
|
||||
|
||||
@staticmethod
|
||||
def get_api_key(tenant_id, model_name):
|
||||
return _MockTableObject(
|
||||
id=1,
|
||||
tenant_id=tenant_id,
|
||||
llm_factory="",
|
||||
model_type="chat",
|
||||
llm_name=model_name,
|
||||
api_key="fake-api-key",
|
||||
api_base="https://api.example.com",
|
||||
max_tokens=8192,
|
||||
used_tokens=0,
|
||||
status=1
|
||||
)
|
||||
|
||||
tenant_llm_service_mod.TenantLLMService = _TenantLLMService
|
||||
monkeypatch.setitem(sys.modules, "api.db.services.tenant_llm_service", tenant_llm_service_mod)
|
||||
|
||||
@@ -253,8 +276,8 @@ def test_set_dialog_branch_matrix_unit(monkeypatch):
|
||||
|
||||
monkeypatch.setattr(module, "duplicate_name", _dup_name)
|
||||
monkeypatch.setattr(module.DialogService, "query", lambda **_kwargs: [SimpleNamespace(name="new dialog")])
|
||||
monkeypatch.setattr(module.TenantService, "get_by_id", lambda _id: (True, SimpleNamespace(llm_id="llm-x")))
|
||||
monkeypatch.setattr(module.KnowledgebaseService, "get_by_ids", lambda _ids: [SimpleNamespace(embd_id="embd-a@builtin")])
|
||||
monkeypatch.setattr(module.TenantService, "get_by_id", lambda _id: (True, SimpleNamespace(llm_id="llm-x", tenant_llm_id=1)))
|
||||
monkeypatch.setattr(module.KnowledgebaseService, "get_by_ids", lambda _ids: [SimpleNamespace(embd_id="embd-a@builtin", tenant_embd_id=2)])
|
||||
monkeypatch.setattr(module.TenantLLMService, "split_model_name_and_factory", lambda embd_id: embd_id.split("@"))
|
||||
monkeypatch.setattr(module.DialogService, "save", lambda **kwargs: captured.update(kwargs) or False)
|
||||
_set_request_json(
|
||||
@@ -301,7 +324,7 @@ def test_set_dialog_branch_matrix_unit(monkeypatch):
|
||||
res = _run(handler())
|
||||
assert res["message"] == "Tenant not found!"
|
||||
|
||||
monkeypatch.setattr(module.TenantService, "get_by_id", lambda _id: (True, SimpleNamespace(llm_id="llm-x")))
|
||||
monkeypatch.setattr(module.TenantService, "get_by_id", lambda _id: (True, SimpleNamespace(llm_id="llm-x", tenant_llm_id=1)))
|
||||
monkeypatch.setattr(
|
||||
module,
|
||||
"get_request_json",
|
||||
@@ -316,7 +339,7 @@ def test_set_dialog_branch_matrix_unit(monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
module.KnowledgebaseService,
|
||||
"get_by_ids",
|
||||
lambda _ids: [SimpleNamespace(embd_id="embd-a@f1"), SimpleNamespace(embd_id="embd-b@f2")],
|
||||
lambda _ids: [SimpleNamespace(embd_id="embd-a@f1", tenant_embd_id=2), SimpleNamespace(embd_id="embd-b@f2", tenant_embd_id=2)],
|
||||
)
|
||||
monkeypatch.setattr(module.TenantLLMService, "split_model_name_and_factory", lambda embd_id: embd_id.split("@"))
|
||||
res = _run(handler())
|
||||
|
||||
@@ -51,11 +51,19 @@ class _DummyTenantLLMModel:
|
||||
llm_factory = _ExprField("llm_factory")
|
||||
llm_name = _ExprField("llm_name")
|
||||
|
||||
def __init__(self, id=None, **kwargs):
|
||||
self.id = id
|
||||
self.api_key = None
|
||||
self.status = None
|
||||
for key, value in kwargs.items():
|
||||
setattr(self, key, value)
|
||||
|
||||
|
||||
class _TenantLLMRow:
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
id,
|
||||
llm_name,
|
||||
llm_factory,
|
||||
model_type,
|
||||
@@ -65,6 +73,7 @@ class _TenantLLMRow:
|
||||
api_base="",
|
||||
max_tokens=8192,
|
||||
):
|
||||
self.id = id
|
||||
self.llm_name = llm_name
|
||||
self.llm_factory = llm_factory
|
||||
self.model_type = model_type
|
||||
@@ -76,6 +85,7 @@ class _TenantLLMRow:
|
||||
|
||||
def to_dict(self):
|
||||
return {
|
||||
"id": self.id,
|
||||
"llm_name": self.llm_name,
|
||||
"llm_factory": self.llm_factory,
|
||||
"model_type": self.model_type,
|
||||
@@ -246,8 +256,8 @@ def test_list_app_grouping_availability_and_merge(monkeypatch):
|
||||
monkeypatch.setattr(module.TenantLLMService, "ensure_mineru_from_env", lambda tenant_id: ensure_calls.append(tenant_id))
|
||||
|
||||
tenant_rows = [
|
||||
_TenantLLMRow(llm_name="fast-emb", llm_factory="FastEmbed", model_type="embedding", api_key="k1", status="1"),
|
||||
_TenantLLMRow(llm_name="tenant-only", llm_factory="CustomFactory", model_type="chat", api_key="k2", status="1"),
|
||||
_TenantLLMRow(id=1, llm_name="fast-emb", llm_factory="FastEmbed", model_type="embedding", api_key="k1", status="1"),
|
||||
_TenantLLMRow(id=2, llm_name="tenant-only", llm_factory="CustomFactory", model_type="chat", api_key="k2", status="1"),
|
||||
]
|
||||
monkeypatch.setattr(module.TenantLLMService, "query", lambda **_kwargs: tenant_rows)
|
||||
|
||||
@@ -263,7 +273,7 @@ def test_list_app_grouping_availability_and_merge(monkeypatch):
|
||||
monkeypatch.setenv("TEI_MODEL", "tei-embed")
|
||||
|
||||
res = _run(module.list_app())
|
||||
assert res["code"] == 0
|
||||
assert res["code"] == 0, res["message"]
|
||||
assert ensure_calls == ["tenant-1"]
|
||||
|
||||
data = res["data"]
|
||||
@@ -291,8 +301,8 @@ def test_list_app_model_type_filter(monkeypatch):
|
||||
module.TenantLLMService,
|
||||
"query",
|
||||
lambda **_kwargs: [
|
||||
_TenantLLMRow(llm_name="fast-emb", llm_factory="FastEmbed", model_type="embedding", api_key="k1", status="1"),
|
||||
_TenantLLMRow(llm_name="tenant-only", llm_factory="CustomFactory", model_type="chat", api_key="k2", status="1"),
|
||||
_TenantLLMRow(id=1, llm_name="fast-emb", llm_factory="FastEmbed", model_type="embedding", api_key="k1", status="1"),
|
||||
_TenantLLMRow(id=2, llm_name="tenant-only", llm_factory="CustomFactory", model_type="chat", api_key="k2", status="1"),
|
||||
],
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
@@ -306,7 +316,7 @@ def test_list_app_model_type_filter(monkeypatch):
|
||||
|
||||
monkeypatch.setattr(module, "request", SimpleNamespace(args={"model_type": "chat"}))
|
||||
res = _run(module.list_app())
|
||||
assert res["code"] == 0
|
||||
assert res["code"] == 0, res["message"]
|
||||
assert list(res["data"].keys()) == ["CustomFactory"]
|
||||
assert res["data"]["CustomFactory"][0]["model_type"] == "chat"
|
||||
|
||||
@@ -799,7 +809,7 @@ def test_add_llm_model_type_probe_and_persistence_matrix_unit(monkeypatch):
|
||||
monkeypatch.setattr(module.TenantLLMService, "filter_update", lambda _filters, _payload: False)
|
||||
monkeypatch.setattr(module.TenantLLMService, "save", lambda **kwargs: saved.append(kwargs) or True)
|
||||
res = _call({"llm_factory": "FChatPass", "llm_name": "m", "model_type": module.LLMType.CHAT.value, "api_key": "k"})
|
||||
assert res["code"] == 0
|
||||
assert res["code"] == 0, res["message"]
|
||||
assert res["data"] is True
|
||||
assert saved
|
||||
assert saved[0]["llm_factory"] == "FChatPass"
|
||||
@@ -841,6 +851,7 @@ def test_my_llms_include_details_and_exception_unit(monkeypatch):
|
||||
"query",
|
||||
lambda **_kwargs: [
|
||||
_TenantLLMRow(
|
||||
id=1,
|
||||
llm_name="chat-model",
|
||||
llm_factory="FactoryX",
|
||||
model_type="chat",
|
||||
|
||||
@@ -32,7 +32,7 @@ def add_memory_func(request, WebApiAuth):
|
||||
payload = {
|
||||
"name": f"test_memory_{i}",
|
||||
"memory_type": ["raw"] + random.choices(["semantic", "episodic", "procedural"], k=random.randint(0, 3)),
|
||||
"embd_id": "BAAI/bge-large-zh-v1.5@SILICONFLOW",
|
||||
"embd_id": "BAAI/bge-small-en-v1.5@Builtin",
|
||||
"llm_id": "glm-4-flash@ZHIPU-AI"
|
||||
}
|
||||
res = create_memory(WebApiAuth, payload)
|
||||
|
||||
@@ -45,7 +45,7 @@ class TestMemoryCreate:
|
||||
payload = {
|
||||
"name": name,
|
||||
"memory_type": ["raw"] + random.choices(["semantic", "episodic", "procedural"], k=random.randint(0, 3)),
|
||||
"embd_id": "BAAI/bge-large-zh-v1.5@SILICONFLOW",
|
||||
"embd_id": "BAAI/bge-small-en-v1.5@Builtin",
|
||||
"llm_id": "glm-4-flash@ZHIPU-AI"
|
||||
}
|
||||
res = create_memory(WebApiAuth, payload)
|
||||
@@ -68,7 +68,7 @@ class TestMemoryCreate:
|
||||
payload = {
|
||||
"name": name,
|
||||
"memory_type": ["raw"] + random.choices(["semantic", "episodic", "procedural"], k=random.randint(0, 3)),
|
||||
"embd_id": "BAAI/bge-large-zh-v1.5@SILICONFLOW",
|
||||
"embd_id": "BAAI/bge-small-en-v1.5@Builtin",
|
||||
"llm_id": "glm-4-flash@ZHIPU-AI"
|
||||
}
|
||||
res = create_memory(WebApiAuth, payload)
|
||||
@@ -80,7 +80,7 @@ class TestMemoryCreate:
|
||||
payload = {
|
||||
"name": name,
|
||||
"memory_type": ["something"],
|
||||
"embd_id": "BAAI/bge-large-zh-v1.5@SILICONFLOW",
|
||||
"embd_id": "BAAI/bge-small-en-v1.5@Builtin",
|
||||
"llm_id": "glm-4-flash@ZHIPU-AI"
|
||||
}
|
||||
res = create_memory(WebApiAuth, payload)
|
||||
@@ -92,7 +92,7 @@ class TestMemoryCreate:
|
||||
payload = {
|
||||
"name": name,
|
||||
"memory_type": ["raw"] + random.choices(["semantic", "episodic", "procedural"], k=random.randint(0, 3)),
|
||||
"embd_id": "BAAI/bge-large-zh-v1.5@SILICONFLOW",
|
||||
"embd_id": "BAAI/bge-small-en-v1.5@Builtin",
|
||||
"llm_id": "glm-4-flash@ZHIPU-AI"
|
||||
}
|
||||
res1 = create_memory(WebApiAuth, payload)
|
||||
|
||||
@@ -216,11 +216,34 @@ def _load_user_app(monkeypatch):
|
||||
|
||||
tenant_llm_service_mod = ModuleType("api.db.services.tenant_llm_service")
|
||||
|
||||
class _MockTableObject:
|
||||
def __init__(self, **kwargs):
|
||||
for key, value in kwargs.items():
|
||||
setattr(self, key, value)
|
||||
|
||||
def to_dict(self):
|
||||
return {k: v for k, v in self.__dict__.items()}
|
||||
|
||||
class _StubTenantLLMService:
|
||||
@staticmethod
|
||||
def insert_many(_payload):
|
||||
return True
|
||||
|
||||
@staticmethod
|
||||
def get_api_key(tenant_id, model_name):
|
||||
return _MockTableObject(
|
||||
id=1,
|
||||
tenant_id=tenant_id,
|
||||
llm_factory="",
|
||||
model_type="chat",
|
||||
llm_name=model_name,
|
||||
api_key="fake-api-key",
|
||||
api_base="https://api.example.com",
|
||||
max_tokens=8192,
|
||||
used_tokens=0,
|
||||
status=1
|
||||
)
|
||||
|
||||
tenant_llm_service_mod.TenantLLMService = _StubTenantLLMService
|
||||
monkeypatch.setitem(sys.modules, "api.db.services.tenant_llm_service", tenant_llm_service_mod)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user