feat(knowledge): 任务队列、独立 worker 与 pgvector 向量索引
知识处理改为数据库任务队列，独立 worker 做切块与向量化，失败可重试。
This commit is contained in:
WangShen007 committed
1 parent
7937de5d7d
commit
c7011040c3
36 files changed
+7590
No files matched your search
@@ -0,0 +1,218 @@
|
||||
"""
|
||||
Copyright 2024-2026 ReplyPier
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
|
||||
Splitting a knowledge page into embedder-sized chunks.
|
||||
|
||||
The search embedder only reads ~512 tokens, so a page stored as one row is
|
||||
searchable by its first screen only, and a single hit can be tens of thousands
|
||||
of chars — enough to overflow the model's context when a few are replayed. Each
|
||||
page is therefore stored as several small rows that together hold every byte of
|
||||
it.
|
||||
|
||||
Stored chunks are disjoint verbatim slices: concatenating a page's chunks gives
|
||||
the page back exactly, which is what the dashboard, the FAQ generator and page
|
||||
edits all rely on. Context across a boundary is handled at embedding time (the
|
||||
tail of the previous chunk is prepended to the text that is embedded, never to
|
||||
the text that is stored) and at search time (neighbouring chunks are returned
|
||||
with a hit).
|
||||
|
||||
Why not agno's chunkers: ``ChunkingStrategy.clean_text`` collapses whitespace,
|
||||
which flattens tables, lists and headings into one line and would persist that
|
||||
way; ``MarkdownChunking`` needs the ``unstructured`` package.
|
||||
|
||||
Chunk ids are ``<page_id>::<n>`` (1-based). The ``_<n>`` form agno uses for PDFs
|
||||
is kept readable by the page-grouping helpers, but it is not used for new chunks
|
||||
because real page URLs end in ``_<n>`` too (``/blog/post_2``) and an upsert on
|
||||
such an id would silently overwrite another page.
|
||||
"""
|
||||
|
||||
import re
|
||||
from dataclasses import dataclass
|
||||
from typing import List, Optional
|
||||
|
||||
from agno.document import Document
|
||||
from agno.document.chunking.strategy import ChunkingStrategy
|
||||
from agno.embedder import Embedder
|
||||
|
||||
from app.core.config import settings
|
||||
|
||||
CHUNK_ID_SEPARATOR = "::"
|
||||
# Legacy agno PDF chunk suffix. Stripped when grouping, never generated.
|
||||
_LEGACY_CHUNK_ID_RE = re.compile(r"_\d+$")
|
||||
_CHUNK_ID_RE = re.compile(re.escape(CHUNK_ID_SEPARATOR) + r"\d+$")
|
||||
|
||||
# Break preference, best first. A chunk ends just after the chosen separator so
|
||||
# the separator itself is never lost.
|
||||
_BREAKS = ("\n\n", "\n", ". ", " ")
|
||||
# Don't take a break that would leave a chunk under this share of the target
|
||||
# size; a hard cut at full size is better than a 40-char chunk.
|
||||
_MIN_FILL = 0.3
|
||||
_TITLE_MAX_CHARS = 120
|
||||
# Metadata that describes a row's place in its page. Set by chunk_page; must be
|
||||
# dropped before a page's content is re-chunked.
|
||||
CHUNK_META_KEYS = frozenset({"page_id", "chunk", "chunk_count", "chunk_size"})
|
||||
|
||||
|
||||
def page_id_of(chunk_id: str) -> str:
|
||||
"""Map a chunk id back to its page id. Mirrors ``PAGE_ID_EXPR`` in SQL and
|
||||
``basePageId`` in the frontend — keep the three in sync."""
|
||||
stripped = _CHUNK_ID_RE.sub("", chunk_id)
|
||||
if stripped != chunk_id:
|
||||
return stripped
|
||||
return _LEGACY_CHUNK_ID_RE.sub("", chunk_id)
|
||||
|
||||
|
||||
def chunk_id(page_id: str, index: int) -> str:
|
||||
return f"{page_id}{CHUNK_ID_SEPARATOR}{index}"
|
||||
|
||||
|
||||
def split_text(text: str, size: Optional[int] = None) -> List[str]:
|
||||
"""Split ``text`` into consecutive pieces of at most ``size`` chars,
|
||||
preferring to break on paragraph, line, sentence then word boundaries.
|
||||
The pieces are disjoint verbatim slices: ``"".join(pieces) == text``."""
|
||||
size = settings.KNOWLEDGE_CHUNK_SIZE if size is None else size
|
||||
if size <= 0:
|
||||
raise ValueError("chunk size must be positive")
|
||||
|
||||
pieces: List[str] = []
|
||||
start = 0
|
||||
length = len(text)
|
||||
while start < length:
|
||||
if length - start <= size:
|
||||
pieces.append(text[start:])
|
||||
break
|
||||
end = _break_point(text, start, size)
|
||||
pieces.append(text[start:end])
|
||||
start = end
|
||||
return pieces
|
||||
|
||||
|
||||
def _break_point(text: str, start: int, size: int) -> int:
|
||||
"""End of the piece starting at ``start``: just after the best separator in
|
||||
the window, else a hard cut at ``size``. Always past ``start``."""
|
||||
window = text[start:start + size]
|
||||
floor = int(size * _MIN_FILL)
|
||||
for sep in _BREAKS:
|
||||
idx = window.rfind(sep)
|
||||
if idx >= floor:
|
||||
return start + idx + len(sep)
|
||||
return start + size
|
||||
|
||||
|
||||
@dataclass
|
||||
class ChunkDocument(Document):
|
||||
"""A chunk whose embedding is computed from more than its stored text: the
|
||||
page title and the tail of the previous chunk are prepended, so a mid-page
|
||||
chunk about pricing still matches "pricing" queries that only name the
|
||||
page, and a sentence cut by the boundary is embedded whole on one side."""
|
||||
|
||||
embedding_text: Optional[str] = None
|
||||
|
||||
def embed(self, embedder: Optional[Embedder] = None) -> None:
|
||||
_embedder = embedder or self.embedder
|
||||
if _embedder is None:
|
||||
raise ValueError("No embedder provided")
|
||||
self.embedding, self.usage = _embedder.get_embedding_and_usage(
|
||||
self.embedding_text or self.content
|
||||
)
|
||||
|
||||
|
||||
def given_title(document: Document) -> Optional[str]:
|
||||
"""The title the source recorded for the page, if any."""
|
||||
title = (document.meta_data or {}).get("title")
|
||||
if isinstance(title, str) and title.strip():
|
||||
return title.strip()[:_TITLE_MAX_CHARS]
|
||||
return None
|
||||
|
||||
|
||||
def page_title(document: Document) -> str:
|
||||
"""What to call the page when embedding: its recorded title, else its first
|
||||
non-empty line. The fallback is for the embedding only — it is never stored
|
||||
as the page's title, since a crawled first line is as likely to be a cookie
|
||||
banner as a heading."""
|
||||
title = given_title(document)
|
||||
if title:
|
||||
return title
|
||||
for line in (document.content or "").splitlines():
|
||||
if line.strip():
|
||||
return line.strip()[:_TITLE_MAX_CHARS]
|
||||
return document.id or ""
|
||||
|
||||
|
||||
def needs_explicit_chunk_id(page_id: Optional[str]) -> bool:
|
||||
"""A page id that ends in ``_<n>`` reads as a legacy chunk suffix to the
|
||||
grouping rule, which would file ``/blog/post_2`` under ``/blog/post`` and
|
||||
let an edit of the sibling delete it. Such pages get a ``::1`` id even
|
||||
when they fit in one chunk."""
|
||||
return bool(page_id) and bool(_LEGACY_CHUNK_ID_RE.search(page_id))
|
||||
|
||||
|
||||
def chunk_page(
|
||||
document: Document,
|
||||
size: Optional[int] = None,
|
||||
overlap: Optional[int] = None,
|
||||
) -> List[Document]:
|
||||
"""Split one page into chunk documents. A page that fits in one chunk is
|
||||
returned unchanged (same id, same metadata), so small pages are stored
|
||||
exactly as before. ``overlap`` is how much of the previous chunk is
|
||||
embedded with each chunk; it never changes what is stored."""
|
||||
size = settings.KNOWLEDGE_CHUNK_SIZE if size is None else size
|
||||
overlap = settings.KNOWLEDGE_CHUNK_OVERLAP if overlap is None else overlap
|
||||
content = document.content or ""
|
||||
if len(content) <= size and not needs_explicit_chunk_id(document.id):
|
||||
return [document]
|
||||
|
||||
pieces = split_text(content, size)
|
||||
title = page_title(document)
|
||||
page_id = document.id
|
||||
base_meta = {k: v for k, v in (document.meta_data or {}).items() if k not in CHUNK_META_KEYS}
|
||||
recorded_title = given_title(document)
|
||||
if recorded_title:
|
||||
base_meta["title"] = recorded_title
|
||||
chunks: List[Document] = []
|
||||
previous = ""
|
||||
for index, piece in enumerate(pieces, start=1):
|
||||
meta = {
|
||||
**base_meta,
|
||||
"page_id": page_id,
|
||||
"chunk": index,
|
||||
"chunk_count": len(pieces),
|
||||
"chunk_size": len(piece),
|
||||
}
|
||||
context = previous[-overlap:] if overlap > 0 else ""
|
||||
chunks.append(
|
||||
ChunkDocument(
|
||||
id=chunk_id(page_id, index),
|
||||
name=document.name,
|
||||
content=piece,
|
||||
meta_data=meta,
|
||||
embedder=document.embedder,
|
||||
embedding_text=f"{title}\n{context}{piece}",
|
||||
)
|
||||
)
|
||||
previous = piece
|
||||
return chunks
|
||||
|
||||
|
||||
class PageChunking(ChunkingStrategy):
|
||||
"""agno ``ChunkingStrategy`` adapter so the PDF readers chunk the same way
|
||||
as the website crawler."""
|
||||
|
||||
def __init__(self, size: Optional[int] = None, overlap: Optional[int] = None):
|
||||
self.size = size
|
||||
self.overlap = overlap
|
||||
|
||||
def chunk(self, document: Document) -> List[Document]:
|
||||
return chunk_page(document, self.size, self.overlap)
|
||||
Reference in new issue
Block a user