Merge branch 'master' of github.com:khoj-ai/khoj into features/advanced-reasoning

tests/conftest.py

@@ -178,6 +178,13 @@ def api_user4(default_user4):
     )
 
 
+@pytest.mark.django_db
+@pytest.fixture
+def default_openai_chat_model_option():
+    chat_model = ChatModelOptionsFactory(chat_model="gpt-4o-mini", model_type="openai")
+    return chat_model
+
+
 @pytest.mark.django_db
 @pytest.fixture
 def offline_agent():

tests/helpers.py

@@ -86,7 +86,7 @@ class SubscriptionFactory(factory.django.DjangoModelFactory):
         model = Subscription
 
     user = factory.SubFactory(UserFactory)
-    type = "standard"
+    type = Subscription.Type.STANDARD
     is_recurring = False
     renewal_date = make_aware(datetime.strptime("2100-04-01", "%Y-%m-%d"))
 

tests/test_agents.py

@@ -0,0 +1,211 @@
+# tests/test_agents.py
+import os
+
+import pytest
+from asgiref.sync import sync_to_async
+
+from khoj.database.adapters import AgentAdapters
+from khoj.database.models import Agent, ChatModelOptions, Entry, KhojUser
+from khoj.routers.api import execute_search
+from khoj.utils.helpers import get_absolute_path
+from tests.helpers import ChatModelOptionsFactory
+
+
+def test_create_default_agent(default_user: KhojUser):
+    ChatModelOptionsFactory()
+
+    agent = AgentAdapters.create_default_agent(default_user)
+    assert agent is not None
+    assert agent.input_tools == []
+    assert agent.output_modes == []
+    assert agent.privacy_level == Agent.PrivacyLevel.PUBLIC
+    assert agent.managed_by_admin == True
+
+
+@pytest.mark.anyio
+@pytest.mark.django_db(transaction=True)
+async def test_create_or_update_agent(default_user: KhojUser, default_openai_chat_model_option: ChatModelOptions):
+    new_agent = await AgentAdapters.aupdate_agent(
+        default_user,
+        "Test Agent",
+        "Test Personality",
+        Agent.PrivacyLevel.PRIVATE,
+        "icon",
+        "color",
+        default_openai_chat_model_option.chat_model,
+        [],
+        [],
+        [],
+    )
+    assert new_agent is not None
+    assert new_agent.name == "Test Agent"
+    assert new_agent.privacy_level == Agent.PrivacyLevel.PRIVATE
+    assert new_agent.creator == default_user
+
+
+@pytest.mark.anyio
+@pytest.mark.django_db(transaction=True)
+async def test_create_or_update_agent_with_knowledge_base(
+    default_user2: KhojUser, default_openai_chat_model_option: ChatModelOptions, chat_client
+):
+    full_filename = get_absolute_path("tests/data/markdown/having_kids.markdown")
+    new_agent = await AgentAdapters.aupdate_agent(
+        default_user2,
+        "Test Agent",
+        "Test Personality",
+        Agent.PrivacyLevel.PRIVATE,
+        "icon",
+        "color",
+        default_openai_chat_model_option.chat_model,
+        [full_filename],
+        [],
+        [],
+    )
+    entries = await sync_to_async(list)(Entry.objects.filter(agent=new_agent))
+    file_names = set()
+    for entry in entries:
+        file_names.add(entry.file_path)
+
+    assert new_agent is not None
+    assert new_agent.name == "Test Agent"
+    assert new_agent.privacy_level == Agent.PrivacyLevel.PRIVATE
+    assert new_agent.creator == default_user2
+    assert len(entries) > 0
+    assert full_filename in file_names
+    assert len(file_names) == 1
+
+
+@pytest.mark.anyio
+@pytest.mark.django_db(transaction=True)
+async def test_create_or_update_agent_with_knowledge_base_and_search(
+    default_user2: KhojUser, default_openai_chat_model_option: ChatModelOptions, chat_client
+):
+    full_filename = get_absolute_path("tests/data/markdown/having_kids.markdown")
+    new_agent = await AgentAdapters.aupdate_agent(
+        default_user2,
+        "Test Agent",
+        "Test Personality",
+        Agent.PrivacyLevel.PRIVATE,
+        "icon",
+        "color",
+        default_openai_chat_model_option.chat_model,
+        [full_filename],
+        [],
+        [],
+    )
+
+    search_result = await execute_search(user=default_user2, q="having kids", agent=new_agent)
+
+    assert len(search_result) == 5
+
+
+@pytest.mark.anyio
+@pytest.mark.django_db(transaction=True)
+async def test_agent_with_knowledge_base_and_search_not_creator(
+    default_user2: KhojUser, default_openai_chat_model_option: ChatModelOptions, chat_client, default_user3: KhojUser
+):
+    full_filename = get_absolute_path("tests/data/markdown/having_kids.markdown")
+    new_agent = await AgentAdapters.aupdate_agent(
+        default_user2,
+        "Test Agent",
+        "Test Personality",
+        Agent.PrivacyLevel.PUBLIC,
+        "icon",
+        "color",
+        default_openai_chat_model_option.chat_model,
+        [full_filename],
+        [],
+        [],
+    )
+
+    search_result = await execute_search(user=default_user3, q="having kids", agent=new_agent)
+
+    assert len(search_result) == 5
+
+
+@pytest.mark.anyio
+@pytest.mark.django_db(transaction=True)
+async def test_agent_with_knowledge_base_and_search_not_creator_and_private(
+    default_user2: KhojUser, default_openai_chat_model_option: ChatModelOptions, chat_client, default_user3: KhojUser
+):
+    full_filename = get_absolute_path("tests/data/markdown/having_kids.markdown")
+    new_agent = await AgentAdapters.aupdate_agent(
+        default_user2,
+        "Test Agent",
+        "Test Personality",
+        Agent.PrivacyLevel.PRIVATE,
+        "icon",
+        "color",
+        default_openai_chat_model_option.chat_model,
+        [full_filename],
+        [],
+        [],
+    )
+
+    search_result = await execute_search(user=default_user3, q="having kids", agent=new_agent)
+
+    assert len(search_result) == 0
+
+
+@pytest.mark.anyio
+@pytest.mark.django_db(transaction=True)
+async def test_agent_with_knowledge_base_and_search_not_creator_and_private_accessible_to_none(
+    default_user2: KhojUser, default_openai_chat_model_option: ChatModelOptions, chat_client
+):
+    full_filename = get_absolute_path("tests/data/markdown/having_kids.markdown")
+    new_agent = await AgentAdapters.aupdate_agent(
+        default_user2,
+        "Test Agent",
+        "Test Personality",
+        Agent.PrivacyLevel.PRIVATE,
+        "icon",
+        "color",
+        default_openai_chat_model_option.chat_model,
+        [full_filename],
+        [],
+        [],
+    )
+
+    search_result = await execute_search(user=None, q="having kids", agent=new_agent)
+
+    assert len(search_result) == 5
+
+
+@pytest.mark.anyio
+@pytest.mark.django_db(transaction=True)
+async def test_multiple_agents_with_knowledge_base_and_users(
+    default_user2: KhojUser, default_openai_chat_model_option: ChatModelOptions, chat_client, default_user3: KhojUser
+):
+    full_filename = get_absolute_path("tests/data/markdown/having_kids.markdown")
+    new_agent = await AgentAdapters.aupdate_agent(
+        default_user2,
+        "Test Agent",
+        "Test Personality",
+        Agent.PrivacyLevel.PUBLIC,
+        "icon",
+        "color",
+        default_openai_chat_model_option.chat_model,
+        [full_filename],
+        [],
+        [],
+    )
+
+    full_filename2 = get_absolute_path("tests/data/markdown/Namita.markdown")
+    new_agent2 = await AgentAdapters.aupdate_agent(
+        default_user2,
+        "Test Agent 2",
+        "Test Personality",
+        Agent.PrivacyLevel.PUBLIC,
+        "icon",
+        "color",
+        default_openai_chat_model_option.chat_model,
+        [full_filename2],
+        [],
+        [],
+    )
+
+    search_result = await execute_search(user=default_user3, q="having kids", agent=new_agent2)
+    search_result2 = await execute_search(user=default_user3, q="Namita", agent=new_agent2)
+
+    assert len(search_result) == 0
+    assert len(search_result2) == 1

tests/test_pdf_to_entries.py

@@ -1,6 +1,8 @@
 import os
 import re
 
+import pytest
+
 from khoj.processor.content.pdf.pdf_to_entries import PdfToEntries
 from khoj.utils.fs_syncer import get_pdf_files
 from khoj.utils.rawconfig import TextContentConfig
@@ -37,6 +39,7 @@ def test_multi_page_pdf_to_jsonl():
     assert len(entries[1]) == 6
 
 
+@pytest.mark.skip(reason="Temporarily disabled OCR due to performance issues")
 def test_ocr_page_pdf_to_jsonl():
     "Convert multiple pages from single PDF file to jsonl."
     # Arrange