kv-ai/tests/test_indexer.py

import tempfile
from pathlib import Path
from unittest.mock import MagicMock, patch

from companion.config import (
    Config,
    VaultConfig,
    IndexingConfig,
    RagConfig,
    EmbeddingConfig,
    VectorStoreConfig,
    SearchConfig,
    HybridSearchConfig,
    FiltersConfig,
    CompanionConfig,
    PersonaConfig,
    MemoryConfig,
    ChatConfig,
    ModelConfig,
    InferenceConfig,
    FineTuningConfig,
    RetrainScheduleConfig,
    ApiConfig,
    AuthConfig,
    UiConfig,
    WebConfig,
    WebFeaturesConfig,
    CliConfig,
    LoggingConfig,
    SecurityConfig,
)
from companion.rag.indexer import Indexer
from companion.rag.vector_store import VectorStore


def _make_config(vault_path: Path, vector_store_path: Path) -> Config:
    return Config(
        companion=CompanionConfig(
            name="SAN",
            persona=PersonaConfig(
                role="companion", tone="reflective", style="questioning", boundaries=[]
            ),
            memory=MemoryConfig(
                session_turns=20, persistent_store="", summarize_after=10
            ),
            chat=ChatConfig(
                streaming=True,
                max_response_tokens=2048,
                default_temperature=0.7,
                allow_temperature_override=True,
            ),
        ),
        vault=VaultConfig(
            path=str(vault_path),
            indexing=IndexingConfig(
                auto_sync=False,
                auto_sync_interval_minutes=1440,
                watch_fs_events=False,
                file_patterns=["*.md"],
                deny_dirs=[".git"],
                deny_patterns=[".*"],
            ),
            chunking_rules={},
        ),
        rag=RagConfig(
            embedding=EmbeddingConfig(
                provider="ollama",
                model="dummy",
                base_url="http://localhost:11434",
                dimensions=4,
                batch_size=2,
            ),
            vector_store=VectorStoreConfig(type="lancedb", path=str(vector_store_path)),
            search=SearchConfig(
                default_top_k=8,
                max_top_k=20,
                similarity_threshold=0.75,
                hybrid_search=HybridSearchConfig(
                    enabled=False, keyword_weight=0.3, semantic_weight=0.7
                ),
                filters=FiltersConfig(
                    date_range_enabled=True,
                    tag_filter_enabled=True,
                    directory_filter_enabled=True,
                ),
            ),
        ),
        model=ModelConfig(
            inference=InferenceConfig(
                backend="llama.cpp",
                model_path="",
                context_length=8192,
                gpu_layers=35,
                batch_size=512,
                threads=8,
            ),
            fine_tuning=FineTuningConfig(
                base_model="",
                output_dir="",
                lora_rank=16,
                lora_alpha=32,
                learning_rate=0.0002,
                batch_size=4,
                gradient_accumulation_steps=4,
                num_epochs=3,
                warmup_steps=100,
                save_steps=500,
                eval_steps=250,
                training_data_path="",
                validation_split=0.1,
            ),
            retrain_schedule=RetrainScheduleConfig(
                auto_reminder=True, default_interval_days=90, reminder_channels=[]
            ),
        ),
        api=ApiConfig(
            host="127.0.0.1", port=7373, cors_origins=[], auth=AuthConfig(enabled=False)
        ),
        ui=UiConfig(
            web=WebConfig(
                enabled=True,
                theme="obsidian",
                features=WebFeaturesConfig(
                    streaming=True, citations=True, source_preview=True
                ),
            ),
            cli=CliConfig(enabled=True, rich_output=True),
        ),
        logging=LoggingConfig(level="INFO", file="", max_size_mb=100, backup_count=5),
        security=SecurityConfig(
            local_only=True,
            vault_path_traversal_check=True,
            sensitive_content_detection=True,
            sensitive_patterns=[],
            require_confirmation_for_external_apis=True,
        ),
    )


@patch("companion.rag.indexer.OllamaEmbedder")
def test_full_index_creates_vectors(mock_embedder_cls):
    mock_embedder = MagicMock()
    mock_embedder.embed.return_value = [[1.0, 0.0, 0.0, 0.0]]
    mock_embedder_cls.return_value = mock_embedder

    with tempfile.TemporaryDirectory() as tmp:
        vault = Path(tmp) / "vault"
        vault.mkdir()
        (vault / "hello.md").write_text("hello world", encoding="utf-8")
        vs_path = Path(tmp) / "vectors"
        config = _make_config(vault, vs_path)
        store = VectorStore(uri=vs_path, dimensions=4)
        indexer = Indexer(config, store)
        indexer.full_index()
        assert store.count() == 1


@patch("companion.rag.indexer.OllamaEmbedder")
def test_sync_updates_changed_files(mock_embedder_cls):
    mock_embedder = MagicMock()
    mock_embedder.embed.return_value = [[1.0, 0.0, 0.0, 0.0]]
    mock_embedder_cls.return_value = mock_embedder

    with tempfile.TemporaryDirectory() as tmp:
        vault = Path(tmp) / "vault"
        vault.mkdir()
        file_path = vault / "hello.md"
        file_path.write_text("hello world", encoding="utf-8")
        vs_path = Path(tmp) / "vectors"
        config = _make_config(vault, vs_path)
        store = VectorStore(uri=vs_path, dimensions=4)
        indexer = Indexer(config, store)

        indexer.full_index()
        assert store.count() == 1
        first_text = store.table.to_pandas()["text"].iloc[0]
        assert first_text == "hello world"

        file_path.write_text("updated content", encoding="utf-8")
        indexer.sync()

        assert store.count() == 1
        updated_text = store.table.to_pandas()["text"].iloc[0]
        assert updated_text == "updated content"