Remove repliqa from test setup

pyproject.toml

@@ -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]",

tests/conftest.py

@@ -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":
+def qa_corpus() -> list[dict[str, str]]:
-    from datasets import Dataset, load_dataset, load_from_disk
+    corpus_path = Path(__file__).parent / "data" / "qa_corpus.json"
-
+    with open(corpus_path) as f:
-    ds_path = Path(__file__).parent / "data" / "dataset"
+        return json.load(f)
-    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
 
 
 @pytest.fixture

tests/data/dataset/dataset_info.json

@@ -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
-  }
-}

tests/data/dataset/state.json

@@ -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"
-}

tests/test_chunk.py

@@ -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

tests/test_chunker.py

@@ -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"

tests/test_client.py

@@ -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

tests/test_document.py

@@ -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:

tests/test_rebuild.py

@@ -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])
+        doc1 = await client.create_document(content=qa_corpus[0]["document_extracted"])
-        doc2 = await client.create_document(content=qa_corpus["document_extracted"][1])
+        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
 

tests/test_search.py

@@ -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)

uv.lock

@@ -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" },