mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-20 00:14:50 +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,385 @@
|
||||
import io
|
||||
import logging
|
||||
import os
|
||||
import tempfile
|
||||
import uuid
|
||||
from typing import Any, Dict, List, Optional, Tuple
|
||||
|
||||
import numpy as np
|
||||
|
||||
from docsgpt.core.settings import settings
|
||||
from docsgpt.storage.storage_creator import StorageCreator
|
||||
from docsgpt.vectorstore.base import BaseVectorStore
|
||||
from docsgpt.vectorstore.document_class import Document
|
||||
from docsgpt.vectorstore.faiss_docstore import (
|
||||
dump_json_sidecar,
|
||||
dump_pickle_sidecar,
|
||||
load_json_sidecar,
|
||||
load_pickle_sidecar,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Sidecar holding chunk text and the row->id mapping. ``index.json`` is what
|
||||
# this version writes; ``index.pkl`` is langchain's historical format, still
|
||||
# read forever (uploads arrive in it) and still written for backward compat.
|
||||
JSON_SIDECAR = "index.json"
|
||||
PICKLE_SIDECAR = "index.pkl"
|
||||
FAISS_INDEX = "index.faiss"
|
||||
|
||||
|
||||
def _dependable_faiss_import():
|
||||
"""Import faiss, with a clearer message than the raw ImportError."""
|
||||
try:
|
||||
import faiss
|
||||
except ImportError as e:
|
||||
raise ImportError(
|
||||
"Could not import faiss. Install it with `pip install faiss-cpu`."
|
||||
) from e
|
||||
return faiss
|
||||
|
||||
|
||||
def get_vectorstore(path: str) -> str:
|
||||
"""Build a safe local path for a FAISS index.
|
||||
|
||||
Args:
|
||||
path: Source identifier provided by the caller.
|
||||
|
||||
Returns:
|
||||
The validated vectorstore path rooted under ``indexes``.
|
||||
|
||||
Raises:
|
||||
ValueError: If ``path`` escapes the ``indexes`` directory.
|
||||
"""
|
||||
base_dir = "indexes"
|
||||
if not path:
|
||||
return base_dir
|
||||
|
||||
normalized = str(path).strip()
|
||||
if "\\" in normalized:
|
||||
raise ValueError("Invalid source_id path")
|
||||
|
||||
candidate = os.path.normpath(os.path.join(base_dir, normalized))
|
||||
base_abs = os.path.abspath(base_dir)
|
||||
candidate_abs = os.path.abspath(candidate)
|
||||
|
||||
if not candidate_abs.startswith(base_abs + os.sep) and candidate_abs != base_abs:
|
||||
raise ValueError("Invalid source_id path")
|
||||
|
||||
return candidate
|
||||
|
||||
|
||||
class FaissStore(BaseVectorStore):
|
||||
"""Vector store backed by a local FAISS index.
|
||||
|
||||
Holds a flat L2 index plus an in-memory docstore mapping chunk ids to
|
||||
their text and metadata, persisted through :class:`StorageCreator`.
|
||||
"""
|
||||
|
||||
# Ranks by L2 distance (lower is better), not cosine — so the number here
|
||||
# is NOT comparable to the ``score_threshold`` the other stores honour,
|
||||
# and must not be shown as one.
|
||||
score_kind = "l2_distance"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
source_id: str,
|
||||
embeddings_key: str,
|
||||
docs_init=None,
|
||||
ids=None,
|
||||
batch_size=None,
|
||||
skip_dimension_check: bool = False,
|
||||
):
|
||||
"""Open or build one source's FAISS index.
|
||||
|
||||
Args:
|
||||
source_id: Source whose index to open.
|
||||
embeddings_key: API key handed to the embeddings provider.
|
||||
docs_init: Documents to build a fresh index from. Loads the stored
|
||||
index instead when omitted.
|
||||
ids: Chunk ids to keep when building. Generated when omitted.
|
||||
batch_size: Documents per embed call when building.
|
||||
skip_dimension_check: Open an index whose width does not match the
|
||||
configured model. Only for a caller that is about to replace
|
||||
that index, such as the re-embed script, which otherwise cannot
|
||||
read the chunks it needs to rebuild from.
|
||||
"""
|
||||
super().__init__()
|
||||
self.source_id = source_id
|
||||
self.path = get_vectorstore(source_id)
|
||||
self.embeddings = self._get_embeddings(settings.EMBEDDINGS_NAME, embeddings_key)
|
||||
self.storage = StorageCreator.get_storage()
|
||||
|
||||
self.index = None
|
||||
self.documents: Dict[str, Dict[str, Any]] = {}
|
||||
self.index_to_docstore_id: Dict[int, str] = {}
|
||||
|
||||
try:
|
||||
if docs_init:
|
||||
self._build_from_documents(docs_init, ids=ids, batch_size=batch_size)
|
||||
else:
|
||||
self._load_from_storage()
|
||||
except Exception as e:
|
||||
raise Exception(f"Error loading FAISS index: {str(e)}")
|
||||
|
||||
if not skip_dimension_check:
|
||||
self.assert_embedding_dimensions(self.embeddings)
|
||||
|
||||
# -- Construction ----------------------------------------------------
|
||||
|
||||
def _build_from_documents(self, docs_init, ids=None, batch_size=None) -> None:
|
||||
"""Create a fresh index seeded with ``docs_init``.
|
||||
|
||||
Args:
|
||||
docs_init: Documents to embed.
|
||||
ids: Chunk ids to keep. Generated when omitted, which renumbers
|
||||
every chunk and orphans anything referencing the old ids.
|
||||
batch_size: Documents per embed call. Without it the whole index
|
||||
goes out in one call, which a remote embeddings server rejects
|
||||
or times out on.
|
||||
"""
|
||||
texts, metadatas = [], []
|
||||
for doc in docs_init:
|
||||
texts.append(getattr(doc, "page_content", None) or getattr(doc, "text", "") or "")
|
||||
metadatas.append(getattr(doc, "metadata", None) or getattr(doc, "extra_info", None) or {})
|
||||
|
||||
faiss = _dependable_faiss_import()
|
||||
ids = list(ids) if ids else None
|
||||
step = batch_size if batch_size and batch_size > 0 else len(texts)
|
||||
for start in range(0, len(texts), max(1, step)):
|
||||
stop = start + max(1, step)
|
||||
vectors = self.embeddings.embed_documents(texts[start:stop])
|
||||
if self.index is None:
|
||||
self.index = faiss.IndexFlatL2(len(vectors[0]))
|
||||
self._append(
|
||||
texts[start:stop],
|
||||
metadatas[start:stop],
|
||||
vectors,
|
||||
ids[start:stop] if ids else None,
|
||||
)
|
||||
|
||||
def _load_from_storage(self) -> None:
|
||||
"""Load the index and its sidecar, preferring JSON over the pickle."""
|
||||
faiss = _dependable_faiss_import()
|
||||
faiss_path = f"{self.path}/{FAISS_INDEX}"
|
||||
json_path = f"{self.path}/{JSON_SIDECAR}"
|
||||
pickle_path = f"{self.path}/{PICKLE_SIDECAR}"
|
||||
|
||||
if not self.storage.file_exists(faiss_path):
|
||||
raise FileNotFoundError(f"Index files not found in storage at {self.path}")
|
||||
|
||||
if self.storage.file_exists(json_path):
|
||||
sidecar, loader = json_path, load_json_sidecar
|
||||
elif self.storage.file_exists(pickle_path):
|
||||
sidecar, loader = pickle_path, load_pickle_sidecar
|
||||
else:
|
||||
raise FileNotFoundError(f"Index files not found in storage at {self.path}")
|
||||
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
local_faiss = os.path.join(temp_dir, FAISS_INDEX)
|
||||
with open(local_faiss, "wb") as f:
|
||||
f.write(self.storage.get_file(faiss_path).read())
|
||||
self.index = faiss.read_index(local_faiss)
|
||||
|
||||
self.documents, self.index_to_docstore_id = loader(
|
||||
self.storage.get_file(sidecar).read()
|
||||
)
|
||||
|
||||
# -- Internals -------------------------------------------------------
|
||||
|
||||
def _append(self, texts, metadatas, vectors, ids=None) -> List[str]:
|
||||
"""Add embedded rows to the index and docstore, returning their ids."""
|
||||
ids = list(ids) if ids else [str(uuid.uuid4()) for _ in texts]
|
||||
self.index.add(np.array(vectors, dtype=np.float32))
|
||||
start = len(self.index_to_docstore_id)
|
||||
for offset, (text, metadata, doc_id) in enumerate(zip(texts, metadatas, ids)):
|
||||
self.documents[doc_id] = {
|
||||
"page_content": text,
|
||||
"metadata": dict(metadata or {}),
|
||||
}
|
||||
self.index_to_docstore_id[start + offset] = doc_id
|
||||
return ids
|
||||
|
||||
def _to_document(self, doc_id: str) -> Optional[Document]:
|
||||
stored = self.documents.get(doc_id)
|
||||
if stored is None:
|
||||
return None
|
||||
return Document(
|
||||
page_content=stored.get("page_content", ""),
|
||||
metadata=stored.get("metadata") or {},
|
||||
)
|
||||
|
||||
# -- Search ----------------------------------------------------------
|
||||
|
||||
def search(self, question: str, k: int = 4, *args, **kwargs) -> List[Document]:
|
||||
"""Return the ``k`` nearest chunks for ``question``."""
|
||||
return [doc for doc, _ in self.search_with_scores(question, k, *args, **kwargs)]
|
||||
|
||||
def search_with_scores(
|
||||
self,
|
||||
question: str,
|
||||
k: int = 4,
|
||||
*args,
|
||||
query_vector: Optional[List[float]] = None,
|
||||
**kwargs,
|
||||
) -> List[Tuple[Document, float]]:
|
||||
"""Same search as :meth:`search`, pairing each hit with its L2 distance.
|
||||
|
||||
Args:
|
||||
query_vector: Precomputed embedding of ``question``; when given the
|
||||
store skips embedding the query itself.
|
||||
"""
|
||||
# FAISS has no relevance-threshold knob; drop it so the per-source
|
||||
# score_threshold is safely ignored rather than crashing the forward.
|
||||
kwargs.pop("score_threshold", None)
|
||||
if self.index is None or self.index.ntotal == 0:
|
||||
return []
|
||||
|
||||
if query_vector is None:
|
||||
query_vector = self.embeddings.embed_query(question)
|
||||
vector = np.array([query_vector], dtype=np.float32)
|
||||
distances, rows = self.index.search(vector, min(k, self.index.ntotal))
|
||||
|
||||
results = []
|
||||
for distance, row in zip(distances[0], rows[0]):
|
||||
if row == -1:
|
||||
continue
|
||||
doc_id = self.index_to_docstore_id.get(int(row))
|
||||
document = self._to_document(doc_id) if doc_id else None
|
||||
if document is not None:
|
||||
results.append((document, float(distance)))
|
||||
return results
|
||||
|
||||
# -- Mutation --------------------------------------------------------
|
||||
|
||||
def add_texts(
|
||||
self,
|
||||
texts: List[str],
|
||||
metadatas: Optional[List[dict]] = None,
|
||||
*args,
|
||||
ids: Optional[List[str]] = None,
|
||||
**kwargs,
|
||||
) -> List[str]:
|
||||
"""Embed and append ``texts`` to the index."""
|
||||
texts = list(texts)
|
||||
if not texts:
|
||||
return []
|
||||
metadatas = list(metadatas or [{} for _ in texts])
|
||||
vectors = self.embeddings.embed_documents(texts)
|
||||
if self.index is None:
|
||||
faiss = _dependable_faiss_import()
|
||||
self.index = faiss.IndexFlatL2(len(vectors[0]))
|
||||
return self._append(texts, metadatas, vectors, ids)
|
||||
|
||||
def delete_index(self, ids: Optional[List[str]] = None, *args, **kwargs):
|
||||
"""Delete the given chunk ids, or the whole index when ids are omitted."""
|
||||
if ids is None:
|
||||
faiss = _dependable_faiss_import()
|
||||
dimension = self.index.d if self.index is not None else None
|
||||
self.index = faiss.IndexFlatL2(dimension) if dimension else None
|
||||
self.documents = {}
|
||||
self.index_to_docstore_id = {}
|
||||
return True
|
||||
|
||||
missing = set(ids) - set(self.documents)
|
||||
if missing:
|
||||
raise ValueError(f"Chunk ids not found in index: {sorted(missing)}")
|
||||
|
||||
rows_by_id = {doc_id: row for row, doc_id in self.index_to_docstore_id.items()}
|
||||
rows_to_drop = {rows_by_id[doc_id] for doc_id in ids}
|
||||
self.index.remove_ids(np.array(sorted(rows_to_drop), dtype=np.int64))
|
||||
for doc_id in ids:
|
||||
self.documents.pop(doc_id, None)
|
||||
|
||||
# remove_ids compacts the index, so the mapping has to be renumbered.
|
||||
remaining = [
|
||||
doc_id
|
||||
for row, doc_id in sorted(self.index_to_docstore_id.items())
|
||||
if row not in rows_to_drop
|
||||
]
|
||||
self.index_to_docstore_id = dict(enumerate(remaining))
|
||||
return True
|
||||
|
||||
def add_chunk(self, text: str, metadata: Optional[Dict[str, Any]] = None) -> str:
|
||||
"""Add a new chunk and save to storage."""
|
||||
ids = self.add_texts([text], [metadata or {}])
|
||||
self._save_to_storage()
|
||||
return ids[0]
|
||||
|
||||
def delete_chunk(self, chunk_id: str) -> bool:
|
||||
"""Delete a chunk and save to storage."""
|
||||
self.delete_index([chunk_id])
|
||||
self._save_to_storage()
|
||||
return True
|
||||
|
||||
# -- Persistence -----------------------------------------------------
|
||||
|
||||
def _write_index_files(self, directory: str) -> None:
|
||||
"""Write index.faiss plus both sidecars into ``directory``."""
|
||||
faiss = _dependable_faiss_import()
|
||||
os.makedirs(directory, exist_ok=True)
|
||||
faiss.write_index(self.index, os.path.join(directory, FAISS_INDEX))
|
||||
with open(os.path.join(directory, JSON_SIDECAR), "wb") as f:
|
||||
f.write(dump_json_sidecar(self.documents, self.index_to_docstore_id))
|
||||
with open(os.path.join(directory, PICKLE_SIDECAR), "wb") as f:
|
||||
f.write(dump_pickle_sidecar(self.documents, self.index_to_docstore_id))
|
||||
|
||||
def _save_to_storage(self) -> bool:
|
||||
"""Persist the index through the configured storage backend.
|
||||
|
||||
Each file is replaced atomically by the backend, so an interrupted save
|
||||
leaves the previous one readable. The three are still written in
|
||||
sequence: a crash between them pairs a new index with an older sidecar,
|
||||
which stays consistent because the ids and row order are preserved.
|
||||
"""
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
self._write_index_files(temp_dir)
|
||||
storage_path = get_vectorstore(self.source_id)
|
||||
for name in (FAISS_INDEX, JSON_SIDECAR, PICKLE_SIDECAR):
|
||||
with open(os.path.join(temp_dir, name), "rb") as f:
|
||||
self.storage.save_file(io.BytesIO(f.read()), f"{storage_path}/{name}")
|
||||
return True
|
||||
|
||||
def save_local(self, path: Optional[str] = None) -> bool:
|
||||
if path:
|
||||
self._write_index_files(path)
|
||||
self._save_to_storage()
|
||||
return True
|
||||
|
||||
# -- Introspection ---------------------------------------------------
|
||||
|
||||
def assert_embedding_dimensions(self, embeddings) -> None:
|
||||
"""Check the index width matches the embedding model's width.
|
||||
|
||||
This used to run only when ``EMBEDDINGS_NAME`` was mpnet, so every
|
||||
other model skipped the check entirely -- exactly the models most
|
||||
likely to differ from an index built earlier. It now runs for any
|
||||
model that reports a width.
|
||||
"""
|
||||
word_embedding_dimension = getattr(embeddings, "dimension", None)
|
||||
if word_embedding_dimension is None:
|
||||
# A remote model of unknown width reports None until its first
|
||||
# call; there is nothing to compare yet.
|
||||
return
|
||||
if self.index is None:
|
||||
return
|
||||
if word_embedding_dimension != self.index.d:
|
||||
raise ValueError(
|
||||
f"Embedding dimension mismatch: {settings.EMBEDDINGS_NAME} produces "
|
||||
f"{word_embedding_dimension}-dim vectors but this FAISS index is "
|
||||
f"{self.index.d}-dim. The index was built with a different "
|
||||
f"embedding model; re-embed it with "
|
||||
f"`python -m docsgpt.scripts.reembed` or point "
|
||||
f"EMBEDDINGS_NAME back at the original model."
|
||||
)
|
||||
|
||||
def get_chunks(self) -> List[Dict[str, Any]]:
|
||||
"""Return every chunk held in the index."""
|
||||
return [
|
||||
{
|
||||
"doc_id": doc_id,
|
||||
"text": stored.get("page_content", ""),
|
||||
"metadata": stored.get("metadata") or {},
|
||||
}
|
||||
for doc_id, stored in self.documents.items()
|
||||
]
|
||||
Reference in new issue
Block a user