From 391e710b6e73af9b87fb86e8500a203f00a60e8d Mon Sep 17 00:00:00 2001 From: Chris Weaver Date: Tue, 8 Jul 2025 21:35:51 -0700 Subject: [PATCH] Slack federated search ux (#4969) * slack_search.py * rem * fix: get elements * feat: better stack message processing * fix: mypy * fix: url parsing * refactor: federated search * feat: proper chunking + source filters * highlighting + source check * feat: forced section insertion * feat: multi slack api queries * slack query expansion * feat: max slack queries env * limit slack search to avoid overloading the search * Initial draft * more * simpify * Improve modal * Fix oauth flow * Fully working versino * More nicities * Improved cascade delete * document set for fed connector UI * Fix source filters + improve document set selection * Improve callback modal * fix: logging error + showing connectors in admin page user settings * better log * Fix mypy * small rei comment * Fix pydantic * Improvements to modals * feat: distributed pruning * random fix * greptile * Encrypt token * respect source filters + revert llm pruning ordering * greptile nit * feat: thread as context in slack search * feat: slack doc ordering * small improvements * rebase * Small improvements * Fix web build * try fix build * Move to seaprate model file * Use default_factory * remove unused model --------- Co-authored-by: Rei Meguro <36625832+Orbital-Web@users.noreply.github.com> --- .github/workflows/pr-python-checks.yml | 2 +- ...326d83aa_add_federated_connector_tables.py | 72 +++ backend/onyx/chat/answer.py | 4 +- backend/onyx/chat/prune_and_merge.py | 81 ++- backend/onyx/configs/app_configs.py | 9 + backend/onyx/configs/constants.py | 9 + .../onyx/context/search/federated/models.py | 18 + .../context/search/federated/slack_search.py | 416 +++++++++++++ backend/onyx/context/search/models.py | 2 + backend/onyx/context/search/pipeline.py | 1 + .../context/search/retrieval/search_runner.py | 84 ++- backend/onyx/context/search/utils.py | 25 + backend/onyx/db/document_set.py | 75 ++- backend/onyx/db/federated.py | 273 ++++++++ backend/onyx/db/models.py | 87 ++- backend/onyx/federated_connectors/__init__.py | 0 backend/onyx/federated_connectors/factory.py | 29 + .../federated_retrieval.py | 96 +++ .../onyx/federated_connectors/interfaces.py | 101 +++ backend/onyx/federated_connectors/models.py | 66 ++ .../onyx/federated_connectors/oauth_utils.py | 169 +++++ .../federated_connectors/slack/__init__.py | 1 + .../slack/federated_connector.py | 236 +++++++ .../onyx/federated_connectors/slack/models.py | 105 ++++ backend/onyx/main.py | 2 + backend/onyx/prompts/chat_prompts.py | 2 +- backend/onyx/prompts/federated_search.py | 27 + .../server/features/document_set/models.py | 42 ++ backend/onyx/server/federated/api.py | 582 +++++++++++++++++ backend/onyx/server/federated/models.py | 81 +++ .../slack/test_slack_federated_connnector.py | 217 +++++++ web/src/app/admin/add-connector/page.tsx | 173 +++++- .../[connector]/AddConnectorPage.tsx | 1 + .../[connector]/ConnectorWrapper.tsx | 85 ++- .../sets/DocumentSetCreationForm.tsx | 67 +- web/src/app/admin/documents/sets/lib.ts | 8 + web/src/app/admin/documents/sets/page.tsx | 95 ++- web/src/app/admin/federated/[id]/page.tsx | 72 +++ .../federated/[id]/useFederatedConnector.ts | 94 +++ .../status/CCPairIndexingStatusTable.tsx | 175 +++++- web/src/app/admin/indexing/status/page.tsx | 28 +- web/src/app/chat/ChatPage.tsx | 150 ++++- web/src/app/chat/input/ChatInputBar.tsx | 14 +- .../app/chat/input/SelectedFilterDisplay.tsx | 5 +- web/src/app/chat/modal/UserSettingsModal.tsx | 298 ++++++++- web/src/app/federated/oauth/callback/page.tsx | 206 ++++++ .../components/FederatedConnectorSelector.tsx | 521 ++++++++++++++++ web/src/components/admin/ClientLayout.tsx | 12 +- .../federated/FederatedConnectorForm.tsx | 586 ++++++++++++++++++ .../components/chat/FederatedOAuthModal.tsx | 130 ++++ web/src/components/filters/SourceSelector.tsx | 242 +------- .../search/filtering/FilterPopup.tsx | 46 +- .../filtering/HorizontalSourceSelector.tsx | 5 +- web/src/lib/connectors/credentials.ts | 1 + web/src/lib/hooks.ts | 23 +- web/src/lib/hooks/useFederatedOAuthStatus.ts | 44 ++ web/src/lib/oauth_utils.ts | 50 ++ web/src/lib/search/interfaces.ts | 4 + web/src/lib/sources.ts | 36 +- web/src/lib/types.ts | 67 +- 60 files changed, 5713 insertions(+), 439 deletions(-) create mode 100644 backend/alembic/versions/0816326d83aa_add_federated_connector_tables.py create mode 100644 backend/onyx/context/search/federated/models.py create mode 100644 backend/onyx/context/search/federated/slack_search.py create mode 100644 backend/onyx/db/federated.py create mode 100644 backend/onyx/federated_connectors/__init__.py create mode 100644 backend/onyx/federated_connectors/factory.py create mode 100644 backend/onyx/federated_connectors/federated_retrieval.py create mode 100644 backend/onyx/federated_connectors/interfaces.py create mode 100644 backend/onyx/federated_connectors/models.py create mode 100644 backend/onyx/federated_connectors/oauth_utils.py create mode 100644 backend/onyx/federated_connectors/slack/__init__.py create mode 100644 backend/onyx/federated_connectors/slack/federated_connector.py create mode 100644 backend/onyx/federated_connectors/slack/models.py create mode 100644 backend/onyx/prompts/federated_search.py create mode 100644 backend/onyx/server/federated/api.py create mode 100644 backend/onyx/server/federated/models.py create mode 100644 backend/tests/unit/federated_connector/slack/test_slack_federated_connnector.py create mode 100644 web/src/app/admin/federated/[id]/page.tsx create mode 100644 web/src/app/admin/federated/[id]/useFederatedConnector.ts create mode 100644 web/src/app/federated/oauth/callback/page.tsx create mode 100644 web/src/components/FederatedConnectorSelector.tsx create mode 100644 web/src/components/admin/federated/FederatedConnectorForm.tsx create mode 100644 web/src/components/chat/FederatedOAuthModal.tsx create mode 100644 web/src/lib/hooks/useFederatedOAuthStatus.ts diff --git a/.github/workflows/pr-python-checks.yml b/.github/workflows/pr-python-checks.yml index d57034b861..cf4a1a8691 100644 --- a/.github/workflows/pr-python-checks.yml +++ b/.github/workflows/pr-python-checks.yml @@ -47,7 +47,7 @@ jobs: -i /local/openapi.json \ -g python \ -o /local/onyx_openapi_client \ - --package-name onyx_openapi_client + --package-name onyx_openapi_client \ - name: Run MyPy run: | diff --git a/backend/alembic/versions/0816326d83aa_add_federated_connector_tables.py b/backend/alembic/versions/0816326d83aa_add_federated_connector_tables.py new file mode 100644 index 0000000000..be1a386fc9 --- /dev/null +++ b/backend/alembic/versions/0816326d83aa_add_federated_connector_tables.py @@ -0,0 +1,72 @@ +"""add federated connector tables + +Revision ID: 0816326d83aa +Revises: 12635f6655b7 +Create Date: 2025-06-29 14:09:45.109518 + +""" + +from alembic import op +import sqlalchemy as sa +from sqlalchemy.dialects import postgresql + + +# revision identifiers, used by Alembic. +revision = "0816326d83aa" +down_revision = "12635f6655b7" +branch_labels = None +depends_on = None + + +def upgrade() -> None: + # Create federated_connector table + op.create_table( + "federated_connector", + sa.Column("id", sa.Integer(), nullable=False), + sa.Column("source", sa.String(), nullable=False), + sa.Column("credentials", sa.LargeBinary(), nullable=False), + sa.PrimaryKeyConstraint("id"), + ) + + # Create federated_connector_oauth_token table + op.create_table( + "federated_connector_oauth_token", + sa.Column("id", sa.Integer(), nullable=False), + sa.Column("federated_connector_id", sa.Integer(), nullable=False), + sa.Column("user_id", postgresql.UUID(as_uuid=True), nullable=False), + sa.Column("token", sa.LargeBinary(), nullable=False), + sa.Column("expires_at", sa.DateTime(), nullable=True), + sa.ForeignKeyConstraint( + ["federated_connector_id"], ["federated_connector.id"], ondelete="CASCADE" + ), + sa.ForeignKeyConstraint(["user_id"], ["user.id"], ondelete="CASCADE"), + sa.PrimaryKeyConstraint("id"), + ) + + # Create federated_connector__document_set table + op.create_table( + "federated_connector__document_set", + sa.Column("id", sa.Integer(), nullable=False), + sa.Column("federated_connector_id", sa.Integer(), nullable=False), + sa.Column("document_set_id", sa.Integer(), nullable=False), + sa.Column("entities", postgresql.JSONB(), nullable=False), + sa.ForeignKeyConstraint( + ["federated_connector_id"], ["federated_connector.id"], ondelete="CASCADE" + ), + sa.ForeignKeyConstraint( + ["document_set_id"], ["document_set.id"], ondelete="CASCADE" + ), + sa.PrimaryKeyConstraint("id"), + sa.UniqueConstraint( + "federated_connector_id", + "document_set_id", + name="uq_federated_connector_document_set", + ), + ) + + +def downgrade() -> None: + # Drop tables in reverse order due to foreign key dependencies + op.drop_table("federated_connector__document_set") + op.drop_table("federated_connector_oauth_token") + op.drop_table("federated_connector") diff --git a/backend/onyx/chat/answer.py b/backend/onyx/chat/answer.py index e083f13b45..b41ff3764e 100644 --- a/backend/onyx/chat/answer.py +++ b/backend/onyx/chat/answer.py @@ -157,9 +157,7 @@ class Answer: else: run_langgraph = run_basic_graph - stream = run_langgraph( - self.graph_config, - ) + stream = run_langgraph(self.graph_config) processed_stream = [] for packet in stream: diff --git a/backend/onyx/chat/prune_and_merge.py b/backend/onyx/chat/prune_and_merge.py index c6af287466..e7389a1660 100644 --- a/backend/onyx/chat/prune_and_merge.py +++ b/backend/onyx/chat/prune_and_merge.py @@ -11,6 +11,7 @@ from onyx.chat.models import ( ) from onyx.chat.models import PromptConfig from onyx.chat.prompt_builder.citations_prompt import compute_max_document_tokens +from onyx.configs.app_configs import MAX_FEDERATED_SECTIONS from onyx.configs.constants import IGNORE_FOR_QA from onyx.configs.model_configs import DOC_EMBEDDING_CONTEXT_SIZE from onyx.context.search.models import InferenceChunk @@ -67,6 +68,38 @@ def merge_chunk_intervals(chunk_ranges: list[ChunkRange]) -> list[ChunkRange]: return combined_ranges +def _separate_federated_sections( + sections: list[InferenceSection], + section_relevance_list: list[bool] | None, +) -> tuple[list[InferenceSection], list[InferenceSection], list[bool] | None]: + """ + Separates out the first NUM_FEDERATED_SECTIONS federated sections to be spared from pruning. + Any remaining federated sections are treated as normal sections, and will get added if it + fits within the allocated context window. This is done as federated sections do not have + a score and would otherwise always get pruned. + """ + federated_sections: list[InferenceSection] = [] + normal_sections: list[InferenceSection] = [] + normal_section_relevance_list: list[bool] = [] + + for i, section in enumerate(sections): + if ( + len(federated_sections) < MAX_FEDERATED_SECTIONS + and section.center_chunk.is_federated + ): + federated_sections.append(section) + continue + normal_sections.append(section) + if section_relevance_list is not None: + normal_section_relevance_list.append(section_relevance_list[i]) + + return ( + federated_sections[:MAX_FEDERATED_SECTIONS], + normal_sections, + normal_section_relevance_list if section_relevance_list is not None else None, + ) + + def _compute_limit( prompt_config: PromptConfig, llm_config: LLMConfig, @@ -105,7 +138,7 @@ def _compute_limit( return int(min(limit_options)) -def reorder_sections( +def _reorder_sections( sections: list[InferenceSection], section_relevance_list: list[bool] | None, ) -> list[InferenceSection]: @@ -113,11 +146,10 @@ def reorder_sections( return sections reordered_sections: list[InferenceSection] = [] - if section_relevance_list is not None: - for selection_target in [True, False]: - for section, is_relevant in zip(sections, section_relevance_list): - if is_relevant == selection_target: - reordered_sections.append(section) + for selection_target in [True, False]: + for section, is_relevant in zip(sections, section_relevance_list): + if is_relevant == selection_target: + reordered_sections.append(section) return reordered_sections @@ -134,6 +166,7 @@ def _remove_sections_to_ignore( def _apply_pruning( sections: list[InferenceSection], section_relevance_list: list[bool] | None, + keep_sections: list[InferenceSection], token_limit: int, is_manually_selected_docs: bool, use_sections: bool, @@ -144,10 +177,22 @@ def _apply_pruning( provider_type=llm_config.model_provider, model_name=llm_config.model_name, ) - sections = deepcopy(sections) # don't modify in place + + # combine the section lists, making sure to add the keep_sections first + sections = deepcopy(keep_sections) + deepcopy(sections) + + # build combined relevance list, treating the keep_sections as relevant + if section_relevance_list is not None: + section_relevance_list = [True] * len(keep_sections) + section_relevance_list + + # map unique_id: relevance for final ordering step + section_id_to_relevance: dict[str, bool] = {} + if section_relevance_list is not None: + for sec, rel in zip(sections, section_relevance_list): + section_id_to_relevance[sec.center_chunk.unique_id] = rel # re-order docs with all the "relevant" docs at the front - sections = reorder_sections( + sections = _reorder_sections( sections=sections, section_relevance_list=section_relevance_list ) # remove docs that are explicitly marked as not for QA @@ -274,6 +319,14 @@ def _apply_pruning( ) sections = [sections[0]] + # sort by relevance, then by score (as we added the keep_sections first) + sections.sort( + key=lambda s: ( + not section_id_to_relevance.get(s.center_chunk.unique_id, True), + -(s.center_chunk.score or 0.0), + ), + ) + return sections @@ -289,9 +342,16 @@ def prune_sections( if section_relevance_list is not None: assert len(sections) == len(section_relevance_list) + # get federated sections (up to NUM_FEDERATED_SECTIONS) + # TODO: if we can somehow score the federated sections well, we don't need this + federated_sections, normal_sections, normal_section_relevance_list = ( + _separate_federated_sections(sections, section_relevance_list) + ) + actual_num_chunks = ( contextual_pruning_config.max_chunks * contextual_pruning_config.num_chunk_multiple + + len(federated_sections) if contextual_pruning_config.max_chunks else None ) @@ -307,8 +367,9 @@ def prune_sections( ) return _apply_pruning( - sections=sections, - section_relevance_list=section_relevance_list, + sections=normal_sections, + section_relevance_list=normal_section_relevance_list, + keep_sections=federated_sections, token_limit=token_limit, is_manually_selected_docs=contextual_pruning_config.is_manually_selected_docs, use_sections=contextual_pruning_config.use_sections, # Now default True diff --git a/backend/onyx/configs/app_configs.py b/backend/onyx/configs/app_configs.py index dd84cd76dc..6d722381ce 100644 --- a/backend/onyx/configs/app_configs.py +++ b/backend/onyx/configs/app_configs.py @@ -481,6 +481,7 @@ LINEAR_CLIENT_SECRET = os.getenv("LINEAR_CLIENT_SECRET") # Slack specific configs SLACK_NUM_THREADS = int(os.getenv("SLACK_NUM_THREADS") or 8) +MAX_SLACK_QUERY_EXPANSIONS = int(os.environ.get("MAX_SLACK_QUERY_EXPANSIONS", "5")) DASK_JOB_CLIENT_ENABLED = ( os.environ.get("DASK_JOB_CLIENT_ENABLED", "").lower() == "true" @@ -654,6 +655,14 @@ except json.JSONDecodeError: # LLM Model Update API endpoint LLM_MODEL_UPDATE_API_URL = os.environ.get("LLM_MODEL_UPDATE_API_URL") +# Federated Search Configs +MAX_FEDERATED_SECTIONS = int( + os.environ.get("MAX_FEDERATED_SECTIONS", "5") +) # max no. of federated sections to always keep +MAX_FEDERATED_CHUNKS = int( + os.environ.get("MAX_FEDERATED_CHUNKS", "5") +) # max no. of chunks to retrieve per federated connector + ##### # Enterprise Edition Configs ##### diff --git a/backend/onyx/configs/constants.py b/backend/onyx/configs/constants.py index 4c26f0a710..18bfb61496 100644 --- a/backend/onyx/configs/constants.py +++ b/backend/onyx/configs/constants.py @@ -192,6 +192,15 @@ class DocumentSource(str, Enum): MOCK_CONNECTOR = "mock_connector" +class FederatedConnectorSource(str, Enum): + FEDERATED_SLACK = "federated_slack" + + def to_non_federated_source(self) -> DocumentSource | None: + if self == FederatedConnectorSource.FEDERATED_SLACK: + return DocumentSource.SLACK + return None + + DocumentSourceRequiringTenantContext: list[DocumentSource] = [DocumentSource.FILE] diff --git a/backend/onyx/context/search/federated/models.py b/backend/onyx/context/search/federated/models.py new file mode 100644 index 0000000000..d453cc50fb --- /dev/null +++ b/backend/onyx/context/search/federated/models.py @@ -0,0 +1,18 @@ +from datetime import datetime + +from pydantic import BaseModel + + +class SlackMessage(BaseModel): + document_id: str + channel_id: str + message_id: str + thread_id: str | None + link: str + metadata: dict[str, str | list[str]] + timestamp: datetime + recency_bias: float + semantic_identifier: str + text: str + highlighted_texts: set[str] + slack_score: float diff --git a/backend/onyx/context/search/federated/slack_search.py b/backend/onyx/context/search/federated/slack_search.py new file mode 100644 index 0000000000..c9bf5a310c --- /dev/null +++ b/backend/onyx/context/search/federated/slack_search.py @@ -0,0 +1,416 @@ +import re +from datetime import datetime +from datetime import timedelta +from typing import Any + +from langchain_core.messages import HumanMessage +from slack_sdk import WebClient +from slack_sdk.errors import SlackApiError +from sqlalchemy.orm import Session + +from onyx.configs.app_configs import ENABLE_CONTEXTUAL_RAG +from onyx.configs.app_configs import MAX_SLACK_QUERY_EXPANSIONS +from onyx.configs.chat_configs import DOC_TIME_DECAY +from onyx.configs.model_configs import DOC_EMBEDDING_CONTEXT_SIZE +from onyx.connectors.models import IndexingDocument +from onyx.connectors.models import TextSection +from onyx.context.search.federated.models import SlackMessage +from onyx.context.search.models import InferenceChunk +from onyx.context.search.models import SearchQuery +from onyx.db.document import DocumentSource +from onyx.db.search_settings import get_current_search_settings +from onyx.document_index.document_index_utils import ( + get_multipass_config, +) +from onyx.indexing.chunker import Chunker +from onyx.indexing.embedder import DefaultIndexingEmbedder +from onyx.indexing.models import DocAwareChunk +from onyx.llm.factory import get_default_llms +from onyx.llm.interfaces import LLM +from onyx.llm.utils import message_to_string +from onyx.prompts.federated_search import SLACK_QUERY_EXPANSION_PROMPT +from onyx.utils.logger import setup_logger +from onyx.utils.threadpool_concurrency import run_functions_tuples_in_parallel +from onyx.utils.timing import log_function_time + +logger = setup_logger() + +HIGHLIGHT_START_CHAR = "\ue000" +HIGHLIGHT_END_CHAR = "\ue001" + + +def build_slack_queries(query: SearchQuery, llm: LLM) -> list[str]: + # get time filter + time_filter = "" + time_cutoff = query.filters.time_cutoff + if time_cutoff is not None: + # slack after: is exclusive, so we need to subtract one day + time_cutoff = time_cutoff - timedelta(days=1) + time_filter = f" after:{time_cutoff.strftime('%Y-%m-%d')}" + + # use llm to generate slack queries + prompt = SLACK_QUERY_EXPANSION_PROMPT.format(query=query.query) + try: + msg = HumanMessage(content=prompt) + response = llm.invoke([msg]) + rephrased_queries = message_to_string(response).split("\n") + except Exception as e: + logger.error(f"Error expanding query: {e}") + rephrased_queries = [query.query] + + return [ + rephrased_query.strip() + time_filter + for rephrased_query in rephrased_queries[:MAX_SLACK_QUERY_EXPANSIONS] + ] + + +def query_slack( + query_string: str, + original_query: SearchQuery, + access_token: str, + limit: int | None = None, +) -> list[SlackMessage]: + # query slack + slack_client = WebClient(token=access_token) + try: + response = slack_client.search_messages( + query=query_string, count=limit, highlight=True + ) + response.validate() + messages: dict[str, Any] = response.get("messages", {}) + matches: list[dict[str, Any]] = messages.get("matches", []) + except SlackApiError as e: + logger.error(f"Slack API error in query_slack: {e}") + return [] + + # convert matches to slack messages + slack_messages: list[SlackMessage] = [] + for match in matches: + text: str | None = match.get("text") + permalink: str | None = match.get("permalink") + message_id: str | None = match.get("ts") + channel_id: str | None = match.get("channel", {}).get("id") + channel_name: str | None = match.get("channel", {}).get("name") + username: str | None = match.get("username") + score: float = match.get("score", 0.0) + if ( # can't use any() because of type checking :( + not text + or not permalink + or not message_id + or not channel_id + or not channel_name + or not username + ): + continue + + # generate thread id and document id + thread_id = ( + permalink.split("?thread_ts=", 1)[1] if "?thread_ts=" in permalink else None + ) + document_id = f"{channel_id}_{message_id}" + + # compute recency bias (parallels vespa calculation) and metadata + decay_factor = DOC_TIME_DECAY * original_query.recency_bias_multiplier + doc_time = datetime.fromtimestamp(float(message_id)) + doc_age_years = (datetime.now() - doc_time).total_seconds() / ( + 365 * 24 * 60 * 60 + ) + recency_bias = max(1 / (1 + decay_factor * doc_age_years), 0.75) + metadata: dict[str, str | list[str]] = { + "channel": channel_name, + "time": doc_time.isoformat(), + } + + # extract out the highlighted texts + highlighted_texts = set( + re.findall( + rf"{re.escape(HIGHLIGHT_START_CHAR)}(.*?){re.escape(HIGHLIGHT_END_CHAR)}", + text, + ) + ) + cleaned_text = text.replace(HIGHLIGHT_START_CHAR, "").replace( + HIGHLIGHT_END_CHAR, "" + ) + + # get the semantic identifier + snippet = ( + cleaned_text[:50].rstrip() + "..." if len(cleaned_text) > 50 else text + ).replace("\n", " ") + doc_sem_id = f"{username} in #{channel_name}: {snippet}" + + slack_messages.append( + SlackMessage( + document_id=document_id, + channel_id=channel_id, + message_id=message_id, + thread_id=thread_id, + link=permalink, + metadata=metadata, + timestamp=doc_time, + recency_bias=recency_bias, + semantic_identifier=doc_sem_id, + text=f"{username}: {cleaned_text}", + highlighted_texts=highlighted_texts, + slack_score=score, + ) + ) + + return slack_messages + + +def merge_slack_messages( + slack_messages: list[list[SlackMessage]], +) -> tuple[list[SlackMessage], dict[str, SlackMessage]]: + merged_messages: list[SlackMessage] = [] + docid_to_message: dict[str, SlackMessage] = {} + + for messages in slack_messages: + for message in messages: + if message.document_id in docid_to_message: + # update the score and highlighted texts, rest should be identical + docid_to_message[message.document_id].slack_score = max( + docid_to_message[message.document_id].slack_score, + message.slack_score, + ) + docid_to_message[message.document_id].highlighted_texts.update( + message.highlighted_texts + ) + continue + + # add the message to the list + docid_to_message[message.document_id] = message + merged_messages.append(message) + + # re-sort by score + merged_messages.sort(key=lambda x: x.slack_score, reverse=True) + + return merged_messages, docid_to_message + + +def get_contextualized_thread_text(message: SlackMessage, access_token: str) -> str: + """ + Retrieves the initial thread message as well as the text following the message + and combines them into a single string. If the slack query fails, returns the + original message text. + + The idea is that the message (the one that actually matched the search), the + initial thread message, and the replies to the message are important in answering + the user's query. + """ + channel_id = message.channel_id + thread_id = message.thread_id + message_id = message.message_id + + # if it's not a thread, return the message text + if thread_id is None: + return message.text + + # get the thread messages + slack_client = WebClient(token=access_token) + try: + response = slack_client.conversations_replies( + channel=channel_id, + ts=thread_id, + ) + response.validate() + messages: list[dict[str, Any]] = response.get("messages", []) + except SlackApiError as e: + logger.error(f"Slack API error in get_contextualized_thread_text: {e}") + return message.text + + # make sure we didn't get an empty response or a single message (not a thread) + if len(messages) <= 1: + return message.text + + # add the initial thread message + msg_text = messages[0].get("text", "") + msg_sender = messages[0].get("user", "") + thread_text = f"<@{msg_sender}>: {msg_text}" + + # add the message (unless it's the initial message) + thread_text += "\n\nReplies:" + if thread_id == message_id: + message_id_idx = 0 + else: + message_id_idx = next( + (i for i, msg in enumerate(messages) if msg.get("ts") == message_id), 0 + ) + if not message_id_idx: + return thread_text + + # add the message + thread_text += "\n..." if message_id_idx > 1 else "" + msg_text = messages[message_id_idx].get("text", "") + msg_sender = messages[message_id_idx].get("user", "") + thread_text += f"\n<@{msg_sender}>: {msg_text}" + + # add the following replies to the thread text + len_replies = 0 + for msg in messages[message_id_idx + 1 :]: + msg_text = msg.get("text", "") + msg_sender = msg.get("user", "") + reply = f"\n\n<@{msg_sender}>: {msg_text}" + thread_text += reply + + # stop if len_replies exceeds chunk_size * 4 chars as the rest likely won't fit + len_replies += len(reply) + if len_replies >= DOC_EMBEDDING_CONTEXT_SIZE * 4: + thread_text += "\n..." + break + + # replace user ids with names in the thread text + userids: set[str] = set(re.findall(r"<@([A-Z0-9]+)>", thread_text)) + for userid in userids: + try: + response = slack_client.users_profile_get(user=userid) + response.validate() + profile: dict[str, Any] = response.get("profile", {}) + name: str | None = profile.get("real_name") or profile.get("email") + except SlackApiError as e: + logger.error(f"Slack API error in get_contextualized_thread_text: {e}") + continue + if not name: + continue + thread_text = thread_text.replace(f"<@{userid}>", name) + + return thread_text + + +def convert_slack_score(slack_score: float) -> float: + """ + Convert slack score to a score between 0 and 1. + Will affect UI ordering and LLM ordering, but not the pruning. + I.e., should have very little effect on the search/answer quality. + """ + return max(0.0, min(1.0, slack_score / 90_000)) + + +@log_function_time(print_only=True) +def slack_retrieval( + query: SearchQuery, + access_token: str, + db_session: Session, + limit: int | None = None, +) -> list[InferenceChunk]: + # query slack + _, fast_llm = get_default_llms() + query_strings = build_slack_queries(query, fast_llm) + + results: list[list[SlackMessage]] = run_functions_tuples_in_parallel( + [ + (query_slack, (query_string, query, access_token, limit)) + for query_string in query_strings + ] + ) + slack_messages, docid_to_message = merge_slack_messages(results) + slack_messages = slack_messages[: limit or len(slack_messages)] + if not slack_messages: + return [] + + # contextualize the slack messages + thread_texts: list[str] = run_functions_tuples_in_parallel( + [ + (get_contextualized_thread_text, (slack_message, access_token)) + for slack_message in slack_messages + ] + ) + for slack_message, thread_text in zip(slack_messages, thread_texts): + slack_message.text = thread_text + + # get the highlighted texts from shortest to longest + highlighted_texts: set[str] = set() + for slack_message in slack_messages: + highlighted_texts.update(slack_message.highlighted_texts) + sorted_highlighted_texts = sorted(highlighted_texts, key=len) + + # convert slack messages to index documents + index_docs: list[IndexingDocument] = [] + for slack_message in slack_messages: + section: TextSection = TextSection( + text=slack_message.text, link=slack_message.link + ) + index_docs.append( + IndexingDocument( + id=slack_message.document_id, + sections=[section], + processed_sections=[section], + source=DocumentSource.SLACK, + title=slack_message.semantic_identifier, + semantic_identifier=slack_message.semantic_identifier, + metadata=slack_message.metadata, + doc_updated_at=slack_message.timestamp, + ) + ) + + # chunk index docs into doc aware chunks + # a single index doc can get split into multiple chunks + search_settings = get_current_search_settings(db_session) + embedder = DefaultIndexingEmbedder.from_db_search_settings( + search_settings=search_settings + ) + multipass_config = get_multipass_config(search_settings) + enable_contextual_rag = ( + search_settings.enable_contextual_rag or ENABLE_CONTEXTUAL_RAG + ) + chunker = Chunker( + tokenizer=embedder.embedding_model.tokenizer, + enable_multipass=multipass_config.multipass_indexing, + enable_large_chunks=multipass_config.enable_large_chunks, + enable_contextual_rag=enable_contextual_rag, + ) + chunks = chunker.chunk(index_docs) + + # prune chunks without any highlighted texts + relevant_chunks: list[DocAwareChunk] = [] + chunkid_to_match_highlight: dict[str, str] = {} + for chunk in chunks: + match_highlight = chunk.content + for highlight in sorted_highlighted_texts: # faster than re sub + match_highlight = match_highlight.replace( + highlight, f"{highlight}" + ) + + # if nothing got replaced, the chunk is irrelevant + if len(match_highlight) == len(chunk.content): + continue + + chunk_id = f"{chunk.source_document.id}__{chunk.chunk_id}" + relevant_chunks.append(chunk) + chunkid_to_match_highlight[chunk_id] = match_highlight + if limit and len(relevant_chunks) >= limit: + break + + # convert to inference chunks + top_chunks: list[InferenceChunk] = [] + for chunk in relevant_chunks: + document_id = chunk.source_document.id + chunk_id = f"{document_id}__{chunk.chunk_id}" + + top_chunks.append( + InferenceChunk( + chunk_id=chunk.chunk_id, + blurb=chunk.blurb, + content=chunk.content, + source_links=chunk.source_links, + image_file_id=chunk.image_file_id, + section_continuation=chunk.section_continuation, + semantic_identifier=docid_to_message[document_id].semantic_identifier, + document_id=document_id, + source_type=DocumentSource.SLACK, + title=chunk.title_prefix, + boost=0, + recency_bias=docid_to_message[document_id].recency_bias, + score=convert_slack_score(docid_to_message[document_id].slack_score), + hidden=False, + is_relevant=None, + relevance_explanation="", + metadata=docid_to_message[document_id].metadata, + match_highlights=[chunkid_to_match_highlight[chunk_id]], + doc_summary="", + chunk_context="", + updated_at=docid_to_message[document_id].timestamp, + is_federated=True, + ) + ) + + return top_chunks diff --git a/backend/onyx/context/search/models.py b/backend/onyx/context/search/models.py index 148a10cbf2..a9e632ef3d 100644 --- a/backend/onyx/context/search/models.py +++ b/backend/onyx/context/search/models.py @@ -252,6 +252,8 @@ class InferenceChunk(BaseChunk): secondary_owners: list[str] | None = None large_chunk_reference_ids: list[int] = Field(default_factory=list) + is_federated: bool = False + @property def unique_id(self) -> str: return f"{self.document_id}__{self.chunk_id}" diff --git a/backend/onyx/context/search/pipeline.py b/backend/onyx/context/search/pipeline.py index 71f71f512a..3967a666e0 100644 --- a/backend/onyx/context/search/pipeline.py +++ b/backend/onyx/context/search/pipeline.py @@ -158,6 +158,7 @@ class SearchPipeline: # These chunks do not include large chunks and have been deduped self._retrieved_chunks = retrieve_chunks( query=self.search_query, + user_id=self.user.id if self.user else None, document_index=self.document_index, db_session=self.db_session, retrieval_metrics_callback=self.retrieval_metrics_callback, diff --git a/backend/onyx/context/search/retrieval/search_runner.py b/backend/onyx/context/search/retrieval/search_runner.py index 5320340567..9d73532f91 100644 --- a/backend/onyx/context/search/retrieval/search_runner.py +++ b/backend/onyx/context/search/retrieval/search_runner.py @@ -1,5 +1,6 @@ import string from collections.abc import Callable +from uuid import UUID import nltk # type:ignore from sqlalchemy.orm import Session @@ -17,15 +18,18 @@ from onyx.context.search.models import SearchQuery from onyx.context.search.postprocessing.postprocessing import cleanup_chunks from onyx.context.search.preprocessing.preprocessing import HYBRID_ALPHA from onyx.context.search.preprocessing.preprocessing import HYBRID_ALPHA_KEYWORD +from onyx.context.search.utils import get_query_embedding +from onyx.context.search.utils import get_query_embeddings from onyx.context.search.utils import inference_section_from_chunks -from onyx.db.search_settings import get_current_search_settings from onyx.db.search_settings import get_multilingual_expansion from onyx.document_index.interfaces import DocumentIndex from onyx.document_index.interfaces import VespaChunkRequest from onyx.document_index.vespa.shared_utils.utils import ( replace_invalid_doc_id_characters, ) -from onyx.natural_language_processing.search_nlp_models import EmbeddingModel +from onyx.federated_connectors.federated_retrieval import ( + get_federated_retrieval_functions, +) from onyx.secondary_llm_flows.query_expansion import multilingual_query_expansion from onyx.utils.logger import setup_logger from onyx.utils.threadpool_concurrency import run_functions_tuples_in_parallel @@ -33,9 +37,6 @@ from onyx.utils.threadpool_concurrency import run_in_background from onyx.utils.threadpool_concurrency import TimeoutThread from onyx.utils.threadpool_concurrency import wait_on_background from onyx.utils.timing import log_function_time -from shared_configs.configs import MODEL_SERVER_HOST -from shared_configs.configs import MODEL_SERVER_PORT -from shared_configs.enums import EmbedTextType from shared_configs.model_server_models import Embedding logger = setup_logger() @@ -115,34 +116,6 @@ def combine_retrieval_results( return sorted_chunks -def get_query_embedding(query: str, db_session: Session) -> Embedding: - search_settings = get_current_search_settings(db_session) - - model = EmbeddingModel.from_db_model( - search_settings=search_settings, - # The below are globally set, this flow always uses the indexing one - server_host=MODEL_SERVER_HOST, - server_port=MODEL_SERVER_PORT, - ) - - query_embedding = model.encode([query], text_type=EmbedTextType.QUERY)[0] - return query_embedding - - -def get_query_embeddings(queries: list[str], db_session: Session) -> list[Embedding]: - search_settings = get_current_search_settings(db_session) - - model = EmbeddingModel.from_db_model( - search_settings=search_settings, - # The below are globally set, this flow always uses the indexing one - server_host=MODEL_SERVER_HOST, - server_port=MODEL_SERVER_PORT, - ) - - query_embedding = model.encode(queries, text_type=EmbedTextType.QUERY) - return query_embedding - - @log_function_time(print_only=True) def doc_index_retrieval( query: SearchQuery, @@ -350,6 +323,7 @@ def _simplify_text(text: str) -> str: def retrieve_chunks( query: SearchQuery, + user_id: UUID | None, document_index: DocumentIndex, db_session: Session, retrieval_metrics_callback: ( @@ -359,14 +333,34 @@ def retrieve_chunks( """Returns a list of the best chunks from an initial keyword/semantic/ hybrid search.""" multilingual_expansion = get_multilingual_expansion(db_session) - # Don't do query expansion on complex queries, rephrasings likely would not work well - if not multilingual_expansion or "\n" in query.query or "\r" in query.query: - top_chunks = doc_index_retrieval( - query=query, document_index=document_index, db_session=db_session - ) - else: + run_queries: list[tuple[Callable, tuple]] = [] + + source_filters = ( + set(query.filters.source_type) if query.filters.source_type else None + ) + + # Federated retrieval + federated_retrieval_infos = get_federated_retrieval_functions( + db_session, user_id, query.filters.source_type, query.filters.document_set + ) + federated_sources = set( + federated_retrieval_info.source.to_non_federated_source() + for federated_retrieval_info in federated_retrieval_infos + ) + for federated_retrieval_info in federated_retrieval_infos: + run_queries.append((federated_retrieval_info.retrieval_function, (query,))) + + # Normal retrieval + normal_search_enabled = (source_filters is None) or ( + len(set(source_filters) - federated_sources) > 0 + ) + if normal_search_enabled and ( + not multilingual_expansion or "\n" in query.query or "\r" in query.query + ): + # Don't do query expansion on complex queries, rephrasings likely would not work well + run_queries.append((doc_index_retrieval, (query, document_index, db_session))) + elif normal_search_enabled: simplified_queries = set() - run_queries: list[tuple[Callable, tuple]] = [] # Currently only uses query expansion on multilingual use cases query_rephrases = multilingual_query_expansion( @@ -393,13 +387,11 @@ def retrieve_chunks( deep=True, ) run_queries.append( - ( - doc_index_retrieval, - (q_copy, document_index, db_session), - ) + (doc_index_retrieval, (q_copy, document_index, db_session)) ) - parallel_search_results = run_functions_tuples_in_parallel(run_queries) - top_chunks = combine_retrieval_results(parallel_search_results) + + parallel_search_results = run_functions_tuples_in_parallel(run_queries) + top_chunks = combine_retrieval_results(parallel_search_results) if not top_chunks: logger.warning( diff --git a/backend/onyx/context/search/utils.py b/backend/onyx/context/search/utils.py index b0fc7c7e2b..607eb877ff 100644 --- a/backend/onyx/context/search/utils.py +++ b/backend/onyx/context/search/utils.py @@ -4,6 +4,7 @@ from typing import TypeVar from nltk.corpus import stopwords # type:ignore from nltk.tokenize import word_tokenize # type:ignore +from sqlalchemy.orm import Session from onyx.chat.models import SectionRelevancePiece from onyx.context.search.models import InferenceChunk @@ -12,7 +13,13 @@ from onyx.context.search.models import SavedSearchDoc from onyx.context.search.models import SavedSearchDocWithContent from onyx.context.search.models import SearchDoc from onyx.db.models import SearchDoc as DBSearchDoc +from onyx.db.search_settings import get_current_search_settings +from onyx.natural_language_processing.search_nlp_models import EmbeddingModel from onyx.utils.logger import setup_logger +from shared_configs.configs import MODEL_SERVER_HOST +from shared_configs.configs import MODEL_SERVER_PORT +from shared_configs.enums import EmbedTextType +from shared_configs.model_server_models import Embedding logger = setup_logger() @@ -160,3 +167,21 @@ def remove_stop_words_and_punctuation(keywords: list[str]) -> list[str]: except Exception as e: logger.warning(f"Error removing stop words and punctuation: {e}") return keywords + + +def get_query_embeddings(queries: list[str], db_session: Session) -> list[Embedding]: + search_settings = get_current_search_settings(db_session) + + model = EmbeddingModel.from_db_model( + search_settings=search_settings, + # The below are globally set, this flow always uses the indexing one + server_host=MODEL_SERVER_HOST, + server_port=MODEL_SERVER_PORT, + ) + + query_embedding = model.encode(queries, text_type=EmbedTextType.QUERY) + return query_embedding + + +def get_query_embedding(query: str, db_session: Session) -> Embedding: + return get_query_embeddings([query], db_session)[0] diff --git a/backend/onyx/db/document_set.py b/backend/onyx/db/document_set.py index 025e48876c..e356ee0069 100644 --- a/backend/onyx/db/document_set.py +++ b/backend/onyx/db/document_set.py @@ -10,6 +10,7 @@ from sqlalchemy import or_ from sqlalchemy import Select from sqlalchemy import select from sqlalchemy.orm import aliased +from sqlalchemy.orm import selectinload from sqlalchemy.orm import Session from onyx.configs.app_configs import DISABLE_AUTH @@ -122,7 +123,11 @@ def get_document_set_by_id_for_user( user: User | None, get_editable: bool = True, ) -> DocumentSetDBModel | None: - stmt = select(DocumentSetDBModel).distinct() + stmt = ( + select(DocumentSetDBModel) + .distinct() + .options(selectinload(DocumentSetDBModel.federated_connectors)) + ) stmt = stmt.where(DocumentSetDBModel.id == document_set_id) stmt = _add_user_filters(stmt=stmt, user=user, get_editable=get_editable) return db_session.scalar(stmt) @@ -145,6 +150,16 @@ def get_document_set_by_name( ) +def get_document_sets_by_name( + db_session: Session, document_set_names: list[str] +) -> Sequence[DocumentSetDBModel]: + return db_session.scalars( + select(DocumentSetDBModel).where( + DocumentSetDBModel.name.in_(document_set_names) + ) + ).all() + + def get_document_sets_by_ids( db_session: Session, document_set_ids: list[int] ) -> Sequence[DocumentSetDBModel]: @@ -210,9 +225,12 @@ def insert_document_set( user_id: UUID | None, db_session: Session, ) -> tuple[DocumentSetDBModel, list[DocumentSet__ConnectorCredentialPair]]: - if not document_set_creation_request.cc_pair_ids: - # It's cc-pairs in actuality but the UI displays this error - raise ValueError("Cannot create a document set with no Connectors") + # Check if we have either CC pairs or federated connectors (or both) + if ( + not document_set_creation_request.cc_pair_ids + and not document_set_creation_request.federated_connectors + ): + raise ValueError("Cannot create a document set with no connectors") if not document_set_creation_request.is_public: _check_if_cc_pairs_are_owned_by_groups( @@ -234,6 +252,7 @@ def insert_document_set( db_session.add(new_document_set_row) db_session.flush() # ensure the new document set gets assigned an ID + # Create CC pair mappings ds_cc_pairs = [ DocumentSet__ConnectorCredentialPair( document_set_id=new_document_set_row.id, @@ -244,6 +263,17 @@ def insert_document_set( ] db_session.add_all(ds_cc_pairs) + # Create federated connector mappings + from onyx.db.federated import create_federated_connector_document_set_mapping + + for fc_config in document_set_creation_request.federated_connectors: + create_federated_connector_document_set_mapping( + db_session=db_session, + federated_connector_id=fc_config.federated_connector_id, + document_set_id=new_document_set_row.id, + entities=fc_config.entities, + ) + versioned_private_doc_set_fn = fetch_versioned_implementation( "onyx.db.document_set", "make_doc_set_private" ) @@ -260,6 +290,7 @@ def insert_document_set( except Exception as e: db_session.rollback() logger.error(f"Error creating document set: {e}") + raise return new_document_set_row, ds_cc_pairs @@ -273,9 +304,12 @@ def update_document_set( That will be processed via Celery in check_for_vespa_sync_task and trigger a long running background sync to Vespa. """ - if not document_set_update_request.cc_pair_ids: - # It's cc-pairs in actuality but the UI displays this error - raise ValueError("Cannot create a document set with no Connectors") + # Check if we have either CC pairs or federated connectors (or both) + if ( + not document_set_update_request.cc_pair_ids + and not document_set_update_request.federated_connectors + ): + raise ValueError("Cannot update a document set with no connectors") if not document_set_update_request.is_public: _check_if_cc_pairs_are_owned_by_groups( @@ -333,6 +367,27 @@ def update_document_set( for cc_pair_id in document_set_update_request.cc_pair_ids ] db_session.add_all(ds_cc_pairs) + + # Update federated connector mappings + from onyx.db.federated import create_federated_connector_document_set_mapping + from onyx.db.models import FederatedConnector__DocumentSet + from sqlalchemy import delete + + # Delete existing federated connector mappings for this document set + delete_stmt = delete(FederatedConnector__DocumentSet).where( + FederatedConnector__DocumentSet.document_set_id == document_set_row.id + ) + db_session.execute(delete_stmt) + + # Create new federated connector mappings + for fc_config in document_set_update_request.federated_connectors: + create_federated_connector_document_set_mapping( + db_session=db_session, + federated_connector_id=fc_config.federated_connector_id, + document_set_id=document_set_row.id, + entities=fc_config.entities, + ) + db_session.commit() except: db_session.rollback() @@ -492,7 +547,11 @@ def fetch_all_document_sets_for_user( user: User | None, get_editable: bool = True, ) -> Sequence[DocumentSetDBModel]: - stmt = select(DocumentSetDBModel).distinct() + stmt = ( + select(DocumentSetDBModel) + .distinct() + .options(selectinload(DocumentSetDBModel.federated_connectors)) + ) stmt = _add_user_filters(stmt, user, get_editable=get_editable) return db_session.scalars(stmt).all() diff --git a/backend/onyx/db/federated.py b/backend/onyx/db/federated.py new file mode 100644 index 0000000000..2555d38866 --- /dev/null +++ b/backend/onyx/db/federated.py @@ -0,0 +1,273 @@ +from datetime import datetime +from typing import Any +from uuid import UUID + +from sqlalchemy import select +from sqlalchemy.orm import joinedload +from sqlalchemy.orm import selectinload +from sqlalchemy.orm import Session + +from onyx.configs.constants import FederatedConnectorSource +from onyx.db.models import DocumentSet +from onyx.db.models import FederatedConnector +from onyx.db.models import FederatedConnector__DocumentSet +from onyx.db.models import FederatedConnectorOAuthToken +from onyx.federated_connectors.factory import get_federated_connector +from onyx.utils.logger import setup_logger + +logger = setup_logger() + + +def fetch_federated_connector_by_id( + federated_connector_id: int, db_session: Session +) -> FederatedConnector | None: + """Fetch a federated connector by its ID.""" + stmt = select(FederatedConnector).where( + FederatedConnector.id == federated_connector_id + ) + result = db_session.execute(stmt) + return result.scalar_one_or_none() + + +def fetch_all_federated_connectors(db_session: Session) -> list[FederatedConnector]: + """Fetch all federated connectors with their OAuth tokens and document sets.""" + stmt = select(FederatedConnector).options( + selectinload(FederatedConnector.oauth_tokens), + selectinload(FederatedConnector.document_sets), + ) + result = db_session.execute(stmt) + return list(result.scalars().all()) + + +def validate_federated_connector_credentials( + source: FederatedConnectorSource, + credentials: dict[str, Any], +) -> bool: + """Validate credentials for a federated connector using the connector's validation logic.""" + try: + # the initialization will fail if the credentials are invalid + get_federated_connector(source, credentials) + return True + except Exception as e: + logger.error(f"Error validating credentials for source {source}: {e}") + return False + + +def create_federated_connector( + db_session: Session, + source: FederatedConnectorSource, + credentials: dict[str, Any], +) -> FederatedConnector: + """Create a new federated connector with credential validation.""" + # Validate credentials before creating + if not validate_federated_connector_credentials(source, credentials): + raise ValueError( + f"Invalid credentials for federated connector source: {source}" + ) + + federated_connector = FederatedConnector( + source=source, + credentials=credentials, + ) + db_session.add(federated_connector) + db_session.commit() + return federated_connector + + +def update_federated_connector_oauth_token( + db_session: Session, + federated_connector_id: int, + user_id: UUID, + token: str, + expires_at: datetime | None = None, +) -> FederatedConnectorOAuthToken: + """Update or create OAuth token for a federated connector and user.""" + # First, try to find existing token for this user and connector + stmt = select(FederatedConnectorOAuthToken).where( + FederatedConnectorOAuthToken.federated_connector_id == federated_connector_id, + FederatedConnectorOAuthToken.user_id == user_id, + ) + existing_token = db_session.execute(stmt).scalar_one_or_none() + + if existing_token: + # Update existing token + existing_token.token = token + existing_token.expires_at = expires_at + db_session.commit() + return existing_token + else: + # Create new token + oauth_token = FederatedConnectorOAuthToken( + federated_connector_id=federated_connector_id, + user_id=user_id, + token=token, + expires_at=expires_at, + ) + db_session.add(oauth_token) + db_session.commit() + return oauth_token + + +def get_federated_connector_oauth_token( + db_session: Session, + federated_connector_id: int, + user_id: UUID, +) -> FederatedConnectorOAuthToken | None: + """Get OAuth token for a federated connector and user.""" + stmt = select(FederatedConnectorOAuthToken).where( + FederatedConnectorOAuthToken.federated_connector_id == federated_connector_id, + FederatedConnectorOAuthToken.user_id == user_id, + ) + result = db_session.execute(stmt) + return result.scalar_one_or_none() + + +def list_federated_connector_oauth_tokens( + db_session: Session, + user_id: UUID, +) -> list[FederatedConnectorOAuthToken]: + """List all OAuth tokens for all federated connectors.""" + stmt = ( + select(FederatedConnectorOAuthToken) + .where( + FederatedConnectorOAuthToken.user_id == user_id, + ) + .options( + joinedload(FederatedConnectorOAuthToken.federated_connector), + ) + ) + result = db_session.scalars(stmt) + return list(result) + + +def create_federated_connector_document_set_mapping( + db_session: Session, + federated_connector_id: int, + document_set_id: int, + entities: dict[str, Any], +) -> FederatedConnector__DocumentSet: + """Create a mapping between federated connector and document set with entities.""" + mapping = FederatedConnector__DocumentSet( + federated_connector_id=federated_connector_id, + document_set_id=document_set_id, + entities=entities, + ) + db_session.add(mapping) + db_session.commit() + return mapping + + +def update_federated_connector_document_set_entities( + db_session: Session, + federated_connector_id: int, + document_set_id: int, + entities: dict[str, Any], +) -> FederatedConnector__DocumentSet | None: + """Update entities for a federated connector document set mapping.""" + stmt = select(FederatedConnector__DocumentSet).where( + FederatedConnector__DocumentSet.federated_connector_id + == federated_connector_id, + FederatedConnector__DocumentSet.document_set_id == document_set_id, + ) + mapping = db_session.execute(stmt).scalar_one_or_none() + + if mapping: + mapping.entities = entities + db_session.commit() + return mapping + + return None + + +def get_federated_connector_document_set_mappings( + db_session: Session, + federated_connector_id: int, +) -> list[FederatedConnector__DocumentSet]: + """Get all document set mappings for a federated connector.""" + stmt = select(FederatedConnector__DocumentSet).where( + FederatedConnector__DocumentSet.federated_connector_id == federated_connector_id + ) + result = db_session.execute(stmt) + return list(result.scalars().all()) + + +def delete_federated_connector_document_set_mapping( + db_session: Session, + federated_connector_id: int, + document_set_id: int, +) -> bool: + """Delete a federated connector document set mapping.""" + stmt = select(FederatedConnector__DocumentSet).where( + FederatedConnector__DocumentSet.federated_connector_id + == federated_connector_id, + FederatedConnector__DocumentSet.document_set_id == document_set_id, + ) + mapping = db_session.execute(stmt).scalar_one_or_none() + + if mapping: + db_session.delete(mapping) + db_session.commit() + return True + + return False + + +def get_federated_connector_document_set_mappings_by_document_set_names( + db_session: Session, + document_set_names: list[str], +) -> list[FederatedConnector__DocumentSet]: + """Get all document set mappings for a federated connector by document set names.""" + stmt = ( + select(FederatedConnector__DocumentSet) + .join( + DocumentSet, + FederatedConnector__DocumentSet.document_set_id == DocumentSet.id, + ) + .where(DocumentSet.name.in_(document_set_names)) + ) + result = db_session.scalars(stmt) + return list(result) + + +def update_federated_connector( + db_session: Session, + federated_connector_id: int, + credentials: dict[str, Any] | None = None, +) -> FederatedConnector | None: + """Update a federated connector with credential validation.""" + federated_connector = fetch_federated_connector_by_id( + federated_connector_id, db_session + ) + if not federated_connector: + return None + + if credentials is not None: + # Validate credentials before updating + if not validate_federated_connector_credentials( + federated_connector.source, credentials + ): + raise ValueError( + f"Invalid credentials for federated connector source: {federated_connector.source}" + ) + federated_connector.credentials = credentials + + db_session.commit() + return federated_connector + + +def delete_federated_connector( + db_session: Session, + federated_connector_id: int, +) -> bool: + """Delete a federated connector and all its related data.""" + federated_connector = fetch_federated_connector_by_id( + federated_connector_id, db_session + ) + if not federated_connector: + return False + + # Delete related OAuth tokens (cascade should handle this) + # Delete related document set mappings (cascade should handle this) + db_session.delete(federated_connector) + db_session.commit() + return True diff --git a/backend/onyx/db/models.py b/backend/onyx/db/models.py index e1432a887e..9edc0db79d 100644 --- a/backend/onyx/db/models.py +++ b/backend/onyx/db/models.py @@ -44,7 +44,11 @@ from sqlalchemy import PrimaryKeyConstraint from onyx.auth.schemas import UserRole from onyx.configs.chat_configs import NUM_POSTPROCESSED_RESULTS -from onyx.configs.constants import DEFAULT_BOOST, MilestoneRecordType +from onyx.configs.constants import ( + DEFAULT_BOOST, + FederatedConnectorSource, + MilestoneRecordType, +) from onyx.configs.constants import DocumentSource from onyx.configs.constants import FileOrigin from onyx.configs.constants import MessageType @@ -1407,6 +1411,80 @@ class Credential(Base): user: Mapped[User | None] = relationship("User", back_populates="credentials") +class FederatedConnector(Base): + __tablename__ = "federated_connector" + + id: Mapped[int] = mapped_column(primary_key=True) + source: Mapped[FederatedConnectorSource] = mapped_column( + Enum(FederatedConnectorSource, native_enum=False) + ) + credentials: Mapped[dict[str, str]] = mapped_column(EncryptedJson(), nullable=False) + + oauth_tokens: Mapped[list["FederatedConnectorOAuthToken"]] = relationship( + "FederatedConnectorOAuthToken", + back_populates="federated_connector", + cascade="all, delete-orphan", + ) + document_sets: Mapped[list["FederatedConnector__DocumentSet"]] = relationship( + "FederatedConnector__DocumentSet", + back_populates="federated_connector", + cascade="all, delete-orphan", + ) + + +class FederatedConnectorOAuthToken(Base): + """NOTE: in the future, can be made more general to support OAuth tokens + for actions.""" + + __tablename__ = "federated_connector_oauth_token" + + id: Mapped[int] = mapped_column(primary_key=True) + federated_connector_id: Mapped[int] = mapped_column( + ForeignKey("federated_connector.id", ondelete="CASCADE"), nullable=False + ) + user_id: Mapped[UUID] = mapped_column( + ForeignKey("user.id", ondelete="CASCADE"), nullable=False + ) + token: Mapped[str] = mapped_column(EncryptedString(), nullable=False) + expires_at: Mapped[datetime.datetime | None] = mapped_column( + DateTime, nullable=True + ) + + federated_connector: Mapped["FederatedConnector"] = relationship( + "FederatedConnector", back_populates="oauth_tokens" + ) + user: Mapped["User"] = relationship("User") + + +class FederatedConnector__DocumentSet(Base): + __tablename__ = "federated_connector__document_set" + + id: Mapped[int] = mapped_column(primary_key=True) + federated_connector_id: Mapped[int] = mapped_column( + ForeignKey("federated_connector.id", ondelete="CASCADE"), nullable=False + ) + document_set_id: Mapped[int] = mapped_column( + ForeignKey("document_set.id", ondelete="CASCADE"), nullable=False + ) + # unique per source type. Validated before insertion. + entities: Mapped[dict[str, Any]] = mapped_column(postgresql.JSONB(), nullable=False) + + federated_connector: Mapped["FederatedConnector"] = relationship( + "FederatedConnector", back_populates="document_sets" + ) + document_set: Mapped["DocumentSet"] = relationship( + "DocumentSet", back_populates="federated_connectors" + ) + + __table_args__ = ( + UniqueConstraint( + "federated_connector_id", + "document_set_id", + name="uq_federated_connector_document_set", + ), + ) + + class SearchSettings(Base): __tablename__ = "search_settings" @@ -2312,6 +2390,13 @@ class DocumentSet(Base): secondary="document_set__user_group", viewonly=True, ) + federated_connectors: Mapped[list["FederatedConnector__DocumentSet"]] = ( + relationship( + "FederatedConnector__DocumentSet", + back_populates="document_set", + cascade="all, delete-orphan", + ) + ) class Prompt(Base): diff --git a/backend/onyx/federated_connectors/__init__.py b/backend/onyx/federated_connectors/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/backend/onyx/federated_connectors/factory.py b/backend/onyx/federated_connectors/factory.py new file mode 100644 index 0000000000..fbbf482fe2 --- /dev/null +++ b/backend/onyx/federated_connectors/factory.py @@ -0,0 +1,29 @@ +"""Factory for creating federated connector instances.""" + +from typing import Any + +from onyx.configs.constants import FederatedConnectorSource +from onyx.federated_connectors.interfaces import FederatedConnector +from onyx.federated_connectors.slack.federated_connector import SlackFederatedConnector +from onyx.utils.logger import setup_logger + +logger = setup_logger() + + +def get_federated_connector( + source: FederatedConnectorSource, + credentials: dict[str, Any], +) -> FederatedConnector: + """Get an instance of the appropriate federated connector.""" + connector_cls = get_federated_connector_cls(source) + return connector_cls(credentials) + + +def get_federated_connector_cls( + source: FederatedConnectorSource, +) -> type[FederatedConnector]: + """Get the class of the appropriate federated connector.""" + if source == FederatedConnectorSource.FEDERATED_SLACK: + return SlackFederatedConnector + else: + raise ValueError(f"Unsupported federated connector source: {source}") diff --git a/backend/onyx/federated_connectors/federated_retrieval.py b/backend/onyx/federated_connectors/federated_retrieval.py new file mode 100644 index 0000000000..4dc930014c --- /dev/null +++ b/backend/onyx/federated_connectors/federated_retrieval.py @@ -0,0 +1,96 @@ +from collections import defaultdict +from collections.abc import Callable +from uuid import UUID + +from pydantic import BaseModel +from pydantic import ConfigDict +from sqlalchemy.orm import Session + +from onyx.configs.app_configs import MAX_FEDERATED_CHUNKS +from onyx.configs.constants import DocumentSource +from onyx.configs.constants import FederatedConnectorSource +from onyx.context.search.models import InferenceChunk +from onyx.context.search.models import SearchQuery +from onyx.db.federated import ( + get_federated_connector_document_set_mappings_by_document_set_names, +) +from onyx.db.federated import list_federated_connector_oauth_tokens +from onyx.db.models import FederatedConnector__DocumentSet +from onyx.federated_connectors.factory import get_federated_connector +from onyx.utils.logger import setup_logger + + +logger = setup_logger() + + +class FederatedRetrievalInfo(BaseModel): + model_config = ConfigDict(arbitrary_types_allowed=True) + + retrieval_function: Callable[[SearchQuery], list[InferenceChunk]] + source: FederatedConnectorSource + + +def get_federated_retrieval_functions( + db_session: Session, + user_id: UUID | None, + source_types: list[DocumentSource] | None, + document_set_names: list[str] | None, +) -> list[FederatedRetrievalInfo]: + if user_id is None: + logger.warning( + "No user ID provided, skipping federated retrieval. Federated retrieval not " + "supported with AUTH_TYPE=disabled." + ) + return [] + + federated_connector__document_set_pairs = ( + ( + get_federated_connector_document_set_mappings_by_document_set_names( + db_session, document_set_names + ) + ) + if document_set_names + else [] + ) + federated_connector_id_to_document_sets: dict[ + int, list[FederatedConnector__DocumentSet] + ] = defaultdict(list) + for pair in federated_connector__document_set_pairs: + federated_connector_id_to_document_sets[pair.federated_connector_id].append( + pair + ) + + federated_retrieval_infos: list[FederatedRetrievalInfo] = [] + federated_oauth_tokens = list_federated_connector_oauth_tokens(db_session, user_id) + for oauth_token in federated_oauth_tokens: + if ( + source_types is not None + and oauth_token.federated_connector.source.to_non_federated_source() + not in source_types + ): + continue + + document_set_associations = federated_connector_id_to_document_sets[ + oauth_token.federated_connector_id + ] + if document_set_associations: + entities = document_set_associations[0].entities + else: + entities = {} + + connector = get_federated_connector( + oauth_token.federated_connector.source, + oauth_token.federated_connector.credentials, + ) + federated_retrieval_infos.append( + FederatedRetrievalInfo( + retrieval_function=lambda query: connector.search( + query, + entities, + access_token=oauth_token.token, + limit=MAX_FEDERATED_CHUNKS, + ), + source=oauth_token.federated_connector.source, + ) + ) + return federated_retrieval_infos diff --git a/backend/onyx/federated_connectors/interfaces.py b/backend/onyx/federated_connectors/interfaces.py new file mode 100644 index 0000000000..80d02be177 --- /dev/null +++ b/backend/onyx/federated_connectors/interfaces.py @@ -0,0 +1,101 @@ +from abc import ABC +from abc import abstractmethod +from typing import Any +from typing import Dict + +from onyx.context.search.models import InferenceChunk +from onyx.context.search.models import SearchQuery +from onyx.federated_connectors.models import CredentialField +from onyx.federated_connectors.models import EntityField +from onyx.federated_connectors.models import OAuthResult + + +class FederatedConnector(ABC): + """Base interface that all federated connectors must implement.""" + + @abstractmethod + def __init__(self, credentials: dict[str, Any]): + """ + Initialize the connector with credentials + validate their structure. + + Args: + credentials: Dictionary of credentials to initialize the connector with + """ + self.credentials = credentials + + @abstractmethod + def validate_entities(self, entities: Dict[str, Any]) -> bool: + """ + Validate that the provided entities match the expected structure. + + Args: + entities: Dictionary of entities to validate + + Returns: + True if entities are valid, False otherwise + """ + + @classmethod + @abstractmethod + def entities_schema(cls) -> Dict[str, EntityField]: + """ + Return the specification of what entities are available for this connector. + + Returns: + Dictionary where keys are entity names and values are EntityField objects + describing the expected structure and constraints. + """ + + @classmethod + @abstractmethod + def credentials_schema(cls) -> Dict[str, CredentialField]: + """ + Return the specification of what credentials are required for this connector. + + Returns: + Dictionary where keys are credential field names and values are CredentialField objects + describing the expected structure, validation rules, and security properties. + """ + + @abstractmethod + def authorize(self, redirect_uri: str) -> str: + """ + Generate the OAuth authorization URL. + + Returns: + The URL where users should be redirected to authorize the application + """ + + @abstractmethod + def callback(self, callback_data: Dict[str, Any], redirect_uri: str) -> OAuthResult: + """ + Handle the OAuth callback and exchange the authorization code for tokens. + + Args: + callback_data: The data received from the OAuth callback (query params, etc.) + redirect_uri: The OAuth redirect URI used in the authorization request + + Returns: + Standardized OAuthResult containing tokens and metadata + """ + + @abstractmethod + def search( + self, + query: SearchQuery, + entities: dict[str, Any], + access_token: str, + limit: int | None = None, + ) -> list[InferenceChunk]: + """ + Perform a federated search using the provided query and entities. + + Args: + query: The search query + entities: The entities to search within (validated by validate()) + access_token: The OAuth access token + limit: Maximum number of results to return + + Returns: + Search results in a standardized format + """ diff --git a/backend/onyx/federated_connectors/models.py b/backend/onyx/federated_connectors/models.py new file mode 100644 index 0000000000..58963f465b --- /dev/null +++ b/backend/onyx/federated_connectors/models.py @@ -0,0 +1,66 @@ +from datetime import datetime +from typing import Any +from typing import Dict +from typing import Optional + +from pydantic import BaseModel +from pydantic import Field + + +class FieldSpec(BaseModel): + """Model for describing a field specification.""" + + type: str = Field( + ..., description="The type of the field (e.g., 'str', 'bool', 'list[str]')" + ) + description: str = Field( + ..., description="Description of what this field represents" + ) + required: bool = Field(default=False, description="Whether this field is required") + default: Optional[Any] = Field( + default=None, description="Default value if not provided" + ) + example: Optional[Any] = Field( + default=None, description="Example value for documentation" + ) + secret: bool = Field( + default=False, description="Whether this field contains sensitive data" + ) + + +class EntityField(FieldSpec): + """Model for describing an entity field in the entities specification.""" + + +class CredentialField(FieldSpec): + """Model for describing a credential field in the credentials specification.""" + + +class OAuthResult(BaseModel): + """Standardized OAuth result that all federated connectors should return from callback.""" + + access_token: Optional[str] = Field( + default=None, description="The access token received" + ) + token_type: Optional[str] = Field( + default=None, description="Token type (usually 'bearer')" + ) + scope: Optional[str] = Field(default=None, description="Granted scopes") + expires_at: Optional[datetime] = Field( + default=None, description="When the token expires" + ) + refresh_token: Optional[str] = Field( + default=None, description="Refresh token if applicable" + ) + + # Additional fields that might be useful + team: Optional[Dict[str, Any]] = Field( + default=None, description="Team/workspace information" + ) + user: Optional[Dict[str, Any]] = Field(default=None, description="User information") + raw_response: Optional[Dict[str, Any]] = Field( + default=None, description="Raw response for debugging" + ) + + class Config: + json_encoders = {datetime: lambda v: v.isoformat()} diff --git a/backend/onyx/federated_connectors/oauth_utils.py b/backend/onyx/federated_connectors/oauth_utils.py new file mode 100644 index 0000000000..7876d3fda8 --- /dev/null +++ b/backend/onyx/federated_connectors/oauth_utils.py @@ -0,0 +1,169 @@ +"""Generic OAuth utilities for federated connectors API layer.""" + +import base64 +import json +import uuid +from typing import Any +from typing import cast +from typing import Dict +from typing import Optional + +from onyx.configs.app_configs import WEB_DOMAIN +from onyx.redis.redis_pool import get_redis_client +from onyx.utils.logger import setup_logger + +logger = setup_logger() + +# Redis key prefix for OAuth state +OAUTH_STATE_PREFIX = "federated_oauth" +# Default TTL for OAuth state (5 minutes) +OAUTH_STATE_TTL = 300 + + +class OAuthSession: + """Represents an OAuth session stored in Redis.""" + + def __init__( + self, + federated_connector_id: int, + user_id: str, + redirect_uri: Optional[str] = None, + additional_data: Optional[Dict[str, Any]] = None, + ): + self.federated_connector_id = federated_connector_id + self.user_id = user_id + self.redirect_uri = redirect_uri + self.additional_data = additional_data or {} + + def to_dict(self) -> Dict[str, Any]: + """Convert to dictionary for Redis storage.""" + return { + "federated_connector_id": self.federated_connector_id, + "user_id": self.user_id, + "redirect_uri": self.redirect_uri, + "additional_data": self.additional_data, + } + + @classmethod + def from_dict(cls, data: Dict[str, Any]) -> "OAuthSession": + """Create from dictionary retrieved from Redis.""" + return cls( + federated_connector_id=data["federated_connector_id"], + user_id=data["user_id"], + redirect_uri=data.get("redirect_uri"), + additional_data=data.get("additional_data", {}), + ) + + +def generate_oauth_state( + federated_connector_id: int, + user_id: str, + redirect_uri: Optional[str] = None, + additional_data: Optional[Dict[str, Any]] = None, + ttl: int = OAUTH_STATE_TTL, +) -> str: + """ + Generate a secure state parameter and store session data in Redis. + + Args: + federated_connector_id: ID of the federated connector + user_id: ID of the user initiating OAuth + redirect_uri: Optional redirect URI after OAuth completion + additional_data: Any additional data to store with the session + ttl: Time-to-live in seconds for the Redis key + + Returns: + Base64-encoded state parameter + """ + # Generate a random UUID for the state + state_uuid = uuid.uuid4() + + # Convert UUID to base64 for URL-safe state parameter + state_bytes = state_uuid.bytes + state_b64 = base64.urlsafe_b64encode(state_bytes).decode("utf-8").rstrip("=") + + # Create session object + session = OAuthSession( + federated_connector_id=federated_connector_id, + user_id=user_id, + redirect_uri=redirect_uri, + additional_data=additional_data, + ) + + # Store in Redis with TTL + redis_client = get_redis_client() + redis_key = f"{OAUTH_STATE_PREFIX}:{state_uuid}" + + redis_client.set( + redis_key, + json.dumps(session.to_dict()), + ex=ttl, + ) + + logger.info( + f"Generated OAuth state for federated_connector_id={federated_connector_id}, " + f"user_id={user_id}, state={state_b64}" + ) + + return state_b64 + + +def verify_oauth_state(state: str) -> OAuthSession: + """ + Verify OAuth state parameter and retrieve session data. + + Args: + state: Base64-encoded state parameter from OAuth callback + + Returns: + OAuthSession if state is valid, None otherwise + """ + # Add padding if needed for base64 decoding + padded_state = state + "=" * (-len(state) % 4) + + # Decode base64 to get UUID bytes + state_bytes = base64.urlsafe_b64decode(padded_state) + state_uuid = uuid.UUID(bytes=state_bytes) + + # Look up in Redis + redis_client = get_redis_client() + redis_key = f"{OAUTH_STATE_PREFIX}:{state_uuid}" + + session_data = cast(bytes, redis_client.get(redis_key)) + if not session_data: + raise ValueError(f"OAuth state not found in Redis: {state}") + + # Delete the key after retrieval (one-time use) + redis_client.delete(redis_key) + + # Parse and return session + session_dict = json.loads(session_data) + return OAuthSession.from_dict(session_dict) + + +def get_oauth_callback_uri() -> str: + """ + Generate the OAuth callback URI for a federated connector. + + Returns: + The callback URI + """ + # Use the frontend callback page as the OAuth redirect URI + # The frontend will then make an API call to process the callback + return f"{WEB_DOMAIN}/federated/oauth/callback" + + +def add_state_to_oauth_url(base_oauth_url: str, state: str) -> str: + """ + Add state parameter to an OAuth URL. + + Args: + base_oauth_url: The base OAuth URL from the connector + state: The state parameter to add + + Returns: + The OAuth URL with state parameter added + """ + # Check if URL already has query parameters + separator = "&" if "?" in base_oauth_url else "?" + return f"{base_oauth_url}{separator}state={state}" diff --git a/backend/onyx/federated_connectors/slack/__init__.py b/backend/onyx/federated_connectors/slack/__init__.py new file mode 100644 index 0000000000..7223eac309 --- /dev/null +++ b/backend/onyx/federated_connectors/slack/__init__.py @@ -0,0 +1 @@ +# Slack federated connector module diff --git a/backend/onyx/federated_connectors/slack/federated_connector.py b/backend/onyx/federated_connectors/slack/federated_connector.py new file mode 100644 index 0000000000..3fd65a34ec --- /dev/null +++ b/backend/onyx/federated_connectors/slack/federated_connector.py @@ -0,0 +1,236 @@ +from datetime import datetime +from datetime import timedelta +from datetime import timezone +from typing import Any +from urllib.parse import urlencode + +import requests +from pydantic import ValidationError +from typing_extensions import override + +from onyx.context.search.federated.slack_search import slack_retrieval +from onyx.context.search.models import InferenceChunk +from onyx.context.search.models import SearchQuery +from onyx.db.engine.sql_engine import get_session_with_current_tenant +from onyx.federated_connectors.interfaces import FederatedConnector +from onyx.federated_connectors.models import CredentialField +from onyx.federated_connectors.models import EntityField +from onyx.federated_connectors.models import OAuthResult +from onyx.federated_connectors.slack.models import SlackCredentials +from onyx.federated_connectors.slack.models import SlackEntities +from onyx.utils.logger import setup_logger + +logger = setup_logger() + + +SCOPES = [ + "search:read", + "channels:history", + "groups:history", + "im:history", + "mpim:history", +] + + +class SlackFederatedConnector(FederatedConnector): + def __init__(self, credentials: dict[str, Any]): + self.slack_credentials = SlackCredentials(**credentials) + + @override + def validate_entities(self, entities: dict[str, Any]) -> bool: + """Check the entities and verify that they match the expected structure/all values are valid. + + For Slack federated search, we expect: + - channels: list[str] (list of channel names or IDs) + - include_dm: bool (whether to include direct messages) + """ + try: + # Use Pydantic model for validation + SlackEntities(**entities) + return True + except ValidationError as e: + logger.warning(f"Validation error for Slack entities: {e}") + return False + except Exception as e: + logger.error(f"Error validating Slack entities: {e}") + return False + + @classmethod + @override + def entities_schema(cls) -> dict[str, EntityField]: + """Return the specifications of what entities are available for this federated search type. + + Returns a specification that tells the caller: + - channels is valid and should be a list[str] + - include_dm is valid and should be a boolean + """ + return { + "channels": EntityField( + type="list[str]", + description="List of Slack channel names or IDs to search in", + required=False, + example=["general", "random", "C1234567890"], + ), + "include_dm": EntityField( + type="bool", + description="Whether to include direct messages in the search", + required=False, + default=False, + example=True, + ), + } + + @classmethod + @override + def credentials_schema(cls) -> dict[str, CredentialField]: + """Return the specification of what credentials are required for Slack connector.""" + return { + "client_id": CredentialField( + type="str", + description="Slack app client ID from your Slack app configuration", + required=True, + example="1234567890.1234567890123", + secret=False, + ), + "client_secret": CredentialField( + type="str", + description="Slack app client secret from your Slack app configuration", + required=True, + example="1a2b3c4d5e6f7g8h9i0j1k2l3m4n5o6p", + secret=True, + ), + } + + @override + def authorize(self, redirect_uri: str) -> str: + """Get back the OAuth URL for Slack authorization. + + Returns the URL where users should be redirected to authorize the application. + Note: State parameter will be added by the API layer. + """ + # Build OAuth URL with proper parameters (no state - handled by API layer) + params = { + "client_id": self.slack_credentials.client_id, + "user_scope": " ".join(SCOPES), + "redirect_uri": redirect_uri, + } + + # Build query string + oauth_url = f"https://slack.com/oauth/v2/authorize?{urlencode(params)}" + + logger.info("Generated Slack OAuth authorization URL") + return oauth_url + + @override + def callback(self, callback_data: dict[str, Any], redirect_uri: str) -> OAuthResult: + """Handle the response from the OAuth flow and return it in a standard format. + + Args: + callback_data: The data received from the OAuth callback (state already validated by API layer) + + Returns: + Standardized OAuthResult + """ + # Extract authorization code from callback + auth_code = callback_data.get("code") + error = callback_data.get("error") + + if error: + raise RuntimeError(f"OAuth error received: {error}") + + if not auth_code: + raise ValueError("No authorization code received") + + # Exchange authorization code for access token + token_response = self._exchange_code_for_token(auth_code, redirect_uri) + + if not token_response.get("ok"): + raise RuntimeError( + f"Failed to exchange authorization code for token: {token_response.get('error')}" + ) + + # Build team info + team_info = None + if "team" in token_response: + team_info = { + "id": token_response["team"]["id"], + "name": token_response["team"]["name"], + } + + # Build user info and extract OAuth tokens + if "authed_user" not in token_response: + raise RuntimeError("Missing authed_user in OAuth response from Slack") + + authed_user = token_response["authed_user"] + user_info = { + "id": authed_user["id"], + "scope": authed_user.get("scope"), + "token_type": authed_user.get("token_type"), + } + + # Extract OAuth tokens from authed_user + access_token = authed_user.get("access_token") + refresh_token = authed_user.get("refresh_token") + token_type = authed_user.get("token_type", "bearer") + scope = authed_user.get("scope") + + # Calculate expires_at from expires_in if present + expires_at = None + if "expires_in" in authed_user: + expires_at = datetime.now(timezone.utc) + timedelta( + seconds=authed_user["expires_in"] + ) + + return OAuthResult( + access_token=access_token, + token_type=token_type, + scope=scope, + expires_at=expires_at, + refresh_token=refresh_token, + team=team_info, + user=user_info, + raw_response=token_response, + ) + + def _exchange_code_for_token(self, code: str, redirect_uri: str) -> dict[str, Any]: + """Exchange authorization code for access token. + + Args: + code: Authorization code from OAuth callback + + Returns: + Token response from Slack API + """ + response = requests.post( + "https://slack.com/api/oauth.v2.access", + data={ + "client_id": self.slack_credentials.client_id, + "client_secret": self.slack_credentials.client_secret, + "code": code, + "redirect_uri": redirect_uri, + }, + ) + response.raise_for_status() + return response.json() + + @override + def search( + self, + query: SearchQuery, + entities: dict[str, Any], + access_token: str, + limit: int | None = None, + ) -> list[InferenceChunk]: + """Perform a federated search on Slack. + + Args: + query: The search query + entities: The entities to search within (validated by validate()) + access_token: The OAuth access token + limit: Maximum number of results to return + + Returns: + Search results in SlackSearchResponse format + """ + with get_session_with_current_tenant() as db_session: + return slack_retrieval(query, access_token, db_session, limit) diff --git a/backend/onyx/federated_connectors/slack/models.py b/backend/onyx/federated_connectors/slack/models.py new file mode 100644 index 0000000000..7f963c9d00 --- /dev/null +++ b/backend/onyx/federated_connectors/slack/models.py @@ -0,0 +1,105 @@ +from typing import Optional + +from pydantic import BaseModel +from pydantic import Field +from pydantic import field_validator + + +class SlackEntities(BaseModel): + """Pydantic model for Slack federated search entities.""" + + channels: Optional[list[str]] = Field( + default=None, description="List of Slack channel names or IDs to search in" + ) + include_dm: Optional[bool] = Field( + default=False, description="Whether to include direct messages in the search" + ) + + @field_validator("channels") + @classmethod + def validate_channels(cls, v: Optional[list[str]]) -> Optional[list[str]]: + if v is not None: + if not isinstance(v, list): + raise ValueError("channels must be a list") + for channel in v: + if not isinstance(channel, str) or not channel.strip(): + raise ValueError("Each channel must be a non-empty string") + return v + + +class SlackCredentials(BaseModel): + """Slack federated connector credentials.""" + + client_id: str = Field(..., description="Slack app client ID") + client_secret: str = Field(..., description="Slack app client secret") + + @field_validator("client_id") + @classmethod + def validate_client_id(cls, v: str) -> str: + if not v or not v.strip(): + raise ValueError("Client ID cannot be empty") + return v.strip() + + @field_validator("client_secret") + @classmethod + def validate_client_secret(cls, v: str) -> str: + if not v or not v.strip(): + raise ValueError("Client secret cannot be empty") + return v.strip() + + +class SlackTeamInfo(BaseModel): + """Information about a Slack team/workspace.""" + + id: str = Field(..., description="Team ID") + name: str = Field(..., description="Team name") + domain: Optional[str] = Field(default=None, description="Team domain") + + +class SlackUserInfo(BaseModel): + """Information about a Slack user.""" + + id: str = Field(..., description="User ID") + team_id: Optional[str] = Field(default=None, description="Team ID") + name: Optional[str] = Field(default=None, description="User name") + email: Optional[str] = Field(default=None, description="User email") + + +class SlackSearchResult(BaseModel): + """Individual search result from Slack.""" + + channel: str = Field(..., description="Channel where the message was found") + timestamp: str = Field(..., description="Message timestamp") + user: Optional[str] = Field(default=None, description="User who sent the message") + text: str = Field(..., description="Message text") + permalink: Optional[str] = Field( + default=None, description="Permalink to the message" + ) + score: Optional[float] = Field(default=None, description="Search relevance score") + + # Additional context + thread_ts: Optional[str] = Field( + default=None, description="Thread timestamp if in a thread" + ) + reply_count: Optional[int] = Field( + default=None, description="Number of replies if it's a thread" + ) + + +class SlackSearchResponse(BaseModel): + """Response from Slack federated search.""" + + query: str = Field(..., description="The search query") + total_count: int = Field(..., description="Total number of results") + results: list[SlackSearchResult] = Field(..., description="Search results") + next_cursor: Optional[str] = Field( + default=None, description="Cursor for pagination" + ) + + # Metadata + channels_searched: Optional[list[str]] = Field( + default=None, description="Channels that were searched" + ) + search_time_ms: Optional[int] = Field( + default=None, description="Time taken to search in milliseconds" + ) diff --git a/backend/onyx/main.py b/backend/onyx/main.py index 85edb8a93f..94297001e8 100644 --- a/backend/onyx/main.py +++ b/backend/onyx/main.py @@ -73,6 +73,7 @@ from onyx.server.features.persona.api import admin_router as admin_persona_route from onyx.server.features.persona.api import basic_router as persona_router from onyx.server.features.tool.api import admin_router as admin_tool_router from onyx.server.features.tool.api import router as tool_router +from onyx.server.federated.api import router as federated_router from onyx.server.gpts.api import router as gpts_router from onyx.server.kg.api import admin_router as kg_admin_router from onyx.server.long_term_logs.long_term_logs_api import ( @@ -370,6 +371,7 @@ def get_application(lifespan_override: Lifespan | None = None) -> FastAPI: include_router_with_global_prefix_prepended(application, long_term_logs_router) include_router_with_global_prefix_prepended(application, api_key_router) include_router_with_global_prefix_prepended(application, standard_oauth_router) + include_router_with_global_prefix_prepended(application, federated_router) if AUTH_TYPE == AuthType.DISABLED: # Server logs this during auth setup verification step diff --git a/backend/onyx/prompts/chat_prompts.py b/backend/onyx/prompts/chat_prompts.py index 5b0eb7ac5a..eb31f1c51d 100644 --- a/backend/onyx/prompts/chat_prompts.py +++ b/backend/onyx/prompts/chat_prompts.py @@ -253,7 +253,7 @@ Please rephrase the following user question/query as a semantic query that would search engine. Note: - - do not change the meaning of the question! Specifically, if the query is a an instruction, keep it \ + - do not change the meaning of the question! Specifically, if the query is an instruction, keep it \ as an instruction! Here is the user question/query: diff --git a/backend/onyx/prompts/federated_search.py b/backend/onyx/prompts/federated_search.py new file mode 100644 index 0000000000..e8effbe8b9 --- /dev/null +++ b/backend/onyx/prompts/federated_search.py @@ -0,0 +1,27 @@ +from onyx.configs.app_configs import MAX_SLACK_QUERY_EXPANSIONS + +# TODO: maybe ask it to generate in and from filters, and fuzzy match them later +SLACK_QUERY_EXPANSION_PROMPT = f""" +Rewrite the user's query and, if helpful, split it into at most {MAX_SLACK_QUERY_EXPANSIONS} \ +keyword-only queries, so that Slack's keyword search yields the best matches. + +Keep in mind the Slack's search behavior: +- Pure keyword AND search (no semantics). +- Word order matters. +- More words = fewer matches, so keep each query concise. + +Guidelines: +1. Remove stop-words and obvious noise. +2. Remove or down-weight meta-instructions (e.g., "show me", "summary of", "how do I") that are \ +unlikely to appear in the target messages. +3. If the query has many keywords, produce several focused queries that keep related words together; \ +never explode into single-word queries. +4. Preserve phrases that belong together (e.g., "performance issues"); a word may appear in multiple queries. +5. When unsure, produce both a broad and a narrow query. +6. If the user asks for X or Y, create separate queries for X and Y. + +Here is the original query: +{{query}} + +Return EXACTLY the new query(ies), one per line, at most {MAX_SLACK_QUERY_EXPANSIONS}. Nothing else. +""" diff --git a/backend/onyx/server/features/document_set/models.py b/backend/onyx/server/features/document_set/models.py index e45ff5314d..7f75b5cdae 100644 --- a/backend/onyx/server/features/document_set/models.py +++ b/backend/onyx/server/features/document_set/models.py @@ -1,3 +1,4 @@ +from typing import Any from uuid import UUID from pydantic import BaseModel @@ -9,6 +10,22 @@ from onyx.server.documents.models import ConnectorSnapshot from onyx.server.documents.models import CredentialSnapshot +class FederatedConnectorConfig(BaseModel): + """Configuration for adding a federated connector to a document set""" + + federated_connector_id: int + entities: dict[str, Any] + + +class FederatedConnectorDescriptor(BaseModel): + """Descriptor for a federated connector in a document set""" + + id: int + name: str + source: str + entities: dict[str, Any] + + class DocumentSetCreationRequest(BaseModel): name: str description: str @@ -17,6 +34,8 @@ class DocumentSetCreationRequest(BaseModel): # For Private Document Sets, who should be able to access these users: list[UUID] = Field(default_factory=list) groups: list[int] = Field(default_factory=list) + # Federated connectors to include in the document set + federated_connectors: list[FederatedConnectorConfig] = Field(default_factory=list) class DocumentSetUpdateRequest(BaseModel): @@ -27,6 +46,8 @@ class DocumentSetUpdateRequest(BaseModel): # For Private Document Sets, who should be able to access these users: list[UUID] groups: list[int] + # Federated connectors to include in the document set + federated_connectors: list[FederatedConnectorConfig] = Field(default_factory=list) class CheckDocSetPublicRequest(BaseModel): @@ -51,6 +72,10 @@ class DocumentSet(BaseModel): # For Private Document Sets, who should be able to access these users: list[UUID] groups: list[int] + # Federated connectors in the document set + federated_connectors: list[FederatedConnectorDescriptor] = Field( + default_factory=list + ) @classmethod def from_model(cls, document_set_model: DocumentSetDBModel) -> "DocumentSet": @@ -76,4 +101,21 @@ class DocumentSet(BaseModel): is_public=document_set_model.is_public, users=[user.id for user in document_set_model.users], groups=[group.id for group in document_set_model.groups], + federated_connectors=[ + FederatedConnectorDescriptor( + id=fc_mapping.federated_connector_id, + name=( + f"{fc_mapping.federated_connector.source.replace('_', ' ').title()}" + if fc_mapping.federated_connector + else "Unknown" + ), + source=( + fc_mapping.federated_connector.source + if fc_mapping.federated_connector + else "unknown" + ), + entities=fc_mapping.entities, + ) + for fc_mapping in document_set_model.federated_connectors + ], ) diff --git a/backend/onyx/server/federated/api.py b/backend/onyx/server/federated/api.py new file mode 100644 index 0000000000..d3c9768118 --- /dev/null +++ b/backend/onyx/server/federated/api.py @@ -0,0 +1,582 @@ +import json +from typing import Any +from uuid import UUID + +from fastapi import APIRouter +from fastapi import Depends +from fastapi import HTTPException +from fastapi import Request +from fastapi import Response +from sqlalchemy.orm import Session + +from onyx.auth.users import current_curator_or_admin_user +from onyx.auth.users import current_user +from onyx.configs.constants import FederatedConnectorSource +from onyx.db.engine.sql_engine import get_session +from onyx.db.federated import ( + create_federated_connector as db_create_federated_connector, +) +from onyx.db.federated import delete_federated_connector +from onyx.db.federated import fetch_all_federated_connectors +from onyx.db.federated import fetch_federated_connector_by_id +from onyx.db.federated import update_federated_connector +from onyx.db.federated import update_federated_connector_oauth_token +from onyx.db.federated import validate_federated_connector_credentials +from onyx.db.models import User +from onyx.federated_connectors.factory import get_federated_connector +from onyx.federated_connectors.factory import get_federated_connector_cls +from onyx.federated_connectors.interfaces import FederatedConnector +from onyx.federated_connectors.oauth_utils import add_state_to_oauth_url +from onyx.federated_connectors.oauth_utils import generate_oauth_state +from onyx.federated_connectors.oauth_utils import get_oauth_callback_uri +from onyx.federated_connectors.oauth_utils import verify_oauth_state +from onyx.server.federated.models import AuthorizeUrlResponse +from onyx.server.federated.models import CredentialSchemaResponse +from onyx.server.federated.models import EntitySpecResponse +from onyx.server.federated.models import FederatedConnectorCredentials +from onyx.server.federated.models import FederatedConnectorDetail +from onyx.server.federated.models import FederatedConnectorRequest +from onyx.server.federated.models import FederatedConnectorResponse +from onyx.server.federated.models import FederatedConnectorStatus +from onyx.server.federated.models import FederatedConnectorUpdateRequest +from onyx.server.federated.models import OAuthCallbackResult +from onyx.server.federated.models import UserOAuthStatus +from onyx.utils.logger import setup_logger +from shared_configs.contextvars import get_current_tenant_id + +logger = setup_logger() + +router = APIRouter(prefix="/federated") + + +def _get_federated_connector_instance( + source: FederatedConnectorSource, + credentials: dict[str, Any], +) -> FederatedConnector: + """Factory function to get the appropriate federated connector instance.""" + try: + return get_federated_connector(source, credentials) + except ValueError as e: + raise HTTPException(status_code=400, detail=str(e)) + + +@router.post("") +def create_federated_connector( + federated_connector_data: FederatedConnectorRequest, + user: User | None = Depends(current_curator_or_admin_user), + db_session: Session = Depends(get_session), +) -> FederatedConnectorResponse: + """Create a new federated connector""" + tenant_id = get_current_tenant_id() + + if user is None: + raise HTTPException( + status_code=400, + detail="Federated retrieval not supported with AUTH_TYPE=disabled.", + ) + + logger.info( + f"Creating federated connector: source={federated_connector_data.source}, " + f"user={user.email}, tenant_id={tenant_id}" + ) + + try: + # Create the federated connector with validation + federated_connector = db_create_federated_connector( + db_session=db_session, + source=federated_connector_data.source, + credentials=federated_connector_data.credentials.model_dump(), + ) + + logger.info( + f"Successfully created federated connector with id={federated_connector.id}" + ) + + return FederatedConnectorResponse( + id=federated_connector.id, + source=federated_connector.source, + ) + + except ValueError as e: + logger.warning(f"Validation error creating federated connector: {e}") + db_session.rollback() + raise HTTPException(status_code=400, detail=str(e)) + except Exception as e: + logger.error(f"Error creating federated connector: {e}") + db_session.rollback() + raise HTTPException(status_code=500, detail=str(e)) + + +@router.get("/{id}/entities") +def get_entities( + id: int, + _: User = Depends(current_curator_or_admin_user), + db_session: Session = Depends(get_session), +) -> EntitySpecResponse: + """Fetch allowed entities for the source type""" + try: + federated_connector = fetch_federated_connector_by_id(id, db_session) + if not federated_connector: + raise HTTPException(status_code=404, detail="Federated connector not found") + + connector_instance = _get_federated_connector_instance( + federated_connector.source, federated_connector.credentials + ) + entities_spec = connector_instance.entities_schema() + + # Convert EntityField objects to a dictionary format for the API response + entities_dict = {} + for key, field in entities_spec.items(): + entities_dict[key] = { + "type": field.type, + "description": field.description, + "required": field.required, + "default": field.default, + "example": field.example, + } + + return EntitySpecResponse(entities=entities_dict) + + except HTTPException: + raise + except Exception as e: + logger.error(f"Error fetching entities for federated connector {id}: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@router.get("/{id}/credentials/schema") +def get_credentials_schema( + id: int, + _: User = Depends(current_curator_or_admin_user), + db_session: Session = Depends(get_session), +) -> CredentialSchemaResponse: + """Fetch credential schema for the source type""" + try: + federated_connector = fetch_federated_connector_by_id(id, db_session) + if not federated_connector: + raise HTTPException(status_code=404, detail="Federated connector not found") + + connector_instance = _get_federated_connector_instance( + federated_connector.source, federated_connector.credentials + ) + credentials_spec = connector_instance.credentials_schema() + + # Convert CredentialField objects to a dictionary format for the API response + credentials_dict = {} + for key, field in credentials_spec.items(): + credentials_dict[key] = { + "type": field.type, + "description": field.description, + "required": field.required, + "default": field.default, + "example": field.example, + "secret": field.secret, + } + + return CredentialSchemaResponse(credentials=credentials_dict) + + except HTTPException: + raise + except Exception as e: + logger.error( + f"Error fetching credentials schema for federated connector {id}: {e}" + ) + raise HTTPException(status_code=500, detail=str(e)) + + +@router.get("/sources/{source}/credentials/schema") +def get_credentials_schema_by_source( + source: FederatedConnectorSource, + _: User = Depends(current_curator_or_admin_user), +) -> CredentialSchemaResponse: + """Fetch credential schema for a specific source type (for setup forms)""" + try: + connector_cls = get_federated_connector_cls(source) + credentials_spec = connector_cls.credentials_schema() + + # Convert CredentialField objects to a dictionary format for the API response + credentials_dict = {} + for key, field in credentials_spec.items(): + credentials_dict[key] = { + "type": field.type, + "description": field.description, + "required": field.required, + "default": field.default, + "example": field.example, + "secret": field.secret, + } + + return CredentialSchemaResponse(credentials=credentials_dict) + + except HTTPException: + raise + except Exception as e: + logger.error(f"Error fetching credentials schema for source {source}: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@router.post("/sources/{source}/credentials/validate") +def validate_credentials( + source: FederatedConnectorSource, + credentials: FederatedConnectorCredentials, + _: User = Depends(current_curator_or_admin_user), +) -> bool: + """Validate credentials for a specific source type""" + try: + is_valid = validate_federated_connector_credentials( + source, credentials.model_dump() + ) + + if not is_valid: + raise HTTPException(status_code=400, detail="Credentials are invalid") + + return is_valid + + except HTTPException: + raise + except Exception as e: + logger.error(f"Error validating credentials for source {source}: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@router.head("/{id}/entities/validate") +def validate_entities( + id: int, + request: Request, + _: User = Depends(current_curator_or_admin_user), + db_session: Session = Depends(get_session), +) -> Response: + """Validate specified entities for source type""" + try: + federated_connector = fetch_federated_connector_by_id(id, db_session) + if not federated_connector: + raise HTTPException(status_code=404, detail="Federated connector not found") + + # For HEAD requests, we'll expect entities as query parameters + # since HEAD requests shouldn't have request bodies + entities_dict = {} + query_params = dict(request.query_params) + if "entities" in query_params: + try: + entities_dict = json.loads(query_params["entities"]) + except json.JSONDecodeError: + logger.warning("Could not parse entities from query parameters") + return Response(status_code=400) + + connector_instance = _get_federated_connector_instance( + federated_connector.source, federated_connector.credentials + ) + is_valid = connector_instance.validate_entities(entities_dict) + + if is_valid: + return Response(status_code=200) + else: + return Response(status_code=400) + + except HTTPException: + raise + except Exception as e: + logger.error(f"Error validating entities for federated connector {id}: {e}") + return Response(status_code=500) + + +@router.get("/{id}/authorize") +def get_authorize_url( + id: int, + user: User = Depends(current_user), + db_session: Session = Depends(get_session), +) -> AuthorizeUrlResponse: + """Get URL to send the user for OAuth""" + # Validate that the ID is not None or invalid + if id is None or id <= 0: + raise HTTPException(status_code=400, detail="Invalid federated connector ID") + + federated_connector = fetch_federated_connector_by_id(id, db_session) + if not federated_connector: + raise HTTPException(status_code=404, detail="Federated connector not found") + + # Update credentials to include the correct redirect URI with the connector ID + updated_credentials = federated_connector.credentials.copy() + if "redirect_uri" in updated_credentials and updated_credentials["redirect_uri"]: + # Replace the {id} placeholder with the actual federated connector ID + updated_credentials["redirect_uri"] = updated_credentials[ + "redirect_uri" + ].replace("{id}", str(id)) + + connector_instance = _get_federated_connector_instance( + federated_connector.source, updated_credentials + ) + base_authorize_url = connector_instance.authorize(get_oauth_callback_uri()) + + # Generate state parameter and store session info + logger.info( + f"Generating OAuth state for federated_connector_id={id}, user_id={user.id}" + ) + state = generate_oauth_state( + federated_connector_id=id, + user_id=str(user.id), + ) + + # Add state to the OAuth URL + authorize_url = add_state_to_oauth_url(base_authorize_url, state) + logger.info(f"Generated OAuth authorize URL with state for connector {id}") + return AuthorizeUrlResponse(authorize_url=authorize_url) + + +@router.post("/callback") +def handle_oauth_callback_generic( + request: Request, + _: User = Depends(current_user), + db_session: Session = Depends(get_session), +) -> OAuthCallbackResult: + """Handle callback for any federated connector using state parameter""" + # Get callback data from request (query parameters) + callback_data = dict(request.query_params) + + # Verify state parameter and get session info + state = callback_data.get("state") + if not state: + raise HTTPException(status_code=400, detail="Missing state parameter") + + try: + oauth_session = verify_oauth_state(state) + except ValueError: + logger.exception("Error verifying OAuth state") + raise HTTPException( + status_code=400, detail="Invalid or expired state parameter" + ) + + if not oauth_session: + raise HTTPException( + status_code=400, detail="Invalid or expired state parameter" + ) + + # Get federated connector ID from the state + federated_connector_id = oauth_session.federated_connector_id + + # Validate federated_connector_id is not None + if federated_connector_id is None: + logger.error("OAuth session has null federated_connector_id") + raise HTTPException( + status_code=400, + detail="Invalid OAuth session: missing federated connector ID", + ) + + federated_connector = fetch_federated_connector_by_id( + federated_connector_id, db_session + ) + if not federated_connector: + raise HTTPException(status_code=404, detail="Federated connector not found") + + connector_instance = _get_federated_connector_instance( + federated_connector.source, federated_connector.credentials + ) + oauth_result = connector_instance.callback(callback_data, get_oauth_callback_uri()) + + # Convert OAuthResult to OAuthCallbackResult for API response + oauth_result_dict = oauth_result.model_dump() + oauth_callback_result = OAuthCallbackResult(**oauth_result_dict) + + # Add source information to the response + oauth_callback_result.source = federated_connector.source + + # Store OAuth token in database if we have an access token + if oauth_result.access_token: + logger.info( + f"Storing OAuth token for federated_connector_id={federated_connector_id}, " + f"user_id={oauth_session.user_id}" + ) + update_federated_connector_oauth_token( + db_session=db_session, + federated_connector_id=federated_connector_id, + user_id=UUID(oauth_session.user_id), + token=oauth_result.access_token, + expires_at=oauth_result.expires_at, + ) + + return oauth_callback_result + + +@router.get("") +def get_federated_connectors( + _: User = Depends(current_user), + db_session: Session = Depends(get_session), +) -> list[FederatedConnectorStatus]: + """Get all federated connectors for display in the status table""" + federated_connectors = fetch_all_federated_connectors(db_session) + + result = [] + for fc in federated_connectors: + status_data = FederatedConnectorStatus( + id=fc.id, + source=fc.source, + name=f"{fc.source.replace('_', ' ').title()}", + ) + result.append(status_data) + + return result + + +@router.get("/oauth-status") +def get_user_oauth_status( + user: User = Depends(current_user), + db_session: Session = Depends(get_session), +) -> list[UserOAuthStatus]: + """Get OAuth status for all federated connectors for the current user""" + federated_connectors = fetch_all_federated_connectors(db_session) + + result = [] + for fc in federated_connectors: + # Check if user has OAuth token for this connector + oauth_token = None + for token in fc.oauth_tokens: + if token.user_id == user.id: + oauth_token = token + break + + # Generate authorize URL if needed + authorize_url = None + if not oauth_token: + connector_instance = _get_federated_connector_instance( + fc.source, fc.credentials + ) + base_authorize_url = connector_instance.authorize(get_oauth_callback_uri()) + + # Generate state parameter and add to URL + state = generate_oauth_state( + federated_connector_id=fc.id, + user_id=str(user.id), + ) + authorize_url = add_state_to_oauth_url(base_authorize_url, state) + + status_data = UserOAuthStatus( + federated_connector_id=fc.id, + source=fc.source, + name=f"{fc.source.replace('_', ' ').title()}", + has_oauth_token=oauth_token is not None, + oauth_token_expires_at=oauth_token.expires_at if oauth_token else None, + authorize_url=authorize_url, + ) + result.append(status_data) + + return result + + +@router.get("/{id}") +def get_federated_connector_detail( + id: int, + user: User = Depends(current_curator_or_admin_user), + db_session: Session = Depends(get_session), +) -> FederatedConnectorDetail: + """Get detailed information about a specific federated connector""" + federated_connector = fetch_federated_connector_by_id(id, db_session) + if not federated_connector: + raise HTTPException(status_code=404, detail="Federated connector not found") + + # Get OAuth token information for the current user + oauth_token = None + for token in federated_connector.oauth_tokens: + if token.user_id == user.id: + oauth_token = token + break + + # Get document set mappings + document_sets = [] + for mapping in federated_connector.document_sets: + document_sets.append( + { + "id": mapping.document_set_id, + "name": ( + mapping.document_set.name if mapping.document_set else "Unknown" + ), + "entities": mapping.entities, + } + ) + + return FederatedConnectorDetail( + id=federated_connector.id, + source=federated_connector.source, + name=f"{federated_connector.source.replace('_', ' ').title()}", + credentials=FederatedConnectorCredentials(**federated_connector.credentials), + oauth_token_exists=oauth_token is not None, + oauth_token_expires_at=oauth_token.expires_at if oauth_token else None, + document_sets=document_sets, + ) + + +@router.put("/{id}") +def update_federated_connector_endpoint( + id: int, + update_request: FederatedConnectorUpdateRequest, + user: User = Depends(current_curator_or_admin_user), + db_session: Session = Depends(get_session), +) -> FederatedConnectorDetail: + """Update a federated connector's configuration""" + try: + # Update the federated connector + updated_connector = update_federated_connector( + db_session=db_session, + federated_connector_id=id, + credentials=( + update_request.credentials.model_dump() + if update_request.credentials + else None + ), + ) + + if not updated_connector: + raise HTTPException(status_code=404, detail="Federated connector not found") + + # Return updated connector details + return get_federated_connector_detail(id, user, db_session) + + except ValueError as e: + logger.warning(f"Validation error updating federated connector {id}: {e}") + raise HTTPException(status_code=400, detail=str(e)) + + +@router.delete("/{id}") +def delete_federated_connector_endpoint( + id: int, + _: User = Depends(current_curator_or_admin_user), + db_session: Session = Depends(get_session), +) -> bool: + """Delete a federated connector""" + success = delete_federated_connector( + db_session=db_session, + federated_connector_id=id, + ) + + if not success: + raise HTTPException(status_code=404, detail="Federated connector not found") + + return True + + +@router.delete("/{id}/oauth") +def disconnect_oauth_token( + id: int, + user: User = Depends(current_user), + db_session: Session = Depends(get_session), +) -> bool: + """Disconnect OAuth token for the current user from a federated connector""" + # Check if the federated connector exists + federated_connector = fetch_federated_connector_by_id(id, db_session) + if not federated_connector: + raise HTTPException(status_code=404, detail="Federated connector not found") + + # Find and delete the user's OAuth token + oauth_token = None + for token in federated_connector.oauth_tokens: + if token.user_id == user.id: + oauth_token = token + break + + if oauth_token: + db_session.delete(oauth_token) + db_session.commit() + return True + else: + raise HTTPException( + status_code=404, detail="No OAuth token found for this user" + ) diff --git a/backend/onyx/server/federated/models.py b/backend/onyx/server/federated/models.py new file mode 100644 index 0000000000..3813872fae --- /dev/null +++ b/backend/onyx/server/federated/models.py @@ -0,0 +1,81 @@ +from datetime import datetime +from typing import Any + +from pydantic import BaseModel +from pydantic import Field + +from onyx.configs.constants import FederatedConnectorSource + + +class FederatedConnectorCredentials(BaseModel): + """Credentials for federated connector""" + + client_id: str | None = None + client_secret: str | None = None + redirect_uri: str | None = None + + +class FederatedConnectorRequest(BaseModel): + source: FederatedConnectorSource + credentials: FederatedConnectorCredentials + + +class FederatedConnectorResponse(BaseModel): + id: int + source: FederatedConnectorSource + + +class AuthorizeUrlResponse(BaseModel): + authorize_url: str + + +class OAuthCallbackResult(BaseModel): + access_token: str | None = None + expires_at: datetime | None = None + refresh_token: str | None = None + token_type: str | None = None + scope: str | None = None + source: FederatedConnectorSource | None = None + + +class FederatedConnectorStatus(BaseModel): + id: int + source: FederatedConnectorSource + name: str + + +class UserOAuthStatus(BaseModel): + """OAuth status for a specific user and federated connector""" + + federated_connector_id: int + source: FederatedConnectorSource + name: str + has_oauth_token: bool + oauth_token_expires_at: datetime | None = None + authorize_url: str | None = None + + +class FederatedConnectorDetail(BaseModel): + id: int + source: FederatedConnectorSource + name: str + credentials: FederatedConnectorCredentials + oauth_token_exists: bool + oauth_token_expires_at: datetime | None = None + document_sets: list[dict[str, Any]] = Field(default_factory=list) + + +class FederatedConnectorUpdateRequest(BaseModel): + credentials: FederatedConnectorCredentials | None = None + + +class EntitySpecResponse(BaseModel): + """Response for entity specification""" + + entities: dict[str, Any] + + +class CredentialSchemaResponse(BaseModel): + """Response for credential schema specification""" + + credentials: dict[str, Any] diff --git a/backend/tests/unit/federated_connector/slack/test_slack_federated_connnector.py b/backend/tests/unit/federated_connector/slack/test_slack_federated_connnector.py new file mode 100644 index 0000000000..3eba64842c --- /dev/null +++ b/backend/tests/unit/federated_connector/slack/test_slack_federated_connnector.py @@ -0,0 +1,217 @@ +from datetime import datetime +from datetime import timedelta +from datetime import timezone +from typing import Any +from unittest.mock import patch + +import pytest + +from onyx.federated_connectors.models import OAuthResult +from onyx.federated_connectors.slack.federated_connector import SlackFederatedConnector + +# Constants for mock Slack OAuth response +MOCK_APP_ID = "A093M5L7Q92" +MOCK_USER_ID = "U05SAH6UGUD" +MOCK_SCOPE = "search:read" +MOCK_ACCESS_TOKEN = ( + "xoxe.xoxp-1-Mi0yLTU5MTAx...MDkwN2U0YjlmZmI4YzA1NTYwZjNlMjRiZDYwNGU0ZA" +) +MOCK_REFRESH_TOKEN = ( + "xoxe-1-My0xLTU5MTAxMz...jcyZjA3NDM3YjdhOTRhYmRhMGJmMGVlMzBjNzQ4Y2I" +) +MOCK_TOKEN_TYPE = "user" +MOCK_EXPIRES_IN = 31659 +MOCK_TEAM_ID = "T05SS40AFAM" +MOCK_TEAM_NAME = "Onyx Team" + + +class TestSlackFederatedConnector: + """Test suite for SlackFederatedConnector""" + + @pytest.fixture + def test_credentials(self) -> dict[str, str]: + """Test credentials for Slack connector""" + return { + "client_id": "test_client_id", + "client_secret": "test_client_secret", + "redirect_uri": "https://test.com/callback", + } + + @pytest.fixture + def slack_connector( + self, test_credentials: dict[str, str] + ) -> SlackFederatedConnector: + """Create a SlackFederatedConnector instance for testing""" + return SlackFederatedConnector(test_credentials) + + @pytest.fixture + def mock_slack_oauth_response(self) -> dict[str, Any]: + """Mock Slack OAuth response based on real example""" + return { + "ok": True, + "app_id": MOCK_APP_ID, + "authed_user": { + "id": MOCK_USER_ID, + "scope": MOCK_SCOPE, + "access_token": MOCK_ACCESS_TOKEN, + "token_type": MOCK_TOKEN_TYPE, + "refresh_token": MOCK_REFRESH_TOKEN, + "expires_in": MOCK_EXPIRES_IN, + }, + "team": {"id": MOCK_TEAM_ID, "name": MOCK_TEAM_NAME}, + "enterprise": None, + "is_enterprise_install": False, + } + + def test_callback_success( + self, + slack_connector: SlackFederatedConnector, + mock_slack_oauth_response: dict[str, Any], + ) -> None: + """Test successful OAuth callback handling""" + # Mock the token exchange method + with patch.object( + slack_connector, + "_exchange_code_for_token", + return_value=mock_slack_oauth_response, + ): + # Simulate callback data with authorization code + callback_data = { + "code": "test_auth_code", + "state": "test_state", + } + redirect_uri = "https://test.com/callback" + + # Call the callback method + result = slack_connector.callback(callback_data, redirect_uri) + + # Assert the result is an OAuthResult + assert isinstance(result, OAuthResult) + + # Assert OAuth token values are correctly extracted + assert result.access_token == MOCK_ACCESS_TOKEN + assert result.refresh_token == MOCK_REFRESH_TOKEN + assert result.token_type == MOCK_TOKEN_TYPE + assert result.scope == MOCK_SCOPE + + # Assert expiration time is calculated correctly + assert result.expires_at is not None + expected_expires_at = datetime.now(timezone.utc) + timedelta( + seconds=MOCK_EXPIRES_IN + ) + # Allow for small time difference due to test execution time + time_diff = abs((result.expires_at - expected_expires_at).total_seconds()) + assert time_diff < 5 # Within 5 seconds + + # Assert team info is extracted correctly + assert result.team is not None + assert result.team["id"] == MOCK_TEAM_ID + assert result.team["name"] == MOCK_TEAM_NAME + + # Assert user info is extracted correctly + assert result.user is not None + assert result.user["id"] == MOCK_USER_ID + assert result.user["scope"] == MOCK_SCOPE + assert result.user["token_type"] == MOCK_TOKEN_TYPE + + # Assert raw response is preserved + assert result.raw_response == mock_slack_oauth_response + + def test_callback_oauth_error( + self, slack_connector: SlackFederatedConnector + ) -> None: + """Test OAuth callback with error response""" + callback_data = { + "error": "access_denied", + "error_description": "User denied access", + } + redirect_uri = "https://test.com/callback" + + with pytest.raises(RuntimeError, match="OAuth error received: access_denied"): + slack_connector.callback(callback_data, redirect_uri) + + def test_callback_missing_code( + self, slack_connector: SlackFederatedConnector + ) -> None: + """Test OAuth callback without authorization code""" + callback_data = {"state": "test_state"} + redirect_uri = "https://test.com/callback" + + with pytest.raises(ValueError, match="No authorization code received"): + slack_connector.callback(callback_data, redirect_uri) + + def test_callback_slack_api_error( + self, slack_connector: SlackFederatedConnector + ) -> None: + """Test OAuth callback when Slack API returns error""" + # Mock failed token exchange + mock_error_response = { + "ok": False, + "error": "invalid_code", + } + + with patch.object( + slack_connector, + "_exchange_code_for_token", + return_value=mock_error_response, + ): + callback_data = {"code": "invalid_code"} + redirect_uri = "https://test.com/callback" + + with pytest.raises( + RuntimeError, match="Failed to exchange authorization code for token" + ): + slack_connector.callback(callback_data, redirect_uri) + + def test_callback_without_authed_user( + self, slack_connector: SlackFederatedConnector + ) -> None: + """Test OAuth callback when authed_user is missing from response""" + # Mock response without authed_user + mock_response = { + "ok": True, + "app_id": MOCK_APP_ID, + "team": {"id": MOCK_TEAM_ID, "name": MOCK_TEAM_NAME}, + } + + with patch.object( + slack_connector, "_exchange_code_for_token", return_value=mock_response + ): + callback_data = {"code": "test_code"} + redirect_uri = "https://test.com/callback" + + with pytest.raises( + RuntimeError, match="Missing authed_user in OAuth response from Slack" + ): + slack_connector.callback(callback_data, redirect_uri) + + def test_callback_with_incomplete_authed_user( + self, slack_connector: SlackFederatedConnector + ) -> None: + """Test OAuth callback when authed_user is missing access_token""" + # Mock response with authed_user but missing access_token + mock_response = { + "ok": True, + "app_id": MOCK_APP_ID, + "authed_user": { + "id": MOCK_USER_ID, + "scope": MOCK_SCOPE, + "token_type": MOCK_TOKEN_TYPE, + # Missing access_token + }, + "team": {"id": MOCK_TEAM_ID, "name": MOCK_TEAM_NAME}, + } + + with patch.object( + slack_connector, "_exchange_code_for_token", return_value=mock_response + ): + callback_data = {"code": "test_code"} + redirect_uri = "https://test.com/callback" + + result = slack_connector.callback(callback_data, redirect_uri) + + # Should handle gracefully - access_token can be None in some edge cases + assert result.access_token is None + assert result.refresh_token is None + assert result.token_type == MOCK_TOKEN_TYPE + assert result.scope == MOCK_SCOPE diff --git a/web/src/app/admin/add-connector/page.tsx b/web/src/app/admin/add-connector/page.tsx index 7fc3d1d0c1..91dccea0ba 100644 --- a/web/src/app/admin/add-connector/page.tsx +++ b/web/src/app/admin/add-connector/page.tsx @@ -1,49 +1,159 @@ "use client"; import { SourceIcon } from "@/components/SourceIcon"; import { AdminPageTitle } from "@/components/admin/Title"; -import { ConnectorIcon } from "@/components/icons/icons"; +import { AlertIcon, ConnectorIcon, InfoIcon } from "@/components/icons/icons"; import { SourceCategory, SourceMetadata } from "@/lib/search/interfaces"; import { listSourceMetadata } from "@/lib/sources"; import Title from "@/components/ui/title"; import { Button } from "@/components/ui/button"; import Link from "next/link"; import { useCallback, useEffect, useMemo, useRef, useState } from "react"; +import { + Tooltip, + TooltipContent, + TooltipProvider, + TooltipTrigger, +} from "@/components/ui/tooltip"; +import { useFederatedConnectors } from "@/lib/hooks"; +import { + FederatedConnectorDetail, + federatedSourceToRegularSource, + ValidSources, +} from "@/lib/types"; +import useSWR from "swr"; +import { errorHandlingFetcher } from "@/lib/fetcher"; +import { buildSimilarCredentialInfoURL } from "@/app/admin/connector/[ccPairId]/lib"; +import { Credential } from "@/lib/connectors/credentials"; function SourceTile({ sourceMetadata, preSelect, + federatedConnectors, + slackCredentials, }: { sourceMetadata: SourceMetadata; preSelect?: boolean; + federatedConnectors?: FederatedConnectorDetail[]; + slackCredentials?: Credential[]; }) { + // Check if there's already a federated connector for this source + const existingFederatedConnector = useMemo(() => { + if (!sourceMetadata.federated || !federatedConnectors) { + return null; + } + + return federatedConnectors.find( + (connector) => + federatedSourceToRegularSource(connector.source) === + sourceMetadata.internalName + ); + }, [sourceMetadata, federatedConnectors]); + + // For Slack specifically, check if there are existing non-federated credentials + const isSlackTile = sourceMetadata.internalName === ValidSources.Slack; + const hasExistingSlackCredentials = useMemo(() => { + return isSlackTile && slackCredentials && slackCredentials.length > 0; + }, [isSlackTile, slackCredentials]); + + // Determine the URL to navigate to + const navigationUrl = useMemo(() => { + // Special logic for Slack: if there are existing credentials, use the old flow + if (isSlackTile && hasExistingSlackCredentials) { + return "/admin/connectors/slack"; + } + + // Otherwise, use the existing logic + if (existingFederatedConnector) { + return `/admin/federated/${existingFederatedConnector.id}`; + } + return sourceMetadata.adminUrl; + }, [ + isSlackTile, + hasExistingSlackCredentials, + existingFederatedConnector, + sourceMetadata.adminUrl, + ]); + return ( - - -

{sourceMetadata.displayName}

- + + + + + {sourceMetadata.federated && !hasExistingSlackCredentials && ( +
+ +
+ )} + +

+ {sourceMetadata.displayName} +

+ + + + {existingFederatedConnector && !hasExistingSlackCredentials ? ( +

+ Federated connector already configured. Click to + edit the existing connector. +

+ ) : hasExistingSlackCredentials ? ( +

+ Existing Slack credentials found. Click to manage + the traditional Slack connector. +

+ ) : sourceMetadata.federated ? ( +

+ {sourceMetadata.federatedTooltip ? ( + sourceMetadata.federatedTooltip + ) : ( + <> + Federated Search. This will result in greater + latency and lower search quality. + + )} +

+ ) : null} + + + ); } + export default function Page() { const sources = useMemo(() => listSourceMetadata(), []); const [searchTerm, setSearchTerm] = useState(""); + const { data: federatedConnectors } = useFederatedConnectors(); + + // Fetch Slack credentials to determine navigation behavior + const { data: slackCredentials } = useSWR[]>( + buildSimilarCredentialInfoURL(ValidSources.Slack), + errorHandlingFetcher + ); const searchInputRef = useRef(null); @@ -93,7 +203,20 @@ export default function Page() { ) { const firstSource = filteredCategories[0][1][0]; if (firstSource) { - window.open(firstSource.adminUrl, "_self"); + // Check if this source has an existing federated connector + const existingFederatedConnector = + firstSource.federated && federatedConnectors + ? federatedConnectors.find( + (connector) => + connector.source === `federated_${firstSource.internalName}` + ) + : null; + + const url = existingFederatedConnector + ? `/admin/federated/${existingFederatedConnector.id}` + : firstSource.adminUrl; + + window.open(url, "_self"); } } } @@ -137,6 +260,8 @@ export default function Page() { } key={source.internalName} sourceMetadata={source} + federatedConnectors={federatedConnectors} + slackCredentials={slackCredentials} /> ))} diff --git a/web/src/app/admin/connectors/[connector]/AddConnectorPage.tsx b/web/src/app/admin/connectors/[connector]/AddConnectorPage.tsx index ec6c2510cd..c50f5e3faf 100644 --- a/web/src/app/admin/connectors/[connector]/AddConnectorPage.tsx +++ b/web/src/app/admin/connectors/[connector]/AddConnectorPage.tsx @@ -511,6 +511,7 @@ export default function AddConnector({ includeDivider={false} icon={} title={displayName} + farRightElement={undefined} /> {formStep == 0 && ( diff --git a/web/src/app/admin/connectors/[connector]/ConnectorWrapper.tsx b/web/src/app/admin/connectors/[connector]/ConnectorWrapper.tsx index e48a587526..9858b714d8 100644 --- a/web/src/app/admin/connectors/[connector]/ConnectorWrapper.tsx +++ b/web/src/app/admin/connectors/[connector]/ConnectorWrapper.tsx @@ -6,22 +6,38 @@ import { FormProvider } from "@/components/context/FormContext"; import Sidebar from "./Sidebar"; import { HeaderTitle } from "@/components/header/HeaderTitle"; import { Button } from "@/components/ui/button"; -import { isValidSource } from "@/lib/sources"; +import { isValidSource, getSourceMetadata } from "@/lib/sources"; +import { FederatedConnectorForm } from "@/components/admin/federated/FederatedConnectorForm"; +import { useSearchParams } from "next/navigation"; +import useSWR from "swr"; +import { errorHandlingFetcher } from "@/lib/fetcher"; +import { buildSimilarCredentialInfoURL } from "@/app/admin/connector/[ccPairId]/lib"; +import { Credential } from "@/lib/connectors/credentials"; export default function ConnectorWrapper({ connector, }: { connector: ConfigurableSources; }) { - return ( - -
- -
- {!isValidSource(connector) ? ( + const searchParams = useSearchParams(); + const mode = searchParams?.get("mode"); // 'federated' or 'regular' + + // Fetch existing credentials for this connector type + const { data: existingCredentials } = useSWR[]>( + buildSimilarCredentialInfoURL(connector), + errorHandlingFetcher + ); + + // Check if the connector is valid + if (!isValidSource(connector)) { + return ( + +
+ +
-

‘{connector}‘ is not a valid Connector Type!

+

‘{connector}’ is not a valid Connector Type!
- ) : ( - - )} +
+
+ + ); + } + + const sourceMetadata = getSourceMetadata(connector); + const supportsFederated = sourceMetadata.federated === true; + const hasExistingCredentials = + existingCredentials && existingCredentials.length > 0; + + // Determine which form to show based on: + // 1. URL parameter mode (takes priority) + // 2. If no mode specified and existing credentials exist, show regular form + // 3. If no mode specified and no credentials, show federated form for federated-supported sources + let showFederatedForm = false; + + if (mode === "federated") { + showFederatedForm = supportsFederated; + } else if (mode === "regular") { + showFederatedForm = false; + } else { + // No mode specified - use default logic + if (hasExistingCredentials) { + // Default to regular form if existing credentials exist + showFederatedForm = false; + } else { + // Default to federated for federated-supported sources with no existing credentials + showFederatedForm = supportsFederated; + } + } + + // For federated form, use the specialized form without FormProvider + if (showFederatedForm) { + return ( +
+
+ +
+
+ ); + } + + // For regular connectors, use the existing flow + return ( + +
+ +
+
diff --git a/web/src/app/admin/documents/sets/DocumentSetCreationForm.tsx b/web/src/app/admin/documents/sets/DocumentSetCreationForm.tsx index cd76712189..3a016c2480 100644 --- a/web/src/app/admin/documents/sets/DocumentSetCreationForm.tsx +++ b/web/src/app/admin/documents/sets/DocumentSetCreationForm.tsx @@ -8,7 +8,13 @@ import { updateDocumentSet, DocumentSetCreationRequest, } from "./lib"; -import { ConnectorStatus, DocumentSet, UserGroup, UserRole } from "@/lib/types"; +import { + ConnectorStatus, + DocumentSet, + UserGroup, + UserRole, + FederatedConnectorConfig, +} from "@/lib/types"; import { TextFormField } from "@/components/Field"; import { Separator } from "@/components/ui/separator"; import { Button } from "@/components/ui/button"; @@ -18,6 +24,8 @@ import React, { useEffect, useState } from "react"; import { useUser } from "@/components/user/UserProvider"; import { ConnectorMultiSelect } from "@/components/ConnectorMultiSelect"; import { NonSelectableConnectors } from "@/components/NonSelectableConnectors"; +import { FederatedConnectorSelector } from "@/components/FederatedConnectorSelector"; +import { useFederatedConnectors } from "@/lib/hooks"; interface SetCreationPopupProps { ccPairs: ConnectorStatus[]; @@ -38,6 +46,7 @@ export const DocumentSetCreationForm = ({ const isUpdate = existingDocumentSet !== undefined; const [localCcPairs, setLocalCcPairs] = useState(ccPairs); const { user } = useUser(); + const { data: federatedConnectors } = useFederatedConnectors(); useEffect(() => { if (existingDocumentSet?.is_public) { @@ -58,14 +67,36 @@ export const DocumentSetCreationForm = ({ is_public: existingDocumentSet?.is_public ?? true, users: existingDocumentSet?.users ?? [], groups: existingDocumentSet?.groups ?? [], + federated_connectors: + existingDocumentSet?.federated_connectors?.map((fc) => ({ + federated_connector_id: fc.id, + entities: fc.entities, + })) ?? [], }} - validationSchema={Yup.object().shape({ - name: Yup.string().required("Please enter a name for the set"), - description: Yup.string().optional(), - cc_pair_ids: Yup.array() - .of(Yup.number().required()) - .required("Please select at least one connector"), - })} + validationSchema={Yup.object() + .shape({ + name: Yup.string().required("Please enter a name for the set"), + description: Yup.string().optional(), + cc_pair_ids: Yup.array().of(Yup.number().required()), + federated_connectors: Yup.array().of( + Yup.object().shape({ + federated_connector_id: Yup.number().required(), + entities: Yup.object().required(), + }) + ), + }) + .test( + "at-least-one-connector", + "Please select at least one connector (regular or federated)", + function (values) { + const hasRegularConnectors = + values.cc_pair_ids && values.cc_pair_ids.length > 0; + const hasFederatedConnectors = + values.federated_connectors && + values.federated_connectors.length > 0; + return hasRegularConnectors || hasFederatedConnectors; + } + )} onSubmit={async (values, formikHelpers) => { formikHelpers.setSubmitting(true); // If the document set is public, then we don't want to send any groups @@ -211,6 +242,26 @@ export const DocumentSetCreationForm = ({ placeholder="Search for connectors..." /> )} + + {/* Federated Connectors Section */} + {federatedConnectors && federatedConnectors.length > 0 && ( + <> + + { + props.setFieldValue( + "federated_connectors", + selectedConfigs + ); + }} + placeholder="Search for federated connectors..." + /> + + )}
diff --git a/web/src/app/admin/documents/sets/lib.ts b/web/src/app/admin/documents/sets/lib.ts index 807016ee00..7a09392c16 100644 --- a/web/src/app/admin/documents/sets/lib.ts +++ b/web/src/app/admin/documents/sets/lib.ts @@ -1,3 +1,5 @@ +import { FederatedConnectorConfig } from "@/lib/types"; + export interface DocumentSetCreationRequest { name: string; description: string; @@ -5,6 +7,7 @@ export interface DocumentSetCreationRequest { is_public: boolean; users: string[]; groups: number[]; + federated_connectors: FederatedConnectorConfig[]; } export const createDocumentSet = async ({ @@ -14,6 +17,7 @@ export const createDocumentSet = async ({ is_public, users, groups, + federated_connectors, }: DocumentSetCreationRequest) => { return fetch("/api/manage/admin/document-set", { method: "POST", @@ -27,6 +31,7 @@ export const createDocumentSet = async ({ is_public, users, groups, + federated_connectors, }), }); }; @@ -38,6 +43,7 @@ interface DocumentSetUpdateRequest { is_public: boolean; users: string[]; groups: number[]; + federated_connectors: FederatedConnectorConfig[]; } export const updateDocumentSet = async ({ @@ -47,6 +53,7 @@ export const updateDocumentSet = async ({ is_public, users, groups, + federated_connectors, }: DocumentSetUpdateRequest) => { return fetch("/api/manage/admin/document-set", { method: "PATCH", @@ -60,6 +67,7 @@ export const updateDocumentSet = async ({ is_public, users, groups, + federated_connectors, }), }); }; diff --git a/web/src/app/admin/documents/sets/page.tsx b/web/src/app/admin/documents/sets/page.tsx index 914c33220d..2900d0d662 100644 --- a/web/src/app/admin/documents/sets/page.tsx +++ b/web/src/app/admin/documents/sets/page.tsx @@ -39,9 +39,68 @@ import { TooltipTrigger, } from "@/components/ui/tooltip"; import CreateButton from "@/components/ui/createButton"; +import { SourceIcon } from "@/components/SourceIcon"; +import Link from "next/link"; const numToDisplay = 50; +// Component to display federated connectors with consistent styling +const FederatedConnectorTitle = ({ + federatedConnector, + showMetadata = true, + isLink = true, +}: { + federatedConnector: any; + showMetadata?: boolean; + isLink?: boolean; +}) => { + const sourceType = federatedConnector.source.replace(/^federated_/, ""); + + const mainSectionClassName = "text-blue-500 dark:text-blue-100 flex w-fit"; + const mainDisplay = ( + <> + +
+ {federatedConnector.name} +
+ + Federated + + + ); + + return ( +
+ {isLink ? ( + + {mainDisplay} + + ) : ( +
{mainDisplay}
+ )} + {showMetadata && Object.keys(federatedConnector.entities).length > 0 && ( +
+ {Object.entries(federatedConnector.entities) + .filter( + ([_, value]) => + value && + (Array.isArray(value) ? value.length > 0 : String(value).trim()) + ) + .map(([key, value]) => ( +
+ {key}:{" "} + {Array.isArray(value) ? value.join(", ") : String(value)} +
+ ))} +
+ )} +
+ ); +}; + const EditRow = ({ documentSet, isEditable, @@ -162,6 +221,7 @@ const DocumentSetTable = ({
+ {/* Regular Connectors */} {documentSet.cc_pair_descriptors.map( (ccPairDescriptor, ind) => { return ( @@ -184,6 +244,37 @@ const DocumentSetTable = ({ ); } )} + + {/* Federated Connectors */} + {documentSet.federated_connectors && + documentSet.federated_connectors.length > 0 && ( + <> + {documentSet.cc_pair_descriptors.length > 0 && ( +
+ )} + {documentSet.federated_connectors.map( + (federatedConnector, ind) => { + return ( +
+ +
+ ); + } + )} + + )}
@@ -191,7 +282,9 @@ const DocumentSetTable = ({ Up to Date - ) : documentSet.cc_pair_descriptors.length > 0 ? ( + ) : documentSet.cc_pair_descriptors.length > 0 || + (documentSet.federated_connectors && + documentSet.federated_connectors.length > 0) ? ( Syncing diff --git a/web/src/app/admin/federated/[id]/page.tsx b/web/src/app/admin/federated/[id]/page.tsx new file mode 100644 index 0000000000..2d024364d1 --- /dev/null +++ b/web/src/app/admin/federated/[id]/page.tsx @@ -0,0 +1,72 @@ +"use client"; + +import { useState, useEffect } from "react"; +import { notFound } from "next/navigation"; +import { Loader2 } from "lucide-react"; +import { useFederatedConnector } from "./useFederatedConnector"; +import { FederatedConnectorForm } from "@/components/admin/federated/FederatedConnectorForm"; + +export default function EditFederatedConnectorPage(props: { + params: Promise<{ id: string }>; +}) { + const [params, setParams] = useState<{ id: string } | null>(null); + + useEffect(() => { + props.params.then(setParams); + }, [props.params]); + + const { sourceType, connectorData, credentialSchema, isLoading, error } = + useFederatedConnector(params?.id ?? ""); + + if (isLoading) { + return ( +
+
+
+ +
+

+ Loading connector configuration... +

+

+ Retrieving connector details and credential schema +

+
+
+
+
+ ); + } + + if (error) { + return ( +
+
+
+

Error

+

{error}

+
+
+
+ ); + } + + if (!sourceType || !params) { + notFound(); + } + + const connectorId = parseInt(params.id); + + return ( +
+
+ +
+
+ ); +} diff --git a/web/src/app/admin/federated/[id]/useFederatedConnector.ts b/web/src/app/admin/federated/[id]/useFederatedConnector.ts new file mode 100644 index 0000000000..2b18776b76 --- /dev/null +++ b/web/src/app/admin/federated/[id]/useFederatedConnector.ts @@ -0,0 +1,94 @@ +import { useState, useEffect } from "react"; +import { + ConfigurableSources, + FederatedConnectorDetail, + CredentialSchemaResponse, +} from "@/lib/types"; + +interface UseFederatedConnectorResult { + sourceType: ConfigurableSources | null; + connectorData: FederatedConnectorDetail | null; + credentialSchema: CredentialSchemaResponse | null; + isLoading: boolean; + error: string | null; +} + +export function useFederatedConnector( + connectorId: string +): UseFederatedConnectorResult { + const [sourceType, setSourceType] = useState( + null + ); + const [connectorData, setConnectorData] = + useState(null); + const [credentialSchema, setCredentialSchema] = + useState(null); + const [isLoading, setIsLoading] = useState(true); + const [error, setError] = useState(null); + + useEffect(() => { + const fetchData = async () => { + try { + setIsLoading(true); + setError(null); + + // First, fetch connector details to get the source type + const connectorResponse = await fetch(`/api/federated/${connectorId}`); + + if (!connectorResponse.ok) { + throw new Error( + `Failed to fetch connector: ${connectorResponse.statusText}` + ); + } + + const connectorData: FederatedConnectorDetail = + await connectorResponse.json(); + + // Extract source type from the federated source string (remove 'federated_' prefix) + const extractedSourceType = connectorData.source.replace( + /^federated_/, + "" + ) as ConfigurableSources; + + // Now fetch credential schema and set state in parallel + const schemaPromise = fetch( + `/api/federated/sources/federated_${extractedSourceType}/credentials/schema` + ); + + // Set the data we already have + setConnectorData(connectorData); + setSourceType(extractedSourceType); + + // Wait for schema fetch to complete + const schemaResponse = await schemaPromise; + + if (!schemaResponse.ok) { + throw new Error( + `Failed to fetch schema: ${schemaResponse.statusText}` + ); + } + + const schemaData: CredentialSchemaResponse = + await schemaResponse.json(); + setCredentialSchema(schemaData); + } catch (error) { + console.error("Error fetching federated connector data:", error); + setError(`Failed to load connector: ${error}`); + } finally { + setIsLoading(false); + } + }; + + if (connectorId) { + fetchData(); + } + }, [connectorId]); + + return { + sourceType, + connectorData, + credentialSchema, + isLoading, + error, + }; +} diff --git a/web/src/app/admin/indexing/status/CCPairIndexingStatusTable.tsx b/web/src/app/admin/indexing/status/CCPairIndexingStatusTable.tsx index 88b317fcd1..b39fb2378f 100644 --- a/web/src/app/admin/indexing/status/CCPairIndexingStatusTable.tsx +++ b/web/src/app/admin/indexing/status/CCPairIndexingStatusTable.tsx @@ -9,7 +9,7 @@ import { } from "@/components/ui/table"; import { Badge } from "@/components/ui/badge"; import { Button } from "@/components/ui/button"; -import { CCPairStatus, IndexAttemptStatus } from "@/components/Status"; +import { CCPairStatus } from "@/components/Status"; import { timeAgo } from "@/lib/time"; import { ConnectorIndexingStatus, @@ -17,6 +17,8 @@ import { GroupedConnectorSummaries, ValidSources, ValidStatuses, + FederatedConnectorDetail, + federatedSourceToRegularSource, } from "@/lib/types"; import { useRouter } from "next/navigation"; import { @@ -26,7 +28,6 @@ import { FiLock, FiUnlock, FiRefreshCw, - FiPauseCircle, } from "react-icons/fi"; import { Tooltip, @@ -36,7 +37,6 @@ import { } from "@/components/ui/tooltip"; import { SourceIcon } from "@/components/SourceIcon"; import { getSourceDisplayName } from "@/lib/sources"; -import { Warning } from "@phosphor-icons/react"; import Cookies from "js-cookie"; import { TOGGLED_CONNECTORS_COOKIE_NAME } from "@/lib/constants"; import { usePaidEnterpriseFeaturesEnabled } from "@/components/settings/usePaidEnterpriseFeaturesEnabled"; @@ -213,12 +213,78 @@ border border-border dark:border-neutral-700 ); } +function FederatedConnectorRow({ + federatedConnector, + invisible, +}: { + federatedConnector: FederatedConnectorDetail; + invisible?: boolean; +}) { + const router = useRouter(); + const isPaidEnterpriseFeaturesEnabled = usePaidEnterpriseFeaturesEnabled(); + + const handleManageClick = (e: any) => { + e.stopPropagation(); + router.push(`/admin/federated/${federatedConnector.id}`); + }; + + return ( + { + router.push(`/admin/federated/${federatedConnector.id}`); + }} + > + +

+ {federatedConnector.name} +

+ + N/A + + Indexed + + {isPaidEnterpriseFeaturesEnabled && ( + + + Federated Access + + + )} + N/A + + + + + + + +

Manage Federated Connector

+ + + + + + ); +} + export function CCPairIndexingStatusTable({ ccPairsIndexingStatuses, editableCcPairsIndexingStatuses, + federatedConnectors, }: { ccPairsIndexingStatuses: ConnectorIndexingStatus[]; editableCcPairsIndexingStatuses: ConnectorIndexingStatus[]; + federatedConnectors: FederatedConnectorDetail[]; }) { const [searchTerm, setSearchTerm] = useState(""); @@ -291,12 +357,17 @@ export function CCPairIndexingStatusTable({ {} as GroupedConnectorSummaries; sorted.forEach((source) => { const statuses = grouped[source]; + const federatedForSource = federatedConnectors.filter( + (fc) => federatedSourceToRegularSource(fc.source) === source + ); + summaries[source] = { - count: statuses.length, - active: statuses.filter( - (status) => - status.cc_pair_status === ConnectorCredentialPairStatus.ACTIVE - ).length, + count: statuses.length + federatedForSource.length, + active: + statuses.filter( + (status) => + status.cc_pair_status === ConnectorCredentialPairStatus.ACTIVE + ).length + federatedForSource.length, // All federated connectors are considered active public: statuses.filter((status) => status.access_type === "public") .length, totalDocsIndexed: statuses.reduce( @@ -372,6 +443,27 @@ export function CCPairIndexingStatusTable({ }; }, [ccPairsIndexingStatuses, editableCcPairsIndexingStatuses, filterOptions]); + // Combine regular connector sources with sources that only have federated connectors + const allSourcesWithFederated = useMemo(() => { + const federatedSources = Array.from( + new Set( + federatedConnectors.map((fc) => + federatedSourceToRegularSource(fc.source) + ) + ) + ) as ValidSources[]; + + // Ensure we keep original ordering for existing sources, then append any new ones + const combined = [...sortedSources]; + federatedSources.forEach((src) => { + if (!combined.includes(src)) { + combined.push(src); + } + }); + + return combined; + }, [sortedSources, federatedConnectors]); + // Determine which sources to display based on filters and search const displaySources = useMemo(() => { const hasActiveFilters = @@ -383,8 +475,8 @@ export function CCPairIndexingStatusTable({ return Object.keys(filteredGroupedStatuses) as ValidSources[]; } - return sortedSources; - }, [sortedSources, filteredGroupedStatuses, filterOptions]); + return allSourcesWithFederated; + }, [allSourcesWithFederated, filteredGroupedStatuses, filterOptions]); const handleFilterChange = (newFilters: FilterOptions) => { setFilterOptions(newFilters); @@ -416,24 +508,6 @@ export function CCPairIndexingStatusTable({ } }; - const clearAllFilters = () => { - const emptyFilters: FilterOptions = { - accessType: null, - docsCountFilter: { - operator: null, - value: null, - }, - lastStatus: null, - }; - - setFilterOptions(emptyFilters); - - // Reset the FilterComponent's internal state - if (filterComponentRef.current) { - filterComponentRef.current.resetFilters(); - } - }; - // Check if filters are active const hasActiveFilters = useMemo(() => { return ( @@ -457,6 +531,7 @@ export function CCPairIndexingStatusTable({ JSON.stringify(newConnectorsToggled) ); }; + const toggleSources = () => { const connectors = sortedSources.reduce( (acc, source) => { @@ -613,7 +688,9 @@ export function CCPairIndexingStatusTable({ .includes(searchTerm.toLowerCase()); const statuses = - filteredGroupedStatuses[source] || groupedStatuses[source]; + filteredGroupedStatuses[source] ?? + groupedStatuses[source] ?? + []; const matchingConnectors = statuses.filter((status) => (status.name || "") @@ -621,13 +698,32 @@ export function CCPairIndexingStatusTable({ .includes(searchTerm.toLowerCase()) ); - if (sourceMatches || matchingConnectors.length > 0) { + const federatedForSource = federatedConnectors.filter( + (fc) => federatedSourceToRegularSource(fc.source) === source + ); + + const hasFederatedMatches = federatedForSource.some((fc) => + fc.name.toLowerCase().includes(searchTerm.toLowerCase()) + ); + + if ( + sourceMatches || + matchingConnectors.length > 0 || + hasFederatedMatches + ) { + const summaryForSource = groupSummaries[source] ?? { + count: federatedForSource.length, + active: federatedForSource.length, // All federated connectors are considered active + public: 0, + totalDocsIndexed: 0, + errors: 0, + }; return (
toggleSource(source)} /> @@ -638,7 +734,7 @@ export function CCPairIndexingStatusTable({ Last Indexed Status {isPaidEnterpriseFeaturesEnabled && ( - Permissions + Permissions / Access )} Total Docs @@ -656,6 +752,21 @@ export function CCPairIndexingStatusTable({ /> ) )} + + {/* Add federated connectors belonging to this source */} + {(sourceMatches + ? federatedForSource + : federatedForSource.filter((fc) => + fc.name + .toLowerCase() + .includes(searchTerm.toLowerCase()) + ) + ).map((federatedConnector) => ( + + ))} )} diff --git a/web/src/app/admin/indexing/status/page.tsx b/web/src/app/admin/indexing/status/page.tsx index e55fdc2a7e..044ce75189 100644 --- a/web/src/app/admin/indexing/status/page.tsx +++ b/web/src/app/admin/indexing/status/page.tsx @@ -6,7 +6,10 @@ import { CCPairIndexingStatusTable } from "./CCPairIndexingStatusTable"; import { AdminPageTitle } from "@/components/admin/Title"; import Link from "next/link"; import Text from "@/components/ui/text"; -import { useConnectorCredentialIndexingStatus } from "@/lib/hooks"; +import { + useConnectorCredentialIndexingStatus, + useFederatedConnectors, +} from "@/lib/hooks"; import { usePopupFromQuery } from "@/components/popup/PopupFromQuery"; import { Button } from "@/components/ui/button"; @@ -23,7 +26,22 @@ function Main() { error: editableIndexAttemptError, } = useConnectorCredentialIndexingStatus(undefined, true); - if (indexAttemptIsLoading || editableIndexAttemptIsLoading) { + const { + data: federatedConnectorsData, + isLoading: federatedConnectorsIsLoading, + error: federatedConnectorsError, + } = useFederatedConnectors(); + + // Debug logging + console.log("Federated connectors data:", federatedConnectorsData); + console.log("Federated connectors error:", federatedConnectorsError); + console.log("Federated connectors loading:", federatedConnectorsIsLoading); + + if ( + indexAttemptIsLoading || + editableIndexAttemptIsLoading || + federatedConnectorsIsLoading + ) { return ; } @@ -42,7 +60,10 @@ function Main() { ); } - if (indexAttemptData.length === 0) { + if ( + indexAttemptData.length === 0 && + (!federatedConnectorsData || federatedConnectorsData.length === 0) + ) { return ( It looks like you don't have any connectors setup yet. Visit the{" "} @@ -69,6 +90,7 @@ function Main() { ); } diff --git a/web/src/app/chat/ChatPage.tsx b/web/src/app/chat/ChatPage.tsx index 18adfe8416..f71f60f18c 100644 --- a/web/src/app/chat/ChatPage.tsx +++ b/web/src/app/chat/ChatPage.tsx @@ -98,6 +98,8 @@ import { ChatInputBar } from "./input/ChatInputBar"; import { useChatContext } from "@/components/context/ChatContext"; import { ChatPopup } from "./ChatPopup"; import FunctionalHeader from "@/components/chat/Header"; +import { FederatedOAuthModal } from "@/components/chat/FederatedOAuthModal"; +import { useFederatedOAuthStatus } from "@/lib/hooks/useFederatedOAuthStatus"; import { useSidebarVisibility } from "@/components/chat/hooks"; import { PRO_SEARCH_TOGGLED_COOKIE_NAME, @@ -127,7 +129,7 @@ import { useSidebarShortcut } from "@/lib/browserUtilities"; import { FilePickerModal } from "./my-documents/components/FilePicker"; import { SourceMetadata } from "@/lib/search/interfaces"; -import { ValidSources } from "@/lib/types"; +import { ValidSources, FederatedConnectorDetail } from "@/lib/types"; import { FileResponse, FolderResponse, @@ -137,6 +139,8 @@ import { ChatSearchModal } from "./chat_search/ChatSearchModal"; import { ErrorBanner } from "./message/Resubmit"; import MinimalMarkdown from "@/components/chat/MinimalMarkdown"; import { WelcomeModal } from "@/components/initialSetup/welcome/WelcomeModal"; +import { useFederatedConnectors } from "@/lib/hooks"; +import { Button } from "@/components/ui/button"; const TEMP_USER_MESSAGE_ID = -1; const TEMP_ASSISTANT_MESSAGE_ID = -2; @@ -191,6 +195,103 @@ export function ChatPage({ setCurrentMessageFiles, } = useDocumentsContext(); + // Federated OAuth status + const { + connectors: federatedConnectors, + hasUnauthenticatedConnectors, + loading: oauthLoading, + refetch: refetchFederatedConnectors, + } = useFederatedOAuthStatus(); + + // Also fetch federated connectors for the sources list + const { data: federatedConnectorsData } = useFederatedConnectors(); + + const MAX_SKIP_COUNT = 1; + + // Check localStorage for previous skip preference and count + const [oAuthModalState, setOAuthModalState] = useState<{ + hidden: boolean; + skipCount: number; + }>(() => { + if (typeof window !== "undefined") { + const skipData = localStorage.getItem("federatedOAuthModalSkipData"); + if (skipData) { + try { + const parsed = JSON.parse(skipData); + // Check if we're still within the hide duration (1 hour) + const now = Date.now(); + const hideUntil = parsed.hideUntil || 0; + const isWithinHideDuration = now < hideUntil; + + return { + hidden: parsed.permanentlyHidden || isWithinHideDuration, + skipCount: parsed.skipCount || 0, + }; + } catch { + return { hidden: false, skipCount: 0 }; + } + } + } + return { hidden: false, skipCount: 0 }; + }); + + const handleOAuthModalSkip = () => { + if (typeof window !== "undefined") { + const newSkipCount = oAuthModalState.skipCount + 1; + + // If we've reached the max skip count, show the "No problem!" modal first + if (newSkipCount >= MAX_SKIP_COUNT) { + // Don't hide immediately - let the "No problem!" modal show + setOAuthModalState({ + hidden: false, + skipCount: newSkipCount, + }); + } else { + // For first skip, hide after a delay to show "No problem!" modal + const oneHourFromNow = Date.now() + 60 * 60 * 1000; // 1 hour in milliseconds + + const skipData = { + skipCount: newSkipCount, + hideUntil: oneHourFromNow, + permanentlyHidden: false, + }; + + localStorage.setItem( + "federatedOAuthModalSkipData", + JSON.stringify(skipData) + ); + + setOAuthModalState({ + hidden: true, + skipCount: newSkipCount, + }); + } + } + }; + + // Handle the final dismissal of the "No problem!" modal + const handleOAuthModalFinalDismiss = () => { + if (typeof window !== "undefined") { + const oneHourFromNow = Date.now() + 60 * 60 * 1000; // 1 hour in milliseconds + + const skipData = { + skipCount: oAuthModalState.skipCount, + hideUntil: oneHourFromNow, + permanentlyHidden: false, + }; + + localStorage.setItem( + "federatedOAuthModalSkipData", + JSON.stringify(skipData) + ); + + setOAuthModalState({ + hidden: true, + skipCount: oAuthModalState.skipCount, + }); + } + }; + const defaultAssistantIdRaw = searchParams?.get( SEARCH_PARAM_NAMES.PERSONA_ID ); @@ -275,7 +376,7 @@ export function ChatPage({ filterManager.buildFiltersFromQueryString( newSearchParams.toString(), - availableSources, + sources, documentSets.map((ds) => ds.name), tags ); @@ -373,8 +474,28 @@ export function ChatPage({ const sources: SourceMetadata[] = useMemo(() => { const uniqueSources = Array.from(new Set(availableSources)); - return uniqueSources.map((source) => getSourceMetadata(source)); - }, [availableSources]); + const regularSources = uniqueSources.map((source) => + getSourceMetadata(source) + ); + + // Add federated connectors as sources + const federatedSources = + federatedConnectorsData?.map((connector: FederatedConnectorDetail) => { + return getSourceMetadata(connector.source); + }) || []; + + // Combine sources and deduplicate based on internalName + const allSources = [...regularSources, ...federatedSources]; + const deduplicatedSources = allSources.reduce((acc, source) => { + const existing = acc.find((s) => s.internalName === source.internalName); + if (!existing) { + acc.push(source); + } + return acc; + }, [] as SourceMetadata[]); + + return deduplicatedSources; + }, [availableSources, federatedConnectorsData]); const stopGenerating = () => { const currentSession = currentSessionId(); @@ -2039,11 +2160,7 @@ export function ChatPage({ Cookies.set( SIDEBAR_TOGGLED_COOKIE_NAME, String(!sidebarVisible).toLocaleLowerCase() - ), - { - path: "/", - }; - + ); toggle(); }; const removeToggle = () => { @@ -2330,6 +2447,18 @@ export function ChatPage({ {shouldShowWelcomeModal && } + {isReady && !oAuthModalState.hidden && hasUnauthenticatedConnectors && ( + = MAX_SKIP_COUNT + ? handleOAuthModalFinalDismiss + : handleOAuthModalSkip + } + skipCount={oAuthModalState.skipCount} + /> + )} + {/* ChatPopup is a custom popup that displays a admin-specified message on initial user visit. Only used in the EE version of the app. */} {popup} @@ -2358,6 +2487,9 @@ export function ChatPage({ setCurrentLlm={(newLlm) => llmManager.updateCurrentLlm(newLlm)} defaultModel={user?.preferences.default_model!} llmProviders={llmProviders} + ccPairs={ccPairs} + federatedConnectors={federatedConnectors} + refetchFederatedConnectors={refetchFederatedConnectors} onClose={() => { setUserSettingsToggled(false); setSettingsToggled(false); diff --git a/web/src/app/chat/input/ChatInputBar.tsx b/web/src/app/chat/input/ChatInputBar.tsx index 3ab2552146..a2c04905cc 100644 --- a/web/src/app/chat/input/ChatInputBar.tsx +++ b/web/src/app/chat/input/ChatInputBar.tsx @@ -867,15 +867,11 @@ export function ChatInputBar({ } availableTags={availableTags} filterManager={filterManager} - trigger={ - - } + trigger={{ + name: "Filters", + Icon: FiFilter, + tooltipContent: "Filter your search", + }} /> )}
diff --git a/web/src/app/chat/input/SelectedFilterDisplay.tsx b/web/src/app/chat/input/SelectedFilterDisplay.tsx index 158dfc5f66..2ead99c0fe 100644 --- a/web/src/app/chat/input/SelectedFilterDisplay.tsx +++ b/web/src/app/chat/input/SelectedFilterDisplay.tsx @@ -98,7 +98,10 @@ export function SelectedFilterDisplay({ } > <> - + {source.displayName} diff --git a/web/src/app/chat/modal/UserSettingsModal.tsx b/web/src/app/chat/modal/UserSettingsModal.tsx index 427b1b332b..e006e901dd 100644 --- a/web/src/app/chat/modal/UserSettingsModal.tsx +++ b/web/src/app/chat/modal/UserSettingsModal.tsx @@ -24,11 +24,15 @@ import { Monitor, Moon, Sun } from "lucide-react"; import { useTheme } from "next-themes"; import { Button } from "@/components/ui/button"; import { Input } from "@/components/ui/input"; -import { FiTrash2 } from "react-icons/fi"; +import { FiTrash2, FiExternalLink } from "react-icons/fi"; import { deleteAllChatSessions } from "../lib"; import { useChatContext } from "@/components/context/ChatContext"; +import { FederatedConnectorOAuthStatus } from "@/components/chat/FederatedOAuthModal"; +import { SourceIcon } from "@/components/SourceIcon"; +import { ValidSources, CCPairBasicInfo } from "@/lib/types"; +import { getSourceMetadata } from "@/lib/sources"; -type SettingsSection = "settings" | "password"; +type SettingsSection = "settings" | "password" | "connectors"; export function UserSettingsModal({ setPopup, @@ -36,12 +40,18 @@ export function UserSettingsModal({ onClose, setCurrentLlm, defaultModel, + ccPairs, + federatedConnectors, + refetchFederatedConnectors, }: { setPopup: (popupSpec: PopupSpec | null) => void; llmProviders: LLMProviderDescriptor[]; setCurrentLlm?: (newLlm: LlmDescriptor) => void; onClose: () => void; defaultModel: string | null; + ccPairs?: CCPairBasicInfo[]; + federatedConnectors?: FederatedConnectorOAuthStatus[]; + refetchFederatedConnectors?: () => void; }) { const { refreshUser, @@ -64,6 +74,11 @@ export function UserSettingsModal({ useState("settings"); const [isDeleteAllLoading, setIsDeleteAllLoading] = useState(false); const [showDeleteConfirmation, setShowDeleteConfirmation] = useState(false); + const [isDisconnecting, setIsDisconnecting] = useState(null); + + const hasConnectors = + (ccPairs && ccPairs.length > 0) || + (federatedConnectors && federatedConnectors.length > 0); useEffect(() => { const container = containerRef.current; @@ -165,6 +180,39 @@ export function UserSettingsModal({ } }; + const handleConnectOAuth = (authorizeUrl: string) => { + // Redirect to OAuth URL in the same window + router.push(authorizeUrl); + }; + + const handleDisconnectOAuth = async (connectorId: number) => { + setIsDisconnecting(connectorId); + try { + const response = await fetch(`/api/federated/${connectorId}/oauth`, { + method: "DELETE", + }); + + if (response.ok) { + setPopup({ + message: "Disconnected successfully", + type: "success", + }); + if (refetchFederatedConnectors) { + refetchFederatedConnectors(); + } + } else { + throw new Error("Failed to disconnect"); + } + } catch (error) { + setPopup({ + message: "Failed to disconnect", + type: "error", + }); + } finally { + setIsDisconnecting(null); + } + }; + const settings = useContext(SettingsContext); const autoScroll = settings?.settings?.auto_scroll; @@ -245,15 +293,15 @@ export function UserSettingsModal({ -
+

User Settings

-
- {showPasswordSection && ( -
+
+ {(showPasswordSection || hasConnectors) && ( +
)} -
+
{activeSection === "settings" && (
@@ -488,6 +558,194 @@ export function UserSettingsModal({
)} + {activeSection === "connectors" && ( +
+
+

+ Connected Services +

+

+ Manage your connected services to search across all your + content. +

+ + {/* Indexed Connectors Section */} + {ccPairs && ccPairs.length > 0 && ( +
+

+ Indexed Connectors +

+ {(() => { + // Group connectors by source + const groupedConnectors = ccPairs.reduce( + (acc, ccPair) => { + const source = ccPair.source; + if (!acc[source]) { + acc[source] = { + source, + count: 0, + hasSuccessfulRun: false, + }; + } + acc[source]!.count++; + if (ccPair.has_successful_run) { + acc[source]!.hasSuccessfulRun = true; + } + return acc; + }, + {} as Record< + string, + { + source: ValidSources; + count: number; + hasSuccessfulRun: boolean; + } + > + ); + + // Helper function to format source names + const formatSourceName = (source: string) => { + return source + .split("_") + .map( + (word) => + word.charAt(0).toUpperCase() + word.slice(1) + ) + .join(" "); + }; + + return Object.values(groupedConnectors).map((group) => ( +
+
+ +
+

+ {formatSourceName(group.source)} +

+

+ {group.count > 1 + ? `${group.count} connectors` + : "Connected"} +

+
+
+
+ Active +
+
+ )); + })()} +
+ )} + + {/* Federated Search Section */} + {federatedConnectors && federatedConnectors.length > 0 && ( +
+

+ Federated Connectors +

+ {(() => { + // Helper function to format source names + const formatSourceName = (source: string) => { + return source + .split("_") + .map( + (word) => + word.charAt(0).toUpperCase() + word.slice(1) + ) + .join(" "); + }; + + return federatedConnectors.map((connector) => { + const sourceMetadata = getSourceMetadata( + connector.source as ValidSources + ); + return ( +
+
+ +
+

+ {formatSourceName( + sourceMetadata.displayName + )} +

+

+ {connector.has_oauth_token + ? "Connected" + : "Not connected"} +

+
+
+
+ {connector.has_oauth_token ? ( + + ) : ( + + )} +
+
+ ); + }); + })()} +
+ )} + + {!hasConnectors && ( +
+

+ No connectors available. +

+
+ )} +
+
+ )}
diff --git a/web/src/app/federated/oauth/callback/page.tsx b/web/src/app/federated/oauth/callback/page.tsx new file mode 100644 index 0000000000..f37f49e449 --- /dev/null +++ b/web/src/app/federated/oauth/callback/page.tsx @@ -0,0 +1,206 @@ +"use client"; + +import { useEffect, useState } from "react"; +import { useRouter, useSearchParams } from "next/navigation"; +import { CheckmarkIcon, TriangleAlertIcon } from "@/components/icons/icons"; +import CardSection from "@/components/admin/CardSection"; +import { Button } from "@/components/ui/button"; +import { getSourceDisplayName } from "@/lib/sources"; + +export default function FederatedOAuthCallbackPage() { + const router = useRouter(); + const searchParams = useSearchParams(); + + const [statusMessage, setStatusMessage] = useState("Processing..."); + const [statusDetails, setStatusDetails] = useState( + "Please wait while we complete the setup." + ); + const [isError, setIsError] = useState(false); + const [isSuccess, setIsSuccess] = useState(false); + const [isLoading, setIsLoading] = useState(true); + + // Extract query parameters + const code = searchParams?.get("code"); + const state = searchParams?.get("state"); + const error = searchParams?.get("error"); + const errorDescription = searchParams?.get("error_description"); + + // Auto-redirect for success cases + useEffect(() => { + if (isSuccess) { + const timer = setTimeout(() => { + router.push("/chat"); + }, 2000); + return () => clearTimeout(timer); + } + }, [isSuccess, router]); + + useEffect(() => { + const handleOAuthCallback = async () => { + // Handle OAuth error from provider + if (error) { + setStatusMessage("Authorization Failed"); + setStatusDetails( + errorDescription || + "The authorization was cancelled or failed. Please try again." + ); + setIsError(true); + setIsLoading(false); + return; + } + + // Validate required parameters + if (!code || !state) { + setStatusMessage("Invalid Request"); + setStatusDetails( + "The authorization request was incomplete. Please try again." + ); + setIsError(true); + setIsLoading(false); + return; + } + + try { + // Use the generic callback endpoint + const url = `/api/federated/callback?code=${encodeURIComponent( + code + )}&state=${encodeURIComponent(state)}`; + + const response = await fetch(url, { + method: "POST", + headers: { + "Content-Type": "application/json", + }, + }); + + if (!response.ok) { + let errorMessage = "Failed to complete authorization"; + try { + const errorData = await response.json(); + if (errorData.detail) { + // Clean up technical error messages for user-friendly display + if (errorData.detail.includes("validation errors")) { + errorMessage = + "Configuration error - please check your connector settings"; + } else if (errorData.detail.includes("client_secret")) { + errorMessage = + "Authentication credentials are missing or invalid"; + } else if (errorData.detail.includes("oauth")) { + errorMessage = "OAuth authorization failed"; + } else { + errorMessage = "Authorization failed - please try again"; + } + } + } catch (parseError) { + console.error("Error parsing response:", parseError); + } + throw new Error(errorMessage); + } + + // Parse the response to get source information + const responseData = await response.json(); + const source = responseData.source; + const displayName = source ? getSourceDisplayName(source) : "connector"; + + setStatusMessage("Success!"); + setStatusDetails( + `Your ${displayName} authorization completed successfully. You can now use this connector for search.` + ); + setIsSuccess(true); + setIsError(false); + setIsLoading(false); + } catch (error) { + console.error("Federated OAuth callback error:", error); + setStatusMessage("Something Went Wrong"); + setStatusDetails( + error instanceof Error + ? error.message + : "An error occurred during the OAuth process. Please try again." + ); + setIsError(true); + setIsLoading(false); + } + }; + + handleOAuthCallback(); + }, [code, state, error, errorDescription]); + + const getStatusIcon = () => { + if (isLoading) { + return ( +
+ ); + } + if (isSuccess) { + return ( + + ); + } + if (isError) { + return ( + + ); + } + return null; + }; + + const getStatusColor = () => { + if (isSuccess) return "text-green-600 dark:text-green-400"; + if (isError) return "text-red-600 dark:text-red-400"; + return "text-gray-600 dark:text-gray-300"; + }; + + return ( +
+
+ +
+ {getStatusIcon()} + +

+ {statusMessage} +

+ +

+ {statusDetails} +

+ + {isSuccess && ( +
+

+ Redirecting to chat in 2 seconds... +

+
+ )} + +
+ {isError && ( +
+ +
+ )} + + {isLoading && ( +

+ This may take a few moments... +

+ )} +
+
+ +
+
+ ); +} diff --git a/web/src/components/FederatedConnectorSelector.tsx b/web/src/components/FederatedConnectorSelector.tsx new file mode 100644 index 0000000000..985d1985d1 --- /dev/null +++ b/web/src/components/FederatedConnectorSelector.tsx @@ -0,0 +1,521 @@ +import React, { useState, useRef, useEffect } from "react"; +import { + FederatedConnectorDetail, + FederatedConnectorConfig, + federatedSourceToRegularSource, + ValidSources, +} from "@/lib/types"; +import { SourceIcon } from "@/components/SourceIcon"; +import { X, Search, Settings } from "lucide-react"; +import { Label } from "@/components/ui/label"; +import { ErrorMessage } from "formik"; +import { Button } from "@/components/ui/button"; +import { Input } from "@/components/ui/input"; +import { + Dialog, + DialogContent, + DialogHeader, + DialogTitle, + DialogTrigger, +} from "@/components/ui/dialog"; + +interface FederatedConnectorSelectorProps { + name: string; + label: string; + federatedConnectors: FederatedConnectorDetail[]; + selectedConfigs: FederatedConnectorConfig[]; + onChange: (selectedConfigs: FederatedConnectorConfig[]) => void; + disabled?: boolean; + placeholder?: string; + showError?: boolean; +} + +interface EntityConfigDialogProps { + connectorId: number; + connectorName: string; + connectorSource: ValidSources | null; + currentEntities: Record; + onSave: (entities: Record) => void; + onClose: () => void; + isOpen: boolean; +} + +const EntityConfigDialog = ({ + connectorId, + connectorName, + connectorSource, + currentEntities, + onSave, + onClose, + isOpen, +}: EntityConfigDialogProps) => { + const [entities, setEntities] = + useState>(currentEntities); + const [entitySchema, setEntitySchema] = useState | null>( + null + ); + const [isLoading, setIsLoading] = useState(false); + const [error, setError] = useState(null); + + useEffect(() => { + if (isOpen && connectorId) { + const fetchEntitySchema = async () => { + setIsLoading(true); + setError(null); + try { + const response = await fetch( + `/api/federated/${connectorId}/entities` + ); + if (!response.ok) { + throw new Error( + `Failed to fetch entity schema: ${response.statusText}` + ); + } + const data = await response.json(); + setEntitySchema(data.entities); + } catch (err) { + setError( + err instanceof Error ? err.message : "Failed to load entity schema" + ); + } finally { + setIsLoading(false); + } + }; + fetchEntitySchema(); + } + }, [isOpen, connectorId]); + + const handleSave = () => { + onSave(entities); + onClose(); + }; + + const handleEntityChange = (key: string, value: any) => { + setEntities((prev) => ({ + ...prev, + [key]: value, + })); + }; + + if (!connectorSource) { + return null; + } + + return ( + + + + + + Configure {connectorName} + + + +
+ {isLoading && ( +
+
+

+ Loading configuration... +

+
+ )} + + {error && ( +
+ {error} +
+ )} + + {entitySchema && !isLoading && ( +
+

+ Configure which entities to include from this connector: +

+ + {Object.entries(entitySchema).map( + ([key, field]: [string, any]) => ( +
+ + + {field.type === "list" ? ( +
+ { + const value = e.target.value; + const list = value + ? value + .split(",") + .map((item) => item.trim()) + .filter(Boolean) + : []; + handleEntityChange(key, list); + }} + /> +

+ {field.description && field.description !== key + ? field.description + : `Enter ${key} separated by commas`} +

+
+ ) : ( +
+ + handleEntityChange(key, e.target.value) + } + /> + {field.description && field.description !== key && ( +

+ {field.description} +

+ )} +
+ )} +
+ ) + )} +
+ )} + +
+ + +
+
+ + + ); +}; + +export const FederatedConnectorSelector = ({ + name, + label, + federatedConnectors, + selectedConfigs, + onChange, + disabled = false, + placeholder = "Search federated connectors...", + showError = false, +}: FederatedConnectorSelectorProps) => { + const [open, setOpen] = useState(false); + const [searchQuery, setSearchQuery] = useState(""); + const [configDialogState, setConfigDialogState] = useState<{ + isOpen: boolean; + connectorId: number | null; + connectorName: string; + connectorSource: ValidSources | null; + currentEntities: Record; + }>({ + isOpen: false, + connectorId: null, + connectorName: "", + connectorSource: null, + currentEntities: {}, + }); + const dropdownRef = useRef(null); + const inputRef = useRef(null); + + const selectedConnectorIds = selectedConfigs.map( + (config) => config.federated_connector_id + ); + + const selectedConnectors = federatedConnectors.filter((connector) => + selectedConnectorIds.includes(connector.id) + ); + + const unselectedConnectors = federatedConnectors.filter( + (connector) => !selectedConnectorIds.includes(connector.id) + ); + + const allConnectorsSelected = unselectedConnectors.length === 0; + + const filteredUnselectedConnectors = unselectedConnectors.filter( + (connector) => { + const connectorName = connector.name; + return connectorName.toLowerCase().includes(searchQuery.toLowerCase()); + } + ); + + useEffect(() => { + if (allConnectorsSelected && open) { + setOpen(false); + inputRef.current?.blur(); + setSearchQuery(""); + } + }, [allConnectorsSelected, open]); + + const selectConnector = (connectorId: number) => { + // Add connector with empty entities configuration + const newConfig: FederatedConnectorConfig = { + federated_connector_id: connectorId, + entities: {}, + }; + + const newSelectedConfigs = [...selectedConfigs, newConfig]; + onChange(newSelectedConfigs); + setSearchQuery(""); + + const willAllBeSelected = + federatedConnectors.length === newSelectedConfigs.length; + + if (!willAllBeSelected) { + setTimeout(() => { + inputRef.current?.focus(); + }, 0); + } + }; + + const removeConnector = (connectorId: number) => { + onChange( + selectedConfigs.filter( + (config) => config.federated_connector_id !== connectorId + ) + ); + }; + + const openConfigDialog = (connectorId: number) => { + const connector = federatedConnectors.find((c) => c.id === connectorId); + const config = selectedConfigs.find( + (c) => c.federated_connector_id === connectorId + ); + + if (connector) { + setConfigDialogState({ + isOpen: true, + connectorId, + connectorName: connector.name, + connectorSource: connector.source, + currentEntities: config?.entities || {}, + }); + } + }; + + const saveEntityConfig = (entities: Record) => { + const updatedConfigs = selectedConfigs.map((config) => { + if (config.federated_connector_id === configDialogState.connectorId) { + return { + ...config, + entities, + }; + } + return config; + }); + onChange(updatedConfigs); + }; + + useEffect(() => { + const handleClickOutside = (event: MouseEvent) => { + if ( + dropdownRef.current && + !dropdownRef.current.contains(event.target as Node) && + inputRef.current !== event.target && + !inputRef.current?.contains(event.target as Node) + ) { + setOpen(false); + } + }; + + document.addEventListener("mousedown", handleClickOutside); + return () => { + document.removeEventListener("mousedown", handleClickOutside); + }; + }, []); + + const handleKeyDown = (e: React.KeyboardEvent) => { + if (e.key === "Escape") { + setOpen(false); + } + }; + + const effectivePlaceholder = allConnectorsSelected + ? "All federated connectors selected" + : placeholder; + + const isInputDisabled = disabled || allConnectorsSelected; + + return ( +
+ {label && } + +

+ Documents from selected federated connectors will be searched in + real-time during queries. +

+
+
+ + { + setSearchQuery(e.target.value); + setOpen(true); + }} + onFocus={() => { + if (!allConnectorsSelected) { + setOpen(true); + } + }} + onKeyDown={handleKeyDown} + placeholder={effectivePlaceholder} + className={`h-9 w-full pl-9 pr-10 py-2 bg-transparent dark:bg-transparent text-sm outline-none disabled:cursor-not-allowed disabled:opacity-50 ${ + allConnectorsSelected + ? "text-neutral-500 dark:text-neutral-400" + : "" + }`} + disabled={isInputDisabled} + /> +
+ + {open && !allConnectorsSelected && ( +
+ {filteredUnselectedConnectors.length === 0 ? ( +
+ {searchQuery + ? "No matching federated connectors found" + : "No more federated connectors available"} +
+ ) : ( +
+ {filteredUnselectedConnectors.map((connector) => ( +
selectConnector(connector.id)} + > +
+
+ +
+ {connector.name} +
+
+ ))} +
+ )} +
+ )} +
+ + {selectedConnectors.length > 0 ? ( +
+
+ {selectedConnectors.map((connector) => { + const config = selectedConfigs.find( + (c) => c.federated_connector_id === connector.id + ); + const hasEntitiesConfigured = + config && Object.keys(config.entities).length > 0; + + return ( +
+
+
+ +
+ + {connector.name} + + {hasEntitiesConfigured && ( +
+ )} +
+
+ + +
+
+ ); + })} +
+
+ ) : ( +
+ No federated connectors selected. Search and select connectors above. +
+ )} + + + setConfigDialogState((prev) => ({ ...prev, isOpen: false })) + } + isOpen={configDialogState.isOpen} + /> + + {showError && ( + + )} +
+ ); +}; diff --git a/web/src/components/admin/ClientLayout.tsx b/web/src/components/admin/ClientLayout.tsx index cf8b8aca9c..2a7b551988 100644 --- a/web/src/components/admin/ClientLayout.tsx +++ b/web/src/components/admin/ClientLayout.tsx @@ -43,6 +43,7 @@ import { Button } from "../ui/button"; import useSWR from "swr"; import { errorHandlingFetcher } from "@/lib/fetcher"; import { useIsKGExposed } from "@/app/admin/kg/utils"; +import { useFederatedOAuthStatus } from "@/lib/hooks/useFederatedOAuthStatus"; const connectors_items = () => [ { @@ -404,9 +405,15 @@ export function ClientLayout({ const toggleUserSettings = () => { setUserSettingsOpen(!userSettingsOpen); }; - const { llmProviders } = useChatContext(); + const { llmProviders, ccPairs } = useChatContext(); const { popup, setPopup } = usePopup(); + // Fetch federated-connector info so the modal can list/refresh them + const { + connectors: federatedConnectors, + refetch: refetchFederatedConnectors, + } = useFederatedOAuthStatus(); + if (isLoading) { return <>; } @@ -428,6 +435,9 @@ export function ClientLayout({ setPopup={setPopup} onClose={() => setUserSettingsOpen(false)} defaultModel={user?.preferences?.default_model!} + ccPairs={ccPairs} + federatedConnectors={federatedConnectors} + refetchFederatedConnectors={refetchFederatedConnectors} /> )} diff --git a/web/src/components/admin/federated/FederatedConnectorForm.tsx b/web/src/components/admin/federated/FederatedConnectorForm.tsx new file mode 100644 index 0000000000..7d7887b825 --- /dev/null +++ b/web/src/components/admin/federated/FederatedConnectorForm.tsx @@ -0,0 +1,586 @@ +"use client"; + +import { useState, useEffect } from "react"; +import { Button } from "@/components/ui/button"; +import { + ConfigurableSources, + CredentialSchemaResponse, + CredentialFieldSpec, + FederatedConnectorCreateRequest, + FederatedConnectorDetail, +} from "@/lib/types"; +import { getSourceMetadata } from "@/lib/sources"; +import { SourceIcon } from "@/components/SourceIcon"; +import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card"; +import { Label } from "@/components/ui/label"; +import { Input } from "@/components/ui/input"; +import { useRouter } from "next/navigation"; +import { AlertTriangle, Check, Loader2, Trash2Icon, Info } from "lucide-react"; +import { BackButton } from "@/components/BackButton"; +import Title from "@/components/ui/title"; +import { EditableStringFieldDisplay } from "@/components/EditableStringFieldDisplay"; +import { + DropdownMenu, + DropdownMenuContent, + DropdownMenuTrigger, +} from "@/components/ui/dropdown-menu"; +import { DropdownMenuItemWithTooltip } from "@/components/ui/dropdown-menu-with-tooltip"; +import { FiSettings } from "react-icons/fi"; +import { usePopup } from "@/components/admin/connectors/Popup"; +import { + Tooltip, + TooltipContent, + TooltipProvider, + TooltipTrigger, +} from "@/components/ui/tooltip"; +import { Badge } from "@/components/ui/badge"; + +export interface FederatedConnectorFormProps { + connector: ConfigurableSources; + connectorId?: number; // Optional ID for editing existing connector + preloadedConnectorData?: FederatedConnectorDetail; + preloadedCredentialSchema?: CredentialSchemaResponse; +} + +interface CredentialForm { + [key: string]: string; +} + +interface FormState { + credentials: CredentialForm; + schema: Record | null; + schemaError: string | null; + connectorError: string | null; +} + +async function validateCredentials( + source: string, + credentials: CredentialForm +): Promise<{ success: boolean; message: string }> { + try { + const response = await fetch( + `/api/federated/sources/federated_${source}/credentials/validate`, + { + method: "POST", + headers: { + "Content-Type": "application/json", + }, + body: JSON.stringify(credentials), + } + ); + + if (!response.ok) { + const errorData = await response.json().catch(() => ({})); + return { + success: false, + message: + errorData.detail || `Validation failed: ${response.statusText}`, + }; + } + + const result = await response.json(); + return { + success: result, + message: result ? "Credentials are valid" : "Credentials are invalid", + }; + } catch (error) { + return { success: false, message: `Validation error: ${error}` }; + } +} + +async function createFederatedConnector( + source: string, + credentials: CredentialForm +): Promise<{ success: boolean; message: string }> { + try { + const response = await fetch("/api/federated", { + method: "POST", + headers: { + "Content-Type": "application/json", + }, + body: JSON.stringify({ + source: `federated_${source}`, + credentials, + } as FederatedConnectorCreateRequest), + }); + + if (response.ok) { + return { + success: true, + message: "Federated connector created successfully!", + }; + } else { + const errorData = await response.json(); + return { + success: false, + message: errorData.detail || "Failed to create federated connector", + }; + } + } catch (error) { + return { success: false, message: `Error: ${error}` }; + } +} + +async function updateFederatedConnector( + id: number, + credentials: CredentialForm +): Promise<{ success: boolean; message: string }> { + try { + const response = await fetch(`/api/federated/${id}`, { + method: "PUT", + headers: { + "Content-Type": "application/json", + }, + body: JSON.stringify({ + credentials, + }), + }); + + if (response.ok) { + return { + success: true, + message: "Federated connector updated successfully!", + }; + } else { + const errorData = await response.json(); + return { + success: false, + message: errorData.detail || "Failed to update federated connector", + }; + } + } catch (error) { + return { success: false, message: `Error: ${error}` }; + } +} + +async function deleteFederatedConnector( + id: number +): Promise<{ success: boolean; message: string }> { + try { + const response = await fetch(`/api/federated/${id}`, { + method: "DELETE", + }); + + if (response.ok) { + return { + success: true, + message: "Federated connector deleted successfully!", + }; + } else { + const errorData = await response.json(); + return { + success: false, + message: errorData.detail || "Failed to delete federated connector", + }; + } + } catch (error) { + return { success: false, message: `Error: ${error}` }; + } +} + +export function FederatedConnectorForm({ + connector, + connectorId, + preloadedConnectorData, + preloadedCredentialSchema, +}: FederatedConnectorFormProps) { + const router = useRouter(); + const sourceMetadata = getSourceMetadata(connector); + const isEditMode = connectorId !== undefined; + const { popup, setPopup } = usePopup(); + + const [formState, setFormState] = useState({ + credentials: preloadedConnectorData?.credentials || {}, + schema: preloadedCredentialSchema?.credentials || null, + schemaError: null, + connectorError: null, + }); + const [isSubmitting, setIsSubmitting] = useState(false); + const [submitMessage, setSubmitMessage] = useState(null); + const [submitSuccess, setSubmitSuccess] = useState(null); + const [isValidating, setIsValidating] = useState(false); + const [isDeleting, setIsDeleting] = useState(false); + const [isLoadingSchema, setIsLoadingSchema] = useState( + !preloadedCredentialSchema + ); + + // Fetch credential schema if not preloaded + useEffect(() => { + const fetchCredentialSchema = async () => { + if (!preloadedCredentialSchema) { + setIsLoadingSchema(true); + try { + const response = await fetch( + `/api/federated/sources/federated_${connector}/credentials/schema` + ); + + if (!response.ok) { + throw new Error( + `Failed to fetch credential schema: ${response.statusText}` + ); + } + + const schemaData: CredentialSchemaResponse = await response.json(); + setFormState((prev) => ({ + ...prev, + schema: schemaData.credentials, + schemaError: null, + })); + } catch (error) { + console.error("Error fetching credential schema:", error); + setFormState((prev) => ({ + ...prev, + schemaError: `Failed to load credential schema: ${error}`, + })); + } finally { + setIsLoadingSchema(false); + } + } + }; + + fetchCredentialSchema(); + }, [connector, preloadedCredentialSchema]); + + // Show loading state at the top level if schema is loading + if (isLoadingSchema) { + return ( +
+ {popup} +
+ +
+

+ Loading credential schema... +

+

+ Retrieving required fields for this connector type +

+
+
+
+ ); + } + + const handleCredentialChange = (key: string, value: string) => { + setFormState((prev) => ({ + ...prev, + credentials: { + ...prev.credentials, + [key]: value, + }, + })); + }; + + const handleValidateCredentials = async () => { + if (!formState.schema) return; + + setIsValidating(true); + setSubmitMessage(null); + setSubmitSuccess(null); + + try { + const result = await validateCredentials( + connector, + formState.credentials + ); + setSubmitMessage(result.message); + setSubmitSuccess(result.success); + } catch (error) { + setSubmitMessage(`Validation error: ${error}`); + setSubmitSuccess(false); + } finally { + setIsValidating(false); + } + }; + + const handleDeleteConnector = async () => { + if (!connectorId) return; + + const confirmed = window.confirm( + "Are you sure you want to delete this federated connector? This action cannot be undone." + ); + + if (!confirmed) return; + + setIsDeleting(true); + + try { + const result = await deleteFederatedConnector(connectorId); + + if (result.success) { + setPopup({ + message: result.message, + type: "success", + }); + // Redirect after a short delay + setTimeout(() => { + router.push("/admin/indexing/status"); + }, 500); + } else { + setPopup({ + message: result.message, + type: "error", + }); + } + } catch (error) { + setPopup({ + message: `Error deleting connector: ${error}`, + type: "error", + }); + } finally { + setIsDeleting(false); + } + }; + + const handleSubmit = async (e: React.FormEvent) => { + e.preventDefault(); + setIsSubmitting(true); + setSubmitMessage(null); + setSubmitSuccess(null); + + try { + // Validate required fields + if (formState.schema) { + const missingRequired = Object.entries(formState.schema) + .filter( + ([key, field]) => field.required && !formState.credentials[key] + ) + .map(([key]) => key); + + if (missingRequired.length > 0) { + setSubmitMessage( + `Missing required fields: ${missingRequired.join(", ")}` + ); + setSubmitSuccess(false); + setIsSubmitting(false); + return; + } + } + + // Validate credentials before creating/updating + const validation = await validateCredentials( + connector, + formState.credentials + ); + if (!validation.success) { + setSubmitMessage(`Credential validation failed: ${validation.message}`); + setSubmitSuccess(false); + setIsSubmitting(false); + return; + } + + // Create or update the connector + const result = + isEditMode && connectorId + ? await updateFederatedConnector(connectorId, formState.credentials) + : await createFederatedConnector(connector, formState.credentials); + + setSubmitMessage(result.message); + setSubmitSuccess(result.success); + setIsSubmitting(false); + + if (result.success) { + // Redirect after a short delay + setTimeout(() => { + router.push("/admin/indexing/status"); + }, 500); + } + } catch (error) { + setSubmitMessage(`Error: ${error}`); + setSubmitSuccess(false); + setIsSubmitting(false); + } + }; + + const renderCredentialFields = () => { + if (formState.schemaError) { + return ( +
+ + {formState.schemaError} +
+ ); + } + + if (formState.connectorError) { + return ( +
+ + {formState.connectorError} +
+ ); + } + + if (!formState.schema) { + return ( +
+ No credential schema available for this connector type. +
+ ); + } + + return ( + <> + {Object.entries(formState.schema).map(([fieldKey, fieldSpec]) => ( +
+ + handleCredentialChange(fieldKey, e.target.value)} + className="w-full" + required={fieldSpec.required} + /> + {fieldSpec.description && ( +

+ {fieldSpec.description} +

+ )} +
+ ))} + + ); + }; + + return ( +
+ {popup} + + +
+
+ +
+ +
+
+ + {isEditMode ? "Edit" : "Setup"} {sourceMetadata.displayName} + + + Federated + + + + + + + +

+ {sourceMetadata.federatedTooltip || + "This is a federated connector. It will result in greater latency and lower search quality compared to regular connectors."} +

+ + + +
+
+ + {isEditMode && ( +
+ + + + + + + + {isDeleting ? "Deleting..." : "Delete"} + + + +
+ )} +
+ + + Connector Configuration + + + + +
+ {renderCredentialFields()} + +
+ {submitMessage && ( +
+ {submitSuccess ? ( + + ) : ( + + )} + {submitMessage} +
+ )} + + + +
+
+ + +
+ ); +} diff --git a/web/src/components/chat/FederatedOAuthModal.tsx b/web/src/components/chat/FederatedOAuthModal.tsx new file mode 100644 index 0000000000..5fa0daa7b9 --- /dev/null +++ b/web/src/components/chat/FederatedOAuthModal.tsx @@ -0,0 +1,130 @@ +"use client"; + +import React, { useContext } from "react"; +import { Modal } from "@/components/Modal"; +import { Button } from "@/components/ui/button"; +import { SourceIcon } from "@/components/SourceIcon"; +import { ValidSources } from "@/lib/types"; +import { SettingsContext } from "@/components/settings/SettingsProvider"; +import { getSourceMetadata } from "@/lib/sources"; +import { useRouter } from "next/navigation"; + +export interface FederatedConnectorOAuthStatus { + federated_connector_id: number; + source: string; + name: string; + has_oauth_token: boolean; + oauth_token_expires_at?: string; + authorize_url?: string; +} + +interface FederatedOAuthModalProps { + connectors: FederatedConnectorOAuthStatus[]; + onSkip: () => void; + skipCount?: number; +} + +const MAX_SKIP_COUNT = 2; + +export function FederatedOAuthModal({ + connectors, + onSkip, + skipCount = 0, +}: FederatedOAuthModalProps) { + const settings = useContext(SettingsContext); + const needsAuth = connectors.filter((c) => !c.has_oauth_token); + const router = useRouter(); + + if (needsAuth.length === 0) { + return null; + } + + const handleAuthorize = (authorizeUrl: string) => { + // Redirect to OAuth URL in the same window + router.push(authorizeUrl); + }; + + const applicationName = + settings?.enterpriseSettings?.application_name || "Onyx"; + + if (skipCount >= MAX_SKIP_COUNT) { + return ( + {}} + hideCloseButton={true} + width="w-full max-w-xl" + > +
+
+

Heads Up!

+

+ You can always connect your apps later by going to the{" "} + User Settings menu (click your profile icon) and + selecting Connectors. +

+
+ +
+ +
+
+ + ); + } + + return ( + +
+

+ Improve answer quality by letting {applicationName} search all your + connected data. +

+ +
+ {needsAuth.map((connector) => { + const sourceMetadata = getSourceMetadata( + connector.source as ValidSources + ); + + return ( +
+
+ + + {sourceMetadata.displayName} + +
+ +
+ ); + })} +
+ + {/* Add visual separation and center modal actions */} +
+
+ +
+
+
+ + ); +} diff --git a/web/src/components/filters/SourceSelector.tsx b/web/src/components/filters/SourceSelector.tsx index 4059190769..a5d1c4ccea 100644 --- a/web/src/components/filters/SourceSelector.tsx +++ b/web/src/components/filters/SourceSelector.tsx @@ -15,19 +15,13 @@ import { SearchDateRangeSelector } from "@/components/dateRangeSelectors/SearchD import { DateRangePickerValue } from "@/components/dateRangeSelectors/AdminDateRangeSelector"; import { listSourceMetadata } from "@/lib/sources"; import { SourceIcon } from "@/components/SourceIcon"; -import { TagFilter } from "@/components/search/filtering/TagFilter"; import { Calendar } from "@/components/ui/calendar"; import { Popover, PopoverTrigger } from "@/components/ui/popover"; import { PopoverContent } from "@radix-ui/react-popover"; import { CalendarIcon } from "lucide-react"; import { getTimeAgoString } from "@/lib/dateUtils"; -import { Separator } from "@/components/ui/separator"; import { FilterDropdown } from "@/components/search/filtering/FilterDropdown"; -const SectionTitle = ({ children }: { children: string }) => ( -
{children}
-); - export interface SourceSelectorProps { timeRange: DateRangePickerValue | null; setTimeRange: React.Dispatch< @@ -48,227 +42,6 @@ export interface SourceSelectorProps { tagsOnLeft: boolean; } -export function SourceSelector({ - timeRange, - setTimeRange, - selectedSources, - setSelectedSources, - selectedDocumentSets, - setSelectedDocumentSets, - selectedTags, - setSelectedTags, - availableDocumentSets, - existingSources, - availableTags, - showDocSidebar, - toggleFilters, - filtersUntoggled, - tagsOnLeft, -}: SourceSelectorProps) { - const handleSelect = (source: SourceMetadata) => { - setSelectedSources((prev: SourceMetadata[]) => { - if ( - prev.map((source) => source.internalName).includes(source.internalName) - ) { - return prev.filter((s) => s.internalName !== source.internalName); - } else { - return [...prev, source]; - } - }); - }; - - const handleDocumentSetSelect = (documentSetName: string) => { - setSelectedDocumentSets((prev: string[]) => { - if (prev.includes(documentSetName)) { - return prev.filter((s) => s !== documentSetName); - } else { - return [...prev, documentSetName]; - } - }); - }; - - let allSourcesSelected = selectedSources.length > 0; - - const toggleAllSources = () => { - if (allSourcesSelected) { - setSelectedSources([]); - } else { - const allSources = listSourceMetadata().filter((source) => - existingSources.includes(source.internalName) - ); - setSelectedSources(allSources); - } - }; - - return ( -
- - {!filtersUntoggled && ( - <> - - - -
-
- Time Range - {true && ( - - )} -
-

- {getTimeAgoString(timeRange?.from!) || "Select a time range"} -

-
- - - { - const initialDate = daterange?.from || new Date(); - const endDate = daterange?.to || new Date(); - setTimeRange({ - from: initialDate, - to: endDate, - selectValue: timeRange?.selectValue || "", - }); - }} - className="rounded-md " - /> - - - - {availableTags.length > 0 && ( - <> -
- Tags -
- - - )} - - {existingSources.length > 0 && ( -
-
-
-

Sources

- -
-
-
- {listSourceMetadata() - .filter((source) => - existingSources.includes(source.internalName) - ) - .map((source) => ( -
source.internalName) - .includes(source.internalName) - ? "bg-accent-background-hovered" - : "hover:bg-accent-background") - } - onClick={() => handleSelect(source)} - > - - - {source.displayName} - -
- ))} -
-
- )} - - {availableDocumentSets.length > 0 && ( - <> -
- Knowledge Sets -
-
- {availableDocumentSets.map((documentSet) => ( -
-
handleDocumentSetSelect(documentSet.name)} - > - - -
- } - popupContent={ -
-
Description
-
- {documentSet.description} -
-
- } - classNameModifications="-ml-2" - /> - {documentSet.name} -
-
- ))} -
- - )} - - )} -
- ); -} - export function SelectedBubble({ children, onClick, @@ -343,7 +116,10 @@ export function HorizontalFilters({ key: source.displayName, display: ( <> - + {source.displayName} ), @@ -404,7 +180,10 @@ export function HorizontalFilters({ onClick={() => handleSourceSelect(source)} > <> - + {source.displayName} @@ -554,7 +333,10 @@ export function HorizontalSourceSelector({ key: source.internalName, display: ( <> - + {source.displayName} ), diff --git a/web/src/components/search/filtering/FilterPopup.tsx b/web/src/components/search/filtering/FilterPopup.tsx index 186ea14cba..514e1a0794 100644 --- a/web/src/components/search/filtering/FilterPopup.tsx +++ b/web/src/components/search/filtering/FilterPopup.tsx @@ -11,6 +11,7 @@ import { FiChevronRight, FiDatabase, FiBook, + FiChevronDown, } from "react-icons/fi"; import { FilterManager } from "@/lib/hooks"; import { DocumentSet, Tag } from "@/lib/types"; @@ -21,10 +22,20 @@ import { Button } from "@/components/ui/button"; import { SourceIcon } from "@/components/SourceIcon"; import { SelectableDropdown, TagFilter } from "./TagFilter"; import { Input } from "@/components/ui/input"; +import { + Tooltip, + TooltipContent, + TooltipProvider, + TooltipTrigger, +} from "@/components/ui/tooltip"; interface FilterPopupProps { filterManager: FilterManager; - trigger: React.ReactNode; + trigger: { + name: string; + Icon: React.ComponentType<{ size?: number; className?: string }>; + tooltipContent?: React.ReactNode; + }; availableSources: SourceMetadata[]; availableDocumentSets: DocumentSet[]; availableTags: Tag[]; @@ -252,7 +263,34 @@ export function FilterPopup({ return ( - + } diff --git a/web/src/components/search/filtering/HorizontalSourceSelector.tsx b/web/src/components/search/filtering/HorizontalSourceSelector.tsx index 7d1209ac4c..b0cefadc6d 100644 --- a/web/src/components/search/filtering/HorizontalSourceSelector.tsx +++ b/web/src/components/search/filtering/HorizontalSourceSelector.tsx @@ -150,7 +150,10 @@ export function HorizontalSourceSelector({ .filter((source) => existingSources.includes(source.internalName)) .map((source) => ({ icon: ( - + ), key: source.internalName, display: ( diff --git a/web/src/lib/connectors/credentials.ts b/web/src/lib/connectors/credentials.ts index b8129dd1c2..99c6713508 100644 --- a/web/src/lib/connectors/credentials.ts +++ b/web/src/lib/connectors/credentials.ts @@ -386,6 +386,7 @@ export const credentialTemplates: Record = { web: null, not_applicable: null, ingestion_api: null, + federated_slack: null, discord: { discord_bot_token: "" } as DiscordCredentialJson, // NOTE: These are Special Cases diff --git a/web/src/lib/hooks.ts b/web/src/lib/hooks.ts index acc4833ddb..23343fbe85 100644 --- a/web/src/lib/hooks.ts +++ b/web/src/lib/hooks.ts @@ -6,7 +6,7 @@ import { UserGroup, ConnectorStatus, CCPairBasicInfo, - ValidSources, + FederatedConnectorDetail, } from "@/lib/types"; import useSWR, { mutate, useSWRConfig } from "swr"; import { errorHandlingFetcher } from "./fetcher"; @@ -121,6 +121,20 @@ export const useBasicConnectorStatus = () => { }; }; +export const useFederatedConnectors = () => { + const { mutate } = useSWRConfig(); + const url = "/api/federated"; + const swrResponse = useSWR( + url, + errorHandlingFetcher + ); + + return { + ...swrResponse, + refreshFederatedConnectors: () => mutate(url), + }; +}; + export const useLabels = () => { const { mutate } = useSWRConfig(); const { data: labels, error } = useSWR( @@ -210,7 +224,7 @@ export interface FilterManager { getFilterString: () => string; buildFiltersFromQueryString: ( filterString: string, - availableSources: ValidSources[], + availableSources: SourceMetadata[], availableDocumentSets: string[], availableTags: Tag[] ) => void; @@ -267,7 +281,7 @@ export function useFilters(): FilterManager { function buildFiltersFromQueryString( filterString: string, - availableSources: ValidSources[], + availableSources: SourceMetadata[], availableDocumentSets: string[], availableTags: Tag[] ): void { @@ -286,12 +300,11 @@ export function useFilters(): FilterManager { } // Parse sources - const availableSourcesMetadata = availableSources.map(getSourceMetadata); let newSelectedSources: SourceMetadata[] = []; const sourcesParam = params.get("sources"); if (sourcesParam) { const sourceNames = sourcesParam.split(",").map(decodeURIComponent); - newSelectedSources = availableSourcesMetadata.filter((source) => + newSelectedSources = availableSources.filter((source) => sourceNames.includes(source.internalName) ); } diff --git a/web/src/lib/hooks/useFederatedOAuthStatus.ts b/web/src/lib/hooks/useFederatedOAuthStatus.ts new file mode 100644 index 0000000000..0599d01255 --- /dev/null +++ b/web/src/lib/hooks/useFederatedOAuthStatus.ts @@ -0,0 +1,44 @@ +import { useState, useEffect } from "react"; +import { FederatedConnectorOAuthStatus } from "@/components/chat/FederatedOAuthModal"; + +export function useFederatedOAuthStatus() { + const [connectors, setConnectors] = useState( + [] + ); + const [loading, setLoading] = useState(true); + const [error, setError] = useState(null); + + const fetchOAuthStatus = async () => { + try { + setLoading(true); + const response = await fetch("/api/federated/oauth-status"); + + if (!response.ok) { + throw new Error("Failed to fetch OAuth status"); + } + + const data = await response.json(); + setConnectors(data); + } catch (err) { + setError(err instanceof Error ? err.message : "An error occurred"); + } finally { + setLoading(false); + } + }; + + useEffect(() => { + fetchOAuthStatus(); + }, []); + + const needsAuth = connectors.filter((c) => !c.has_oauth_token); + const hasUnauthenticatedConnectors = needsAuth.length > 0; + + return { + connectors, + needsAuth, + hasUnauthenticatedConnectors, + loading, + error, + refetch: fetchOAuthStatus, + }; +} diff --git a/web/src/lib/oauth_utils.ts b/web/src/lib/oauth_utils.ts index db3342efed..2236fa45f0 100644 --- a/web/src/lib/oauth_utils.ts +++ b/web/src/lib/oauth_utils.ts @@ -62,6 +62,56 @@ export async function handleOAuthAuthorizationResponse( return; } +// Handler for federated connector OAuth callbacks +export async function handleFederatedOAuthCallback( + federatedConnectorId: string, + code: string, + state: string +): Promise { + // Use the generic callback endpoint - the connector ID will be extracted from the state parameter + const url = `/api/federated/callback?code=${encodeURIComponent( + code + )}&state=${encodeURIComponent(state)}`; + + const response = await fetch(url, { + method: "POST", + headers: { + "Content-Type": "application/json", + }, + }); + + if (!response.ok) { + let errorDetails = `Failed to handle federated OAuth callback: ${response.status}`; + + try { + const responseBody = await response.text(); + errorDetails += `\nResponse Body: ${responseBody}`; + } catch (err) { + if (err instanceof Error) { + errorDetails += `\nUnable to read response body: ${err.message}`; + } else { + errorDetails += `\nUnable to read response body: Unknown error type`; + } + } + + throw new Error(errorDetails); + } + + // Parse the JSON response and extract the data field + const result = await response.json(); + + if (!result.success) { + throw new Error(result.message || "OAuth callback failed"); + } + + return { + success: true, + message: result.message || "OAuth authorization successful", + redirect_on_success: `/admin/federated/${federatedConnectorId}`, + finalize_url: null, + }; +} + // server side handler to process the oauth redirect callback // https://api.slack.com/authentication/oauth-v2#exchanging export async function handleOAuthSlackAuthorizationResponse( diff --git a/web/src/lib/search/interfaces.ts b/web/src/lib/search/interfaces.ts index faf2e2579f..1fde64f8d9 100644 --- a/web/src/lib/search/interfaces.ts +++ b/web/src/lib/search/interfaces.ts @@ -180,6 +180,10 @@ export interface SourceMetadata { internalName: ValidSources; adminUrl: string; oauthSupported?: boolean; + federated?: boolean; + federatedTooltip?: string; + // For federated connectors, this stores the base source type for the icon + baseSourceType?: ValidSources; } export interface SearchDefaultOverrides { diff --git a/web/src/lib/sources.ts b/web/src/lib/sources.ts index 32249535bc..ce0bd90da6 100644 --- a/web/src/lib/sources.ts +++ b/web/src/lib/sources.ts @@ -54,10 +54,29 @@ interface PartialSourceMetadata { displayName: string; category: SourceCategory; docs?: string; + oauthSupported?: boolean; + federated?: boolean; + federatedTooltip?: string; + // federated connectors store the base source type if it's a source + // that has both indexed connectors and federated connectors + baseSourceType?: ValidSources; } type SourceMap = { - [K in ValidSources]: PartialSourceMetadata; + [K in ValidSources | "federated_slack"]: PartialSourceMetadata; +}; + +const slackMetadata = { + icon: ColorSlackIcon, + displayName: "Slack", + category: SourceCategory.Messaging, + docs: "https://docs.onyx.app/connectors/slack", + oauthSupported: true, + federated: true, + federatedTooltip: + "⚠️ WARNING: Due to Slack's rate limit and ToS changes, Slack is now federated. " + + "This will result in significantly greater latency and lower search quality.", + baseSourceType: "slack", }; export const SOURCE_METADATA_MAP: SourceMap = { @@ -73,13 +92,8 @@ export const SOURCE_METADATA_MAP: SourceMap = { category: SourceCategory.Storage, docs: "https://docs.onyx.app/connectors/file", }, - slack: { - icon: ColorSlackIcon, - displayName: "Slack", - category: SourceCategory.Messaging, - docs: "https://docs.onyx.app/connectors/slack", - oauthSupported: true, - }, + slack: slackMetadata, + federated_slack: slackMetadata, discord: { icon: ColorDiscordIcon, displayName: "Discord", @@ -352,7 +366,7 @@ function fillSourceMetadata( internalName: ValidSources ): SourceMetadata { return { - internalName: internalName, + internalName: partialMetadata.baseSourceType || internalName, ...partialMetadata, adminUrl: `/admin/connectors/${internalName}`, }; @@ -375,7 +389,9 @@ export function listSourceMetadata(): SourceMetadata[] { ([source, _]) => source !== "not_applicable" && source !== "ingestion_api" && - source !== "mock_connector" + source !== "mock_connector" && + // use the "regular" slack connector when listing + source !== "federated_slack" ) .map(([source, metadata]) => { return fillSourceMetadata(metadata, source as ValidSources); diff --git a/web/src/lib/types.ts b/web/src/lib/types.ts index 3d9b2b4017..4f94ab2122 100644 --- a/web/src/lib/types.ts +++ b/web/src/lib/types.ts @@ -175,6 +175,20 @@ export interface ConnectorIndexingStatus< docs_indexed: number; } +export interface FederatedConnectorDetail { + id: number; + source: ValidSources.FederatedSlack; + name: string; + credentials: Record; + oauth_token_exists: boolean; + oauth_token_expires_at: string | null; + document_sets: Array<{ + id: number; + name: string; + entities: Record; + }>; +} + export interface OAuthPrepareAuthorizationResponse { url: string; } @@ -243,6 +257,18 @@ export interface CCPairDescriptor { access_type: AccessType; } +export interface FederatedConnectorConfig { + federated_connector_id: number; + entities: Record; +} + +export interface FederatedConnectorDescriptor { + id: number; + name: string; + source: string; + entities: Record; +} + export interface DocumentSet { id: number; name: string; @@ -252,6 +278,7 @@ export interface DocumentSet { is_public: boolean; users: string[]; groups: number[]; + federated_connectors: FederatedConnectorDescriptor[]; } export interface Tag { @@ -392,8 +419,20 @@ export enum ValidSources { Airtable = "airtable", Gitbook = "gitbook", Highspot = "highspot", + + // Federated Connectors + FederatedSlack = "federated_slack", } +export const federatedSourceToRegularSource = ( + maybeFederatedSource: ValidSources +): ValidSources => { + if (maybeFederatedSource === ValidSources.FederatedSlack) { + return ValidSources.Slack; + } + return maybeFederatedSource; +}; + export const validAutoSyncSources = [ ValidSources.Confluence, ValidSources.GoogleDrive, @@ -407,7 +446,9 @@ export type ValidAutoSyncSource = (typeof validAutoSyncSources)[number]; export type ConfigurableSources = Exclude< ValidSources, - ValidSources.NotApplicable | ValidSources.IngestionApi + | ValidSources.NotApplicable + | ValidSources.IngestionApi + | ValidSources.FederatedSlack // is part of ValiedSources.Slack >; export const oauthSupportedSources: ConfigurableSources[] = [ @@ -418,3 +459,27 @@ export const oauthSupportedSources: ConfigurableSources[] = [ ]; export type OAuthSupportedSource = (typeof oauthSupportedSources)[number]; + +// Federated Connector Types +export interface CredentialFieldSpec { + type: string; + description: string; + required: boolean; + default?: any; + example?: any; + secret: boolean; +} + +export interface CredentialSchemaResponse { + credentials: Record; +} + +export interface FederatedConnectorCreateRequest { + source: string; + credentials: Record; +} + +export interface FederatedConnectorCreateResponse { + id: number; + source: string; +}