From 1b72899f246ba46ab8f7cd6aee2d56ef69217d82 Mon Sep 17 00:00:00 2001 From: Timothy Jaeryang Baek Date: Mon, 10 Aug 2026 00:26:44 -0600 Subject: [PATCH] refac --- .env.example | 3 ++ backend/open_webui/retrieval/loaders/main.py | 54 +++++++++++++++++++- 2 files changed, 56 insertions(+), 1 deletion(-) diff --git a/.env.example b/.env.example index 6f674e6e50..40f5b61fc5 100644 --- a/.env.example +++ b/.env.example @@ -16,6 +16,9 @@ CORS_ALLOW_ORIGIN='*' # Set to false to keep memory tools enabled without adding memory context to the system context. ENABLE_MEMORY_SYSTEM_CONTEXT=true +# Set to true to add compact row/column stats to parsed CSV retrieval context. +ENABLE_RAG_CSV_SUMMARY=false + # Set to false to disable workspace Tools and Functions. ENABLE_PLUGINS=true diff --git a/backend/open_webui/retrieval/loaders/main.py b/backend/open_webui/retrieval/loaders/main.py index 874085f87f..c86590583b 100644 --- a/backend/open_webui/retrieval/loaders/main.py +++ b/backend/open_webui/retrieval/loaders/main.py @@ -1,5 +1,7 @@ import asyncio +import csv import logging +import os import sys import ftfy @@ -111,6 +113,52 @@ class ExcelLoader: ] +def get_csv_summary(filename: str, file_path: str, encoding: str) -> str | None: + try: + with open(file_path, newline='', encoding=encoding) as f: + sample = f.read(4096) + f.seek(0) + try: + dialect = csv.Sniffer().sniff(sample) + except csv.Error: + dialect = csv.excel + + total_rows = 0 + max_columns = 0 + headers = [] + for row in csv.reader(f, dialect): + total_rows += 1 + max_columns = max(max_columns, len(row)) + if total_rows == 1: + headers = [header.lstrip('\ufeff') for header in row] + except Exception: + return None + + if total_rows == 0: + return None + + return ( + f'Table: {total_rows} rows incl. header; ' + f'{max(total_rows - 1, 0)} data rows; ' + f'{max_columns} columns: {", ".join(headers)}.' + ) + + +class CSVLoaderWithSummary: + def __init__(self, file_path: str, filename: str, encoding: str): + self.file_path = file_path + self.filename = filename + self.encoding = encoding + + def load(self) -> list[Document]: + docs = CSVLoader(self.file_path, encoding=self.encoding).load() + if os.getenv('ENABLE_RAG_CSV_SUMMARY', 'False').lower() == 'true': + summary = get_csv_summary(self.filename, self.file_path, self.encoding) + if summary: + docs.insert(0, Document(page_content=summary, metadata={'source': self.file_path, 'row': -1})) + return docs + + class PptxLoader: """Fallback PowerPoint loader using python-pptx when unstructured is not installed.""" @@ -594,7 +642,11 @@ class Loader: mode=self.kwargs.get('PDF_LOADER_MODE', 'page'), ) elif file_ext == 'csv': - loader = CSVLoader(file_path, encoding=self._detect_text_encoding(file_path)) + loader = CSVLoaderWithSummary( + file_path, + filename, + self._detect_text_encoding(file_path), + ) elif file_ext == 'rst': try: from langchain_community.document_loaders import UnstructuredRSTLoader