mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-07-28 03:38:11 +08:00
Refa: GraphRAG to use async chat methods instead of thread pool execution (#14002)
### What problem does this PR solve? GraphRAG _async_chat. ### Type of change - [x] Refactoring - [x] Performance Improvement <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Refactor** * Unified chat calls to an async invocation across extractors, improving timeout handling and ensuring task IDs propagate reliably. * **Tests** * Added and expanded unit tests and mocks to cover extractor behavior, timeout scenarios, and safe test-package imports, reducing regression risk. <!-- end of auto-generated comment: release notes by coderabbit.ai -->
This commit is contained in:
@@ -32,7 +32,6 @@ from rag.graphrag.utils import perform_variable_replacements, chat_limiter, Grap
|
||||
from api.db.services.task_service import has_canceled
|
||||
from common.exceptions import TaskCanceledException
|
||||
|
||||
from common.misc_utils import thread_pool_exec
|
||||
|
||||
DEFAULT_RECORD_DELIMITER = "##"
|
||||
DEFAULT_ENTITY_INDEX_DELIMITER = "<|>"
|
||||
@@ -213,21 +212,15 @@ class EntityResolution(Extractor):
|
||||
timeout_seconds = 280 if os.environ.get("ENABLE_TIMEOUT_ASSERTION") else 1000000000
|
||||
try:
|
||||
response = await asyncio.wait_for(
|
||||
thread_pool_exec(
|
||||
self._chat,
|
||||
text,
|
||||
[{"role": "user", "content": "Output:"}],
|
||||
{},
|
||||
task_id
|
||||
),
|
||||
self._async_chat(text, [{"role": "user", "content": "Output:"}], {}, task_id),
|
||||
timeout=timeout_seconds,
|
||||
)
|
||||
|
||||
except asyncio.TimeoutError:
|
||||
logging.warning("_resolve_candidate._chat timeout, skipping...")
|
||||
logging.warning("_resolve_candidate._async_chat timeout, skipping...")
|
||||
return
|
||||
except Exception as e:
|
||||
logging.error(f"_resolve_candidate._chat failed: {e}")
|
||||
logging.error(f"_resolve_candidate._async_chat failed: {e}")
|
||||
return
|
||||
|
||||
logging.debug(f"_resolve_candidate chat prompt: {text}\nchat response: {response}")
|
||||
|
||||
Reference in New Issue
Block a user