Explorar el Código

Fix/vdb index issue (#1776)

Co-authored-by: jyong <jyong@dify.ai>
tags/0.3.34
Jyong hace 1 año
padre
commit
2fd56cb01c
No account linked to committer's email address

+ 0
- 1
api/controllers/console/app/annotation.py Ver fichero

@setup_required @setup_required
@login_required @login_required
@account_initialization_required @account_initialization_required
@cloud_edition_billing_resource_check('annotation')
def delete(self, app_id, annotation_id): def delete(self, app_id, annotation_id):
# The role of the current user in the ta table must be admin or owner # The role of the current user in the ta table must be admin or owner
if current_user.current_tenant.current_role not in ['admin', 'owner']: if current_user.current_tenant.current_role not in ['admin', 'owner']:

+ 1
- 1
api/controllers/console/wraps.py Ver fichero

abort(403, error_msg) abort(403, error_msg)
elif resource == 'workspace_custom' and not billing_info['can_replace_logo']: elif resource == 'workspace_custom' and not billing_info['can_replace_logo']:
abort(403, error_msg) abort(403, error_msg)
elif resource == 'annotation' and 0 < annotation_quota_limit['limit'] <= annotation_quota_limit['size']:
elif resource == 'annotation' and 0 < annotation_quota_limit['limit'] < annotation_quota_limit['size']:
abort(403, error_msg) abort(403, error_msg)
else: else:
return view(*args, **kwargs) return view(*args, **kwargs)

+ 2
- 1
api/core/completion.py Ver fichero

vector_index = VectorIndex( vector_index = VectorIndex(
dataset=dataset, dataset=dataset,
config=current_app.config, config=current_app.config,
embeddings=embeddings
embeddings=embeddings,
attributes=['doc_id', 'annotation_id', 'app_id']
) )


documents = vector_index.search( documents = vector_index.search(

+ 0
- 1
api/core/index/vector_index/milvus_vector_index.py Ver fichero

"""Only for created index.""" """Only for created index."""
if self._vector_store: if self._vector_store:
return self._vector_store return self._vector_store
attributes = ['doc_id', 'dataset_id', 'document_id']


return MilvusVectorStore( return MilvusVectorStore(
collection_name=self.get_index_name(self.dataset), collection_name=self.get_index_name(self.dataset),

+ 10
- 4
api/core/index/vector_index/vector_index.py Ver fichero





class VectorIndex: class VectorIndex:
def __init__(self, dataset: Dataset, config: dict, embeddings: Embeddings):
def __init__(self, dataset: Dataset, config: dict, embeddings: Embeddings,
attributes: list = None):
if attributes is None:
attributes = ['doc_id', 'dataset_id', 'document_id', 'doc_hash']
self._dataset = dataset self._dataset = dataset
self._embeddings = embeddings self._embeddings = embeddings
self._vector_index = self._init_vector_index(dataset, config, embeddings)
self._vector_index = self._init_vector_index(dataset, config, embeddings, attributes)
self._attributes = attributes


def _init_vector_index(self, dataset: Dataset, config: dict, embeddings: Embeddings) -> BaseVectorIndex:
def _init_vector_index(self, dataset: Dataset, config: dict, embeddings: Embeddings,
attributes: list) -> BaseVectorIndex:
vector_type = config.get('VECTOR_STORE') vector_type = config.get('VECTOR_STORE')


if self._dataset.index_struct_dict: if self._dataset.index_struct_dict:
api_key=config.get('WEAVIATE_API_KEY'), api_key=config.get('WEAVIATE_API_KEY'),
batch_size=int(config.get('WEAVIATE_BATCH_SIZE')) batch_size=int(config.get('WEAVIATE_BATCH_SIZE'))
), ),
embeddings=embeddings
embeddings=embeddings,
attributes=attributes
) )
elif vector_type == "qdrant": elif vector_type == "qdrant":
from core.index.vector_index.qdrant_vector_index import QdrantVectorIndex, QdrantConfig from core.index.vector_index.qdrant_vector_index import QdrantVectorIndex, QdrantConfig

+ 3
- 2
api/core/index/vector_index/weaviate_vector_index.py Ver fichero



class WeaviateVectorIndex(BaseVectorIndex): class WeaviateVectorIndex(BaseVectorIndex):


def __init__(self, dataset: Dataset, config: WeaviateConfig, embeddings: Embeddings):
def __init__(self, dataset: Dataset, config: WeaviateConfig, embeddings: Embeddings, attributes: list):
super().__init__(dataset, embeddings) super().__init__(dataset, embeddings)
self._client = self._init_client(config) self._client = self._init_client(config)
self._attributes = attributes


def _init_client(self, config: WeaviateConfig) -> weaviate.Client: def _init_client(self, config: WeaviateConfig) -> weaviate.Client:
auth_config = weaviate.auth.AuthApiKey(api_key=config.api_key) auth_config = weaviate.auth.AuthApiKey(api_key=config.api_key)
if self._vector_store: if self._vector_store:
return self._vector_store return self._vector_store


attributes = ['doc_id', 'dataset_id', 'document_id', 'doc_hash']
attributes = self._attributes
if self._is_origin(): if self._is_origin():
attributes = ['doc_id'] attributes = ['doc_id']



+ 2
- 0
api/tasks/annotation/add_annotation_to_index_task.py Ver fichero

id=app_id, id=app_id,
tenant_id=tenant_id, tenant_id=tenant_id,
indexing_technique='high_quality', indexing_technique='high_quality',
embedding_model_provider=dataset_collection_binding.provider_name,
embedding_model=dataset_collection_binding.model_name,
collection_binding_id=dataset_collection_binding.id collection_binding_id=dataset_collection_binding.id
) )



Cargando…
Cancelar
Guardar