Make get_config the only configuration lookup
haiku.rag.config exported two configuration instances: the lazy _config behind get_config/set_config, and Config, loaded at import time. Nothing linked them, and eleven signatures captured Config as a default argument, so set_config could not reach the factories, the client, the store or the MCP server. reranking/base.py went further and snapshotted the configured reranker name into a class attribute at import. Config is removed. Internal defaults are config: AppConfig | None = None, resolved through get_config() per call. RerankerBase._model is None and CohereReranker takes its model name as an argument, like every other reranker. The suite patched attributes on Config while production read the instance get_config() returns, a different object, so those patches were no-ops waiting to happen. They now go through get_config().
This commit is contained in:
parent
7465026b6a
commit
e8f00fcff4
36 changed files with 260 additions and 183 deletions
|
|
@ -25,6 +25,7 @@
|
||||||
|
|
||||||
### Removed
|
### Removed
|
||||||
|
|
||||||
|
- `haiku.rag.config.Config`, the eagerly loaded configuration instance. Use `get_config()` for the current global config, or pass an `AppConfig`. Every internal default (`get_embedder`, `get_converter`, `get_chunker`, `get_reranker`, `embed_chunks`, `HaikuRAG`, `Store`, `HaikuRAGApp`, `create_mcp_server`) now takes `config: AppConfig | None = None` and resolves it per call, so `set_config()` reaches them. `RerankerBase._model` no longer defaults to the configured reranker name; `CohereReranker` takes its model name as an argument.
|
||||||
- `wix` evaluation dataset and its reference config `evaluations/configs/wix.yaml`.
|
- `wix` evaluation dataset and its reference config `evaluations/configs/wix.yaml`.
|
||||||
|
|
||||||
### Fixed
|
### Fixed
|
||||||
|
|
|
||||||
|
|
@ -44,16 +44,16 @@ logger = logging.getLogger(__name__)
|
||||||
config_path = Path("/app/haiku.rag.yaml")
|
config_path = Path("/app/haiku.rag.yaml")
|
||||||
if config_path.exists():
|
if config_path.exists():
|
||||||
yaml_data = load_yaml_config(config_path)
|
yaml_data = load_yaml_config(config_path)
|
||||||
Config = AppConfig.model_validate(yaml_data)
|
config = AppConfig.model_validate(yaml_data)
|
||||||
else:
|
else:
|
||||||
Config = AppConfig()
|
config = AppConfig()
|
||||||
|
|
||||||
# Get DB path from environment
|
# Get DB path from environment
|
||||||
db_path_str = os.getenv("DB_PATH", "haiku_rag.lancedb")
|
db_path_str = os.getenv("DB_PATH", "haiku_rag.lancedb")
|
||||||
db_path = Path(db_path_str)
|
db_path = Path(db_path_str)
|
||||||
|
|
||||||
logger.info(f"Database path: {db_path}")
|
logger.info(f"Database path: {db_path}")
|
||||||
logger.info(f"QA Provider: {Config.qa.model.provider}, Model: {Config.qa.model.name}")
|
logger.info(f"QA Provider: {config.qa.model.provider}, Model: {config.qa.model.name}")
|
||||||
|
|
||||||
# Only HaikuRAG client is a singleton (expensive to create)
|
# Only HaikuRAG client is a singleton (expensive to create)
|
||||||
_client: HaikuRAG | None = None
|
_client: HaikuRAG | None = None
|
||||||
|
|
@ -71,7 +71,7 @@ async def get_client() -> HaikuRAG:
|
||||||
if _client is None:
|
if _client is None:
|
||||||
async with _client_lock:
|
async with _client_lock:
|
||||||
if _client is None:
|
if _client is None:
|
||||||
client = HaikuRAG(db_path=db_path, config=Config, create=True)
|
client = HaikuRAG(db_path=db_path, config=config, create=True)
|
||||||
await client.__aenter__()
|
await client.__aenter__()
|
||||||
_client = client
|
_client = client
|
||||||
return _client
|
return _client
|
||||||
|
|
@ -82,10 +82,10 @@ class AppDeps:
|
||||||
state: dict[str, Any] = field(default_factory=dict)
|
state: dict[str, Any] = field(default_factory=dict)
|
||||||
|
|
||||||
|
|
||||||
capability = create_capability(db_path=db_path, config=Config, defer_loading=False)
|
capability = create_capability(db_path=db_path, config=config, defer_loading=False)
|
||||||
|
|
||||||
agent = Agent(
|
agent = Agent(
|
||||||
get_model(Config.qa.model, Config),
|
get_model(config.qa.model, config),
|
||||||
instructions=AGENT_PREAMBLE,
|
instructions=AGENT_PREAMBLE,
|
||||||
# Conversations here are multi-turn, so earlier questions are reduced to the
|
# Conversations here are multi-turn, so earlier questions are reduced to the
|
||||||
# evidence they cited rather than carried whole, and every answer declares
|
# evidence they cited rather than carried whole, and every answer declares
|
||||||
|
|
@ -136,8 +136,8 @@ async def health_check(_: Request) -> JSONResponse:
|
||||||
return JSONResponse(
|
return JSONResponse(
|
||||||
{
|
{
|
||||||
"status": "healthy",
|
"status": "healthy",
|
||||||
"qa_provider": Config.qa.model.provider,
|
"qa_provider": config.qa.model.provider,
|
||||||
"qa_model": Config.qa.model.name,
|
"qa_model": config.qa.model.name,
|
||||||
"db_path": str(db_path),
|
"db_path": str(db_path),
|
||||||
"db_exists": db_path.exists(),
|
"db_exists": db_path.exists(),
|
||||||
}
|
}
|
||||||
|
|
|
||||||
|
|
@ -16,7 +16,7 @@ from rich.progress import (
|
||||||
)
|
)
|
||||||
|
|
||||||
from haiku.rag.client import HaikuRAG, RebuildMode
|
from haiku.rag.client import HaikuRAG, RebuildMode
|
||||||
from haiku.rag.config import AppConfig, Config
|
from haiku.rag.config import AppConfig, get_config
|
||||||
from haiku.rag.mcp import create_mcp_server
|
from haiku.rag.mcp import create_mcp_server
|
||||||
from haiku.rag.store.models.chunk import SearchType
|
from haiku.rag.store.models.chunk import SearchType
|
||||||
from haiku.rag.store.models.document import Document
|
from haiku.rag.store.models.document import Document
|
||||||
|
|
@ -33,11 +33,11 @@ class HaikuRAGApp: # pragma: no cover
|
||||||
def __init__(
|
def __init__(
|
||||||
self,
|
self,
|
||||||
db_path: Path,
|
db_path: Path,
|
||||||
config: AppConfig = Config,
|
config: AppConfig | None = None,
|
||||||
read_only: bool = False,
|
read_only: bool = False,
|
||||||
):
|
):
|
||||||
self.db_path = db_path
|
self.db_path = db_path
|
||||||
self.config = config
|
self.config = config if config is not None else get_config()
|
||||||
self.read_only = read_only
|
self.read_only = read_only
|
||||||
self.console = Console()
|
self.console = Console()
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -1,16 +1,16 @@
|
||||||
"""Document chunker abstraction for haiku.rag."""
|
"""Document chunker abstraction for haiku.rag."""
|
||||||
|
|
||||||
from haiku.rag.chunkers.base import DocumentChunker
|
from haiku.rag.chunkers.base import DocumentChunker
|
||||||
from haiku.rag.config import AppConfig, Config
|
from haiku.rag.config import AppConfig, get_config
|
||||||
|
|
||||||
__all__ = ["DocumentChunker", "get_chunker"]
|
__all__ = ["DocumentChunker", "get_chunker"]
|
||||||
|
|
||||||
|
|
||||||
def get_chunker(config: AppConfig = Config) -> DocumentChunker:
|
def get_chunker(config: AppConfig | None = None) -> DocumentChunker:
|
||||||
"""Get a document chunker instance based on configuration.
|
"""Get a document chunker instance based on configuration.
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
config: Configuration to use. Defaults to global Config.
|
config: Configuration to use. Defaults to the current global config.
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
DocumentChunker instance configured according to the config.
|
DocumentChunker instance configured according to the config.
|
||||||
|
|
@ -18,6 +18,7 @@ def get_chunker(config: AppConfig = Config) -> DocumentChunker:
|
||||||
Raises:
|
Raises:
|
||||||
ValueError: If the chunker provider is not recognized.
|
ValueError: If the chunker provider is not recognized.
|
||||||
"""
|
"""
|
||||||
|
config = config if config is not None else get_config()
|
||||||
if config.processing.chunker == "docling-local":
|
if config.processing.chunker == "docling-local":
|
||||||
from haiku.rag.chunkers.docling_local import DoclingLocalChunker
|
from haiku.rag.chunkers.docling_local import DoclingLocalChunker
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -3,7 +3,7 @@ from functools import cache
|
||||||
from typing import TYPE_CHECKING, cast
|
from typing import TYPE_CHECKING, cast
|
||||||
|
|
||||||
from haiku.rag.chunkers.base import DocumentChunker
|
from haiku.rag.chunkers.base import DocumentChunker
|
||||||
from haiku.rag.config import AppConfig, Config
|
from haiku.rag.config import AppConfig, get_config
|
||||||
from haiku.rag.store.models.chunk import Chunk, ChunkMetadata
|
from haiku.rag.store.models.chunk import Chunk, ChunkMetadata
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
|
|
@ -68,7 +68,7 @@ class DoclingLocalChunker(DocumentChunker):
|
||||||
config: Application configuration.
|
config: Application configuration.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
def __init__(self, config: AppConfig = Config):
|
def __init__(self, config: AppConfig | None = None):
|
||||||
from docling_core.transforms.chunker.hierarchical_chunker import (
|
from docling_core.transforms.chunker.hierarchical_chunker import (
|
||||||
HierarchicalChunker,
|
HierarchicalChunker,
|
||||||
)
|
)
|
||||||
|
|
@ -77,10 +77,10 @@ class DoclingLocalChunker(DocumentChunker):
|
||||||
HuggingFaceTokenizer,
|
HuggingFaceTokenizer,
|
||||||
)
|
)
|
||||||
|
|
||||||
self.config = config
|
self.config = config if config is not None else get_config()
|
||||||
self.chunk_size = config.processing.chunk_size
|
self.chunk_size = self.config.processing.chunk_size
|
||||||
self.chunker_type = config.processing.chunker_type
|
self.chunker_type = self.config.processing.chunker_type
|
||||||
self.tokenizer_name = config.processing.chunking_tokenizer
|
self.tokenizer_name = self.config.processing.chunking_tokenizer
|
||||||
|
|
||||||
if self.chunker_type == "hybrid":
|
if self.chunker_type == "hybrid":
|
||||||
hf_tokenizer = _get_tokenizer(self.tokenizer_name)
|
hf_tokenizer = _get_tokenizer(self.tokenizer_name)
|
||||||
|
|
@ -88,16 +88,16 @@ class DoclingLocalChunker(DocumentChunker):
|
||||||
tokenizer=hf_tokenizer, max_tokens=self.chunk_size
|
tokenizer=hf_tokenizer, max_tokens=self.chunk_size
|
||||||
)
|
)
|
||||||
serializer_provider = _create_markdown_serializer_provider(
|
serializer_provider = _create_markdown_serializer_provider(
|
||||||
use_markdown_tables=config.processing.chunking_use_markdown_tables
|
use_markdown_tables=self.config.processing.chunking_use_markdown_tables
|
||||||
)
|
)
|
||||||
self.chunker = HybridChunker(
|
self.chunker = HybridChunker(
|
||||||
tokenizer=tokenizer,
|
tokenizer=tokenizer,
|
||||||
merge_peers=config.processing.chunking_merge_peers,
|
merge_peers=self.config.processing.chunking_merge_peers,
|
||||||
serializer_provider=serializer_provider,
|
serializer_provider=serializer_provider,
|
||||||
)
|
)
|
||||||
elif self.chunker_type == "hierarchical":
|
elif self.chunker_type == "hierarchical":
|
||||||
serializer_provider = _create_markdown_serializer_provider(
|
serializer_provider = _create_markdown_serializer_provider(
|
||||||
use_markdown_tables=config.processing.chunking_use_markdown_tables
|
use_markdown_tables=self.config.processing.chunking_use_markdown_tables
|
||||||
)
|
)
|
||||||
self.chunker = HierarchicalChunker(serializer_provider=serializer_provider)
|
self.chunker = HierarchicalChunker(serializer_provider=serializer_provider)
|
||||||
else:
|
else:
|
||||||
|
|
|
||||||
|
|
@ -4,7 +4,7 @@ from io import BytesIO
|
||||||
from typing import TYPE_CHECKING
|
from typing import TYPE_CHECKING
|
||||||
|
|
||||||
from haiku.rag.chunkers.base import DocumentChunker
|
from haiku.rag.chunkers.base import DocumentChunker
|
||||||
from haiku.rag.config import AppConfig, Config
|
from haiku.rag.config import AppConfig, get_config
|
||||||
from haiku.rag.providers.docling_serve import DoclingServeClient
|
from haiku.rag.providers.docling_serve import DoclingServeClient
|
||||||
from haiku.rag.store.models.chunk import Chunk, ChunkMetadata
|
from haiku.rag.store.models.chunk import Chunk, ChunkMetadata
|
||||||
|
|
||||||
|
|
@ -54,10 +54,12 @@ class DoclingServeChunker(DocumentChunker):
|
||||||
config: Application configuration containing docling-serve settings.
|
config: Application configuration containing docling-serve settings.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
def __init__(self, config: AppConfig = Config):
|
def __init__(self, config: AppConfig | None = None):
|
||||||
self.config = config
|
self.config = config if config is not None else get_config()
|
||||||
self.client = DoclingServeClient.from_config(config.providers.docling_serve)
|
self.client = DoclingServeClient.from_config(
|
||||||
self.chunker_type = config.processing.chunker_type
|
self.config.providers.docling_serve
|
||||||
|
)
|
||||||
|
self.chunker_type = self.config.processing.chunker_type
|
||||||
|
|
||||||
def _build_chunking_data(self) -> dict[str, str | list[str]]:
|
def _build_chunking_data(self) -> dict[str, str | list[str]]:
|
||||||
"""Build form data for chunking request."""
|
"""Build form data for chunking request."""
|
||||||
|
|
|
||||||
|
|
@ -15,7 +15,7 @@ from urllib.parse import urlparse
|
||||||
import httpx
|
import httpx
|
||||||
|
|
||||||
from haiku.rag.client.documents import DocumentImport
|
from haiku.rag.client.documents import DocumentImport
|
||||||
from haiku.rag.config import AppConfig, Config
|
from haiku.rag.config import AppConfig, get_config
|
||||||
from haiku.rag.converters import get_converter
|
from haiku.rag.converters import get_converter
|
||||||
from haiku.rag.reranking import get_reranker
|
from haiku.rag.reranking import get_reranker
|
||||||
from haiku.rag.store.engine import Store
|
from haiku.rag.store.engine import Store
|
||||||
|
|
@ -67,7 +67,7 @@ class HaikuRAG:
|
||||||
def __init__(
|
def __init__(
|
||||||
self,
|
self,
|
||||||
db_path: Path | None = None,
|
db_path: Path | None = None,
|
||||||
config: AppConfig = Config,
|
config: AppConfig | None = None,
|
||||||
skip_validation: bool = False,
|
skip_validation: bool = False,
|
||||||
create: bool = False,
|
create: bool = False,
|
||||||
read_only: bool = False,
|
read_only: bool = False,
|
||||||
|
|
@ -76,12 +76,12 @@ class HaikuRAG:
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
db_path: Path to the database file. If None, uses config.storage.data_dir.
|
db_path: Path to the database file. If None, uses config.storage.data_dir.
|
||||||
config: Configuration to use. Defaults to global Config.
|
config: Configuration to use. Defaults to the current global config.
|
||||||
skip_validation: Whether to skip configuration validation on database load.
|
skip_validation: Whether to skip configuration validation on database load.
|
||||||
create: Whether to create the database if it doesn't exist.
|
create: Whether to create the database if it doesn't exist.
|
||||||
read_only: Whether to open the database in read-only mode.
|
read_only: Whether to open the database in read-only mode.
|
||||||
"""
|
"""
|
||||||
self._config = config
|
self._config = config if config is not None else get_config()
|
||||||
if db_path is None:
|
if db_path is None:
|
||||||
db_path = self._config.storage.data_dir / "haiku.rag.lancedb"
|
db_path = self._config.storage.data_dir / "haiku.rag.lancedb"
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -36,7 +36,6 @@ from haiku.rag.config.models import (
|
||||||
)
|
)
|
||||||
|
|
||||||
__all__ = [
|
__all__ = [
|
||||||
"Config",
|
|
||||||
"APIConfig",
|
"APIConfig",
|
||||||
"AppConfig",
|
"AppConfig",
|
||||||
"CircuitBreakerConfig",
|
"CircuitBreakerConfig",
|
||||||
|
|
@ -78,7 +77,7 @@ _config: AppConfig | None = None
|
||||||
|
|
||||||
|
|
||||||
def _load_default_config() -> AppConfig:
|
def _load_default_config() -> AppConfig:
|
||||||
"""Load config from default locations (used at import time)."""
|
"""Load config from the default locations."""
|
||||||
config_path = find_config_file(None)
|
config_path = find_config_file(None)
|
||||||
if config_path:
|
if config_path:
|
||||||
yaml_data = load_yaml_config(config_path)
|
yaml_data = load_yaml_config(config_path)
|
||||||
|
|
@ -93,12 +92,8 @@ def set_config(config: AppConfig) -> None:
|
||||||
|
|
||||||
|
|
||||||
def get_config() -> AppConfig:
|
def get_config() -> AppConfig:
|
||||||
"""Get the current config instance."""
|
"""Get the current config instance, loading it on first use."""
|
||||||
global _config
|
global _config
|
||||||
if _config is None:
|
if _config is None:
|
||||||
_config = _load_default_config()
|
_config = _load_default_config()
|
||||||
return _config
|
return _config
|
||||||
|
|
||||||
|
|
||||||
# Legacy compatibility - Config is the default instance
|
|
||||||
Config = _load_default_config()
|
|
||||||
|
|
|
||||||
|
|
@ -1,16 +1,16 @@
|
||||||
"""Document converter abstraction for haiku.rag."""
|
"""Document converter abstraction for haiku.rag."""
|
||||||
|
|
||||||
from haiku.rag.config import AppConfig, Config
|
from haiku.rag.config import AppConfig, get_config
|
||||||
from haiku.rag.converters.base import DocumentConverter
|
from haiku.rag.converters.base import DocumentConverter
|
||||||
|
|
||||||
__all__ = ["DocumentConverter", "get_converter"]
|
__all__ = ["DocumentConverter", "get_converter"]
|
||||||
|
|
||||||
|
|
||||||
def get_converter(config: AppConfig = Config) -> DocumentConverter:
|
def get_converter(config: AppConfig | None = None) -> DocumentConverter:
|
||||||
"""Get a document converter instance based on configuration.
|
"""Get a document converter instance based on configuration.
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
config: Configuration to use. Defaults to global Config.
|
config: Configuration to use. Defaults to the current global config.
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
DocumentConverter instance configured according to the config.
|
DocumentConverter instance configured according to the config.
|
||||||
|
|
@ -18,6 +18,7 @@ def get_converter(config: AppConfig = Config) -> DocumentConverter:
|
||||||
Raises:
|
Raises:
|
||||||
ValueError: If the converter provider is not recognized.
|
ValueError: If the converter provider is not recognized.
|
||||||
"""
|
"""
|
||||||
|
config = config if config is not None else get_config()
|
||||||
if config.processing.converter == "docling-local":
|
if config.processing.converter == "docling-local":
|
||||||
from haiku.rag.converters.docling_local import DoclingLocalConverter
|
from haiku.rag.converters.docling_local import DoclingLocalConverter
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -7,7 +7,7 @@ from pydantic_ai.embeddings.openai import OpenAIEmbeddingModel
|
||||||
from pydantic_ai.providers.ollama import OllamaProvider
|
from pydantic_ai.providers.ollama import OllamaProvider
|
||||||
from pydantic_ai.providers.openai import OpenAIProvider
|
from pydantic_ai.providers.openai import OpenAIProvider
|
||||||
|
|
||||||
from haiku.rag.config import AppConfig, Config
|
from haiku.rag.config import AppConfig, get_config
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
from PIL import Image as PILImage
|
from PIL import Image as PILImage
|
||||||
|
|
@ -118,7 +118,9 @@ def contextualize(chunks: list["Chunk"]) -> list[str]:
|
||||||
|
|
||||||
|
|
||||||
async def embed_chunks(
|
async def embed_chunks(
|
||||||
chunks: list["Chunk"], embedder: "EmbedderWrapper", config: AppConfig = Config
|
chunks: list["Chunk"],
|
||||||
|
embedder: "EmbedderWrapper",
|
||||||
|
config: AppConfig | None = None,
|
||||||
) -> list["Chunk"]:
|
) -> list["Chunk"]:
|
||||||
"""Generate embeddings for chunks, dispatching text vs picture variants.
|
"""Generate embeddings for chunks, dispatching text vs picture variants.
|
||||||
|
|
||||||
|
|
@ -127,6 +129,7 @@ async def embed_chunks(
|
||||||
are routed through ``embed_images`` and require a multimodal embedder.
|
are routed through ``embed_images`` and require a multimodal embedder.
|
||||||
Vectors land in the original chunk order.
|
Vectors land in the original chunk order.
|
||||||
"""
|
"""
|
||||||
|
config = config if config is not None else get_config()
|
||||||
if not chunks:
|
if not chunks:
|
||||||
return []
|
return []
|
||||||
|
|
||||||
|
|
@ -181,15 +184,16 @@ async def embed_chunks(
|
||||||
]
|
]
|
||||||
|
|
||||||
|
|
||||||
def get_embedder(config: AppConfig = Config) -> EmbedderWrapper:
|
def get_embedder(config: AppConfig | None = None) -> EmbedderWrapper:
|
||||||
"""Factory function to get the appropriate embedder based on the configuration.
|
"""Factory function to get the appropriate embedder based on the configuration.
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
config: Configuration to use. Defaults to global Config.
|
config: Configuration to use. Defaults to the current global config.
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
An embedder instance configured according to the config.
|
An embedder instance configured according to the config.
|
||||||
"""
|
"""
|
||||||
|
config = config if config is not None else get_config()
|
||||||
embedding_model = config.embeddings.model
|
embedding_model = config.embeddings.model
|
||||||
provider = embedding_model.provider
|
provider = embedding_model.provider
|
||||||
model_name = embedding_model.name
|
model_name = embedding_model.name
|
||||||
|
|
|
||||||
|
|
@ -7,7 +7,7 @@ from typing import Any
|
||||||
from fastmcp import FastMCP
|
from fastmcp import FastMCP
|
||||||
|
|
||||||
from haiku.rag.client import HaikuRAG
|
from haiku.rag.client import HaikuRAG
|
||||||
from haiku.rag.config import AppConfig, Config
|
from haiku.rag.config import AppConfig, get_config
|
||||||
from haiku.rag.store.models import Document, SearchResult
|
from haiku.rag.store.models import Document, SearchResult
|
||||||
from haiku.rag.tools.document import DocumentInfo
|
from haiku.rag.tools.document import DocumentInfo
|
||||||
from haiku.rag.utils import format_citations
|
from haiku.rag.utils import format_citations
|
||||||
|
|
@ -22,7 +22,7 @@ def _decode_images(images_base64: list[str] | None) -> list[bytes] | None:
|
||||||
|
|
||||||
|
|
||||||
def create_mcp_server(
|
def create_mcp_server(
|
||||||
db_path: Path, config: AppConfig = Config, read_only: bool = False
|
db_path: Path, config: AppConfig | None = None, read_only: bool = False
|
||||||
) -> FastMCP:
|
) -> FastMCP:
|
||||||
"""Create an MCP server with the specified database path.
|
"""Create an MCP server with the specified database path.
|
||||||
|
|
||||||
|
|
@ -31,6 +31,7 @@ def create_mcp_server(
|
||||||
config: Configuration to use.
|
config: Configuration to use.
|
||||||
read_only: If True, write tools (add_document_*, delete_document) are not registered.
|
read_only: If True, write tools (add_document_*, delete_document) are not registered.
|
||||||
"""
|
"""
|
||||||
|
config = config if config is not None else get_config()
|
||||||
client: HaikuRAG | None = None
|
client: HaikuRAG | None = None
|
||||||
stack = AsyncExitStack()
|
stack = AsyncExitStack()
|
||||||
client_lock = asyncio.Lock()
|
client_lock = asyncio.Lock()
|
||||||
|
|
|
||||||
|
|
@ -1,10 +1,11 @@
|
||||||
from haiku.rag.config import AppConfig, Config
|
from haiku.rag.config import AppConfig, get_config
|
||||||
from haiku.rag.reranking.base import RerankerBase
|
from haiku.rag.reranking.base import RerankerBase
|
||||||
|
|
||||||
|
|
||||||
def get_reranker(config: AppConfig = Config) -> RerankerBase | None:
|
def get_reranker(config: AppConfig | None = None) -> RerankerBase | None:
|
||||||
"""Build the configured reranker, or None if reranking is disabled or its
|
"""Build the configured reranker, or None if reranking is disabled or its
|
||||||
optional dependency is not installed."""
|
optional dependency is not installed."""
|
||||||
|
config = config if config is not None else get_config()
|
||||||
model = config.reranking.model
|
model = config.reranking.model
|
||||||
if model is None:
|
if model is None:
|
||||||
return None
|
return None
|
||||||
|
|
@ -16,7 +17,7 @@ def get_reranker(config: AppConfig = Config) -> RerankerBase | None:
|
||||||
if model.provider == "cohere":
|
if model.provider == "cohere":
|
||||||
from haiku.rag.reranking.cohere import CohereReranker
|
from haiku.rag.reranking.cohere import CohereReranker
|
||||||
|
|
||||||
return CohereReranker()
|
return CohereReranker(model.name)
|
||||||
|
|
||||||
if model.provider == "vllm":
|
if model.provider == "vllm":
|
||||||
if not model.base_url:
|
if not model.base_url:
|
||||||
|
|
|
||||||
|
|
@ -1,9 +1,8 @@
|
||||||
from haiku.rag.config import Config
|
|
||||||
from haiku.rag.store.models.chunk import Chunk
|
from haiku.rag.store.models.chunk import Chunk
|
||||||
|
|
||||||
|
|
||||||
class RerankerBase:
|
class RerankerBase:
|
||||||
_model: str | None = Config.reranking.model.name if Config.reranking.model else None
|
_model: str | None = None
|
||||||
|
|
||||||
async def rerank(
|
async def rerank(
|
||||||
self, query: str, chunks: list[Chunk], top_n: int = 10
|
self, query: str, chunks: list[Chunk], top_n: int = 10
|
||||||
|
|
|
||||||
|
|
@ -10,7 +10,8 @@ except ImportError as e: # pragma: no cover
|
||||||
|
|
||||||
|
|
||||||
class CohereReranker(RerankerBase): # pragma: no cover
|
class CohereReranker(RerankerBase): # pragma: no cover
|
||||||
def __init__(self):
|
def __init__(self, model: str | None = None):
|
||||||
|
self._model = model
|
||||||
# Cohere SDK reads CO_API_KEY from environment by default
|
# Cohere SDK reads CO_API_KEY from environment by default
|
||||||
self._client = cohere.AsyncClientV2()
|
self._client = cohere.AsyncClientV2()
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -18,7 +18,7 @@ from lancedb.pydantic import LanceModel, Vector
|
||||||
from packaging.version import parse
|
from packaging.version import parse
|
||||||
from pydantic import BaseModel, Field
|
from pydantic import BaseModel, Field
|
||||||
|
|
||||||
from haiku.rag.config import AppConfig, Config
|
from haiku.rag.config import AppConfig, get_config
|
||||||
from haiku.rag.embeddings import get_embedder
|
from haiku.rag.embeddings import get_embedder
|
||||||
from haiku.rag.store.exceptions import MigrationRequiredError, ReadOnlyError
|
from haiku.rag.store.exceptions import MigrationRequiredError, ReadOnlyError
|
||||||
|
|
||||||
|
|
@ -497,14 +497,14 @@ class Store:
|
||||||
def __init__(
|
def __init__(
|
||||||
self,
|
self,
|
||||||
db_path: Path,
|
db_path: Path,
|
||||||
config: AppConfig = Config,
|
config: AppConfig | None = None,
|
||||||
skip_validation: bool = False,
|
skip_validation: bool = False,
|
||||||
create: bool = False,
|
create: bool = False,
|
||||||
read_only: bool = False,
|
read_only: bool = False,
|
||||||
skip_migration_check: bool = False,
|
skip_migration_check: bool = False,
|
||||||
):
|
):
|
||||||
self.db_path: Path = db_path
|
self.db_path: Path = db_path
|
||||||
self._config = config
|
self._config = config if config is not None else get_config()
|
||||||
self._read_only = read_only
|
self._read_only = read_only
|
||||||
self._create = create
|
self._create = create
|
||||||
self._skip_validation = skip_validation
|
self._skip_validation = skip_validation
|
||||||
|
|
|
||||||
|
|
@ -119,7 +119,7 @@ def get_model(
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
model_config: ModelConfig with provider, model, and settings
|
model_config: ModelConfig with provider, model, and settings
|
||||||
app_config: AppConfig for provider base URLs (defaults to global Config)
|
app_config: AppConfig for provider base URLs (defaults to the current global config)
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
A configured model instance
|
A configured model instance
|
||||||
|
|
@ -129,9 +129,9 @@ def get_model(
|
||||||
from pydantic_ai.providers.openai import OpenAIProvider
|
from pydantic_ai.providers.openai import OpenAIProvider
|
||||||
|
|
||||||
if app_config is None:
|
if app_config is None:
|
||||||
from haiku.rag.config import Config
|
from haiku.rag.config import get_config
|
||||||
|
|
||||||
app_config = Config
|
app_config = get_config()
|
||||||
|
|
||||||
provider = model_config.provider
|
provider = model_config.provider
|
||||||
model = model_config.name
|
model = model_config.name
|
||||||
|
|
|
||||||
|
|
@ -1,7 +1,7 @@
|
||||||
import pytest
|
import pytest
|
||||||
|
|
||||||
from haiku.rag.client import HaikuRAG
|
from haiku.rag.client import HaikuRAG
|
||||||
from haiku.rag.config import Config
|
from haiku.rag.config import get_config
|
||||||
from haiku.rag.store.models.chunk import Chunk, ChunkMetadata, SearchResult
|
from haiku.rag.store.models.chunk import Chunk, ChunkMetadata, SearchResult
|
||||||
from tests.conftest import capture_logs
|
from tests.conftest import capture_logs
|
||||||
|
|
||||||
|
|
@ -11,7 +11,9 @@ async def test_chunk_repository_operations(
|
||||||
qa_corpus: list[dict[str, str]], temp_db_path
|
qa_corpus: list[dict[str, str]], temp_db_path
|
||||||
):
|
):
|
||||||
"""Test ChunkRepository operations."""
|
"""Test ChunkRepository operations."""
|
||||||
async with HaikuRAG(db_path=temp_db_path, config=Config, create=True) as client:
|
async with HaikuRAG(
|
||||||
|
db_path=temp_db_path, config=get_config(), create=True
|
||||||
|
) as client:
|
||||||
# Get the first document from the corpus
|
# Get the first document from the corpus
|
||||||
first_doc = qa_corpus[0]
|
first_doc = qa_corpus[0]
|
||||||
document_text = first_doc["document_extracted"]
|
document_text = first_doc["document_extracted"]
|
||||||
|
|
@ -52,7 +54,9 @@ async def test_chunk_repository_pagination(
|
||||||
qa_corpus: list[dict[str, str]], temp_db_path
|
qa_corpus: list[dict[str, str]], temp_db_path
|
||||||
):
|
):
|
||||||
"""Test ChunkRepository pagination with get_by_document_id and count_by_document_id."""
|
"""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:
|
async with HaikuRAG(
|
||||||
|
db_path=temp_db_path, config=get_config(), create=True
|
||||||
|
) as client:
|
||||||
# Get the first document from the corpus (should produce multiple chunks)
|
# Get the first document from the corpus (should produce multiple chunks)
|
||||||
first_doc = qa_corpus[0]
|
first_doc = qa_corpus[0]
|
||||||
document_text = first_doc["document_extracted"]
|
document_text = first_doc["document_extracted"]
|
||||||
|
|
@ -443,7 +447,9 @@ async def test_chunk_content_fts(temp_db_path, metadata, content, expected_conte
|
||||||
"""content_fts holds the contextualized content while content stays raw."""
|
"""content_fts holds the contextualized content while content stays raw."""
|
||||||
from haiku.rag.embeddings import get_embedder
|
from haiku.rag.embeddings import get_embedder
|
||||||
|
|
||||||
async with HaikuRAG(db_path=temp_db_path, config=Config, create=True) as client:
|
async with HaikuRAG(
|
||||||
|
db_path=temp_db_path, config=get_config(), create=True
|
||||||
|
) as client:
|
||||||
chunk = Chunk(
|
chunk = Chunk(
|
||||||
document_id="test-doc",
|
document_id="test-doc",
|
||||||
content=content,
|
content=content,
|
||||||
|
|
@ -451,7 +457,7 @@ async def test_chunk_content_fts(temp_db_path, metadata, content, expected_conte
|
||||||
order=0,
|
order=0,
|
||||||
)
|
)
|
||||||
|
|
||||||
embedder = get_embedder(Config)
|
embedder = get_embedder(get_config())
|
||||||
embedding = (await embedder.embed_documents([chunk.content]))[0]
|
embedding = (await embedder.embed_documents([chunk.content]))[0]
|
||||||
chunk.embedding = embedding
|
chunk.embedding = embedding
|
||||||
|
|
||||||
|
|
@ -477,7 +483,9 @@ async def test_ensure_fts_index_warns_on_failure(temp_db_path):
|
||||||
|
|
||||||
from haiku.rag.store.repositories import chunk as chunk_module
|
from haiku.rag.store.repositories import chunk as chunk_module
|
||||||
|
|
||||||
async with HaikuRAG(db_path=temp_db_path, config=Config, create=True) as client:
|
async with HaikuRAG(
|
||||||
|
db_path=temp_db_path, config=get_config(), create=True
|
||||||
|
) as client:
|
||||||
repo = client.chunk_repository
|
repo = client.chunk_repository
|
||||||
|
|
||||||
async def _boom(*_args, **_kwargs):
|
async def _boom(*_args, **_kwargs):
|
||||||
|
|
@ -497,7 +505,9 @@ async def test_chunk_repository_get_by_id_and_list_all_pagination(
|
||||||
qa_corpus: list[dict[str, str]], temp_db_path
|
qa_corpus: list[dict[str, str]], temp_db_path
|
||||||
):
|
):
|
||||||
"""get_by_id resolves a stored chunk; list_all honours limit and offset."""
|
"""get_by_id resolves a stored chunk; list_all honours limit and offset."""
|
||||||
async with HaikuRAG(db_path=temp_db_path, config=Config, create=True) as client:
|
async with HaikuRAG(
|
||||||
|
db_path=temp_db_path, config=get_config(), create=True
|
||||||
|
) as client:
|
||||||
# A corpus document is long enough to chunk more than once, which is
|
# A corpus document is long enough to chunk more than once, which is
|
||||||
# what makes the offset assertion below meaningful.
|
# what makes the offset assertion below meaningful.
|
||||||
doc = await client.create_document(content=qa_corpus[0]["document_extracted"])
|
doc = await client.create_document(content=qa_corpus[0]["document_extracted"])
|
||||||
|
|
@ -531,7 +541,9 @@ async def test_chunk_repository_get_by_id_and_list_all_pagination(
|
||||||
@pytest.mark.vcr()
|
@pytest.mark.vcr()
|
||||||
async def test_chunk_search_returns_empty_for_blank_query(temp_db_path):
|
async def test_chunk_search_returns_empty_for_blank_query(temp_db_path):
|
||||||
"""A blank query with no precomputed vector short-circuits before searching."""
|
"""A blank query with no precomputed vector short-circuits before searching."""
|
||||||
async with HaikuRAG(db_path=temp_db_path, config=Config, create=True) as client:
|
async with HaikuRAG(
|
||||||
|
db_path=temp_db_path, config=get_config(), create=True
|
||||||
|
) as client:
|
||||||
await client.create_document(content="Searchable body about elections.")
|
await client.create_document(content="Searchable body about elections.")
|
||||||
|
|
||||||
# Positive control: the corpus is non-empty, so [] is a real decision
|
# Positive control: the corpus is non-empty, so [] is a real decision
|
||||||
|
|
@ -543,7 +555,9 @@ async def test_chunk_search_returns_empty_for_blank_query(temp_db_path):
|
||||||
@pytest.mark.vcr()
|
@pytest.mark.vcr()
|
||||||
async def test_chunk_search_with_precomputed_vector_skips_text_query(temp_db_path):
|
async def test_chunk_search_with_precomputed_vector_skips_text_query(temp_db_path):
|
||||||
"""The image-as-query path searches vector-only using a stored embedding."""
|
"""The image-as-query path searches vector-only using a stored embedding."""
|
||||||
async with HaikuRAG(db_path=temp_db_path, config=Config, create=True) as client:
|
async with HaikuRAG(
|
||||||
|
db_path=temp_db_path, config=get_config(), create=True
|
||||||
|
) as client:
|
||||||
doc = await client.create_document(content="Vector-only search target.")
|
doc = await client.create_document(content="Vector-only search target.")
|
||||||
assert doc.id is not None
|
assert doc.id is not None
|
||||||
|
|
||||||
|
|
@ -557,7 +571,9 @@ async def test_chunk_search_with_precomputed_vector_skips_text_query(temp_db_pat
|
||||||
|
|
||||||
|
|
||||||
async def test_get_chunk_ids_by_self_ref_grouped_without_documents(temp_db_path):
|
async def test_get_chunk_ids_by_self_ref_grouped_without_documents(temp_db_path):
|
||||||
async with HaikuRAG(db_path=temp_db_path, config=Config, create=True) as client:
|
async with HaikuRAG(
|
||||||
|
db_path=temp_db_path, config=get_config(), create=True
|
||||||
|
) as client:
|
||||||
assert await client.chunk_repository.get_chunk_ids_by_self_ref_grouped([]) == {}
|
assert await client.chunk_repository.get_chunk_ids_by_self_ref_grouped([]) == {}
|
||||||
|
|
||||||
|
|
||||||
|
|
@ -565,7 +581,9 @@ async def test_process_search_results_rejects_unknown_score_column(temp_db_path)
|
||||||
"""A result frame with no recognised score column is a programming error."""
|
"""A result frame with no recognised score column is a programming error."""
|
||||||
import pandas as pd
|
import pandas as pd
|
||||||
|
|
||||||
async with HaikuRAG(db_path=temp_db_path, config=Config, create=True) as client:
|
async with HaikuRAG(
|
||||||
|
db_path=temp_db_path, config=get_config(), create=True
|
||||||
|
) as client:
|
||||||
|
|
||||||
class _Frame:
|
class _Frame:
|
||||||
async def to_pandas(self):
|
async def to_pandas(self):
|
||||||
|
|
|
||||||
|
|
@ -7,7 +7,7 @@ from transformers import AutoTokenizer
|
||||||
from haiku.rag.chunkers import get_chunker
|
from haiku.rag.chunkers import get_chunker
|
||||||
from haiku.rag.chunkers.docling_local import DoclingLocalChunker
|
from haiku.rag.chunkers.docling_local import DoclingLocalChunker
|
||||||
from haiku.rag.chunkers.docling_serve import DoclingServeChunker
|
from haiku.rag.chunkers.docling_serve import DoclingServeChunker
|
||||||
from haiku.rag.config import AppConfig, Config
|
from haiku.rag.config import AppConfig, get_config
|
||||||
from haiku.rag.converters import get_converter
|
from haiku.rag.converters import get_converter
|
||||||
|
|
||||||
|
|
||||||
|
|
@ -23,7 +23,7 @@ async def test_local_chunker(qa_corpus: list[dict[str, str]]):
|
||||||
doc_text = qa_corpus[0]["document_extracted"]
|
doc_text = qa_corpus[0]["document_extracted"]
|
||||||
|
|
||||||
# Convert text to DoclingDocument
|
# Convert text to DoclingDocument
|
||||||
converter = get_converter(Config)
|
converter = get_converter(get_config())
|
||||||
doc = await converter.convert_text(doc_text, name="test.md")
|
doc = await converter.convert_text(doc_text, name="test.md")
|
||||||
|
|
||||||
chunks = await chunker.chunk(doc)
|
chunks = await chunker.chunk(doc)
|
||||||
|
|
@ -79,7 +79,7 @@ async def test_local_chunker_runs_off_event_loop_thread():
|
||||||
called_from.append(threading.current_thread())
|
called_from.append(threading.current_thread())
|
||||||
return original(document)
|
return original(document)
|
||||||
|
|
||||||
converter = get_converter(Config)
|
converter = get_converter(get_config())
|
||||||
doc = await converter.convert_text("# Hello\n\nWorld", name="test.md")
|
doc = await converter.convert_text("# Hello\n\nWorld", name="test.md")
|
||||||
|
|
||||||
with patch.object(DoclingLocalChunker, "_chunk_sync", recording_chunk_sync):
|
with patch.object(DoclingLocalChunker, "_chunk_sync", recording_chunk_sync):
|
||||||
|
|
@ -149,7 +149,7 @@ async def test_local_chunker_hierarchical(qa_corpus: list[dict[str, str]]):
|
||||||
chunker = DoclingLocalChunker(config)
|
chunker = DoclingLocalChunker(config)
|
||||||
|
|
||||||
doc_text = qa_corpus[0]["document_extracted"]
|
doc_text = qa_corpus[0]["document_extracted"]
|
||||||
converter = get_converter(Config)
|
converter = get_converter(get_config())
|
||||||
doc = await converter.convert_text(doc_text, name="test.md")
|
doc = await converter.convert_text(doc_text, name="test.md")
|
||||||
|
|
||||||
chunks = await chunker.chunk(doc)
|
chunks = await chunker.chunk(doc)
|
||||||
|
|
@ -180,7 +180,7 @@ async def test_local_chunker_markdown_tables():
|
||||||
| Value D | Value E |
|
| Value D | Value E |
|
||||||
"""
|
"""
|
||||||
|
|
||||||
converter = get_converter(Config)
|
converter = get_converter(get_config())
|
||||||
doc = await converter.convert_text(markdown_with_table, name="test.md")
|
doc = await converter.convert_text(markdown_with_table, name="test.md")
|
||||||
|
|
||||||
# Test with markdown tables enabled
|
# Test with markdown tables enabled
|
||||||
|
|
@ -222,7 +222,7 @@ Second paragraph.
|
||||||
|
|
||||||
Third paragraph.
|
Third paragraph.
|
||||||
"""
|
"""
|
||||||
converter = get_converter(Config)
|
converter = get_converter(get_config())
|
||||||
doc = await converter.convert_text(sample_md, name="test.md")
|
doc = await converter.convert_text(sample_md, name="test.md")
|
||||||
|
|
||||||
chunker = DoclingLocalChunker()
|
chunker = DoclingLocalChunker()
|
||||||
|
|
@ -249,7 +249,7 @@ Here is some background information.
|
||||||
|----------|----------|
|
|----------|----------|
|
||||||
| Value 1 | Value 2 |
|
| Value 1 | Value 2 |
|
||||||
"""
|
"""
|
||||||
converter = get_converter(Config)
|
converter = get_converter(get_config())
|
||||||
doc = await converter.convert_text(sample_md, name="test.md")
|
doc = await converter.convert_text(sample_md, name="test.md")
|
||||||
|
|
||||||
chunker = DoclingLocalChunker()
|
chunker = DoclingLocalChunker()
|
||||||
|
|
@ -347,7 +347,7 @@ class TestDoclingServeChunker:
|
||||||
mock_client_class.return_value.__aenter__.return_value = mock_client
|
mock_client_class.return_value.__aenter__.return_value = mock_client
|
||||||
|
|
||||||
# Create a simple document
|
# Create a simple document
|
||||||
converter = get_converter(Config)
|
converter = get_converter(get_config())
|
||||||
doc = await converter.convert_text("# Test\n\nContent", name="test.md")
|
doc = await converter.convert_text("# Test\n\nContent", name="test.md")
|
||||||
|
|
||||||
chunks = await chunker.chunk(doc)
|
chunks = await chunker.chunk(doc)
|
||||||
|
|
@ -371,7 +371,7 @@ class TestDoclingServeChunker:
|
||||||
mock_client.get = AsyncMock(side_effect=[poll_resp, result_resp])
|
mock_client.get = AsyncMock(side_effect=[poll_resp, result_resp])
|
||||||
mock_client_class.return_value.__aenter__.return_value = mock_client
|
mock_client_class.return_value.__aenter__.return_value = mock_client
|
||||||
|
|
||||||
converter = get_converter(Config)
|
converter = get_converter(get_config())
|
||||||
doc = await converter.convert_text("# Test", name="test.md")
|
doc = await converter.convert_text("# Test", name="test.md")
|
||||||
await chunker.chunk(doc)
|
await chunker.chunk(doc)
|
||||||
|
|
||||||
|
|
@ -394,7 +394,7 @@ class TestDoclingServeChunker:
|
||||||
mock_client.get = AsyncMock(side_effect=[poll_resp, result_resp])
|
mock_client.get = AsyncMock(side_effect=[poll_resp, result_resp])
|
||||||
mock_client_class.return_value.__aenter__.return_value = mock_client
|
mock_client_class.return_value.__aenter__.return_value = mock_client
|
||||||
|
|
||||||
converter = get_converter(Config)
|
converter = get_converter(get_config())
|
||||||
doc = await converter.convert_text("# Test", name="test.md")
|
doc = await converter.convert_text("# Test", name="test.md")
|
||||||
await chunker.chunk(doc)
|
await chunker.chunk(doc)
|
||||||
|
|
||||||
|
|
@ -422,7 +422,7 @@ class TestDoclingServeChunker:
|
||||||
mock_client.get = AsyncMock(side_effect=[poll_resp, result_resp])
|
mock_client.get = AsyncMock(side_effect=[poll_resp, result_resp])
|
||||||
mock_client_class.return_value.__aenter__.return_value = mock_client
|
mock_client_class.return_value.__aenter__.return_value = mock_client
|
||||||
|
|
||||||
converter = get_converter(Config)
|
converter = get_converter(get_config())
|
||||||
doc = await converter.convert_text("# Test", name="test.md")
|
doc = await converter.convert_text("# Test", name="test.md")
|
||||||
await chunker.chunk(doc)
|
await chunker.chunk(doc)
|
||||||
|
|
||||||
|
|
@ -453,7 +453,7 @@ class TestDoclingServeChunker:
|
||||||
mock_client.get = AsyncMock(side_effect=[poll_resp, result_resp])
|
mock_client.get = AsyncMock(side_effect=[poll_resp, result_resp])
|
||||||
mock_client_class.return_value.__aenter__.return_value = mock_client
|
mock_client_class.return_value.__aenter__.return_value = mock_client
|
||||||
|
|
||||||
converter = get_converter(Config)
|
converter = get_converter(get_config())
|
||||||
doc = await converter.convert_text("# Test", name="test.md")
|
doc = await converter.convert_text("# Test", name="test.md")
|
||||||
await chunker.chunk(doc)
|
await chunker.chunk(doc)
|
||||||
|
|
||||||
|
|
@ -476,7 +476,7 @@ class TestDoclingServeChunker:
|
||||||
mock_client.post.side_effect = httpx.ConnectError("Connection failed")
|
mock_client.post.side_effect = httpx.ConnectError("Connection failed")
|
||||||
mock_client_class.return_value.__aenter__.return_value = mock_client
|
mock_client_class.return_value.__aenter__.return_value = mock_client
|
||||||
|
|
||||||
converter = get_converter(Config)
|
converter = get_converter(get_config())
|
||||||
doc = await converter.convert_text("# Test", name="test.md")
|
doc = await converter.convert_text("# Test", name="test.md")
|
||||||
|
|
||||||
with pytest.raises(httpx.ConnectError):
|
with pytest.raises(httpx.ConnectError):
|
||||||
|
|
@ -492,7 +492,7 @@ class TestDoclingServeChunker:
|
||||||
mock_client.post.side_effect = httpx.TimeoutException("Timeout")
|
mock_client.post.side_effect = httpx.TimeoutException("Timeout")
|
||||||
mock_client_class.return_value.__aenter__.return_value = mock_client
|
mock_client_class.return_value.__aenter__.return_value = mock_client
|
||||||
|
|
||||||
converter = get_converter(Config)
|
converter = get_converter(get_config())
|
||||||
doc = await converter.convert_text("# Test", name="test.md")
|
doc = await converter.convert_text("# Test", name="test.md")
|
||||||
|
|
||||||
with pytest.raises(httpx.TimeoutException):
|
with pytest.raises(httpx.TimeoutException):
|
||||||
|
|
@ -517,7 +517,7 @@ class TestDoclingServeChunker:
|
||||||
mock_client.post.return_value = mock_response
|
mock_client.post.return_value = mock_response
|
||||||
mock_client_class.return_value.__aenter__.return_value = mock_client
|
mock_client_class.return_value.__aenter__.return_value = mock_client
|
||||||
|
|
||||||
converter = get_converter(Config)
|
converter = get_converter(get_config())
|
||||||
doc = await converter.convert_text("# Test", name="test.md")
|
doc = await converter.convert_text("# Test", name="test.md")
|
||||||
|
|
||||||
with pytest.raises(httpx.HTTPStatusError) as exc_info:
|
with pytest.raises(httpx.HTTPStatusError) as exc_info:
|
||||||
|
|
@ -545,7 +545,7 @@ class TestDoclingServeChunker:
|
||||||
mock_client.get = AsyncMock(side_effect=[poll_resp, result_resp])
|
mock_client.get = AsyncMock(side_effect=[poll_resp, result_resp])
|
||||||
mock_client_class.return_value.__aenter__.return_value = mock_client
|
mock_client_class.return_value.__aenter__.return_value = mock_client
|
||||||
|
|
||||||
converter = get_converter(Config)
|
converter = get_converter(get_config())
|
||||||
doc = await converter.convert_text("# Test", name="test.md")
|
doc = await converter.convert_text("# Test", name="test.md")
|
||||||
|
|
||||||
with pytest.raises(ValueError, match="Chunking failed"):
|
with pytest.raises(ValueError, match="Chunking failed"):
|
||||||
|
|
@ -574,7 +574,7 @@ class TestDoclingServeChunker:
|
||||||
mock_client.get = AsyncMock(side_effect=[poll_resp, result_resp])
|
mock_client.get = AsyncMock(side_effect=[poll_resp, result_resp])
|
||||||
mock_client_class.return_value.__aenter__.return_value = mock_client
|
mock_client_class.return_value.__aenter__.return_value = mock_client
|
||||||
|
|
||||||
converter = get_converter(Config)
|
converter = get_converter(get_config())
|
||||||
doc = await converter.convert_text("# Test", name="test.md")
|
doc = await converter.convert_text("# Test", name="test.md")
|
||||||
|
|
||||||
chunks = await chunker.chunk(doc)
|
chunks = await chunker.chunk(doc)
|
||||||
|
|
@ -612,7 +612,7 @@ class TestDoclingServeChunker:
|
||||||
mock_client_class.return_value.__aenter__.return_value = mock_client
|
mock_client_class.return_value.__aenter__.return_value = mock_client
|
||||||
|
|
||||||
# Create a document with texts and tables that match the mocked refs
|
# Create a document with texts and tables that match the mocked refs
|
||||||
converter = get_converter(Config)
|
converter = get_converter(get_config())
|
||||||
doc = await converter.convert_text(
|
doc = await converter.convert_text(
|
||||||
"""# Chapter 1
|
"""# Chapter 1
|
||||||
|
|
||||||
|
|
@ -700,7 +700,7 @@ async def test_local_and_serve_chunkers_produce_same_output(doclaynet_first_page
|
||||||
from haiku.rag.converters.docling_serve import DoclingServeConverter
|
from haiku.rag.converters.docling_serve import DoclingServeConverter
|
||||||
|
|
||||||
# Use docling-serve to convert the PDF (ensures same conversion for both chunkers)
|
# Use docling-serve to convert the PDF (ensures same conversion for both chunkers)
|
||||||
converter = DoclingServeConverter(Config)
|
converter = DoclingServeConverter(get_config())
|
||||||
pdf_path = doclaynet_first_page_pdf
|
pdf_path = doclaynet_first_page_pdf
|
||||||
doc = await converter.convert_file(pdf_path)
|
doc = await converter.convert_file(pdf_path)
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -17,7 +17,7 @@ from haiku.rag.client.documents import (
|
||||||
_write_fetch_body,
|
_write_fetch_body,
|
||||||
check_source_accessible,
|
check_source_accessible,
|
||||||
)
|
)
|
||||||
from haiku.rag.config import Config
|
from haiku.rag.config import get_config
|
||||||
from haiku.rag.embeddings import EmbedderWrapper
|
from haiku.rag.embeddings import EmbedderWrapper
|
||||||
from haiku.rag.ingester.sources.base import FetchResult
|
from haiku.rag.ingester.sources.base import FetchResult
|
||||||
from haiku.rag.store.compression import decompress_json
|
from haiku.rag.store.compression import decompress_json
|
||||||
|
|
@ -813,7 +813,7 @@ async def test_client_import_document_with_custom_chunks(temp_db_path):
|
||||||
Chunk(
|
Chunk(
|
||||||
content="This is the second chunk",
|
content="This is the second chunk",
|
||||||
metadata={"custom": "metadata2"},
|
metadata={"custom": "metadata2"},
|
||||||
embedding=[0.1] * Config.embeddings.model.vector_dim,
|
embedding=[0.1] * get_config().embeddings.model.vector_dim,
|
||||||
order=1,
|
order=1,
|
||||||
), # With embedding
|
), # With embedding
|
||||||
Chunk(
|
Chunk(
|
||||||
|
|
@ -856,7 +856,7 @@ def _docling_doc(name: str, text: str):
|
||||||
|
|
||||||
def _import(name: str, text: str, **overrides) -> "DocumentImport":
|
def _import(name: str, text: str, **overrides) -> "DocumentImport":
|
||||||
"""Build a DocumentImport with one pre-embedded chunk (no embedder call)."""
|
"""Build a DocumentImport with one pre-embedded chunk (no embedder call)."""
|
||||||
dim = Config.embeddings.model.vector_dim
|
dim = get_config().embeddings.model.vector_dim
|
||||||
chunk = Chunk(content=text, embedding=[0.1] * dim, order=0)
|
chunk = Chunk(content=text, embedding=[0.1] * dim, order=0)
|
||||||
return DocumentImport(
|
return DocumentImport(
|
||||||
docling_document=_docling_doc(name, text),
|
docling_document=_docling_doc(name, text),
|
||||||
|
|
@ -868,7 +868,7 @@ def _import(name: str, text: str, **overrides) -> "DocumentImport":
|
||||||
async def test_client_import_documents_single_version_per_table(temp_db_path):
|
async def test_client_import_documents_single_version_per_table(temp_db_path):
|
||||||
"""import_documents writes documents/chunks/document_items once for the
|
"""import_documents writes documents/chunks/document_items once for the
|
||||||
whole batch (issue #287)."""
|
whole batch (issue #287)."""
|
||||||
config = Config.model_copy(deep=True)
|
config = get_config().model_copy(deep=True)
|
||||||
config.storage.auto_vacuum = False
|
config.storage.auto_vacuum = False
|
||||||
|
|
||||||
async with HaikuRAG(temp_db_path, config=config, create=True) as client:
|
async with HaikuRAG(temp_db_path, config=config, create=True) as client:
|
||||||
|
|
@ -902,7 +902,7 @@ async def test_client_import_documents_single_version_per_table(temp_db_path):
|
||||||
|
|
||||||
async def test_client_import_documents_rolls_back_on_failure(temp_db_path):
|
async def test_client_import_documents_rolls_back_on_failure(temp_db_path):
|
||||||
"""A failure mid-batch restores all tables: nothing is persisted."""
|
"""A failure mid-batch restores all tables: nothing is persisted."""
|
||||||
dim = Config.embeddings.model.vector_dim
|
dim = get_config().embeddings.model.vector_dim
|
||||||
|
|
||||||
async with HaikuRAG(temp_db_path, create=True) as client:
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
||||||
good = _import("good", "Good document body", uri="mem://good")
|
good = _import("good", "Good document body", uri="mem://good")
|
||||||
|
|
@ -943,7 +943,7 @@ async def test_client_import_documents_batches_embeddings(temp_db_path):
|
||||||
"""Chunks missing embeddings are embedded in one pass across the whole
|
"""Chunks missing embeddings are embedded in one pass across the whole
|
||||||
batch, not one embedder call per document. Duplicate chunk texts across
|
batch, not one embedder call per document. Duplicate chunk texts across
|
||||||
documents keep their per-document embeddings."""
|
documents keep their per-document embeddings."""
|
||||||
dim = Config.embeddings.model.vector_dim
|
dim = get_config().embeddings.model.vector_dim
|
||||||
embedder = _CountingEmbedder(dim)
|
embedder = _CountingEmbedder(dim)
|
||||||
|
|
||||||
async with HaikuRAG(temp_db_path, create=True) as client:
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
||||||
|
|
@ -977,7 +977,7 @@ async def test_client_import_documents_batches_embeddings(temp_db_path):
|
||||||
async def test_client_import_documents_mixed_embeddings(temp_db_path):
|
async def test_client_import_documents_mixed_embeddings(temp_db_path):
|
||||||
"""Pre-embedded chunks keep their vectors; only the unembedded ones go
|
"""Pre-embedded chunks keep their vectors; only the unembedded ones go
|
||||||
through the embedder, in one batch."""
|
through the embedder, in one batch."""
|
||||||
dim = Config.embeddings.model.vector_dim
|
dim = get_config().embeddings.model.vector_dim
|
||||||
embedder = _CountingEmbedder(dim)
|
embedder = _CountingEmbedder(dim)
|
||||||
|
|
||||||
async with HaikuRAG(temp_db_path, create=True) as client:
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
||||||
|
|
@ -1024,8 +1024,8 @@ async def test_client_update_document_replaces_rows_with_bounded_versions(
|
||||||
|
|
||||||
auto_vacuum is off: its writes would land inside the measured window.
|
auto_vacuum is off: its writes would land inside the measured window.
|
||||||
"""
|
"""
|
||||||
dim = Config.embeddings.model.vector_dim
|
dim = get_config().embeddings.model.vector_dim
|
||||||
config = Config.model_copy(deep=True)
|
config = get_config().model_copy(deep=True)
|
||||||
config.storage.auto_vacuum = False
|
config.storage.auto_vacuum = False
|
||||||
|
|
||||||
async with HaikuRAG(temp_db_path, config=config, create=True) as client:
|
async with HaikuRAG(temp_db_path, config=config, create=True) as client:
|
||||||
|
|
@ -1076,8 +1076,8 @@ async def test_metadata_only_update_does_not_advance_documents_table(temp_db_pat
|
||||||
live in document_meta, so the documents table version must stay frozen while
|
live in document_meta, so the documents table version must stay frozen while
|
||||||
only metadata/title change — and reads must still hydrate the full Document.
|
only metadata/title change — and reads must still hydrate the full Document.
|
||||||
"""
|
"""
|
||||||
dim = Config.embeddings.model.vector_dim
|
dim = get_config().embeddings.model.vector_dim
|
||||||
config = Config.model_copy(deep=True)
|
config = get_config().model_copy(deep=True)
|
||||||
config.storage.auto_vacuum = False
|
config.storage.auto_vacuum = False
|
||||||
|
|
||||||
async with HaikuRAG(temp_db_path, config=config, create=True) as client:
|
async with HaikuRAG(temp_db_path, config=config, create=True) as client:
|
||||||
|
|
@ -1121,7 +1121,7 @@ async def test_metadata_only_update_does_not_advance_documents_table(temp_db_pat
|
||||||
async def test_delete_marks_vacuum_dirty(temp_db_path):
|
async def test_delete_marks_vacuum_dirty(temp_db_path):
|
||||||
"""A delete adds tombstone/table versions, so it must enter the auto-vacuum
|
"""A delete adds tombstone/table versions, so it must enter the auto-vacuum
|
||||||
lifecycle — otherwise a delete-only run closes without a final vacuum."""
|
lifecycle — otherwise a delete-only run closes without a final vacuum."""
|
||||||
dim = Config.embeddings.model.vector_dim
|
dim = get_config().embeddings.model.vector_dim
|
||||||
async with HaikuRAG(temp_db_path, create=True) as client:
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
||||||
doc = await client.import_document(
|
doc = await client.import_document(
|
||||||
_docling_doc("d", "body"),
|
_docling_doc("d", "body"),
|
||||||
|
|
@ -1138,8 +1138,8 @@ async def test_delete_marks_vacuum_dirty(temp_db_path):
|
||||||
async def test_delete_rolls_back_on_partial_failure(temp_db_path, monkeypatch):
|
async def test_delete_rolls_back_on_partial_failure(temp_db_path, monkeypatch):
|
||||||
"""A multi-table delete is atomic: if a later table delete fails, the write
|
"""A multi-table delete is atomic: if a later table delete fails, the write
|
||||||
lock + version restore bring every table back, leaving no orphaned rows."""
|
lock + version restore bring every table back, leaving no orphaned rows."""
|
||||||
dim = Config.embeddings.model.vector_dim
|
dim = get_config().embeddings.model.vector_dim
|
||||||
config = Config.model_copy(deep=True)
|
config = get_config().model_copy(deep=True)
|
||||||
config.storage.auto_vacuum = False
|
config.storage.auto_vacuum = False
|
||||||
|
|
||||||
async with HaikuRAG(temp_db_path, config=config, create=True) as client:
|
async with HaikuRAG(temp_db_path, config=config, create=True) as client:
|
||||||
|
|
@ -1175,8 +1175,8 @@ async def test_cascade_delete_is_atomic(temp_db_path, monkeypatch):
|
||||||
"""Deleting a parent cascades to children under one lock + snapshot. If any
|
"""Deleting a parent cascades to children under one lock + snapshot. If any
|
||||||
delete in the subtree fails, the whole subtree is restored — a child isn't
|
delete in the subtree fails, the whole subtree is restored — a child isn't
|
||||||
left deleted while its parent survives."""
|
left deleted while its parent survives."""
|
||||||
dim = Config.embeddings.model.vector_dim
|
dim = get_config().embeddings.model.vector_dim
|
||||||
config = Config.model_copy(deep=True)
|
config = get_config().model_copy(deep=True)
|
||||||
config.storage.auto_vacuum = False
|
config.storage.auto_vacuum = False
|
||||||
|
|
||||||
async with HaikuRAG(temp_db_path, config=config, create=True) as client:
|
async with HaikuRAG(temp_db_path, config=config, create=True) as client:
|
||||||
|
|
@ -2390,7 +2390,7 @@ async def test_metadata_only_update_waits_for_write_lock(temp_db_path):
|
||||||
create_tag's version snapshot and its per-table tag creation)."""
|
create_tag's version snapshot and its per-table tag creation)."""
|
||||||
import asyncio
|
import asyncio
|
||||||
|
|
||||||
dim = Config.embeddings.model.vector_dim
|
dim = get_config().embeddings.model.vector_dim
|
||||||
docling_doc = DoclingDocument(name="d")
|
docling_doc = DoclingDocument(name="d")
|
||||||
docling_doc.add_text(label=DocItemLabel.TEXT, text="body")
|
docling_doc.add_text(label=DocItemLabel.TEXT, text="body")
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -3,7 +3,7 @@ from pathlib import Path
|
||||||
import pytest
|
import pytest
|
||||||
import yaml
|
import yaml
|
||||||
|
|
||||||
from haiku.rag.config import AppConfig
|
from haiku.rag.config import AppConfig, set_config
|
||||||
from haiku.rag.config.loader import (
|
from haiku.rag.config.loader import (
|
||||||
find_config_file,
|
find_config_file,
|
||||||
generate_default_config,
|
generate_default_config,
|
||||||
|
|
@ -577,3 +577,48 @@ def test_example_configs_are_present():
|
||||||
)
|
)
|
||||||
def test_example_config_validates(path: Path):
|
def test_example_config_validates(path: Path):
|
||||||
AppConfig.model_validate(yaml.safe_load(path.read_text()) or {})
|
AppConfig.model_validate(yaml.safe_load(path.read_text()) or {})
|
||||||
|
|
||||||
|
|
||||||
|
def _use_config(monkeypatch, cfg):
|
||||||
|
"""Install cfg as the global config, restored by monkeypatch teardown."""
|
||||||
|
import haiku.rag.config as config_module
|
||||||
|
|
||||||
|
monkeypatch.setattr(config_module, "_config", None)
|
||||||
|
set_config(cfg)
|
||||||
|
|
||||||
|
|
||||||
|
def test_set_config_reaches_the_factories(monkeypatch, tmp_path):
|
||||||
|
"""A config installed after import must reach every factory and
|
||||||
|
constructor that resolves the global config itself."""
|
||||||
|
from haiku.rag.chunkers import get_chunker
|
||||||
|
from haiku.rag.chunkers.docling_serve import DoclingServeChunker
|
||||||
|
from haiku.rag.client import HaikuRAG
|
||||||
|
from haiku.rag.config.models import ModelConfig, RerankingConfig
|
||||||
|
from haiku.rag.converters import get_converter
|
||||||
|
from haiku.rag.converters.docling_serve import DoclingServeConverter
|
||||||
|
from haiku.rag.embeddings import get_embedder
|
||||||
|
from haiku.rag.reranking import get_reranker
|
||||||
|
from haiku.rag.reranking.vllm import VLLMReranker
|
||||||
|
from haiku.rag.store.engine import Store
|
||||||
|
|
||||||
|
cfg = AppConfig()
|
||||||
|
cfg.processing.converter = "docling-serve"
|
||||||
|
cfg.processing.chunker = "docling-serve"
|
||||||
|
cfg.embeddings.model.vector_dim = 7
|
||||||
|
cfg.reranking.model = ModelConfig(
|
||||||
|
provider="vllm", name="reranker-x", base_url="http://localhost:9/v1"
|
||||||
|
)
|
||||||
|
assert isinstance(cfg.reranking, RerankingConfig)
|
||||||
|
|
||||||
|
_use_config(monkeypatch, cfg)
|
||||||
|
|
||||||
|
assert isinstance(get_converter(), DoclingServeConverter)
|
||||||
|
assert isinstance(get_chunker(), DoclingServeChunker)
|
||||||
|
assert get_embedder().vector_dim == 7
|
||||||
|
|
||||||
|
reranker = get_reranker()
|
||||||
|
assert isinstance(reranker, VLLMReranker)
|
||||||
|
assert reranker._model == "reranker-x"
|
||||||
|
|
||||||
|
assert HaikuRAG(tmp_path / "db")._config is cfg
|
||||||
|
assert Store(tmp_path / "db", create=True)._config is cfg
|
||||||
|
|
|
||||||
|
|
@ -69,7 +69,7 @@ def create_code_document() -> DoclingDocument:
|
||||||
# Multiple adjacent code blocks - type-aware expansion should group them
|
# Multiple adjacent code blocks - type-aware expansion should group them
|
||||||
doc.add_text(label=DocItemLabel.CODE, text="# Part 1: Setup\nimport os\nimport sys")
|
doc.add_text(label=DocItemLabel.CODE, text="# Part 1: Setup\nimport os\nimport sys")
|
||||||
doc.add_text(
|
doc.add_text(
|
||||||
label=DocItemLabel.CODE, text='# Part 2: Config\nCONFIG = {"debug": True}'
|
label=DocItemLabel.CODE, text='# Part 2: get_config()\nCONFIG = {"debug": True}'
|
||||||
)
|
)
|
||||||
doc.add_text(
|
doc.add_text(
|
||||||
label=DocItemLabel.CODE, text="# Part 3: Main\ndef main():\n print(CONFIG)"
|
label=DocItemLabel.CODE, text="# Part 3: Main\ndef main():\n print(CONFIG)"
|
||||||
|
|
@ -81,7 +81,7 @@ def create_code_document() -> DoclingDocument:
|
||||||
|
|
||||||
@pytest.fixture
|
@pytest.fixture
|
||||||
def small_chunk_config() -> AppConfig:
|
def small_chunk_config() -> AppConfig:
|
||||||
"""Config with small chunk size to force splitting."""
|
"""get_config() with small chunk size to force splitting."""
|
||||||
config = AppConfig()
|
config = AppConfig()
|
||||||
config.processing.chunk_size = 32
|
config.processing.chunk_size = 32
|
||||||
config.search.max_context_chars = 10000
|
config.search.max_context_chars = 10000
|
||||||
|
|
|
||||||
|
|
@ -5,7 +5,7 @@ import httpx
|
||||||
import pytest
|
import pytest
|
||||||
|
|
||||||
from haiku.rag.client.downloads import download_models
|
from haiku.rag.client.downloads import download_models
|
||||||
from haiku.rag.config import Config
|
from haiku.rag.config import get_config
|
||||||
from haiku.rag.config.models import ModelConfig
|
from haiku.rag.config.models import ModelConfig
|
||||||
|
|
||||||
|
|
||||||
|
|
@ -39,7 +39,7 @@ async def test_download_models_ollama_connect_error(mock_to_thread):
|
||||||
with pytest.raises(
|
with pytest.raises(
|
||||||
ConnectionError, match="Cannot connect to Ollama"
|
ConnectionError, match="Cannot connect to Ollama"
|
||||||
) as exc_info:
|
) as exc_info:
|
||||||
async for _ in download_models(Config):
|
async for _ in download_models(get_config()):
|
||||||
pass
|
pass
|
||||||
|
|
||||||
assert "ollama serve" in str(exc_info.value)
|
assert "ollama serve" in str(exc_info.value)
|
||||||
|
|
@ -74,7 +74,7 @@ async def test_download_models_ollama_pulls_models(mock_to_thread):
|
||||||
return_value=_mock_httpx_client(mock_stream),
|
return_value=_mock_httpx_client(mock_stream),
|
||||||
):
|
):
|
||||||
events = []
|
events = []
|
||||||
async for progress in download_models(Config):
|
async for progress in download_models(get_config()):
|
||||||
events.append(progress)
|
events.append(progress)
|
||||||
|
|
||||||
# Default config has embeddings=qwen3-embedding:4b, qa=gpt-oss
|
# Default config has embeddings=qwen3-embedding:4b, qa=gpt-oss
|
||||||
|
|
|
||||||
|
|
@ -5,9 +5,9 @@ import pytest
|
||||||
|
|
||||||
from haiku.rag.config import (
|
from haiku.rag.config import (
|
||||||
AppConfig,
|
AppConfig,
|
||||||
Config,
|
|
||||||
EmbeddingModelConfig,
|
EmbeddingModelConfig,
|
||||||
EmbeddingsConfig,
|
EmbeddingsConfig,
|
||||||
|
get_config,
|
||||||
)
|
)
|
||||||
from haiku.rag.embeddings import (
|
from haiku.rag.embeddings import (
|
||||||
EmbedderWrapper,
|
EmbedderWrapper,
|
||||||
|
|
@ -197,7 +197,7 @@ async def test_embed_chunks_basic(allow_model_requests):
|
||||||
),
|
),
|
||||||
]
|
]
|
||||||
|
|
||||||
embedded_chunks = await embed_chunks(chunks, get_embedder(Config))
|
embedded_chunks = await embed_chunks(chunks, get_embedder(get_config()))
|
||||||
|
|
||||||
assert len(embedded_chunks) == 2
|
assert len(embedded_chunks) == 2
|
||||||
# Check that all original fields are preserved
|
# Check that all original fields are preserved
|
||||||
|
|
@ -216,7 +216,7 @@ async def test_embed_chunks_basic(allow_model_requests):
|
||||||
async def test_embed_chunks_returns_new_objects(allow_model_requests):
|
async def test_embed_chunks_returns_new_objects(allow_model_requests):
|
||||||
"""Test that embed_chunks returns new Chunk objects (immutable pattern)."""
|
"""Test that embed_chunks returns new Chunk objects (immutable pattern)."""
|
||||||
original = Chunk(id="orig", content="Test content.")
|
original = Chunk(id="orig", content="Test content.")
|
||||||
embedded = await embed_chunks([original], get_embedder(Config))
|
embedded = await embed_chunks([original], get_embedder(get_config()))
|
||||||
|
|
||||||
# Original should be unchanged
|
# Original should be unchanged
|
||||||
assert original.embedding is None
|
assert original.embedding is None
|
||||||
|
|
@ -302,7 +302,7 @@ async def test_embed_chunks_preserves_all_fields(allow_model_requests):
|
||||||
document_meta={"author": "Test"},
|
document_meta={"author": "Test"},
|
||||||
)
|
)
|
||||||
|
|
||||||
embedded = await embed_chunks([chunk], get_embedder(Config))
|
embedded = await embed_chunks([chunk], get_embedder(get_config()))
|
||||||
|
|
||||||
assert embedded[0].id == "test-id"
|
assert embedded[0].id == "test-id"
|
||||||
assert embedded[0].document_id == "doc-id"
|
assert embedded[0].document_id == "doc-id"
|
||||||
|
|
|
||||||
|
|
@ -4,7 +4,7 @@ from unittest.mock import AsyncMock, patch
|
||||||
import pytest
|
import pytest
|
||||||
from pydantic import ValidationError
|
from pydantic import ValidationError
|
||||||
|
|
||||||
from haiku.rag.config import Config
|
from haiku.rag.config import get_config
|
||||||
from haiku.rag.config.models import AppConfig, LanceDBConfig
|
from haiku.rag.config.models import AppConfig, LanceDBConfig
|
||||||
from haiku.rag.store.engine import ConnectionMode, Store, connect_lancedb
|
from haiku.rag.store.engine import ConnectionMode, Store, connect_lancedb
|
||||||
|
|
||||||
|
|
@ -139,16 +139,16 @@ class TestStoreConnectionMode:
|
||||||
async def test_store_connection_mode_cloud(self, temp_db_path):
|
async def test_store_connection_mode_cloud(self, temp_db_path):
|
||||||
async with Store(temp_db_path, create=True) as store:
|
async with Store(temp_db_path, create=True) as store:
|
||||||
with (
|
with (
|
||||||
patch.object(Config.lancedb, "uri", "db://test-database"),
|
patch.object(get_config().lancedb, "uri", "db://test-database"),
|
||||||
patch.object(Config.lancedb, "api_key", "test-api-key"),
|
patch.object(get_config().lancedb, "api_key", "test-api-key"),
|
||||||
patch.object(Config.lancedb, "region", "us-east-1"),
|
patch.object(get_config().lancedb, "region", "us-east-1"),
|
||||||
):
|
):
|
||||||
assert store._connection_mode == ConnectionMode.CLOUD
|
assert store._connection_mode == ConnectionMode.CLOUD
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_store_connection_mode_object_storage(self, temp_db_path):
|
async def test_store_connection_mode_object_storage(self, temp_db_path):
|
||||||
async with Store(temp_db_path, create=True) as store:
|
async with Store(temp_db_path, create=True) as store:
|
||||||
with patch.object(Config.lancedb, "uri", "s3://bucket/path"):
|
with patch.object(get_config().lancedb, "uri", "s3://bucket/path"):
|
||||||
assert store._connection_mode == ConnectionMode.OBJECT_STORAGE
|
assert store._connection_mode == ConnectionMode.OBJECT_STORAGE
|
||||||
|
|
||||||
|
|
||||||
|
|
@ -157,9 +157,9 @@ class TestVacuumByConnectionMode:
|
||||||
async def test_cloud_skips_vacuum(self, temp_db_path):
|
async def test_cloud_skips_vacuum(self, temp_db_path):
|
||||||
async with Store(temp_db_path, create=True) as store:
|
async with Store(temp_db_path, create=True) as store:
|
||||||
with (
|
with (
|
||||||
patch.object(Config.lancedb, "uri", "db://test-database"),
|
patch.object(get_config().lancedb, "uri", "db://test-database"),
|
||||||
patch.object(Config.lancedb, "api_key", "test-api-key"),
|
patch.object(get_config().lancedb, "api_key", "test-api-key"),
|
||||||
patch.object(Config.lancedb, "region", "us-east-1"),
|
patch.object(get_config().lancedb, "region", "us-east-1"),
|
||||||
):
|
):
|
||||||
with patch.object(
|
with patch.object(
|
||||||
store.chunks_table, "optimize", new_callable=AsyncMock
|
store.chunks_table, "optimize", new_callable=AsyncMock
|
||||||
|
|
@ -170,7 +170,7 @@ class TestVacuumByConnectionMode:
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_object_storage_runs_vacuum(self, temp_db_path):
|
async def test_object_storage_runs_vacuum(self, temp_db_path):
|
||||||
async with Store(temp_db_path, create=True) as store:
|
async with Store(temp_db_path, create=True) as store:
|
||||||
with patch.object(Config.lancedb, "uri", "s3://bucket/path"):
|
with patch.object(get_config().lancedb, "uri", "s3://bucket/path"):
|
||||||
with patch.object(
|
with patch.object(
|
||||||
store.chunks_table, "optimize", new_callable=AsyncMock
|
store.chunks_table, "optimize", new_callable=AsyncMock
|
||||||
) as mock_optimize:
|
) as mock_optimize:
|
||||||
|
|
@ -180,7 +180,7 @@ class TestVacuumByConnectionMode:
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_local_runs_vacuum(self, temp_db_path):
|
async def test_local_runs_vacuum(self, temp_db_path):
|
||||||
async with Store(temp_db_path, create=True) as store:
|
async with Store(temp_db_path, create=True) as store:
|
||||||
with patch.object(Config.lancedb, "uri", ""):
|
with patch.object(get_config().lancedb, "uri", ""):
|
||||||
with patch.object(
|
with patch.object(
|
||||||
store.chunks_table, "optimize", new_callable=AsyncMock
|
store.chunks_table, "optimize", new_callable=AsyncMock
|
||||||
) as mock_optimize:
|
) as mock_optimize:
|
||||||
|
|
@ -193,9 +193,9 @@ class TestVectorIndexByConnectionMode:
|
||||||
async def test_cloud_skips_index_creation(self, temp_db_path):
|
async def test_cloud_skips_index_creation(self, temp_db_path):
|
||||||
async with Store(temp_db_path, create=True) as store:
|
async with Store(temp_db_path, create=True) as store:
|
||||||
with (
|
with (
|
||||||
patch.object(Config.lancedb, "uri", "db://test-database"),
|
patch.object(get_config().lancedb, "uri", "db://test-database"),
|
||||||
patch.object(Config.lancedb, "api_key", "test-api-key"),
|
patch.object(get_config().lancedb, "api_key", "test-api-key"),
|
||||||
patch.object(Config.lancedb, "region", "us-east-1"),
|
patch.object(get_config().lancedb, "region", "us-east-1"),
|
||||||
):
|
):
|
||||||
with patch.object(
|
with patch.object(
|
||||||
store.chunks_table, "count_rows", new_callable=AsyncMock
|
store.chunks_table, "count_rows", new_callable=AsyncMock
|
||||||
|
|
@ -206,7 +206,7 @@ class TestVectorIndexByConnectionMode:
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_object_storage_runs_index_creation(self, temp_db_path):
|
async def test_object_storage_runs_index_creation(self, temp_db_path):
|
||||||
async with Store(temp_db_path, create=True) as store:
|
async with Store(temp_db_path, create=True) as store:
|
||||||
with patch.object(Config.lancedb, "uri", "s3://bucket/path"):
|
with patch.object(get_config().lancedb, "uri", "s3://bucket/path"):
|
||||||
with patch.object(
|
with patch.object(
|
||||||
store.chunks_table,
|
store.chunks_table,
|
||||||
"count_rows",
|
"count_rows",
|
||||||
|
|
|
||||||
|
|
@ -644,10 +644,10 @@ class TestMCPClientLifetime:
|
||||||
"""Validation is unchanged: same-dimension identity drift warns in
|
"""Validation is unchanged: same-dimension identity drift warns in
|
||||||
read-only mode and raises in writable mode. The MCP server no longer
|
read-only mode and raises in writable mode. The MCP server no longer
|
||||||
opts out of it for deletion."""
|
opts out of it for deletion."""
|
||||||
from haiku.rag.config import Config
|
from haiku.rag.config import get_config
|
||||||
from haiku.rag.store.repositories.settings import ConfigMismatchError
|
from haiku.rag.store.repositories.settings import ConfigMismatchError
|
||||||
|
|
||||||
drifted = Config.model_copy(deep=True)
|
drifted = get_config().model_copy(deep=True)
|
||||||
drifted.embeddings.model.name = "a-different-model"
|
drifted.embeddings.model.name = "a-different-model"
|
||||||
|
|
||||||
async with create_mcp_server(
|
async with create_mcp_server(
|
||||||
|
|
|
||||||
|
|
@ -15,7 +15,7 @@ from pydantic_ai.usage import RunUsage
|
||||||
from haiku.rag.capabilities.rag import RAGState, create_capability
|
from haiku.rag.capabilities.rag import RAGState, create_capability
|
||||||
from haiku.rag.client import HaikuRAG
|
from haiku.rag.client import HaikuRAG
|
||||||
from haiku.rag.client.search import _populate_image_data
|
from haiku.rag.client.search import _populate_image_data
|
||||||
from haiku.rag.config import AppConfig, Config
|
from haiku.rag.config import AppConfig, get_config
|
||||||
from haiku.rag.store.models.chunk import Chunk, SearchResult
|
from haiku.rag.store.models.chunk import Chunk, SearchResult
|
||||||
from haiku.rag.store.models.document_item import DocumentItem
|
from haiku.rag.store.models.document_item import DocumentItem
|
||||||
from haiku.rag.tools.search import create_search_toolset
|
from haiku.rag.tools.search import create_search_toolset
|
||||||
|
|
@ -980,7 +980,7 @@ async def test_search_tool_returns_plain_string_when_no_pictures():
|
||||||
fake_client.search = AsyncMock(return_value=[text_result])
|
fake_client.search = AsyncMock(return_value=[text_result])
|
||||||
fake_client.expand_context = AsyncMock(return_value=[text_result])
|
fake_client.expand_context = AsyncMock(return_value=[text_result])
|
||||||
|
|
||||||
toolset = create_search_toolset(Config, expand_context=False)
|
toolset = create_search_toolset(get_config(), expand_context=False)
|
||||||
func = toolset.tools["search"].function
|
func = toolset.tools["search"].function
|
||||||
|
|
||||||
ctx = RunContext(
|
ctx = RunContext(
|
||||||
|
|
|
||||||
|
|
@ -3,7 +3,7 @@ from pathlib import Path
|
||||||
|
|
||||||
import pytest
|
import pytest
|
||||||
|
|
||||||
from haiku.rag.config import Config
|
from haiku.rag.config import get_config
|
||||||
from haiku.rag.converters import get_converter
|
from haiku.rag.converters import get_converter
|
||||||
|
|
||||||
|
|
||||||
|
|
@ -19,7 +19,7 @@ async def test_code_file_wrapped_in_code_block():
|
||||||
f.flush()
|
f.flush()
|
||||||
temp_path = Path(f.name)
|
temp_path = Path(f.name)
|
||||||
|
|
||||||
converter = get_converter(Config)
|
converter = get_converter(get_config())
|
||||||
document = await converter.convert_file(temp_path)
|
document = await converter.convert_file(temp_path)
|
||||||
result = document.export_to_markdown()
|
result = document.export_to_markdown()
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -6,7 +6,7 @@ from typing import TypedDict
|
||||||
import pytest
|
import pytest
|
||||||
|
|
||||||
from haiku.rag.client import HaikuRAG, RebuildMode
|
from haiku.rag.client import HaikuRAG, RebuildMode
|
||||||
from haiku.rag.config import Config
|
from haiku.rag.config import get_config
|
||||||
from tests.conftest import capture_logs
|
from tests.conftest import capture_logs
|
||||||
|
|
||||||
|
|
||||||
|
|
@ -194,7 +194,7 @@ async def test_rebuild_resumes_phase2_from_staging_after_crash(
|
||||||
# crash, where no background vacuum would be in flight. Leaving it on lets
|
# crash, where no background vacuum would be in flight. Leaving it on lets
|
||||||
# create_document's scheduled optimize race the raw drop_table ("Directory
|
# create_document's scheduled optimize race the raw drop_table ("Directory
|
||||||
# not empty").
|
# not empty").
|
||||||
config = Config.model_copy(deep=True)
|
config = get_config().model_copy(deep=True)
|
||||||
config.storage.auto_vacuum = False
|
config.storage.auto_vacuum = False
|
||||||
|
|
||||||
async with HaikuRAG(temp_db_path, config=config, create=True) as client:
|
async with HaikuRAG(temp_db_path, config=config, create=True) as client:
|
||||||
|
|
@ -1093,7 +1093,7 @@ async def test_rebuild_blocks_tag_operations(temp_db_path, monkeypatch):
|
||||||
|
|
||||||
monkeypatch.setattr(EmbedderWrapper, "embed_documents", fake_embed_documents)
|
monkeypatch.setattr(EmbedderWrapper, "embed_documents", fake_embed_documents)
|
||||||
|
|
||||||
dim = Config.embeddings.model.vector_dim
|
dim = get_config().embeddings.model.vector_dim
|
||||||
docling_doc = DoclingDocument(name="d")
|
docling_doc = DoclingDocument(name="d")
|
||||||
docling_doc.add_text(label=DocItemLabel.TEXT, text="body")
|
docling_doc.add_text(label=DocItemLabel.TEXT, text="body")
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -38,11 +38,9 @@ chunks = [
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_reranker_base():
|
async def test_reranker_base():
|
||||||
from haiku.rag.config import Config
|
|
||||||
|
|
||||||
reranker = RerankerBase()
|
reranker = RerankerBase()
|
||||||
expected_model = Config.reranking.model.name if Config.reranking.model else None
|
# The base carries no model: each reranker takes its own from the factory.
|
||||||
assert reranker._model == expected_model
|
assert reranker._model is None
|
||||||
|
|
||||||
# Empty input short-circuits in the base class without dispatching to _rerank.
|
# Empty input short-circuits in the base class without dispatching to _rerank.
|
||||||
assert await reranker.rerank("query", []) == []
|
assert await reranker.rerank("query", []) == []
|
||||||
|
|
@ -156,12 +154,12 @@ class TestGetReranker:
|
||||||
[
|
[
|
||||||
(
|
(
|
||||||
"cohere",
|
"cohere",
|
||||||
"rerank-v3.5",
|
"rerank-english-v3.0",
|
||||||
"haiku.rag.reranking.cohere",
|
"haiku.rag.reranking.cohere",
|
||||||
"CohereReranker",
|
"CohereReranker",
|
||||||
{},
|
{},
|
||||||
{},
|
{"_model": "rerank-english-v3.0"},
|
||||||
{},
|
{"CO_API_KEY": "test-api-key"},
|
||||||
),
|
),
|
||||||
(
|
(
|
||||||
"vllm",
|
"vllm",
|
||||||
|
|
|
||||||
|
|
@ -1,14 +1,16 @@
|
||||||
import pytest
|
import pytest
|
||||||
|
|
||||||
from haiku.rag.client import HaikuRAG
|
from haiku.rag.client import HaikuRAG
|
||||||
from haiku.rag.config import Config
|
from haiku.rag.config import get_config
|
||||||
from haiku.rag.store.models import SearchResult
|
from haiku.rag.store.models import SearchResult
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.vcr()
|
@pytest.mark.vcr()
|
||||||
async def test_search_qa_corpus(qa_corpus: list[dict[str, str]], 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."""
|
"""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:
|
async with HaikuRAG(
|
||||||
|
db_path=temp_db_path, config=get_config(), create=True
|
||||||
|
) as client:
|
||||||
# Load unique documents (limited to 10)
|
# Load unique documents (limited to 10)
|
||||||
seen_documents = set()
|
seen_documents = set()
|
||||||
documents = []
|
documents = []
|
||||||
|
|
@ -57,7 +59,9 @@ async def test_search_qa_corpus(qa_corpus: list[dict[str, str]], temp_db_path):
|
||||||
@pytest.mark.vcr()
|
@pytest.mark.vcr()
|
||||||
async def test_search_chunk_includes_document_provenance(temp_db_path):
|
async def test_search_chunk_includes_document_provenance(temp_db_path):
|
||||||
"""Test that raw chunk search results include document URI, metadata, and ID."""
|
"""Test that raw chunk search results include document URI, metadata, and ID."""
|
||||||
async with HaikuRAG(db_path=temp_db_path, config=Config, create=True) as client:
|
async with HaikuRAG(
|
||||||
|
db_path=temp_db_path, config=get_config(), create=True
|
||||||
|
) as client:
|
||||||
# Create a document with URI and metadata but no title
|
# Create a document with URI and metadata but no title
|
||||||
created_document = await client.create_document(
|
created_document = await client.create_document(
|
||||||
content="This is a test document with some content for searching.",
|
content="This is a test document with some content for searching.",
|
||||||
|
|
@ -92,7 +96,9 @@ async def test_search_chunk_includes_document_provenance(temp_db_path):
|
||||||
@pytest.mark.vcr()
|
@pytest.mark.vcr()
|
||||||
async def test_search_score_types(temp_db_path):
|
async def test_search_score_types(temp_db_path):
|
||||||
"""Test that different search types return appropriate score ranges."""
|
"""Test that different search types return appropriate score ranges."""
|
||||||
async with HaikuRAG(db_path=temp_db_path, config=Config, create=True) as client:
|
async with HaikuRAG(
|
||||||
|
db_path=temp_db_path, config=get_config(), create=True
|
||||||
|
) as client:
|
||||||
# Create multiple documents with different content
|
# Create multiple documents with different content
|
||||||
documents_content = [
|
documents_content = [
|
||||||
"Machine learning algorithms are powerful tools for data analysis and pattern recognition.",
|
"Machine learning algorithms are powerful tools for data analysis and pattern recognition.",
|
||||||
|
|
@ -159,7 +165,9 @@ async def test_search_score_types(temp_db_path):
|
||||||
@pytest.mark.vcr()
|
@pytest.mark.vcr()
|
||||||
async def test_search_returns_search_result(temp_db_path):
|
async def test_search_returns_search_result(temp_db_path):
|
||||||
"""Test that client.search() returns SearchResult with provenance info."""
|
"""Test that client.search() returns SearchResult with provenance info."""
|
||||||
async with HaikuRAG(db_path=temp_db_path, config=Config, create=True) as client:
|
async with HaikuRAG(
|
||||||
|
db_path=temp_db_path, config=get_config(), create=True
|
||||||
|
) as client:
|
||||||
await client.create_document(
|
await client.create_document(
|
||||||
content="Machine learning models can classify images with high accuracy.",
|
content="Machine learning models can classify images with high accuracy.",
|
||||||
uri="https://example.com/ml.html",
|
uri="https://example.com/ml.html",
|
||||||
|
|
@ -188,7 +196,9 @@ async def test_search_graceful_degradation(temp_db_path):
|
||||||
"""Test search works when docling data is unavailable."""
|
"""Test search works when docling data is unavailable."""
|
||||||
from haiku.rag.store.models import Chunk
|
from haiku.rag.store.models import Chunk
|
||||||
|
|
||||||
async with HaikuRAG(db_path=temp_db_path, config=Config, create=True) as client:
|
async with HaikuRAG(
|
||||||
|
db_path=temp_db_path, config=get_config(), create=True
|
||||||
|
) as client:
|
||||||
# Import document with custom chunks (no docling document)
|
# Import document with custom chunks (no docling document)
|
||||||
custom_chunks = [
|
custom_chunks = [
|
||||||
Chunk(content="Custom chunk without docling metadata", metadata={}),
|
Chunk(content="Custom chunk without docling metadata", metadata={}),
|
||||||
|
|
|
||||||
|
|
@ -2,7 +2,7 @@ import logging
|
||||||
|
|
||||||
import pytest
|
import pytest
|
||||||
|
|
||||||
from haiku.rag.config import AppConfig, Config
|
from haiku.rag.config import AppConfig, get_config
|
||||||
from haiku.rag.store.repositories.settings import ConfigMismatchError
|
from haiku.rag.store.repositories.settings import ConfigMismatchError
|
||||||
|
|
||||||
|
|
||||||
|
|
@ -16,7 +16,7 @@ async def test_settings_table_populated_on_store_init(temp_db_path):
|
||||||
settings_repo = SettingsRepository(store)
|
settings_repo = SettingsRepository(store)
|
||||||
|
|
||||||
db_settings = await settings_repo.get_current_settings()
|
db_settings = await settings_repo.get_current_settings()
|
||||||
config_dict = Config.model_dump(mode="json")
|
config_dict = get_config().model_dump(mode="json")
|
||||||
|
|
||||||
# Remove version from db_settings since it's added automatically
|
# Remove version from db_settings since it's added automatically
|
||||||
db_settings_without_version = {
|
db_settings_without_version = {
|
||||||
|
|
@ -34,14 +34,14 @@ async def test_settings_save_and_retrieve(temp_db_path):
|
||||||
async with Store(temp_db_path, create=True) as store:
|
async with Store(temp_db_path, create=True) as store:
|
||||||
settings_repo = SettingsRepository(store)
|
settings_repo = SettingsRepository(store)
|
||||||
|
|
||||||
original_chunk_size = Config.processing.chunk_size
|
original_chunk_size = get_config().processing.chunk_size
|
||||||
Config.processing.chunk_size = 2 * original_chunk_size
|
get_config().processing.chunk_size = 2 * original_chunk_size
|
||||||
|
|
||||||
await settings_repo.save_current_settings()
|
await settings_repo.save_current_settings()
|
||||||
retrieved_settings = await settings_repo.get_current_settings()
|
retrieved_settings = await settings_repo.get_current_settings()
|
||||||
assert retrieved_settings["processing"]["chunk_size"] == 2 * original_chunk_size
|
assert retrieved_settings["processing"]["chunk_size"] == 2 * original_chunk_size
|
||||||
|
|
||||||
Config.processing.chunk_size = original_chunk_size
|
get_config().processing.chunk_size = original_chunk_size
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
|
|
@ -52,7 +52,7 @@ async def test_set_haiku_version_recreates_row_from_store_config(temp_db_path):
|
||||||
from haiku.rag.store.repositories.settings import SettingsRepository
|
from haiku.rag.store.repositories.settings import SettingsRepository
|
||||||
|
|
||||||
config = AppConfig()
|
config = AppConfig()
|
||||||
config.processing.chunk_size = Config.processing.chunk_size + 512
|
config.processing.chunk_size = get_config().processing.chunk_size + 512
|
||||||
|
|
||||||
async with Store(temp_db_path, config=config, create=True) as store:
|
async with Store(temp_db_path, config=config, create=True) as store:
|
||||||
settings_repo = SettingsRepository(store)
|
settings_repo = SettingsRepository(store)
|
||||||
|
|
|
||||||
|
|
@ -3,7 +3,7 @@ import importlib.util
|
||||||
import pytest
|
import pytest
|
||||||
from pydantic_ai.models.openai import OpenAIChatModel
|
from pydantic_ai.models.openai import OpenAIChatModel
|
||||||
|
|
||||||
from haiku.rag.config import Config
|
from haiku.rag.config import get_config
|
||||||
from haiku.rag.config.models import ModelConfig
|
from haiku.rag.config.models import ModelConfig
|
||||||
from haiku.rag.converters import get_converter
|
from haiku.rag.converters import get_converter
|
||||||
from haiku.rag.utils import get_model
|
from haiku.rag.utils import get_model
|
||||||
|
|
@ -20,7 +20,7 @@ async def test_text_to_docling_document():
|
||||||
"""Test text to DoclingDocument conversion."""
|
"""Test text to DoclingDocument conversion."""
|
||||||
# Test basic text conversion
|
# Test basic text conversion
|
||||||
simple_text = "This is a simple text document."
|
simple_text = "This is a simple text document."
|
||||||
converter = get_converter(Config)
|
converter = get_converter(get_config())
|
||||||
doc = await converter.convert_text(simple_text)
|
doc = await converter.convert_text(simple_text)
|
||||||
|
|
||||||
# Verify it returns a DoclingDocument
|
# Verify it returns a DoclingDocument
|
||||||
|
|
@ -44,7 +44,7 @@ def hello():
|
||||||
return True
|
return True
|
||||||
```"""
|
```"""
|
||||||
|
|
||||||
converter = get_converter(Config)
|
converter = get_converter(get_config())
|
||||||
doc = await converter.convert_text(code_text, name="hello.md")
|
doc = await converter.convert_text(code_text, name="hello.md")
|
||||||
|
|
||||||
# Verify it's a valid DoclingDocument
|
# Verify it's a valid DoclingDocument
|
||||||
|
|
@ -77,7 +77,7 @@ def test():
|
||||||
|
|
||||||
**Bold text** and *italic text*."""
|
**Bold text** and *italic text*."""
|
||||||
|
|
||||||
converter = get_converter(Config)
|
converter = get_converter(get_config())
|
||||||
doc = await converter.convert_text(markdown_text, name="test.md")
|
doc = await converter.convert_text(markdown_text, name="test.md")
|
||||||
|
|
||||||
# Verify it's a DoclingDocument
|
# Verify it's a DoclingDocument
|
||||||
|
|
@ -95,7 +95,7 @@ def test():
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_text_to_docling_document_empty_content():
|
async def test_text_to_docling_document_empty_content():
|
||||||
"""Test text to DoclingDocument conversion with empty content."""
|
"""Test text to DoclingDocument conversion with empty content."""
|
||||||
converter = get_converter(Config)
|
converter = get_converter(get_config())
|
||||||
doc = await converter.convert_text("")
|
doc = await converter.convert_text("")
|
||||||
|
|
||||||
# Should still create a valid DoclingDocument
|
# Should still create a valid DoclingDocument
|
||||||
|
|
@ -124,7 +124,7 @@ function saludar() {
|
||||||
|
|
||||||
Emoji test: 🚀 ✅ 📝"""
|
Emoji test: 🚀 ✅ 📝"""
|
||||||
|
|
||||||
converter = get_converter(Config)
|
converter = get_converter(get_config())
|
||||||
doc = await converter.convert_text(unicode_text, name="unicode.md")
|
doc = await converter.convert_text(unicode_text, name="unicode.md")
|
||||||
|
|
||||||
# Verify it's a DoclingDocument
|
# Verify it's a DoclingDocument
|
||||||
|
|
|
||||||
|
|
@ -5,7 +5,7 @@ import pytest
|
||||||
import haiku.rag.client as client_mod
|
import haiku.rag.client as client_mod
|
||||||
from haiku.rag.client import HaikuRAG
|
from haiku.rag.client import HaikuRAG
|
||||||
from haiku.rag.client.documents import _refresh_doc_metadata
|
from haiku.rag.client.documents import _refresh_doc_metadata
|
||||||
from haiku.rag.config import Config
|
from haiku.rag.config import get_config
|
||||||
from haiku.rag.store.models.chunk import Chunk
|
from haiku.rag.store.models.chunk import Chunk
|
||||||
|
|
||||||
|
|
||||||
|
|
@ -73,7 +73,7 @@ async def test_metadata_refresh_sweep_schedules_vacuum(temp_db_path):
|
||||||
"""A source re-sweep that only rolls source_revision (MD5/revision
|
"""A source re-sweep that only rolls source_revision (MD5/revision
|
||||||
short-circuit) writes document_meta and must still schedule the (debounced)
|
short-circuit) writes document_meta and must still schedule the (debounced)
|
||||||
vacuum, so that tiny churn gets reclaimed instead of accumulating."""
|
vacuum, so that tiny churn gets reclaimed instead of accumulating."""
|
||||||
dim = Config.embeddings.model.vector_dim
|
dim = get_config().embeddings.model.vector_dim
|
||||||
async with HaikuRAG(temp_db_path, create=True) as client:
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
||||||
doc = await client.import_document(
|
doc = await client.import_document(
|
||||||
_docling_doc("d", "body"),
|
_docling_doc("d", "body"),
|
||||||
|
|
@ -99,7 +99,7 @@ async def test_metadata_refresh_waits_for_write_lock(temp_db_path):
|
||||||
"""The revision/MD5 short-circuit write serializes with other writers so
|
"""The revision/MD5 short-circuit write serializes with other writers so
|
||||||
it cannot land inside another writer's critical section (e.g. between
|
it cannot land inside another writer's critical section (e.g. between
|
||||||
create_tag's version snapshot and its per-table tag creation)."""
|
create_tag's version snapshot and its per-table tag creation)."""
|
||||||
dim = Config.embeddings.model.vector_dim
|
dim = get_config().embeddings.model.vector_dim
|
||||||
async with HaikuRAG(temp_db_path, create=True) as client:
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
||||||
doc = await client.import_document(
|
doc = await client.import_document(
|
||||||
_docling_doc("d", "body"),
|
_docling_doc("d", "body"),
|
||||||
|
|
|
||||||
|
|
@ -271,10 +271,10 @@ async def test_vacuum_with_retention_threshold(temp_db_path):
|
||||||
@pytest.mark.vcr()
|
@pytest.mark.vcr()
|
||||||
async def test_vacuum_completes_before_context_exit(temp_db_path, monkeypatch):
|
async def test_vacuum_completes_before_context_exit(temp_db_path, monkeypatch):
|
||||||
"""Test that background vacuum completes when context manager exits."""
|
"""Test that background vacuum completes when context manager exits."""
|
||||||
from haiku.rag.config import Config
|
from haiku.rag.config import get_config
|
||||||
|
|
||||||
# Set aggressive vacuum retention for this test
|
# Set aggressive vacuum retention for this test
|
||||||
monkeypatch.setattr(Config.storage, "vacuum_retention_seconds", 0)
|
monkeypatch.setattr(get_config().storage, "vacuum_retention_seconds", 0)
|
||||||
|
|
||||||
async with HaikuRAG(db_path=temp_db_path, create=True) as client:
|
async with HaikuRAG(db_path=temp_db_path, create=True) as client:
|
||||||
# Create multiple documents - each creation triggers automatic vacuum with retention=0
|
# Create multiple documents - each creation triggers automatic vacuum with retention=0
|
||||||
|
|
@ -302,9 +302,9 @@ async def test_aexit_awaits_background_vacuum(temp_db_path, monkeypatch):
|
||||||
it yet when __aexit__ runs. Simply acquiring the vacuum lock (which is free until
|
it yet when __aexit__ runs. Simply acquiring the vacuum lock (which is free until
|
||||||
the task actually starts) would let close() proceed before vacuum runs.
|
the task actually starts) would let close() proceed before vacuum runs.
|
||||||
"""
|
"""
|
||||||
from haiku.rag.config import Config
|
from haiku.rag.config import get_config
|
||||||
|
|
||||||
monkeypatch.setattr(Config.storage, "auto_vacuum", True)
|
monkeypatch.setattr(get_config().storage, "auto_vacuum", True)
|
||||||
|
|
||||||
vacuum_started = asyncio.Event()
|
vacuum_started = asyncio.Event()
|
||||||
vacuum_completed = asyncio.Event()
|
vacuum_completed = asyncio.Event()
|
||||||
|
|
@ -339,9 +339,9 @@ async def test_aexit_awaits_all_background_vacuums(temp_db_path, monkeypatch):
|
||||||
__aexit__ awaits the fast no-op B and closes the connection while Task A
|
__aexit__ awaits the fast no-op B and closes the connection while Task A
|
||||||
is still running.
|
is still running.
|
||||||
"""
|
"""
|
||||||
from haiku.rag.config import Config
|
from haiku.rag.config import get_config
|
||||||
|
|
||||||
monkeypatch.setattr(Config.storage, "auto_vacuum", True)
|
monkeypatch.setattr(get_config().storage, "auto_vacuum", True)
|
||||||
|
|
||||||
first_vacuum_completed = asyncio.Event()
|
first_vacuum_completed = asyncio.Event()
|
||||||
|
|
||||||
|
|
@ -378,10 +378,10 @@ async def test_aexit_awaits_all_background_vacuums(temp_db_path, monkeypatch):
|
||||||
@pytest.mark.vcr()
|
@pytest.mark.vcr()
|
||||||
async def test_auto_vacuum_disabled_skips_vacuum(temp_db_path, monkeypatch):
|
async def test_auto_vacuum_disabled_skips_vacuum(temp_db_path, monkeypatch):
|
||||||
"""Test that auto_vacuum=False prevents automatic vacuum after operations."""
|
"""Test that auto_vacuum=False prevents automatic vacuum after operations."""
|
||||||
from haiku.rag.config import Config
|
from haiku.rag.config import get_config
|
||||||
|
|
||||||
# Disable auto-vacuum
|
# Disable auto-vacuum
|
||||||
monkeypatch.setattr(Config.storage, "auto_vacuum", False)
|
monkeypatch.setattr(get_config().storage, "auto_vacuum", False)
|
||||||
|
|
||||||
async with HaikuRAG(db_path=temp_db_path, create=True) as client:
|
async with HaikuRAG(db_path=temp_db_path, create=True) as client:
|
||||||
# Create multiple documents
|
# Create multiple documents
|
||||||
|
|
@ -404,11 +404,11 @@ async def test_auto_vacuum_disabled_skips_vacuum(temp_db_path, monkeypatch):
|
||||||
@pytest.mark.vcr()
|
@pytest.mark.vcr()
|
||||||
async def test_auto_vacuum_enabled_triggers_vacuum(temp_db_path, monkeypatch):
|
async def test_auto_vacuum_enabled_triggers_vacuum(temp_db_path, monkeypatch):
|
||||||
"""Test that auto_vacuum=True (default) triggers vacuum after operations."""
|
"""Test that auto_vacuum=True (default) triggers vacuum after operations."""
|
||||||
from haiku.rag.config import Config
|
from haiku.rag.config import get_config
|
||||||
|
|
||||||
# Enable auto-vacuum with aggressive retention
|
# Enable auto-vacuum with aggressive retention
|
||||||
monkeypatch.setattr(Config.storage, "auto_vacuum", True)
|
monkeypatch.setattr(get_config().storage, "auto_vacuum", True)
|
||||||
monkeypatch.setattr(Config.storage, "vacuum_retention_seconds", 0)
|
monkeypatch.setattr(get_config().storage, "vacuum_retention_seconds", 0)
|
||||||
|
|
||||||
async with HaikuRAG(db_path=temp_db_path, create=True) as client:
|
async with HaikuRAG(db_path=temp_db_path, create=True) as client:
|
||||||
# Create multiple documents
|
# Create multiple documents
|
||||||
|
|
|
||||||
|
|
@ -238,6 +238,6 @@ async def doc_client(temp_db_path):
|
||||||
@pytest.fixture
|
@pytest.fixture
|
||||||
def doc_config():
|
def doc_config():
|
||||||
"""Default AppConfig for document tests."""
|
"""Default AppConfig for document tests."""
|
||||||
from haiku.rag.config import Config
|
from haiku.rag.config import get_config
|
||||||
|
|
||||||
return Config
|
return get_config()
|
||||||
|
|
|
||||||
|
|
@ -205,9 +205,9 @@ async def search_client(temp_db_path):
|
||||||
@pytest.fixture
|
@pytest.fixture
|
||||||
def search_config():
|
def search_config():
|
||||||
"""Default AppConfig for search tests."""
|
"""Default AppConfig for search tests."""
|
||||||
from haiku.rag.config import Config
|
from haiku.rag.config import get_config
|
||||||
|
|
||||||
return Config
|
return get_config()
|
||||||
|
|
||||||
|
|
||||||
def _png_b64():
|
def _png_b64():
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue