Initial commit: Pegas Memory (PG + Qdrant + Memgraph, MCP server, 14 tools)

pegas_memory/ingest.py

@@ -0,0 +1,84 @@
+"""Парсеры источников: 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