mirror of
https://github.com/khoaliber/khoj.git
synced 2026-03-02 13:18:18 +00:00
[Multi-User Part 1]: Enable storage of settings for plaintext files based on user account (#498)
- Partition configuration for indexing local data based on user accounts - Store indexed data in an underlying postgres db using the `pgvector` extension - Add migrations for all relevant user data and embeddings generation. Very little performance optimization has been done for the lookup time - Apply filters using SQL queries - Start removing many server-level configuration settings - Configure GitHub test actions to run during any PR. Update the test action to run in a containerized environment with a DB. - Update the Docker image and docker-compose.yml to work with the new application design
This commit is contained in:
sabaimran
@@ -2,22 +2,21 @@
|
||||
from io import BytesIO
|
||||
from PIL import Image
|
||||
from urllib.parse import quote
|
||||
|
||||
import pytest
|
||||
|
||||
# External Packages
|
||||
from fastapi.testclient import TestClient
|
||||
import pytest
|
||||
from fastapi import FastAPI
|
||||
|
||||
# Internal Packages
|
||||
from app.main import app
|
||||
from khoj.configure import configure_routes, configure_search_types
|
||||
from khoj.utils import state
|
||||
from khoj.utils.state import search_models, content_index, config
|
||||
from khoj.search_type import text_search, image_search
|
||||
from khoj.utils.rawconfig import ContentConfig, SearchConfig
|
||||
from khoj.processor.org_mode.org_to_jsonl import OrgToJsonl
|
||||
from khoj.search_filter.word_filter import WordFilter
|
||||
from khoj.search_filter.file_filter import FileFilter
|
||||
from database.models import KhojUser
|
||||
from database.adapters import EmbeddingsAdapters
|
||||
|
||||
|
||||
# Test
|
||||
@@ -35,7 +34,7 @@ def test_search_with_invalid_content_type(client):
|
||||
|
||||
# ----------------------------------------------------------------------------------------------------
|
||||
def test_search_with_valid_content_type(client):
|
||||
for content_type in ["all", "org", "markdown", "image", "pdf", "github", "notion", "plugin1"]:
|
||||
for content_type in ["all", "org", "markdown", "image", "pdf", "github", "notion"]:
|
||||
# Act
|
||||
response = client.get(f"/api/search?q=random&t={content_type}")
|
||||
# Assert
|
||||
@@ -75,7 +74,7 @@ def test_index_update(client):
|
||||
|
||||
# ----------------------------------------------------------------------------------------------------
|
||||
def test_regenerate_with_valid_content_type(client):
|
||||
for content_type in ["all", "org", "markdown", "image", "pdf", "notion", "plugin1"]:
|
||||
for content_type in ["all", "org", "markdown", "image", "pdf", "notion"]:
|
||||
# Arrange
|
||||
files = get_sample_files_data()
|
||||
headers = {"x-api-key": "secret"}
|
||||
@@ -102,60 +101,42 @@ def test_regenerate_with_github_fails_without_pat(client):
|
||||
|
||||
|
||||
# ----------------------------------------------------------------------------------------------------
|
||||
@pytest.mark.django_db
|
||||
@pytest.mark.skip(reason="Flaky test on parallel test runs")
|
||||
def test_get_configured_types_via_api(client):
|
||||
def test_get_configured_types_via_api(client, sample_org_data):
|
||||
# Act
|
||||
response = client.get(f"/api/config/types")
|
||||
text_search.setup(OrgToJsonl, sample_org_data, regenerate=False)
|
||||
|
||||
enabled_types = EmbeddingsAdapters.get_unique_file_types(user=None).all().values_list("file_type", flat=True)
|
||||
|
||||
# Assert
|
||||
assert response.status_code == 200
|
||||
assert response.json() == ["all", "org", "image", "plaintext", "plugin1"]
|
||||
assert list(enabled_types) == ["org"]
|
||||
|
||||
|
||||
# ----------------------------------------------------------------------------------------------------
|
||||
def test_get_configured_types_with_only_plugin_content_config(content_config):
|
||||
@pytest.mark.django_db(transaction=True)
|
||||
def test_get_api_config_types(client, search_config: SearchConfig, sample_org_data):
|
||||
# Arrange
|
||||
config.content_type = ContentConfig()
|
||||
config.content_type.plugins = content_config.plugins
|
||||
state.SearchType = configure_search_types(config)
|
||||
|
||||
configure_routes(app)
|
||||
client = TestClient(app)
|
||||
text_search.setup(OrgToJsonl, sample_org_data, regenerate=False)
|
||||
|
||||
# Act
|
||||
response = client.get(f"/api/config/types")
|
||||
|
||||
# Assert
|
||||
assert response.status_code == 200
|
||||
assert response.json() == ["all", "plugin1"]
|
||||
assert response.json() == ["all", "org", "markdown", "image"]
|
||||
|
||||
|
||||
# ----------------------------------------------------------------------------------------------------
|
||||
def test_get_configured_types_with_no_plugin_content_config(content_config):
|
||||
@pytest.mark.django_db(transaction=True)
|
||||
def test_get_configured_types_with_no_content_config(fastapi_app: FastAPI):
|
||||
# Arrange
|
||||
config.content_type = content_config
|
||||
config.content_type.plugins = None
|
||||
state.SearchType = configure_search_types(config)
|
||||
original_config = state.config.content_type
|
||||
state.config.content_type = None
|
||||
|
||||
configure_routes(app)
|
||||
client = TestClient(app)
|
||||
|
||||
# Act
|
||||
response = client.get(f"/api/config/types")
|
||||
|
||||
# Assert
|
||||
assert response.status_code == 200
|
||||
assert "plugin1" not in response.json()
|
||||
|
||||
|
||||
# ----------------------------------------------------------------------------------------------------
|
||||
def test_get_configured_types_with_no_content_config():
|
||||
# Arrange
|
||||
config.content_type = ContentConfig()
|
||||
state.SearchType = configure_search_types(config)
|
||||
|
||||
configure_routes(app)
|
||||
client = TestClient(app)
|
||||
configure_routes(fastapi_app)
|
||||
client = TestClient(fastapi_app)
|
||||
|
||||
# Act
|
||||
response = client.get(f"/api/config/types")
|
||||
@@ -164,6 +145,9 @@ def test_get_configured_types_with_no_content_config():
|
||||
assert response.status_code == 200
|
||||
assert response.json() == ["all"]
|
||||
|
||||
# Restore
|
||||
state.config.content_type = original_config
|
||||
|
||||
|
||||
# ----------------------------------------------------------------------------------------------------
|
||||
def test_image_search(client, content_config: ContentConfig, search_config: SearchConfig):
|
||||
@@ -192,12 +176,10 @@ def test_image_search(client, content_config: ContentConfig, search_config: Sear
|
||||
|
||||
|
||||
# ----------------------------------------------------------------------------------------------------
|
||||
def test_notes_search(client, content_config: ContentConfig, search_config: SearchConfig, sample_org_data):
|
||||
@pytest.mark.django_db(transaction=True)
|
||||
def test_notes_search(client, search_config: SearchConfig, sample_org_data):
|
||||
# Arrange
|
||||
search_models.text_search = text_search.initialize_model(search_config.asymmetric)
|
||||
content_index.org = text_search.setup(
|
||||
OrgToJsonl, sample_org_data, content_config.org, search_models.text_search.bi_encoder, regenerate=False
|
||||
)
|
||||
text_search.setup(OrgToJsonl, sample_org_data, regenerate=False)
|
||||
user_query = quote("How to git install application?")
|
||||
|
||||
# Act
|
||||
@@ -211,19 +193,15 @@ def test_notes_search(client, content_config: ContentConfig, search_config: Sear
|
||||
|
||||
|
||||
# ----------------------------------------------------------------------------------------------------
|
||||
@pytest.mark.django_db(transaction=True)
|
||||
def test_notes_search_with_only_filters(
|
||||
client, content_config: ContentConfig, search_config: SearchConfig, sample_org_data
|
||||
):
|
||||
# Arrange
|
||||
filters = [WordFilter(), FileFilter()]
|
||||
search_models.text_search = text_search.initialize_model(search_config.asymmetric)
|
||||
content_index.org = text_search.setup(
|
||||
text_search.setup(
|
||||
OrgToJsonl,
|
||||
sample_org_data,
|
||||
content_config.org,
|
||||
search_models.text_search.bi_encoder,
|
||||
regenerate=False,
|
||||
filters=filters,
|
||||
)
|
||||
user_query = quote('+"Emacs" file:"*.org"')
|
||||
|
||||
@@ -238,15 +216,10 @@ def test_notes_search_with_only_filters(
|
||||
|
||||
|
||||
# ----------------------------------------------------------------------------------------------------
|
||||
def test_notes_search_with_include_filter(
|
||||
client, content_config: ContentConfig, search_config: SearchConfig, sample_org_data
|
||||
):
|
||||
@pytest.mark.django_db(transaction=True)
|
||||
def test_notes_search_with_include_filter(client, sample_org_data):
|
||||
# Arrange
|
||||
filters = [WordFilter()]
|
||||
search_models.text_search = text_search.initialize_model(search_config.asymmetric)
|
||||
content_index.org = text_search.setup(
|
||||
OrgToJsonl, sample_org_data, content_config.org, search_models.text_search, regenerate=False, filters=filters
|
||||
)
|
||||
text_search.setup(OrgToJsonl, sample_org_data, regenerate=False)
|
||||
user_query = quote('How to git install application? +"Emacs"')
|
||||
|
||||
# Act
|
||||
@@ -260,19 +233,13 @@ def test_notes_search_with_include_filter(
|
||||
|
||||
|
||||
# ----------------------------------------------------------------------------------------------------
|
||||
def test_notes_search_with_exclude_filter(
|
||||
client, content_config: ContentConfig, search_config: SearchConfig, sample_org_data
|
||||
):
|
||||
@pytest.mark.django_db(transaction=True)
|
||||
def test_notes_search_with_exclude_filter(client, sample_org_data):
|
||||
# Arrange
|
||||
filters = [WordFilter()]
|
||||
search_models.text_search = text_search.initialize_model(search_config.asymmetric)
|
||||
content_index.org = text_search.setup(
|
||||
text_search.setup(
|
||||
OrgToJsonl,
|
||||
sample_org_data,
|
||||
content_config.org,
|
||||
search_models.text_search.bi_encoder,
|
||||
regenerate=False,
|
||||
filters=filters,
|
||||
)
|
||||
user_query = quote('How to git install application? -"clone"')
|
||||
|
||||
@@ -286,6 +253,22 @@ def test_notes_search_with_exclude_filter(
|
||||
assert "clone" not in search_result
|
||||
|
||||
|
||||
# ----------------------------------------------------------------------------------------------------
|
||||
@pytest.mark.django_db(transaction=True)
|
||||
def test_different_user_data_not_accessed(client, sample_org_data, default_user: KhojUser):
|
||||
# Arrange
|
||||
text_search.setup(OrgToJsonl, sample_org_data, regenerate=False, user=default_user)
|
||||
user_query = quote("How to git install application?")
|
||||
|
||||
# Act
|
||||
response = client.get(f"/api/search?q={user_query}&n=1&t=org")
|
||||
|
||||
# Assert
|
||||
assert response.status_code == 200
|
||||
# assert actual response has no data as the default_user is different from the user making the query (anonymous)
|
||||
assert len(response.json()) == 0
|
||||
|
||||
|
||||
def get_sample_files_data():
|
||||
return {
|
||||
"files": ("path/to/filename.org", "* practicing piano", "text/org"),
|
||||
|
||||
Reference in New Issue
Block a user