diff --git a/packages/graphrag/graphrag/config/enums.py b/packages/graphrag/graphrag/config/enums.py index 5084f2154b..3238d4b8c8 100644 --- a/packages/graphrag/graphrag/config/enums.py +++ b/packages/graphrag/graphrag/config/enums.py @@ -8,6 +8,19 @@ from enum import Enum +class ModelType(str, Enum): + """The type of language model to use (deprecated: kept for backward compatibility).""" + + Chat = "chat" + """Chat-based language model.""" + Embedding = "embedding" + """Embedding-based language model.""" + + def __repr__(self): + """Get a string representation.""" + return f'"{self.value}"' + + class ReportingType(str, Enum): """The reporting configuration type for the pipeline.""" diff --git a/packages/graphrag/graphrag/query/indexer_adapters.py b/packages/graphrag/graphrag/query/indexer_adapters.py index 7119ad842c..a6ded1f9f7 100644 --- a/packages/graphrag/graphrag/query/indexer_adapters.py +++ b/packages/graphrag/graphrag/query/indexer_adapters.py @@ -221,5 +221,5 @@ def _filter_under_community_level( ) -> pd.DataFrame: return cast( "pd.DataFrame", - df[df.level <= community_level], + df[(df.level <= community_level) | df.level.isna()], )