| @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']: |
| 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) |
| 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( |
| """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), |
| 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 |
| 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'] | ||||
| 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 | ||||
| ) | ) | ||||