auto-append /v1 to per-model Ollama base_url, fix flaky tests

This commit is contained in:
Yiorgis Gozadinos 2026-05-04 13:12:19 +03:00
parent 65d9c74224
commit aa3e9406cf
No known key found for this signature in database
8 changed files with 161 additions and 20 deletions

View file

@ -167,7 +167,9 @@ def get_embedder(config: AppConfig = Config) -> EmbedderWrapper:
if provider == "ollama":
# Use model-level base_url if set, otherwise fall back to providers config
base_url = embedding_model.base_url or f"{config.providers.ollama.base_url}/v1"
base_url = embedding_model.base_url or config.providers.ollama.base_url
if not base_url.rstrip("/").endswith("/v1"):
base_url = base_url.rstrip("/") + "/v1"
model = OpenAIEmbeddingModel(
model_name,
provider=OllamaProvider(base_url=base_url),

View file

@ -158,8 +158,11 @@ def get_model(
model_settings, OpenAIChatModelSettings, model_config
)
# Use model-level base_url if set, otherwise fall back to providers config
base_url = model_config.base_url or f"{app_config.providers.ollama.base_url}/v1"
# Ollama's OpenAI-compatible API lives under /v1. Append it if the
# configured base_url doesn't already include it.
base_url = model_config.base_url or app_config.providers.ollama.base_url
if not base_url.rstrip("/").endswith("/v1"):
base_url = base_url.rstrip("/") + "/v1"
return OpenAIChatModel(
model_name=model,

File diff suppressed because one or more lines are too long

View file

@ -5,6 +5,7 @@ from haiku.rag.client.documents import (
_store_document_with_chunks,
_update_document_with_chunks,
)
from haiku.rag.config import AppConfig
from haiku.rag.store.engine import Store
from haiku.rag.store.models.document_item import (
DocumentItem,
@ -646,9 +647,11 @@ class TestPictureDataPreservedThroughRoundTrip:
docling_doc = _docling_doc_with_picture()
async with HaikuRAG(temp_db_path, create=True) as rag:
# Preservation only kicks in under modes that retain picture bytes.
rag._config.processing.pictures = "image"
config = AppConfig()
# Preservation only kicks in under modes that retain picture bytes.
config.processing.pictures = "image"
async with HaikuRAG(temp_db_path, config=config, create=True) as rag:
document = Document(content="Hello world", uri="test://doc")
document.set_docling(docling_doc)
created = await _store_document_with_chunks(rag, document, [], docling_doc)

View file

@ -1,3 +1,5 @@
import logging
import pytest
import yaml
@ -7,6 +9,16 @@ from haiku.rag.config.loader import (
generate_default_config,
load_yaml_config,
)
from haiku.rag.config.loader import logger as loader_logger
class _ListHandler(logging.Handler):
def __init__(self) -> None:
super().__init__(level=logging.WARNING)
self.records: list[logging.LogRecord] = []
def emit(self, record: logging.LogRecord) -> None:
self.records.append(record)
def test_load_yaml_config(tmp_path):
@ -237,7 +249,7 @@ def _write(tmp_path, body: str):
return p
def test_load_yaml_legacy_picture_description_maps_to_description(tmp_path, caplog):
def test_load_yaml_legacy_picture_description_maps_to_description(tmp_path):
"""`picture_description.enabled=true` (with or without the image flag)
maps to `processing.pictures: description`."""
config_file = _write(
@ -251,15 +263,21 @@ processing:
timeout: 120
""",
)
with caplog.at_level("WARNING", logger="haiku.rag.config.loader"):
handler = _ListHandler()
loader_logger.addHandler(handler)
try:
data = load_yaml_config(config_file)
finally:
loader_logger.removeHandler(handler)
cfg = AppConfig.model_validate(data)
assert cfg.processing.pictures == "description"
assert cfg.processing.conversion_options.picture_description.timeout == 120
assert any("picture_description.enabled=true" in m.message for m in caplog.records)
assert any(
"picture_description.enabled=true" in r.getMessage() for r in handler.records
)
def test_load_yaml_legacy_generate_picture_images_maps_to_image(tmp_path, caplog):
def test_load_yaml_legacy_generate_picture_images_maps_to_image(tmp_path):
"""`generate_picture_images=true` alone maps to `pictures: image`."""
config_file = _write(
tmp_path,
@ -269,14 +287,20 @@ processing:
generate_picture_images: true
""",
)
with caplog.at_level("WARNING", logger="haiku.rag.config.loader"):
handler = _ListHandler()
loader_logger.addHandler(handler)
try:
data = load_yaml_config(config_file)
finally:
loader_logger.removeHandler(handler)
cfg = AppConfig.model_validate(data)
assert cfg.processing.pictures == "image"
assert any("generate_picture_images=true" in m.message for m in caplog.records)
assert any(
"generate_picture_images=true" in r.getMessage() for r in handler.records
)
def test_load_yaml_no_legacy_fields_keeps_default_none(tmp_path, caplog):
def test_load_yaml_no_legacy_fields_keeps_default_none(tmp_path):
"""Empty processing block leaves the default `none` mode untouched and
does not warn."""
config_file = _write(
@ -286,11 +310,15 @@ processing:
chunk_size: 256
""",
)
with caplog.at_level("WARNING", logger="haiku.rag.config.loader"):
handler = _ListHandler()
loader_logger.addHandler(handler)
try:
data = load_yaml_config(config_file)
finally:
loader_logger.removeHandler(handler)
cfg = AppConfig.model_validate(data)
assert cfg.processing.pictures == "none"
assert not caplog.records
assert not handler.records
def test_load_yaml_explicit_pictures_wins_over_legacy(tmp_path):

View file

@ -75,3 +75,39 @@ def test_unsupported_provider_raises():
with pytest.raises(ValueError, match="Unsupported embedding provider"):
get_embedder(custom_config)
def test_ollama_embedder_appends_v1_when_missing():
"""Per-model base_url without /v1 should get it appended for Ollama."""
config = AppConfig(
embeddings=EmbeddingsConfig(
model=EmbeddingModelConfig(
provider="ollama",
name="qwen3-embedding:4b",
vector_dim=2560,
base_url="http://my-ollama:11434",
),
),
)
embedder = get_embedder(config)
pa_model = embedder._embedder._model # type: ignore[union-attr] # ty: ignore[unresolved-attribute]
assert str(pa_model.base_url).rstrip("/").endswith("/v1") # type: ignore[union-attr] # ty: ignore[unresolved-attribute]
def test_ollama_embedder_does_not_double_append_v1():
"""If the user already includes /v1 we leave it alone."""
config = AppConfig(
embeddings=EmbeddingsConfig(
model=EmbeddingModelConfig(
provider="ollama",
name="qwen3-embedding:4b",
vector_dim=2560,
base_url="http://my-ollama:11434/v1",
),
),
)
embedder = get_embedder(config)
pa_model = embedder._embedder._model # type: ignore[union-attr] # ty: ignore[unresolved-attribute]
url = str(pa_model.base_url).rstrip("/") # type: ignore[union-attr] # ty: ignore[unresolved-attribute]
assert url.endswith("/v1")
assert not url.endswith("/v1/v1")

View file

@ -12,7 +12,7 @@ from pydantic_ai.usage import RunUsage
from haiku.rag.client import HaikuRAG
from haiku.rag.client.search import _populate_image_data
from haiku.rag.config import Config
from haiku.rag.config import AppConfig, Config
from haiku.rag.context import expand_with_items
from haiku.rag.store.models.chunk import SearchResult
from haiku.rag.store.models.document_item import DocumentItem
@ -167,6 +167,7 @@ async def test_expand_context_preserves_picture_refs_with_empty_text(temp_db_pat
assert "picture" in out.labels
@pytest.mark.vcr()
@pytest.mark.asyncio
async def test_rechunk_preserves_picture_data(temp_db_path):
"""``rebuild --rechunk`` keeps ``picture_data`` for every picture row."""
@ -177,8 +178,10 @@ async def test_rechunk_preserves_picture_data(temp_db_path):
docling_doc = _docling_doc_with_picture()
async with HaikuRAG(temp_db_path, create=True) as rag:
rag._config.processing.pictures = "image"
config = AppConfig()
config.processing.pictures = "image"
async with HaikuRAG(temp_db_path, config=config, create=True) as rag:
document = Document(content="x", uri="test://doc")
document.set_docling(docling_doc)
created = await _store_document_with_chunks(rag, document, [], docling_doc)
@ -193,6 +196,7 @@ async def test_rechunk_preserves_picture_data(temp_db_path):
assert after.get("#/pictures/0") == before.get("#/pictures/0")
@pytest.mark.vcr()
@pytest.mark.asyncio
async def test_update_clears_picture_data_when_mode_none(temp_db_path):
"""Switching to ``pictures="none"`` and re-running update_document
@ -207,9 +211,11 @@ async def test_update_clears_picture_data_when_mode_none(temp_db_path):
docling_doc = _docling_doc_with_picture()
async with HaikuRAG(temp_db_path, create=True) as rag:
config = AppConfig()
config.processing.pictures = "image"
async with HaikuRAG(temp_db_path, config=config, create=True) as rag:
# Ingest under "image" so picture bytes land in document_items.
rag._config.processing.pictures = "image"
document = Document(content="x", uri="test://doc")
document.set_docling(docling_doc)
created = await _store_document_with_chunks(rag, document, [], docling_doc)

View file

@ -162,6 +162,27 @@ def test_get_model_ollama_with_settings():
assert isinstance(result, OpenAIChatModel)
def test_get_model_ollama_appends_v1_to_per_model_base_url():
"""Per-model base_url without /v1 should get it appended."""
model_config = ModelConfig(
provider="ollama", name="qwen3.6", base_url="http://my-ollama:11434"
)
result = get_model(model_config)
assert isinstance(result, OpenAIChatModel)
assert str(result.client.base_url).rstrip("/").endswith("/v1")
def test_get_model_ollama_does_not_double_append_v1():
"""If the per-model base_url already ends with /v1, leave it alone."""
model_config = ModelConfig(
provider="ollama", name="qwen3.6", base_url="http://my-ollama:11434/v1"
)
result = get_model(model_config)
url = str(result.client.base_url).rstrip("/")
assert url.endswith("/v1")
assert not url.endswith("/v1/v1")
def test_get_model_openai():
"""Test get_model returns OpenAIChatModel for OpenAI."""
model_config = ModelConfig(provider="openai", name="gpt-4o")