danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Merge pull request #481 from ggozad/chore/clean-up

Browse source 
Clean up reranker/converter/embedder layers and cut test runtime
This commit is contained in:
Yiorgis Gozadinos 2026-06-29 15:26:07 +03:00 committed by GitHub
commit 07098d55e4
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
22 changed files with 32219 additions and 226444 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
CHANGELOG.md
View file

@ -1,6 +1,10 @@
# Changelog
## [Unreleased]
### Changed
- Custom rerankers override `RerankerBase._rerank` instead of `rerank`; the base `rerank` handles the empty-input short-circuit.
## [0.63.1] - 2026-06-29
### Changed

17
haiku_rag_slim/haiku/rag/converters/base.py
View file

@ -7,6 +7,23 @@ from typing import TYPE_CHECKING
if TYPE_CHECKING:
from docling_core.types.doc.document import DoclingDocument
from haiku.rag.config import AppConfig
from haiku.rag.config.models import ModelConfig
def vlm_api_url(config: "AppConfig", model: "ModelConfig") -> str:
"""Construct the VLM chat-completions URL for a picture-description model."""
if model.base_url:
return f"{model.base_url.rstrip('/')}/v1/chat/completions"
if model.provider == "ollama":
return f"{config.providers.ollama.base_url.rstrip('/')}/v1/chat/completions"
if model.provider == "openai":
return "https://api.openai.com/v1/chat/completions"
raise ValueError(f"Unsupported VLM provider: {model.provider}")
class DocumentConverter(ABC):
"""Abstract base class for document converters.

21
haiku_rag_slim/haiku/rag/converters/docling_local.py
View file

@ -5,7 +5,7 @@ from pathlib import Path
from typing import TYPE_CHECKING, ClassVar
from haiku.rag.config import AppConfig
from haiku.rag.converters.base import DocumentConverter
from haiku.rag.converters.base import DocumentConverter, vlm_api_url
from haiku.rag.converters.text_utils import TextFileHandler
if TYPE_CHECKING:
@ -13,7 +13,7 @@ if TYPE_CHECKING:
from docling.document_converter import FormatOption
from docling_core.types.doc.document import DoclingDocument
from haiku.rag.config.models import ConversionOptions, ModelConfig
from haiku.rag.config.models import ConversionOptions
class DoclingLocalConverter(DocumentConverter):
@ -62,21 +62,6 @@ class DoclingLocalConverter(DocumentConverter):
"""Return list of file extensions supported by this converter."""
return self.docling_extensions + TextFileHandler.text_extensions
def _get_vlm_api_url(self, model: "ModelConfig") -> str:
"""Construct VLM API URL from model config."""
if model.base_url:
base = model.base_url.rstrip("/")
return f"{base}/v1/chat/completions"
if model.provider == "ollama":
base = self.config.providers.ollama.base_url.rstrip("/")
return f"{base}/v1/chat/completions"
if model.provider == "openai":
return "https://api.openai.com/v1/chat/completions"
raise ValueError(f"Unsupported VLM provider: {model.provider}")
def _get_ocr_options(self, opts: "ConversionOptions"):
"""Get OCR options based on configuration."""
from docling.datamodel.pipeline_options import (
@ -145,7 +130,7 @@ class DoclingLocalConverter(DocumentConverter):
pipeline_options.enable_remote_services = True
pipeline_options.picture_description_options = PictureDescriptionApiOptions(
url=AnyUrl(self._get_vlm_api_url(pic_desc.model)),
url=AnyUrl(vlm_api_url(self.config, pic_desc.model)),
params=dict(
model=pic_desc.model.name,
max_completion_tokens=pic_desc.max_tokens,

21
haiku_rag_slim/haiku/rag/converters/docling_serve.py
View file

@ -6,15 +6,13 @@ from pathlib import Path
from typing import TYPE_CHECKING, ClassVar
from haiku.rag.config import AppConfig
from haiku.rag.converters.base import DocumentConverter
from haiku.rag.converters.base import DocumentConverter, vlm_api_url
from haiku.rag.converters.text_utils import TextFileHandler
from haiku.rag.providers.docling_serve import DoclingServeClient
if TYPE_CHECKING:
from docling_core.types.doc.document import DoclingDocument
from haiku.rag.config.models import ModelConfig
class DoclingServeConverter(DocumentConverter):
"""Converter that uses docling-serve for document conversion.
@ -69,21 +67,6 @@ class DoclingServeConverter(DocumentConverter):
"""Return list of file extensions supported by this converter."""
return self.docling_serve_extensions + TextFileHandler.text_extensions
def _get_vlm_api_url(self, model: "ModelConfig") -> str:
"""Construct VLM API URL from model config."""
if model.base_url:
base = model.base_url.rstrip("/")
return f"{base}/v1/chat/completions"
if model.provider == "ollama":
base = self.config.providers.ollama.base_url.rstrip("/")
return f"{base}/v1/chat/completions"
if model.provider == "openai":
return "https://api.openai.com/v1/chat/completions"
raise ValueError(f"Unsupported VLM provider: {model.provider}")
def _build_conversion_data(self) -> dict[str, str | list[str]]:
"""Build form data for conversion request.
@ -124,7 +107,7 @@ class DoclingServeConverter(DocumentConverter):
if runs_vlm:
prompt = self.config.prompts.picture_description
picture_description_api = {
"url": self._get_vlm_api_url(pic_desc.model),
"url": vlm_api_url(self.config, pic_desc.model),
"params": {
"model": pic_desc.model.name,
"max_completion_tokens": pic_desc.max_tokens,

3
haiku_rag_slim/haiku/rag/embeddings/__init__.py
View file

@ -54,6 +54,9 @@ class EmbedderWrapper:
"""Embed documents/chunks for indexing."""
if not texts:
return []
return await self._embed_documents(texts)
async def _embed_documents(self, texts: list[str]) -> list[list[float]]:
assert self._embedder is not None
result = await self._embedder.embed_documents(texts)
return [list(e) for e in result.embeddings]

4
haiku_rag_slim/haiku/rag/embeddings/cohere.py
View file

@ -43,9 +43,7 @@ class CohereMultimodalEmbedder(EmbedderWrapper):
rows = await self._embed_texts([text], "search_query")
return rows[0]
async def embed_documents(self, texts: list[str]) -> list[list[float]]:
if not texts:
return []
async def _embed_documents(self, texts: list[str]) -> list[list[float]]:
return await self._embed_texts(texts, "search_document")
async def embed_image(self, image: "bytes | PILImage.Image") -> list[float]:

4
haiku_rag_slim/haiku/rag/embeddings/vllm.py
View file

@ -86,9 +86,7 @@ class VLLMMultimodalEmbedder(EmbedderWrapper):
)
return rows[0]
async def embed_documents(self, texts: list[str]) -> list[list[float]]:
if not texts:
return []
async def _embed_documents(self, texts: list[str]) -> list[list[float]]:
return await self._post(
{
"model": self._model_name,

4
haiku_rag_slim/haiku/rag/embeddings/voyageai.py
View file

@ -37,9 +37,7 @@ class VoyageMultimodalEmbedder(EmbedderWrapper):
)
return list(result.embeddings[0])
async def embed_documents(self, texts: list[str]) -> list[list[float]]:
if not texts:
return []
async def _embed_documents(self, texts: list[str]) -> list[list[float]]:
result = await self._client.multimodal_embed(
inputs=[[text] for text in texts],
model=self._model_name,

76
haiku_rag_slim/haiku/rag/reranking/__init__.py
View file

@ -5,79 +5,55 @@ from haiku.rag.reranking.base import RerankerBase
def get_reranker(config: AppConfig = Config) -> RerankerBase | None:
"""
Factory function to get the appropriate reranker based on the configuration.
Returns None if reranking is disabled.
"""Build the configured reranker, or None if reranking is disabled or its
optional dependency is not installed."""
model = config.reranking.model
if model is None:
return None
Args:
config: Configuration to use. Defaults to global Config.
Returns:
A reranker instance if configured, None otherwise.
"""
if config.reranking.model and config.reranking.model.provider == "mxbai":
try:
try:
if model.provider == "mxbai":
from haiku.rag.reranking.mxbai import MxBAIReranker
os.environ["TOKENIZERS_PARALLELISM"] = "true"
return MxBAIReranker()
except ImportError: # pragma: no cover
return None
if config.reranking.model and config.reranking.model.provider == "cohere":
try:
if model.provider == "cohere":
from haiku.rag.reranking.cohere import CohereReranker
return CohereReranker()
except ImportError: # pragma: no cover
return None
if config.reranking.model and config.reranking.model.provider == "vllm":
try:
if model.provider == "vllm":
if not model.base_url:
raise ValueError("vLLM reranker requires base_url in reranking.model")
from haiku.rag.reranking.vllm import VLLMReranker
base_url = config.reranking.model.base_url
if not base_url:
raise ValueError("vLLM reranker requires base_url in reranking.model")
return VLLMReranker(config.reranking.model.name, base_url)
except ImportError: # pragma: no cover
return None
return VLLMReranker(model.name, model.base_url)
if config.reranking.model and config.reranking.model.provider == "zeroentropy":
try:
if model.provider == "zeroentropy":
from haiku.rag.reranking.zeroentropy import ZeroEntropyReranker
model = config.reranking.model.name or "zerank-1"
return ZeroEntropyReranker(model)
except ImportError: # pragma: no cover
return None
return ZeroEntropyReranker(model.name or "zerank-1")
if config.reranking.model and config.reranking.model.provider == "jina":
from haiku.rag.reranking.jina import JinaReranker
if model.provider == "jina":
from haiku.rag.reranking.jina import JinaReranker
model = config.reranking.model.name or "jina-reranker-v3"
return JinaReranker(model)
return JinaReranker(model.name or "jina-reranker-v3")
if config.reranking.model and config.reranking.model.provider == "jina-local":
try:
if model.provider == "jina-local":
from haiku.rag.reranking.jina_local import JinaLocalReranker
model = config.reranking.model.name or "jinaai/jina-reranker-v3"
return JinaLocalReranker(model)
except ImportError: # pragma: no cover
return None
return JinaLocalReranker(model.name or "jinaai/jina-reranker-v3")
if config.reranking.model and config.reranking.model.provider == "cross-encoder":
try:
from haiku.rag.reranking.cross_encoder import CrossEncoderReranker
name = config.reranking.model.name
if not name:
if model.provider == "cross-encoder":
if not model.name:
raise ValueError(
"cross-encoder reranker requires name in reranking.model"
)
return CrossEncoderReranker(name)
except ImportError: # pragma: no cover
return None
from haiku.rag.reranking.cross_encoder import CrossEncoderReranker
return CrossEncoderReranker(model.name)
except ImportError: # pragma: no cover
return None
return None

9
haiku_rag_slim/haiku/rag/reranking/base.py
View file

@ -7,7 +7,14 @@ class RerankerBase:
async def rerank(
self, query: str, chunks: list[Chunk], top_n: int = 10
) -> list[tuple[Chunk, float]]:
if not chunks:
return []
return await self._rerank(query, chunks, top_n)
async def _rerank(
self, query: str, chunks: list[Chunk], top_n: int = 10
) -> list[tuple[Chunk, float]]:
raise NotImplementedError(
"Reranker is an abstract class. Please implement the rerank method in a subclass."
"Reranker is an abstract class. Please implement the _rerank method in a subclass."
)

5
haiku_rag_slim/haiku/rag/reranking/cohere.py
View file

@ -14,12 +14,9 @@ class CohereReranker(RerankerBase): # pragma: no cover
# Cohere SDK reads CO_API_KEY from environment by default
self._client = cohere.AsyncClientV2()
async def rerank(
async def _rerank(
self, query: str, chunks: list[Chunk], top_n: int = 10
) -> list[tuple[Chunk, float]]:
if not chunks:
return []
documents = [chunk.content for chunk in chunks]
model_name = self._model or "rerank-v3.5"

5
haiku_rag_slim/haiku/rag/reranking/cross_encoder.py
View file

@ -26,12 +26,9 @@ class CrossEncoderReranker(RerankerBase):
self._model = model
self._reranker = CrossEncoder(model)
async def rerank(
async def _rerank(
self, query: str, chunks: list[Chunk], top_n: int = 10
) -> list[tuple[Chunk, float]]:
if not chunks:
return []
documents = [chunk.content for chunk in chunks]
rankings = await asyncio.to_thread(
lambda: self._reranker.rank(query, documents, top_k=top_n)

5
haiku_rag_slim/haiku/rag/reranking/jina.py
View file

@ -15,12 +15,9 @@ class JinaReranker(RerankerBase):
if not self._api_key:
raise ValueError("JINA_API_KEY environment variable required")
async def rerank(
async def _rerank(
self, query: str, chunks: list[Chunk], top_n: int = 10
) -> list[tuple[Chunk, float]]:
if not chunks:
return []
documents = [chunk.content for chunk in chunks]
async with httpx.AsyncClient() as client:

5
haiku_rag_slim/haiku/rag/reranking/jina_local.py
View file

@ -26,12 +26,9 @@ class JinaLocalReranker(RerankerBase): # pragma: no cover
self._reranker = AutoModel.from_pretrained(model, trust_remote_code=True)
self._reranker.eval()
async def rerank(
async def _rerank(
self, query: str, chunks: list[Chunk], top_n: int = 10
) -> list[tuple[Chunk, float]]:
if not chunks:
return []
documents = [chunk.content for chunk in chunks]
results = await asyncio.to_thread(

5
haiku_rag_slim/haiku/rag/reranking/mxbai.py
View file

@ -25,12 +25,9 @@ class MxBAIReranker(RerankerBase):
)
self._client = MxbaiRerankV2(model_name, disable_transformers_warnings=True)
async def rerank(
async def _rerank(
self, query: str, chunks: list[Chunk], top_n: int = 10
) -> list[tuple[Chunk, float]]:
if not chunks:
return []
documents = [chunk.content for chunk in chunks]
results = await asyncio.to_thread(

5
haiku_rag_slim/haiku/rag/reranking/vllm.py
View file

@ -9,12 +9,9 @@ class VLLMReranker(RerankerBase): # pragma: no cover
self._model = model
self._base_url = base_url
async def rerank(
async def _rerank(
self, query: str, chunks: list[Chunk], top_n: int = 10
) -> list[tuple[Chunk, float]]:
if not chunks:
return []
# Prepare documents for reranking
documents = [chunk.content for chunk in chunks]

21
haiku_rag_slim/haiku/rag/reranking/zeroentropy.py
View file

@ -17,22 +17,9 @@ class ZeroEntropyReranker(RerankerBase): # pragma: no cover
# Zero Entropy SDK reads ZEROENTROPY_API_KEY from environment by default
self._client = AsyncZeroEntropy()
async def rerank(
async def _rerank(
self, query: str, chunks: list[Chunk], top_n: int = 10
) -> list[tuple[Chunk, float]]:
"""Rerank the given chunks based on relevance to the query.
Args:
query: The query to rank against
chunks: The chunks to rerank
top_n: The number of top results to return
Returns:
A list of (chunk, score) tuples, sorted by relevance
"""
if not chunks:
return []
# Prepare documents for Zero Entropy API
documents = [chunk.content for chunk in chunks]
@ -44,13 +31,9 @@ class ZeroEntropyReranker(RerankerBase): # pragma: no cover
documents=documents,
)
# Extract results and map back to chunks
# Zero Entropy returns results sorted by relevance with scores
reranked_results = []
# Get top_n results
for i, result in enumerate(response.results[:top_n]):
# Zero Entropy returns index and score for each document
for result in response.results[:top_n]:
chunk_index = result.index
score = result.relevance_score

258200
tests/cassettes/test_context_enhancement/test_expand_context_no_base64_images_docling_serve.yaml
View file

File diff suppressed because one or more lines are too long

76
tests/skills/conftest.py
View file

@ -1,5 +1,5 @@
import random
from unittest.mock import MagicMock
from unittest.mock import MagicMock, patch
import pytest
from pydantic_ai import RunContext
@ -12,6 +12,19 @@ from haiku.rag.skills._deps import AnalysisRunDeps, RAGRunDeps
VECTOR_DIM = 2560
def _seeded_vector(text: str) -> list[float]:
random.seed(hash(text) % (2**32))
return [random.random() for _ in range(VECTOR_DIM)]
async def _fake_embed_query(self, text: str) -> list[float]:
return _seeded_vector(text)
async def _fake_embed_documents(self, texts: list[str]) -> list[list[float]]:
return [_seeded_vector(t) for t in texts]
def _make_ctx(state=None, rag=None, sandbox=None):
"""Create a mock RunContext with RAGRunDeps (or AnalysisRunDeps when state is AnalysisState)."""
from haiku.rag.skills.analysis import AnalysisState
@ -35,20 +48,8 @@ def _get_tool(skill, name):
@pytest.fixture(autouse=True)
def mock_embedder(monkeypatch):
"""Monkeypatch the embedder to return deterministic vectors."""
async def fake_embed_query(self, text):
random.seed(hash(text) % (2**32))
return [random.random() for _ in range(VECTOR_DIM)]
async def fake_embed_documents(self, texts):
result = []
for t in texts:
random.seed(hash(t) % (2**32))
result.append([random.random() for _ in range(VECTOR_DIM)])
return result
monkeypatch.setattr(EmbedderWrapper, "embed_query", fake_embed_query)
monkeypatch.setattr(EmbedderWrapper, "embed_documents", fake_embed_documents)
monkeypatch.setattr(EmbedderWrapper, "embed_query", _fake_embed_query)
monkeypatch.setattr(EmbedderWrapper, "embed_documents", _fake_embed_documents)
@pytest.fixture
@ -56,23 +57,34 @@ def test_app_config():
return AppConfig(environment="skills-test")
@pytest.fixture
async def rag_db(temp_db_path):
"""Create a test database with sample documents."""
async with HaikuRAG(temp_db_path, create=True) as rag:
await rag.create_document(
"Artificial intelligence is transforming industries worldwide. "
"Deep learning models are used in healthcare, finance, and transportation.",
title="AI Overview",
uri="test://ai-overview",
)
await rag.create_document(
"Machine learning is a subset of artificial intelligence. "
"It includes supervised learning, unsupervised learning, and reinforcement learning.",
title="ML Basics",
uri="test://ml-basics",
)
return temp_db_path
@pytest.fixture(scope="session")
async def rag_db(tmp_path_factory):
"""Sample database with two documents, built once and shared read-only.
Consumers (``rag_client``, ``sandbox_factory``) only read, so the docling
conversion + ingest is paid once per session instead of per test. Document
vectors use the same seeded fakes as ``mock_embedder`` so search stays
consistent with query-time embeddings.
"""
db_path = tmp_path_factory.mktemp("skills_rag_db") / "rag.lancedb"
with (
patch.object(EmbedderWrapper, "embed_query", _fake_embed_query),
patch.object(EmbedderWrapper, "embed_documents", _fake_embed_documents),
):
async with HaikuRAG(db_path, create=True) as rag:
await rag.create_document(
"Artificial intelligence is transforming industries worldwide. "
"Deep learning models are used in healthcare, finance, and transportation.",
title="AI Overview",
uri="test://ai-overview",
)
await rag.create_document(
"Machine learning is a subset of artificial intelligence. "
"It includes supervised learning, unsupervised learning, and reinforcement learning.",
title="ML Basics",
uri="test://ml-basics",
)
return db_path
@pytest.fixture

10
tests/test_context_enhancement.py
View file

@ -558,19 +558,19 @@ async def test_expand_context_no_base64_images_docling_local(
@pytest.mark.vcr()
async def test_expand_context_no_base64_images_docling_serve(temp_db_path):
async def test_expand_context_no_base64_images_docling_serve(
temp_db_path, doclaynet_first_page_pdf
):
"""Ensure expanded context from real PDF does not contain base64 image data.
Tests end-to-end with doclaynet.pdf using docling-serve converter.
Tests end-to-end with a single page of doclaynet.pdf using docling-serve converter.
"""
from pathlib import Path
config = AppConfig()
config.processing.converter = "docling-serve"
config.processing.chunker = "docling-serve"
async with HaikuRAG(temp_db_path, config=config, create=True) as client:
pdf_path = Path(__file__).parent / "data" / "doclaynet.pdf"
pdf_path = doclaynet_first_page_pdf
result = await client.create_document_from_source(pdf_path)
doc = result if not isinstance(result, list) else result[0]
assert doc.id is not None

107
tests/test_converters.py
View file

@ -9,12 +9,43 @@ import pytest
from docling_core.types.doc.document import DoclingDocument
from haiku.rag.config import AppConfig
from haiku.rag.config.models import ModelConfig
from haiku.rag.converters import get_converter
from haiku.rag.converters.base import vlm_api_url
from haiku.rag.converters.docling_local import DoclingLocalConverter
from haiku.rag.converters.docling_serve import DoclingServeConverter
from haiku.rag.converters.text_utils import TextFileHandler
class TestVlmApiUrl:
"""URL construction for picture-description VLM models (shared by both converters)."""
def test_ollama_uses_provider_base_url(self):
url = vlm_api_url(
AppConfig(), ModelConfig(provider="ollama", name="ministral-3")
)
assert url == "http://localhost:11434/v1/chat/completions"
def test_custom_base_url_takes_precedence(self):
url = vlm_api_url(
AppConfig(),
ModelConfig(
provider="openai", name="gpt-4-vision", base_url="http://my-vllm:8000"
),
)
assert url == "http://my-vllm:8000/v1/chat/completions"
def test_openai_uses_public_endpoint(self):
url = vlm_api_url(
AppConfig(), ModelConfig(provider="openai", name="gpt-4-vision")
)
assert url == "https://api.openai.com/v1/chat/completions"
def test_unsupported_provider_raises(self):
with pytest.raises(ValueError, match="Unsupported VLM provider"):
vlm_api_url(AppConfig(), ModelConfig(provider="unsupported", name="test"))
@pytest.fixture(scope="module")
def vcr_cassette_dir():
return str(Path(__file__).parent / "cassettes" / "test_converters")
@ -763,44 +794,6 @@ class TestDoclingLocalConverter:
"Pages should have image data when generate_page_images=True"
)
def test_get_vlm_api_url_with_ollama(self, config):
"""Test VLM API URL construction for Ollama provider."""
converter = DoclingLocalConverter(config)
from haiku.rag.config.models import ModelConfig
model = ModelConfig(provider="ollama", name="ministral-3")
url = converter._get_vlm_api_url(model)
assert url == "http://localhost:11434/v1/chat/completions"
def test_get_vlm_api_url_with_custom_base_url(self, config):
"""Test VLM API URL construction with custom base_url."""
converter = DoclingLocalConverter(config)
from haiku.rag.config.models import ModelConfig
model = ModelConfig(
provider="openai", name="gpt-4-vision", base_url="http://my-vllm:8000"
)
url = converter._get_vlm_api_url(model)
assert url == "http://my-vllm:8000/v1/chat/completions"
def test_get_vlm_api_url_with_openai(self, config):
"""Test VLM API URL construction for OpenAI provider."""
converter = DoclingLocalConverter(config)
from haiku.rag.config.models import ModelConfig
model = ModelConfig(provider="openai", name="gpt-4-vision")
url = converter._get_vlm_api_url(model)
assert url == "https://api.openai.com/v1/chat/completions"
def test_get_vlm_api_url_unsupported_provider(self, config):
"""Test VLM API URL construction raises error for unsupported provider."""
converter = DoclingLocalConverter(config)
from haiku.rag.config.models import ModelConfig
model = ModelConfig(provider="unsupported", name="test")
with pytest.raises(ValueError, match="Unsupported VLM provider"):
converter._get_vlm_api_url(model)
def test_ocr_engine_config_applied(self, config):
"""Test that ocr_engine config is stored correctly."""
config.processing.conversion_options.ocr_engine = "rapidocr"
@ -1266,44 +1259,6 @@ class TestDoclingServeConverterPictureDescription:
config.providers.docling_serve.api_key = ""
return config
def test_get_vlm_api_url_with_ollama(self, config):
"""Test VLM API URL construction for Ollama provider."""
converter = DoclingServeConverter(config)
from haiku.rag.config.models import ModelConfig
model = ModelConfig(provider="ollama", name="ministral-3")
url = converter._get_vlm_api_url(model)
assert url == "http://localhost:11434/v1/chat/completions"
def test_get_vlm_api_url_with_custom_base_url(self, config):
"""Test VLM API URL construction with custom base_url."""
converter = DoclingServeConverter(config)
from haiku.rag.config.models import ModelConfig
model = ModelConfig(
provider="openai", name="gpt-4-vision", base_url="http://my-vllm:8000"
)
url = converter._get_vlm_api_url(model)
assert url == "http://my-vllm:8000/v1/chat/completions"
def test_get_vlm_api_url_with_openai(self, config):
"""Test VLM API URL construction for OpenAI provider."""
converter = DoclingServeConverter(config)
from haiku.rag.config.models import ModelConfig
model = ModelConfig(provider="openai", name="gpt-4-vision")
url = converter._get_vlm_api_url(model)
assert url == "https://api.openai.com/v1/chat/completions"
def test_get_vlm_api_url_unsupported_provider(self, config):
"""Test VLM API URL construction raises error for unsupported provider."""
converter = DoclingServeConverter(config)
from haiku.rag.config.models import ModelConfig
model = ModelConfig(provider="unsupported", name="test")
with pytest.raises(ValueError, match="Unsupported VLM provider"):
converter._get_vlm_api_url(model)
@pytest.mark.asyncio
async def test_picture_description_options_passed_to_api(self, config):
"""Picture-description options reach the docling-serve API when the

56
tests/test_reranker.py
View file

@ -1,4 +1,5 @@
from pathlib import Path
from unittest.mock import MagicMock
import pytest
@ -7,6 +8,14 @@ from haiku.rag.reranking import get_reranker
from haiku.rag.reranking.base import RerankerBase
from haiku.rag.store.models.chunk import Chunk
# Providers whose constructor loads a model in-process. Factory-routing tests
# patch the loader so they assert dispatch without paying the model load.
HEAVY_LOADERS = {
"mxbai": "MxbaiRerankV2",
"jina-local": "AutoModel",
"cross-encoder": "CrossEncoder",
}
@pytest.fixture(scope="module")
def vcr_cassette_dir():
@ -36,11 +45,16 @@ async def test_reranker_base():
expected_model = Config.reranking.model.name if Config.reranking.model else None
assert reranker._model == expected_model
# Empty input short-circuits in the base class without dispatching to _rerank.
assert await reranker.rerank("query", []) == []
# The actual rerank step is abstract.
with pytest.raises(NotImplementedError):
await reranker.rerank("query", [])
await reranker.rerank("query", chunks)
@pytest.mark.asyncio
@pytest.mark.integration
async def test_mxbai_reranker():
try:
from haiku.rag.config import Config
@ -62,18 +76,6 @@ async def test_mxbai_reranker():
pytest.skip("MxBAI package not installed")
@pytest.mark.asyncio
async def test_mxbai_reranker_empty_chunks():
try:
from haiku.rag.reranking.mxbai import MxBAIReranker
reranker = MxBAIReranker()
result = await reranker.rerank("query", [], top_n=2)
assert result == []
except ImportError:
pytest.skip("MxBAI package not installed")
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_cohere_reranker():
@ -251,6 +253,10 @@ class TestGetReranker:
mod = pytest.importorskip(class_module)
expected_class = getattr(mod, class_name)
loader_attr = HEAVY_LOADERS.get(provider)
if loader_attr:
monkeypatch.setattr(mod, loader_attr, MagicMock())
for key, value in env_vars.items():
monkeypatch.setenv(key, value)
@ -277,17 +283,6 @@ def test_jina_reranker_missing_api_key(monkeypatch):
JinaReranker("jina-reranker-v3")
@pytest.mark.asyncio
async def test_jina_reranker_empty_chunks(monkeypatch):
monkeypatch.setenv("JINA_API_KEY", "test-api-key")
from haiku.rag.reranking.jina import JinaReranker
reranker = JinaReranker("jina-reranker-v3")
result = await reranker.rerank("query", [], top_n=2)
assert result == []
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_jina_reranker(monkeypatch):
@ -332,6 +327,7 @@ async def test_jina_local_reranker():
@pytest.mark.asyncio
@pytest.mark.integration
async def test_cross_encoder_reranker():
try:
from haiku.rag.reranking.cross_encoder import CrossEncoderReranker
@ -347,15 +343,3 @@ async def test_cross_encoder_reranker():
assert "0" in top_ids or "2" in top_ids
except ImportError:
pytest.skip("sentence-transformers not installed")
@pytest.mark.asyncio
async def test_cross_encoder_reranker_empty_chunks():
try:
from haiku.rag.reranking.cross_encoder import CrossEncoderReranker
reranker = CrossEncoderReranker("cross-encoder/ms-marco-MiniLM-L-6-v2")
result = await reranker.rerank("query", [], top_n=2)
assert result == []
except ImportError:
pytest.skip("sentence-transformers not installed")