486 lines
16 KiB
Python
486 lines
16 KiB
Python
"""Text extraction, chunking, embeddings, ingest and search for RAG."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import re
|
|
from html.parser import HTMLParser
|
|
from pathlib import Path
|
|
from typing import Any, Optional
|
|
from urllib.parse import urlparse
|
|
|
|
import httpx
|
|
from sqlalchemy import delete, or_, select, text
|
|
from sqlalchemy.orm import Session
|
|
|
|
from app.core.config import get_settings
|
|
from app.services.rag.models import EMBEDDING_DIM, RagChunk, RagSource
|
|
|
|
CHUNK_SIZE = 1200
|
|
CHUNK_OVERLAP = 200
|
|
|
|
|
|
class _HTMLStripper(HTMLParser):
|
|
def __init__(self) -> None:
|
|
super().__init__()
|
|
self.parts: list[str] = []
|
|
|
|
def handle_data(self, data: str) -> None:
|
|
if data:
|
|
self.parts.append(data)
|
|
|
|
def get_text(self) -> str:
|
|
return re.sub(r"\s+", " ", " ".join(self.parts)).strip()
|
|
|
|
|
|
def strip_html(html: str) -> str:
|
|
parser = _HTMLStripper()
|
|
try:
|
|
parser.feed(html or "")
|
|
parser.close()
|
|
except Exception:
|
|
return re.sub(r"<[^>]+>", " ", html or "")
|
|
return parser.get_text()
|
|
|
|
|
|
def extract_text_from_bytes(data: bytes, filename: str = "", content_type: str = "") -> str:
|
|
name = (filename or "").lower()
|
|
ctype = (content_type or "").lower()
|
|
if name.endswith(".pdf") or "pdf" in ctype:
|
|
try:
|
|
from pypdf import PdfReader
|
|
import io
|
|
|
|
reader = PdfReader(io.BytesIO(data))
|
|
return "\n\n".join((p.extract_text() or "") for p in reader.pages).strip()
|
|
except Exception as e:
|
|
raise ValueError(f"Не удалось извлечь текст из PDF: {e}") from e
|
|
if name.endswith(".docx") or "wordprocessingml" in ctype or name.endswith(".doc"):
|
|
try:
|
|
import io
|
|
from docx import Document
|
|
|
|
doc = Document(io.BytesIO(data))
|
|
return "\n\n".join(p.text for p in doc.paragraphs if p.text.strip()).strip()
|
|
except Exception as e:
|
|
raise ValueError(f"Не удалось извлечь текст из DOCX: {e}") from e
|
|
if name.endswith((".html", ".htm")) or "html" in ctype:
|
|
return strip_html(data.decode("utf-8", errors="ignore"))
|
|
# plain / markdown / txt
|
|
try:
|
|
return data.decode("utf-8")
|
|
except UnicodeDecodeError:
|
|
return data.decode("latin-1", errors="ignore")
|
|
|
|
|
|
def chunk_text(text: str, size: int = CHUNK_SIZE, overlap: int = CHUNK_OVERLAP) -> list[str]:
|
|
clean = re.sub(r"\n{3,}", "\n\n", (text or "").strip())
|
|
if not clean:
|
|
return []
|
|
if len(clean) <= size:
|
|
return [clean]
|
|
chunks: list[str] = []
|
|
start = 0
|
|
while start < len(clean):
|
|
end = min(len(clean), start + size)
|
|
if end < len(clean):
|
|
# prefer break at paragraph / sentence
|
|
window = clean[start:end]
|
|
br = max(window.rfind("\n\n"), window.rfind(". "), window.rfind("\n"))
|
|
if br > size // 3:
|
|
end = start + br + 1
|
|
piece = clean[start:end].strip()
|
|
if piece:
|
|
chunks.append(piece)
|
|
if end >= len(clean):
|
|
break
|
|
start = max(0, end - overlap)
|
|
return chunks
|
|
|
|
|
|
async def embed_texts(texts: list[str]) -> list[list[float]]:
|
|
if not texts:
|
|
return []
|
|
settings = get_settings()
|
|
api_key = (settings.openrouter_api_key or "").strip()
|
|
if not api_key:
|
|
# deterministic pseudo-embedding for offline/dev (cosine still works within corpus)
|
|
return [_hash_embed(t) for t in texts]
|
|
|
|
model = getattr(settings, "openrouter_embedding_model", None) or "openai/text-embedding-3-small"
|
|
headers = {
|
|
"Authorization": f"Bearer {api_key}",
|
|
"Content-Type": "application/json",
|
|
}
|
|
base = settings.openrouter_base_url.rstrip("/")
|
|
proxy = (settings.openrouter_proxy or "").strip() or None
|
|
async with httpx.AsyncClient(proxy=proxy, timeout=120.0) as client:
|
|
resp = await client.post(
|
|
f"{base}/embeddings",
|
|
headers=headers,
|
|
json={"model": model, "input": texts},
|
|
)
|
|
resp.raise_for_status()
|
|
data = resp.json()
|
|
items = sorted(data.get("data") or [], key=lambda x: x.get("index", 0))
|
|
vectors = [list(it.get("embedding") or []) for it in items]
|
|
# pad/truncate
|
|
out: list[list[float]] = []
|
|
for v in vectors:
|
|
if len(v) < EMBEDDING_DIM:
|
|
v = v + [0.0] * (EMBEDDING_DIM - len(v))
|
|
out.append(v[:EMBEDDING_DIM])
|
|
return out
|
|
|
|
|
|
def _hash_embed(text: str) -> list[float]:
|
|
vec = [0.0] * EMBEDDING_DIM
|
|
tokens = re.findall(r"\w+", (text or "").lower())
|
|
for tok in tokens:
|
|
h = int(hashlib.md5(tok.encode("utf-8")).hexdigest(), 16)
|
|
vec[h % EMBEDDING_DIM] += 1.0
|
|
vec[(h // 7) % EMBEDDING_DIM] += 0.5
|
|
norm = sum(x * x for x in vec) ** 0.5 or 1.0
|
|
return [x / norm for x in vec]
|
|
|
|
|
|
def rag_storage_dir(scope: str, project_id: int | None, source_id: int) -> Path:
|
|
settings = get_settings()
|
|
root = Path(settings.data_dir) / "rag"
|
|
if scope == "project" and project_id is not None:
|
|
return root / "projects" / str(project_id) / str(source_id)
|
|
return root / "global" / str(source_id)
|
|
|
|
|
|
async def fetch_url_bytes(url: str) -> tuple[bytes, str, str]:
|
|
parsed = urlparse(url)
|
|
if parsed.scheme not in ("http", "https"):
|
|
raise ValueError("Разрешены только http/https URL")
|
|
async with httpx.AsyncClient(follow_redirects=True, timeout=60.0) as client:
|
|
resp = await client.get(url)
|
|
resp.raise_for_status()
|
|
ctype = resp.headers.get("content-type", "")
|
|
name = Path(parsed.path).name or "page.html"
|
|
return resp.content, name, ctype
|
|
|
|
|
|
def source_to_dict(src: RagSource, chunk_count: int | None = None) -> dict[str, Any]:
|
|
return {
|
|
"id": src.id,
|
|
"scope": src.scope,
|
|
"project_id": src.project_id,
|
|
"kind": src.kind,
|
|
"title": src.title,
|
|
"uri": src.uri,
|
|
"path": src.path,
|
|
"gost_id": src.gost_id,
|
|
"status": src.status,
|
|
"error_message": src.error_message,
|
|
"meta": src.meta,
|
|
"chunk_count": chunk_count,
|
|
"created_at": src.created_at.isoformat() if src.created_at else None,
|
|
"updated_at": src.updated_at.isoformat() if src.updated_at else None,
|
|
}
|
|
|
|
|
|
async def index_source(db: Session, source_id: int) -> RagSource:
|
|
src = db.get(RagSource, source_id)
|
|
if not src:
|
|
raise ValueError("Источник не найден")
|
|
src.status = "indexing"
|
|
src.error_message = None
|
|
db.commit()
|
|
|
|
try:
|
|
data: bytes
|
|
filename = src.title or "document"
|
|
ctype = ""
|
|
if src.kind == "upload" and src.path:
|
|
data = Path(src.path).read_bytes()
|
|
filename = Path(src.path).name
|
|
elif src.kind == "url" and src.uri:
|
|
data, filename, ctype = await fetch_url_bytes(src.uri)
|
|
# optionally cache
|
|
store = rag_storage_dir(src.scope, src.project_id, src.id)
|
|
store.mkdir(parents=True, exist_ok=True)
|
|
cached = store / (filename or "page.html")
|
|
cached.write_bytes(data)
|
|
src.path = str(cached)
|
|
else:
|
|
raise ValueError("Нет данных для индексации")
|
|
|
|
text_body = extract_text_from_bytes(data, filename=filename, content_type=ctype)
|
|
pieces = chunk_text(text_body)
|
|
db.execute(delete(RagChunk).where(RagChunk.source_id == src.id))
|
|
db.flush()
|
|
|
|
embeddings: list[list[float]] = []
|
|
batch = 16
|
|
for i in range(0, len(pieces), batch):
|
|
embeddings.extend(await embed_texts(pieces[i : i + batch]))
|
|
|
|
for idx, piece in enumerate(pieces):
|
|
emb = embeddings[idx] if idx < len(embeddings) else None
|
|
db.add(
|
|
RagChunk(
|
|
source_id=src.id,
|
|
chunk_index=idx,
|
|
text=piece,
|
|
embedding=emb,
|
|
meta={"chars": len(piece)},
|
|
)
|
|
)
|
|
src.status = "ready"
|
|
src.meta = {**(src.meta or {}), "chunks": len(pieces), "chars": len(text_body)}
|
|
db.commit()
|
|
db.refresh(src)
|
|
return src
|
|
except Exception as e:
|
|
src.status = "error"
|
|
src.error_message = str(e)[:1000]
|
|
db.commit()
|
|
db.refresh(src)
|
|
raise
|
|
|
|
|
|
async def create_upload_source(
|
|
db: Session,
|
|
*,
|
|
scope: str,
|
|
project_id: int | None,
|
|
filename: str,
|
|
data: bytes,
|
|
title: str | None = None,
|
|
gost_id: str | None = None,
|
|
) -> RagSource:
|
|
src = RagSource(
|
|
scope=scope,
|
|
project_id=project_id if scope == "project" else None,
|
|
kind="upload",
|
|
title=title or filename,
|
|
gost_id=gost_id or None,
|
|
status="pending",
|
|
)
|
|
db.add(src)
|
|
db.flush()
|
|
store = rag_storage_dir(scope, project_id, src.id)
|
|
store.mkdir(parents=True, exist_ok=True)
|
|
path = store / filename
|
|
path.write_bytes(data)
|
|
src.path = str(path)
|
|
src.uri = filename
|
|
db.commit()
|
|
db.refresh(src)
|
|
await index_source(db, src.id)
|
|
db.refresh(src)
|
|
return src
|
|
|
|
|
|
async def create_url_source(
|
|
db: Session,
|
|
*,
|
|
scope: str,
|
|
project_id: int | None,
|
|
url: str,
|
|
title: str | None = None,
|
|
gost_id: str | None = None,
|
|
) -> RagSource:
|
|
src = RagSource(
|
|
scope=scope,
|
|
project_id=project_id if scope == "project" else None,
|
|
kind="url",
|
|
title=title or url,
|
|
uri=url,
|
|
gost_id=gost_id or None,
|
|
status="pending",
|
|
)
|
|
db.add(src)
|
|
db.commit()
|
|
db.refresh(src)
|
|
await index_source(db, src.id)
|
|
db.refresh(src)
|
|
return src
|
|
|
|
|
|
def delete_source(db: Session, source_id: int) -> None:
|
|
src = db.get(RagSource, source_id)
|
|
if not src:
|
|
return
|
|
if src.path:
|
|
try:
|
|
p = Path(src.path)
|
|
if p.is_file():
|
|
p.unlink(missing_ok=True)
|
|
parent = p.parent
|
|
if parent.exists() and parent.name == str(source_id):
|
|
for child in parent.iterdir():
|
|
child.unlink(missing_ok=True)
|
|
parent.rmdir()
|
|
except Exception:
|
|
pass
|
|
db.delete(src)
|
|
db.commit()
|
|
|
|
|
|
async def search_chunks(
|
|
db: Session,
|
|
query: str,
|
|
*,
|
|
project_id: int | None = None,
|
|
top_k: int = 5,
|
|
gost_only: bool = False,
|
|
) -> list[dict[str, Any]]:
|
|
q = (query or "").strip()
|
|
if not q:
|
|
return []
|
|
top_k = max(1, min(int(top_k or 5), 20))
|
|
|
|
# candidate sources: global + optional project
|
|
src_q = select(RagSource).where(RagSource.status == "ready")
|
|
if project_id is not None:
|
|
src_q = src_q.where(
|
|
or_(RagSource.scope == "global", RagSource.project_id == project_id)
|
|
)
|
|
else:
|
|
src_q = src_q.where(RagSource.scope == "global")
|
|
if gost_only:
|
|
src_q = src_q.where(
|
|
or_(RagSource.gost_id.isnot(None), RagSource.scope == "global")
|
|
)
|
|
sources = db.scalars(src_q).all()
|
|
if not sources:
|
|
return []
|
|
source_ids = [s.id for s in sources]
|
|
source_map = {s.id: s for s in sources}
|
|
|
|
vectors = await embed_texts([q])
|
|
query_vec = vectors[0] if vectors else None
|
|
|
|
results: list[dict[str, Any]] = []
|
|
if query_vec is not None:
|
|
try:
|
|
# pgvector cosine distance
|
|
literal = "[" + ",".join(str(float(x)) for x in query_vec) + "]"
|
|
from sqlalchemy import bindparam
|
|
|
|
stmt = text(
|
|
"""
|
|
SELECT c.id, c.source_id, c.chunk_index, c.text,
|
|
(c.embedding <=> CAST(:emb AS vector)) AS distance
|
|
FROM rag_chunks c
|
|
WHERE c.source_id IN :sids AND c.embedding IS NOT NULL
|
|
ORDER BY c.embedding <=> CAST(:emb AS vector)
|
|
LIMIT :lim
|
|
"""
|
|
).bindparams(bindparam("sids", expanding=True))
|
|
rows = db.execute(
|
|
stmt,
|
|
{"emb": literal, "sids": source_ids, "lim": top_k},
|
|
).fetchall()
|
|
for row in rows:
|
|
src = source_map.get(row.source_id)
|
|
results.append(
|
|
{
|
|
"chunk_id": row.id,
|
|
"source_id": row.source_id,
|
|
"chunk_index": row.chunk_index,
|
|
"text": row.text,
|
|
"score": float(1.0 / (1.0 + float(row.distance or 0))),
|
|
"title": src.title if src else "",
|
|
"gost_id": src.gost_id if src else None,
|
|
"scope": src.scope if src else "",
|
|
"uri": src.uri if src else None,
|
|
}
|
|
)
|
|
if results:
|
|
return results
|
|
except Exception:
|
|
pass
|
|
|
|
# fallback: ILIKE
|
|
like = f"%{q[:80]}%"
|
|
chunks = db.scalars(
|
|
select(RagChunk)
|
|
.where(RagChunk.source_id.in_(source_ids), RagChunk.text.ilike(like))
|
|
.limit(top_k)
|
|
).all()
|
|
if not chunks:
|
|
# last resort: first chunks from matching sources
|
|
chunks = db.scalars(
|
|
select(RagChunk)
|
|
.where(RagChunk.source_id.in_(source_ids))
|
|
.order_by(RagChunk.id.asc())
|
|
.limit(top_k)
|
|
).all()
|
|
for ch in chunks:
|
|
src = source_map.get(ch.source_id)
|
|
results.append(
|
|
{
|
|
"chunk_id": ch.id,
|
|
"source_id": ch.source_id,
|
|
"chunk_index": ch.chunk_index,
|
|
"text": ch.text,
|
|
"score": 0.3,
|
|
"title": src.title if src else "",
|
|
"gost_id": src.gost_id if src else None,
|
|
"scope": src.scope if src else "",
|
|
"uri": src.uri if src else None,
|
|
}
|
|
)
|
|
return results
|
|
|
|
|
|
async def gost_check(
|
|
db: Session,
|
|
text_body: str,
|
|
focus: str = "",
|
|
project_id: int | None = None,
|
|
) -> dict[str, Any]:
|
|
from app.services.llm import openrouter_chat
|
|
|
|
query = (focus or "") + "\n" + (text_body or "")[:1500]
|
|
hits = await search_chunks(db, query, project_id=project_id, top_k=8, gost_only=True)
|
|
citations = [
|
|
{
|
|
"title": h.get("title"),
|
|
"gost_id": h.get("gost_id"),
|
|
"excerpt": (h.get("text") or "")[:500],
|
|
"source_id": h.get("source_id"),
|
|
"uri": h.get("uri"),
|
|
}
|
|
for h in hits
|
|
]
|
|
context = "\n\n".join(
|
|
f"[{i+1}] {c['title'] or 'Источник'}"
|
|
+ (f" ({c['gost_id']})" if c.get("gost_id") else "")
|
|
+ f":\n{c['excerpt']}"
|
|
for i, c in enumerate(citations)
|
|
)
|
|
prompt = (
|
|
"Проверь фрагмент документа ЕСПД на соответствие нормам. "
|
|
"Ссылайся на цитаты по номерам [n]. Перечисли замечания списком.\n\n"
|
|
f"Фокус: {focus or 'общее соответствие'}\n\n"
|
|
f"Текст документа:\n{(text_body or '')[:6000]}\n\n"
|
|
f"Нормы из базы знаний:\n{context or '(пусто — скажи, что корпус пуст)'}"
|
|
)
|
|
try:
|
|
msg = await openrouter_chat(
|
|
[
|
|
{
|
|
"role": "system",
|
|
"content": "Ты эксперт по ГОСТам ЕСПД. Отвечай по-русски, с ссылками [n].",
|
|
},
|
|
{"role": "user", "content": prompt},
|
|
],
|
|
temperature=0.2,
|
|
max_tokens=2000,
|
|
)
|
|
reply = (msg.get("content") or "").strip() or "Нет ответа модели."
|
|
except Exception as e:
|
|
reply = f"Не удалось вызвать LLM для проверки: {e}"
|
|
if citations:
|
|
reply += "\n\nНайденные фрагменты норм:\n" + "\n".join(
|
|
f"- {c['title']}: {c['excerpt'][:200]}" for c in citations
|
|
)
|
|
return {"reply": reply, "sources": citations, "hits": hits}
|