Kevin Hu commited on
Commit
eb7da20
·
1 Parent(s): b57c1d6

Fix API retrieval error. (#4408)

Browse files

### What problem does this PR solve?

#4403

### Type of change

- [x] Bug Fix (non-breaking change which fixes an issue)

Files changed (1) hide show
  1. api/apps/sdk/doc.py +4 -8
api/apps/sdk/doc.py CHANGED
@@ -20,7 +20,7 @@ from api.db.services.dialog_service import keyword_extraction
20
  from rag.app.qa import rmPrefix, beAdoc
21
  from rag.nlp import rag_tokenizer
22
  from api.db import LLMType, ParserType
23
- from api.db.services.llm_service import TenantLLMService
24
  from api import settings
25
  import xxhash
26
  import re
@@ -1331,18 +1331,14 @@ def retrieval_test(tenant_id):
1331
  e, kb = KnowledgebaseService.get_by_id(kb_ids[0])
1332
  if not e:
1333
  return get_error_data_result(message="Dataset not found!")
1334
- embd_mdl = TenantLLMService.model_instance(
1335
- kb.tenant_id, LLMType.EMBEDDING.value, llm_name=kb.embd_id
1336
- )
1337
 
1338
  rerank_mdl = None
1339
  if req.get("rerank_id"):
1340
- rerank_mdl = TenantLLMService.model_instance(
1341
- kb.tenant_id, LLMType.RERANK.value, llm_name=req["rerank_id"]
1342
- )
1343
 
1344
  if req.get("keyword", False):
1345
- chat_mdl = TenantLLMService.model_instance(kb.tenant_id, LLMType.CHAT)
1346
  question += keyword_extraction(chat_mdl, question)
1347
 
1348
  retr = settings.retrievaler if kb.parser_id != ParserType.KG else settings.kg_retrievaler
 
20
  from rag.app.qa import rmPrefix, beAdoc
21
  from rag.nlp import rag_tokenizer
22
  from api.db import LLMType, ParserType
23
+ from api.db.services.llm_service import TenantLLMService, LLMBundle
24
  from api import settings
25
  import xxhash
26
  import re
 
1331
  e, kb = KnowledgebaseService.get_by_id(kb_ids[0])
1332
  if not e:
1333
  return get_error_data_result(message="Dataset not found!")
1334
+ embd_mdl = LLMBundle(kb.tenant_id, LLMType.EMBEDDING, llm_name=kb.embd_id)
 
 
1335
 
1336
  rerank_mdl = None
1337
  if req.get("rerank_id"):
1338
+ rerank_mdl = LLMBundle(kb.tenant_id, LLMType.RERANK, llm_name=req["rerank_id"])
 
 
1339
 
1340
  if req.get("keyword", False):
1341
+ chat_mdl = LLMBundle(kb.tenant_id, LLMType.CHAT)
1342
  question += keyword_extraction(chat_mdl, question)
1343
 
1344
  retr = settings.retrievaler if kb.parser_id != ParserType.KG else settings.kg_retrievaler