mirror of
https://github.com/microsoft/graphrag.git
synced 2026-02-08 12:12:59 +08:00
48 lines
1.2 KiB
Python
48 lines
1.2 KiB
Python
# Copyright (c) 2024 Microsoft Corporation.
|
|
# Licensed under the MIT License
|
|
|
|
"""Base classes for global and local context builders."""
|
|
|
|
from abc import ABC, abstractmethod
|
|
|
|
import pandas as pd
|
|
|
|
from graphrag.query.context_builder.conversation_history import (
|
|
ConversationHistory,
|
|
)
|
|
|
|
|
|
class GlobalContextBuilder(ABC):
|
|
"""Base class for global-search context builders."""
|
|
|
|
@abstractmethod
|
|
def build_context(
|
|
self, conversation_history: ConversationHistory | None = None, **kwargs
|
|
) -> tuple[str | list[str], dict[str, pd.DataFrame]]:
|
|
"""Build the context for the global search mode."""
|
|
|
|
|
|
class LocalContextBuilder(ABC):
|
|
"""Base class for local-search context builders."""
|
|
|
|
@abstractmethod
|
|
def build_context(
|
|
self,
|
|
query: str,
|
|
conversation_history: ConversationHistory | None = None,
|
|
**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_context(
|
|
self,
|
|
query: str,
|
|
**kwargs,
|
|
) -> pd.DataFrame:
|
|
"""Build the context for the primer search actions."""
|