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"]