# System Packages import logging from pathlib import Path import os # External Packages import pytest import torch from khoj.utils.config import SearchModels # Internal Packages from khoj.utils.state import content_index, search_models from khoj.search_type import text_search from khoj.utils.rawconfig import ContentConfig, SearchConfig, TextContentConfig from khoj.processor.org_mode.org_to_jsonl import OrgToJsonl from khoj.processor.github.github_to_jsonl import GithubToJsonl # Test # ---------------------------------------------------------------------------------------------------- def test_text_search_setup_with_missing_file_raises_error( org_config_with_only_new_file: TextContentConfig, search_config: SearchConfig ): # Arrange # Ensure file mentioned in org.input-files is missing single_new_file = Path(org_config_with_only_new_file.input_files[0]) single_new_file.unlink() # Act # Generate notes embeddings during asymmetric setup with pytest.raises(FileNotFoundError): text_search.setup(OrgToJsonl, org_config_with_only_new_file, search_config.asymmetric, regenerate=True) # ---------------------------------------------------------------------------------------------------- def test_text_search_setup_with_empty_file_raises_error( org_config_with_only_new_file: TextContentConfig, search_config: SearchConfig ): # Act # Generate notes embeddings during asymmetric setup with pytest.raises(ValueError, match=r"^No valid entries found*"): text_search.setup(OrgToJsonl, org_config_with_only_new_file, search_config.asymmetric, regenerate=True) # ---------------------------------------------------------------------------------------------------- def test_text_search_setup(content_config: ContentConfig, search_models: SearchModels): # Act # Regenerate notes embeddings during asymmetric setup notes_model = text_search.setup( OrgToJsonl, content_config.org, search_models.text_search.bi_encoder, regenerate=True ) # Assert assert len(notes_model.entries) == 10 assert len(notes_model.corpus_embeddings) == 10 # ---------------------------------------------------------------------------------------------------- def test_text_index_same_if_content_unchanged(content_config: ContentConfig, search_models: SearchModels, caplog): # Arrange caplog.set_level(logging.INFO, logger="khoj") # Act # Generate initial notes embeddings during asymmetric setup text_search.setup(OrgToJsonl, content_config.org, search_models.text_search.bi_encoder, regenerate=True) initial_logs = caplog.text caplog.clear() # Clear logs # Run asymmetric setup again with no changes to data source. Ensure index is not updated text_search.setup(OrgToJsonl, content_config.org, search_models.text_search.bi_encoder, regenerate=False) final_logs = caplog.text # Assert assert "Creating index from scratch." in initial_logs assert "Creating index from scratch." not in final_logs # ---------------------------------------------------------------------------------------------------- @pytest.mark.anyio async def test_text_search(content_config: ContentConfig, search_config: SearchConfig): # Arrange search_models.text_search = text_search.initialize_model(search_config.asymmetric) content_index.org = text_search.setup( OrgToJsonl, content_config.org, search_models.text_search.bi_encoder, regenerate=True ) query = "How to git install application?" # Act hits, entries = await text_search.query( query, search_model=search_models.text_search, content=content_index.org, rank_results=True ) results = text_search.collate_results(hits, entries, count=1) # Assert # search results should contain "git clone" entry search_result = results[0].entry assert "git clone" in search_result # ---------------------------------------------------------------------------------------------------- def test_entry_chunking_by_max_tokens(org_config_with_only_new_file: TextContentConfig, search_models: SearchModels): # Arrange # Insert org-mode entry with size exceeding max token limit to new org file max_tokens = 256 new_file_to_index = Path(org_config_with_only_new_file.input_files[0]) with open(new_file_to_index, "w") as f: f.write(f"* Entry more than {max_tokens} words\n") for index in range(max_tokens + 1): f.write(f"{index} ") # Act # reload embeddings, entries, notes model after adding new org-mode file initial_notes_model = text_search.setup( OrgToJsonl, org_config_with_only_new_file, search_models.text_search.bi_encoder, regenerate=False ) # Assert # verify newly added org-mode entry is split by max tokens assert len(initial_notes_model.entries) == 2 assert len(initial_notes_model.corpus_embeddings) == 2 # ---------------------------------------------------------------------------------------------------- def test_regenerate_index_with_new_entry( content_config: ContentConfig, search_models: SearchModels, new_org_file: Path ): # Arrange initial_notes_model = text_search.setup( OrgToJsonl, content_config.org, search_models.text_search.bi_encoder, regenerate=True ) assert len(initial_notes_model.entries) == 10 assert len(initial_notes_model.corpus_embeddings) == 10 # append org-mode entry to first org input file in config content_config.org.input_files = [f"{new_org_file}"] with open(new_org_file, "w") as f: f.write("\n* A Chihuahua doing Tango\n- Saw a super cute video of a chihuahua doing the Tango on Youtube\n") # Act # regenerate notes jsonl, model embeddings and model to include entry from new file regenerated_notes_model = text_search.setup( OrgToJsonl, content_config.org, search_models.text_search.bi_encoder, regenerate=True ) # Assert assert len(regenerated_notes_model.entries) == 11 assert len(regenerated_notes_model.corpus_embeddings) == 11 # verify new entry appended to index, without disrupting order or content of existing entries error_details = compare_index(initial_notes_model, regenerated_notes_model) if error_details: pytest.fail(error_details, False) # Cleanup # reset input_files in config to empty list content_config.org.input_files = [] # ---------------------------------------------------------------------------------------------------- def test_update_index_with_duplicate_entries_in_stable_order( org_config_with_only_new_file: TextContentConfig, search_models: SearchModels ): # Arrange new_file_to_index = Path(org_config_with_only_new_file.input_files[0]) # Insert org-mode entries with same compiled form into new org file new_entry = "* TODO A Chihuahua doing Tango\n- Saw a super cute video of a chihuahua doing the Tango on Youtube\n" with open(new_file_to_index, "w") as f: f.write(f"{new_entry}{new_entry}") # Act # load embeddings, entries, notes model after adding new org-mode file initial_index = text_search.setup( OrgToJsonl, org_config_with_only_new_file, search_models.text_search.bi_encoder, regenerate=True ) # update embeddings, entries, notes model after adding new org-mode file updated_index = text_search.setup( OrgToJsonl, org_config_with_only_new_file, search_models.text_search.bi_encoder, regenerate=False ) # Assert # verify only 1 entry added even if there are multiple duplicate entries assert len(initial_index.entries) == len(updated_index.entries) == 1 assert len(initial_index.corpus_embeddings) == len(updated_index.corpus_embeddings) == 1 # verify the same entry is added even when there are multiple duplicate entries error_details = compare_index(initial_index, updated_index) if error_details: pytest.fail(error_details) # ---------------------------------------------------------------------------------------------------- def test_update_index_with_deleted_entry(org_config_with_only_new_file: TextContentConfig, search_models: SearchModels): # Arrange new_file_to_index = Path(org_config_with_only_new_file.input_files[0]) # Insert org-mode entries with same compiled form into new org file new_entry = "* TODO A Chihuahua doing Tango\n- Saw a super cute video of a chihuahua doing the Tango on Youtube\n" with open(new_file_to_index, "w") as f: f.write(f"{new_entry}{new_entry} -- Tatooine") # load embeddings, entries, notes model after adding new org file with 2 entries initial_index = text_search.setup( OrgToJsonl, org_config_with_only_new_file, search_models.text_search.bi_encoder, regenerate=True ) # update embeddings, entries, notes model after removing an entry from the org file with open(new_file_to_index, "w") as f: f.write(f"{new_entry}") # Act updated_index = text_search.setup( OrgToJsonl, org_config_with_only_new_file, search_models.text_search.bi_encoder, regenerate=False ) # Assert # verify only 1 entry added even if there are multiple duplicate entries assert len(initial_index.entries) == len(updated_index.entries) + 1 assert len(initial_index.corpus_embeddings) == len(updated_index.corpus_embeddings) + 1 # verify the same entry is added even when there are multiple duplicate entries error_details = compare_index(updated_index, initial_index) if error_details: pytest.fail(error_details) # ---------------------------------------------------------------------------------------------------- def test_update_index_with_new_entry(content_config: ContentConfig, search_models: SearchModels, new_org_file: Path): # Arrange initial_notes_model = text_search.setup( OrgToJsonl, content_config.org, search_models.text_search.bi_encoder, regenerate=True, normalize=False ) # append org-mode entry to first org input file in config with open(new_org_file, "w") as f: new_entry = "\n* A Chihuahua doing Tango\n- Saw a super cute video of a chihuahua doing the Tango on Youtube\n" f.write(new_entry) # Act # update embeddings, entries with the newly added note content_config.org.input_files = [f"{new_org_file}"] final_notes_model = text_search.setup( OrgToJsonl, content_config.org, search_models.text_search.bi_encoder, regenerate=False, normalize=False ) # Assert assert len(final_notes_model.entries) == len(initial_notes_model.entries) + 1 assert len(final_notes_model.corpus_embeddings) == len(initial_notes_model.corpus_embeddings) + 1 # verify new entry appended to index, without disrupting order or content of existing entries error_details = compare_index(initial_notes_model, final_notes_model) if error_details: pytest.fail(error_details, False) # Cleanup # reset input_files in config to empty list content_config.org.input_files = [] # ---------------------------------------------------------------------------------------------------- @pytest.mark.skipif(os.getenv("GITHUB_PAT_TOKEN") is None, reason="GITHUB_PAT_TOKEN not set") def test_text_search_setup_github(content_config: ContentConfig, search_models: SearchModels): # Act # Regenerate github embeddings to test asymmetric setup without caching github_model = text_search.setup( GithubToJsonl, content_config.github, search_models.text_search.bi_encoder, regenerate=True ) # Assert assert len(github_model.entries) > 1 def compare_index(initial_notes_model, final_notes_model): mismatched_entries, mismatched_embeddings = [], [] for index in range(len(initial_notes_model.entries)): if initial_notes_model.entries[index].to_json() != final_notes_model.entries[index].to_json(): mismatched_entries.append(index) # verify new entry embedding appended to embeddings tensor, without disrupting order or content of existing embeddings for index in range(len(initial_notes_model.corpus_embeddings)): if not torch.equal(final_notes_model.corpus_embeddings[index], initial_notes_model.corpus_embeddings[index]): mismatched_embeddings.append(index) error_details = "" if mismatched_entries: mismatched_entries_str = ",".join(map(str, mismatched_entries)) error_details += f"Entries at {mismatched_entries_str} not equal\n" if mismatched_embeddings: mismatched_embeddings_str = ", ".join(map(str, mismatched_embeddings)) error_details += f"Embeddings at {mismatched_embeddings_str} not equal\n" return error_details