feat(rag): index markdown KB — chunker, embed client, delta importer, Sources page
Phase 02 (story: import documents):
- fence-aware markdown chunker (heading sections, 200-char overlap,
heading anchor on every chunk, 1200-char hard cap, fence blocks
kept atomic and split under the cap)
- LLMClient over aipi (LiteLLM) reusing the openai client's httpx
transport to send a clean {model, input} payload — the openai SDK
injects encoding_format, which aipi's openai_like group rejects;
token-budget batching + halving retry for the endpoint's
~1024-token per-request input cap
- two-phase per-file upsert importer: sha256 delta (unchanged skip),
atomic commit, A9 exclusion walk, per-source prune, per-file error
tolerance (rollback + log + continue, non-zero CLI exit), adaptive
re-chunk at half target for URL-dense files the endpoint rejects
- scripts/import_docs CLI (repeatable --source, --prune, --limit,
defaults ~/Homelab + ~/Deployments)
- GET /api/docs with per-doc chunk counts; Sources page wired to the
real endpoint (stat cards, full-width a11y table, designed empty
state, DOM-built rows — no innerHTML)
- tests: 63 passed (chunker/llm/importer units, docs API + importer
integration), story E2E 3/3 (real endpoints, in-thread import);
app/ coverage 98%
- real KB imported: 672 docs / 8969 chunks in ~3m, idempotent
re-run (672 unchanged, 0 batches)
- harness: .agent/validate.sh now gates through uv (pytest +
coverage >90% + ruff + pyright) instead of system python3
This commit is contained in:
@@ -0,0 +1,159 @@
|
||||
"""Unit tests: markdown-aware chunker (PLAN §5 policy)."""
|
||||
from __future__ import annotations
|
||||
|
||||
from itertools import pairwise
|
||||
|
||||
import pytest
|
||||
|
||||
from app.rag.chunker import HARD_MAX_CHARS, chunk_markdown, extract_title
|
||||
|
||||
ANCHOR = "## Big"
|
||||
ANCHOR_PREFIX = f"{ANCHOR}\n\n"
|
||||
|
||||
|
||||
def _paras(n: int, char: str = "l", width: int = 300) -> list[str]:
|
||||
return [f"paragraph {i} " + char * (width - 12) for i in range(n)]
|
||||
|
||||
|
||||
def test_short_document_is_single_chunk() -> None:
|
||||
doc = "# Title\n\nJust some intro, no section headings at all."
|
||||
chunks = chunk_markdown(doc)
|
||||
assert chunks == [doc.strip()]
|
||||
|
||||
|
||||
def test_empty_and_whitespace_only_content() -> None:
|
||||
assert chunk_markdown("") == []
|
||||
assert chunk_markdown(" \n\n \n") == []
|
||||
|
||||
|
||||
def test_invalid_params_raise() -> None:
|
||||
with pytest.raises(ValueError):
|
||||
chunk_markdown("# x", target_chars=0)
|
||||
with pytest.raises(ValueError):
|
||||
chunk_markdown("# x", overlap_chars=-1)
|
||||
|
||||
|
||||
def test_splits_on_headings_and_keeps_nearest_heading() -> None:
|
||||
doc = (
|
||||
"# Title\n"
|
||||
"intro line\n"
|
||||
"## Alpha\n"
|
||||
"alpha body\n"
|
||||
"### Beta\n"
|
||||
"beta body\n"
|
||||
"## Gamma\n"
|
||||
"gamma body\n"
|
||||
)
|
||||
chunks = chunk_markdown(doc)
|
||||
assert chunks[0] == "# Title\nintro line"
|
||||
assert chunks[1] == "## Alpha\nalpha body"
|
||||
assert chunks[2] == "### Beta\nbeta body"
|
||||
assert chunks[3] == "## Gamma\ngamma body"
|
||||
|
||||
|
||||
def test_document_without_h1_starts_at_first_section() -> None:
|
||||
chunks = chunk_markdown("## Only\n\nbody")
|
||||
assert chunks == ["## Only\n\nbody"]
|
||||
|
||||
|
||||
def test_long_section_splits_with_overlap_and_anchor_on_every_chunk() -> None:
|
||||
body = "\n\n".join(_paras(10, width=138))
|
||||
doc = f"{ANCHOR}\n\n{body}"
|
||||
chunks = chunk_markdown(doc, target_chars=800, overlap_chars=100)
|
||||
|
||||
assert len(chunks) == 3
|
||||
# Every chunk keeps its nearest preceding heading (the section anchor).
|
||||
assert all(c.startswith(ANCHOR) for c in chunks)
|
||||
# All chunks respect the target budget (anchor + packed body).
|
||||
assert all(len(c) <= 800 for c in chunks)
|
||||
# Overlap: the tail of each chunk is at the start of the next one.
|
||||
for prev, nxt in pairwise(chunks):
|
||||
assert nxt[len(ANCHOR_PREFIX) :].startswith(prev[-100:])
|
||||
|
||||
|
||||
def test_overlap_zero_disables_tail_carryover() -> None:
|
||||
body = "\n\n".join(_paras(8, width=200))
|
||||
chunks = chunk_markdown(f"{ANCHOR}\n\n{body}", target_chars=800, overlap_chars=0)
|
||||
assert len(chunks) >= 2
|
||||
for prev, nxt in pairwise(chunks):
|
||||
assert not nxt[len(ANCHOR_PREFIX) :].startswith(prev[-50:])
|
||||
|
||||
|
||||
def test_code_fences_stay_intact() -> None:
|
||||
doc = (
|
||||
"## Section\n"
|
||||
"before fence\n"
|
||||
"```\n"
|
||||
"## fake heading inside fence\n"
|
||||
"\n"
|
||||
"still in fence\n"
|
||||
"```\n"
|
||||
"after fence\n"
|
||||
"## Other\n"
|
||||
"other body\n"
|
||||
)
|
||||
chunks = chunk_markdown(doc)
|
||||
assert any(c.startswith("## Other") for c in chunks)
|
||||
# The fake heading inside the fence never opens a section…
|
||||
assert not any(c.startswith("## fake heading") for c in chunks)
|
||||
# …and the fence itself is whole in the chunk that contains it.
|
||||
fenced = [c for c in chunks if "still in fence" in c]
|
||||
assert len(fenced) == 1
|
||||
assert "## fake heading inside fence" in fenced[0]
|
||||
assert fenced[0].count("```") == 2
|
||||
# Blank lines inside the fence did not create extra paragraph chunks.
|
||||
assert not any(c.startswith("before fence\n\n") for c in chunks)
|
||||
|
||||
|
||||
def test_fence_block_is_atomic_across_forced_split() -> None:
|
||||
fence = "```\n" + "\n".join(f"code line {i}" for i in range(60)) + "\n```"
|
||||
doc = (
|
||||
f"{ANCHOR}\n\npara A "
|
||||
+ "a" * 300
|
||||
+ f"\n\n{fence}\n\npara B "
|
||||
+ "b" * 300
|
||||
+ "\n\npara C "
|
||||
+ "c" * 300
|
||||
)
|
||||
chunks = chunk_markdown(doc, target_chars=1000, overlap_chars=100)
|
||||
assert len(chunks) >= 2
|
||||
# The whole fence (first and last code line) lives in one chunk — a
|
||||
# chunk boundary never falls inside a code block.
|
||||
assert any("code line 0" in c and "code line 59" in c for c in chunks)
|
||||
|
||||
|
||||
def test_oversized_fence_block_is_split_to_stay_under_hard_cap() -> None:
|
||||
"""aipi's embedding endpoint caps requests at ~1024 input tokens — a
|
||||
multi-KB fenced code block must not survive chunking as one piece."""
|
||||
code = "\n".join(f"int value_{i:03d} = {i}; // padding to grow the line" for i in range(160))
|
||||
doc = (
|
||||
"# Big Doc\n\n"
|
||||
"## Usage Example\n\n"
|
||||
f"```cpp\n{code}\n```\n\n"
|
||||
"## After\n\nDone.\n"
|
||||
)
|
||||
chunks = chunk_markdown(doc)
|
||||
assert len(chunks) >= 3
|
||||
# No chunk exceeds the hard cap (heading anchor adds a little).
|
||||
assert all(len(c) <= HARD_MAX_CHARS + 60 for c in chunks)
|
||||
# Content survives the split, and later sections are untouched.
|
||||
joined = "\n".join(chunks)
|
||||
assert "value_000" in joined
|
||||
assert "value_159" in joined
|
||||
assert any(c.startswith("## After") for c in chunks)
|
||||
|
||||
|
||||
def test_unclosed_fence_does_not_break_sections() -> None:
|
||||
doc = "## A\n\n```\nunterminated fence\n\n## B\n\nbody\n"
|
||||
chunks = chunk_markdown(doc)
|
||||
# "## B" is inside the unterminated fence → not a real heading.
|
||||
assert len(chunks) == 1
|
||||
assert "## B" in chunks[0]
|
||||
|
||||
|
||||
def test_extract_title_prefers_h1() -> None:
|
||||
assert extract_title("# My Title\n\nbody") == "My Title"
|
||||
assert extract_title(" # Indented H1\nbody") == "" # ATX must be at col 0
|
||||
assert extract_title("## not a title\n\nbody") == ""
|
||||
assert extract_title("## sub only", fallback="stem") == "stem"
|
||||
assert extract_title("", fallback="fallback") == "fallback"
|
||||
@@ -0,0 +1,279 @@
|
||||
"""Unit tests: importer directory walk + sha256 delta logic.
|
||||
|
||||
The walk tests are pure filesystem (``tmp_path``); the delta tests run
|
||||
against the local compose Postgres (preferred — a real vector table),
|
||||
skipping with clear instructions when the stack is not up.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
from sqlalchemy import func, select
|
||||
|
||||
from app.models import Chunk, Document
|
||||
from app.rag.importer import (
|
||||
EXCLUDED_DIRS,
|
||||
import_sources,
|
||||
iter_markdown_files,
|
||||
)
|
||||
from app.rag.llm import EmbeddingError
|
||||
from tests.fakes import FakeEmbedder
|
||||
|
||||
|
||||
class _PoisonEmbedder(FakeEmbedder):
|
||||
"""Fails (like a real endpoint) on any text containing 'poison'."""
|
||||
|
||||
async def embed(self, texts: list[str]) -> list[list[float]]:
|
||||
if any("poison" in t for t in texts):
|
||||
raise EmbeddingError("embeddings endpoint refused the input (simulated)")
|
||||
return await super().embed(texts)
|
||||
|
||||
|
||||
class _CapEmbedder(FakeEmbedder):
|
||||
"""Simulates the endpoint's ~1024-token input cap at ~1.1 chars/token:
|
||||
any single text over 1000 chars is rejected (URL-dense worst case)."""
|
||||
|
||||
async def embed(self, texts: list[str]) -> list[list[float]]:
|
||||
if any(len(t) > 1000 for t in texts):
|
||||
raise EmbeddingError(
|
||||
"a single 1100-char chunk exceeded the endpoint's per-request "
|
||||
"input token cap — lower BOR_CHUNK_TARGET_CHARS and re-import"
|
||||
)
|
||||
return await super().embed(texts)
|
||||
|
||||
|
||||
def _cleanup_source(db, source: str) -> None:
|
||||
for doc in db.scalars(select(Document).where(Document.source == source)).all():
|
||||
db.delete(doc)
|
||||
db.commit()
|
||||
|
||||
|
||||
def test_iter_markdown_files_excludes_noncontent_dirs(tmp_path: Path) -> None:
|
||||
root = tmp_path / "proj"
|
||||
for d in (
|
||||
"notes/sub",
|
||||
".venv/lib",
|
||||
"node_modules/x",
|
||||
".git",
|
||||
"__pycache__",
|
||||
".pytest_cache",
|
||||
"dist",
|
||||
"build",
|
||||
):
|
||||
(root / d).mkdir(parents=True)
|
||||
files = {
|
||||
"README.md": "readme",
|
||||
"notes/sub/deep.md": "deep",
|
||||
".venv/lib/junk.md": "junk",
|
||||
"node_modules/x/j.md": "j",
|
||||
".git/c.md": "g",
|
||||
"__pycache__/c.md": "p",
|
||||
".pytest_cache/c.md": "pc",
|
||||
"dist/d.md": "d",
|
||||
"build/b.md": "b",
|
||||
}
|
||||
for rel, text in files.items():
|
||||
(root / rel).write_text(text)
|
||||
(root / "notes" / "not-md.txt").write_text("skip me")
|
||||
|
||||
found = {p.relative_to(root).as_posix() for p in iter_markdown_files(root)}
|
||||
assert found == {"README.md", "notes/sub/deep.md"}
|
||||
|
||||
|
||||
def test_iter_markdown_files_missing_dir_yields_nothing(tmp_path: Path) -> None:
|
||||
assert iter_markdown_files(tmp_path / "definitely-missing") == []
|
||||
|
||||
|
||||
def test_excluded_dirs_match_plan_anchor_a9() -> None:
|
||||
assert {
|
||||
".venv", "node_modules", ".git", "__pycache__", ".pytest_cache", "dist", "build"
|
||||
} == EXCLUDED_DIRS
|
||||
|
||||
|
||||
def test_added_then_unchanged_then_updated_then_pruned(db, tmp_path: Path) -> None:
|
||||
root = tmp_path / "src"
|
||||
root.mkdir()
|
||||
(root / "a.md").write_text("# A\n\nalpha\n\n## Sub\n\nmore alpha\n")
|
||||
(root / "b.md").write_text("# B\n\nbeta\n")
|
||||
llm = FakeEmbedder()
|
||||
try:
|
||||
s1 = asyncio.run(import_sources([root], llm, session=db))
|
||||
assert (s1.files, s1.added, s1.unchanged, s1.updated, s1.pruned) == (2, 2, 0, 0, 0)
|
||||
# a.md has two sections (2 chunks), b.md one (1 chunk).
|
||||
assert s1.chunks == 3
|
||||
# Embeddings are stored with the configured dimension.
|
||||
n = db.scalar(
|
||||
select(func.count())
|
||||
.select_from(Chunk)
|
||||
.join(Document, Document.id == Chunk.document_id)
|
||||
.where(Document.source == root.name)
|
||||
)
|
||||
assert n == 3
|
||||
for c in db.scalars(
|
||||
select(Chunk)
|
||||
.join(Document, Document.id == Chunk.document_id)
|
||||
.where(Document.source == root.name)
|
||||
).all():
|
||||
assert c.embedding is not None and len(c.embedding) == 768
|
||||
|
||||
s2 = asyncio.run(import_sources([root], llm, session=db))
|
||||
assert s2.added == 0 and s2.unchanged == 2
|
||||
|
||||
(root / "a.md").write_text("# A\n\nalpha CHANGED\n")
|
||||
s3 = asyncio.run(import_sources([root], llm, session=db))
|
||||
assert s3.updated == 1 and s3.unchanged == 1
|
||||
doc = db.scalar(
|
||||
select(Document).where(Document.source == root.name, Document.path == "a.md")
|
||||
)
|
||||
assert doc is not None and "CHANGED" in doc.content
|
||||
|
||||
(root / "a.md").unlink()
|
||||
s4 = asyncio.run(import_sources([root], llm, session=db, prune=True))
|
||||
assert s4.pruned == 1
|
||||
assert db.scalar(
|
||||
select(Document).where(Document.source == root.name, Document.path == "a.md")
|
||||
) is None
|
||||
# Chunks of the pruned document are gone (FK cascade).
|
||||
n_after = db.scalar(
|
||||
select(func.count())
|
||||
.select_from(Chunk)
|
||||
.join(Document, Document.id == Chunk.document_id)
|
||||
.where(Document.source == root.name)
|
||||
)
|
||||
assert n_after == 1
|
||||
finally:
|
||||
_cleanup_source(db, root.name)
|
||||
|
||||
|
||||
def test_embedding_failure_is_logged_and_import_continues(db, tmp_path: Path) -> None:
|
||||
"""A file the embedding endpoint refuses must not abort the whole KB:
|
||||
its rows are rolled back, the error is counted, and other files import."""
|
||||
root = tmp_path / "mixed"
|
||||
root.mkdir()
|
||||
(root / "bad.md").write_text("# Bad\n\npoison content that the endpoint refuses\n")
|
||||
(root / "good.md").write_text("# Good\n\nperfectly fine content\n")
|
||||
try:
|
||||
summary = asyncio.run(import_sources([root], _PoisonEmbedder(), session=db))
|
||||
assert summary.files == 2
|
||||
assert summary.errors == 1
|
||||
assert summary.added == 1 # only good.md
|
||||
# bad.md left no row and no orphan chunks behind (rolled back).
|
||||
assert db.scalar(
|
||||
select(Document).where(Document.source == root.name, Document.path == "bad.md")
|
||||
) is None
|
||||
assert db.scalar(
|
||||
select(func.count())
|
||||
.select_from(Chunk)
|
||||
.join(Document, Document.id == Chunk.document_id)
|
||||
.where(Document.source == root.name, Document.path == "bad.md")
|
||||
) == 0
|
||||
assert db.scalar(
|
||||
select(Document).where(Document.source == root.name, Document.path == "good.md")
|
||||
) is not None
|
||||
finally:
|
||||
_cleanup_source(db, root.name)
|
||||
|
||||
|
||||
def test_oversized_chunk_triggers_adaptive_rechunk(db, tmp_path: Path) -> None:
|
||||
"""A URL-dense paragraph the endpoint rejects must be re-chunked smaller
|
||||
for that file only — the import still succeeds."""
|
||||
root = tmp_path / "dense"
|
||||
root.mkdir()
|
||||
# One ~1165-char paragraph: under the 1200-char hard cap, over the
|
||||
# simulated token cap. The retry at 600 chars must split it.
|
||||
para = "see https://example.com/" + "a" * 1100
|
||||
(root / "dense.md").write_text(f"# D\n\n{para}\n")
|
||||
try:
|
||||
summary = asyncio.run(import_sources([root], _CapEmbedder(), session=db))
|
||||
assert summary.errors == 0
|
||||
assert summary.added == 1
|
||||
doc = db.scalar(
|
||||
select(Document).where(Document.source == root.name, Document.path == "dense.md")
|
||||
)
|
||||
assert doc is not None
|
||||
assert len(doc.chunks) >= 2 # re-chunked smaller than the hard cap
|
||||
assert all(len(c.content) <= 1000 for c in doc.chunks)
|
||||
assert all(c.embedding is not None for c in doc.chunks)
|
||||
# The content survives the split.
|
||||
assert "".join(c.content for c in doc.chunks).count("a" * 500) >= 1
|
||||
finally:
|
||||
_cleanup_source(db, root.name)
|
||||
|
||||
|
||||
def test_missing_source_dir_is_skipped(db, tmp_path: Path) -> None:
|
||||
llm = FakeEmbedder()
|
||||
summary = asyncio.run(import_sources([tmp_path / "missing"], llm, session=db))
|
||||
assert summary.files == 0 and summary.added == 0
|
||||
|
||||
|
||||
def test_limit_caps_files_and_disables_prune(db, tmp_path: Path) -> None:
|
||||
root = tmp_path / "limited"
|
||||
root.mkdir()
|
||||
for name in ("a.md", "b.md", "c.md"):
|
||||
(root / name).write_text(f"# {name}\n\nbody {name}\n")
|
||||
llm = FakeEmbedder()
|
||||
try:
|
||||
summary = asyncio.run(import_sources([root], llm, limit=2, session=db, prune=True))
|
||||
assert summary.files == 2 and summary.added == 2
|
||||
# c.md was never walked, so it must NOT be pruned (prune disabled
|
||||
# under --limit) — and nothing else disappears either.
|
||||
assert summary.pruned == 0
|
||||
assert db.scalar(
|
||||
select(func.count()).select_from(Document).where(Document.source == root.name)
|
||||
) == 2
|
||||
finally:
|
||||
_cleanup_source(db, root.name)
|
||||
|
||||
|
||||
def test_limit_must_be_positive(db, tmp_path: Path) -> None:
|
||||
with pytest.raises(ValueError):
|
||||
asyncio.run(import_sources([tmp_path], FakeEmbedder(), limit=0, session=db))
|
||||
|
||||
|
||||
def test_prune_is_scoped_to_the_given_sources(db, tmp_path: Path) -> None:
|
||||
src_x = tmp_path / "SourceX"
|
||||
src_y = tmp_path / "SourceY"
|
||||
src_x.mkdir()
|
||||
src_y.mkdir()
|
||||
(src_x / "x.md").write_text("# X\n\nx body\n")
|
||||
(src_y / "y.md").write_text("# Y\n\ny body\n")
|
||||
llm = FakeEmbedder()
|
||||
try:
|
||||
asyncio.run(import_sources([src_x, src_y], llm, session=db))
|
||||
# Re-import ONLY source Y (y.md removed) with prune: source X's doc
|
||||
# must survive — prune never touches sources not passed to this run.
|
||||
(src_y / "y.md").unlink()
|
||||
summary = asyncio.run(import_sources([src_y], llm, session=db, prune=True))
|
||||
assert summary.pruned == 1
|
||||
assert db.scalar(
|
||||
select(Document).where(Document.source == "SourceX", Document.path == "x.md")
|
||||
) is not None
|
||||
finally:
|
||||
_cleanup_source(db, "SourceX")
|
||||
_cleanup_source(db, "SourceY")
|
||||
|
||||
|
||||
def test_chunk_positions_and_titles(db, tmp_path: Path) -> None:
|
||||
root = tmp_path / "titled"
|
||||
root.mkdir()
|
||||
(root / "multi.md").write_text("# Real Title\n\n## One\n\na\n\n## Two\n\nb\n")
|
||||
(root / "noh1.md").write_text("## Only heading\n\nbody\n")
|
||||
llm = FakeEmbedder()
|
||||
try:
|
||||
asyncio.run(import_sources([root], llm, session=db))
|
||||
titles = {
|
||||
d.path: d.title
|
||||
for d in db.scalars(select(Document).where(Document.source == root.name)).all()
|
||||
}
|
||||
assert titles["multi.md"] == "Real Title" # H1 wins
|
||||
assert titles["noh1.md"] == "noh1" # …else the file stem
|
||||
doc = db.scalar(
|
||||
select(Document).where(Document.source == root.name, Document.path == "multi.md")
|
||||
)
|
||||
assert doc is not None
|
||||
positions = sorted(c.position for c in doc.chunks)
|
||||
assert positions == list(range(len(doc.chunks))) and len(doc.chunks) >= 2
|
||||
finally:
|
||||
_cleanup_source(db, root.name)
|
||||
@@ -0,0 +1,232 @@
|
||||
"""Unit tests: LLMClient embeddings (batching, order, loud dim failure).
|
||||
|
||||
The fakes stand in at the httpx-transport layer — that is where LLMClient
|
||||
actually talks to the endpoint (see ``LLMClient._embed_batch`` in
|
||||
``app/rag/llm.py`` for why the openai SDK's own ``embeddings.create`` is
|
||||
bypassed: it injects ``encoding_format``, which aipi's litellm proxy
|
||||
rejects).
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
|
||||
from app.config import Settings
|
||||
from app.rag.llm import EmbeddingDimensionError, EmbeddingError, LLMClient
|
||||
|
||||
|
||||
def _settings(**kwargs: Any) -> Settings:
|
||||
kwargs.setdefault("_env_file", None)
|
||||
return Settings(**kwargs) # pyright: ignore[reportCallIssue] (kwarg exists at runtime)
|
||||
|
||||
|
||||
class _Row:
|
||||
def __init__(self, index: int, embedding: list[float]) -> None:
|
||||
self.index = index
|
||||
self.embedding = embedding
|
||||
|
||||
|
||||
class _Response:
|
||||
def __init__(self, rows: list[_Row]) -> None:
|
||||
self.data = rows
|
||||
|
||||
|
||||
class _FakeEmbeddingsService:
|
||||
"""Simulates the /embeddings endpoint; records calls; can fail."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dim: int = 768,
|
||||
fail: Exception | None = None,
|
||||
drop_index: int = -1,
|
||||
http_error: int | None = None,
|
||||
too_large_min: int | None = None,
|
||||
) -> None:
|
||||
self.dim = dim
|
||||
self.fail = fail
|
||||
self.drop_index = drop_index
|
||||
self.http_error = http_error
|
||||
self.too_large_min = too_large_min
|
||||
self.calls: list[list[str]] = []
|
||||
|
||||
async def create(self, *, model: str, input: list[str]) -> _Response:
|
||||
self.calls.append(list(input))
|
||||
if self.fail is not None:
|
||||
raise self.fail
|
||||
return _Response(
|
||||
[
|
||||
_Row(i, [0.5] * self.dim)
|
||||
for i in range(len(input))
|
||||
if i != self.drop_index
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
class _FakeHttpResponse:
|
||||
def __init__(
|
||||
self, status_code: int, payload: dict[str, Any] | None = None, text: str = ""
|
||||
) -> None:
|
||||
self.status_code = status_code
|
||||
self._payload = payload
|
||||
self.text = text or (json.dumps(payload) if payload is not None else "boom")
|
||||
|
||||
def json(self) -> Any:
|
||||
if self._payload is None:
|
||||
raise ValueError("no json body")
|
||||
return self._payload
|
||||
|
||||
|
||||
#: The endpoint's real error phrasing (litellm) — the client keys off it.
|
||||
_TOO_LARGE_TEXT = 'input (9999 tokens) is too large to process. increase the physical batch size'
|
||||
|
||||
|
||||
class _FakeHttp:
|
||||
"""Stands in for the httpx transport the openai client owns."""
|
||||
|
||||
def __init__(self, service: _FakeEmbeddingsService) -> None:
|
||||
self.service = service
|
||||
self.bodies: list[dict[str, Any]] = []
|
||||
|
||||
async def post(
|
||||
self, url: str, *, json: dict[str, Any], headers: dict[str, str] | None = None
|
||||
) -> _FakeHttpResponse:
|
||||
self.bodies.append(json)
|
||||
assert "Authorization" in (headers or {})
|
||||
if self.service.http_error is not None:
|
||||
return _FakeHttpResponse(self.service.http_error)
|
||||
if (
|
||||
self.service.too_large_min is not None
|
||||
and len(json["input"]) >= self.service.too_large_min
|
||||
):
|
||||
return _FakeHttpResponse(500, None, _TOO_LARGE_TEXT)
|
||||
rows = await self.service.create(model=json["model"], input=json["input"])
|
||||
payload = {"data": [{"index": r.index, "embedding": r.embedding} for r in rows.data]}
|
||||
return _FakeHttpResponse(200, payload)
|
||||
|
||||
|
||||
class _FakeClient:
|
||||
"""Stands in for the openai AsyncOpenAI object (only its transport is used)."""
|
||||
|
||||
def __init__(self, http: _FakeHttp) -> None:
|
||||
self._client = http
|
||||
|
||||
|
||||
def _make_client(service: _FakeEmbeddingsService, **kwargs: Any) -> tuple[LLMClient, _FakeHttp]:
|
||||
kwargs.setdefault("embed_batch_size", 2)
|
||||
llm = LLMClient(_settings(**kwargs))
|
||||
http = _FakeHttp(service)
|
||||
llm._client = _FakeClient(http) # pyright: ignore[reportAttributeAccessIssue]
|
||||
return llm, http
|
||||
|
||||
|
||||
def test_embed_batches_by_batch_size_and_keeps_order() -> None:
|
||||
service = _FakeEmbeddingsService()
|
||||
llm, http = _make_client(service)
|
||||
texts = [f"t{i}" for i in range(5)]
|
||||
vecs = asyncio.run(llm.embed(texts))
|
||||
|
||||
assert [len(c) for c in service.calls] == [2, 2, 1]
|
||||
assert [t for call in service.calls for t in call] == texts
|
||||
assert len(vecs) == 5
|
||||
assert all(len(v) == 768 for v in vecs)
|
||||
assert llm.embed_batches == 3
|
||||
# aipi (litellm) rejects the SDK's injected "encoding_format" — the
|
||||
# payload must stay a minimal {model, input} body.
|
||||
assert all(set(b) == {"model", "input"} for b in http.bodies)
|
||||
|
||||
|
||||
def test_embed_empty_returns_empty_without_calling_endpoint() -> None:
|
||||
service = _FakeEmbeddingsService()
|
||||
llm, http = _make_client(service)
|
||||
assert asyncio.run(llm.embed([])) == []
|
||||
assert http.bodies == []
|
||||
assert service.calls == []
|
||||
assert llm.embed_batches == 0
|
||||
|
||||
|
||||
def test_embed_one_returns_single_vector() -> None:
|
||||
llm, _ = _make_client(_FakeEmbeddingsService())
|
||||
vec = asyncio.run(llm.embed_one("hello"))
|
||||
assert len(vec) == 768
|
||||
|
||||
|
||||
def test_dim_mismatch_fails_loudly_with_actionable_message() -> None:
|
||||
llm, _ = _make_client(_FakeEmbeddingsService(dim=512))
|
||||
with pytest.raises(EmbeddingDimensionError) as exc:
|
||||
asyncio.run(llm.embed(["hello"]))
|
||||
msg = str(exc.value)
|
||||
assert "512" in msg and "768" in msg
|
||||
assert "BOR_EMBEDDING_DIM" in msg
|
||||
assert "llm_probe" in msg
|
||||
|
||||
|
||||
def test_endpoint_error_is_wrapped() -> None:
|
||||
llm, _ = _make_client(_FakeEmbeddingsService(fail=RuntimeError("connection refused")))
|
||||
with pytest.raises(EmbeddingError, match="connection refused"):
|
||||
asyncio.run(llm.embed(["hello"]))
|
||||
assert llm.embed_batches == 0
|
||||
|
||||
|
||||
def test_http_error_surfaces_status() -> None:
|
||||
llm, _ = _make_client(_FakeEmbeddingsService(http_error=502))
|
||||
with pytest.raises(EmbeddingError, match="HTTP 502"):
|
||||
asyncio.run(llm.embed(["hello"]))
|
||||
assert llm.embed_batches == 0
|
||||
|
||||
|
||||
def test_missing_vector_row_is_rejected() -> None:
|
||||
llm, _ = _make_client(_FakeEmbeddingsService(drop_index=1))
|
||||
with pytest.raises(EmbeddingError, match="returned 1 vectors for 2 inputs"):
|
||||
asyncio.run(llm.embed(["a", "b"]))
|
||||
|
||||
|
||||
def test_batch_size_one_forces_one_call_per_text() -> None:
|
||||
service = _FakeEmbeddingsService()
|
||||
llm, _ = _make_client(service, embed_batch_size=1)
|
||||
asyncio.run(llm.embed(["a", "b", "c"]))
|
||||
assert [len(c) for c in service.calls] == [1, 1, 1]
|
||||
|
||||
|
||||
def test_token_budget_limits_texts_per_request() -> None:
|
||||
"""~2000-char chunks must not stack up past aipi's ~1024-token cap."""
|
||||
service = _FakeEmbeddingsService()
|
||||
llm, _ = _make_client(service, embed_batch_size=16) # high count cap
|
||||
texts = ["x" * 2000 for _ in range(4)]
|
||||
vecs = asyncio.run(llm.embed(texts))
|
||||
# 2000 + 2000 chars > 3600-char (≈900-token) budget ⇒ one chunk per request
|
||||
assert [len(c) for c in service.calls] == [1, 1, 1, 1]
|
||||
assert len(vecs) == 4
|
||||
|
||||
|
||||
def test_small_chunks_pack_up_to_count_cap() -> None:
|
||||
service = _FakeEmbeddingsService()
|
||||
llm, _ = _make_client(service, embed_batch_size=4) # count cap binds
|
||||
texts = ["short text" for _ in range(9)]
|
||||
vecs = asyncio.run(llm.embed(texts))
|
||||
assert [len(c) for c in service.calls] == [4, 4, 1]
|
||||
assert len(vecs) == 9
|
||||
|
||||
|
||||
def test_too_large_response_halves_batch_until_it_fits() -> None:
|
||||
"""The tokenizer estimate can be wrong for dense content — the client
|
||||
must halve an over-large request and preserve order."""
|
||||
service = _FakeEmbeddingsService(too_large_min=3)
|
||||
llm, http = _make_client(service, embed_batch_size=16) # all 8 fit one request
|
||||
texts = [f"t{i}" for i in range(8)]
|
||||
vecs = asyncio.run(llm.embed(texts))
|
||||
# http.bodies sees every request (including the rejected ones); the
|
||||
# left-half recursion completes before the right half starts.
|
||||
assert [len(b["input"]) for b in http.bodies] == [8, 4, 2, 2, 4, 2, 2]
|
||||
assert len(vecs) == 8
|
||||
assert all(len(v) == 768 for v in vecs)
|
||||
|
||||
|
||||
def test_single_oversized_text_fails_actionably() -> None:
|
||||
service = _FakeEmbeddingsService(too_large_min=1)
|
||||
llm, _ = _make_client(service)
|
||||
with pytest.raises(EmbeddingError, match="token cap"):
|
||||
asyncio.run(llm.embed(["x" * 3000]))
|
||||
assert llm.embed_batches == 0
|
||||
Reference in New Issue
Block a user