mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-04 12:13:05 +00:00
Alex
574f96341e
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.
390 lines
15 KiB
Python
390 lines
15 KiB
Python
"""Boot-time ownership of the pgvector schema.
|
|
|
|
``PGVectorStore`` used to run its DDL on every instantiation — once per source
|
|
per request. The schema is created here instead, once per process, and the boot
|
|
hook is also the only place that can catch an embedding-dimension mismatch:
|
|
a ``documents`` table built for a different model silently retrieves garbage.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import ast
|
|
from pathlib import Path
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import pytest
|
|
|
|
from docsgpt.storage.db.bootstrap import ensure_vector_schema
|
|
|
|
_APP_PY = Path(__file__).resolve().parents[3] / "docsgpt" / "app.py"
|
|
|
|
|
|
@pytest.fixture
|
|
def vector_settings(monkeypatch):
|
|
"""Settings configured for a pgvector deployment."""
|
|
from docsgpt.core import settings as settings_module
|
|
|
|
settings = settings_module.settings
|
|
monkeypatch.setattr(settings, "VECTOR_STORE", "pgvector", raising=False)
|
|
monkeypatch.setattr(
|
|
settings,
|
|
"PGVECTOR_CONNECTION_STRING",
|
|
"postgresql://user:pass@localhost/db",
|
|
raising=False,
|
|
)
|
|
monkeypatch.setattr(settings, "GRAPHRAG_ENABLED", False, raising=False)
|
|
monkeypatch.setattr(settings, "EMBEDDINGS_NAME", "test-model", raising=False)
|
|
return settings
|
|
|
|
|
|
def _embeddings(dimension):
|
|
stub = MagicMock()
|
|
stub.dimension = dimension
|
|
return stub
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestEnsureVectorSchemaSkips:
|
|
def test_skips_when_the_store_is_not_pgvector(self, vector_settings, monkeypatch):
|
|
monkeypatch.setattr(vector_settings, "VECTOR_STORE", "faiss", raising=False)
|
|
|
|
with patch("psycopg.connect") as connect:
|
|
ensure_vector_schema()
|
|
|
|
connect.assert_not_called()
|
|
|
|
def test_skips_when_no_connection_string_is_configured(
|
|
self, vector_settings, monkeypatch
|
|
):
|
|
monkeypatch.setattr(
|
|
vector_settings, "PGVECTOR_CONNECTION_STRING", None, raising=False
|
|
)
|
|
monkeypatch.setattr(vector_settings, "POSTGRES_URI", None, raising=False)
|
|
|
|
with patch("psycopg.connect") as connect:
|
|
ensure_vector_schema()
|
|
|
|
connect.assert_not_called()
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestEnsureVectorSchemaCreates:
|
|
def _run(self, dimension=768, table_dimension=768):
|
|
conn = MagicMock()
|
|
cursor = MagicMock()
|
|
conn.cursor.return_value = cursor
|
|
with patch("psycopg.connect", return_value=conn) as connect, patch(
|
|
"docsgpt.vectorstore.model_registry.dimension_for",
|
|
return_value=dimension,
|
|
), patch(
|
|
"docsgpt.vectorstore.pgvector.PGVectorStore.create_schema"
|
|
) as vector_schema, patch(
|
|
"docsgpt.vectorstore.pgvector.PGVectorStore.table_dimension",
|
|
return_value=table_dimension,
|
|
), patch(
|
|
"docsgpt.graphrag.store.GraphStore.create_schema"
|
|
) as graph_schema, patch(
|
|
"docsgpt.vectorstore.pgvector._pool_for"
|
|
) as pool_for:
|
|
ensure_vector_schema()
|
|
return connect, conn, cursor, vector_schema, graph_schema, pool_for
|
|
|
|
def test_creates_the_vector_schema_under_the_advisory_lock(self, vector_settings):
|
|
connect, conn, cursor, vector_schema, graph_schema, _ = self._run()
|
|
|
|
# A bounded connect: an unreachable or suspended vector DB must fail the
|
|
# hook, not hang boot until a liveness probe kills the process.
|
|
connect.assert_called_once_with(
|
|
"postgresql://user:pass@localhost/db", connect_timeout=10
|
|
)
|
|
statements = " ".join(str(c) for c in cursor.execute.call_args_list)
|
|
assert "pg_advisory_xact_lock" in statements
|
|
vector_schema.assert_called_once()
|
|
assert vector_schema.call_args.kwargs["dimension"] == 768
|
|
graph_schema.assert_not_called()
|
|
conn.commit.assert_called_once()
|
|
conn.close.assert_called_once()
|
|
|
|
def test_creates_the_graph_schema_only_when_graphrag_is_enabled(
|
|
self, vector_settings, monkeypatch
|
|
):
|
|
monkeypatch.setattr(vector_settings, "GRAPHRAG_ENABLED", True, raising=False)
|
|
|
|
_, _, _, _, graph_schema, _ = self._run()
|
|
|
|
graph_schema.assert_called_once()
|
|
assert graph_schema.call_args.kwargs["dimension"] == 768
|
|
|
|
def test_uses_a_direct_connection_never_the_store_pool(self, vector_settings):
|
|
# The hook can run pre-fork under ``gunicorn --preload``; a pooled
|
|
# socket inherited by a worker is a corrupted connection.
|
|
_, _, _, _, _, pool_for = self._run()
|
|
|
|
pool_for.assert_not_called()
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestEnsureVectorSchemaDimensionCheck:
|
|
def test_raises_when_the_table_width_disagrees_with_the_model(
|
|
self, vector_settings
|
|
):
|
|
conn = MagicMock()
|
|
with patch("psycopg.connect", return_value=conn), patch(
|
|
"docsgpt.vectorstore.model_registry.dimension_for",
|
|
return_value=1536,
|
|
), patch("docsgpt.vectorstore.pgvector.PGVectorStore.create_schema"), patch(
|
|
"docsgpt.vectorstore.pgvector.PGVectorStore.table_dimension",
|
|
return_value=768,
|
|
):
|
|
with pytest.raises(RuntimeError) as excinfo:
|
|
ensure_vector_schema()
|
|
|
|
message = str(excinfo.value)
|
|
assert "768" in message and "1536" in message
|
|
assert "test-model" in message
|
|
conn.close.assert_called_once()
|
|
|
|
def test_skips_the_check_when_the_embeddings_expose_no_dimension(
|
|
self, vector_settings
|
|
):
|
|
conn = MagicMock()
|
|
stub = MagicMock()
|
|
del stub.dimension
|
|
with patch("psycopg.connect", return_value=conn), patch(
|
|
"docsgpt.vectorstore.base.get_embeddings", return_value=stub
|
|
), patch(
|
|
"docsgpt.vectorstore.pgvector.PGVectorStore.create_schema"
|
|
) as vector_schema, patch(
|
|
"docsgpt.vectorstore.pgvector.PGVectorStore.table_dimension",
|
|
return_value=1536,
|
|
):
|
|
ensure_vector_schema() # must not raise
|
|
|
|
assert vector_schema.call_args.kwargs["dimension"] == 768
|
|
|
|
def test_skips_the_check_when_the_model_cannot_be_loaded(self, vector_settings):
|
|
conn = MagicMock()
|
|
with patch("psycopg.connect", return_value=conn), patch(
|
|
"docsgpt.vectorstore.base.get_embeddings",
|
|
side_effect=RuntimeError("no model"),
|
|
), patch(
|
|
"docsgpt.vectorstore.pgvector.PGVectorStore.create_schema"
|
|
) as vector_schema, patch(
|
|
"docsgpt.vectorstore.pgvector.PGVectorStore.table_dimension",
|
|
return_value=1536,
|
|
):
|
|
ensure_vector_schema() # must not raise
|
|
|
|
assert vector_schema.call_args.kwargs["dimension"] == 768
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestBootGating:
|
|
"""The boot hook must stay behind AUTO_VECTOR_SCHEMA.
|
|
|
|
Asserted on the source because importing ``docsgpt.app`` runs the hook,
|
|
and a test that imports it cannot observe the gate it just executed.
|
|
"""
|
|
|
|
def test_app_calls_the_hook_behind_the_setting(self):
|
|
source = _APP_PY.read_text()
|
|
|
|
assert "if settings.AUTO_VECTOR_SCHEMA:" in source
|
|
assert "ensure_vector_schema(" in source
|
|
|
|
def test_the_boot_hook_cannot_take_the_process_down(self):
|
|
"""A vector-DB fault must degrade retrieval, not crash-loop the app.
|
|
|
|
The hook runs at import time, so an exception escaping it stops
|
|
gunicorn and every Celery worker from booting -- taking auth, chat
|
|
history and webhooks down with retrieval. ``PGVectorStore`` re-checks
|
|
the schema on its write path, so failing soft here loses nothing.
|
|
"""
|
|
tree = ast.parse(_APP_PY.read_text())
|
|
|
|
def _calls_hook(node) -> bool:
|
|
return any(
|
|
isinstance(n, ast.Call)
|
|
and isinstance(n.func, ast.Name)
|
|
and n.func.id == "ensure_vector_schema"
|
|
for n in ast.walk(node)
|
|
)
|
|
|
|
assert any(
|
|
isinstance(node, ast.Try) and node.handlers and _calls_hook(node)
|
|
for node in ast.walk(tree)
|
|
), "ensure_vector_schema() at import time must be wrapped in try/except"
|
|
|
|
def test_setting_defaults_on(self):
|
|
from docsgpt.core.settings import Settings
|
|
|
|
assert Settings.model_fields["AUTO_VECTOR_SCHEMA"].default is True
|
|
|
|
def test_test_suite_opts_out_by_default(self):
|
|
conftest = Path(__file__).resolve().parents[2] / "conftest.py"
|
|
|
|
assert 'os.environ.setdefault("AUTO_VECTOR_SCHEMA", "false")' in (
|
|
conftest.read_text()
|
|
)
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestBootDoesNotLoadTheModel:
|
|
"""The hook needs an integer, not an inference session.
|
|
|
|
It used to build the embeddings instance to read ``.dimension`` off it,
|
|
loading several hundred MB of ONNX into every API and worker process at
|
|
import. For a model the registry describes that is a lookup.
|
|
"""
|
|
|
|
def _run(self, registry_dim, loader):
|
|
conn = MagicMock()
|
|
conn.cursor.return_value = MagicMock()
|
|
with patch("psycopg.connect", return_value=conn), patch(
|
|
"docsgpt.vectorstore.model_registry.dimension_for",
|
|
return_value=registry_dim,
|
|
), patch(
|
|
"docsgpt.vectorstore.base.build_local_embeddings", loader
|
|
), patch(
|
|
"docsgpt.vectorstore.pgvector.PGVectorStore.create_schema"
|
|
) as vector_schema, patch(
|
|
"docsgpt.vectorstore.pgvector.PGVectorStore.table_dimension",
|
|
return_value=None,
|
|
):
|
|
ensure_vector_schema()
|
|
return vector_schema
|
|
|
|
def test_a_registered_model_is_never_constructed(self, vector_settings):
|
|
loader = MagicMock()
|
|
vector_schema = self._run(768, loader)
|
|
loader.assert_not_called()
|
|
assert vector_schema.call_args.kwargs["dimension"] == 768
|
|
|
|
def test_an_unregistered_model_still_falls_back_to_loading(self, vector_settings):
|
|
loader = MagicMock(return_value=_embeddings(1024))
|
|
vector_schema = self._run(None, loader)
|
|
loader.assert_called_once()
|
|
assert vector_schema.call_args.kwargs["dimension"] == 1024
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestUnknownWidthIsProbed:
|
|
"""A remote server's width is only knowable by asking it.
|
|
|
|
``RemoteEmbeddings`` reports ``None`` until its first call, so sizing the
|
|
table from the attribute alone fell back to 768 and skipped the mismatch
|
|
check — the silent ``vector(768)`` column this hook exists to prevent.
|
|
"""
|
|
|
|
def _run(self, remote, table_dimension=1024):
|
|
conn = MagicMock()
|
|
conn.cursor.return_value = MagicMock()
|
|
with patch("psycopg.connect", return_value=conn), patch(
|
|
"docsgpt.vectorstore.model_registry.dimension_for", return_value=None
|
|
), patch(
|
|
"docsgpt.vectorstore.base.build_local_embeddings", return_value=remote
|
|
), patch(
|
|
"docsgpt.vectorstore.pgvector.PGVectorStore.create_schema"
|
|
) as vector_schema, patch(
|
|
"docsgpt.vectorstore.pgvector.PGVectorStore.table_dimension",
|
|
return_value=table_dimension,
|
|
):
|
|
try:
|
|
ensure_vector_schema()
|
|
raised = False
|
|
except RuntimeError:
|
|
raised = True
|
|
return vector_schema, raised
|
|
|
|
@staticmethod
|
|
def _remote(width=None, error=None):
|
|
remote = MagicMock()
|
|
remote.dimension = None
|
|
remote.embed_query.side_effect = error or (lambda _text: [0.0] * width)
|
|
return remote
|
|
|
|
def test_the_table_is_sized_from_the_probe(self, vector_settings):
|
|
remote = self._remote(width=1024)
|
|
vector_schema, _ = self._run(remote, table_dimension=1024)
|
|
remote.embed_query.assert_called_once()
|
|
assert vector_schema.call_args.kwargs["dimension"] == 1024
|
|
|
|
def test_the_probe_restores_the_mismatch_check(self, vector_settings):
|
|
_, raised = self._run(self._remote(width=768), table_dimension=1024)
|
|
assert raised, "a 768-dim model against a vector(1024) table must fail loudly"
|
|
|
|
def test_an_unreachable_server_does_not_block_boot(self, vector_settings):
|
|
vector_schema, raised = self._run(
|
|
self._remote(error=ConnectionError("server down")), table_dimension=1024
|
|
)
|
|
assert not raised
|
|
assert vector_schema.call_args.kwargs["dimension"] == 768
|
|
|
|
def test_a_model_that_knows_its_width_is_not_probed(self, vector_settings):
|
|
local = MagicMock()
|
|
local.dimension = 384
|
|
vector_schema, _ = self._run(local, table_dimension=384)
|
|
local.embed_query.assert_not_called()
|
|
assert vector_schema.call_args.kwargs["dimension"] == 384
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestBootLoadedModelIsReleased:
|
|
"""The width probe must not leave a model resident in a delegating process.
|
|
|
|
Reading ``.dimension`` off an unregistered model means loading it, and
|
|
``EmbeddingsSingleton`` caches what it builds. In an API that delegates
|
|
every embed to the worker that cached copy is never called again — it is
|
|
several hundred megabytes held for the life of the process, which is the
|
|
cost ``EMBEDDINGS_DELEGATE_TO_WORKER`` exists to avoid.
|
|
"""
|
|
|
|
def _run(self, vector_settings, *, delegate, base_url=None):
|
|
from docsgpt.vectorstore.base import EmbeddingsSingleton
|
|
|
|
monkeyed = _embeddings(1024)
|
|
conn = MagicMock()
|
|
conn.cursor.return_value = MagicMock()
|
|
EmbeddingsSingleton._instances.pop("test-model", None)
|
|
|
|
def _build(*_args, **_kwargs):
|
|
EmbeddingsSingleton._instances["test-model"] = monkeyed
|
|
return monkeyed
|
|
|
|
with patch.object(
|
|
vector_settings, "EMBEDDINGS_DELEGATE_TO_WORKER", delegate
|
|
), patch.object(
|
|
vector_settings, "EMBEDDINGS_BASE_URL", base_url
|
|
), patch("psycopg.connect", return_value=conn), patch(
|
|
"docsgpt.vectorstore.model_registry.dimension_for", return_value=None
|
|
), patch(
|
|
"docsgpt.vectorstore.base.build_local_embeddings", side_effect=_build
|
|
), patch(
|
|
"docsgpt.vectorstore.pgvector.PGVectorStore.create_schema"
|
|
) as vector_schema, patch(
|
|
"docsgpt.vectorstore.pgvector.PGVectorStore.table_dimension",
|
|
return_value=1024,
|
|
):
|
|
ensure_vector_schema()
|
|
try:
|
|
return vector_schema, "test-model" in EmbeddingsSingleton._instances
|
|
finally:
|
|
EmbeddingsSingleton._instances.pop("test-model", None)
|
|
|
|
def test_a_delegating_process_does_not_retain_it(self, vector_settings):
|
|
vector_schema, retained = self._run(vector_settings, delegate=True)
|
|
assert not retained, "a delegating API must not hold the model it probed"
|
|
assert vector_schema.call_args.kwargs["dimension"] == 1024
|
|
|
|
def test_a_process_that_embeds_locally_keeps_it(self, vector_settings):
|
|
_, retained = self._run(vector_settings, delegate=False)
|
|
assert retained, "without delegation the model is used, so evicting it "\
|
|
"would only force a rebuild on the first query"
|
|
|
|
def test_a_remote_client_is_kept(self, vector_settings):
|
|
_, retained = self._run(
|
|
vector_settings, delegate=True, base_url="http://embeddings:8080"
|
|
)
|
|
assert retained, "a RemoteEmbeddings holds no model and is what the "\
|
|
"process goes on to use"
|