zlei6/ragflow
1
0
Fork 0
mirror of https://github.com/infiniflow/ragflow.git synced 2026-09-08 10:14:35 +08:00
Code Issues Packages Projects Releases Wiki Activity

Refactor: reformat all code for lefthook using ruff and gofmt (#16585)

Browse Source
This commit is contained in:
Wang Qi
2026-07-03 12:53:39 +08:00
committed by GitHub
parent 19fcb4a981
commit 6a4b9be426
588 changed files with 11123 additions and 15412 deletions
Download Patch File Download Diff File Expand all files Collapse all files
rag/svr/task_executor_refactor/chunk_builder.py
+1 -4
View File
@@ -126,10 +126,7 @@ async def extract_outline(cks: List[Dict], ctx: TaskContext) -> None:
ctx.write_interceptor.intercept("DocMetadataService.update_document_metadata")
else:
temp_doc = DocMetadataService.get_document_metadata(ctx.doc_id) or {}
DocMetadataService.update_document_metadata(
ctx.doc_id,
update_metadata_to({"outline": outline}, temp_doc)
)
DocMetadataService.update_document_metadata(ctx.doc_id, update_metadata_to({"outline": outline}, temp_doc))
logging.info("Persisted PDF outline (%d entries) for doc %s", len(outline), ctx.doc_id)
except Exception as e:
rag/svr/task_executor_refactor/chunk_post_processor.py
+2 -6
View File
@@ -1114,15 +1114,11 @@ async def run_document_post_chunking_if_last(
return False
chunking_aborted = is_doc_chunking_aborted(task_doc_id)
remaining_chunking_tasks = (
0 if ctx.write_interceptor
else credit_doc_chunking_task(task_doc_id, task_id)
)
remaining_chunking_tasks = 0 if ctx.write_interceptor else credit_doc_chunking_task(task_doc_id, task_id)
if remaining_chunking_tasks != 0:
if chunking_aborted:
logging.info(
"Chunking for doc %s was aborted before task %s reached post-processing; "
"skip document finalizers.",
"Chunking for doc %s was aborted before task %s reached post-processing; skip document finalizers.",
task_doc_id,
task_id,
)
rag/svr/task_executor_refactor/chunk_service.py
+20 -56
View File
@@ -109,8 +109,7 @@ class ChunkService:
ctx = self._task_context
# Validate file size
if ctx.size > settings.DOC_MAXIMUM_SIZE:
self._progress(prog=-1, msg="File size exceeds( <= %dMb )" %
(int(settings.DOC_MAXIMUM_SIZE / 1024 / 1024)))
self._progress(prog=-1, msg="File size exceeds( <= %dMb )" % (int(settings.DOC_MAXIMUM_SIZE / 1024 / 1024)))
self._task_context.recording_context.record("file_size_exceeded", True)
return []
ctx.recording_context.record("file_size_exceeded", False)
@@ -123,9 +122,7 @@ class ChunkService:
chunk_config = {
"parser_id": ctx.parser_id,
"chunk_token_num": ctx.parser_config.get("chunk_token_num", 128),
"overlapped_percent": normalize_overlapped_percent(
ctx.parser_config.get("overlapped_percent", 0)
),
"overlapped_percent": normalize_overlapped_percent(ctx.parser_config.get("overlapped_percent", 0)),
"delimiter": ctx.parser_config.get("delimiter", "\n!?。；！？"),
"from_page": ctx.from_page,
"to_page": ctx.to_page,
@@ -160,9 +157,7 @@ class ChunkService:
questions = [d for d in docs if d.get("question_kwd")]
self._task_context.recording_context.record("questions_generated", questions)
if ctx.parser_config.get("enable_metadata", False) and (
ctx.parser_config.get("metadata") or ctx.parser_config.get("built_in_metadata")
):
if ctx.parser_config.get("enable_metadata", False) and (ctx.parser_config.get("metadata") or ctx.parser_config.get("built_in_metadata")):
await generate_metadata(docs, ctx)
metadata_list = [d for d in docs if d.get("metadata_obj")]
self._task_context.recording_context.record("metadata_list_generated", metadata_list)
@@ -183,10 +178,7 @@ class ChunkService:
"""Prepare docs and upload images to MinIO."""
ctx = self._task_context
docs = []
doc = {
"doc_id": ctx.doc_id,
"kb_id": str(ctx.kb_id)
}
doc = {"doc_id": ctx.doc_id, "kb_id": str(ctx.kb_id)}
if ctx.pagerank:
doc[PAGERANK_FLD] = int(ctx.pagerank)
@@ -197,8 +189,7 @@ class ChunkService:
try:
d = copy.deepcopy(document)
d.update(chunk)
d["id"] = xxhash.xxh64(
(chunk["content_with_weight"] + str(d["doc_id"])).encode("utf-8", "surrogatepass")).hexdigest()
d["id"] = xxhash.xxh64((chunk["content_with_weight"] + str(d["doc_id"])).encode("utf-8", "surrogatepass")).hexdigest()
d["create_time"] = str(datetime.now()).replace("T", " ")[:19]
d["create_timestamp_flt"] = datetime.now().timestamp()
@@ -215,8 +206,7 @@ class ChunkService:
await image2id(d, partial(settings.STORAGE_IMPL.put, tenant_id=ctx.tenant_id), d["id"], ctx.kb_id)
docs.append(d)
except Exception:
logging.exception(
"Saving image of chunk {}/{}/{} got exception".format(ctx.location, ctx.name, d["id"]))
logging.exception("Saving image of chunk {}/{}/{} got exception".format(ctx.location, ctx.name, d["id"]))
raise
tasks = []
@@ -303,11 +293,7 @@ class ChunkService:
mom_ck["available_int"] = 0
# Keep only essential fields
allowed_fields = [
"id", "content_with_weight", "doc_id", "docnm_kwd",
"kb_id", "available_int", "position_int",
"create_timestamp_flt", "page_num_int", "top_int"
]
allowed_fields = ["id", "content_with_weight", "doc_id", "docnm_kwd", "kb_id", "available_int", "position_int", "create_timestamp_flt", "page_num_int", "top_int"]
for fld in list(mom_ck.keys()):
if fld not in allowed_fields:
del mom_ck[fld]
@@ -326,11 +312,7 @@ class ChunkService:
) -> bool:
"""Insert mother chunks in batches."""
for b in range(0, len(mothers), doc_bulk_size):
await self._intercept_doc_store_insert(
mothers[b:b + doc_bulk_size],
search.index_name(task_tenant_id),
task_dataset_id
)
await self._intercept_doc_store_insert(mothers[b : b + doc_bulk_size], search.index_name(task_tenant_id), task_dataset_id)
if self._task_context.has_canceled_func(task_id):
self._task_context.progress_cb(-1, msg="Task has been canceled.")
@@ -346,7 +328,7 @@ class ChunkService:
async def _intercept_doc_store_insert(self, chunks: list, index_name: str, task_dataset_id: str) -> Any:
if self._task_context.write_interceptor:
if self._task_context.doc_id == GRAPH_RAPTOR_FAKE_DOC_ID: # raptor - non-determinisic
if self._task_context.doc_id == GRAPH_RAPTOR_FAKE_DOC_ID: # raptor - non-determinisic
return self._task_context.write_interceptor.intercept("docStoreConn.insert", [])
return self._task_context.write_interceptor.intercept("docStoreConn.insert")
else:
@@ -362,40 +344,28 @@ class ChunkService:
) -> bool:
"""Insert main chunks in batches with cancellation handling."""
for b in range(0, len(chunks), doc_bulk_size):
doc_store_result = await self._intercept_doc_store_insert(
chunks[b:b + doc_bulk_size],
search.index_name(task_tenant_id),
task_dataset_id
)
doc_store_result = await self._intercept_doc_store_insert(chunks[b : b + doc_bulk_size], search.index_name(task_tenant_id), task_dataset_id)
if self._task_context.has_canceled_func(task_id):
# Roll back partial RAPTOR summary inserts
await self._rollback_raptor_chunks(
task_id, task_tenant_id, task_dataset_id, chunks, b, doc_bulk_size
)
await self._rollback_raptor_chunks(task_id, task_tenant_id, task_dataset_id, chunks, b, doc_bulk_size)
self._task_context.progress_cb(-1, msg="Task has been canceled.")
return False
if b % 128 == 0:
self._task_context.progress_cb(prog=0.8 + 0.1 * (b + 1) / len(chunks),msg="")
self._task_context.progress_cb(prog=0.8 + 0.1 * (b + 1) / len(chunks), msg="")
if doc_store_result:
error_message = (
f"Insert chunk error: {doc_store_result}, "
"please check log file and Elasticsearch/Infinity status!"
)
error_message = f"Insert chunk error: {doc_store_result}, please check log file and Elasticsearch/Infinity status!"
self._task_context.progress_cb(-1, msg=error_message)
raise Exception(error_message)
# Update chunk IDs in task
chunk_ids = [chunk["id"] for chunk in chunks[:b + doc_bulk_size]]
chunk_ids = [chunk["id"] for chunk in chunks[: b + doc_bulk_size]]
if not await self._update_task_chunk_ids(task_id, chunk_ids):
# Roll back on failure
await self._rollback_insertion(task_tenant_id, task_dataset_id, chunk_ids)
self._task_context.progress_cb(
-1,
msg=f"Chunk updates failed since task {task_id} is unknown."
)
self._task_context.progress_cb(-1, msg=f"Chunk updates failed since task {task_id} is unknown.")
return False
return True
@@ -410,19 +380,15 @@ class ChunkService:
doc_bulk_size: int,
):
"""Roll back partial RAPTOR summary inserts after cancellation."""
raptor_ids = [
c["id"] for c in chunks[:up_to_batch + doc_bulk_size]
if c.get("raptor_kwd") == "raptor"
]
raptor_ids = [c["id"] for c in chunks[: up_to_batch + doc_bulk_size] if c.get("raptor_kwd") == "raptor"]
if raptor_ids:
try:
await self._intercept_doc_store_delete(
{"id": raptor_ids}, search.index_name(task_tenant_id), task_dataset_id
)
await self._intercept_doc_store_delete({"id": raptor_ids}, search.index_name(task_tenant_id), task_dataset_id)
logging.info(
"insert_chunks: rolled back %d partial RAPTOR chunks after cancellation (task=%s)",
len(raptor_ids), task_id,
len(raptor_ids),
task_id,
)
except Exception:
logging.exception(
@@ -454,9 +420,7 @@ class ChunkService:
chunk_ids: List[str],
):
"""Roll back an insertion by deleting chunks and images."""
await self._intercept_doc_store_delete(
{"id": chunk_ids}, search.index_name(task_tenant_id), task_dataset_id
)
await self._intercept_doc_store_delete({"id": chunk_ids}, search.index_name(task_tenant_id), task_dataset_id)
# Delete associated images
tasks = []
rag/svr/task_executor_refactor/comparator.py
+3 -5
View File
@@ -102,19 +102,17 @@ class ContextComparator:
A new dictionary with non-deterministic fields removed.
"""
import copy
result = copy.copy(data)
for key, value in result.items():
if isinstance(value, dict):
# Create a new dict without the non-deterministic keys
cleaned = {
k: v for k, v in value.items()
if k not in self.DICT_KEYS_TO_STRIP
}
cleaned = {k: v for k, v in value.items() if k not in self.DICT_KEYS_TO_STRIP}
result[key] = cleaned
return result
@staticmethod
def _get_key_values_to_compare(prod_data_all:dict):
def _get_key_values_to_compare(prod_data_all: dict):
prod_data = dict()
for key, value in prod_data_all.items():
if key in ALLOWED_METHOD_NAMES:
rag/svr/task_executor_refactor/dataflow_service.py
+16 -44
View File
@@ -109,10 +109,7 @@ class DataflowService:
dataflow_id = corrected_id
# Run pipeline
pipeline = Pipeline(
dsl, tenant_id=ctx.tenant_id, doc_id=doc_id,
task_id=task_id, flow_id=dataflow_id
)
pipeline = Pipeline(dsl, tenant_id=ctx.tenant_id, doc_id=doc_id, task_id=task_id, flow_id=dataflow_id)
chunks = await pipeline.run(file=ctx.file) if ctx.file else await pipeline.run()
if doc_id == CANVAS_DEBUG_DOC_ID:
@@ -140,9 +137,7 @@ class DataflowService:
# Embed chunks if needed
keys = [k for o in chunks for k in list(o.keys())]
if not any([re.match(r"q_[0-9]+_vec", k) for k in keys]):
chunks, embedding_token_consumption = await self._embed_chunks(
chunks, embedding_token_consumption
)
chunks, embedding_token_consumption = await self._embed_chunks(chunks, embedding_token_consumption)
if chunks is None:
self._record_pipeline_log(doc_id, dataflow_id, pipeline)
return
@@ -157,33 +152,22 @@ class DataflowService:
# Insert chunks
start_ts = timer()
self._progress(prog=0.82, msg="[DOC Engine]:\nStart to index...")
e = await self._insert_chunks(
task_id, ctx.tenant_id, ctx.kb_id, chunks
)
e = await self._insert_chunks(task_id, ctx.tenant_id, ctx.kb_id, chunks)
if not e:
self._record_pipeline_log(doc_id, dataflow_id, pipeline)
return
time_cost = timer() - start_ts
task_time_cost = timer() - task_start_ts
self._progress(
prog=1.,
msg="Indexing done ({:.2f}s). Task done ({:.2f}s)".format(time_cost, task_time_cost)
)
self._progress(prog=1.0, msg="Indexing done ({:.2f}s). Task done ({:.2f}s)".format(time_cost, task_time_cost))
# Update document stats
if ctx.write_interceptor:
ctx.write_interceptor.intercept("DocumentService.increment_chunk_num")
else:
DocumentService.increment_chunk_num(
doc_id, task_dataset_id, embedding_token_consumption, len(chunks), task_time_cost
)
DocumentService.increment_chunk_num(doc_id, task_dataset_id, embedding_token_consumption, len(chunks), task_time_cost)
logging.info(
"[Done], chunks({}), token({}), elapsed:{:.2f}".format(
len(chunks), embedding_token_consumption, task_time_cost
)
)
logging.info("[Done], chunks({}), token({}), elapsed:{:.2f}".format(len(chunks), embedding_token_consumption, task_time_cost))
ctx.recording_context.record("dataflow_chunks", chunks)
self._record_pipeline_log(doc_id, dataflow_id, pipeline)
@@ -244,21 +228,17 @@ class DataflowService:
return []
@timeout(60)
async def _embed_chunks(
self, chunks: List[Dict], token_consumption: int
) -> Tuple[Optional[List[Dict]], int]:
async def _embed_chunks(self, chunks: List[Dict], token_consumption: int) -> Tuple[Optional[List[Dict]], int]:
"""Embed chunks using the embedding model."""
ctx = self._task_context
try:
self._progress(prog=0.82, msg="\n-------------------------------------\nStart to embedding...")
e, kb = self._get_kb_by_id(ctx.kb_id)
embedding_id = kb.embd_id
embd_model_config = get_model_config_from_provider_instance(
ctx.tenant_id, LLMType.EMBEDDING, embedding_id
)
embd_model_config = get_model_config_from_provider_instance(ctx.tenant_id, LLMType.EMBEDDING, embedding_id)
from api.db.services.llm_service import LLMBundle
with LLMBundle(ctx.tenant_id, embd_model_config) as embedding_model:
with LLMBundle(ctx.tenant_id, embd_model_config) as embedding_model:
# Prepare texts for embedding using EmbeddingUtils
texts = EmbeddingUtils.prepare_texts_for_dataflow_embedding(chunks)
delta = 0.20 / (len(texts) // self._embedding_batch_size + 1)
@@ -267,19 +247,14 @@ class DataflowService:
# Batch encode using EmbeddingUtils
vects_batches = []
for i in range(0, len(texts), self._embedding_batch_size):
batch = texts[i: i + self._embedding_batch_size]
batch = texts[i : i + self._embedding_batch_size]
async with ctx.embed_limiter:
vts, c = await thread_pool_exec(
self._encode_batch, batch, embedding_model
)
vts, c = await thread_pool_exec(self._encode_batch, batch, embedding_model)
vects_batches.append(vts)
token_consumption += c
prog += delta
if i % (len(texts) // self._embedding_batch_size / 100 + 1) == 1:
self._progress(
prog=prog,
msg=f"{i + 1} / {len(texts) // self._embedding_batch_size}"
)
self._progress(prog=prog, msg=f"{i + 1} / {len(texts) // self._embedding_batch_size}")
# Stack vectors using EmbeddingUtils
vects = EmbeddingUtils.stack_vectors(vects_batches)
@@ -358,11 +333,10 @@ class DataflowService:
else:
DocMetadataService.update_document_metadata(doc_id, metadata)
async def _insert_chunks(
self, task_id: str, tenant_id: str, kb_id: str, chunks: List[Dict]
) -> bool:
async def _insert_chunks(self, task_id: str, tenant_id: str, kb_id: str, chunks: List[Dict]) -> bool:
"""Insert chunks into document store."""
from rag.svr.task_executor_refactor.chunk_service import ChunkService
chunk_service = ChunkService(self._task_context)
return await chunk_service.insert_chunks(task_id, tenant_id, kb_id, chunks)
@@ -371,15 +345,13 @@ class DataflowService:
if self._task_context.write_interceptor:
self._task_context.write_interceptor.intercept("PipelineOperationLogService.create")
else:
PipelineOperationLogService.create(
document_id=doc_id, pipeline_id=dataflow_id,
task_type=PipelineTaskType.PARSE, dsl=str(pipeline)
)
PipelineOperationLogService.create(document_id=doc_id, pipeline_id=dataflow_id, task_type=PipelineTaskType.PARSE, dsl=str(pipeline))
@classmethod
def _get_kb_by_id(cls, kb_id: str):
"""Get knowledge base by ID."""
from api.db.services.knowledgebase_service import KnowledgebaseService
return KnowledgebaseService.get_by_id(kb_id)
def _progress(self, prog=None, msg=None):
rag/svr/task_executor_refactor/embedding_service.py
+1 -1
View File
@@ -90,7 +90,7 @@ class EmbeddingService:
# Batch encode contents using EmbeddingUtils
vects_batches = []
for i in range(0, len(contents), self._embedding_batch_size):
batch = contents[i: i + self._embedding_batch_size]
batch = contents[i : i + self._embedding_batch_size]
async with self._task_context.embed_limiter:
vts, c = await thread_pool_exec(
self._batch_encode_wrapper,
rag/svr/task_executor_refactor/embedding_utils.py
+1 -5
View File
@@ -180,11 +180,7 @@ class EmbeddingUtils:
if not title_weight:
title_weight = cls.DEFAULT_TITLE_WEIGHT
if (
title_vecs is not None
and content_vecs.ndim == 2
and title_vecs.shape == content_vecs.shape
):
if title_vecs is not None and content_vecs.ndim == 2 and title_vecs.shape == content_vecs.shape:
return title_weight * title_vecs + (1 - title_weight) * content_vecs
return content_vecs
rag/svr/task_executor_refactor/recording_context.py
+1 -1
View File
@@ -416,4 +416,4 @@ def timed_with_recording(
return wrapper
return decorator
return decorator
rag/svr/task_executor_refactor/report_generator.py
+1 -4
View File
@@ -83,10 +83,7 @@ class ComparisonReport:
if self.total_keys == 0:
return f"Task {self.task_id}: No keys to compare"
match_rate = (self.matched_keys / self.total_keys) * 100
return (
f"Task {self.task_id}: {self.matched_keys}/{self.total_keys} "
f"keys matched ({match_rate:.1f}%)"
)
return f"Task {self.task_id}: {self.matched_keys}/{self.total_keys} keys matched ({match_rate:.1f}%)"
def to_dict(self) -> dict:
"""Convert to dictionary for serialization.
rag/svr/task_executor_refactor/task_context.py
+1 -1
View File
@@ -143,6 +143,7 @@ class TaskDict(TypedDict, total=False):
message_dict: Dict[str, Any]
"""Message dictionary for memory tasks."""
# ============================================================================
# Data Classes
# ============================================================================
@@ -270,7 +271,6 @@ class TaskContext:
self._write_interceptor = write_interceptor
self._recording_context = recording_context
# Prepare progress callback and set it on the context
progress_cb = partial(
callbacks.progress,
rag/svr/task_executor_refactor/task_manager.py
+8 -5
View File
@@ -30,7 +30,8 @@ from rag.svr.task_executor_refactor.recording_context import (
BaseRecordingContext,
RecordingContext,
_NULL_RECORDING_CONTEXT,
set_recording_context, recording_context_manager,
set_recording_context,
recording_context_manager,
)
from rag.svr.task_executor_refactor.task_context import TaskContext
from rag.svr.task_executor_refactor.task_handler import TaskHandler
@@ -166,12 +167,14 @@ class TaskManager:
comp_result = comp.compare(task_context.id, recording_ctx1, recording_ctx2)
logging.info(f"-------{task_context.name}, compare result:{comp_result.to_markdown()}")
if interceptor.remaining_values_count() > 0 or comp_result.mismatched_keys > 0:
logging.info(f"------task:{task_context.id} {task_context.name} differs, "
f"interceptor.remaining_values_count():{interceptor.remaining_values_count()}, "
f"mismatched_keys:{comp_result.mismatched_keys}")
logging.info(
f"------task:{task_context.id} {task_context.name} differs, "
f"interceptor.remaining_values_count():{interceptor.remaining_values_count()}, "
f"mismatched_keys:{comp_result.mismatched_keys}"
)
if interceptor.remaining_values_count() > 0:
logging.info(f"------task:{task_context.id}, remaining values:{interceptor.remaining_values()}")
if comp_result.mismatched_keys > 0:
logging.info(f"-------compare result:{comp_result.details}")
else:
logging.info(f"------task:{task_context.id} {task_context.name} same result for prod and dry run ")
logging.info(f"------task:{task_context.id} {task_context.name} same result for prod and dry run ")
rag/svr/task_executor_refactor/write_operation_interceptor.py
+4 -8
View File
@@ -20,6 +20,7 @@ Provides a mechanism to intercept write operations during comparison mode.
The interceptor consumes pre-recorded return values (from production execution)
and returns them one by one when the corresponding methods are called.
"""
import logging
from typing import Any, Dict, List
@@ -34,7 +35,7 @@ ALLOWED_METHOD_NAMES = {
"delete_raptor_chunks",
"handle_save_to_memory_task",
"docStoreConn.insert",
"docStoreConn.delete"
"docStoreConn.delete",
}
_NO_DEFAULT = object()
@@ -80,7 +81,7 @@ class WriteOperationInterceptor:
for key in ALLOWED_METHOD_NAMES:
self._recorded_values[key] = list(recorded_values.get(key, []))
def intercept(self, method_name: str, default_value = _NO_DEFAULT) -> Any:
def intercept(self, method_name: str, default_value=_NO_DEFAULT) -> Any:
"""Intercept a method call and return the next pre-recorded value.
Args:
@@ -96,10 +97,7 @@ class WriteOperationInterceptor:
IndexError: If the recorded values list for method_name is empty.
"""
if method_name not in ALLOWED_METHOD_NAMES:
raise ValueError(
f"Cannot intercept method '{method_name}'. "
f"Allowed method names: {ALLOWED_METHOD_NAMES}"
)
raise ValueError(f"Cannot intercept method '{method_name}'. Allowed method names: {ALLOWED_METHOD_NAMES}")
if method_name not in self._recorded_values:
raise KeyError(f"No recorded values found for method '{method_name}'")
@@ -113,7 +111,6 @@ class WriteOperationInterceptor:
return values_list.pop(0)
def remaining_count(self, method_name: str) -> int:
"""Get the number of remaining recorded values for a method.
@@ -127,7 +124,6 @@ class WriteOperationInterceptor:
return 0
return len(self._recorded_values[method_name])
def remaining_values(self):
return {k: list(v) for k, v in self._recorded_values.items()}