mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-13 22:15:10 +00:00
refactor: rename the application package to docsgpt
The backend import package is now docsgpt, the name it will carry on PyPI; application was far too generic to install into anyone's site-packages. git mv plus a mechanical rewrite of every import, dotted string and path reference: 734 Python files, the compose files, Dockerfile, workflows, docs, setup scripts, devcontainer, k8s manifests, vscode config, pytest and coverage config, .gitignore. Behaviour is unchanged. Kept for one release: - A top-level application package whose meta-path finder resolves application.x.y to the already-imported docsgpt.x.y object, so old imports and entry points (celery -A application.app.celery, uvicorn application.asgi:asgi_app) keep working with a FutureWarning. - Celery registers every application.* task name as an alias of its docsgpt.* task on start-up, so messages queued by the previous release still run. The redbeat key prefix moves to redbeat:docsgpt:v2: so schedule entries the previous release wrote are left unread instead of firing twice. The backend image builds from the repository root (docker build -f docsgpt/Dockerfile .) so it can ship the alias package; a root .dockerignore allow-lists docsgpt/ and application/ and keeps caches, local data, .env files, the sample index files and the Dockerfile out. Compose and the image workflows point at the new context.
This commit is contained in:
1 parent
3e7f1f91b4
commit
574f96341e
984 files changed
+9847
-9682
No files matched your search
@@ -0,0 +1,542 @@
|
||||
"""Re-embed an existing index with the currently configured embedding model.
|
||||
|
||||
Changing ``EMBEDDINGS_NAME`` does not change vectors already stored. Because
|
||||
several models share a width -- granite-311m and all-mpnet-base-v2 are both
|
||||
768-dimensional -- a swapped model does not fail any dimension check. It simply
|
||||
retrieves against vectors that mean something else, which shows up as bad
|
||||
answers rather than as an error.
|
||||
|
||||
This script rebuilds the vectors from the chunk text already held in the store,
|
||||
so it never re-downloads, re-parses or re-chunks anything: no source files, no
|
||||
crawl, no docling. It works against pgvector and FAISS. With ``GRAPHRAG_ENABLED``
|
||||
it also rewrites ``graph_nodes.name_embedding``, which seeds graph traversal and
|
||||
would otherwise be left behind in the previous model's space.
|
||||
|
||||
Usage::
|
||||
|
||||
python -m docsgpt.scripts.reembed --dry-run # report, change nothing
|
||||
python -m docsgpt.scripts.reembed # re-embed everything
|
||||
python -m docsgpt.scripts.reembed --sources a,b # only these sources
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import logging
|
||||
import sys
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple
|
||||
|
||||
from psycopg import sql
|
||||
|
||||
from docsgpt.core.settings import settings
|
||||
from docsgpt.vectorstore.model_registry import resolve
|
||||
from docsgpt.vectorstore.vector_creator import VectorCreator
|
||||
|
||||
logger = logging.getLogger("reembed")
|
||||
|
||||
#: Stores this script knows how to rewrite in place.
|
||||
SUPPORTED_STORES = ("pgvector", "faiss")
|
||||
|
||||
#: Chunks embedded (and, for pgvector, written) per transaction.
|
||||
DEFAULT_BATCH_SIZE = 64
|
||||
|
||||
|
||||
@dataclass
|
||||
class _RebuildDoc:
|
||||
"""Minimal document shape ``FaissStore(docs_init=...)`` accepts."""
|
||||
|
||||
page_content: str
|
||||
metadata: Dict[str, Any]
|
||||
|
||||
|
||||
class ReembedError(RuntimeError):
|
||||
"""Raised for a misconfiguration the user needs to fix before running."""
|
||||
|
||||
|
||||
def _log_setup(verbose: bool) -> None:
|
||||
logging.basicConfig(
|
||||
level=logging.DEBUG if verbose else logging.INFO,
|
||||
format="%(asctime)s %(levelname)s %(message)s",
|
||||
)
|
||||
|
||||
|
||||
def _batched(items: Sequence[Any], size: int) -> Iterable[Sequence[Any]]:
|
||||
for start in range(0, len(items), size):
|
||||
yield items[start : start + size]
|
||||
|
||||
|
||||
def list_source_ids(store_type: str) -> List[str]:
|
||||
"""Every source id that has vectors in the configured store.
|
||||
|
||||
Args:
|
||||
store_type: ``"pgvector"`` or ``"faiss"``.
|
||||
|
||||
Returns:
|
||||
Source ids, sorted for a stable and resumable ordering.
|
||||
"""
|
||||
if store_type == "pgvector":
|
||||
return _pgvector_source_ids()
|
||||
return _faiss_source_ids()
|
||||
|
||||
|
||||
def _pgvector_source_ids() -> List[str]:
|
||||
store = VectorCreator.create_vectorstore("pgvector", source_id="")
|
||||
# ``_get_connection`` hands back a pooled connection, not a context
|
||||
# manager: closing it via ``with`` would cost the pool a slot.
|
||||
conn = store._get_connection()
|
||||
cursor = conn.cursor()
|
||||
try:
|
||||
# Table and column names cannot be bound as parameters, so they go
|
||||
# through ``sql.Identifier``, which quotes them. They are internal
|
||||
# constants rather than user input, but composing SQL by f-string is
|
||||
# the habit worth not having.
|
||||
cursor.execute(
|
||||
sql.SQL(
|
||||
"SELECT DISTINCT source_id FROM {table} "
|
||||
"WHERE source_id IS NOT NULL ORDER BY source_id"
|
||||
).format(table=sql.Identifier(store._table_name))
|
||||
)
|
||||
return [row[0] for row in cursor.fetchall()]
|
||||
finally:
|
||||
cursor.close()
|
||||
store.close()
|
||||
|
||||
|
||||
def _faiss_source_ids() -> List[str]:
|
||||
"""Directories under the FAISS root that hold an index."""
|
||||
from docsgpt.storage.storage_creator import StorageCreator
|
||||
|
||||
storage = StorageCreator.get_storage()
|
||||
root = "indexes"
|
||||
ids: List[str] = []
|
||||
try:
|
||||
for path in storage.list_files(root):
|
||||
# ``indexes/<source_id>/index.faiss``
|
||||
parts = str(path).split("/")
|
||||
if parts and parts[-1].endswith(".faiss") and len(parts) >= 2:
|
||||
ids.append(parts[-2])
|
||||
except Exception as exc: # noqa: BLE001 -- surfaced to the user below
|
||||
raise ReembedError(
|
||||
f"Could not list FAISS indexes under {root!r}: {exc}"
|
||||
) from exc
|
||||
return sorted(set(ids))
|
||||
|
||||
|
||||
def _graph_nodes_exist(conn) -> bool:
|
||||
"""True when the GraphRAG node table is present in this database."""
|
||||
cursor = conn.cursor()
|
||||
try:
|
||||
cursor.execute("SELECT to_regclass('public.graph_nodes')")
|
||||
row = cursor.fetchone()
|
||||
return bool(row and row[0])
|
||||
finally:
|
||||
cursor.close()
|
||||
|
||||
|
||||
def reembed_graph_nodes(store, conn, source_id: str, batch_size: int, dry_run: bool) -> int:
|
||||
"""Rewrite one source's GraphRAG node name vectors.
|
||||
|
||||
``graph_nodes.name_embedding`` is what every graph traversal starts from:
|
||||
the retriever embeds the query and takes the nearest node names as seeds.
|
||||
It is written once at extraction time and never revisited, so re-embedding
|
||||
only the chunk table leaves the graph seeding against the previous model.
|
||||
Nothing errors -- mpnet and granite-311m are both 768-dimensional, so the
|
||||
column accepts the mismatch -- the graph just walks from the wrong nodes.
|
||||
|
||||
The embedded text is ``graph_nodes.name``, which is already stored, so this
|
||||
needs no LLM re-extraction.
|
||||
|
||||
Args:
|
||||
store: The pgvector store, for its embeddings client.
|
||||
conn: Open connection to the same database.
|
||||
source_id: Source whose nodes to re-embed.
|
||||
batch_size: Names per embed call and per transaction.
|
||||
dry_run: When true, count the work and change nothing.
|
||||
|
||||
Returns:
|
||||
Number of node rows seen (dry run) or rewritten.
|
||||
"""
|
||||
if not _graph_nodes_exist(conn):
|
||||
return 0
|
||||
|
||||
cursor = conn.cursor()
|
||||
try:
|
||||
cursor.execute(
|
||||
"SELECT id, name FROM graph_nodes "
|
||||
"WHERE source_id = %s AND name_embedding IS NOT NULL ORDER BY id",
|
||||
(source_id,),
|
||||
)
|
||||
rows = cursor.fetchall()
|
||||
finally:
|
||||
cursor.close()
|
||||
|
||||
if dry_run or not rows:
|
||||
return len(rows)
|
||||
|
||||
written = 0
|
||||
for batch in _batched(rows, batch_size):
|
||||
vectors = store._embedding.embed_documents([row[1] or "" for row in batch])
|
||||
cursor = conn.cursor()
|
||||
try:
|
||||
cursor.executemany(
|
||||
"UPDATE graph_nodes SET name_embedding = %s::vector WHERE id = %s",
|
||||
[
|
||||
(str(list(vector)), row[0])
|
||||
for vector, row in zip(vectors, batch)
|
||||
],
|
||||
)
|
||||
conn.commit()
|
||||
except Exception:
|
||||
conn.rollback()
|
||||
raise
|
||||
finally:
|
||||
cursor.close()
|
||||
written += len(batch)
|
||||
return written
|
||||
|
||||
|
||||
def _count_source_chunks(conn, table: str, source_id: str) -> int:
|
||||
"""How many chunk rows one source holds."""
|
||||
cursor = conn.cursor()
|
||||
try:
|
||||
cursor.execute(
|
||||
sql.SQL("SELECT count(*) FROM {table} WHERE source_id = %s").format(
|
||||
table=sql.Identifier(table)
|
||||
),
|
||||
(source_id,),
|
||||
)
|
||||
return int(cursor.fetchone()[0])
|
||||
finally:
|
||||
cursor.close()
|
||||
|
||||
|
||||
def _fetch_chunk_page(
|
||||
conn, table: str, source_id: str, after_id: Optional[Any], limit: int
|
||||
) -> List[Tuple[Any, str]]:
|
||||
"""One page of ``(id, text)`` ordered by id, starting after ``after_id``.
|
||||
|
||||
Keyed off ``id`` rather than ``OFFSET`` so each page costs an index seek,
|
||||
and the page boundaries stay stable across the updates the caller makes
|
||||
between pages -- only the vector column changes, never ``id``.
|
||||
"""
|
||||
cursor = conn.cursor()
|
||||
try:
|
||||
if after_id is None:
|
||||
cursor.execute(
|
||||
sql.SQL(
|
||||
"SELECT id, text FROM {table} WHERE source_id = %s "
|
||||
"ORDER BY id LIMIT %s"
|
||||
).format(table=sql.Identifier(table)),
|
||||
(source_id, limit),
|
||||
)
|
||||
else:
|
||||
cursor.execute(
|
||||
sql.SQL(
|
||||
"SELECT id, text FROM {table} WHERE source_id = %s AND id > %s "
|
||||
"ORDER BY id LIMIT %s"
|
||||
).format(table=sql.Identifier(table)),
|
||||
(source_id, after_id, limit),
|
||||
)
|
||||
return cursor.fetchall()
|
||||
finally:
|
||||
cursor.close()
|
||||
|
||||
|
||||
def _write_chunk_page(conn, table: str, vector_column: str, store, page) -> None:
|
||||
"""Embed one page and commit its vectors."""
|
||||
vectors = store._embedding.embed_documents([row[1] or "" for row in page])
|
||||
cursor = conn.cursor()
|
||||
try:
|
||||
cursor.executemany(
|
||||
sql.SQL("UPDATE {table} SET {column} = %s::vector WHERE id = %s").format(
|
||||
table=sql.Identifier(table),
|
||||
column=sql.Identifier(vector_column),
|
||||
),
|
||||
[(str(list(vector)), row[0]) for vector, row in zip(vectors, page)],
|
||||
)
|
||||
conn.commit()
|
||||
except Exception:
|
||||
conn.rollback()
|
||||
raise
|
||||
finally:
|
||||
cursor.close()
|
||||
|
||||
|
||||
def reembed_pgvector(source_id: str, batch_size: int, dry_run: bool) -> Tuple[int, int]:
|
||||
"""Rewrite one source's vectors in place.
|
||||
|
||||
Rows are updated, never deleted and re-inserted, so an interrupted run
|
||||
leaves every chunk present and simply re-does its batch next time.
|
||||
|
||||
Chunks are read a page at a time rather than all at once. A source's text is
|
||||
the whole corpus: at the default chunk size 200k chunks materialise to about
|
||||
1.6 GB of Python strings, several times that for non-Latin scripts, and the
|
||||
result set is held alongside them until the cursor closes. That is enough to
|
||||
be OOM-killed inside a 4 GiB container while also holding the model -- and a
|
||||
kill here is what a torn index costs the most.
|
||||
|
||||
Args:
|
||||
source_id: Source whose chunks to re-embed.
|
||||
batch_size: Chunks per page, per embed call and per transaction.
|
||||
dry_run: When true, count the work and change nothing.
|
||||
|
||||
Returns:
|
||||
``(chunks_seen, chunks_written)``.
|
||||
"""
|
||||
store = VectorCreator.create_vectorstore("pgvector", source_id=source_id)
|
||||
table, vector_column = store._table_name, store._vector_column
|
||||
conn = store._get_connection()
|
||||
|
||||
seen = written = 0
|
||||
try:
|
||||
if dry_run:
|
||||
seen = _count_source_chunks(conn, table, source_id)
|
||||
else:
|
||||
after_id: Optional[Any] = None
|
||||
while True:
|
||||
page = _fetch_chunk_page(
|
||||
conn, table, source_id, after_id, batch_size
|
||||
)
|
||||
if not page:
|
||||
break
|
||||
_write_chunk_page(conn, table, vector_column, store, page)
|
||||
after_id = page[-1][0]
|
||||
seen += len(page)
|
||||
written += len(page)
|
||||
|
||||
# The graph seeds every traversal from its own vectors, so leaving them
|
||||
# in the old model's space is the same silent mismatch this script
|
||||
# exists to remove -- and at equal widths nothing would report it.
|
||||
if getattr(settings, "GRAPHRAG_ENABLED", False):
|
||||
nodes = reembed_graph_nodes(store, conn, source_id, batch_size, dry_run)
|
||||
if nodes:
|
||||
logger.info(
|
||||
" %s: %d graph node name(s)%s",
|
||||
source_id,
|
||||
nodes,
|
||||
"" if dry_run else " re-embedded",
|
||||
)
|
||||
finally:
|
||||
store.close()
|
||||
return seen, written
|
||||
|
||||
|
||||
def reembed_faiss(source_id: str, batch_size: int, dry_run: bool) -> Tuple[int, int]:
|
||||
"""Rebuild one FAISS index from the chunk text in its sidecar.
|
||||
|
||||
A FAISS index is a flat array, so this rebuilds rather than updates. Every
|
||||
new vector is computed *before* the old index is touched, so an interrupt
|
||||
during embedding leaves the existing index intact -- and ``save_local``
|
||||
writes each file to a temporary path and moves it into place, so an
|
||||
interrupt during the write does too.
|
||||
|
||||
Args:
|
||||
source_id: Source whose index to rebuild.
|
||||
batch_size: Chunks per embed call.
|
||||
dry_run: When true, count the work and change nothing.
|
||||
|
||||
Returns:
|
||||
``(chunks_seen, chunks_written)``.
|
||||
"""
|
||||
# A width change is the main reason to run this, and it is exactly what
|
||||
# ``assert_embedding_dimensions`` refuses to open. The chunk text lives in
|
||||
# the sidecar rather than the index, so reading it needs no matching width,
|
||||
# and the index this reads is replaced below.
|
||||
store = VectorCreator.create_vectorstore(
|
||||
"faiss",
|
||||
source_id=source_id,
|
||||
embeddings_key=settings.EMBEDDINGS_KEY,
|
||||
skip_dimension_check=True,
|
||||
)
|
||||
chunks: List[Dict[str, Any]] = store.get_chunks() or []
|
||||
if dry_run or not chunks:
|
||||
return len(chunks), 0
|
||||
|
||||
docs = [
|
||||
_RebuildDoc(chunk.get("text") or "", chunk.get("metadata") or {})
|
||||
for chunk in chunks
|
||||
]
|
||||
|
||||
# Constructing with ``docs_init`` embeds everything into a fresh in-memory
|
||||
# index and touches nothing on disk. The existing index is only replaced by
|
||||
# ``save_local`` below, so a failure while embedding leaves it intact.
|
||||
# Keep the existing ids: GraphRAG's ``graph_node_chunks`` rows and any id a
|
||||
# client already holds point at them.
|
||||
ids = [chunk.get("doc_id") for chunk in chunks]
|
||||
rebuilt = VectorCreator.create_vectorstore(
|
||||
"faiss",
|
||||
source_id=source_id,
|
||||
embeddings_key=settings.EMBEDDINGS_KEY,
|
||||
docs_init=docs,
|
||||
ids=ids if all(ids) else None,
|
||||
batch_size=batch_size,
|
||||
)
|
||||
rebuilt.save_local()
|
||||
return len(chunks), len(docs)
|
||||
|
||||
|
||||
def record_source_model(source_id: str) -> None:
|
||||
"""Stamp ``sources.model`` with the model these vectors were just built by.
|
||||
|
||||
``sources`` lives in the user-data database while the vectors may not, so
|
||||
this takes its own session. Left unwritten, the column keeps naming the old
|
||||
model and every consumer that trusts it -- the boot mismatch check above
|
||||
all -- reports a source as stale immediately after it was migrated.
|
||||
"""
|
||||
from sqlalchemy import text
|
||||
|
||||
from docsgpt.storage.db.session import db_session
|
||||
|
||||
try:
|
||||
with db_session() as conn:
|
||||
conn.execute(
|
||||
text("UPDATE sources SET model = :model WHERE id = :id"),
|
||||
{"model": settings.EMBEDDINGS_NAME, "id": source_id},
|
||||
)
|
||||
except Exception as exc: # noqa: BLE001 — the vectors are already rewritten
|
||||
logger.warning(
|
||||
" %s: re-embedded, but could not update sources.model (%s). Retrieval "
|
||||
"is correct; the mismatch warning may persist until it is.",
|
||||
source_id,
|
||||
exc,
|
||||
)
|
||||
|
||||
|
||||
def run(
|
||||
store_type: str,
|
||||
source_ids: Optional[Sequence[str]],
|
||||
batch_size: int,
|
||||
dry_run: bool,
|
||||
) -> int:
|
||||
"""Re-embed the requested sources. Returns a process exit code."""
|
||||
model = settings.EMBEDDINGS_NAME
|
||||
spec = resolve(model)
|
||||
logger.info(
|
||||
"Re-embedding %s with %s (%s)",
|
||||
store_type,
|
||||
model,
|
||||
f"{spec.dimension} dims" if spec else "width determined at runtime",
|
||||
)
|
||||
if dry_run:
|
||||
logger.info("Dry run: nothing will be written.")
|
||||
|
||||
ids = list(source_ids) if source_ids else list_source_ids(store_type)
|
||||
if not ids:
|
||||
logger.warning("No sources found; nothing to do.")
|
||||
return 0
|
||||
logger.info("%d source(s) to process", len(ids))
|
||||
|
||||
handler = reembed_pgvector if store_type == "pgvector" else reembed_faiss
|
||||
total_seen = total_written = 0
|
||||
failures: List[str] = []
|
||||
started = time.monotonic()
|
||||
|
||||
for position, source_id in enumerate(ids, start=1):
|
||||
try:
|
||||
seen, written = handler(source_id, batch_size, dry_run)
|
||||
if written and not dry_run:
|
||||
record_source_model(source_id)
|
||||
total_seen += seen
|
||||
total_written += written
|
||||
logger.info(
|
||||
"[%d/%d] %s: %d chunk(s)%s",
|
||||
position,
|
||||
len(ids),
|
||||
source_id,
|
||||
seen,
|
||||
"" if dry_run else f", {written} re-embedded",
|
||||
)
|
||||
except Exception as exc: # noqa: BLE001 -- one bad source must not stop the run
|
||||
failures.append(source_id)
|
||||
logger.error("[%d/%d] %s FAILED: %s", position, len(ids), source_id, exc)
|
||||
|
||||
elapsed = time.monotonic() - started
|
||||
logger.info(
|
||||
"Done in %.1fs: %d chunk(s) seen, %d re-embedded, %d source(s) failed",
|
||||
elapsed,
|
||||
total_seen,
|
||||
total_written,
|
||||
len(failures),
|
||||
)
|
||||
if failures:
|
||||
logger.error("Failed sources: %s", ", ".join(failures))
|
||||
logger.error("Re-run with --sources %s to retry just those.", ",".join(failures))
|
||||
return 1
|
||||
if dry_run:
|
||||
logger.info("Re-run without --dry-run to apply.")
|
||||
return 0
|
||||
|
||||
|
||||
def build_parser() -> argparse.ArgumentParser:
|
||||
parser = argparse.ArgumentParser(
|
||||
prog="python -m docsgpt.scripts.reembed",
|
||||
description=__doc__.split("\n\n")[0],
|
||||
)
|
||||
parser.add_argument(
|
||||
"--dry-run", action="store_true", help="report what would change, write nothing"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--sources", help="comma-separated source ids; default is every source"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--batch-size",
|
||||
type=int,
|
||||
default=DEFAULT_BATCH_SIZE,
|
||||
help=f"chunks per embed call (default {DEFAULT_BATCH_SIZE})",
|
||||
)
|
||||
parser.add_argument("-v", "--verbose", action="store_true", help="debug logging")
|
||||
return parser
|
||||
|
||||
|
||||
def main(argv: Optional[Sequence[str]] = None) -> int:
|
||||
args = build_parser().parse_args(argv)
|
||||
_log_setup(args.verbose)
|
||||
|
||||
# Which model this installation uses may live in ``app_metadata`` rather
|
||||
# than the environment -- ``docsgpt.app`` resolves it at boot, and this
|
||||
# script never imports that. Without this, an install pinned to granite
|
||||
# with no EMBEDDINGS_NAME set (every stock Kubernetes deployment: the
|
||||
# manifests carry no embedding config at all) would re-embed its whole
|
||||
# index with the *legacy* code default and stamp ``sources.model`` to
|
||||
# match -- the silent cross-model index this script exists to repair.
|
||||
from docsgpt.storage.db.embeddings_pin import resolve_embeddings_pin
|
||||
|
||||
resolve_embeddings_pin(logger)
|
||||
|
||||
# Embed in this process. ``EMBEDDINGS_DELEGATE_TO_WORKER`` exists to keep a
|
||||
# model out of the API, which serves one query at a time and holds the
|
||||
# model for nothing in between. This is the opposite case: a batch job that
|
||||
# embeds every chunk in the index, where a broker round trip per batch adds
|
||||
# latency and a dependency on a worker running. Loading the model here also
|
||||
# means the script reports a real failure for a model it cannot load,
|
||||
# instead of timing out against an empty queue.
|
||||
if getattr(settings, "EMBEDDINGS_DELEGATE_TO_WORKER", False):
|
||||
logger.info("Embedding in-process; worker delegation does not apply here.")
|
||||
settings.EMBEDDINGS_DELEGATE_TO_WORKER = False
|
||||
|
||||
store_type = (settings.VECTOR_STORE or "").lower()
|
||||
if store_type not in SUPPORTED_STORES:
|
||||
logger.error(
|
||||
"VECTOR_STORE is %r; this script supports %s. For other stores, "
|
||||
"delete and re-ingest the affected sources.",
|
||||
settings.VECTOR_STORE,
|
||||
" and ".join(SUPPORTED_STORES),
|
||||
)
|
||||
return 2
|
||||
|
||||
sources = (
|
||||
[s.strip() for s in args.sources.split(",") if s.strip()]
|
||||
if args.sources
|
||||
else None
|
||||
)
|
||||
try:
|
||||
return run(store_type, sources, max(1, args.batch_size), args.dry_run)
|
||||
except ReembedError as exc:
|
||||
logger.error("%s", exc)
|
||||
return 2
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
Reference in new issue
Block a user