Skip to content

Commit

Permalink
Merge branch 'onyx-dot-app:main' into bugfix/gmail-date-parse
Browse files Browse the repository at this point in the history
  • Loading branch information
sam-w authored Feb 6, 2025
2 parents 145a2a3 + 58fdc86 commit 2759cb9
Show file tree
Hide file tree
Showing 89 changed files with 1,687 additions and 771 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -67,6 +67,7 @@ jobs:
NEXT_PUBLIC_SENTRY_DSN=${{ secrets.SENTRY_DSN }}
NEXT_PUBLIC_GTM_ENABLED=true
NEXT_PUBLIC_FORGOT_PASSWORD_ENABLED=true
NEXT_PUBLIC_INCLUDE_ERROR_POPUP_SUPPORT_LINK=true
NODE_OPTIONS=--max-old-space-size=8192
# needed due to weird interactions with the builds for different platforms
no-cache: true
Expand Down
3 changes: 2 additions & 1 deletion backend/Dockerfile
Original file line number Diff line number Diff line change
Expand Up @@ -101,7 +101,8 @@ COPY ./alembic_tenants /app/alembic_tenants
COPY ./alembic.ini /app/alembic.ini
COPY supervisord.conf /usr/etc/supervisord.conf

# Escape hatch
# Escape hatch scripts
COPY ./scripts/debugging /app/scripts/debugging
COPY ./scripts/force_delete_connector_by_id.py /app/scripts/force_delete_connector_by_id.py

# Put logo in assets
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -52,7 +52,11 @@ def upgrade() -> None:
slack_bot_id, persona_id, channel_config, enable_auto_filters, is_default
) VALUES (
:bot_id, NULL,
'{"channel_name": null, "respond_member_group_list": [], "answer_filters": [], "follow_up_tags": []}',
'{"channel_name": null, '
'"respond_member_group_list": [], '
'"answer_filters": [], '
'"follow_up_tags": [], '
'"respond_tag_only": true}',
FALSE, TRUE
)
"""
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,53 @@
"""delete non-search assistants
Revision ID: f5437cc136c5
Revises: eaa3b5593925
Create Date: 2025-02-04 16:17:15.677256
"""
from alembic import op


# revision identifiers, used by Alembic.
revision = "f5437cc136c5"
down_revision = "eaa3b5593925"
branch_labels = None
depends_on = None


def upgrade() -> None:
pass


def downgrade() -> None:
# Fix: split the statements into multiple op.execute() calls
op.execute(
"""
WITH personas_without_search AS (
SELECT p.id
FROM persona p
LEFT JOIN persona__tool pt ON p.id = pt.persona_id
LEFT JOIN tool t ON pt.tool_id = t.id
GROUP BY p.id
HAVING COUNT(CASE WHEN t.in_code_tool_id = 'run_search' THEN 1 END) = 0
)
UPDATE slack_channel_config
SET persona_id = NULL
WHERE is_default = TRUE AND persona_id IN (SELECT id FROM personas_without_search)
"""
)

op.execute(
"""
WITH personas_without_search AS (
SELECT p.id
FROM persona p
LEFT JOIN persona__tool pt ON p.id = pt.persona_id
LEFT JOIN tool t ON pt.tool_id = t.id
GROUP BY p.id
HAVING COUNT(CASE WHEN t.in_code_tool_id = 'run_search' THEN 1 END) = 0
)
DELETE FROM slack_channel_config
WHERE is_default = FALSE AND persona_id IN (SELECT id FROM personas_without_search)
"""
)
11 changes: 11 additions & 0 deletions backend/ee/onyx/db/persona.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,8 +2,11 @@

from sqlalchemy.orm import Session

from onyx.configs.constants import NotificationType
from onyx.db.models import Persona__User
from onyx.db.models import Persona__UserGroup
from onyx.db.notification import create_notification
from onyx.server.features.persona.models import PersonaSharedNotificationData


def make_persona_private(
Expand All @@ -23,6 +26,14 @@ def make_persona_private(
for user_uuid in user_ids:
db_session.add(Persona__User(persona_id=persona_id, user_id=user_uuid))

create_notification(
user_id=user_uuid,
notif_type=NotificationType.PERSONA_SHARED,
db_session=db_session,
additional_data=PersonaSharedNotificationData(
persona_id=persona_id,
).model_dump(),
)
if group_ids:
for group_id in group_ids:
db_session.add(
Expand Down
4 changes: 2 additions & 2 deletions backend/ee/onyx/db/user_group.py
Original file line number Diff line number Diff line change
Expand Up @@ -218,14 +218,14 @@ def fetch_user_groups_for_user(
return db_session.scalars(stmt).all()


def construct_document_select_by_usergroup(
def construct_document_id_select_by_usergroup(
user_group_id: int,
) -> Select:
"""This returns a statement that should be executed using
.yield_per() to minimize overhead. The primary consumers of this function
are background processing task generators."""
stmt = (
select(Document)
select(Document.id)
.join(
DocumentByConnectorCredentialPair,
Document.id == DocumentByConnectorCredentialPair.id,
Expand Down
5 changes: 3 additions & 2 deletions backend/ee/onyx/server/middleware/tenant_tracking.py
Original file line number Diff line number Diff line change
Expand Up @@ -87,13 +87,14 @@ async def _get_tenant_id_from_request(
if not is_valid_schema_name(tenant_id):
raise HTTPException(status_code=400, detail="Invalid tenant ID format")

return tenant_id

except Exception as e:
logger.error(f"Unexpected error in _get_tenant_id_from_request: {str(e)}")
raise HTTPException(status_code=500, detail="Internal server error")

finally:
if tenant_id:
return tenant_id

# As a final step, check for explicit tenant_id cookie
tenant_id_cookie = request.cookies.get(TENANT_ID_COOKIE_NAME)
if tenant_id_cookie and is_valid_schema_name(tenant_id_cookie):
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -12,8 +12,9 @@
from onyx.agents.agent_search.deep_search.main.models import (
AgentRefinedMetrics,
)
from onyx.agents.agent_search.deep_search.main.operations import dispatch_subquestion
from onyx.agents.agent_search.deep_search.main.operations import (
dispatch_subquestion,
dispatch_subquestion_sep,
)
from onyx.agents.agent_search.deep_search.main.states import (
InitialQuestionDecompositionUpdate,
Expand Down Expand Up @@ -109,9 +110,12 @@ def decompose_orig_question(
),
writer,
)

# dispatches custom events for subquestion tokens, adding in subquestion ids.
streamed_tokens = dispatch_separated(
model.stream(msg), dispatch_subquestion(0, writer)
model.stream(msg),
dispatch_subquestion(0, writer),
sep_callback=dispatch_subquestion_sep(0, writer),
)

stop_event = StreamStopInfo(
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -9,8 +9,9 @@
from onyx.agents.agent_search.deep_search.main.models import (
RefinementSubQuestion,
)
from onyx.agents.agent_search.deep_search.main.operations import dispatch_subquestion
from onyx.agents.agent_search.deep_search.main.operations import (
dispatch_subquestion,
dispatch_subquestion_sep,
)
from onyx.agents.agent_search.deep_search.main.states import MainState
from onyx.agents.agent_search.deep_search.main.states import (
Expand Down Expand Up @@ -96,7 +97,9 @@ def create_refined_sub_questions(
model = graph_config.tooling.fast_llm

streamed_tokens = dispatch_separated(
model.stream(msg), dispatch_subquestion(1, writer)
model.stream(msg),
dispatch_subquestion(1, writer),
sep_callback=dispatch_subquestion_sep(1, writer),
)
response = merge_content(*streamed_tokens)

Expand Down
19 changes: 19 additions & 0 deletions backend/onyx/agents/agent_search/deep_search/main/operations.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,9 @@
SubQuestionAnswerResults,
)
from onyx.agents.agent_search.shared_graph_utils.utils import write_custom_event
from onyx.chat.models import StreamStopInfo
from onyx.chat.models import StreamStopReason
from onyx.chat.models import StreamType
from onyx.chat.models import SubQuestionPiece
from onyx.context.search.models import IndexFilters
from onyx.tools.models import SearchQueryInfo
Expand All @@ -34,6 +37,22 @@ def _helper(sub_question_part: str, sep_num: int) -> None:
return _helper


def dispatch_subquestion_sep(level: int, writer: StreamWriter) -> Callable[[int], None]:
def _helper(sep_num: int) -> None:
write_custom_event(
"stream_finished",
StreamStopInfo(
stop_reason=StreamStopReason.FINISHED,
stream_type=StreamType.SUB_QUESTIONS,
level=level,
level_question_num=sep_num,
),
writer,
)

return _helper


def calculate_initial_agent_stats(
decomp_answer_results: list[SubQuestionAnswerResults],
original_question_stats: AgentChunkRetrievalStats,
Expand Down
8 changes: 8 additions & 0 deletions backend/onyx/agents/agent_search/shared_graph_utils/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -295,6 +295,7 @@ def _dispatch_nonempty(
def dispatch_separated(
tokens: Iterator[BaseMessage],
dispatch_event: Callable[[str, int], None],
sep_callback: Callable[[int], None] | None = None,
sep: str = DISPATCH_SEP_CHAR,
) -> list[BaseMessage_Content]:
num = 1
Expand All @@ -304,6 +305,10 @@ def dispatch_separated(
if sep in content:
sub_question_parts = content.split(sep)
_dispatch_nonempty(sub_question_parts[0], dispatch_event, num)

if sep_callback:
sep_callback(num)

num += 1
_dispatch_nonempty(
"".join(sub_question_parts[1:]).strip(), dispatch_event, num
Expand All @@ -312,6 +317,9 @@ def dispatch_separated(
_dispatch_nonempty(content, dispatch_event, num)
streamed_tokens.append(content)

if sep_callback:
sep_callback(num)

return streamed_tokens


Expand Down
13 changes: 8 additions & 5 deletions backend/onyx/background/celery/tasks/connector_deletion/tasks.py
Original file line number Diff line number Diff line change
Expand Up @@ -179,11 +179,14 @@ def try_generate_document_cc_pair_cleanup_tasks(
if tasks_generated is None:
raise ValueError("RedisConnectorDeletion.generate_tasks returned None")

insert_sync_record(
db_session=db_session,
entity_id=cc_pair_id,
sync_type=SyncType.CONNECTOR_DELETION,
)
try:
insert_sync_record(
db_session=db_session,
entity_id=cc_pair_id,
sync_type=SyncType.CONNECTOR_DELETION,
)
except Exception:
pass

except TaskDependencyError:
redis_connector.delete.set_fence(None)
Expand Down
26 changes: 26 additions & 0 deletions backend/onyx/db/document.py
Original file line number Diff line number Diff line change
Expand Up @@ -105,6 +105,32 @@ def construct_document_select_for_connector_credential_pair_by_needs_sync(
return stmt


def construct_document_id_select_for_connector_credential_pair_by_needs_sync(
connector_id: int, credential_id: int
) -> Select:
initial_doc_ids_stmt = select(DocumentByConnectorCredentialPair.id).where(
and_(
DocumentByConnectorCredentialPair.connector_id == connector_id,
DocumentByConnectorCredentialPair.credential_id == credential_id,
)
)

stmt = (
select(DbDocument.id)
.where(
DbDocument.id.in_(initial_doc_ids_stmt),
or_(
DbDocument.last_modified
> DbDocument.last_synced, # last_modified is newer than last_synced
DbDocument.last_synced.is_(None), # never synced
),
)
.distinct()
)

return stmt


def get_all_documents_needing_vespa_sync_for_cc_pair(
db_session: Session, cc_pair_id: int
) -> list[DbDocument]:
Expand Down
6 changes: 3 additions & 3 deletions backend/onyx/db/document_set.py
Original file line number Diff line number Diff line change
Expand Up @@ -198,7 +198,7 @@ def _check_if_cc_pairs_are_owned_by_groups(
ids=missing_cc_pair_ids,
)
for cc_pair in cc_pairs:
if cc_pair.access_type != AccessType.PUBLIC:
if cc_pair.access_type == AccessType.PRIVATE:
raise ValueError(
f"Connector Credential Pair with ID: '{cc_pair.id}'"
" is not owned by the specified groups"
Expand Down Expand Up @@ -545,7 +545,7 @@ def fetch_documents_for_document_set_paginated(
return documents, documents[-1].id if documents else None


def construct_document_select_by_docset(
def construct_document_id_select_by_docset(
document_set_id: int,
current_only: bool = True,
) -> Select:
Expand All @@ -554,7 +554,7 @@ def construct_document_select_by_docset(
are background processing task generators."""

stmt = (
select(Document)
select(Document.id)
.join(
DocumentByConnectorCredentialPair,
DocumentByConnectorCredentialPair.id == Document.id,
Expand Down
25 changes: 25 additions & 0 deletions backend/onyx/db/persona.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@
from sqlalchemy import select
from sqlalchemy import update
from sqlalchemy.orm import aliased
from sqlalchemy.orm import joinedload
from sqlalchemy.orm import selectinload
from sqlalchemy.orm import Session

Expand All @@ -19,6 +20,7 @@
from onyx.configs.chat_configs import BING_API_KEY
from onyx.configs.chat_configs import CONTEXT_CHUNKS_ABOVE
from onyx.configs.chat_configs import CONTEXT_CHUNKS_BELOW
from onyx.configs.constants import NotificationType
from onyx.context.search.enums import RecencyBiasSetting
from onyx.db.constants import SLACK_BOT_PERSONA_PREFIX
from onyx.db.models import DocumentSet
Expand All @@ -32,6 +34,8 @@
from onyx.db.models import User
from onyx.db.models import User__UserGroup
from onyx.db.models import UserGroup
from onyx.db.notification import create_notification
from onyx.server.features.persona.models import PersonaSharedNotificationData
from onyx.server.features.persona.models import PersonaSnapshot
from onyx.server.features.persona.models import PersonaUpsertRequest
from onyx.utils.logger import setup_logger
Expand Down Expand Up @@ -169,6 +173,15 @@ def make_persona_private(
for user_uuid in user_ids:
db_session.add(Persona__User(persona_id=persona_id, user_id=user_uuid))

create_notification(
user_id=user_uuid,
notif_type=NotificationType.PERSONA_SHARED,
db_session=db_session,
additional_data=PersonaSharedNotificationData(
persona_id=persona_id,
).model_dump(),
)

db_session.commit()

# May cause error if someone switches down to MIT from EE
Expand Down Expand Up @@ -708,3 +721,15 @@ def update_persona_label(
def delete_persona_label(label_id: int, db_session: Session) -> None:
db_session.query(PersonaLabel).filter(PersonaLabel.id == label_id).delete()
db_session.commit()


def persona_has_search_tool(persona_id: int, db_session: Session) -> bool:
persona = (
db_session.query(Persona)
.options(joinedload(Persona.tools))
.filter(Persona.id == persona_id)
.one_or_none()
)
if persona is None:
raise ValueError(f"Persona with ID {persona_id} does not exist")
return any(tool.in_code_tool_id == "run_search" for tool in persona.tools)
Loading

0 comments on commit 2759cb9

Please sign in to comment.