mirror of
https://github.com/langgenius/dify.git
synced 2024-11-16 19:59:50 +08:00
96 lines
3.5 KiB
Python
96 lines
3.5 KiB
Python
|
|
from typing import Optional
|
|
from flask import current_app, Flask
|
|
from langchain.embeddings.base import Embeddings
|
|
from core.index.vector_index.vector_index import VectorIndex
|
|
from core.model_providers.model_factory import ModelFactory
|
|
from extensions.ext_database import db
|
|
from models.dataset import Dataset
|
|
|
|
default_retrieval_model = {
|
|
'search_method': 'semantic_search',
|
|
'reranking_enable': False,
|
|
'reranking_model': {
|
|
'reranking_provider_name': '',
|
|
'reranking_model_name': ''
|
|
},
|
|
'top_k': 2,
|
|
'score_threshold_enabled': False
|
|
}
|
|
|
|
|
|
class RetrievalService:
|
|
|
|
@classmethod
|
|
def embedding_search(cls, flask_app: Flask, dataset_id: str, query: str,
|
|
top_k: int, score_threshold: Optional[float], reranking_model: Optional[dict],
|
|
all_documents: list, search_method: str, embeddings: Embeddings):
|
|
with flask_app.app_context():
|
|
dataset = db.session.query(Dataset).filter(
|
|
Dataset.id == dataset_id
|
|
).first()
|
|
|
|
vector_index = VectorIndex(
|
|
dataset=dataset,
|
|
config=current_app.config,
|
|
embeddings=embeddings
|
|
)
|
|
|
|
documents = vector_index.search(
|
|
query,
|
|
search_type='similarity_score_threshold',
|
|
search_kwargs={
|
|
'k': top_k,
|
|
'score_threshold': score_threshold,
|
|
'filter': {
|
|
'group_id': [dataset.id]
|
|
}
|
|
}
|
|
)
|
|
|
|
if documents:
|
|
if reranking_model and search_method == 'semantic_search':
|
|
rerank = ModelFactory.get_reranking_model(
|
|
tenant_id=dataset.tenant_id,
|
|
model_provider_name=reranking_model['reranking_provider_name'],
|
|
model_name=reranking_model['reranking_model_name']
|
|
)
|
|
all_documents.extend(rerank.rerank(query, documents, score_threshold, len(documents)))
|
|
else:
|
|
all_documents.extend(documents)
|
|
|
|
@classmethod
|
|
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],
|
|
all_documents: list, search_method: str, embeddings: Embeddings):
|
|
with flask_app.app_context():
|
|
dataset = db.session.query(Dataset).filter(
|
|
Dataset.id == dataset_id
|
|
).first()
|
|
|
|
vector_index = VectorIndex(
|
|
dataset=dataset,
|
|
config=current_app.config,
|
|
embeddings=embeddings
|
|
)
|
|
|
|
documents = vector_index.search_by_full_text_index(
|
|
query,
|
|
search_type='similarity_score_threshold',
|
|
top_k=top_k
|
|
)
|
|
if documents:
|
|
if reranking_model and search_method == 'full_text_search':
|
|
rerank = ModelFactory.get_reranking_model(
|
|
tenant_id=dataset.tenant_id,
|
|
model_provider_name=reranking_model['reranking_provider_name'],
|
|
model_name=reranking_model['reranking_model_name']
|
|
)
|
|
all_documents.extend(rerank.rerank(query, documents, score_threshold, len(documents)))
|
|
else:
|
|
all_documents.extend(documents)
|
|
|
|
|
|
|
|
|