add qdrant migration (#1046)

Co-authored-by: jyong <jyong@dify.ai>
This commit is contained in:
Jyong
2023-08-29 10:37:04 +08:00
committed by GitHub
parent ad5f27bc5f
commit a43e80dd9c
2 changed files with 91 additions and 5 deletions

View File

@@ -15,12 +15,12 @@ from models.dataset import Document as DatasetDocument
class BaseVectorIndex(BaseIndex):
def __init__(self, dataset: Dataset, embeddings: Embeddings):
super().__init__(dataset)
self._embeddings = embeddings
self._vector_store = None
def get_type(self) -> str:
raise NotImplementedError
@@ -143,7 +143,7 @@ class BaseVectorIndex(BaseIndex):
DocumentSegment.status == 'completed',
DocumentSegment.enabled == True
).all()
for segment in segments:
document = Document(
page_content=segment.content,
@@ -218,4 +218,28 @@ class BaseVectorIndex(BaseIndex):
except Exception as e:
raise e
logging.info(f"Dataset {dataset.id} recreate successfully.")
logging.info(f"Dataset {dataset.id} recreate successfully.")
def update_qdrant_dataset(self, dataset: Dataset):
logging.info(f"update_qdrant_dataset {dataset.id}")
segment = db.session.query(DocumentSegment).filter(
DocumentSegment.dataset_id == dataset.id,
DocumentSegment.status == 'completed',
DocumentSegment.enabled == True
).first()
if segment:
try:
exist = self.text_exists(segment.index_node_id)
if exist:
index_struct = {
"type": 'qdrant',
"vector_store": {"class_prefix": dataset.index_struct_dict['vector_store']['class_prefix']}
}
dataset.index_struct = json.dumps(index_struct)
db.session.commit()
except Exception as e:
raise e
logging.info(f"Dataset {dataset.id} recreate successfully.")