Skip to content

Commit

Permalink
Merge branch 'main' of https://github.com/onyx-dot-app/onyx into feat…
Browse files Browse the repository at this point in the history
…ure/more_celery_fanout
  • Loading branch information
Richard Kuo (Danswer) committed Jan 22, 2025
2 parents 4ee7bc5 + b9c29f2 commit 5608553
Show file tree
Hide file tree
Showing 10 changed files with 267 additions and 23 deletions.
11 changes: 7 additions & 4 deletions backend/ee/onyx/db/query_history.py
Original file line number Diff line number Diff line change
Expand Up @@ -98,10 +98,9 @@ def get_page_of_chat_sessions(
conditions = _build_filter_conditions(start_time, end_time, feedback_filter)

subquery = (
select(ChatSession.id, ChatSession.time_created)
select(ChatSession.id)
.filter(*conditions)
.order_by(ChatSession.id, desc(ChatSession.time_created))
.distinct(ChatSession.id)
.order_by(desc(ChatSession.time_created), ChatSession.id)
.limit(page_size)
.offset(page_num * page_size)
.subquery()
Expand All @@ -118,7 +117,11 @@ def get_page_of_chat_sessions(
ChatMessage.chat_message_feedbacks
),
)
.order_by(desc(ChatSession.time_created), asc(ChatMessage.id))
.order_by(
desc(ChatSession.time_created),
ChatSession.id,
asc(ChatMessage.id), # Ensure chronological message order
)
)

return db_session.scalars(stmt).unique().all()
Expand Down
2 changes: 1 addition & 1 deletion backend/onyx/db/index_attempt.py
Original file line number Diff line number Diff line change
Expand Up @@ -432,7 +432,7 @@ def get_paginated_index_attempts_for_cc_pair_id(
stmt = stmt.order_by(IndexAttempt.time_started.desc())

# Apply pagination
stmt = stmt.offset((page - 1) * page_size).limit(page_size)
stmt = stmt.offset(page * page_size).limit(page_size)

return list(db_session.execute(stmt).scalars().all())

Expand Down
4 changes: 2 additions & 2 deletions backend/onyx/server/documents/cc_pair.py
Original file line number Diff line number Diff line change
Expand Up @@ -62,7 +62,7 @@
@router.get("/admin/cc-pair/{cc_pair_id}/index-attempts")
def get_cc_pair_index_attempts(
cc_pair_id: int,
page: int = Query(1, ge=1),
page_num: int = Query(0, ge=0),
page_size: int = Query(10, ge=1, le=1000),
user: User | None = Depends(current_curator_or_admin_user),
db_session: Session = Depends(get_session),
Expand All @@ -81,7 +81,7 @@ def get_cc_pair_index_attempts(
index_attempts = get_paginated_index_attempts_for_cc_pair_id(
db_session=db_session,
connector_id=cc_pair.connector_id,
page=page,
page=page_num,
page_size=page_size,
)
return PaginatedReturn(
Expand Down
106 changes: 106 additions & 0 deletions backend/tests/integration/common_utils/managers/index_attempt.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,106 @@
from datetime import datetime
from datetime import timedelta
from urllib.parse import urlencode

import requests

from onyx.db.engine import get_session_context_manager
from onyx.db.enums import IndexModelStatus
from onyx.db.models import IndexAttempt
from onyx.db.models import IndexingStatus
from onyx.db.search_settings import get_current_search_settings
from onyx.server.documents.models import IndexAttemptSnapshot
from onyx.server.documents.models import PaginatedReturn
from tests.integration.common_utils.constants import API_SERVER_URL
from tests.integration.common_utils.constants import GENERAL_HEADERS
from tests.integration.common_utils.test_models import DATestIndexAttempt
from tests.integration.common_utils.test_models import DATestUser


class IndexAttemptManager:
@staticmethod
def create_test_index_attempts(
num_attempts: int,
cc_pair_id: int,
from_beginning: bool = False,
status: IndexingStatus = IndexingStatus.SUCCESS,
new_docs_indexed: int = 10,
total_docs_indexed: int = 10,
docs_removed_from_index: int = 0,
error_msg: str | None = None,
base_time: datetime | None = None,
) -> list[DATestIndexAttempt]:
if base_time is None:
base_time = datetime.now()

attempts = []
with get_session_context_manager() as db_session:
# Get the current search settings
search_settings = get_current_search_settings(db_session)
if (
not search_settings
or search_settings.status != IndexModelStatus.PRESENT
):
raise ValueError("No current search settings found with PRESENT status")

for i in range(num_attempts):
time_created = base_time - timedelta(hours=i)

index_attempt = IndexAttempt(
connector_credential_pair_id=cc_pair_id,
from_beginning=from_beginning,
status=status,
new_docs_indexed=new_docs_indexed,
total_docs_indexed=total_docs_indexed,
docs_removed_from_index=docs_removed_from_index,
error_msg=error_msg,
time_created=time_created,
time_started=time_created,
time_updated=time_created,
search_settings_id=search_settings.id,
)

db_session.add(index_attempt)
db_session.flush() # To get the ID

attempts.append(
DATestIndexAttempt(
id=index_attempt.id,
status=index_attempt.status,
new_docs_indexed=index_attempt.new_docs_indexed,
total_docs_indexed=index_attempt.total_docs_indexed,
docs_removed_from_index=index_attempt.docs_removed_from_index,
error_msg=index_attempt.error_msg,
time_started=index_attempt.time_started,
time_updated=index_attempt.time_updated,
)
)

db_session.commit()

return attempts

@staticmethod
def get_index_attempt_page(
cc_pair_id: int,
page: int = 0,
page_size: int = 10,
user_performing_action: DATestUser | None = None,
) -> PaginatedReturn[IndexAttemptSnapshot]:
query_params: dict[str, str | int] = {
"page_num": page,
"page_size": page_size,
}

response = requests.get(
url=f"{API_SERVER_URL}/manage/admin/cc-pair/{cc_pair_id}/index-attempts?{urlencode(query_params, doseq=True)}",
headers=user_performing_action.headers
if user_performing_action
else GENERAL_HEADERS,
)
response.raise_for_status()
data = response.json()
return PaginatedReturn(
items=[IndexAttemptSnapshot(**item) for item in data["items"]],
total_items=data["total_items"],
)
33 changes: 33 additions & 0 deletions backend/tests/integration/common_utils/test_models.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
from dataclasses import dataclass
from datetime import datetime
from enum import Enum
from typing import Any
from uuid import UUID
Expand All @@ -10,6 +12,8 @@
from onyx.context.search.enums import RecencyBiasSetting
from onyx.db.enums import AccessType
from onyx.server.documents.models import DocumentSource
from onyx.server.documents.models import IndexAttemptSnapshot
from onyx.server.documents.models import IndexingStatus
from onyx.server.documents.models import InputType

"""
Expand Down Expand Up @@ -171,3 +175,32 @@ class DATestSettings(BaseModel):
gpu_enabled: bool | None = None
product_gating: DATestGatingType = DATestGatingType.NONE
anonymous_user_enabled: bool | None = None


@dataclass
class DATestIndexAttempt:
id: int
status: IndexingStatus | None
new_docs_indexed: int | None
total_docs_indexed: int | None
docs_removed_from_index: int | None
error_msg: str | None
time_started: datetime | None
time_updated: datetime | None

@classmethod
def from_index_attempt_snapshot(
cls, index_attempt: IndexAttemptSnapshot
) -> "DATestIndexAttempt":
return cls(
id=index_attempt.id,
status=index_attempt.status,
new_docs_indexed=index_attempt.new_docs_indexed,
total_docs_indexed=index_attempt.total_docs_indexed,
docs_removed_from_index=index_attempt.docs_removed_from_index,
error_msg=index_attempt.error_msg,
time_started=datetime.fromisoformat(index_attempt.time_started)
if index_attempt.time_started
else None,
time_updated=datetime.fromisoformat(index_attempt.time_updated),
)
Original file line number Diff line number Diff line change
@@ -0,0 +1,90 @@
from datetime import datetime

from onyx.db.models import IndexingStatus
from tests.integration.common_utils.managers.cc_pair import CCPairManager
from tests.integration.common_utils.managers.index_attempt import IndexAttemptManager
from tests.integration.common_utils.managers.user import UserManager
from tests.integration.common_utils.test_models import DATestIndexAttempt
from tests.integration.common_utils.test_models import DATestUser


def _verify_index_attempt_pagination(
cc_pair_id: int,
index_attempts: list[DATestIndexAttempt],
page_size: int = 5,
user_performing_action: DATestUser | None = None,
) -> None:
retrieved_attempts: list[int] = []
last_time_started = None # Track the last time_started seen

for i in range(0, len(index_attempts), page_size):
paginated_result = IndexAttemptManager.get_index_attempt_page(
cc_pair_id=cc_pair_id,
page=(i // page_size),
page_size=page_size,
user_performing_action=user_performing_action,
)

# Verify that the total items is equal to the length of the index attempts list
assert paginated_result.total_items == len(index_attempts)
# Verify that the number of items in the page is equal to the page size
assert len(paginated_result.items) == min(page_size, len(index_attempts) - i)

# Verify time ordering within the page (descending order)
for attempt in paginated_result.items:
if last_time_started is not None:
assert (
attempt.time_started <= last_time_started
), "Index attempts not in descending time order"
last_time_started = attempt.time_started

# Add the retrieved index attempts to the list of retrieved attempts
retrieved_attempts.extend([attempt.id for attempt in paginated_result.items])

# Create a set of all the expected index attempt IDs
all_expected_attempts = set(attempt.id for attempt in index_attempts)
# Create a set of all the retrieved index attempt IDs
all_retrieved_attempts = set(retrieved_attempts)

# Verify that the set of retrieved attempts is equal to the set of expected attempts
assert all_expected_attempts == all_retrieved_attempts


def test_index_attempt_pagination(reset: None) -> None:
# Create an admin user to perform actions
user_performing_action: DATestUser = UserManager.create(
name="admin_performing_action",
is_first_user=True,
)

# Create a CC pair to attach index attempts to
cc_pair = CCPairManager.create_from_scratch(
user_performing_action=user_performing_action,
)

# Create 300 successful index attempts
base_time = datetime.now()
all_attempts = IndexAttemptManager.create_test_index_attempts(
num_attempts=300,
cc_pair_id=cc_pair.id,
status=IndexingStatus.SUCCESS,
base_time=base_time,
)

# Verify basic pagination with different page sizes
print("Verifying basic pagination with page size 5")
_verify_index_attempt_pagination(
cc_pair_id=cc_pair.id,
index_attempts=all_attempts,
page_size=5,
user_performing_action=user_performing_action,
)

# Test with a larger page size
print("Verifying pagination with page size 100")
_verify_index_attempt_pagination(
cc_pair_id=cc_pair.id,
index_attempts=all_attempts,
page_size=100,
user_performing_action=user_performing_action,
)
24 changes: 16 additions & 8 deletions backend/tests/integration/tests/llm_provider/test_llm_provider.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
import requests

from tests.integration.common_utils.constants import API_SERVER_URL
from tests.integration.common_utils.managers.user import UserManager
from tests.integration.common_utils.test_models import DATestUser


Expand All @@ -20,14 +21,15 @@ def _get_provider_by_id(admin_user: DATestUser, provider_id: str) -> dict | None
return next((p for p in providers if p["id"] == provider_id), None)


def test_create_llm_provider_without_display_model_names(
admin_user: DATestUser,
) -> None:
def test_create_llm_provider_without_display_model_names(reset: None) -> None:
"""Test creating an LLM provider without specifying
display_model_names and verify it's null in response"""
# Create admin user
admin_user = UserManager.create(name="admin_user")

# Create LLM provider without model_names
response = requests.put(
f"{API_SERVER_URL}/admin/llm/provider",
f"{API_SERVER_URL}/admin/llm/provider?is_creation=true",
headers=admin_user.headers,
json={
"name": str(uuid.uuid4()),
Expand All @@ -49,12 +51,15 @@ def test_create_llm_provider_without_display_model_names(
assert provider_data["display_model_names"] is None


def test_update_llm_provider_model_names(admin_user: DATestUser) -> None:
def test_update_llm_provider_model_names(reset: None) -> None:
"""Test updating an LLM provider's model_names"""
# Create admin user
admin_user = UserManager.create(name="admin_user")

# First create provider without model_names
name = str(uuid.uuid4())
response = requests.put(
f"{API_SERVER_URL}/admin/llm/provider",
f"{API_SERVER_URL}/admin/llm/provider?is_creation=true",
headers=admin_user.headers,
json={
"name": name,
Expand Down Expand Up @@ -90,11 +95,14 @@ def test_update_llm_provider_model_names(admin_user: DATestUser) -> None:
assert provider_data["model_names"] == _DEFAULT_MODELS


def test_delete_llm_provider(admin_user: DATestUser) -> None:
def test_delete_llm_provider(reset: None) -> None:
"""Test deleting an LLM provider"""
# Create admin user
admin_user = UserManager.create(name="admin_user")

# Create a provider
response = requests.put(
f"{API_SERVER_URL}/admin/llm/provider",
f"{API_SERVER_URL}/admin/llm/provider?is_creation=true",
headers=admin_user.headers,
json={
"name": "test-provider-delete",
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -46,7 +46,9 @@ function QueryHistoryTableRow({
>
<TableCell>
<Text className="whitespace-normal line-clamp-5">
{chatSessionMinimal.first_user_message || "-"}
{chatSessionMinimal.first_user_message ||
chatSessionMinimal.name ||
"-"}
</Text>
</TableCell>
<TableCell>
Expand Down
Loading

0 comments on commit 5608553

Please sign in to comment.