From 78d05856d23c809dd4073e300af4ae7bba20ae47 Mon Sep 17 00:00:00 2001 From: Julian Whiting Date: Tue, 8 Oct 2024 11:55:06 -0400 Subject: [PATCH] abstract class for drift context --- graphrag/query/context_builder/builders.py | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/graphrag/query/context_builder/builders.py b/graphrag/query/context_builder/builders.py index 7a4ba277..23887c03 100644 --- a/graphrag/query/context_builder/builders.py +++ b/graphrag/query/context_builder/builders.py @@ -4,7 +4,7 @@ """Base classes for global and local context builders.""" from abc import ABC, abstractmethod - +from typing import Any import pandas as pd from graphrag.query.context_builder.conversation_history import ( @@ -33,3 +33,15 @@ class LocalContextBuilder(ABC): **kwargs, ) -> tuple[str | list[str], dict[str, pd.DataFrame]]: """Build the context for the local search mode.""" + + +class DRIFTContextBuilder(ABC): + """Base class for DRIFT-search context builders.""" + + @abstractmethod + def build_primer_context( + self, + query: str, + **kwargs, + ) -> pd.DataFrame: + """Build the context for the primer search actions"""