mirror of
https://github.com/dataelement/bisheng.git
synced 2026-08-30 17:58:00 +08:00
fix: normalize legacy chunks for shared recovery
This commit is contained in:
GuoQing Zhang
+21
-4
@@ -511,17 +511,34 @@ class KnowledgeDocumentProjectionService:
|
||||
"shared content chunk loader is unavailable"
|
||||
)
|
||||
chunks = await self.shared_content_chunk_loader(content_file)
|
||||
source_file_id = int(content_file.id)
|
||||
if not chunks:
|
||||
tried_file_ids = {source_file_id}
|
||||
for candidate in entries:
|
||||
tried_file_ids = {int(content_file.id)}
|
||||
original_knowledge_id = int(
|
||||
content_file.original_knowledge_id or 0
|
||||
)
|
||||
entry_type_priority = {
|
||||
KnowledgeFileEntryType.PUBLISH.value: 0,
|
||||
KnowledgeFileEntryType.MANAGER.value: 1,
|
||||
KnowledgeFileEntryType.SHARE.value: 2,
|
||||
}
|
||||
candidates = sorted(
|
||||
entries,
|
||||
key=lambda candidate: (
|
||||
0
|
||||
if original_knowledge_id
|
||||
and int(candidate.knowledge_id) == original_knowledge_id
|
||||
else 1,
|
||||
entry_type_priority.get(str(candidate.entry_type), 99),
|
||||
int(candidate.id),
|
||||
),
|
||||
)
|
||||
for candidate in candidates:
|
||||
candidate_id = int(candidate.id)
|
||||
if candidate_id in tried_file_ids:
|
||||
continue
|
||||
tried_file_ids.add(candidate_id)
|
||||
chunks = await self.shared_content_chunk_loader(candidate)
|
||||
if chunks:
|
||||
source_file_id = candidate_id
|
||||
logger.info(
|
||||
"shared content projection recovered chunks from active entry "
|
||||
"tenant_id=%s document_id=%s content_file_id=%s "
|
||||
|
||||
@@ -151,22 +151,94 @@ async def load_shared_content_chunks_from_legacy(
|
||||
|
||||
if not rows:
|
||||
continue
|
||||
chunks: list[SharedContentChunk] = []
|
||||
|
||||
desired_generation = int(
|
||||
content_file.desired_content_generation or 0
|
||||
)
|
||||
|
||||
def row_generation(row: dict) -> int | None:
|
||||
user_metadata = row.get("user_metadata")
|
||||
if not isinstance(user_metadata, dict):
|
||||
return None
|
||||
value = user_metadata.get("content_generation")
|
||||
try:
|
||||
return int(value) if value is not None else None
|
||||
except (TypeError, ValueError):
|
||||
return None
|
||||
|
||||
generations = {
|
||||
generation
|
||||
for row in rows
|
||||
if (generation := row_generation(row)) is not None
|
||||
}
|
||||
selected_generation = (
|
||||
desired_generation
|
||||
if desired_generation in generations
|
||||
else max(generations, default=None)
|
||||
)
|
||||
if selected_generation is not None:
|
||||
rows = [
|
||||
row
|
||||
for row in rows
|
||||
if row_generation(row) == selected_generation
|
||||
]
|
||||
|
||||
chunks_by_index: dict[int, SharedContentChunk] = {}
|
||||
corrupted = False
|
||||
for offset, row in enumerate(rows):
|
||||
metadata = {
|
||||
key: value
|
||||
for key, value in row.items()
|
||||
if key not in {"pk", "text", "vector"}
|
||||
}
|
||||
chunks.append(
|
||||
SharedContentChunk(
|
||||
chunk_index=int(row.get("chunk_index", offset) or offset),
|
||||
text=str(row.get("text", "")),
|
||||
vector=row.get("vector"),
|
||||
metadata=metadata,
|
||||
)
|
||||
raw_chunk_index = row.get("chunk_index")
|
||||
chunk_index = (
|
||||
offset
|
||||
if raw_chunk_index is None
|
||||
else int(raw_chunk_index)
|
||||
)
|
||||
return sorted(chunks, key=lambda chunk: chunk.chunk_index)
|
||||
chunk = SharedContentChunk(
|
||||
chunk_index=chunk_index,
|
||||
text=str(row.get("text", "")),
|
||||
vector=row.get("vector"),
|
||||
metadata=metadata,
|
||||
)
|
||||
existing = chunks_by_index.get(chunk_index)
|
||||
if existing is None:
|
||||
chunks_by_index[chunk_index] = chunk
|
||||
continue
|
||||
existing_vector = (
|
||||
tuple(existing.vector)
|
||||
if existing.vector is not None
|
||||
else None
|
||||
)
|
||||
chunk_vector = (
|
||||
tuple(chunk.vector) if chunk.vector is not None else None
|
||||
)
|
||||
existing_sparse = tuple(
|
||||
sorted((existing.sparse_vector or {}).items())
|
||||
)
|
||||
chunk_sparse = tuple(sorted((chunk.sparse_vector or {}).items()))
|
||||
if (
|
||||
existing.text != chunk.text
|
||||
or existing_vector != chunk_vector
|
||||
or existing_sparse != chunk_sparse
|
||||
):
|
||||
corrupted = True
|
||||
break
|
||||
indexes = sorted(chunks_by_index)
|
||||
if corrupted or indexes != list(range(len(indexes))):
|
||||
logger.warning(
|
||||
"skip corrupted legacy chunk source file_id=%s knowledge_id=%s "
|
||||
"generation=%s rows=%s unique_indexes=%s",
|
||||
content_file.id,
|
||||
knowledge.id,
|
||||
selected_generation,
|
||||
len(rows),
|
||||
len(indexes),
|
||||
)
|
||||
continue
|
||||
return [chunks_by_index[index] for index in indexes]
|
||||
return []
|
||||
|
||||
return await asyncio.to_thread(_load)
|
||||
|
||||
@@ -295,6 +295,85 @@ async def test_chunk_loader_falls_back_to_original_space_after_publish_move():
|
||||
assert chunks[0].vector == [0.1, 0.2]
|
||||
|
||||
|
||||
async def test_chunk_loader_selects_current_generation_and_deduplicates():
|
||||
content_file = KnowledgeFile(
|
||||
id=100,
|
||||
tenant_id=TENANT,
|
||||
knowledge_id=20,
|
||||
original_knowledge_id=10,
|
||||
desired_content_generation=1,
|
||||
file_name="doc.pdf",
|
||||
)
|
||||
collection = SimpleNamespace(
|
||||
schema=SimpleNamespace(
|
||||
fields=[
|
||||
SimpleNamespace(name=name)
|
||||
for name in (
|
||||
"pk",
|
||||
"document_id",
|
||||
"chunk_index",
|
||||
"text",
|
||||
"vector",
|
||||
"user_metadata",
|
||||
)
|
||||
]
|
||||
),
|
||||
query=lambda **_kwargs: [
|
||||
{
|
||||
"document_id": 100,
|
||||
"chunk_index": 0,
|
||||
"text": "stale zero",
|
||||
"vector": [0.0],
|
||||
"user_metadata": {"content_generation": 0},
|
||||
},
|
||||
{
|
||||
"document_id": 100,
|
||||
"chunk_index": 0,
|
||||
"text": "current zero",
|
||||
"vector": [1.0],
|
||||
"user_metadata": {"content_generation": 1},
|
||||
},
|
||||
{
|
||||
"document_id": 100,
|
||||
"chunk_index": 1,
|
||||
"text": "current one",
|
||||
"vector": [2.0],
|
||||
"user_metadata": {"content_generation": 1},
|
||||
},
|
||||
{
|
||||
"document_id": 100,
|
||||
"chunk_index": 0,
|
||||
"text": "current zero",
|
||||
"vector": [1.0],
|
||||
"user_metadata": {"content_generation": 1},
|
||||
},
|
||||
{
|
||||
"document_id": 100,
|
||||
"chunk_index": 1,
|
||||
"text": "current one",
|
||||
"vector": [2.0],
|
||||
"user_metadata": {"content_generation": 1},
|
||||
},
|
||||
],
|
||||
)
|
||||
|
||||
with (
|
||||
patch(
|
||||
"bisheng.knowledge.domain.models.knowledge.KnowledgeDao.aquery_by_id",
|
||||
new=AsyncMock(return_value=SimpleNamespace(id=10)),
|
||||
),
|
||||
patch(
|
||||
"bisheng.knowledge.domain.knowledge_rag.KnowledgeRag."
|
||||
"init_knowledge_milvus_vectorstore_sync",
|
||||
return_value=SimpleNamespace(col=collection),
|
||||
),
|
||||
):
|
||||
chunks = await load_shared_content_chunks_from_legacy(content_file)
|
||||
|
||||
assert [chunk.chunk_index for chunk in chunks] == [0, 1]
|
||||
assert [chunk.text for chunk in chunks] == ["current zero", "current one"]
|
||||
|
||||
|
||||
async def test_content_generation_requeue_resets_exhausted_retry_state(
|
||||
async_db_session: AsyncSession,
|
||||
):
|
||||
|
||||
Reference in New Issue
Block a user