From 9c4619af432ef552220007afaf8185faf374ad30 Mon Sep 17 00:00:00 2001
From: Julian Whiting <j2whitin@gmail.com>
Date: Tue, 8 Oct 2024 11:53:48 -0400
Subject: [PATCH] accept drift context in search base

---
 graphrag/query/structured_search/base.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/graphrag/query/structured_search/base.py b/graphrag/query/structured_search/base.py
index e74a03f6..099dea39 100644
--- a/graphrag/query/structured_search/base.py
+++ b/graphrag/query/structured_search/base.py
@@ -14,6 +14,7 @@ import tiktoken
 from graphrag.query.context_builder.builders import (
     GlobalContextBuilder,
     LocalContextBuilder,
+    DRIFTContextBuilder,
 )
 from graphrag.query.context_builder.conversation_history import (
     ConversationHistory,
@@ -40,7 +41,7 @@ class BaseSearch(ABC):
     def __init__(
         self,
         llm: BaseLLM,
-        context_builder: GlobalContextBuilder | LocalContextBuilder,
+        context_builder: GlobalContextBuilder | LocalContextBuilder | DRIFTContextBuilder,
         token_encoder: tiktoken.Encoding | None = None,
         llm_params: dict[str, Any] | None = None,
         context_builder_params: dict[str, Any] | None = None,