mirror of
https://github.com/khoj-ai/khoj.git
synced 2024-11-27 17:35:07 +01:00
Setup Search with Search_Config to Fix Tests
- Rename pytest fixture search_config to more appropriate content_config - Create search_config pytest fixture - Use search_config where search being setup, used in tests
This commit is contained in:
Debanjum Singh Solanky
parent
c64e0c2965
commit
ed144f7984
4 changed files with 80 additions and 46 deletions
|
|
@ -1,51 +1,78 @@
|
|||
# Standard Packages
|
||||
import pytest
|
||||
from pathlib import Path
|
||||
from src import search_type
|
||||
|
||||
# Internal Packages
|
||||
from src.search_type import asymmetric, image_search
|
||||
from src.utils.rawconfig import ContentTypeConfig, ImageSearchConfig, TextSearchConfig
|
||||
from src.utils.rawconfig import AsymmetricConfig, ContentTypeConfig, ImageSearchConfig, ImageSearchTypeConfig, SearchTypeConfig, SymmetricConfig, TextSearchConfig
|
||||
|
||||
|
||||
@pytest.fixture(scope='session')
|
||||
def model_dir(tmp_path_factory):
|
||||
def search_config(tmp_path_factory):
|
||||
model_dir = tmp_path_factory.mktemp('data')
|
||||
|
||||
search_config = SearchTypeConfig()
|
||||
|
||||
search_config.asymmetric = SymmetricConfig(
|
||||
encoder = "sentence-transformers/paraphrase-MiniLM-L6-v2",
|
||||
cross_encoder = "cross-encoder/ms-marco-MiniLM-L-6-v2",
|
||||
model_directory = model_dir
|
||||
)
|
||||
|
||||
search_config.asymmetric = AsymmetricConfig(
|
||||
encoder = "sentence-transformers/msmarco-MiniLM-L-6-v3",
|
||||
cross_encoder = "cross-encoder/ms-marco-MiniLM-L-6-v2",
|
||||
model_directory = model_dir
|
||||
)
|
||||
|
||||
search_config.image = ImageSearchTypeConfig(
|
||||
encoder = "clip-ViT-B-32",
|
||||
model_directory = model_dir
|
||||
)
|
||||
|
||||
return search_config
|
||||
|
||||
|
||||
@pytest.fixture(scope='session')
|
||||
def model_dir(search_config):
|
||||
model_dir = search_config.asymmetric.model_directory
|
||||
|
||||
# Generate Image Embeddings from Test Images
|
||||
search_config = ContentTypeConfig()
|
||||
search_config.image = ImageSearchConfig(
|
||||
content_config = ContentTypeConfig()
|
||||
content_config.image = ImageSearchConfig(
|
||||
input_directory = 'tests/data',
|
||||
embeddings_file = model_dir.joinpath('.image_embeddings.pt'),
|
||||
batch_size = 10,
|
||||
use_xmp_metadata = False)
|
||||
|
||||
image_search.setup(search_config.image, regenerate=False, verbose=True)
|
||||
image_search.setup(content_config.image, search_config.image, regenerate=False, verbose=True)
|
||||
|
||||
# Generate Notes Embeddings from Test Notes
|
||||
search_config.org = TextSearchConfig(
|
||||
content_config.org = TextSearchConfig(
|
||||
input_files = ['tests/data/main_readme.org', 'tests/data/interface_emacs_readme.org'],
|
||||
input_filter = None,
|
||||
compressed_jsonl = model_dir.joinpath('.notes.jsonl.gz'),
|
||||
embeddings_file = model_dir.joinpath('.note_embeddings.pt'))
|
||||
|
||||
asymmetric.setup(search_config.org, regenerate=False, verbose=True)
|
||||
asymmetric.setup(content_config.org, search_config.asymmetric, regenerate=False, verbose=True)
|
||||
|
||||
return model_dir
|
||||
|
||||
|
||||
@pytest.fixture(scope='session')
|
||||
def search_config(model_dir):
|
||||
search_config = ContentTypeConfig()
|
||||
search_config.org = TextSearchConfig(
|
||||
def content_config(model_dir):
|
||||
content_config = ContentTypeConfig()
|
||||
content_config.org = TextSearchConfig(
|
||||
input_files = ['tests/data/main_readme.org', 'tests/data/interface_emacs_readme.org'],
|
||||
input_filter = None,
|
||||
compressed_jsonl = model_dir.joinpath('.notes.jsonl.gz'),
|
||||
embeddings_file = model_dir.joinpath('.note_embeddings.pt'))
|
||||
|
||||
search_config.image = ImageSearchConfig(
|
||||
content_config.image = ImageSearchConfig(
|
||||
input_directory = 'tests/data',
|
||||
embeddings_file = 'tests/data/.image_embeddings.pt',
|
||||
embeddings_file = model_dir.joinpath('.image_embeddings.pt'),
|
||||
batch_size = 10,
|
||||
use_xmp_metadata = False)
|
||||
|
||||
return search_config
|
||||
return content_config
|
||||
|
|
|
|||
|
|
@ -1,14 +1,15 @@
|
|||
# Internal Packages
|
||||
from src.main import model
|
||||
from src.search_type import asymmetric
|
||||
from src.utils.rawconfig import ContentTypeConfig, SearchTypeConfig
|
||||
|
||||
|
||||
# Test
|
||||
# ----------------------------------------------------------------------------------------------------
|
||||
def test_asymmetric_setup(search_config):
|
||||
def test_asymmetric_setup(content_config: ContentTypeConfig, search_config: SearchTypeConfig):
|
||||
# Act
|
||||
# Regenerate notes embeddings during asymmetric setup
|
||||
notes_model = asymmetric.setup(search_config.org, regenerate=True)
|
||||
notes_model = asymmetric.setup(content_config.org, search_config.asymmetric, regenerate=True)
|
||||
|
||||
# Assert
|
||||
assert len(notes_model.entries) == 10
|
||||
|
|
@ -16,9 +17,9 @@ def test_asymmetric_setup(search_config):
|
|||
|
||||
|
||||
# ----------------------------------------------------------------------------------------------------
|
||||
def test_asymmetric_search(search_config):
|
||||
def test_asymmetric_search(content_config: ContentTypeConfig, search_config: SearchTypeConfig):
|
||||
# Arrange
|
||||
model.notes_search = asymmetric.setup(search_config.org, regenerate=False)
|
||||
model.notes_search = asymmetric.setup(content_config.org, search_config.asymmetric, regenerate=False)
|
||||
query = "How to git install application?"
|
||||
|
||||
# Act
|
||||
|
|
|
|||
|
|
@ -9,7 +9,7 @@ import pytest
|
|||
from src.main import app, model, config
|
||||
from src.search_type import asymmetric, image_search
|
||||
from src.utils.helpers import resolve_absolute_path
|
||||
from src.utils.rawconfig import ContentTypeConfig
|
||||
from src.utils.rawconfig import ContentTypeConfig, SearchTypeConfig
|
||||
|
||||
|
||||
# Arrange
|
||||
|
|
@ -18,55 +18,60 @@ client = TestClient(app)
|
|||
|
||||
# Test
|
||||
# ----------------------------------------------------------------------------------------------------
|
||||
def test_search_with_invalid_search_type():
|
||||
def test_search_with_invalid_content_type():
|
||||
# Arrange
|
||||
user_query = "How to call semantic search from Emacs?"
|
||||
|
||||
# Act
|
||||
response = client.get(f"/search?q={user_query}&t=invalid_search_type")
|
||||
response = client.get(f"/search?q={user_query}&t=invalid_content_type")
|
||||
|
||||
# Assert
|
||||
assert response.status_code == 422
|
||||
|
||||
|
||||
# ----------------------------------------------------------------------------------------------------
|
||||
def test_search_with_valid_search_type(search_config: ContentTypeConfig):
|
||||
def test_search_with_valid_content_type(content_config: ContentTypeConfig, search_config: SearchTypeConfig):
|
||||
# Arrange
|
||||
config.content_type = search_config
|
||||
config.content_type = content_config
|
||||
config.search_type = search_config
|
||||
|
||||
# config.content_type.image = search_config.image
|
||||
for search_type in ["notes", "ledger", "music", "image"]:
|
||||
for content_type in ["notes", "ledger", "music", "image"]:
|
||||
# Act
|
||||
response = client.get(f"/search?q=random&t={search_type}")
|
||||
response = client.get(f"/search?q=random&t={content_type}")
|
||||
# Assert
|
||||
assert response.status_code == 200
|
||||
|
||||
|
||||
# ----------------------------------------------------------------------------------------------------
|
||||
def test_regenerate_with_invalid_search_type():
|
||||
def test_regenerate_with_invalid_content_type():
|
||||
# Act
|
||||
response = client.get(f"/regenerate?t=invalid_search_type")
|
||||
response = client.get(f"/regenerate?t=invalid_content_type")
|
||||
|
||||
# Assert
|
||||
assert response.status_code == 422
|
||||
|
||||
|
||||
# ----------------------------------------------------------------------------------------------------
|
||||
def test_regenerate_with_valid_search_type(search_config: ContentTypeConfig):
|
||||
def test_regenerate_with_valid_content_type(content_config: ContentTypeConfig, search_config: SearchTypeConfig):
|
||||
# Arrange
|
||||
config.content_type = search_config
|
||||
for search_type in ["notes", "ledger", "music", "image"]:
|
||||
config.content_type = content_config
|
||||
config.search_type = search_config
|
||||
|
||||
for content_type in ["notes", "ledger", "music", "image"]:
|
||||
# Act
|
||||
response = client.get(f"/regenerate?t={search_type}")
|
||||
response = client.get(f"/regenerate?t={content_type}")
|
||||
# Assert
|
||||
assert response.status_code == 200
|
||||
|
||||
|
||||
# ----------------------------------------------------------------------------------------------------
|
||||
@pytest.mark.skip(reason="Flaky test. Search doesn't always return expected image path.")
|
||||
def test_image_search(search_config: ContentTypeConfig):
|
||||
def test_image_search(content_config: ContentTypeConfig, search_config: SearchTypeConfig):
|
||||
# Arrange
|
||||
config.content_type = search_config
|
||||
model.image_search = image_search.setup(search_config.image, regenerate=False)
|
||||
config.content_type = content_config
|
||||
config.search_type = search_config
|
||||
model.image_search = image_search.setup(content_config.image, search_config.image, regenerate=False)
|
||||
query_expected_image_pairs = [("brown kitten next to fallen plant", "kitten_park.jpg"),
|
||||
("a horse and dog on a leash", "horse_dog.jpg"),
|
||||
("A guinea pig eating grass", "guineapig_grass.jpg")]
|
||||
|
|
@ -78,16 +83,16 @@ def test_image_search(search_config: ContentTypeConfig):
|
|||
# Assert
|
||||
assert response.status_code == 200
|
||||
actual_image = Path(response.json()[0]["Entry"])
|
||||
expected_image = resolve_absolute_path(search_config.image.input_directory.joinpath(expected_image_name))
|
||||
expected_image = resolve_absolute_path(content_config.image.input_directory.joinpath(expected_image_name))
|
||||
|
||||
# Assert
|
||||
assert expected_image == actual_image
|
||||
|
||||
|
||||
# ----------------------------------------------------------------------------------------------------
|
||||
def test_notes_search(search_config: ContentTypeConfig):
|
||||
def test_notes_search(content_config: ContentTypeConfig, search_config: SearchTypeConfig):
|
||||
# Arrange
|
||||
model.notes_search = asymmetric.setup(search_config.org, regenerate=False)
|
||||
model.notes_search = asymmetric.setup(content_config.org, search_config.asymmetric, regenerate=False)
|
||||
user_query = "How to git install application?"
|
||||
|
||||
# Act
|
||||
|
|
@ -101,9 +106,9 @@ def test_notes_search(search_config: ContentTypeConfig):
|
|||
|
||||
|
||||
# ----------------------------------------------------------------------------------------------------
|
||||
def test_notes_search_with_include_filter(search_config: ContentTypeConfig):
|
||||
def test_notes_search_with_include_filter(content_config: ContentTypeConfig, search_config: SearchTypeConfig):
|
||||
# Arrange
|
||||
model.notes_search = asymmetric.setup(search_config.org, regenerate=False)
|
||||
model.notes_search = asymmetric.setup(content_config.org, search_config.asymmetric, regenerate=False)
|
||||
user_query = "How to git install application? +Emacs"
|
||||
|
||||
# Act
|
||||
|
|
@ -117,9 +122,9 @@ def test_notes_search_with_include_filter(search_config: ContentTypeConfig):
|
|||
|
||||
|
||||
# ----------------------------------------------------------------------------------------------------
|
||||
def test_notes_search_with_exclude_filter(search_config: ContentTypeConfig):
|
||||
def test_notes_search_with_exclude_filter(content_config: ContentTypeConfig, search_config: SearchTypeConfig):
|
||||
# Arrange
|
||||
model.notes_search = asymmetric.setup(search_config.org, regenerate=False)
|
||||
model.notes_search = asymmetric.setup(content_config.org, search_config.asymmetric, regenerate=False)
|
||||
user_query = "How to git install application? -clone"
|
||||
|
||||
# Act
|
||||
|
|
|
|||
|
|
@ -5,14 +5,15 @@ import pytest
|
|||
from src.main import model
|
||||
from src.search_type import image_search
|
||||
from src.utils.helpers import resolve_absolute_path
|
||||
from src.utils.rawconfig import ContentTypeConfig, SearchTypeConfig
|
||||
|
||||
|
||||
# Test
|
||||
# ----------------------------------------------------------------------------------------------------
|
||||
def test_image_search_setup(search_config):
|
||||
def test_image_search_setup(content_config: ContentTypeConfig, search_config: SearchTypeConfig):
|
||||
# Act
|
||||
# Regenerate image search embeddings during image setup
|
||||
image_search_model = image_search.setup(search_config.image, regenerate=True)
|
||||
image_search_model = image_search.setup(content_config.image, search_config.image, regenerate=True)
|
||||
|
||||
# Assert
|
||||
assert len(image_search_model.image_names) == 3
|
||||
|
|
@ -21,9 +22,9 @@ def test_image_search_setup(search_config):
|
|||
|
||||
# ----------------------------------------------------------------------------------------------------
|
||||
@pytest.mark.skip(reason="results inconsistent currently")
|
||||
def test_image_search(search_config):
|
||||
def test_image_search(content_config: ContentTypeConfig, search_config: SearchTypeConfig):
|
||||
# Arrange
|
||||
model.image_search = image_search.setup(search_config.image, regenerate=False)
|
||||
model.image_search = image_search.setup(content_config.image, search_config.image, regenerate=False)
|
||||
query_expected_image_pairs = [("brown kitten next to plant", "kitten_park.jpg"),
|
||||
("horse and dog in a farm", "horse_dog.jpg"),
|
||||
("A guinea pig eating grass", "guineapig_grass.jpg")]
|
||||
|
|
@ -38,11 +39,11 @@ def test_image_search(search_config):
|
|||
results = image_search.collate_results(
|
||||
hits,
|
||||
model.image_search.image_names,
|
||||
search_config.image.input_directory,
|
||||
content_config.image.input_directory,
|
||||
count=1)
|
||||
|
||||
actual_image = results[0]["Entry"]
|
||||
expected_image = resolve_absolute_path(search_config.image.input_directory.joinpath(expected_image_name))
|
||||
expected_image = resolve_absolute_path(content_config.image.input_directory.joinpath(expected_image_name))
|
||||
|
||||
# Assert
|
||||
assert expected_image == actual_image
|
||||
|
|
|
|||
Loading…
Reference in a new issue