graphgen/models/reader/csv_reader.py

from typing import Any, Dict, List

import pandas as pd

from graphgen.bases.base_reader import BaseReader


class CSVReader(BaseReader):
    def read(self, file_path: str) -> List[Dict[str, Any]]:

        df = pd.read_csv(file_path)
        for _, row in df.iterrows():
            if "type" in row and row["type"] == "text" and self.text_column not in row:
                raise ValueError(
                    f"Missing '{self.text_column}' in document: {row.to_dict()}"
                )
        return self.filter(df.to_dict(orient="records"))