88 lines
3.4 KiB
Python
88 lines
3.4 KiB
Python
import base64
|
|
import os
|
|
import subprocess
|
|
import logging
|
|
import pandas as pd
|
|
|
|
logger = logging.getLogger("FILE_PARSER")
|
|
|
|
def extract_text_from_file(file_bytes: bytes, filename: str) -> dict:
|
|
ext = os.path.splitext(filename)[1].lower()
|
|
temp_filepath = f"/tmp/upload_{os.getpid()}_{filename}"
|
|
|
|
with open(temp_filepath, "wb") as f:
|
|
f.write(file_bytes)
|
|
|
|
try:
|
|
# 1. ИЗОБРАЖЕНИЯ (.png, .jpg, .jpeg, .bmp, .webp) -> Кодируем в Base64 для Vision LLM
|
|
if ext in ['.png', '.jpg', '.jpeg', '.bmp', '.webp']:
|
|
b64_str = base64.b64encode(file_bytes).decode('utf-8')
|
|
return {
|
|
"text": f"[ПРИКРЕПЛЕНО ИЗОБРАЖЕНИЕ: {filename}]",
|
|
"image_b64": b64_str
|
|
}
|
|
|
|
# 2. PDF ДОКУМЕНТЫ (Конвертируем 1-ю страницу в картинку для Vision LLM)
|
|
elif ext == '.pdf':
|
|
cmd = ['pdftotext', temp_filepath, '-']
|
|
res = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
|
|
pdf_text = res.stdout.strip()
|
|
|
|
img_prefix = f"/tmp/pdf_preview_{os.getpid()}"
|
|
subprocess.run(['pdftoppm', '-png', '-r', '200', '-f', '1', '-l', '1', temp_filepath, img_prefix], check=True)
|
|
|
|
page_png = f"{img_prefix}-1.png"
|
|
b64_str = None
|
|
if os.path.exists(page_png):
|
|
with open(page_png, "rb") as pf:
|
|
b64_str = base64.b64encode(pf.read()).decode('utf-8')
|
|
os.remove(page_png)
|
|
|
|
context_text = f"[ПРИКРЕПЛЕН ДОКУМЕНТ PDF: {filename}]"
|
|
if pdf_text:
|
|
context_text += f"\n\n[ЭЛЕКТРОННЫЙ ТЕКСТОВЫЙ СЛОЙ PDF]:\n{pdf_text}"
|
|
|
|
return {
|
|
"text": context_text,
|
|
"image_b64": b64_str
|
|
}
|
|
|
|
# 3. ЭЛЕКТРОННЫЕ ТАБЛИЦЫ (.xlsx, .xls, .csv)
|
|
elif ext in ['.xlsx', '.xls', '.csv']:
|
|
if ext == '.csv':
|
|
df = pd.read_csv(temp_filepath)
|
|
else:
|
|
df = pd.read_excel(temp_filepath)
|
|
|
|
total_rows = len(df)
|
|
df_preview = df.head(100)
|
|
table_str = df_preview.to_string(index=False)
|
|
note = f"\n(Показано первых 100 строк из {total_rows})" if total_rows > 100 else ""
|
|
return {
|
|
"text": f"[СОДЕРЖИМОЕ ТАБЛИЦЫ {filename}]:\n{table_str}{note}",
|
|
"image_b64": None
|
|
}
|
|
|
|
# 4. ТЕКСТОВЫЕ ФАЙЛЫ
|
|
elif ext in ['.txt', '.log', '.json', '.xml', '.md']:
|
|
with open(temp_filepath, 'r', encoding='utf-8', errors='replace') as tf:
|
|
return {
|
|
"text": tf.read().strip(),
|
|
"image_b64": None
|
|
}
|
|
|
|
else:
|
|
return {
|
|
"text": f"[ОШИБКА: Формат {ext} не поддерживается]",
|
|
"image_b64": None
|
|
}
|
|
|
|
except Exception as e:
|
|
logger.error(f"Ошибка при анализе файла {filename}: {e}")
|
|
return {
|
|
"text": f"[ОШИБКА ОБРАБОТКИ ФАЙЛА: {str(e)}]",
|
|
"image_b64": None
|
|
}
|
|
finally:
|
|
if os.path.exists(temp_filepath):
|
|
os.remove(temp_filepath) |