85 lines
2.9 KiB
Python
85 lines
2.9 KiB
Python
"""Парсеры источников: text/pdf/epub/web → чанки."""
|
||
from __future__ import annotations
|
||
import hashlib
|
||
from dataclasses import dataclass, field
|
||
|
||
@dataclass
|
||
class ParsedChunk:
|
||
seq: int
|
||
content: str
|
||
locator: dict = field(default_factory=dict)
|
||
|
||
@dataclass
|
||
class ParsedDoc:
|
||
kind: str
|
||
uri: str | None
|
||
title: str | None
|
||
chunks: list[ParsedChunk]
|
||
content_hash: str = ""
|
||
|
||
def parse_text(text: str, *, chunk_size: int = 2000, overlap: int = 200,
|
||
title: str | None = None, uri: str | None = None) -> ParsedDoc:
|
||
"""Разбить сырой текст на чанки с перекрытием."""
|
||
text = text.strip()
|
||
chunks: list[ParsedChunk] = []
|
||
start = 0
|
||
seq = 0
|
||
while start < len(text):
|
||
end = min(start + chunk_size, len(text))
|
||
# не резать посреди слова
|
||
if end < len(text):
|
||
space = text.rfind(" ", start, end)
|
||
if space > start + chunk_size // 2:
|
||
end = space
|
||
chunks.append(ParsedChunk(seq=seq, content=text[start:end].strip()))
|
||
seq += 1
|
||
if end >= len(text):
|
||
break
|
||
start = end - overlap
|
||
return ParsedDoc(kind="text", uri=uri, title=title, chunks=chunks,
|
||
content_hash=hashlib.sha256(text.encode()).hexdigest())
|
||
|
||
def parse_pdf(path: str) -> ParsedDoc:
|
||
from pypdf import PdfReader
|
||
reader = PdfReader(path)
|
||
text = "\n\n".join((page.extract_text() or "") for page in reader.pages)
|
||
doc = parse_text(text, title=path.split("/")[-1], uri=path)
|
||
doc.kind = "pdf"
|
||
return doc
|
||
|
||
def parse_epub(path: str) -> ParsedDoc:
|
||
import ebooklib
|
||
from ebooklib import epub
|
||
book = epub.read_epub(path)
|
||
parts = []
|
||
for item in book.get_items_of_type(ebooklib.ITEM_DOCUMENT):
|
||
parts.append(item.get_content().decode("utf-8", errors="ignore"))
|
||
text = "\n\n".join(parts)
|
||
doc = parse_text(text, title=path.split("/")[-1], uri=path)
|
||
doc.kind = "epub"
|
||
return doc
|
||
|
||
def _http_client():
|
||
"""httpx.Client с опциональным прокси и единым User-Agent."""
|
||
import httpx
|
||
from . import config
|
||
kwargs = dict(timeout=30, follow_redirects=True,
|
||
headers={"User-Agent": config.USER_AGENT})
|
||
if config.PROXY_URL:
|
||
kwargs["proxy"] = config.PROXY_URL
|
||
return httpx.Client(**kwargs)
|
||
|
||
def parse_web(url: str) -> ParsedDoc:
|
||
import trafilatura
|
||
# fetch через httpx с таймаутом + прокси (trafilatura.fetch_url не имеет таймаута и виснет)
|
||
with _http_client() as client:
|
||
r = client.get(url)
|
||
r.raise_for_status()
|
||
text = trafilatura.extract(r.text, include_comments=False, include_tables=True,
|
||
url=url)
|
||
if not text:
|
||
raise ValueError(f"trafilatura: no main content for {url}")
|
||
doc = parse_text(text, uri=url, title=url)
|
||
doc.kind = "web"
|
||
return doc
|