mirror of
https://github.com/onyx-dot-app/onyx.git
synced 2026-04-22 01:46:47 +00:00
Compare commits
21 Commits
v3.0.9
...
nikg/std-e
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
435ac4d06b | ||
|
|
6041773f60 | ||
|
|
63d3efd380 | ||
|
|
ec978d9a3f | ||
|
|
d4d98a6cd0 | ||
|
|
dc40e86dac | ||
|
|
e495f7a13e | ||
|
|
4761e4b132 | ||
|
|
6b5ab54b85 | ||
|
|
959cf444f8 | ||
|
|
2ebccea6d6 | ||
|
|
5fe7a474db | ||
|
|
9d7dc3da21 | ||
|
|
2899be4c5e | ||
|
|
64ee7fc23f | ||
|
|
e07764285d | ||
|
|
cc2e6ffa8a | ||
|
|
d3ee5c9b59 | ||
|
|
dfa0efc093 | ||
|
|
9aad4077f1 | ||
|
|
29d9ebf7b3 |
2
.github/workflows/deployment.yml
vendored
2
.github/workflows/deployment.yml
vendored
@@ -151,7 +151,7 @@ jobs:
|
||||
fetch-depth: 0
|
||||
|
||||
- name: Setup uv
|
||||
uses: astral-sh/setup-uv@61cb8a9741eeb8a550a1b8544337180c0fc8476b # ratchet:astral-sh/setup-uv@v7
|
||||
uses: astral-sh/setup-uv@5a095e7a2014a4212f075830d4f7277575a9d098 # ratchet:astral-sh/setup-uv@v7
|
||||
with:
|
||||
version: "0.9.9"
|
||||
# NOTE: This isn't caching much and zizmor suggests this could be poisoned, so disable.
|
||||
|
||||
@@ -70,7 +70,7 @@ jobs:
|
||||
|
||||
- name: Install the latest version of uv
|
||||
if: steps.gate.outputs.should_cherrypick == 'true'
|
||||
uses: astral-sh/setup-uv@61cb8a9741eeb8a550a1b8544337180c0fc8476b # ratchet:astral-sh/setup-uv@v7
|
||||
uses: astral-sh/setup-uv@5a095e7a2014a4212f075830d4f7277575a9d098 # ratchet:astral-sh/setup-uv@v7
|
||||
with:
|
||||
enable-cache: false
|
||||
version: "0.9.9"
|
||||
|
||||
4
.github/workflows/pr-playwright-tests.yml
vendored
4
.github/workflows/pr-playwright-tests.yml
vendored
@@ -468,7 +468,7 @@ jobs:
|
||||
|
||||
- name: Install the latest version of uv
|
||||
if: always()
|
||||
uses: astral-sh/setup-uv@61cb8a9741eeb8a550a1b8544337180c0fc8476b # ratchet:astral-sh/setup-uv@v7
|
||||
uses: astral-sh/setup-uv@5a095e7a2014a4212f075830d4f7277575a9d098 # ratchet:astral-sh/setup-uv@v7
|
||||
with:
|
||||
enable-cache: false
|
||||
version: "0.9.9"
|
||||
@@ -707,7 +707,7 @@ jobs:
|
||||
pull-requests: write
|
||||
steps:
|
||||
- name: Download visual diff summaries
|
||||
uses: actions/download-artifact@37930b1c2abaa49bbe596cd826c3c89aef350131
|
||||
uses: actions/download-artifact@70fc10c6e5e1ce46ad2ea6f2b72d43f7d47b13c3
|
||||
with:
|
||||
pattern: screenshot-diff-summary-*
|
||||
path: summaries/
|
||||
|
||||
2
.github/workflows/pr-quality-checks.yml
vendored
2
.github/workflows/pr-quality-checks.yml
vendored
@@ -28,7 +28,7 @@ jobs:
|
||||
with:
|
||||
python-version: "3.11"
|
||||
- name: Setup Terraform
|
||||
uses: hashicorp/setup-terraform@b9cd54a3c349d3f38e8881555d616ced269862dd # ratchet:hashicorp/setup-terraform@v3
|
||||
uses: hashicorp/setup-terraform@5e8dbf3c6d9deaf4193ca7a8fb23f2ac83bb6c85 # ratchet:hashicorp/setup-terraform@v4.0.0
|
||||
- name: Setup node
|
||||
uses: actions/setup-node@6044e13b5dc448c55e2357c09f80417699197238 # ratchet:actions/setup-node@v6
|
||||
with: # zizmor: ignore[cache-poisoning]
|
||||
|
||||
2
.github/workflows/release-devtools.yml
vendored
2
.github/workflows/release-devtools.yml
vendored
@@ -26,7 +26,7 @@ jobs:
|
||||
- uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # ratchet:actions/checkout@v6
|
||||
with:
|
||||
persist-credentials: false
|
||||
- uses: astral-sh/setup-uv@61cb8a9741eeb8a550a1b8544337180c0fc8476b # ratchet:astral-sh/setup-uv@v7
|
||||
- uses: astral-sh/setup-uv@5a095e7a2014a4212f075830d4f7277575a9d098 # ratchet:astral-sh/setup-uv@v7
|
||||
with:
|
||||
enable-cache: false
|
||||
version: "0.9.9"
|
||||
|
||||
2
.github/workflows/zizmor.yml
vendored
2
.github/workflows/zizmor.yml
vendored
@@ -24,7 +24,7 @@ jobs:
|
||||
persist-credentials: false
|
||||
|
||||
- name: Install the latest version of uv
|
||||
uses: astral-sh/setup-uv@61cb8a9741eeb8a550a1b8544337180c0fc8476b # ratchet:astral-sh/setup-uv@v7
|
||||
uses: astral-sh/setup-uv@5a095e7a2014a4212f075830d4f7277575a9d098 # ratchet:astral-sh/setup-uv@v7
|
||||
with:
|
||||
enable-cache: false
|
||||
version: "0.9.9"
|
||||
|
||||
@@ -141,6 +141,7 @@ COPY --chown=onyx:onyx ./scripts/debugging /app/scripts/debugging
|
||||
COPY --chown=onyx:onyx ./scripts/force_delete_connector_by_id.py /app/scripts/force_delete_connector_by_id.py
|
||||
COPY --chown=onyx:onyx ./scripts/supervisord_entrypoint.sh /app/scripts/supervisord_entrypoint.sh
|
||||
COPY --chown=onyx:onyx ./scripts/setup_craft_templates.sh /app/scripts/setup_craft_templates.sh
|
||||
COPY --chown=onyx:onyx ./scripts/reencrypt_secrets.py /app/scripts/reencrypt_secrets.py
|
||||
RUN chmod +x /app/scripts/supervisord_entrypoint.sh /app/scripts/setup_craft_templates.sh
|
||||
|
||||
# Run Craft template setup at build time when ENABLE_CRAFT=true
|
||||
|
||||
@@ -11,7 +11,6 @@ from sqlalchemy import text
|
||||
from alembic import op
|
||||
from onyx.configs.app_configs import DB_READONLY_PASSWORD
|
||||
from onyx.configs.app_configs import DB_READONLY_USER
|
||||
from shared_configs.configs import MULTI_TENANT
|
||||
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
@@ -22,59 +21,52 @@ depends_on = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
if MULTI_TENANT:
|
||||
# Enable pg_trgm extension if not already enabled
|
||||
op.execute("CREATE EXTENSION IF NOT EXISTS pg_trgm")
|
||||
|
||||
# Enable pg_trgm extension if not already enabled
|
||||
op.execute("CREATE EXTENSION IF NOT EXISTS pg_trgm")
|
||||
# Create the read-only db user if it does not already exist.
|
||||
if not (DB_READONLY_USER and DB_READONLY_PASSWORD):
|
||||
raise Exception("DB_READONLY_USER or DB_READONLY_PASSWORD is not set")
|
||||
|
||||
# Create read-only db user here only in multi-tenant mode. For single-tenant mode,
|
||||
# the user is created in the standard migration.
|
||||
if not (DB_READONLY_USER and DB_READONLY_PASSWORD):
|
||||
raise Exception("DB_READONLY_USER or DB_READONLY_PASSWORD is not set")
|
||||
|
||||
op.execute(
|
||||
text(
|
||||
f"""
|
||||
DO $$
|
||||
BEGIN
|
||||
-- Check if the read-only user already exists
|
||||
IF NOT EXISTS (SELECT FROM pg_catalog.pg_roles WHERE rolname = '{DB_READONLY_USER}') THEN
|
||||
-- Create the read-only user with the specified password
|
||||
EXECUTE format('CREATE USER %I WITH PASSWORD %L', '{DB_READONLY_USER}', '{DB_READONLY_PASSWORD}');
|
||||
-- First revoke all privileges to ensure a clean slate
|
||||
EXECUTE format('REVOKE ALL ON DATABASE %I FROM %I', current_database(), '{DB_READONLY_USER}');
|
||||
-- Grant only the CONNECT privilege to allow the user to connect to the database
|
||||
-- but not perform any operations without additional specific grants
|
||||
EXECUTE format('GRANT CONNECT ON DATABASE %I TO %I', current_database(), '{DB_READONLY_USER}');
|
||||
END IF;
|
||||
END
|
||||
$$;
|
||||
"""
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
if MULTI_TENANT:
|
||||
# Drop read-only db user here only in single tenant mode. For multi-tenant mode,
|
||||
# the user is dropped in the alembic_tenants migration.
|
||||
|
||||
op.execute(
|
||||
text(
|
||||
f"""
|
||||
op.execute(
|
||||
text(
|
||||
f"""
|
||||
DO $$
|
||||
BEGIN
|
||||
IF EXISTS (SELECT FROM pg_catalog.pg_roles WHERE rolname = '{DB_READONLY_USER}') THEN
|
||||
-- First revoke all privileges from the database
|
||||
-- Check if the read-only user already exists
|
||||
IF NOT EXISTS (SELECT FROM pg_catalog.pg_roles WHERE rolname = '{DB_READONLY_USER}') THEN
|
||||
-- Create the read-only user with the specified password
|
||||
EXECUTE format('CREATE USER %I WITH PASSWORD %L', '{DB_READONLY_USER}', '{DB_READONLY_PASSWORD}');
|
||||
-- First revoke all privileges to ensure a clean slate
|
||||
EXECUTE format('REVOKE ALL ON DATABASE %I FROM %I', current_database(), '{DB_READONLY_USER}');
|
||||
-- Then revoke all privileges from the public schema
|
||||
EXECUTE format('REVOKE ALL ON SCHEMA public FROM %I', '{DB_READONLY_USER}');
|
||||
-- Then drop the user
|
||||
EXECUTE format('DROP USER %I', '{DB_READONLY_USER}');
|
||||
-- Grant only the CONNECT privilege to allow the user to connect to the database
|
||||
-- but not perform any operations without additional specific grants
|
||||
EXECUTE format('GRANT CONNECT ON DATABASE %I TO %I', current_database(), '{DB_READONLY_USER}');
|
||||
END IF;
|
||||
END
|
||||
$$;
|
||||
"""
|
||||
)
|
||||
"""
|
||||
)
|
||||
op.execute(text("DROP EXTENSION IF EXISTS pg_trgm"))
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
op.execute(
|
||||
text(
|
||||
f"""
|
||||
DO $$
|
||||
BEGIN
|
||||
IF EXISTS (SELECT FROM pg_catalog.pg_roles WHERE rolname = '{DB_READONLY_USER}') THEN
|
||||
-- First revoke all privileges from the database
|
||||
EXECUTE format('REVOKE ALL ON DATABASE %I FROM %I', current_database(), '{DB_READONLY_USER}');
|
||||
-- Then revoke all privileges from the public schema
|
||||
EXECUTE format('REVOKE ALL ON SCHEMA public FROM %I', '{DB_READONLY_USER}');
|
||||
-- Then drop the user
|
||||
EXECUTE format('DROP USER %I', '{DB_READONLY_USER}');
|
||||
END IF;
|
||||
END
|
||||
$$;
|
||||
"""
|
||||
)
|
||||
)
|
||||
op.execute(text("DROP EXTENSION IF EXISTS pg_trgm"))
|
||||
|
||||
@@ -68,6 +68,7 @@ def get_external_access_for_raw_gdrive_file(
|
||||
company_domain: str,
|
||||
retriever_drive_service: GoogleDriveService | None,
|
||||
admin_drive_service: GoogleDriveService,
|
||||
fallback_user_email: str,
|
||||
add_prefix: bool = False,
|
||||
) -> ExternalAccess:
|
||||
"""
|
||||
@@ -79,6 +80,11 @@ def get_external_access_for_raw_gdrive_file(
|
||||
set add_prefix to True so group IDs are prefixed with the source type.
|
||||
When invoked from doc_sync (permission sync), use the default (False)
|
||||
since upsert_document_external_perms handles prefixing.
|
||||
fallback_user_email: When we cannot retrieve any permission info for a file
|
||||
(e.g. externally-owned files where the API returns no permissions
|
||||
and permissions.list returns 403), fall back to granting access
|
||||
to this user. This is typically the impersonated org user whose
|
||||
drive contained the file.
|
||||
"""
|
||||
doc_id = file.get("id")
|
||||
if not doc_id:
|
||||
@@ -117,6 +123,26 @@ def get_external_access_for_raw_gdrive_file(
|
||||
[permissions_list, backup_permissions_list]
|
||||
)
|
||||
|
||||
# For externally-owned files, the Drive API may return no permissions
|
||||
# and permissions.list may return 403. In this case, fall back to
|
||||
# granting access to the user who found the file in their drive.
|
||||
# Note, even if other users also have access to this file,
|
||||
# they will not be granted access in Onyx.
|
||||
# We check permissions_list (the final result after all fetch attempts)
|
||||
# rather than the raw fields, because permission_ids may be present
|
||||
# but the actual fetch can still return empty due to a 403.
|
||||
if not permissions_list:
|
||||
logger.info(
|
||||
f"No permission info available for file {doc_id} "
|
||||
f"(likely owned by a user outside of your organization). "
|
||||
f"Falling back to granting access to retriever user: {fallback_user_email}"
|
||||
)
|
||||
return ExternalAccess(
|
||||
external_user_emails={fallback_user_email},
|
||||
external_user_group_ids=set(),
|
||||
is_public=False,
|
||||
)
|
||||
|
||||
folder_ids_to_inherit_permissions_from: set[str] = set()
|
||||
user_emails: set[str] = set()
|
||||
group_emails: set[str] = set()
|
||||
|
||||
@@ -14,67 +14,91 @@ from onyx.utils.variable_functionality import fetch_versioned_implementation
|
||||
logger = setup_logger()
|
||||
|
||||
|
||||
@lru_cache(maxsize=1)
|
||||
@lru_cache(maxsize=2)
|
||||
def _get_trimmed_key(key: str) -> bytes:
|
||||
encoded_key = key.encode()
|
||||
key_length = len(encoded_key)
|
||||
if key_length < 16:
|
||||
raise RuntimeError("Invalid ENCRYPTION_KEY_SECRET - too short")
|
||||
elif key_length > 32:
|
||||
key = key[:32]
|
||||
elif key_length not in (16, 24, 32):
|
||||
valid_lengths = [16, 24, 32]
|
||||
key = key[: min(valid_lengths, key=lambda x: abs(x - key_length))]
|
||||
|
||||
return encoded_key
|
||||
# Trim to the largest valid AES key size that fits
|
||||
valid_lengths = [32, 24, 16]
|
||||
for size in valid_lengths:
|
||||
if key_length >= size:
|
||||
return encoded_key[:size]
|
||||
|
||||
raise AssertionError("unreachable")
|
||||
|
||||
|
||||
def _encrypt_string(input_str: str) -> bytes:
|
||||
if not ENCRYPTION_KEY_SECRET:
|
||||
def _encrypt_string(input_str: str, key: str | None = None) -> bytes:
|
||||
effective_key = key if key is not None else ENCRYPTION_KEY_SECRET
|
||||
if not effective_key:
|
||||
return input_str.encode()
|
||||
|
||||
key = _get_trimmed_key(ENCRYPTION_KEY_SECRET)
|
||||
trimmed = _get_trimmed_key(effective_key)
|
||||
iv = urandom(16)
|
||||
padder = padding.PKCS7(algorithms.AES.block_size).padder()
|
||||
padded_data = padder.update(input_str.encode()) + padder.finalize()
|
||||
|
||||
cipher = Cipher(algorithms.AES(key), modes.CBC(iv), backend=default_backend())
|
||||
cipher = Cipher(algorithms.AES(trimmed), modes.CBC(iv), backend=default_backend())
|
||||
encryptor = cipher.encryptor()
|
||||
encrypted_data = encryptor.update(padded_data) + encryptor.finalize()
|
||||
|
||||
return iv + encrypted_data
|
||||
|
||||
|
||||
def _decrypt_bytes(input_bytes: bytes) -> str:
|
||||
if not ENCRYPTION_KEY_SECRET:
|
||||
def _decrypt_bytes(input_bytes: bytes, key: str | None = None) -> str:
|
||||
effective_key = key if key is not None else ENCRYPTION_KEY_SECRET
|
||||
if not effective_key:
|
||||
return input_bytes.decode()
|
||||
|
||||
key = _get_trimmed_key(ENCRYPTION_KEY_SECRET)
|
||||
iv = input_bytes[:16]
|
||||
encrypted_data = input_bytes[16:]
|
||||
trimmed = _get_trimmed_key(effective_key)
|
||||
try:
|
||||
iv = input_bytes[:16]
|
||||
encrypted_data = input_bytes[16:]
|
||||
|
||||
cipher = Cipher(algorithms.AES(key), modes.CBC(iv), backend=default_backend())
|
||||
decryptor = cipher.decryptor()
|
||||
decrypted_padded_data = decryptor.update(encrypted_data) + decryptor.finalize()
|
||||
cipher = Cipher(
|
||||
algorithms.AES(trimmed), modes.CBC(iv), backend=default_backend()
|
||||
)
|
||||
decryptor = cipher.decryptor()
|
||||
decrypted_padded_data = decryptor.update(encrypted_data) + decryptor.finalize()
|
||||
|
||||
unpadder = padding.PKCS7(algorithms.AES.block_size).unpadder()
|
||||
decrypted_data = unpadder.update(decrypted_padded_data) + unpadder.finalize()
|
||||
unpadder = padding.PKCS7(algorithms.AES.block_size).unpadder()
|
||||
decrypted_data = unpadder.update(decrypted_padded_data) + unpadder.finalize()
|
||||
|
||||
return decrypted_data.decode()
|
||||
return decrypted_data.decode()
|
||||
except (ValueError, UnicodeDecodeError):
|
||||
if key is not None:
|
||||
# Explicit key was provided — don't fall back silently
|
||||
raise
|
||||
# Read path: attempt raw UTF-8 decode as a fallback for legacy data.
|
||||
# Does NOT handle data encrypted with a different key — that
|
||||
# ciphertext is not valid UTF-8 and will raise below.
|
||||
logger.warning(
|
||||
"AES decryption failed — falling back to raw decode. "
|
||||
"Run the re-encrypt secrets script to rotate to the current key."
|
||||
)
|
||||
try:
|
||||
return input_bytes.decode()
|
||||
except UnicodeDecodeError:
|
||||
raise ValueError(
|
||||
"Data is not valid UTF-8 — likely encrypted with a different key. "
|
||||
"Run the re-encrypt secrets script to rotate to the current key."
|
||||
) from None
|
||||
|
||||
|
||||
def encrypt_string_to_bytes(input_str: str) -> bytes:
|
||||
def encrypt_string_to_bytes(input_str: str, key: str | None = None) -> bytes:
|
||||
versioned_encryption_fn = fetch_versioned_implementation(
|
||||
"onyx.utils.encryption", "_encrypt_string"
|
||||
)
|
||||
return versioned_encryption_fn(input_str)
|
||||
return versioned_encryption_fn(input_str, key=key)
|
||||
|
||||
|
||||
def decrypt_bytes_to_string(input_bytes: bytes) -> str:
|
||||
def decrypt_bytes_to_string(input_bytes: bytes, key: str | None = None) -> str:
|
||||
versioned_decryption_fn = fetch_versioned_implementation(
|
||||
"onyx.utils.encryption", "_decrypt_bytes"
|
||||
)
|
||||
return versioned_decryption_fn(input_bytes)
|
||||
return versioned_decryption_fn(input_bytes, key=key)
|
||||
|
||||
|
||||
def test_encryption() -> None:
|
||||
|
||||
@@ -4,10 +4,11 @@ from typing import Any
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from fastapi import APIRouter
|
||||
from fastapi import HTTPException
|
||||
from fastapi import Request
|
||||
|
||||
from model_server.utils import simple_log_function_time
|
||||
from onyx.error_handling.error_codes import OnyxErrorCode
|
||||
from onyx.error_handling.exceptions import OnyxError
|
||||
from onyx.utils.logger import setup_logger
|
||||
from shared_configs.enums import EmbedTextType
|
||||
from shared_configs.model_server_models import Embedding
|
||||
@@ -188,7 +189,7 @@ async def process_embed_request(
|
||||
)
|
||||
|
||||
if not embed_request.texts:
|
||||
raise HTTPException(status_code=400, detail="No texts to be embedded")
|
||||
raise OnyxError(OnyxErrorCode.VALIDATION_ERROR, "No texts to be embedded")
|
||||
|
||||
if not all(embed_request.texts):
|
||||
raise ValueError("Empty strings are not allowed for embedding.")
|
||||
@@ -211,14 +212,12 @@ async def process_embed_request(
|
||||
)
|
||||
return EmbedResponse(embeddings=embeddings)
|
||||
except RateLimitError as e:
|
||||
raise HTTPException(
|
||||
status_code=429,
|
||||
detail=str(e),
|
||||
)
|
||||
raise OnyxError(OnyxErrorCode.RATE_LIMITED, str(e))
|
||||
except Exception as e:
|
||||
logger.exception(
|
||||
f"Error during embedding process: provider={embed_request.provider_type} model={embed_request.model_name}"
|
||||
)
|
||||
raise HTTPException(
|
||||
status_code=500, detail=f"Error during embedding process: {e}"
|
||||
raise OnyxError(
|
||||
OnyxErrorCode.INTERNAL_ERROR,
|
||||
f"Error during embedding process: {e}",
|
||||
)
|
||||
|
||||
@@ -18,6 +18,7 @@ from model_server.encoders import router as encoders_router
|
||||
from model_server.management_endpoints import router as management_router
|
||||
from model_server.utils import get_gpu_type
|
||||
from onyx import __version__
|
||||
from onyx.error_handling.exceptions import register_onyx_exception_handlers
|
||||
from onyx.utils.logger import setup_logger
|
||||
from onyx.utils.logger import setup_uvicorn_logger
|
||||
from onyx.utils.middleware import add_onyx_request_id_middleware
|
||||
@@ -108,6 +109,8 @@ def get_model_app() -> FastAPI:
|
||||
application.include_router(management_router)
|
||||
application.include_router(encoders_router)
|
||||
|
||||
register_onyx_exception_handlers(application)
|
||||
|
||||
request_id_prefix = "INF"
|
||||
if INDEXING_ONLY:
|
||||
request_id_prefix = "IDX"
|
||||
|
||||
@@ -15,6 +15,7 @@ from onyx.chat.citation_processor import DynamicCitationProcessor
|
||||
from onyx.chat.emitter import Emitter
|
||||
from onyx.chat.models import ChatMessageSimple
|
||||
from onyx.chat.models import LlmStepResult
|
||||
from onyx.chat.tool_call_args_streaming import maybe_emit_argument_delta
|
||||
from onyx.configs.app_configs import LOG_ONYX_MODEL_INTERACTIONS
|
||||
from onyx.configs.app_configs import PROMPT_CACHE_CHAT_HISTORY
|
||||
from onyx.configs.constants import MessageType
|
||||
@@ -54,6 +55,7 @@ from onyx.server.query_and_chat.streaming_models import ReasoningStart
|
||||
from onyx.tools.models import ToolCallKickoff
|
||||
from onyx.tracing.framework.create import generation_span
|
||||
from onyx.utils.b64 import get_image_type_from_bytes
|
||||
from onyx.utils.jsonriver import Parser
|
||||
from onyx.utils.logger import setup_logger
|
||||
from onyx.utils.postgres_sanitization import sanitize_string
|
||||
from onyx.utils.text_processing import find_all_json_objects
|
||||
@@ -1009,6 +1011,7 @@ def run_llm_step_pkt_generator(
|
||||
)
|
||||
|
||||
id_to_tool_call_map: dict[int, dict[str, Any]] = {}
|
||||
arg_parsers: dict[int, Parser] = {}
|
||||
reasoning_start = False
|
||||
answer_start = False
|
||||
accumulated_reasoning = ""
|
||||
@@ -1215,7 +1218,14 @@ def run_llm_step_pkt_generator(
|
||||
yield from _close_reasoning_if_active()
|
||||
|
||||
for tool_call_delta in delta.tool_calls:
|
||||
# maybe_emit depends and update being called first and attaching the delta
|
||||
_update_tool_call_with_delta(id_to_tool_call_map, tool_call_delta)
|
||||
yield from maybe_emit_argument_delta(
|
||||
tool_calls_in_progress=id_to_tool_call_map,
|
||||
tool_call_delta=tool_call_delta,
|
||||
placement=_current_placement(),
|
||||
parsers=arg_parsers,
|
||||
)
|
||||
|
||||
# Flush any tail text buffered while checking for split "<function_calls" markers.
|
||||
filtered_content_tail = xml_tool_call_content_filter.flush()
|
||||
|
||||
77
backend/onyx/chat/tool_call_args_streaming.py
Normal file
77
backend/onyx/chat/tool_call_args_streaming.py
Normal file
@@ -0,0 +1,77 @@
|
||||
from collections.abc import Generator
|
||||
from collections.abc import Mapping
|
||||
from typing import Any
|
||||
from typing import Type
|
||||
|
||||
from onyx.llm.model_response import ChatCompletionDeltaToolCall
|
||||
from onyx.server.query_and_chat.placement import Placement
|
||||
from onyx.server.query_and_chat.streaming_models import Packet
|
||||
from onyx.server.query_and_chat.streaming_models import ToolCallArgumentDelta
|
||||
from onyx.tools.built_in_tools import TOOL_NAME_TO_CLASS
|
||||
from onyx.tools.interface import Tool
|
||||
from onyx.utils.jsonriver import Parser
|
||||
|
||||
|
||||
def _get_tool_class(
|
||||
tool_calls_in_progress: Mapping[int, Mapping[str, Any]],
|
||||
tool_call_delta: ChatCompletionDeltaToolCall,
|
||||
) -> Type[Tool] | None:
|
||||
"""Look up the Tool subclass for a streaming tool call delta."""
|
||||
tool_name = tool_calls_in_progress.get(tool_call_delta.index, {}).get("name")
|
||||
if not tool_name:
|
||||
return None
|
||||
return TOOL_NAME_TO_CLASS.get(tool_name)
|
||||
|
||||
|
||||
def maybe_emit_argument_delta(
|
||||
tool_calls_in_progress: Mapping[int, Mapping[str, Any]],
|
||||
tool_call_delta: ChatCompletionDeltaToolCall,
|
||||
placement: Placement,
|
||||
parsers: dict[int, Parser],
|
||||
) -> Generator[Packet, None, None]:
|
||||
"""Emit decoded tool-call argument deltas to the frontend.
|
||||
|
||||
Uses a ``jsonriver.Parser`` per tool-call index to incrementally parse
|
||||
the JSON argument string and extract only the newly-appended content
|
||||
for each string-valued argument.
|
||||
|
||||
NOTE: Non-string arguments (numbers, booleans, null, arrays, objects)
|
||||
are skipped — they are available in the final tool-call kickoff packet.
|
||||
|
||||
``parsers`` is a mutable dict keyed by tool-call index. A new
|
||||
``Parser`` is created automatically for each new index.
|
||||
"""
|
||||
tool_cls = _get_tool_class(tool_calls_in_progress, tool_call_delta)
|
||||
if not tool_cls or not tool_cls.should_emit_argument_deltas():
|
||||
return
|
||||
|
||||
fn = tool_call_delta.function
|
||||
delta_fragment = fn.arguments if fn else None
|
||||
if not delta_fragment:
|
||||
return
|
||||
|
||||
idx = tool_call_delta.index
|
||||
if idx not in parsers:
|
||||
parsers[idx] = Parser()
|
||||
parser = parsers[idx]
|
||||
|
||||
deltas = parser.feed(delta_fragment)
|
||||
|
||||
argument_deltas: dict[str, str] = {}
|
||||
for delta in deltas:
|
||||
if isinstance(delta, dict):
|
||||
for key, value in delta.items():
|
||||
if isinstance(value, str):
|
||||
argument_deltas[key] = argument_deltas.get(key, "") + value
|
||||
|
||||
if not argument_deltas:
|
||||
return
|
||||
|
||||
tc_data = tool_calls_in_progress[tool_call_delta.index]
|
||||
yield Packet(
|
||||
placement=placement,
|
||||
obj=ToolCallArgumentDelta(
|
||||
tool_type=tc_data.get("name", ""),
|
||||
argument_deltas=argument_deltas,
|
||||
),
|
||||
)
|
||||
@@ -68,6 +68,10 @@ FILE_TOKEN_COUNT_THRESHOLD = int(
|
||||
os.environ.get("FILE_TOKEN_COUNT_THRESHOLD", str(_DEFAULT_FILE_TOKEN_LIMIT))
|
||||
)
|
||||
|
||||
# Maximum upload size for a single user file (chat/projects) in MB.
|
||||
USER_FILE_MAX_UPLOAD_SIZE_MB = int(os.environ.get("USER_FILE_MAX_UPLOAD_SIZE_MB") or 50)
|
||||
USER_FILE_MAX_UPLOAD_SIZE_BYTES = USER_FILE_MAX_UPLOAD_SIZE_MB * 1024 * 1024
|
||||
|
||||
# If set to true, will show extra/uncommon connectors in the "Other" category
|
||||
SHOW_EXTRA_CONNECTORS = os.environ.get("SHOW_EXTRA_CONNECTORS", "").lower() == "true"
|
||||
|
||||
|
||||
@@ -1722,6 +1722,7 @@ class GoogleDriveConnector(
|
||||
primary_admin_email=self.primary_admin_email,
|
||||
google_domain=self.google_domain,
|
||||
),
|
||||
retriever_email=file.user_email,
|
||||
):
|
||||
slim_batch.append(doc)
|
||||
|
||||
|
||||
@@ -476,6 +476,7 @@ def _get_external_access_for_raw_gdrive_file(
|
||||
company_domain: str,
|
||||
retriever_drive_service: GoogleDriveService | None,
|
||||
admin_drive_service: GoogleDriveService,
|
||||
fallback_user_email: str,
|
||||
add_prefix: bool = False,
|
||||
) -> ExternalAccess:
|
||||
"""
|
||||
@@ -484,6 +485,8 @@ def _get_external_access_for_raw_gdrive_file(
|
||||
add_prefix: When True, prefix group IDs with source type (for indexing path).
|
||||
When False (default), leave unprefixed (for permission sync path
|
||||
where upsert_document_external_perms handles prefixing).
|
||||
fallback_user_email: When permission info can't be retrieved (e.g. externally-owned
|
||||
files), fall back to granting access to this user.
|
||||
"""
|
||||
external_access_fn = cast(
|
||||
Callable[
|
||||
@@ -492,6 +495,7 @@ def _get_external_access_for_raw_gdrive_file(
|
||||
str,
|
||||
GoogleDriveService | None,
|
||||
GoogleDriveService,
|
||||
str,
|
||||
bool,
|
||||
],
|
||||
ExternalAccess,
|
||||
@@ -507,6 +511,7 @@ def _get_external_access_for_raw_gdrive_file(
|
||||
company_domain,
|
||||
retriever_drive_service,
|
||||
admin_drive_service,
|
||||
fallback_user_email,
|
||||
add_prefix,
|
||||
)
|
||||
|
||||
@@ -672,6 +677,7 @@ def _convert_drive_item_to_document(
|
||||
creds, user_email=permission_sync_context.primary_admin_email
|
||||
),
|
||||
add_prefix=True, # Indexing path - prefix here
|
||||
fallback_user_email=retriever_email,
|
||||
)
|
||||
if permission_sync_context
|
||||
else None
|
||||
@@ -753,6 +759,7 @@ def build_slim_document(
|
||||
# if not specified, we will not sync permissions
|
||||
# will also be a no-op if EE is not enabled
|
||||
permission_sync_context: PermissionSyncContext | None,
|
||||
retriever_email: str,
|
||||
) -> SlimDocument | None:
|
||||
if file.get("mimeType") in [DRIVE_FOLDER_TYPE, DRIVE_SHORTCUT_TYPE]:
|
||||
return None
|
||||
@@ -774,6 +781,7 @@ def build_slim_document(
|
||||
creds,
|
||||
user_email=permission_sync_context.primary_admin_email,
|
||||
),
|
||||
fallback_user_email=retriever_email,
|
||||
)
|
||||
if permission_sync_context
|
||||
else None
|
||||
|
||||
@@ -36,9 +36,11 @@ from sqlalchemy import Text
|
||||
from sqlalchemy import text
|
||||
from sqlalchemy import UniqueConstraint
|
||||
from sqlalchemy.dialects import postgresql
|
||||
from sqlalchemy import event
|
||||
from sqlalchemy.engine.interfaces import Dialect
|
||||
from sqlalchemy.orm import DeclarativeBase
|
||||
from sqlalchemy.orm import Mapped
|
||||
from sqlalchemy.orm import Mapper
|
||||
from sqlalchemy.orm import mapped_column
|
||||
from sqlalchemy.orm import relationship
|
||||
from sqlalchemy.types import LargeBinary
|
||||
@@ -117,10 +119,50 @@ class Base(DeclarativeBase):
|
||||
__abstract__ = True
|
||||
|
||||
|
||||
class EncryptedString(TypeDecorator):
|
||||
class _EncryptedBase(TypeDecorator):
|
||||
"""Base for encrypted column types that wrap values in SensitiveValue."""
|
||||
|
||||
impl = LargeBinary
|
||||
# This type's behavior is fully deterministic and doesn't depend on any external factors.
|
||||
cache_ok = True
|
||||
_is_json: bool = False
|
||||
|
||||
def wrap_raw(self, value: Any) -> SensitiveValue:
|
||||
"""Encrypt a raw value and wrap it in SensitiveValue.
|
||||
|
||||
Called by the attribute set event so the Python-side type is always
|
||||
SensitiveValue, regardless of whether the value was loaded from the DB
|
||||
or assigned in application code.
|
||||
"""
|
||||
if self._is_json:
|
||||
if not isinstance(value, dict):
|
||||
raise TypeError(
|
||||
f"EncryptedJson column expected dict, got {type(value).__name__}"
|
||||
)
|
||||
raw_str = json.dumps(value)
|
||||
else:
|
||||
if not isinstance(value, str):
|
||||
raise TypeError(
|
||||
f"EncryptedString column expected str, got {type(value).__name__}"
|
||||
)
|
||||
raw_str = value
|
||||
return SensitiveValue(
|
||||
encrypted_bytes=encrypt_string_to_bytes(raw_str),
|
||||
decrypt_fn=decrypt_bytes_to_string,
|
||||
is_json=self._is_json,
|
||||
)
|
||||
|
||||
def compare_values(self, x: Any, y: Any) -> bool:
|
||||
if x is None or y is None:
|
||||
return x == y
|
||||
if isinstance(x, SensitiveValue):
|
||||
x = x.get_value(apply_mask=False)
|
||||
if isinstance(y, SensitiveValue):
|
||||
y = y.get_value(apply_mask=False)
|
||||
return x == y
|
||||
|
||||
|
||||
class EncryptedString(_EncryptedBase):
|
||||
_is_json: bool = False
|
||||
|
||||
def process_bind_param(
|
||||
self, value: str | SensitiveValue[str] | None, dialect: Dialect # noqa: ARG002
|
||||
@@ -144,20 +186,9 @@ class EncryptedString(TypeDecorator):
|
||||
)
|
||||
return None
|
||||
|
||||
def compare_values(self, x: Any, y: Any) -> bool:
|
||||
if x is None or y is None:
|
||||
return x == y
|
||||
if isinstance(x, SensitiveValue):
|
||||
x = x.get_value(apply_mask=False)
|
||||
if isinstance(y, SensitiveValue):
|
||||
y = y.get_value(apply_mask=False)
|
||||
return x == y
|
||||
|
||||
|
||||
class EncryptedJson(TypeDecorator):
|
||||
impl = LargeBinary
|
||||
# This type's behavior is fully deterministic and doesn't depend on any external factors.
|
||||
cache_ok = True
|
||||
class EncryptedJson(_EncryptedBase):
|
||||
_is_json: bool = True
|
||||
|
||||
def process_bind_param(
|
||||
self,
|
||||
@@ -165,9 +196,7 @@ class EncryptedJson(TypeDecorator):
|
||||
dialect: Dialect, # noqa: ARG002
|
||||
) -> bytes | None:
|
||||
if value is not None:
|
||||
# Handle both raw dicts and SensitiveValue wrappers
|
||||
if isinstance(value, SensitiveValue):
|
||||
# Get raw value for storage
|
||||
value = value.get_value(apply_mask=False)
|
||||
json_str = json.dumps(value)
|
||||
return encrypt_string_to_bytes(json_str)
|
||||
@@ -184,14 +213,40 @@ class EncryptedJson(TypeDecorator):
|
||||
)
|
||||
return None
|
||||
|
||||
def compare_values(self, x: Any, y: Any) -> bool:
|
||||
if x is None or y is None:
|
||||
return x == y
|
||||
if isinstance(x, SensitiveValue):
|
||||
x = x.get_value(apply_mask=False)
|
||||
if isinstance(y, SensitiveValue):
|
||||
y = y.get_value(apply_mask=False)
|
||||
return x == y
|
||||
|
||||
_REGISTERED_ATTRS: set[str] = set()
|
||||
|
||||
|
||||
@event.listens_for(Mapper, "mapper_configured")
|
||||
def _register_sensitive_value_set_events(
|
||||
mapper: Mapper,
|
||||
class_: type,
|
||||
) -> None:
|
||||
"""Auto-wrap raw values in SensitiveValue when assigned to encrypted columns."""
|
||||
for prop in mapper.column_attrs:
|
||||
for col in prop.columns:
|
||||
if isinstance(col.type, _EncryptedBase):
|
||||
col_type = col.type
|
||||
attr = getattr(class_, prop.key)
|
||||
|
||||
# Guard against double-registration (e.g. if mapper is
|
||||
# re-configured in test setups)
|
||||
attr_key = f"{class_.__qualname__}.{prop.key}"
|
||||
if attr_key in _REGISTERED_ATTRS:
|
||||
continue
|
||||
_REGISTERED_ATTRS.add(attr_key)
|
||||
|
||||
@event.listens_for(attr, "set", retval=True)
|
||||
def _wrap_value(
|
||||
target: Any, # noqa: ARG001
|
||||
value: Any,
|
||||
oldvalue: Any, # noqa: ARG001
|
||||
initiator: Any, # noqa: ARG001
|
||||
_col_type: _EncryptedBase = col_type,
|
||||
) -> Any:
|
||||
if value is not None and not isinstance(value, SensitiveValue):
|
||||
return _col_type.wrap_raw(value)
|
||||
return value
|
||||
|
||||
|
||||
class NullFilteredString(TypeDecorator):
|
||||
|
||||
161
backend/onyx/db/rotate_encryption_key.py
Normal file
161
backend/onyx/db/rotate_encryption_key.py
Normal file
@@ -0,0 +1,161 @@
|
||||
"""Rotate encryption key for all encrypted columns.
|
||||
|
||||
Dynamically discovers all columns using EncryptedString / EncryptedJson,
|
||||
decrypts each value with the old key, and re-encrypts with the current
|
||||
ENCRYPTION_KEY_SECRET.
|
||||
|
||||
The operation is idempotent: rows already encrypted with the current key
|
||||
are skipped. Commits are made in batches so a crash mid-rotation can be
|
||||
safely resumed by re-running.
|
||||
"""
|
||||
|
||||
import json
|
||||
from typing import Any
|
||||
|
||||
from sqlalchemy import LargeBinary
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy import update
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from onyx.configs.app_configs import ENCRYPTION_KEY_SECRET
|
||||
from onyx.db.models import Base
|
||||
from onyx.db.models import EncryptedJson
|
||||
from onyx.db.models import EncryptedString
|
||||
from onyx.utils.encryption import decrypt_bytes_to_string
|
||||
from onyx.utils.logger import setup_logger
|
||||
from onyx.utils.variable_functionality import global_version
|
||||
|
||||
logger = setup_logger()
|
||||
|
||||
_BATCH_SIZE = 500
|
||||
|
||||
|
||||
def _can_decrypt_with_current_key(data: bytes) -> bool:
|
||||
"""Check if data is already encrypted with the current key.
|
||||
|
||||
Passes the key explicitly so the fallback-to-raw-decode path in
|
||||
_decrypt_bytes is NOT triggered — a clean success/failure signal.
|
||||
"""
|
||||
try:
|
||||
decrypt_bytes_to_string(data, key=ENCRYPTION_KEY_SECRET)
|
||||
return True
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
|
||||
def _discover_encrypted_columns() -> list[tuple[type, str, list[str], bool]]:
|
||||
"""Walk all ORM models and find columns using EncryptedString/EncryptedJson.
|
||||
|
||||
Returns list of (ModelClass, column_attr_name, [pk_attr_names], is_json).
|
||||
"""
|
||||
results: list[tuple[type, str, list[str], bool]] = []
|
||||
|
||||
for mapper in Base.registry.mappers:
|
||||
model_cls = mapper.class_
|
||||
pk_names = [col.key for col in mapper.primary_key]
|
||||
|
||||
for prop in mapper.column_attrs:
|
||||
for col in prop.columns:
|
||||
if isinstance(col.type, EncryptedJson):
|
||||
results.append((model_cls, prop.key, pk_names, True))
|
||||
elif isinstance(col.type, EncryptedString):
|
||||
results.append((model_cls, prop.key, pk_names, False))
|
||||
|
||||
return results
|
||||
|
||||
|
||||
def rotate_encryption_key(
|
||||
db_session: Session,
|
||||
old_key: str | None,
|
||||
dry_run: bool = False,
|
||||
) -> dict[str, int]:
|
||||
"""Decrypt all encrypted columns with old_key and re-encrypt with the current key.
|
||||
|
||||
Args:
|
||||
db_session: Active database session.
|
||||
old_key: The previous encryption key. Pass None or "" if values were
|
||||
not previously encrypted with a key.
|
||||
dry_run: If True, count rows that need rotation without modifying data.
|
||||
|
||||
Returns:
|
||||
Dict of "table.column" -> number of rows re-encrypted (or would be).
|
||||
|
||||
Commits every _BATCH_SIZE rows so that locks are held briefly and progress
|
||||
is preserved on crash. Already-rotated rows are detected and skipped,
|
||||
making the operation safe to re-run.
|
||||
"""
|
||||
if not global_version.is_ee_version():
|
||||
raise RuntimeError("EE mode is not enabled — rotation requires EE encryption.")
|
||||
|
||||
if not ENCRYPTION_KEY_SECRET:
|
||||
raise RuntimeError(
|
||||
"ENCRYPTION_KEY_SECRET is not set — cannot rotate. "
|
||||
"Set the target encryption key in the environment before running."
|
||||
)
|
||||
|
||||
encrypted_columns = _discover_encrypted_columns()
|
||||
totals: dict[str, int] = {}
|
||||
|
||||
for model_cls, col_name, pk_names, is_json in encrypted_columns:
|
||||
table_name: str = model_cls.__tablename__ # type: ignore[attr-defined]
|
||||
col_attr = getattr(model_cls, col_name)
|
||||
pk_attrs = [getattr(model_cls, pk) for pk in pk_names]
|
||||
|
||||
# Read raw bytes directly, bypassing the TypeDecorator
|
||||
raw_col = col_attr.property.columns[0]
|
||||
|
||||
stmt = select(*pk_attrs, raw_col.cast(LargeBinary)).where(col_attr.is_not(None))
|
||||
rows = db_session.execute(stmt).all()
|
||||
|
||||
reencrypted = 0
|
||||
batch_pending = 0
|
||||
for row in rows:
|
||||
raw_bytes: bytes | None = row[-1]
|
||||
if raw_bytes is None:
|
||||
continue
|
||||
|
||||
if _can_decrypt_with_current_key(raw_bytes):
|
||||
continue
|
||||
|
||||
try:
|
||||
if not old_key:
|
||||
decrypted_str = raw_bytes.decode("utf-8")
|
||||
else:
|
||||
decrypted_str = decrypt_bytes_to_string(raw_bytes, key=old_key)
|
||||
|
||||
# For EncryptedJson, parse back to dict so the TypeDecorator
|
||||
# can json.dumps() it cleanly (avoids double-encoding).
|
||||
value: Any = json.loads(decrypted_str) if is_json else decrypted_str
|
||||
except (ValueError, UnicodeDecodeError) as e:
|
||||
pk_vals = [row[i] for i in range(len(pk_names))]
|
||||
logger.warning(
|
||||
f"Could not decrypt/parse {table_name}.{col_name} "
|
||||
f"row {pk_vals} — skipping: {e}"
|
||||
)
|
||||
continue
|
||||
|
||||
if not dry_run:
|
||||
pk_filters = [pk_attr == row[i] for i, pk_attr in enumerate(pk_attrs)]
|
||||
update_stmt = (
|
||||
update(model_cls).where(*pk_filters).values({col_name: value})
|
||||
)
|
||||
db_session.execute(update_stmt)
|
||||
batch_pending += 1
|
||||
|
||||
if batch_pending >= _BATCH_SIZE:
|
||||
db_session.commit()
|
||||
batch_pending = 0
|
||||
reencrypted += 1
|
||||
|
||||
# Flush remaining rows in this column
|
||||
if batch_pending > 0:
|
||||
db_session.commit()
|
||||
|
||||
if reencrypted > 0:
|
||||
totals[f"{table_name}.{col_name}"] = reencrypted
|
||||
logger.info(
|
||||
f"{'[DRY RUN] Would re-encrypt' if dry_run else 'Re-encrypted'} "
|
||||
f"{reencrypted} value(s) in {table_name}.{col_name}"
|
||||
)
|
||||
|
||||
return totals
|
||||
@@ -10,6 +10,7 @@ from onyx.mcp_server.utils import get_indexed_sources
|
||||
from onyx.mcp_server.utils import require_access_token
|
||||
from onyx.utils.logger import setup_logger
|
||||
from onyx.utils.variable_functionality import build_api_server_url_for_http_requests
|
||||
from onyx.utils.variable_functionality import global_version
|
||||
|
||||
logger = setup_logger()
|
||||
|
||||
@@ -26,6 +27,14 @@ async def search_indexed_documents(
|
||||
Use this tool for information that is not public knowledge and specific to the user,
|
||||
their team, their work, or their organization/company.
|
||||
|
||||
Note: In CE mode, this tool uses the chat endpoint internally which invokes an LLM
|
||||
on every call, consuming tokens and adding latency.
|
||||
Additionally, CE callers receive a truncated snippet (blurb) instead of a full document chunk,
|
||||
but this should still be sufficient for most use cases. CE mode functionality should be swapped
|
||||
when a dedicated CE search endpoint is implemented.
|
||||
|
||||
In EE mode, the dedicated search endpoint is used instead.
|
||||
|
||||
To find a list of available sources, use the `indexed_sources` resource.
|
||||
Returns chunks of text as search results with snippets, scores, and metadata.
|
||||
|
||||
@@ -111,48 +120,73 @@ async def search_indexed_documents(
|
||||
if time_cutoff_dt:
|
||||
filters["time_cutoff"] = time_cutoff_dt.isoformat()
|
||||
|
||||
# Build the search request using the new SendSearchQueryRequest format
|
||||
search_request = {
|
||||
"search_query": query,
|
||||
"filters": filters,
|
||||
"num_docs_fed_to_llm_selection": limit,
|
||||
"run_query_expansion": False,
|
||||
"include_content": True,
|
||||
"stream": False,
|
||||
}
|
||||
is_ee = global_version.is_ee_version()
|
||||
base_url = build_api_server_url_for_http_requests(respect_env_override_if_set=True)
|
||||
auth_headers = {"Authorization": f"Bearer {access_token.token}"}
|
||||
|
||||
search_request: dict[str, Any]
|
||||
if is_ee:
|
||||
# EE: use the dedicated search endpoint (no LLM invocation)
|
||||
search_request = {
|
||||
"search_query": query,
|
||||
"filters": filters,
|
||||
"num_docs_fed_to_llm_selection": limit,
|
||||
"run_query_expansion": False,
|
||||
"include_content": True,
|
||||
"stream": False,
|
||||
}
|
||||
endpoint = f"{base_url}/search/send-search-message"
|
||||
error_key = "error"
|
||||
docs_key = "search_docs"
|
||||
content_field = "content"
|
||||
else:
|
||||
# CE: fall back to the chat endpoint (invokes LLM, consumes tokens)
|
||||
search_request = {
|
||||
"message": query,
|
||||
"stream": False,
|
||||
"chat_session_info": {},
|
||||
}
|
||||
if filters:
|
||||
search_request["internal_search_filters"] = filters
|
||||
endpoint = f"{base_url}/chat/send-chat-message"
|
||||
error_key = "error_msg"
|
||||
docs_key = "top_documents"
|
||||
content_field = "blurb"
|
||||
|
||||
# Call the API server using the new send-search-message route
|
||||
try:
|
||||
response = await get_http_client().post(
|
||||
f"{build_api_server_url_for_http_requests(respect_env_override_if_set=True)}/search/send-search-message",
|
||||
endpoint,
|
||||
json=search_request,
|
||||
headers={"Authorization": f"Bearer {access_token.token}"},
|
||||
headers=auth_headers,
|
||||
)
|
||||
response.raise_for_status()
|
||||
result = response.json()
|
||||
|
||||
# Check for error in response
|
||||
if result.get("error"):
|
||||
if result.get(error_key):
|
||||
return {
|
||||
"documents": [],
|
||||
"total_results": 0,
|
||||
"query": query,
|
||||
"error": result.get("error"),
|
||||
"error": result.get(error_key),
|
||||
}
|
||||
|
||||
# Return simplified format for MCP clients
|
||||
fields_to_return = [
|
||||
"semantic_identifier",
|
||||
"content",
|
||||
"source_type",
|
||||
"link",
|
||||
"score",
|
||||
]
|
||||
documents = [
|
||||
{key: doc.get(key) for key in fields_to_return}
|
||||
for doc in result.get("search_docs", [])
|
||||
{
|
||||
"semantic_identifier": doc.get("semantic_identifier"),
|
||||
"content": doc.get(content_field),
|
||||
"source_type": doc.get("source_type"),
|
||||
"link": doc.get("link"),
|
||||
"score": doc.get("score"),
|
||||
}
|
||||
for doc in result.get(docs_key, [])
|
||||
]
|
||||
|
||||
# NOTE: search depth is controlled by the backend persona defaults, not `limit`.
|
||||
# `limit` only caps the returned list; fewer results may be returned if the
|
||||
# backend retrieves fewer documents than requested.
|
||||
documents = documents[:limit]
|
||||
|
||||
logger.info(
|
||||
f"Onyx MCP Server: Internal search returned {len(documents)} results"
|
||||
)
|
||||
@@ -160,7 +194,6 @@ async def search_indexed_documents(
|
||||
"documents": documents,
|
||||
"total_results": len(documents),
|
||||
"query": query,
|
||||
"executed_queries": result.get("all_executed_queries", [query]),
|
||||
}
|
||||
except Exception as e:
|
||||
logger.error(f"Onyx MCP Server: Document search error: {e}", exc_info=True)
|
||||
|
||||
@@ -10,6 +10,8 @@ from pydantic import Field
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from onyx.configs.app_configs import FILE_TOKEN_COUNT_THRESHOLD
|
||||
from onyx.configs.app_configs import USER_FILE_MAX_UPLOAD_SIZE_BYTES
|
||||
from onyx.configs.app_configs import USER_FILE_MAX_UPLOAD_SIZE_MB
|
||||
from onyx.db.llm import fetch_default_llm_model
|
||||
from onyx.file_processing.extract_file_text import extract_file_text
|
||||
from onyx.file_processing.extract_file_text import get_file_ext
|
||||
@@ -35,6 +37,38 @@ def get_safe_filename(upload: UploadFile) -> str:
|
||||
return upload.filename
|
||||
|
||||
|
||||
def get_upload_size_bytes(upload: UploadFile) -> int | None:
|
||||
"""Best-effort file size in bytes without consuming the stream."""
|
||||
if upload.size is not None:
|
||||
return upload.size
|
||||
|
||||
try:
|
||||
current_pos = upload.file.tell()
|
||||
upload.file.seek(0, 2)
|
||||
size = upload.file.tell()
|
||||
upload.file.seek(current_pos)
|
||||
return size
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
"Could not determine upload size via stream seek "
|
||||
f"(filename='{get_safe_filename(upload)}', "
|
||||
f"error_type={type(e).__name__}, error={e})"
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
def is_upload_too_large(upload: UploadFile, max_bytes: int) -> bool:
|
||||
"""Return True when upload size is known and exceeds max_bytes."""
|
||||
size_bytes = get_upload_size_bytes(upload)
|
||||
if size_bytes is None:
|
||||
logger.warning(
|
||||
"Could not determine upload size; skipping size-limit check for "
|
||||
f"'{get_safe_filename(upload)}'"
|
||||
)
|
||||
return False
|
||||
return size_bytes > max_bytes
|
||||
|
||||
|
||||
# Guard against extremely large images
|
||||
Image.MAX_IMAGE_PIXELS = 12000 * 12000
|
||||
|
||||
@@ -159,6 +193,18 @@ def categorize_uploaded_files(
|
||||
for upload in files:
|
||||
try:
|
||||
filename = get_safe_filename(upload)
|
||||
|
||||
# Size limit is a hard safety cap and is enforced even when token
|
||||
# threshold checks are skipped via SKIP_USERFILE_THRESHOLD settings.
|
||||
if is_upload_too_large(upload, USER_FILE_MAX_UPLOAD_SIZE_BYTES):
|
||||
results.rejected.append(
|
||||
RejectedFile(
|
||||
filename=filename,
|
||||
reason=f"Exceeds {USER_FILE_MAX_UPLOAD_SIZE_MB} MB file size limit",
|
||||
)
|
||||
)
|
||||
continue
|
||||
|
||||
extension = get_file_ext(filename)
|
||||
|
||||
# If image, estimate tokens via dedicated method first
|
||||
|
||||
@@ -41,6 +41,7 @@ class StreamingType(Enum):
|
||||
REASONING_DONE = "reasoning_done"
|
||||
CITATION_INFO = "citation_info"
|
||||
TOOL_CALL_DEBUG = "tool_call_debug"
|
||||
TOOL_CALL_ARGUMENT_DELTA = "tool_call_argument_delta"
|
||||
|
||||
MEMORY_TOOL_START = "memory_tool_start"
|
||||
MEMORY_TOOL_DELTA = "memory_tool_delta"
|
||||
@@ -259,6 +260,15 @@ class CustomToolDelta(BaseObj):
|
||||
file_ids: list[str] | None = None
|
||||
|
||||
|
||||
class ToolCallArgumentDelta(BaseObj):
|
||||
type: Literal["tool_call_argument_delta"] = (
|
||||
StreamingType.TOOL_CALL_ARGUMENT_DELTA.value
|
||||
)
|
||||
|
||||
tool_type: str
|
||||
argument_deltas: dict[str, Any]
|
||||
|
||||
|
||||
################################################
|
||||
# File Reader Packets
|
||||
################################################
|
||||
@@ -379,6 +389,7 @@ PacketObj = Union[
|
||||
# Citation Packets
|
||||
CitationInfo,
|
||||
ToolCallDebug,
|
||||
ToolCallArgumentDelta,
|
||||
# Deep Research Packets
|
||||
DeepResearchPlanStart,
|
||||
DeepResearchPlanDelta,
|
||||
|
||||
@@ -78,6 +78,7 @@ class Settings(BaseModel):
|
||||
|
||||
# User Knowledge settings
|
||||
user_knowledge_enabled: bool | None = True
|
||||
user_file_max_upload_size_mb: int | None = None
|
||||
|
||||
# Connector settings
|
||||
show_extra_connectors: bool | None = True
|
||||
|
||||
@@ -3,6 +3,7 @@ from onyx.configs.app_configs import DISABLE_USER_KNOWLEDGE
|
||||
from onyx.configs.app_configs import ENABLE_OPENSEARCH_INDEXING_FOR_ONYX
|
||||
from onyx.configs.app_configs import ONYX_QUERY_HISTORY_TYPE
|
||||
from onyx.configs.app_configs import SHOW_EXTRA_CONNECTORS
|
||||
from onyx.configs.app_configs import USER_FILE_MAX_UPLOAD_SIZE_MB
|
||||
from onyx.configs.constants import KV_SETTINGS_KEY
|
||||
from onyx.configs.constants import OnyxRedisLocks
|
||||
from onyx.key_value_store.factory import get_kv_store
|
||||
@@ -50,6 +51,7 @@ def load_settings() -> Settings:
|
||||
if DISABLE_USER_KNOWLEDGE:
|
||||
settings.user_knowledge_enabled = False
|
||||
|
||||
settings.user_file_max_upload_size_mb = USER_FILE_MAX_UPLOAD_SIZE_MB
|
||||
settings.show_extra_connectors = SHOW_EXTRA_CONNECTORS
|
||||
settings.opensearch_indexing_enabled = ENABLE_OPENSEARCH_INDEXING_FOR_ONYX
|
||||
return settings
|
||||
|
||||
@@ -56,3 +56,23 @@ def get_built_in_tool_ids() -> list[str]:
|
||||
|
||||
def get_built_in_tool_by_id(in_code_tool_id: str) -> Type[BUILT_IN_TOOL_TYPES]:
|
||||
return BUILT_IN_TOOL_MAP[in_code_tool_id]
|
||||
|
||||
|
||||
def _build_tool_name_to_class() -> dict[str, Type[BUILT_IN_TOOL_TYPES]]:
|
||||
"""Build a mapping from LLM-facing tool name to tool class."""
|
||||
result: dict[str, Type[BUILT_IN_TOOL_TYPES]] = {}
|
||||
for cls in BUILT_IN_TOOL_MAP.values():
|
||||
name_attr = cls.__dict__.get("name")
|
||||
if isinstance(name_attr, property) and name_attr.fget is not None:
|
||||
tool_name = name_attr.fget(cls)
|
||||
elif isinstance(name_attr, str):
|
||||
tool_name = name_attr
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Built-in tool {cls.__name__} must define a valid LLM-facing tool name"
|
||||
)
|
||||
result[tool_name] = cls
|
||||
return result
|
||||
|
||||
|
||||
TOOL_NAME_TO_CLASS: dict[str, Type[BUILT_IN_TOOL_TYPES]] = _build_tool_name_to_class()
|
||||
|
||||
@@ -92,3 +92,7 @@ class Tool(abc.ABC, Generic[TOverride]):
|
||||
**llm_kwargs: Any,
|
||||
) -> ToolResponse:
|
||||
raise NotImplementedError
|
||||
|
||||
@classmethod
|
||||
def should_emit_argument_deltas(cls) -> bool:
|
||||
return False
|
||||
|
||||
@@ -376,3 +376,8 @@ class PythonTool(Tool[PythonToolOverrideKwargs]):
|
||||
rich_response=None,
|
||||
llm_facing_response=llm_response,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
@override
|
||||
def should_emit_argument_deltas(cls) -> bool:
|
||||
return True
|
||||
|
||||
@@ -11,16 +11,20 @@ logger = setup_logger()
|
||||
|
||||
|
||||
# IMPORTANT DO NOT DELETE, THIS IS USED BY fetch_versioned_implementation
|
||||
def _encrypt_string(input_str: str) -> bytes:
|
||||
def _encrypt_string(input_str: str, key: str | None = None) -> bytes: # noqa: ARG001
|
||||
if ENCRYPTION_KEY_SECRET:
|
||||
logger.warning("MIT version of Onyx does not support encryption of secrets.")
|
||||
elif key is not None:
|
||||
logger.debug("MIT encrypt called with explicit key — key ignored.")
|
||||
return input_str.encode()
|
||||
|
||||
|
||||
# IMPORTANT DO NOT DELETE, THIS IS USED BY fetch_versioned_implementation
|
||||
def _decrypt_bytes(input_bytes: bytes) -> str:
|
||||
# No need to double warn. If you wish to learn more about encryption features
|
||||
# refer to the Onyx EE code
|
||||
def _decrypt_bytes(input_bytes: bytes, key: str | None = None) -> str: # noqa: ARG001
|
||||
if ENCRYPTION_KEY_SECRET:
|
||||
logger.warning("MIT version of Onyx does not support decryption of secrets.")
|
||||
elif key is not None:
|
||||
logger.debug("MIT decrypt called with explicit key — key ignored.")
|
||||
return input_bytes.decode()
|
||||
|
||||
|
||||
@@ -86,15 +90,15 @@ def _mask_list(items: list[Any]) -> list[Any]:
|
||||
return masked
|
||||
|
||||
|
||||
def encrypt_string_to_bytes(intput_str: str) -> bytes:
|
||||
def encrypt_string_to_bytes(intput_str: str, key: str | None = None) -> bytes:
|
||||
versioned_encryption_fn = fetch_versioned_implementation(
|
||||
"onyx.utils.encryption", "_encrypt_string"
|
||||
)
|
||||
return versioned_encryption_fn(intput_str)
|
||||
return versioned_encryption_fn(intput_str, key=key)
|
||||
|
||||
|
||||
def decrypt_bytes_to_string(intput_bytes: bytes) -> str:
|
||||
def decrypt_bytes_to_string(intput_bytes: bytes, key: str | None = None) -> str:
|
||||
versioned_decryption_fn = fetch_versioned_implementation(
|
||||
"onyx.utils.encryption", "_decrypt_bytes"
|
||||
)
|
||||
return versioned_decryption_fn(intput_bytes)
|
||||
return versioned_decryption_fn(intput_bytes, key=key)
|
||||
|
||||
@@ -128,6 +128,8 @@ class SensitiveValue(Generic[T]):
|
||||
value = self._decrypt()
|
||||
|
||||
if not apply_mask:
|
||||
# Callers must not mutate the returned dict — doing so would
|
||||
# desync the cache from the encrypted bytes and the DB.
|
||||
return value
|
||||
|
||||
# Apply masking
|
||||
@@ -174,18 +176,20 @@ class SensitiveValue(Generic[T]):
|
||||
)
|
||||
|
||||
def __eq__(self, other: Any) -> bool:
|
||||
"""Prevent direct comparison which might expose value."""
|
||||
if isinstance(other, SensitiveValue):
|
||||
# Compare encrypted bytes for equality check
|
||||
return self._encrypted_bytes == other._encrypted_bytes
|
||||
raise SensitiveAccessError(
|
||||
"Cannot compare SensitiveValue with non-SensitiveValue. "
|
||||
"Use .get_value(apply_mask=True/False) to access the value for comparison."
|
||||
)
|
||||
"""Compare SensitiveValues by their decrypted content."""
|
||||
# NOTE: if you attempt to compare a string/dict to a SensitiveValue,
|
||||
# this comparison will return NotImplemented, which then evaluates to False.
|
||||
# This is the convention and required for SQLAlchemy's attribute tracking.
|
||||
if not isinstance(other, SensitiveValue):
|
||||
return NotImplemented
|
||||
return self._decrypt() == other._decrypt()
|
||||
|
||||
def __hash__(self) -> int:
|
||||
"""Allow hashing based on encrypted bytes."""
|
||||
return hash(self._encrypted_bytes)
|
||||
"""Hash based on decrypted content."""
|
||||
value = self._decrypt()
|
||||
if isinstance(value, dict):
|
||||
return hash(json.dumps(value, sort_keys=True))
|
||||
return hash(value)
|
||||
|
||||
# Prevent JSON serialization
|
||||
def __json__(self) -> Any:
|
||||
|
||||
@@ -1,48 +1,93 @@
|
||||
"""Decrypt a raw hex-encoded credential value.
|
||||
|
||||
Usage:
|
||||
python -m scripts.decrypt <hex_value>
|
||||
python -m scripts.decrypt <hex_value> --key "my-encryption-key"
|
||||
python -m scripts.decrypt <hex_value> --key ""
|
||||
|
||||
Pass --key "" to skip decryption and just decode the raw bytes as UTF-8.
|
||||
Omit --key to use the current ENCRYPTION_KEY_SECRET from the environment.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import binascii
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
|
||||
from onyx.utils.encryption import decrypt_bytes_to_string
|
||||
parent_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
sys.path.append(parent_dir)
|
||||
|
||||
from onyx.utils.encryption import decrypt_bytes_to_string # noqa: E402
|
||||
from onyx.utils.variable_functionality import global_version # noqa: E402
|
||||
|
||||
|
||||
def decrypt_raw_credential(encrypted_value: str) -> None:
|
||||
"""Decrypt and display a raw encrypted credential value
|
||||
def decrypt_raw_credential(encrypted_value: str, key: str | None = None) -> None:
|
||||
"""Decrypt and display a raw encrypted credential value.
|
||||
|
||||
Args:
|
||||
encrypted_value: The hex encoded encrypted credential value
|
||||
encrypted_value: The hex-encoded encrypted credential value.
|
||||
key: Encryption key to use. None means use ENCRYPTION_KEY_SECRET,
|
||||
empty string means just decode as UTF-8.
|
||||
"""
|
||||
# Strip common hex prefixes
|
||||
if encrypted_value.startswith("\\x"):
|
||||
encrypted_value = encrypted_value[2:]
|
||||
elif encrypted_value.startswith("x"):
|
||||
encrypted_value = encrypted_value[1:]
|
||||
print(encrypted_value)
|
||||
|
||||
try:
|
||||
# If string starts with 'x', remove it as it's just a prefix indicating hex
|
||||
if encrypted_value.startswith("x"):
|
||||
encrypted_value = encrypted_value[1:]
|
||||
elif encrypted_value.startswith("\\x"):
|
||||
encrypted_value = encrypted_value[2:]
|
||||
|
||||
# Convert hex string to bytes
|
||||
encrypted_bytes = binascii.unhexlify(encrypted_value)
|
||||
|
||||
# Decrypt the bytes
|
||||
decrypted_str = decrypt_bytes_to_string(encrypted_bytes)
|
||||
|
||||
# Parse and pretty print the decrypted JSON
|
||||
decrypted_json = json.loads(decrypted_str)
|
||||
print("Decrypted credential value:")
|
||||
print(json.dumps(decrypted_json, indent=2))
|
||||
|
||||
raw_bytes = binascii.unhexlify(encrypted_value)
|
||||
except binascii.Error:
|
||||
print("Error: Invalid hex encoded string")
|
||||
print("Error: Invalid hex-encoded string")
|
||||
sys.exit(1)
|
||||
|
||||
except json.JSONDecodeError as e:
|
||||
print(f"Decrypted raw value (not JSON): {e}")
|
||||
if key == "":
|
||||
# Empty key → just decode as UTF-8, no decryption
|
||||
try:
|
||||
decrypted_str = raw_bytes.decode("utf-8")
|
||||
except UnicodeDecodeError as e:
|
||||
print(f"Error decoding bytes as UTF-8: {e}")
|
||||
sys.exit(1)
|
||||
else:
|
||||
print(key)
|
||||
try:
|
||||
decrypted_str = decrypt_bytes_to_string(raw_bytes, key=key)
|
||||
except Exception as e:
|
||||
print(f"Error decrypting value: {e}")
|
||||
sys.exit(1)
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error decrypting value: {e}")
|
||||
# Try to pretty-print as JSON, otherwise print raw
|
||||
try:
|
||||
parsed = json.loads(decrypted_str)
|
||||
print(json.dumps(parsed, indent=2))
|
||||
except json.JSONDecodeError:
|
||||
print(decrypted_str)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Decrypt a hex-encoded credential value."
|
||||
)
|
||||
parser.add_argument(
|
||||
"value",
|
||||
help="Hex-encoded encrypted value to decrypt.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--key",
|
||||
default=None,
|
||||
help=(
|
||||
"Encryption key. Omit to use ENCRYPTION_KEY_SECRET from env. "
|
||||
'Pass "" (empty) to just decode as UTF-8 without decryption.'
|
||||
),
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
global_version.set_ee()
|
||||
decrypt_raw_credential(args.value, key=args.key)
|
||||
global_version.unset_ee()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
if len(sys.argv) != 2:
|
||||
print("Usage: python decrypt.py <hex_encoded_encrypted_value>")
|
||||
sys.exit(1)
|
||||
|
||||
encrypted_value = sys.argv[1]
|
||||
decrypt_raw_credential(encrypted_value)
|
||||
main()
|
||||
|
||||
107
backend/scripts/reencrypt_secrets.py
Normal file
107
backend/scripts/reencrypt_secrets.py
Normal file
@@ -0,0 +1,107 @@
|
||||
"""Re-encrypt secrets under the current ENCRYPTION_KEY_SECRET.
|
||||
|
||||
Decrypts all encrypted columns using the old key (or raw decode if the old key
|
||||
is empty), then re-encrypts them with the current ENCRYPTION_KEY_SECRET.
|
||||
|
||||
Usage (docker):
|
||||
docker exec -it onyx-api_server-1 \
|
||||
python -m scripts.reencrypt_secrets --old-key "previous-key"
|
||||
|
||||
Usage (kubernetes):
|
||||
kubectl exec -it <pod> -- \
|
||||
python -m scripts.reencrypt_secrets --old-key "previous-key"
|
||||
|
||||
Omit --old-key (or pass "") if secrets were not previously encrypted.
|
||||
|
||||
For multi-tenant deployments, pass --tenant-id to target a specific tenant,
|
||||
or --all-tenants to iterate every tenant.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import sys
|
||||
|
||||
parent_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
sys.path.append(parent_dir)
|
||||
|
||||
from onyx.db.rotate_encryption_key import rotate_encryption_key # noqa: E402
|
||||
from onyx.db.engine.sql_engine import get_session_with_tenant # noqa: E402
|
||||
from onyx.db.engine.sql_engine import SqlEngine # noqa: E402
|
||||
from onyx.db.engine.tenant_utils import get_all_tenant_ids # noqa: E402
|
||||
from onyx.utils.variable_functionality import global_version # noqa: E402
|
||||
from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA # noqa: E402
|
||||
|
||||
|
||||
def _run_for_tenant(tenant_id: str, old_key: str | None, dry_run: bool = False) -> None:
|
||||
print(f"Re-encrypting secrets for tenant: {tenant_id}")
|
||||
with get_session_with_tenant(tenant_id=tenant_id) as db_session:
|
||||
results = rotate_encryption_key(db_session, old_key=old_key, dry_run=dry_run)
|
||||
|
||||
if results:
|
||||
for col, count in results.items():
|
||||
print(
|
||||
f" {col}: {count} row(s) {'would be ' if dry_run else ''}re-encrypted"
|
||||
)
|
||||
else:
|
||||
print("No rows needed re-encryption.")
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Re-encrypt secrets under the current encryption key."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--old-key",
|
||||
default=None,
|
||||
help="Previous encryption key. Omit or pass empty string if not applicable.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--dry-run",
|
||||
action="store_true",
|
||||
help="Show what would be re-encrypted without making changes.",
|
||||
)
|
||||
|
||||
tenant_group = parser.add_mutually_exclusive_group()
|
||||
tenant_group.add_argument(
|
||||
"--tenant-id",
|
||||
default=None,
|
||||
help="Target a specific tenant schema.",
|
||||
)
|
||||
tenant_group.add_argument(
|
||||
"--all-tenants",
|
||||
action="store_true",
|
||||
help="Iterate all tenants.",
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
old_key = args.old_key if args.old_key else None
|
||||
|
||||
global_version.set_ee()
|
||||
SqlEngine.init_engine(pool_size=5, max_overflow=2)
|
||||
|
||||
if args.dry_run:
|
||||
print("DRY RUN — no changes will be made")
|
||||
|
||||
if args.all_tenants:
|
||||
tenant_ids = get_all_tenant_ids()
|
||||
print(f"Found {len(tenant_ids)} tenant(s)")
|
||||
failed_tenants: list[str] = []
|
||||
for tid in tenant_ids:
|
||||
try:
|
||||
_run_for_tenant(tid, old_key, dry_run=args.dry_run)
|
||||
except Exception as e:
|
||||
print(f" ERROR for tenant {tid}: {e}")
|
||||
failed_tenants.append(tid)
|
||||
if failed_tenants:
|
||||
print(f"FAILED tenants ({len(failed_tenants)}): {failed_tenants}")
|
||||
sys.exit(1)
|
||||
else:
|
||||
tenant_id = args.tenant_id or POSTGRES_DEFAULT_SCHEMA
|
||||
_run_for_tenant(tenant_id, old_key, dry_run=args.dry_run)
|
||||
|
||||
print("Done.")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,90 @@
|
||||
"""Test that Credential with nested JSON round-trips through SensitiveValue correctly.
|
||||
|
||||
Exercises the full encrypt → store → read → decrypt → SensitiveValue path
|
||||
with realistic nested OAuth credential data, and verifies SQLAlchemy dirty
|
||||
tracking works with nested dict comparison.
|
||||
|
||||
Requires a running Postgres instance.
|
||||
"""
|
||||
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from onyx.configs.constants import DocumentSource
|
||||
from onyx.db.models import Credential
|
||||
from onyx.utils.sensitive import SensitiveValue
|
||||
|
||||
# NOTE: this is not the real shape of a Drive credential,
|
||||
# but it is intended to test nested JSON credential handling
|
||||
|
||||
_NESTED_CRED_JSON = {
|
||||
"oauth_tokens": {
|
||||
"access_token": "ya29.abc123",
|
||||
"refresh_token": "1//xEg-def456",
|
||||
},
|
||||
"scopes": ["read", "write", "admin"],
|
||||
"client_config": {
|
||||
"client_id": "123.apps.googleusercontent.com",
|
||||
"client_secret": "GOCSPX-secret",
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def test_nested_credential_json_round_trip(db_session: Session) -> None:
|
||||
"""Nested OAuth credential survives encrypt → store → read → decrypt."""
|
||||
credential = Credential(
|
||||
source=DocumentSource.GOOGLE_DRIVE,
|
||||
credential_json=_NESTED_CRED_JSON,
|
||||
)
|
||||
db_session.add(credential)
|
||||
db_session.flush()
|
||||
|
||||
# Immediate read (no DB round-trip) — tests the set event wrapping
|
||||
assert isinstance(credential.credential_json, SensitiveValue)
|
||||
assert credential.credential_json.get_value(apply_mask=False) == _NESTED_CRED_JSON
|
||||
|
||||
# DB round-trip — tests process_result_value
|
||||
db_session.expire(credential)
|
||||
reloaded = credential.credential_json
|
||||
assert isinstance(reloaded, SensitiveValue)
|
||||
assert reloaded.get_value(apply_mask=False) == _NESTED_CRED_JSON
|
||||
|
||||
db_session.rollback()
|
||||
|
||||
|
||||
def test_reassign_same_nested_json_not_dirty(db_session: Session) -> None:
|
||||
"""Re-assigning the same nested dict should not mark the session dirty."""
|
||||
credential = Credential(
|
||||
source=DocumentSource.GOOGLE_DRIVE,
|
||||
credential_json=_NESTED_CRED_JSON,
|
||||
)
|
||||
db_session.add(credential)
|
||||
db_session.flush()
|
||||
|
||||
# Clear dirty state from the insert
|
||||
db_session.expire(credential)
|
||||
_ = credential.credential_json # force reload
|
||||
|
||||
# Re-assign identical value
|
||||
credential.credential_json = _NESTED_CRED_JSON # type: ignore[assignment]
|
||||
assert not db_session.is_modified(credential)
|
||||
|
||||
db_session.rollback()
|
||||
|
||||
|
||||
def test_assign_different_nested_json_is_dirty(db_session: Session) -> None:
|
||||
"""Assigning a different nested dict should mark the session dirty."""
|
||||
credential = Credential(
|
||||
source=DocumentSource.GOOGLE_DRIVE,
|
||||
credential_json=_NESTED_CRED_JSON,
|
||||
)
|
||||
db_session.add(credential)
|
||||
db_session.flush()
|
||||
|
||||
db_session.expire(credential)
|
||||
_ = credential.credential_json # force reload
|
||||
|
||||
modified_cred = {**_NESTED_CRED_JSON, "scopes": ["read"]}
|
||||
credential.credential_json = modified_cred # type: ignore[assignment]
|
||||
assert db_session.is_modified(credential)
|
||||
|
||||
db_session.rollback()
|
||||
@@ -0,0 +1,305 @@
|
||||
"""Tests for rotate_encryption_key against real Postgres.
|
||||
|
||||
Uses real ORM models (Credential, InternetSearchProvider) and the actual
|
||||
Postgres database. Discovery is mocked in rotation tests to scope mutations
|
||||
to only the test rows — the real _discover_encrypted_columns walk is tested
|
||||
separately in TestDiscoverEncryptedColumns.
|
||||
|
||||
Requires a running Postgres instance. Run with::
|
||||
|
||||
python -m dotenv -f .vscode/.env run -- pytest tests/external_dependency_unit/db/test_rotate_encryption_key.py
|
||||
"""
|
||||
|
||||
import json
|
||||
from collections.abc import Generator
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
from sqlalchemy import LargeBinary
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy import text
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from ee.onyx.utils.encryption import _decrypt_bytes
|
||||
from ee.onyx.utils.encryption import _encrypt_string
|
||||
from ee.onyx.utils.encryption import _get_trimmed_key
|
||||
from onyx.configs.constants import DocumentSource
|
||||
from onyx.db.models import Credential
|
||||
from onyx.db.models import EncryptedJson
|
||||
from onyx.db.models import EncryptedString
|
||||
from onyx.db.models import InternetSearchProvider
|
||||
from onyx.db.rotate_encryption_key import _discover_encrypted_columns
|
||||
from onyx.db.rotate_encryption_key import rotate_encryption_key
|
||||
from onyx.utils.variable_functionality import fetch_versioned_implementation
|
||||
from onyx.utils.variable_functionality import global_version
|
||||
|
||||
EE_MODULE = "ee.onyx.utils.encryption"
|
||||
ROTATE_MODULE = "onyx.db.rotate_encryption_key"
|
||||
|
||||
OLD_KEY = "o" * 16
|
||||
NEW_KEY = "n" * 16
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _enable_ee() -> Generator[None, None, None]:
|
||||
prev = global_version._is_ee
|
||||
global_version.set_ee()
|
||||
fetch_versioned_implementation.cache_clear()
|
||||
yield
|
||||
global_version._is_ee = prev
|
||||
fetch_versioned_implementation.cache_clear()
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _clear_key_cache() -> None:
|
||||
_get_trimmed_key.cache_clear()
|
||||
|
||||
|
||||
def _raw_credential_bytes(db_session: Session, credential_id: int) -> bytes | None:
|
||||
"""Read raw bytes from credential_json, bypassing the TypeDecorator."""
|
||||
col = Credential.__table__.c.credential_json
|
||||
stmt = select(col.cast(LargeBinary)).where(
|
||||
Credential.__table__.c.id == credential_id
|
||||
)
|
||||
return db_session.execute(stmt).scalar()
|
||||
|
||||
|
||||
def _raw_isp_bytes(db_session: Session, isp_id: int) -> bytes | None:
|
||||
"""Read raw bytes from InternetSearchProvider.api_key."""
|
||||
col = InternetSearchProvider.__table__.c.api_key
|
||||
stmt = select(col.cast(LargeBinary)).where(
|
||||
InternetSearchProvider.__table__.c.id == isp_id
|
||||
)
|
||||
return db_session.execute(stmt).scalar()
|
||||
|
||||
|
||||
class TestDiscoverEncryptedColumns:
|
||||
"""Verify _discover_encrypted_columns finds real production models."""
|
||||
|
||||
def test_discovers_credential_json(self) -> None:
|
||||
results = _discover_encrypted_columns()
|
||||
found = {
|
||||
(model_cls.__tablename__, col_name, is_json) # type: ignore[attr-defined]
|
||||
for model_cls, col_name, _, is_json in results
|
||||
}
|
||||
assert ("credential", "credential_json", True) in found
|
||||
|
||||
def test_discovers_internet_search_provider_api_key(self) -> None:
|
||||
results = _discover_encrypted_columns()
|
||||
found = {
|
||||
(model_cls.__tablename__, col_name, is_json) # type: ignore[attr-defined]
|
||||
for model_cls, col_name, _, is_json in results
|
||||
}
|
||||
assert ("internet_search_provider", "api_key", False) in found
|
||||
|
||||
def test_all_encrypted_string_columns_are_not_json(self) -> None:
|
||||
results = _discover_encrypted_columns()
|
||||
for model_cls, col_name, _, is_json in results:
|
||||
col = getattr(model_cls, col_name).property.columns[0]
|
||||
if isinstance(col.type, EncryptedString):
|
||||
assert not is_json, (
|
||||
f"{model_cls.__tablename__}.{col_name} is EncryptedString " # type: ignore[attr-defined]
|
||||
f"but is_json={is_json}"
|
||||
)
|
||||
|
||||
def test_all_encrypted_json_columns_are_json(self) -> None:
|
||||
results = _discover_encrypted_columns()
|
||||
for model_cls, col_name, _, is_json in results:
|
||||
col = getattr(model_cls, col_name).property.columns[0]
|
||||
if isinstance(col.type, EncryptedJson):
|
||||
assert is_json, (
|
||||
f"{model_cls.__tablename__}.{col_name} is EncryptedJson " # type: ignore[attr-defined]
|
||||
f"but is_json={is_json}"
|
||||
)
|
||||
|
||||
|
||||
class TestRotateCredential:
|
||||
"""Test rotation against the real Credential table (EncryptedJson).
|
||||
|
||||
Discovery is scoped to only the Credential model to avoid mutating
|
||||
other tables in the test database.
|
||||
"""
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _limit_discovery(self) -> Generator[None, None, None]:
|
||||
with patch(
|
||||
f"{ROTATE_MODULE}._discover_encrypted_columns",
|
||||
return_value=[(Credential, "credential_json", ["id"], True)],
|
||||
):
|
||||
yield
|
||||
|
||||
@pytest.fixture()
|
||||
def credential_id(
|
||||
self, db_session: Session, tenant_context: None # noqa: ARG002
|
||||
) -> Generator[int, None, None]:
|
||||
"""Insert a Credential row with raw encrypted bytes, clean up after."""
|
||||
config = {"api_key": "sk-test-1234", "endpoint": "https://example.com"}
|
||||
encrypted = _encrypt_string(json.dumps(config), key=OLD_KEY)
|
||||
|
||||
result = db_session.execute(
|
||||
text(
|
||||
"INSERT INTO credential "
|
||||
"(source, credential_json, admin_public, curator_public) "
|
||||
"VALUES (:source, :cred_json, true, false) "
|
||||
"RETURNING id"
|
||||
),
|
||||
{"source": DocumentSource.INGESTION_API.value, "cred_json": encrypted},
|
||||
)
|
||||
cred_id = result.scalar_one()
|
||||
db_session.commit()
|
||||
|
||||
yield cred_id
|
||||
|
||||
db_session.execute(
|
||||
text("DELETE FROM credential WHERE id = :id"), {"id": cred_id}
|
||||
)
|
||||
db_session.commit()
|
||||
|
||||
def test_rotates_credential_json(
|
||||
self, db_session: Session, credential_id: int
|
||||
) -> None:
|
||||
with (
|
||||
patch(f"{ROTATE_MODULE}.ENCRYPTION_KEY_SECRET", NEW_KEY),
|
||||
patch(f"{EE_MODULE}.ENCRYPTION_KEY_SECRET", NEW_KEY),
|
||||
):
|
||||
totals = rotate_encryption_key(db_session, old_key=OLD_KEY)
|
||||
|
||||
assert totals.get("credential.credential_json", 0) >= 1
|
||||
|
||||
raw = _raw_credential_bytes(db_session, credential_id)
|
||||
assert raw is not None
|
||||
decrypted = json.loads(_decrypt_bytes(raw, key=NEW_KEY))
|
||||
assert decrypted["api_key"] == "sk-test-1234"
|
||||
assert decrypted["endpoint"] == "https://example.com"
|
||||
|
||||
def test_skips_already_rotated(
|
||||
self, db_session: Session, credential_id: int
|
||||
) -> None:
|
||||
with (
|
||||
patch(f"{ROTATE_MODULE}.ENCRYPTION_KEY_SECRET", NEW_KEY),
|
||||
patch(f"{EE_MODULE}.ENCRYPTION_KEY_SECRET", NEW_KEY),
|
||||
):
|
||||
rotate_encryption_key(db_session, old_key=OLD_KEY)
|
||||
_ = rotate_encryption_key(db_session, old_key=OLD_KEY)
|
||||
|
||||
raw = _raw_credential_bytes(db_session, credential_id)
|
||||
assert raw is not None
|
||||
decrypted = json.loads(_decrypt_bytes(raw, key=NEW_KEY))
|
||||
assert decrypted["api_key"] == "sk-test-1234"
|
||||
|
||||
def test_dry_run_does_not_modify(
|
||||
self, db_session: Session, credential_id: int
|
||||
) -> None:
|
||||
original = _raw_credential_bytes(db_session, credential_id)
|
||||
|
||||
with (
|
||||
patch(f"{ROTATE_MODULE}.ENCRYPTION_KEY_SECRET", NEW_KEY),
|
||||
patch(f"{EE_MODULE}.ENCRYPTION_KEY_SECRET", NEW_KEY),
|
||||
):
|
||||
totals = rotate_encryption_key(db_session, old_key=OLD_KEY, dry_run=True)
|
||||
|
||||
assert totals.get("credential.credential_json", 0) >= 1
|
||||
|
||||
raw_after = _raw_credential_bytes(db_session, credential_id)
|
||||
assert raw_after == original
|
||||
|
||||
|
||||
class TestRotateInternetSearchProvider:
|
||||
"""Test rotation against the real InternetSearchProvider table (EncryptedString).
|
||||
|
||||
Discovery is scoped to only the InternetSearchProvider model to avoid
|
||||
mutating other tables in the test database.
|
||||
"""
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _limit_discovery(self) -> Generator[None, None, None]:
|
||||
with patch(
|
||||
f"{ROTATE_MODULE}._discover_encrypted_columns",
|
||||
return_value=[
|
||||
(InternetSearchProvider, "api_key", ["id"], False),
|
||||
],
|
||||
):
|
||||
yield
|
||||
|
||||
@pytest.fixture()
|
||||
def isp_id(
|
||||
self, db_session: Session, tenant_context: None # noqa: ARG002
|
||||
) -> Generator[int, None, None]:
|
||||
"""Insert an InternetSearchProvider row with raw encrypted bytes."""
|
||||
encrypted = _encrypt_string("sk-secret-api-key", key=OLD_KEY)
|
||||
|
||||
result = db_session.execute(
|
||||
text(
|
||||
"INSERT INTO internet_search_provider "
|
||||
"(name, provider_type, api_key, is_active) "
|
||||
"VALUES (:name, :ptype, :api_key, false) "
|
||||
"RETURNING id"
|
||||
),
|
||||
{
|
||||
"name": f"test-rotation-{id(self)}",
|
||||
"ptype": "test",
|
||||
"api_key": encrypted,
|
||||
},
|
||||
)
|
||||
isp_id = result.scalar_one()
|
||||
db_session.commit()
|
||||
|
||||
yield isp_id
|
||||
|
||||
db_session.execute(
|
||||
text("DELETE FROM internet_search_provider WHERE id = :id"),
|
||||
{"id": isp_id},
|
||||
)
|
||||
db_session.commit()
|
||||
|
||||
def test_rotates_api_key(self, db_session: Session, isp_id: int) -> None:
|
||||
with (
|
||||
patch(f"{ROTATE_MODULE}.ENCRYPTION_KEY_SECRET", NEW_KEY),
|
||||
patch(f"{EE_MODULE}.ENCRYPTION_KEY_SECRET", NEW_KEY),
|
||||
):
|
||||
totals = rotate_encryption_key(db_session, old_key=OLD_KEY)
|
||||
|
||||
assert totals.get("internet_search_provider.api_key", 0) >= 1
|
||||
|
||||
raw = _raw_isp_bytes(db_session, isp_id)
|
||||
assert raw is not None
|
||||
assert _decrypt_bytes(raw, key=NEW_KEY) == "sk-secret-api-key"
|
||||
|
||||
def test_rotates_from_unencrypted(
|
||||
self, db_session: Session, tenant_context: None # noqa: ARG002
|
||||
) -> None:
|
||||
"""Test rotating data that was stored without any encryption key."""
|
||||
result = db_session.execute(
|
||||
text(
|
||||
"INSERT INTO internet_search_provider "
|
||||
"(name, provider_type, api_key, is_active) "
|
||||
"VALUES (:name, :ptype, :api_key, false) "
|
||||
"RETURNING id"
|
||||
),
|
||||
{
|
||||
"name": f"test-raw-{id(self)}",
|
||||
"ptype": "test",
|
||||
"api_key": b"raw-api-key",
|
||||
},
|
||||
)
|
||||
isp_id = result.scalar_one()
|
||||
db_session.commit()
|
||||
|
||||
try:
|
||||
with (
|
||||
patch(f"{ROTATE_MODULE}.ENCRYPTION_KEY_SECRET", NEW_KEY),
|
||||
patch(f"{EE_MODULE}.ENCRYPTION_KEY_SECRET", NEW_KEY),
|
||||
):
|
||||
totals = rotate_encryption_key(db_session, old_key=None)
|
||||
|
||||
assert totals.get("internet_search_provider.api_key", 0) >= 1
|
||||
|
||||
raw = _raw_isp_bytes(db_session, isp_id)
|
||||
assert raw is not None
|
||||
assert _decrypt_bytes(raw, key=NEW_KEY) == "raw-api-key"
|
||||
finally:
|
||||
db_session.execute(
|
||||
text("DELETE FROM internet_search_provider WHERE id = :id"),
|
||||
{"id": isp_id},
|
||||
)
|
||||
db_session.commit()
|
||||
@@ -0,0 +1,85 @@
|
||||
"""Tests that SlackBot CRUD operations return properly typed SensitiveValue fields.
|
||||
|
||||
Regression test for the bug where insert_slack_bot/update_slack_bot returned
|
||||
objects with raw string tokens instead of SensitiveValue wrappers, causing
|
||||
'str object has no attribute get_value' errors in SlackBot.from_model().
|
||||
"""
|
||||
|
||||
from uuid import uuid4
|
||||
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from onyx.db.slack_bot import insert_slack_bot
|
||||
from onyx.db.slack_bot import update_slack_bot
|
||||
from onyx.server.manage.models import SlackBot
|
||||
from onyx.utils.sensitive import SensitiveValue
|
||||
|
||||
|
||||
def _unique(prefix: str) -> str:
|
||||
return f"{prefix}-{uuid4().hex[:8]}"
|
||||
|
||||
|
||||
def test_insert_slack_bot_returns_sensitive_values(db_session: Session) -> None:
|
||||
bot_token = _unique("xoxb-insert")
|
||||
app_token = _unique("xapp-insert")
|
||||
user_token = _unique("xoxp-insert")
|
||||
|
||||
slack_bot = insert_slack_bot(
|
||||
db_session=db_session,
|
||||
name=_unique("test-bot-insert"),
|
||||
enabled=True,
|
||||
bot_token=bot_token,
|
||||
app_token=app_token,
|
||||
user_token=user_token,
|
||||
)
|
||||
|
||||
assert isinstance(slack_bot.bot_token, SensitiveValue)
|
||||
assert isinstance(slack_bot.app_token, SensitiveValue)
|
||||
assert isinstance(slack_bot.user_token, SensitiveValue)
|
||||
|
||||
assert slack_bot.bot_token.get_value(apply_mask=False) == bot_token
|
||||
assert slack_bot.app_token.get_value(apply_mask=False) == app_token
|
||||
assert slack_bot.user_token.get_value(apply_mask=False) == user_token
|
||||
|
||||
# Verify from_model works without error
|
||||
pydantic_bot = SlackBot.from_model(slack_bot)
|
||||
assert pydantic_bot.bot_token # masked, but not empty
|
||||
assert pydantic_bot.app_token
|
||||
|
||||
|
||||
def test_update_slack_bot_returns_sensitive_values(db_session: Session) -> None:
|
||||
slack_bot = insert_slack_bot(
|
||||
db_session=db_session,
|
||||
name=_unique("test-bot-update"),
|
||||
enabled=True,
|
||||
bot_token=_unique("xoxb-update"),
|
||||
app_token=_unique("xapp-update"),
|
||||
)
|
||||
|
||||
new_bot_token = _unique("xoxb-update-new")
|
||||
new_app_token = _unique("xapp-update-new")
|
||||
new_user_token = _unique("xoxp-update-new")
|
||||
|
||||
updated = update_slack_bot(
|
||||
db_session=db_session,
|
||||
slack_bot_id=slack_bot.id,
|
||||
name=_unique("test-bot-updated"),
|
||||
enabled=False,
|
||||
bot_token=new_bot_token,
|
||||
app_token=new_app_token,
|
||||
user_token=new_user_token,
|
||||
)
|
||||
|
||||
assert isinstance(updated.bot_token, SensitiveValue)
|
||||
assert isinstance(updated.app_token, SensitiveValue)
|
||||
assert isinstance(updated.user_token, SensitiveValue)
|
||||
|
||||
assert updated.bot_token.get_value(apply_mask=False) == new_bot_token
|
||||
assert updated.app_token.get_value(apply_mask=False) == new_app_token
|
||||
assert updated.user_token.get_value(apply_mask=False) == new_user_token
|
||||
|
||||
# Verify from_model works without error
|
||||
pydantic_bot = SlackBot.from_model(updated)
|
||||
assert pydantic_bot.bot_token
|
||||
assert pydantic_bot.app_token
|
||||
assert pydantic_bot.user_token is not None
|
||||
@@ -148,8 +148,16 @@ class TestOAuthConfigCRUD:
|
||||
)
|
||||
|
||||
# Secrets should be preserved
|
||||
assert updated_config.client_id == original_client_id
|
||||
assert updated_config.client_secret == original_client_secret
|
||||
assert updated_config.client_id is not None
|
||||
assert original_client_id is not None
|
||||
assert updated_config.client_id.get_value(
|
||||
apply_mask=False
|
||||
) == original_client_id.get_value(apply_mask=False)
|
||||
assert updated_config.client_secret is not None
|
||||
assert original_client_secret is not None
|
||||
assert updated_config.client_secret.get_value(
|
||||
apply_mask=False
|
||||
) == original_client_secret.get_value(apply_mask=False)
|
||||
# But name should be updated
|
||||
assert updated_config.name == new_name
|
||||
|
||||
@@ -173,9 +181,14 @@ class TestOAuthConfigCRUD:
|
||||
)
|
||||
|
||||
# client_id should be cleared (empty string)
|
||||
assert updated_config.client_id == ""
|
||||
assert updated_config.client_id is not None
|
||||
assert updated_config.client_id.get_value(apply_mask=False) == ""
|
||||
# client_secret should be preserved
|
||||
assert updated_config.client_secret == original_client_secret
|
||||
assert updated_config.client_secret is not None
|
||||
assert original_client_secret is not None
|
||||
assert updated_config.client_secret.get_value(
|
||||
apply_mask=False
|
||||
) == original_client_secret.get_value(apply_mask=False)
|
||||
|
||||
def test_update_oauth_config_clear_client_secret(self, db_session: Session) -> None:
|
||||
"""Test clearing client_secret while preserving client_id"""
|
||||
@@ -190,9 +203,14 @@ class TestOAuthConfigCRUD:
|
||||
)
|
||||
|
||||
# client_secret should be cleared (empty string)
|
||||
assert updated_config.client_secret == ""
|
||||
assert updated_config.client_secret is not None
|
||||
assert updated_config.client_secret.get_value(apply_mask=False) == ""
|
||||
# client_id should be preserved
|
||||
assert updated_config.client_id == original_client_id
|
||||
assert updated_config.client_id is not None
|
||||
assert original_client_id is not None
|
||||
assert updated_config.client_id.get_value(
|
||||
apply_mask=False
|
||||
) == original_client_id.get_value(apply_mask=False)
|
||||
|
||||
def test_update_oauth_config_clear_both_secrets(self, db_session: Session) -> None:
|
||||
"""Test clearing both client_id and client_secret"""
|
||||
@@ -207,8 +225,10 @@ class TestOAuthConfigCRUD:
|
||||
)
|
||||
|
||||
# Both should be cleared (empty strings)
|
||||
assert updated_config.client_id == ""
|
||||
assert updated_config.client_secret == ""
|
||||
assert updated_config.client_id is not None
|
||||
assert updated_config.client_id.get_value(apply_mask=False) == ""
|
||||
assert updated_config.client_secret is not None
|
||||
assert updated_config.client_secret.get_value(apply_mask=False) == ""
|
||||
|
||||
def test_update_oauth_config_authorization_url(self, db_session: Session) -> None:
|
||||
"""Test updating authorization_url"""
|
||||
@@ -275,7 +295,8 @@ class TestOAuthConfigCRUD:
|
||||
assert updated_config.token_url == new_token_url
|
||||
assert updated_config.scopes == new_scopes
|
||||
assert updated_config.additional_params == new_params
|
||||
assert updated_config.client_id == new_client_id
|
||||
assert updated_config.client_id is not None
|
||||
assert updated_config.client_id.get_value(apply_mask=False) == new_client_id
|
||||
|
||||
def test_delete_oauth_config(self, db_session: Session) -> None:
|
||||
"""Test deleting an OAuth configuration"""
|
||||
@@ -416,7 +437,8 @@ class TestOAuthUserTokenCRUD:
|
||||
assert user_token.id is not None
|
||||
assert user_token.oauth_config_id == oauth_config.id
|
||||
assert user_token.user_id == user.id
|
||||
assert user_token.token_data == token_data
|
||||
assert user_token.token_data is not None
|
||||
assert user_token.token_data.get_value(apply_mask=False) == token_data
|
||||
assert user_token.created_at is not None
|
||||
assert user_token.updated_at is not None
|
||||
|
||||
@@ -446,8 +468,13 @@ class TestOAuthUserTokenCRUD:
|
||||
|
||||
# Should be the same token record (updated, not inserted)
|
||||
assert updated_token.id == initial_token_id
|
||||
assert updated_token.token_data == updated_token_data
|
||||
assert updated_token.token_data != initial_token_data
|
||||
assert updated_token.token_data is not None
|
||||
assert (
|
||||
updated_token.token_data.get_value(apply_mask=False) == updated_token_data
|
||||
)
|
||||
assert (
|
||||
updated_token.token_data.get_value(apply_mask=False) != initial_token_data
|
||||
)
|
||||
|
||||
def test_get_user_oauth_token(self, db_session: Session) -> None:
|
||||
"""Test retrieving a user's OAuth token"""
|
||||
@@ -463,7 +490,8 @@ class TestOAuthUserTokenCRUD:
|
||||
|
||||
assert retrieved_token is not None
|
||||
assert retrieved_token.id == created_token.id
|
||||
assert retrieved_token.token_data == token_data
|
||||
assert retrieved_token.token_data is not None
|
||||
assert retrieved_token.token_data.get_value(apply_mask=False) == token_data
|
||||
|
||||
def test_get_user_oauth_token_not_found(self, db_session: Session) -> None:
|
||||
"""Test retrieving a non-existent user token returns None"""
|
||||
@@ -519,7 +547,8 @@ class TestOAuthUserTokenCRUD:
|
||||
retrieved_token = get_user_oauth_token(oauth_config.id, user.id, db_session)
|
||||
assert retrieved_token is not None
|
||||
assert retrieved_token.id == updated_token.id
|
||||
assert retrieved_token.token_data == token_data2
|
||||
assert retrieved_token.token_data is not None
|
||||
assert retrieved_token.token_data.get_value(apply_mask=False) == token_data2
|
||||
|
||||
def test_cascade_delete_user_tokens_on_config_deletion(
|
||||
self, db_session: Session
|
||||
|
||||
@@ -374,8 +374,14 @@ class TestOAuthTokenManagerCodeExchange:
|
||||
assert call_args[0][0] == oauth_config.token_url
|
||||
assert call_args[1]["data"]["grant_type"] == "authorization_code"
|
||||
assert call_args[1]["data"]["code"] == "auth_code_123"
|
||||
assert call_args[1]["data"]["client_id"] == oauth_config.client_id
|
||||
assert call_args[1]["data"]["client_secret"] == oauth_config.client_secret
|
||||
assert oauth_config.client_id is not None
|
||||
assert oauth_config.client_secret is not None
|
||||
assert call_args[1]["data"]["client_id"] == oauth_config.client_id.get_value(
|
||||
apply_mask=False
|
||||
)
|
||||
assert call_args[1]["data"][
|
||||
"client_secret"
|
||||
] == oauth_config.client_secret.get_value(apply_mask=False)
|
||||
assert call_args[1]["data"]["redirect_uri"] == "https://example.com/callback"
|
||||
|
||||
@patch("onyx.auth.oauth_token_manager.requests.post")
|
||||
|
||||
@@ -950,6 +950,7 @@ from onyx.server.query_and_chat.streaming_models import Packet
|
||||
from onyx.server.query_and_chat.streaming_models import PythonToolDelta
|
||||
from onyx.server.query_and_chat.streaming_models import PythonToolStart
|
||||
from onyx.server.query_and_chat.streaming_models import SectionEnd
|
||||
from onyx.server.query_and_chat.streaming_models import ToolCallArgumentDelta
|
||||
from onyx.tools.tool_implementations.python.python_tool import PythonTool
|
||||
from tests.external_dependency_unit.answer.stream_test_builder import StreamTestBuilder
|
||||
from tests.external_dependency_unit.answer.stream_test_utils import create_chat_session
|
||||
@@ -1294,9 +1295,18 @@ def test_code_interpreter_replay_packets_include_code_and_output(
|
||||
).expect(
|
||||
Packet(
|
||||
placement=create_placement(0),
|
||||
obj=PythonToolStart(code=code),
|
||||
obj=ToolCallArgumentDelta(
|
||||
tool_type="python",
|
||||
argument_deltas={"code": code},
|
||||
),
|
||||
),
|
||||
forward=2,
|
||||
).expect(
|
||||
Packet(
|
||||
placement=create_placement(0),
|
||||
obj=PythonToolStart(code=code),
|
||||
),
|
||||
forward=False,
|
||||
).expect(
|
||||
Packet(
|
||||
placement=create_placement(0),
|
||||
|
||||
@@ -64,7 +64,8 @@ class TestBotConfigAPI:
|
||||
db_session.commit()
|
||||
|
||||
assert config is not None
|
||||
assert config.bot_token == "test_token_123"
|
||||
assert config.bot_token is not None
|
||||
assert config.bot_token.get_value(apply_mask=False) == "test_token_123"
|
||||
|
||||
# Cleanup
|
||||
delete_discord_bot_config(db_session)
|
||||
|
||||
165
backend/tests/unit/ee/onyx/utils/test_encryption.py
Normal file
165
backend/tests/unit/ee/onyx/utils/test_encryption.py
Normal file
@@ -0,0 +1,165 @@
|
||||
"""Tests for EE AES-CBC encryption/decryption with explicit key support.
|
||||
|
||||
With EE mode enabled (via conftest), fetch_versioned_implementation resolves
|
||||
to the EE implementations, so no patching of the MIT layer is needed.
|
||||
"""
|
||||
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
from ee.onyx.utils.encryption import _decrypt_bytes
|
||||
from ee.onyx.utils.encryption import _encrypt_string
|
||||
from ee.onyx.utils.encryption import _get_trimmed_key
|
||||
from ee.onyx.utils.encryption import decrypt_bytes_to_string
|
||||
from ee.onyx.utils.encryption import encrypt_string_to_bytes
|
||||
|
||||
EE_MODULE = "ee.onyx.utils.encryption"
|
||||
|
||||
# Keys must be exactly 16, 24, or 32 bytes for AES
|
||||
KEY_16 = "a" * 16
|
||||
KEY_16_ALT = "b" * 16
|
||||
KEY_24 = "d" * 24
|
||||
KEY_32 = "c" * 32
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _clear_key_cache() -> None:
|
||||
_get_trimmed_key.cache_clear()
|
||||
|
||||
|
||||
class TestEncryptDecryptRoundTrip:
|
||||
def test_roundtrip_with_env_key(self) -> None:
|
||||
with patch(f"{EE_MODULE}.ENCRYPTION_KEY_SECRET", KEY_16):
|
||||
encrypted = _encrypt_string("hello world")
|
||||
assert encrypted != b"hello world"
|
||||
assert _decrypt_bytes(encrypted) == "hello world"
|
||||
|
||||
def test_roundtrip_with_explicit_key(self) -> None:
|
||||
encrypted = _encrypt_string("secret data", key=KEY_32)
|
||||
assert encrypted != b"secret data"
|
||||
assert _decrypt_bytes(encrypted, key=KEY_32) == "secret data"
|
||||
|
||||
def test_roundtrip_no_key(self) -> None:
|
||||
"""Without any key, data is raw-encoded (no encryption)."""
|
||||
with patch(f"{EE_MODULE}.ENCRYPTION_KEY_SECRET", ""):
|
||||
encrypted = _encrypt_string("plain text")
|
||||
assert encrypted == b"plain text"
|
||||
assert _decrypt_bytes(encrypted) == "plain text"
|
||||
|
||||
def test_explicit_key_overrides_env(self) -> None:
|
||||
with patch(f"{EE_MODULE}.ENCRYPTION_KEY_SECRET", KEY_16):
|
||||
encrypted = _encrypt_string("data", key=KEY_16_ALT)
|
||||
with pytest.raises(ValueError):
|
||||
_decrypt_bytes(encrypted, key=KEY_16)
|
||||
assert _decrypt_bytes(encrypted, key=KEY_16_ALT) == "data"
|
||||
|
||||
def test_different_encryptions_produce_different_bytes(self) -> None:
|
||||
"""Each encryption uses a random IV, so results differ."""
|
||||
a = _encrypt_string("same", key=KEY_16)
|
||||
b = _encrypt_string("same", key=KEY_16)
|
||||
assert a != b
|
||||
|
||||
def test_roundtrip_empty_string(self) -> None:
|
||||
encrypted = _encrypt_string("", key=KEY_16)
|
||||
assert encrypted != b""
|
||||
assert _decrypt_bytes(encrypted, key=KEY_16) == ""
|
||||
|
||||
def test_roundtrip_unicode(self) -> None:
|
||||
text = "日本語テスト 🔐 émojis"
|
||||
encrypted = _encrypt_string(text, key=KEY_16)
|
||||
assert _decrypt_bytes(encrypted, key=KEY_16) == text
|
||||
|
||||
|
||||
class TestDecryptFallbackBehavior:
|
||||
def test_wrong_env_key_falls_back_to_raw_decode(self) -> None:
|
||||
"""Default key path: AES fails on non-AES data → fallback to raw decode."""
|
||||
raw = "readable text".encode()
|
||||
with patch(f"{EE_MODULE}.ENCRYPTION_KEY_SECRET", KEY_16):
|
||||
assert _decrypt_bytes(raw) == "readable text"
|
||||
|
||||
def test_explicit_wrong_key_raises(self) -> None:
|
||||
"""Explicit key path: AES fails → raises, no fallback."""
|
||||
encrypted = _encrypt_string("secret", key=KEY_16)
|
||||
with pytest.raises(ValueError):
|
||||
_decrypt_bytes(encrypted, key=KEY_16_ALT)
|
||||
|
||||
def test_explicit_none_key_with_no_env(self) -> None:
|
||||
"""key=None with empty env → raw decode."""
|
||||
with patch(f"{EE_MODULE}.ENCRYPTION_KEY_SECRET", ""):
|
||||
assert _decrypt_bytes(b"hello", key=None) == "hello"
|
||||
|
||||
def test_explicit_empty_string_key(self) -> None:
|
||||
"""key='' means no encryption."""
|
||||
encrypted = _encrypt_string("test", key="")
|
||||
assert encrypted == b"test"
|
||||
assert _decrypt_bytes(encrypted, key="") == "test"
|
||||
|
||||
|
||||
class TestKeyValidation:
|
||||
def test_key_too_short_raises(self) -> None:
|
||||
with pytest.raises(RuntimeError, match="too short"):
|
||||
_encrypt_string("data", key="short")
|
||||
|
||||
def test_16_byte_key(self) -> None:
|
||||
encrypted = _encrypt_string("data", key=KEY_16)
|
||||
assert _decrypt_bytes(encrypted, key=KEY_16) == "data"
|
||||
|
||||
def test_24_byte_key(self) -> None:
|
||||
encrypted = _encrypt_string("data", key=KEY_24)
|
||||
assert _decrypt_bytes(encrypted, key=KEY_24) == "data"
|
||||
|
||||
def test_32_byte_key(self) -> None:
|
||||
encrypted = _encrypt_string("data", key=KEY_32)
|
||||
assert _decrypt_bytes(encrypted, key=KEY_32) == "data"
|
||||
|
||||
def test_long_key_truncated_to_32(self) -> None:
|
||||
"""Keys longer than 32 bytes are truncated to 32."""
|
||||
long_key = "e" * 64
|
||||
encrypted = _encrypt_string("data", key=long_key)
|
||||
assert _decrypt_bytes(encrypted, key=long_key) == "data"
|
||||
|
||||
def test_20_byte_key_trimmed_to_16(self) -> None:
|
||||
"""A 20-byte key is trimmed to the largest valid AES size that fits (16)."""
|
||||
key_20 = "f" * 20
|
||||
encrypted = _encrypt_string("data", key=key_20)
|
||||
assert _decrypt_bytes(encrypted, key=key_20) == "data"
|
||||
|
||||
# Verify it was trimmed to 16 by checking that the first 16 bytes
|
||||
# of the key can also decrypt it
|
||||
key_16_same_prefix = "f" * 16
|
||||
assert _decrypt_bytes(encrypted, key=key_16_same_prefix) == "data"
|
||||
|
||||
def test_25_byte_key_trimmed_to_24(self) -> None:
|
||||
"""A 25-byte key is trimmed to the largest valid AES size that fits (24)."""
|
||||
key_25 = "g" * 25
|
||||
encrypted = _encrypt_string("data", key=key_25)
|
||||
assert _decrypt_bytes(encrypted, key=key_25) == "data"
|
||||
|
||||
key_24_same_prefix = "g" * 24
|
||||
assert _decrypt_bytes(encrypted, key=key_24_same_prefix) == "data"
|
||||
|
||||
def test_30_byte_key_trimmed_to_24(self) -> None:
|
||||
"""A 30-byte key is trimmed to the largest valid AES size that fits (24)."""
|
||||
key_30 = "h" * 30
|
||||
encrypted = _encrypt_string("data", key=key_30)
|
||||
assert _decrypt_bytes(encrypted, key=key_30) == "data"
|
||||
|
||||
key_24_same_prefix = "h" * 24
|
||||
assert _decrypt_bytes(encrypted, key=key_24_same_prefix) == "data"
|
||||
|
||||
|
||||
class TestWrapperFunctions:
|
||||
"""Test encrypt_string_to_bytes / decrypt_bytes_to_string pass key through.
|
||||
|
||||
With EE mode enabled, the wrappers resolve to EE implementations automatically.
|
||||
"""
|
||||
|
||||
def test_wrapper_passes_key(self) -> None:
|
||||
encrypted = encrypt_string_to_bytes("payload", key=KEY_16)
|
||||
assert decrypt_bytes_to_string(encrypted, key=KEY_16) == "payload"
|
||||
|
||||
def test_wrapper_no_key_uses_env(self) -> None:
|
||||
with patch(f"{EE_MODULE}.ENCRYPTION_KEY_SECRET", KEY_32):
|
||||
encrypted = encrypt_string_to_bytes("payload")
|
||||
assert decrypt_bytes_to_string(encrypted) == "payload"
|
||||
630
backend/tests/unit/onyx/chat/test_argument_delta_streaming.py
Normal file
630
backend/tests/unit/onyx/chat/test_argument_delta_streaming.py
Normal file
@@ -0,0 +1,630 @@
|
||||
from typing import Any
|
||||
from unittest.mock import MagicMock
|
||||
from unittest.mock import patch
|
||||
|
||||
from onyx.chat.tool_call_args_streaming import maybe_emit_argument_delta
|
||||
from onyx.server.query_and_chat.placement import Placement
|
||||
from onyx.server.query_and_chat.streaming_models import ToolCallArgumentDelta
|
||||
from onyx.utils.jsonriver import Parser
|
||||
|
||||
|
||||
def _make_tool_call_delta(
|
||||
index: int = 0,
|
||||
name: str | None = None,
|
||||
arguments: str | None = None,
|
||||
function_is_none: bool = False,
|
||||
) -> MagicMock:
|
||||
"""Create a mock tool_call_delta matching the LiteLLM streaming shape."""
|
||||
delta = MagicMock()
|
||||
delta.index = index
|
||||
if function_is_none:
|
||||
delta.function = None
|
||||
else:
|
||||
delta.function = MagicMock()
|
||||
delta.function.name = name
|
||||
delta.function.arguments = arguments
|
||||
return delta
|
||||
|
||||
|
||||
def _make_placement() -> Placement:
|
||||
return Placement(turn_index=0, tab_index=0)
|
||||
|
||||
|
||||
def _mock_tool_class(emit: bool = True) -> MagicMock:
|
||||
cls = MagicMock()
|
||||
cls.should_emit_argument_deltas.return_value = emit
|
||||
return cls
|
||||
|
||||
|
||||
def _collect(
|
||||
tc_map: dict[int, dict[str, Any]],
|
||||
delta: MagicMock,
|
||||
placement: Placement | None = None,
|
||||
parsers: dict[int, Parser] | None = None,
|
||||
) -> list[Any]:
|
||||
"""Run maybe_emit_argument_delta and return the yielded packets."""
|
||||
return list(
|
||||
maybe_emit_argument_delta(
|
||||
tc_map,
|
||||
delta,
|
||||
placement or _make_placement(),
|
||||
parsers if parsers is not None else {},
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def _stream_fragments(
|
||||
fragments: list[str],
|
||||
tc_map: dict[int, dict[str, Any]],
|
||||
placement: Placement | None = None,
|
||||
) -> list[str]:
|
||||
"""Feed fragments into maybe_emit_argument_delta one by one, returning
|
||||
all emitted content values concatenated per-key as a flat list."""
|
||||
pl = placement or _make_placement()
|
||||
parsers: dict[int, Parser] = {}
|
||||
emitted: list[str] = []
|
||||
for frag in fragments:
|
||||
tc_map[0]["arguments"] += frag
|
||||
delta = _make_tool_call_delta(arguments=frag)
|
||||
for packet in maybe_emit_argument_delta(tc_map, delta, pl, parsers=parsers):
|
||||
obj = packet.obj
|
||||
assert isinstance(obj, ToolCallArgumentDelta)
|
||||
for value in obj.argument_deltas.values():
|
||||
emitted.append(value)
|
||||
return emitted
|
||||
|
||||
|
||||
class TestMaybeEmitArgumentDeltaGuards:
|
||||
"""Tests for conditions that cause no packet to be emitted."""
|
||||
|
||||
@patch("onyx.chat.tool_call_args_streaming._get_tool_class")
|
||||
def test_no_emission_when_tool_does_not_opt_in(
|
||||
self, mock_get_tool: MagicMock
|
||||
) -> None:
|
||||
"""Tools that return False from should_emit_argument_deltas emit nothing."""
|
||||
mock_get_tool.return_value = _mock_tool_class(emit=False)
|
||||
|
||||
tc_map: dict[int, dict[str, Any]] = {
|
||||
0: {"id": "tc_1", "name": "python", "arguments": '{"code": "x'}
|
||||
}
|
||||
assert _collect(tc_map, _make_tool_call_delta(arguments="x")) == []
|
||||
|
||||
@patch("onyx.chat.tool_call_args_streaming._get_tool_class")
|
||||
def test_no_emission_when_tool_class_unknown(
|
||||
self, mock_get_tool: MagicMock
|
||||
) -> None:
|
||||
mock_get_tool.return_value = None
|
||||
|
||||
tc_map: dict[int, dict[str, Any]] = {
|
||||
0: {"id": "tc_1", "name": "unknown", "arguments": '{"code": "x'}
|
||||
}
|
||||
assert _collect(tc_map, _make_tool_call_delta(arguments="x")) == []
|
||||
|
||||
@patch("onyx.chat.tool_call_args_streaming._get_tool_class")
|
||||
def test_no_emission_when_no_argument_fragment(
|
||||
self, mock_get_tool: MagicMock
|
||||
) -> None:
|
||||
mock_get_tool.return_value = _mock_tool_class()
|
||||
|
||||
tc_map: dict[int, dict[str, Any]] = {
|
||||
0: {"id": "tc_1", "name": "python", "arguments": '{"code": "x'}
|
||||
}
|
||||
assert _collect(tc_map, _make_tool_call_delta(arguments=None)) == []
|
||||
|
||||
@patch("onyx.chat.tool_call_args_streaming._get_tool_class")
|
||||
def test_no_emission_when_key_value_not_started(
|
||||
self, mock_get_tool: MagicMock
|
||||
) -> None:
|
||||
"""Key exists in JSON but its string value hasn't begun yet."""
|
||||
mock_get_tool.return_value = _mock_tool_class()
|
||||
|
||||
tc_map: dict[int, dict[str, Any]] = {
|
||||
0: {"id": "tc_1", "name": "python", "arguments": '{"code":'}
|
||||
}
|
||||
assert _collect(tc_map, _make_tool_call_delta(arguments=":")) == []
|
||||
|
||||
@patch("onyx.chat.tool_call_args_streaming._get_tool_class")
|
||||
def test_no_emission_before_any_key(self, mock_get_tool: MagicMock) -> None:
|
||||
"""Only the opening brace has arrived — no key to stream yet."""
|
||||
mock_get_tool.return_value = _mock_tool_class()
|
||||
|
||||
tc_map: dict[int, dict[str, Any]] = {
|
||||
0: {"id": "tc_1", "name": "python", "arguments": "{"}
|
||||
}
|
||||
assert _collect(tc_map, _make_tool_call_delta(arguments="{")) == []
|
||||
|
||||
|
||||
class TestMaybeEmitArgumentDeltaBasic:
|
||||
"""Tests for correct packet content and incremental emission."""
|
||||
|
||||
@patch("onyx.chat.tool_call_args_streaming._get_tool_class")
|
||||
def test_emits_packet_with_correct_fields(self, mock_get_tool: MagicMock) -> None:
|
||||
mock_get_tool.return_value = _mock_tool_class()
|
||||
|
||||
tc_map: dict[int, dict[str, Any]] = {
|
||||
0: {"id": "tc_1", "name": "python", "arguments": ""}
|
||||
}
|
||||
fragments = ['{"code": "', "print(1)", '"}']
|
||||
|
||||
pl = _make_placement()
|
||||
parsers: dict[int, Parser] = {}
|
||||
all_packets = []
|
||||
for frag in fragments:
|
||||
tc_map[0]["arguments"] += frag
|
||||
packets = _collect(
|
||||
tc_map, _make_tool_call_delta(arguments=frag), pl, parsers
|
||||
)
|
||||
all_packets.extend(packets)
|
||||
|
||||
assert len(all_packets) >= 1
|
||||
# Verify packet structure
|
||||
obj = all_packets[0].obj
|
||||
assert isinstance(obj, ToolCallArgumentDelta)
|
||||
assert obj.tool_type == "python"
|
||||
# All emitted content should reconstruct the value
|
||||
full_code = ""
|
||||
for p in all_packets:
|
||||
assert isinstance(p.obj, ToolCallArgumentDelta)
|
||||
if "code" in p.obj.argument_deltas:
|
||||
full_code += p.obj.argument_deltas["code"]
|
||||
assert full_code == "print(1)"
|
||||
|
||||
@patch("onyx.chat.tool_call_args_streaming._get_tool_class")
|
||||
def test_emits_only_new_content_on_subsequent_call(
|
||||
self, mock_get_tool: MagicMock
|
||||
) -> None:
|
||||
"""After a first emission, subsequent calls emit only the diff."""
|
||||
mock_get_tool.return_value = _mock_tool_class()
|
||||
|
||||
tc_map: dict[int, dict[str, Any]] = {
|
||||
0: {"id": "tc_1", "name": "python", "arguments": ""}
|
||||
}
|
||||
parsers: dict[int, Parser] = {}
|
||||
pl = _make_placement()
|
||||
|
||||
# First fragment opens the string
|
||||
tc_map[0]["arguments"] = '{"code": "abc'
|
||||
packets_1 = _collect(
|
||||
tc_map, _make_tool_call_delta(arguments='{"code": "abc'), pl, parsers
|
||||
)
|
||||
code_1 = ""
|
||||
for p in packets_1:
|
||||
assert isinstance(p.obj, ToolCallArgumentDelta)
|
||||
code_1 += p.obj.argument_deltas.get("code", "")
|
||||
assert code_1 == "abc"
|
||||
|
||||
# Second fragment appends more
|
||||
tc_map[0]["arguments"] = '{"code": "abcdef'
|
||||
packets_2 = _collect(
|
||||
tc_map, _make_tool_call_delta(arguments="def"), pl, parsers
|
||||
)
|
||||
code_2 = ""
|
||||
for p in packets_2:
|
||||
assert isinstance(p.obj, ToolCallArgumentDelta)
|
||||
code_2 += p.obj.argument_deltas.get("code", "")
|
||||
assert code_2 == "def"
|
||||
|
||||
@patch("onyx.chat.tool_call_args_streaming._get_tool_class")
|
||||
def test_handles_multiple_keys_sequentially(self, mock_get_tool: MagicMock) -> None:
|
||||
"""When a second key starts, emissions switch to that key."""
|
||||
mock_get_tool.return_value = _mock_tool_class()
|
||||
|
||||
tc_map: dict[int, dict[str, Any]] = {
|
||||
0: {"id": "tc_1", "name": "python", "arguments": ""}
|
||||
}
|
||||
fragments = [
|
||||
'{"code": "x',
|
||||
'", "output": "hello',
|
||||
'"}',
|
||||
]
|
||||
|
||||
emitted = _stream_fragments(fragments, tc_map)
|
||||
full = "".join(emitted)
|
||||
assert "x" in full
|
||||
assert "hello" in full
|
||||
|
||||
@patch("onyx.chat.tool_call_args_streaming._get_tool_class")
|
||||
def test_delta_spans_key_boundary(self, mock_get_tool: MagicMock) -> None:
|
||||
"""A single delta contains the end of one value and the start of the next key."""
|
||||
mock_get_tool.return_value = _mock_tool_class()
|
||||
|
||||
tc_map: dict[int, dict[str, Any]] = {
|
||||
0: {"id": "tc_1", "name": "python", "arguments": ""}
|
||||
}
|
||||
fragments = [
|
||||
'{"code": "x',
|
||||
'y", "lang": "py',
|
||||
'"}',
|
||||
]
|
||||
|
||||
emitted = _stream_fragments(fragments, tc_map)
|
||||
full = "".join(emitted)
|
||||
assert "xy" in full
|
||||
assert "py" in full
|
||||
|
||||
@patch("onyx.chat.tool_call_args_streaming._get_tool_class")
|
||||
def test_empty_value_emits_nothing(self, mock_get_tool: MagicMock) -> None:
|
||||
"""An empty string value has nothing to emit."""
|
||||
mock_get_tool.return_value = _mock_tool_class()
|
||||
|
||||
tc_map: dict[int, dict[str, Any]] = {
|
||||
0: {"id": "tc_1", "name": "python", "arguments": ""}
|
||||
}
|
||||
# Opening quote just arrived, value is empty
|
||||
tc_map[0]["arguments"] = '{"code": "'
|
||||
packets = _collect(tc_map, _make_tool_call_delta(arguments='{"code": "'))
|
||||
# No string content yet, so either no packet or empty deltas
|
||||
for p in packets:
|
||||
assert isinstance(p.obj, ToolCallArgumentDelta)
|
||||
assert p.obj.argument_deltas.get("code", "") == ""
|
||||
|
||||
|
||||
class TestMaybeEmitArgumentDeltaDecoding:
|
||||
"""Tests verifying that JSON escape sequences are properly decoded."""
|
||||
|
||||
@patch("onyx.chat.tool_call_args_streaming._get_tool_class")
|
||||
def test_decodes_newlines(self, mock_get_tool: MagicMock) -> None:
|
||||
mock_get_tool.return_value = _mock_tool_class()
|
||||
|
||||
tc_map: dict[int, dict[str, Any]] = {
|
||||
0: {"id": "tc_1", "name": "python", "arguments": ""}
|
||||
}
|
||||
fragments = ['{"code": "line1\\nline2"}']
|
||||
|
||||
emitted = _stream_fragments(fragments, tc_map)
|
||||
assert "".join(emitted) == "line1\nline2"
|
||||
|
||||
@patch("onyx.chat.tool_call_args_streaming._get_tool_class")
|
||||
def test_decodes_tabs(self, mock_get_tool: MagicMock) -> None:
|
||||
mock_get_tool.return_value = _mock_tool_class()
|
||||
|
||||
tc_map: dict[int, dict[str, Any]] = {
|
||||
0: {"id": "tc_1", "name": "python", "arguments": ""}
|
||||
}
|
||||
fragments = ['{"code": "\\tindented"}']
|
||||
|
||||
emitted = _stream_fragments(fragments, tc_map)
|
||||
assert "".join(emitted) == "\tindented"
|
||||
|
||||
@patch("onyx.chat.tool_call_args_streaming._get_tool_class")
|
||||
def test_decodes_escaped_quotes(self, mock_get_tool: MagicMock) -> None:
|
||||
mock_get_tool.return_value = _mock_tool_class()
|
||||
|
||||
tc_map: dict[int, dict[str, Any]] = {
|
||||
0: {"id": "tc_1", "name": "python", "arguments": ""}
|
||||
}
|
||||
fragments = ['{"code": "say \\"hi\\""}']
|
||||
|
||||
emitted = _stream_fragments(fragments, tc_map)
|
||||
assert "".join(emitted) == 'say "hi"'
|
||||
|
||||
@patch("onyx.chat.tool_call_args_streaming._get_tool_class")
|
||||
def test_decodes_escaped_backslashes(self, mock_get_tool: MagicMock) -> None:
|
||||
mock_get_tool.return_value = _mock_tool_class()
|
||||
|
||||
tc_map: dict[int, dict[str, Any]] = {
|
||||
0: {"id": "tc_1", "name": "python", "arguments": ""}
|
||||
}
|
||||
fragments = ['{"code": "path\\\\dir"}']
|
||||
|
||||
emitted = _stream_fragments(fragments, tc_map)
|
||||
assert "".join(emitted) == "path\\dir"
|
||||
|
||||
@patch("onyx.chat.tool_call_args_streaming._get_tool_class")
|
||||
def test_decodes_unicode_escape(self, mock_get_tool: MagicMock) -> None:
|
||||
mock_get_tool.return_value = _mock_tool_class()
|
||||
|
||||
tc_map: dict[int, dict[str, Any]] = {
|
||||
0: {"id": "tc_1", "name": "python", "arguments": ""}
|
||||
}
|
||||
fragments = ['{"code": "\\u0041"}']
|
||||
|
||||
emitted = _stream_fragments(fragments, tc_map)
|
||||
assert "".join(emitted) == "A"
|
||||
|
||||
@patch("onyx.chat.tool_call_args_streaming._get_tool_class")
|
||||
def test_incomplete_escape_at_end_decoded_on_next_chunk(
|
||||
self, mock_get_tool: MagicMock
|
||||
) -> None:
|
||||
"""A trailing backslash (incomplete escape) is completed in the next chunk."""
|
||||
mock_get_tool.return_value = _mock_tool_class()
|
||||
|
||||
tc_map: dict[int, dict[str, Any]] = {
|
||||
0: {"id": "tc_1", "name": "python", "arguments": ""}
|
||||
}
|
||||
fragments = ['{"code": "hello\\', 'n"}']
|
||||
|
||||
emitted = _stream_fragments(fragments, tc_map)
|
||||
assert "".join(emitted) == "hello\n"
|
||||
|
||||
@patch("onyx.chat.tool_call_args_streaming._get_tool_class")
|
||||
def test_incomplete_unicode_escape_completed_on_next_chunk(
|
||||
self, mock_get_tool: MagicMock
|
||||
) -> None:
|
||||
"""A partial \\uXX sequence is completed in the next chunk."""
|
||||
mock_get_tool.return_value = _mock_tool_class()
|
||||
|
||||
tc_map: dict[int, dict[str, Any]] = {
|
||||
0: {"id": "tc_1", "name": "python", "arguments": ""}
|
||||
}
|
||||
fragments = ['{"code": "hello\\u00', '41"}']
|
||||
|
||||
emitted = _stream_fragments(fragments, tc_map)
|
||||
assert "".join(emitted) == "helloA"
|
||||
|
||||
|
||||
class TestArgumentDeltaStreamingE2E:
|
||||
"""Simulates realistic sequences of LLM argument deltas to verify
|
||||
the full pipeline produces correct decoded output."""
|
||||
|
||||
@patch("onyx.chat.tool_call_args_streaming._get_tool_class")
|
||||
def test_realistic_python_code_streaming(self, mock_get_tool: MagicMock) -> None:
|
||||
"""Streams: {"code": "print('hello')\\nprint('world')"}"""
|
||||
mock_get_tool.return_value = _mock_tool_class()
|
||||
|
||||
tc_map: dict[int, dict[str, Any]] = {
|
||||
0: {"id": "tc_1", "name": "python", "arguments": ""}
|
||||
}
|
||||
fragments = [
|
||||
'{"',
|
||||
"code",
|
||||
'": "',
|
||||
"print(",
|
||||
"'hello')",
|
||||
"\\n",
|
||||
"print(",
|
||||
"'world')",
|
||||
'"}',
|
||||
]
|
||||
|
||||
full = "".join(_stream_fragments(fragments, tc_map))
|
||||
assert full == "print('hello')\nprint('world')"
|
||||
|
||||
@patch("onyx.chat.tool_call_args_streaming._get_tool_class")
|
||||
def test_streaming_with_tabs_and_newlines(self, mock_get_tool: MagicMock) -> None:
|
||||
"""Streams code with tabs and newlines."""
|
||||
mock_get_tool.return_value = _mock_tool_class()
|
||||
|
||||
tc_map: dict[int, dict[str, Any]] = {
|
||||
0: {"id": "tc_1", "name": "python", "arguments": ""}
|
||||
}
|
||||
fragments = [
|
||||
'{"code": "',
|
||||
"if True:",
|
||||
"\\n",
|
||||
"\\t",
|
||||
"pass",
|
||||
'"}',
|
||||
]
|
||||
|
||||
full = "".join(_stream_fragments(fragments, tc_map))
|
||||
assert full == "if True:\n\tpass"
|
||||
|
||||
@patch("onyx.chat.tool_call_args_streaming._get_tool_class")
|
||||
def test_split_escape_sequence(self, mock_get_tool: MagicMock) -> None:
|
||||
"""An escape sequence split across two fragments (backslash in one,
|
||||
'n' in the next) should still decode correctly."""
|
||||
mock_get_tool.return_value = _mock_tool_class()
|
||||
|
||||
tc_map: dict[int, dict[str, Any]] = {
|
||||
0: {"id": "tc_1", "name": "python", "arguments": ""}
|
||||
}
|
||||
fragments = [
|
||||
'{"code": "hello',
|
||||
"\\",
|
||||
"n",
|
||||
'world"}',
|
||||
]
|
||||
|
||||
full = "".join(_stream_fragments(fragments, tc_map))
|
||||
assert full == "hello\nworld"
|
||||
|
||||
@patch("onyx.chat.tool_call_args_streaming._get_tool_class")
|
||||
def test_multiple_newlines_and_indentation(self, mock_get_tool: MagicMock) -> None:
|
||||
"""Streams a multi-line function with multiple escape sequences."""
|
||||
mock_get_tool.return_value = _mock_tool_class()
|
||||
|
||||
tc_map: dict[int, dict[str, Any]] = {
|
||||
0: {"id": "tc_1", "name": "python", "arguments": ""}
|
||||
}
|
||||
fragments = [
|
||||
'{"code": "',
|
||||
"def foo():",
|
||||
"\\n",
|
||||
"\\t",
|
||||
"x = 1",
|
||||
"\\n",
|
||||
"\\t",
|
||||
"return x",
|
||||
'"}',
|
||||
]
|
||||
|
||||
full = "".join(_stream_fragments(fragments, tc_map))
|
||||
assert full == "def foo():\n\tx = 1\n\treturn x"
|
||||
|
||||
@patch("onyx.chat.tool_call_args_streaming._get_tool_class")
|
||||
def test_two_keys_streamed_sequentially(self, mock_get_tool: MagicMock) -> None:
|
||||
"""Streams code first, then a second key (language) — both decoded."""
|
||||
mock_get_tool.return_value = _mock_tool_class()
|
||||
|
||||
tc_map: dict[int, dict[str, Any]] = {
|
||||
0: {"id": "tc_1", "name": "python", "arguments": ""}
|
||||
}
|
||||
fragments = [
|
||||
'{"code": "',
|
||||
"x = 1",
|
||||
'", "language": "',
|
||||
"python",
|
||||
'"}',
|
||||
]
|
||||
|
||||
emitted = _stream_fragments(fragments, tc_map)
|
||||
# Should have emissions for both keys
|
||||
full = "".join(emitted)
|
||||
assert "x = 1" in full
|
||||
assert "python" in full
|
||||
|
||||
@patch("onyx.chat.tool_call_args_streaming._get_tool_class")
|
||||
def test_code_containing_dict_literal(self, mock_get_tool: MagicMock) -> None:
|
||||
"""Python code like `x = {"key": "val"}` contains JSON-like patterns.
|
||||
The escaped quotes inside the *outer* JSON value should prevent the
|
||||
inner `"key":` from being mistaken for a top-level JSON key."""
|
||||
mock_get_tool.return_value = _mock_tool_class()
|
||||
|
||||
tc_map: dict[int, dict[str, Any]] = {
|
||||
0: {"id": "tc_1", "name": "python", "arguments": ""}
|
||||
}
|
||||
# The LLM sends: {"code": "x = {\"key\": \"val\"}"}
|
||||
# The inner quotes are escaped as \" in the JSON value.
|
||||
fragments = [
|
||||
'{"code": "',
|
||||
"x = {",
|
||||
'\\"key\\"',
|
||||
": ",
|
||||
'\\"val\\"',
|
||||
"}",
|
||||
'"}',
|
||||
]
|
||||
|
||||
full = "".join(_stream_fragments(fragments, tc_map))
|
||||
assert full == 'x = {"key": "val"}'
|
||||
|
||||
@patch("onyx.chat.tool_call_args_streaming._get_tool_class")
|
||||
def test_code_with_colon_in_value(self, mock_get_tool: MagicMock) -> None:
|
||||
"""Colons inside the string value should not confuse key detection."""
|
||||
mock_get_tool.return_value = _mock_tool_class()
|
||||
|
||||
tc_map: dict[int, dict[str, Any]] = {
|
||||
0: {"id": "tc_1", "name": "python", "arguments": ""}
|
||||
}
|
||||
fragments = [
|
||||
'{"code": "',
|
||||
"url = ",
|
||||
'\\"https://example.com\\"',
|
||||
'"}',
|
||||
]
|
||||
|
||||
full = "".join(_stream_fragments(fragments, tc_map))
|
||||
assert full == 'url = "https://example.com"'
|
||||
|
||||
|
||||
class TestMaybeEmitArgumentDeltaEdgeCases:
|
||||
"""Edge cases not covered by the standard test classes."""
|
||||
|
||||
@patch("onyx.chat.tool_call_args_streaming._get_tool_class")
|
||||
def test_no_emission_when_function_is_none(self, mock_get_tool: MagicMock) -> None:
|
||||
"""Some delta chunks have function=None (e.g. role-only deltas)."""
|
||||
mock_get_tool.return_value = _mock_tool_class()
|
||||
|
||||
tc_map: dict[int, dict[str, Any]] = {
|
||||
0: {"id": "tc_1", "name": "python", "arguments": '{"code": "x'}
|
||||
}
|
||||
delta = _make_tool_call_delta(arguments=None, function_is_none=True)
|
||||
assert _collect(tc_map, delta) == []
|
||||
|
||||
@patch("onyx.chat.tool_call_args_streaming._get_tool_class")
|
||||
def test_multiple_concurrent_tool_calls(self, mock_get_tool: MagicMock) -> None:
|
||||
"""Two tool calls streaming at different indices in parallel."""
|
||||
mock_get_tool.return_value = _mock_tool_class()
|
||||
|
||||
tc_map: dict[int, dict[str, Any]] = {
|
||||
0: {"id": "tc_1", "name": "python", "arguments": ""},
|
||||
1: {"id": "tc_2", "name": "python", "arguments": ""},
|
||||
}
|
||||
|
||||
parsers: dict[int, Parser] = {}
|
||||
pl = _make_placement()
|
||||
|
||||
# Feed full JSON to index 0
|
||||
tc_map[0]["arguments"] = '{"code": "aaa"}'
|
||||
packets_0 = _collect(
|
||||
tc_map,
|
||||
_make_tool_call_delta(index=0, arguments='{"code": "aaa"}'),
|
||||
pl,
|
||||
parsers,
|
||||
)
|
||||
code_0 = ""
|
||||
for p in packets_0:
|
||||
assert isinstance(p.obj, ToolCallArgumentDelta)
|
||||
code_0 += p.obj.argument_deltas.get("code", "")
|
||||
assert code_0 == "aaa"
|
||||
|
||||
# Feed full JSON to index 1
|
||||
tc_map[1]["arguments"] = '{"code": "bbb"}'
|
||||
packets_1 = _collect(
|
||||
tc_map,
|
||||
_make_tool_call_delta(index=1, arguments='{"code": "bbb"}'),
|
||||
pl,
|
||||
parsers,
|
||||
)
|
||||
code_1 = ""
|
||||
for p in packets_1:
|
||||
assert isinstance(p.obj, ToolCallArgumentDelta)
|
||||
code_1 += p.obj.argument_deltas.get("code", "")
|
||||
assert code_1 == "bbb"
|
||||
|
||||
@patch("onyx.chat.tool_call_args_streaming._get_tool_class")
|
||||
def test_delta_with_four_arguments(self, mock_get_tool: MagicMock) -> None:
|
||||
"""A single delta contains four complete key-value pairs."""
|
||||
mock_get_tool.return_value = _mock_tool_class()
|
||||
|
||||
full = '{"a": "one", "b": "two", "c": "three", "d": "four"}'
|
||||
tc_map: dict[int, dict[str, Any]] = {
|
||||
0: {"id": "tc_1", "name": "python", "arguments": ""}
|
||||
}
|
||||
tc_map[0]["arguments"] = full
|
||||
parsers: dict[int, Parser] = {}
|
||||
packets = _collect(
|
||||
tc_map, _make_tool_call_delta(arguments=full), parsers=parsers
|
||||
)
|
||||
|
||||
# Collect all argument deltas across packets
|
||||
all_deltas: dict[str, str] = {}
|
||||
for p in packets:
|
||||
assert isinstance(p.obj, ToolCallArgumentDelta)
|
||||
for k, v in p.obj.argument_deltas.items():
|
||||
all_deltas[k] = all_deltas.get(k, "") + v
|
||||
|
||||
assert all_deltas == {
|
||||
"a": "one",
|
||||
"b": "two",
|
||||
"c": "three",
|
||||
"d": "four",
|
||||
}
|
||||
|
||||
@patch("onyx.chat.tool_call_args_streaming._get_tool_class")
|
||||
def test_delta_on_second_arg_after_first_complete(
|
||||
self, mock_get_tool: MagicMock
|
||||
) -> None:
|
||||
"""First argument is fully complete; delta only adds to the second."""
|
||||
mock_get_tool.return_value = _mock_tool_class()
|
||||
|
||||
tc_map: dict[int, dict[str, Any]] = {
|
||||
0: {"id": "tc_1", "name": "python", "arguments": ""}
|
||||
}
|
||||
|
||||
fragments = [
|
||||
'{"code": "print(1)", "lang": "py',
|
||||
'"}',
|
||||
]
|
||||
|
||||
emitted = _stream_fragments(fragments, tc_map)
|
||||
full = "".join(emitted)
|
||||
assert "print(1)" in full
|
||||
assert "py" in full
|
||||
|
||||
@patch("onyx.chat.tool_call_args_streaming._get_tool_class")
|
||||
def test_non_string_values_skipped(self, mock_get_tool: MagicMock) -> None:
|
||||
"""Non-string values (numbers, booleans, null) are skipped — they are
|
||||
available in the final tool-call kickoff packet. String arguments
|
||||
following them are still emitted."""
|
||||
mock_get_tool.return_value = _mock_tool_class()
|
||||
|
||||
tc_map: dict[int, dict[str, Any]] = {
|
||||
0: {"id": "tc_1", "name": "python", "arguments": ""}
|
||||
}
|
||||
fragments = ['{"timeout": 30, "code": "hello"}']
|
||||
|
||||
emitted = _stream_fragments(fragments, tc_map)
|
||||
full = "".join(emitted)
|
||||
assert full == "hello"
|
||||
188
backend/tests/unit/onyx/server/test_projects_file_utils.py
Normal file
188
backend/tests/unit/onyx/server/test_projects_file_utils.py
Normal file
@@ -0,0 +1,188 @@
|
||||
from io import BytesIO
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
from fastapi import UploadFile
|
||||
|
||||
from onyx.server.features.projects import projects_file_utils as utils
|
||||
|
||||
|
||||
class _Tokenizer:
|
||||
def encode(self, text: str) -> list[int]:
|
||||
return [1] * len(text)
|
||||
|
||||
|
||||
class _NonSeekableFile(BytesIO):
|
||||
def tell(self) -> int:
|
||||
raise OSError("tell not supported")
|
||||
|
||||
def seek(self, *_args: object, **_kwargs: object) -> int:
|
||||
raise OSError("seek not supported")
|
||||
|
||||
|
||||
def _make_upload(filename: str, size: int, content: bytes | None = None) -> UploadFile:
|
||||
payload = content if content is not None else (b"x" * size)
|
||||
return UploadFile(filename=filename, file=BytesIO(payload), size=size)
|
||||
|
||||
|
||||
def _make_upload_no_size(filename: str, content: bytes) -> UploadFile:
|
||||
return UploadFile(filename=filename, file=BytesIO(content), size=None)
|
||||
|
||||
|
||||
def _patch_common_dependencies(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(utils, "fetch_default_llm_model", lambda _db: None)
|
||||
monkeypatch.setattr(utils, "get_tokenizer", lambda **_kwargs: _Tokenizer())
|
||||
monkeypatch.setattr(utils, "is_file_password_protected", lambda **_kwargs: False)
|
||||
|
||||
|
||||
def test_get_upload_size_bytes_falls_back_to_stream_size() -> None:
|
||||
upload = UploadFile(filename="example.txt", file=BytesIO(b"abcdef"), size=None)
|
||||
upload.file.seek(2)
|
||||
|
||||
size = utils.get_upload_size_bytes(upload)
|
||||
|
||||
assert size == 6
|
||||
assert upload.file.tell() == 2
|
||||
|
||||
|
||||
def test_get_upload_size_bytes_logs_warning_when_stream_size_unavailable(
|
||||
caplog: pytest.LogCaptureFixture,
|
||||
) -> None:
|
||||
upload = UploadFile(filename="non_seekable.txt", file=_NonSeekableFile(), size=None)
|
||||
|
||||
caplog.set_level("WARNING")
|
||||
size = utils.get_upload_size_bytes(upload)
|
||||
|
||||
assert size is None
|
||||
assert "Could not determine upload size via stream seek" in caplog.text
|
||||
assert "non_seekable.txt" in caplog.text
|
||||
|
||||
|
||||
def test_is_upload_too_large_logs_warning_when_size_unknown(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
caplog: pytest.LogCaptureFixture,
|
||||
) -> None:
|
||||
upload = _make_upload("size_unknown.txt", size=1)
|
||||
monkeypatch.setattr(utils, "get_upload_size_bytes", lambda _upload: None)
|
||||
|
||||
caplog.set_level("WARNING")
|
||||
is_too_large = utils.is_upload_too_large(upload, max_bytes=100)
|
||||
|
||||
assert is_too_large is False
|
||||
assert "Could not determine upload size; skipping size-limit check" in caplog.text
|
||||
assert "size_unknown.txt" in caplog.text
|
||||
|
||||
|
||||
def test_categorize_uploaded_files_accepts_size_under_limit(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
_patch_common_dependencies(monkeypatch)
|
||||
monkeypatch.setattr(utils, "USER_FILE_MAX_UPLOAD_SIZE_BYTES", 100)
|
||||
monkeypatch.setattr(utils, "USER_FILE_MAX_UPLOAD_SIZE_MB", 1)
|
||||
monkeypatch.setattr(utils, "estimate_image_tokens_for_upload", lambda _upload: 10)
|
||||
|
||||
upload = _make_upload("small.png", size=99)
|
||||
result = utils.categorize_uploaded_files([upload], MagicMock())
|
||||
|
||||
assert len(result.acceptable) == 1
|
||||
assert len(result.rejected) == 0
|
||||
|
||||
|
||||
def test_categorize_uploaded_files_uses_seek_fallback_when_upload_size_missing(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
_patch_common_dependencies(monkeypatch)
|
||||
monkeypatch.setattr(utils, "USER_FILE_MAX_UPLOAD_SIZE_BYTES", 100)
|
||||
monkeypatch.setattr(utils, "USER_FILE_MAX_UPLOAD_SIZE_MB", 1)
|
||||
monkeypatch.setattr(utils, "estimate_image_tokens_for_upload", lambda _upload: 10)
|
||||
|
||||
upload = _make_upload_no_size("small.png", content=b"x" * 99)
|
||||
result = utils.categorize_uploaded_files([upload], MagicMock())
|
||||
|
||||
assert len(result.acceptable) == 1
|
||||
assert len(result.rejected) == 0
|
||||
|
||||
|
||||
def test_categorize_uploaded_files_accepts_size_at_limit(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
_patch_common_dependencies(monkeypatch)
|
||||
monkeypatch.setattr(utils, "USER_FILE_MAX_UPLOAD_SIZE_BYTES", 100)
|
||||
monkeypatch.setattr(utils, "USER_FILE_MAX_UPLOAD_SIZE_MB", 1)
|
||||
monkeypatch.setattr(utils, "estimate_image_tokens_for_upload", lambda _upload: 10)
|
||||
|
||||
upload = _make_upload("edge.png", size=100)
|
||||
result = utils.categorize_uploaded_files([upload], MagicMock())
|
||||
|
||||
assert len(result.acceptable) == 1
|
||||
assert len(result.rejected) == 0
|
||||
|
||||
|
||||
def test_categorize_uploaded_files_rejects_size_over_limit_with_reason(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
_patch_common_dependencies(monkeypatch)
|
||||
monkeypatch.setattr(utils, "USER_FILE_MAX_UPLOAD_SIZE_BYTES", 100)
|
||||
monkeypatch.setattr(utils, "USER_FILE_MAX_UPLOAD_SIZE_MB", 1)
|
||||
monkeypatch.setattr(utils, "estimate_image_tokens_for_upload", lambda _upload: 10)
|
||||
|
||||
upload = _make_upload("large.png", size=101)
|
||||
result = utils.categorize_uploaded_files([upload], MagicMock())
|
||||
|
||||
assert len(result.acceptable) == 0
|
||||
assert len(result.rejected) == 1
|
||||
assert result.rejected[0].reason == "Exceeds 1 MB file size limit"
|
||||
|
||||
|
||||
def test_categorize_uploaded_files_mixed_batch_keeps_valid_and_rejects_oversized(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
_patch_common_dependencies(monkeypatch)
|
||||
monkeypatch.setattr(utils, "USER_FILE_MAX_UPLOAD_SIZE_BYTES", 100)
|
||||
monkeypatch.setattr(utils, "USER_FILE_MAX_UPLOAD_SIZE_MB", 1)
|
||||
monkeypatch.setattr(utils, "estimate_image_tokens_for_upload", lambda _upload: 10)
|
||||
|
||||
small = _make_upload("small.png", size=50)
|
||||
large = _make_upload("large.png", size=101)
|
||||
|
||||
result = utils.categorize_uploaded_files([small, large], MagicMock())
|
||||
|
||||
assert [file.filename for file in result.acceptable] == ["small.png"]
|
||||
assert len(result.rejected) == 1
|
||||
assert result.rejected[0].filename == "large.png"
|
||||
assert result.rejected[0].reason == "Exceeds 1 MB file size limit"
|
||||
|
||||
|
||||
def test_categorize_uploaded_files_enforces_size_limit_even_when_threshold_is_skipped(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
_patch_common_dependencies(monkeypatch)
|
||||
monkeypatch.setattr(utils, "SKIP_USERFILE_THRESHOLD", True)
|
||||
monkeypatch.setattr(utils, "USER_FILE_MAX_UPLOAD_SIZE_BYTES", 100)
|
||||
monkeypatch.setattr(utils, "USER_FILE_MAX_UPLOAD_SIZE_MB", 1)
|
||||
|
||||
upload = _make_upload("oversized.pdf", size=101)
|
||||
result = utils.categorize_uploaded_files([upload], MagicMock())
|
||||
|
||||
assert len(result.acceptable) == 0
|
||||
assert len(result.rejected) == 1
|
||||
assert result.rejected[0].reason == "Exceeds 1 MB file size limit"
|
||||
|
||||
|
||||
def test_categorize_uploaded_files_checks_size_before_text_extraction(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
_patch_common_dependencies(monkeypatch)
|
||||
monkeypatch.setattr(utils, "USER_FILE_MAX_UPLOAD_SIZE_BYTES", 100)
|
||||
monkeypatch.setattr(utils, "USER_FILE_MAX_UPLOAD_SIZE_MB", 1)
|
||||
|
||||
extract_mock = MagicMock(return_value="this should not run")
|
||||
monkeypatch.setattr(utils, "extract_file_text", extract_mock)
|
||||
|
||||
oversized_doc = _make_upload("oversized.pdf", size=101)
|
||||
result = utils.categorize_uploaded_files([oversized_doc], MagicMock())
|
||||
|
||||
extract_mock.assert_not_called()
|
||||
assert len(result.acceptable) == 0
|
||||
assert len(result.rejected) == 1
|
||||
assert result.rejected[0].reason == "Exceeds 1 MB file size limit"
|
||||
32
backend/tests/unit/onyx/server/test_settings_store.py
Normal file
32
backend/tests/unit/onyx/server/test_settings_store.py
Normal file
@@ -0,0 +1,32 @@
|
||||
import pytest
|
||||
|
||||
from onyx.key_value_store.interface import KvKeyNotFoundError
|
||||
from onyx.server.settings import store as settings_store
|
||||
|
||||
|
||||
class _FakeKvStore:
|
||||
def load(self, _key: str) -> dict:
|
||||
raise KvKeyNotFoundError()
|
||||
|
||||
|
||||
class _FakeCache:
|
||||
def __init__(self) -> None:
|
||||
self._vals: dict[str, bytes] = {}
|
||||
|
||||
def get(self, key: str) -> bytes | None:
|
||||
return self._vals.get(key)
|
||||
|
||||
def set(self, key: str, value: str, ex: int | None = None) -> None: # noqa: ARG002
|
||||
self._vals[key] = value.encode("utf-8")
|
||||
|
||||
|
||||
def test_load_settings_includes_user_file_max_upload_size_mb(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
monkeypatch.setattr(settings_store, "get_kv_store", lambda: _FakeKvStore())
|
||||
monkeypatch.setattr(settings_store, "get_cache_backend", lambda: _FakeCache())
|
||||
monkeypatch.setattr(settings_store, "USER_FILE_MAX_UPLOAD_SIZE_MB", 77)
|
||||
|
||||
settings = settings_store.load_settings()
|
||||
|
||||
assert settings.user_file_max_upload_size_mb == 77
|
||||
@@ -147,15 +147,18 @@ class TestSensitiveValueString:
|
||||
)
|
||||
assert sensitive1 != sensitive2
|
||||
|
||||
def test_equality_with_non_sensitive_raises(self) -> None:
|
||||
"""Test that comparing with non-SensitiveValue raises error."""
|
||||
def test_equality_with_non_sensitive_returns_not_equal(self) -> None:
|
||||
"""Test that comparing with non-SensitiveValue is always not-equal.
|
||||
|
||||
Returns NotImplemented so Python falls back to identity comparison.
|
||||
This is required for compatibility with SQLAlchemy's attribute tracking.
|
||||
"""
|
||||
sensitive = SensitiveValue(
|
||||
encrypted_bytes=_encrypt_string("secret"),
|
||||
decrypt_fn=_decrypt_string,
|
||||
is_json=False,
|
||||
)
|
||||
with pytest.raises(SensitiveAccessError):
|
||||
_ = sensitive == "secret"
|
||||
assert not (sensitive == "secret")
|
||||
|
||||
|
||||
class TestSensitiveValueJson:
|
||||
|
||||
@@ -1261,3 +1261,5 @@ configMap:
|
||||
SKIP_USERFILE_THRESHOLD: ""
|
||||
# For multi-tenant: comma-separated list of tenant IDs to skip threshold
|
||||
SKIP_USERFILE_THRESHOLD_TENANT_IDS: ""
|
||||
# Maximum user upload file size in MB for chat/projects uploads
|
||||
USER_FILE_MAX_UPLOAD_SIZE_MB: ""
|
||||
|
||||
@@ -156,6 +156,7 @@ module.exports = {
|
||||
"**/src/app/**/*.test.tsx",
|
||||
"**/src/components/**/*.test.tsx",
|
||||
"**/src/lib/**/*.test.tsx",
|
||||
"**/src/providers/**/*.test.tsx",
|
||||
"**/src/refresh-components/**/*.test.tsx",
|
||||
"**/src/hooks/**/*.test.tsx",
|
||||
"**/src/sections/**/*.test.tsx",
|
||||
|
||||
@@ -39,7 +39,7 @@ type ButtonProps = InteractiveStatelessProps &
|
||||
/** Tooltip text shown on hover. */
|
||||
tooltip?: string;
|
||||
|
||||
/** Width preset. `"auto"` shrink-wraps, `"full"` stretches to parent width. */
|
||||
/** Width preset. `"fit"` shrink-wraps, `"full"` stretches to parent width. */
|
||||
width?: WidthVariant;
|
||||
|
||||
/** Which side the tooltip appears on. */
|
||||
|
||||
@@ -0,0 +1,83 @@
|
||||
# LineItemButton
|
||||
|
||||
**Import:** `import { LineItemButton, type LineItemButtonProps } from "@opal/components";`
|
||||
|
||||
A composite component that wraps `Interactive.Stateful > Interactive.Container > ContentAction` into a single API. Use it for selectable list rows such as model pickers, menu items, or any row that acts like a button.
|
||||
|
||||
## Architecture
|
||||
|
||||
```
|
||||
Interactive.Stateful <- selectVariant, state, interaction, onClick, href, ref
|
||||
└─ Interactive.Container <- type, width, roundingVariant
|
||||
└─ ContentAction <- withInteractive, paddingVariant="lg"
|
||||
├─ Content <- icon, title, description, sizePreset, variant, ...
|
||||
└─ rightChildren
|
||||
```
|
||||
|
||||
`paddingVariant` is hardcoded to `"lg"` and `withInteractive` is always `true`. These are not exposed as props.
|
||||
|
||||
## Props
|
||||
|
||||
### Interactive surface
|
||||
|
||||
| Prop | Type | Default | Description |
|
||||
|------|------|---------|-------------|
|
||||
| `selectVariant` | `"select-light" \| "select-heavy"` | `"select-light"` | Interactive select variant |
|
||||
| `state` | `InteractiveStatefulState` | `"empty"` | Value state (`"empty"`, `"filled"`, `"selected"`) |
|
||||
| `interaction` | `InteractiveStatefulInteraction` | `"rest"` | JS-controlled interaction state override |
|
||||
| `onClick` | `MouseEventHandler<HTMLElement>` | — | Click handler |
|
||||
| `href` | `string` | — | Renders an anchor instead of a div |
|
||||
| `target` | `string` | — | Anchor target (e.g. `"_blank"`) |
|
||||
| `group` | `string` | — | Interactive group key |
|
||||
| `ref` | `React.Ref<HTMLElement>` | — | Forwarded ref |
|
||||
|
||||
### Sizing
|
||||
|
||||
| Prop | Type | Default | Description |
|
||||
|------|------|---------|-------------|
|
||||
| `roundingVariant` | `InteractiveContainerRoundingVariant` | `"default"` | Corner rounding preset (height is content-driven) |
|
||||
| `width` | `WidthVariant` | `"full"` | Container width |
|
||||
| `type` | `"submit" \| "button" \| "reset"` | `"button"` | HTML button type |
|
||||
| `tooltip` | `string` | — | Tooltip text shown on hover |
|
||||
| `tooltipSide` | `TooltipSide` | `"top"` | Tooltip side |
|
||||
|
||||
### Content (pass-through to ContentAction)
|
||||
|
||||
| Prop | Type | Default | Description |
|
||||
|------|------|---------|-------------|
|
||||
| `title` | `string` | **(required)** | Row label |
|
||||
| `icon` | `IconFunctionComponent` | — | Left icon |
|
||||
| `description` | `string` | — | Description below the title |
|
||||
| `sizePreset` | `SizePreset` | `"headline"` | Content size preset |
|
||||
| `variant` | `ContentVariant` | `"heading"` | Content layout variant |
|
||||
| `rightChildren` | `ReactNode` | — | Content after the label (e.g. action button) |
|
||||
|
||||
All other `ContentAction` / `Content` props (`editable`, `onTitleChange`, `optional`, `auxIcon`, `tag`, etc.) are also passed through. Note: `withInteractive` is always `true` inside `LineItemButton` and cannot be overridden.
|
||||
|
||||
## Usage
|
||||
|
||||
```tsx
|
||||
import { LineItemButton } from "@opal/components";
|
||||
|
||||
// Simple selectable row
|
||||
<LineItemButton
|
||||
selectVariant="select-heavy"
|
||||
state={isSelected ? "selected" : "empty"}
|
||||
roundingVariant="compact"
|
||||
onClick={handleClick}
|
||||
title="gpt-4o"
|
||||
sizePreset="main-ui"
|
||||
variant="section"
|
||||
/>
|
||||
|
||||
// With right-side action
|
||||
<LineItemButton
|
||||
selectVariant="select-heavy"
|
||||
state={isSelected ? "selected" : "empty"}
|
||||
onClick={handleClick}
|
||||
title="claude-opus-4"
|
||||
sizePreset="main-ui"
|
||||
variant="section"
|
||||
rightChildren={<Tag title="Default" color="blue" />}
|
||||
/>
|
||||
```
|
||||
@@ -0,0 +1,137 @@
|
||||
import "@opal/components/tooltip.css";
|
||||
import {
|
||||
Interactive,
|
||||
type InteractiveStatefulState,
|
||||
type InteractiveStatefulInteraction,
|
||||
type InteractiveStatefulProps,
|
||||
InteractiveContainerRoundingVariant,
|
||||
} from "@opal/core";
|
||||
import { type WidthVariant } from "@opal/shared";
|
||||
import type { TooltipSide } from "@opal/components";
|
||||
import type { DistributiveOmit } from "@opal/types";
|
||||
import type { ContentActionProps } from "@opal/layouts/content-action/components";
|
||||
import { ContentAction } from "@opal/layouts";
|
||||
import * as TooltipPrimitive from "@radix-ui/react-tooltip";
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Types
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
type ContentPassthroughProps = DistributiveOmit<
|
||||
ContentActionProps,
|
||||
"paddingVariant" | "widthVariant" | "ref" | "withInteractive"
|
||||
>;
|
||||
|
||||
type LineItemButtonOwnProps = {
|
||||
/** Interactive select variant. @default "select-light" */
|
||||
selectVariant?: "select-light" | "select-heavy";
|
||||
|
||||
/** Value state. @default "empty" */
|
||||
state?: InteractiveStatefulState;
|
||||
|
||||
/** JS-controllable interaction state override. @default "rest" */
|
||||
interaction?: InteractiveStatefulInteraction;
|
||||
|
||||
/** Click handler. */
|
||||
onClick?: InteractiveStatefulProps["onClick"];
|
||||
|
||||
/** When provided, renders an anchor instead of a div. */
|
||||
href?: string;
|
||||
|
||||
/** Anchor target (e.g. "_blank"). */
|
||||
target?: string;
|
||||
|
||||
/** Interactive group key. */
|
||||
group?: string;
|
||||
|
||||
/** Forwarded ref. */
|
||||
ref?: React.Ref<HTMLElement>;
|
||||
|
||||
/** Corner rounding preset (height is always content-driven). @default "default" */
|
||||
roundingVariant?: InteractiveContainerRoundingVariant;
|
||||
|
||||
/** Container width. @default "full" */
|
||||
width?: WidthVariant;
|
||||
|
||||
/** HTML button type. @default "button" */
|
||||
type?: "submit" | "button" | "reset";
|
||||
|
||||
/** Tooltip text shown on hover. */
|
||||
tooltip?: string;
|
||||
|
||||
/** Which side the tooltip appears on. @default "top" */
|
||||
tooltipSide?: TooltipSide;
|
||||
};
|
||||
|
||||
type LineItemButtonProps = ContentPassthroughProps & LineItemButtonOwnProps;
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// LineItemButton
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
function LineItemButton({
|
||||
// Interactive surface
|
||||
selectVariant = "select-light",
|
||||
state,
|
||||
interaction,
|
||||
onClick,
|
||||
href,
|
||||
target,
|
||||
group,
|
||||
ref,
|
||||
|
||||
// Sizing
|
||||
roundingVariant = "default",
|
||||
width = "full",
|
||||
type = "button",
|
||||
tooltip,
|
||||
tooltipSide = "top",
|
||||
|
||||
// ContentAction pass-through
|
||||
...contentActionProps
|
||||
}: LineItemButtonProps) {
|
||||
const item = (
|
||||
<Interactive.Stateful
|
||||
variant={selectVariant}
|
||||
state={state}
|
||||
interaction={interaction}
|
||||
onClick={onClick}
|
||||
href={href}
|
||||
target={target}
|
||||
group={group}
|
||||
ref={ref}
|
||||
>
|
||||
<Interactive.Container
|
||||
type={type}
|
||||
widthVariant={width}
|
||||
heightVariant="lg"
|
||||
roundingVariant={roundingVariant}
|
||||
>
|
||||
<ContentAction
|
||||
{...(contentActionProps as ContentActionProps)}
|
||||
withInteractive
|
||||
paddingVariant="fit"
|
||||
/>
|
||||
</Interactive.Container>
|
||||
</Interactive.Stateful>
|
||||
);
|
||||
|
||||
if (!tooltip) return item;
|
||||
|
||||
return (
|
||||
<TooltipPrimitive.Root>
|
||||
<TooltipPrimitive.Trigger asChild>{item}</TooltipPrimitive.Trigger>
|
||||
<TooltipPrimitive.Portal>
|
||||
<TooltipPrimitive.Content
|
||||
className="opal-tooltip"
|
||||
side={tooltipSide}
|
||||
sideOffset={4}
|
||||
>
|
||||
{tooltip}
|
||||
</TooltipPrimitive.Content>
|
||||
</TooltipPrimitive.Portal>
|
||||
</TooltipPrimitive.Root>
|
||||
);
|
||||
}
|
||||
|
||||
export { LineItemButton, type LineItemButtonProps };
|
||||
@@ -56,7 +56,7 @@ type SelectButtonProps = InteractiveStatefulProps &
|
||||
/** Tooltip text shown on hover. */
|
||||
tooltip?: string;
|
||||
|
||||
/** Width preset. `"auto"` shrink-wraps, `"full"` stretches to parent width. */
|
||||
/** Width preset. `"fit"` shrink-wraps, `"full"` stretches to parent width. */
|
||||
width?: WidthVariant;
|
||||
|
||||
/** Which side the tooltip appears on. */
|
||||
|
||||
@@ -19,6 +19,12 @@ export {
|
||||
type OpenButtonProps,
|
||||
} from "@opal/components/buttons/open-button/components";
|
||||
|
||||
/* LineItemButton */
|
||||
export {
|
||||
LineItemButton,
|
||||
type LineItemButtonProps,
|
||||
} from "@opal/components/buttons/line-item-button/components";
|
||||
|
||||
/* Tag */
|
||||
export {
|
||||
Tag,
|
||||
|
||||
@@ -2,6 +2,7 @@ import "@opal/core/animations/styles.css";
|
||||
import React, { createContext, useContext, useState, useCallback } from "react";
|
||||
import { cn } from "@opal/utils";
|
||||
import type { WithoutStyles } from "@opal/types";
|
||||
import { widthVariants, type WidthVariant } from "@opal/shared";
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Context-per-group registry
|
||||
@@ -38,6 +39,10 @@ interface HoverableRootProps
|
||||
extends WithoutStyles<React.HTMLAttributes<HTMLDivElement>> {
|
||||
children: React.ReactNode;
|
||||
group: string;
|
||||
/** Width preset. @default "auto" */
|
||||
widthVariant?: WidthVariant;
|
||||
/** Ref forwarded to the root `<div>`. */
|
||||
ref?: React.Ref<HTMLDivElement>;
|
||||
}
|
||||
|
||||
type HoverableItemVariant = "opacity-on-hover";
|
||||
@@ -47,6 +52,8 @@ interface HoverableItemProps
|
||||
children: React.ReactNode;
|
||||
group?: string;
|
||||
variant?: HoverableItemVariant;
|
||||
/** Ref forwarded to the item `<div>`. */
|
||||
ref?: React.Ref<HTMLDivElement>;
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
@@ -77,6 +84,8 @@ interface HoverableItemProps
|
||||
function HoverableRoot({
|
||||
group,
|
||||
children,
|
||||
widthVariant = "auto",
|
||||
ref,
|
||||
onMouseEnter: consumerMouseEnter,
|
||||
onMouseLeave: consumerMouseLeave,
|
||||
...props
|
||||
@@ -103,7 +112,13 @@ function HoverableRoot({
|
||||
|
||||
return (
|
||||
<GroupContext.Provider value={hovered}>
|
||||
<div {...props} onMouseEnter={onMouseEnter} onMouseLeave={onMouseLeave}>
|
||||
<div
|
||||
{...props}
|
||||
ref={ref}
|
||||
className={cn(widthVariants[widthVariant])}
|
||||
onMouseEnter={onMouseEnter}
|
||||
onMouseLeave={onMouseLeave}
|
||||
>
|
||||
{children}
|
||||
</div>
|
||||
</GroupContext.Provider>
|
||||
@@ -147,6 +162,7 @@ function HoverableItem({
|
||||
group,
|
||||
variant = "opacity-on-hover",
|
||||
children,
|
||||
ref,
|
||||
...props
|
||||
}: HoverableItemProps) {
|
||||
const contextValue = useContext(
|
||||
@@ -165,6 +181,7 @@ function HoverableItem({
|
||||
return (
|
||||
<div
|
||||
{...props}
|
||||
ref={ref}
|
||||
className={cn("hoverable-item")}
|
||||
data-hoverable-variant={variant}
|
||||
data-hoverable-active={
|
||||
|
||||
@@ -10,7 +10,7 @@ Structural container shared by both `Interactive.Stateless` and `Interactive.Sta
|
||||
|------|------|---------|-------------|
|
||||
| `heightVariant` | `SizeVariant` | `"lg"` | Height preset (`2xs`–`lg`, `fit`) |
|
||||
| `roundingVariant` | `"default" \| "compact" \| "mini"` | `"default"` | Border-radius preset |
|
||||
| `widthVariant` | `WidthVariant` | — | Width preset (`auto`, `full`) |
|
||||
| `widthVariant` | `WidthVariant` | — | Width preset (`"auto"`, `"fit"`, `"full"`) |
|
||||
| `border` | `boolean` | `false` | Renders a 1px border |
|
||||
| `type` | `"submit" \| "button" \| "reset"` | — | When set, renders a `<button>` element |
|
||||
|
||||
|
||||
@@ -78,7 +78,7 @@ interface InteractiveContainerProps
|
||||
/**
|
||||
* Width preset controlling the container's horizontal size.
|
||||
*
|
||||
* @default "auto"
|
||||
* @default "fit"
|
||||
*/
|
||||
widthVariant?: WidthVariant;
|
||||
}
|
||||
@@ -101,7 +101,7 @@ function InteractiveContainer({
|
||||
border,
|
||||
roundingVariant = "default",
|
||||
heightVariant = "lg",
|
||||
widthVariant = "auto",
|
||||
widthVariant = "fit",
|
||||
...props
|
||||
}: InteractiveContainerProps) {
|
||||
const { allowClick } = useDisabled();
|
||||
|
||||
@@ -59,8 +59,8 @@
|
||||
--------------------------------------------------------------------------- */
|
||||
.interactive[data-interactive-variant="select-heavy"][data-interactive-state="filled"] {
|
||||
@apply bg-transparent;
|
||||
--interactive-foreground: var(--text-04);
|
||||
--interactive-foreground-icon: var(--text-04);
|
||||
--interactive-foreground: var(--action-link-05);
|
||||
--interactive-foreground-icon: var(--action-link-05);
|
||||
}
|
||||
.interactive[data-interactive-variant="select-heavy"][data-interactive-state="filled"]:hover:not(
|
||||
[data-disabled]
|
||||
@@ -76,9 +76,7 @@
|
||||
.interactive[data-interactive-variant="select-heavy"][data-interactive-state="filled"][data-interaction="active"]:not(
|
||||
[data-disabled]
|
||||
) {
|
||||
@apply bg-background-neutral-00;
|
||||
--interactive-foreground: var(--text-05);
|
||||
--interactive-foreground-icon: var(--text-05);
|
||||
@apply bg-background-tint-00;
|
||||
}
|
||||
.interactive[data-interactive-variant="select-heavy"][data-interactive-state="filled"][data-disabled] {
|
||||
@apply bg-transparent;
|
||||
@@ -158,8 +156,8 @@
|
||||
--------------------------------------------------------------------------- */
|
||||
.interactive[data-interactive-variant="select-light"][data-interactive-state="filled"] {
|
||||
@apply bg-transparent;
|
||||
--interactive-foreground: var(--text-04);
|
||||
--interactive-foreground-icon: var(--text-04);
|
||||
--interactive-foreground: var(--action-link-05);
|
||||
--interactive-foreground-icon: var(--action-link-05);
|
||||
}
|
||||
.interactive[data-interactive-variant="select-light"][data-interactive-state="filled"]:hover:not(
|
||||
[data-disabled]
|
||||
@@ -175,9 +173,7 @@
|
||||
.interactive[data-interactive-variant="select-light"][data-interactive-state="filled"][data-interaction="active"]:not(
|
||||
[data-disabled]
|
||||
) {
|
||||
@apply bg-background-neutral-00;
|
||||
--interactive-foreground: var(--text-05);
|
||||
--interactive-foreground-icon: var(--text-05);
|
||||
@apply bg-background-tint-00;
|
||||
}
|
||||
.interactive[data-interactive-variant="select-light"][data-interactive-state="filled"][data-disabled] {
|
||||
@apply bg-transparent;
|
||||
|
||||
@@ -40,6 +40,9 @@ interface BodyLayoutProps {
|
||||
|
||||
/** Title prominence. Default: `"default"`. */
|
||||
prominence?: BodyProminence;
|
||||
|
||||
/** Ref forwarded to the root `<div>`. */
|
||||
ref?: React.Ref<HTMLDivElement>;
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
@@ -80,6 +83,7 @@ function BodyLayout({
|
||||
sizePreset = "main-ui",
|
||||
orientation = "inline",
|
||||
prominence = "default",
|
||||
ref,
|
||||
}: BodyLayoutProps) {
|
||||
const config = BODY_PRESETS[sizePreset];
|
||||
const titleColorClass =
|
||||
@@ -87,6 +91,7 @@ function BodyLayout({
|
||||
|
||||
return (
|
||||
<div
|
||||
ref={ref}
|
||||
className="opal-content-body"
|
||||
data-orientation={orientation}
|
||||
style={{ gap: config.gap }}
|
||||
|
||||
@@ -48,6 +48,12 @@ interface ContentLgProps {
|
||||
|
||||
/** Size preset. Default: `"headline"`. */
|
||||
sizePreset?: ContentLgSizePreset;
|
||||
|
||||
/** When `true`, the title color hooks into `Interactive`'s `--interactive-foreground` variable. */
|
||||
withInteractive?: boolean;
|
||||
|
||||
/** Ref forwarded to the root `<div>`. */
|
||||
ref?: React.Ref<HTMLDivElement>;
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
@@ -86,6 +92,8 @@ function ContentLg({
|
||||
description,
|
||||
editable,
|
||||
onTitleChange,
|
||||
withInteractive,
|
||||
ref,
|
||||
}: ContentLgProps) {
|
||||
const [editing, setEditing] = useState(false);
|
||||
const [editValue, setEditValue] = useState(title);
|
||||
@@ -104,7 +112,12 @@ function ContentLg({
|
||||
}
|
||||
|
||||
return (
|
||||
<div className="opal-content-lg" style={{ gap: config.gap }}>
|
||||
<div
|
||||
ref={ref}
|
||||
className="opal-content-lg"
|
||||
data-interactive={withInteractive || undefined}
|
||||
style={{ gap: config.gap }}
|
||||
>
|
||||
{Icon && (
|
||||
<div
|
||||
className={cn(
|
||||
|
||||
@@ -61,6 +61,12 @@ interface ContentMdProps {
|
||||
|
||||
/** Size preset. Default: `"main-ui"`. */
|
||||
sizePreset?: ContentMdSizePreset;
|
||||
|
||||
/** When `true`, the title color hooks into `Interactive`'s `--interactive-foreground` variable. */
|
||||
withInteractive?: boolean;
|
||||
|
||||
/** Ref forwarded to the root `<div>`. */
|
||||
ref?: React.Ref<HTMLDivElement>;
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
@@ -130,6 +136,8 @@ function ContentMd({
|
||||
auxIcon,
|
||||
tag,
|
||||
sizePreset = "main-ui",
|
||||
withInteractive,
|
||||
ref,
|
||||
}: ContentMdProps) {
|
||||
const [editing, setEditing] = useState(false);
|
||||
const [editValue, setEditValue] = useState(title);
|
||||
@@ -149,7 +157,12 @@ function ContentMd({
|
||||
}
|
||||
|
||||
return (
|
||||
<div className="opal-content-md" style={{ gap: config.gap }}>
|
||||
<div
|
||||
ref={ref}
|
||||
className="opal-content-md"
|
||||
data-interactive={withInteractive || undefined}
|
||||
style={{ gap: config.gap }}
|
||||
>
|
||||
{Icon && (
|
||||
<div
|
||||
className={cn(
|
||||
|
||||
@@ -40,6 +40,12 @@ interface ContentSmProps {
|
||||
|
||||
/** Title prominence. Default: `"default"`. */
|
||||
prominence?: ContentSmProminence;
|
||||
|
||||
/** When `true`, the title color hooks into `Interactive`'s `--interactive-foreground` variable. */
|
||||
withInteractive?: boolean;
|
||||
|
||||
/** Ref forwarded to the root `<div>`. */
|
||||
ref?: React.Ref<HTMLDivElement>;
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
@@ -80,14 +86,18 @@ function ContentSm({
|
||||
sizePreset = "main-ui",
|
||||
orientation = "inline",
|
||||
prominence = "default",
|
||||
withInteractive,
|
||||
ref,
|
||||
}: ContentSmProps) {
|
||||
const config = CONTENT_SM_PRESETS[sizePreset];
|
||||
|
||||
return (
|
||||
<div
|
||||
ref={ref}
|
||||
className="opal-content-sm"
|
||||
data-orientation={orientation}
|
||||
data-prominence={prominence}
|
||||
data-interactive={withInteractive || undefined}
|
||||
style={{ gap: config.gap }}
|
||||
>
|
||||
{Icon && (
|
||||
|
||||
@@ -60,6 +60,12 @@ interface ContentXlProps {
|
||||
|
||||
/** Optional tertiary icon rendered in the icon row. */
|
||||
moreIcon2?: IconFunctionComponent;
|
||||
|
||||
/** When `true`, the title color hooks into `Interactive`'s `--interactive-foreground` variable. */
|
||||
withInteractive?: boolean;
|
||||
|
||||
/** Ref forwarded to the root `<div>`. */
|
||||
ref?: React.Ref<HTMLDivElement>;
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
@@ -106,6 +112,8 @@ function ContentXl({
|
||||
onTitleChange,
|
||||
moreIcon1: MoreIcon1,
|
||||
moreIcon2: MoreIcon2,
|
||||
withInteractive,
|
||||
ref,
|
||||
}: ContentXlProps) {
|
||||
const [editing, setEditing] = useState(false);
|
||||
const [editValue, setEditValue] = useState(title);
|
||||
@@ -124,7 +132,11 @@ function ContentXl({
|
||||
}
|
||||
|
||||
return (
|
||||
<div className="opal-content-xl">
|
||||
<div
|
||||
ref={ref}
|
||||
className="opal-content-xl"
|
||||
data-interactive={withInteractive || undefined}
|
||||
>
|
||||
{(Icon || MoreIcon1 || MoreIcon2) && (
|
||||
<div className="opal-content-xl-icon-row">
|
||||
{Icon && (
|
||||
|
||||
@@ -52,6 +52,9 @@ interface HeadingLayoutProps {
|
||||
|
||||
/** Variant controls icon placement. `"heading"` = top, `"section"` = inline. Default: `"heading"`. */
|
||||
variant?: HeadingVariant;
|
||||
|
||||
/** Ref forwarded to the root `<div>`. */
|
||||
ref?: React.Ref<HTMLDivElement>;
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
@@ -91,6 +94,7 @@ function HeadingLayout({
|
||||
description,
|
||||
editable,
|
||||
onTitleChange,
|
||||
ref,
|
||||
}: HeadingLayoutProps) {
|
||||
const [editing, setEditing] = useState(false);
|
||||
const [editValue, setEditValue] = useState(title);
|
||||
@@ -112,6 +116,7 @@ function HeadingLayout({
|
||||
|
||||
return (
|
||||
<div
|
||||
ref={ref}
|
||||
className="opal-content-heading"
|
||||
data-icon-placement={iconPlacement}
|
||||
style={{ gap: iconPlacement === "left" ? config.gap : undefined }}
|
||||
|
||||
@@ -61,6 +61,9 @@ interface LabelLayoutProps {
|
||||
|
||||
/** Size preset. Default: `"main-ui"`. */
|
||||
sizePreset?: LabelSizePreset;
|
||||
|
||||
/** Ref forwarded to the root `<div>`. */
|
||||
ref?: React.Ref<HTMLDivElement>;
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
@@ -130,6 +133,7 @@ function LabelLayout({
|
||||
auxIcon,
|
||||
tag,
|
||||
sizePreset = "main-ui",
|
||||
ref,
|
||||
}: LabelLayoutProps) {
|
||||
const [editing, setEditing] = useState(false);
|
||||
const [editValue, setEditValue] = useState(title);
|
||||
@@ -149,7 +153,7 @@ function LabelLayout({
|
||||
}
|
||||
|
||||
return (
|
||||
<div className="opal-content-label" style={{ gap: config.gap }}>
|
||||
<div ref={ref} className="opal-content-label" style={{ gap: config.gap }}>
|
||||
{Icon && (
|
||||
<div
|
||||
className={cn(
|
||||
|
||||
@@ -54,11 +54,18 @@ interface ContentBaseProps {
|
||||
* Uses the shared `WidthVariant` scale from `@opal/shared`.
|
||||
*
|
||||
* - `"auto"` — Shrink-wraps to content width
|
||||
* - `"fit"` — Shrink-wraps to content width
|
||||
* - `"full"` — Stretches to fill the parent's width
|
||||
*
|
||||
* @default "auto"
|
||||
* @default "fit"
|
||||
*/
|
||||
widthVariant?: WidthVariant;
|
||||
|
||||
/** When `true`, the title color hooks into `Interactive.Stateful`/`Interactive.Stateless`'s `--interactive-foreground` variable. */
|
||||
withInteractive?: boolean;
|
||||
|
||||
/** Ref forwarded to the root `<div>` of the resolved layout. */
|
||||
ref?: React.Ref<HTMLDivElement>;
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
@@ -122,11 +129,11 @@ function Content(props: ContentProps) {
|
||||
sizePreset = "headline",
|
||||
variant = "heading",
|
||||
widthVariant = "auto",
|
||||
withInteractive,
|
||||
ref,
|
||||
...rest
|
||||
} = props;
|
||||
|
||||
const widthClass = widthVariants[widthVariant];
|
||||
|
||||
let layout: React.ReactNode = null;
|
||||
|
||||
// ContentXl / ContentLg: headline/section presets
|
||||
@@ -135,6 +142,8 @@ function Content(props: ContentProps) {
|
||||
layout = (
|
||||
<ContentXl
|
||||
sizePreset={sizePreset}
|
||||
withInteractive={withInteractive}
|
||||
ref={ref}
|
||||
{...(rest as Omit<ContentXlProps, "sizePreset">)}
|
||||
/>
|
||||
);
|
||||
@@ -142,6 +151,8 @@ function Content(props: ContentProps) {
|
||||
layout = (
|
||||
<ContentLg
|
||||
sizePreset={sizePreset}
|
||||
withInteractive={withInteractive}
|
||||
ref={ref}
|
||||
{...(rest as Omit<ContentLgProps, "sizePreset">)}
|
||||
/>
|
||||
);
|
||||
@@ -154,6 +165,8 @@ function Content(props: ContentProps) {
|
||||
layout = (
|
||||
<ContentMd
|
||||
sizePreset={sizePreset}
|
||||
withInteractive={withInteractive}
|
||||
ref={ref}
|
||||
{...(rest as Omit<ContentMdProps, "sizePreset">)}
|
||||
/>
|
||||
);
|
||||
@@ -164,6 +177,8 @@ function Content(props: ContentProps) {
|
||||
layout = (
|
||||
<ContentSm
|
||||
sizePreset={sizePreset}
|
||||
withInteractive={withInteractive}
|
||||
ref={ref}
|
||||
{...(rest as Omit<
|
||||
React.ComponentProps<typeof ContentSm>,
|
||||
"sizePreset"
|
||||
@@ -178,11 +193,7 @@ function Content(props: ContentProps) {
|
||||
`Content: no layout matched for sizePreset="${sizePreset}" variant="${variant}"`
|
||||
);
|
||||
|
||||
// "auto" → return layout directly (a block div with w-auto still
|
||||
// stretches to its parent, defeating shrink-to-content).
|
||||
if (widthVariant === "auto") return layout;
|
||||
|
||||
return <div className={widthClass}>{layout}</div>;
|
||||
return <div className={widthVariants[widthVariant]}>{layout}</div>;
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
@@ -22,6 +22,7 @@
|
||||
|
||||
.opal-content-xl-icon-row {
|
||||
@apply flex flex-row items-center;
|
||||
gap: 0.25rem;
|
||||
}
|
||||
|
||||
/* ---------------------------------------------------------------------------
|
||||
@@ -63,7 +64,7 @@
|
||||
}
|
||||
|
||||
.opal-content-xl-title {
|
||||
@apply text-left overflow-hidden;
|
||||
@apply text-left overflow-hidden text-text-04;
|
||||
display: -webkit-box;
|
||||
-webkit-box-orient: vertical;
|
||||
-webkit-line-clamp: 1;
|
||||
@@ -162,7 +163,7 @@
|
||||
}
|
||||
|
||||
.opal-content-lg-title {
|
||||
@apply text-left overflow-hidden;
|
||||
@apply text-left overflow-hidden text-text-04;
|
||||
display: -webkit-box;
|
||||
-webkit-box-orient: vertical;
|
||||
-webkit-line-clamp: 1;
|
||||
@@ -255,7 +256,7 @@
|
||||
}
|
||||
|
||||
.opal-content-md-title {
|
||||
@apply text-left overflow-hidden;
|
||||
@apply text-left overflow-hidden text-text-04;
|
||||
display: -webkit-box;
|
||||
-webkit-box-orient: vertical;
|
||||
-webkit-line-clamp: 1;
|
||||
@@ -361,8 +362,12 @@
|
||||
@apply text-text-03;
|
||||
}
|
||||
|
||||
.opal-content-sm[data-prominence="muted"] .opal-content-sm-icon {
|
||||
@apply text-text-02;
|
||||
}
|
||||
|
||||
.opal-content-sm[data-prominence="muted-2x"] .opal-content-sm-icon {
|
||||
@apply text-text-02 stroke-text-02;
|
||||
@apply text-text-02;
|
||||
}
|
||||
|
||||
/* ---------------------------------------------------------------------------
|
||||
@@ -385,3 +390,44 @@
|
||||
.opal-content-sm[data-prominence="muted-2x"] .opal-content-sm-title {
|
||||
@apply text-text-02;
|
||||
}
|
||||
|
||||
/* ===========================================================================
|
||||
Interactive-foreground opt-in
|
||||
|
||||
When a Content variant is nested inside an Interactive and
|
||||
`withInteractive` is set, the title and icon delegate their color to the
|
||||
`--interactive-foreground` / `--interactive-foreground-icon` CSS variables
|
||||
controlled by the ancestor Interactive variant.
|
||||
=========================================================================== */
|
||||
|
||||
.opal-content-xl[data-interactive] .opal-content-xl-title {
|
||||
color: var(--interactive-foreground);
|
||||
}
|
||||
|
||||
.opal-content-xl[data-interactive] .opal-content-xl-icon {
|
||||
color: var(--interactive-foreground-icon);
|
||||
}
|
||||
|
||||
.opal-content-lg[data-interactive] .opal-content-lg-title {
|
||||
color: var(--interactive-foreground);
|
||||
}
|
||||
|
||||
.opal-content-lg[data-interactive] .opal-content-lg-icon {
|
||||
color: var(--interactive-foreground-icon);
|
||||
}
|
||||
|
||||
.opal-content-md[data-interactive] .opal-content-md-title {
|
||||
color: var(--interactive-foreground);
|
||||
}
|
||||
|
||||
.opal-content-md[data-interactive] .opal-content-md-icon {
|
||||
color: var(--interactive-foreground-icon);
|
||||
}
|
||||
|
||||
.opal-content-sm[data-interactive] .opal-content-sm-title {
|
||||
color: var(--interactive-foreground);
|
||||
}
|
||||
|
||||
.opal-content-sm[data-interactive] .opal-content-sm-icon {
|
||||
color: var(--interactive-foreground-icon);
|
||||
}
|
||||
|
||||
@@ -67,10 +67,12 @@ type SizeVariant = keyof typeof sizeVariants;
|
||||
* | Key | Tailwind class |
|
||||
* |--------|----------------|
|
||||
* | `auto` | `w-auto` |
|
||||
* | `fit` | `w-fit` |
|
||||
* | `full` | `w-full` |
|
||||
*/
|
||||
const widthVariants = {
|
||||
auto: "w-auto",
|
||||
fit: "w-fit",
|
||||
full: "w-full",
|
||||
} as const;
|
||||
|
||||
|
||||
@@ -30,6 +30,11 @@ export interface IconProps extends SVGProps<SVGSVGElement> {
|
||||
/** Strips `className` and `style` from a props type to enforce design-system styling. */
|
||||
export type WithoutStyles<T> = Omit<T, "className" | "style">;
|
||||
|
||||
/** Like `Omit` but distributes over union types, preserving discriminated unions. */
|
||||
export type DistributiveOmit<T, K extends keyof any> = T extends any
|
||||
? Omit<T, K>
|
||||
: never;
|
||||
|
||||
/**
|
||||
* A React function component that accepts {@link IconProps}.
|
||||
*
|
||||
|
||||
@@ -8,7 +8,8 @@ const cspHeader = `
|
||||
base-uri 'self';
|
||||
form-action 'self';
|
||||
${
|
||||
process.env.NEXT_PUBLIC_CLOUD_ENABLED === "true"
|
||||
process.env.NEXT_PUBLIC_CLOUD_ENABLED === "true" &&
|
||||
process.env.NODE_ENV !== "development"
|
||||
? "upgrade-insecure-requests;"
|
||||
: ""
|
||||
}
|
||||
|
||||
@@ -7,7 +7,7 @@ import { SlackTokensForm } from "./SlackTokensForm";
|
||||
import * as SettingsLayouts from "@/layouts/settings-layouts";
|
||||
import { SvgSlack } from "@opal/icons";
|
||||
|
||||
export const NewSlackBotForm = () => {
|
||||
export function NewSlackBotForm() {
|
||||
const [formValues] = useState({
|
||||
name: "",
|
||||
enabled: true,
|
||||
@@ -19,7 +19,12 @@ export const NewSlackBotForm = () => {
|
||||
|
||||
return (
|
||||
<SettingsLayouts.Root>
|
||||
<SettingsLayouts.Header icon={SvgSlack} title="New Slack Bot" separator />
|
||||
<SettingsLayouts.Header
|
||||
icon={SvgSlack}
|
||||
title="New Slack Bot"
|
||||
separator
|
||||
backButton
|
||||
/>
|
||||
<SettingsLayouts.Body>
|
||||
<CardSection>
|
||||
<div className="p-4">
|
||||
@@ -33,4 +38,4 @@ export const NewSlackBotForm = () => {
|
||||
</SettingsLayouts.Body>
|
||||
</SettingsLayouts.Root>
|
||||
);
|
||||
};
|
||||
}
|
||||
|
||||
@@ -1,12 +1,12 @@
|
||||
"use client";
|
||||
|
||||
import { toast } from "@/hooks/useToast";
|
||||
import { SlackBot, ValidSources } from "@/lib/types";
|
||||
import { SlackBot } from "@/lib/types";
|
||||
import { useRouter } from "next/navigation";
|
||||
import { useState, useEffect, useRef } from "react";
|
||||
import { updateSlackBotField } from "@/lib/updateSlackBotField";
|
||||
import { SlackTokensForm } from "./SlackTokensForm";
|
||||
import { SourceIcon } from "@/components/SourceIcon";
|
||||
|
||||
import { EditableStringFieldDisplay } from "@/components/EditableStringFieldDisplay";
|
||||
import { deleteSlackBot } from "./new/lib";
|
||||
import GenericConfirmModal from "@/components/modals/GenericConfirmModal";
|
||||
@@ -90,10 +90,7 @@ export const ExistingSlackBotForm = ({
|
||||
<div>
|
||||
<div className="flex items-center justify-between h-14">
|
||||
<div className="flex items-center gap-2">
|
||||
<div className="my-auto">
|
||||
<SourceIcon iconSize={32} sourceType={ValidSources.Slack} />
|
||||
</div>
|
||||
<div className="ml-1">
|
||||
<div>
|
||||
<EditableStringFieldDisplay
|
||||
value={formValues.name}
|
||||
isEditable={true}
|
||||
|
||||
@@ -1,100 +1,122 @@
|
||||
import { SlackChannelConfigCreationForm } from "../SlackChannelConfigCreationForm";
|
||||
import { fetchSS } from "@/lib/utilsSS";
|
||||
"use client";
|
||||
|
||||
import { use } from "react";
|
||||
import { SlackChannelConfigCreationForm } from "@/app/admin/bots/[bot-id]/channels/SlackChannelConfigCreationForm";
|
||||
import { ErrorCallout } from "@/components/ErrorCallout";
|
||||
import { DocumentSetSummary, SlackChannelConfig } from "@/lib/types";
|
||||
import { InstantSSRAutoRefresh } from "@/components/SSRAutoRefresh";
|
||||
import SimpleLoader from "@/refresh-components/loaders/SimpleLoader";
|
||||
import * as SettingsLayouts from "@/layouts/settings-layouts";
|
||||
import { SvgSlack } from "@opal/icons";
|
||||
import { FetchAgentsResponse, fetchAgentsSS } from "@/lib/agentsSS";
|
||||
import { getStandardAnswerCategoriesIfEE } from "@/components/standardAnswers/getStandardAnswerCategoriesIfEE";
|
||||
import { useSlackChannelConfigs } from "@/app/admin/bots/[bot-id]/hooks";
|
||||
import { useDocumentSets } from "@/app/admin/documents/sets/hooks";
|
||||
import { useAgents } from "@/hooks/useAgents";
|
||||
import { useStandardAnswerCategories } from "@/app/ee/admin/standard-answer/hooks";
|
||||
import { usePaidEnterpriseFeaturesEnabled } from "@/components/settings/usePaidEnterpriseFeaturesEnabled";
|
||||
import type { StandardAnswerCategoryResponse } from "@/components/standardAnswers/getStandardAnswerCategoriesIfEE";
|
||||
|
||||
async function EditslackChannelConfigPage(props: {
|
||||
params: Promise<{ id: number }>;
|
||||
}) {
|
||||
const params = await props.params;
|
||||
const tasks = [
|
||||
fetchSS("/manage/admin/slack-app/channel"),
|
||||
fetchSS("/manage/document-set"),
|
||||
fetchAgentsSS(),
|
||||
];
|
||||
function EditSlackChannelConfigContent({ id }: { id: string }) {
|
||||
const isPaidEnterprise = usePaidEnterpriseFeaturesEnabled();
|
||||
|
||||
const [
|
||||
slackChannelsResponse,
|
||||
documentSetsResponse,
|
||||
[assistants, agentsFetchError],
|
||||
] = (await Promise.all(tasks)) as [Response, Response, FetchAgentsResponse];
|
||||
const {
|
||||
data: slackChannelConfigs,
|
||||
isLoading: isChannelsLoading,
|
||||
error: channelsError,
|
||||
} = useSlackChannelConfigs();
|
||||
|
||||
const eeStandardAnswerCategoryResponse =
|
||||
await getStandardAnswerCategoriesIfEE();
|
||||
const {
|
||||
data: documentSets,
|
||||
isLoading: isDocSetsLoading,
|
||||
error: docSetsError,
|
||||
} = useDocumentSets();
|
||||
|
||||
if (!slackChannelsResponse.ok) {
|
||||
return (
|
||||
<ErrorCallout
|
||||
errorTitle="Something went wrong :("
|
||||
errorMsg={`Failed to fetch Slack Channels - ${await slackChannelsResponse.text()}`}
|
||||
/>
|
||||
);
|
||||
}
|
||||
const allslackChannelConfigs =
|
||||
(await slackChannelsResponse.json()) as SlackChannelConfig[];
|
||||
const {
|
||||
agents,
|
||||
isLoading: isAgentsLoading,
|
||||
error: agentsError,
|
||||
} = useAgents();
|
||||
|
||||
const slackChannelConfig = allslackChannelConfigs.find(
|
||||
(config) => config.id === Number(params.id)
|
||||
const {
|
||||
data: standardAnswerCategories,
|
||||
isLoading: isStdAnswerLoading,
|
||||
error: stdAnswerError,
|
||||
} = useStandardAnswerCategories();
|
||||
|
||||
const isLoading =
|
||||
isChannelsLoading ||
|
||||
isDocSetsLoading ||
|
||||
isAgentsLoading ||
|
||||
(isPaidEnterprise && isStdAnswerLoading);
|
||||
|
||||
const slackChannelConfig = slackChannelConfigs?.find(
|
||||
(config) => config.id === Number(id)
|
||||
);
|
||||
|
||||
if (!slackChannelConfig) {
|
||||
return (
|
||||
<ErrorCallout
|
||||
errorTitle="Something went wrong :("
|
||||
errorMsg={`Did not find Slack Channel config with ID: ${params.id}`}
|
||||
/>
|
||||
);
|
||||
}
|
||||
|
||||
if (!documentSetsResponse.ok) {
|
||||
return (
|
||||
<ErrorCallout
|
||||
errorTitle="Something went wrong :("
|
||||
errorMsg={`Failed to fetch document sets - ${await documentSetsResponse.text()}`}
|
||||
/>
|
||||
);
|
||||
}
|
||||
const response = await documentSetsResponse.json();
|
||||
const documentSets = response as DocumentSetSummary[];
|
||||
|
||||
if (agentsFetchError) {
|
||||
return (
|
||||
<ErrorCallout
|
||||
errorTitle="Something went wrong :("
|
||||
errorMsg={`Failed to fetch personas - ${agentsFetchError}`}
|
||||
/>
|
||||
);
|
||||
}
|
||||
const title = slackChannelConfig?.is_default
|
||||
? "Edit Default Slack Config"
|
||||
: "Edit Slack Channel Config";
|
||||
|
||||
return (
|
||||
<SettingsLayouts.Root>
|
||||
<InstantSSRAutoRefresh />
|
||||
<SettingsLayouts.Header
|
||||
icon={SvgSlack}
|
||||
title={
|
||||
slackChannelConfig.is_default
|
||||
? "Edit Default Slack Config"
|
||||
: "Edit Slack Channel Config"
|
||||
}
|
||||
title={title}
|
||||
separator
|
||||
backButton
|
||||
/>
|
||||
<SettingsLayouts.Body>
|
||||
<SlackChannelConfigCreationForm
|
||||
slack_bot_id={slackChannelConfig.slack_bot_id}
|
||||
documentSets={documentSets}
|
||||
personas={assistants}
|
||||
standardAnswerCategoryResponse={eeStandardAnswerCategoryResponse}
|
||||
existingSlackChannelConfig={slackChannelConfig}
|
||||
/>
|
||||
{isLoading ? (
|
||||
<SimpleLoader />
|
||||
) : channelsError || !slackChannelConfigs ? (
|
||||
<ErrorCallout
|
||||
errorTitle="Something went wrong :("
|
||||
errorMsg={`Failed to fetch Slack Channels - ${
|
||||
channelsError?.message ?? "unknown error"
|
||||
}`}
|
||||
/>
|
||||
) : !slackChannelConfig ? (
|
||||
<ErrorCallout
|
||||
errorTitle="Something went wrong :("
|
||||
errorMsg={`Did not find Slack Channel config with ID: ${id}`}
|
||||
/>
|
||||
) : docSetsError || !documentSets ? (
|
||||
<ErrorCallout
|
||||
errorTitle="Something went wrong :("
|
||||
errorMsg={`Failed to fetch document sets - ${
|
||||
docSetsError?.message ?? "unknown error"
|
||||
}`}
|
||||
/>
|
||||
) : agentsError ? (
|
||||
<ErrorCallout
|
||||
errorTitle="Something went wrong :("
|
||||
errorMsg={`Failed to fetch agents - ${
|
||||
agentsError?.message ?? "unknown error"
|
||||
}`}
|
||||
/>
|
||||
) : (
|
||||
<SlackChannelConfigCreationForm
|
||||
slack_bot_id={slackChannelConfig.slack_bot_id}
|
||||
documentSets={documentSets}
|
||||
personas={agents}
|
||||
standardAnswerCategoryResponse={
|
||||
isPaidEnterprise
|
||||
? {
|
||||
paidEnterpriseFeaturesEnabled: true,
|
||||
categories: standardAnswerCategories ?? [],
|
||||
...(stdAnswerError
|
||||
? { error: { message: String(stdAnswerError) } }
|
||||
: {}),
|
||||
}
|
||||
: { paidEnterpriseFeaturesEnabled: false }
|
||||
}
|
||||
existingSlackChannelConfig={slackChannelConfig}
|
||||
/>
|
||||
)}
|
||||
</SettingsLayouts.Body>
|
||||
</SettingsLayouts.Root>
|
||||
);
|
||||
}
|
||||
|
||||
export default EditslackChannelConfigPage;
|
||||
export default function Page(props: { params: Promise<{ id: string }> }) {
|
||||
const params = use(props.params);
|
||||
|
||||
return <EditSlackChannelConfigContent id={params.id} />;
|
||||
}
|
||||
|
||||
@@ -1,53 +1,109 @@
|
||||
import { SlackChannelConfigCreationForm } from "../SlackChannelConfigCreationForm";
|
||||
import { fetchSS } from "@/lib/utilsSS";
|
||||
"use client";
|
||||
|
||||
import { use, useEffect } from "react";
|
||||
import { SlackChannelConfigCreationForm } from "@/app/admin/bots/[bot-id]/channels/SlackChannelConfigCreationForm";
|
||||
import { ErrorCallout } from "@/components/ErrorCallout";
|
||||
import { DocumentSetSummary } from "@/lib/types";
|
||||
import { fetchAgentsSS } from "@/lib/agentsSS";
|
||||
import { getStandardAnswerCategoriesIfEE } from "@/components/standardAnswers/getStandardAnswerCategoriesIfEE";
|
||||
import { redirect } from "next/navigation";
|
||||
import SimpleLoader from "@/refresh-components/loaders/SimpleLoader";
|
||||
import * as SettingsLayouts from "@/layouts/settings-layouts";
|
||||
import { SvgSlack } from "@opal/icons";
|
||||
import { useDocumentSets } from "@/app/admin/documents/sets/hooks";
|
||||
import { useAgents } from "@/hooks/useAgents";
|
||||
import { useStandardAnswerCategories } from "@/app/ee/admin/standard-answer/hooks";
|
||||
import { usePaidEnterpriseFeaturesEnabled } from "@/components/settings/usePaidEnterpriseFeaturesEnabled";
|
||||
import type { StandardAnswerCategoryResponse } from "@/components/standardAnswers/getStandardAnswerCategoriesIfEE";
|
||||
import { useRouter } from "next/navigation";
|
||||
|
||||
function NewChannelConfigContent({ slackBotId }: { slackBotId: number }) {
|
||||
const isPaidEnterprise = usePaidEnterpriseFeaturesEnabled();
|
||||
|
||||
const {
|
||||
data: documentSets,
|
||||
isLoading: isDocSetsLoading,
|
||||
error: docSetsError,
|
||||
} = useDocumentSets();
|
||||
|
||||
const {
|
||||
agents,
|
||||
isLoading: isAgentsLoading,
|
||||
error: agentsError,
|
||||
} = useAgents();
|
||||
|
||||
const {
|
||||
data: standardAnswerCategories,
|
||||
isLoading: isStdAnswerLoading,
|
||||
error: stdAnswerError,
|
||||
} = useStandardAnswerCategories();
|
||||
|
||||
if (
|
||||
isDocSetsLoading ||
|
||||
isAgentsLoading ||
|
||||
(isPaidEnterprise && isStdAnswerLoading)
|
||||
) {
|
||||
return <SimpleLoader />;
|
||||
}
|
||||
|
||||
if (docSetsError || !documentSets) {
|
||||
return (
|
||||
<ErrorCallout
|
||||
errorTitle="Something went wrong :("
|
||||
errorMsg={`Failed to fetch document sets - ${
|
||||
docSetsError?.message ?? "unknown error"
|
||||
}`}
|
||||
/>
|
||||
);
|
||||
}
|
||||
|
||||
if (agentsError) {
|
||||
return (
|
||||
<ErrorCallout
|
||||
errorTitle="Something went wrong :("
|
||||
errorMsg={`Failed to fetch agents - ${
|
||||
agentsError?.message ?? "unknown error"
|
||||
}`}
|
||||
/>
|
||||
);
|
||||
}
|
||||
|
||||
const standardAnswerCategoryResponse: StandardAnswerCategoryResponse =
|
||||
isPaidEnterprise
|
||||
? {
|
||||
paidEnterpriseFeaturesEnabled: true,
|
||||
categories: standardAnswerCategories ?? [],
|
||||
...(stdAnswerError
|
||||
? { error: { message: String(stdAnswerError) } }
|
||||
: {}),
|
||||
}
|
||||
: { paidEnterpriseFeaturesEnabled: false };
|
||||
|
||||
return (
|
||||
<SlackChannelConfigCreationForm
|
||||
slack_bot_id={slackBotId}
|
||||
documentSets={documentSets}
|
||||
personas={agents}
|
||||
standardAnswerCategoryResponse={standardAnswerCategoryResponse}
|
||||
/>
|
||||
);
|
||||
}
|
||||
|
||||
export default function Page(props: { params: Promise<{ "bot-id": string }> }) {
|
||||
const unwrappedParams = use(props.params);
|
||||
const router = useRouter();
|
||||
|
||||
async function NewChannelConfigPage(props: {
|
||||
params: Promise<{ "bot-id": string }>;
|
||||
}) {
|
||||
const unwrappedParams = await props.params;
|
||||
const slack_bot_id_raw = unwrappedParams?.["bot-id"] || null;
|
||||
const slack_bot_id = slack_bot_id_raw
|
||||
? parseInt(slack_bot_id_raw as string, 10)
|
||||
: null;
|
||||
|
||||
useEffect(() => {
|
||||
if (!slack_bot_id || isNaN(slack_bot_id)) {
|
||||
router.replace("/admin/bots");
|
||||
}
|
||||
}, [slack_bot_id, router]);
|
||||
|
||||
if (!slack_bot_id || isNaN(slack_bot_id)) {
|
||||
redirect("/admin/bots");
|
||||
return null;
|
||||
}
|
||||
|
||||
const [documentSetsResponse, agentsResponse, standardAnswerCategoryResponse] =
|
||||
await Promise.all([
|
||||
fetchSS("/manage/document-set") as Promise<Response>,
|
||||
fetchAgentsSS(),
|
||||
getStandardAnswerCategoriesIfEE(),
|
||||
]);
|
||||
|
||||
if (!documentSetsResponse.ok) {
|
||||
return (
|
||||
<ErrorCallout
|
||||
errorTitle="Something went wrong :("
|
||||
errorMsg={`Failed to fetch document sets - ${await documentSetsResponse.text()}`}
|
||||
/>
|
||||
);
|
||||
}
|
||||
const documentSets =
|
||||
(await documentSetsResponse.json()) as DocumentSetSummary[];
|
||||
|
||||
if (agentsResponse[1]) {
|
||||
return (
|
||||
<ErrorCallout
|
||||
errorTitle="Something went wrong :("
|
||||
errorMsg={`Failed to fetch agents - ${agentsResponse[1]}`}
|
||||
/>
|
||||
);
|
||||
}
|
||||
|
||||
return (
|
||||
<SettingsLayouts.Root>
|
||||
<SettingsLayouts.Header
|
||||
@@ -57,15 +113,8 @@ async function NewChannelConfigPage(props: {
|
||||
backButton
|
||||
/>
|
||||
<SettingsLayouts.Body>
|
||||
<SlackChannelConfigCreationForm
|
||||
slack_bot_id={slack_bot_id}
|
||||
documentSets={documentSets}
|
||||
personas={agentsResponse[0]}
|
||||
standardAnswerCategoryResponse={standardAnswerCategoryResponse}
|
||||
/>
|
||||
<NewChannelConfigContent slackBotId={slack_bot_id} />
|
||||
</SettingsLayouts.Body>
|
||||
</SettingsLayouts.Root>
|
||||
);
|
||||
}
|
||||
|
||||
export default NewChannelConfigPage;
|
||||
|
||||
@@ -1,82 +1,62 @@
|
||||
"use client";
|
||||
|
||||
import { use } from "react";
|
||||
import BackButton from "@/refresh-components/buttons/BackButton";
|
||||
import { ErrorCallout } from "@/components/ErrorCallout";
|
||||
import { ThreeDotsLoader } from "@/components/Loading";
|
||||
import { InstantSSRAutoRefresh } from "@/components/SSRAutoRefresh";
|
||||
import SimpleLoader from "@/refresh-components/loaders/SimpleLoader";
|
||||
import SlackChannelConfigsTable from "./SlackChannelConfigsTable";
|
||||
import { useSlackBot, useSlackChannelConfigsByBot } from "./hooks";
|
||||
import { ExistingSlackBotForm } from "../SlackBotUpdateForm";
|
||||
import Separator from "@/refresh-components/Separator";
|
||||
|
||||
function SlackBotEditPage({
|
||||
params,
|
||||
}: {
|
||||
params: Promise<{ "bot-id": string }>;
|
||||
}) {
|
||||
// Unwrap the params promise
|
||||
const unwrappedParams = use(params);
|
||||
import * as SettingsLayouts from "@/layouts/settings-layouts";
|
||||
import { SvgSlack } from "@opal/icons";
|
||||
import { getErrorMsg } from "@/lib/error";
|
||||
|
||||
function SlackBotEditContent({ botId }: { botId: string }) {
|
||||
const {
|
||||
data: slackBot,
|
||||
isLoading: isSlackBotLoading,
|
||||
error: slackBotError,
|
||||
refreshSlackBot,
|
||||
} = useSlackBot(Number(unwrappedParams["bot-id"]));
|
||||
} = useSlackBot(Number(botId));
|
||||
|
||||
const {
|
||||
data: slackChannelConfigs,
|
||||
isLoading: isSlackChannelConfigsLoading,
|
||||
error: slackChannelConfigsError,
|
||||
refreshSlackChannelConfigs,
|
||||
} = useSlackChannelConfigsByBot(Number(unwrappedParams["bot-id"]));
|
||||
} = useSlackChannelConfigsByBot(Number(botId));
|
||||
|
||||
if (isSlackBotLoading || isSlackChannelConfigsLoading) {
|
||||
return (
|
||||
<div className="flex justify-center items-center h-screen">
|
||||
<ThreeDotsLoader />
|
||||
</div>
|
||||
);
|
||||
return <SimpleLoader />;
|
||||
}
|
||||
|
||||
if (slackBotError || !slackBot) {
|
||||
const errorMsg =
|
||||
slackBotError?.info?.message ||
|
||||
slackBotError?.info?.detail ||
|
||||
"An unknown error occurred";
|
||||
return (
|
||||
<ErrorCallout
|
||||
errorTitle="Something went wrong :("
|
||||
errorMsg={`Failed to fetch Slack Bot ${unwrappedParams["bot-id"]}: ${errorMsg}`}
|
||||
errorMsg={`Failed to fetch Slack Bot ${botId}: ${getErrorMsg(
|
||||
slackBotError
|
||||
)}`}
|
||||
/>
|
||||
);
|
||||
}
|
||||
|
||||
if (slackChannelConfigsError || !slackChannelConfigs) {
|
||||
const errorMsg =
|
||||
slackChannelConfigsError?.info?.message ||
|
||||
slackChannelConfigsError?.info?.detail ||
|
||||
"An unknown error occurred";
|
||||
return (
|
||||
<ErrorCallout
|
||||
errorTitle="Something went wrong :("
|
||||
errorMsg={`Failed to fetch Slack Bot ${unwrappedParams["bot-id"]}: ${errorMsg}`}
|
||||
errorMsg={`Failed to fetch Slack Bot ${botId}: ${getErrorMsg(
|
||||
slackChannelConfigsError
|
||||
)}`}
|
||||
/>
|
||||
);
|
||||
}
|
||||
|
||||
return (
|
||||
<>
|
||||
<InstantSSRAutoRefresh />
|
||||
|
||||
<BackButton routerOverride="/admin/bots" />
|
||||
|
||||
<ExistingSlackBotForm
|
||||
existingSlackBot={slackBot}
|
||||
refreshSlackBot={refreshSlackBot}
|
||||
/>
|
||||
<Separator />
|
||||
|
||||
<div className="mt-8">
|
||||
<SlackChannelConfigsTable
|
||||
@@ -94,9 +74,19 @@ export default function Page({
|
||||
}: {
|
||||
params: Promise<{ "bot-id": string }>;
|
||||
}) {
|
||||
const unwrappedParams = use(params);
|
||||
|
||||
return (
|
||||
<>
|
||||
<SlackBotEditPage params={params} />
|
||||
</>
|
||||
<SettingsLayouts.Root>
|
||||
<SettingsLayouts.Header
|
||||
icon={SvgSlack}
|
||||
title="Edit Slack Bot"
|
||||
backButton
|
||||
separator
|
||||
/>
|
||||
<SettingsLayouts.Body>
|
||||
<SlackBotEditContent botId={unwrappedParams["bot-id"]} />
|
||||
</SettingsLayouts.Body>
|
||||
</SettingsLayouts.Root>
|
||||
);
|
||||
}
|
||||
|
||||
@@ -1,12 +1,7 @@
|
||||
import BackButton from "@/refresh-components/buttons/BackButton";
|
||||
"use client";
|
||||
|
||||
import { NewSlackBotForm } from "../SlackBotCreationForm";
|
||||
|
||||
export default async function NewSlackBotPage() {
|
||||
return (
|
||||
<>
|
||||
<BackButton routerOverride="/admin/bots" />
|
||||
|
||||
<NewSlackBotForm />
|
||||
</>
|
||||
);
|
||||
export default function Page() {
|
||||
return <NewSlackBotForm />;
|
||||
}
|
||||
|
||||
@@ -1,11 +1,14 @@
|
||||
import React, { JSX, memo } from "react";
|
||||
import {
|
||||
ChatPacket,
|
||||
CODE_INTERPRETER_TOOL_TYPES,
|
||||
ImageGenerationToolPacket,
|
||||
Packet,
|
||||
PacketType,
|
||||
ReasoningPacket,
|
||||
SearchToolStart,
|
||||
StopReason,
|
||||
ToolCallArgumentDelta,
|
||||
} from "../../services/streamingModels";
|
||||
import {
|
||||
FullChatState,
|
||||
@@ -26,7 +29,6 @@ import { DeepResearchPlanRenderer } from "./timeline/renderers/deepresearch/Deep
|
||||
import { ResearchAgentRenderer } from "./timeline/renderers/deepresearch/ResearchAgentRenderer";
|
||||
import { WebSearchToolRenderer } from "./timeline/renderers/search/WebSearchToolRenderer";
|
||||
import { InternalSearchToolRenderer } from "./timeline/renderers/search/InternalSearchToolRenderer";
|
||||
import { SearchToolStart } from "../../services/streamingModels";
|
||||
|
||||
// Different types of chat packets using discriminated unions
|
||||
interface GroupedPackets {
|
||||
@@ -56,7 +58,12 @@ function isImageToolPacket(packet: Packet) {
|
||||
}
|
||||
|
||||
function isPythonToolPacket(packet: Packet) {
|
||||
return packet.obj.type === PacketType.PYTHON_TOOL_START;
|
||||
return (
|
||||
packet.obj.type === PacketType.PYTHON_TOOL_START ||
|
||||
(packet.obj.type === PacketType.TOOL_CALL_ARGUMENT_DELTA &&
|
||||
(packet.obj as ToolCallArgumentDelta).tool_type ===
|
||||
CODE_INTERPRETER_TOOL_TYPES.PYTHON)
|
||||
);
|
||||
}
|
||||
|
||||
function isCustomToolPacket(packet: Packet) {
|
||||
|
||||
@@ -10,6 +10,8 @@ import {
|
||||
Stop,
|
||||
ImageGenerationToolDelta,
|
||||
MessageStart,
|
||||
ToolCallArgumentDelta,
|
||||
CODE_INTERPRETER_TOOL_TYPES,
|
||||
} from "@/app/app/services/streamingModels";
|
||||
import { CitationMap } from "@/app/app/interfaces";
|
||||
import { OnyxDocument } from "@/lib/search/interfaces";
|
||||
@@ -138,6 +140,7 @@ const CONTENT_PACKET_TYPES_SET = new Set<PacketType>([
|
||||
PacketType.SEARCH_TOOL_START,
|
||||
PacketType.IMAGE_GENERATION_TOOL_START,
|
||||
PacketType.PYTHON_TOOL_START,
|
||||
PacketType.TOOL_CALL_ARGUMENT_DELTA,
|
||||
PacketType.CUSTOM_TOOL_START,
|
||||
PacketType.FILE_READER_START,
|
||||
PacketType.FETCH_TOOL_START,
|
||||
@@ -149,9 +152,16 @@ const CONTENT_PACKET_TYPES_SET = new Set<PacketType>([
|
||||
]);
|
||||
|
||||
function hasContentPackets(packets: Packet[]): boolean {
|
||||
return packets.some((packet) =>
|
||||
CONTENT_PACKET_TYPES_SET.has(packet.obj.type as PacketType)
|
||||
);
|
||||
return packets.some((packet) => {
|
||||
const type = packet.obj.type as PacketType;
|
||||
if (type === PacketType.TOOL_CALL_ARGUMENT_DELTA) {
|
||||
return (
|
||||
(packet.obj as ToolCallArgumentDelta).tool_type ===
|
||||
CODE_INTERPRETER_TOOL_TYPES.PYTHON
|
||||
);
|
||||
}
|
||||
return CONTENT_PACKET_TYPES_SET.has(type);
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
|
||||
@@ -1,6 +1,13 @@
|
||||
import { Packet, PacketType } from "@/app/app/services/streamingModels";
|
||||
import {
|
||||
CODE_INTERPRETER_TOOL_TYPES,
|
||||
Packet,
|
||||
PacketType,
|
||||
ToolCallArgumentDelta,
|
||||
} from "@/app/app/services/streamingModels";
|
||||
|
||||
// Packet types with renderers supporting collapsed streaming mode
|
||||
// Packet types with renderers supporting collapsed streaming mode.
|
||||
// TOOL_CALL_ARGUMENT_DELTA is intentionally excluded here because it requires
|
||||
// a tool_type check — it's handled separately in stepSupportsCollapsedStreaming.
|
||||
export const COLLAPSED_STREAMING_PACKET_TYPES = new Set<PacketType>([
|
||||
PacketType.SEARCH_TOOL_START,
|
||||
PacketType.FETCH_TOOL_START,
|
||||
@@ -21,7 +28,13 @@ export const isSearchToolPackets = (packets: Packet[]): boolean =>
|
||||
|
||||
// Check if packets belong to a python tool
|
||||
export const isPythonToolPackets = (packets: Packet[]): boolean =>
|
||||
packets.some((p) => p.obj.type === PacketType.PYTHON_TOOL_START);
|
||||
packets.some(
|
||||
(p) =>
|
||||
p.obj.type === PacketType.PYTHON_TOOL_START ||
|
||||
(p.obj.type === PacketType.TOOL_CALL_ARGUMENT_DELTA &&
|
||||
(p.obj as ToolCallArgumentDelta).tool_type ===
|
||||
CODE_INTERPRETER_TOOL_TYPES.PYTHON)
|
||||
);
|
||||
|
||||
// Check if packets belong to reasoning
|
||||
export const isReasoningPackets = (packets: Packet[]): boolean =>
|
||||
@@ -29,8 +42,12 @@ export const isReasoningPackets = (packets: Packet[]): boolean =>
|
||||
|
||||
// Check if step supports collapsed streaming rendering mode
|
||||
export const stepSupportsCollapsedStreaming = (packets: Packet[]): boolean =>
|
||||
packets.some((p) =>
|
||||
COLLAPSED_STREAMING_PACKET_TYPES.has(p.obj.type as PacketType)
|
||||
packets.some(
|
||||
(p) =>
|
||||
COLLAPSED_STREAMING_PACKET_TYPES.has(p.obj.type as PacketType) ||
|
||||
(p.obj.type === PacketType.TOOL_CALL_ARGUMENT_DELTA &&
|
||||
(p.obj as ToolCallArgumentDelta).tool_type ===
|
||||
CODE_INTERPRETER_TOOL_TYPES.PYTHON)
|
||||
);
|
||||
|
||||
// Check if packets have content worth rendering in collapsed streaming mode.
|
||||
@@ -67,7 +84,13 @@ export const stepHasCollapsedStreamingContent = (
|
||||
// Python tool renders code/output from the start packet onward
|
||||
if (
|
||||
packetTypes.has(PacketType.PYTHON_TOOL_START) ||
|
||||
packetTypes.has(PacketType.PYTHON_TOOL_DELTA)
|
||||
packetTypes.has(PacketType.PYTHON_TOOL_DELTA) ||
|
||||
packets.some(
|
||||
(p) =>
|
||||
p.obj.type === PacketType.TOOL_CALL_ARGUMENT_DELTA &&
|
||||
(p.obj as ToolCallArgumentDelta).tool_type ===
|
||||
CODE_INTERPRETER_TOOL_TYPES.PYTHON
|
||||
)
|
||||
) {
|
||||
return true;
|
||||
}
|
||||
|
||||
@@ -4,7 +4,9 @@ import {
|
||||
PythonToolPacket,
|
||||
PythonToolStart,
|
||||
PythonToolDelta,
|
||||
ToolCallArgumentDelta,
|
||||
SectionEnd,
|
||||
CODE_INTERPRETER_TOOL_TYPES,
|
||||
} from "@/app/app/services/streamingModels";
|
||||
import {
|
||||
MessageRenderer,
|
||||
@@ -39,6 +41,18 @@ function HighlightedPythonCode({ code }: { code: string }) {
|
||||
|
||||
// Helper function to construct current Python execution state
|
||||
function constructCurrentPythonState(packets: PythonToolPacket[]) {
|
||||
// Accumulate streaming code from argument deltas (arrives before PythonToolStart)
|
||||
const streamingCode = packets
|
||||
.filter(
|
||||
(packet) =>
|
||||
packet.obj.type === PacketType.TOOL_CALL_ARGUMENT_DELTA &&
|
||||
(packet.obj as ToolCallArgumentDelta).tool_type ===
|
||||
CODE_INTERPRETER_TOOL_TYPES.PYTHON
|
||||
)
|
||||
.map((packet) =>
|
||||
String((packet.obj as ToolCallArgumentDelta).argument_deltas.code ?? "")
|
||||
)
|
||||
.join("");
|
||||
const pythonStart = packets.find(
|
||||
(packet) => packet.obj.type === PacketType.PYTHON_TOOL_START
|
||||
)?.obj as PythonToolStart | null;
|
||||
@@ -51,7 +65,8 @@ function constructCurrentPythonState(packets: PythonToolPacket[]) {
|
||||
packet.obj.type === PacketType.ERROR
|
||||
)?.obj as SectionEnd | null;
|
||||
|
||||
const code = pythonStart?.code || "";
|
||||
// Use complete code from PythonToolStart if available, else use streamed code.
|
||||
const code = pythonStart?.code || streamingCode;
|
||||
const stdout = pythonDeltas
|
||||
.map((delta) => delta?.stdout || "")
|
||||
.filter((s) => s)
|
||||
@@ -61,6 +76,7 @@ function constructCurrentPythonState(packets: PythonToolPacket[]) {
|
||||
.filter((s) => s)
|
||||
.join("");
|
||||
const fileIds = pythonDeltas.flatMap((delta) => delta?.file_ids || []);
|
||||
const isStreaming = !pythonStart && streamingCode.length > 0;
|
||||
const isExecuting = pythonStart && !pythonEnd;
|
||||
const isComplete = pythonStart && pythonEnd;
|
||||
const hasError = stderr.length > 0;
|
||||
@@ -70,6 +86,7 @@ function constructCurrentPythonState(packets: PythonToolPacket[]) {
|
||||
stdout,
|
||||
stderr,
|
||||
fileIds,
|
||||
isStreaming,
|
||||
isExecuting,
|
||||
isComplete,
|
||||
hasError,
|
||||
@@ -82,8 +99,16 @@ export const PythonToolRenderer: MessageRenderer<PythonToolPacket, {}> = ({
|
||||
renderType,
|
||||
children,
|
||||
}) => {
|
||||
const { code, stdout, stderr, fileIds, isExecuting, isComplete, hasError } =
|
||||
constructCurrentPythonState(packets);
|
||||
const {
|
||||
code,
|
||||
stdout,
|
||||
stderr,
|
||||
fileIds,
|
||||
isStreaming,
|
||||
isExecuting,
|
||||
isComplete,
|
||||
hasError,
|
||||
} = constructCurrentPythonState(packets);
|
||||
|
||||
useEffect(() => {
|
||||
if (isComplete) {
|
||||
@@ -92,6 +117,9 @@ export const PythonToolRenderer: MessageRenderer<PythonToolPacket, {}> = ({
|
||||
}, [isComplete, onComplete]);
|
||||
|
||||
const status = useMemo(() => {
|
||||
if (isStreaming) {
|
||||
return "Writing code...";
|
||||
}
|
||||
if (isExecuting) {
|
||||
return "Executing Python code...";
|
||||
}
|
||||
@@ -102,13 +130,13 @@ export const PythonToolRenderer: MessageRenderer<PythonToolPacket, {}> = ({
|
||||
return "Python execution completed";
|
||||
}
|
||||
return "Python execution";
|
||||
}, [isComplete, isExecuting, hasError]);
|
||||
}, [isStreaming, isComplete, isExecuting, hasError]);
|
||||
|
||||
// Shared content for all states - used by both FULL and compact modes
|
||||
const content = (
|
||||
<div className="flex flex-col mb-1 space-y-2">
|
||||
{/* Loading indicator when executing */}
|
||||
{isExecuting && (
|
||||
{/* Loading indicator when streaming or executing */}
|
||||
{(isStreaming || isExecuting) && (
|
||||
<div className="flex items-center gap-2 text-sm text-muted-foreground">
|
||||
<div className="flex gap-0.5">
|
||||
<div className="w-1 h-1 bg-current rounded-full animate-pulse"></div>
|
||||
@@ -121,7 +149,7 @@ export const PythonToolRenderer: MessageRenderer<PythonToolPacket, {}> = ({
|
||||
style={{ animationDelay: "0.2s" }}
|
||||
></div>
|
||||
</div>
|
||||
<span>Running code...</span>
|
||||
<span>{isStreaming ? "Writing code..." : "Running code..."}</span>
|
||||
</div>
|
||||
)}
|
||||
|
||||
|
||||
@@ -308,15 +308,15 @@ export async function getProjectTokenCount(projectId: number): Promise<number> {
|
||||
|
||||
export async function getMaxSelectedDocumentTokens(
|
||||
personaId: number
|
||||
): Promise<number> {
|
||||
): Promise<number | null> {
|
||||
const response = await fetch(
|
||||
`/api/chat/max-selected-document-tokens?persona_id=${personaId}`
|
||||
);
|
||||
if (!response.ok) {
|
||||
return 128_000;
|
||||
return null;
|
||||
}
|
||||
const json = await response.json();
|
||||
return (json?.max_tokens as number) ?? 128_000;
|
||||
return (json?.max_tokens as number) ?? null;
|
||||
}
|
||||
|
||||
export async function moveChatSession(
|
||||
|
||||
@@ -288,15 +288,15 @@ export async function deleteAllChatSessions() {
|
||||
|
||||
export async function getAvailableContextTokens(
|
||||
chatSessionId: string
|
||||
): Promise<number> {
|
||||
): Promise<number | null> {
|
||||
const response = await fetch(
|
||||
`/api/chat/available-context-tokens/${chatSessionId}`
|
||||
);
|
||||
if (!response.ok) {
|
||||
return 0;
|
||||
return null;
|
||||
}
|
||||
const data = (await response.json()) as { available_tokens: number };
|
||||
return data?.available_tokens ?? 0;
|
||||
return data?.available_tokens ?? null;
|
||||
}
|
||||
|
||||
export function processRawChatHistory(
|
||||
|
||||
@@ -16,6 +16,7 @@ export function isToolPacket(
|
||||
PacketType.SEARCH_TOOL_DOCUMENTS_DELTA,
|
||||
PacketType.PYTHON_TOOL_START,
|
||||
PacketType.PYTHON_TOOL_DELTA,
|
||||
PacketType.TOOL_CALL_ARGUMENT_DELTA,
|
||||
PacketType.CUSTOM_TOOL_START,
|
||||
PacketType.CUSTOM_TOOL_DELTA,
|
||||
PacketType.FILE_READER_START,
|
||||
|
||||
@@ -27,6 +27,9 @@ export enum PacketType {
|
||||
FETCH_TOOL_URLS = "open_url_urls",
|
||||
FETCH_TOOL_DOCUMENTS = "open_url_documents",
|
||||
|
||||
// Tool call argument delta (streams tool args before tool executes)
|
||||
TOOL_CALL_ARGUMENT_DELTA = "tool_call_argument_delta",
|
||||
|
||||
// Custom tool packets
|
||||
CUSTOM_TOOL_START = "custom_tool_start",
|
||||
CUSTOM_TOOL_DELTA = "custom_tool_delta",
|
||||
@@ -59,6 +62,10 @@ export enum PacketType {
|
||||
INTERMEDIATE_REPORT_CITED_DOCS = "intermediate_report_cited_docs",
|
||||
}
|
||||
|
||||
export const CODE_INTERPRETER_TOOL_TYPES = {
|
||||
PYTHON: "python",
|
||||
} as const;
|
||||
|
||||
// Basic Message Packets
|
||||
export interface MessageStart extends BaseObj {
|
||||
id: string;
|
||||
@@ -149,6 +156,13 @@ export interface PythonToolDelta extends BaseObj {
|
||||
file_ids: string[];
|
||||
}
|
||||
|
||||
export interface ToolCallArgumentDelta extends BaseObj {
|
||||
type: "tool_call_argument_delta";
|
||||
tool_type: string;
|
||||
tool_id: string;
|
||||
argument_deltas: Record<string, unknown>;
|
||||
}
|
||||
|
||||
export interface FetchToolStart extends BaseObj {
|
||||
type: "open_url_start";
|
||||
}
|
||||
@@ -294,6 +308,7 @@ export type ImageGenerationToolObj =
|
||||
export type PythonToolObj =
|
||||
| PythonToolStart
|
||||
| PythonToolDelta
|
||||
| ToolCallArgumentDelta
|
||||
| SectionEnd
|
||||
| PacketError;
|
||||
export type FetchToolObj =
|
||||
|
||||
@@ -2,9 +2,12 @@
|
||||
|
||||
import {
|
||||
buildChatUrl,
|
||||
getAvailableContextTokens,
|
||||
nameChatSession,
|
||||
updateLlmOverrideForChatSession,
|
||||
} from "@/app/app/services/lib";
|
||||
import { getMaxSelectedDocumentTokens } from "@/app/app/projects/projectsService";
|
||||
import { DEFAULT_CONTEXT_TOKENS } from "@/lib/constants";
|
||||
import { StreamStopInfo } from "@/lib/search/interfaces";
|
||||
import { useCallback, useEffect, useMemo, useRef, useState } from "react";
|
||||
import type { Route } from "next";
|
||||
@@ -194,9 +197,6 @@ export default function useChatController({
|
||||
|
||||
const navigatingAway = useRef(false);
|
||||
|
||||
// Local state that doesn't need to be in the store
|
||||
const [_maxTokens, setMaxTokens] = useState<number>(4096);
|
||||
|
||||
// Sync store state changes
|
||||
useEffect(() => {
|
||||
if (currentSessionId) {
|
||||
@@ -1067,21 +1067,59 @@ export default function useChatController({
|
||||
handleSlackChatRedirect();
|
||||
}, [searchParams, router]);
|
||||
|
||||
// fetch # of allowed document tokens for the selected Persona
|
||||
useEffect(() => {
|
||||
if (!liveAgent?.id) return; // avoid calling with undefined persona id
|
||||
// Available context tokens: if a chat session exists, fetch from the session
|
||||
// API (dynamic per session/model). Otherwise derive from the persona's max
|
||||
// document tokens. The backend already accounts for system prompt, tools,
|
||||
// and user-message reservations.
|
||||
const [availableContextTokens, setAvailableContextTokens] = useState<number>(
|
||||
DEFAULT_CONTEXT_TOKENS
|
||||
);
|
||||
|
||||
async function fetchMaxTokens() {
|
||||
const response = await fetch(
|
||||
`/api/chat/max-selected-document-tokens?persona_id=${liveAgent?.id}`
|
||||
);
|
||||
if (response.ok) {
|
||||
const maxTokens = (await response.json()).max_tokens as number;
|
||||
setMaxTokens(maxTokens);
|
||||
useEffect(() => {
|
||||
if (!llmManager.hasAnyProvider) return;
|
||||
|
||||
let cancelled = false;
|
||||
|
||||
const setIfActive = (tokens: number) => {
|
||||
if (!cancelled) setAvailableContextTokens(tokens);
|
||||
};
|
||||
|
||||
// Prefer the Zustand session ID, but fall back to the URL-derived prop
|
||||
// so we don't incorrectly take the persona path while the store is
|
||||
// still initialising on navigation to an existing chat.
|
||||
const sessionId = currentSessionId || existingChatSessionId;
|
||||
|
||||
(async () => {
|
||||
try {
|
||||
if (sessionId) {
|
||||
const available = await getAvailableContextTokens(sessionId);
|
||||
setIfActive(available ?? DEFAULT_CONTEXT_TOKENS);
|
||||
return;
|
||||
}
|
||||
|
||||
const personaId = liveAgent?.id;
|
||||
if (personaId == null) {
|
||||
setIfActive(DEFAULT_CONTEXT_TOKENS);
|
||||
return;
|
||||
}
|
||||
|
||||
const maxTokens = await getMaxSelectedDocumentTokens(personaId);
|
||||
setIfActive(maxTokens ?? DEFAULT_CONTEXT_TOKENS);
|
||||
} catch (e) {
|
||||
console.error("Failed to fetch available context tokens:", e);
|
||||
setIfActive(DEFAULT_CONTEXT_TOKENS);
|
||||
}
|
||||
}
|
||||
fetchMaxTokens();
|
||||
}, [liveAgent]);
|
||||
})();
|
||||
|
||||
return () => {
|
||||
cancelled = true;
|
||||
};
|
||||
}, [
|
||||
currentSessionId,
|
||||
existingChatSessionId,
|
||||
liveAgent?.id,
|
||||
llmManager.hasAnyProvider,
|
||||
]);
|
||||
|
||||
// check if there's an image file in the message history so that we know
|
||||
// which LLMs are available to use
|
||||
@@ -1110,5 +1148,7 @@ export default function useChatController({
|
||||
onSubmit,
|
||||
stopGenerating,
|
||||
handleMessageSpecificFileUpload,
|
||||
// data
|
||||
availableContextTokens,
|
||||
};
|
||||
}
|
||||
|
||||
@@ -36,6 +36,7 @@ export interface Settings {
|
||||
|
||||
// User Knowledge settings
|
||||
user_knowledge_enabled?: boolean;
|
||||
user_file_max_upload_size_mb?: number | null;
|
||||
|
||||
// Connector settings
|
||||
show_extra_connectors?: boolean;
|
||||
|
||||
10
web/src/lib/error.ts
Normal file
10
web/src/lib/error.ts
Normal file
@@ -0,0 +1,10 @@
|
||||
/**
|
||||
* Extract a human-readable error message from an SWR error object.
|
||||
* SWR errors from `errorHandlingFetcher` attach `info.message` or `info.detail`.
|
||||
*/
|
||||
export function getErrorMsg(
|
||||
error: { info?: { message?: string; detail?: string } } | null | undefined,
|
||||
fallback = "An unknown error occurred"
|
||||
): string {
|
||||
return error?.info?.message || error?.info?.detail || fallback;
|
||||
}
|
||||
@@ -43,6 +43,7 @@ import { useAppRouter } from "@/hooks/appNavigation";
|
||||
import { ChatFileType } from "@/app/app/interfaces";
|
||||
import { toast } from "@/hooks/useToast";
|
||||
import { useProjects } from "@/lib/hooks/useProjects";
|
||||
import { SettingsContext } from "@/providers/SettingsProvider";
|
||||
|
||||
export type { Project, ProjectFile } from "@/app/app/projects/projectsService";
|
||||
|
||||
@@ -84,6 +85,8 @@ function buildFileKey(file: File): string {
|
||||
return `${file.size}|${namePrefix}`;
|
||||
}
|
||||
|
||||
const DEFAULT_USER_FILE_MAX_UPLOAD_SIZE_MB = 50;
|
||||
|
||||
interface ProjectsContextType {
|
||||
projects: Project[];
|
||||
recentFiles: ProjectFile[];
|
||||
@@ -157,6 +160,7 @@ export function ProjectsProvider({ children }: ProjectsProviderProps) {
|
||||
new Map()
|
||||
);
|
||||
const route = useAppRouter();
|
||||
const settingsContext = useContext(SettingsContext);
|
||||
|
||||
// Use SWR's mutate to refresh projects - returns the new data
|
||||
const fetchProjects = useCallback(async (): Promise<Project[]> => {
|
||||
@@ -336,16 +340,40 @@ export function ProjectsProvider({ children }: ProjectsProviderProps) {
|
||||
onSuccess?: (uploaded: CategorizedFiles) => void,
|
||||
onFailure?: (failedTempIds: string[]) => void
|
||||
): Promise<ProjectFile[]> => {
|
||||
const optimisticFiles = files.map((f) =>
|
||||
const rawMax = settingsContext?.settings?.user_file_max_upload_size_mb;
|
||||
const maxUploadSizeMb =
|
||||
rawMax && rawMax > 0 ? rawMax : DEFAULT_USER_FILE_MAX_UPLOAD_SIZE_MB;
|
||||
const maxUploadSizeBytes = maxUploadSizeMb * 1024 * 1024;
|
||||
|
||||
const oversizedFiles = files.filter(
|
||||
(file) => file.size > maxUploadSizeBytes
|
||||
);
|
||||
const validFiles = files.filter(
|
||||
(file) => file.size <= maxUploadSizeBytes
|
||||
);
|
||||
|
||||
if (oversizedFiles.length > 0) {
|
||||
const skippedNames = oversizedFiles.map((file) => file.name).join(", ");
|
||||
toast.warning(
|
||||
`Skipped ${oversizedFiles.length} oversized file(s) (>${maxUploadSizeMb} MB): ${skippedNames}`
|
||||
);
|
||||
}
|
||||
|
||||
if (validFiles.length === 0) {
|
||||
onFailure?.([]);
|
||||
return [];
|
||||
}
|
||||
|
||||
const optimisticFiles = validFiles.map((f) =>
|
||||
createOptimisticFile(f, projectId)
|
||||
);
|
||||
const tempIdMap = getTempIdMap(files, optimisticFiles);
|
||||
const tempIdMap = getTempIdMap(validFiles, optimisticFiles);
|
||||
setAllRecentFiles((prev) => [...optimisticFiles, ...prev]);
|
||||
if (projectId) {
|
||||
setAllCurrentProjectFiles((prev) => [...optimisticFiles, ...prev]);
|
||||
projectToUploadFilesMapRef.current.set(projectId, optimisticFiles);
|
||||
}
|
||||
svcUploadFiles(files, projectId, tempIdMap)
|
||||
svcUploadFiles(validFiles, projectId, tempIdMap)
|
||||
.then((uploaded) => {
|
||||
const uploadedFiles = uploaded.user_files || [];
|
||||
const tempIdToUploadedFileMap = new Map(
|
||||
@@ -445,6 +473,7 @@ export function ProjectsProvider({ children }: ProjectsProviderProps) {
|
||||
refreshCurrentProjectDetails,
|
||||
refreshRecentFiles,
|
||||
removeOptimisticFilesByTempIds,
|
||||
settingsContext,
|
||||
]
|
||||
);
|
||||
|
||||
|
||||
166
web/src/providers/__tests__/ProjectsContext.test.tsx
Normal file
166
web/src/providers/__tests__/ProjectsContext.test.tsx
Normal file
@@ -0,0 +1,166 @@
|
||||
import React, { PropsWithChildren } from "react";
|
||||
import { act, renderHook } from "@testing-library/react";
|
||||
import {
|
||||
ProjectsProvider,
|
||||
useProjectsContext,
|
||||
} from "@/providers/ProjectsContext";
|
||||
import { SettingsContext } from "@/providers/SettingsProvider";
|
||||
import { CombinedSettings } from "@/interfaces/settings";
|
||||
import type { ProjectFile } from "@/app/app/projects/projectsService";
|
||||
|
||||
const mockUploadFiles = jest.fn();
|
||||
const mockGetRecentFiles = jest.fn();
|
||||
const mockToastWarning = jest.fn();
|
||||
|
||||
jest.mock("next/navigation", () => ({
|
||||
useSearchParams: () => ({
|
||||
get: () => null,
|
||||
}),
|
||||
}));
|
||||
|
||||
jest.mock("@/hooks/appNavigation", () => ({
|
||||
useAppRouter: () => jest.fn(),
|
||||
}));
|
||||
|
||||
jest.mock("@/lib/hooks/useProjects", () => ({
|
||||
useProjects: () => ({
|
||||
projects: [],
|
||||
refreshProjects: jest.fn().mockResolvedValue([]),
|
||||
}),
|
||||
}));
|
||||
|
||||
jest.mock("@/hooks/useToast", () => ({
|
||||
toast: {
|
||||
warning: (...args: unknown[]) => mockToastWarning(...args),
|
||||
error: jest.fn(),
|
||||
success: jest.fn(),
|
||||
},
|
||||
}));
|
||||
|
||||
jest.mock("@/app/app/projects/projectsService", () => {
|
||||
const actual = jest.requireActual("@/app/app/projects/projectsService");
|
||||
return {
|
||||
...actual,
|
||||
fetchProjects: jest.fn().mockResolvedValue([]),
|
||||
createProject: jest.fn(),
|
||||
uploadFiles: (...args: unknown[]) => mockUploadFiles(...args),
|
||||
getRecentFiles: (...args: unknown[]) => mockGetRecentFiles(...args),
|
||||
getFilesInProject: jest.fn().mockResolvedValue([]),
|
||||
getProject: jest.fn(),
|
||||
getProjectInstructions: jest.fn(),
|
||||
upsertProjectInstructions: jest.fn(),
|
||||
getProjectDetails: jest.fn(),
|
||||
renameProject: jest.fn(),
|
||||
deleteProject: jest.fn(),
|
||||
deleteUserFile: jest.fn(),
|
||||
getUserFileStatuses: jest.fn().mockResolvedValue([]),
|
||||
unlinkFileFromProject: jest.fn(),
|
||||
linkFileToProject: jest.fn(),
|
||||
};
|
||||
});
|
||||
|
||||
const settingsValue: CombinedSettings = {
|
||||
settings: {
|
||||
user_file_max_upload_size_mb: 1,
|
||||
} as CombinedSettings["settings"],
|
||||
enterpriseSettings: null,
|
||||
customAnalyticsScript: null,
|
||||
webVersion: null,
|
||||
webDomain: null,
|
||||
isSearchModeAvailable: true,
|
||||
};
|
||||
|
||||
const wrapper = ({ children }: PropsWithChildren) => (
|
||||
<SettingsContext.Provider value={settingsValue}>
|
||||
<ProjectsProvider>{children}</ProjectsProvider>
|
||||
</SettingsContext.Provider>
|
||||
);
|
||||
|
||||
describe("ProjectsContext beginUpload size precheck", () => {
|
||||
beforeEach(() => {
|
||||
mockUploadFiles.mockReset();
|
||||
mockGetRecentFiles.mockReset();
|
||||
mockToastWarning.mockReset();
|
||||
|
||||
mockUploadFiles.mockResolvedValue({
|
||||
user_files: [],
|
||||
rejected_files: [],
|
||||
});
|
||||
mockGetRecentFiles.mockResolvedValue([]);
|
||||
});
|
||||
|
||||
it("only sends valid files to the upload API when oversized files are present", async () => {
|
||||
const { result } = renderHook(() => useProjectsContext(), { wrapper });
|
||||
|
||||
const valid = new File(["small"], "small.txt", { type: "text/plain" });
|
||||
const oversized = new File([new Uint8Array(2 * 1024 * 1024)], "big.txt", {
|
||||
type: "text/plain",
|
||||
});
|
||||
|
||||
let optimisticFiles: ProjectFile[] = [];
|
||||
await act(async () => {
|
||||
optimisticFiles = await result.current.beginUpload(
|
||||
[valid, oversized],
|
||||
null
|
||||
);
|
||||
});
|
||||
|
||||
expect(mockUploadFiles).toHaveBeenCalledTimes(1);
|
||||
const [uploadedFiles] = mockUploadFiles.mock.calls[0];
|
||||
expect((uploadedFiles as File[]).map((f) => f.name)).toEqual(["small.txt"]);
|
||||
expect(optimisticFiles.map((f) => f.name)).toEqual(["small.txt"]);
|
||||
expect(mockToastWarning).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
|
||||
it("uploads all files when none are oversized", async () => {
|
||||
const { result } = renderHook(() => useProjectsContext(), { wrapper });
|
||||
|
||||
const first = new File(["small"], "first.txt", { type: "text/plain" });
|
||||
const second = new File(["small"], "second.txt", { type: "text/plain" });
|
||||
|
||||
let optimisticFiles: ProjectFile[] = [];
|
||||
await act(async () => {
|
||||
optimisticFiles = await result.current.beginUpload([first, second], null);
|
||||
});
|
||||
|
||||
expect(mockUploadFiles).toHaveBeenCalledTimes(1);
|
||||
const [uploadedFiles] = mockUploadFiles.mock.calls[0];
|
||||
expect((uploadedFiles as File[]).map((f) => f.name)).toEqual([
|
||||
"first.txt",
|
||||
"second.txt",
|
||||
]);
|
||||
expect(mockToastWarning).not.toHaveBeenCalled();
|
||||
expect(optimisticFiles.map((f) => f.name)).toEqual([
|
||||
"first.txt",
|
||||
"second.txt",
|
||||
]);
|
||||
});
|
||||
|
||||
it("does not call upload API when all files are oversized", async () => {
|
||||
const { result } = renderHook(() => useProjectsContext(), { wrapper });
|
||||
|
||||
const oversized = new File(
|
||||
[new Uint8Array(2 * 1024 * 1024)],
|
||||
"too-big.txt",
|
||||
{ type: "text/plain" }
|
||||
);
|
||||
const onSuccess = jest.fn();
|
||||
const onFailure = jest.fn();
|
||||
|
||||
let optimisticFiles: ProjectFile[] = [];
|
||||
await act(async () => {
|
||||
optimisticFiles = await result.current.beginUpload(
|
||||
[oversized],
|
||||
null,
|
||||
onSuccess,
|
||||
onFailure
|
||||
);
|
||||
});
|
||||
|
||||
expect(mockUploadFiles).not.toHaveBeenCalled();
|
||||
expect(optimisticFiles).toEqual([]);
|
||||
expect(mockToastWarning).toHaveBeenCalledTimes(1);
|
||||
expect(onSuccess).not.toHaveBeenCalled();
|
||||
expect(onFailure).toHaveBeenCalledWith([]);
|
||||
});
|
||||
});
|
||||
@@ -1,10 +1,7 @@
|
||||
"use client";
|
||||
|
||||
import { redirect, useRouter, useSearchParams } from "next/navigation";
|
||||
import {
|
||||
personaIncludesRetrieval,
|
||||
getAvailableContextTokens,
|
||||
} from "@/app/app/services/lib";
|
||||
import { personaIncludesRetrieval } from "@/app/app/services/lib";
|
||||
import { useCallback, useEffect, useMemo, useRef, useState } from "react";
|
||||
import { toast } from "@/hooks/useToast";
|
||||
import { SEARCH_PARAM_NAMES } from "@/app/app/services/searchParams";
|
||||
@@ -56,10 +53,7 @@ import ChatScrollContainer, {
|
||||
} from "@/sections/chat/ChatScrollContainer";
|
||||
import ProjectContextPanel from "@/app/app/components/projects/ProjectContextPanel";
|
||||
import { useProjectsContext } from "@/providers/ProjectsContext";
|
||||
import {
|
||||
getProjectTokenCount,
|
||||
getMaxSelectedDocumentTokens,
|
||||
} from "@/app/app/projects/projectsService";
|
||||
import { getProjectTokenCount } from "@/app/app/projects/projectsService";
|
||||
import ProjectChatSessionList from "@/app/app/components/projects/ProjectChatSessionList";
|
||||
import { cn } from "@/lib/utils";
|
||||
import Suggestions from "@/sections/Suggestions";
|
||||
@@ -70,7 +64,6 @@ import * as AppLayouts from "@/layouts/app-layouts";
|
||||
import { SvgChevronDown, SvgFileText } from "@opal/icons";
|
||||
import { Button } from "@opal/components";
|
||||
import Spacer from "@/refresh-components/Spacer";
|
||||
import { DEFAULT_CONTEXT_TOKENS } from "@/lib/constants";
|
||||
import useAppFocus from "@/hooks/useAppFocus";
|
||||
import { useQueryController } from "@/providers/QueryControllerProvider";
|
||||
import WelcomeMessage from "@/app/app/components/WelcomeMessage";
|
||||
@@ -367,18 +360,22 @@ export default function AppPage({ firstMessage }: ChatPageProps) {
|
||||
const autoScrollEnabled = user?.preferences?.auto_scroll !== false;
|
||||
const isStreaming = currentChatState === "streaming";
|
||||
|
||||
const { onSubmit, stopGenerating, handleMessageSpecificFileUpload } =
|
||||
useChatController({
|
||||
filterManager,
|
||||
llmManager,
|
||||
availableAgents: agents,
|
||||
liveAgent,
|
||||
existingChatSessionId: currentChatSessionId,
|
||||
selectedDocuments,
|
||||
searchParams,
|
||||
resetInputBar,
|
||||
setSelectedAgentFromId,
|
||||
});
|
||||
const {
|
||||
onSubmit,
|
||||
stopGenerating,
|
||||
handleMessageSpecificFileUpload,
|
||||
availableContextTokens,
|
||||
} = useChatController({
|
||||
filterManager,
|
||||
llmManager,
|
||||
availableAgents: agents,
|
||||
liveAgent,
|
||||
existingChatSessionId: currentChatSessionId,
|
||||
selectedDocuments,
|
||||
searchParams,
|
||||
resetInputBar,
|
||||
setSelectedAgentFromId,
|
||||
});
|
||||
|
||||
const { onMessageSelection, currentSessionFileTokenCount } =
|
||||
useChatSessionController({
|
||||
@@ -595,43 +592,6 @@ export default function AppPage({ firstMessage }: ChatPageProps) {
|
||||
};
|
||||
}, [currentChatSessionId, currentProjectId, currentProjectDetails?.files]);
|
||||
|
||||
// Available context tokens source of truth:
|
||||
// - If a chat session exists, fetch from session API (dynamic per session/model)
|
||||
// - If no session, derive from the default/current persona's max document tokens
|
||||
const [availableContextTokens, setAvailableContextTokens] = useState<number>(
|
||||
DEFAULT_CONTEXT_TOKENS * 0.5
|
||||
);
|
||||
useEffect(() => {
|
||||
let cancelled = false;
|
||||
async function run() {
|
||||
try {
|
||||
if (currentChatSessionId) {
|
||||
const available =
|
||||
await getAvailableContextTokens(currentChatSessionId);
|
||||
const capped_context_tokens =
|
||||
(available ?? DEFAULT_CONTEXT_TOKENS) * 0.5;
|
||||
if (!cancelled) setAvailableContextTokens(capped_context_tokens);
|
||||
} else {
|
||||
const personaId = (selectedAgent || liveAgent)?.id;
|
||||
if (personaId !== undefined && personaId !== null) {
|
||||
const maxTokens = await getMaxSelectedDocumentTokens(personaId);
|
||||
const capped_context_tokens =
|
||||
(maxTokens ?? DEFAULT_CONTEXT_TOKENS) * 0.5;
|
||||
if (!cancelled) setAvailableContextTokens(capped_context_tokens);
|
||||
} else if (!cancelled) {
|
||||
setAvailableContextTokens(DEFAULT_CONTEXT_TOKENS * 0.5);
|
||||
}
|
||||
}
|
||||
} catch (e) {
|
||||
if (!cancelled) setAvailableContextTokens(DEFAULT_CONTEXT_TOKENS * 0.5);
|
||||
}
|
||||
}
|
||||
run();
|
||||
return () => {
|
||||
cancelled = true;
|
||||
};
|
||||
}, [currentChatSessionId, selectedAgent?.id, liveAgent?.id]);
|
||||
|
||||
// handle error case where no assistants are available
|
||||
// Only show this after agents have loaded to prevent flash during initial load
|
||||
if (noAgents && !isLoadingAgents) {
|
||||
|
||||
@@ -21,10 +21,47 @@ type DefaultModelInfo = {
|
||||
model_name: string;
|
||||
} | null;
|
||||
|
||||
type ProviderModelConfig = {
|
||||
name: string;
|
||||
is_visible: boolean;
|
||||
};
|
||||
|
||||
function uniqueName(prefix: string): string {
|
||||
return `${prefix}-${Date.now()}-${Math.random().toString(36).slice(2, 8)}`;
|
||||
}
|
||||
|
||||
function normalizeAlphaNum(input: string): string {
|
||||
return input.toLowerCase().replace(/[^a-z0-9]/g, "");
|
||||
}
|
||||
|
||||
function modelTokenVariants(modelName: string): string[][] {
|
||||
return modelName
|
||||
.toLowerCase()
|
||||
.split(/[^a-z0-9]+/)
|
||||
.filter((token) => token.length > 0)
|
||||
.map((token) => {
|
||||
// Display names may shorten long numeric segments to suffixes.
|
||||
if (/^\d+$/.test(token) && token.length > 5) {
|
||||
return [token, token.slice(-5)];
|
||||
}
|
||||
return [token];
|
||||
});
|
||||
}
|
||||
|
||||
function textMatchesModel(modelName: string, candidateText: string): boolean {
|
||||
const normalizedCandidate = normalizeAlphaNum(candidateText);
|
||||
if (!normalizedCandidate) {
|
||||
return false;
|
||||
}
|
||||
|
||||
const tokenVariants = modelTokenVariants(modelName);
|
||||
return tokenVariants.every((variants) =>
|
||||
variants.some((variant) =>
|
||||
normalizedCandidate.includes(normalizeAlphaNum(variant))
|
||||
)
|
||||
);
|
||||
}
|
||||
|
||||
async function getAdminLLMProviderResponse(page: Page) {
|
||||
const response = await page.request.get(`${BASE_URL}/api/admin/llm/provider`);
|
||||
expect(response.ok()).toBeTruthy();
|
||||
@@ -50,6 +87,18 @@ async function createPublicProvider(
|
||||
providerName: string,
|
||||
modelName: string = "gpt-4o"
|
||||
): Promise<number> {
|
||||
return createPublicProviderWithModels(page, providerName, [
|
||||
{ name: modelName, is_visible: true },
|
||||
]);
|
||||
}
|
||||
|
||||
async function createPublicProviderWithModels(
|
||||
page: Page,
|
||||
providerName: string,
|
||||
modelConfigurations: ProviderModelConfig[]
|
||||
): Promise<number> {
|
||||
expect(modelConfigurations.length).toBeGreaterThan(0);
|
||||
|
||||
const response = await page.request.put(
|
||||
`${BASE_URL}/api/admin/llm/provider?is_creation=true`,
|
||||
{
|
||||
@@ -60,7 +109,7 @@ async function createPublicProvider(
|
||||
is_public: true,
|
||||
groups: [],
|
||||
personas: [],
|
||||
model_configurations: [{ name: modelName, is_visible: true }],
|
||||
model_configurations: modelConfigurations,
|
||||
},
|
||||
}
|
||||
);
|
||||
@@ -69,6 +118,86 @@ async function createPublicProvider(
|
||||
return data.id;
|
||||
}
|
||||
|
||||
async function navigateToAdminLlmPageFromChat(page: Page): Promise<void> {
|
||||
await page.goto(LLM_SETUP_URL);
|
||||
await page.waitForURL("**/admin/configuration/llm**");
|
||||
await expect(page.getByLabel("admin-page-title")).toHaveText(
|
||||
/^Language Models/
|
||||
);
|
||||
}
|
||||
|
||||
async function exitAdminToChat(page: Page): Promise<void> {
|
||||
await page.goto("/app");
|
||||
await page.waitForURL("**/app**");
|
||||
await page
|
||||
.locator("#onyx-chat-input-textarea")
|
||||
.waitFor({ state: "visible", timeout: 15000 });
|
||||
}
|
||||
|
||||
async function isModelVisibleInChatProviders(
|
||||
page: Page,
|
||||
modelName: string
|
||||
): Promise<boolean> {
|
||||
const response = await page.request.get(`${BASE_URL}/api/llm/provider`);
|
||||
expect(response.ok()).toBeTruthy();
|
||||
|
||||
const data = (await response.json()) as {
|
||||
providers: {
|
||||
model_configurations: { name: string; is_visible: boolean }[];
|
||||
}[];
|
||||
};
|
||||
|
||||
return data.providers.some((provider) =>
|
||||
provider.model_configurations.some(
|
||||
(model) => model.name === modelName && model.is_visible
|
||||
)
|
||||
);
|
||||
}
|
||||
|
||||
async function expectModelVisibilityInChatProviders(
|
||||
page: Page,
|
||||
modelName: string,
|
||||
expectedVisible: boolean
|
||||
): Promise<void> {
|
||||
await expect
|
||||
.poll(() => isModelVisibleInChatProviders(page, modelName), {
|
||||
timeout: 30000,
|
||||
})
|
||||
.toBe(expectedVisible);
|
||||
}
|
||||
|
||||
async function getModelCountInChatSelector(
|
||||
page: Page,
|
||||
modelName: string
|
||||
): Promise<number> {
|
||||
const dialog = page.locator('[role="dialog"]').first();
|
||||
|
||||
// When used in expect.poll retries, a previous attempt may leave the
|
||||
// popover open. Ensure a clean state before toggling it.
|
||||
if (await dialog.isVisible()) {
|
||||
await page.keyboard.press("Escape");
|
||||
await dialog.waitFor({ state: "hidden", timeout: 5000 });
|
||||
}
|
||||
|
||||
await page.getByTestId("AppInputBar/llm-popover-trigger").click();
|
||||
await dialog.waitFor({ state: "visible", timeout: 10000 });
|
||||
|
||||
await dialog.getByPlaceholder("Search models...").fill(modelName);
|
||||
const optionButtons = dialog.getByRole("button");
|
||||
const optionTexts = await optionButtons.allTextContents();
|
||||
const uniqueOptionTexts = Array.from(
|
||||
new Set(optionTexts.map((text) => text.trim()))
|
||||
);
|
||||
const count = uniqueOptionTexts.filter((text) =>
|
||||
textMatchesModel(modelName, text)
|
||||
).length;
|
||||
|
||||
await page.keyboard.press("Escape");
|
||||
await dialog.waitFor({ state: "hidden", timeout: 10000 });
|
||||
|
||||
return count;
|
||||
}
|
||||
|
||||
async function getProviderByName(
|
||||
page: Page,
|
||||
providerName: string
|
||||
@@ -272,4 +401,132 @@ test.describe("LLM Provider Setup @exclusive", () => {
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
test("adding a hidden model on an existing provider shows it in chat after one save", async ({
|
||||
page,
|
||||
}) => {
|
||||
await page.route("**/api/admin/llm/test", async (route) => {
|
||||
await route.fulfill({
|
||||
status: 200,
|
||||
contentType: "application/json",
|
||||
body: JSON.stringify({ success: true }),
|
||||
});
|
||||
});
|
||||
|
||||
const providerName = uniqueName("PW Provider Add Model");
|
||||
const ts = Date.now();
|
||||
const alwaysVisibleModel = `pw-visible-${ts}-base`;
|
||||
const modelToEnable = `pw-hidden-${ts}-to-enable`;
|
||||
|
||||
const providerId = await createPublicProviderWithModels(
|
||||
page,
|
||||
providerName,
|
||||
[
|
||||
{ name: alwaysVisibleModel, is_visible: true },
|
||||
{ name: modelToEnable, is_visible: false },
|
||||
]
|
||||
);
|
||||
providersToCleanup.push(providerId);
|
||||
await expectModelVisibilityInChatProviders(page, modelToEnable, false);
|
||||
|
||||
await page.goto("/app");
|
||||
await page.waitForLoadState("networkidle");
|
||||
await page
|
||||
.locator("#onyx-chat-input-textarea")
|
||||
.waitFor({ state: "visible", timeout: 15000 });
|
||||
|
||||
await expect
|
||||
.poll(() => getModelCountInChatSelector(page, modelToEnable), {
|
||||
timeout: 15000,
|
||||
})
|
||||
.toBe(0);
|
||||
|
||||
await navigateToAdminLlmPageFromChat(page);
|
||||
|
||||
const editModal = await openProviderEditModal(page, providerName);
|
||||
await editModal.getByText(modelToEnable, { exact: true }).click();
|
||||
|
||||
const updateButton = editModal.getByRole("button", { name: "Update" });
|
||||
const providerUpdateResponsePromise = page.waitForResponse(
|
||||
(response) =>
|
||||
response.url().includes("/api/admin/llm/provider") &&
|
||||
response.request().method() === "PUT"
|
||||
);
|
||||
await expect(updateButton).toBeEnabled({ timeout: 10000 });
|
||||
await updateButton.click();
|
||||
await providerUpdateResponsePromise;
|
||||
await expect(editModal).not.toBeVisible({ timeout: 30000 });
|
||||
await expectModelVisibilityInChatProviders(page, modelToEnable, true);
|
||||
|
||||
await exitAdminToChat(page);
|
||||
await expect
|
||||
.poll(() => getModelCountInChatSelector(page, modelToEnable), {
|
||||
timeout: 15000,
|
||||
})
|
||||
.toBe(1);
|
||||
});
|
||||
|
||||
test("removing a visible model on an existing provider hides it in chat after one save", async ({
|
||||
page,
|
||||
}) => {
|
||||
await page.route("**/api/admin/llm/test", async (route) => {
|
||||
await route.fulfill({
|
||||
status: 200,
|
||||
contentType: "application/json",
|
||||
body: JSON.stringify({ success: true }),
|
||||
});
|
||||
});
|
||||
|
||||
const providerName = uniqueName("PW Provider Remove Model");
|
||||
const ts = Date.now();
|
||||
const alwaysVisibleModel = `pw-visible-${ts}-base`;
|
||||
const modelToDisable = `pw-visible-${ts}-to-disable`;
|
||||
|
||||
const providerId = await createPublicProviderWithModels(
|
||||
page,
|
||||
providerName,
|
||||
[
|
||||
{ name: alwaysVisibleModel, is_visible: true },
|
||||
{ name: modelToDisable, is_visible: true },
|
||||
]
|
||||
);
|
||||
providersToCleanup.push(providerId);
|
||||
await expectModelVisibilityInChatProviders(page, modelToDisable, true);
|
||||
|
||||
await page.goto("/app");
|
||||
await page.waitForLoadState("networkidle");
|
||||
await page
|
||||
.locator("#onyx-chat-input-textarea")
|
||||
.waitFor({ state: "visible", timeout: 15000 });
|
||||
|
||||
await expect
|
||||
.poll(() => getModelCountInChatSelector(page, modelToDisable), {
|
||||
timeout: 15000,
|
||||
})
|
||||
.toBe(1);
|
||||
|
||||
await navigateToAdminLlmPageFromChat(page);
|
||||
|
||||
const editModal = await openProviderEditModal(page, providerName);
|
||||
await editModal.getByText(modelToDisable, { exact: true }).click();
|
||||
|
||||
const updateButton = editModal.getByRole("button", { name: "Update" });
|
||||
const providerUpdateResponsePromise = page.waitForResponse(
|
||||
(response) =>
|
||||
response.url().includes("/api/admin/llm/provider") &&
|
||||
response.request().method() === "PUT"
|
||||
);
|
||||
await expect(updateButton).toBeEnabled({ timeout: 10000 });
|
||||
await updateButton.click();
|
||||
await providerUpdateResponsePromise;
|
||||
await expect(editModal).not.toBeVisible({ timeout: 30000 });
|
||||
await expectModelVisibilityInChatProviders(page, modelToDisable, false);
|
||||
|
||||
await exitAdminToChat(page);
|
||||
await expect
|
||||
.poll(() => getModelCountInChatSelector(page, modelToDisable), {
|
||||
timeout: 15000,
|
||||
})
|
||||
.toBe(0);
|
||||
});
|
||||
});
|
||||
|
||||
Reference in New Issue
Block a user