rag_template / src /document_processing.py
Guilherme Favaron
Initial commit of local project
f5eb34f
"""
Processamento e extração de texto de documentos
"""
import io
import os
from typing import Union, Optional, Tuple
from pypdf import PdfReader
def extract_text_from_pdf(pdf_data: Union[bytes, io.BytesIO]) -> str:
"""
Extrai texto de arquivo PDF
Args:
pdf_data: Dados do PDF em bytes ou BytesIO
Returns:
Texto extraído
"""
try:
if isinstance(pdf_data, bytes):
pdf_data = io.BytesIO(pdf_data)
reader = PdfReader(pdf_data)
text = ""
for page in reader.pages:
page_text = page.extract_text()
if page_text:
text += page_text + "\n"
return text.strip()
except Exception as e:
return f"Erro ao extrair PDF: {str(e)}"
def extract_text_from_txt(txt_data: Union[bytes, str]) -> str:
"""
Extrai texto de arquivo TXT
Args:
txt_data: Dados do TXT em bytes ou string
Returns:
Texto extraído
"""
try:
if isinstance(txt_data, bytes):
return txt_data.decode("utf-8", errors="ignore")
return txt_data
except Exception as e:
return f"Erro ao extrair TXT: {str(e)}"
def process_uploaded_file(file_obj) -> Tuple[str, str]:
"""
Processa arquivo enviado via Gradio
Args:
file_obj: Objeto de arquivo do Gradio
Returns:
Tupla (nome_arquivo, texto_extraído)
"""
# Extrai nome do arquivo
name = "arquivo_desconhecido"
if hasattr(file_obj, "name"):
name = os.path.basename(file_obj.name)
elif isinstance(file_obj, dict) and "name" in file_obj:
name = os.path.basename(file_obj["name"])
elif isinstance(file_obj, str):
name = os.path.basename(file_obj)
# Extrai dados
data = None
path = None
if hasattr(file_obj, "read"):
# Objeto file-like
try:
data = file_obj.read()
except Exception:
data = None
elif isinstance(file_obj, dict):
# Dicionário com caminho ou dados
path = file_obj.get("path") or file_obj.get("name")
if path and os.path.exists(path):
with open(path, "rb") as f:
data = f.read()
elif "data" in file_obj:
data = file_obj["data"]
elif isinstance(file_obj, str):
# Caminho de arquivo
path = file_obj
if os.path.exists(path):
with open(path, "rb") as f:
data = f.read()
# Processa baseado no tipo
if data is None:
return name, "Erro: não foi possível ler o arquivo"
is_pdf = name.lower().endswith(".pdf")
if is_pdf:
text = extract_text_from_pdf(data)
else:
text = extract_text_from_txt(data)
return name, text
def get_document_preview(text: str, max_chars: int = 500) -> str:
"""
Retorna preview do documento
Args:
text: Texto completo
max_chars: Quantidade máxima de caracteres
Returns:
Preview do texto
"""
if len(text) <= max_chars:
return text
return text[:max_chars] + "..."
def get_document_stats(text: str) -> dict:
"""
Calcula estatísticas do documento
Args:
text: Texto do documento
Returns:
Dicionário com estatísticas
"""
words = text.split()
lines = text.split("\n")
return {
"total_chars": len(text),
"total_words": len(words),
"total_lines": len(lines),
"avg_word_length": sum(len(w) for w in words) / len(words) if words else 0
}