Compare commits

..

18 Commits

Author SHA1 Message Date
pablonyx
ddaaaeeb40 k 2025-02-27 20:28:49 -08:00
pablonyx
4e31bc19dc k 2025-02-27 17:02:11 -08:00
pablonyx
1c0c80116e fix function 2025-02-27 17:00:05 -08:00
pablonyx
27891ad34d k 2025-02-27 16:51:06 -08:00
pablonyx
511341fd7c k 2025-02-27 16:49:23 -08:00
pablonyx
88d4e7defa fix handling 2025-02-27 16:14:02 -08:00
pablonyx
a7ba0da8cc Lowercase multi tenant email mapping (#4141) 2025-02-27 15:33:40 -08:00
Richard Kuo (Danswer)
aaced6d551 scan images 2025-02-27 15:25:29 -08:00
Richard Kuo (Danswer)
4c230f92ea trivy test 2025-02-27 15:05:03 -08:00
Richard Kuo (Danswer)
07d75b04d1 enable trivy scan 2025-02-27 14:22:44 -08:00
evan-danswer
a8d10750c1 fix propagation of is_agentic (#4150) 2025-02-27 11:56:51 -08:00
pablonyx
85e3ed57f1 Order chat sessions by time updated, not created (#4143)
* order chat sessions by time updated, not created

* quick update

* k
2025-02-27 17:35:42 +00:00
pablonyx
e10cc8ccdb Multi tenant user google auth fix (#4145) 2025-02-27 10:35:38 -08:00
pablonyx
7018bc974b Better looking errors (#4050)
* add error handling

* fix

* k
2025-02-27 04:58:25 +00:00
pablonyx
9c9075d71d Minor improvements to provisioning (#4109)
* quick fix

* k

* nit
2025-02-27 04:57:31 +00:00
pablonyx
338e084062 Improved tenant handling for slack bot (#4099) 2025-02-27 04:06:26 +00:00
pablonyx
2f64031f5c Improved tenant handling for slack bot1 (#4104) 2025-02-27 03:40:50 +00:00
pablonyx
abb74f2eaa Improved chat search (#4137)
* functional + fast

* k

* adapt

* k

* nit

* k

* k

* fix typing

* k
2025-02-27 02:27:45 +00:00
45 changed files with 1245 additions and 568 deletions

View File

@@ -53,24 +53,90 @@ jobs:
exclude: '(?i)^(pylint|aio[-_]*).*'
- name: Print report
if: ${{ always() }}
if: always()
run: echo "${{ steps.license_check_report.outputs.report }}"
- name: Install npm dependencies
working-directory: ./web
run: npm ci
- name: Run Trivy vulnerability scanner in repo mode
uses: aquasecurity/trivy-action@0.28.0
with:
scan-type: fs
scanners: license
format: table
# format: sarif
# output: trivy-results.sarif
severity: HIGH,CRITICAL
# - name: Upload Trivy scan results to GitHub Security tab
# uses: github/codeql-action/upload-sarif@v3
# be careful enabling the sarif and upload as it may spam the security tab
# with a huge amount of items. Work out the issues before enabling upload.
# - name: Run Trivy vulnerability scanner in repo mode
# if: always()
# uses: aquasecurity/trivy-action@0.29.0
# with:
# sarif_file: trivy-results.sarif
# scan-type: fs
# scan-ref: .
# scanners: license
# format: table
# severity: HIGH,CRITICAL
# # format: sarif
# # output: trivy-results.sarif
#
# # - name: Upload Trivy scan results to GitHub Security tab
# # uses: github/codeql-action/upload-sarif@v3
# # with:
# # sarif_file: trivy-results.sarif
scan-trivy:
# See https://runs-on.com/runners/linux/
runs-on: [runs-on,runner=2cpu-linux-x64,"run-id=${{ github.run_id }}"]
steps:
- name: Set up Docker Buildx
uses: docker/setup-buildx-action@v3
- name: Login to Docker Hub
uses: docker/login-action@v3
with:
username: ${{ secrets.DOCKER_USERNAME }}
password: ${{ secrets.DOCKER_TOKEN }}
# Backend
- name: Pull backend docker image
run: docker pull onyxdotapp/onyx-backend:latest
- name: Run Trivy vulnerability scanner on backend
uses: aquasecurity/trivy-action@0.29.0
env:
TRIVY_DB_REPOSITORY: 'public.ecr.aws/aquasecurity/trivy-db:2'
TRIVY_JAVA_DB_REPOSITORY: 'public.ecr.aws/aquasecurity/trivy-java-db:1'
with:
image-ref: onyxdotapp/onyx-backend:latest
scanners: license
severity: HIGH,CRITICAL
vuln-type: library
exit-code: 0 # Set to 1 if we want a failed scan to fail the workflow
# Web server
- name: Pull web server docker image
run: docker pull onyxdotapp/onyx-web-server:latest
- name: Run Trivy vulnerability scanner on web server
uses: aquasecurity/trivy-action@0.29.0
env:
TRIVY_DB_REPOSITORY: 'public.ecr.aws/aquasecurity/trivy-db:2'
TRIVY_JAVA_DB_REPOSITORY: 'public.ecr.aws/aquasecurity/trivy-java-db:1'
with:
image-ref: onyxdotapp/onyx-web-server:latest
scanners: license
severity: HIGH,CRITICAL
vuln-type: library
exit-code: 0
# Model server
- name: Pull model server docker image
run: docker pull onyxdotapp/onyx-model-server:latest
- name: Run Trivy vulnerability scanner
uses: aquasecurity/trivy-action@0.29.0
env:
TRIVY_DB_REPOSITORY: 'public.ecr.aws/aquasecurity/trivy-db:2'
TRIVY_JAVA_DB_REPOSITORY: 'public.ecr.aws/aquasecurity/trivy-java-db:1'
with:
image-ref: onyxdotapp/onyx-model-server:latest
scanners: license
severity: HIGH,CRITICAL
vuln-type: library
exit-code: 0

View File

@@ -0,0 +1,84 @@
"""improved index
Revision ID: 3bd4c84fe72f
Revises: 8f43500ee275
Create Date: 2025-02-26 13:07:56.217791
"""
from alembic import op
# revision identifiers, used by Alembic.
revision = "3bd4c84fe72f"
down_revision = "8f43500ee275"
branch_labels = None
depends_on = None
# NOTE:
# This migration addresses issues with the previous migration (8f43500ee275) which caused
# an outage by creating an index without using CONCURRENTLY. This migration:
#
# 1. Creates more efficient full-text search capabilities using tsvector columns and GIN indexes
# 2. Uses CONCURRENTLY for all index creation to prevent table locking
# 3. Explicitly manages transactions with COMMIT statements to allow CONCURRENTLY to work
# (see: https://www.postgresql.org/docs/9.4/sql-createindex.html#SQL-CREATEINDEX-CONCURRENTLY)
# (see: https://github.com/sqlalchemy/alembic/issues/277)
# 4. Adds indexes to both chat_message and chat_session tables for comprehensive search
def upgrade() -> None:
# Create a GIN index for full-text search on chat_message.message
op.execute(
"""
ALTER TABLE chat_message
ADD COLUMN message_tsv tsvector
GENERATED ALWAYS AS (to_tsvector('english', message)) STORED;
"""
)
# Commit the current transaction before creating concurrent indexes
op.execute("COMMIT")
op.execute(
"""
CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_chat_message_tsv
ON chat_message
USING GIN (message_tsv)
"""
)
# Also add a stored tsvector column for chat_session.description
op.execute(
"""
ALTER TABLE chat_session
ADD COLUMN description_tsv tsvector
GENERATED ALWAYS AS (to_tsvector('english', coalesce(description, ''))) STORED;
"""
)
# Commit again before creating the second concurrent index
op.execute("COMMIT")
op.execute(
"""
CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_chat_session_desc_tsv
ON chat_session
USING GIN (description_tsv)
"""
)
def downgrade() -> None:
# Drop the indexes first (use CONCURRENTLY for dropping too)
op.execute("COMMIT")
op.execute("DROP INDEX CONCURRENTLY IF EXISTS idx_chat_message_tsv;")
op.execute("COMMIT")
op.execute("DROP INDEX CONCURRENTLY IF EXISTS idx_chat_session_desc_tsv;")
# Then drop the columns
op.execute("ALTER TABLE chat_message DROP COLUMN IF EXISTS message_tsv;")
op.execute("ALTER TABLE chat_session DROP COLUMN IF EXISTS description_tsv;")
op.execute("DROP INDEX IF EXISTS idx_chat_message_message_lower;")

View File

@@ -0,0 +1,36 @@
"""force lowercase all users
Revision ID: f11b408e39d3
Revises: 3bd4c84fe72f
Create Date: 2025-02-26 17:04:55.683500
"""
# revision identifiers, used by Alembic.
revision = "f11b408e39d3"
down_revision = "3bd4c84fe72f"
branch_labels = None
depends_on = None
def upgrade() -> None:
# 1) Convert all existing user emails to lowercase
from alembic import op
op.execute(
"""
UPDATE "user"
SET email = LOWER(email)
"""
)
# 2) Add a check constraint to ensure emails are always lowercase
op.create_check_constraint("ensure_lowercase_email", "user", "email = LOWER(email)")
def downgrade() -> None:
# Drop the check constraint
from alembic import op
op.drop_constraint("ensure_lowercase_email", "user", type_="check")

View File

@@ -0,0 +1,42 @@
"""lowercase multi-tenant user auth
Revision ID: 34e3630c7f32
Revises: a4f6ee863c47
Create Date: 2025-02-26 15:03:01.211894
"""
from alembic import op
# revision identifiers, used by Alembic.
revision = "34e3630c7f32"
down_revision = "a4f6ee863c47"
branch_labels = None
depends_on = None
def upgrade() -> None:
# 1) Convert all existing rows to lowercase
op.execute(
"""
UPDATE user_tenant_mapping
SET email = LOWER(email)
"""
)
# 2) Add a check constraint so that emails cannot be written in uppercase
op.create_check_constraint(
"ensure_lowercase_email",
"user_tenant_mapping",
"email = LOWER(email)",
schema="public",
)
def downgrade() -> None:
# Drop the check constraint
op.drop_constraint(
"ensure_lowercase_email",
"user_tenant_mapping",
schema="public",
type_="check",
)

View File

@@ -138,6 +138,7 @@ def get_user_chat_sessions(
name=chat.description,
persona_id=chat.persona_id,
time_created=chat.time_created.isoformat(),
time_updated=chat.time_updated.isoformat(),
shared_status=chat.shared_status,
folder_id=chat.folder_id,
current_alternate_model=chat.current_alternate_model,

View File

@@ -7,6 +7,7 @@ from ee.onyx.configs.app_configs import STRIPE_PRICE_ID
from ee.onyx.configs.app_configs import STRIPE_SECRET_KEY
from ee.onyx.server.tenants.access import generate_data_plane_token
from ee.onyx.server.tenants.models import BillingInformation
from ee.onyx.server.tenants.models import SubscriptionStatusResponse
from onyx.configs.app_configs import CONTROL_PLANE_API_BASE_URL
from onyx.utils.logger import setup_logger
@@ -41,7 +42,9 @@ def fetch_tenant_stripe_information(tenant_id: str) -> dict:
return response.json()
def fetch_billing_information(tenant_id: str) -> BillingInformation:
def fetch_billing_information(
tenant_id: str,
) -> BillingInformation | SubscriptionStatusResponse:
logger.info("Fetching billing information")
token = generate_data_plane_token()
headers = {
@@ -52,8 +55,19 @@ def fetch_billing_information(tenant_id: str) -> BillingInformation:
params = {"tenant_id": tenant_id}
response = requests.get(url, headers=headers, params=params)
response.raise_for_status()
billing_info = BillingInformation(**response.json())
return billing_info
response_data = response.json()
# Check if the response indicates no subscription
if (
isinstance(response_data, dict)
and "subscribed" in response_data
and not response_data["subscribed"]
):
return SubscriptionStatusResponse(**response_data)
# Otherwise, parse as BillingInformation
return BillingInformation(**response_data)
def register_tenant_users(tenant_id: str, number_of_users: int) -> stripe.Subscription:

View File

@@ -200,25 +200,6 @@ async def rollback_tenant_provisioning(tenant_id: str) -> None:
def configure_default_api_keys(db_session: Session) -> None:
if OPENAI_DEFAULT_API_KEY:
open_provider = LLMProviderUpsertRequest(
name="OpenAI",
provider=OPENAI_PROVIDER_NAME,
api_key=OPENAI_DEFAULT_API_KEY,
default_model_name="gpt-4",
fast_default_model_name="gpt-4o-mini",
model_names=OPEN_AI_MODEL_NAMES,
)
try:
full_provider = upsert_llm_provider(open_provider, db_session)
update_default_provider(full_provider.id, db_session)
except Exception as e:
logger.error(f"Failed to configure OpenAI provider: {e}")
else:
logger.error(
"OPENAI_DEFAULT_API_KEY not set, skipping OpenAI provider configuration"
)
if ANTHROPIC_DEFAULT_API_KEY:
anthropic_provider = LLMProviderUpsertRequest(
name="Anthropic",
@@ -227,6 +208,7 @@ def configure_default_api_keys(db_session: Session) -> None:
default_model_name="claude-3-7-sonnet-20250219",
fast_default_model_name="claude-3-5-sonnet-20241022",
model_names=ANTHROPIC_MODEL_NAMES,
display_model_names=["claude-3-5-sonnet-20241022"],
)
try:
full_provider = upsert_llm_provider(anthropic_provider, db_session)
@@ -238,6 +220,26 @@ def configure_default_api_keys(db_session: Session) -> None:
"ANTHROPIC_DEFAULT_API_KEY not set, skipping Anthropic provider configuration"
)
if OPENAI_DEFAULT_API_KEY:
open_provider = LLMProviderUpsertRequest(
name="OpenAI",
provider=OPENAI_PROVIDER_NAME,
api_key=OPENAI_DEFAULT_API_KEY,
default_model_name="gpt-4o",
fast_default_model_name="gpt-4o-mini",
model_names=OPEN_AI_MODEL_NAMES,
display_model_names=["o1", "o3-mini", "gpt-4o", "gpt-4o-mini"],
)
try:
full_provider = upsert_llm_provider(open_provider, db_session)
update_default_provider(full_provider.id, db_session)
except Exception as e:
logger.error(f"Failed to configure OpenAI provider: {e}")
else:
logger.error(
"OPENAI_DEFAULT_API_KEY not set, skipping OpenAI provider configuration"
)
if COHERE_DEFAULT_API_KEY:
cloud_embedding_provider = CloudEmbeddingProviderCreationRequest(
provider_type=EmbeddingProvider.COHERE,

View File

@@ -411,7 +411,7 @@ class UserManager(UUIDIDMixin, BaseUserManager[User, uuid.UUID]):
"refresh_token": refresh_token,
}
user: User
user: User | None = None
try:
# Attempt to get user by OAuth account
@@ -420,15 +420,20 @@ class UserManager(UUIDIDMixin, BaseUserManager[User, uuid.UUID]):
except exceptions.UserNotExists:
try:
# Attempt to get user by email
user = cast(User, await self.user_db.get_by_email(account_email))
user = await self.user_db.get_by_email(account_email)
if not associate_by_email:
raise exceptions.UserAlreadyExists()
user = await self.user_db.add_oauth_account(
user, oauth_account_dict
)
# Make sure user is not None before adding OAuth account
if user is not None:
user = await self.user_db.add_oauth_account(
user, oauth_account_dict
)
else:
# This shouldn't happen since get_by_email would raise UserNotExists
# but adding as a safeguard
raise exceptions.UserNotExists()
# If user not found by OAuth account or email, create a new user
except exceptions.UserNotExists:
password = self.password_helper.generate()
user_dict = {
@@ -439,26 +444,36 @@ class UserManager(UUIDIDMixin, BaseUserManager[User, uuid.UUID]):
user = await self.user_db.create(user_dict)
# Explicitly set the Postgres schema for this session to ensure
# OAuth account creation happens in the correct tenant schema
# Add OAuth account
await self.user_db.add_oauth_account(user, oauth_account_dict)
await self.on_after_register(user, request)
# Add OAuth account only if user creation was successful
if user is not None:
await self.user_db.add_oauth_account(user, oauth_account_dict)
await self.on_after_register(user, request)
else:
raise HTTPException(
status_code=500, detail="Failed to create user account"
)
else:
for existing_oauth_account in user.oauth_accounts:
if (
existing_oauth_account.account_id == account_id
and existing_oauth_account.oauth_name == oauth_name
):
user = await self.user_db.update_oauth_account(
user,
# NOTE: OAuthAccount DOES implement the OAuthAccountProtocol
# but the type checker doesn't know that :(
existing_oauth_account, # type: ignore
oauth_account_dict,
)
# User exists, update OAuth account if needed
if user is not None: # Add explicit check
for existing_oauth_account in user.oauth_accounts:
if (
existing_oauth_account.account_id == account_id
and existing_oauth_account.oauth_name == oauth_name
):
user = await self.user_db.update_oauth_account(
user,
# NOTE: OAuthAccount DOES implement the OAuthAccountProtocol
# but the type checker doesn't know that :(
existing_oauth_account, # type: ignore
oauth_account_dict,
)
# Ensure user is not None before proceeding
if user is None:
raise HTTPException(
status_code=500, detail="Failed to authenticate or create user"
)
# NOTE: Most IdPs have very short expiry times, and we don't want to force the user to
# re-authenticate that frequently, so by default this is disabled

View File

@@ -16,7 +16,7 @@ from typing import Optional
from onyx.configs.constants import POSTGRES_CELERY_WORKER_INDEXING_CHILD_APP_NAME
from onyx.db.engine import SqlEngine
from onyx.utils.logger import setup_logger
from onyx.setup import setup_logger
from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA
from shared_configs.configs import TENANT_ID_PREFIX
from shared_configs.contextvars import CURRENT_TENANT_ID_CONTEXTVAR

View File

@@ -168,7 +168,7 @@ def get_chat_sessions_by_user(
if not include_onyxbot_flows:
stmt = stmt.where(ChatSession.onyxbot_flow.is_(False))
stmt = stmt.order_by(desc(ChatSession.time_created))
stmt = stmt.order_by(desc(ChatSession.time_updated))
if deleted is not None:
stmt = stmt.where(ChatSession.deleted == deleted)
@@ -962,6 +962,7 @@ def translate_db_message_to_chat_message_detail(
chat_message.sub_questions
),
refined_answer_improvement=chat_message.refined_answer_improvement,
is_agentic=chat_message.is_agentic,
error=chat_message.error,
)

View File

@@ -3,14 +3,13 @@ from typing import Optional
from typing import Tuple
from uuid import UUID
from sqlalchemy import column
from sqlalchemy import desc
from sqlalchemy import func
from sqlalchemy import literal
from sqlalchemy import Select
from sqlalchemy import select
from sqlalchemy import union_all
from sqlalchemy.orm import joinedload
from sqlalchemy.orm import Session
from sqlalchemy.sql.expression import ColumnClause
from onyx.db.models import ChatMessage
from onyx.db.models import ChatSession
@@ -26,127 +25,87 @@ def search_chat_sessions(
include_onyxbot_flows: bool = False,
) -> Tuple[List[ChatSession], bool]:
"""
Search for chat sessions based on the provided query.
If no query is provided, returns recent chat sessions.
Fast full-text search on ChatSession + ChatMessage using tsvectors.
Returns a tuple of (chat_sessions, has_more)
If no query is provided, returns the most recent chat sessions.
Otherwise, searches both chat messages and session descriptions.
Returns a tuple of (sessions, has_more) where has_more indicates if
there are additional results beyond the requested page.
"""
offset = (page - 1) * page_size
offset_val = (page - 1) * page_size
# If no search query, we use standard SQLAlchemy pagination
# If no query, just return the most recent sessions
if not query or not query.strip():
stmt = select(ChatSession)
if user_id:
stmt = (
select(ChatSession)
.order_by(desc(ChatSession.time_created))
.offset(offset_val)
.limit(page_size + 1)
)
if user_id is not None:
stmt = stmt.where(ChatSession.user_id == user_id)
if not include_onyxbot_flows:
stmt = stmt.where(ChatSession.onyxbot_flow.is_(False))
if not include_deleted:
stmt = stmt.where(ChatSession.deleted.is_(False))
stmt = stmt.order_by(desc(ChatSession.time_created))
# Apply pagination
stmt = stmt.offset(offset).limit(page_size + 1)
result = db_session.execute(stmt.options(joinedload(ChatSession.persona)))
chat_sessions = result.scalars().all()
sessions = result.scalars().all()
has_more = len(chat_sessions) > page_size
has_more = len(sessions) > page_size
if has_more:
chat_sessions = chat_sessions[:page_size]
sessions = sessions[:page_size]
return list(chat_sessions), has_more
return list(sessions), has_more
words = query.lower().strip().split()
# Otherwise, proceed with full-text search
query = query.strip()
# Message mach subquery
message_matches = []
for word in words:
word_like = f"%{word}%"
message_match: Select = (
select(ChatMessage.chat_session_id, literal(1.0).label("search_rank"))
.join(ChatSession, ChatSession.id == ChatMessage.chat_session_id)
.where(func.lower(ChatMessage.message).like(word_like))
)
if user_id:
message_match = message_match.where(ChatSession.user_id == user_id)
message_matches.append(message_match)
if message_matches:
message_matches_query = union_all(*message_matches).alias("message_matches")
else:
return [], False
# Description matches
description_match: Select = select(
ChatSession.id.label("chat_session_id"), literal(0.5).label("search_rank")
).where(func.lower(ChatSession.description).like(f"%{query.lower()}%"))
if user_id:
description_match = description_match.where(ChatSession.user_id == user_id)
base_conditions = []
if user_id is not None:
base_conditions.append(ChatSession.user_id == user_id)
if not include_onyxbot_flows:
description_match = description_match.where(ChatSession.onyxbot_flow.is_(False))
base_conditions.append(ChatSession.onyxbot_flow.is_(False))
if not include_deleted:
description_match = description_match.where(ChatSession.deleted.is_(False))
base_conditions.append(ChatSession.deleted.is_(False))
# Combine all match sources
combined_matches = union_all(
message_matches_query.select(), description_match
).alias("combined_matches")
message_tsv: ColumnClause = column("message_tsv")
description_tsv: ColumnClause = column("description_tsv")
# Use CTE to group and get max rank
session_ranks = (
select(
combined_matches.c.chat_session_id,
func.max(combined_matches.c.search_rank).label("rank"),
)
.group_by(combined_matches.c.chat_session_id)
.alias("session_ranks")
ts_query = func.plainto_tsquery("english", query)
description_session_ids = (
select(ChatSession.id)
.where(*base_conditions)
.where(description_tsv.op("@@")(ts_query))
)
# Get ranked sessions with pagination
ranked_query = (
db_session.query(session_ranks.c.chat_session_id, session_ranks.c.rank)
.order_by(desc(session_ranks.c.rank), session_ranks.c.chat_session_id)
.offset(offset)
message_session_ids = (
select(ChatMessage.chat_session_id)
.join(ChatSession, ChatMessage.chat_session_id == ChatSession.id)
.where(*base_conditions)
.where(message_tsv.op("@@")(ts_query))
)
combined_ids = description_session_ids.union(message_session_ids).alias(
"combined_ids"
)
final_stmt = (
select(ChatSession)
.join(combined_ids, ChatSession.id == combined_ids.c.id)
.order_by(desc(ChatSession.time_created))
.distinct()
.offset(offset_val)
.limit(page_size + 1)
.options(joinedload(ChatSession.persona))
)
result = ranked_query.all()
session_objs = db_session.execute(final_stmt).scalars().all()
# Extract session IDs and ranks
session_ids_with_ranks = {row.chat_session_id: row.rank for row in result}
session_ids = list(session_ids_with_ranks.keys())
if not session_ids:
return [], False
# Now, let's query the actual ChatSession objects using the IDs
stmt = select(ChatSession).where(ChatSession.id.in_(session_ids))
if user_id:
stmt = stmt.where(ChatSession.user_id == user_id)
if not include_onyxbot_flows:
stmt = stmt.where(ChatSession.onyxbot_flow.is_(False))
if not include_deleted:
stmt = stmt.where(ChatSession.deleted.is_(False))
# Full objects with eager loading
result = db_session.execute(stmt.options(joinedload(ChatSession.persona)))
chat_sessions = result.scalars().all()
# Sort based on above ranking
chat_sessions = sorted(
chat_sessions,
key=lambda session: (
-session_ids_with_ranks.get(session.id, 0), # Rank (higher first)
session.time_created.timestamp() * -1, # Then by time (newest first)
),
)
has_more = len(chat_sessions) > page_size
has_more = len(session_objs) > page_size
if has_more:
chat_sessions = chat_sessions[:page_size]
session_objs = session_objs[:page_size]
return chat_sessions, has_more
return list(session_objs), has_more

View File

@@ -360,18 +360,13 @@ def backend_update_credential_json(
db_session.commit()
def delete_credential(
def _delete_credential_internal(
credential: Credential,
credential_id: int,
user: User | None,
db_session: Session,
force: bool = False,
) -> None:
credential = fetch_credential_by_id_for_user(credential_id, user, db_session)
if credential is None:
raise ValueError(
f"Credential by provided id {credential_id} does not exist or does not belong to user"
)
"""Internal utility function to handle the actual deletion of a credential"""
associated_connectors = (
db_session.query(ConnectorCredentialPair)
.filter(ConnectorCredentialPair.credential_id == credential_id)
@@ -416,6 +411,35 @@ def delete_credential(
db_session.commit()
def delete_credential_for_user(
credential_id: int,
user: User,
db_session: Session,
force: bool = False,
) -> None:
"""Delete a credential that belongs to a specific user"""
credential = fetch_credential_by_id_for_user(credential_id, user, db_session)
if credential is None:
raise ValueError(
f"Credential by provided id {credential_id} does not exist or does not belong to user"
)
_delete_credential_internal(credential, credential_id, db_session, force)
def delete_credential(
credential_id: int,
db_session: Session,
force: bool = False,
) -> None:
"""Delete a credential regardless of ownership (admin function)"""
credential = fetch_credential_by_id(credential_id, db_session)
if credential is None:
raise ValueError(f"Credential by provided id {credential_id} does not exist")
_delete_credential_internal(credential, credential_id, db_session, force)
def create_initial_public_credential(db_session: Session) -> None:
error_msg = (
"DB is not in a valid initial state."

View File

@@ -7,6 +7,7 @@ from typing import Optional
from uuid import uuid4
from pydantic import BaseModel
from sqlalchemy.orm import validates
from typing_extensions import TypedDict # noreorder
from uuid import UUID
@@ -25,6 +26,7 @@ from sqlalchemy import ForeignKey
from sqlalchemy import func
from sqlalchemy import Index
from sqlalchemy import Integer
from sqlalchemy import Sequence
from sqlalchemy import String
from sqlalchemy import Text
@@ -205,6 +207,10 @@ class User(SQLAlchemyBaseUserTableUUID, Base):
primaryjoin="User.id == foreign(ConnectorCredentialPair.creator_id)",
)
@validates("email")
def validate_email(self, key: str, value: str) -> str:
return value.lower() if value else value
@property
def password_configured(self) -> bool:
"""
@@ -2269,6 +2275,10 @@ class UserTenantMapping(Base):
email: Mapped[str] = mapped_column(String, nullable=False, primary_key=True)
tenant_id: Mapped[str] = mapped_column(String, nullable=False)
@validates("email")
def validate_email(self, key: str, value: str) -> str:
return value.lower() if value else value
# This is a mapping from tenant IDs to anonymous user paths
class TenantAnonymousUserPath(Base):

View File

@@ -23,7 +23,7 @@ from onyx.configs.constants import SearchFeedbackType
from onyx.configs.onyxbot_configs import DANSWER_BOT_NUM_DOCS_TO_DISPLAY
from onyx.context.search.models import SavedSearchDoc
from onyx.db.chat import get_chat_session_by_message_id
from onyx.db.engine import get_session_with_tenant
from onyx.db.engine import get_session_with_current_tenant
from onyx.db.models import ChannelConfig
from onyx.onyxbot.slack.constants import CONTINUE_IN_WEB_UI_ACTION_ID
from onyx.onyxbot.slack.constants import DISLIKE_BLOCK_ACTION_ID
@@ -410,12 +410,11 @@ def _build_qa_response_blocks(
def _build_continue_in_web_ui_block(
tenant_id: str,
message_id: int | None,
) -> Block:
if message_id is None:
raise ValueError("No message id provided to build continue in web ui block")
with get_session_with_tenant(tenant_id=tenant_id) as db_session:
with get_session_with_current_tenant() as db_session:
chat_session = get_chat_session_by_message_id(
db_session=db_session,
message_id=message_id,
@@ -482,7 +481,6 @@ def build_follow_up_resolved_blocks(
def build_slack_response_blocks(
answer: ChatOnyxBotResponse,
tenant_id: str,
message_info: SlackMessageInfo,
channel_conf: ChannelConfig | None,
use_citations: bool,
@@ -517,7 +515,6 @@ def build_slack_response_blocks(
if channel_conf and channel_conf.get("show_continue_in_web_ui"):
web_follow_up_block.append(
_build_continue_in_web_ui_block(
tenant_id=tenant_id,
message_id=answer.chat_message_id,
)
)

View File

@@ -11,7 +11,7 @@ from onyx.configs.constants import SearchFeedbackType
from onyx.configs.onyxbot_configs import DANSWER_FOLLOWUP_EMOJI
from onyx.connectors.slack.utils import expert_info_from_slack_id
from onyx.connectors.slack.utils import make_slack_api_rate_limited
from onyx.db.engine import get_session_with_tenant
from onyx.db.engine import get_session_with_current_tenant
from onyx.db.feedback import create_chat_message_feedback
from onyx.db.feedback import create_doc_retrieval_feedback
from onyx.onyxbot.slack.blocks import build_follow_up_resolved_blocks
@@ -114,7 +114,7 @@ def handle_generate_answer_button(
thread_ts=thread_ts,
)
with get_session_with_tenant(tenant_id=client.tenant_id) as db_session:
with get_session_with_current_tenant() as db_session:
slack_channel_config = get_slack_channel_config_for_bot_and_channel(
db_session=db_session,
slack_bot_id=client.slack_bot_id,
@@ -136,7 +136,6 @@ def handle_generate_answer_button(
slack_channel_config=slack_channel_config,
receiver_ids=None,
client=client.web_client,
tenant_id=client.tenant_id,
channel=channel_id,
logger=logger,
feedback_reminder_id=None,
@@ -151,11 +150,10 @@ def handle_slack_feedback(
user_id_to_post_confirmation: str,
channel_id_to_post_confirmation: str,
thread_ts_to_post_confirmation: str,
tenant_id: str,
) -> None:
message_id, doc_id, doc_rank = decompose_action_id(feedback_id)
with get_session_with_tenant(tenant_id=tenant_id) as db_session:
with get_session_with_current_tenant() as db_session:
if feedback_type in [LIKE_BLOCK_ACTION_ID, DISLIKE_BLOCK_ACTION_ID]:
create_chat_message_feedback(
is_positive=feedback_type == LIKE_BLOCK_ACTION_ID,
@@ -246,7 +244,7 @@ def handle_followup_button(
tag_ids: list[str] = []
group_ids: list[str] = []
with get_session_with_tenant(tenant_id=client.tenant_id) as db_session:
with get_session_with_current_tenant() as db_session:
channel_name, is_dm = get_channel_name_from_id(
client=client.web_client, channel_id=channel_id
)

View File

@@ -5,7 +5,7 @@ from slack_sdk.errors import SlackApiError
from onyx.configs.onyxbot_configs import DANSWER_BOT_FEEDBACK_REMINDER
from onyx.configs.onyxbot_configs import DANSWER_REACT_EMOJI
from onyx.db.engine import get_session_with_tenant
from onyx.db.engine import get_session_with_current_tenant
from onyx.db.models import SlackChannelConfig
from onyx.db.users import add_slack_user_if_not_exists
from onyx.onyxbot.slack.blocks import get_feedback_reminder_blocks
@@ -109,7 +109,6 @@ def handle_message(
slack_channel_config: SlackChannelConfig,
client: WebClient,
feedback_reminder_id: str | None,
tenant_id: str,
) -> bool:
"""Potentially respond to the user message depending on filters and if an answer was generated
@@ -135,9 +134,7 @@ def handle_message(
action = "slack_tag_message"
elif is_bot_dm:
action = "slack_dm_message"
slack_usage_report(
action=action, sender_id=sender_id, client=client, tenant_id=tenant_id
)
slack_usage_report(action=action, sender_id=sender_id, client=client)
document_set_names: list[str] | None = None
persona = slack_channel_config.persona if slack_channel_config else None
@@ -218,7 +215,7 @@ def handle_message(
except SlackApiError as e:
logger.error(f"Was not able to react to user message due to: {e}")
with get_session_with_tenant(tenant_id=tenant_id) as db_session:
with get_session_with_current_tenant() as db_session:
if message_info.email:
add_slack_user_if_not_exists(db_session, message_info.email)
@@ -244,6 +241,5 @@ def handle_message(
channel=channel,
logger=logger,
feedback_reminder_id=feedback_reminder_id,
tenant_id=tenant_id,
)
return issue_with_regular_answer

View File

@@ -24,7 +24,6 @@ from onyx.context.search.enums import OptionalSearchSetting
from onyx.context.search.models import BaseFilters
from onyx.context.search.models import RetrievalDetails
from onyx.db.engine import get_session_with_current_tenant
from onyx.db.engine import get_session_with_tenant
from onyx.db.models import SlackChannelConfig
from onyx.db.models import User
from onyx.db.persona import get_persona_by_id
@@ -72,7 +71,6 @@ def handle_regular_answer(
channel: str,
logger: OnyxLoggingAdapter,
feedback_reminder_id: str | None,
tenant_id: str,
num_retries: int = DANSWER_BOT_NUM_RETRIES,
thread_context_percent: float = MAX_THREAD_CONTEXT_PERCENTAGE,
should_respond_with_error_msgs: bool = DANSWER_BOT_DISPLAY_ERROR_MSGS,
@@ -87,7 +85,7 @@ def handle_regular_answer(
user = None
if message_info.is_bot_dm:
if message_info.email:
with get_session_with_tenant(tenant_id=tenant_id) as db_session:
with get_session_with_current_tenant() as db_session:
user = get_user_by_email(message_info.email, db_session)
document_set_names: list[str] | None = None
@@ -96,7 +94,7 @@ def handle_regular_answer(
# This way slack flow always has a persona
persona = slack_channel_config.persona
if not persona:
with get_session_with_tenant(tenant_id=tenant_id) as db_session:
with get_session_with_current_tenant() as db_session:
persona = get_persona_by_id(DEFAULT_PERSONA_ID, user, db_session)
document_set_names = [
document_set.name for document_set in persona.document_sets
@@ -157,7 +155,7 @@ def handle_regular_answer(
def _get_slack_answer(
new_message_request: CreateChatMessageRequest, onyx_user: User | None
) -> ChatOnyxBotResponse:
with get_session_with_tenant(tenant_id=tenant_id) as db_session:
with get_session_with_current_tenant() as db_session:
packets = stream_chat_message_objects(
new_msg_req=new_message_request,
user=onyx_user,
@@ -197,7 +195,7 @@ def handle_regular_answer(
enable_auto_detect_filters=auto_detect_filters,
)
with get_session_with_tenant(tenant_id=tenant_id) as db_session:
with get_session_with_current_tenant() as db_session:
answer_request = prepare_chat_message_request(
message_text=user_message.message,
user=user,
@@ -361,7 +359,6 @@ def handle_regular_answer(
return True
all_blocks = build_slack_response_blocks(
tenant_id=tenant_id,
message_info=message_info,
answer=answer,
channel_conf=channel_conf,

View File

@@ -37,6 +37,7 @@ from onyx.context.search.retrieval.search_runner import (
download_nltk_data,
)
from onyx.db.engine import get_all_tenant_ids
from onyx.db.engine import get_session_with_current_tenant
from onyx.db.engine import get_session_with_tenant
from onyx.db.models import SlackBot
from onyx.db.search_settings import get_current_search_settings
@@ -92,6 +93,7 @@ from shared_configs.configs import MODEL_SERVER_PORT
from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA
from shared_configs.configs import SLACK_CHANNEL_ID
from shared_configs.contextvars import CURRENT_TENANT_ID_CONTEXTVAR
from shared_configs.contextvars import get_current_tenant_id
logger = setup_logger()
@@ -347,7 +349,7 @@ class SlackbotHandler:
redis_client = get_redis_client(tenant_id=tenant_id)
try:
with get_session_with_tenant(tenant_id=tenant_id) as db_session:
with get_session_with_current_tenant() as db_session:
# Attempt to fetch Slack bots
try:
bots = list(fetch_slack_bots(db_session=db_session))
@@ -586,7 +588,7 @@ def prefilter_requests(req: SocketModeRequest, client: TenantSocketModeClient) -
channel_name, _ = get_channel_name_from_id(
client=client.web_client, channel_id=channel
)
with get_session_with_tenant(tenant_id=client.tenant_id) as db_session:
with get_session_with_current_tenant() as db_session:
slack_channel_config = get_slack_channel_config_for_bot_and_channel(
db_session=db_session,
slack_bot_id=client.slack_bot_id,
@@ -680,7 +682,6 @@ def process_feedback(req: SocketModeRequest, client: TenantSocketModeClient) ->
user_id_to_post_confirmation=user_id,
channel_id_to_post_confirmation=channel_id,
thread_ts_to_post_confirmation=thread_ts,
tenant_id=client.tenant_id,
)
query_event_id, _, _ = decompose_action_id(feedback_id)
@@ -796,8 +797,9 @@ def process_message(
respond_every_channel: bool = DANSWER_BOT_RESPOND_EVERY_CHANNEL,
notify_no_answer: bool = NOTIFY_SLACKBOT_NO_ANSWER,
) -> None:
tenant_id = get_current_tenant_id()
logger.debug(
f"Received Slack request of type: '{req.type}' for tenant, {client.tenant_id}"
f"Received Slack request of type: '{req.type}' for tenant, {tenant_id}"
)
# Throw out requests that can't or shouldn't be handled
@@ -810,50 +812,39 @@ def process_message(
client=client.web_client, channel_id=channel
)
token: Token[str | None] | None = None
# Set the current tenant ID at the beginning for all DB calls within this thread
if client.tenant_id:
logger.info(f"Setting tenant ID to {client.tenant_id}")
token = CURRENT_TENANT_ID_CONTEXTVAR.set(client.tenant_id)
try:
with get_session_with_tenant(tenant_id=client.tenant_id) as db_session:
slack_channel_config = get_slack_channel_config_for_bot_and_channel(
db_session=db_session,
slack_bot_id=client.slack_bot_id,
channel_name=channel_name,
)
with get_session_with_current_tenant() as db_session:
slack_channel_config = get_slack_channel_config_for_bot_and_channel(
db_session=db_session,
slack_bot_id=client.slack_bot_id,
channel_name=channel_name,
)
follow_up = bool(
slack_channel_config.channel_config
and slack_channel_config.channel_config.get("follow_up_tags")
is not None
)
follow_up = bool(
slack_channel_config.channel_config
and slack_channel_config.channel_config.get("follow_up_tags") is not None
)
feedback_reminder_id = schedule_feedback_reminder(
details=details, client=client.web_client, include_followup=follow_up
)
feedback_reminder_id = schedule_feedback_reminder(
details=details, client=client.web_client, include_followup=follow_up
)
failed = handle_message(
message_info=details,
slack_channel_config=slack_channel_config,
client=client.web_client,
feedback_reminder_id=feedback_reminder_id,
tenant_id=client.tenant_id,
)
failed = handle_message(
message_info=details,
slack_channel_config=slack_channel_config,
client=client.web_client,
feedback_reminder_id=feedback_reminder_id,
)
if failed:
if feedback_reminder_id:
remove_scheduled_feedback_reminder(
client=client.web_client,
channel=details.sender_id,
msg_id=feedback_reminder_id,
)
# Skipping answering due to pre-filtering is not considered a failure
if notify_no_answer:
apologize_for_fail(details, client)
finally:
if token:
CURRENT_TENANT_ID_CONTEXTVAR.reset(token)
if failed:
if feedback_reminder_id:
remove_scheduled_feedback_reminder(
client=client.web_client,
channel=details.sender_id,
msg_id=feedback_reminder_id,
)
# Skipping answering due to pre-filtering is not considered a failure
if notify_no_answer:
apologize_for_fail(details, client)
def acknowledge_message(req: SocketModeRequest, client: TenantSocketModeClient) -> None:

View File

@@ -4,6 +4,8 @@ import re
import string
import time
import uuid
from collections.abc import Generator
from contextlib import contextmanager
from typing import Any
from typing import cast
@@ -30,7 +32,7 @@ from onyx.configs.onyxbot_configs import (
)
from onyx.connectors.slack.utils import make_slack_api_rate_limited
from onyx.connectors.slack.utils import SlackTextCleaner
from onyx.db.engine import get_session_with_tenant
from onyx.db.engine import get_session_with_current_tenant
from onyx.db.users import get_user_by_email
from onyx.llm.exceptions import GenAIDisabledException
from onyx.llm.factory import get_default_llms
@@ -43,6 +45,7 @@ from onyx.utils.logger import setup_logger
from onyx.utils.telemetry import optional_telemetry
from onyx.utils.telemetry import RecordType
from onyx.utils.text_processing import replace_whitespaces_w_space
from shared_configs.contextvars import CURRENT_TENANT_ID_CONTEXTVAR
logger = setup_logger()
@@ -569,9 +572,7 @@ def read_slack_thread(
return thread_messages
def slack_usage_report(
action: str, sender_id: str | None, client: WebClient, tenant_id: str
) -> None:
def slack_usage_report(action: str, sender_id: str | None, client: WebClient) -> None:
if DISABLE_TELEMETRY:
return
@@ -583,14 +584,13 @@ def slack_usage_report(
logger.warning("Unable to find sender email")
if sender_email is not None:
with get_session_with_tenant(tenant_id=tenant_id) as db_session:
with get_session_with_current_tenant() as db_session:
onyx_user = get_user_by_email(email=sender_email, db_session=db_session)
optional_telemetry(
record_type=RecordType.USAGE,
data={"action": action},
user_id=str(onyx_user.id) if onyx_user else "Non-Onyx-Or-No-Auth-User",
tenant_id=tenant_id,
)
@@ -665,5 +665,28 @@ def get_feedback_visibility() -> FeedbackVisibility:
class TenantSocketModeClient(SocketModeClient):
def __init__(self, tenant_id: str, slack_bot_id: int, *args: Any, **kwargs: Any):
super().__init__(*args, **kwargs)
self.tenant_id = tenant_id
self._tenant_id = tenant_id
self.slack_bot_id = slack_bot_id
@contextmanager
def _set_tenant_context(self) -> Generator[None, None, None]:
token = None
try:
if self._tenant_id:
token = CURRENT_TENANT_ID_CONTEXTVAR.set(self._tenant_id)
yield
finally:
if token:
CURRENT_TENANT_ID_CONTEXTVAR.reset(token)
def enqueue_message(self, message: str) -> None:
with self._set_tenant_context():
super().enqueue_message(message)
def process_message(self) -> None:
with self._set_tenant_context():
super().process_message()
def run_message_listeners(self, message: dict, raw_message: str) -> None:
with self._set_tenant_context():
super().run_message_listeners(message, raw_message)

View File

@@ -646,7 +646,6 @@ def associate_credential_to_connector(
)
return response
except ValidationError as e:
# If validation fails, delete the connector and commit the changes
# Ensures we don't leave invalid connectors in the database
@@ -660,10 +659,14 @@ def associate_credential_to_connector(
)
except IntegrityError as e:
logger.error(f"IntegrityError: {e}")
delete_connector(db_session, connector_id)
db_session.commit()
raise HTTPException(status_code=400, detail="Name must be unique")
except Exception as e:
logger.exception(f"Unexpected error: {e}")
raise HTTPException(status_code=500, detail="Unexpected error")

View File

@@ -13,6 +13,7 @@ from onyx.db.credentials import cleanup_gmail_credentials
from onyx.db.credentials import create_credential
from onyx.db.credentials import CREDENTIAL_PERMISSIONS_TO_IGNORE
from onyx.db.credentials import delete_credential
from onyx.db.credentials import delete_credential_for_user
from onyx.db.credentials import fetch_credential_by_id_for_user
from onyx.db.credentials import fetch_credentials_by_source_for_user
from onyx.db.credentials import fetch_credentials_for_user
@@ -88,7 +89,7 @@ def delete_credential_by_id_admin(
db_session: Session = Depends(get_session),
) -> StatusResponse:
"""Same as the user endpoint, but can delete any credential (not just the user's own)"""
delete_credential(db_session=db_session, credential_id=credential_id, user=None)
delete_credential(db_session=db_session, credential_id=credential_id)
return StatusResponse(
success=True, message="Credential deleted successfully", data=credential_id
)
@@ -242,7 +243,7 @@ def delete_credential_by_id(
user: User = Depends(current_user),
db_session: Session = Depends(get_session),
) -> StatusResponse:
delete_credential(
delete_credential_for_user(
credential_id,
user,
db_session,
@@ -259,7 +260,7 @@ def force_delete_credential_by_id(
user: User = Depends(current_user),
db_session: Session = Depends(get_session),
) -> StatusResponse:
delete_credential(credential_id, user, db_session, True)
delete_credential_for_user(credential_id, user, db_session, True)
return StatusResponse(
success=True, message="Credential deleted successfully", data=credential_id

View File

@@ -49,6 +49,7 @@ def get_folders(
name=chat_session.description,
persona_id=chat_session.persona_id,
time_created=chat_session.time_created.isoformat(),
time_updated=chat_session.time_updated.isoformat(),
shared_status=chat_session.shared_status,
folder_id=folder.id,
)

View File

@@ -147,9 +147,11 @@ def list_threads(
name=chat.description,
persona_id=chat.persona_id,
time_created=chat.time_created.isoformat(),
time_updated=chat.time_updated.isoformat(),
shared_status=chat.shared_status,
folder_id=chat.folder_id,
current_alternate_model=chat.current_alternate_model,
current_temperature_override=chat.temperature_override,
)
for chat in chat_sessions
]

View File

@@ -119,6 +119,7 @@ def get_user_chat_sessions(
name=chat.description,
persona_id=chat.persona_id,
time_created=chat.time_created.isoformat(),
time_updated=chat.time_updated.isoformat(),
shared_status=chat.shared_status,
folder_id=chat.folder_id,
current_alternate_model=chat.current_alternate_model,

View File

@@ -181,6 +181,7 @@ class ChatSessionDetails(BaseModel):
name: str | None
persona_id: int | None = None
time_created: str
time_updated: str
shared_status: ChatSessionSharedStatus
folder_id: int | None = None
current_alternate_model: str | None = None
@@ -241,6 +242,7 @@ class ChatMessageDetail(BaseModel):
files: list[FileDescriptor]
tool_call: ToolCallFinalResult | None
refined_answer_improvement: bool | None = None
is_agentic: bool | None = None
error: str | None = None
def model_dump(self, *args: list, **kwargs: dict[str, Any]) -> dict[str, Any]: # type: ignore

View File

@@ -159,6 +159,7 @@ def get_user_search_sessions(
name=sessions_with_documents_dict[search.id],
persona_id=search.persona_id,
time_created=search.time_created.isoformat(),
time_updated=search.time_updated.isoformat(),
shared_status=search.shared_status,
folder_id=search.folder_id,
current_alternate_model=search.current_alternate_model,

View File

@@ -1,6 +1,6 @@
import { Button } from "@/components/Button";
import { PopupSpec } from "@/components/admin/connectors/Popup";
import { useState } from "react";
import React, { useState, useEffect } from "react";
import { useSWRConfig } from "swr";
import * as Yup from "yup";
import { useRouter } from "next/navigation";
@@ -17,13 +17,18 @@ import {
GoogleDriveCredentialJson,
GoogleDriveServiceAccountCredentialJson,
} from "@/lib/connectors/credentials";
import { refreshAllGoogleData } from "@/lib/googleConnector";
import { ValidSources } from "@/lib/types";
import { buildSimilarCredentialInfoURL } from "@/app/admin/connector/[ccPairId]/lib";
type GoogleDriveCredentialJsonTypes = "authorized_user" | "service_account";
export const DriveJsonUpload = ({
setPopup,
onSuccess,
}: {
setPopup: (popupSpec: PopupSpec | null) => void;
onSuccess?: () => void;
}) => {
const { mutate } = useSWRConfig();
const [credentialJsonStr, setCredentialJsonStr] = useState<
@@ -62,7 +67,6 @@ export const DriveJsonUpload = ({
<Button
disabled={!credentialJsonStr}
onClick={async () => {
// check if the JSON is a app credential or a service account credential
let credentialFileType: GoogleDriveCredentialJsonTypes;
try {
const appCredentialJson = JSON.parse(credentialJsonStr!);
@@ -99,6 +103,10 @@ export const DriveJsonUpload = ({
message: "Successfully uploaded app credentials",
type: "success",
});
mutate("/api/manage/admin/connector/google-drive/app-credential");
if (onSuccess) {
onSuccess();
}
} else {
const errorMsg = await response.text();
setPopup({
@@ -106,7 +114,6 @@ export const DriveJsonUpload = ({
type: "error",
});
}
mutate("/api/manage/admin/connector/google-drive/app-credential");
}
if (credentialFileType === "service_account") {
@@ -122,19 +129,22 @@ export const DriveJsonUpload = ({
);
if (response.ok) {
setPopup({
message: "Successfully uploaded app credentials",
message: "Successfully uploaded service account key",
type: "success",
});
mutate(
"/api/manage/admin/connector/google-drive/service-account-key"
);
if (onSuccess) {
onSuccess();
}
} else {
const errorMsg = await response.text();
setPopup({
message: `Failed to upload app credentials - ${errorMsg}`,
message: `Failed to upload service account key - ${errorMsg}`,
type: "error",
});
}
mutate(
"/api/manage/admin/connector/google-drive/service-account-key"
);
}
}}
>
@@ -149,6 +159,7 @@ interface DriveJsonUploadSectionProps {
appCredentialData?: { client_id: string };
serviceAccountCredentialData?: { service_account_email: string };
isAdmin: boolean;
onSuccess?: () => void;
}
export const DriveJsonUploadSection = ({
@@ -156,17 +167,36 @@ export const DriveJsonUploadSection = ({
appCredentialData,
serviceAccountCredentialData,
isAdmin,
onSuccess,
}: DriveJsonUploadSectionProps) => {
const { mutate } = useSWRConfig();
const router = useRouter();
const [localServiceAccountData, setLocalServiceAccountData] = useState(
serviceAccountCredentialData
);
const [localAppCredentialData, setLocalAppCredentialData] =
useState(appCredentialData);
if (serviceAccountCredentialData?.service_account_email) {
useEffect(() => {
setLocalServiceAccountData(serviceAccountCredentialData);
setLocalAppCredentialData(appCredentialData);
}, [serviceAccountCredentialData, appCredentialData]);
const handleSuccess = () => {
if (onSuccess) {
onSuccess();
} else {
refreshAllGoogleData(ValidSources.GoogleDrive);
}
};
if (localServiceAccountData?.service_account_email) {
return (
<div className="mt-2 text-sm">
<div>
Found existing service account key with the following <b>Email:</b>
<p className="italic mt-1">
{serviceAccountCredentialData.service_account_email}
{localServiceAccountData.service_account_email}
</p>
</div>
{isAdmin ? (
@@ -188,11 +218,15 @@ export const DriveJsonUploadSection = ({
mutate(
"/api/manage/admin/connector/google-drive/service-account-key"
);
mutate(
buildSimilarCredentialInfoURL(ValidSources.GoogleDrive)
);
setPopup({
message: "Successfully deleted service account key",
type: "success",
});
router.refresh();
setLocalServiceAccountData(undefined);
handleSuccess();
} else {
const errorMsg = await response.text();
setPopup({
@@ -216,12 +250,12 @@ export const DriveJsonUploadSection = ({
);
}
if (appCredentialData?.client_id) {
if (localAppCredentialData?.client_id) {
return (
<div className="mt-2 text-sm">
<div>
Found existing app credentials with the following <b>Client ID:</b>
<p className="italic mt-1">{appCredentialData.client_id}</p>
<p className="italic mt-1">{localAppCredentialData.client_id}</p>
</div>
{isAdmin ? (
<>
@@ -242,10 +276,15 @@ export const DriveJsonUploadSection = ({
mutate(
"/api/manage/admin/connector/google-drive/app-credential"
);
mutate(
buildSimilarCredentialInfoURL(ValidSources.GoogleDrive)
);
setPopup({
message: "Successfully deleted app credentials",
type: "success",
});
setLocalAppCredentialData(undefined);
handleSuccess();
} else {
const errorMsg = await response.text();
setPopup({
@@ -297,7 +336,7 @@ export const DriveJsonUploadSection = ({
Download the credentials JSON if choosing option (1) or the Service
Account key JSON if chooosing option (2), and upload it here.
</p>
<DriveJsonUpload setPopup={setPopup} />
<DriveJsonUpload setPopup={setPopup} onSuccess={handleSuccess} />
</div>
);
};
@@ -348,13 +387,41 @@ export const DriveAuthSection = ({
appCredentialData,
setPopup,
refreshCredentials,
connectorAssociated, // don't allow revoke if a connector / credential pair is active with the uploaded credential
connectorAssociated,
user,
}: DriveCredentialSectionProps) => {
const router = useRouter();
const [localServiceAccountData, setLocalServiceAccountData] = useState(
serviceAccountKeyData
);
const [localAppCredentialData, setLocalAppCredentialData] =
useState(appCredentialData);
const [
localGoogleDrivePublicCredential,
setLocalGoogleDrivePublicCredential,
] = useState(googleDrivePublicUploadedCredential);
const [
localGoogleDriveServiceAccountCredential,
setLocalGoogleDriveServiceAccountCredential,
] = useState(googleDriveServiceAccountCredential);
useEffect(() => {
setLocalServiceAccountData(serviceAccountKeyData);
setLocalAppCredentialData(appCredentialData);
setLocalGoogleDrivePublicCredential(googleDrivePublicUploadedCredential);
setLocalGoogleDriveServiceAccountCredential(
googleDriveServiceAccountCredential
);
}, [
serviceAccountKeyData,
appCredentialData,
googleDrivePublicUploadedCredential,
googleDriveServiceAccountCredential,
]);
const existingCredential =
googleDrivePublicUploadedCredential || googleDriveServiceAccountCredential;
localGoogleDrivePublicCredential ||
localGoogleDriveServiceAccountCredential;
if (existingCredential) {
return (
<>
@@ -377,7 +444,7 @@ export const DriveAuthSection = ({
);
}
if (serviceAccountKeyData?.service_account_email) {
if (localServiceAccountData?.service_account_email) {
return (
<div>
<Formik
@@ -438,7 +505,7 @@ export const DriveAuthSection = ({
);
}
if (appCredentialData?.client_id) {
if (localAppCredentialData?.client_id) {
return (
<div className="text-sm mb-4">
<p className="mb-2">

View File

@@ -1,8 +1,7 @@
"use client";
import React, { useEffect, useState } from "react";
import useSWR, { mutate } from "swr";
import { FetchError, errorHandlingFetcher } from "@/lib/fetcher";
import React from "react";
import { FetchError } from "@/lib/fetcher";
import { ErrorCallout } from "@/components/ErrorCallout";
import { LoadingAnimation } from "@/components/Loading";
import { PopupSpec, usePopup } from "@/components/admin/connectors/Popup";
@@ -15,22 +14,17 @@ import {
GoogleDriveCredentialJson,
GoogleDriveServiceAccountCredentialJson,
} from "@/lib/connectors/credentials";
import { ConnectorSnapshot } from "@/lib/connectors/connectors";
import { useUser } from "@/components/user/UserProvider";
import { buildSimilarCredentialInfoURL } from "@/app/admin/connector/[ccPairId]/lib";
const useConnectorsByCredentialId = (credential_id: number | null) => {
let url: string | null = null;
if (credential_id !== null) {
url = `/api/manage/admin/connector?credential=${credential_id}`;
}
const swrResponse = useSWR<ConnectorSnapshot[]>(url, errorHandlingFetcher);
return {
...swrResponse,
refreshConnectorsByCredentialId: () => mutate(url),
};
};
import {
useGoogleAppCredential,
useGoogleServiceAccountKey,
useGoogleCredentials,
useConnectorsByCredentialId,
checkCredentialsFetched,
filterUploadedCredentials,
checkConnectorsExist,
refreshAllGoogleData,
} from "@/lib/googleConnector";
const GDriveMain = ({
setPopup,
@@ -39,27 +33,20 @@ const GDriveMain = ({
}) => {
const { isAdmin, user } = useUser();
// tries getting the uploaded credential json
// Get app credential and service account key
const {
data: appCredentialData,
isLoading: isAppCredentialLoading,
error: isAppCredentialError,
} = useSWR<{ client_id: string }, FetchError>(
"/api/manage/admin/connector/google-drive/app-credential",
errorHandlingFetcher
);
} = useGoogleAppCredential("google_drive");
// tries getting the uploaded service account key
const {
data: serviceAccountKeyData,
isLoading: isServiceAccountKeyLoading,
error: isServiceAccountKeyError,
} = useSWR<{ service_account_email: string }, FetchError>(
"/api/manage/admin/connector/google-drive/service-account-key",
errorHandlingFetcher
);
} = useGoogleServiceAccountKey("google_drive");
// gets all public credentials
// Get all public credentials
const {
data: credentialsData,
isLoading: isCredentialsLoading,
@@ -67,33 +54,19 @@ const GDriveMain = ({
refreshCredentials,
} = usePublicCredentials();
// gets all credentials for source type google drive
// Get Google Drive-specific credentials
const {
data: googleDriveCredentials,
isLoading: isGoogleDriveCredentialsLoading,
error: googleDriveCredentialsError,
} = useSWR<Credential<any>[]>(
buildSimilarCredentialInfoURL(ValidSources.GoogleDrive),
errorHandlingFetcher,
{ refreshInterval: 5000 }
} = useGoogleCredentials(ValidSources.GoogleDrive);
// Filter uploaded credentials and get credential ID
const { credential_id, uploadedCredentials } = filterUploadedCredentials(
googleDriveCredentials
);
// filters down to just credentials that were created via upload (there should be only one)
let credential_id = null;
if (googleDriveCredentials) {
const googleDriveUploadedCredentials: Credential<GoogleDriveCredentialJson>[] =
googleDriveCredentials.filter(
(googleDriveCredential) =>
googleDriveCredential.credential_json.authentication_method !==
"oauth_interactive"
);
if (googleDriveUploadedCredentials.length > 0) {
credential_id = googleDriveUploadedCredentials[0].id;
}
}
// retrieves all connectors for that credential id
// Get connectors for the credential ID
const {
data: googleDriveConnectors,
isLoading: isGoogleDriveConnectorsLoading,
@@ -101,13 +74,25 @@ const GDriveMain = ({
refreshConnectorsByCredentialId,
} = useConnectorsByCredentialId(credential_id);
const appCredentialSuccessfullyFetched =
appCredentialData ||
(isAppCredentialError && isAppCredentialError.status === 404);
const serviceAccountKeySuccessfullyFetched =
serviceAccountKeyData ||
(isServiceAccountKeyError && isServiceAccountKeyError.status === 404);
// Check if credentials were successfully fetched
const {
appCredentialSuccessfullyFetched,
serviceAccountKeySuccessfullyFetched,
} = checkCredentialsFetched(
appCredentialData,
isAppCredentialError,
serviceAccountKeyData,
isServiceAccountKeyError
);
// Handle refresh of all data
const handleRefresh = () => {
refreshCredentials();
refreshConnectorsByCredentialId();
refreshAllGoogleData(ValidSources.GoogleDrive);
};
// Loading state
if (
(!appCredentialSuccessfullyFetched && isAppCredentialLoading) ||
(!serviceAccountKeySuccessfullyFetched && isServiceAccountKeyLoading) ||
@@ -122,6 +107,7 @@ const GDriveMain = ({
);
}
// Error states
if (credentialsError || !credentialsData) {
return <ErrorCallout errorTitle="Failed to load credentials." />;
}
@@ -141,7 +127,16 @@ const GDriveMain = ({
);
}
// get the actual uploaded oauth or service account credentials
if (googleDriveConnectorsError) {
return (
<ErrorCallout errorTitle="Failed to load Google Drive associated connectors." />
);
}
// Check if connectors exist
const connectorAssociated = checkConnectorsExist(googleDriveConnectors);
// Get the uploaded OAuth credential
const googleDrivePublicUploadedCredential:
| Credential<GoogleDriveCredentialJson>
| undefined = credentialsData.find(
@@ -152,6 +147,7 @@ const GDriveMain = ({
credential.credential_json.authentication_method !== "oauth_interactive"
);
// Get the service account credential
const googleDriveServiceAccountCredential:
| Credential<GoogleDriveServiceAccountCredentialJson>
| undefined = credentialsData.find(
@@ -160,19 +156,6 @@ const GDriveMain = ({
credential.source === "google_drive"
);
if (googleDriveConnectorsError) {
return (
<ErrorCallout errorTitle="Failed to load Google Drive associated connectors." />
);
}
let connectorAssociated = false;
if (googleDriveConnectors) {
if (googleDriveConnectors.length > 0) {
connectorAssociated = true;
}
}
return (
<>
<Title className="mb-2 mt-6">Step 1: Provide your Credentials</Title>
@@ -181,27 +164,30 @@ const GDriveMain = ({
appCredentialData={appCredentialData}
serviceAccountCredentialData={serviceAccountKeyData}
isAdmin={isAdmin}
onSuccess={handleRefresh}
/>
{isAdmin && (
<>
<Title className="mb-2 mt-6">Step 2: Authenticate with Onyx</Title>
<DriveAuthSection
setPopup={setPopup}
refreshCredentials={refreshCredentials}
googleDrivePublicUploadedCredential={
googleDrivePublicUploadedCredential
}
googleDriveServiceAccountCredential={
googleDriveServiceAccountCredential
}
appCredentialData={appCredentialData}
serviceAccountKeyData={serviceAccountKeyData}
connectorAssociated={connectorAssociated}
user={user}
/>
</>
)}
{isAdmin &&
(appCredentialData?.client_id ||
serviceAccountKeyData?.service_account_email) && (
<>
<Title className="mb-2 mt-6">Step 2: Authenticate with Onyx</Title>
<DriveAuthSection
setPopup={setPopup}
refreshCredentials={handleRefresh}
googleDrivePublicUploadedCredential={
googleDrivePublicUploadedCredential
}
googleDriveServiceAccountCredential={
googleDriveServiceAccountCredential
}
appCredentialData={appCredentialData}
serviceAccountKeyData={serviceAccountKeyData}
connectorAssociated={connectorAssociated}
user={user}
/>
</>
)}
</>
);
};

View File

@@ -1,6 +1,6 @@
import { Button } from "@/components/Button";
import { PopupSpec } from "@/components/admin/connectors/Popup";
import { useState } from "react";
import React, { useState, useEffect } from "react";
import { useSWRConfig } from "swr";
import * as Yup from "yup";
import { useRouter } from "next/navigation";
@@ -17,13 +17,18 @@ import {
GmailCredentialJson,
GmailServiceAccountCredentialJson,
} from "@/lib/connectors/credentials";
import { refreshAllGoogleData } from "@/lib/googleConnector";
import { ValidSources } from "@/lib/types";
import { buildSimilarCredentialInfoURL } from "@/app/admin/connector/[ccPairId]/lib";
type GmailCredentialJsonTypes = "authorized_user" | "service_account";
const DriveJsonUpload = ({
setPopup,
onSuccess,
}: {
setPopup: (popupSpec: PopupSpec | null) => void;
onSuccess?: () => void;
}) => {
const { mutate } = useSWRConfig();
const [credentialJsonStr, setCredentialJsonStr] = useState<
@@ -72,7 +77,7 @@ const DriveJsonUpload = ({
credentialFileType = "service_account";
} else {
throw new Error(
"Unknown credential type, expected 'OAuth Web application'"
"Unknown credential type, expected one of 'OAuth Web application' or 'Service Account'"
);
}
} catch (e) {
@@ -99,6 +104,10 @@ const DriveJsonUpload = ({
message: "Successfully uploaded app credentials",
type: "success",
});
mutate("/api/manage/admin/connector/gmail/app-credential");
if (onSuccess) {
onSuccess();
}
} else {
const errorMsg = await response.text();
setPopup({
@@ -106,7 +115,6 @@ const DriveJsonUpload = ({
type: "error",
});
}
mutate("/api/manage/admin/connector/gmail/app-credential");
}
if (credentialFileType === "service_account") {
@@ -122,17 +130,20 @@ const DriveJsonUpload = ({
);
if (response.ok) {
setPopup({
message: "Successfully uploaded app credentials",
message: "Successfully uploaded service account key",
type: "success",
});
mutate("/api/manage/admin/connector/gmail/service-account-key");
if (onSuccess) {
onSuccess();
}
} else {
const errorMsg = await response.text();
setPopup({
message: `Failed to upload app credentials - ${errorMsg}`,
message: `Failed to upload service account key - ${errorMsg}`,
type: "error",
});
}
mutate("/api/manage/admin/connector/gmail/service-account-key");
}
}}
>
@@ -147,6 +158,7 @@ interface DriveJsonUploadSectionProps {
appCredentialData?: { client_id: string };
serviceAccountCredentialData?: { service_account_email: string };
isAdmin: boolean;
onSuccess?: () => void;
}
export const GmailJsonUploadSection = ({
@@ -154,16 +166,37 @@ export const GmailJsonUploadSection = ({
appCredentialData,
serviceAccountCredentialData,
isAdmin,
onSuccess,
}: DriveJsonUploadSectionProps) => {
const { mutate } = useSWRConfig();
const router = useRouter();
const [localServiceAccountData, setLocalServiceAccountData] = useState(
serviceAccountCredentialData
);
const [localAppCredentialData, setLocalAppCredentialData] =
useState(appCredentialData);
if (serviceAccountCredentialData?.service_account_email) {
// Update local state when props change
useEffect(() => {
setLocalServiceAccountData(serviceAccountCredentialData);
setLocalAppCredentialData(appCredentialData);
}, [serviceAccountCredentialData, appCredentialData]);
const handleSuccess = () => {
if (onSuccess) {
onSuccess();
} else {
refreshAllGoogleData(ValidSources.Gmail);
}
};
if (localServiceAccountData?.service_account_email) {
return (
<div className="mt-2 text-sm">
<div>
Found existing service account key with the following <b>Email:</b>
<p className="italic mt-1">
{serviceAccountCredentialData.service_account_email}
{localServiceAccountData.service_account_email}
</p>
</div>
{isAdmin ? (
@@ -185,10 +218,15 @@ export const GmailJsonUploadSection = ({
mutate(
"/api/manage/admin/connector/gmail/service-account-key"
);
// Also mutate the credential endpoints to ensure Step 2 is reset
mutate(buildSimilarCredentialInfoURL(ValidSources.Gmail));
setPopup({
message: "Successfully deleted service account key",
type: "success",
});
// Immediately update local state
setLocalServiceAccountData(undefined);
handleSuccess();
} else {
const errorMsg = await response.text();
setPopup({
@@ -212,43 +250,56 @@ export const GmailJsonUploadSection = ({
);
}
if (appCredentialData?.client_id) {
if (localAppCredentialData?.client_id) {
return (
<div className="mt-2 text-sm">
<div>
Found existing app credentials with the following <b>Client ID:</b>
<p className="italic mt-1">{appCredentialData.client_id}</p>
<p className="italic mt-1">{localAppCredentialData.client_id}</p>
</div>
<div className="mt-4 mb-1">
If you want to update these credentials, delete the existing
credentials through the button below, and then upload a new
credentials JSON.
</div>
<Button
onClick={async () => {
const response = await fetch(
"/api/manage/admin/connector/gmail/app-credential",
{
method: "DELETE",
}
);
if (response.ok) {
mutate("/api/manage/admin/connector/gmail/app-credential");
setPopup({
message: "Successfully deleted service account key",
type: "success",
});
} else {
const errorMsg = await response.text();
setPopup({
message: `Failed to delete app credential - ${errorMsg}`,
type: "error",
});
}
}}
>
Delete
</Button>
{isAdmin ? (
<>
<div className="mt-4 mb-1">
If you want to update these credentials, delete the existing
credentials through the button below, and then upload a new
credentials JSON.
</div>
<Button
onClick={async () => {
const response = await fetch(
"/api/manage/admin/connector/gmail/app-credential",
{
method: "DELETE",
}
);
if (response.ok) {
mutate("/api/manage/admin/connector/gmail/app-credential");
// Also mutate the credential endpoints to ensure Step 2 is reset
mutate(buildSimilarCredentialInfoURL(ValidSources.Gmail));
setPopup({
message: "Successfully deleted app credentials",
type: "success",
});
// Immediately update local state
setLocalAppCredentialData(undefined);
handleSuccess();
} else {
const errorMsg = await response.text();
setPopup({
message: `Failed to delete app credential - ${errorMsg}`,
type: "error",
});
}
}}
>
Delete
</Button>
</>
) : (
<div className="mt-4 mb-1">
To change these credentials, please contact an administrator.
</div>
)}
</div>
);
}
@@ -276,14 +327,14 @@ export const GmailJsonUploadSection = ({
>
here
</a>{" "}
to either (1) setup a google OAuth App in your company workspace or (2)
to either (1) setup a Google OAuth App in your company workspace or (2)
create a Service Account.
<br />
<br />
Download the credentials JSON if choosing option (1) or the Service
Account key JSON if chooosing option (2), and upload it here.
Account key JSON if choosing option (2), and upload it here.
</p>
<DriveJsonUpload setPopup={setPopup} />
<DriveJsonUpload setPopup={setPopup} onSuccess={handleSuccess} />
</div>
);
};
@@ -299,6 +350,34 @@ interface DriveCredentialSectionProps {
user: User | null;
}
async function handleRevokeAccess(
connectorExists: boolean,
setPopup: (popupSpec: PopupSpec | null) => void,
existingCredential:
| Credential<GmailCredentialJson>
| Credential<GmailServiceAccountCredentialJson>,
refreshCredentials: () => void
) {
if (connectorExists) {
const message =
"Cannot revoke the Gmail credential while any connector is still associated with the credential. " +
"Please delete all associated connectors, then try again.";
setPopup({
message: message,
type: "error",
});
return;
}
await adminDeleteCredential(existingCredential.id);
setPopup({
message: "Successfully revoked the Gmail credential!",
type: "success",
});
refreshCredentials();
}
export const GmailAuthSection = ({
gmailPublicCredential,
gmailServiceAccountCredential,
@@ -310,31 +389,49 @@ export const GmailAuthSection = ({
user,
}: DriveCredentialSectionProps) => {
const router = useRouter();
const [isAuthenticating, setIsAuthenticating] = useState(false);
const [localServiceAccountData, setLocalServiceAccountData] = useState(
serviceAccountKeyData
);
const [localAppCredentialData, setLocalAppCredentialData] =
useState(appCredentialData);
const [localGmailPublicCredential, setLocalGmailPublicCredential] = useState(
gmailPublicCredential
);
const [
localGmailServiceAccountCredential,
setLocalGmailServiceAccountCredential,
] = useState(gmailServiceAccountCredential);
// Update local state when props change
useEffect(() => {
setLocalServiceAccountData(serviceAccountKeyData);
setLocalAppCredentialData(appCredentialData);
setLocalGmailPublicCredential(gmailPublicCredential);
setLocalGmailServiceAccountCredential(gmailServiceAccountCredential);
}, [
serviceAccountKeyData,
appCredentialData,
gmailPublicCredential,
gmailServiceAccountCredential,
]);
const existingCredential =
gmailPublicCredential || gmailServiceAccountCredential;
localGmailPublicCredential || localGmailServiceAccountCredential;
if (existingCredential) {
return (
<>
<p className="mb-2 text-sm">
<i>Existing credential already set up!</i>
<i>Uploaded and authenticated credential already exists!</i>
</p>
<Button
onClick={async () => {
if (connectorExists) {
setPopup({
message:
"Cannot revoke access to Gmail while any connector is still set up. Please delete all connectors, then try again.",
type: "error",
});
return;
}
await adminDeleteCredential(existingCredential.id);
setPopup({
message: "Successfully revoked access to Gmail!",
type: "success",
});
refreshCredentials();
handleRevokeAccess(
connectorExists,
setPopup,
existingCredential,
refreshCredentials
);
}}
>
Revoke Access
@@ -343,20 +440,21 @@ export const GmailAuthSection = ({
);
}
if (serviceAccountKeyData?.service_account_email) {
if (localServiceAccountData?.service_account_email) {
return (
<div>
<CardSection>
<Formik
initialValues={{
google_primary_admin: user?.email || "",
}}
validationSchema={Yup.object().shape({
google_primary_admin: Yup.string().required(),
})}
onSubmit={async (values, formikHelpers) => {
formikHelpers.setSubmitting(true);
<Formik
initialValues={{
google_primary_admin: user?.email || "",
}}
validationSchema={Yup.object().shape({
google_primary_admin: Yup.string()
.email("Must be a valid email")
.required("Required"),
})}
onSubmit={async (values, formikHelpers) => {
formikHelpers.setSubmitting(true);
try {
const response = await fetch(
"/api/manage/admin/connector/gmail/service-account-credential",
{
@@ -375,6 +473,7 @@ export const GmailAuthSection = ({
message: "Successfully created service account credential",
type: "success",
});
refreshCredentials();
} else {
const errorMsg = await response.text();
setPopup({
@@ -382,65 +481,73 @@ export const GmailAuthSection = ({
type: "error",
});
}
refreshCredentials();
}}
>
{({ isSubmitting }) => (
<Form>
<TextFormField
name="google_primary_admin"
label="Primary Admin Email:"
subtext="You must provide an admin/owner account to retrieve all org emails."
/>
<div className="flex">
<button
type="submit"
disabled={isSubmitting}
className={
"bg-slate-500 hover:bg-slate-700 text-white " +
"font-bold py-2 px-4 rounded focus:outline-none " +
"focus:shadow-outline w-full max-w-sm mx-auto"
}
>
Submit
</button>
</div>
</Form>
)}
</Formik>
</CardSection>
} catch (error) {
setPopup({
message: `Failed to create service account credential - ${error}`,
type: "error",
});
} finally {
formikHelpers.setSubmitting(false);
}
}}
>
{({ isSubmitting }) => (
<Form>
<TextFormField
name="google_primary_admin"
label="Primary Admin Email:"
subtext="Enter the email of an admin/owner of the Google Organization that owns the Gmail account(s) you want to index."
/>
<div className="flex">
<Button type="submit" disabled={isSubmitting}>
Create Credential
</Button>
</div>
</Form>
)}
</Formik>
</div>
);
}
if (appCredentialData?.client_id) {
if (localAppCredentialData?.client_id) {
return (
<div className="text-sm mb-4">
<p className="mb-2">
Next, you must provide credentials via OAuth. This gives us read
access to the docs you have access to in your gmail account.
access to the emails you have access to in your Gmail account.
</p>
<Button
onClick={async () => {
const [authUrl, errorMsg] = await setupGmailOAuth({
isAdmin: true,
});
if (authUrl) {
// cookie used by callback to determine where to finally redirect to
setIsAuthenticating(true);
try {
Cookies.set(GMAIL_AUTH_IS_ADMIN_COOKIE_NAME, "true", {
path: "/",
});
router.push(authUrl);
return;
}
const [authUrl, errorMsg] = await setupGmailOAuth({
isAdmin: true,
});
setPopup({
message: errorMsg,
type: "error",
});
if (authUrl) {
router.push(authUrl);
} else {
setPopup({
message: errorMsg,
type: "error",
});
setIsAuthenticating(false);
}
} catch (error) {
setPopup({
message: `Failed to authenticate with Gmail - ${error}`,
type: "error",
});
setIsAuthenticating(false);
}
}}
disabled={isAuthenticating}
>
Authenticate with Gmail
{isAuthenticating ? "Authenticating..." : "Authenticate with Gmail"}
</Button>
</div>
);
@@ -449,8 +556,8 @@ export const GmailAuthSection = ({
// case where no keys have been uploaded in step 1
return (
<p className="text-sm">
Please upload an OAuth or Service Account Credential JSON in Step 1 before
moving onto Step 2.
Please upload either a OAuth Client Credential JSON or a Gmail Service
Account Key JSON in Step 1 before moving onto Step 2.
</p>
);
};

View File

@@ -1,10 +1,11 @@
"use client";
import useSWR from "swr";
import { errorHandlingFetcher } from "@/lib/fetcher";
import React from "react";
import { FetchError } from "@/lib/fetcher";
import { ErrorCallout } from "@/components/ErrorCallout";
import { LoadingAnimation } from "@/components/Loading";
import { usePopup } from "@/components/admin/connectors/Popup";
import { CCPairBasicInfo } from "@/lib/types";
import { PopupSpec, usePopup } from "@/components/admin/connectors/Popup";
import { CCPairBasicInfo, ValidSources } from "@/lib/types";
import {
Credential,
GmailCredentialJson,
@@ -14,26 +15,33 @@ import { GmailAuthSection, GmailJsonUploadSection } from "./Credential";
import { usePublicCredentials, useBasicConnectorStatus } from "@/lib/hooks";
import Title from "@/components/ui/title";
import { useUser } from "@/components/user/UserProvider";
import {
useGoogleAppCredential,
useGoogleServiceAccountKey,
useGoogleCredentials,
useConnectorsByCredentialId,
checkCredentialsFetched,
filterUploadedCredentials,
checkConnectorsExist,
refreshAllGoogleData,
} from "@/lib/googleConnector";
export const GmailMain = () => {
const { isAdmin, user } = useUser();
const { popup, setPopup } = usePopup();
const {
data: appCredentialData,
isLoading: isAppCredentialLoading,
error: isAppCredentialError,
} = useSWR<{ client_id: string }>(
"/api/manage/admin/connector/gmail/app-credential",
errorHandlingFetcher
);
} = useGoogleAppCredential("gmail");
const {
data: serviceAccountKeyData,
isLoading: isServiceAccountKeyLoading,
error: isServiceAccountKeyError,
} = useSWR<{ service_account_email: string }>(
"/api/manage/admin/connector/gmail/service-account-key",
errorHandlingFetcher
);
} = useGoogleServiceAccountKey("gmail");
const {
data: connectorIndexingStatuses,
isLoading: isConnectorIndexingStatusesLoading,
@@ -47,20 +55,45 @@ export const GmailMain = () => {
refreshCredentials,
} = usePublicCredentials();
const { popup, setPopup } = usePopup();
const {
data: gmailCredentials,
isLoading: isGmailCredentialsLoading,
error: gmailCredentialsError,
} = useGoogleCredentials(ValidSources.Gmail);
const appCredentialSuccessfullyFetched =
appCredentialData ||
(isAppCredentialError && isAppCredentialError.status === 404);
const serviceAccountKeySuccessfullyFetched =
serviceAccountKeyData ||
(isServiceAccountKeyError && isServiceAccountKeyError.status === 404);
const { credential_id, uploadedCredentials } =
filterUploadedCredentials(gmailCredentials);
const {
data: gmailConnectors,
isLoading: isGmailConnectorsLoading,
error: gmailConnectorsError,
refreshConnectorsByCredentialId,
} = useConnectorsByCredentialId(credential_id);
const {
appCredentialSuccessfullyFetched,
serviceAccountKeySuccessfullyFetched,
} = checkCredentialsFetched(
appCredentialData,
isAppCredentialError,
serviceAccountKeyData,
isServiceAccountKeyError
);
const handleRefresh = () => {
refreshCredentials();
refreshConnectorsByCredentialId();
refreshAllGoogleData(ValidSources.Gmail);
};
if (
(!appCredentialSuccessfullyFetched && isAppCredentialLoading) ||
(!serviceAccountKeySuccessfullyFetched && isServiceAccountKeyLoading) ||
(!connectorIndexingStatuses && isConnectorIndexingStatusesLoading) ||
(!credentialsData && isCredentialsLoading)
(!credentialsData && isCredentialsLoading) ||
(!gmailCredentials && isGmailCredentialsLoading) ||
(!gmailConnectors && isGmailConnectorsLoading)
) {
return (
<div className="mx-auto">
@@ -70,19 +103,15 @@ export const GmailMain = () => {
}
if (credentialsError || !credentialsData) {
return (
<div className="mx-auto">
<div className="text-red-500">Failed to load credentials.</div>
</div>
);
return <ErrorCallout errorTitle="Failed to load credentials." />;
}
if (gmailCredentialsError || !gmailCredentials) {
return <ErrorCallout errorTitle="Failed to load Gmail credentials." />;
}
if (connectorIndexingStatusesError || !connectorIndexingStatuses) {
return (
<div className="mx-auto">
<div className="text-red-500">Failed to load connectors.</div>
</div>
);
return <ErrorCallout errorTitle="Failed to load connectors." />;
}
if (
@@ -90,21 +119,28 @@ export const GmailMain = () => {
!serviceAccountKeySuccessfullyFetched
) {
return (
<div className="mx-auto">
<div className="text-red-500">
Error loading Gmail app credentials. Contact an administrator.
</div>
</div>
<ErrorCallout errorTitle="Error loading Gmail app credentials. Contact an administrator." />
);
}
const gmailPublicCredential: Credential<GmailCredentialJson> | undefined =
credentialsData.find(
(credential) =>
(credential.credential_json?.google_service_account_key ||
credential.credential_json?.google_tokens) &&
credential.admin_public
if (gmailConnectorsError) {
return (
<ErrorCallout errorTitle="Failed to load Gmail associated connectors." />
);
}
const connectorExistsFromCredential = checkConnectorsExist(gmailConnectors);
const gmailPublicUploadedCredential:
| Credential<GmailCredentialJson>
| undefined = credentialsData.find(
(credential) =>
credential.credential_json?.google_tokens &&
credential.admin_public &&
credential.source === "gmail" &&
credential.credential_json.authentication_method !== "oauth_interactive"
);
const gmailServiceAccountCredential:
| Credential<GmailServiceAccountCredentialJson>
| undefined = credentialsData.find(
@@ -118,6 +154,13 @@ export const GmailMain = () => {
(connectorIndexingStatus) => connectorIndexingStatus.source === "gmail"
);
const connectorExists =
connectorExistsFromCredential || gmailConnectorIndexingStatuses.length > 0;
const hasUploadedCredentials =
Boolean(appCredentialData?.client_id) ||
Boolean(serviceAccountKeyData?.service_account_email);
return (
<>
{popup}
@@ -129,21 +172,22 @@ export const GmailMain = () => {
appCredentialData={appCredentialData}
serviceAccountCredentialData={serviceAccountKeyData}
isAdmin={isAdmin}
onSuccess={handleRefresh}
/>
{isAdmin && (
{isAdmin && hasUploadedCredentials && (
<>
<Title className="mb-2 mt-6 ml-auto mr-auto">
Step 2: Authenticate with Onyx
</Title>
<GmailAuthSection
setPopup={setPopup}
refreshCredentials={refreshCredentials}
gmailPublicCredential={gmailPublicCredential}
refreshCredentials={handleRefresh}
gmailPublicCredential={gmailPublicUploadedCredential}
gmailServiceAccountCredential={gmailServiceAccountCredential}
appCredentialData={appCredentialData}
serviceAccountKeyData={serviceAccountKeyData}
connectorExists={gmailConnectorIndexingStatuses.length > 0}
connectorExists={connectorExists}
user={user}
/>
</>

View File

@@ -36,11 +36,14 @@ export function EmailPasswordForm({
{popup}
<Formik
initialValues={{
email: defaultEmail || "",
email: defaultEmail ? defaultEmail.toLowerCase() : "",
password: "",
}}
validationSchema={Yup.object().shape({
email: Yup.string().email().required(),
email: Yup.string()
.email()
.required()
.transform((value) => value.toLowerCase()),
password: Yup.string().required(),
})}
onSubmit={async (values) => {

View File

@@ -132,18 +132,12 @@ import {
import { getSourceMetadata } from "@/lib/sources";
import { UserSettingsModal } from "./modal/UserSettingsModal";
import { AlignStartVertical } from "lucide-react";
import { AgenticMessage } from "./message/AgenticMessage";
import AssistantModal from "../assistants/mine/AssistantModal";
import {
OperatingSystem,
useOperatingSystem,
useSidebarShortcut,
} from "@/lib/browserUtilities";
import { Button } from "@/components/ui/button";
import { useSidebarShortcut } from "@/lib/browserUtilities";
import { ConfirmEntityModal } from "@/components/modals/ConfirmEntityModal";
import { MessageChannel } from "node:worker_threads";
import { ChatSearchModal } from "./chat_search/ChatSearchModal";
import { ErrorBanner } from "./message/Resubmit";
const TEMP_USER_MESSAGE_ID = -1;
const TEMP_ASSISTANT_MESSAGE_ID = -2;
@@ -1169,6 +1163,7 @@ export function ChatPage({
navigatingAway.current = false;
let frozenSessionId = currentSessionId();
updateCanContinue(false, frozenSessionId);
setUncaughtError(null);
// Mark that we've sent a message for this session in the current page load
markSessionMessageSent(frozenSessionId);
@@ -1319,6 +1314,7 @@ export function ChatPage({
let isStreamingQuestions = true;
let includeAgentic = false;
let secondLevelMessageId: number | null = null;
let isAgentic: boolean = false;
let initialFetchDetails: null | {
user_message_id: number;
@@ -1481,6 +1477,9 @@ export function ChatPage({
second_level_generating = true;
}
}
if (Object.hasOwn(packet, "is_agentic")) {
isAgentic = (packet as any).is_agentic;
}
if (Object.hasOwn(packet, "refined_answer_improvement")) {
isImprovement = (packet as RefinedAnswerImprovement)
@@ -1514,6 +1513,7 @@ export function ChatPage({
);
} else if (Object.hasOwn(packet, "sub_question")) {
updateChatState("toolBuilding", frozenSessionId);
isAgentic = true;
is_generating = true;
sub_questions = constructSubQuestions(
sub_questions,
@@ -1714,6 +1714,7 @@ export function ChatPage({
sub_questions: sub_questions,
second_level_generating: second_level_generating,
agentic_docs: agenticDocs,
is_agentic: isAgentic,
},
...(includeAgentic
? [
@@ -2062,6 +2063,26 @@ export function ChatPage({
const [sharedChatSession, setSharedChatSession] =
useState<ChatSession | null>();
const handleResubmitLastMessage = () => {
// Grab the last user-type message
const lastUserMsg = messageHistory
.slice()
.reverse()
.find((m) => m.type === "user");
if (!lastUserMsg) {
setPopup({
message: "No previously-submitted user message found.",
type: "error",
});
return;
}
// We call onSubmit, passing a `messageOverride`
onSubmit({
messageIdToResend: lastUserMsg.messageId,
messageOverride: lastUserMsg.message,
});
};
const showShareModal = (chatSession: ChatSession) => {
setSharedChatSession(chatSession);
};
@@ -2644,9 +2665,9 @@ export function ChatPage({
: null
}
>
{message.sub_questions &&
message.sub_questions.length > 0 ? (
{message.is_agentic ? (
<AgenticMessage
resubmit={handleResubmitLastMessage}
error={uncaughtError}
isStreamingQuestions={
message.isStreamingQuestions ?? false
@@ -2994,21 +3015,18 @@ export function ChatPage({
currentPersona={liveAssistant}
messageId={message.messageId}
content={
<p className="text-red-700 text-sm my-auto">
{message.message}
{message.stackTrace && (
<span
onClick={() =>
setStackTraceModalContent(
message.stackTrace!
)
}
className="ml-2 cursor-pointer underline"
>
Show stack trace.
</span>
)}
</p>
<ErrorBanner
resubmit={handleResubmitLastMessage}
error={message.message}
showStackTrace={
message.stackTrace
? () =>
setStackTraceModalContent(
message.stackTrace!
)
: undefined
}
/>
}
/>
</div>

View File

@@ -15,8 +15,8 @@ export function ChatSearchGroup({
}: ChatSearchGroupProps) {
return (
<div className="mb-4">
<div className="sticky -top-1 mt-1 z-10 bg-[#fff]/90 dark:bg-gray-800/90 py-2 px-4 px-4">
<div className="text-xs font-medium leading-4 text-gray-600 dark:text-gray-400">
<div className="sticky -top-1 mt-1 z-10 bg-[#fff]/90 dark:bg-neutral-800/90 py-2 px-4 px-4">
<div className="text-xs font-medium leading-4 text-neutral-600 dark:text-neutral-400">
{title}
</div>
</div>

View File

@@ -1,6 +1,7 @@
import React from "react";
import { MessageSquare } from "lucide-react";
import { ChatSessionSummary } from "./interfaces";
import { truncateString } from "@/lib/utils";
interface ChatSearchItemProps {
chat: ChatSessionSummary;
@@ -11,12 +12,12 @@ export function ChatSearchItem({ chat, onSelect }: ChatSearchItemProps) {
return (
<li>
<div className="cursor-pointer" onClick={() => onSelect(chat.id)}>
<div className="group relative flex flex-col rounded-lg px-4 py-3 hover:bg-neutral-100 dark:hover:bg-neutral-800">
<div className="flex items-center">
<div className="group relative flex flex-col rounded-lg px-4 py-3 hover:bg-neutral-100 dark:hover:bg-neutral-700">
<div className="flex max-w-full mx-2 items-center">
<MessageSquare className="h-5 w-5 text-neutral-600 dark:text-neutral-400" />
<div className="relative grow overflow-hidden whitespace-nowrap pl-4">
<div className="text-sm dark:text-neutral-200">
{chat.name || "Untitled Chat"}
<div className="relative max-w-full grow overflow-hidden whitespace-nowrap pl-4">
<div className="text-sm max-w-full dark:text-neutral-200">
{truncateString(chat.name || "Untitled Chat", 90)}
</div>
</div>
<div className="opacity-0 group-hover:opacity-100 transition-opacity text-xs text-neutral-500 dark:text-neutral-400">

View File

@@ -168,7 +168,7 @@ const FolderItem = ({
};
const folders = folder.chat_sessions.sort((a, b) => {
return a.time_created.localeCompare(b.time_created);
return a.time_updated.localeCompare(b.time_updated);
});
// Determine whether to show the trash can icon

View File

@@ -70,6 +70,7 @@ export interface ChatSession {
name: string;
persona_id: number;
time_created: string;
time_updated: string;
shared_status: ChatSessionSharedStatus;
folder_id: number | null;
current_alternate_model: string;
@@ -103,6 +104,7 @@ export interface Message {
overridden_model?: string;
stopReason?: StreamStopReason | null;
sub_questions?: SubQuestionDetail[] | null;
is_agentic?: boolean | null;
// Streaming only
second_level_generating?: boolean;
@@ -122,6 +124,7 @@ export interface BackendChatSession {
persona_icon_shape: number | null;
messages: BackendMessage[];
time_created: string;
time_updated: string;
shared_status: ChatSessionSharedStatus;
current_temperature_override: number | null;
current_alternate_model?: string;
@@ -148,6 +151,7 @@ export interface BackendMessage {
comments: any;
parentMessageId: number | null;
refined_answer_improvement: boolean | null;
is_agentic: boolean | null;
}
export interface MessageResponseIDInfo {

View File

@@ -48,10 +48,10 @@ export function getChatRetentionInfo(
): ChatRetentionInfo {
// If `maximum_chat_retention_days` isn't set- never display retention warning.
const chatRetentionDays = settings.maximum_chat_retention_days || 10000;
const createdDate = new Date(chatSession.time_created);
const updatedDate = new Date(chatSession.time_updated);
const today = new Date();
const daysFromCreation = Math.ceil(
(today.getTime() - createdDate.getTime()) / (1000 * 3600 * 24)
(today.getTime() - updatedDate.getTime()) / (1000 * 3600 * 24)
);
const daysUntilExpiration = chatRetentionDays - daysFromCreation;
const showRetentionWarning =
@@ -419,7 +419,7 @@ export function groupSessionsByDateRange(chatSessions: ChatSession[]) {
};
chatSessions.forEach((chatSession) => {
const chatSessionDate = new Date(chatSession.time_created);
const chatSessionDate = new Date(chatSession.time_updated);
const diffTime = today.getTime() - chatSessionDate.getTime();
const diffDays = diffTime / (1000 * 3600 * 24); // Convert time difference to days
@@ -501,6 +501,7 @@ export function processRawChatHistory(
sub_questions: subQuestions,
isImprovement:
(messageInfo.refined_answer_improvement as unknown as boolean) || false,
is_agentic: messageInfo.is_agentic,
};
messages.set(messageInfo.message_id, message);

View File

@@ -50,6 +50,9 @@ import "katex/dist/katex.min.css";
import SubQuestionsDisplay from "./SubQuestionsDisplay";
import { StatusRefinement } from "../Refinement";
import { copyAll, handleCopy } from "./copyingUtils";
import { Button } from "@/components/ui/button";
import { RefreshCw } from "lucide-react";
import { ErrorBanner, Resubmit } from "./Resubmit";
export const AgenticMessage = ({
isStreamingQuestions,
@@ -84,7 +87,9 @@ export const AgenticMessage = ({
secondLevelSubquestions,
toggleDocDisplay,
error,
resubmit,
}: {
resubmit?: () => void;
isStreamingQuestions: boolean;
isGenerating: boolean;
docSidebarToggled?: boolean;
@@ -455,7 +460,6 @@ export const AgenticMessage = ({
finalContent.length > 8) ||
(files && files.length > 0) ? (
<>
{/* <FileDisplay files={files || []} /> */}
<div className="w-full py-4 flex flex-col gap-4">
<div className="flex items-center gap-x-2 px-4">
<div className="text-black text-lg font-medium">
@@ -503,9 +507,7 @@ export const AgenticMessage = ({
content
)}
{error && (
<p className="mt-2 text-red-700 text-sm my-auto">
{error}
</p>
<ErrorBanner error={error} resubmit={resubmit} />
)}
</div>
</div>
@@ -513,15 +515,13 @@ export const AgenticMessage = ({
) : isComplete ? (
error && (
<p className="mt-2 mx-4 text-red-700 text-sm my-auto">
{error}
<ErrorBanner error={error} resubmit={resubmit} />
</p>
)
) : (
<>
{error && (
<p className="mt-2 mx-4 text-red-700 text-sm my-auto">
{error}
</p>
<ErrorBanner error={error} resubmit={resubmit} />
)}
</>
)}

View File

@@ -0,0 +1,58 @@
import { Alert, AlertDescription, AlertTitle } from "@/components/ui/alert";
import { AlertCircle } from "lucide-react";
import { Button } from "@/components/ui/button";
import { RefreshCw } from "lucide-react";
interface ResubmitProps {
resubmit: () => void;
}
export const Resubmit: React.FC<ResubmitProps> = ({ resubmit }) => {
return (
<div className="flex flex-col items-center justify-center gap-y-2 mt-4">
<p className="text-sm text-neutral-700 dark:text-neutral-300">
There was an error with the response.
</p>
<Button
onClick={resubmit}
variant="agent"
size="sm"
className="flex items-center gap-2 text-white font-medium py-2 px-4 rounded"
>
<RefreshCw className="w-4 h-4" />
Regenerate
</Button>
</div>
);
};
export const ErrorBanner = ({
error,
showStackTrace,
resubmit,
}: {
error: string;
showStackTrace?: () => void;
resubmit?: () => void;
}) => {
return (
<div className="text-red-700 mt-4 text-sm my-auto">
<Alert variant="broken">
<AlertCircle className="h-4 w-4" />
<AlertTitle>Error</AlertTitle>
<AlertDescription className="flex gap-x-2">
{error}
{showStackTrace && (
<span
className="text-red-600 hover:text-red-800 cursor-pointer underline"
onClick={showStackTrace}
>
Show stack trace
</span>
)}
</AlertDescription>
</Alert>
{resubmit && <Resubmit resubmit={resubmit} />}
</div>
);
};

View File

@@ -206,7 +206,7 @@ export function SharedChatDisplay({
{chatSession.description || `Unnamed Chat`}
</h1>
<p className=" text-text-darker">
{humanReadableFormat(chatSession.time_created)}
{humanReadableFormat(chatSession.time_updated)}
</p>
<div
className={`

View File

@@ -3,7 +3,6 @@ interface Props {
onClick?: React.MouseEventHandler<HTMLButtonElement>;
type?: "button" | "submit" | "reset";
disabled?: boolean;
fullWidth?: boolean;
className?: string;
}
@@ -12,14 +11,12 @@ export const Button = ({
onClick,
type = "submit",
disabled = false,
fullWidth = false,
className = "",
}: Props) => {
return (
<button
className={
"group relative " +
(fullWidth ? "w-full " : "") +
"py-1 px-2 border border-transparent text-sm " +
"font-medium rounded-md text-white " +
"focus:outline-none focus:ring-2 " +

View File

@@ -8,8 +8,10 @@ const alertVariants = cva(
{
variants: {
variant: {
broken:
"border-red-500/50 text-red-500 dark:border-red-500 [&>svg]:text-red-500 dark:border-red-900/50 dark:text-red-100 dark:dark:border-red-900 dark:[&>svg]:text-red-700 bg-red-50 dark:bg-red-950",
ark: "border-amber-500/50 text-amber-500 dark:border-amber-500 [&>svg]:text-amber-500 dark:border-amber-900/50 dark:text-amber-900 dark:dark:border-amber-900 dark:[&>svg]:text-amber-900 bg-amber-50 dark:bg-amber-950",
info: "border-black/50 dark:border-black dark:border-black/50 dark:dark:border-black",
default:
"bg-neutral-50 text-neutral-darker dark:bg-neutral-950 dark:text-text",
destructive:

View File

@@ -9,6 +9,8 @@ const buttonVariants = cva(
{
variants: {
variant: {
agent:
"bg-agent text-white hover:bg-agent-hovered dark:bg-agent dark:text-white dark:hover:bg-agent/90",
success:
"bg-green-100 text-green-600 hover:bg-green-500/90 dark:bg-green-700 dark:text-green-100 dark:hover:bg-green-600/90",
"success-reverse":

View File

@@ -148,7 +148,7 @@ export async function fetchChatData(searchParams: {
chatSessions.sort(
(a, b) =>
new Date(b.time_created).getTime() - new Date(a.time_created).getTime()
new Date(b.time_updated).getTime() - new Date(a.time_updated).getTime()
);
let documentSets: DocumentSet[] = [];

View File

@@ -0,0 +1,120 @@
import useSWR, { mutate } from "swr";
import { FetchError, errorHandlingFetcher } from "@/lib/fetcher";
import { Credential } from "@/lib/connectors/credentials";
import { ConnectorSnapshot } from "@/lib/connectors/connectors";
import { ValidSources } from "@/lib/types";
import { buildSimilarCredentialInfoURL } from "@/app/admin/connector/[ccPairId]/lib";
// Constants for service names to avoid typos
export const GOOGLE_SERVICES = {
GMAIL: "gmail",
GOOGLE_DRIVE: "google-drive",
} as const;
export const useGoogleAppCredential = (service: "gmail" | "google_drive") => {
const endpoint = `/api/manage/admin/connector/${
service === "gmail" ? GOOGLE_SERVICES.GMAIL : GOOGLE_SERVICES.GOOGLE_DRIVE
}/app-credential`;
return useSWR<{ client_id: string }, FetchError>(
endpoint,
errorHandlingFetcher
);
};
export const useGoogleServiceAccountKey = (
service: "gmail" | "google_drive"
) => {
const endpoint = `/api/manage/admin/connector/${
service === "gmail" ? GOOGLE_SERVICES.GMAIL : GOOGLE_SERVICES.GOOGLE_DRIVE
}/service-account-key`;
return useSWR<{ service_account_email: string }, FetchError>(
endpoint,
errorHandlingFetcher
);
};
export const useGoogleCredentials = (
source: ValidSources.Gmail | ValidSources.GoogleDrive
) => {
return useSWR<Credential<any>[]>(
buildSimilarCredentialInfoURL(source),
errorHandlingFetcher,
{ refreshInterval: 5000 }
);
};
export const useConnectorsByCredentialId = (credential_id: number | null) => {
let url: string | null = null;
if (credential_id !== null) {
url = `/api/manage/admin/connector?credential=${credential_id}`;
}
const swrResponse = useSWR<ConnectorSnapshot[]>(url, errorHandlingFetcher);
return {
...swrResponse,
refreshConnectorsByCredentialId: () => mutate(url),
};
};
export const checkCredentialsFetched = (
appCredentialData: any,
appCredentialError: FetchError | undefined,
serviceAccountKeyData: any,
serviceAccountKeyError: FetchError | undefined
) => {
const appCredentialSuccessfullyFetched =
appCredentialData ||
(appCredentialError && appCredentialError.status === 404);
const serviceAccountKeySuccessfullyFetched =
serviceAccountKeyData ||
(serviceAccountKeyError && serviceAccountKeyError.status === 404);
return {
appCredentialSuccessfullyFetched,
serviceAccountKeySuccessfullyFetched,
};
};
export const filterUploadedCredentials = <
T extends { authentication_method?: string },
>(
credentials: Credential<T>[] | undefined
): { credential_id: number | null; uploadedCredentials: Credential<T>[] } => {
let credential_id = null;
let uploadedCredentials: Credential<T>[] = [];
if (credentials) {
uploadedCredentials = credentials.filter(
(credential) =>
credential.credential_json.authentication_method !== "oauth_interactive"
);
if (uploadedCredentials.length > 0) {
credential_id = uploadedCredentials[0].id;
}
}
return { credential_id, uploadedCredentials };
};
export const checkConnectorsExist = (
connectors: ConnectorSnapshot[] | undefined
): boolean => {
return !!connectors && connectors.length > 0;
};
export const refreshAllGoogleData = (
source: ValidSources.Gmail | ValidSources.GoogleDrive
) => {
mutate(buildSimilarCredentialInfoURL(source));
const service =
source === ValidSources.Gmail
? GOOGLE_SERVICES.GMAIL
: GOOGLE_SERVICES.GOOGLE_DRIVE;
mutate(`/api/manage/admin/connector/${service}/app-credential`);
mutate(`/api/manage/admin/connector/${service}/service-account-key`);
};