Feat/support parent child chunk (#12092)

This commit is contained in:
Jyong
2024-12-25 19:49:07 +08:00
committed by GitHub
parent 017d7538ae
commit 9231fdbf4c
54 changed files with 2578 additions and 808 deletions

View File

@@ -7,7 +7,7 @@ from core.rag.models.document import Document
from core.rag.retrieval.retrieval_methods import RetrievalMethod
from extensions.ext_database import db
from models.account import Account
from models.dataset import Dataset, DatasetQuery, DocumentSegment
from models.dataset import Dataset, DatasetQuery
default_retrieval_model = {
"search_method": RetrievalMethod.SEMANTIC_SEARCH.value,
@@ -69,7 +69,7 @@ class HitTestingService:
db.session.add(dataset_query)
db.session.commit()
return dict(cls.compact_retrieve_response(dataset, query, all_documents))
return cls.compact_retrieve_response(query, all_documents)
@classmethod
def external_retrieve(
@@ -106,41 +106,14 @@ class HitTestingService:
return dict(cls.compact_external_retrieve_response(dataset, query, all_documents))
@classmethod
def compact_retrieve_response(cls, dataset: Dataset, query: str, documents: list[Document]):
records = []
for document in documents:
if document.metadata is None:
continue
index_node_id = document.metadata["doc_id"]
segment = (
db.session.query(DocumentSegment)
.filter(
DocumentSegment.dataset_id == dataset.id,
DocumentSegment.enabled == True,
DocumentSegment.status == "completed",
DocumentSegment.index_node_id == index_node_id,
)
.first()
)
if not segment:
continue
record = {
"segment": segment,
"score": document.metadata.get("score", None),
}
records.append(record)
def compact_retrieve_response(cls, query: str, documents: list[Document]):
records = RetrievalService.format_retrieval_documents(documents)
return {
"query": {
"content": query,
},
"records": records,
"records": [record.model_dump() for record in records],
}
@classmethod