Source code for langchain_community.document_loaders.csv_loader

import csv
from io import TextIOWrapper
from pathlib import Path
from typing import Any, Dict, Iterator, List, Optional, Sequence, Union

from langchain_core.documents import Document

from langchain_community.document_loaders.base import BaseLoader
from langchain_community.document_loaders.helpers import detect_file_encodings
from langchain_community.document_loaders.unstructured import (
    UnstructuredFileLoader,
    validate_unstructured_version,
)


class CSVLoader(BaseLoader):
    """将`CSV`文件加载到文档列表中。

    每个文档代表CSV文件的一行。每一行被转换为键/值对，并输出到文档的page_content中的新行。

    从CSV加载的每个文档的源默认设置为`file_path`参数的值。
    您可以通过将`source_column`参数设置为CSV文件中的列名来覆盖此设置。
    然后，每个文档的源将设置为指定在`source_column`中的列的值。

    输出示例:
        .. code-block:: txt

            column1: value1
            column2: value2
            column3: value3"""

    def __init__(
        self,
        file_path: Union[str, Path],
        source_column: Optional[str] = None,
        metadata_columns: Sequence[str] = (),
        csv_args: Optional[Dict] = None,
        encoding: Optional[str] = None,
        autodetect_encoding: bool = False,
    ):
        """参数：
    file_path：CSV文件的路径。
    source_column：CSV文件中用作源的列的名称。可选。默认为None。
    metadata_columns：要用作元数据的列名称序列。可选。
    csv_args：要传递给csv.DictReader的参数字典。可选。默认为None。
    encoding：CSV文件的编码。可选。默认为None。
    autodetect_encoding：是否尝试自动检测文件编码。
"""
        self.file_path = file_path
        self.source_column = source_column
        self.metadata_columns = metadata_columns
        self.encoding = encoding
        self.csv_args = csv_args or {}
        self.autodetect_encoding = autodetect_encoding

    def lazy_load(self) -> Iterator[Document]:
        try:
            with open(self.file_path, newline="", encoding=self.encoding) as csvfile:
                yield from self.__read_file(csvfile)
        except UnicodeDecodeError as e:
            if self.autodetect_encoding:
                detected_encodings = detect_file_encodings(self.file_path)
                for encoding in detected_encodings:
                    try:
                        with open(
                            self.file_path, newline="", encoding=encoding.encoding
                        ) as csvfile:
                            yield from self.__read_file(csvfile)
                            break
                    except UnicodeDecodeError:
                        continue
            else:
                raise RuntimeError(f"Error loading {self.file_path}") from e
        except Exception as e:
            raise RuntimeError(f"Error loading {self.file_path}") from e

    def __read_file(self, csvfile: TextIOWrapper) -> Iterator[Document]:
        csv_reader = csv.DictReader(csvfile, **self.csv_args)
        for i, row in enumerate(csv_reader):
            try:
                source = (
                    row[self.source_column]
                    if self.source_column is not None
                    else str(self.file_path)
                )
            except KeyError:
                raise ValueError(
                    f"Source column '{self.source_column}' not found in CSV file."
                )
            content = "\n".join(
                f"{k.strip()}: {v.strip() if v is not None else v}"
                for k, v in row.items()
                if k not in self.metadata_columns
            )
            metadata = {"source": source, "row": i}
            for col in self.metadata_columns:
                try:
                    metadata[col] = row[col]
                except KeyError:
                    raise ValueError(f"Metadata column '{col}' not found in CSV file.")
            yield Document(page_content=content, metadata=metadata)


class UnstructuredCSVLoader(UnstructuredFileLoader):
    """使用`Unstructured`加载`CSV`文件。

    像其他Unstructured加载器一样，UnstructuredCSVLoader可以在“single”和“elements”模式下使用。如果您在“elements”模式下使用加载器，则CSV文件将成为单个Unstructured表元素。如果您在“elements”模式下使用加载器，则表的HTML表示将在文档元数据中的“text_as_html”键中可用。

    示例
    --------
    from langchain_community.document_loaders.csv_loader import UnstructuredCSVLoader

    loader = UnstructuredCSVLoader("stanley-cups.csv", mode="elements")
    docs = loader.load()"""

    def __init__(
        self, file_path: str, mode: str = "single", **unstructured_kwargs: Any
    ):
        """参数：
    file_path: CSV文件的路径。
    mode: 加载CSV文件时要使用的模式。
      可选。默认为"single"。
    **unstructured_kwargs: 要传递给unstructured的关键字参数。
"""
        validate_unstructured_version(min_unstructured_version="0.6.8")
        super().__init__(file_path=file_path, mode=mode, **unstructured_kwargs)

    def _get_elements(self) -> List:
        from unstructured.partition.csv import partition_csv

        return partition_csv(filename=self.file_path, **self.unstructured_kwargs)