Remove indexing embedding and input configs

Author: natoverse

graphrag/api/query.py

@@ -24,12 +24,12 @@
 import pandas as pd
 from pydantic import validate_call
 
-from graphrag.config.models.graph_rag_config import GraphRagConfig
-from graphrag.index.config.embeddings import (
+from graphrag.config.embeddings import (
     community_full_content_embedding,
     entity_description_embedding,
     text_unit_text_embedding,
 )
+from graphrag.config.models.graph_rag_config import GraphRagConfig
 from graphrag.logger.print_progress import PrintProgressLogger
 from graphrag.query.factory import (
     get_basic_search_engine,

graphrag/index/config/embeddings.py graphrag/config/embeddings.pygraphrag/index/config/embeddings.py renamed to graphrag/config/embeddings.py

@@ -9,7 +9,7 @@
 
 from graphrag.cache.pipeline_cache import PipelineCache
 from graphrag.callbacks.workflow_callbacks import WorkflowCallbacks
-from graphrag.index.config.embeddings import (
+from graphrag.config.embeddings import (
     community_full_content_embedding,
     community_summary_embedding,
     community_title_embedding,

graphrag/index/config/__init__.py

@@ -6,11 +6,10 @@
 import logging
 import re
 from io import BytesIO
-from typing import cast
 
 import pandas as pd
 
-from graphrag.index.config.input import PipelineCSVInputConfig, PipelineInputConfig
+from graphrag.config.models.input_config import InputConfig
 from graphrag.index.utils.hashing import gen_sha512_hash
 from graphrag.logger.base import ProgressLogger
 from graphrag.storage.pipeline_storage import PipelineStorage
@@ -23,13 +22,12 @@
 
 
 async def load(
-    config: PipelineInputConfig,
+    config: InputConfig,
     progress: ProgressLogger | None,
     storage: PipelineStorage,
 ) -> pd.DataFrame:
     """Load csv inputs from a directory."""
-    csv_config = cast("PipelineCSVInputConfig", config)
-    log.info("Loading csv files from %s", csv_config.base_dir)
+    log.info("Loading csv files from %s", config.base_dir)
 
     async def load_file(path: str, group: dict | None) -> pd.DataFrame:
         if group is None:
@@ -43,51 +41,49 @@ async def load_file(path: str, group: dict | None) -> pd.DataFrame:
             )
         if "id" not in data.columns:
             data["id"] = data.apply(lambda x: gen_sha512_hash(x, x.keys()), axis=1)
-        if csv_config.source_column is not None and "source" not in data.columns:
-            if csv_config.source_column not in data.columns:
+        if config.source_column is not None and "source" not in data.columns:
+            if config.source_column not in data.columns:
                 log.warning(
                     "source_column %s not found in csv file %s",
-                    csv_config.source_column,
+                    config.source_column,
                     path,
                 )
             else:
-                data["source"] = data.apply(
-                    lambda x: x[csv_config.source_column], axis=1
-                )
-        if csv_config.text_column is not None and "text" not in data.columns:
-            if csv_config.text_column not in data.columns:
+                data["source"] = data.apply(lambda x: x[config.source_column], axis=1)
+        if config.text_column is not None and "text" not in data.columns:
+            if config.text_column not in data.columns:
                 log.warning(
                     "text_column %s not found in csv file %s",
-                    csv_config.text_column,
+                    config.text_column,
                     path,
                 )
             else:
-                data["text"] = data.apply(lambda x: x[csv_config.text_column], axis=1)
-        if csv_config.title_column is not None and "title" not in data.columns:
-            if csv_config.title_column not in data.columns:
+                data["text"] = data.apply(lambda x: x[config.text_column], axis=1)
+        if config.title_column is not None and "title" not in data.columns:
+            if config.title_column not in data.columns:
                 log.warning(
                     "title_column %s not found in csv file %s",
-                    csv_config.title_column,
+                    config.title_column,
                     path,
                 )
             else:
-                data["title"] = data.apply(lambda x: x[csv_config.title_column], axis=1)
+                data["title"] = data.apply(lambda x: x[config.title_column], axis=1)
 
-        if csv_config.timestamp_column is not None:
-            fmt = csv_config.timestamp_format
+        if config.timestamp_column is not None:
+            fmt = config.timestamp_format
             if fmt is None:
                 msg = "Must specify timestamp_format if timestamp_column is specified"
                 raise ValueError(msg)
 
-            if csv_config.timestamp_column not in data.columns:
+            if config.timestamp_column not in data.columns:
                 log.warning(
                     "timestamp_column %s not found in csv file %s",
-                    csv_config.timestamp_column,
+                    config.timestamp_column,
                     path,
                 )
             else:
                 data["timestamp"] = pd.to_datetime(
-                    data[csv_config.timestamp_column], format=fmt
+                    data[config.timestamp_column], format=fmt
                 )
 
             # TODO: Theres probably a less gross way to do this

graphrag/index/config/input.py

@@ -12,7 +12,6 @@
 
 from graphrag.config.enums import InputType
 from graphrag.config.models.input_config import InputConfig
-from graphrag.index.config.input import PipelineInputConfig
 from graphrag.index.input.csv import input_type as csv
 from graphrag.index.input.csv import load as load_csv
 from graphrag.index.input.text import input_type as text
@@ -30,7 +29,7 @@
 
 
 async def create_input(
-    config: PipelineInputConfig | InputConfig,
+    config: InputConfig,
     progress_reporter: ProgressLogger | None = None,
     root_dir: str | None = None,
 ) -> pd.DataFrame:

graphrag/index/errors.py

@@ -10,7 +10,7 @@
 
 import pandas as pd
 
-from graphrag.index.config.input import PipelineInputConfig
+from graphrag.config.models.input_config import InputConfig
 from graphrag.index.utils.hashing import gen_sha512_hash
 from graphrag.logger.base import ProgressLogger
 from graphrag.storage.pipeline_storage import PipelineStorage
@@ -23,7 +23,7 @@
 
 
 async def load(
-    config: PipelineInputConfig,
+    config: InputConfig,
     progress: ProgressLogger | None,
     storage: PipelineStorage,
 ) -> pd.DataFrame:

graphrag/index/flows/generate_text_embeddings.py

@@ -10,8 +10,8 @@
 
 from graphrag.cache.pipeline_cache import PipelineCache
 from graphrag.callbacks.workflow_callbacks import WorkflowCallbacks
+from graphrag.config.embeddings import get_embedded_fields, get_embedding_settings
 from graphrag.config.models.graph_rag_config import GraphRagConfig
-from graphrag.index.config.embeddings import get_embedded_fields, get_embedding_settings
 from graphrag.index.flows.generate_text_embeddings import generate_text_embeddings
 from graphrag.index.update.communities import (
     _merge_and_resolve_nodes,