mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-13 12:14:00 +00:00
fix: embeddings registry edge cases and chunk-size accounting
Follow-up to the embeddings work, from a review pass over the branch. - Route the OpenAI/Azure key handling through the model registry instead of matching the canonical name literally, so the `text-embedding-ada-002` alias the registry now accepts also reaches the Azure deployment name rather than failing every embed with DeploymentNotFound. - Fall back to a default width where the embeddings model reports no dimension. A model outside the registry returns None rather than no attribute, so `getattr` with a default did not catch it and the width reached the DDL as `vector(None)` / `list_size=None`. - Point HF_HUB_CACHE at the prefetch directory. Chunking loads the tokenizer through `tokenizers`, which reads the hub cache, so a fresh container fetched over the network on first ingest and an offline one silently fell back to cl100k. - Charge a token that collapses a long unbroken run by its character span. WordPiece emits one [UNK] for any word over its character limit, which made base64 and minified content count as near-zero tokens, so nothing split it and oversized chunks reached the embedding server. - Preserve chunk ids and honour --batch-size when rebuilding a FAISS index. Fresh uuids orphaned GraphRAG's graph_node_chunks rows, and the whole index went out in a single embed call on remote servers. - Document that granite runs an int8-quantised graph, and scope the SentenceTransformer parity claim to mpnet's fp32 graph, which is where it was measured. - Correct the embeddings docs: a matching dimension is not a matching model, so a same-width swap raises nothing and silently degrades retrieval.
This commit is contained in:
1 parent
869d87452b
commit
25e07f5cee
15 files changed
+296
-20
No files matched your search
@@ -70,7 +70,12 @@ COPY . /app/application
|
||||
# does not download on first ingest and an air-gapped install works at all.
|
||||
# Both defaults are baked: an upgraded deployment keeps using mpnet until it
|
||||
# runs the re-embed script, while a new one starts on granite.
|
||||
ENV EMBEDDINGS_CACHE_DIR=/app/models
|
||||
# The prefetch writes hub-layout snapshots (including tokenizer.json) here, so
|
||||
# HF_HUB_CACHE has to point at the same directory: chunking loads the tokenizer
|
||||
# through ``tokenizers``, which reads the hub cache and would otherwise fetch
|
||||
# over the network on first ingest -- and fall back to cl100k when offline.
|
||||
ENV EMBEDDINGS_CACHE_DIR=/app/models \
|
||||
HF_HUB_CACHE=/app/models
|
||||
ARG EMBEDDINGS_PREFETCH=""
|
||||
RUN PYTHONPATH=/app /venv/bin/python -m application.scripts.prefetch_models ${EMBEDDINGS_PREFETCH}
|
||||
|
||||
|
||||
@@ -157,12 +157,14 @@ class GraphStore:
|
||||
"""Dimension of the configured embeddings model, matching ``PGVectorStore``.
|
||||
|
||||
Falls back to ``DEFAULT_NAME_EMBEDDING_DIM`` so the graph table and the
|
||||
pgvector ``documents`` table always agree on the configured model.
|
||||
pgvector ``documents`` table always agree on the configured model. A
|
||||
model outside the registry reports ``None`` rather than no attribute,
|
||||
so the fallback cannot be left to ``getattr``.
|
||||
"""
|
||||
from application.vectorstore.base import get_embeddings
|
||||
|
||||
embedding = get_embeddings()
|
||||
return getattr(embedding, "dimension", DEFAULT_NAME_EMBEDDING_DIM)
|
||||
return getattr(embedding, "dimension", None) or DEFAULT_NAME_EMBEDDING_DIM
|
||||
|
||||
@staticmethod
|
||||
def create_schema(conn, *, dimension: int = DEFAULT_NAME_EMBEDDING_DIM) -> None:
|
||||
|
||||
@@ -49,6 +49,37 @@ def _windows(total: int, first: int, rest: int) -> Iterator[Tuple[int, int]]:
|
||||
budget = max(1, rest)
|
||||
|
||||
|
||||
# A token spanning more characters than this collapsed a run the tokenizer
|
||||
# could not break up: WordPiece emits a single ``[UNK]`` for any word longer
|
||||
# than ``max_input_chars_per_word``. Charging that once makes a base64 blob or
|
||||
# a minified bundle look tiny, so nothing splits it and an oversized chunk
|
||||
# reaches the embedding server. Real tokens are a few characters, so prose
|
||||
# never reaches this bound.
|
||||
_MAX_CHARS_PER_TOKEN = 16
|
||||
|
||||
|
||||
def _token_weight(start: int, end: int) -> int:
|
||||
"""Tokens a span costs, charging a collapsed run by its length."""
|
||||
span = end - start
|
||||
return max(1, -(-span // _MAX_CHARS_PER_TOKEN))
|
||||
|
||||
|
||||
def _cap_piece_chars(pieces: List[str], first: int, rest: int) -> List[str]:
|
||||
"""Cut any piece holding more characters than its budget can cover."""
|
||||
capped: List[str] = []
|
||||
budget = first
|
||||
for piece in pieces:
|
||||
limit = max(1, budget * _MAX_CHARS_PER_TOKEN)
|
||||
while len(piece) > limit:
|
||||
capped.append(piece[:limit])
|
||||
piece = piece[limit:]
|
||||
budget = rest
|
||||
limit = max(1, budget * _MAX_CHARS_PER_TOKEN)
|
||||
capped.append(piece)
|
||||
budget = rest
|
||||
return capped
|
||||
|
||||
|
||||
class TokenCounter:
|
||||
"""Counts and splits text in one tokenizer's units.
|
||||
|
||||
@@ -123,7 +154,10 @@ class HuggingFaceCounter(TokenCounter):
|
||||
def count(self, text: str) -> int:
|
||||
if not text:
|
||||
return 0
|
||||
return len(self._encode(text).ids)
|
||||
encoding = self._encode(text)
|
||||
if not encoding.offsets:
|
||||
return len(encoding.ids)
|
||||
return sum(_token_weight(start, end) for start, end in encoding.offsets)
|
||||
|
||||
def split(
|
||||
self, text: str, max_tokens: int, first_max_tokens: Optional[int] = None
|
||||
@@ -133,7 +167,7 @@ class HuggingFaceCounter(TokenCounter):
|
||||
rest = max(1, max_tokens)
|
||||
first = max(1, first_max_tokens if first_max_tokens is not None else rest)
|
||||
offsets = self._encode(text).offsets
|
||||
if len(offsets) <= first:
|
||||
if self.count(text) <= first:
|
||||
return [text]
|
||||
|
||||
pieces: List[str] = []
|
||||
@@ -155,7 +189,7 @@ class HuggingFaceCounter(TokenCounter):
|
||||
pieces[-1] = pieces[-1] + text[cursor:]
|
||||
else:
|
||||
pieces.append(text[cursor:])
|
||||
return pieces
|
||||
return _cap_piece_chars(pieces, first, rest)
|
||||
|
||||
|
||||
def _load_hf_counter(repo: str) -> Optional[HuggingFaceCounter]:
|
||||
|
||||
@@ -216,11 +216,16 @@ def reembed_faiss(source_id: str, batch_size: int, dry_run: bool) -> Tuple[int,
|
||||
# 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)
|
||||
|
||||
@@ -285,7 +285,11 @@ def get_embeddings(
|
||||
)
|
||||
return EmbeddingsSingleton._remote_instance(embeddings_name, embeddings_key)
|
||||
|
||||
if embeddings_name == "openai_text-embedding-ada-002":
|
||||
# Match through the registry, not on the canonical string: the bare
|
||||
# ``text-embedding-ada-002`` alias resolves here too, and skipping this
|
||||
# branch would drop the Azure deployment name and the API key.
|
||||
spec = resolve(embeddings_name)
|
||||
if spec is not None and spec.provider == "openai":
|
||||
if _azure_configured():
|
||||
embedding_instance = EmbeddingsSingleton.get_instance(
|
||||
embeddings_name, model=settings.AZURE_EMBEDDINGS_DEPLOYMENT_NAME
|
||||
|
||||
@@ -6,9 +6,11 @@ removes torch, transformers and sentence-transformers from the dependency set
|
||||
and measurably reduces both resident memory and import cost.
|
||||
|
||||
The swap is numerically transparent for existing indexes: embedding the same
|
||||
text with SentenceTransformer and with FastEmbed's ONNX graph of the same
|
||||
model returns vectors at cosine 1.0, so an mpnet index built before this
|
||||
change keeps working unchanged.
|
||||
text with SentenceTransformer and with FastEmbed's fp32 ONNX graph of mpnet
|
||||
returns vectors at cosine 1.0, so an mpnet index built before this change
|
||||
keeps working unchanged. That result is specific to the fp32 graph -- the
|
||||
granite entries run an int8-quantised one and are not bit-comparable to a
|
||||
fp32 index of the same model.
|
||||
|
||||
Models are described in :mod:`application.vectorstore.model_registry`. A name
|
||||
the registry does not know is treated as a Hugging Face repository, which is
|
||||
|
||||
@@ -81,7 +81,14 @@ class FaissStore(BaseVectorStore):
|
||||
# and must not be shown as one.
|
||||
score_kind = "l2_distance"
|
||||
|
||||
def __init__(self, source_id: str, embeddings_key: str, docs_init=None):
|
||||
def __init__(
|
||||
self,
|
||||
source_id: str,
|
||||
embeddings_key: str,
|
||||
docs_init=None,
|
||||
ids=None,
|
||||
batch_size=None,
|
||||
):
|
||||
super().__init__()
|
||||
self.source_id = source_id
|
||||
self.path = get_vectorstore(source_id)
|
||||
@@ -94,7 +101,7 @@ class FaissStore(BaseVectorStore):
|
||||
|
||||
try:
|
||||
if docs_init:
|
||||
self._build_from_documents(docs_init)
|
||||
self._build_from_documents(docs_init, ids=ids, batch_size=batch_size)
|
||||
else:
|
||||
self._load_from_storage()
|
||||
except Exception as e:
|
||||
@@ -104,17 +111,36 @@ class FaissStore(BaseVectorStore):
|
||||
|
||||
# -- Construction ----------------------------------------------------
|
||||
|
||||
def _build_from_documents(self, docs_init) -> None:
|
||||
"""Create a fresh index seeded with ``docs_init``."""
|
||||
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()
|
||||
vectors = self.embeddings.embed_documents(texts)
|
||||
self.index = faiss.IndexFlatL2(len(vectors[0]))
|
||||
self._append(texts, metadatas, vectors)
|
||||
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."""
|
||||
|
||||
@@ -2,6 +2,7 @@ from typing import List, Optional
|
||||
import importlib
|
||||
from application.vectorstore.base import BaseVectorStore
|
||||
from application.core.settings import settings
|
||||
from application.vectorstore.model_registry import DEFAULT_EMBEDDING_DIMENSION
|
||||
|
||||
class LanceDBVectorStore(BaseVectorStore):
|
||||
"""Class for LanceDB Vector Store integration."""
|
||||
@@ -54,8 +55,11 @@ class LanceDBVectorStore(BaseVectorStore):
|
||||
"""Ensure the table exists before performing operations."""
|
||||
if self.table is None:
|
||||
embeddings = self._get_embeddings(settings.EMBEDDINGS_NAME, self.embeddings_key)
|
||||
# A model outside the registry reports no width until it has run;
|
||||
# ``list_size=None`` is a TypeError, not a permissive schema.
|
||||
dimension = getattr(embeddings, "dimension", None) or DEFAULT_EMBEDDING_DIMENSION
|
||||
schema = self.pa.schema([
|
||||
self.pa.field("vector", self.pa.list_(self.pa.float32(), list_size=embeddings.dimension)),
|
||||
self.pa.field("vector", self.pa.list_(self.pa.float32(), list_size=dimension)),
|
||||
self.pa.field("text", self.pa.string()),
|
||||
self.pa.field("metadata", self.pa.struct([
|
||||
self.pa.field("key", self.pa.string()),
|
||||
|
||||
@@ -70,6 +70,12 @@ MPNET = EmbeddingModel(
|
||||
|
||||
#: Default for new installs: same 768 dimensions as mpnet, an 8x wider
|
||||
#: effective input, and multilingual retrieval.
|
||||
#:
|
||||
#: Runs the int8-quantised graph, which is ~4x smaller than the fp32 one
|
||||
#: (313 MB against 1247 MB) and keeps the image shippable now that both
|
||||
#: defaults are baked in. It costs some numeric fidelity against fp32 --
|
||||
#: measured at cosine 0.96 on short texts, with retrieval rank order
|
||||
#: unchanged -- so it is not bit-comparable to a fp32 granite index.
|
||||
GRANITE_311M = EmbeddingModel(
|
||||
name="ibm-granite/granite-embedding-311m-multilingual-r2",
|
||||
dimension=768,
|
||||
@@ -82,6 +88,7 @@ GRANITE_311M = EmbeddingModel(
|
||||
)
|
||||
|
||||
#: Smaller granite. Half the vector width, roughly three times the speed.
|
||||
#: Int8-quantised on the same terms as GRANITE_311M above.
|
||||
GRANITE_97M = EmbeddingModel(
|
||||
name="ibm-granite/granite-embedding-97m-multilingual-r2",
|
||||
dimension=384,
|
||||
|
||||
@@ -89,7 +89,11 @@ Some remote servers (notably llama.cpp) reject any single input larger than thei
|
||||
EMBEDDINGS_MAX_INPUT_TOKENS=512
|
||||
```
|
||||
|
||||
When set, each input string is truncated to that many tokens and the overflow is dropped (lossy by design). Token counts use DocsGPT's shared tiktoken encoding, which differs from your server's tokenizer, so choose a limit with some headroom below the server's true limit to absorb tokenizer skew. Leave the setting unset (or `0`) to disable truncation.
|
||||
When set, each input string is truncated to that many tokens and the overflow is dropped (lossy by design).
|
||||
|
||||
You usually do not need to set it. When `EMBEDDINGS_NAME` names a model DocsGPT knows, its context window is used automatically, and counting uses that model's own tokenizer, so the limit and the count are in the same unit. Set `EMBEDDINGS_MAX_INPUT_TOKENS` when your server serves a model DocsGPT does not know, or when you deliberately want a limit below the model's own.
|
||||
|
||||
Leaving `EMBEDDINGS_NAME` unset imposes no limit: the name is only forwarded as the `model` field in each request, so a default nobody chose is not taken as a description of your server. When no model tokenizer is available, counting falls back to tiktoken — pick a limit with headroom below the server's true limit to absorb the skew between the two tokenizers.
|
||||
|
||||
## Important: Embedding Dimensions Must Stay Consistent
|
||||
|
||||
@@ -97,6 +101,18 @@ Each embedding model produces vectors of a fixed dimension, and your vector stor
|
||||
|
||||
If you need to switch embedding models, you must re-ingest your sources so the index is rebuilt with the new dimension. This also applies to the [GraphRAG](/Sources/GraphRAG) graph tables, which are sized to the embedding dimension at creation time.
|
||||
|
||||
### A matching dimension is not a matching model
|
||||
|
||||
The dimension check is a guard against a corrupt index, not a guarantee that a swap is safe. Two models of the *same* width — `all-mpnet-base-v2` and `granite-embedding-311m-multilingual-r2` are both 768 — raise nothing at all, and every query is then embedded by a different model than the stored vectors were. Nothing fails; retrieval quality simply degrades.
|
||||
|
||||
Switching between same-width models therefore still requires re-embedding:
|
||||
|
||||
```bash
|
||||
python -m application.scripts.reembed
|
||||
```
|
||||
|
||||
Run it after changing `EMBEDDINGS_NAME` and before serving queries. See [Upgrading](/upgrading) for the granite migration specifically. Changing to a model of a *different* width is not supported by the script — re-ingest those sources instead.
|
||||
|
||||
## Adding Support for Other Embedding Models
|
||||
|
||||
If you wish to use an embedding model that is not supported out-of-the-box, a good starting point for adding custom embedding model support is to examine the `base.py` file located in the `application/vectorstore` directory.
|
||||
|
||||
@@ -514,6 +514,27 @@ class TestEmbeddingDim:
|
||||
store = GraphStore.__new__(GraphStore)
|
||||
assert store._embedding_dim() == 1536
|
||||
|
||||
def test_none_dimension_falls_back_to_default(self, monkeypatch):
|
||||
"""A remote model outside the registry reports ``None``, not nothing.
|
||||
|
||||
``getattr`` with a default cannot catch that -- the attribute exists --
|
||||
so the width reached the DDL as ``vector(None)``.
|
||||
"""
|
||||
from application.vectorstore import base as base_module
|
||||
|
||||
monkeypatch.setattr(base_module.settings, "EMBEDDINGS_BASE_URL", None)
|
||||
fake_embedding = MagicMock()
|
||||
fake_embedding.dimension = None
|
||||
monkeypatch.setattr(
|
||||
base_module.EmbeddingsSingleton,
|
||||
"get_instance",
|
||||
staticmethod(lambda *a, **k: fake_embedding),
|
||||
)
|
||||
monkeypatch.setattr(GraphStore, "_embedding_dim", _REAL_EMBEDDING_DIM)
|
||||
|
||||
store = GraphStore.__new__(GraphStore)
|
||||
assert store._embedding_dim() == store_module.DEFAULT_NAME_EMBEDDING_DIM
|
||||
|
||||
def test_falls_back_to_default_dimension(self, monkeypatch):
|
||||
from application.vectorstore import base as base_module
|
||||
|
||||
|
||||
@@ -202,3 +202,50 @@ class TestOffsetsWithoutSpans:
|
||||
def test_all_span_less_offsets_still_return_the_text(self):
|
||||
counter = self._counter([(0, 0), (0, 0), (0, 0)])
|
||||
assert "".join(counter.split("abc", 1)) == "abc"
|
||||
|
||||
|
||||
class TestUnknownTokenCollapse:
|
||||
"""A tokenizer that folds a long unbroken run into one ``[UNK]``.
|
||||
|
||||
WordPiece gives up on any word longer than ``max_input_chars_per_word``
|
||||
and emits a single unknown token for it. Counting that as one token makes
|
||||
a base64 blob or a minified bundle look tiny, so the chunker never splits
|
||||
it and an oversized chunk reaches the embedding server.
|
||||
"""
|
||||
|
||||
class _CollapsingEncoding:
|
||||
"""One token per whitespace-separated word, however long the word."""
|
||||
|
||||
def __init__(self, text):
|
||||
self.ids = []
|
||||
self.offsets = []
|
||||
cursor = 0
|
||||
for word in text.split(" "):
|
||||
if word:
|
||||
self.ids.append(0)
|
||||
self.offsets.append((cursor, cursor + len(word)))
|
||||
cursor += len(word) + 1
|
||||
|
||||
class _CollapsingTokenizer:
|
||||
def encode(self, text, add_special_tokens=False):
|
||||
return TestUnknownTokenCollapse._CollapsingEncoding(text)
|
||||
|
||||
def _counter(self):
|
||||
return tokenization.HuggingFaceCounter(self._CollapsingTokenizer(), "stub")
|
||||
|
||||
def test_long_unbroken_run_is_charged_by_its_span(self):
|
||||
counter = self._counter()
|
||||
assert counter.count("a" * 9000) > 100
|
||||
|
||||
def test_ordinary_prose_is_unaffected(self):
|
||||
counter = self._counter()
|
||||
text = "the quick brown fox jumps over the lazy dog"
|
||||
assert counter.count(text) == 9
|
||||
|
||||
def test_split_bounds_a_collapsed_run(self):
|
||||
counter = self._counter()
|
||||
text = "a" * 9000
|
||||
pieces = counter.split(text, 20)
|
||||
assert "".join(pieces) == text, "split must not lose or alter text"
|
||||
assert len(pieces) > 1, "a collapsed run must still be cut into pieces"
|
||||
assert all(counter.count(p) <= 20 for p in pieces)
|
||||
@@ -107,6 +107,22 @@ class TestFaissRebuild:
|
||||
keys = [call.kwargs.get("embeddings_key") for call in factory.call_args_list]
|
||||
assert keys == ["sk-real", "sk-real"]
|
||||
|
||||
def test_chunk_ids_are_preserved(self, stores):
|
||||
"""Re-embedding must not renumber chunks.
|
||||
|
||||
Fresh ids orphan every GraphRAG ``graph_node_chunks`` row for the
|
||||
source and invalidate any id a client already holds.
|
||||
"""
|
||||
_, _, factory = stores
|
||||
reembed.reembed_faiss("s1", batch_size=8, dry_run=False)
|
||||
assert factory.call_args_list[1].kwargs["ids"] == ["1", "2"]
|
||||
|
||||
def test_batch_size_is_forwarded_to_the_rebuild(self, stores):
|
||||
"""On a remote embeddings server the whole index is otherwise one POST."""
|
||||
_, _, factory = stores
|
||||
reembed.reembed_faiss("s1", batch_size=8, dry_run=False)
|
||||
assert factory.call_args_list[1].kwargs["batch_size"] == 8
|
||||
|
||||
def test_existing_index_is_not_deleted(self, stores):
|
||||
"""The rebuild must not destroy the old index before the new one exists."""
|
||||
existing, rebuilt, _ = stores
|
||||
|
||||
@@ -591,6 +591,49 @@ class TestGetEmbeddingsResolver:
|
||||
"openai_text-embedding-ada-002", model="embed-deploy"
|
||||
)
|
||||
|
||||
@patch("application.vectorstore.base.settings")
|
||||
@patch("application.vectorstore.base.EmbeddingsSingleton.get_instance")
|
||||
def test_openai_alias_also_reaches_the_azure_deployment(
|
||||
self, mock_get_instance, mock_settings
|
||||
):
|
||||
"""The registry accepts the bare alias, so the key handling must too.
|
||||
|
||||
Matching on the canonical string alone sent the alias down the generic
|
||||
branch, where the deployment name is never passed and Azure answers
|
||||
every embed with DeploymentNotFound.
|
||||
"""
|
||||
mock_settings.EMBEDDINGS_BASE_URL = None
|
||||
mock_settings.OPENAI_API_BASE = "https://azure.openai.com"
|
||||
mock_settings.OPENAI_API_VERSION = "2023-05-15"
|
||||
mock_settings.AZURE_DEPLOYMENT_NAME = "deploy"
|
||||
mock_settings.AZURE_EMBEDDINGS_DEPLOYMENT_NAME = "embed-deploy"
|
||||
mock_settings.EMBEDDINGS_NAME = "text-embedding-ada-002"
|
||||
mock_settings.EMBEDDINGS_KEY = "sk-key"
|
||||
|
||||
get_embeddings()
|
||||
|
||||
mock_get_instance.assert_called_once_with(
|
||||
"text-embedding-ada-002", model="embed-deploy"
|
||||
)
|
||||
|
||||
@patch("application.vectorstore.base.settings")
|
||||
@patch("application.vectorstore.base.EmbeddingsSingleton.get_instance")
|
||||
def test_openai_name_is_matched_case_insensitively(
|
||||
self, mock_get_instance, mock_settings
|
||||
):
|
||||
mock_settings.EMBEDDINGS_BASE_URL = None
|
||||
mock_settings.OPENAI_API_BASE = None
|
||||
mock_settings.OPENAI_API_VERSION = None
|
||||
mock_settings.AZURE_DEPLOYMENT_NAME = None
|
||||
mock_settings.EMBEDDINGS_NAME = "OpenAI_Text-Embedding-Ada-002"
|
||||
mock_settings.EMBEDDINGS_KEY = "sk-from-settings"
|
||||
|
||||
get_embeddings()
|
||||
|
||||
mock_get_instance.assert_called_once_with(
|
||||
"OpenAI_Text-Embedding-Ada-002", openai_api_key="sk-from-settings"
|
||||
)
|
||||
|
||||
@patch("application.vectorstore.base.settings")
|
||||
@patch("application.vectorstore.base.EmbeddingsSingleton.get_instance")
|
||||
def test_explicit_arguments_win_over_settings(
|
||||
|
||||
@@ -5,11 +5,13 @@ than asserting that calls were forwarded to a mock.
|
||||
"""
|
||||
|
||||
import json
|
||||
from unittest.mock import Mock, patch
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import MagicMock, Mock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from application.storage.local import LocalStorage
|
||||
from application.vectorstore.faiss import FaissStore
|
||||
|
||||
|
||||
class _FakeEmbeddings:
|
||||
@@ -243,3 +245,45 @@ class TestGetVectorstore:
|
||||
|
||||
with pytest.raises(ValueError, match="Invalid source_id path"):
|
||||
get_vectorstore(bad)
|
||||
|
||||
|
||||
class TestBuildFromDocumentsBatching:
|
||||
"""The full rebuild path: honour caller ids and embed in batches."""
|
||||
|
||||
def _docs(self, n):
|
||||
return [
|
||||
SimpleNamespace(page_content=f"chunk {i}", metadata={"i": i})
|
||||
for i in range(n)
|
||||
]
|
||||
|
||||
def _store(self, embed):
|
||||
store = FaissStore.__new__(FaissStore)
|
||||
store.embeddings = MagicMock()
|
||||
store.embeddings.embed_documents.side_effect = embed
|
||||
store.documents = {}
|
||||
store.index_to_docstore_id = {}
|
||||
store.index = None
|
||||
return store
|
||||
|
||||
def test_supplied_ids_are_used(self):
|
||||
store = self._store(lambda texts: [[float(len(texts))] * 2 for _ in texts])
|
||||
store._build_from_documents(self._docs(3), ids=["a", "b", "c"])
|
||||
assert list(store.documents) == ["a", "b", "c"]
|
||||
|
||||
def test_embedding_is_split_into_batches(self):
|
||||
sizes = []
|
||||
|
||||
def embed(texts):
|
||||
sizes.append(len(texts))
|
||||
return [[1.0, 2.0] for _ in texts]
|
||||
|
||||
store = self._store(embed)
|
||||
store._build_from_documents(self._docs(5), batch_size=2)
|
||||
assert sizes == [2, 2, 1]
|
||||
assert len(store.index_to_docstore_id) == 5
|
||||
|
||||
def test_defaults_are_unchanged(self):
|
||||
store = self._store(lambda texts: [[1.0, 2.0] for _ in texts])
|
||||
store._build_from_documents(self._docs(3))
|
||||
assert len(store.documents) == 3
|
||||
assert all(len(k) == 36 for k in store.documents), "uuid4 ids by default"
|
||||
Reference in new issue
Block a user