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

Remove repliqa from test setup

Browse source
This commit is contained in:
Yiorgis Gozadinos 2026-05-29 15:06:51 +03:00
parent e653c49fce
commit 339d231a4e
No known key found for this signature in database
13 changed files with 302 additions and 177 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
pyproject.toml
View file

@ -71,7 +71,6 @@ members = ["haiku_rag_slim", "evaluations"]
dev = [
"haiku.rag-evals",
"haiku.rag-slim[ingester]",
"datasets>=4.8.4",
"zensical",
"pre-commit>=4.5.1",
"pydantic-ai-slim[anthropic]",

21
tests/conftest.py
View file

@ -1,8 +1,9 @@
import json
import logging
import os
import tempfile
from pathlib import Path
from typing import TYPE_CHECKING, Any, cast
from typing import TYPE_CHECKING, Any
# Prevent tests from loading user's local haiku.rag.yaml by setting env var
# to a test config file BEFORE any haiku.rag imports.
@ -17,14 +18,12 @@ embeddings:
vector_dim: 2560
""")
os.environ["HAIKU_RAG_CONFIG_PATH"] = str(_test_config_path)
os.environ["HF_HUB_OFFLINE"] = "1"
import pydantic_ai.models # noqa: E402
import pytest # noqa: E402
import yaml # noqa: E402
if TYPE_CHECKING:
from datasets import Dataset
from vcr import VCR
setattr(pydantic_ai.models, "ALLOW_MODEL_REQUESTS", False)
@ -32,18 +31,10 @@ logging.getLogger("vcr.cassette").setLevel(logging.WARNING)
@pytest.fixture(scope="session")
def qa_corpus() -> "Dataset":
from datasets import Dataset, load_dataset, load_from_disk
ds_path = Path(__file__).parent / "data" / "dataset"
ds_path.mkdir(parents=True, exist_ok=True)
try:
return cast(Dataset, load_from_disk(ds_path))
except FileNotFoundError:
ds: Dataset = load_dataset("ServiceNow/repliqa")["repliqa_3"]
corpus = ds.filter(lambda doc: doc["document_topic"] == "News Stories")
corpus.save_to_disk(ds_path)
return corpus
def qa_corpus() -> list[dict[str, str]]:
corpus_path = Path(__file__).parent / "data" / "qa_corpus.json"
with open(corpus_path) as f:
return json.load(f)
@pytest.fixture

BIN
tests/data/dataset/data-00000-of-00001.arrow
View file

Binary file not shown.

106
tests/data/dataset/dataset_info.json
View file

@ -1,106 +0,0 @@
{
"builder_name": "parquet",
"citation": "",
"config_name": "default",
"dataset_name": "repliqa",
"dataset_size": 648387648,
"description": "",
"download_checksums": {
"hf://datasets/ServiceNow/repliqa@bc880adc948fd3a70d5f8b2b3a1d1ee90d820dbd/data/repliqa_0-00000-of-00001.parquet": {
"num_bytes": 20855947,
"checksum": null
},
"hf://datasets/ServiceNow/repliqa@bc880adc948fd3a70d5f8b2b3a1d1ee90d820dbd/data/repliqa_1-00000-of-00001.parquet": {
"num_bytes": 20903799,
"checksum": null
},
"hf://datasets/ServiceNow/repliqa@bc880adc948fd3a70d5f8b2b3a1d1ee90d820dbd/data/repliqa_2-00000-of-00001.parquet": {
"num_bytes": 20876034,
"checksum": null
},
"hf://datasets/ServiceNow/repliqa@bc880adc948fd3a70d5f8b2b3a1d1ee90d820dbd/data/repliqa_3-00000-of-00001.parquet": {
"num_bytes": 20919214,
"checksum": null
},
"hf://datasets/ServiceNow/repliqa@bc880adc948fd3a70d5f8b2b3a1d1ee90d820dbd/data/repliqa_4-00000-of-00001.parquet": {
"num_bytes": 20866527,
"checksum": null
}
},
"download_size": 104421521,
"features": {
"document_id": {
"dtype": "string",
"_type": "Value"
},
"document_topic": {
"dtype": "string",
"_type": "Value"
},
"document_path": {
"dtype": "string",
"_type": "Value"
},
"document_extracted": {
"dtype": "string",
"_type": "Value"
},
"question_id": {
"dtype": "string",
"_type": "Value"
},
"question": {
"dtype": "string",
"_type": "Value"
},
"answer": {
"dtype": "string",
"_type": "Value"
},
"long_answer": {
"dtype": "string",
"_type": "Value"
}
},
"homepage": "",
"license": "",
"size_in_bytes": 752809169,
"splits": {
"repliqa_0": {
"name": "repliqa_0",
"num_bytes": 129579219,
"num_examples": 17955,
"dataset_name": "repliqa"
},
"repliqa_1": {
"name": "repliqa_1",
"num_bytes": 129844840,
"num_examples": 17955,
"dataset_name": "repliqa"
},
"repliqa_2": {
"name": "repliqa_2",
"num_bytes": 129338522,
"num_examples": 17955,
"dataset_name": "repliqa"
},
"repliqa_3": {
"name": "repliqa_3",
"num_bytes": 129981017,
"num_examples": 17955,
"dataset_name": "repliqa"
},
"repliqa_4": {
"name": "repliqa_4",
"num_bytes": 129644050,
"num_examples": 17950,
"dataset_name": "repliqa"
}
},
"version": {
"version_str": "0.0.0",
"major": 0,
"minor": 0,
"patch": 0
}
}

13
tests/data/dataset/state.json
View file

@ -1,13 +0,0 @@
{
"_data_files": [
{
"filename": "data-00000-of-00001.arrow"
}
],
"_fingerprint": "2e29f63d782f12f0",
"_format_columns": null,
"_format_kwargs": {},
"_format_type": null,
"_output_all_columns": false,
"_split": "repliqa_3"
}

252
tests/data/qa_corpus.json Normal file
View file

File diff suppressed because one or more lines are too long

11
tests/test_chunk.py
View file

@ -1,5 +1,4 @@
import pytest
from datasets import Dataset
from haiku.rag.client import HaikuRAG
from haiku.rag.config import Config
@ -7,7 +6,9 @@ from haiku.rag.store.models.chunk import Chunk, ChunkMetadata, SearchResult
@pytest.mark.vcr()
async def test_chunk_repository_operations(qa_corpus: Dataset, temp_db_path):
async def test_chunk_repository_operations(
qa_corpus: list[dict[str, str]], temp_db_path
):
"""Test ChunkRepository operations."""
async with HaikuRAG(db_path=temp_db_path, config=Config, create=True) as client:
# Get the first document from the corpus
@ -46,7 +47,9 @@ async def test_chunk_repository_operations(qa_corpus: Dataset, temp_db_path):
@pytest.mark.vcr()
async def test_chunk_repository_pagination(qa_corpus: Dataset, temp_db_path):
async def test_chunk_repository_pagination(
qa_corpus: list[dict[str, str]], temp_db_path
):
"""Test ChunkRepository pagination with get_by_document_id and count_by_document_id."""
async with HaikuRAG(db_path=temp_db_path, config=Config, create=True) as client:
# Get the first document from the corpus (should produce multiple chunks)
@ -95,7 +98,7 @@ async def test_chunk_repository_pagination(qa_corpus: Dataset, temp_db_path):
@pytest.mark.vcr()
async def test_chunking_pipeline(qa_corpus: Dataset, temp_db_path):
async def test_chunking_pipeline(qa_corpus: list[dict[str, str]], temp_db_path):
"""Test document chunking using client primitives."""
from haiku.rag.client import HaikuRAG
from haiku.rag.embeddings import embed_chunks

5
tests/test_chunker.py
View file

@ -2,7 +2,6 @@ from pathlib import Path
from unittest.mock import AsyncMock, Mock, patch
import pytest
from datasets import Dataset
from transformers import AutoTokenizer
from haiku.rag.chunkers import get_chunker
@ -18,7 +17,7 @@ def vcr_cassette_dir():
@pytest.mark.asyncio
async def test_local_chunker(qa_corpus: Dataset):
async def test_local_chunker(qa_corpus: list[dict[str, str]]):
"""Test DoclingLocalChunker with real document."""
chunker = DoclingLocalChunker()
doc_text = qa_corpus[0]["document_extracted"]
@ -106,7 +105,7 @@ def test_tokenizer_cached_across_chunker_instances():
@pytest.mark.asyncio
async def test_local_chunker_hierarchical(qa_corpus: Dataset):
async def test_local_chunker_hierarchical(qa_corpus: list[dict[str, str]]):
"""Test DoclingLocalChunker with hierarchical chunking."""
config = AppConfig()
config.processing.chunker_type = "hierarchical"

5
tests/test_client.py
View file

@ -5,7 +5,6 @@ from unittest.mock import AsyncMock, patch
import httpx
import pytest
from datasets import Dataset
from haiku.rag.client import HaikuRAG
from haiku.rag.config import Config
@ -20,7 +19,7 @@ def vcr_cassette_dir():
@pytest.mark.vcr()
async def test_client_document_crud(qa_corpus: Dataset, temp_db_path):
async def test_client_document_crud(qa_corpus: list[dict[str, str]], temp_db_path):
"""Test HaikuRAG CRUD operations for documents."""
async with HaikuRAG(temp_db_path, create=True) as client:
# Get test data
@ -123,7 +122,7 @@ async def test_client_resolve_document(temp_db_path):
@pytest.mark.vcr()
async def test_client_update_document(qa_corpus: Dataset, temp_db_path):
async def test_client_update_document(qa_corpus: list[dict[str, str]], temp_db_path):
"""Test updating document with individual parameters."""
async with HaikuRAG(temp_db_path, create=True) as client:
# Get test data

17
tests/test_document.py
View file

@ -1,5 +1,4 @@
import pytest
from datasets import Dataset
from haiku.rag.store.engine import Store
from haiku.rag.store.models.document import Document
@ -8,7 +7,7 @@ from haiku.rag.store.repositories.document import DocumentRepository
@pytest.mark.asyncio
async def test_document_list_excludes_content_by_default(
qa_corpus: Dataset, temp_db_path
qa_corpus: list[dict[str, str]], temp_db_path
):
"""list_all excludes content and docling_document by default."""
async with Store(temp_db_path, create=True) as store:
@ -34,7 +33,7 @@ async def test_document_list_excludes_content_by_default(
@pytest.mark.asyncio
async def test_document_list_includes_content_when_requested(
qa_corpus: Dataset, temp_db_path
qa_corpus: list[dict[str, str]], temp_db_path
):
"""list_all returns content when include_content=True."""
async with Store(temp_db_path, create=True) as store:
@ -51,7 +50,7 @@ async def test_document_list_includes_content_when_requested(
@pytest.mark.asyncio
async def test_document_list_with_filter(qa_corpus: Dataset, temp_db_path):
async def test_document_list_with_filter(qa_corpus: list[dict[str, str]], temp_db_path):
"""Test listing documents with filter clause."""
async with Store(temp_db_path, create=True) as store:
doc_repo = DocumentRepository(store)
@ -229,7 +228,7 @@ def test_get_page_images():
@pytest.mark.asyncio
async def test_get_docling_data_loads_only_docling_columns(
qa_corpus: Dataset, temp_db_path
qa_corpus: list[dict[str, str]], temp_db_path
):
"""get_docling_data returns docling blob without loading content."""
import json
@ -276,7 +275,9 @@ async def test_get_docling_data_loads_only_docling_columns(
@pytest.mark.asyncio
async def test_get_pages_data_loads_only_pages_column(qa_corpus: Dataset, temp_db_path):
async def test_get_pages_data_loads_only_pages_column(
qa_corpus: list[dict[str, str]], temp_db_path
):
"""get_pages_data returns only page image data for a document."""
import json
@ -309,7 +310,7 @@ async def test_get_pages_data_loads_only_pages_column(qa_corpus: Dataset, temp_d
@pytest.mark.asyncio
async def test_get_pages_data_none_for_markdown_document(
qa_corpus: Dataset, temp_db_path
qa_corpus: list[dict[str, str]], temp_db_path
):
"""Markdown documents have no page images — get_pages_data returns None pages."""
async with Store(temp_db_path, create=True) as store:
@ -330,7 +331,7 @@ async def test_get_pages_data_none_for_markdown_document(
@pytest.mark.asyncio
async def test_document_get_by_uri_with_special_characters(
qa_corpus: Dataset, temp_db_path
qa_corpus: list[dict[str, str]], temp_db_path
):
"""Test get_by_uri handles URIs with special characters like single quotes."""
async with Store(temp_db_path, create=True) as store:

43
tests/test_rebuild.py
View file

@ -4,7 +4,6 @@ from pathlib import Path
from typing import TypedDict
import pytest
from datasets import Dataset
from haiku.rag.client import HaikuRAG, RebuildMode
@ -19,10 +18,10 @@ class ChunkData(TypedDict):
@pytest.mark.vcr()
async def test_rebuild_full(qa_corpus: Dataset, temp_db_path):
async def test_rebuild_full(qa_corpus: list[dict[str, str]], temp_db_path):
"""Test full rebuild: converts, chunks, and embeds all documents."""
async with HaikuRAG(temp_db_path, create=True) as client:
doc = await client.create_document(content=qa_corpus["document_extracted"][0])
doc = await client.create_document(content=qa_corpus[0]["document_extracted"])
assert doc.id is not None
assert doc.docling_document is not None
@ -49,10 +48,10 @@ async def test_rebuild_full(qa_corpus: Dataset, temp_db_path):
@pytest.mark.vcr()
async def test_rebuild_embed_only(qa_corpus: Dataset, temp_db_path):
async def test_rebuild_embed_only(qa_corpus: list[dict[str, str]], temp_db_path):
"""Test embed-only rebuild: keeps chunks, only regenerates embeddings."""
async with HaikuRAG(temp_db_path, create=True) as client:
doc = await client.create_document(content=qa_corpus["document_extracted"][0])
doc = await client.create_document(content=qa_corpus[0]["document_extracted"])
assert doc.id is not None
original_docling_json = doc.docling_document
@ -86,7 +85,7 @@ async def test_rebuild_embed_only(qa_corpus: Dataset, temp_db_path):
@pytest.mark.vcr()
async def test_rebuild_embed_only_multi_doc_streams_via_staging(
qa_corpus: Dataset, temp_db_path
qa_corpus: list[dict[str, str]], temp_db_path
):
"""Embed-only rebuild with multiple documents preserves chunks via staging.
@ -100,8 +99,8 @@ async def test_rebuild_embed_only_multi_doc_streams_via_staging(
- the rebuild yields every document with chunks.
"""
async with HaikuRAG(temp_db_path, create=True) as client:
doc1 = await client.create_document(content=qa_corpus["document_extracted"][0])
doc2 = await client.create_document(content=qa_corpus["document_extracted"][1])
doc1 = await client.create_document(content=qa_corpus[0]["document_extracted"])
doc2 = await client.create_document(content=qa_corpus[1]["document_extracted"])
assert doc1.id is not None and doc2.id is not None
chunks_before_1 = await client.chunk_repository.get_by_document_id(doc1.id)
@ -133,7 +132,9 @@ async def test_rebuild_embed_only_multi_doc_streams_via_staging(
@pytest.mark.vcr()
async def test_rebuild_drops_leftover_staging_table(qa_corpus: Dataset, temp_db_path):
async def test_rebuild_drops_leftover_staging_table(
qa_corpus: list[dict[str, str]], temp_db_path
):
"""Staging table without marker is treated as partial phase 1 and dropped.
Simulates a phase-1 interruption by creating only the staging table (no
@ -143,7 +144,7 @@ async def test_rebuild_drops_leftover_staging_table(qa_corpus: Dataset, temp_db_
from haiku.rag.client.rebuild import _StagingChunkRecord
async with HaikuRAG(temp_db_path, create=True) as client:
doc = await client.create_document(content=qa_corpus["document_extracted"][0])
doc = await client.create_document(content=qa_corpus[0]["document_extracted"])
assert doc.id is not None
# Simulate a partial phase 1 (staging exists, marker absent).
@ -167,7 +168,7 @@ async def test_rebuild_drops_leftover_staging_table(qa_corpus: Dataset, temp_db_
@pytest.mark.vcr()
async def test_rebuild_resumes_phase2_from_staging_after_crash(
qa_corpus: Dataset, temp_db_path
qa_corpus: list[dict[str, str]], temp_db_path
):
"""Marker + staging present → phase 2 resumes from staging instead of
redoing phase 1.
@ -184,7 +185,7 @@ async def test_rebuild_resumes_phase2_from_staging_after_crash(
)
async with HaikuRAG(temp_db_path, create=True) as client:
doc = await client.create_document(content=qa_corpus["document_extracted"][0])
doc = await client.create_document(content=qa_corpus[0]["document_extracted"])
assert doc.id is not None
original_chunks = await client.chunk_repository.get_by_document_id(doc.id)
assert original_chunks
@ -280,7 +281,7 @@ async def test_rebuild_drops_orphan_marker(temp_db_path):
@pytest.mark.vcr()
async def test_rebuild_non_embed_mode_drops_staging_recovery_state(
qa_corpus: Dataset, temp_db_path
qa_corpus: list[dict[str, str]], temp_db_path
):
"""Staging + marker from a prior embed-only crash → dropped on RECHUNK.
@ -295,7 +296,7 @@ async def test_rebuild_non_embed_mode_drops_staging_recovery_state(
)
async with HaikuRAG(temp_db_path, create=True) as client:
doc = await client.create_document(content=qa_corpus["document_extracted"][0])
doc = await client.create_document(content=qa_corpus[0]["document_extracted"])
assert doc.id is not None
await client.store.db.create_table(
@ -317,10 +318,12 @@ async def test_rebuild_non_embed_mode_drops_staging_recovery_state(
@pytest.mark.vcr()
async def test_rebuild_embed_only_skips_unchanged(qa_corpus: Dataset, temp_db_path):
async def test_rebuild_embed_only_skips_unchanged(
qa_corpus: list[dict[str, str]], temp_db_path
):
"""Test embed-only rebuild skips chunks with unchanged embeddings."""
async with HaikuRAG(temp_db_path, create=True) as client:
doc = await client.create_document(content=qa_corpus["document_extracted"][0])
doc = await client.create_document(content=qa_corpus[0]["document_extracted"])
assert doc.id is not None
# Get embeddings before rebuild
@ -354,7 +357,7 @@ async def test_rebuild_embed_only_skips_unchanged(qa_corpus: Dataset, temp_db_pa
@pytest.mark.vcr()
async def test_rebuild_embed_only_with_changed_vector_dim(
qa_corpus: Dataset, temp_db_path
qa_corpus: list[dict[str, str]], temp_db_path
):
"""Test embed-only rebuild when vector dimension changes.
@ -373,7 +376,7 @@ async def test_rebuild_embed_only_with_changed_vector_dim(
# Step 1: Create a database with normal 2560-dim embeddings
async with HaikuRAG(temp_db_path, create=True) as client:
doc = await client.create_document(content=qa_corpus["document_extracted"][0])
doc = await client.create_document(content=qa_corpus[0]["document_extracted"])
assert doc.id is not None
chunks_before = await client.chunk_repository.get_by_document_id(doc.id)
@ -463,10 +466,10 @@ async def test_rebuild_embed_only_with_changed_vector_dim(
@pytest.mark.vcr()
async def test_rebuild_rechunk(qa_corpus: Dataset, temp_db_path):
async def test_rebuild_rechunk(qa_corpus: list[dict[str, str]], temp_db_path):
"""Test rechunk rebuild: re-chunks from content without accessing source files."""
async with HaikuRAG(temp_db_path, create=True) as client:
doc = await client.create_document(content=qa_corpus["document_extracted"][0])
doc = await client.create_document(content=qa_corpus[0]["document_extracted"])
assert doc.id is not None
assert doc.docling_document is not None

3
tests/test_search.py
View file

@ -1,5 +1,4 @@
import pytest
from datasets import Dataset
from haiku.rag.client import HaikuRAG
from haiku.rag.config import Config
@ -7,7 +6,7 @@ from haiku.rag.store.models import SearchResult
@pytest.mark.vcr()
async def test_search_qa_corpus(qa_corpus: Dataset, temp_db_path):
async def test_search_qa_corpus(qa_corpus: list[dict[str, str]], temp_db_path):
"""Test that documents can be found by searching with their associated questions."""
async with HaikuRAG(db_path=temp_db_path, config=Config, create=True) as client:
# Load unique documents (limited to 10)

2
uv.lock
View file

@ -1479,7 +1479,6 @@ tui = [
[package.dev-dependencies]
dev = [
{ name = "datasets" },
{ name = "haiku-rag-evals" },
{ name = "haiku-rag-slim", extra = ["ingester"] },
{ name = "pre-commit" },
@ -1506,7 +1505,6 @@ provides-extras = ["tui", "s3", "cross-encoder", "ingester"]
[package.metadata.requires-dev]
dev = [
{ name = "datasets", specifier = ">=4.8.4" },
{ name = "haiku-rag-evals", editable = "evaluations" },
{ name = "haiku-rag-slim", extras = ["ingester"], editable = "haiku_rag_slim" },
{ name = "pre-commit", specifier = ">=4.5.1" },