feat(slack federated search scoping - 3/4): Add connector-level config support (#6178)

Co-authored-by: Nikolas Garza <nikolas@Nikolass-MacBook-Pro.local>

Author: nmgarza5

backend/onyx/db/federated.py

@@ -63,17 +63,31 @@ def create_federated_connector(
     db_session: Session,
     source: FederatedConnectorSource,
     credentials: dict[str, Any],
+    config: dict[str, Any] | None = None,
 ) -> FederatedConnector:
-    """Create a new federated connector with credential validation."""
+    """Create a new federated connector with credential and config validation."""
     # Validate credentials before creating
     if not validate_federated_connector_credentials(source, credentials):
         raise ValueError(
             f"Invalid credentials for federated connector source: {source}"
         )
 
+    # Validate config using connector-specific validation
+    if config:
+        try:
+            # Get connector instance to access validate_config method
+            connector = get_federated_connector(source, credentials)
+            if not connector.validate_config(config):
+                raise ValueError(
+                    f"Invalid config for federated connector source: {source}"
+                )
+        except Exception as e:
+            raise ValueError(f"Config validation failed for {source}: {str(e)}")
+
     federated_connector = FederatedConnector(
         source=source,
         credentials=credentials,
+        config=config or {},
     )
     db_session.add(federated_connector)
     db_session.commit()
@@ -239,14 +253,21 @@ def update_federated_connector(
     db_session: Session,
     federated_connector_id: int,
     credentials: dict[str, Any] | None = None,
+    config: dict[str, Any] | None = None,
 ) -> FederatedConnector | None:
-    """Update a federated connector with credential validation."""
+    """Update a federated connector with credential and config validation."""
     federated_connector = fetch_federated_connector_by_id(
         federated_connector_id, db_session
     )
     if not federated_connector:
         return None
 
+    # Use provided credentials if updating them, otherwise use existing credentials
+    # This is needed to instantiate the connector for config validation when only config is being updated
+    creds_to_use = (
+        credentials if credentials is not None else federated_connector.credentials
+    )
+
     if credentials is not None:
         # Validate credentials before updating
         if not validate_federated_connector_credentials(
@@ -257,6 +278,23 @@ def update_federated_connector(
             )
         federated_connector.credentials = credentials
 
+    if config is not None:
+        # Validate config using connector-specific validation
+        try:
+            # Get connector instance to access validate_config method
+            connector = get_federated_connector(
+                federated_connector.source, creds_to_use
+            )
+            if not connector.validate_config(config):
+                raise ValueError(
+                    f"Invalid config for federated connector source: {federated_connector.source}"
+                )
+        except Exception as e:
+            raise ValueError(
+                f"Config validation failed for {federated_connector.source}: {str(e)}"
+            )
+        federated_connector.config = config
+
     db_session.commit()
     return federated_connector
 

backend/onyx/federated_connectors/federated_retrieval.py

@@ -1,12 +1,12 @@
 from collections import defaultdict
 from collections.abc import Callable
+from typing import Any
 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
@@ -18,6 +18,7 @@
 from onyx.db.models import FederatedConnector__DocumentSet
 from onyx.db.slack_bot import fetch_slack_bots
 from onyx.federated_connectors.factory import get_federated_connector
+from onyx.federated_connectors.interfaces import FederatedConnector
 from onyx.onyxbot.slack.models import SlackContext
 from onyx.utils.logger import setup_logger
 
@@ -86,15 +87,31 @@ def get_federated_retrieval_functions(
                     credentials,
                 )
 
-                federated_retrieval_infos_slack.append(
-                    FederatedRetrievalInfo(
-                        retrieval_function=lambda query: connector.search(
+                # Capture variables by value to avoid lambda closure issues
+                bot_token = tenant_slack_bot.bot_token
+
+                def create_slack_retrieval_function(
+                    conn: FederatedConnector,
+                    token: str,
+                    ctx: SlackContext,
+                    bot_tok: str,
+                ) -> Callable[[SearchQuery], list[InferenceChunk]]:
+                    def retrieval_fn(query: SearchQuery) -> list[InferenceChunk]:
+                        return conn.search(
                             query,
                             {},  # Empty entities for Slack context
-                            access_token=access_token,
-                            limit=MAX_FEDERATED_CHUNKS,
-                            slack_event_context=slack_context,
-                            bot_token=tenant_slack_bot.bot_token,
+                            access_token=token,
+                            limit=None,  # Let connector use its own max_messages_per_query config
+                            slack_event_context=ctx,
+                            bot_token=bot_tok,
+                        )
+
+                    return retrieval_fn
+
+                federated_retrieval_infos_slack.append(
+                    FederatedRetrievalInfo(
+                        retrieval_function=create_slack_retrieval_function(
+                            connector, access_token, slack_context, bot_token
                         ),
                         source=FederatedConnectorSource.FEDERATED_SLACK,
                     )
@@ -158,22 +175,33 @@ def get_federated_retrieval_functions(
         if document_set_names and not document_set_associations:
             continue
 
-        if document_set_associations:
-            entities = document_set_associations[0].entities
-        else:
-            entities = {}
+        # Only use connector-level config (no junction table entities)
+        entities = oauth_token.federated_connector.config or {}
 
         connector = get_federated_connector(
             oauth_token.federated_connector.source,
             oauth_token.federated_connector.credentials,
         )
+
+        # Capture variables by value to avoid lambda closure issues
+        access_token = oauth_token.token
+
+        def create_retrieval_function(
+            conn: FederatedConnector,
+            ent: dict[str, Any],
+            token: str,
+        ) -> Callable[[SearchQuery], list[InferenceChunk]]:
+            return lambda query: conn.search(
+                query,
+                ent,
+                access_token=token,
+                limit=None,  # Let connector use its own max_messages_per_query config
+            )
+
         federated_retrieval_infos.append(
             FederatedRetrievalInfo(
-                retrieval_function=lambda query: connector.search(
-                    query,
-                    entities,
-                    access_token=oauth_token.token,
-                    limit=MAX_FEDERATED_CHUNKS,
+                retrieval_function=create_retrieval_function(
+                    connector, entities, access_token
                 ),
                 source=oauth_token.federated_connector.source,
             )

backend/onyx/federated_connectors/interfaces.py

@@ -34,16 +34,34 @@ def validate_entities(self, entities: Dict[str, Any]) -> bool:
 
         Returns:
             True if entities are valid, False otherwise
+
+        Note: This method is used for backward compatibility with document-set level entities.
+        For connector-level config validation, use validate_config() instead.
+        """
+
+    def validate_config(self, config: Dict[str, Any]) -> bool:
+        """
+        Validate that the provided config matches the expected structure.
+
+        This is an alias for validate_entities() to provide clearer semantics
+        when validating connector-level configuration.
+
+        Args:
+            config: Dictionary of configuration to validate
+
+        Returns:
+            True if config is valid, False otherwise
         """
+        return self.validate_entities(config)
 
     @classmethod
     @abstractmethod
-    def entities_schema(cls) -> Dict[str, EntityField]:
+    def configuration_schema(cls) -> Dict[str, EntityField]:
         """
-        Return the specification of what entities are available for this connector.
+        Return the specification of what configuration fields are available for this connector.
 
         Returns:
-            Dictionary where keys are entity names and values are EntityField objects
+            Dictionary where keys are configuration field names and values are EntityField objects
             describing the expected structure and constraints.
         """
 
@@ -96,7 +114,7 @@ def search(
 
         Args:
             query: The search query
-            entities: The entities to search within (validated by validate())
+            entities: Connector-level config (entity filtering configuration)
             access_token: The OAuth access token
             limit: Maximum number of results to return
             slack_event_context: Slack-specific context (only used by Slack bot)

backend/onyx/federated_connectors/slack/federated_connector.py

@@ -281,7 +281,7 @@ def search(
 
         Args:
             query: The search query
-            entities: The entities to search within (validated by validate())
+            entities: Connector-level config (entity filtering configuration)
             access_token: The OAuth access token
             limit: Maximum number of results to return
             slack_event_context: Optional Slack context for slack bot
@@ -311,9 +311,9 @@ def search(
                 query,
                 access_token,
                 db_session,
+                entities=entities,
                 limit=limit,
                 slack_event_context=slack_event_context,
                 bot_token=bot_token,
-                entities=entities,
                 team_id=team_id,
             )

backend/onyx/server/federated/api.py

@@ -31,6 +31,7 @@
 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 ConfigurationSchemaResponse
 from onyx.server.federated.models import CredentialSchemaResponse
 from onyx.server.federated.models import EntitySpecResponse
 from onyx.server.federated.models import FederatedConnectorCredentials
@@ -86,6 +87,7 @@ def create_federated_connector(
             db_session=db_session,
             source=federated_connector_data.source,
             credentials=federated_connector_data.credentials.model_dump(),
+            config=federated_connector_data.config,
         )
 
         logger.info(
@@ -122,7 +124,7 @@ def get_entities(
         connector_instance = _get_federated_connector_instance(
             federated_connector.source, federated_connector.credentials
         )
-        entities_spec = connector_instance.entities_schema()
+        entities_spec = connector_instance.configuration_schema()
 
         # Convert EntityField objects to a dictionary format for the API response
         entities_dict = {}
@@ -184,6 +186,34 @@ def get_credentials_schema(
         raise HTTPException(status_code=500, detail=str(e))
 
 
+@router.get("/sources/{source}/configuration/schema")
+def get_configuration_schema_by_source(
+    source: FederatedConnectorSource,
+    _: User = Depends(current_curator_or_admin_user),
+) -> ConfigurationSchemaResponse:
+    """Fetch configuration schema for a specific source type (for setup/edit forms)"""
+    try:
+        connector_cls = get_federated_connector_cls(source)
+        entities_spec = connector_cls.configuration_schema()
+
+        # Convert EntityField objects to a dictionary format for the API response
+        configuration_dict = {}
+        for key, field in entities_spec.items():
+            configuration_dict[key] = {
+                "type": field.type,
+                "description": field.description,
+                "required": field.required,
+                "default": field.default,
+                "example": field.example,
+            }
+
+        return ConfigurationSchemaResponse(configuration=configuration_dict)
+
+    except Exception as e:
+        logger.error(f"Error fetching configuration schema for source {source}: {e}")
+        raise HTTPException(status_code=500, detail=str(e))
+
+
 @router.get("/sources/{source}/credentials/schema")
 def get_credentials_schema_by_source(
     source: FederatedConnectorSource,
@@ -498,6 +528,7 @@ def get_federated_connector_detail(
         source=federated_connector.source,
         name=f"{federated_connector.source.replace('_', ' ').title()}",
         credentials=FederatedConnectorCredentials(**federated_connector.credentials),
+        config=federated_connector.config,
         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,
@@ -522,6 +553,7 @@ def update_federated_connector_endpoint(
                 if update_request.credentials
                 else None
             ),
+            config=update_request.config,
         )
 
         if not updated_connector:

backend/onyx/server/federated/models.py

@@ -18,6 +18,7 @@ class FederatedConnectorCredentials(BaseModel):
 class FederatedConnectorRequest(BaseModel):
     source: FederatedConnectorSource
     credentials: FederatedConnectorCredentials
+    config: dict[str, Any] = Field(default_factory=dict)
 
 
 class FederatedConnectorResponse(BaseModel):
@@ -60,6 +61,7 @@ class FederatedConnectorDetail(BaseModel):
     source: FederatedConnectorSource
     name: str
     credentials: FederatedConnectorCredentials
+    config: dict[str, Any] = Field(default_factory=dict)
     oauth_token_exists: bool
     oauth_token_expires_at: datetime | None = None
     document_sets: list[dict[str, Any]] = Field(default_factory=list)
@@ -87,6 +89,7 @@ def from_federated_connector_detail(
 
 class FederatedConnectorUpdateRequest(BaseModel):
     credentials: FederatedConnectorCredentials | None = None
+    config: dict[str, Any] | None = None
 
 
 class EntitySpecResponse(BaseModel):
@@ -95,6 +98,12 @@ class EntitySpecResponse(BaseModel):
     entities: dict[str, Any]
 
 
+class ConfigurationSchemaResponse(BaseModel):
+    """Response for configuration schema specification"""
+
+    configuration: dict[str, Any]
+
+
 class CredentialSchemaResponse(BaseModel):
     """Response for credential schema specification"""
 

web/src/lib/types.ts

@@ -560,6 +560,10 @@ export interface CredentialSchemaResponse {
   credentials: Record<string, CredentialFieldSpec>;
 }
 
+export interface ConfigurationSchemaResponse {
+  configuration: Record<string, CredentialFieldSpec>;
+}
+
 export interface FederatedConnectorCreateRequest {
   source: string;
   credentials: Record<string, any>;