Ruff

graphrag/query/structured_search/drift_search/action.py

@@ -86,9 +86,9 @@ class DriftAction:
             generated_tokens = 0
         else:
             generated_tokens = num_tokens(self.answer, search_engine.token_encoder)
-        self.metadata.update({
-            "token_ct": search_result.prompt_tokens + generated_tokens
-        })
+        self.metadata.update(
+            {"token_ct": search_result.prompt_tokens + generated_tokens}
+        )
 
         self.follow_ups = response.pop("follow_up_queries", [])
         if not self.follow_ups:
@@ -203,7 +203,7 @@ class DriftAction:
                 raise ValueError(error_message)
             except json.JSONDecodeError as e:
                 error_message = f"Failed to parse response string: {e}. Parsed response must be a dictionary."
-                raise ValueError(error_message)
+                raise ValueError(error_message) from e
 
         error_message = f"Unsupported response type: {type(response).__name__}. Expected a dictionary or JSON string."
         raise ValueError(error_message)

graphrag/query/structured_search/drift_search/primer.py

@@ -170,7 +170,7 @@ class DRIFTPrimer:
         return SearchResult(
             response=[response for response, _ in results_with_tokens],
             context_data={"top_k_reports": top_k_reports},
-            context_text=top_k_reports.to_json(),
+            context_text=top_k_reports.to_json() or "",
             completion_time=completion_time,
             llm_calls=len(results_with_tokens),
             prompt_tokens=sum(tokens for _, tokens in results_with_tokens),

graphrag/query/structured_search/drift_search/state.py

@@ -5,11 +5,8 @@
 
 import logging
 import random
-from typing import TYPE_CHECKING, Any
-
-if TYPE_CHECKING:
-    from typing import Callable
-
+from collections.abc import Callable
+from typing import Any
 
 import networkx as nx