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 4 Packages Projects Releases Wiki Activity

fix segment enable service api (#10445)

Browse Source
This commit is contained in:
Jyong 2024-11-08 17:09:05 +08:00 committed by GitHub
parent 919275cc58
commit 888d7e6422
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 24 additions and 25 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

39
api/services/dataset_service.py
View File

@ -14,8 +14,6 @@ from configs import dify_config
from core.errors.error import LLMBadRequestError, ProviderTokenNotInitError from core.errors.error import LLMBadRequestError, ProviderTokenNotInitError
from core.model_manager import ModelManager from core.model_manager import ModelManager
from core.model_runtime.entities.model_entities import ModelType from core.model_runtime.entities.model_entities import ModelType
from core.rag.datasource.keyword.keyword_factory import Keyword
from core.rag.models.document import Document as RAGDocument
from core.rag.retrieval.retrieval_methods import RetrievalMethod from core.rag.retrieval.retrieval_methods import RetrievalMethod
from events.dataset_event import dataset_was_deleted from events.dataset_event import dataset_was_deleted
from events.document_event import document_was_deleted from events.document_event import document_was_deleted
@ -37,6 +35,7 @@ from models.dataset import (
) )
from models.model import UploadFile from models.model import UploadFile
from models.source import DataSourceOauthBinding from models.source import DataSourceOauthBinding
from services.entities.knowledge_entities.knowledge_entities import SegmentUpdateEntity
from services.errors.account import NoPermissionError from services.errors.account import NoPermissionError
from services.errors.dataset import DatasetNameDuplicateError from services.errors.dataset import DatasetNameDuplicateError
from services.errors.document import DocumentIndexingError from services.errors.document import DocumentIndexingError
@ -1503,12 +1502,13 @@ class SegmentService:
@classmethod @classmethod
def update_segment(cls, args: dict, segment: DocumentSegment, document: Document, dataset: Dataset): def update_segment(cls, args: dict, segment: DocumentSegment, document: Document, dataset: Dataset):
segment_update_entity = SegmentUpdateEntity(**args)
indexing_cache_key = "segment_{}_indexing".format(segment.id) indexing_cache_key = "segment_{}_indexing".format(segment.id)
cache_result = redis_client.get(indexing_cache_key) cache_result = redis_client.get(indexing_cache_key)
if cache_result is not None: if cache_result is not None:
raise ValueError("Segment is indexing, please try again later") raise ValueError("Segment is indexing, please try again later")
if "enabled" in args and args["enabled"] is not None: if segment_update_entity.enabled is not None:
action = args["enabled"] action = segment_update_entity.enabled
if segment.enabled != action: if segment.enabled != action:
if not action: if not action:
segment.enabled = action segment.enabled = action
@ -1521,37 +1521,26 @@ class SegmentService:
disable_segment_from_index_task.delay(segment.id) disable_segment_from_index_task.delay(segment.id)
return segment return segment
if not segment.enabled: if not segment.enabled:
if "enabled" in args and args["enabled"] is not None: if segment_update_entity.enabled is not None:
if not args["enabled"]: if not segment_update_entity.enabled:
raise ValueError("Can't update disabled segment") raise ValueError("Can't update disabled segment")
else: else:
raise ValueError("Can't update disabled segment") raise ValueError("Can't update disabled segment")
try: try:
content = args["content"] content = segment_update_entity.content
if segment.content == content: if segment.content == content:
if document.doc_form == "qa_model": if document.doc_form == "qa_model":
segment.answer = args["answer"] segment.answer = segment_update_entity.answer
if args.get("keywords"): if segment_update_entity.keywords:
segment.keywords = args["keywords"] segment.keywords = segment_update_entity.keywords
segment.enabled = True segment.enabled = True
segment.disabled_at = None segment.disabled_at = None
segment.disabled_by = None segment.disabled_by = None
db.session.add(segment) db.session.add(segment)
db.session.commit() db.session.commit()
# update segment index task # update segment index task
if "keywords" in args: if segment_update_entity.enabled:
keyword = Keyword(dataset) VectorService.create_segments_vector([segment_update_entity.keywords], [segment], dataset)
keyword.delete_by_ids([segment.index_node_id])
document = RAGDocument(
page_content=segment.content,
metadata={
"doc_id": segment.index_node_id,
"doc_hash": segment.index_node_hash,
"document_id": segment.document_id,
"dataset_id": segment.dataset_id,
},
)
keyword.add_texts([document], keywords_list=[args["keywords"]])
else: else:
segment_hash = helper.generate_text_hash(content) segment_hash = helper.generate_text_hash(content)
tokens = 0 tokens = 0
@ -1579,11 +1568,11 @@ class SegmentService:
segment.disabled_at = None segment.disabled_at = None
segment.disabled_by = None segment.disabled_by = None
if document.doc_form == "qa_model": if document.doc_form == "qa_model":
segment.answer = args["answer"] segment.answer = segment_update_entity.answer
db.session.add(segment) db.session.add(segment)
db.session.commit() db.session.commit()
# update segment vector index # update segment vector index
VectorService.update_segment_vector(args["keywords"], segment, dataset) VectorService.update_segment_vector(segment_update_entity.keywords, segment, dataset)
except Exception as e: except Exception as e:
logging.exception("update segment index failed") logging.exception("update segment index failed")

10
api/services/entities/knowledge_entities/knowledge_entities.py Normal file
View File

@ -0,0 +1,10 @@
from typing import Optional
from pydantic import BaseModel
class SegmentUpdateEntity(BaseModel):
content: str
answer: Optional[str] = None
keywords: Optional[list[str]] = None
enabled: Optional[bool] = None