2022-06-29 19:09:48 +00:00
|
|
|
# System Packages
|
2023-05-12 09:37:34 +00:00
|
|
|
import logging
|
2022-06-29 19:09:48 +00:00
|
|
|
from pathlib import Path
|
2023-06-13 23:55:58 +00:00
|
|
|
import os
|
2023-10-26 16:42:29 +00:00
|
|
|
import asyncio
|
2022-06-29 19:09:48 +00:00
|
|
|
|
2022-09-10 10:05:21 +00:00
|
|
|
# External Packages
|
|
|
|
import pytest
|
|
|
|
|
2021-10-03 03:28:33 +00:00
|
|
|
# Internal Packages
|
2023-02-14 20:50:51 +00:00
|
|
|
from khoj.search_type import text_search
|
2023-10-26 16:42:29 +00:00
|
|
|
from khoj.utils.rawconfig import ContentConfig, SearchConfig
|
2023-11-01 21:51:33 +00:00
|
|
|
from khoj.processor.org_mode.org_to_entries import OrgToEntries
|
|
|
|
from khoj.processor.github.github_to_entries import GithubToEntries
|
2023-10-26 18:37:41 +00:00
|
|
|
from khoj.utils.fs_syncer import collect_files, get_org_files
|
2023-11-01 01:50:54 +00:00
|
|
|
from database.models import LocalOrgConfig, KhojUser, Entry, GithubConfig
|
2023-10-26 16:42:29 +00:00
|
|
|
|
|
|
|
logger = logging.getLogger(__name__)
|
2021-10-03 03:28:33 +00:00
|
|
|
|
|
|
|
|
|
|
|
# Test
|
2022-09-10 22:09:24 +00:00
|
|
|
# ----------------------------------------------------------------------------------------------------
|
2023-10-26 16:42:29 +00:00
|
|
|
@pytest.mark.django_db
|
2023-11-01 00:59:53 +00:00
|
|
|
def test_text_search_setup_with_missing_file_raises_error(org_config_with_only_new_file: LocalOrgConfig):
|
2022-09-10 22:09:24 +00:00
|
|
|
# Arrange
|
2023-01-09 19:17:36 +00:00
|
|
|
# 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()
|
2022-09-10 22:09:24 +00:00
|
|
|
|
|
|
|
# Act
|
|
|
|
# Generate notes embeddings during asymmetric setup
|
2023-08-31 19:55:17 +00:00
|
|
|
with pytest.raises(FileNotFoundError):
|
2023-10-18 04:31:15 +00:00
|
|
|
get_org_files(org_config_with_only_new_file)
|
2022-09-10 22:09:24 +00:00
|
|
|
|
|
|
|
|
2023-10-18 05:59:10 +00:00
|
|
|
# ----------------------------------------------------------------------------------------------------
|
2023-10-26 16:42:29 +00:00
|
|
|
@pytest.mark.django_db
|
|
|
|
def test_get_org_files_with_org_suffixed_dir_doesnt_raise_error(tmp_path, default_user: KhojUser):
|
2023-10-18 05:59:10 +00:00
|
|
|
# Arrange
|
|
|
|
orgfile = tmp_path / "directory.org" / "file.org"
|
|
|
|
orgfile.parent.mkdir()
|
|
|
|
with open(orgfile, "w") as f:
|
|
|
|
f.write("* Heading\n- List item\n")
|
2023-10-26 16:42:29 +00:00
|
|
|
|
|
|
|
LocalOrgConfig.objects.create(
|
|
|
|
input_filter=[f"{tmp_path}/**/*"],
|
|
|
|
input_files=None,
|
|
|
|
user=default_user,
|
2023-10-18 05:59:10 +00:00
|
|
|
)
|
|
|
|
|
2023-10-26 16:42:29 +00:00
|
|
|
org_files = collect_files(user=default_user)["org"]
|
|
|
|
|
2023-10-18 05:59:10 +00:00
|
|
|
# Act
|
|
|
|
# should not raise IsADirectoryError and return orgfile
|
2023-10-26 16:42:29 +00:00
|
|
|
assert org_files == {f"{orgfile}": "* Heading\n- List item\n"}
|
2023-10-18 05:59:10 +00:00
|
|
|
|
|
|
|
|
2022-09-10 10:05:21 +00:00
|
|
|
# ----------------------------------------------------------------------------------------------------
|
2023-10-26 16:42:29 +00:00
|
|
|
@pytest.mark.django_db
|
2023-07-16 08:23:22 +00:00
|
|
|
def test_text_search_setup_with_empty_file_raises_error(
|
2023-10-26 16:42:29 +00:00
|
|
|
org_config_with_only_new_file: LocalOrgConfig, default_user: KhojUser, caplog
|
2023-02-17 16:04:26 +00:00
|
|
|
):
|
2023-08-31 19:55:17 +00:00
|
|
|
# Arrange
|
|
|
|
data = get_org_files(org_config_with_only_new_file)
|
2022-09-10 10:05:21 +00:00
|
|
|
# Act
|
|
|
|
# Generate notes embeddings during asymmetric setup
|
2023-10-26 16:42:29 +00:00
|
|
|
with caplog.at_level(logging.INFO):
|
2023-11-01 21:51:33 +00:00
|
|
|
text_search.setup(OrgToEntries, data, regenerate=True, user=default_user)
|
2023-10-26 16:42:29 +00:00
|
|
|
|
2023-11-01 00:59:53 +00:00
|
|
|
assert "Created 0 new embeddings. Deleted 3 embeddings for user " in caplog.records[-1].message
|
2023-10-26 16:42:29 +00:00
|
|
|
verify_embeddings(0, default_user)
|
2022-09-10 10:05:21 +00:00
|
|
|
|
|
|
|
|
2021-10-03 03:28:33 +00:00
|
|
|
# ----------------------------------------------------------------------------------------------------
|
2023-10-26 16:42:29 +00:00
|
|
|
@pytest.mark.django_db
|
2023-11-01 00:59:53 +00:00
|
|
|
def test_text_indexer_deletes_embedding_before_regenerate(
|
|
|
|
content_config: ContentConfig, default_user: KhojUser, caplog
|
|
|
|
):
|
2023-08-31 19:55:17 +00:00
|
|
|
# Arrange
|
2023-10-26 16:42:29 +00:00
|
|
|
org_config = LocalOrgConfig.objects.filter(user=default_user).first()
|
|
|
|
data = get_org_files(org_config)
|
2023-11-01 00:59:53 +00:00
|
|
|
with caplog.at_level(logging.DEBUG):
|
2023-11-01 21:51:33 +00:00
|
|
|
text_search.setup(OrgToEntries, data, regenerate=True, user=default_user)
|
2023-11-01 00:59:53 +00:00
|
|
|
|
|
|
|
# Assert
|
|
|
|
assert "Deleting all embeddings for file type org" in caplog.text
|
|
|
|
assert "Created 10 new embeddings. Deleted 3 embeddings for user " in caplog.records[-1].message
|
|
|
|
|
|
|
|
|
|
|
|
# ----------------------------------------------------------------------------------------------------
|
|
|
|
@pytest.mark.django_db
|
|
|
|
def test_text_search_setup_batch_processes(content_config: ContentConfig, default_user: KhojUser, caplog):
|
|
|
|
# Arrange
|
|
|
|
org_config = LocalOrgConfig.objects.filter(user=default_user).first()
|
|
|
|
data = get_org_files(org_config)
|
|
|
|
with caplog.at_level(logging.DEBUG):
|
2023-11-01 21:51:33 +00:00
|
|
|
text_search.setup(OrgToEntries, data, regenerate=True, user=default_user)
|
2021-10-03 03:28:33 +00:00
|
|
|
|
|
|
|
# Assert
|
2023-11-01 00:59:53 +00:00
|
|
|
assert "Created 4 new embeddings" in caplog.text
|
|
|
|
assert "Created 6 new embeddings" in caplog.text
|
|
|
|
assert "Created 10 new embeddings. Deleted 3 embeddings for user " in caplog.records[-1].message
|
2021-10-03 03:28:33 +00:00
|
|
|
|
|
|
|
|
2023-05-12 09:37:34 +00:00
|
|
|
# ----------------------------------------------------------------------------------------------------
|
2023-10-26 16:42:29 +00:00
|
|
|
@pytest.mark.django_db
|
|
|
|
def test_text_index_same_if_content_unchanged(content_config: ContentConfig, default_user: KhojUser, caplog):
|
2023-05-12 09:37:34 +00:00
|
|
|
# Arrange
|
2023-10-26 16:42:29 +00:00
|
|
|
org_config = LocalOrgConfig.objects.filter(user=default_user).first()
|
|
|
|
data = get_org_files(org_config)
|
2023-08-31 19:55:17 +00:00
|
|
|
|
2023-05-12 09:37:34 +00:00
|
|
|
# Act
|
|
|
|
# Generate initial notes embeddings during asymmetric setup
|
2023-11-01 00:59:53 +00:00
|
|
|
with caplog.at_level(logging.DEBUG):
|
2023-11-01 21:51:33 +00:00
|
|
|
text_search.setup(OrgToEntries, data, regenerate=True, user=default_user)
|
2023-05-12 09:37:34 +00:00
|
|
|
initial_logs = caplog.text
|
|
|
|
caplog.clear() # Clear logs
|
|
|
|
|
|
|
|
# Run asymmetric setup again with no changes to data source. Ensure index is not updated
|
2023-11-01 00:59:53 +00:00
|
|
|
with caplog.at_level(logging.DEBUG):
|
2023-11-01 21:51:33 +00:00
|
|
|
text_search.setup(OrgToEntries, data, regenerate=False, user=default_user)
|
2023-05-12 09:37:34 +00:00
|
|
|
final_logs = caplog.text
|
|
|
|
|
|
|
|
# Assert
|
2023-10-26 16:42:29 +00:00
|
|
|
assert "Deleting all embeddings for file type org" in initial_logs
|
|
|
|
assert "Deleting all embeddings for file type org" not in final_logs
|
2023-05-12 09:37:34 +00:00
|
|
|
|
|
|
|
|
2021-10-03 03:28:33 +00:00
|
|
|
# ----------------------------------------------------------------------------------------------------
|
2023-10-26 16:42:29 +00:00
|
|
|
@pytest.mark.django_db
|
2023-06-29 03:11:26 +00:00
|
|
|
@pytest.mark.anyio
|
2023-10-26 16:42:29 +00:00
|
|
|
# @pytest.mark.asyncio
|
|
|
|
async def test_text_search(search_config: SearchConfig):
|
|
|
|
# Arrange
|
|
|
|
default_user = await KhojUser.objects.acreate(
|
|
|
|
username="test_user", password="test_password", email="test@example.com"
|
|
|
|
)
|
2021-10-03 03:28:33 +00:00
|
|
|
# Arrange
|
2023-10-26 16:42:29 +00:00
|
|
|
org_config = await LocalOrgConfig.objects.acreate(
|
|
|
|
input_files=None,
|
|
|
|
input_filter=["tests/data/org/*.org"],
|
|
|
|
index_heading_entries=False,
|
|
|
|
user=default_user,
|
|
|
|
)
|
|
|
|
data = get_org_files(org_config)
|
2023-08-31 19:55:17 +00:00
|
|
|
|
2023-10-26 16:42:29 +00:00
|
|
|
loop = asyncio.get_event_loop()
|
|
|
|
await loop.run_in_executor(
|
|
|
|
None,
|
|
|
|
text_search.setup,
|
2023-11-01 21:51:33 +00:00
|
|
|
OrgToEntries,
|
2023-10-26 16:42:29 +00:00
|
|
|
data,
|
|
|
|
True,
|
|
|
|
True,
|
|
|
|
default_user,
|
2023-07-14 08:19:38 +00:00
|
|
|
)
|
2023-10-26 16:42:29 +00:00
|
|
|
|
2021-10-03 03:28:33 +00:00
|
|
|
query = "How to git install application?"
|
|
|
|
|
|
|
|
# Act
|
2023-10-26 16:42:29 +00:00
|
|
|
hits = await text_search.query(default_user, query)
|
2021-10-03 03:28:33 +00:00
|
|
|
|
|
|
|
# Assert
|
2023-10-26 16:42:29 +00:00
|
|
|
results = text_search.collate_results(hits)
|
|
|
|
results = sorted(results, key=lambda x: float(x.score))[:1]
|
2023-07-28 22:33:22 +00:00
|
|
|
# search results should contain "git clone" entry
|
2022-09-15 20:34:43 +00:00
|
|
|
search_result = results[0].entry
|
2021-10-03 03:28:33 +00:00
|
|
|
assert "git clone" in search_result
|
2022-06-29 19:09:48 +00:00
|
|
|
|
|
|
|
|
2022-12-23 18:52:02 +00:00
|
|
|
# ----------------------------------------------------------------------------------------------------
|
2023-10-26 16:42:29 +00:00
|
|
|
@pytest.mark.django_db
|
|
|
|
def test_entry_chunking_by_max_tokens(org_config_with_only_new_file: LocalOrgConfig, default_user: KhojUser, caplog):
|
2022-12-23 18:52:02 +00:00
|
|
|
# Arrange
|
2022-12-26 00:45:40 +00:00
|
|
|
# Insert org-mode entry with size exceeding max token limit to new org file
|
2022-12-23 18:52:02 +00:00
|
|
|
max_tokens = 256
|
2023-01-09 19:17:36 +00:00
|
|
|
new_file_to_index = Path(org_config_with_only_new_file.input_files[0])
|
|
|
|
with open(new_file_to_index, "w") as f:
|
2022-12-23 18:52:02 +00:00
|
|
|
f.write(f"* Entry more than {max_tokens} words\n")
|
2023-02-17 16:04:26 +00:00
|
|
|
for index in range(max_tokens + 1):
|
2022-12-23 18:52:02 +00:00
|
|
|
f.write(f"{index} ")
|
|
|
|
|
2023-08-31 19:55:17 +00:00
|
|
|
data = get_org_files(org_config_with_only_new_file)
|
|
|
|
|
2022-12-23 18:52:02 +00:00
|
|
|
# Act
|
|
|
|
# reload embeddings, entries, notes model after adding new org-mode file
|
2023-10-26 16:42:29 +00:00
|
|
|
with caplog.at_level(logging.INFO):
|
2023-11-01 21:51:33 +00:00
|
|
|
text_search.setup(OrgToEntries, data, regenerate=False, user=default_user)
|
2022-12-23 18:52:02 +00:00
|
|
|
|
|
|
|
# Assert
|
|
|
|
# verify newly added org-mode entry is split by max tokens
|
2023-11-01 00:59:53 +00:00
|
|
|
assert "Created 2 new embeddings. Deleted 0 embeddings for user " in caplog.records[-1].message
|
2022-12-23 18:52:02 +00:00
|
|
|
|
|
|
|
|
2023-09-06 19:04:18 +00:00
|
|
|
# ----------------------------------------------------------------------------------------------------
|
2023-10-26 16:42:29 +00:00
|
|
|
@pytest.mark.django_db
|
2023-09-06 19:04:18 +00:00
|
|
|
def test_entry_chunking_by_max_tokens_not_full_corpus(
|
2023-10-26 16:42:29 +00:00
|
|
|
org_config_with_only_new_file: LocalOrgConfig, default_user: KhojUser, caplog
|
2023-09-06 19:04:18 +00:00
|
|
|
):
|
|
|
|
# Arrange
|
|
|
|
# Insert org-mode entry with size exceeding max token limit to new org file
|
|
|
|
data = {
|
|
|
|
"readme.org": """
|
|
|
|
* Khoj
|
2023-10-26 16:42:29 +00:00
|
|
|
/Allow natural language search on user content like notes, images using transformer based models/
|
2023-09-06 19:04:18 +00:00
|
|
|
|
2023-10-26 16:42:29 +00:00
|
|
|
All data is processed locally. User can interface with khoj app via [[./interface/emacs/khoj.el][Emacs]], API or Commandline
|
2023-09-06 19:04:18 +00:00
|
|
|
|
|
|
|
** Dependencies
|
2023-10-26 16:42:29 +00:00
|
|
|
- Python3
|
|
|
|
- [[https://docs.conda.io/en/latest/miniconda.html#latest-miniconda-installer-links][Miniconda]]
|
2023-09-06 19:04:18 +00:00
|
|
|
|
|
|
|
** Install
|
2023-10-26 16:42:29 +00:00
|
|
|
#+begin_src shell
|
|
|
|
git clone https://github.com/khoj-ai/khoj && cd khoj
|
|
|
|
conda env create -f environment.yml
|
|
|
|
conda activate khoj
|
|
|
|
#+end_src"""
|
2023-09-06 19:04:18 +00:00
|
|
|
}
|
|
|
|
text_search.setup(
|
2023-11-01 21:51:33 +00:00
|
|
|
OrgToEntries,
|
2023-09-06 19:04:18 +00:00
|
|
|
data,
|
|
|
|
regenerate=False,
|
2023-10-26 16:42:29 +00:00
|
|
|
user=default_user,
|
2023-09-06 19:04:18 +00:00
|
|
|
)
|
|
|
|
|
|
|
|
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} ")
|
|
|
|
|
|
|
|
data = get_org_files(org_config_with_only_new_file)
|
|
|
|
|
|
|
|
# Act
|
|
|
|
# reload embeddings, entries, notes model after adding new org-mode file
|
2023-10-26 16:42:29 +00:00
|
|
|
with caplog.at_level(logging.INFO):
|
|
|
|
text_search.setup(
|
2023-11-01 21:51:33 +00:00
|
|
|
OrgToEntries,
|
2023-10-26 16:42:29 +00:00
|
|
|
data,
|
|
|
|
regenerate=False,
|
|
|
|
full_corpus=False,
|
|
|
|
user=default_user,
|
|
|
|
)
|
|
|
|
|
2023-09-06 19:04:18 +00:00
|
|
|
# Assert
|
|
|
|
# verify newly added org-mode entry is split by max tokens
|
2023-11-01 00:59:53 +00:00
|
|
|
assert "Created 2 new embeddings. Deleted 0 embeddings for user " in caplog.records[-1].message
|
2023-09-06 19:04:18 +00:00
|
|
|
|
|
|
|
|
2022-06-29 19:09:48 +00:00
|
|
|
# ----------------------------------------------------------------------------------------------------
|
2023-10-26 16:42:29 +00:00
|
|
|
@pytest.mark.django_db
|
2023-07-16 07:47:11 +00:00
|
|
|
def test_regenerate_index_with_new_entry(
|
2023-10-26 16:42:29 +00:00
|
|
|
content_config: ContentConfig, new_org_file: Path, default_user: KhojUser, caplog
|
2023-07-16 07:47:11 +00:00
|
|
|
):
|
2022-06-29 19:09:48 +00:00
|
|
|
# Arrange
|
2023-10-26 16:42:29 +00:00
|
|
|
org_config = LocalOrgConfig.objects.filter(user=default_user).first()
|
|
|
|
data = get_org_files(org_config)
|
|
|
|
|
|
|
|
with caplog.at_level(logging.INFO):
|
2023-11-01 21:51:33 +00:00
|
|
|
text_search.setup(OrgToEntries, data, regenerate=True, user=default_user)
|
2022-06-29 19:09:48 +00:00
|
|
|
|
2023-11-01 00:59:53 +00:00
|
|
|
assert "Created 10 new embeddings. Deleted 3 embeddings for user " in caplog.records[-1].message
|
2022-06-29 19:09:48 +00:00
|
|
|
|
2023-01-09 19:17:36 +00:00
|
|
|
# append org-mode entry to first org input file in config
|
2023-10-26 16:42:29 +00:00
|
|
|
org_config.input_files = [f"{new_org_file}"]
|
2023-01-09 19:17:36 +00:00
|
|
|
with open(new_org_file, "w") as f:
|
2022-06-29 19:09:48 +00:00
|
|
|
f.write("\n* A Chihuahua doing Tango\n- Saw a super cute video of a chihuahua doing the Tango on Youtube\n")
|
|
|
|
|
2023-10-26 16:42:29 +00:00
|
|
|
data = get_org_files(org_config)
|
2023-08-31 19:55:17 +00:00
|
|
|
|
2023-07-16 07:47:11 +00:00
|
|
|
# Act
|
2022-06-29 19:09:48 +00:00
|
|
|
# regenerate notes jsonl, model embeddings and model to include entry from new file
|
2023-10-26 16:42:29 +00:00
|
|
|
with caplog.at_level(logging.INFO):
|
2023-11-01 21:51:33 +00:00
|
|
|
text_search.setup(OrgToEntries, data, regenerate=True, user=default_user)
|
2022-06-29 19:09:48 +00:00
|
|
|
|
|
|
|
# Assert
|
2023-10-26 16:42:29 +00:00
|
|
|
assert "Created 11 new embeddings. Deleted 10 embeddings for user " in caplog.records[-1].message
|
|
|
|
verify_embeddings(11, default_user)
|
2022-09-07 00:06:29 +00:00
|
|
|
|
|
|
|
|
2023-07-16 00:18:07 +00:00
|
|
|
# ----------------------------------------------------------------------------------------------------
|
2023-10-26 16:42:29 +00:00
|
|
|
@pytest.mark.django_db
|
2023-07-16 00:18:07 +00:00
|
|
|
def test_update_index_with_duplicate_entries_in_stable_order(
|
2023-10-26 16:42:29 +00:00
|
|
|
org_config_with_only_new_file: LocalOrgConfig, default_user: KhojUser, caplog
|
2023-07-16 00:18:07 +00:00
|
|
|
):
|
|
|
|
# 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}")
|
|
|
|
|
2023-08-31 19:55:17 +00:00
|
|
|
data = get_org_files(org_config_with_only_new_file)
|
|
|
|
|
2023-07-16 00:18:07 +00:00
|
|
|
# Act
|
2023-11-01 00:59:53 +00:00
|
|
|
# generate embeddings, entries, notes model from scratch after adding new org-mode file
|
2023-10-26 16:42:29 +00:00
|
|
|
with caplog.at_level(logging.INFO):
|
2023-11-01 21:51:33 +00:00
|
|
|
text_search.setup(OrgToEntries, data, regenerate=True, user=default_user)
|
2023-11-01 00:59:53 +00:00
|
|
|
initial_logs = caplog.text
|
|
|
|
caplog.clear() # Clear logs
|
2023-07-16 00:18:07 +00:00
|
|
|
|
2023-08-31 19:55:17 +00:00
|
|
|
data = get_org_files(org_config_with_only_new_file)
|
|
|
|
|
2023-11-01 00:59:53 +00:00
|
|
|
# update embeddings, entries, notes model with no new changes
|
2023-10-26 16:42:29 +00:00
|
|
|
with caplog.at_level(logging.INFO):
|
2023-11-01 21:51:33 +00:00
|
|
|
text_search.setup(OrgToEntries, data, regenerate=False, user=default_user)
|
2023-11-01 00:59:53 +00:00
|
|
|
final_logs = caplog.text
|
2023-07-16 00:18:07 +00:00
|
|
|
|
|
|
|
# Assert
|
|
|
|
# verify only 1 entry added even if there are multiple duplicate entries
|
2023-11-01 00:59:53 +00:00
|
|
|
assert "Created 1 new embeddings. Deleted 3 embeddings for user " in initial_logs
|
|
|
|
assert "Created 0 new embeddings. Deleted 0 embeddings for user " in final_logs
|
2023-07-16 00:18:07 +00:00
|
|
|
|
2023-10-26 16:42:29 +00:00
|
|
|
verify_embeddings(1, default_user)
|
2023-07-16 00:18:07 +00:00
|
|
|
|
|
|
|
|
2023-07-16 10:47:05 +00:00
|
|
|
# ----------------------------------------------------------------------------------------------------
|
2023-10-26 16:42:29 +00:00
|
|
|
@pytest.mark.django_db
|
|
|
|
def test_update_index_with_deleted_entry(org_config_with_only_new_file: LocalOrgConfig, default_user: KhojUser, caplog):
|
2023-07-16 10:47:05 +00:00
|
|
|
# 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")
|
2023-08-31 19:55:17 +00:00
|
|
|
data = get_org_files(org_config_with_only_new_file)
|
2023-07-16 10:47:05 +00:00
|
|
|
|
|
|
|
# load embeddings, entries, notes model after adding new org file with 2 entries
|
2023-10-26 16:42:29 +00:00
|
|
|
with caplog.at_level(logging.INFO):
|
2023-11-01 21:51:33 +00:00
|
|
|
text_search.setup(OrgToEntries, data, regenerate=True, user=default_user)
|
2023-11-01 00:59:53 +00:00
|
|
|
initial_logs = caplog.text
|
|
|
|
caplog.clear() # Clear logs
|
2023-07-16 10:47:05 +00:00
|
|
|
|
|
|
|
# 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}")
|
|
|
|
|
2023-08-31 19:55:17 +00:00
|
|
|
data = get_org_files(org_config_with_only_new_file)
|
|
|
|
|
2023-07-16 10:47:05 +00:00
|
|
|
# Act
|
2023-10-26 16:42:29 +00:00
|
|
|
with caplog.at_level(logging.INFO):
|
2023-11-01 21:51:33 +00:00
|
|
|
text_search.setup(OrgToEntries, data, regenerate=False, user=default_user)
|
2023-11-01 00:59:53 +00:00
|
|
|
final_logs = caplog.text
|
2023-07-16 10:47:05 +00:00
|
|
|
|
|
|
|
# Assert
|
|
|
|
# verify only 1 entry added even if there are multiple duplicate entries
|
2023-11-01 00:59:53 +00:00
|
|
|
assert "Created 2 new embeddings. Deleted 3 embeddings for user " in initial_logs
|
|
|
|
assert "Created 0 new embeddings. Deleted 1 embeddings for user " in final_logs
|
2023-07-16 10:47:05 +00:00
|
|
|
|
2023-10-26 16:42:29 +00:00
|
|
|
verify_embeddings(1, default_user)
|
2023-07-16 10:47:05 +00:00
|
|
|
|
|
|
|
|
2022-09-07 00:06:29 +00:00
|
|
|
# ----------------------------------------------------------------------------------------------------
|
2023-10-26 16:42:29 +00:00
|
|
|
@pytest.mark.django_db
|
|
|
|
def test_update_index_with_new_entry(content_config: ContentConfig, new_org_file: Path, default_user: KhojUser, caplog):
|
2022-09-07 00:06:29 +00:00
|
|
|
# Arrange
|
2023-10-26 16:42:29 +00:00
|
|
|
org_config = LocalOrgConfig.objects.filter(user=default_user).first()
|
|
|
|
data = get_org_files(org_config)
|
|
|
|
with caplog.at_level(logging.INFO):
|
2023-11-01 21:51:33 +00:00
|
|
|
text_search.setup(OrgToEntries, data, regenerate=True, user=default_user)
|
2023-11-01 00:59:53 +00:00
|
|
|
initial_logs = caplog.text
|
|
|
|
caplog.clear() # Clear logs
|
2022-09-07 00:06:29 +00:00
|
|
|
|
2023-01-09 19:17:36 +00:00
|
|
|
# append org-mode entry to first org input file in config
|
|
|
|
with open(new_org_file, "w") as f:
|
2023-07-16 08:24:03 +00:00
|
|
|
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)
|
2022-09-07 00:06:29 +00:00
|
|
|
|
2023-10-26 16:42:29 +00:00
|
|
|
data = get_org_files(org_config)
|
2023-08-31 19:55:17 +00:00
|
|
|
|
2022-09-07 00:06:29 +00:00
|
|
|
# Act
|
|
|
|
# update embeddings, entries with the newly added note
|
2023-10-26 16:42:29 +00:00
|
|
|
with caplog.at_level(logging.INFO):
|
2023-11-01 21:51:33 +00:00
|
|
|
text_search.setup(OrgToEntries, data, regenerate=False, user=default_user)
|
2023-11-01 00:59:53 +00:00
|
|
|
final_logs = caplog.text
|
2022-09-07 00:06:29 +00:00
|
|
|
|
2023-01-09 19:17:36 +00:00
|
|
|
# Assert
|
2023-11-01 00:59:53 +00:00
|
|
|
assert "Created 10 new embeddings. Deleted 3 embeddings for user " in initial_logs
|
|
|
|
assert "Created 1 new embeddings. Deleted 0 embeddings for user " in final_logs
|
2022-09-07 00:06:29 +00:00
|
|
|
|
2023-10-26 16:42:29 +00:00
|
|
|
verify_embeddings(11, default_user)
|
2023-06-13 23:55:58 +00:00
|
|
|
|
|
|
|
|
|
|
|
# ----------------------------------------------------------------------------------------------------
|
|
|
|
@pytest.mark.skipif(os.getenv("GITHUB_PAT_TOKEN") is None, reason="GITHUB_PAT_TOKEN not set")
|
2023-10-26 16:42:29 +00:00
|
|
|
def test_text_search_setup_github(content_config: ContentConfig, default_user: KhojUser):
|
|
|
|
# Arrange
|
|
|
|
github_config = GithubConfig.objects.filter(user=default_user).first()
|
2023-06-13 23:55:58 +00:00
|
|
|
# Act
|
2023-06-18 09:24:25 +00:00
|
|
|
# Regenerate github embeddings to test asymmetric setup without caching
|
2023-10-26 16:42:29 +00:00
|
|
|
text_search.setup(
|
2023-11-01 21:51:33 +00:00
|
|
|
GithubToEntries,
|
2023-10-26 16:42:29 +00:00
|
|
|
{},
|
|
|
|
regenerate=True,
|
|
|
|
user=default_user,
|
|
|
|
config=github_config,
|
2023-07-14 08:19:38 +00:00
|
|
|
)
|
2023-06-13 23:55:58 +00:00
|
|
|
|
|
|
|
# Assert
|
2023-11-01 01:50:54 +00:00
|
|
|
embeddings = Entry.objects.filter(user=default_user, file_type="github").count()
|
2023-10-26 16:42:29 +00:00
|
|
|
assert embeddings > 1
|
2023-07-15 21:33:15 +00:00
|
|
|
|
|
|
|
|
2023-10-26 16:42:29 +00:00
|
|
|
def verify_embeddings(expected_count, user):
|
2023-11-01 01:50:54 +00:00
|
|
|
embeddings = Entry.objects.filter(user=user, file_type="org").count()
|
2023-10-26 16:42:29 +00:00
|
|
|
assert embeddings == expected_count
|