Spaces:
Sleeping
Sleeping
File size: 2,813 Bytes
c96b98a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 | import csv
import os
from pathlib import Path
from typing import List, Union, IO, Any
from urllib.parse import urlparse
from phi.document.base import Document
from phi.document.reader.base import Reader
from phi.utils.log import logger
import io
class CSVReader(Reader):
"""Reader for CSV files"""
def read(self, file: Union[Path, IO[Any]], delimiter: str = ",", quotechar: str = '"') -> List[Document]:
if not file:
raise ValueError("No file provided")
try:
if isinstance(file, Path):
if not file.exists():
raise FileNotFoundError(f"Could not find file: {file}")
logger.info(f"Reading: {file}")
file_content = file.open(newline="", mode="r", encoding="utf-8")
else:
logger.info(f"Reading uploaded file: {file.name}")
file.seek(0)
file_content = io.StringIO(file.read().decode("utf-8")) # type: ignore
csv_name = Path(file.name).stem if isinstance(file, Path) else file.name.split(".")[0]
csv_content = ""
with file_content as csvfile:
csv_reader = csv.reader(csvfile, delimiter=delimiter, quotechar=quotechar)
for row in csv_reader:
csv_content += ", ".join(row) + "\n"
documents = [
Document(
name=csv_name,
id=csv_name,
content=csv_content,
)
]
if self.chunk:
chunked_documents = []
for document in documents:
chunked_documents.extend(self.chunk_document(document))
return chunked_documents
return documents
except Exception as e:
logger.error(f"Error reading: {file.name if isinstance(file, IO) else file}: {e}")
return []
class CSVUrlReader(Reader):
"""Reader for CSV files"""
def read(self, url: str) -> List[Document]:
if not url:
raise ValueError("No URL provided")
try:
import httpx
except ImportError:
raise ImportError("`httpx` not installed")
logger.info(f"Reading: {url}")
response = httpx.get(url)
try:
response.raise_for_status()
except httpx.HTTPStatusError as e:
logger.error(f"HTTP error occurred: {e.response.status_code} - {e.response.text}")
raise
parsed_url = urlparse(url)
filename = os.path.basename(parsed_url.path) or "data.csv"
file_obj = io.BytesIO(response.content)
file_obj.name = filename
documents = CSVReader().read(file=file_obj)
file_obj.close()
return documents
|