Underground-Digital/Workflow-Engine

1"""Abstract interface for document loader implementations."""2 3import csv4from typing import Optional5 6import pandas as pd7 8from core.rag.extractor.extractor_base import BaseExtractor9from core.rag.extractor.helpers import detect_file_encodings10from core.rag.models.document import Document11 12 13class CSVExtractor(BaseExtractor):14    """Load CSV files.15 16 17    Args:18        file_path: Path to the file to load.19    """20 21    def __init__(22        self,23        file_path: str,24        encoding: Optional[str] = None,25        autodetect_encoding: bool = False,26        source_column: Optional[str] = None,27        csv_args: Optional[dict] = None,28    ):29        """Initialize with file path."""30        self._file_path = file_path31        self._encoding = encoding32        self._autodetect_encoding = autodetect_encoding33        self.source_column = source_column34        self.csv_args = csv_args or {}35 36    def extract(self) -> list[Document]:37        """Load data into document objects."""38        docs = []39        try:40            with open(self._file_path, newline="", encoding=self._encoding) as csvfile:41                docs = self._read_from_file(csvfile)42        except UnicodeDecodeError as e:43            if self._autodetect_encoding:44                detected_encodings = detect_file_encodings(self._file_path)45                for encoding in detected_encodings:46                    try:47                        with open(self._file_path, newline="", encoding=encoding.encoding) as csvfile:48                            docs = self._read_from_file(csvfile)49                        break50                    except UnicodeDecodeError:51                        continue52            else:53                raise RuntimeError(f"Error loading {self._file_path}") from e54 55        return docs56 57    def _read_from_file(self, csvfile) -> list[Document]:58        docs = []59        try:60            # load csv file into pandas dataframe61            df = pd.read_csv(csvfile, on_bad_lines="skip", **self.csv_args)62 63            # check source column exists64            if self.source_column and self.source_column not in df.columns:65                raise ValueError(f"Source column '{self.source_column}' not found in CSV file.")66 67            # create document objects68 69            for i, row in df.iterrows():70                content = ";".join(f"{col.strip()}: {str(row[col]).strip()}" for col in df.columns)71                source = row[self.source_column] if self.source_column else ""72                metadata = {"source": source, "row": i}73                doc = Document(page_content=content, metadata=metadata)74                docs.append(doc)75        except csv.Error as e:76            raise e77 78        return docs79