From 0f51ec48f1110c3d4fb35e78abce5d4ef4200e2d Mon Sep 17 00:00:00 2001 From: yangdx Date: Sat, 28 Jun 2025 09:18:06 +0800 Subject: [PATCH] fix: streaming error when only_need_context=True returns empty results Prevents NoneType async iteration error by handling None responses in stream_generator and ensuring kg_query returns valid strings. --- lightrag/api/routers/query_routes.py | 3 +++ lightrag/operate.py | 2 +- 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/lightrag/api/routers/query_routes.py b/lightrag/api/routers/query_routes.py index b23dbd32..69aa32d8 100644 --- a/lightrag/api/routers/query_routes.py +++ b/lightrag/api/routers/query_routes.py @@ -183,6 +183,9 @@ def create_query_routes(rag, api_key: Optional[str] = None, top_k: int = 60): if isinstance(response, str): # If it's a string, send it all at once yield f"{json.dumps({'response': response})}\n" + elif response is None: + # Handle None response (e.g., when only_need_context=True but no context found) + yield f"{json.dumps({'response': 'No relevant context found for the query.'})}\n" else: # If it's an async generator, send chunks one by one try: diff --git a/lightrag/operate.py b/lightrag/operate.py index 77568161..ad26f857 100644 --- a/lightrag/operate.py +++ b/lightrag/operate.py @@ -1390,7 +1390,7 @@ async def kg_query( ) if query_param.only_need_context: - return context + return context if context is not None else PROMPTS["fail_response"] if context is None: return PROMPTS["fail_response"]