diff --git a/README.md b/README.md index fd52aaff..ab68136d 100644 --- a/README.md +++ b/README.md @@ -357,7 +357,7 @@ Bootstrap settings use the `BIGRAG_` prefix as environment variables, or configu | Variable | Description | Default | |----------|-------------|---------| | `BIGRAG_CHAT_PROVIDER` | Chat provider | `openai` | -| `BIGRAG_CHAT_MODEL` | Default chat model | `gpt-4o-mini` | +| `BIGRAG_CHAT_MODEL` | Default chat model | `gpt-4.1` | | `BIGRAG_CHAT_BASE_URL` | Base URL for OpenAI-compatible chat endpoints | — | | `BIGRAG_CHAT_TEMPERATURE` | Default chat temperature | `0.2` | | `BIGRAG_CHAT_MAX_CONTEXT_CHARS` | Max retrieved-context characters per chat call | `120000` | diff --git a/api/bigrag/app_factory/routers.py b/api/bigrag/app_factory/routers.py index 552d3cdb..2996962a 100644 --- a/api/bigrag/app_factory/routers.py +++ b/api/bigrag/app_factory/routers.py @@ -27,7 +27,7 @@ def include_all_routers(app: FastAPI) -> None: from bigrag.routers.mcp_servers import router as mcp_servers_router from bigrag.routers.preferences import router as preferences_router from bigrag.routers.query import router as query_router - from bigrag.routers.realtime import router as realtime_router + from bigrag.routers.status import router as status_router from bigrag.routers.upload_sessions import router as upload_sessions_router from bigrag.routers.usage import router as usage_router from bigrag.routers.vectors import router as vectors_router @@ -43,7 +43,7 @@ def include_all_routers(app: FastAPI) -> None: app.include_router(admin_vector_storage_router) app.include_router(mcp_servers_router) app.include_router(admin_audit_router) - app.include_router(realtime_router) + app.include_router(status_router) app.include_router(embedding_presets_router) app.include_router(collections_router) app.include_router(connectors_router) diff --git a/api/bigrag/config.py b/api/bigrag/config.py index 065ddb19..48dc770a 100644 --- a/api/bigrag/config.py +++ b/api/bigrag/config.py @@ -68,7 +68,7 @@ class Settings(BaseSettings): allow_local_webhooks: bool = False chat_provider: str = "openai" - chat_model: str = "gpt-4o-mini" + chat_model: str = "gpt-4.1" chat_base_url: str | None = None chat_temperature: float = 0.2 chat_max_context_chars: int = 120_000 diff --git a/api/bigrag/models/status.py b/api/bigrag/models/status.py new file mode 100644 index 00000000..3ed9ea28 --- /dev/null +++ b/api/bigrag/models/status.py @@ -0,0 +1,20 @@ +from __future__ import annotations + +from pydantic import BaseModel + + +class OverviewStatusResponse(BaseModel): + platform: dict[str, object] + readiness: dict[str, object] + + +class CollectionsStatusResponse(BaseModel): + collections_total: int + documents_total: int + documents_ready: int + documents_pending: int + documents_processing: int + documents_failed: int + total_chunks: int + total_tokens: int + total_size_bytes: int diff --git a/api/bigrag/routers/realtime.py b/api/bigrag/routers/realtime.py deleted file mode 100644 index a50433e4..00000000 --- a/api/bigrag/routers/realtime.py +++ /dev/null @@ -1,47 +0,0 @@ -from __future__ import annotations - -import sqlalchemy as sa -from fastapi import APIRouter, Depends, HTTPException, WebSocket -from sqlalchemy.ext.asyncio import AsyncSession - -from bigrag.db.models import Collection -from bigrag.db.session import get_session -from bigrag.middleware.auth import get_current_user -from bigrag.services.error_sanitize import safe_error_detail -from bigrag.services.realtime import RealtimeConnection -from bigrag.services.realtime.auth import connection_principal, cookie_origin_allowed -from bigrag.services.realtime_tokens import ( - REALTIME_TOKEN_TTL_SECONDS, - create_realtime_token, -) - -router = APIRouter(tags=["realtime"]) - - -@router.websocket("/v1/realtime") -async def realtime_socket(websocket: WebSocket): - if not await cookie_origin_allowed(websocket): - await websocket.close(code=1008) - return - principal = await connection_principal(websocket) - await websocket.accept() - await RealtimeConnection(websocket, principal).run() - - -@router.post("/v1/collections/{name}/realtime-token", response_model=dict[str, str | int]) -async def create_collection_realtime_token( - name: str, - user: dict = Depends(get_current_user), - session: AsyncSession = Depends(get_session), -) -> dict[str, str | int]: - exists = await session.scalar(sa.select(Collection.id).where(Collection.name == name)) - if exists is None: - raise HTTPException(status_code=404, detail="Collection not found") - try: - token = await create_realtime_token(user, name) - except RuntimeError as exc: - raise HTTPException( - status_code=503, - detail=safe_error_detail(exc, "Realtime tokens are unavailable; check Redis."), - ) from exc - return {"token": token, "expires_in": REALTIME_TOKEN_TTL_SECONDS} diff --git a/api/bigrag/routers/status.py b/api/bigrag/routers/status.py new file mode 100644 index 00000000..fafafecb --- /dev/null +++ b/api/bigrag/routers/status.py @@ -0,0 +1,54 @@ +from __future__ import annotations + +from fastapi import APIRouter, Depends, Query, Request +from sqlalchemy.ext.asyncio import AsyncSession + +from bigrag.db.session import get_session +from bigrag.middleware.auth import get_current_user, require_admin_session +from bigrag.models.access import AccessLogOverviewResponse +from bigrag.models.status import CollectionsStatusResponse, OverviewStatusResponse +from bigrag.services.access_log.queries import access_overview_payload +from bigrag.services.health import readiness_payload +from bigrag.services.platform_stats import platform_stats_payload +from bigrag.services.status import collections_status_payload +from bigrag.services.usage import UsageResponse, usage_payload + +router = APIRouter(tags=["status"]) + + +@router.get("/v1/status/overview", response_model=OverviewStatusResponse) +async def overview_status( + request: Request, + _: dict = Depends(get_current_user), + session: AsyncSession = Depends(get_session), +) -> OverviewStatusResponse: + return OverviewStatusResponse( + platform=await platform_stats_payload(request.app.state.queue, session), + readiness=await readiness_payload(request.app.state.vector_store, request.app.state.queue), + ) + + +@router.get("/v1/status/collections", response_model=CollectionsStatusResponse) +async def collections_status( + _: dict = Depends(get_current_user), + session: AsyncSession = Depends(get_session), +) -> CollectionsStatusResponse: + return await collections_status_payload(session) + + +@router.get("/v1/status/usage", response_model=UsageResponse) +async def usage_status( + window_days: int = Query(default=30, ge=1, le=365), + _: dict = Depends(get_current_user), + session: AsyncSession = Depends(get_session), +) -> UsageResponse: + return await usage_payload(session, window_days=window_days) + + +@router.get("/v1/admin/status/access", response_model=AccessLogOverviewResponse) +async def access_status( + window_days: int = Query(default=7, ge=1, le=90), + _: dict = Depends(require_admin_session), + session: AsyncSession = Depends(get_session), +) -> AccessLogOverviewResponse: + return await access_overview_payload(session, window_days=window_days) diff --git a/api/bigrag/services/access_log/queries.py b/api/bigrag/services/access_log/queries.py index 3fac5ebf..79af3881 100644 --- a/api/bigrag/services/access_log/queries.py +++ b/api/bigrag/services/access_log/queries.py @@ -5,7 +5,6 @@ import sqlalchemy as sa from fastapi import HTTPException -from fastapi.encoders import jsonable_encoder from sqlalchemy.ext.asyncio import AsyncSession from bigrag.db.models import AccessLog @@ -16,12 +15,10 @@ AccessLogOverviewResponse, AccessLogTimelinePoint, ) -from bigrag.services import redis_cache from bigrag.services.access_log.middleware import RAG_ACCESS_ACTIONS from bigrag.services.pagination import paginate _RAG_ACTION_FILTER = AccessLog.action.in_(tuple(sorted(RAG_ACCESS_ACTIONS))) -_ACCESS_OVERVIEW_TTL = 15 def access_log_entry(row: AccessLog) -> AccessLogEntry: @@ -155,11 +152,6 @@ async def access_overview_payload( *, window_days: int, ) -> AccessLogOverviewResponse: - cache_key = f"access:overview:{window_days}" - cached = await redis_cache.get(cache_key) - if cached: - return AccessLogOverviewResponse.model_validate(cached) - filters = [_RAG_ACTION_FILTER, _window_filter(window_days)] summary = ( @@ -204,7 +196,7 @@ async def access_overview_payload( total = int(summary.total or 0) errors = int(summary.errors or 0) successes = int(summary.successes or 0) - response = AccessLogOverviewResponse( + return AccessLogOverviewResponse( window_days=window_days, total_events=total, success_rate=round((successes / total) * 100, 2) if total else 0, @@ -241,5 +233,3 @@ async def access_overview_payload( ], recent=[access_log_entry(row) for row in recent], ) - await redis_cache.set(cache_key, jsonable_encoder(response), ttl=_ACCESS_OVERVIEW_TTL) - return response diff --git a/api/bigrag/services/collection_scope.py b/api/bigrag/services/collection_scope.py index dd4250ed..1f893f1a 100644 --- a/api/bigrag/services/collection_scope.py +++ b/api/bigrag/services/collection_scope.py @@ -13,6 +13,9 @@ ("GET", "/v1/collections"), ("GET", "/v1/usage"), ("GET", "/v1/stats"), + ("GET", "/v1/status/overview"), + ("GET", "/v1/status/collections"), + ("GET", "/v1/status/usage"), ("GET", "/v1/embeddings/models"), ) diff --git a/api/bigrag/services/connectors/progress.py b/api/bigrag/services/connectors/progress.py index 271585f2..52e10583 100644 --- a/api/bigrag/services/connectors/progress.py +++ b/api/bigrag/services/connectors/progress.py @@ -3,7 +3,6 @@ from typing import Any from bigrag.db.models import ConnectorDocument, ConnectorSource, ConnectorSyncJob -from bigrag.services.connectors.realtime import notify_connector_sync_jobs from bigrag.services.connectors.types import ConnectorSyncCounters, RemoteConnectorFile SYNC_PROGRESS_FIXED_PERCENT = { @@ -98,5 +97,3 @@ async def update_sync_progress( ), } await session.commit() - if source is not None: - notify_connector_sync_jobs(job.provider, source.collection_name, str(source.id)) diff --git a/api/bigrag/services/connectors/realtime.py b/api/bigrag/services/connectors/realtime.py deleted file mode 100644 index 805c5b24..00000000 --- a/api/bigrag/services/connectors/realtime.py +++ /dev/null @@ -1,48 +0,0 @@ -from __future__ import annotations - -from bigrag.services.event_bus import event_bus - - -def connector_sources_event_key(provider: str, collection_name: str | None = None) -> str: - return f"connector:{provider}:sources:{collection_name or 'all'}" - - -def connector_sync_jobs_event_key( - provider: str, - collection_name: str | None = None, - source_id: str | None = None, -) -> str: - return f"connector:{provider}:sync-jobs:{collection_name or 'all'}:{source_id or 'all'}" - - -def notify_connector_sources(provider: str, collection_name: str | None = None) -> None: - keys = {connector_sources_event_key(provider)} - if collection_name: - keys.add(connector_sources_event_key(provider, collection_name)) - for key in keys: - event_bus.notify(key) - - -def notify_connector_sync_jobs( - provider: str, - collection_name: str | None = None, - source_id: str | None = None, -) -> None: - keys = {connector_sync_jobs_event_key(provider)} - if collection_name: - keys.add(connector_sync_jobs_event_key(provider, collection_name)) - if source_id: - keys.add(connector_sync_jobs_event_key(provider, None, source_id)) - if collection_name and source_id: - keys.add(connector_sync_jobs_event_key(provider, collection_name, source_id)) - for key in keys: - event_bus.notify(key) - - -def notify_connector_state( - provider: str, - collection_name: str | None = None, - source_id: str | None = None, -) -> None: - notify_connector_sources(provider, collection_name) - notify_connector_sync_jobs(provider, collection_name, source_id) diff --git a/api/bigrag/services/connectors/scheduler.py b/api/bigrag/services/connectors/scheduler.py index 30e62aab..35098b21 100644 --- a/api/bigrag/services/connectors/scheduler.py +++ b/api/bigrag/services/connectors/scheduler.py @@ -7,7 +7,6 @@ from bigrag.db.engine import session_factory from bigrag.db.models import ConnectorSource, ConnectorSyncJob -from bigrag.services.connectors.realtime import notify_connector_state from bigrag.services.connectors.sources import create_sync_job from bigrag.services.connectors.time import next_sync_at, utcnow @@ -58,7 +57,6 @@ async def run_due_syncs( async with session_factory()() as session: await reap_stale_syncs(session, provider=provider) job_ids: list[str] = [] - notifications: list[tuple[str, str, str]] = [] async with session_factory()() as session: rows = ( await session.scalars( @@ -84,12 +82,9 @@ async def run_due_syncs( commit=False, ) await session.flush() - notifications.append((provider, source.collection_name, str(source.id))) if job.status == "pending" and job.started_at is None: job_ids.append(str(job.id)) await session.commit() - for event in notifications: - notify_connector_state(*event) for job_id in job_ids: start_sync_job(job_id) return len(job_ids) diff --git a/api/bigrag/services/connectors/sources/sources_mutations.py b/api/bigrag/services/connectors/sources/sources_mutations.py index 221d8587..fb62fb82 100644 --- a/api/bigrag/services/connectors/sources/sources_mutations.py +++ b/api/bigrag/services/connectors/sources/sources_mutations.py @@ -15,7 +15,6 @@ ) from bigrag.services import collection_cache from bigrag.services.connectors.progress import sync_progress_details -from bigrag.services.connectors.realtime import notify_connector_sources, notify_connector_state from bigrag.services.connectors.sources.sources_credentials import upsert_source_credential from bigrag.services.connectors.sources.sources_queries import source_by_id from bigrag.services.connectors.time import next_sync_at, utcnow @@ -146,7 +145,6 @@ async def create_source( ) await session.commit() await session.refresh(existing) - notify_connector_state(provider, existing.collection_name, str(existing.id)) if job.status == "pending" and job.started_at is None: start_sync_job(str(job.id)) return existing, job @@ -164,7 +162,6 @@ async def create_source( await session.commit() await session.refresh(source) await session.refresh(job) - notify_connector_state(provider, source.collection_name, str(source.id)) if job.status == "pending" and job.started_at is None: start_sync_job(str(job.id)) return source, job @@ -197,7 +194,6 @@ async def trigger_sync( ) if job.status == "pending" and job.started_at is None: start_sync_job(str(job.id)) - notify_connector_state(provider, source.collection_name, str(source.id)) return job @@ -235,7 +231,6 @@ async def update_source( source.next_sync_at = next_sync_at(source) await session.commit() await session.refresh(source) - notify_connector_sources(provider, source.collection_name) return source @@ -253,7 +248,6 @@ async def delete_source( not_found_message=not_found_message, ) collection_name = source.collection_name - source_uuid = str(source.id) manifests = ( await session.scalars( sa.select(ConnectorDocument).where(ConnectorDocument.source_id == source.id) @@ -277,7 +271,6 @@ async def delete_source( await session.delete(source) await session.commit() await collection_cache.invalidate(collection_name) - notify_connector_state(provider, collection_name, source_uuid) def _tenant_id(collection: Collection, metadata: dict[str, Any]) -> str | None: diff --git a/api/bigrag/services/connectors/status.py b/api/bigrag/services/connectors/status.py index b9532e19..2b42391f 100644 --- a/api/bigrag/services/connectors/status.py +++ b/api/bigrag/services/connectors/status.py @@ -5,7 +5,6 @@ from bigrag.db.models import ConnectorSource, ConnectorSyncJob from bigrag.services.connectors.manifest import apply_counters from bigrag.services.connectors.progress import sync_counter_details, update_sync_progress -from bigrag.services.connectors.realtime import notify_connector_sources from bigrag.services.connectors.time import next_sync_at, utcnow from bigrag.services.connectors.types import ConnectorSyncCounters from bigrag.services.webhook import enqueue_webhook_event @@ -37,7 +36,6 @@ async def fail_sync( phase="failed", message=message, ) - notify_connector_sources(job.provider, source.collection_name) data = { "provider": job.provider, "source_id": str(source.id), diff --git a/api/bigrag/services/connectors/sync/finalize.py b/api/bigrag/services/connectors/sync/finalize.py index b97acca0..ef765489 100644 --- a/api/bigrag/services/connectors/sync/finalize.py +++ b/api/bigrag/services/connectors/sync/finalize.py @@ -5,7 +5,6 @@ from bigrag.services import collection_cache from bigrag.services.connectors.manifest import apply_counters from bigrag.services.connectors.progress import sync_counter_details, update_sync_progress -from bigrag.services.connectors.realtime import notify_connector_sources from bigrag.services.connectors.time import next_sync_at, utcnow from bigrag.services.connectors.types import ConnectorSyncAdapter, ConnectorSyncCounters from bigrag.services.documents import recount_collection_documents @@ -58,7 +57,6 @@ async def finalize_sync( processed_items=counters.found + counters.deleted, total_items=counters.found + missing_count, ) - notify_connector_sources(adapter.provider, source.collection_name) await collection_cache.invalidate(source.collection_name) await invalidate_collection_query_cache(source.collection_name) webhook_event = ( diff --git a/api/bigrag/services/connectors/sync/lifecycle.py b/api/bigrag/services/connectors/sync/lifecycle.py index d7b7fe86..47a083cb 100644 --- a/api/bigrag/services/connectors/sync/lifecycle.py +++ b/api/bigrag/services/connectors/sync/lifecycle.py @@ -4,7 +4,6 @@ from bigrag.services.connectors.manifest import apply_counters from bigrag.services.connectors.progress import update_sync_progress -from bigrag.services.connectors.realtime import notify_connector_sources from bigrag.services.connectors.time import next_sync_at, utcnow from bigrag.services.connectors.types import ConnectorSyncAdapter, ConnectorSyncCounters from bigrag.services.webhook import enqueue_webhook_event @@ -70,4 +69,3 @@ async def mark_sync_deferred_by_queue( phase="complete", message="Ingestion queue full; remaining files will sync on the next run.", ) - notify_connector_sources(adapter.provider, source.collection_name) diff --git a/api/bigrag/services/embedding/__init__.py b/api/bigrag/services/embedding/__init__.py index 01bf96cd..1a00ae1e 100644 --- a/api/bigrag/services/embedding/__init__.py +++ b/api/bigrag/services/embedding/__init__.py @@ -2,7 +2,6 @@ from bigrag.services.embedding.base import ( EmbeddingModel, - reset_embedding_semaphores, truncate_to_tokens, ) from bigrag.services.embedding.cohere import CohereEmbedding @@ -13,6 +12,7 @@ get_embedding_model, ) from bigrag.services.embedding.voyage import VoyageEmbedding +from bigrag.services.embedding_gate import reset_embedding_limiters __all__ = [ "AVAILABLE_MODELS", @@ -22,6 +22,6 @@ "VoyageEmbedding", "close_embedding_models", "get_embedding_model", - "reset_embedding_semaphores", + "reset_embedding_limiters", "truncate_to_tokens", ] diff --git a/api/bigrag/services/embedding/base.py b/api/bigrag/services/embedding/base.py index 78729e07..607170eb 100644 --- a/api/bigrag/services/embedding/base.py +++ b/api/bigrag/services/embedding/base.py @@ -1,15 +1,11 @@ from __future__ import annotations -import asyncio from abc import ABC, abstractmethod from bigrag.logging import get_logger logger = get_logger("bigrag.embedding") -_embed_semaphores: dict[str, asyncio.Semaphore] = {} -_embed_semaphores_lock = asyncio.Lock() - _TOKEN_LIMITS: dict[str, int] = { "text-embedding-3-small": 8000, "text-embedding-3-large": 8000, @@ -27,21 +23,6 @@ } -async def get_semaphore(key: str) -> asyncio.Semaphore: - if key in _embed_semaphores: - return _embed_semaphores[key] - async with _embed_semaphores_lock: - if key not in _embed_semaphores: - from bigrag.services.runtime_settings import sync_value - - _embed_semaphores[key] = asyncio.Semaphore(sync_value("embedding_concurrency")) - return _embed_semaphores[key] - - -def reset_embedding_semaphores() -> None: - _embed_semaphores.clear() - - def truncate_to_tokens( texts: list[str], model: str | None, diff --git a/api/bigrag/services/embedding/cohere.py b/api/bigrag/services/embedding/cohere.py index 103cb613..c5c45a46 100644 --- a/api/bigrag/services/embedding/cohere.py +++ b/api/bigrag/services/embedding/cohere.py @@ -2,14 +2,8 @@ import asyncio -from bigrag.services.embedding.base import EmbeddingModel, get_semaphore, logger, truncate_to_tokens -from bigrag.services.embedding_rate_limit import ( - is_rate_limit_error, - rate_limit_cooldown_key, - rate_limit_delay, - record_rate_limit_cooldown, - wait_for_rate_limit_cooldown, -) +from bigrag.services.embedding.base import EmbeddingModel, logger, truncate_to_tokens +from bigrag.services.embedding_gate import embedding_gate class CohereEmbedding(EmbeddingModel): @@ -35,7 +29,6 @@ def __init__( self._model_name = model_name self._dimension = dimension - self._semaphore_key = "cohere" self._cache_identity = f"cohere:{model_name}:{dimension}" self._client = cohere.AsyncClient(api_key=api_key) logger.info("initialized cohere embedding", model=model_name, dimension=dimension) @@ -66,25 +59,16 @@ async def embed(self, texts: list[str], *, input_type: str = "document") -> list return await self._embed_single(texts, cohere_input_type) async def _embed_single(self, texts: list[str], cohere_input_type: str) -> list[list[float]]: - cooldown_key = rate_limit_cooldown_key( - self._cache_identity, self.provider, self._model_name, self._dimension - ) - async with await get_semaphore(self._semaphore_key): - await wait_for_rate_limit_cooldown(cooldown_key, self.provider, self._model_name) - try: - response = await asyncio.wait_for( - self._client.embed( - texts=texts, - model=self._model_name, - input_type=cohere_input_type, - embedding_types=["float"], - ), - timeout=60, - ) - except Exception as exc: - if is_rate_limit_error(exc): - await record_rate_limit_cooldown(cooldown_key, rate_limit_delay(exc, 1.0)) - raise + async with embedding_gate(self._cache_identity, self.provider, self._model_name): + response = await asyncio.wait_for( + self._client.embed( + texts=texts, + model=self._model_name, + input_type=cohere_input_type, + embedding_types=["float"], + ), + timeout=60, + ) vectors = [list(e) for e in response.embeddings.float_] for vector in vectors: if len(vector) != self._dimension: diff --git a/api/bigrag/services/embedding/openai.py b/api/bigrag/services/embedding/openai.py index bf7c21c1..a3ed1cf9 100644 --- a/api/bigrag/services/embedding/openai.py +++ b/api/bigrag/services/embedding/openai.py @@ -3,14 +3,8 @@ import asyncio import hashlib -from bigrag.services.embedding.base import EmbeddingModel, get_semaphore, logger, truncate_to_tokens -from bigrag.services.embedding_rate_limit import ( - is_rate_limit_error, - rate_limit_cooldown_key, - rate_limit_delay, - record_rate_limit_cooldown, - wait_for_rate_limit_cooldown, -) +from bigrag.services.embedding.base import EmbeddingModel, logger, truncate_to_tokens +from bigrag.services.embedding_gate import embedding_gate from bigrag.services.url_security import ( pinned_async_client, resolve_and_pin_sync, @@ -43,7 +37,6 @@ def __init__( self._model_name = model_name self._dimension = dimension self._base_url = validate_embedding_base_url_sync(base_url) - self._semaphore_key = f"openai:{self._base_url or 'default'}" base_tag = hashlib.sha256((self._base_url or "").encode()).hexdigest()[:12] self._cache_identity = f"openai:{model_name}:{dimension}:{base_tag}" from bigrag.services.runtime_settings import sync_value @@ -91,23 +84,14 @@ async def embed(self, texts: list[str], *, input_type: str = "document") -> list return await self._embed_single(texts) async def _embed_single(self, texts: list[str]) -> list[list[float]]: - cooldown_key = rate_limit_cooldown_key( - self._cache_identity, self.provider, self._model_name, self._dimension - ) kwargs: dict = {"input": texts, "model": self._model_name} if self._supports_dimensions(self._model_name): kwargs["dimensions"] = self._dimension - async with await get_semaphore(self._semaphore_key): - await wait_for_rate_limit_cooldown(cooldown_key, self.provider, self._model_name) - try: - response = await asyncio.wait_for( - self._client.embeddings.create(**kwargs), - timeout=60, - ) - except Exception as exc: - if is_rate_limit_error(exc): - await record_rate_limit_cooldown(cooldown_key, rate_limit_delay(exc, 1.0)) - raise + async with embedding_gate(self._cache_identity, self.provider, self._model_name): + response = await asyncio.wait_for( + self._client.embeddings.create(**kwargs), + timeout=60, + ) vectors = [item.embedding for item in response.data] for vector in vectors: if len(vector) != self._dimension: diff --git a/api/bigrag/services/embedding/voyage.py b/api/bigrag/services/embedding/voyage.py index e710d84c..bc9b4fb3 100644 --- a/api/bigrag/services/embedding/voyage.py +++ b/api/bigrag/services/embedding/voyage.py @@ -4,14 +4,8 @@ import httpx -from bigrag.services.embedding.base import EmbeddingModel, get_semaphore, logger, truncate_to_tokens -from bigrag.services.embedding_rate_limit import ( - is_rate_limit_error, - rate_limit_cooldown_key, - rate_limit_delay, - record_rate_limit_cooldown, - wait_for_rate_limit_cooldown, -) +from bigrag.services.embedding.base import EmbeddingModel, logger, truncate_to_tokens +from bigrag.services.embedding_gate import embedding_gate from bigrag.services.url_security import pin_embedding_base_url, pinned_async_client @@ -39,7 +33,6 @@ def __init__( self._model_name = model_name self._dimension = dimension self._api_key = api_key - self._semaphore_key = "voyage" self._cache_identity = f"voyage:{model_name}:{dimension}" self._client: httpx.AsyncClient | None = None self._client_lock = asyncio.Lock() @@ -100,22 +93,13 @@ async def _embed_single(self, texts: list[str], voyage_input_type: str) -> list[ "Authorization": f"Bearer {self._api_key}", "Content-Type": "application/json", } - cooldown_key = rate_limit_cooldown_key( - self._cache_identity, self.provider, self._model_name, self._dimension - ) client = await self._get_client() - async with await get_semaphore(self._semaphore_key): - await wait_for_rate_limit_cooldown(cooldown_key, self.provider, self._model_name) - try: - response = await client.post( - f"{self._DEFAULT_BASE_URL}{self._EMBEDDINGS_PATH}", - json=payload, - headers=headers, - ) - except Exception as exc: - if is_rate_limit_error(exc): - await record_rate_limit_cooldown(cooldown_key, rate_limit_delay(exc, 1.0)) - raise + async with embedding_gate(self._cache_identity, self.provider, self._model_name): + response = await client.post( + f"{self._DEFAULT_BASE_URL}{self._EMBEDDINGS_PATH}", + json=payload, + headers=headers, + ) if response.status_code >= 400: logger.warning( "voyage embed http error", @@ -128,8 +112,6 @@ async def _embed_single(self, texts: list[str], voyage_input_type: str) -> list[ f"Voyage embed failed ({response.status_code})", ) exc.headers = response.headers - if is_rate_limit_error(exc): - await record_rate_limit_cooldown(cooldown_key, rate_limit_delay(exc, 1.0)) raise exc data = response.json() vectors = [item["embedding"] for item in data["data"]] diff --git a/api/bigrag/services/embedding_gate.py b/api/bigrag/services/embedding_gate.py new file mode 100644 index 00000000..7a34d2b2 --- /dev/null +++ b/api/bigrag/services/embedding_gate.py @@ -0,0 +1,269 @@ +from __future__ import annotations + +import asyncio +import hashlib +import random +import time +import uuid +from contextlib import asynccontextmanager + +from bigrag.logging import get_logger +from bigrag.services import redis_cache +from bigrag.services.embedding_rate_limit import ( + RATE_LIMIT_COOLDOWN_KEY_PREFIX, + is_rate_limit_error, + rate_limit_delay, + record_rate_limit_cooldown, + wait_for_rate_limit_cooldown, +) + +logger = get_logger("bigrag.embedding_gate") + +MIN_LIMIT = 1.0 +DECREASE_GUARD_MS = 1000 +ACQUIRE_RETRY_MIN_MS = 25 +ACQUIRE_RETRY_MAX_MS = 100 +LEASE_SECONDS = 60 +LIMIT_TTL_MS = 3_600_000 + +INFLIGHT_PREFIX = "bigrag:embedding:inflight:" +LIMIT_PREFIX = "bigrag:embedding:limit:" +LIMIT_DEC_PREFIX = "bigrag:embedding:limit-dec:" + +_LOCAL_TOKEN = "__local__" +_FAILOPEN_TOKEN = "__failopen__" + +_ACQUIRE_LUA = """ +redis.call('ZREMRANGEBYSCORE', KEYS[1], '-inf', ARGV[1]) +local limit = tonumber(redis.call('GET', KEYS[2])) +if limit == nil then limit = tonumber(ARGV[4]); redis.call('SET', KEYS[2], limit, 'PX', ARGV[6]) end +local count = redis.call('ZCARD', KEYS[1]) +if count < math.floor(limit) then + redis.call('ZADD', KEYS[1], ARGV[2], ARGV[3]) + redis.call('PEXPIRE', KEYS[1], ARGV[5]) + return 1 +end +return 0 +""" + +_SUCCESS_LUA = """ +local limit = tonumber(redis.call('GET', KEYS[1])) +if limit == nil then limit = tonumber(ARGV[1]) end +limit = limit + 1.0 / limit +if limit > tonumber(ARGV[1]) then limit = tonumber(ARGV[1]) end +redis.call('SET', KEYS[1], limit, 'PX', ARGV[2]) +return tostring(limit) +""" + +_DECREASE_LUA = """ +local last = tonumber(redis.call('GET', KEYS[2])) or 0 +local limit = tonumber(redis.call('GET', KEYS[1])) +if limit == nil then limit = tonumber(ARGV[2]) end +local changed = 0 +if tonumber(ARGV[1]) - last > tonumber(ARGV[4]) then + limit = limit * 0.5 + if limit < tonumber(ARGV[3]) then limit = tonumber(ARGV[3]) end + redis.call('SET', KEYS[1], limit, 'PX', ARGV[5]) + redis.call('SET', KEYS[2], ARGV[1], 'PX', ARGV[5]) + changed = 1 +end +return {tostring(limit), changed} +""" + + +class _LocalLimiter: + def __init__(self, ceiling: float) -> None: + self.limit = ceiling + self.inflight = 0 + self.last_decrease = 0.0 + self.cond = asyncio.Condition() + + async def acquire(self) -> None: + async with self.cond: + while self.inflight >= int(self.limit): + await self.cond.wait() + self.inflight += 1 + + async def release(self) -> None: + async with self.cond: + self.inflight = max(0, self.inflight - 1) + self.cond.notify(1) + + async def on_success(self, ceiling: float) -> float: + async with self.cond: + self.limit = min(self.limit + 1.0 / self.limit, ceiling) + self.cond.notify(1) + return self.limit + + async def on_rate_limited(self) -> tuple[float, bool]: + async with self.cond: + now = time.monotonic() + changed = False + if now - self.last_decrease > DECREASE_GUARD_MS / 1000: + self.limit = max(self.limit * 0.5, MIN_LIMIT) + self.last_decrease = now + changed = True + return self.limit, changed + + +_local_limiters: dict[str, _LocalLimiter] = {} +_scripts: dict[str, tuple] = {} + + +def _ceiling() -> float: + from bigrag.services.runtime_settings import sync_value + + return max(float(sync_value("embedding_concurrency")), MIN_LIMIT) + + +def _digest(cache_identity: str) -> str: + return hashlib.sha256(str(cache_identity).encode()).hexdigest()[:24] + + +def _local(digest: str) -> _LocalLimiter: + limiter = _local_limiters.get(digest) + if limiter is None: + limiter = _LocalLimiter(_ceiling()) + _local_limiters[digest] = limiter + return limiter + + +def _script(redis, name: str, body: str): + cached = _scripts.get(name) + if cached is not None and cached[0] is redis: + return cached[1] + script = redis.register_script(body) + _scripts[name] = (redis, script) + return script + + +def _as_float(raw) -> float: + if isinstance(raw, (bytes, bytearray)): + return float(raw.decode()) + return float(raw) + + +async def _delete_keys_by_prefix(redis, prefix: str) -> None: + batch: list[bytes | str] = [] + async for key in redis.scan_iter(prefix + "*", count=500): + batch.append(key) + if len(batch) >= 500: + await redis.delete(*batch) + batch = [] + if batch: + await redis.delete(*batch) + + +async def reset_embedding_limiters() -> None: + _local_limiters.clear() + redis = redis_cache.get_redis() + if redis is None: + return + try: + await _delete_keys_by_prefix(redis, LIMIT_PREFIX) + await _delete_keys_by_prefix(redis, LIMIT_DEC_PREFIX) + except Exception as exc: + logger.debug("embedding gate reset failed", error=repr(exc)) + + +async def _acquire(redis, digest: str) -> str: + if redis is None: + await _local(digest).acquire() + return _LOCAL_TOKEN + inflight_key = INFLIGHT_PREFIX + digest + limit_key = LIMIT_PREFIX + digest + ceiling = _ceiling() + script = _script(redis, "acquire", _ACQUIRE_LUA) + while True: + token = uuid.uuid4().hex + now_ms = int(time.time() * 1000) + try: + ok = await script( + keys=[inflight_key, limit_key], + args=[ + now_ms, + now_ms + LEASE_SECONDS * 1000, + token, + ceiling, + LEASE_SECONDS * 1000 * 2, + LIMIT_TTL_MS, + ], + ) + except Exception as exc: + logger.debug("embedding gate acquire fell back", error=repr(exc)) + return _FAILOPEN_TOKEN + if int(ok) == 1: + return token + await asyncio.sleep(random.uniform(ACQUIRE_RETRY_MIN_MS, ACQUIRE_RETRY_MAX_MS) / 1000) + + +async def _release(redis, digest: str, token: str) -> None: + if token == _LOCAL_TOKEN: + await _local(digest).release() + return + if token == _FAILOPEN_TOKEN or redis is None: + return + try: + await redis.zrem(INFLIGHT_PREFIX + digest, token) + except Exception as exc: + logger.debug("embedding gate release failed", error=repr(exc)) + + +async def _on_success(redis, digest: str) -> None: + if redis is None: + await _local(digest).on_success(_ceiling()) + return + try: + script = _script(redis, "success", _SUCCESS_LUA) + await script(keys=[LIMIT_PREFIX + digest], args=[_ceiling(), LIMIT_TTL_MS]) + except Exception as exc: + logger.debug("embedding gate success update failed", error=repr(exc)) + + +async def _on_rate_limited( + redis, digest: str, cooldown_key: str, exc: Exception, provider: str, model_name: str +) -> None: + await record_rate_limit_cooldown(cooldown_key, rate_limit_delay(exc, 1.0)) + if redis is None: + new_limit, changed = await _local(digest).on_rate_limited() + else: + try: + script = _script(redis, "decrease", _DECREASE_LUA) + now_ms = int(time.time() * 1000) + raw = await script( + keys=[LIMIT_PREFIX + digest, LIMIT_DEC_PREFIX + digest], + args=[now_ms, _ceiling(), MIN_LIMIT, DECREASE_GUARD_MS, LIMIT_TTL_MS], + ) + new_limit = _as_float(raw[0]) + changed = bool(int(raw[1])) + except Exception as update_exc: + logger.debug("embedding gate decrease failed", error=repr(update_exc)) + return + if changed: + logger.warning( + "embedding limit decreased", + provider=provider, + model=model_name, + new_limit=round(new_limit, 2), + ) + + +@asynccontextmanager +async def embedding_gate(cache_identity: str, provider: str, model_name: str): + digest = _digest(cache_identity) + cooldown_key = RATE_LIMIT_COOLDOWN_KEY_PREFIX + digest + await wait_for_rate_limit_cooldown(cooldown_key, provider, model_name) + redis = redis_cache.get_redis() + token = await _acquire(redis, digest) + err: BaseException | None = None + try: + yield + except BaseException as exc: + err = exc + raise + finally: + await _release(redis, digest, token) + if err is None: + await _on_success(redis, digest) + elif is_rate_limit_error(err): + await _on_rate_limited(redis, digest, cooldown_key, err, provider, model_name) diff --git a/api/bigrag/services/health.py b/api/bigrag/services/health.py index 35718a1c..d47bed9c 100644 --- a/api/bigrag/services/health.py +++ b/api/bigrag/services/health.py @@ -14,7 +14,7 @@ logger = get_logger("bigrag.services.health") _EMBEDDING_HEALTH_TTL = 60 -READINESS_TTL = 10 +READINESS_TTL = 5 READINESS_CACHE_KEY = "health:readiness" READINESS_CHECK_TIMEOUT = 5 diff --git a/api/bigrag/services/platform_stats.py b/api/bigrag/services/platform_stats.py index 36b216b6..25f3651f 100644 --- a/api/bigrag/services/platform_stats.py +++ b/api/bigrag/services/platform_stats.py @@ -8,24 +8,13 @@ from sqlalchemy.ext.asyncio import AsyncSession from bigrag.db.models import Collection, Document, Webhook -from bigrag.services.health import cache_get, cache_set from bigrag.services.jobs.broker import INGESTION_QUEUE, worker_heartbeat_key -PLATFORM_STATS_CACHE_KEY = "stats:platform" -PLATFORM_STATS_TTL = 15 - async def platform_stats_payload( queue: Any, session: AsyncSession, - *, - use_cache: bool = True, ) -> dict[str, object]: - if use_cache: - cached = await cache_get(PLATFORM_STATS_CACHE_KEY) - if cached: - return cached - async def db_stats(): cols = await session.scalar(sa.select(sa.func.count()).select_from(Collection)) doc_row = ( @@ -108,8 +97,6 @@ async def worker_stats(): "queue_health": queue_health, "workers": worker_data, } - if use_cache: - await cache_set(PLATFORM_STATS_CACHE_KEY, result, ttl=PLATFORM_STATS_TTL) return result diff --git a/api/bigrag/services/queue_embedding/__init__.py b/api/bigrag/services/queue_embedding/__init__.py index 282e7f3b..b03df88b 100644 --- a/api/bigrag/services/queue_embedding/__init__.py +++ b/api/bigrag/services/queue_embedding/__init__.py @@ -1,7 +1,6 @@ from __future__ import annotations from bigrag.services.queue_embedding.embed import ( - EMBEDDING_TIMEOUT_SECONDS, PERMANENT_ERRORS, delete_document_vectors_after_failure, embed_with_cache, @@ -9,7 +8,6 @@ from bigrag.services.queue_embedding.insert import chunk_and_embed __all__ = [ - "EMBEDDING_TIMEOUT_SECONDS", "PERMANENT_ERRORS", "chunk_and_embed", "delete_document_vectors_after_failure", diff --git a/api/bigrag/services/queue_embedding/embed.py b/api/bigrag/services/queue_embedding/embed.py index d8211437..a8d288de 100644 --- a/api/bigrag/services/queue_embedding/embed.py +++ b/api/bigrag/services/queue_embedding/embed.py @@ -1,23 +1,14 @@ from __future__ import annotations -import asyncio import math import time from bigrag.logging import get_logger from bigrag.services import embedding_cache from bigrag.services.embedding import truncate_to_tokens -from bigrag.services.embedding_rate_limit import ( - is_rate_limit_error, - rate_limit_cooldown_key, - rate_limit_delay, - record_rate_limit_cooldown, - wait_for_rate_limit_cooldown, -) logger = get_logger("bigrag.queue") -EMBEDDING_TIMEOUT_SECONDS = 60 PERMANENT_ERRORS = (ValueError, UnicodeDecodeError, KeyError) @@ -55,8 +46,6 @@ async def embed_with_cache( provider_idx = list(missing_by_cache_text.values()) missing_texts = [texts[i] for i in provider_idx] missing_cache_texts = [cache_texts[i] for i in provider_idx] - cooldown_key = rate_limit_cooldown_key(model, provider, model_name, dimension) - await wait_for_rate_limit_cooldown(cooldown_key, provider, model_name) t0 = time.monotonic() logger.debug( "embedding provider request", @@ -64,15 +53,7 @@ async def embed_with_cache( model=model_name, inputs=len(missing_texts), ) - try: - fresh = await asyncio.wait_for( - model.embed(missing_texts, input_type=input_type), - timeout=EMBEDDING_TIMEOUT_SECONDS, - ) - except Exception as exc: - if is_rate_limit_error(exc): - await record_rate_limit_cooldown(cooldown_key, rate_limit_delay(exc, 1.0)) - raise + fresh = await model.embed(missing_texts, input_type=input_type) logger.debug( "embedding provider response", provider=provider, diff --git a/api/bigrag/services/queue_embedding/embed_batches.py b/api/bigrag/services/queue_embedding/embed_batches.py index aec0fcc3..c9da8d0b 100644 --- a/api/bigrag/services/queue_embedding/embed_batches.py +++ b/api/bigrag/services/queue_embedding/embed_batches.py @@ -8,7 +8,6 @@ MAX_RATE_LIMIT_RETRIES, is_rate_limit_error, rate_limit_delay, - record_rate_limit_cooldown, ) from bigrag.services.ingestion_job import IngestionJob from bigrag.services.queue_embedding.embed import PERMANENT_ERRORS, embed_with_cache @@ -24,7 +23,6 @@ async def embed_all_batches( prefix: str, *, embedding_model, - cooldown_key: str, batches: list[tuple[int, int, int, list]], total_batches: int, ) -> list[tuple[int, int, int, list, list[list[float]], float]]: @@ -77,7 +75,6 @@ async def _embed_batch( raise fallback_delay = BATCH_BACKOFF_BASE ** min(rate_limit_attempt, 5) delay = rate_limit_delay(exc, float(fallback_delay)) - await record_rate_limit_cooldown(cooldown_key, delay) logger.warning( "batch rate limited", prefix=prefix, diff --git a/api/bigrag/services/queue_embedding/insert.py b/api/bigrag/services/queue_embedding/insert.py index 71b279c4..169a028a 100644 --- a/api/bigrag/services/queue_embedding/insert.py +++ b/api/bigrag/services/queue_embedding/insert.py @@ -32,7 +32,6 @@ async def chunk_and_embed( job, prefix, embedding_model=plan.embedding_model, - cooldown_key=plan.cooldown_key, batches=plan.batches, total_batches=plan.total_batches, ) diff --git a/api/bigrag/services/queue_embedding/plan.py b/api/bigrag/services/queue_embedding/plan.py index 6bd13308..863f78d7 100644 --- a/api/bigrag/services/queue_embedding/plan.py +++ b/api/bigrag/services/queue_embedding/plan.py @@ -6,7 +6,6 @@ from bigrag.logging import get_logger from bigrag.services.document_elements import ParsedDocument -from bigrag.services.embedding_rate_limit import rate_limit_cooldown_key from bigrag.services.ingestion_job import IngestionJob logger = get_logger("bigrag.queue") @@ -16,7 +15,6 @@ class EmbedPlan: collection: object embedding_model: object - cooldown_key: str text: str elements: list include_elements: bool @@ -60,12 +58,6 @@ async def build_plan( model=job.embedding_model, elapsed=round(elapsed, 2), ) - cooldown_key = rate_limit_cooldown_key( - embedding_model, - job.embedding_provider, - job.embedding_model, - job.embedding_dimension, - ) emit( job.document_id, "model_loaded", @@ -132,7 +124,6 @@ async def build_plan( return EmbedPlan( collection=collection, embedding_model=embedding_model, - cooldown_key=cooldown_key, text=text, elements=elements, include_elements=include_elements, diff --git a/api/bigrag/services/realtime/__init__.py b/api/bigrag/services/realtime/__init__.py deleted file mode 100644 index 71423eb9..00000000 --- a/api/bigrag/services/realtime/__init__.py +++ /dev/null @@ -1,5 +0,0 @@ -from __future__ import annotations - -from bigrag.services.realtime.connection import RealtimeConnection - -__all__ = ["RealtimeConnection"] diff --git a/api/bigrag/services/realtime/admin_topics/__init__.py b/api/bigrag/services/realtime/admin_topics/__init__.py deleted file mode 100644 index 21d5c7ba..00000000 --- a/api/bigrag/services/realtime/admin_topics/__init__.py +++ /dev/null @@ -1,5 +0,0 @@ -from __future__ import annotations - -from bigrag.services.realtime.admin_topics.registry import admin_topic - -__all__ = ["admin_topic"] diff --git a/api/bigrag/services/realtime/admin_topics/_session.py b/api/bigrag/services/realtime/admin_topics/_session.py deleted file mode 100644 index 475c93ea..00000000 --- a/api/bigrag/services/realtime/admin_topics/_session.py +++ /dev/null @@ -1,16 +0,0 @@ -from __future__ import annotations - -from collections.abc import Awaitable, Callable -from typing import Any - -from fastapi import WebSocket - -from bigrag.db.engine import session_factory -from bigrag.services.realtime.specs import SnapshotTopic - -TopicBuilder = Callable[[WebSocket, dict, dict[str, Any]], SnapshotTopic] - - -async def with_session(load: Callable[[Any], Awaitable[Any]]) -> Any: - async with session_factory()() as session: - return await load(session) diff --git a/api/bigrag/services/realtime/admin_topics/access.py b/api/bigrag/services/realtime/admin_topics/access.py deleted file mode 100644 index 4c375722..00000000 --- a/api/bigrag/services/realtime/admin_topics/access.py +++ /dev/null @@ -1,111 +0,0 @@ -from __future__ import annotations - -from typing import Any - -from fastapi import WebSocket - -from bigrag.services.access_log.queries import access_logs_payload, access_overview_payload -from bigrag.services.audit import audit_log_payload -from bigrag.services.realtime.admin_topics._session import with_session -from bigrag.services.realtime.params import boolean, integer, string -from bigrag.services.realtime.specs import SnapshotTopic, TopicError, fixed -from bigrag.services.usage import usage_payload - - -def _access_overview_topic( - websocket: WebSocket, user: dict, params: dict[str, Any] -) -> SnapshotTopic: - window_days = integer(params, "window_days", default=7, minimum=1, maximum=90) - snapshot_topic = f"access:overview:{window_days}" - - async def load(): - return await with_session( - lambda session: access_overview_payload(session, window_days=window_days) - ) - - return SnapshotTopic(snapshot_topic, load, fixed(20.0)) - - -def _access_logs_topic(websocket: WebSocket, user: dict, params: dict[str, Any]) -> SnapshotTopic: - action = string(params, "action", max_length=100) - actor_id = string(params, "actor_id", max_length=120) - collection = string(params, "collection", max_length=120) - method = string(params, "method", max_length=10) - path = string(params, "path", max_length=300) - status_family = string(params, "status_family", max_length=3) - success = boolean(params, "success") - limit = integer(params, "limit", default=100, minimum=1, maximum=1000) - offset = integer(params, "offset", default=0, minimum=0) - if status_family is not None and status_family not in {"1xx", "2xx", "3xx", "4xx", "5xx"}: - raise TopicError("status_family must be one of 1xx, 2xx, 3xx, 4xx, or 5xx") - snapshot_topic = ":".join( - [ - "access:logs", - action or "*", - actor_id or "*", - collection or "*", - method or "*", - path or "*", - status_family or "*", - str(success), - str(limit), - str(offset), - ] - ) - - async def load(): - return await with_session( - lambda session: access_logs_payload( - session, - action=action, - actor_id=actor_id, - collection=collection, - method=method, - path=path, - status_family=status_family, - success=success, - limit=limit, - offset=offset, - cursor=None, - include_total=False, - ) - ) - - return SnapshotTopic(snapshot_topic, load, fixed(20.0)) - - -def _audit_topic(websocket: WebSocket, user: dict, params: dict[str, Any]) -> SnapshotTopic: - action = string(params, "action", max_length=100) - actor_id = string(params, "actor_id", max_length=120) - resource_type = string(params, "resource_type", max_length=50) - limit = integer(params, "limit", default=100, minimum=1, maximum=1000) - offset = integer(params, "offset", default=0, minimum=0) - snapshot_topic = ( - f"audit:{action or '*'}:{actor_id or '*'}:{resource_type or '*'}:{limit}:{offset}" - ) - - async def load(): - return await with_session( - lambda session: audit_log_payload( - session, - action=action, - actor_id=actor_id, - resource_type=resource_type, - limit=limit, - offset=offset, - cursor=None, - include_total=False, - ) - ) - - return SnapshotTopic(snapshot_topic, load, fixed(60.0)) - - -def _usage_topic(websocket: WebSocket, user: dict, params: dict[str, Any]) -> SnapshotTopic: - window_days = integer(params, "window_days", default=30, minimum=1, maximum=365) - snapshot_topic = f"usage:{window_days}" - - async def load(): - return await with_session(lambda session: usage_payload(session, window_days=window_days)) - - return SnapshotTopic(snapshot_topic, load, fixed(60.0)) diff --git a/api/bigrag/services/realtime/admin_topics/connectors.py b/api/bigrag/services/realtime/admin_topics/connectors.py deleted file mode 100644 index c80fcbc7..00000000 --- a/api/bigrag/services/realtime/admin_topics/connectors.py +++ /dev/null @@ -1,85 +0,0 @@ -from __future__ import annotations - -from typing import Any - -from fastapi import WebSocket - -from bigrag.services.connectors.realtime import ( - connector_sources_event_key, - connector_sync_jobs_event_key, -) -from bigrag.services.connectors.views import ( - connector_sources_payload, - connector_sync_jobs_payload, -) -from bigrag.services.realtime.admin_topics._session import with_session -from bigrag.services.realtime.params import integer, string -from bigrag.services.realtime.specs import SnapshotTopic - -ACTIVE_SYNC_JOB_STATUSES = {"pending", "running"} - - -def _connector_sources_topic( - websocket: WebSocket, user: dict, params: dict[str, Any] -) -> SnapshotTopic: - provider = string(params, "provider", required=True, max_length=80) - collection = string(params, "collection", max_length=120) - snapshot_topic = f"{provider}:sources:{collection or 'all'}" - - async def load(): - return await with_session( - lambda session: connector_sources_payload( - session, - provider_slug=provider, - collection=collection, - ) - ) - - return SnapshotTopic( - snapshot_topic, - load, - _connector_sources_interval, - connector_sources_event_key(provider, collection), - ) - - -def _connector_jobs_topic( - websocket: WebSocket, user: dict, params: dict[str, Any] -) -> SnapshotTopic: - provider = string(params, "provider", required=True, max_length=80) - collection = string(params, "collection", max_length=120) - source_id = string(params, "source_id", max_length=120) - limit = integer(params, "limit", default=20, minimum=1, maximum=100) - snapshot_topic = f"{provider}:sync-jobs:{collection or 'all'}:{source_id or 'all'}" - - async def load(): - return await with_session( - lambda session: connector_sync_jobs_payload( - session, - provider_slug=provider, - collection=collection, - source_id=source_id, - limit=limit, - ) - ) - - return SnapshotTopic( - snapshot_topic, - load, - _connector_jobs_interval, - connector_sync_jobs_event_key(provider, collection, source_id), - ) - - -def _connector_sources_interval(payload: Any | None) -> float: - sources = getattr(payload, "sources", []) if payload is not None else [] - return 2.5 if any(getattr(source, "status", None) == "syncing" for source in sources) else 10.0 - - -def _connector_jobs_interval(payload: Any | None) -> float: - jobs = getattr(payload, "jobs", []) if payload is not None else [] - return ( - 2.5 - if any(getattr(job, "status", None) in ACTIVE_SYNC_JOB_STATUSES for job in jobs) - else 10.0 - ) diff --git a/api/bigrag/services/realtime/admin_topics/documents.py b/api/bigrag/services/realtime/admin_topics/documents.py deleted file mode 100644 index f68a6780..00000000 --- a/api/bigrag/services/realtime/admin_topics/documents.py +++ /dev/null @@ -1,138 +0,0 @@ -from __future__ import annotations - -from typing import Any - -from fastapi import WebSocket - -from bigrag.services.collections.stats import collection_stats_payload -from bigrag.services.document_batch import batch_status_payload -from bigrag.services.document_progress import TERMINAL_DOCUMENT_STATUSES -from bigrag.services.documents import get_document_payload, list_documents_payload -from bigrag.services.realtime.admin_topics._session import with_session -from bigrag.services.realtime.params import boolean, document_ids, integer, string -from bigrag.services.realtime.specs import SnapshotTopic, fixed -from bigrag.services.upload_sessions import upload_session_payload - - -def _documents_topic(websocket: WebSocket, user: dict, params: dict[str, Any]) -> SnapshotTopic: - collection = string(params, "collection", required=True, max_length=120) - q = string(params, "q", max_length=200) - status = string(params, "status") - sort = string(params, "sort", default="created_at") or "created_at" - order = string(params, "order", default="desc") or "desc" - include_total = boolean(params, "include_total") or False - limit = integer(params, "limit", default=100, minimum=1, maximum=1000) - offset = integer(params, "offset", default=0, minimum=0) - snapshot_topic = f"documents:list:{collection}" - - async def load(): - return await with_session( - lambda session: list_documents_payload( - session, - collection_name=collection, - q=q, - status=status, - sort=sort, - order=order, - limit=limit, - offset=offset, - cursor=None, - include_total=include_total, - ) - ) - - return SnapshotTopic(snapshot_topic, load, fixed(5.0), f"collection:{collection}") - - -def _batch_status_topic(websocket: WebSocket, user: dict, params: dict[str, Any]) -> SnapshotTopic: - collection = string(params, "collection", required=True, max_length=120) - ids = document_ids(params.get("document_ids")) - snapshot_topic = f"documents:batch:{collection}:{','.join(ids)}" - - async def load(): - return await with_session( - lambda session: batch_status_payload( - session, - user=user, - collection_name=collection, - document_ids=ids, - ) - ) - - return SnapshotTopic( - snapshot_topic, - load, - fixed(2.0), - f"collection:{collection}", - lambda payload: _batch_done(payload, len(ids)), - ) - - -def _document_topic(websocket: WebSocket, user: dict, params: dict[str, Any]) -> SnapshotTopic: - collection = string(params, "collection", required=True, max_length=120) - document_id = string(params, "document_id", required=True, max_length=120) - snapshot_topic = f"documents:detail:{collection}:{document_id}" - - async def load(): - return await with_session( - lambda session: get_document_payload( - session, - user=user, - collection_name=collection, - document_id=document_id, - ) - ) - - return SnapshotTopic(snapshot_topic, load, fixed(2.0), document_id, _document_done) - - -def _upload_session_topic( - websocket: WebSocket, user: dict, params: dict[str, Any] -) -> SnapshotTopic: - collection = string(params, "collection", required=True, max_length=120) - session_id = string(params, "session_id", required=True, max_length=120) - snapshot_topic = f"upload-session:{collection}:{session_id}" - - async def load(): - return await with_session( - lambda session: upload_session_payload( - session, - user=user, - collection_name=collection, - session_id=session_id, - ) - ) - - return SnapshotTopic(snapshot_topic, load, fixed(2.0), None, _upload_session_done) - - -def _collection_stats_topic( - websocket: WebSocket, user: dict, params: dict[str, Any] -) -> SnapshotTopic: - collection = string(params, "collection", required=True, max_length=120) - snapshot_topic = f"collections:stats:{collection}" - - async def load(): - return await with_session( - lambda session: collection_stats_payload(session, name=collection, use_cache=False) - ) - - return SnapshotTopic(snapshot_topic, load, fixed(10.0), f"collection:{collection}") - - -def _document_done(payload: Any) -> bool: - return getattr(payload, "status", None) in TERMINAL_DOCUMENT_STATUSES - - -def _batch_done(payload: Any, expected_count: int) -> bool: - documents = getattr(payload, "documents", []) - terminal = all( - getattr(document, "status", None) in TERMINAL_DOCUMENT_STATUSES for document in documents - ) - if len(documents) < expected_count: - return terminal - return bool(documents) and terminal - - -def _upload_session_done(payload: Any) -> bool: - return getattr(payload, "status", None) in {"complete", "failed", "canceled"} diff --git a/api/bigrag/services/realtime/admin_topics/platform.py b/api/bigrag/services/realtime/admin_topics/platform.py deleted file mode 100644 index 30694d25..00000000 --- a/api/bigrag/services/realtime/admin_topics/platform.py +++ /dev/null @@ -1,38 +0,0 @@ -from __future__ import annotations - -from typing import Any - -from fastapi import WebSocket - -from bigrag.services.event_bus import INGESTION_EVENTS_KEY -from bigrag.services.health import readiness_payload -from bigrag.services.platform_stats import platform_stats_payload -from bigrag.services.realtime.admin_topics._session import with_session -from bigrag.services.realtime.specs import SnapshotTopic, fixed - - -def _platform_stats_topic( - websocket: WebSocket, user: dict, params: dict[str, Any] -) -> SnapshotTopic: - async def load(): - return await with_session( - lambda session: platform_stats_payload( - websocket.app.state.queue, - session, - use_cache=False, - ) - ) - - return SnapshotTopic("platform:stats", load, fixed(5.0), INGESTION_EVENTS_KEY) - - -def _platform_readiness_topic( - websocket: WebSocket, user: dict, params: dict[str, Any] -) -> SnapshotTopic: - async def load(): - return await readiness_payload( - websocket.app.state.vector_store, - websocket.app.state.queue, - ) - - return SnapshotTopic("platform:readiness", load, fixed(10.0)) diff --git a/api/bigrag/services/realtime/admin_topics/registry.py b/api/bigrag/services/realtime/admin_topics/registry.py deleted file mode 100644 index 76ca10ab..00000000 --- a/api/bigrag/services/realtime/admin_topics/registry.py +++ /dev/null @@ -1,52 +0,0 @@ -from __future__ import annotations - -from typing import Any - -from fastapi import WebSocket - -from bigrag.services.realtime.admin_topics._session import TopicBuilder -from bigrag.services.realtime.admin_topics.access import ( - _access_logs_topic, - _access_overview_topic, - _audit_topic, - _usage_topic, -) -from bigrag.services.realtime.admin_topics.connectors import ( - _connector_jobs_topic, - _connector_sources_topic, -) -from bigrag.services.realtime.admin_topics.documents import ( - _batch_status_topic, - _collection_stats_topic, - _document_topic, - _documents_topic, - _upload_session_topic, -) -from bigrag.services.realtime.admin_topics.platform import ( - _platform_readiness_topic, - _platform_stats_topic, -) -from bigrag.services.realtime.specs import TopicError - -_TOPIC_BUILDERS: dict[str, TopicBuilder] = { - "admin.collections.documents": _documents_topic, - "admin.collections.documents.batch_status": _batch_status_topic, - "admin.collections.documents.detail": _document_topic, - "admin.collections.upload_session": _upload_session_topic, - "admin.collections.stats": _collection_stats_topic, - "admin.connectors.sources": _connector_sources_topic, - "admin.connectors.sync_jobs": _connector_jobs_topic, - "admin.access.overview": _access_overview_topic, - "admin.access.logs": _access_logs_topic, - "admin.audit": _audit_topic, - "admin.usage": _usage_topic, - "admin.platform.stats": _platform_stats_topic, - "admin.platform.readiness": _platform_readiness_topic, -} - - -def admin_topic(websocket: WebSocket, user: dict, topic: str, params: dict[str, Any]): - builder = _TOPIC_BUILDERS.get(topic) - if builder is None: - raise TopicError("Unknown realtime topic") - return builder(websocket, user, params) diff --git a/api/bigrag/services/realtime/auth.py b/api/bigrag/services/realtime/auth.py deleted file mode 100644 index 6c30bb87..00000000 --- a/api/bigrag/services/realtime/auth.py +++ /dev/null @@ -1,75 +0,0 @@ -from __future__ import annotations - -from urllib.parse import urlparse - -from fastapi import WebSocket - -from bigrag import config as _config -from bigrag.db.engine import session_factory -from bigrag.middleware import auth as auth_middleware -from bigrag.services import runtime_settings -from bigrag.services.collection_scope import assert_collection_matches_pin -from bigrag.services.realtime_tokens import validate_realtime_token -from bigrag.services.scopes import has_scope - - -async def connection_principal(websocket: WebSocket) -> dict | None: - async with session_factory()() as session: - principal = await auth_middleware._user_from_session(websocket, session) - if principal is None: - principal = await auth_middleware._user_from_api_key(websocket, session) - if principal is not None: - websocket.state.user = principal - return principal - - -async def cookie_origin_allowed(websocket: WebSocket) -> bool: - if _config.settings.session_cookie_name not in websocket.cookies: - return True - origin = websocket.headers.get("origin") - if not origin: - return False - try: - cors_origins = await runtime_settings.get_value("cors_origins") - if not isinstance(cors_origins, list): - cors_origins = [] - except Exception: - cors_origins = websocket.app.state.settings.cors_origins - if "*" in cors_origins or origin in cors_origins: - return True - parsed = urlparse(origin) - if not parsed.scheme or not parsed.netloc: - return False - websocket_scheme = websocket.url.scheme - expected_scheme = "https" if websocket_scheme == "wss" else "http" - return parsed.scheme == expected_scheme and parsed.netloc == websocket.headers.get("host") - - -def is_admin_session(principal: dict | None) -> bool: - return bool( - principal and principal.get("auth_method") == "session" and principal.get("role") == "admin" - ) - - -async def authorize_admin(principal: dict | None) -> None: - if not is_admin_session(principal): - raise PermissionError("Admin session required") - - -async def authorize_collection_events( - websocket: WebSocket, - principal: dict | None, - collection_name: str, - token: str | None, -) -> None: - pinned = principal.get("collection") if principal else None - if pinned: - assert_collection_matches_pin(pinned, collection_name) - if token and await validate_realtime_token(token, collection_name, principal): - return - if principal is None: - raise PermissionError("Authentication required") - if principal.get("auth_method") == "api_key" and not has_scope( - principal.get("scopes"), "collection:read" - ): - raise PermissionError("API key missing required scope: collection:read") diff --git a/api/bigrag/services/realtime/collection_topics.py b/api/bigrag/services/realtime/collection_topics.py deleted file mode 100644 index dc5d7d04..00000000 --- a/api/bigrag/services/realtime/collection_topics.py +++ /dev/null @@ -1,41 +0,0 @@ -from __future__ import annotations - -import re -from typing import Any - -import sqlalchemy as sa -from fastapi import WebSocket - -from bigrag.db.engine import session_factory -from bigrag.db.models import Collection -from bigrag.services.realtime.auth import authorize_collection_events -from bigrag.services.realtime.params import string -from bigrag.services.realtime.specs import EventTopic, TopicError - -_COLLECTION_NAME = re.compile(r"^[a-zA-Z][a-zA-Z0-9_]*$") - - -async def collection_events_topic( - websocket: WebSocket, - principal: dict | None, - params: dict[str, Any], -) -> EventTopic: - collection = string( - params, "collection", required=True, max_length=120, pattern=_COLLECTION_NAME - ) - token = string(params, "token", max_length=500) - await authorize_collection_events(websocket, principal, collection, token) - async with session_factory()() as session: - exists = await session.scalar(sa.select(Collection.id).where(Collection.name == collection)) - if exists is None: - raise TopicError("Collection not found") - return EventTopic( - topic=f"collection.events:{collection}", - event_key=f"collection:{collection}", - connected_payload={ - "step": "connected", - "status": "connected", - "message": f"Listening for events on {collection}", - "progress": 0, - }, - ) diff --git a/api/bigrag/services/realtime/connection.py b/api/bigrag/services/realtime/connection.py deleted file mode 100644 index 091096f7..00000000 --- a/api/bigrag/services/realtime/connection.py +++ /dev/null @@ -1,298 +0,0 @@ -from __future__ import annotations - -import asyncio -import json -import logging -import time -from collections.abc import Awaitable, Callable -from dataclasses import dataclass -from typing import Any - -from fastapi import HTTPException, WebSocket -from starlette.websockets import WebSocketDisconnect - -from bigrag.services.error_sanitize import sanitize_message_text -from bigrag.services.event_bus import event_bus -from bigrag.services.realtime import messages -from bigrag.services.realtime.specs import ( - EventTopic, - RealtimeTopic, - SnapshotTopic, - TopicError, -) -from bigrag.services.realtime.topics import ( - resolve_topic, - topic_error_message, -) - -logger = logging.getLogger(__name__) - -HEARTBEAT_SECONDS = 30.0 -SEND_TIMEOUT_SECONDS = 30.0 -MAX_MESSAGE_BYTES = 64 * 1024 - -SendMessage = Callable[[dict[str, Any]], Awaitable[None]] - - -@dataclass -class Subscription: - id: str - topic: RealtimeTopic - task: asyncio.Task - - -class RealtimeConnection: - def __init__(self, websocket: WebSocket, principal: dict | None) -> None: - self.websocket = websocket - self.principal = principal - self.outgoing: asyncio.Queue[dict[str, Any] | None] = asyncio.Queue(maxsize=1024) - self.subscriptions: dict[str, Subscription] = {} - self._closed = asyncio.Event() - - async def run(self) -> None: - writer = asyncio.create_task(self._write_loop()) - heartbeat = asyncio.create_task(self._heartbeat_loop()) - closed = asyncio.create_task(self._closed.wait()) - try: - while not self._closed.is_set(): - receive = asyncio.create_task(self.websocket.receive_text()) - done, _ = await asyncio.wait({receive, closed}, return_when=asyncio.FIRST_COMPLETED) - if receive not in done: - receive.cancel() - break - try: - raw = receive.result() - except WebSocketDisconnect: - break - except Exception as exc: - logger.debug("realtime receive failed: %s", exc) - break - await self._handle_raw(raw) - finally: - closed.cancel() - await self.close() - heartbeat.cancel() - writer.cancel() - await asyncio.gather(heartbeat, writer, closed, return_exceptions=True) - - async def close(self) -> None: - self._closed.set() - for subscription_id in list(self.subscriptions): - await self.unsubscribe(subscription_id, notify=False) - await self.outgoing.put(None) - - async def send(self, message: dict[str, Any]) -> None: - if self._closed.is_set(): - return - try: - await asyncio.wait_for(self.outgoing.put(message), timeout=SEND_TIMEOUT_SECONDS) - except TimeoutError: - self._closed.set() - - async def _write_loop(self) -> None: - while True: - message = await self.outgoing.get() - if message is None: - return - try: - await asyncio.wait_for( - self.websocket.send_json(messages.encode_message(message)), - timeout=SEND_TIMEOUT_SECONDS, - ) - except (WebSocketDisconnect, RuntimeError, TimeoutError) as exc: - logger.debug("realtime write failed: %s", exc) - self._closed.set() - return - - async def _heartbeat_loop(self) -> None: - while not self._closed.is_set(): - await asyncio.sleep(HEARTBEAT_SECONDS) - await self.send(messages.heartbeat()) - - async def _handle_raw(self, raw: str) -> None: - if len(raw) > MAX_MESSAGE_BYTES: - await self.send(messages.error(None, None, "Message too large")) - return - try: - message = json.loads(raw) - except json.JSONDecodeError: - await self.send(messages.error(None, None, "Invalid JSON message")) - return - if not isinstance(message, dict): - await self.send(messages.error(None, None, "Message must be a JSON object")) - return - message_type = message.get("type") - if message_type == "ping": - await self.send(messages.pong()) - return - if message_type == "subscribe": - await self._subscribe(message) - return - if message_type == "unsubscribe": - subscription_id = _subscription_id(message) - if subscription_id is None: - await self.send(messages.error(None, None, "id is required")) - return - await self.unsubscribe(subscription_id) - return - await self.send(messages.error(None, None, "Unknown realtime message type")) - - async def _subscribe(self, message: dict[str, Any]) -> None: - subscription_id = _subscription_id(message) - topic_name = message.get("topic") - params = message.get("params") or {} - if subscription_id is None: - await self.send(messages.error(None, None, "id is required")) - return - if not isinstance(topic_name, str) or not topic_name: - await self.send(messages.error(subscription_id, None, "topic is required")) - return - if not isinstance(params, dict): - await self.send(messages.error(subscription_id, topic_name, "params must be an object")) - return - if subscription_id in self.subscriptions: - await self.unsubscribe(subscription_id, notify=False) - try: - topic = await resolve_topic(self.websocket, self.principal, topic_name, params) - except Exception as exc: - if not isinstance(exc, TopicError | HTTPException | PermissionError): - logger.exception("realtime topic resolution failed: %s", topic_name) - await self.send(messages.error(subscription_id, topic_name, topic_error_message(exc))) - return - task = asyncio.create_task(self._run_subscription(subscription_id, topic)) - self.subscriptions[subscription_id] = Subscription(subscription_id, topic, task) - await self.send(messages.subscribed(subscription_id, topic.topic)) - - async def unsubscribe(self, subscription_id: str, *, notify: bool = True) -> None: - subscription = self.subscriptions.pop(subscription_id, None) - if subscription is None: - return - subscription.task.cancel() - await asyncio.gather(subscription.task, return_exceptions=True) - if notify: - await self.send(messages.complete(subscription_id, subscription.topic.topic)) - - async def _run_subscription(self, subscription_id: str, topic: RealtimeTopic) -> None: - try: - if isinstance(topic, SnapshotTopic): - await run_snapshot_subscription(subscription_id, topic, self.send) - else: - await run_event_subscription(subscription_id, topic, self.send) - except asyncio.CancelledError: - raise - except Exception: - logger.exception("realtime subscription failed: %s", topic.topic) - await self.send( - messages.error(subscription_id, topic.topic, "Realtime subscription failed") - ) - finally: - current = self.subscriptions.get(subscription_id) - if current is not None and current.task is asyncio.current_task(): - del self.subscriptions[subscription_id] - - -async def run_snapshot_subscription( - subscription_id: str, - topic: SnapshotTopic, - send: SendMessage, -) -> None: - queue = event_bus.subscribe(topic.event_key) if topic.event_key else None - payload: Any | None = None - last_snapshot = 0.0 - try: - payload = await _load_snapshot(subscription_id, topic, send) - last_snapshot = time.monotonic() - if payload is not None and topic.done is not None and topic.done(payload): - await send(messages.complete(subscription_id, topic.topic)) - return - while True: - interval = max(1.0, topic.interval_for(payload)) - if queue is None: - await asyncio.sleep(interval) - elif await _wait_for_refresh(queue, last_snapshot + interval): - if payload is not None and topic.done is not None and topic.done(payload): - await send(messages.complete(subscription_id, topic.topic)) - return - payload = await _load_snapshot(subscription_id, topic, send) - last_snapshot = time.monotonic() - if payload is not None and topic.done is not None and topic.done(payload): - await send(messages.complete(subscription_id, topic.topic)) - return - finally: - if queue is not None and topic.event_key is not None: - event_bus.unsubscribe(topic.event_key, queue) - - -async def _wait_for_refresh(queue: asyncio.Queue, deadline: float) -> bool: - while True: - remaining = deadline - time.monotonic() - if remaining <= 0: - return False - try: - await asyncio.wait_for(queue.get(), timeout=remaining) - except TimeoutError: - return False - return True - - -async def run_event_subscription( - subscription_id: str, - topic: EventTopic, - send: SendMessage, -) -> None: - queue = event_bus.subscribe(topic.event_key) - try: - if topic.connected_payload is not None: - await send(messages.event(subscription_id, topic.topic, topic.connected_payload)) - while True: - event = await queue.get() - if event is None: - await send(messages.complete(subscription_id, topic.topic)) - return - await send( - messages.event( - subscription_id, - topic.topic, - { - **event.detail, - "document_id": event.document_id, - "step": event.step, - "status": event.status, - "message": event.message, - "progress": event.progress, - }, - ) - ) - finally: - event_bus.unsubscribe(topic.event_key, queue) - - -async def _load_snapshot( - subscription_id: str, - topic: SnapshotTopic, - send: SendMessage, -) -> Any | None: - try: - payload = await topic.load() - except asyncio.CancelledError: - raise - except Exception as exc: - message = _load_error_message(exc) - await send(messages.error(subscription_id, topic.topic, message)) - return None - await send(messages.snapshot(subscription_id, topic.topic, payload)) - return payload - - -def _load_error_message(exc: Exception) -> str: - detail = getattr(exc, "detail", None) - if isinstance(detail, str): - return detail - return sanitize_message_text(f"{type(exc).__name__}") or "snapshot error" - - -def _subscription_id(message: dict[str, Any]) -> str | None: - value = message.get("id") - if isinstance(value, str) and value: - return value - return None diff --git a/api/bigrag/services/realtime/messages.py b/api/bigrag/services/realtime/messages.py deleted file mode 100644 index 732c1f64..00000000 --- a/api/bigrag/services/realtime/messages.py +++ /dev/null @@ -1,53 +0,0 @@ -from __future__ import annotations - -from datetime import UTC, datetime -from typing import Any - -from fastapi.encoders import jsonable_encoder - -PROTOCOL_VERSION = 1 - - -def now_iso() -> str: - return datetime.now(UTC).isoformat() - - -def encode_message(message: dict[str, Any]) -> dict[str, Any]: - return jsonable_encoder(message) - - -def _envelope(message_type: str, **fields: Any) -> dict[str, Any]: - return { - "type": message_type, - "version": PROTOCOL_VERSION, - **fields, - "generated_at": now_iso(), - } - - -def subscribed(subscription_id: str, topic: str) -> dict[str, Any]: - return _envelope("subscribed", id=subscription_id, topic=topic) - - -def snapshot(subscription_id: str, topic: str, payload: Any) -> dict[str, Any]: - return _envelope("snapshot", id=subscription_id, topic=topic, payload=payload) - - -def event(subscription_id: str, topic: str, payload: Any) -> dict[str, Any]: - return _envelope("event", id=subscription_id, topic=topic, payload=payload) - - -def error(subscription_id: str | None, topic: str | None, message: str) -> dict[str, Any]: - return _envelope("error", id=subscription_id, topic=topic, message=message) - - -def complete(subscription_id: str, topic: str) -> dict[str, Any]: - return _envelope("complete", id=subscription_id, topic=topic) - - -def heartbeat() -> dict[str, Any]: - return _envelope("heartbeat") - - -def pong() -> dict[str, Any]: - return _envelope("pong") diff --git a/api/bigrag/services/realtime/params.py b/api/bigrag/services/realtime/params.py deleted file mode 100644 index 3bcf0a62..00000000 --- a/api/bigrag/services/realtime/params.py +++ /dev/null @@ -1,80 +0,0 @@ -from __future__ import annotations - -import re -from typing import Any - -from bigrag.services.realtime.specs import TopicError - - -def string( - params: dict[str, Any], - key: str, - *, - default: str | None = None, - required: bool = False, - max_length: int | None = None, - pattern: re.Pattern[str] | None = None, -) -> str | None: - value = params.get(key, default) - if value is None or value == "": - if required: - raise TopicError(f"{key} is required") - return None - if not isinstance(value, str): - raise TopicError(f"{key} must be a string") - if max_length is not None and len(value) > max_length: - raise TopicError(f"{key} is too long") - if pattern is not None and not pattern.fullmatch(value): - raise TopicError(f"{key} is invalid") - return value - - -def integer( - params: dict[str, Any], - key: str, - *, - default: int, - minimum: int | None = None, - maximum: int | None = None, -) -> int: - value = params.get(key, default) - if isinstance(value, bool): - raise TopicError(f"{key} must be an integer") - try: - parsed = int(value) - except (TypeError, ValueError) as exc: - raise TopicError(f"{key} must be an integer") from exc - if minimum is not None and parsed < minimum: - raise TopicError(f"{key} is too small") - if maximum is not None and parsed > maximum: - raise TopicError(f"{key} is too large") - return parsed - - -def boolean(params: dict[str, Any], key: str) -> bool | None: - value = params.get(key) - if value is None or value == "": - return None - if isinstance(value, bool): - return value - if isinstance(value, str): - if value.lower() == "true": - return True - if value.lower() == "false": - return False - raise TopicError(f"{key} must be a boolean") - - -def document_ids(value: Any) -> list[str]: - if isinstance(value, str): - raw_values = [value] - elif isinstance(value, list): - raw_values = [str(item) for item in value] - else: - raise TopicError("document_ids is required") - parsed = [item.strip() for raw in raw_values for item in raw.split(",") if item.strip()] - if not parsed: - raise TopicError("document_ids is required") - if len(parsed) > 100: - raise TopicError("Maximum 100 documents per subscription") - return parsed diff --git a/api/bigrag/services/realtime/specs.py b/api/bigrag/services/realtime/specs.py deleted file mode 100644 index 6de98eb4..00000000 --- a/api/bigrag/services/realtime/specs.py +++ /dev/null @@ -1,36 +0,0 @@ -from __future__ import annotations - -from collections.abc import Awaitable, Callable -from dataclasses import dataclass -from typing import Any - -SnapshotLoader = Callable[[], Awaitable[Any]] -SnapshotInterval = Callable[[Any | None], float] -SnapshotDone = Callable[[Any], bool] - - -class TopicError(Exception): - pass - - -@dataclass -class SnapshotTopic: - topic: str - load: SnapshotLoader - interval_for: SnapshotInterval - event_key: str | None = None - done: SnapshotDone | None = None - - -@dataclass -class EventTopic: - topic: str - event_key: str - connected_payload: dict[str, Any] | None = None - - -RealtimeTopic = SnapshotTopic | EventTopic - - -def fixed(seconds: float) -> SnapshotInterval: - return lambda _payload: seconds diff --git a/api/bigrag/services/realtime/topics.py b/api/bigrag/services/realtime/topics.py deleted file mode 100644 index b14f9bdb..00000000 --- a/api/bigrag/services/realtime/topics.py +++ /dev/null @@ -1,31 +0,0 @@ -from __future__ import annotations - -from typing import Any - -from fastapi import HTTPException, WebSocket - -from bigrag.services.realtime.admin_topics import admin_topic -from bigrag.services.realtime.auth import authorize_admin -from bigrag.services.realtime.collection_topics import collection_events_topic -from bigrag.services.realtime.specs import RealtimeTopic, TopicError - - -async def resolve_topic( - websocket: WebSocket, - principal: dict | None, - topic: str, - params: dict[str, Any], -) -> RealtimeTopic: - if topic == "collection.events": - return await collection_events_topic(websocket, principal, params) - await authorize_admin(principal) - return admin_topic(websocket, principal or {}, topic, params) - - -def topic_error_message(exc: Exception) -> str: - if isinstance(exc, TopicError | HTTPException | PermissionError): - detail = getattr(exc, "detail", None) - if isinstance(detail, str): - return detail - return str(exc) or "Realtime subscription failed" - return "Realtime subscription failed" diff --git a/api/bigrag/services/realtime_tokens.py b/api/bigrag/services/realtime_tokens.py deleted file mode 100644 index 692d21bf..00000000 --- a/api/bigrag/services/realtime_tokens.py +++ /dev/null @@ -1,45 +0,0 @@ -from __future__ import annotations - -import secrets - -from bigrag.services.redis_cache import get_redis - -REALTIME_TOKEN_TTL_SECONDS = 300 -_PREFIX = "bigrag:realtime_token:" - - -async def create_realtime_token(user: dict, collection_name: str) -> str: - token = secrets.token_urlsafe(32) - redis = get_redis() - if redis is None: - raise RuntimeError("Redis is required for realtime tokens") - payload = f"{user['id']}|{collection_name}" - await redis.set( - f"{_PREFIX}{token}", - payload.encode("utf-8"), - ex=REALTIME_TOKEN_TTL_SECONDS, - ) - return token - - -async def validate_realtime_token( - token: str | None, collection_name: str, user: dict | None = None -) -> bool: - if not token: - return False - redis = get_redis() - if redis is None: - return False - key = f"{_PREFIX}{token}" - raw = await redis.get(key) - if raw is None: - return False - decoded = raw.decode("utf-8") - token_user_id, sep, token_collection = decoded.partition("|") - if not sep: - return False - if token_collection != collection_name: - return False - if user is not None and token_user_id != str(user.get("id", "")): - return False - return True diff --git a/api/bigrag/services/runtime_setting_specs/chat.py b/api/bigrag/services/runtime_setting_specs/chat.py index 6060c48d..49305e24 100644 --- a/api/bigrag/services/runtime_setting_specs/chat.py +++ b/api/bigrag/services/runtime_setting_specs/chat.py @@ -17,7 +17,7 @@ group="chat", label="Default chat model", kind="string", - default="gpt-4o-mini", + default="gpt-4.1", description="Default model for chat answers.", ), SettingSpec( diff --git a/api/bigrag/services/runtime_setting_specs/search.py b/api/bigrag/services/runtime_setting_specs/search.py index e14ecb7f..cfd7d9f7 100644 --- a/api/bigrag/services/runtime_setting_specs/search.py +++ b/api/bigrag/services/runtime_setting_specs/search.py @@ -9,7 +9,10 @@ label="Embedding concurrency", kind="int", default=8, - description="Maximum concurrent embedding requests per provider endpoint.", + description=( + "Global ceiling on concurrent embedding requests per endpoint across all " + "workers; the limiter backs off below this on rate limits and recovers toward it." + ), min=1, max=1024, ), diff --git a/api/bigrag/services/runtime_settings_apply.py b/api/bigrag/services/runtime_settings_apply.py index ce08eda7..4f6f9a5f 100644 --- a/api/bigrag/services/runtime_settings_apply.py +++ b/api/bigrag/services/runtime_settings_apply.py @@ -7,7 +7,7 @@ from bigrag import config as config_module from bigrag.logging import get_logger from bigrag.services import runtime_settings -from bigrag.services.embedding import reset_embedding_semaphores +from bigrag.services.embedding import reset_embedding_limiters from bigrag.services.runtime_setting_specs import REGISTRY from bigrag.services.vector_store import VectorStore, vector_store @@ -75,7 +75,7 @@ async def apply_prepared_runtime_settings(app: Any, prepared: PreparedRuntimeSet app.state.vector_store = vector_store prepared.vector_backend = None if "embedding_concurrency" in keyset: - reset_embedding_semaphores() + await reset_embedding_limiters() logger.info("runtime settings applied", keys=prepared.keys) diff --git a/api/bigrag/services/scopes.py b/api/bigrag/services/scopes.py index 82db3a3d..99afd182 100644 --- a/api/bigrag/services/scopes.py +++ b/api/bigrag/services/scopes.py @@ -31,7 +31,6 @@ ("POST", "/v1/collections/{name}/vectors/upsert", "vector:write"), ("POST", "/v1/collections/{name}/vectors/delete", "vector:delete"), ("GET", "/v1/collections/{name}/analytics", "collection:read"), - ("POST", "/v1/collections/{name}/realtime-token", "collection:read"), ("GET", "/v1/documents/{id}/chunks", "document:read"), ("GET", "/v1/documents/{id}", "document:read"), ("GET", "/v1/documents/", "document:read"), @@ -45,7 +44,11 @@ ("GET", "/v1/admin/webhooks/{id}", "webhook:read"), ("GET", "/v1/admin/webhooks/{id}/deliveries", "webhook:read"), ("GET", "/v1/usage", "audit:read"), + ("GET", "/v1/status/overview", "collection:read"), + ("GET", "/v1/status/collections", "collection:read"), + ("GET", "/v1/status/usage", "audit:read"), ("GET", "/v1/admin/audit", "audit:read"), + ("GET", "/v1/admin/status/access", "audit:read"), ("GET", "/v1/admin/access/overview", "audit:read"), ("GET", "/v1/admin/access/logs", "audit:read"), ("POST", "/v1/collections", "collection:write"), diff --git a/api/bigrag/services/status.py b/api/bigrag/services/status.py new file mode 100644 index 00000000..5c9f25be --- /dev/null +++ b/api/bigrag/services/status.py @@ -0,0 +1,37 @@ +from __future__ import annotations + +import sqlalchemy as sa +from sqlalchemy.ext.asyncio import AsyncSession + +from bigrag.db.models import Collection, Document +from bigrag.models.status import CollectionsStatusResponse + + +async def collections_status_payload(session: AsyncSession) -> CollectionsStatusResponse: + collections_total = await session.scalar(sa.select(sa.func.count()).select_from(Collection)) + documents = ( + await session.execute( + sa.select( + sa.func.count().label("total"), + sa.func.coalesce(sa.func.sum(Document.chunk_count), 0).label("chunks"), + sa.func.coalesce(sa.func.sum(Document.token_count), 0).label("tokens"), + sa.func.coalesce(sa.func.sum(Document.file_size), 0).label("size_bytes"), + sa.func.count().filter(Document.status == "ready").label("ready"), + sa.func.count().filter(Document.status == "pending").label("pending"), + sa.func.count().filter(Document.status == "processing").label("processing"), + sa.func.count().filter(Document.status == "failed").label("failed"), + ) + ) + ).one() + + return CollectionsStatusResponse( + collections_total=collections_total or 0, + documents_total=documents.total, + documents_ready=documents.ready, + documents_pending=documents.pending, + documents_processing=documents.processing, + documents_failed=documents.failed, + total_chunks=int(documents.chunks), + total_tokens=int(documents.tokens), + total_size_bytes=int(documents.size_bytes), + ) diff --git a/app/src/config/runtime.ts b/app/src/config/runtime.ts index 8d88afcd..523c5f12 100644 --- a/app/src/config/runtime.ts +++ b/app/src/config/runtime.ts @@ -18,10 +18,3 @@ export const bigragApiUrl = trimSlash( ); export const apiUrl = (path: string) => `${bigragApiUrl}/${path.replace(/^\/+/, "")}`; - -export const realtimeUrl = () => { - const url = new URL(apiUrl("v1/realtime")); - if (url.protocol === "https:") url.protocol = "wss:"; - else if (url.protocol === "http:") url.protocol = "ws:"; - return url.toString(); -}; diff --git a/app/src/features/audit/audit-page.tsx b/app/src/features/audit/audit-page.tsx index f4a1f2b5..25f4fa3a 100644 --- a/app/src/features/audit/audit-page.tsx +++ b/app/src/features/audit/audit-page.tsx @@ -1,3 +1,4 @@ +import { useQuery } from "@tanstack/react-query"; import { ChevronLeft, ChevronRight } from "lucide-react"; import { useMemo, useState } from "react"; import { Button } from "@/components/ui/button"; @@ -6,7 +7,6 @@ import { Empty } from "@/components/ui/empty"; import { Input } from "@/components/ui/input"; import { Page } from "@/components/ui/page"; import { Spinner } from "@/components/ui/spinner"; -import { useRealtimeSnapshotQuery } from "@/hooks/use-realtime-snapshot-query"; import { apiClient } from "@/lib/api"; import { formatNumber, formatRelative } from "@/lib/format"; import { queryKeys } from "@/lib/query-keys"; @@ -50,11 +50,9 @@ export const AuditPage = () => { }), [action, offset, resourceType], ); - const { data, isPending, error, realtimeUnavailable } = useRealtimeSnapshotQuery({ + const { data, isPending, error } = useQuery({ queryKey, queryFn: () => apiClient.get("v1/admin/audit", params), - topic: "admin.audit", - params, }); const total = data?.total ?? 0; const pageCount = Math.max(1, Math.ceil(total / PAGE_SIZE)); @@ -186,7 +184,6 @@ export const AuditPage = () => {
Showing {formatNumber(firstEntry)}-{formatNumber(lastEntry)} of{" "} {formatNumber(total)} entries - {realtimeUnavailable ? " · realtime unavailable, polling" : ""}
diff --git a/app/src/features/chat/chat-input-controls.tsx b/app/src/features/chat/chat-input-controls.tsx index 7988e67f..7507a1b4 100644 --- a/app/src/features/chat/chat-input-controls.tsx +++ b/app/src/features/chat/chat-input-controls.tsx @@ -6,10 +6,10 @@ import { cn } from "@/lib/cn"; import type { Collection } from "@/types/bigrag"; export const OPENAI_MODELS = [ + { value: "gpt-4.1", label: "GPT-4.1" }, + { value: "gpt-4.1-mini", label: "GPT-4.1 mini" }, { value: "gpt-4o-mini", label: "GPT-4o mini" }, { value: "gpt-4o", label: "GPT-4o" }, - { value: "gpt-4.1-mini", label: "GPT-4.1 mini" }, - { value: "gpt-4.1", label: "GPT-4.1" }, { value: "gpt-3.5-turbo", label: "GPT-3.5 turbo" }, ]; diff --git a/app/src/features/chat/chat-page-defaults.ts b/app/src/features/chat/chat-page-defaults.ts index 3578c0cf..c40aa54d 100644 --- a/app/src/features/chat/chat-page-defaults.ts +++ b/app/src/features/chat/chat-page-defaults.ts @@ -12,7 +12,7 @@ export const defaultSystemPrompt = export const defaultChatState: ChatState = { hasOpenAIKey: false, - model: "gpt-4o-mini", + model: "gpt-4.1", topK: 5, temperature: 0.2, searchMode: "semantic", diff --git a/app/src/features/collections/document-detail-route.tsx b/app/src/features/collections/document-detail-route.tsx index 258c4ea7..2b623ae9 100644 --- a/app/src/features/collections/document-detail-route.tsx +++ b/app/src/features/collections/document-detail-route.tsx @@ -93,7 +93,7 @@ export const DocumentDetail = () => { const navigate = useNavigate(); const [deleteOpen, setDeleteOpen] = useState(false); - const { data: doc, dataUpdatedAt, isPending, streaming } = useDocument(name, docId); + const { data: doc, dataUpdatedAt, isFetching, isPending } = useDocument(name, docId); const { data: chunks, refetch: refetchChunks } = useChunks(name, docId); const remove = useDeleteDocument(name); @@ -150,7 +150,7 @@ export const DocumentDetail = () => {

{progress.message}

- {streaming ? ( + {isFetching && doc.status !== "ready" && doc.status !== "failed" ? ( ) : ( {progressPct}% diff --git a/app/src/features/collections/documents/documents-tab.tsx b/app/src/features/collections/documents/documents-tab.tsx index 79989068..c3427438 100644 --- a/app/src/features/collections/documents/documents-tab.tsx +++ b/app/src/features/collections/documents/documents-tab.tsx @@ -62,6 +62,7 @@ export const DocumentsTab = ({ filters, name, onFiltersChange }: DocumentsTabPro fetchNextPage, hasNextPage, isError, + isFetching, isFetchingNextPage, isPending, refetch, @@ -154,7 +155,6 @@ export const DocumentsTab = ({ filters, name, onFiltersChange }: DocumentsTabPro onCancel={() => cancelSession.mutate(activeSessionId)} onDismiss={() => clearActiveSessionId(name)} session={uploadSession.data} - streaming={uploadSession.streaming} /> )} @@ -184,6 +184,10 @@ export const DocumentsTab = ({ filters, name, onFiltersChange }: DocumentsTabPro selectedCount={selected.size} onClearSelection={clearSelection} onBulkDelete={() => setBulkDeleteOpen(true)} + onRefresh={() => { + void refetch(); + }} + refreshPending={isFetching && !isFetchingNextPage} /> {isPending ? ( diff --git a/app/src/features/collections/documents/documents-toolbar.tsx b/app/src/features/collections/documents/documents-toolbar.tsx index c6daa871..dac6f42f 100644 --- a/app/src/features/collections/documents/documents-toolbar.tsx +++ b/app/src/features/collections/documents/documents-toolbar.tsx @@ -1,9 +1,11 @@ -import { Search, Trash2 } from "lucide-react"; +import { RefreshCw, Search, Trash2 } from "lucide-react"; import { useEffect, useState } from "react"; import { Button } from "@/components/ui/button"; import { Card, CardContent } from "@/components/ui/card"; import { Input } from "@/components/ui/input"; import { Select } from "@/components/ui/select"; +import { Spinner } from "@/components/ui/spinner"; +import { Tooltip } from "@/components/ui/tooltip"; import type { DocumentListOrder, DocumentListSort } from "@/hooks/use-documents"; const statusOptions = [ @@ -38,6 +40,8 @@ interface DocumentsToolbarProps { readonly selectedCount: number; readonly onClearSelection: () => void; readonly onBulkDelete: () => void; + readonly onRefresh: () => void; + readonly refreshPending: boolean; } export const DocumentsToolbar = ({ @@ -50,6 +54,8 @@ export const DocumentsToolbar = ({ selectedCount, onClearSelection, onBulkDelete, + onRefresh, + refreshPending, }: DocumentsToolbarProps) => { const [qDraft, setQDraft] = useState(q); @@ -110,7 +116,20 @@ export const DocumentsToolbar = ({
- {loadedLabel} +
+ {loadedLabel} + + + +
{selectedCount > 0 && (
{selectedCount} selected diff --git a/app/src/features/collections/documents/upload-session-panel.tsx b/app/src/features/collections/documents/upload-session-panel.tsx index 3dd2e50a..49e17973 100644 --- a/app/src/features/collections/documents/upload-session-panel.tsx +++ b/app/src/features/collections/documents/upload-session-panel.tsx @@ -3,7 +3,6 @@ import { Badge } from "@/components/ui/badge"; import { Button } from "@/components/ui/button"; import { Card, CardContent } from "@/components/ui/card"; import { ProgressBar } from "@/components/ui/progress-bar"; -import { Spinner } from "@/components/ui/spinner"; import { FileType } from "@/features/collections/documents/file-type"; import { cn } from "@/lib/cn"; import { formatBytes } from "@/lib/format"; @@ -45,7 +44,6 @@ interface UploadSessionPanelProps { readonly onCancel: () => void; readonly onDismiss: () => void; readonly session: UploadSession; - readonly streaming: boolean; } export const UploadSessionPanel = ({ @@ -53,7 +51,6 @@ export const UploadSessionPanel = ({ onCancel, onDismiss, session, - streaming, }: UploadSessionPanelProps) => { const progressPct = sessionProgress(session); const remaining = Math.max(session.total_files - session.uploaded_files, 0); @@ -142,7 +139,6 @@ export const UploadSessionPanel = ({
- {!terminal && streaming && } )} diff --git a/app/src/features/collections/s3/s3-connector-panel.tsx b/app/src/features/collections/s3/s3-connector-panel.tsx index 7812c8d3..3efa0c4e 100644 --- a/app/src/features/collections/s3/s3-connector-panel.tsx +++ b/app/src/features/collections/s3/s3-connector-panel.tsx @@ -74,11 +74,7 @@ export const S3ConnectorPanel = ({ collection }: { collection: string }) => { workerOffline={workerOffline} /> {addSourceOpen && ( diff --git a/app/src/features/collections/s3/sync-monitor.tsx b/app/src/features/collections/s3/sync-monitor.tsx index f20b8ff6..b79f7073 100644 --- a/app/src/features/collections/s3/sync-monitor.tsx +++ b/app/src/features/collections/s3/sync-monitor.tsx @@ -14,16 +14,13 @@ import type { S3SyncJob } from "@/types/bigrag"; export const SyncMonitor = ({ isPending, job, - streaming, }: { isPending: boolean; job: S3SyncJob | undefined; - streaming: boolean; }) => (

Sync monitor

- {streaming && live}
{isPending ? (
diff --git a/app/src/features/overview/overview-page.tsx b/app/src/features/overview/overview-page.tsx index 8ef2f5ed..04dc583d 100644 --- a/app/src/features/overview/overview-page.tsx +++ b/app/src/features/overview/overview-page.tsx @@ -17,18 +17,19 @@ import { getWorkerAvailability } from "@/features/workers/worker-status"; import { useAccessOverview } from "@/hooks/use-access-logs"; import { useSession } from "@/hooks/use-auth"; import { useCollections } from "@/hooks/use-collections"; -import { usePlatformStats, useReadiness } from "@/hooks/use-platform"; +import { useOverviewStatus } from "@/hooks/use-platform"; import { cn } from "@/lib/cn"; import { formatNumber, formatRelative } from "@/lib/format"; export const OverviewPage = () => { const { data: session } = useSession(); - const { data: stats, isPending: statsPending } = usePlatformStats(); - const { data: readiness } = useReadiness(); + const { data: overviewStatus, isPending: statsPending } = useOverviewStatus(); const { data: collectionsData } = useCollections(); const canSeeAccess = session?.user.role === "admin"; const { data: accessOverview, isPending: accessPending } = useAccessOverview(canSeeAccess); + const stats = overviewStatus?.platform; + const readiness = overviewStatus?.readiness; const collections = collectionsData?.collections ?? []; const firstName = session?.user.display_name?.split(" ")[0] || session?.user.email || "there"; const docs = stats?.documents; @@ -68,7 +69,8 @@ export const OverviewPage = () => { Good to see you, {firstName}

- Live readout of retrieval coverage, ingestion health, and the systems behind bigRAG. + Current readout of retrieval coverage, ingestion health, and the systems behind + bigRAG.

diff --git a/app/src/features/usage/usage-page.tsx b/app/src/features/usage/usage-page.tsx index 2c30de2f..34c44d6d 100644 --- a/app/src/features/usage/usage-page.tsx +++ b/app/src/features/usage/usage-page.tsx @@ -1,9 +1,9 @@ +import { useQuery } from "@tanstack/react-query"; import { useMemo, useState } from "react"; import { Card, CardContent, CardDescription, CardHeader, CardTitle } from "@/components/ui/card"; import { Page } from "@/components/ui/page"; import { Select } from "@/components/ui/select"; import { Spinner } from "@/components/ui/spinner"; -import { useRealtimeSnapshotQuery } from "@/hooks/use-realtime-snapshot-query"; import { apiClient } from "@/lib/api"; import { queryKeys } from "@/lib/query-keys"; @@ -42,11 +42,10 @@ const humanBytes = (n: number): string => { export const UsagePage = () => { const [windowDays, setWindowDays] = useState(30); const queryKey = useMemo(() => queryKeys.usage({ windowDays }), [windowDays]); - const { data, isPending, error, realtimeUnavailable } = useRealtimeSnapshotQuery({ + const { data, isPending, error } = useQuery({ queryKey, - queryFn: () => apiClient.get("v1/usage", { window_days: windowDays }), - topic: "admin.usage", - params: { window_days: windowDays }, + queryFn: () => apiClient.get("v1/status/usage", { window_days: windowDays }), + refetchInterval: 5_000, }); const timelineMax = useMemo( @@ -87,9 +86,6 @@ export const UsagePage = () => { onChange={(v) => setWindowDays(Number(v))} />
- {realtimeUnavailable ? ( - realtime unavailable, polling - ) : null} {isPending ? ( diff --git a/app/src/hooks/use-access-logs.ts b/app/src/hooks/use-access-logs.ts index 6d0bb0e9..6636c1ee 100644 --- a/app/src/hooks/use-access-logs.ts +++ b/app/src/hooks/use-access-logs.ts @@ -1,11 +1,13 @@ +import { useQuery } from "@tanstack/react-query"; import { useMemo } from "react"; -import { useRealtimeSnapshotQuery } from "@/hooks/use-realtime-snapshot-query"; import { apiClient } from "@/lib/api"; import { queryKeys } from "@/lib/query-keys"; import type { AccessLogFilters, AccessLogListResponse, AccessLogOverview } from "@/types/bigrag"; export type { AccessLogFilters }; +const statusPollMs = 5_000; + const compactFilters = (filters: AccessLogFilters & { include_total?: boolean }) => Object.fromEntries( Object.entries(filters).filter(([, value]) => value !== undefined && value !== ""), @@ -13,13 +15,12 @@ const compactFilters = (filters: AccessLogFilters & { include_total?: boolean }) export const useAccessOverview = (enabled: boolean, windowDays = 7) => { const queryKey = useMemo(() => queryKeys.access.overview({ windowDays }), [windowDays]); - return useRealtimeSnapshotQuery({ + return useQuery({ queryKey, queryFn: () => - apiClient.get("v1/admin/access/overview", { window_days: windowDays }), + apiClient.get("v1/admin/status/access", { window_days: windowDays }), enabled, - topic: "admin.access.overview", - params: { window_days: windowDays }, + refetchInterval: enabled ? statusPollMs : false, }); }; @@ -52,11 +53,9 @@ export const useAccessLogs = (filters: AccessLogFilters, enabled = true) => { [action, actor_id, collection, limit, method, offset, pathFilter, status_family, success], ); const queryKey = useMemo(() => queryKeys.access.logs(searchParams), [searchParams]); - return useRealtimeSnapshotQuery({ + return useQuery({ queryKey, queryFn: () => apiClient.get("v1/admin/access/logs", searchParams), enabled, - topic: "admin.access.logs", - params: searchParams, }); }; diff --git a/app/src/hooks/use-auth.ts b/app/src/hooks/use-auth.ts index 69a93971..7e07479b 100644 --- a/app/src/hooks/use-auth.ts +++ b/app/src/hooks/use-auth.ts @@ -1,7 +1,6 @@ import { APIError, type User as CurrentUser, type SessionResponse } from "@bigrag/client/browser"; import { useMutation, useQuery, useQueryClient } from "@tanstack/react-query"; import { toast } from "sonner"; -import { closeAllRealtimeStreams } from "@/hooks/use-realtime-snapshot-query"; import { AUTH_TIMEOUT_MS, apiClient } from "@/lib/api"; import { queryKeys } from "@/lib/query-keys"; @@ -62,7 +61,6 @@ export const useLogout = () => { return useMutation({ mutationFn: () => apiClient.post("v1/auth/logout"), onSuccess: () => { - closeAllRealtimeStreams(); qc.clear(); qc.setQueryData(queryKeys.auth.session(), null); toast.success("Signed out"); @@ -75,7 +73,6 @@ export const useLogoutAll = () => { return useMutation({ mutationFn: () => apiClient.post("v1/auth/logout-all"), onSuccess: () => { - closeAllRealtimeStreams(); qc.clear(); qc.setQueryData(queryKeys.auth.session(), null); toast.success("Signed out of all devices"); diff --git a/app/src/hooks/use-collections.ts b/app/src/hooks/use-collections.ts index f237de63..204662d5 100644 --- a/app/src/hooks/use-collections.ts +++ b/app/src/hooks/use-collections.ts @@ -1,7 +1,6 @@ import { type QueryClient, useMutation, useQuery, useQueryClient } from "@tanstack/react-query"; import { useMemo } from "react"; import { toast } from "sonner"; -import { useRealtimeSnapshotQuery } from "@/hooks/use-realtime-snapshot-query"; import { apiClient } from "@/lib/api"; import { errorToast } from "@/lib/mutation-toast"; import { queryKeys } from "@/lib/query-keys"; @@ -10,6 +9,8 @@ import type { Paginated } from "@/types/pagination"; type ListResponse = Paginated<"collections", Collection>; +const statusPollMs = 5_000; + const invalidateCollectionData = (queryClient: QueryClient, name: string) => { queryClient.invalidateQueries({ queryKey: queryKeys.collections.all() }); queryClient.invalidateQueries({ queryKey: queryKeys.collections.one({ name }) }); @@ -34,13 +35,12 @@ export const useCollection = (name: string) => export const useCollectionStats = (name: string) => { const queryKey = useMemo(() => queryKeys.collections.stats({ name }), [name]); - return useRealtimeSnapshotQuery({ + return useQuery({ queryKey, queryFn: () => apiClient.get(`v1/collections/${encodeURIComponent(name)}/stats`), enabled: !!name, - topic: "admin.collections.stats", - params: { collection: name }, + refetchInterval: statusPollMs, }); }; diff --git a/app/src/hooks/use-documents.ts b/app/src/hooks/use-documents.ts index 8e2dd7b6..995adf61 100644 --- a/app/src/hooks/use-documents.ts +++ b/app/src/hooks/use-documents.ts @@ -1,26 +1,8 @@ -import { useInfiniteQuery, useQuery, useQueryClient } from "@tanstack/react-query"; +import { useInfiniteQuery, useQuery } from "@tanstack/react-query"; import { useMemo } from "react"; -import { useRealtimeSnapshotQuery } from "@/hooks/use-realtime-snapshot-query"; -import { - type RealtimeSnapshotSubscription, - useRealtimeSnapshotSubscriptions, -} from "@/hooks/use-realtime-subscriptions"; import { apiClient } from "@/lib/api"; -import { - type BatchStatusResponse, - type DocListResponse, - type DocumentPageParam, - type InfiniteDocumentsData, - mergeDocumentListSnapshot, - mergeDocumentStatusUpdates, - watchedDocumentIds, -} from "@/lib/document-cache"; -import { - documentListSubscription, - documentStatusSubscriptions, - subscriptionDocumentIds, -} from "@/lib/document-subscriptions"; -import { fetchBatchStatus, fetchDocumentList } from "@/lib/documents-api"; +import type { DocumentPageParam } from "@/lib/document-cache"; +import { fetchDocumentList } from "@/lib/documents-api"; import { queryKeys } from "@/lib/query-keys"; import type { Chunk, @@ -40,6 +22,8 @@ export type { DocumentListFilters, DocumentListOrder, DocumentListSort }; const documentListLimit = 1000; const chunkListLimit = 1000; +const statusPollMs = 5_000; +const terminalDocumentStatuses = new Set(["ready", "failed"]); export const useDocuments = (collection: string, filters: DocumentListFilters = {}) => { const limit = filters.limit ?? documentListLimit; @@ -52,11 +36,7 @@ export const useDocuments = (collection: string, filters: DocumentListFilters = () => queryKeys.documents.list({ collection, q, status, sort, order, limit, offset }), [collection, limit, offset, order, q, sort, status], ); - const realtimeParams = useMemo( - () => ({ collection, limit, offset, order, q, sort, status }), - [collection, limit, offset, order, q, sort, status], - ); - return useRealtimeSnapshotQuery({ + return useQuery({ queryKey, queryFn: ({ signal }) => fetchDocumentList( @@ -65,102 +45,6 @@ export const useDocuments = (collection: string, filters: DocumentListFilters = signal, ), enabled: !!collection, - topic: "admin.collections.documents", - params: realtimeParams, - }); -}; - -const useInfiniteDocumentsRealtime = ({ - collection, - initialPageParam, - limit, - order, - q, - queryData, - queryKey, - sort, - status, -}: { - collection: string; - initialPageParam: DocumentPageParam; - limit: number; - order: DocumentListOrder; - q?: string; - queryData: InfiniteDocumentsData | undefined; - queryKey: ReturnType; - sort: DocumentListSort; - status?: string; -}) => { - const queryClient = useQueryClient(); - const subscriptions = useMemo( - () => documentStatusSubscriptions(collection, watchedDocumentIds(queryData)), - [collection, queryData], - ); - const listSubscription = useMemo( - () => documentListSubscription({ collection, limit, order, q, sort, status }), - [collection, limit, order, q, sort, status], - ); - const applyListPayload = (payload: DocListResponse) => { - queryClient.setQueryData(queryKey, (current) => - mergeDocumentListSnapshot(current, payload, initialPageParam, limit), - ); - }; - const applyStatusPayload = ( - payload: BatchStatusResponse, - subscription: RealtimeSnapshotSubscription, - ) => { - const documentIds = subscriptionDocumentIds(subscription); - queryClient.setQueryData(queryKey, (current) => - mergeDocumentStatusUpdates(current, payload.documents), - ); - if (documentIds.length > payload.documents.length) { - void queryClient.invalidateQueries({ queryKey }); - } - }; - - useRealtimeSnapshotSubscriptions({ - enabled: !!collection && subscriptions.length > 0, - pollIntervalMs: 5_000, - subscriptions, - onSnapshot: (payload, subscription) => { - applyStatusPayload(payload, subscription); - }, - onUnavailable: (subscription) => { - const documentIds = subscriptionDocumentIds(subscription); - if (documentIds.length === 0) { - void queryClient.invalidateQueries({ queryKey }); - return; - } - void fetchBatchStatus(collection, documentIds) - .then((payload) => { - applyStatusPayload(payload, subscription); - }) - .catch(() => { - void queryClient.invalidateQueries({ queryKey }); - }); - }, - }); - - useRealtimeSnapshotSubscriptions({ - enabled: !!collection, - pollIntervalMs: 5_000, - subscriptions: [listSubscription], - onSnapshot: applyListPayload, - onUnavailable: () => { - void fetchDocumentList(collection, { - include_total: true, - limit, - offset: 0, - order, - q, - sort, - status, - }) - .then(applyListPayload) - .catch(() => { - void queryClient.invalidateQueries({ queryKey }); - }); - }, }); }; @@ -211,18 +95,6 @@ export const useInfiniteDocuments = (collection: string, filters: DocumentListFi retry: false, }); - useInfiniteDocumentsRealtime({ - collection, - initialPageParam, - limit, - order, - q, - queryData: query.data, - queryKey, - sort, - status, - }); - return query; }; @@ -231,7 +103,7 @@ export const useDocument = (collection: string, docId: string) => { () => queryKeys.documents.one({ collection, id: docId }), [collection, docId], ); - return useRealtimeSnapshotQuery({ + return useQuery({ queryKey, queryFn: ({ signal }) => apiClient.get( @@ -239,9 +111,10 @@ export const useDocument = (collection: string, docId: string) => { { signal }, ), enabled: !!collection && !!docId, - topic: "admin.collections.documents.detail", - params: { collection, document_id: docId }, - closeWhen: (doc) => doc.status === "ready" || doc.status === "failed", + refetchInterval: (query) => { + const doc = query.state.data; + return doc && terminalDocumentStatuses.has(doc.status) ? false : statusPollMs; + }, }); }; diff --git a/app/src/hooks/use-platform.ts b/app/src/hooks/use-platform.ts index 842898c5..69268344 100644 --- a/app/src/hooks/use-platform.ts +++ b/app/src/hooks/use-platform.ts @@ -1,23 +1,36 @@ import { useQuery } from "@tanstack/react-query"; import { useMemo } from "react"; import { apiUrl } from "@/config/runtime"; -import { useRealtimeSnapshotQuery } from "@/hooks/use-realtime-snapshot-query"; import { apiClient } from "@/lib/api"; import { queryKeys } from "@/lib/query-keys"; import type { PlatformStats, ReadinessReport } from "@/types/bigrag"; +const statusPollMs = 5_000; + +export type OverviewStatus = { + platform: PlatformStats; + readiness: ReadinessReport; +}; + +export const useOverviewStatus = () => + useQuery({ + queryKey: queryKeys.platform.overviewStatus(), + queryFn: () => apiClient.get("v1/status/overview"), + refetchInterval: statusPollMs, + }); + export const usePlatformStats = () => { const queryKey = useMemo(() => queryKeys.platform.stats(), []); - return useRealtimeSnapshotQuery({ + return useQuery({ queryKey, queryFn: () => apiClient.get("v1/stats"), - topic: "admin.platform.stats", + refetchInterval: statusPollMs, }); }; export const useReadiness = () => { const queryKey = useMemo(() => queryKeys.platform.readiness(), []); - return useRealtimeSnapshotQuery({ + return useQuery({ queryKey, queryFn: async (): Promise => { const res = await fetch(apiUrl("health/ready"), { @@ -29,7 +42,7 @@ export const useReadiness = () => { } return (await res.json()) as ReadinessReport; }, - topic: "admin.platform.readiness", + refetchInterval: statusPollMs, }); }; diff --git a/app/src/hooks/use-realtime-snapshot-query.ts b/app/src/hooks/use-realtime-snapshot-query.ts deleted file mode 100644 index 53367117..00000000 --- a/app/src/hooks/use-realtime-snapshot-query.ts +++ /dev/null @@ -1,164 +0,0 @@ -import { - hashKey, - type QueryClient, - type QueryFunction, - type QueryKey, - useQuery, - useQueryClient, -} from "@tanstack/react-query"; -import { useEffect, useMemo, useRef, useState } from "react"; -import { - type RealtimeSnapshotSubscription, - useRealtimeSnapshotSubscriptions, -} from "@/hooks/use-realtime-subscriptions"; -import { - closeAllRealtimeStreams, - compactParams, - DEFAULT_FIRST_SNAPSHOT_TIMEOUT_MS, - DEFAULT_POLL_INTERVAL_MS, - type ErrorKind, - MAX_RECONNECT_ATTEMPTS, - type RealtimeMessage, - type RealtimeParams, - type SnapshotEvent, - subscribeStream, -} from "@/lib/realtime-socket"; - -export type { RealtimeParams, RealtimeSnapshotSubscription }; -export { closeAllRealtimeStreams, useRealtimeSnapshotSubscriptions }; - -type RealtimeSnapshotQueryOptions = { - closeWhen?: (payload: T) => boolean; - enabled?: boolean; - firstSnapshotTimeoutMs?: number; - params?: RealtimeParams; - pollIntervalMs?: number; - queryFn: QueryFunction; - queryKey: QueryKey; - topic: string; -}; - -const jitterInterval = (ms: number) => Math.round(ms * (0.85 + Math.random() * 0.3)); - -export const useRealtimeSnapshotQuery = ({ - closeWhen, - enabled = true, - firstSnapshotTimeoutMs = DEFAULT_FIRST_SNAPSHOT_TIMEOUT_MS, - params, - pollIntervalMs = DEFAULT_POLL_INTERVAL_MS, - queryFn, - queryKey, - topic, -}: RealtimeSnapshotQueryOptions) => { - const queryClient = useQueryClient(); - const queryClientRef = useRef(queryClient); - const closeWhenRef = useRef(closeWhen); - const queryFnRef = useRef(queryFn); - const fallbackStartedRef = useRef(false); - const [realtimeUnavailable, setRealtimeUnavailable] = useState(false); - const [streaming, setStreaming] = useState(false); - const queryKeyHash = useMemo(() => hashKey(queryKey), [queryKey]); - const queryKeyRef = useRef(queryKey); - const paramsHash = useMemo(() => JSON.stringify(compactParams(params)), [params]); - - useEffect(() => { - queryClientRef.current = queryClient; - }, [queryClient]); - - useEffect(() => { - closeWhenRef.current = closeWhen; - }, [closeWhen]); - - useEffect(() => { - queryFnRef.current = queryFn; - }, [queryFn]); - - useEffect(() => { - queryKeyRef.current = queryKey; - }, [queryKey]); - - const query = useQuery({ - enabled, - queryFn: (context) => queryFnRef.current(context), - queryKey, - refetchInterval: (q) => { - if (q.state.data != null && closeWhenRef.current?.(q.state.data as T)) return false; - return realtimeUnavailable ? jitterInterval(pollIntervalMs) : false; - }, - retry: false, - }); - - useEffect(() => { - void queryKeyHash; - if (!enabled) { - setStreaming(false); - setRealtimeUnavailable(false); - return; - } - - fallbackStartedRef.current = false; - setRealtimeUnavailable(false); - setStreaming(true); - - let sawSnapshot = false; - let unsubscribed = false; - let failureCount = 0; - - const fetchFallback = () => { - if (fallbackStartedRef.current) return; - fallbackStartedRef.current = true; - setRealtimeUnavailable(true); - void queryClientRef.current.invalidateQueries({ queryKey: queryKeyRef.current }); - }; - - const firstSnapshotTimer = window.setTimeout(() => { - if (!sawSnapshot) fetchFallback(); - }, firstSnapshotTimeoutMs); - - const handleMessage = (message: RealtimeMessage) => { - if (message.type === "complete") { - setStreaming(false); - return; - } - if (message.type !== "snapshot") return; - const snapshot = message as SnapshotEvent; - sawSnapshot = true; - failureCount = 0; - setRealtimeUnavailable(false); - fallbackStartedRef.current = false; - queryClientRef.current.setQueryData(queryKeyRef.current, snapshot.payload); - if (snapshot.payload != null && closeWhenRef.current?.(snapshot.payload)) { - unsubscribe(); - unsubscribed = true; - setStreaming(false); - } - }; - - const handleError = (kind: ErrorKind) => { - if (kind === "transport") { - fetchFallback(); - return; - } - failureCount += 1; - if (failureCount >= MAX_RECONNECT_ATTEMPTS) { - fetchFallback(); - setStreaming(false); - } - }; - - const unsubscribe = subscribeStream( - topic, - JSON.parse(paramsHash) as RealtimeParams, - handleMessage, - handleError, - ); - - return () => { - window.clearTimeout(firstSnapshotTimer); - if (!unsubscribed) unsubscribe(); - setStreaming(false); - }; - }, [enabled, firstSnapshotTimeoutMs, paramsHash, queryKeyHash, topic]); - - return { ...query, realtimeUnavailable, streaming }; -}; diff --git a/app/src/hooks/use-realtime-subscriptions.ts b/app/src/hooks/use-realtime-subscriptions.ts deleted file mode 100644 index ef7373b2..00000000 --- a/app/src/hooks/use-realtime-subscriptions.ts +++ /dev/null @@ -1,155 +0,0 @@ -import { useEffect, useMemo, useRef, useState } from "react"; -import { - compactParams, - DEFAULT_FIRST_SNAPSHOT_TIMEOUT_MS, - DEFAULT_POLL_INTERVAL_MS, - type ErrorKind, - MAX_RECONNECT_ATTEMPTS, - type RealtimeMessage, - type RealtimeParams, - type SnapshotEvent, - subscribeStream, -} from "@/lib/realtime-socket"; - -export type RealtimeSnapshotSubscription = { - key: string; - params?: RealtimeParams; - topic: string; -}; - -type RealtimeSnapshotSubscriptionsOptions = { - enabled?: boolean; - firstSnapshotTimeoutMs?: number; - onSnapshot: (payload: T, subscription: RealtimeSnapshotSubscription) => void; - onUnavailable?: (subscription: RealtimeSnapshotSubscription) => void; - pollIntervalMs?: number; - subscriptions: RealtimeSnapshotSubscription[]; -}; - -export const useRealtimeSnapshotSubscriptions = ({ - enabled = true, - firstSnapshotTimeoutMs = DEFAULT_FIRST_SNAPSHOT_TIMEOUT_MS, - onSnapshot, - onUnavailable, - pollIntervalMs = DEFAULT_POLL_INTERVAL_MS, - subscriptions, -}: RealtimeSnapshotSubscriptionsOptions) => { - const onSnapshotRef = useRef(onSnapshot); - const onUnavailableRef = useRef(onUnavailable); - const [realtimeUnavailable, setRealtimeUnavailable] = useState(false); - const [streaming, setStreaming] = useState(false); - const normalizedSubscriptions = useMemo( - () => - subscriptions.map((subscription) => ({ - ...subscription, - params: compactParams(subscription.params), - })), - [subscriptions], - ); - const subscriptionsHash = useMemo( - () => JSON.stringify(normalizedSubscriptions), - [normalizedSubscriptions], - ); - - useEffect(() => { - onSnapshotRef.current = onSnapshot; - }, [onSnapshot]); - - useEffect(() => { - onUnavailableRef.current = onUnavailable; - }, [onUnavailable]); - - useEffect(() => { - const activeSubscriptions = JSON.parse(subscriptionsHash) as RealtimeSnapshotSubscription[]; - if (!enabled || activeSubscriptions.length === 0) { - setStreaming(false); - setRealtimeUnavailable(false); - return; - } - - setRealtimeUnavailable(false); - setStreaming(true); - - const activeKeys = new Set(activeSubscriptions.map((subscription) => subscription.key)); - const fallbackKeys = new Set(); - const failureCounts = new Map(); - const pollTimers = new Map(); - const snapshotKeys = new Set(); - const timers: number[] = []; - const unsubscribes: (() => void)[] = []; - - const clearPoll = (key: string) => { - const timer = pollTimers.get(key); - if (timer === undefined) return; - window.clearInterval(timer); - pollTimers.delete(key); - }; - - const markComplete = (key: string) => { - activeKeys.delete(key); - clearPoll(key); - if (activeKeys.size === 0) setStreaming(false); - }; - - const fallback = (subscription: RealtimeSnapshotSubscription) => { - if (fallbackKeys.has(subscription.key)) return; - fallbackKeys.add(subscription.key); - setRealtimeUnavailable(true); - onUnavailableRef.current?.(subscription); - if (pollIntervalMs > 0) { - pollTimers.set( - subscription.key, - window.setInterval(() => onUnavailableRef.current?.(subscription), pollIntervalMs), - ); - } - }; - - for (const subscription of activeSubscriptions) { - const timer = window.setTimeout(() => { - if (!snapshotKeys.has(subscription.key)) fallback(subscription); - }, firstSnapshotTimeoutMs); - timers.push(timer); - - const handleMessage = (message: RealtimeMessage) => { - if (message.type === "complete") { - markComplete(subscription.key); - return; - } - if (message.type !== "snapshot") return; - const snapshot = message as SnapshotEvent; - snapshotKeys.add(subscription.key); - failureCounts.set(subscription.key, 0); - fallbackKeys.delete(subscription.key); - clearPoll(subscription.key); - if (fallbackKeys.size === 0) setRealtimeUnavailable(false); - onSnapshotRef.current(snapshot.payload, subscription); - }; - - const handleError = (kind: ErrorKind) => { - if (kind === "transport") { - fallback(subscription); - return; - } - const failureCount = (failureCounts.get(subscription.key) ?? 0) + 1; - failureCounts.set(subscription.key, failureCount); - if (failureCount >= MAX_RECONNECT_ATTEMPTS) { - fallback(subscription); - markComplete(subscription.key); - } - }; - - unsubscribes.push( - subscribeStream(subscription.topic, subscription.params, handleMessage, handleError), - ); - } - - return () => { - for (const timer of timers) window.clearTimeout(timer); - for (const timer of pollTimers.values()) window.clearInterval(timer); - for (const unsubscribe of unsubscribes) unsubscribe(); - setStreaming(false); - }; - }, [enabled, firstSnapshotTimeoutMs, pollIntervalMs, subscriptionsHash]); - - return { realtimeUnavailable, streaming }; -}; diff --git a/app/src/hooks/use-s3-connector.ts b/app/src/hooks/use-s3-connector.ts index e0239383..1c1e4c2d 100644 --- a/app/src/hooks/use-s3-connector.ts +++ b/app/src/hooks/use-s3-connector.ts @@ -1,7 +1,6 @@ -import { type QueryClient, useMutation, useQueryClient } from "@tanstack/react-query"; +import { type QueryClient, useMutation, useQuery, useQueryClient } from "@tanstack/react-query"; import { useMemo } from "react"; import { toast } from "sonner"; -import { useRealtimeSnapshotQuery } from "@/hooks/use-realtime-snapshot-query"; import { apiClient } from "@/lib/api"; import { errorToast } from "@/lib/mutation-toast"; import { queryKeys } from "@/lib/query-keys"; @@ -14,6 +13,8 @@ import type { UpdateS3SourceBody, } from "@/types/bigrag"; +const connectorPollMs = 2_000; + const updateS3SourcesCache = ( queryClient: QueryClient, collection: string | undefined, @@ -35,15 +36,13 @@ const invalidateS3SyncJobs = (queryClient: QueryClient) => { export const useS3Sources = (collection?: string) => { const queryKey = useMemo(() => queryKeys.connectors.s3Sources({ collection }), [collection]); - const params = useMemo(() => ({ provider: "s3", collection }), [collection]); - return useRealtimeSnapshotQuery({ + return useQuery({ queryKey, queryFn: () => apiClient.get("v1/connectors/s3/sources", { ...(collection ? { collection } : {}), }), - topic: "admin.connectors.sources", - params, + refetchInterval: connectorPollMs, }); }; @@ -60,11 +59,7 @@ export const useS3SyncJobs = ({ () => queryKeys.connectors.s3SyncJobs({ collection, limit, sourceId }), [collection, limit, sourceId], ); - const params = useMemo( - () => ({ provider: "s3", collection, limit, source_id: sourceId }), - [collection, limit, sourceId], - ); - return useRealtimeSnapshotQuery({ + return useQuery({ queryKey, queryFn: () => apiClient.get("v1/connectors/s3/sync-jobs", { @@ -72,8 +67,7 @@ export const useS3SyncJobs = ({ ...(collection ? { collection } : {}), ...(sourceId ? { source_id: sourceId } : {}), }), - topic: "admin.connectors.sync_jobs", - params, + refetchInterval: connectorPollMs, }); }; diff --git a/app/src/hooks/use-upload-session-documents.ts b/app/src/hooks/use-upload-session-documents.ts index 1424b0fd..365a3573 100644 --- a/app/src/hooks/use-upload-session-documents.ts +++ b/app/src/hooks/use-upload-session-documents.ts @@ -1,7 +1,6 @@ -import { useMutation, useQueryClient } from "@tanstack/react-query"; +import { useMutation, useQuery, useQueryClient } from "@tanstack/react-query"; import { useMemo } from "react"; import { toast } from "sonner"; -import { useRealtimeSnapshotQuery } from "@/hooks/use-realtime-snapshot-query"; import { apiClient } from "@/lib/api"; import { runWithConcurrency } from "@/lib/concurrency"; import { errorToast } from "@/lib/mutation-toast"; @@ -9,6 +8,12 @@ import { queryKeys } from "@/lib/query-keys"; import type { UploadSession, UploadSessionFileResponse } from "@/types/bigrag"; const uploadConcurrency = 4; +const uploadSessionPollMs = 2_000; +const terminalUploadSessionStatuses = new Set([ + "complete", + "failed", + "canceled", +]); const uploadSessionFileName = (file: File) => (file as File & { webkitRelativePath?: string }).webkitRelativePath || file.name; @@ -22,7 +27,7 @@ export const useUploadSession = (collection: string, sessionId: string | null) = [collection, sessionId], ); const enabled = Boolean(collection && sessionId); - return useRealtimeSnapshotQuery({ + return useQuery({ queryKey, queryFn: ({ signal }) => apiClient.get( @@ -30,11 +35,12 @@ export const useUploadSession = (collection: string, sessionId: string | null) = { signal }, ), enabled, - topic: "admin.collections.upload_session", - params: { collection, session_id: sessionId ?? "" }, - pollIntervalMs: 2_000, - closeWhen: (session) => - session.status === "complete" || session.status === "failed" || session.status === "canceled", + refetchInterval: (query) => { + const session = query.state.data; + return session && terminalUploadSessionStatuses.has(session.status) + ? false + : uploadSessionPollMs; + }, }); }; diff --git a/app/src/lib/document-cache.ts b/app/src/lib/document-cache.ts index 93567630..6932be0f 100644 --- a/app/src/lib/document-cache.ts +++ b/app/src/lib/document-cache.ts @@ -1,4 +1,3 @@ -import type { InfiniteData } from "@tanstack/react-query"; import type { Document } from "@/types/bigrag"; import type { Paginated } from "@/types/pagination"; @@ -19,96 +18,3 @@ export type BatchStatusResponse = { documents: DocumentStatusUpdate[]; total: number; }; - -export type InfiniteDocumentsData = InfiniteData; - -const documentStatusBatchSize = 100; -const activeDocumentStatuses = new Set(["pending", "processing"]); - -export const chunkDocumentIds = (ids: string[]) => { - const chunks: string[][] = []; - for (let index = 0; index < ids.length; index += documentStatusBatchSize) { - chunks.push(ids.slice(index, index + documentStatusBatchSize)); - } - return chunks; -}; - -export const watchedDocumentIds = (data: InfiniteDocumentsData | undefined) => { - const ids: string[] = []; - const seen = new Set(); - for (const page of data?.pages ?? []) { - for (const document of page.documents) { - if (!activeDocumentStatuses.has(document.status) || seen.has(document.id)) continue; - seen.add(document.id); - ids.push(document.id); - } - } - return ids; -}; - -export const mergeDocumentStatusUpdates = ( - current: InfiniteDocumentsData | undefined, - updates: DocumentStatusUpdate[], -): InfiniteDocumentsData | undefined => { - if (!current || updates.length === 0) return current; - const updatesById = new Map(updates.map((document) => [document.id, document])); - let changed = false; - const pages = current.pages.map((page) => { - let pageChanged = false; - const documents = page.documents.map((document) => { - const update = updatesById.get(document.id); - if (!update) return document; - const unchanged = - document.status === update.status && - document.error_message === update.error_message && - document.chunk_count === update.chunk_count && - document.multimodal_element_count === update.multimodal_element_count && - document.progress === update.progress; - if (unchanged) return document; - pageChanged = true; - return { - ...document, - chunk_count: update.chunk_count, - error_message: update.error_message, - multimodal_element_count: update.multimodal_element_count, - progress: update.progress, - status: update.status, - }; - }); - if (!pageChanged) return page; - changed = true; - return { ...page, documents }; - }); - return changed ? { ...current, pages } : current; -}; - -export const mergeDocumentListSnapshot = ( - current: InfiniteDocumentsData | undefined, - snapshot: DocListResponse, - firstPageParam: DocumentPageParam, - limit: number, -): InfiniteDocumentsData => { - if (!current || current.pages.length === 0) { - return { pageParams: [firstPageParam], pages: [snapshot] }; - } - const existingTotal = current.pages.find((page) => page.total !== null)?.total ?? null; - const firstPage = { ...snapshot, total: snapshot.total ?? existingTotal }; - const firstPageIds = new Set(firstPage.documents.map((document) => document.id)); - const previousDocuments = current.pages - .flatMap((page) => page.documents) - .filter((document) => !firstPageIds.has(document.id)); - const pages = [firstPage]; - let previousIndex = 0; - for (let pageIndex = 1; pageIndex < current.pages.length; pageIndex += 1) { - const pageSize = Math.min(limit, current.pages[pageIndex]?.documents.length ?? limit); - const documents = previousDocuments.slice(previousIndex, previousIndex + pageSize); - previousIndex += pageSize; - if (documents.length === 0) break; - pages.push({ ...current.pages[pageIndex], documents }); - } - return { - ...current, - pageParams: current.pageParams.slice(0, pages.length), - pages, - }; -}; diff --git a/app/src/lib/document-subscriptions.ts b/app/src/lib/document-subscriptions.ts deleted file mode 100644 index 6882a3ec..00000000 --- a/app/src/lib/document-subscriptions.ts +++ /dev/null @@ -1,54 +0,0 @@ -import type { RealtimeSnapshotSubscription } from "@/hooks/use-realtime-subscriptions"; -import { chunkDocumentIds } from "@/lib/document-cache"; -import type { DocumentListOrder, DocumentListSort } from "@/types/bigrag"; - -export const documentStatusSubscriptions = ( - collection: string, - ids: string[], -): RealtimeSnapshotSubscription[] => - chunkDocumentIds(ids).map((documentIds, index) => ({ - key: `${collection}:${index}:${documentIds.join(",")}`, - topic: "admin.collections.documents.batch_status", - params: { collection, document_ids: documentIds }, - })); - -export const documentListSubscription = ({ - collection, - limit, - order, - q, - sort, - status, -}: { - collection: string; - limit: number; - order: DocumentListOrder; - q?: string; - sort: DocumentListSort; - status?: string; -}): RealtimeSnapshotSubscription => ({ - key: `${collection}:${limit}:${order}:${q ?? ""}:${sort}:${status ?? ""}`, - topic: "admin.collections.documents", - params: { - collection, - include_total: true, - limit, - offset: 0, - order, - q, - sort, - status, - }, -}); - -export const subscriptionDocumentIds = (subscription: RealtimeSnapshotSubscription) => { - const value = subscription.params?.document_ids; - if (Array.isArray(value)) return value; - if (typeof value === "string") { - return value - .split(",") - .map((id) => id.trim()) - .filter(Boolean); - } - return []; -}; diff --git a/app/src/lib/query-keys.ts b/app/src/lib/query-keys.ts index 1c745e11..5ff381c9 100644 --- a/app/src/lib/query-keys.ts +++ b/app/src/lib/query-keys.ts @@ -126,6 +126,7 @@ export const queryKeys = { ["documents", "upload-session", { collection, id }] as const, }, platform: { + overviewStatus: () => ["status", "overview"] as const, stats: () => ["platform", "stats"] as const, readiness: () => ["platform", "readiness"] as const, embeddingModels: () => ["platform", "embedding-models"] as const, diff --git a/app/src/lib/realtime-socket.ts b/app/src/lib/realtime-socket.ts deleted file mode 100644 index 621e1c2a..00000000 --- a/app/src/lib/realtime-socket.ts +++ /dev/null @@ -1,248 +0,0 @@ -import { realtimeUrl } from "@/config/runtime"; - -export type RealtimeParams = Record< - string, - string | number | boolean | string[] | null | undefined ->; - -export type SnapshotEvent = { - generated_at: string; - id: string; - payload: T; - topic: string; - type: "snapshot"; -}; - -export type RealtimeMessage = - | SnapshotEvent - | { - generated_at?: string; - id?: string; - message?: string; - topic?: string; - type: string; - version?: number; - }; - -export const DEFAULT_FIRST_SNAPSHOT_TIMEOUT_MS = 5_000; -export const DEFAULT_POLL_INTERVAL_MS = 30_000; -export const MAX_RECONNECT_ATTEMPTS = 5; -const MAX_BACKOFF_MS = 30_000; -const BASE_BACKOFF_MS = 1_000; -const WS_POLICY_VIOLATION = 1008; - -export type ErrorKind = "transport" | "subscription"; -export type MessageListener = (event: RealtimeMessage) => void; -export type ErrorListener = (kind: ErrorKind) => void; - -type StreamEntry = { - completed: boolean; - errorListeners: Set; - id: string; - listeners: Set; - params: Record; - refcount: number; - topic: string; -}; - -let socket: WebSocket | null = null; -let socketGeneration = 0; -let reconnectAttempt = 0; -let reconnectTimer: ReturnType | null = null; -const streams = new Map(); -const streamsById = new Map(); - -const randomId = () => { - if (globalThis.crypto?.randomUUID) return globalThis.crypto.randomUUID(); - return `${Date.now().toString(36)}-${Math.random().toString(36).slice(2)}-${Math.random().toString(36).slice(2)}`; -}; - -const jitter = (ms: number) => ms * (0.75 + Math.random() * 0.5); - -const backoffDelay = (attempt: number) => - Math.round(jitter(Math.min(BASE_BACKOFF_MS * 2 ** attempt, MAX_BACKOFF_MS))); - -export const compactParams = (params?: RealtimeParams) => - Object.fromEntries( - Object.entries(params ?? {}).filter(([, value]) => value !== undefined && value !== null), - ) as Record; - -const streamKey = (topic: string, params?: RealtimeParams) => - `${topic}:${JSON.stringify(compactParams(params))}`; - -const send = (message: unknown) => { - if (!socket || socket.readyState !== WebSocket.OPEN) return; - socket.send(JSON.stringify(message)); -}; - -const subscribeEntry = (entry: StreamEntry) => { - if (entry.completed) return; - send({ type: "subscribe", id: entry.id, topic: entry.topic, params: entry.params }); -}; - -const dispatchError = (kind: ErrorKind) => { - for (const entry of streams.values()) { - for (const listener of entry.errorListeners) listener(kind); - } -}; - -const closeSocket = () => { - const current = socket; - socket = null; - socketGeneration += 1; - current?.close(); -}; - -const openSocket = () => { - if (socket && socket.readyState !== WebSocket.CLOSING && socket.readyState !== WebSocket.CLOSED) { - return; - } - const currentSocket = new WebSocket(realtimeUrl()); - socket = currentSocket; - const generation = socketGeneration + 1; - socketGeneration = generation; - const isCurrentSocket = () => socket === currentSocket && socketGeneration === generation; - - currentSocket.addEventListener("open", () => { - if (!isCurrentSocket()) return; - reconnectAttempt = 0; - for (const entry of streams.values()) subscribeEntry(entry); - }); - - currentSocket.addEventListener("message", (event) => { - if (!isCurrentSocket()) return; - let message: RealtimeMessage; - try { - message = JSON.parse(String(event.data)) as RealtimeMessage; - } catch { - return; - } - if (message.type === "heartbeat" || message.type === "pong" || message.type === "subscribed") { - return; - } - if (message.type === "error") { - const entry = message.id ? streamsById.get(message.id) : undefined; - const targets = entry ? [entry] : Array.from(streams.values()); - for (const target of targets) { - for (const listener of target.errorListeners) listener("subscription"); - } - return; - } - if (!message.id) return; - const entry = streamsById.get(message.id); - if (!entry) return; - if (message.type === "complete") { - entry.completed = true; - } - for (const listener of entry.listeners) listener(message); - }); - - currentSocket.addEventListener("close", (event) => { - if (!isCurrentSocket()) return; - socket = null; - if (streams.size === 0) return; - dispatchError("transport"); - if (event.code === WS_POLICY_VIOLATION) return; - scheduleReconnect(); - }); - - currentSocket.addEventListener("error", () => { - if (!isCurrentSocket()) return; - dispatchError("transport"); - }); -}; - -const scheduleReconnect = () => { - if (reconnectTimer) return; - if (streams.size === 0) return; - if (reconnectAttempt >= MAX_RECONNECT_ATTEMPTS) return; - const delay = backoffDelay(reconnectAttempt); - reconnectAttempt += 1; - reconnectTimer = setTimeout(() => { - reconnectTimer = null; - openSocket(); - }, delay); -}; - -const resumeRealtime = () => { - if (streams.size === 0) return; - if (socket && socket.readyState !== WebSocket.CLOSING && socket.readyState !== WebSocket.CLOSED) { - return; - } - reconnectAttempt = 0; - if (reconnectTimer) { - clearTimeout(reconnectTimer); - reconnectTimer = null; - } - openSocket(); -}; - -if (typeof window !== "undefined") { - window.addEventListener("online", resumeRealtime); - document.addEventListener("visibilitychange", () => { - if (document.visibilityState === "visible") resumeRealtime(); - }); -} - -export const subscribeStream = ( - topic: string, - params: RealtimeParams | undefined, - onMessage: MessageListener, - onError: ErrorListener, -): (() => void) => { - const key = streamKey(topic, params); - let entry = streams.get(key); - let created = false; - if (!entry) { - entry = { - id: randomId(), - topic, - params: compactParams(params), - listeners: new Set(), - errorListeners: new Set(), - refcount: 0, - completed: false, - }; - streams.set(key, entry); - streamsById.set(entry.id, entry); - created = true; - } - entry.listeners.add(onMessage); - entry.errorListeners.add(onError); - entry.refcount += 1; - reconnectAttempt = 0; - openSocket(); - if (created && socket?.readyState === WebSocket.OPEN) subscribeEntry(entry); - - return () => { - const current = streams.get(key); - if (!current) return; - current.listeners.delete(onMessage); - current.errorListeners.delete(onError); - current.refcount = Math.max(0, current.refcount - 1); - if (current.refcount === 0) { - send({ type: "unsubscribe", id: current.id }); - streams.delete(key); - streamsById.delete(current.id); - if (streams.size === 0 && reconnectTimer) { - clearTimeout(reconnectTimer); - reconnectTimer = null; - } - } - }; -}; - -export const closeAllRealtimeStreams = () => { - if (reconnectTimer) { - clearTimeout(reconnectTimer); - reconnectTimer = null; - } - for (const entry of streams.values()) { - entry.listeners.clear(); - entry.errorListeners.clear(); - entry.refcount = 0; - } - streams.clear(); - streamsById.clear(); - closeSocket(); -}; diff --git a/app/src/routes/_dashboard.collections.$name.search.tsx b/app/src/routes/_dashboard.collections.$name.search.tsx index 170ea9d7..a6af2132 100644 --- a/app/src/routes/_dashboard.collections.$name.search.tsx +++ b/app/src/routes/_dashboard.collections.$name.search.tsx @@ -215,7 +215,7 @@ const SearchTab = () => { ) : ( )} - {run.data.timings.cache_hit ? "cache hit" : "live"} + {run.data.timings.cache_hit ? "cache hit" : "uncached"} )} diff --git a/bigrag.toml b/bigrag.toml index a305c038..ca09fcb2 100644 --- a/bigrag.toml +++ b/bigrag.toml @@ -52,7 +52,7 @@ # allow_private_embedding_base_urls = false # chat_provider = "openai" # openai | openai_compatible -# chat_model = "gpt-4o-mini" +# chat_model = "gpt-4.1" # chat_base_url = "" # required for openai_compatible endpoints # chat_temperature = 0.2 # chat_max_context_chars = 120000 diff --git a/docs/superpowers/plans/2026-05-24-adaptive-embedding-rate-limiting.md b/docs/superpowers/plans/2026-05-24-adaptive-embedding-rate-limiting.md new file mode 100644 index 00000000..13e0647b --- /dev/null +++ b/docs/superpowers/plans/2026-05-24-adaptive-embedding-rate-limiting.md @@ -0,0 +1,1073 @@ +# Adaptive Embedding Rate Limiting Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Replace the per-process, reactive embedding rate-limit handling with a single Redis-coordinated, self-tuning (AIMD) concurrency limiter so the 5 workers stop blasting the provider in lockstep and 429s become rare. + +**Architecture:** A new `embedding_gate` async context manager wraps every provider HTTP call. It (1) waits out any active Redis 429 cooldown, (2) acquires a permit from a shared, dynamically-sized concurrency limit stored in Redis, (3) on success nudges the limit up toward the `embedding_concurrency` ceiling, and (4) on a 429 halves the limit and records the cooldown. The limit lives in Redis (a ZSET of leased permits plus a float limit key), with a process-local fallback when Redis is absent. The previously per-process semaphores and the duplicated cooldown logic in each provider are removed. + +**Tech Stack:** Python 3.12, asyncio, `redis.asyncio` (Lua via `register_script`), structlog. No new dependencies. Lint via `api/.venv/bin/ruff` (pre-commit also runs ruff/biome on commit). No unit-test framework exists in this repo — verification is ruff + the existing load test (`tests/load`). + +**Conventions (from user):** Conventional commit messages (`feat: ...`, not `feat(scope): ...`). **No co-author trailer.** **No explanatory comments in code** — the diff must speak for itself. Commit + push after each task. + +**Key bounds:** +- Ceiling = `sync_value("embedding_concurrency")` clamped to `>= 1` (now global across workers, was per-process). +- `MIN_LIMIT = 1.0`, `DECREASE_GUARD_MS = 1000`, `ACQUIRE_RETRY_MIN_MS = 25`, `ACQUIRE_RETRY_MAX_MS = 100`, `LEASE_SECONDS = 60`, `LIMIT_TTL_MS = 3_600_000`. +- The Redis client uses `decode_responses=False` (`redis_cache.py:22`), and Lua numeric returns truncate to int — so the limit-math scripts `return tostring(limit)` and Python parses the returned **bytes** back to `float` via `float(raw.decode())`. + +--- + +### Task 1: Create the `embedding_gate` module + +**Files:** +- Create: `api/bigrag/services/embedding_gate.py` + +- [ ] **Step 1: Write the module** + +Create `api/bigrag/services/embedding_gate.py` with exactly this content: + +```python +from __future__ import annotations + +import asyncio +import hashlib +import random +import time +import uuid +from contextlib import asynccontextmanager + +from bigrag.logging import get_logger +from bigrag.services import redis_cache +from bigrag.services.embedding_rate_limit import ( + RATE_LIMIT_COOLDOWN_KEY_PREFIX, + is_rate_limit_error, + rate_limit_delay, + record_rate_limit_cooldown, + wait_for_rate_limit_cooldown, +) + +logger = get_logger("bigrag.embedding_gate") + +MIN_LIMIT = 1.0 +DECREASE_GUARD_MS = 1000 +ACQUIRE_RETRY_MIN_MS = 25 +ACQUIRE_RETRY_MAX_MS = 100 +LEASE_SECONDS = 60 +LIMIT_TTL_MS = 3_600_000 + +INFLIGHT_PREFIX = "bigrag:embedding:inflight:" +LIMIT_PREFIX = "bigrag:embedding:limit:" +LIMIT_DEC_PREFIX = "bigrag:embedding:limit-dec:" + +_LOCAL_TOKEN = "__local__" +_FAILOPEN_TOKEN = "__failopen__" + +_ACQUIRE_LUA = """ +redis.call('ZREMRANGEBYSCORE', KEYS[1], '-inf', ARGV[1]) +local limit = tonumber(redis.call('GET', KEYS[2])) +if limit == nil then limit = tonumber(ARGV[4]); redis.call('SET', KEYS[2], limit, 'PX', ARGV[6]) end +local count = redis.call('ZCARD', KEYS[1]) +if count < math.floor(limit) then + redis.call('ZADD', KEYS[1], ARGV[2], ARGV[3]) + redis.call('PEXPIRE', KEYS[1], ARGV[5]) + return 1 +end +return 0 +""" + +_SUCCESS_LUA = """ +local limit = tonumber(redis.call('GET', KEYS[1])) +if limit == nil then limit = tonumber(ARGV[1]) end +limit = limit + 1.0 / limit +if limit > tonumber(ARGV[1]) then limit = tonumber(ARGV[1]) end +redis.call('SET', KEYS[1], limit, 'PX', ARGV[2]) +return tostring(limit) +""" + +_DECREASE_LUA = """ +local last = tonumber(redis.call('GET', KEYS[2])) or 0 +local limit = tonumber(redis.call('GET', KEYS[1])) +if limit == nil then limit = tonumber(ARGV[2]) end +if tonumber(ARGV[1]) - last > tonumber(ARGV[4]) then + limit = limit * 0.5 + if limit < tonumber(ARGV[3]) then limit = tonumber(ARGV[3]) end + redis.call('SET', KEYS[1], limit, 'PX', ARGV[5]) + redis.call('SET', KEYS[2], ARGV[1], 'PX', ARGV[5]) +end +return tostring(limit) +""" + + +class _LocalLimiter: + def __init__(self, ceiling: float) -> None: + self.limit = ceiling + self.inflight = 0 + self.last_decrease = 0.0 + self.cond = asyncio.Condition() + + async def acquire(self) -> None: + async with self.cond: + while self.inflight >= int(self.limit): + await self.cond.wait() + self.inflight += 1 + + async def release(self) -> None: + async with self.cond: + self.inflight = max(0, self.inflight - 1) + self.cond.notify(1) + + async def on_success(self, ceiling: float) -> float: + async with self.cond: + self.limit = min(self.limit + 1.0 / self.limit, ceiling) + self.cond.notify(1) + return self.limit + + async def on_rate_limited(self) -> float: + async with self.cond: + now = time.monotonic() + if now - self.last_decrease > DECREASE_GUARD_MS / 1000: + self.limit = max(self.limit * 0.5, MIN_LIMIT) + self.last_decrease = now + return self.limit + + +_local_limiters: dict[str, _LocalLimiter] = {} +_scripts: dict[str, tuple] = {} + + +def _ceiling() -> float: + from bigrag.services.runtime_settings import sync_value + + return max(float(sync_value("embedding_concurrency")), MIN_LIMIT) + + +def _digest(cache_identity: str) -> str: + return hashlib.sha256(str(cache_identity).encode()).hexdigest()[:24] + + +def _local(digest: str) -> _LocalLimiter: + limiter = _local_limiters.get(digest) + if limiter is None: + limiter = _LocalLimiter(_ceiling()) + _local_limiters[digest] = limiter + return limiter + + +def _script(redis, name: str, body: str): + cached = _scripts.get(name) + if cached is not None and cached[0] is redis: + return cached[1] + script = redis.register_script(body) + _scripts[name] = (redis, script) + return script + + +def _as_float(raw) -> float: + if isinstance(raw, (bytes, bytearray)): + return float(raw.decode()) + return float(raw) + + +def reset_embedding_limiters() -> None: + _local_limiters.clear() + + +async def _acquire(redis, digest: str) -> str: + if redis is None: + await _local(digest).acquire() + return _LOCAL_TOKEN + inflight_key = INFLIGHT_PREFIX + digest + limit_key = LIMIT_PREFIX + digest + ceiling = _ceiling() + script = _script(redis, "acquire", _ACQUIRE_LUA) + while True: + token = uuid.uuid4().hex + now_ms = int(time.time() * 1000) + try: + ok = await script( + keys=[inflight_key, limit_key], + args=[ + now_ms, + now_ms + LEASE_SECONDS * 1000, + token, + ceiling, + LEASE_SECONDS * 1000 * 2, + LIMIT_TTL_MS, + ], + ) + except Exception as exc: + logger.debug("embedding gate acquire fell back", error=repr(exc)) + return _FAILOPEN_TOKEN + if int(ok) == 1: + return token + await asyncio.sleep(random.uniform(ACQUIRE_RETRY_MIN_MS, ACQUIRE_RETRY_MAX_MS) / 1000) + + +async def _release(redis, digest: str, token: str) -> None: + if token == _LOCAL_TOKEN: + await _local(digest).release() + return + if token == _FAILOPEN_TOKEN or redis is None: + return + try: + await redis.zrem(INFLIGHT_PREFIX + digest, token) + except Exception as exc: + logger.debug("embedding gate release failed", error=repr(exc)) + + +async def _on_success(redis, digest: str) -> None: + if redis is None: + await _local(digest).on_success(_ceiling()) + return + try: + script = _script(redis, "success", _SUCCESS_LUA) + await script(keys=[LIMIT_PREFIX + digest], args=[_ceiling(), LIMIT_TTL_MS]) + except Exception as exc: + logger.debug("embedding gate success update failed", error=repr(exc)) + + +async def _on_rate_limited( + redis, digest: str, cooldown_key: str, exc: Exception, provider: str, model_name: str +) -> None: + await record_rate_limit_cooldown(cooldown_key, rate_limit_delay(exc, 1.0)) + if redis is None: + new_limit = await _local(digest).on_rate_limited() + else: + try: + script = _script(redis, "decrease", _DECREASE_LUA) + now_ms = int(time.time() * 1000) + raw = await script( + keys=[LIMIT_PREFIX + digest, LIMIT_DEC_PREFIX + digest], + args=[now_ms, _ceiling(), MIN_LIMIT, DECREASE_GUARD_MS, LIMIT_TTL_MS], + ) + new_limit = _as_float(raw) + except Exception as update_exc: + logger.debug("embedding gate decrease failed", error=repr(update_exc)) + return + logger.warning( + "embedding limit decreased", + provider=provider, + model=model_name, + new_limit=round(new_limit, 2), + ) + + +@asynccontextmanager +async def embedding_gate(cache_identity: str, provider: str, model_name: str): + digest = _digest(cache_identity) + cooldown_key = RATE_LIMIT_COOLDOWN_KEY_PREFIX + digest + await wait_for_rate_limit_cooldown(cooldown_key, provider, model_name) + redis = redis_cache.get_redis() + token = await _acquire(redis, digest) + err: Exception | None = None + try: + yield + except Exception as exc: + err = exc + raise + finally: + await _release(redis, digest, token) + if err is None: + await _on_success(redis, digest) + elif is_rate_limit_error(err): + await _on_rate_limited(redis, digest, cooldown_key, err, provider, model_name) +``` + +- [ ] **Step 2: Lint the new module** + +Run: `api/.venv/bin/ruff check api/bigrag/services/embedding_gate.py` +Expected: `All checks passed!` + +- [ ] **Step 3: Import smoke test** + +Run: `api/.venv/bin/python -c "import bigrag.services.embedding_gate as g; print(g.embedding_gate, g.reset_embedding_limiters)"` +Expected: prints the two function objects, no traceback. (If it fails due to missing env/settings, note it and rely on ruff — the module has no import-time side effects beyond importing `redis_cache`.) + +- [ ] **Step 4: Commit** + +```bash +git add api/bigrag/services/embedding_gate.py +git commit -m "feat: add adaptive embedding rate-limit gate" +``` + +--- + +### Task 2: Route all three providers through the gate + +**Files:** +- Modify: `api/bigrag/services/embedding/openai.py:1-13,46,93-118` +- Modify: `api/bigrag/services/embedding/cohere.py:1-13,38,68-95` +- Modify: `api/bigrag/services/embedding/voyage.py:1-15,42,92-142` + +- [ ] **Step 1: Update `openai.py` imports** + +Replace `api/bigrag/services/embedding/openai.py:6-13`: + +```python +from bigrag.services.embedding.base import EmbeddingModel, get_semaphore, logger, truncate_to_tokens +from bigrag.services.embedding_rate_limit import ( + is_rate_limit_error, + rate_limit_cooldown_key, + rate_limit_delay, + record_rate_limit_cooldown, + wait_for_rate_limit_cooldown, +) +``` + +with: + +```python +from bigrag.services.embedding.base import EmbeddingModel, logger, truncate_to_tokens +from bigrag.services.embedding_gate import embedding_gate +``` + +- [ ] **Step 2: Delete the unused semaphore-key line in `openai.py`** + +Delete `api/bigrag/services/embedding/openai.py:46`: + +```python + self._semaphore_key = f"openai:{self._base_url or 'default'}" +``` + +- [ ] **Step 3: Replace `openai.py` `_embed_single`** + +Replace `api/bigrag/services/embedding/openai.py:93-118`: + +```python + async def _embed_single(self, texts: list[str]) -> list[list[float]]: + cooldown_key = rate_limit_cooldown_key( + self._cache_identity, self.provider, self._model_name, self._dimension + ) + kwargs: dict = {"input": texts, "model": self._model_name} + if self._supports_dimensions(self._model_name): + kwargs["dimensions"] = self._dimension + async with await get_semaphore(self._semaphore_key): + await wait_for_rate_limit_cooldown(cooldown_key, self.provider, self._model_name) + try: + response = await asyncio.wait_for( + self._client.embeddings.create(**kwargs), + timeout=60, + ) + except Exception as exc: + if is_rate_limit_error(exc): + await record_rate_limit_cooldown(cooldown_key, rate_limit_delay(exc, 1.0)) + raise + vectors = [item.embedding for item in response.data] + for vector in vectors: + if len(vector) != self._dimension: + raise ValueError( + f"openai returned vector of length {len(vector)}, " + f"expected {self._dimension} for model {self._model_name}" + ) + return vectors +``` + +with: + +```python + async def _embed_single(self, texts: list[str]) -> list[list[float]]: + kwargs: dict = {"input": texts, "model": self._model_name} + if self._supports_dimensions(self._model_name): + kwargs["dimensions"] = self._dimension + async with embedding_gate(self._cache_identity, self.provider, self._model_name): + response = await asyncio.wait_for( + self._client.embeddings.create(**kwargs), + timeout=60, + ) + vectors = [item.embedding for item in response.data] + for vector in vectors: + if len(vector) != self._dimension: + raise ValueError( + f"openai returned vector of length {len(vector)}, " + f"expected {self._dimension} for model {self._model_name}" + ) + return vectors +``` + +- [ ] **Step 4: Update `cohere.py` imports** + +Replace `api/bigrag/services/embedding/cohere.py:5-12`: + +```python +from bigrag.services.embedding.base import EmbeddingModel, get_semaphore, logger, truncate_to_tokens +from bigrag.services.embedding_rate_limit import ( + is_rate_limit_error, + rate_limit_cooldown_key, + rate_limit_delay, + record_rate_limit_cooldown, + wait_for_rate_limit_cooldown, +) +``` + +with: + +```python +from bigrag.services.embedding.base import EmbeddingModel, logger, truncate_to_tokens +from bigrag.services.embedding_gate import embedding_gate +``` + +- [ ] **Step 5: Delete the unused semaphore-key line in `cohere.py`** + +Delete `api/bigrag/services/embedding/cohere.py:38`: + +```python + self._semaphore_key = "cohere" +``` + +- [ ] **Step 6: Replace `cohere.py` `_embed_single`** + +Replace `api/bigrag/services/embedding/cohere.py:68-95`: + +```python + async def _embed_single(self, texts: list[str], cohere_input_type: str) -> list[list[float]]: + cooldown_key = rate_limit_cooldown_key( + self._cache_identity, self.provider, self._model_name, self._dimension + ) + async with await get_semaphore(self._semaphore_key): + await wait_for_rate_limit_cooldown(cooldown_key, self.provider, self._model_name) + try: + response = await asyncio.wait_for( + self._client.embed( + texts=texts, + model=self._model_name, + input_type=cohere_input_type, + embedding_types=["float"], + ), + timeout=60, + ) + except Exception as exc: + if is_rate_limit_error(exc): + await record_rate_limit_cooldown(cooldown_key, rate_limit_delay(exc, 1.0)) + raise + vectors = [list(e) for e in response.embeddings.float_] + for vector in vectors: + if len(vector) != self._dimension: + raise ValueError( + f"cohere returned vector of length {len(vector)}, " + f"expected {self._dimension} for model {self._model_name}" + ) + return vectors +``` + +with: + +```python + async def _embed_single(self, texts: list[str], cohere_input_type: str) -> list[list[float]]: + async with embedding_gate(self._cache_identity, self.provider, self._model_name): + response = await asyncio.wait_for( + self._client.embed( + texts=texts, + model=self._model_name, + input_type=cohere_input_type, + embedding_types=["float"], + ), + timeout=60, + ) + vectors = [list(e) for e in response.embeddings.float_] + for vector in vectors: + if len(vector) != self._dimension: + raise ValueError( + f"cohere returned vector of length {len(vector)}, " + f"expected {self._dimension} for model {self._model_name}" + ) + return vectors +``` + +- [ ] **Step 7: Update `voyage.py` imports** + +Replace `api/bigrag/services/embedding/voyage.py:7-14`: + +```python +from bigrag.services.embedding.base import EmbeddingModel, get_semaphore, logger, truncate_to_tokens +from bigrag.services.embedding_rate_limit import ( + is_rate_limit_error, + rate_limit_cooldown_key, + rate_limit_delay, + record_rate_limit_cooldown, + wait_for_rate_limit_cooldown, +) +``` + +with (all `embedding_rate_limit` imports are dropped — Voyage no longer records cooldowns itself; the gate inspects the raised `VoyageHTTPError` and handles the 429): + +```python +from bigrag.services.embedding.base import EmbeddingModel, logger, truncate_to_tokens +from bigrag.services.embedding_gate import embedding_gate +``` + +- [ ] **Step 8: Delete the unused semaphore-key line in `voyage.py`** + +Delete `api/bigrag/services/embedding/voyage.py:42`: + +```python + self._semaphore_key = "voyage" +``` + +- [ ] **Step 9: Replace `voyage.py` `_embed_single`** + +Replace `api/bigrag/services/embedding/voyage.py:92-142`: + +```python + async def _embed_single(self, texts: list[str], voyage_input_type: str) -> list[list[float]]: + payload = { + "input": texts, + "model": self._model_name, + "input_type": voyage_input_type, + "output_dimension": self._dimension, + } + headers = { + "Authorization": f"Bearer {self._api_key}", + "Content-Type": "application/json", + } + cooldown_key = rate_limit_cooldown_key( + self._cache_identity, self.provider, self._model_name, self._dimension + ) + client = await self._get_client() + async with await get_semaphore(self._semaphore_key): + await wait_for_rate_limit_cooldown(cooldown_key, self.provider, self._model_name) + try: + response = await client.post( + f"{self._DEFAULT_BASE_URL}{self._EMBEDDINGS_PATH}", + json=payload, + headers=headers, + ) + except Exception as exc: + if is_rate_limit_error(exc): + await record_rate_limit_cooldown(cooldown_key, rate_limit_delay(exc, 1.0)) + raise + if response.status_code >= 400: + logger.warning( + "voyage embed http error", + status=response.status_code, + body_preview=response.text[:500], + model=self._model_name, + ) + exc = VoyageHTTPError( + response.status_code, + f"Voyage embed failed ({response.status_code})", + ) + exc.headers = response.headers + if is_rate_limit_error(exc): + await record_rate_limit_cooldown(cooldown_key, rate_limit_delay(exc, 1.0)) + raise exc + data = response.json() + vectors = [item["embedding"] for item in data["data"]] + for vector in vectors: + if len(vector) != self._dimension: + raise ValueError( + f"voyage returned vector of length {len(vector)}, " + f"expected {self._dimension} for model {self._model_name}" + ) + return vectors +``` + +with: + +```python + async def _embed_single(self, texts: list[str], voyage_input_type: str) -> list[list[float]]: + payload = { + "input": texts, + "model": self._model_name, + "input_type": voyage_input_type, + "output_dimension": self._dimension, + } + headers = { + "Authorization": f"Bearer {self._api_key}", + "Content-Type": "application/json", + } + client = await self._get_client() + async with embedding_gate(self._cache_identity, self.provider, self._model_name): + response = await client.post( + f"{self._DEFAULT_BASE_URL}{self._EMBEDDINGS_PATH}", + json=payload, + headers=headers, + ) + if response.status_code >= 400: + logger.warning( + "voyage embed http error", + status=response.status_code, + body_preview=response.text[:500], + model=self._model_name, + ) + exc = VoyageHTTPError( + response.status_code, + f"Voyage embed failed ({response.status_code})", + ) + exc.headers = response.headers + raise exc + data = response.json() + vectors = [item["embedding"] for item in data["data"]] + for vector in vectors: + if len(vector) != self._dimension: + raise ValueError( + f"voyage returned vector of length {len(vector)}, " + f"expected {self._dimension} for model {self._model_name}" + ) + return vectors +``` + +Why this is correct: the `VoyageHTTPError` raised inside the `async with embedding_gate(...)` block propagates through the context manager. The gate's `finally` calls `is_rate_limit_error(err)` on it; because `VoyageHTTPError` carries `.status_code` and `.headers`, a 429 is detected there and the cooldown + AIMD decrease are applied — replacing the manual `record_rate_limit_cooldown` calls removed here. + +- [ ] **Step 10: Lint changed providers** + +Run: `api/.venv/bin/ruff check api/bigrag/services/embedding/openai.py api/bigrag/services/embedding/cohere.py api/bigrag/services/embedding/voyage.py` +Expected: `All checks passed!` (no unused-import warnings — confirms the imports were trimmed correctly). + +- [ ] **Step 11: Commit** + +```bash +git add api/bigrag/services/embedding/openai.py api/bigrag/services/embedding/cohere.py api/bigrag/services/embedding/voyage.py +git commit -m "feat: route embedding providers through the adaptive gate" +``` + +--- + +### Task 3: Remove the old per-process semaphores + +**Files:** +- Modify: `api/bigrag/services/embedding/base.py:1-11,30-42` +- Modify: `api/bigrag/services/embedding/__init__.py:3-7,17-27` +- Modify: `api/bigrag/services/runtime_settings_apply.py:10,77-78` + +- [ ] **Step 1: Strip semaphore code from `base.py`** + +In `api/bigrag/services/embedding/base.py`, delete: +- the `import asyncio` line (line 3) — becomes unused; +- the module globals `_embed_semaphores` and `_embed_semaphores_lock` (lines 10-11); +- the `get_semaphore` function (lines 30-38); +- the `reset_embedding_semaphores` function (lines 41-42). + +After the edit the top of the file reads exactly: + +```python +from __future__ import annotations + +from abc import ABC, abstractmethod + +from bigrag.logging import get_logger + +logger = get_logger("bigrag.embedding") + +_TOKEN_LIMITS: dict[str, int] = { +``` + +and `truncate_to_tokens` follows directly after the `_TOKEN_LIMITS` dict (nothing between the dict and that function). + +- [ ] **Step 2: Update `embedding/__init__.py`** + +Replace `api/bigrag/services/embedding/__init__.py:3-7`: + +```python +from bigrag.services.embedding.base import ( + EmbeddingModel, + reset_embedding_semaphores, + truncate_to_tokens, +) +``` + +with: + +```python +from bigrag.services.embedding.base import ( + EmbeddingModel, + truncate_to_tokens, +) +from bigrag.services.embedding_gate import reset_embedding_limiters +``` + +Then in `__all__`, replace `"reset_embedding_semaphores",` with `"reset_embedding_limiters",`. + +- [ ] **Step 3: Update `runtime_settings_apply.py`** + +At `api/bigrag/services/runtime_settings_apply.py:10`, replace: + +```python +from bigrag.services.embedding import reset_embedding_semaphores +``` + +with: + +```python +from bigrag.services.embedding import reset_embedding_limiters +``` + +At `api/bigrag/services/runtime_settings_apply.py:77-78`, replace: + +```python + if "embedding_concurrency" in keyset: + reset_embedding_semaphores() +``` + +with: + +```python + if "embedding_concurrency" in keyset: + reset_embedding_limiters() +``` + +- [ ] **Step 4: Confirm no dangling references** + +Run: `grep -rn "reset_embedding_semaphores\|get_semaphore" api/bigrag/services/embedding api/bigrag/services/runtime_settings_apply.py` +Expected: no output. (The only other `get_semaphore` is in `api/bigrag/services/webhook/`, unrelated and untouched.) + +- [ ] **Step 5: Lint** + +Run: `api/.venv/bin/ruff check api/bigrag/services/embedding/base.py api/bigrag/services/embedding/__init__.py api/bigrag/services/runtime_settings_apply.py` +Expected: `All checks passed!` + +- [ ] **Step 6: Commit** + +```bash +git add api/bigrag/services/embedding/base.py api/bigrag/services/embedding/__init__.py api/bigrag/services/runtime_settings_apply.py +git commit -m "refactor: drop per-process embedding semaphores" +``` + +--- + +### Task 4: Remove now-redundant cooldown logic from the ingestion path + +The gate owns cooldown waiting and recording. Remove the duplicate handling and the vestigial `cooldown_key` threading. + +**Files:** +- Modify: `api/bigrag/services/queue_embedding/embed.py:1-21,51-96` +- Modify: `api/bigrag/services/queue_embedding/embed_batches.py:1-19,22-30,80` +- Modify: `api/bigrag/services/queue_embedding/plan.py:9,19,63-68,135` +- Modify: `api/bigrag/services/queue_embedding/insert.py:31-38` + +- [ ] **Step 1: Simplify `embed.py` imports/header** + +Replace `api/bigrag/services/queue_embedding/embed.py:1-21`: + +```python +from __future__ import annotations + +import asyncio +import math +import time + +from bigrag.logging import get_logger +from bigrag.services import embedding_cache +from bigrag.services.embedding import truncate_to_tokens +from bigrag.services.embedding_rate_limit import ( + is_rate_limit_error, + rate_limit_cooldown_key, + rate_limit_delay, + record_rate_limit_cooldown, + wait_for_rate_limit_cooldown, +) + +logger = get_logger("bigrag.queue") + +EMBEDDING_TIMEOUT_SECONDS = 60 +PERMANENT_ERRORS = (ValueError, UnicodeDecodeError, KeyError) +``` + +with: + +```python +from __future__ import annotations + +import math +import time + +from bigrag.logging import get_logger +from bigrag.services import embedding_cache +from bigrag.services.embedding import truncate_to_tokens + +logger = get_logger("bigrag.queue") + +PERMANENT_ERRORS = (ValueError, UnicodeDecodeError, KeyError) +``` + +(`asyncio`, `EMBEDDING_TIMEOUT_SECONDS`, and all `embedding_rate_limit` imports become unused — the gate handles cooldowns and the per-call timeout already lives inside each provider's `_embed_single`.) + +- [ ] **Step 2: Simplify the provider call in `embed.py`** + +Replace `api/bigrag/services/queue_embedding/embed.py:51-96` (the `if missing_idx:` block through the function's final `return`): + +```python + if missing_idx: + missing_by_cache_text: dict[str, int] = {} + for idx in missing_idx: + missing_by_cache_text.setdefault(cache_texts[idx], idx) + provider_idx = list(missing_by_cache_text.values()) + missing_texts = [texts[i] for i in provider_idx] + missing_cache_texts = [cache_texts[i] for i in provider_idx] + cooldown_key = rate_limit_cooldown_key(model, provider, model_name, dimension) + await wait_for_rate_limit_cooldown(cooldown_key, provider, model_name) + t0 = time.monotonic() + logger.debug( + "embedding provider request", + provider=provider, + model=model_name, + inputs=len(missing_texts), + ) + try: + fresh = await asyncio.wait_for( + model.embed(missing_texts, input_type=input_type), + timeout=EMBEDDING_TIMEOUT_SECONDS, + ) + except Exception as exc: + if is_rate_limit_error(exc): + await record_rate_limit_cooldown(cooldown_key, rate_limit_delay(exc, 1.0)) + raise + logger.debug( + "embedding provider response", + provider=provider, + model=model_name, + inputs=len(missing_texts), + elapsed=round(time.monotonic() - t0, 2), + ) + if len(fresh) != len(missing_texts): + raise ValueError( + f"embedding provider returned {len(fresh)} vectors for {len(missing_texts)} inputs" + ) + for vec in fresh: + if any(not math.isfinite(v) for v in vec): + raise ValueError("embedding provider returned non-finite values") + await embedding_cache.put_many( + missing_cache_texts, fresh, model.cache_identity, dimension, input_type + ) + fresh_by_cache_text = dict(zip(missing_cache_texts, fresh, strict=False)) + for idx in missing_idx: + cached[idx] = fresh_by_cache_text[cache_texts[idx]] + return [cached[i] for i in range(len(texts))] +``` + +with: + +```python + if missing_idx: + missing_by_cache_text: dict[str, int] = {} + for idx in missing_idx: + missing_by_cache_text.setdefault(cache_texts[idx], idx) + provider_idx = list(missing_by_cache_text.values()) + missing_texts = [texts[i] for i in provider_idx] + missing_cache_texts = [cache_texts[i] for i in provider_idx] + t0 = time.monotonic() + logger.debug( + "embedding provider request", + provider=provider, + model=model_name, + inputs=len(missing_texts), + ) + fresh = await model.embed(missing_texts, input_type=input_type) + logger.debug( + "embedding provider response", + provider=provider, + model=model_name, + inputs=len(missing_texts), + elapsed=round(time.monotonic() - t0, 2), + ) + if len(fresh) != len(missing_texts): + raise ValueError( + f"embedding provider returned {len(fresh)} vectors for {len(missing_texts)} inputs" + ) + for vec in fresh: + if any(not math.isfinite(v) for v in vec): + raise ValueError("embedding provider returned non-finite values") + await embedding_cache.put_many( + missing_cache_texts, fresh, model.cache_identity, dimension, input_type + ) + fresh_by_cache_text = dict(zip(missing_cache_texts, fresh, strict=False)) + for idx in missing_idx: + cached[idx] = fresh_by_cache_text[cache_texts[idx]] + return [cached[i] for i in range(len(texts))] +``` + +- [ ] **Step 3: Trim `embed_batches.py` imports** + +Replace `api/bigrag/services/queue_embedding/embed_batches.py:7-12`: + +```python +from bigrag.services.embedding_rate_limit import ( + MAX_RATE_LIMIT_RETRIES, + is_rate_limit_error, + rate_limit_delay, + record_rate_limit_cooldown, +) +``` + +with: + +```python +from bigrag.services.embedding_rate_limit import ( + MAX_RATE_LIMIT_RETRIES, + is_rate_limit_error, + rate_limit_delay, +) +``` + +- [ ] **Step 4: Drop `cooldown_key` from `embed_all_batches` signature** + +Replace `api/bigrag/services/queue_embedding/embed_batches.py:22-30`: + +```python +async def embed_all_batches( + job: IngestionJob, + prefix: str, + *, + embedding_model, + cooldown_key: str, + batches: list[tuple[int, int, int, list]], + total_batches: int, +) -> list[tuple[int, int, int, list, list[list[float]], float]]: +``` + +with: + +```python +async def embed_all_batches( + job: IngestionJob, + prefix: str, + *, + embedding_model, + batches: list[tuple[int, int, int, list]], + total_batches: int, +) -> list[tuple[int, int, int, list, list[list[float]], float]]: +``` + +- [ ] **Step 5: Remove the redundant cooldown write in the retry loop** + +In `api/bigrag/services/queue_embedding/embed_batches.py`, delete the single line (line 80): + +```python + await record_rate_limit_cooldown(cooldown_key, delay) +``` + +Leave the `fallback_delay` / `delay` computation, the `logger.warning(...)`, and `await asyncio.sleep(delay)` exactly as they are — the retry/backoff loop remains the safety net; only the cooldown *write* is removed (the gate already recorded it). + +- [ ] **Step 6: Drop `cooldown_key` from `plan.py`** + +In `api/bigrag/services/queue_embedding/plan.py`: +- delete the import at line 9: `from bigrag.services.embedding_rate_limit import rate_limit_cooldown_key`; +- delete the `cooldown_key: str` field from the `EmbedPlan` dataclass (line 19); +- delete the cooldown_key computation block (lines 63-68): + +```python + cooldown_key = rate_limit_cooldown_key( + embedding_model, + job.embedding_provider, + job.embedding_model, + job.embedding_dimension, + ) +``` + +- delete the `cooldown_key=cooldown_key,` line inside the `EmbedPlan(...)` return (line 135). + +- [ ] **Step 7: Drop `cooldown_key` from the `insert.py` call** + +Replace `api/bigrag/services/queue_embedding/insert.py:31-38`: + +```python + embed_results = await embed_all_batches( + job, + prefix, + embedding_model=plan.embedding_model, + cooldown_key=plan.cooldown_key, + batches=plan.batches, + total_batches=plan.total_batches, + ) +``` + +with: + +```python + embed_results = await embed_all_batches( + job, + prefix, + embedding_model=plan.embedding_model, + batches=plan.batches, + total_batches=plan.total_batches, + ) +``` + +- [ ] **Step 8: Lint the ingestion path** + +Run: `api/.venv/bin/ruff check api/bigrag/services/queue_embedding/` +Expected: `All checks passed!` (verifies no unused imports / undefined names remain). + +- [ ] **Step 9: Commit** + +```bash +git add api/bigrag/services/queue_embedding/embed.py api/bigrag/services/queue_embedding/embed_batches.py api/bigrag/services/queue_embedding/plan.py api/bigrag/services/queue_embedding/insert.py +git commit -m "refactor: let the gate own embedding cooldowns" +``` + +--- + +### Task 5: Document the semantic change of `embedding_concurrency` + +**Files:** +- Modify: `api/bigrag/services/runtime_setting_specs/search.py:12` + +- [ ] **Step 1: Update the setting description** + +In `api/bigrag/services/runtime_setting_specs/search.py`, replace line 12: + +```python + description="Maximum concurrent embedding requests per provider endpoint.", +``` + +with: + +```python + description="Global ceiling on concurrent embedding requests per endpoint across all workers; the limiter backs off below this on rate limits and recovers toward it.", +``` + +- [ ] **Step 2: Lint** + +Run: `api/.venv/bin/ruff check api/bigrag/services/runtime_setting_specs/search.py` +Expected: `All checks passed!` + +- [ ] **Step 3: Commit** + +```bash +git add api/bigrag/services/runtime_setting_specs/search.py +git commit -m "docs: clarify embedding_concurrency is a global adaptive ceiling" +``` + +--- + +### Task 6: Full-package lint + manual verification + +**Files:** none (verification only). + +- [ ] **Step 1: Lint the whole touched surface** + +Run: `api/.venv/bin/ruff check api/bigrag` +Expected: `All checks passed!` + +- [ ] **Step 2: Format check** + +Run: `api/.venv/bin/ruff format --check api/bigrag/services/embedding_gate.py api/bigrag/services/embedding api/bigrag/services/queue_embedding` +Expected: `N files already formatted`. If it lists files that would be reformatted, run `api/.venv/bin/ruff format `, review the diff, and amend the relevant commit. + +- [ ] **Step 3: Manual load-test verification (requires a running stack + Redis)** + +With the API + workers running against a local Redis, drive ingestion using the existing load test in `tests/load/` (see `tests/load/test.sh`). While it runs, watch the worker logs: + +- Confirm `embedding limit decreased` warnings appear *when* a 429 occurs (the gate detects and backs off). +- Confirm the prior flood of `batch rate limited` / repeated 429 warnings drops sharply after the first backoff (the herd is tamed). +- Optionally inspect Redis: `redis-cli --scan --pattern 'bigrag:embedding:limit:*'`, then `GET ` — the live limit should settle below `embedding_concurrency` after 429s and recover toward it. + +Tuning: if 429s remain frequent, lower `embedding_concurrency` (now global); if throughput is low and there are no 429s, raise it. + +- [ ] **Step 4: Push the branch** + +```bash +git push +``` + +--- + +## Notes for the implementer + +- **No tests:** this repo has no unit-test framework (only `tests/load`). Verification is ruff + the load test, per project convention. +- **Lua numeric returns truncate to int** and the Redis client uses `decode_responses=False`, so `_SUCCESS_LUA` / `_DECREASE_LUA` `return tostring(limit)` and Python parses the resulting **bytes** via `_as_float` (`float(raw.decode())`). Do not change these to return numbers. +- **`register_script` not direct script execution:** scripts are registered once per client (`_script` helper, cached in `_scripts`) so the redis client uses EVALSHA under the hood. `_acquire` returns int `1`/`0` directly. +- **Fail-open is deliberate:** if a Redis script call raises, `_acquire` returns `_FAILOPEN_TOKEN` and admits the request rather than stalling ingestion. Release/success/decrease for that token are no-ops. +- **Permit leaks self-heal:** each permit is a ZSET member scored with its expiry; `_acquire` reaps expired members first, so a crashed worker frees its slots after `LEASE_SECONDS`. +- **Per-job `EMBED_CONCURRENCY` stays:** the local per-job fan-out bound in `embed_batches.py` is left intact; it is now subordinate to the global gate. +- **The webhook `get_semaphore`** (`api/bigrag/services/webhook/http.py`) is a different, unrelated function — do not touch it. +``` diff --git a/docs/superpowers/specs/2026-05-24-adaptive-embedding-rate-limiting-design.md b/docs/superpowers/specs/2026-05-24-adaptive-embedding-rate-limiting-design.md new file mode 100644 index 00000000..dcea994e --- /dev/null +++ b/docs/superpowers/specs/2026-05-24-adaptive-embedding-rate-limiting-design.md @@ -0,0 +1,258 @@ +# Adaptive Embedding Rate Limiting + +**Date:** 2026-05-24 +**Status:** Design — pending review + +## Problem + +Production logs show frequent embedding `429` hits. Repeatedly hammering a +provider after `429`s risks an abuse ban, so the goal is to **minimize how often +we hit the limit**, not just survive it. + +The current handling is purely reactive and uncoordinated: + +1. **No proactive pacing.** Requests fire at full speed until a `429` returns. + Only then does `record_rate_limit_cooldown` write a Redis cooldown + (`embed.py:74`, `openai.py:109`, `cohere.py:86`, `voyage.py:117/132`). Steady + state is a sawtooth: blast → `429` → everyone waits `Retry-After` → cooldown + expires → everyone blasts again in lockstep → immediate re-`429`. + +2. **Concurrency limits are per-process.** `EMBED_CONCURRENCY = 8` + (`embed_batches.py:19`, a fresh semaphore per job) and `get_semaphore` + (`base.py:30`, default 8) live in process-local memory bound to each thread's + event loop. With **5 worker processes × 8 threads**, true in-flight + concurrency to the provider is dozens, not 8. The `embedding_concurrency` + runtime setting is effectively meaningless at this scale. + +3. **TPM is never modeled.** OpenAI embeddings are gated mainly by + tokens-per-minute. A batch can be 512 chunks → up to 2048 inputs/request. + Nothing counts tokens, so even modest concurrency trips TPM. + +## Constraints (from brainstorming) + +- **Must work for all providers** (OpenAI, Cohere, Voyage, OpenAI-compatible). + Testing with OpenAI first. +- **No hardcoded or fetched RPM/TPM limits.** Configuring or fetching quotas is + considered brittle and wasteful. The only source of truth for "how fast can we + go" is the `429` signal itself. +- Keep the existing `Retry-After` cooldown behavior ("wait the time, continue"). +- Must coordinate across the **5 workers** that share one API key. + +## Approach: self-tuning adaptive concurrency (AIMD) + +Because we refuse to configure or fetch the limit, the safe rate must be +**learned from feedback**. We use AIMD (additive-increase / multiplicative +decrease, the TCP congestion-control pattern): a shared, dynamically-sized +concurrency limit per endpoint, stored in Redis so all 5 workers obey the same +ceiling. + +- **On success:** nudge the allowed concurrency up. +- **On `429`:** cut it multiplicatively *and* set the existing `Retry-After` + cooldown. + +This replaces the fixed per-process semaphores with one cross-worker limiter. +**No token counting is needed:** large (high-token) requests trip `429`s sooner, +which lowers concurrency automatically — so AIMD adapts to TPM pressure +implicitly. The net behavior changes from `blast → 429 → wait → blast → 429` +to `ramp up → find the edge → one 429 → back off → settle just under the edge`. + +## Architecture + +### New module: `services/embedding_gate.py` + +A single async context manager that every provider call site uses. It absorbs +the duplicated `get_semaphore` + `wait_for_rate_limit_cooldown` + +`record_rate_limit_cooldown` logic currently copy-pasted across the three +providers. + +```python +@asynccontextmanager +async def embedding_gate(endpoint_key: str, provider: str, model_name: str): + await wait_for_rate_limit_cooldown(cooldown_key, provider, model_name) # existing + token = await _acquire(endpoint_key) # block until an adaptive slot frees + try: + yield # caller runs the HTTP request here + except Exception as exc: + if is_rate_limit_error(exc): + await _on_rate_limited(endpoint_key, exc) # AIMD decrease + record cooldown + raise + else: + await _on_success(endpoint_key) # AIMD increase + finally: + await _release(endpoint_key, token) # always free the slot +``` + +`endpoint_key` is derived from `cache_identity` (already +`provider:model:dimension[:base_tag]`), so each distinct provider/model/endpoint +gets its own independent limiter — matching the existing cooldown key. + +### Redis state (per endpoint) + +| Key | Type | Purpose | +|-----|------|---------| +| `bigrag:embedding:inflight:{hash}` | ZSET | in-flight permits; member = unique token, score = lease expiry (ms) | +| `bigrag:embedding:limit:{hash}` | string (float) | current AIMD concurrency limit | +| `bigrag:embedding:limit-dec:{hash}` | string (ms) | timestamp of last decrease (decrease guard) | +| `bigrag:embedding:rate-limit:{hash}` | string | **existing** `Retry-After` cooldown (unchanged) | + +The ZSET doubles as a **self-healing lease**: each acquire reaps members whose +score (expiry) is in the past, so a crashed worker's permits free themselves +after the request timeout (`EMBEDDING_TIMEOUT_SECONDS = 60`). This prevents +permanent permit leaks / deadlock. + +### Atomic operations (Lua via `redis.eval`) + +**Acquire** — reap expired permits, init limit if missing, admit only if under +limit: + +```lua +redis.call('ZREMRANGEBYSCORE', KEYS[1], '-inf', ARGV[1]) -- reap (now) +local limit = tonumber(redis.call('GET', KEYS[2])) +if limit == nil then limit = tonumber(ARGV[4]); redis.call('SET', KEYS[2], limit) end +local count = redis.call('ZCARD', KEYS[1]) +if count < math.floor(limit) then + redis.call('ZADD', KEYS[1], ARGV[2], ARGV[3]) -- score = lease expiry + redis.call('PEXPIRE', KEYS[1], ARGV[5]) + return 1 +end +return 0 +``` + +Caller loop: if `0`, sleep a short jittered interval (e.g. 25–100 ms) and retry. +This is the proactive throttle — workers wait *before* sending instead of after +a `429`. + +**On success** — release token, additive increase (TCP-style: `+1/limit` per +success, so it takes ~`limit` successes to add one slot): + +```lua +redis.call('ZREM', KEYS[1], ARGV[3]) +local limit = tonumber(redis.call('GET', KEYS[2])) or tonumber(ARGV[4]) +limit = math.min(limit + 1.0/limit, tonumber(ARGV[5])) -- cap at ceiling +redis.call('SET', KEYS[2], limit) +``` + +**On `429`** — release token, multiplicative decrease with a guard so a burst of +simultaneous `429`s applies **one** halving, not many: + +```lua +redis.call('ZREM', KEYS[1], ARGV[3]) +local last = tonumber(redis.call('GET', KEYS[3])) or 0 +if tonumber(ARGV[1]) - last > tonumber(ARGV[6]) then -- guard window + local limit = tonumber(redis.call('GET', KEYS[2])) or tonumber(ARGV[4]) + limit = math.max(limit * 0.5, tonumber(ARGV[7])) -- floor at MIN + redis.call('SET', KEYS[2], limit) + redis.call('SET', KEYS[3], ARGV[1]) +end +``` + +The existing `record_rate_limit_cooldown` is still called on `429` so all workers +also pause for the provider's `Retry-After`. + +### Bounds — reuse the existing `embedding_concurrency` setting + +We do **not** add new limit constants. The existing `embedding_concurrency` +runtime setting (`config.py:65`, default 8, `runtime_setting_specs/search.py:6`) +becomes the adaptive **ceiling**, with one important semantic change: + +> **It is now global across all workers, not per-process.** Today each of the 5 +> worker processes independently allows 8 concurrent calls (~40 effective). After +> this change, `embedding_concurrency` caps total concurrent calls *across the +> whole fleet* via the shared Redis limiter. Operators should expect to raise +> this value from 8 to preserve aggregate throughput. + +- **Ceiling = `sync_value("embedding_concurrency")`**, clamped to `>= 1`. Used as + both the initial seed and the recovery cap. AIMD only ever moves the live limit + *below* this and recovers back up to it. +- `MIN_LIMIT = 1` (constant floor). +- `DECREASE_GUARD_MS = 1000` (collapse a burst of concurrent 429s to one halving). +- `ACQUIRE_RETRY_MS = 25–100` jittered (poll interval when at capacity). +- lease = `EMBEDDING_TIMEOUT_SECONDS` (60s) — the self-healing permit TTL. + +So steady state with no 429s sits at the ceiling; a 429 halves the live limit +(floored at 1) and it recovers additively toward the ceiling. The ceiling is a +user knob, not a hardcoded provider limit. + +### Local fallback (no Redis) + +Mirror the cooldown module's existing pattern: a process-local +`{endpoint_key: AdaptiveLimiter}` using an `asyncio` counter + condition variable +with the same AIMD math. Keeps single-process/dev working. (The 5-worker +coordination benefit only applies with Redis, which production has.) + +## Call-site changes + +Replace the duplicated block in each provider's `_embed_single` with the gate: + +- `services/embedding/openai.py:93` — `OpenAIEmbedding._embed_single` +- `services/embedding/cohere.py:68` — `CohereEmbedding._embed_single` +- `services/embedding/voyage.py:92` — `VoyageEmbedding._embed_single` + +Before: +```python +async with await get_semaphore(self._semaphore_key): + await wait_for_rate_limit_cooldown(cooldown_key, ...) + try: + response = await asyncio.wait_for(call(...), timeout=60) + except Exception as exc: + if is_rate_limit_error(exc): + await record_rate_limit_cooldown(cooldown_key, rate_limit_delay(exc, 1.0)) + raise +``` + +After: +```python +async with embedding_gate(self._cache_identity, self.provider, self._model_name): + response = await asyncio.wait_for(call(...), timeout=60) +``` + +`get_semaphore` / `reset_embedding_semaphores` in `base.py` become unused and are +removed. `embed.py:embed_with_cache` wraps `model.embed`, which fans out to +`_embed_single` where the gate now lives — so its pre-call +`wait_for_rate_limit_cooldown` and its on-`429` `record_rate_limit_cooldown` +(`embed.py:59,74`) are removed to keep the gate the single point of admission and +avoid double-handling. The per-job `EMBED_CONCURRENCY` semaphore in +`embed_batches.py` stays as a local per-job fan-out bound but is now subordinate +to the global gate. + +Query-time embedding (`retrieval/cache.py:90,127`) calls `model.embed` → +`_embed_single`, so it inherits the gate automatically. No change there. + +## Observability + +Structured logs (existing `get_logger` style) when the limit changes: +`embedding limit decreased` / `embedding limit recovered`, with `endpoint`, +`old_limit`, `new_limit`, `inflight`. This makes "are we still hitting 429s" +directly visible and shows the limit settling at the discovered ceiling. + +## Error handling + +- **Redis unavailable:** fall back to the local adaptive limiter; never block + embedding on Redis errors. Wrap Lua calls; on exception, log and admit (fail + open) so a Redis blip can't stall ingestion. +- **Permit leak on crash:** healed by the lease TTL reap in `_acquire`. +- **Burst of concurrent 429s:** collapsed to a single halving by the decrease + guard, so `limit` can't crater to 1 from one bad moment. +- **`MAX_RATE_LIMIT_RETRIES` / backoff in `embed_batches.py`:** unchanged; still + the safety net if the learned limit lags a sudden quota cut. + +## Testing + +- **AIMD math:** unit tests for increase (`+1/limit`), decrease (`*0.5` floored), + and the decrease guard (N concurrent 429s → one halving). +- **Lua acquire/release:** against a real/fake Redis — admission under limit, + rejection at limit, expired-lease reaping. +- **Concurrency simulation:** many coroutines through the gate never exceed + `floor(limit)` in flight. +- **429 adaptation:** simulate provider 429s; assert limit halves, cooldown is + recorded, then recovers on sustained success. +- **Local fallback:** same behaviors with Redis disabled. +- **Provider integration:** each provider's `_embed_single` admits/blocks via the + gate (mock the HTTP client). + +## Out of scope + +- Configuring or fetching provider RPM/TPM (explicitly rejected). +- Token-budget accounting (handled implicitly by AIMD). +- Changing batch sizes or the Dramatiq worker topology. +``` diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index e6282600..15fd1581 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -111,10 +111,6 @@ importers: version: 8.0.13(@types/node@25.9.0)(esbuild@0.28.0)(jiti@2.7.0)(tsx@4.22.1) sdks/typescript: - dependencies: - ws: - specifier: ^8.18.3 - version: 8.20.1 devDependencies: '@types/node': specifier: ^25.9.0 @@ -3131,18 +3127,6 @@ packages: resolution: {integrity: sha512-YVGIj2kamLSTxw6NsZjoBxfSwsn0ycdesmc4p+Q21c5zPuZ1pl+NfxVdxPtdHvmNVOQ6XSYG4AUtyt/Fi7D16Q==} engines: {node: '>=10'} - ws@8.20.1: - resolution: {integrity: sha512-It4dO0K5v//JtTXuPkfEOaI3uUN87iYPnqo/ZzqCoG3g8uhA66QUMs/SrM0YK7/NAu+r4LMh/9dq2A7k+rHs+w==} - engines: {node: '>=10.0.0'} - peerDependencies: - bufferutil: ^4.0.1 - utf-8-validate: '>=5.0.2' - peerDependenciesMeta: - bufferutil: - optional: true - utf-8-validate: - optional: true - y18n@5.0.8: resolution: {integrity: sha512-0pfFzegeDWJHJIAmTLRP2DwHjdF5s7jo9tuztdQxAhINCdvS+3nGINqPd00AphqJR/0LhANUS6/+7SCb98YOfA==} engines: {node: '>=10'} @@ -6352,8 +6336,6 @@ snapshots: string-width: 4.2.3 strip-ansi: 6.0.1 - ws@8.20.1: {} - y18n@5.0.8: {} yallist@3.1.1: {} diff --git a/sdks/python/README.md b/sdks/python/README.md index 137447b1..ac5edb19 100644 --- a/sdks/python/README.md +++ b/sdks/python/README.md @@ -26,14 +26,14 @@ The SDK is fully typed, ships `py.typed`, and uses CalVer releases like `2026.5. ## Namespaces -- `client.collections` for collection CRUD, stats, realtime tokens, and event streams. +- `client.collections` for collection CRUD, stats, and analytics. - `client.documents` for uploads, batch operations, chunks, elements, and status polling. - `client.queries` for single, multi-collection, and batch retrieval queries. - `client.vectors` for raw vector upsert and delete. - `client.webhooks` for webhook management and delivery replay. - `client.auth` for session login, setup, preferences, and identity. -- `client.admin` for users, API keys, access logs, audit logs, runtime settings, vector storage overview, admin realtime helpers, connector config, embedding presets, and MCP server keys. -- `client.realtime` for explicit WebSocket connect, subscribe, and unsubscribe control. +- `client.admin` for users, API keys, access logs, audit logs, runtime settings, vector storage overview, connector config, embedding presets, and MCP server keys. +- Top-level status helpers (`get_overview_status`, `get_collections_status`, `get_usage_status`, `get_access_status`) for pollable admin UI aggregates. - `client.connectors.s3` for S3-compatible bucket-prefix sources and sync jobs. - `client.evaluations` for golden-set retrieval evaluations. diff --git a/sdks/python/pyproject.toml b/sdks/python/pyproject.toml index 5ccc7c21..19a3789b 100644 --- a/sdks/python/pyproject.toml +++ b/sdks/python/pyproject.toml @@ -22,7 +22,6 @@ classifiers = [ ] dependencies = [ "httpx>=0.28.0", - "websockets>=15.0,<17", ] [project.urls] diff --git a/sdks/python/src/bigrag/__init__.py b/sdks/python/src/bigrag/__init__.py index e84a2aa9..e245fcbb 100644 --- a/sdks/python/src/bigrag/__init__.py +++ b/sdks/python/src/bigrag/__init__.py @@ -22,7 +22,6 @@ error_for_status, ) from bigrag._files import FileInput, normalize_file_input -from bigrag._realtime import RealtimeConnection, RealtimeResource from bigrag._sse import parse_sse_frames from bigrag._version import __version__ from bigrag.resources import ( @@ -46,8 +45,6 @@ "BigRAG", "BigRAGCore", "CollectionClient", - "RealtimeConnection", - "RealtimeResource", "__version__", "APIConnectionError", "APIError", diff --git a/sdks/python/src/bigrag/_client.py b/sdks/python/src/bigrag/_client.py index c8743007..62df77b2 100644 --- a/sdks/python/src/bigrag/_client.py +++ b/sdks/python/src/bigrag/_client.py @@ -1,6 +1,5 @@ from __future__ import annotations -from collections.abc import AsyncGenerator from typing import Any from bigrag._core import BigRAGCore @@ -14,17 +13,16 @@ DocumentsResource, EvaluationsResource, QueryResource, - RealtimeResource, VectorsResource, WebhooksResource, ) +from bigrag.types.access import AccessLogOverviewResponse from bigrag.types.analytics import AnalyticsResponse -from bigrag.types.collections import ( - CollectionRealtimeTokenResponse, - CollectionStatsResponse, -) +from bigrag.types.collections import CollectionStatsResponse from bigrag.types.common import ( + CollectionsStatusResponse, HealthResponse, + OverviewStatusResponse, PlatformStatsResponse, ReadinessResponse, StatusResponse, @@ -45,7 +43,6 @@ QueryBody, QueryResponse, ) -from bigrag.types.realtime import ProgressEvent from bigrag.types.usage import UsageResponse @@ -60,7 +57,6 @@ class BigRAG(BigRAGCore): auth: AuthResource admin: AdminResource evaluations: EvaluationsResource - realtime: RealtimeResource def __init__(self, **kwargs: Any) -> None: super().__init__(**kwargs) @@ -74,7 +70,6 @@ def __init__(self, **kwargs: Any) -> None: self.auth = AuthResource(self) self.admin = AdminResource(self) self.evaluations = EvaluationsResource(self) - self.realtime = RealtimeResource(self) async def health(self) -> HealthResponse: return await self._request("GET", "/health") @@ -94,6 +89,26 @@ async def get_usage(self, *, window_days: int | None = None) -> UsageResponse: params["window_days"] = str(window_days) return await self._request("GET", "/v1/usage", params=params) + async def get_overview_status(self) -> OverviewStatusResponse: + return await self._request("GET", "/v1/status/overview") + + async def get_collections_status(self) -> CollectionsStatusResponse: + return await self._request("GET", "/v1/status/collections") + + async def get_usage_status(self, *, window_days: int | None = None) -> UsageResponse: + params: dict[str, str] = {} + if window_days is not None: + params["window_days"] = str(window_days) + return await self._request("GET", "/v1/status/usage", params=params) + + async def get_access_status( + self, *, window_days: int | None = None + ) -> AccessLogOverviewResponse: + params: dict[str, str] = {} + if window_days is not None: + params["window_days"] = str(window_days) + return await self._request("GET", "/v1/admin/status/access", params=params) + def collection(self, name: str) -> CollectionClient: return CollectionClient(self, name) @@ -206,10 +221,3 @@ async def query(self, body: QueryBody) -> QueryResponse: async def analytics(self) -> AnalyticsResponse: return await self._client.collections.analytics(self._name) - - async def stream_events(self) -> AsyncGenerator[ProgressEvent, None]: - async for event in self._client.collections.stream_events(self._name): - yield event - - async def create_realtime_token(self) -> CollectionRealtimeTokenResponse: - return await self._client.collections.create_realtime_token(self._name) diff --git a/sdks/python/src/bigrag/_realtime.py b/sdks/python/src/bigrag/_realtime.py deleted file mode 100644 index be2bbb4f..00000000 --- a/sdks/python/src/bigrag/_realtime.py +++ /dev/null @@ -1,150 +0,0 @@ -from __future__ import annotations - -import asyncio -import json -from collections.abc import AsyncGenerator -from typing import Any -from uuid import uuid4 - -from websockets.asyncio.client import connect -from websockets.exceptions import ConnectionClosedOK - -from bigrag.types.realtime import RealtimeMessage - - -class RealtimeConnection: - def __init__(self, client) -> None: - self._client = client - self._socket = None - self._reader: asyncio.Task | None = None - self._queues: dict[str, asyncio.Queue[RealtimeMessage | None]] = {} - self._closed_exc: Exception | None = None - - async def __aenter__(self) -> RealtimeConnection: - await self.connect() - return self - - async def __aexit__(self, *_exc: object) -> None: - await self.close() - - async def connect(self) -> None: - if self._socket is not None: - return - self._socket = await connect( - _realtime_url(self._client.base_url), - additional_headers=self._client._headers(), - open_timeout=self._client.timeout, - ping_interval=20, - ping_timeout=20, - ) - self._reader = asyncio.create_task(self._read_loop()) - - async def close(self) -> None: - for queue in self._queues.values(): - queue.put_nowait(None) - self._queues.clear() - if self._reader is not None: - self._reader.cancel() - await asyncio.gather(self._reader, return_exceptions=True) - self._reader = None - if self._socket is not None: - await self._socket.close() - self._socket = None - - async def subscribe( - self, topic: str, params: dict[str, Any] | None = None - ) -> AsyncGenerator[RealtimeMessage, None]: - await self.connect() - subscription_id = uuid4().hex - queue: asyncio.Queue[RealtimeMessage | None] = asyncio.Queue() - self._queues[subscription_id] = queue - await self._send( - { - "type": "subscribe", - "id": subscription_id, - "topic": topic, - "params": _compact_params(params or {}), - } - ) - try: - while True: - message = await queue.get() - if message is None: - if self._closed_exc is not None: - raise self._closed_exc - return - yield message - if message.get("type") in {"complete", "error"}: - return - finally: - try: - await self._send({"type": "unsubscribe", "id": subscription_id}) - except Exception: - pass - self._queues.pop(subscription_id, None) - queue.put_nowait(None) - - async def _send(self, message: dict[str, Any]) -> None: - if self._socket is None: - raise RuntimeError("realtime connection is not open") - await self._socket.send(json.dumps(message)) - - async def _read_loop(self) -> None: - try: - while True: - raw = await self._socket.recv() - if isinstance(raw, bytes): - raw = raw.decode() - try: - message = json.loads(raw) - except json.JSONDecodeError: - continue - if not isinstance(message, dict): - continue - if message.get("type") in {"heartbeat", "pong", "subscribed"}: - continue - subscription_id = message.get("id") - if not isinstance(subscription_id, str): - continue - queue = self._queues.get(subscription_id) - if queue is not None: - queue.put_nowait(message) - except asyncio.CancelledError: - raise - except ConnectionClosedOK: - pass - except Exception as exc: - self._closed_exc = exc - finally: - for queue in self._queues.values(): - queue.put_nowait(None) - - -class RealtimeResource: - def __init__(self, client) -> None: - self._client = client - - def connect(self) -> RealtimeConnection: - return RealtimeConnection(self._client) - - async def subscribe( - self, topic: str, params: dict[str, Any] | None = None - ) -> AsyncGenerator[RealtimeMessage, None]: - connection = self.connect() - try: - async for message in connection.subscribe(topic, params): - yield message - finally: - await connection.close() - - -def _realtime_url(base_url: str) -> str: - if base_url.startswith("https://"): - return f"wss://{base_url.removeprefix('https://').rstrip('/')}/v1/realtime" - if base_url.startswith("http://"): - return f"ws://{base_url.removeprefix('http://').rstrip('/')}/v1/realtime" - return f"{base_url.rstrip('/')}/v1/realtime" - - -def _compact_params(params: dict[str, Any]) -> dict[str, Any]: - return {key: value for key, value in params.items() if value is not None} diff --git a/sdks/python/src/bigrag/resources/__init__.py b/sdks/python/src/bigrag/resources/__init__.py index 2b6041fc..e55b0a56 100644 --- a/sdks/python/src/bigrag/resources/__init__.py +++ b/sdks/python/src/bigrag/resources/__init__.py @@ -1,13 +1,11 @@ from __future__ import annotations -from bigrag._realtime import RealtimeResource from bigrag.resources.admin import ( AdminAccessResource, AdminApiKeysResource, AdminAuditResource, AdminEmbeddingPresetsResource, AdminMcpServersResource, - AdminRealtimeResource, AdminResource, AdminSettingsResource, AdminUsersResource, @@ -29,7 +27,6 @@ "AdminAuditResource", "AdminEmbeddingPresetsResource", "AdminMcpServersResource", - "AdminRealtimeResource", "AdminResource", "AdminSettingsResource", "AdminUsersResource", @@ -41,7 +38,6 @@ "DocumentsResource", "EvaluationsResource", "QueryResource", - "RealtimeResource", "S3ConnectorResource", "VectorsResource", "WebhooksResource", diff --git a/sdks/python/src/bigrag/resources/admin/__init__.py b/sdks/python/src/bigrag/resources/admin/__init__.py index 8a2c2c02..c27e988b 100644 --- a/sdks/python/src/bigrag/resources/admin/__init__.py +++ b/sdks/python/src/bigrag/resources/admin/__init__.py @@ -7,7 +7,6 @@ from bigrag.resources.admin.audit import AdminAuditResource from bigrag.resources.admin.embedding_presets import AdminEmbeddingPresetsResource from bigrag.resources.admin.mcp_servers import AdminMcpServersResource -from bigrag.resources.admin.realtime import AdminRealtimeResource from bigrag.resources.admin.settings import AdminSettingsResource from bigrag.resources.admin.users import AdminUsersResource from bigrag.resources.admin.vector_storage import AdminVectorStorageResource @@ -23,7 +22,6 @@ class AdminResource: audit: AdminAuditResource embedding_presets: AdminEmbeddingPresetsResource mcp_servers: AdminMcpServersResource - realtime: AdminRealtimeResource settings: AdminSettingsResource vector_storage: AdminVectorStorageResource @@ -34,7 +32,6 @@ def __init__(self, client: BigRAGCore) -> None: self.audit = AdminAuditResource(client) self.embedding_presets = AdminEmbeddingPresetsResource(client) self.mcp_servers = AdminMcpServersResource(client) - self.realtime = AdminRealtimeResource(client) self.settings = AdminSettingsResource(client) self.vector_storage = AdminVectorStorageResource(client) @@ -45,7 +42,6 @@ def __init__(self, client: BigRAGCore) -> None: "AdminAuditResource", "AdminEmbeddingPresetsResource", "AdminMcpServersResource", - "AdminRealtimeResource", "AdminResource", "AdminSettingsResource", "AdminUsersResource", diff --git a/sdks/python/src/bigrag/resources/admin/realtime.py b/sdks/python/src/bigrag/resources/admin/realtime.py deleted file mode 100644 index b6c4cb76..00000000 --- a/sdks/python/src/bigrag/resources/admin/realtime.py +++ /dev/null @@ -1,173 +0,0 @@ -from __future__ import annotations - -from collections.abc import AsyncGenerator -from typing import TYPE_CHECKING, Any - -from bigrag._realtime import RealtimeConnection -from bigrag.types.admin import AdminRealtimeEvent - -if TYPE_CHECKING: - from bigrag._core import BigRAGCore - - -class AdminRealtimeResource: - def __init__(self, client: BigRAGCore) -> None: - self._client = client - - def documents( - self, - collection: str, - *, - status: str | None = None, - limit: int | None = None, - offset: int | None = None, - order: str | None = None, - q: str | None = None, - sort: str | None = None, - ) -> AsyncGenerator[AdminRealtimeEvent, None]: - return self._stream( - "admin.collections.documents", - { - "collection": collection, - "status": status, - "limit": limit, - "offset": offset, - "order": order, - "q": q, - "sort": sort, - }, - ) - - def document_batch_status( - self, collection: str, document_ids: list[str] - ) -> AsyncGenerator[AdminRealtimeEvent, None]: - return self._stream( - "admin.collections.documents.batch_status", - {"collection": collection, "document_ids": document_ids}, - ) - - def document( - self, collection: str, document_id: str - ) -> AsyncGenerator[AdminRealtimeEvent, None]: - return self._stream( - "admin.collections.documents.detail", - {"collection": collection, "document_id": document_id}, - ) - - def upload_session( - self, collection: str, session_id: str - ) -> AsyncGenerator[AdminRealtimeEvent, None]: - return self._stream( - "admin.collections.upload_session", - {"collection": collection, "session_id": session_id}, - ) - - def collection_stats(self, collection: str) -> AsyncGenerator[AdminRealtimeEvent, None]: - return self._stream("admin.collections.stats", {"collection": collection}) - - def connector_sources( - self, provider: str, *, collection: str | None = None - ) -> AsyncGenerator[AdminRealtimeEvent, None]: - return self._stream( - "admin.connectors.sources", - {"provider": provider, "collection": collection}, - ) - - def connector_sync_jobs( - self, - provider: str, - *, - collection: str | None = None, - source_id: str | None = None, - limit: int | None = None, - ) -> AsyncGenerator[AdminRealtimeEvent, None]: - return self._stream( - "admin.connectors.sync_jobs", - { - "provider": provider, - "collection": collection, - "source_id": source_id, - "limit": limit, - }, - ) - - def access_overview( - self, *, window_days: int | None = None - ) -> AsyncGenerator[AdminRealtimeEvent, None]: - return self._stream("admin.access.overview", {"window_days": window_days}) - - def access_logs( - self, - *, - action: str | None = None, - actor_id: str | None = None, - collection: str | None = None, - method: str | None = None, - path: str | None = None, - status_family: str | None = None, - success: bool | None = None, - limit: int | None = None, - offset: int | None = None, - ) -> AsyncGenerator[AdminRealtimeEvent, None]: - return self._stream( - "admin.access.logs", - { - "action": action, - "actor_id": actor_id, - "collection": collection, - "method": method, - "path": path, - "status_family": status_family, - "success": success, - "limit": limit, - "offset": offset, - }, - ) - - def audit( - self, - *, - action: str | None = None, - actor_id: str | None = None, - resource_type: str | None = None, - limit: int | None = None, - offset: int | None = None, - ) -> AsyncGenerator[AdminRealtimeEvent, None]: - return self._stream( - "admin.audit", - { - "action": action, - "actor_id": actor_id, - "resource_type": resource_type, - "limit": limit, - "offset": offset, - }, - ) - - def usage(self, *, window_days: int | None = None) -> AsyncGenerator[AdminRealtimeEvent, None]: - return self._stream("admin.usage", {"window_days": window_days}) - - def platform_stats(self) -> AsyncGenerator[AdminRealtimeEvent, None]: - return self._stream("admin.platform.stats") - - def platform_readiness(self) -> AsyncGenerator[AdminRealtimeEvent, None]: - return self._stream("admin.platform.readiness") - - def custom( - self, topic: str, params: dict[str, Any] | None = None - ) -> AsyncGenerator[AdminRealtimeEvent, None]: - if not topic.startswith("admin."): - raise ValueError("admin.realtime.custom topic must start with admin.") - return self._stream(topic, params or {}) - - async def _stream( - self, topic: str, params: dict[str, Any] | None = None - ) -> AsyncGenerator[AdminRealtimeEvent, None]: - connection = RealtimeConnection(self._client) - try: - async for message in connection.subscribe(topic, params or {}): - yield message - if message["type"] == "complete": - return - finally: - await connection.close() diff --git a/sdks/python/src/bigrag/resources/collections.py b/sdks/python/src/bigrag/resources/collections.py index 1bd6eb9c..b5c31c8f 100644 --- a/sdks/python/src/bigrag/resources/collections.py +++ b/sdks/python/src/bigrag/resources/collections.py @@ -4,18 +4,15 @@ from typing import TYPE_CHECKING from urllib.parse import quote -from bigrag._realtime import RealtimeConnection from bigrag.types.analytics import AnalyticsResponse from bigrag.types.collections import ( Collection, CollectionListResponse, - CollectionRealtimeTokenResponse, CollectionStatsResponse, CreateCollectionBody, UpdateCollectionBody, ) from bigrag.types.common import StatusResponse -from bigrag.types.realtime import ProgressEvent if TYPE_CHECKING: from bigrag._core import BigRAGCore @@ -89,25 +86,3 @@ async def truncate(self, name: str) -> StatusResponse: return await self._client._request( "POST", f"/v1/collections/{quote(name, safe='')}/truncate" ) - - async def create_realtime_token(self, name: str) -> CollectionRealtimeTokenResponse: - return await self._client._request( - "POST", f"/v1/collections/{quote(name, safe='')}/realtime-token" - ) - - async def stream_events( - self, name: str, *, token: str | None = None - ) -> AsyncGenerator[ProgressEvent, None]: - connection = RealtimeConnection(self._client) - try: - async for message in connection.subscribe( - "collection.events", {"collection": name, "token": token} - ): - if message["type"] == "event": - yield message["payload"] - elif message["type"] == "error": - raise RuntimeError(message["message"]) - elif message["type"] == "complete": - return - finally: - await connection.close() diff --git a/sdks/python/src/bigrag/types/__init__.py b/sdks/python/src/bigrag/types/__init__.py index 5088d39c..eb3a3643 100644 --- a/sdks/python/src/bigrag/types/__init__.py +++ b/sdks/python/src/bigrag/types/__init__.py @@ -6,7 +6,6 @@ AccessLogTimelinePoint, ) from bigrag.types.admin import ( - AdminRealtimeEvent, ApiKey, ApiKeyListResponse, AuditLogEntry, @@ -63,14 +62,15 @@ from bigrag.types.collections import ( Collection, CollectionListResponse, - CollectionRealtimeTokenResponse, CollectionStatsResponse, CreateCollectionBody, UpdateCollectionBody, ) from bigrag.types.common import ( + CollectionsStatusResponse, DocumentStats, HealthResponse, + OverviewStatusResponse, PlatformStatsResponse, QueueStatsResponse, ReadinessResponse, @@ -129,15 +129,6 @@ QueryTimings, SearchMode, ) -from bigrag.types.realtime import ( - ProgressEvent, - RealtimeComplete, - RealtimeControlMessage, - RealtimeError, - RealtimeEvent, - RealtimeMessage, - RealtimeSnapshot, -) from bigrag.types.usage import CollectionUsage, UsageResponse from bigrag.types.vectors import DeleteResponse, UpsertResponse, VectorEntry from bigrag.types.webhooks import ( @@ -160,6 +151,8 @@ "WorkerStatsResponse", "DocumentStats", "PlatformStatsResponse", + "OverviewStatusResponse", + "CollectionsStatusResponse", "SetupStatusResponse", "LoginBody", "SetupBody", @@ -203,7 +196,6 @@ "ResetInstanceSettingsBody", "TestInstanceSettingsBody", "InstanceSettingsTestResponse", - "AdminRealtimeEvent", "VectorStorageHealth", "VectorStorageCollection", "VectorStorageTotals", @@ -214,7 +206,6 @@ "AccessLogTimelinePoint", "AccessLogOverviewResponse", "Collection", - "CollectionRealtimeTokenResponse", "CollectionListResponse", "CollectionStatsResponse", "CreateCollectionBody", @@ -286,11 +277,4 @@ "EmbeddingModelListResponse", "CollectionUsage", "UsageResponse", - "ProgressEvent", - "RealtimeSnapshot", - "RealtimeEvent", - "RealtimeError", - "RealtimeComplete", - "RealtimeControlMessage", - "RealtimeMessage", ] diff --git a/sdks/python/src/bigrag/types/admin.py b/sdks/python/src/bigrag/types/admin.py index bddbc5c2..ec235fd7 100644 --- a/sdks/python/src/bigrag/types/admin.py +++ b/sdks/python/src/bigrag/types/admin.py @@ -3,7 +3,6 @@ from typing import Any, Literal, NotRequired, TypedDict from bigrag.types.auth import User -from bigrag.types.realtime import RealtimeMessage InstanceSettingKind = Literal[ "bool", @@ -73,9 +72,6 @@ class InstanceSettingsTestResponse(TypedDict): message: str -AdminRealtimeEvent = RealtimeMessage - - class VectorStorageHealth(TypedDict): status: Literal["ok", "error"] error: str | None diff --git a/sdks/python/src/bigrag/types/collections.py b/sdks/python/src/bigrag/types/collections.py index 3fcb858e..93132b9d 100644 --- a/sdks/python/src/bigrag/types/collections.py +++ b/sdks/python/src/bigrag/types/collections.py @@ -46,11 +46,6 @@ class CollectionStatsResponse(TypedDict): status_counts: dict[str, int] -class CollectionRealtimeTokenResponse(TypedDict): - token: str - expires_in: int - - class CreateCollectionBody(TypedDict): name: str description: NotRequired[str] diff --git a/sdks/python/src/bigrag/types/common.py b/sdks/python/src/bigrag/types/common.py index 58a19631..df4d1c27 100644 --- a/sdks/python/src/bigrag/types/common.py +++ b/sdks/python/src/bigrag/types/common.py @@ -59,3 +59,20 @@ class PlatformStatsResponse(TypedDict): webhooks: int queue: QueueStatsResponse workers: NotRequired[WorkerStatsResponse] + + +class OverviewStatusResponse(TypedDict): + platform: PlatformStatsResponse + readiness: ReadinessResponse + + +class CollectionsStatusResponse(TypedDict): + collections_total: int + documents_total: int + documents_ready: int + documents_pending: int + documents_processing: int + documents_failed: int + total_chunks: int + total_tokens: int + total_size_bytes: int diff --git a/sdks/python/src/bigrag/types/realtime.py b/sdks/python/src/bigrag/types/realtime.py deleted file mode 100644 index fd5bdc26..00000000 --- a/sdks/python/src/bigrag/types/realtime.py +++ /dev/null @@ -1,55 +0,0 @@ -from __future__ import annotations - -from typing import Any, Literal, NotRequired, TypedDict - - -class RealtimeSnapshot(TypedDict): - type: Literal["snapshot"] - id: str - topic: str - payload: Any - generated_at: str - - -class RealtimeEvent(TypedDict): - type: Literal["event"] - id: str - topic: str - payload: Any - generated_at: str - - -class RealtimeError(TypedDict): - type: Literal["error"] - message: str - id: NotRequired[str] - topic: NotRequired[str] - generated_at: str - - -class RealtimeComplete(TypedDict): - type: Literal["complete"] - id: str - topic: str - generated_at: str - - -class RealtimeControlMessage(TypedDict): - type: Literal["subscribed", "heartbeat", "pong"] - id: NotRequired[str] - topic: NotRequired[str] - generated_at: NotRequired[str] - - -RealtimeMessage = ( - RealtimeSnapshot | RealtimeEvent | RealtimeError | RealtimeComplete | RealtimeControlMessage -) - - -class ProgressEvent(TypedDict): - step: str - message: str - progress: float - document_id: NotRequired[str] - status: NotRequired[str] - detail: NotRequired[dict[str, Any]] diff --git a/sdks/python/src/bigrag/types/usage.py b/sdks/python/src/bigrag/types/usage.py index 2f7cd414..b332f168 100644 --- a/sdks/python/src/bigrag/types/usage.py +++ b/sdks/python/src/bigrag/types/usage.py @@ -23,4 +23,6 @@ class UsageResponse(TypedDict): storage_bytes_total: int embedding_tokens_total: int embedding_cost_usd_estimate: float + avg_latency_ms: float + timeline: list[dict[str, int | float | str]] by_collection: list[CollectionUsage] diff --git a/sdks/python/uv.lock b/sdks/python/uv.lock index 10e7f5bb..31a7c820 100644 --- a/sdks/python/uv.lock +++ b/sdks/python/uv.lock @@ -21,14 +21,10 @@ version = "2026.5.23" source = { editable = "." } dependencies = [ { name = "httpx" }, - { name = "websockets" }, ] [package.metadata] -requires-dist = [ - { name = "httpx", specifier = ">=0.28.0" }, - { name = "websockets", specifier = ">=15.0,<17" }, -] +requires-dist = [{ name = "httpx", specifier = ">=0.28.0" }] [[package]] name = "certifi" @@ -93,62 +89,3 @@ sdist = { url = "https://files.pythonhosted.org/packages/72/94/1a15dd82efb362ac8 wheels = [ { url = "https://files.pythonhosted.org/packages/18/67/36e9267722cc04a6b9f15c7f3441c2363321a3ea07da7ae0c0707beb2a9c/typing_extensions-4.15.0-py3-none-any.whl", hash = "sha256:f0fa19c6845758ab08074a0cfa8b7aecb71c999ca73d62883bc25cc018c4e548", size = 44614, upload-time = "2025-08-25T13:49:24.86Z" }, ] - -[[package]] -name = "websockets" -version = "16.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/04/24/4b2031d72e840ce4c1ccb255f693b15c334757fc50023e4db9537080b8c4/websockets-16.0.tar.gz", hash = "sha256:5f6261a5e56e8d5c42a4497b364ea24d94d9563e8fbd44e78ac40879c60179b5", size = 179346, upload-time = "2026-01-10T09:23:47.181Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/f2/db/de907251b4ff46ae804ad0409809504153b3f30984daf82a1d84a9875830/websockets-16.0-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:31a52addea25187bde0797a97d6fc3d2f92b6f72a9370792d65a6e84615ac8a8", size = 177340, upload-time = "2026-01-10T09:22:34.539Z" }, - { url = "https://files.pythonhosted.org/packages/f3/fa/abe89019d8d8815c8781e90d697dec52523fb8ebe308bf11664e8de1877e/websockets-16.0-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:417b28978cdccab24f46400586d128366313e8a96312e4b9362a4af504f3bbad", size = 175022, upload-time = "2026-01-10T09:22:36.332Z" }, - { url = "https://files.pythonhosted.org/packages/58/5d/88ea17ed1ded2079358b40d31d48abe90a73c9e5819dbcde1606e991e2ad/websockets-16.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:af80d74d4edfa3cb9ed973a0a5ba2b2a549371f8a741e0800cb07becdd20f23d", size = 175319, upload-time = "2026-01-10T09:22:37.602Z" }, - { url = "https://files.pythonhosted.org/packages/d2/ae/0ee92b33087a33632f37a635e11e1d99d429d3d323329675a6022312aac2/websockets-16.0-cp311-cp311-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:08d7af67b64d29823fed316505a89b86705f2b7981c07848fb5e3ea3020c1abe", size = 184631, upload-time = "2026-01-10T09:22:38.789Z" }, - { url = "https://files.pythonhosted.org/packages/c8/c5/27178df583b6c5b31b29f526ba2da5e2f864ecc79c99dae630a85d68c304/websockets-16.0-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:7be95cfb0a4dae143eaed2bcba8ac23f4892d8971311f1b06f3c6b78952ee70b", size = 185870, upload-time = "2026-01-10T09:22:39.893Z" }, - { url = "https://files.pythonhosted.org/packages/87/05/536652aa84ddc1c018dbb7e2c4cbcd0db884580bf8e95aece7593fde526f/websockets-16.0-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:d6297ce39ce5c2e6feb13c1a996a2ded3b6832155fcfc920265c76f24c7cceb5", size = 185361, upload-time = "2026-01-10T09:22:41.016Z" }, - { url = "https://files.pythonhosted.org/packages/6d/e2/d5332c90da12b1e01f06fb1b85c50cfc489783076547415bf9f0a659ec19/websockets-16.0-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:1c1b30e4f497b0b354057f3467f56244c603a79c0d1dafce1d16c283c25f6e64", size = 184615, upload-time = "2026-01-10T09:22:42.442Z" }, - { url = "https://files.pythonhosted.org/packages/77/fb/d3f9576691cae9253b51555f841bc6600bf0a983a461c79500ace5a5b364/websockets-16.0-cp311-cp311-win32.whl", hash = "sha256:5f451484aeb5cafee1ccf789b1b66f535409d038c56966d6101740c1614b86c6", size = 178246, upload-time = "2026-01-10T09:22:43.654Z" }, - { url = "https://files.pythonhosted.org/packages/54/67/eaff76b3dbaf18dcddabc3b8c1dba50b483761cccff67793897945b37408/websockets-16.0-cp311-cp311-win_amd64.whl", hash = "sha256:8d7f0659570eefb578dacde98e24fb60af35350193e4f56e11190787bee77dac", size = 178684, upload-time = "2026-01-10T09:22:44.941Z" }, - { url = "https://files.pythonhosted.org/packages/84/7b/bac442e6b96c9d25092695578dda82403c77936104b5682307bd4deb1ad4/websockets-16.0-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:71c989cbf3254fbd5e84d3bff31e4da39c43f884e64f2551d14bb3c186230f00", size = 177365, upload-time = "2026-01-10T09:22:46.787Z" }, - { url = "https://files.pythonhosted.org/packages/b0/fe/136ccece61bd690d9c1f715baaeefd953bb2360134de73519d5df19d29ca/websockets-16.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:8b6e209ffee39ff1b6d0fa7bfef6de950c60dfb91b8fcead17da4ee539121a79", size = 175038, upload-time = "2026-01-10T09:22:47.999Z" }, - { url = "https://files.pythonhosted.org/packages/40/1e/9771421ac2286eaab95b8575b0cb701ae3663abf8b5e1f64f1fd90d0a673/websockets-16.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:86890e837d61574c92a97496d590968b23c2ef0aeb8a9bc9421d174cd378ae39", size = 175328, upload-time = "2026-01-10T09:22:49.809Z" }, - { url = "https://files.pythonhosted.org/packages/18/29/71729b4671f21e1eaa5d6573031ab810ad2936c8175f03f97f3ff164c802/websockets-16.0-cp312-cp312-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:9b5aca38b67492ef518a8ab76851862488a478602229112c4b0d58d63a7a4d5c", size = 184915, upload-time = "2026-01-10T09:22:51.071Z" }, - { url = "https://files.pythonhosted.org/packages/97/bb/21c36b7dbbafc85d2d480cd65df02a1dc93bf76d97147605a8e27ff9409d/websockets-16.0-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:e0334872c0a37b606418ac52f6ab9cfd17317ac26365f7f65e203e2d0d0d359f", size = 186152, upload-time = "2026-01-10T09:22:52.224Z" }, - { url = "https://files.pythonhosted.org/packages/4a/34/9bf8df0c0cf88fa7bfe36678dc7b02970c9a7d5e065a3099292db87b1be2/websockets-16.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:a0b31e0b424cc6b5a04b8838bbaec1688834b2383256688cf47eb97412531da1", size = 185583, upload-time = "2026-01-10T09:22:53.443Z" }, - { url = "https://files.pythonhosted.org/packages/47/88/4dd516068e1a3d6ab3c7c183288404cd424a9a02d585efbac226cb61ff2d/websockets-16.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:485c49116d0af10ac698623c513c1cc01c9446c058a4e61e3bf6c19dff7335a2", size = 184880, upload-time = "2026-01-10T09:22:55.033Z" }, - { url = "https://files.pythonhosted.org/packages/91/d6/7d4553ad4bf1c0421e1ebd4b18de5d9098383b5caa1d937b63df8d04b565/websockets-16.0-cp312-cp312-win32.whl", hash = "sha256:eaded469f5e5b7294e2bdca0ab06becb6756ea86894a47806456089298813c89", size = 178261, upload-time = "2026-01-10T09:22:56.251Z" }, - { url = "https://files.pythonhosted.org/packages/c3/f0/f3a17365441ed1c27f850a80b2bc680a0fa9505d733fe152fdf5e98c1c0b/websockets-16.0-cp312-cp312-win_amd64.whl", hash = "sha256:5569417dc80977fc8c2d43a86f78e0a5a22fee17565d78621b6bb264a115d4ea", size = 178693, upload-time = "2026-01-10T09:22:57.478Z" }, - { url = "https://files.pythonhosted.org/packages/cc/9c/baa8456050d1c1b08dd0ec7346026668cbc6f145ab4e314d707bb845bf0d/websockets-16.0-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:878b336ac47938b474c8f982ac2f7266a540adc3fa4ad74ae96fea9823a02cc9", size = 177364, upload-time = "2026-01-10T09:22:59.333Z" }, - { url = "https://files.pythonhosted.org/packages/7e/0c/8811fc53e9bcff68fe7de2bcbe75116a8d959ac699a3200f4847a8925210/websockets-16.0-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:52a0fec0e6c8d9a784c2c78276a48a2bdf099e4ccc2a4cad53b27718dbfd0230", size = 175039, upload-time = "2026-01-10T09:23:01.171Z" }, - { url = "https://files.pythonhosted.org/packages/aa/82/39a5f910cb99ec0b59e482971238c845af9220d3ab9fa76dd9162cda9d62/websockets-16.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:e6578ed5b6981005df1860a56e3617f14a6c307e6a71b4fff8c48fdc50f3ed2c", size = 175323, upload-time = "2026-01-10T09:23:02.341Z" }, - { url = "https://files.pythonhosted.org/packages/bd/28/0a25ee5342eb5d5f297d992a77e56892ecb65e7854c7898fb7d35e9b33bd/websockets-16.0-cp313-cp313-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:95724e638f0f9c350bb1c2b0a7ad0e83d9cc0c9259f3ea94e40d7b02a2179ae5", size = 184975, upload-time = "2026-01-10T09:23:03.756Z" }, - { url = "https://files.pythonhosted.org/packages/f9/66/27ea52741752f5107c2e41fda05e8395a682a1e11c4e592a809a90c6a506/websockets-16.0-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:c0204dc62a89dc9d50d682412c10b3542d748260d743500a85c13cd1ee4bde82", size = 186203, upload-time = "2026-01-10T09:23:05.01Z" }, - { url = "https://files.pythonhosted.org/packages/37/e5/8e32857371406a757816a2b471939d51c463509be73fa538216ea52b792a/websockets-16.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:52ac480f44d32970d66763115edea932f1c5b1312de36df06d6b219f6741eed8", size = 185653, upload-time = "2026-01-10T09:23:06.301Z" }, - { url = "https://files.pythonhosted.org/packages/9b/67/f926bac29882894669368dc73f4da900fcdf47955d0a0185d60103df5737/websockets-16.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:6e5a82b677f8f6f59e8dfc34ec06ca6b5b48bc4fcda346acd093694cc2c24d8f", size = 184920, upload-time = "2026-01-10T09:23:07.492Z" }, - { url = "https://files.pythonhosted.org/packages/3c/a1/3d6ccdcd125b0a42a311bcd15a7f705d688f73b2a22d8cf1c0875d35d34a/websockets-16.0-cp313-cp313-win32.whl", hash = "sha256:abf050a199613f64c886ea10f38b47770a65154dc37181bfaff70c160f45315a", size = 178255, upload-time = "2026-01-10T09:23:09.245Z" }, - { url = "https://files.pythonhosted.org/packages/6b/ae/90366304d7c2ce80f9b826096a9e9048b4bb760e44d3b873bb272cba696b/websockets-16.0-cp313-cp313-win_amd64.whl", hash = "sha256:3425ac5cf448801335d6fdc7ae1eb22072055417a96cc6b31b3861f455fbc156", size = 178689, upload-time = "2026-01-10T09:23:10.483Z" }, - { url = "https://files.pythonhosted.org/packages/f3/1d/e88022630271f5bd349ed82417136281931e558d628dd52c4d8621b4a0b2/websockets-16.0-cp314-cp314-macosx_10_15_universal2.whl", hash = "sha256:8cc451a50f2aee53042ac52d2d053d08bf89bcb31ae799cb4487587661c038a0", size = 177406, upload-time = "2026-01-10T09:23:12.178Z" }, - { url = "https://files.pythonhosted.org/packages/f2/78/e63be1bf0724eeb4616efb1ae1c9044f7c3953b7957799abb5915bffd38e/websockets-16.0-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:daa3b6ff70a9241cf6c7fc9e949d41232d9d7d26fd3522b1ad2b4d62487e9904", size = 175085, upload-time = "2026-01-10T09:23:13.511Z" }, - { url = "https://files.pythonhosted.org/packages/bb/f4/d3c9220d818ee955ae390cf319a7c7a467beceb24f05ee7aaaa2414345ba/websockets-16.0-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:fd3cb4adb94a2a6e2b7c0d8d05cb94e6f1c81a0cf9dc2694fb65c7e8d94c42e4", size = 175328, upload-time = "2026-01-10T09:23:14.727Z" }, - { url = "https://files.pythonhosted.org/packages/63/bc/d3e208028de777087e6fb2b122051a6ff7bbcca0d6df9d9c2bf1dd869ae9/websockets-16.0-cp314-cp314-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:781caf5e8eee67f663126490c2f96f40906594cb86b408a703630f95550a8c3e", size = 185044, upload-time = "2026-01-10T09:23:15.939Z" }, - { url = "https://files.pythonhosted.org/packages/ad/6e/9a0927ac24bd33a0a9af834d89e0abc7cfd8e13bed17a86407a66773cc0e/websockets-16.0-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:caab51a72c51973ca21fa8a18bd8165e1a0183f1ac7066a182ff27107b71e1a4", size = 186279, upload-time = "2026-01-10T09:23:17.148Z" }, - { url = "https://files.pythonhosted.org/packages/b9/ca/bf1c68440d7a868180e11be653c85959502efd3a709323230314fda6e0b3/websockets-16.0-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:19c4dc84098e523fd63711e563077d39e90ec6702aff4b5d9e344a60cb3c0cb1", size = 185711, upload-time = "2026-01-10T09:23:18.372Z" }, - { url = "https://files.pythonhosted.org/packages/c4/f8/fdc34643a989561f217bb477cbc47a3a07212cbda91c0e4389c43c296ebf/websockets-16.0-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:a5e18a238a2b2249c9a9235466b90e96ae4795672598a58772dd806edc7ac6d3", size = 184982, upload-time = "2026-01-10T09:23:19.652Z" }, - { url = "https://files.pythonhosted.org/packages/dd/d1/574fa27e233764dbac9c52730d63fcf2823b16f0856b3329fc6268d6ae4f/websockets-16.0-cp314-cp314-win32.whl", hash = "sha256:a069d734c4a043182729edd3e9f247c3b2a4035415a9172fd0f1b71658a320a8", size = 177915, upload-time = "2026-01-10T09:23:21.458Z" }, - { url = "https://files.pythonhosted.org/packages/8a/f1/ae6b937bf3126b5134ce1f482365fde31a357c784ac51852978768b5eff4/websockets-16.0-cp314-cp314-win_amd64.whl", hash = "sha256:c0ee0e63f23914732c6d7e0cce24915c48f3f1512ec1d079ed01fc629dab269d", size = 178381, upload-time = "2026-01-10T09:23:22.715Z" }, - { url = "https://files.pythonhosted.org/packages/06/9b/f791d1db48403e1f0a27577a6beb37afae94254a8c6f08be4a23e4930bc0/websockets-16.0-cp314-cp314t-macosx_10_15_universal2.whl", hash = "sha256:a35539cacc3febb22b8f4d4a99cc79b104226a756aa7400adc722e83b0d03244", size = 177737, upload-time = "2026-01-10T09:23:24.523Z" }, - { url = "https://files.pythonhosted.org/packages/bd/40/53ad02341fa33b3ce489023f635367a4ac98b73570102ad2cdd770dacc9a/websockets-16.0-cp314-cp314t-macosx_10_15_x86_64.whl", hash = "sha256:b784ca5de850f4ce93ec85d3269d24d4c82f22b7212023c974c401d4980ebc5e", size = 175268, upload-time = "2026-01-10T09:23:25.781Z" }, - { url = "https://files.pythonhosted.org/packages/74/9b/6158d4e459b984f949dcbbb0c5d270154c7618e11c01029b9bbd1bb4c4f9/websockets-16.0-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:569d01a4e7fba956c5ae4fc988f0d4e187900f5497ce46339c996dbf24f17641", size = 175486, upload-time = "2026-01-10T09:23:27.033Z" }, - { url = "https://files.pythonhosted.org/packages/e5/2d/7583b30208b639c8090206f95073646c2c9ffd66f44df967981a64f849ad/websockets-16.0-cp314-cp314t-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:50f23cdd8343b984957e4077839841146f67a3d31ab0d00e6b824e74c5b2f6e8", size = 185331, upload-time = "2026-01-10T09:23:28.259Z" }, - { url = "https://files.pythonhosted.org/packages/45/b0/cce3784eb519b7b5ad680d14b9673a31ab8dcb7aad8b64d81709d2430aa8/websockets-16.0-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:152284a83a00c59b759697b7f9e9cddf4e3c7861dd0d964b472b70f78f89e80e", size = 186501, upload-time = "2026-01-10T09:23:29.449Z" }, - { url = "https://files.pythonhosted.org/packages/19/60/b8ebe4c7e89fb5f6cdf080623c9d92789a53636950f7abacfc33fe2b3135/websockets-16.0-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:bc59589ab64b0022385f429b94697348a6a234e8ce22544e3681b2e9331b5944", size = 186062, upload-time = "2026-01-10T09:23:31.368Z" }, - { url = "https://files.pythonhosted.org/packages/88/a8/a080593f89b0138b6cba1b28f8df5673b5506f72879322288b031337c0b8/websockets-16.0-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:32da954ffa2814258030e5a57bc73a3635463238e797c7375dc8091327434206", size = 185356, upload-time = "2026-01-10T09:23:32.627Z" }, - { url = "https://files.pythonhosted.org/packages/c2/b6/b9afed2afadddaf5ebb2afa801abf4b0868f42f8539bfe4b071b5266c9fe/websockets-16.0-cp314-cp314t-win32.whl", hash = "sha256:5a4b4cc550cb665dd8a47f868c8d04c8230f857363ad3c9caf7a0c3bf8c61ca6", size = 178085, upload-time = "2026-01-10T09:23:33.816Z" }, - { url = "https://files.pythonhosted.org/packages/9f/3e/28135a24e384493fa804216b79a6a6759a38cc4ff59118787b9fb693df93/websockets-16.0-cp314-cp314t-win_amd64.whl", hash = "sha256:b14dc141ed6d2dde437cddb216004bcac6a1df0935d79656387bd41632ba0bbd", size = 178531, upload-time = "2026-01-10T09:23:35.016Z" }, - { url = "https://files.pythonhosted.org/packages/72/07/c98a68571dcf256e74f1f816b8cc5eae6eb2d3d5cfa44d37f801619d9166/websockets-16.0-pp311-pypy311_pp73-macosx_10_15_x86_64.whl", hash = "sha256:349f83cd6c9a415428ee1005cadb5c2c56f4389bc06a9af16103c3bc3dcc8b7d", size = 174947, upload-time = "2026-01-10T09:23:36.166Z" }, - { url = "https://files.pythonhosted.org/packages/7e/52/93e166a81e0305b33fe416338be92ae863563fe7bce446b0f687b9df5aea/websockets-16.0-pp311-pypy311_pp73-macosx_11_0_arm64.whl", hash = "sha256:4a1aba3340a8dca8db6eb5a7986157f52eb9e436b74813764241981ca4888f03", size = 175260, upload-time = "2026-01-10T09:23:37.409Z" }, - { url = "https://files.pythonhosted.org/packages/56/0c/2dbf513bafd24889d33de2ff0368190a0e69f37bcfa19009ef819fe4d507/websockets-16.0-pp311-pypy311_pp73-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:f4a32d1bd841d4bcbffdcb3d2ce50c09c3909fbead375ab28d0181af89fd04da", size = 176071, upload-time = "2026-01-10T09:23:39.158Z" }, - { url = "https://files.pythonhosted.org/packages/a5/8f/aea9c71cc92bf9b6cc0f7f70df8f0b420636b6c96ef4feee1e16f80f75dd/websockets-16.0-pp311-pypy311_pp73-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:0298d07ee155e2e9fda5be8a9042200dd2e3bb0b8a38482156576f863a9d457c", size = 176968, upload-time = "2026-01-10T09:23:41.031Z" }, - { url = "https://files.pythonhosted.org/packages/9a/3f/f70e03f40ffc9a30d817eef7da1be72ee4956ba8d7255c399a01b135902a/websockets-16.0-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:a653aea902e0324b52f1613332ddf50b00c06fdaf7e92624fbf8c77c78fa5767", size = 178735, upload-time = "2026-01-10T09:23:42.259Z" }, - { url = "https://files.pythonhosted.org/packages/6f/28/258ebab549c2bf3e64d2b0217b973467394a9cea8c42f70418ca2c5d0d2e/websockets-16.0-py3-none-any.whl", hash = "sha256:1637db62fad1dc833276dded54215f2c7fa46912301a24bd94d45d46a011ceec", size = 171598, upload-time = "2026-01-10T09:23:45.395Z" }, -] diff --git a/sdks/typescript/README.md b/sdks/typescript/README.md index 3d368e39..1c0c89f2 100644 --- a/sdks/typescript/README.md +++ b/sdks/typescript/README.md @@ -2,7 +2,7 @@ TypeScript client for [bigRAG](https://github.com/bigint/rag.computer) — a self-hostable RAG platform. -Works in Node.js 18+, browsers, Deno, Bun, and edge runtimes. Browser realtime uses native `WebSocket`; Node realtime uses the SDK's `ws` dependency. +Works in Node.js 18+, browsers, Deno, Bun, and edge runtimes. ## Installation @@ -55,15 +55,15 @@ while (current.status === "pending" || current.status === "processing") { ## Namespaces -- `client.collections` for collection CRUD, stats, analytics, realtime tokens, and event streams. +- `client.collections` for collection CRUD, stats, and analytics. - `client.documents` for uploads, batch operations, chunks, elements, and status polling. - `client.queries` for single, multi-collection, and batch retrieval queries. - `client.chat` for generated answers, question suggestions, and streaming. - `client.vectors` for raw vector upsert and delete. - `client.webhooks` for webhook management and delivery replay. - `client.auth` for setup, login, identity, password, and preferences. -- `client.admin` for users, API keys, access logs, audit logs, runtime settings, vector storage overview, admin realtime helpers, connectors, embedding presets, and MCP server keys. -- `client.realtime` for explicit WebSocket connect, subscribe, and unsubscribe control. +- `client.admin` for users, API keys, access logs, audit logs, runtime settings, vector storage overview, connectors, embedding presets, and MCP server keys. +- Top-level status helpers (`getOverviewStatus`, `getCollectionsStatus`, `getUsageStatus`, `getAccessStatus`) for pollable admin UI aggregates. - `client.connectors.s3` for S3-compatible bucket-prefix sources and sync jobs. - `client.evaluations` for golden-set retrieval evaluations. diff --git a/sdks/typescript/package.json b/sdks/typescript/package.json index 6b2d3314..f2b2e334 100644 --- a/sdks/typescript/package.json +++ b/sdks/typescript/package.json @@ -21,9 +21,6 @@ "engines": { "node": ">=18" }, - "dependencies": { - "ws": "^8.18.3" - }, "scripts": { "clean": "node -e \"require('node:fs').rmSync('dist',{recursive:true,force:true})\"", "build": "pnpm run clean && tsc", diff --git a/sdks/typescript/src/client.ts b/sdks/typescript/src/client.ts index 33c47da5..18b11b59 100644 --- a/sdks/typescript/src/client.ts +++ b/sdks/typescript/src/client.ts @@ -1,6 +1,5 @@ import type { BigRAGOptions } from "./core.js"; import { BigRAGCore } from "./core.js"; -import { RealtimeResource } from "./realtime.js"; import { AdminResource, AuthResource, @@ -14,8 +13,11 @@ import { WebhooksResource, } from "./resources/index.js"; import type { + AccessLogOverviewResponse, + CollectionsStatusResponse, EmbeddingModelListResponse, HealthResponse, + OverviewStatusResponse, PlatformStatsResponse, ReadinessResponse, UsageResponse, @@ -30,7 +32,6 @@ export class BigRAG extends BigRAGCore { readonly documents: DocumentsResource; readonly evaluations: EvaluationsResource; readonly queries: QueryResource; - readonly realtime: RealtimeResource; readonly vectors: VectorsResource; readonly webhooks: WebhooksResource; @@ -44,7 +45,6 @@ export class BigRAG extends BigRAGCore { this.documents = new DocumentsResource(this); this.evaluations = new EvaluationsResource(this); this.queries = new QueryResource(this); - this.realtime = new RealtimeResource(this); this.vectors = new VectorsResource(this); this.webhooks = new WebhooksResource(this); } @@ -70,4 +70,24 @@ export class BigRAG extends BigRAGCore { if (options.windowDays !== undefined) params.window_days = String(options.windowDays); return this._request("GET", "/v1/usage", { params }); } + + getOverviewStatus(): Promise { + return this._request("GET", "/v1/status/overview"); + } + + getCollectionsStatus(): Promise { + return this._request("GET", "/v1/status/collections"); + } + + getUsageStatus(options: { windowDays?: number } = {}): Promise { + const params: Record = {}; + if (options.windowDays !== undefined) params.window_days = String(options.windowDays); + return this._request("GET", "/v1/status/usage", { params }); + } + + getAccessStatus(options: { windowDays?: number } = {}): Promise { + const params: Record = {}; + if (options.windowDays !== undefined) params.window_days = String(options.windowDays); + return this._request("GET", "/v1/admin/status/access", { params }); + } } diff --git a/sdks/typescript/src/index.ts b/sdks/typescript/src/index.ts index 13452aca..48b549ca 100644 --- a/sdks/typescript/src/index.ts +++ b/sdks/typescript/src/index.ts @@ -3,7 +3,6 @@ export type { BigRAGOptions, RequestClient } from "./core.js"; export { BigRAGCore } from "./core.js"; export * from "./errors.js"; export { normalizeFileInput } from "./files.js"; -export { BigRAGRealtimeConnection, RealtimeResource } from "./realtime.js"; export { AdminResource, AuthResource, diff --git a/sdks/typescript/src/realtime.ts b/sdks/typescript/src/realtime.ts deleted file mode 100644 index 7d3ec47b..00000000 --- a/sdks/typescript/src/realtime.ts +++ /dev/null @@ -1,247 +0,0 @@ -import { type RequestClient, USER_AGENT } from "./core.js"; -import { APIConnectionError } from "./errors.js"; -import type { RealtimeMessage } from "./types/index.js"; - -type AnySocket = { - addEventListener?: (event: string, handler: (event: unknown) => void, options?: unknown) => void; - close: () => void; - on?: (event: string, handler: (...args: unknown[]) => void) => void; - readyState?: number; - send: (data: string) => void; -}; - -const OPEN = 1; - -class AsyncQueue { - private readonly items: T[] = []; - private readonly waiters: ((value: T | undefined) => void)[] = []; - private closed = false; - - push(item: T): void { - const waiter = this.waiters.shift(); - if (waiter) { - waiter(item); - return; - } - this.items.push(item); - } - - close(): void { - this.closed = true; - for (const waiter of this.waiters.splice(0)) waiter(undefined); - } - - next(): Promise { - const item = this.items.shift(); - if (item !== undefined) return Promise.resolve(item); - if (this.closed) return Promise.resolve(undefined); - return new Promise((resolve) => this.waiters.push(resolve)); - } -} - -export class BigRAGRealtimeConnection { - private socket: AnySocket | null = null; - private connecting: Promise | null = null; - private closing = false; - private closedError: Error | null = null; - private readonly queues = new Map>>(); - - constructor(private readonly client: RequestClient) {} - - async connect(): Promise { - if (this.socket?.readyState === OPEN) return; - if (this.connecting) return this.connecting; - this.connecting = this.open(); - try { - await this.connecting; - } finally { - this.connecting = null; - } - } - - async close(): Promise { - this.closing = true; - for (const queue of this.queues.values()) queue.close(); - this.queues.clear(); - this.socket?.close(); - this.socket = null; - } - - async *subscribe( - topic: string, - params: Record = {}, - ): AsyncGenerator> { - await this.connect(); - const id = randomId(); - const queue = new AsyncQueue>(); - this.queues.set(id, queue); - this.send({ type: "subscribe", id, topic, params: compactParams(params) }); - try { - while (true) { - const message = await queue.next(); - if (!message) { - if (this.closedError) throw this.closedError; - return; - } - yield message as RealtimeMessage; - if (message.type === "complete" || message.type === "error") return; - } - } finally { - this.send({ type: "unsubscribe", id }); - this.queues.delete(id); - queue.close(); - } - } - - private async open(): Promise { - const { Socket, supportsHeaders } = await websocketConstructor(this.client.apiKey !== ""); - if (this.client.apiKey && !supportsHeaders) { - throw new APIConnectionError( - "Browser WebSocket cannot send the Authorization header required for API-key auth; use session-cookie auth or a collection realtime token.", - ); - } - const url = realtimeUrl(this.client.baseUrl); - const headers: Record = { "User-Agent": USER_AGENT }; - if (this.client.apiKey) headers.Authorization = `Bearer ${this.client.apiKey}`; - const socket = ( - supportsHeaders ? new Socket(url, undefined, { headers }) : new Socket(url) - ) as AnySocket; - this.socket = socket; - await new Promise((resolve, reject) => { - let settled = false; - const done = (fn: () => void) => { - if (settled) return; - settled = true; - fn(); - }; - on(socket, "open", () => done(resolve)); - on(socket, "error", (event) => - done(() => reject(new APIConnectionError(websocketErrorMessage(event)))), - ); - on(socket, "message", (event) => this.handleMessage(event)); - on(socket, "close", () => this.handleClose()); - }); - } - - private send(message: unknown): void { - if (!this.socket || this.socket.readyState !== OPEN) return; - this.socket.send(JSON.stringify(message)); - } - - private handleMessage(event: unknown): void { - const raw = messageData(event); - if (raw === undefined) return; - let message: RealtimeMessage; - try { - message = JSON.parse(raw) as RealtimeMessage; - } catch { - return; - } - if (message.type === "heartbeat" || message.type === "pong" || message.type === "subscribed") { - return; - } - if (!message.id) return; - this.queues.get(message.id)?.push(message); - } - - private handleClose(): void { - this.socket = null; - if (!this.closing) { - this.closedError = new APIConnectionError("Realtime connection closed unexpectedly"); - } - for (const queue of this.queues.values()) queue.close(); - this.queues.clear(); - } -} - -export class RealtimeResource { - constructor(private readonly client: RequestClient) {} - - connect(): BigRAGRealtimeConnection { - return new BigRAGRealtimeConnection(this.client); - } - - async *subscribe( - topic: string, - params: Record = {}, - ): AsyncGenerator> { - const connection = this.connect(); - try { - yield* connection.subscribe(topic, params); - } finally { - await connection.close(); - } - } -} - -const randomId = () => { - const crypto = (globalThis as { crypto?: { randomUUID?: () => string } }).crypto; - if (crypto?.randomUUID) return crypto.randomUUID(); - return `${Date.now().toString(36)}-${Math.random().toString(36).slice(2)}-${Math.random().toString(36).slice(2)}`; -}; - -const websocketErrorMessage = (event: unknown): string => { - if (event && typeof event === "object") { - const candidate = event as { message?: unknown; error?: { message?: unknown } }; - if (typeof candidate.message === "string") return candidate.message; - if (candidate.error && typeof candidate.error.message === "string") { - return candidate.error.message; - } - } - return "WebSocket connection failed"; -}; - -const compactParams = (params: Record) => - Object.fromEntries(Object.entries(params).filter(([, value]) => value !== undefined)); - -const realtimeUrl = (baseUrl: string) => { - const url = new URL("/v1/realtime", `${baseUrl}/`); - url.protocol = url.protocol === "https:" ? "wss:" : "ws:"; - return url.toString(); -}; - -const websocketConstructor = async (needsHeaders: boolean) => { - if (needsHeaders && !isBrowserRuntime()) return websocketModuleConstructor(); - const globalSocket = (globalThis as { WebSocket?: unknown }).WebSocket; - if (globalSocket) { - return { - Socket: globalSocket as new (...args: unknown[]) => AnySocket, - supportsHeaders: false, - }; - } - return websocketModuleConstructor(); -}; - -const websocketModuleConstructor = async () => { - const dynamicImport = new Function("specifier", "return import(specifier)") as ( - specifier: string, - ) => Promise<{ WebSocket?: new (...args: unknown[]) => AnySocket; default?: unknown }>; - const module = await dynamicImport("ws"); - return { - Socket: (module.WebSocket ?? module.default) as new (...args: unknown[]) => AnySocket, - supportsHeaders: true, - }; -}; - -const isBrowserRuntime = () => - (globalThis as { window?: unknown }).window !== undefined && - (globalThis as { window?: unknown }).window === globalThis; - -const on = (socket: AnySocket, event: string, handler: (...args: unknown[]) => void) => { - if (socket.addEventListener) { - socket.addEventListener(event, (payload) => handler(payload)); - return; - } - socket.on?.(event, handler); -}; - -const messageData = (event: unknown): string | undefined => { - if (typeof event === "string") return event; - if (event instanceof Uint8Array) return new TextDecoder().decode(event); - if (event && typeof event === "object" && "data" in event) { - const data = (event as { data?: unknown }).data; - if (typeof data === "string") return data; - if (data instanceof Uint8Array) return new TextDecoder().decode(data); - } - return undefined; -}; diff --git a/sdks/typescript/src/resources/admin/index.ts b/sdks/typescript/src/resources/admin/index.ts index 670b3893..892d8c8b 100644 --- a/sdks/typescript/src/resources/admin/index.ts +++ b/sdks/typescript/src/resources/admin/index.ts @@ -4,7 +4,6 @@ import { AdminApiKeysResource } from "./api_keys.js"; import { AdminAuditResource } from "./audit.js"; import { AdminEmbeddingPresetsResource } from "./embedding_presets.js"; import { AdminMcpServersResource } from "./mcp_servers.js"; -import { AdminRealtimeResource } from "./realtime.js"; import { AdminSettingsResource } from "./settings.js"; import { AdminUsersResource } from "./users.js"; import { AdminVectorStorageResource } from "./vector_storage.js"; @@ -14,7 +13,6 @@ export { AdminApiKeysResource } from "./api_keys.js"; export { AdminAuditResource } from "./audit.js"; export { AdminEmbeddingPresetsResource } from "./embedding_presets.js"; export { AdminMcpServersResource } from "./mcp_servers.js"; -export { AdminRealtimeResource } from "./realtime.js"; export { AdminSettingsResource } from "./settings.js"; export { AdminUsersResource } from "./users.js"; export { AdminVectorStorageResource } from "./vector_storage.js"; @@ -26,7 +24,6 @@ export class AdminResource { readonly audit: AdminAuditResource; readonly embeddingPresets: AdminEmbeddingPresetsResource; readonly mcpServers: AdminMcpServersResource; - readonly realtime: AdminRealtimeResource; readonly settings: AdminSettingsResource; readonly vectorStorage: AdminVectorStorageResource; @@ -37,7 +34,6 @@ export class AdminResource { this.audit = new AdminAuditResource(client); this.embeddingPresets = new AdminEmbeddingPresetsResource(client); this.mcpServers = new AdminMcpServersResource(client); - this.realtime = new AdminRealtimeResource(client); this.settings = new AdminSettingsResource(client); this.vectorStorage = new AdminVectorStorageResource(client); } diff --git a/sdks/typescript/src/resources/admin/realtime.ts b/sdks/typescript/src/resources/admin/realtime.ts deleted file mode 100644 index 20749771..00000000 --- a/sdks/typescript/src/resources/admin/realtime.ts +++ /dev/null @@ -1,185 +0,0 @@ -import type { RequestClient } from "../../core.js"; -import { BigRAGRealtimeConnection } from "../../realtime.js"; -import type { - AccessLogListResponse, - AccessLogOverviewResponse, - AdminRealtimeEvent, - AuditLogListResponse, - BatchStatusResponse, - CollectionStatsResponse, - Document, - DocumentListResponse, - PlatformStatsResponse, - ReadinessResponse, - S3SourceListResponse, - S3SyncJobListResponse, - UploadSession, - UsageResponse, -} from "../../types/index.js"; - -export class AdminRealtimeResource { - constructor(private readonly _client: RequestClient) {} - - documents( - collection: string, - options: { - limit?: number; - offset?: number; - order?: "asc" | "desc"; - q?: string; - sort?: "created_at" | "updated_at" | "filename" | "file_size" | "chunk_count" | "status"; - status?: string; - } = {}, - ): AsyncGenerator> { - return this._stream("admin.collections.documents", { - collection, - limit: options.limit, - offset: options.offset, - order: options.order, - q: options.q, - sort: options.sort, - status: options.status, - }); - } - - documentBatchStatus( - collection: string, - documentIds: string[], - ): AsyncGenerator> { - return this._stream("admin.collections.documents.batch_status", { - collection, - document_ids: documentIds, - }); - } - - document(collection: string, documentId: string): AsyncGenerator> { - return this._stream("admin.collections.documents.detail", { - collection, - document_id: documentId, - }); - } - - uploadSession( - collection: string, - sessionId: string, - ): AsyncGenerator> { - return this._stream("admin.collections.upload_session", { collection, session_id: sessionId }); - } - - collectionStats(collection: string): AsyncGenerator> { - return this._stream("admin.collections.stats", { collection }); - } - - connectorSources( - provider: string, - options: { collection?: string } = {}, - ): AsyncGenerator> { - return this._stream("admin.connectors.sources", { - provider, - collection: options.collection, - }); - } - - connectorSyncJobs( - provider: string, - options: { collection?: string; sourceId?: string; limit?: number } = {}, - ): AsyncGenerator> { - return this._stream("admin.connectors.sync_jobs", { - provider, - collection: options.collection, - source_id: options.sourceId, - limit: options.limit, - }); - } - - accessOverview( - options: { windowDays?: number } = {}, - ): AsyncGenerator> { - return this._stream("admin.access.overview", { - window_days: options.windowDays, - }); - } - - accessLogs( - options: { - action?: string; - actorId?: string; - collection?: string; - method?: string; - path?: string; - statusFamily?: string; - success?: boolean; - limit?: number; - offset?: number; - } = {}, - ): AsyncGenerator> { - return this._stream("admin.access.logs", { - action: options.action, - actor_id: options.actorId, - collection: options.collection, - method: options.method, - path: options.path, - status_family: options.statusFamily, - success: options.success, - limit: options.limit, - offset: options.offset, - }); - } - - audit( - options: { - action?: string; - actorId?: string; - resourceType?: string; - limit?: number; - offset?: number; - } = {}, - ): AsyncGenerator> { - return this._stream("admin.audit", { - action: options.action, - actor_id: options.actorId, - resource_type: options.resourceType, - limit: options.limit, - offset: options.offset, - }); - } - - usage(options: { windowDays?: number } = {}): AsyncGenerator> { - return this._stream("admin.usage", { - window_days: options.windowDays, - }); - } - - platformStats(): AsyncGenerator> { - return this._stream("admin.platform.stats"); - } - - platformReadiness(): AsyncGenerator> { - return this._stream("admin.platform.readiness"); - } - - custom( - topic: string, - params: Record = {}, - ): AsyncGenerator> { - if (!topic.startsWith("admin.")) { - throw new Error("admin.realtime.custom topic must start with admin."); - } - return this._stream(topic, params); - } - - private async *_stream( - topic: string, - params: Record = {}, - ): AsyncGenerator> { - const connection = new BigRAGRealtimeConnection(this._client); - try { - for await (const message of connection.subscribe(topic, params)) { - yield message as AdminRealtimeEvent; - if (message.type === "complete") return; - } - } finally { - await connection.close(); - } - } -} diff --git a/sdks/typescript/src/resources/collections.ts b/sdks/typescript/src/resources/collections.ts index bd96138b..9b5c0c83 100644 --- a/sdks/typescript/src/resources/collections.ts +++ b/sdks/typescript/src/resources/collections.ts @@ -1,14 +1,11 @@ import type { RequestClient } from "../core.js"; -import { BigRAGRealtimeConnection } from "../realtime.js"; import type { AnalyticsResponse, Collection, CollectionListOptions, CollectionListResponse, - CollectionRealtimeTokenResponse, CollectionStatsResponse, CreateCollectionBody, - ProgressEvent, StatusResponse, UpdateCollectionBody, } from "../types/index.js"; @@ -68,30 +65,4 @@ export class CollectionsResource { analytics(name: string): Promise { return this._client._request("GET", `/v1/collections/${encodeURIComponent(name)}/analytics`); } - - createRealtimeToken(name: string): Promise { - return this._client._request( - "POST", - `/v1/collections/${encodeURIComponent(name)}/realtime-token`, - ); - } - - async *streamEvents( - name: string, - options: { token?: string } = {}, - ): AsyncGenerator { - const connection = new BigRAGRealtimeConnection(this._client); - try { - for await (const message of connection.subscribe("collection.events", { - collection: name, - token: options.token, - })) { - if (message.type === "event") yield message.payload; - if (message.type === "error") throw new Error(message.message); - if (message.type === "complete") return; - } - } finally { - await connection.close(); - } - } } diff --git a/sdks/typescript/src/types/admin.ts b/sdks/typescript/src/types/admin.ts index 34583f05..d5805c95 100644 --- a/sdks/typescript/src/types/admin.ts +++ b/sdks/typescript/src/types/admin.ts @@ -1,5 +1,4 @@ import type { User } from "./auth.js"; -import type { RealtimeMessage } from "./realtime.js"; export type InstanceSettingKind = | "bool" @@ -66,8 +65,6 @@ export interface InstanceSettingsTestResponse { message: string; } -export type AdminRealtimeEvent = RealtimeMessage; - export interface VectorStorageHealth { status: "ok" | "error"; error: string | null; diff --git a/sdks/typescript/src/types/collections.ts b/sdks/typescript/src/types/collections.ts index cc3d7c66..af05fd2f 100644 --- a/sdks/typescript/src/types/collections.ts +++ b/sdks/typescript/src/types/collections.ts @@ -46,11 +46,6 @@ export interface CollectionStatsResponse { status_counts: Record; } -export interface CollectionRealtimeTokenResponse { - token: string; - expires_in: number; -} - export interface CreateCollectionBody { name: string; description?: string; diff --git a/sdks/typescript/src/types/common.ts b/sdks/typescript/src/types/common.ts index 14e42601..e8081bfd 100644 --- a/sdks/typescript/src/types/common.ts +++ b/sdks/typescript/src/types/common.ts @@ -55,3 +55,20 @@ export interface PlatformStatsResponse { queue: QueueStatsResponse; workers?: WorkerStatsResponse; } + +export interface OverviewStatusResponse { + platform: PlatformStatsResponse; + readiness: ReadinessResponse; +} + +export interface CollectionsStatusResponse { + collections_total: number; + documents_total: number; + documents_ready: number; + documents_pending: number; + documents_processing: number; + documents_failed: number; + total_chunks: number; + total_tokens: number; + total_size_bytes: number; +} diff --git a/sdks/typescript/src/types/index.ts b/sdks/typescript/src/types/index.ts index 36ad5fc2..8522e476 100644 --- a/sdks/typescript/src/types/index.ts +++ b/sdks/typescript/src/types/index.ts @@ -10,7 +10,6 @@ export * from "./documents.js"; export * from "./embeddings.js"; export * from "./evaluations.js"; export * from "./query.js"; -export * from "./realtime.js"; export * from "./usage.js"; export * from "./vectors.js"; export * from "./webhooks.js"; diff --git a/sdks/typescript/src/types/realtime.ts b/sdks/typescript/src/types/realtime.ts deleted file mode 100644 index a3d3b02c..00000000 --- a/sdks/typescript/src/types/realtime.ts +++ /dev/null @@ -1,53 +0,0 @@ -export interface RealtimeSnapshot { - type: "snapshot"; - id: string; - topic: string; - payload: T; - generated_at: string; -} - -export interface RealtimeEvent { - type: "event"; - id: string; - topic: string; - payload: T; - generated_at: string; -} - -export interface RealtimeError { - type: "error"; - id?: string; - topic?: string; - message: string; - generated_at: string; -} - -export interface RealtimeComplete { - type: "complete"; - id: string; - topic: string; - generated_at: string; -} - -export interface RealtimeControlMessage { - type: "subscribed" | "heartbeat" | "pong"; - id?: string; - topic?: string; - generated_at?: string; -} - -export type RealtimeMessage = - | RealtimeSnapshot - | RealtimeEvent - | RealtimeError - | RealtimeComplete - | RealtimeControlMessage; - -export interface ProgressEvent { - step: string; - message: string; - progress: number; - document_id?: string; - status?: string; - detail?: Record; -} diff --git a/sdks/typescript/src/types/usage.ts b/sdks/typescript/src/types/usage.ts index cd77a739..7831cce7 100644 --- a/sdks/typescript/src/types/usage.ts +++ b/sdks/typescript/src/types/usage.ts @@ -18,5 +18,7 @@ export interface UsageResponse { storage_bytes_total: number; embedding_tokens_total: number; embedding_cost_usd_estimate: number; + avg_latency_ms: number; + timeline: { date: string; queries: number; avg_latency_ms: number }[]; by_collection: CollectionUsage[]; } diff --git a/tests/load/streaming.py b/tests/load/streaming.py index 6b26822f..a2ec46e3 100644 --- a/tests/load/streaming.py +++ b/tests/load/streaming.py @@ -17,7 +17,7 @@ from bigrag import APIError, BigRAG -PAYLOAD_BYTES = 512 +PAYLOAD_BYTES = 1000 * 1024 PAYLOAD_BODY_ALPHABET = "abcdefghijklmnopqrstuvwxyz0123456789 " diff --git a/website/content/docs/admin-ui.mdx b/website/content/docs/admin-ui.mdx index 45ad14d3..0e753dac 100644 --- a/website/content/docs/admin-ui.mdx +++ b/website/content/docs/admin-ui.mdx @@ -33,8 +33,8 @@ For cross-site production deployments, enable secure session cookies and set `Sa | `/overview` | Platform-wide stats — collection count, document states, stored token count, queue depth, readiness, and worker heartbeat. | | `/collections` | List, create, delete, and search collections. | | `/collections/[name]` | Collection overview with configuration badges, status counts, recent documents, and shortcuts. | -| `/collections/[name]/documents` | Upload files or folders through resumable upload sessions, filter/search/sort documents with infinite scrolling, inspect status, bulk delete with typed confirmation, and delete individual documents. Uploads are disabled when `/v1/stats` reports no live `bigrag-worker`. | -| `/collections/[name]/documents/[docId]` | Document detail — chunk viewer, live progress, failed-ingest error, and delete. | +| `/collections/[name]/documents` | Upload files or folders through resumable upload sessions, filter/search/sort documents with infinite scrolling, inspect status, bulk delete with typed confirmation, and delete individual documents. Uploads are disabled when `/v1/stats` reports no running `bigrag-worker`. | +| `/collections/[name]/documents/[docId]` | Document detail — chunk viewer, ingestion progress, failed-ingest error, and delete. | | `/collections/[name]/connectors` | Redirects to the first available collection connector from the provider catalog. | | `/collections/[name]/connectors/s3` | Open the collection-scoped S3/R2 connector workspace with an inline source list, an Add source modal for bucket-prefix credentials, sync progress, manual resync, and schedule controls. Source creation and manual sync are disabled when the worker is offline. | | `/collections/[name]/search` | Query tester scoped to the collection — mode segments, top-k control, skip-cache runs, cancel/retry states, result metadata, document filenames, and chunk deep links. | @@ -68,30 +68,32 @@ Settings accepts tab values via query param: `/settings?tab=account`, `/settings | Group | Section | Reads from | Writes to | |-------|---------|-----------|-----------| | **Personal** | Account | `GET /v1/auth/me` | `PATCH /v1/admin/users/{user_id}`, `POST /v1/auth/password`, `POST /v1/auth/logout-all` | -| **Operate** | Health | `GET /health/ready`, `admin.platform.readiness` realtime topic | — | +| **Operate** | Health | `GET /health/ready`, `GET /v1/status/overview` | — | | **Operate** | Security | `GET /v1/admin/settings` | `PUT /v1/admin/settings`, `POST /v1/admin/settings/embedding-cache/purge` | | **Operate** | Data | `GET /v1/admin/settings` | `PUT /v1/admin/settings` | | **Operate** | Vector storage | `GET /v1/admin/settings` | `PUT /v1/admin/settings` | -## Realtime updates +## Polling updates -The admin UI uses one session-cookie [realtime WebSocket](/docs/api-reference/realtime) for live document progress, batch upload progress, collection stats, connector sync state, access logs, paginated audit log pages, usage, platform stats, and readiness. Each subscription sends an immediate full `snapshot` payload matching the related REST response, then sends later snapshots as data changes. The socket stays open across dashboard navigation and only swaps topic subscriptions as pages mount and unmount. Overview platform stats refresh from ingestion events instead of waiting on the REST stats cache. The UI replays active subscriptions after reconnect and falls back to REST polling when the socket cannot deliver an initial snapshot. +The admin UI does not open sockets. Page header numbers use lightweight REST status endpoints and React Query polling: -The collection documents page keeps pagination on REST fetches for initial load, filter changes, and loading more rows. The first page stays live through `admin.collections.documents`, so newly created matching documents appear without a manual refresh. Loaded `pending` and `processing` rows are watched through `admin.collections.documents.batch_status` in batches of up to 100 document IDs on the shared WebSocket; newly paginated active rows are added to the watched set, and terminal rows are removed. If realtime is unavailable, the page falls back to first-page and status refreshes for the affected ID batch instead of polling every loaded page. +| Surface | Endpoint | Interval | +|---------|----------|----------| +| Overview | `GET /v1/status/overview` | 5 seconds | +| Usage | `GET /v1/status/usage?window_days=...` | 5 seconds | +| Access log header cards | `GET /v1/admin/status/access?window_days=...` | 5 seconds | -Expensive overview surfaces keep short-lived backend caches behind the readiness and access-overview loaders so several open admin sessions do not rerun the same aggregate or dependency checks on every interval tick. +Platform, collection, usage, and access-log stat endpoints are uncached so reloads and manual refreshes reflect the latest totals. Readiness and provider health checks may still cache dependency probes because they are not document or traffic stats. -**Worker-offline UX:** The same platform stats stream drives the offline indicator. When `workers.online` is `false`, the admin UI shows `bigrag-worker is offline`, includes the last heartbeat age when available, and disables actions that would enqueue uploads, S3 source sync, or manual S3 resync. Missing worker stats are treated as unknown, so actions remain available until an offline heartbeat is confirmed. +The collection documents page keeps pagination on REST fetches for initial load, filter changes, loading more rows, and the toolbar refresh button. It does not auto-refresh the document list. Document detail pages may poll the single-document endpoint every 5 seconds while a document is still `pending` or `processing`, then stop once the document is `ready` or `failed`. -**Large upload progress:** After a large local upload starts, the collection documents page subscribes to `admin.collections.upload_session` with the upload session ID. The snapshot returns received, queued, ingesting, completed, failed, and canceled counts. The panel renders the active summary plus a small item window that prioritizes active and failed files before older completed rows, so 10,000-file uploads do not require a 10,000-ID status stream. If the browser cannot keep the WebSocket healthy, the panel falls back to two-second polling. If a linked document is deleted while the upload session remains visible, the item is treated as canceled rather than active queued work. Terminal sessions are removed automatically once no active or failed work remains. +**Worker-offline UX:** The overview status response drives the offline indicator. When `workers.online` is `false`, the admin UI shows `bigrag-worker is offline`, includes the last heartbeat age when available, and disables actions that would enqueue uploads, S3 source sync, or manual S3 resync. Missing worker stats are treated as unknown, so actions remain available until an offline heartbeat is confirmed. -**S3/R2:** The S3 connector page subscribes to `admin.connectors.sources` and `admin.connectors.sync_jobs` for the selected collection. The sync monitor shows object scanning, concurrent download/update progress, remote deletion cleanup, and queued-for-ingestion completion; document conversion and indexing continue through the normal document progress surfaces. +**Large upload progress:** After a large local upload starts, the collection documents page polls `GET /v1/collections/{collection}/upload-sessions/{session_id}` every 2 seconds while the upload session is active and visible. The panel renders the active summary plus a small item window that prioritizes active and failed files before older completed rows, so 10,000-file uploads do not require 10,000 document status reads. If a linked document is deleted while the upload session remains visible, the item is treated as canceled rather than active queued work. Terminal sessions are removed automatically once no active or failed work remains. -Upload progress restoration is browser-local admin UI state. Upload session IDs are persisted in local storage so a collection's upload panel can reconnect after navigation; stale upload-session IDs are cleared without a failure alert when the API returns `404`. Chat turns are in-memory only and are not persisted in browser storage or the REST API. +**S3/R2:** The S3 connector page polls source and sync-job REST endpoints every 2 seconds. The sync monitor shows object scanning, concurrent download/update progress, remote deletion cleanup, and queued-for-ingestion completion; document conversion and indexing continue through the normal document progress surfaces. -REST document status endpoints remain available for API clients and SDKs that prefer polling. Collection-wide ingestion events are available through the `collection.events` realtime topic. - -Realtime pages only evaluate terminal-state predicates after the first snapshot or poll response exists, so stale browser-local upload-session state cannot crash the documents route during initial load. +Upload progress restoration is browser-local admin UI state. Upload session IDs are persisted in local storage so a collection's upload panel can resume polling after navigation; stale upload-session IDs are cleared without a failure alert when the API returns `404`. Chat turns are in-memory only and are not persisted in browser storage or the REST API. ## Chat specifics diff --git a/website/content/docs/api-reference/authentication.mdx b/website/content/docs/api-reference/authentication.mdx index 9a0f8b2d..0b7619b0 100644 --- a/website/content/docs/api-reference/authentication.mdx +++ b/website/content/docs/api-reference/authentication.mdx @@ -4,7 +4,6 @@ description: Admin accounts with session cookies and minted API keys. --- import { Callout } from "fumadocs-ui/components/callout"; -import { Tabs, Tab } from "fumadocs-ui/components/tabs"; Every request is authenticated as either: @@ -75,24 +74,12 @@ curl -X POST http://localhost:4000/v1/admin/api-keys \ # "key": "bigrag_sk_abcd…_XXXXXXXXXXXXXXXXXXXXXXXX"} ``` -Send the key as a bearer token. For browser WebSocket clients that cannot set bearer headers, mint a short-lived collection realtime token first. +Send the key as a bearer token: - - ```bash curl http://localhost:4000/v1/collections \ -H "Authorization: Bearer bigrag_sk_…" ``` - - -```bash -REALTIME_TOKEN=$(curl -s -X POST http://localhost:4000/v1/collections/docs/realtime-token \ - -H "Authorization: Bearer bigrag_sk_…" | jq -r .token) -``` - - - -The `collection.events` realtime topic accepts bearer headers or short-lived realtime tokens. Long-lived API keys are not accepted in query strings or WebSocket message params. ### Key fields @@ -124,12 +111,12 @@ Scopes use `resource:action` with `*` as a wildcard. | `chat:write` | Generate chat answers and question suggestions | | `vector:write` | Direct raw vector upsert | | `vector:delete` | Direct raw vector delete | -| `audit:read` | Read the API-key-capable usage rollup at `/v1/usage` | +| `audit:read` | Read the API-key-capable usage rollups at `/v1/usage` and `/v1/status/usage` | | `*:*` | Full access (equivalent to no `scopes`) | A key missing a required scope returns `403` with `{"detail": "API key missing required scope: "}`. -If a key is pinned to a single `collection`, that pin is enforced even on the global read helpers like `/v1/documents/{id}` by resolving the document's owning collection before returning it. +If a key is pinned to a single `collection`, that pin is enforced even on the global read helpers like `/v1/documents/{id}` by resolving the document's owning collection before returning it. Cross-collection status endpoints (`/v1/status/overview`, `/v1/status/collections`, and `/v1/status/usage`) are blocked for pinned keys. ## Which endpoints accept which auth @@ -139,10 +126,9 @@ If a key is pinned to a single `collection`, that pin is enforced even on the gl | `GET /v1/auth/setup-status`, `POST /v1/auth/setup` (first run only) | — | — | ✓ | | `POST /v1/auth/login` | — | — | ✓ | | `/v1/auth/*` (logout, me, whoami, password) | ✓ | ✓* | — | -| `/v1/collections`, `/v1/collections/**`, `/v1/query`, `/v1/batch/query`, `/v1/stats`, `/v1/usage`, `/v1/embeddings/models` | ✓ | ✓ | — | -| `WS /v1/realtime` collection topics, `POST /v1/collections/{name}/realtime-token` | ✓ | ✓ | — | +| `/v1/collections`, `/v1/collections/**`, `/v1/query`, `/v1/batch/query`, `/v1/stats`, `/v1/usage`, `/v1/status/overview`, `/v1/status/collections`, `/v1/status/usage`, `/v1/embeddings/models` | ✓ | ✓ | — | | `/v1/admin/api-keys`, `/v1/admin/users`, `/v1/admin/settings`, `/v1/admin/vector-storage`, `/v1/admin/connectors`, `/v1/admin/webhooks`, `/v1/admin/embedding-presets`, `/v1/admin/mcp-servers` | ✓ | — | — | -| `/v1/admin/access`, `/v1/admin/audit`, `WS /v1/realtime` admin topics, `/v1/auth/preferences` | ✓ | — | — | +| `/v1/admin/access`, `/v1/admin/audit`, `/v1/admin/status/access`, `/v1/auth/preferences` | ✓ | — | — | \* Read-only auth endpoints (`/me`) accept either; state-changing ones (`/password`, `/logout-all`) require a session. diff --git a/website/content/docs/api-reference/chat.mdx b/website/content/docs/api-reference/chat.mdx index ba55c7f0..7684f834 100644 --- a/website/content/docs/api-reference/chat.mdx +++ b/website/content/docs/api-reference/chat.mdx @@ -36,7 +36,7 @@ POST /v1/chat/question-suggestions ```json { "collection": "handbook", - "model": "gpt-4o-mini", + "model": "gpt-4.1", "temperature": 0.7 } ``` @@ -54,7 +54,7 @@ POST /v1/chat/question-suggestions "How do reimbursement limits work?" ], "generated_at": "2026-05-16T04:45:00+00:00", - "model": "gpt-4o-mini" + "model": "gpt-4.1" } ``` @@ -72,7 +72,7 @@ POST /v1/chat { "collection": "handbook", "message": "What is the PTO approval policy?", - "model": "gpt-4o-mini", + "model": "gpt-4.1", "top_k": 8, "search_mode": "hybrid", "rerank": true, diff --git a/website/content/docs/api-reference/collections.mdx b/website/content/docs/api-reference/collections.mdx index 135e2d9f..3769c977 100644 --- a/website/content/docs/api-reference/collections.mdx +++ b/website/content/docs/api-reference/collections.mdx @@ -206,48 +206,3 @@ Deletes every document, vector, and remaining staged ingestion file in the colle - -## Realtime - - - - -Creates a short-lived token for the `collection.events` WebSocket topic. Use it when a browser client cannot send an API-key bearer header during the WebSocket handshake. - -```http -POST /v1/collections/{name}/realtime-token -``` - -**Response** `200`: - -```json -{ "token": "...", "expires_in": 300 } -``` - - - - - -The `collection.events` topic streams every document event in the collection. It accepts a session cookie, an API-key bearer header, or the short-lived token from `POST /v1/collections/{name}/realtime-token`. - -```json -{ "type": "subscribe", "id": "events", "topic": "collection.events", "params": { "collection": "research_papers", "token": "..." } } -``` - -**Event payload:** - -```json -{ - "document_id": "abc-123", - "step": "embedded", - "status": "in_progress", - "message": "Embedded 128/482 chunks", - "progress": 26, - "detail": {} -} -``` - -`step` values: `received`, `converted`, `embedded`, `indexed`, `completed`, `failed`. - - - diff --git a/website/content/docs/api-reference/documents.mdx b/website/content/docs/api-reference/documents.mdx index 922403fb..46a31ace 100644 --- a/website/content/docs/api-reference/documents.mdx +++ b/website/content/docs/api-reference/documents.mdx @@ -388,13 +388,13 @@ Check status of up to 100 documents. { "document_ids": ["doc-id-1", "doc-id-2"] } ``` -Returns `id`, `status`, `chunk_count`, `error_message`, and the latest `progress` snapshot for each found document. API clients can poll this endpoint after batch upload until every requested document is `ready` or `failed`; the admin UI uses the `admin.collections.documents.batch_status` realtime topic instead. +Returns `id`, `status`, `chunk_count`, `error_message`, and the latest `progress` snapshot for each found document. API clients can poll this endpoint after batch upload until every requested document is `ready` or `failed`. -Use `GET /v1/collections/{collection_name}/documents/{document_id}` for a single document. Poll while `status` is `pending` or `processing`, then stop once it reaches `ready` or `failed`. The admin UI uses [realtime WebSocket subscriptions](/docs/api-reference/realtime) for the same full payloads. +Use `GET /v1/collections/{collection_name}/documents/{document_id}` for a single document. Poll while `status` is `pending` or `processing`, then stop once it reaches `ready` or `failed`. Document status responses include a `progress` object with the latest ingestion event snapshot: diff --git a/website/content/docs/api-reference/health.mdx b/website/content/docs/api-reference/health.mdx index 99debcd6..ef015dd1 100644 --- a/website/content/docs/api-reference/health.mdx +++ b/website/content/docs/api-reference/health.mdx @@ -63,7 +63,7 @@ Returns HTTP `503` with `"status": "degraded"` when any dependency is unhealthy. -Platform-wide statistics including collections, documents, and queue status. +Platform-wide statistics including collections, documents, and queue status. This endpoint is not cached. **Response** `200`: diff --git a/website/content/docs/api-reference/meta.json b/website/content/docs/api-reference/meta.json index 122917d5..8eeff056 100644 --- a/website/content/docs/api-reference/meta.json +++ b/website/content/docs/api-reference/meta.json @@ -5,7 +5,7 @@ "users", "api-keys", "audit", - "realtime", + "status", "instance-settings", "collections", "documents", diff --git a/website/content/docs/api-reference/realtime.mdx b/website/content/docs/api-reference/realtime.mdx deleted file mode 100644 index 7a58bafe..00000000 --- a/website/content/docs/api-reference/realtime.mdx +++ /dev/null @@ -1,75 +0,0 @@ ---- -title: Realtime -description: One WebSocket endpoint for subscribed realtime snapshots and events. ---- - -import { Callout } from "fumadocs-ui/components/callout"; - -Realtime subscriptions use one WebSocket connection at `WS /v1/realtime`. Clients subscribe and unsubscribe topics over that socket instead of opening one HTTP stream per page or resource. - -Chat streaming remains separate because `POST /v1/chat` is request-response generation streaming, not a reusable subscription surface. - -## Client Messages - -```json -{ "type": "subscribe", "id": "sub_1", "topic": "admin.collections.documents", "params": { "collection": "docs" } } -{ "type": "unsubscribe", "id": "sub_1" } -{ "type": "ping" } -``` - -Subscription IDs are client-generated and scoped to the WebSocket connection. Reusing an ID replaces the existing subscription with the new topic. - -## Server Messages - -```json -{ "type": "subscribed", "version": 1, "id": "sub_1", "topic": "documents:list:docs", "generated_at": "2026-05-23T00:00:00Z" } -{ "type": "snapshot", "version": 1, "id": "sub_1", "topic": "documents:list:docs", "payload": { "...": "..." }, "generated_at": "2026-05-23T00:00:00Z" } -{ "type": "event", "version": 1, "id": "sub_1", "topic": "collection.events:docs", "payload": { "...": "..." }, "generated_at": "2026-05-23T00:00:00Z" } -{ "type": "error", "version": 1, "id": "sub_1", "topic": "collection.events", "message": "Collection not found", "generated_at": "2026-05-23T00:00:00Z" } -{ "type": "complete", "version": 1, "id": "sub_1", "topic": "documents:detail:docs:doc_1", "generated_at": "2026-05-23T00:00:00Z" } -{ "type": "pong", "version": 1, "generated_at": "2026-05-23T00:00:00Z" } -{ "type": "heartbeat", "version": 1, "generated_at": "2026-05-23T00:00:00Z" } -``` - -Every server message carries a `version` for the protocol envelope. The `topic` field is the backend's resolved topic key (handy for debugging); clients demultiplex on the subscription `id`, not the topic. Snapshot payloads match the related REST response shape. - -## Authentication - -Admin topics require an admin session cookie. Cookie-authenticated WebSocket connections validate the `Origin` header against the configured CORS origins or the backend's own origin. - -Collection event topics accept any of: -- Admin or member session cookie. -- API-key bearer auth with `collection:read`. -- Short-lived realtime token from `POST /v1/collections/{name}/realtime-token`. - -Realtime tokens are useful when a browser cannot send bearer headers during the WebSocket handshake. Tokens are short-lived (5 minutes) and collection-scoped, and stay valid for the duration of the connection so they survive automatic reconnects. - -```bash -TOKEN=$(curl -s -X POST http://localhost:4000/v1/collections/docs/realtime-token \ - -H "Authorization: Bearer $BIGRAG_API_KEY" | jq -r .token) -``` - -## Topics - -| Topic | Params | Payload shape | Lifecycle | -|-------|--------|---------------|-----------| -| `admin.collections.documents` | `collection`, `q`, `status`, `sort`, `order`, `limit`, `offset`, `include_total` | `GET /v1/collections/{name}/documents` | stays open | -| `admin.collections.documents.detail` | `collection`, `document_id` | `GET /v1/collections/{name}/documents/{document_id}` | completes when `ready` or `failed` | -| `admin.collections.documents.batch_status` | `collection`, `document_ids` | `POST /v1/collections/{name}/documents/batch/status` | completes when watched documents are terminal or missing/deleted with no active returned documents | -| `admin.collections.upload_session` | `collection`, `session_id` | `GET /v1/collections/{name}/upload-sessions/{session_id}` | completes when `complete`, `failed`, or `canceled` | -| `admin.collections.stats` | `collection` | `GET /v1/collections/{name}/stats` | stays open | -| `admin.connectors.sources` | `provider`, `collection` | `GET /v1/connectors/{provider}/sources` | stays open | -| `admin.connectors.sync_jobs` | `provider`, `collection`, `source_id`, `limit` | `GET /v1/connectors/{provider}/sync-jobs` | stays open | -| `admin.access.overview` | `window_days` | `GET /v1/admin/access/overview` | stays open | -| `admin.access.logs` | access-log filters, `limit`, `offset` | `GET /v1/admin/access/logs` | stays open | -| `admin.audit` | audit filters, `limit`, `offset` | `GET /v1/admin/audit` | stays open | -| `admin.usage` | `window_days` | `GET /v1/usage` | stays open | -| `admin.platform.stats` | none | `GET /v1/stats` | stays open | -| `admin.platform.readiness` | none | `GET /health/ready` | stays open | -| `collection.events` | `collection`, optional `token` | ingestion progress event | stays open | - -The backend uses event-bus wakeups for ingestion-backed topics, connector source changes, and connector sync progress. Paginated or aggregate admin surfaces still send full snapshots so reconnects and reloads see the same shape as the REST endpoint. - - - REST endpoints remain available for polling fallback. Realtime SSE endpoints were removed in favor of `WS /v1/realtime`. - diff --git a/website/content/docs/api-reference/status.mdx b/website/content/docs/api-reference/status.mdx new file mode 100644 index 00000000..ac443436 --- /dev/null +++ b/website/content/docs/api-reference/status.mdx @@ -0,0 +1,115 @@ +--- +title: Status +description: Lightweight status endpoints for admin UI polling. +--- + +import { Callout } from "fumadocs-ui/components/callout"; +import { Accordions, Accordion } from "fumadocs-ui/components/accordion"; + +Status endpoints return page-specific aggregate numbers for dashboards and operators. They are designed for short-interval polling and do not cache stat totals. Readiness and provider health checks may still cache dependency probes because they are not document or traffic stats. + + +Collection-pinned API keys cannot use the cross-collection status endpoints. Use collection-specific document, stats, query, and upload-session endpoints for pinned integrations. + + + + + + +Authenticated overview numbers for platform stats and readiness. + +```http +GET /v1/status/overview +``` + +**Response** `200`: + +```json +{ + "platform": { + "collections": 4, + "documents": { + "total": 124, + "ready": 120, + "pending": 1, + "processing": 2, + "failed": 1, + "total_chunks": 9200, + "total_tokens": 1820432, + "total_size_bytes": 52428800 + }, + "webhooks": 2, + "queue": {}, + "workers": { + "online": true, + "heartbeat_at": "2026-05-25T06:00:00Z", + "heartbeat_age_seconds": 4 + } + }, + "readiness": { + "status": "ok", + "version": "2026.5.23", + "postgres": true, + "vector_store": true, + "redis": true + } +} +``` + +Requires a session or an API key with `collection:read`. Collection-pinned API keys are blocked. + + + + + +Authenticated collection and document aggregate numbers. + +```http +GET /v1/status/collections +``` + +**Response** `200`: + +```json +{ + "collections_total": 4, + "documents_total": 124, + "documents_ready": 120, + "documents_pending": 1, + "documents_processing": 2, + "documents_failed": 1, + "total_chunks": 9200, + "total_tokens": 1820432, + "total_size_bytes": 52428800 +} +``` + +Requires a session or an API key with `collection:read`. Collection-pinned API keys are blocked. + + + + + +Authenticated usage report for the usage page. + +```http +GET /v1/status/usage?window_days=30 +``` + +Uses the same response shape as [`GET /v1/usage`](/docs/api-reference/usage). Requires a session or an API key with `audit:read`. Collection-pinned API keys are blocked. + + + + + +Admin-only access-log numbers for the access-log page header cards. + +```http +GET /v1/admin/status/access?window_days=7 +``` + +Uses the same response shape as `GET /v1/admin/access/overview`. Requires an admin session. + + + + diff --git a/website/content/docs/api-reference/usage.mdx b/website/content/docs/api-reference/usage.mdx index f83c662c..b8582419 100644 --- a/website/content/docs/api-reference/usage.mdx +++ b/website/content/docs/api-reference/usage.mdx @@ -11,6 +11,8 @@ GET /v1/usage Accepts a session cookie or an API key with `audit:read`. Returns the per-collection resource footprint over a trailing window — useful for quota enforcement, billing, and capacity planning. +`GET /v1/status/usage` returns the same response shape for admin UI polling. Collection-pinned API keys cannot use the status variant because it aggregates across collections. + ## Query parameters | Parameter | Type | Default | Notes | diff --git a/website/content/docs/comparison.mdx b/website/content/docs/comparison.mdx index 137c14cd..33ec1dff 100644 --- a/website/content/docs/comparison.mdx +++ b/website/content/docs/comparison.mdx @@ -46,7 +46,7 @@ These are the closest alternatives to bigRAG — self-hostable platforms that ha - Multi-collection architecture with per-collection embedding configuration (OpenRAG uses a single index) - Turbopuffer namespaces per collection for scoped vector, keyword, hybrid, export, and delete operations - Mature webhook system with HMAC signatures, retry logic, and delivery tracking -- Batch operations — 100-file upload, WebSocket admin progress, 20-query batch, bulk status/delete +- Batch operations — 100-file upload, pollable admin progress, 20-query batch, bulk status/delete - Query analytics with 24h/7d/30d windows and top queries per collection - Raw vector upsert API for bring-your-own-embeddings workflows - Broader file format support (18+ formats including XLSX, CSV, TSV, XML, JSON) diff --git a/website/content/docs/concepts/embeddings.mdx b/website/content/docs/concepts/embeddings.mdx index e0265e61..0392f524 100644 --- a/website/content/docs/concepts/embeddings.mdx +++ b/website/content/docs/concepts/embeddings.mdx @@ -161,13 +161,15 @@ Redis query caches are also encrypted when `BIGRAG_MASTER_KEY` is configured. Th ## Concurrency & throughput -Embedding requests are guarded by a semaphore (`BIGRAG_EMBEDDING_CONCURRENCY`, default 8): +Embedding requests are guarded by a Redis-backed adaptive limiter with `embedding_concurrency` as its ceiling (`BIGRAG_EMBEDDING_CONCURRENCY`, default 8): - Raise it for high-QPS providers. - Lower it if your embedding provider throttles requests. When a provider returns a rate limit with `Retry-After`, `retry-after-ms`, or a message such as `Please try again in 37ms`, ingestion records a short Redis cooldown for that provider/model, waits for the hint, and retries the same batch without consuming the generic transient retry budget. Repeated rate limits are capped per batch, so sustained TPM pressure should be handled by lowering concurrency or batch size. +Changing `embedding_concurrency` in `/settings` resets the local and Redis limiter state so workers use the new ceiling immediately. + Token counting uses `tiktoken` where the provider ships a tokenizer and falls back to a 4-character-per-token heuristic otherwise. ## Reranking diff --git a/website/content/docs/sdks/mcp.mdx b/website/content/docs/sdks/mcp.mdx index 7f4c1872..95789d09 100644 --- a/website/content/docs/sdks/mcp.mdx +++ b/website/content/docs/sdks/mcp.mdx @@ -156,5 +156,5 @@ Once configured: ## Limitations - Ingestion (upload and webhooks) is intentionally not exposed — MCP tools target retrieval workflows. Use the HTTP API, a language SDK, or the admin UI for writes. -- WebSocket realtime subscriptions are not wrapped — MCP tools are request/response only. +- Operator status polling endpoints are not wrapped — MCP tools are request/response retrieval helpers. - OAuth 2.0 flows are not implemented. Remote HTTP clients must send `Authorization: Bearer ...`; URL query-token authentication is not accepted on `/mcp`. diff --git a/website/content/docs/sdks/python.mdx b/website/content/docs/sdks/python.mdx index d0316bb8..2db246cc 100644 --- a/website/content/docs/sdks/python.mdx +++ b/website/content/docs/sdks/python.mdx @@ -65,15 +65,14 @@ The client reads `BIGRAG_API_KEY` from the environment if `api_key` is not passe | Namespace | Description | |---|---| -| `client.collections` | Collection CRUD, stats, realtime tokens, and event streams | +| `client.collections` | Collection CRUD, stats, and analytics | | `client.documents` | Document upload, list, delete, batch ops | | `client.chat` | Generated answers, question suggestions, and streaming | | `client.queries` | Single, multi-collection, and batch queries | | `client.vectors` | Raw vector upsert and delete | | `client.webhooks` | Webhook management | | `client.auth` | Setup, login, logout, identity, password, and preferences | -| `client.admin` | Users, API keys, access logs, audit logs, runtime settings, vector storage overview, admin realtime helpers, embedding presets, and MCP server keys | -| `client.realtime` | Explicit WebSocket connect, subscribe, and unsubscribe control | +| `client.admin` | Users, API keys, access logs, audit logs, runtime settings, vector storage overview, embedding presets, and MCP server keys | | `client.connectors.s3` | S3-compatible bucket prefixes, sources, and sync jobs | | `client.evaluations` | Golden-set retrieval evaluation runs | @@ -117,11 +116,6 @@ stats = await client.collections.stats("docs") # Truncate (delete all documents, keep the collection) await client.collections.truncate("docs") - -# Stream collection realtime events -token = await client.collections.create_realtime_token("docs") -async for event in client.collections.stream_events("docs", token=token["token"]): - print(event) ``` @@ -366,10 +360,6 @@ await client.admin.settings.reset({"keys": ["trusted_proxies"]}) await client.admin.settings.purge_embedding_cache() overview = await client.admin.vector_storage.overview() - -async for event in client.admin.realtime.platform_readiness(): - if event["type"] == "snapshot": - print(event["payload"]) ``` ## Platform Endpoints @@ -387,6 +377,12 @@ stats = await client.get_stats() # Usage rollup usage = await client.get_usage(window_days=30) +# Pollable admin UI status helpers +overview_status = await client.get_overview_status() +collections_status = await client.get_collections_status() +usage_status = await client.get_usage_status(window_days=30) +access_status = await client.get_access_status(window_days=7) + # Available embedding models models = await client.list_embedding_models() diff --git a/website/content/docs/sdks/typescript.mdx b/website/content/docs/sdks/typescript.mdx index 20231daa..c2c8611d 100644 --- a/website/content/docs/sdks/typescript.mdx +++ b/website/content/docs/sdks/typescript.mdx @@ -13,7 +13,7 @@ import { Accordions, Accordion } from "fumadocs-ui/components/accordion"; npm install @bigrag/client ``` -Works in Node.js 18+, browsers, Deno, Bun, and edge runtimes. Browser realtime uses native `WebSocket`; Node realtime uses the SDK's `ws` dependency. +Works in Node.js 18+, browsers, Deno, Bun, and edge runtimes. Published npm releases use CalVer (`YYYY.M.D`), for example `@bigrag/client@2026.5.23`. Browser applications that only need SDK error classes and shared response types can import the browser-safe subpath: @@ -106,10 +106,6 @@ client.collections.update(name, { client.collections.delete(name) client.collections.truncate(name) // delete all documents, keep the collection client.collections.analytics(name) -client.collections.createRealtimeToken(name) -client.collections.streamEvents(name, { token? }) // AsyncIterable of collection realtime events -client.realtime.connect() -client.realtime.subscribe(topic, params?) ``` @@ -293,22 +289,6 @@ client.admin.mcpServers.rotate(serverId) client.admin.mcpServers.delete(serverId) ``` - - - - -```typescript -for await (const event of client.admin.realtime.platformReadiness()) { - if (event.type === "snapshot") { - console.log(event.payload) - } -} -``` - - - - - Admin endpoints (`/v1/admin/*` and most of `/v1/auth/*`) are session-only. API keys remain the right fit for read/query/ingest workloads. @@ -370,6 +350,10 @@ client.readiness() // GET /health/ready client.getStats() // GET /v1/stats client.listEmbeddingModels() // GET /v1/embeddings/models client.getUsage({ windowDays? }) // GET /v1/usage +client.getOverviewStatus() // GET /v1/status/overview +client.getCollectionsStatus() // GET /v1/status/collections +client.getUsageStatus({ windowDays? }) // GET /v1/status/usage +client.getAccessStatus({ windowDays? }) // GET /v1/admin/status/access client.collections.analytics(collection) // GET /v1/collections/{name}/analytics ```