Austin/dify
1
0
Fork 0
mirror of https://github.com/langgenius/dify.git synced 2024-11-16 11:42:29 +08:00
Code Issues Actions 3 Packages Projects Releases Wiki Activity

add error msg for hit test (#4704)

Browse Source
This commit is contained in:
Jyong 2024-05-28 14:54:53 +08:00 committed by GitHub
parent e6f6a59f3b
commit 1b2d862973
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 81 additions and 64 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

145
api/core/rag/datasource/retrieval_service.py
View File

@ -33,6 +33,7 @@ class RetrievalService:
return [] return []
all_documents = [] all_documents = []
threads = [] threads = []
exceptions = []
# retrieval_model source with keyword # retrieval_model source with keyword
if retrival_method == 'keyword_search': if retrival_method == 'keyword_search':
keyword_thread = threading.Thread(target=RetrievalService.keyword_search, kwargs={ keyword_thread = threading.Thread(target=RetrievalService.keyword_search, kwargs={
@ -40,7 +41,8 @@ class RetrievalService:
'dataset_id': dataset_id, 'dataset_id': dataset_id,
'query': query, 'query': query,
'top_k': top_k, 'top_k': top_k,
'all_documents': all_documents 'all_documents': all_documents,
'exceptions': exceptions,
}) })
threads.append(keyword_thread) threads.append(keyword_thread)
keyword_thread.start() keyword_thread.start()
@ -54,7 +56,8 @@ class RetrievalService:
'score_threshold': score_threshold, 'score_threshold': score_threshold,
'reranking_model': reranking_model, 'reranking_model': reranking_model,
'all_documents': all_documents, 'all_documents': all_documents,
'retrival_method': retrival_method 'retrival_method': retrival_method,
'exceptions': exceptions,
}) })
threads.append(embedding_thread) threads.append(embedding_thread)
embedding_thread.start() embedding_thread.start()
@ -69,7 +72,8 @@ class RetrievalService:
'score_threshold': score_threshold, 'score_threshold': score_threshold,
'top_k': top_k, 'top_k': top_k,
'reranking_model': reranking_model, 'reranking_model': reranking_model,
'all_documents': all_documents 'all_documents': all_documents,
'exceptions': exceptions,
}) })
threads.append(full_text_index_thread) threads.append(full_text_index_thread)
full_text_index_thread.start() full_text_index_thread.start()
@ -77,6 +81,10 @@ class RetrievalService:
for thread in threads: for thread in threads:
thread.join() thread.join()
if exceptions:
exception_message = ';\n'.join(exceptions)
raise Exception(exception_message)
if retrival_method == 'hybrid_search': if retrival_method == 'hybrid_search':
data_post_processor = DataPostProcessor(str(dataset.tenant_id), reranking_model, False) data_post_processor = DataPostProcessor(str(dataset.tenant_id), reranking_model, False)
all_documents = data_post_processor.invoke( all_documents = data_post_processor.invoke(
@ -89,82 +97,91 @@ class RetrievalService:
@classmethod @classmethod
def keyword_search(cls, flask_app: Flask, dataset_id: str, query: str, def keyword_search(cls, flask_app: Flask, dataset_id: str, query: str,
top_k: int, all_documents: list): top_k: int, all_documents: list, exceptions: list):
with flask_app.app_context(): with flask_app.app_context():
dataset = db.session.query(Dataset).filter( try:
Dataset.id == dataset_id dataset = db.session.query(Dataset).filter(
).first() Dataset.id == dataset_id
).first()
keyword = Keyword( keyword = Keyword(
dataset=dataset dataset=dataset
) )
documents = keyword.search( documents = keyword.search(
query, query,
top_k=top_k top_k=top_k
) )
all_documents.extend(documents) all_documents.extend(documents)
except Exception as e:
exceptions.append(str(e))
@classmethod @classmethod
def embedding_search(cls, flask_app: Flask, dataset_id: str, query: str, def embedding_search(cls, flask_app: Flask, dataset_id: str, query: str,
top_k: int, score_threshold: Optional[float], reranking_model: Optional[dict], top_k: int, score_threshold: Optional[float], reranking_model: Optional[dict],
all_documents: list, retrival_method: str): all_documents: list, retrival_method: str, exceptions: list):
with flask_app.app_context(): with flask_app.app_context():
dataset = db.session.query(Dataset).filter( try:
Dataset.id == dataset_id dataset = db.session.query(Dataset).filter(
).first() Dataset.id == dataset_id
).first()
vector = Vector( vector = Vector(
dataset=dataset dataset=dataset
) )
documents = vector.search_by_vector( documents = vector.search_by_vector(
query, query,
search_type='similarity_score_threshold', search_type='similarity_score_threshold',
top_k=top_k, top_k=top_k,
score_threshold=score_threshold, score_threshold=score_threshold,
filter={ filter={
'group_id': [dataset.id] 'group_id': [dataset.id]
} }
) )
if documents: if documents:
if reranking_model and retrival_method == 'semantic_search': if reranking_model and retrival_method == 'semantic_search':
data_post_processor = DataPostProcessor(str(dataset.tenant_id), reranking_model, False) data_post_processor = DataPostProcessor(str(dataset.tenant_id), reranking_model, False)
all_documents.extend(data_post_processor.invoke( all_documents.extend(data_post_processor.invoke(
query=query, query=query,
documents=documents, documents=documents,
score_threshold=score_threshold, score_threshold=score_threshold,
top_n=len(documents) top_n=len(documents)
)) ))
else: else:
all_documents.extend(documents) all_documents.extend(documents)
except Exception as e:
exceptions.append(str(e))
@classmethod @classmethod
def full_text_index_search(cls, flask_app: Flask, dataset_id: str, query: str, def full_text_index_search(cls, flask_app: Flask, dataset_id: str, query: str,
top_k: int, score_threshold: Optional[float], reranking_model: Optional[dict], top_k: int, score_threshold: Optional[float], reranking_model: Optional[dict],
all_documents: list, retrival_method: str): all_documents: list, retrival_method: str, exceptions: list):
with flask_app.app_context(): with flask_app.app_context():
dataset = db.session.query(Dataset).filter( try:
Dataset.id == dataset_id dataset = db.session.query(Dataset).filter(
).first() Dataset.id == dataset_id
).first()
vector_processor = Vector( vector_processor = Vector(
dataset=dataset, dataset=dataset,
) )
documents = vector_processor.search_by_full_text( documents = vector_processor.search_by_full_text(
query, query,
top_k=top_k top_k=top_k
) )
if documents: if documents:
if reranking_model and retrival_method == 'full_text_search': if reranking_model and retrival_method == 'full_text_search':
data_post_processor = DataPostProcessor(str(dataset.tenant_id), reranking_model, False) data_post_processor = DataPostProcessor(str(dataset.tenant_id), reranking_model, False)
all_documents.extend(data_post_processor.invoke( all_documents.extend(data_post_processor.invoke(
query=query, query=query,
documents=documents, documents=documents,
score_threshold=score_threshold, score_threshold=score_threshold,
top_n=len(documents) top_n=len(documents)
)) ))
else: else:
all_documents.extend(documents) all_documents.extend(documents)
except Exception as e:
exceptions.append(str(e))