diff --git a/.gitmodules b/.gitmodules new file mode 100644 index 000000000..a176df463 --- /dev/null +++ b/.gitmodules @@ -0,0 +1,3 @@ +[submodule "vendor/pg-llm-batch"] + path = vendor/pg-llm-batch + url = https://github.com/ContextualWisdomLab/pg-llm-batch.git diff --git a/README.md b/README.md index a5cc6252f..037478f7a 100644 --- a/README.md +++ b/README.md @@ -609,6 +609,29 @@ vector counts, unsupported embedding model names, fake quality totals, or inert permanent ready-soon controls; use source-backed rows or explicit pending states. +## Batch-tolerant LLM embeddings (optional submodule) + +`vendor/pg-llm-batch` is a git **submodule** (the org's standalone +`pg_tiktoken` Postgres batch engine, Apache-2.0). It powers batch-tolerant +hotspots such as bulk email-import embeddings. It is fully optional: naruon runs +normally with the submodule uninitialized — the batch path degrades to the +existing per-item embedding path whenever it is unconfigured or unavailable. + +To enable it (human actions): + +1. Initialize the submodule: `git submodule update --init vendor/pg-llm-batch` + and install it on the backend path (`pip install -e vendor/pg-llm-batch`). +2. Bring up the batch Postgres: + `docker compose -f docker-compose.yml -f docker-compose.pg-llm-batch.yml up`. +3. Seed per-tenant batch config in the Fernet DB (`tenant_configs`, never via + `os.getenv`): set `batch_embedding_enabled = true` and the Fernet-encrypted + `batch_embedding_dsn` (plus `batch_embedding_endpoint` / `batch_embedding_model`). + Provider credentials continue to resolve through `resolve_runtime_llm_provider`. + +Batch runs are recorded in the `llm_batch_jobs` / `llm_batch_items` control-plane +tables (migration `0010_llm_batch_embedding`). The submodule can also run +standalone — see `vendor/pg-llm-batch/README.md` and its own `docker-compose.yml`. + ## Operations and release docs - `docs/operations/release-deployment-architecture.md`: release, CI, GHCR, and diff --git a/backend/alembic/versions/0010_llm_batch_embedding.py b/backend/alembic/versions/0010_llm_batch_embedding.py new file mode 100644 index 000000000..b96832bcb --- /dev/null +++ b/backend/alembic/versions/0010_llm_batch_embedding.py @@ -0,0 +1,128 @@ +"""add llm batch embedding job/item tables and tenant batch config + +Revision ID: 0010_llm_batch_embedding +Revises: 0009_project_graph_projection +Create Date: 2026-07-08 00:00:00.000000 +""" + +from alembic import op +import sqlalchemy as sa + +revision = "0010_llm_batch_embedding" +down_revision = "0009_project_graph_projection" + +_JOBS_TABLE = "llm_batch_jobs" +_ITEMS_TABLE = "llm_batch_items" +_TENANT_TABLE = "tenant_configs" + + +def upgrade() -> None: + connection = op.get_bind() + inspector = sa.inspect(connection) + + if not inspector.has_table(_JOBS_TABLE): + op.create_table( + _JOBS_TABLE, + sa.Column("batch_job_uid", sa.String(), nullable=False), + sa.Column("organization_id", sa.String(), nullable=False), + sa.Column("user_id", sa.String(), nullable=False), + sa.Column("job_status", sa.String(), nullable=False), + sa.Column("model_name", sa.String(), nullable=False), + sa.Column("endpoint_alias", sa.String(), nullable=True), + sa.Column("total_items", sa.Integer(), nullable=False), + sa.Column("completed_items", sa.Integer(), nullable=False), + sa.Column("failed_items", sa.Integer(), nullable=False), + sa.Column("total_tokens", sa.Integer(), nullable=False), + sa.Column("part_count", sa.Integer(), nullable=False), + sa.Column("error_code", sa.String(), nullable=True), + sa.Column("created_at", sa.DateTime(timezone=True), nullable=False), + sa.Column("updated_at", sa.DateTime(timezone=True), nullable=False), + sa.PrimaryKeyConstraint("batch_job_uid"), + ) + + if not inspector.has_table(_ITEMS_TABLE): + op.create_table( + _ITEMS_TABLE, + sa.Column("batch_item_uid", sa.String(), nullable=False), + sa.Column("batch_job_uid", sa.String(), nullable=False), + sa.Column("sequence_no", sa.Integer(), nullable=False), + sa.Column("part_index", sa.Integer(), nullable=False), + sa.Column("token_count", sa.Integer(), nullable=False), + sa.Column("item_status", sa.String(), nullable=False), + sa.Column("error_code", sa.String(), nullable=True), + sa.Column("created_at", sa.DateTime(timezone=True), nullable=False), + sa.Column("updated_at", sa.DateTime(timezone=True), nullable=False), + sa.ForeignKeyConstraint( + ["batch_job_uid"], + [f"{_JOBS_TABLE}.batch_job_uid"], + ondelete="CASCADE", + ), + sa.PrimaryKeyConstraint("batch_item_uid"), + ) + + for table_name, index_name, columns in _batch_indexes(): + op.create_index(index_name, table_name, columns, if_not_exists=True) + + if inspector.has_table(_TENANT_TABLE): + for column in _tenant_batch_columns(): + if not _has_column(inspector, _TENANT_TABLE, column.name): + op.add_column(_TENANT_TABLE, column) + + +def downgrade() -> None: + connection = op.get_bind() + inspector = sa.inspect(connection) + + if inspector.has_table(_TENANT_TABLE): + for column in reversed(_tenant_batch_columns()): + if _has_column(inspector, _TENANT_TABLE, column.name): + op.drop_column(_TENANT_TABLE, column.name) + + for table_name, index_name, _columns in reversed(_batch_indexes()): + if inspector.has_table(table_name): + op.drop_index(index_name, table_name=table_name, if_exists=True) + + if inspector.has_table(_ITEMS_TABLE): + op.drop_table(_ITEMS_TABLE) + if inspector.has_table(_JOBS_TABLE): + op.drop_table(_JOBS_TABLE) + + +def _batch_indexes() -> list[tuple[str, str, list[str]]]: + return [ + (_JOBS_TABLE, "ix_llm_batch_jobs_organization_id", ["organization_id"]), + (_JOBS_TABLE, "ix_llm_batch_jobs_user_id", ["user_id"]), + (_JOBS_TABLE, "ix_llm_batch_jobs_job_status", ["job_status"]), + ( + _JOBS_TABLE, + "ix_llm_batch_jobs_scope_status", + ["organization_id", "user_id", "job_status"], + ), + (_ITEMS_TABLE, "ix_llm_batch_items_batch_job_uid", ["batch_job_uid"]), + (_ITEMS_TABLE, "ix_llm_batch_items_item_status", ["item_status"]), + ( + _ITEMS_TABLE, + "ix_llm_batch_items_job_sequence", + ["batch_job_uid", "sequence_no"], + ), + ] + + +def _tenant_batch_columns() -> list["sa.Column"]: + return [ + sa.Column( + "batch_embedding_enabled", + sa.Boolean(), + nullable=False, + server_default=sa.false(), + ), + sa.Column("batch_embedding_dsn", sa.String(), nullable=True), + sa.Column("batch_embedding_endpoint", sa.String(), nullable=True), + sa.Column("batch_embedding_model", sa.String(), nullable=True), + ] + + +def _has_column(inspector, table_name: str, column_name: str) -> bool: + return any( + column["name"] == column_name for column in inspector.get_columns(table_name) + ) diff --git a/backend/db/models.py b/backend/db/models.py index b68bb23f9..e20216e67 100644 --- a/backend/db/models.py +++ b/backend/db/models.py @@ -274,6 +274,116 @@ class ProviderWritebackRetryItem(Base): ) +class LlmBatchJob(Base): + """A batch-tolerant embedding/completion job routed via pg-llm-batch. + + naruon-side control-plane mirror of the component's ``llm_batches`` table. + Records one job per bulk embedding run (e.g. an email import batch) so the + batched work has a durable audit trail even though the JSONL assembly lives + in the batch engine's own Postgres. Modeled on + :class:`ProviderWritebackRetryItem` (string uid PK, scope indexes). + """ + + __tablename__ = "llm_batch_jobs" + + batch_job_uid: Mapped[str] = mapped_column( + String, + primary_key=True, + default=lambda: f"llm_batch_{uuid.uuid4().hex}", + ) + organization_id: Mapped[str] = mapped_column(String, index=True, nullable=False) + user_id: Mapped[str] = mapped_column(String, index=True, nullable=False) + job_status: Mapped[str] = mapped_column( + String, + index=True, + default="preparing", + nullable=False, + ) + model_name: Mapped[str] = mapped_column(String, nullable=False) + endpoint_alias: Mapped[str | None] = mapped_column(String, nullable=True) + total_items: Mapped[int] = mapped_column(Integer, default=0, nullable=False) + completed_items: Mapped[int] = mapped_column(Integer, default=0, nullable=False) + failed_items: Mapped[int] = mapped_column(Integer, default=0, nullable=False) + total_tokens: Mapped[int] = mapped_column(Integer, default=0, nullable=False) + part_count: Mapped[int] = mapped_column(Integer, default=0, nullable=False) + error_code: Mapped[str | None] = mapped_column(String, nullable=True) + created_at: Mapped[datetime.datetime] = mapped_column( + DateTime(timezone=True), + default=lambda: datetime.datetime.now(datetime.timezone.utc), + nullable=False, + ) + updated_at: Mapped[datetime.datetime] = mapped_column( + DateTime(timezone=True), + default=lambda: datetime.datetime.now(datetime.timezone.utc), + onupdate=lambda: datetime.datetime.now(datetime.timezone.utc), + nullable=False, + ) + items: Mapped[list["LlmBatchItem"]] = relationship( + back_populates="job", + cascade="all, delete-orphan", + ) + __table_args__ = ( + Index( + "ix_llm_batch_jobs_scope_status", + "organization_id", + "user_id", + "job_status", + ), + ) + + +class LlmBatchItem(Base): + """A single request within an :class:`LlmBatchJob`. + + naruon-side mirror of the component's ``llm_requests`` rows. One item per + input text, carrying its token count and the partition (batch file part) it + was assigned to by the engine's token/byte/record accumulator. + """ + + __tablename__ = "llm_batch_items" + + batch_item_uid: Mapped[str] = mapped_column( + String, + primary_key=True, + default=lambda: f"llm_batch_item_{uuid.uuid4().hex}", + ) + batch_job_uid: Mapped[str] = mapped_column( + String, + ForeignKey("llm_batch_jobs.batch_job_uid", ondelete="CASCADE"), + index=True, + nullable=False, + ) + sequence_no: Mapped[int] = mapped_column(Integer, nullable=False) + part_index: Mapped[int] = mapped_column(Integer, default=0, nullable=False) + token_count: Mapped[int] = mapped_column(Integer, default=0, nullable=False) + item_status: Mapped[str] = mapped_column( + String, + index=True, + default="queued", + nullable=False, + ) + error_code: Mapped[str | None] = mapped_column(String, nullable=True) + created_at: Mapped[datetime.datetime] = mapped_column( + DateTime(timezone=True), + default=lambda: datetime.datetime.now(datetime.timezone.utc), + nullable=False, + ) + updated_at: Mapped[datetime.datetime] = mapped_column( + DateTime(timezone=True), + default=lambda: datetime.datetime.now(datetime.timezone.utc), + onupdate=lambda: datetime.datetime.now(datetime.timezone.utc), + nullable=False, + ) + job: Mapped["LlmBatchJob"] = relationship(back_populates="items") + __table_args__ = ( + Index( + "ix_llm_batch_items_job_sequence", + "batch_job_uid", + "sequence_no", + ), + ) + + class Organization(Base): __tablename__ = "organization_entities" @@ -1042,6 +1152,20 @@ class TenantConfig(Base): EncryptedString, nullable=True ) + # Batch-tolerant embedding routing (pg-llm-batch submodule). All config here + # lives in the Fernet DB, never in os.getenv. The batch Postgres DSN is a + # connection secret, so it is stored EncryptedString (Fernet at rest). + batch_embedding_enabled: Mapped[bool] = mapped_column( + Boolean, default=False, nullable=False + ) + batch_embedding_dsn: Mapped[str | None] = mapped_column( + EncryptedString, nullable=True + ) + batch_embedding_endpoint: Mapped[str | None] = mapped_column( + String, nullable=True + ) + batch_embedding_model: Mapped[str | None] = mapped_column(String, nullable=True) + def __repr__(self) -> str: return ( f" ModuleType | None: + """Return the ``pg_llm_batch`` module, or ``None`` when unavailable. + + The submodule is an *optional* dependency: naruon must run with it + uninitialized. When the package cannot be imported we log once at debug + level and signal fallback by returning ``None``. + """ + if _ENGINE_CACHE: + return _ENGINE_CACHE[0] + try: + import pg_llm_batch # type: ignore + except ImportError: + logger.debug( + "pg_llm_batch submodule not importable; batch embedding disabled, " + "falling back to per-item path" + ) + _ENGINE_CACHE.append(None) + return None + _ENGINE_CACHE.append(pg_llm_batch) + return pg_llm_batch + + +@dataclass(frozen=True) +class BatchEmbeddingSettings: + """Per-tenant batch configuration resolved from the Fernet DB (never env).""" + + dsn: str + endpoint_alias: str | None + model: str + + +async def resolve_batch_embedding_settings( + session: AsyncSession, + *, + user_id: str, + organization_id: str | None, +) -> BatchEmbeddingSettings | None: + """Resolve batch settings from the per-tenant Fernet-encrypted config. + + Returns ``None`` (i.e. "route the normal per-item path") unless the tenant + has both enabled batching and stored a batch Postgres DSN. The DSN is read + back through the ``EncryptedString`` column, so it is decrypted from the + Fernet DB rather than read from the process environment. + """ + tenant_config = await get_scoped_tenant_config(session, user_id, organization_id) + if tenant_config is None: + return None + if not getattr(tenant_config, "batch_embedding_enabled", False): + return None + dsn = getattr(tenant_config, "batch_embedding_dsn", None) + if not dsn: + return None + return BatchEmbeddingSettings( + dsn=dsn, + endpoint_alias=getattr(tenant_config, "batch_embedding_endpoint", None), + model=(getattr(tenant_config, "batch_embedding_model", None) or "").strip() + or None, # resolved against the provider model below + ) + + +def _plan_partitions( + engine: ModuleType, dsn: str, model: str, texts: list[str] +) -> tuple[list[list[int]], list[int]]: + """Group text indices into token/byte/record-bounded partitions. + + Runs the component's ``TokenCounter`` + ``BatchAccumulator`` (which count + tokens inside Postgres via ``pg_tiktoken``) to split the batch the same way + the standalone engine would. Pure planning: returns index partitions plus + the per-text token counts. Synchronous (psycopg) — call via a thread. + """ + config = engine.PostgresConfigStore(dsn) + try: + counter = engine.TokenCounter(dsn, config=config) + accumulator = engine.BatchAccumulator(counter, model) + partitions: list[list[int]] = [] + current: list[int] = [] + token_counts: list[int] = [] + for index, text in enumerate(texts): + total_tokens, _system, _user = accumulator.compute_tokens("", text) + token_counts.append(int(total_tokens)) + byte_size = engine.BatchAccumulator.compute_byte_size(text) + if current and accumulator.would_exceed(total_tokens, byte_size): + partitions.append(current) + current = [] + accumulator.reset() + accumulator.add_entry(str(index), text, total_tokens, byte_size) + current.append(index) + if current: + partitions.append(current) + return partitions, token_counts + finally: + close = getattr(config, "close", None) + if callable(close): + try: + close() + except Exception: # pragma: no cover - defensive + pass + + +async def try_batch_import_embeddings( + session: AsyncSession, + texts: list[str], + *, + embedding_provider: "EmailImportEmbeddingProvider", + user_id: str, + organization_id: str | None, + dimension: int = STORAGE_EMBEDDING_DIMENSION, +) -> list[list[float]] | None: + """Route bulk embeddings through the batch engine, or ``None`` to fall back. + + On success returns one fitted vector per input text (original order). Any + failure — batch disabled, submodule missing, batch DB unreachable, embedding + error — returns ``None`` so the caller uses its per-item path. The batch run + is recorded in ``llm_batch_jobs`` / ``llm_batch_items`` for observability. + """ + if not texts: + return None + + settings = await resolve_batch_embedding_settings( + session, user_id=user_id, organization_id=organization_id + ) + if settings is None: + return None + + engine = load_batch_engine() + if engine is None: + return None + + model = settings.model or embedding_provider.embedding_model + + try: + partitions, token_counts = await asyncio.to_thread( + _plan_partitions, engine, settings.dsn, model, texts + ) + except Exception as exc: + logger.warning( + "Batch embedding planning unavailable; falling back to per-item path: " + "error_type=%s text_count=%s", + type(exc).__name__, + len(texts), + ) + return None + + job = _new_batch_job( + organization_id=organization_id, + user_id=user_id, + model=model, + endpoint_alias=settings.endpoint_alias, + text_count=len(texts), + total_tokens=sum(token_counts), + part_count=len(partitions), + ) + items = _new_batch_items(job.batch_job_uid, partitions, token_counts) + session.add(job) + for item in items: + session.add(item) + + results: list[list[float] | None] = [None] * len(texts) + try: + for index_group in partitions: + part_texts = [texts[i] for i in index_group] + vectors = await generate_embeddings( + part_texts, + embedding_provider.api_key, + base_url=embedding_provider.base_url, + model=model, + ) + for offset, text_index in enumerate(index_group): + if offset < len(vectors): + results[text_index] = fit_embedding_vector( + vectors[offset], dimension + ) + except (EmbeddingGenerationError, ValueError, TypeError) as exc: + logger.warning( + "Batch embedding generation failed; falling back to per-item path: " + "error_type=%s text_count=%s part_count=%s", + type(exc).__name__, + len(texts), + len(partitions), + ) + _mark_job_failed(job, items, error_code=type(exc).__name__) + return None + + if any(vector is None for vector in results): + # A partition returned fewer vectors than inputs — treat as incomplete + # and fall back rather than persisting zero vectors silently. + logger.warning( + "Batch embedding returned incomplete vectors; falling back: " + "text_count=%s", + len(texts), + ) + _mark_job_failed(job, items, error_code="incomplete_vectors") + return None + + _mark_job_completed(job, items) + return [vector for vector in results if vector is not None] + + +def _new_batch_job( + *, + organization_id: str | None, + user_id: str, + model: str, + endpoint_alias: str | None, + text_count: int, + total_tokens: int, + part_count: int, +) -> LlmBatchJob: + return LlmBatchJob( + batch_job_uid=f"llm_batch_{uuid.uuid4().hex}", + organization_id=organization_id or "", + user_id=user_id, + job_status="preparing", + model_name=model, + endpoint_alias=endpoint_alias, + total_items=text_count, + completed_items=0, + failed_items=0, + total_tokens=total_tokens, + part_count=part_count, + ) + + +def _new_batch_items( + batch_job_uid: str, + partitions: list[list[int]], + token_counts: list[int], +) -> list[LlmBatchItem]: + items: list[LlmBatchItem] = [] + for part_index, index_group in enumerate(partitions): + for text_index in index_group: + items.append( + LlmBatchItem( + batch_item_uid=f"llm_batch_item_{uuid.uuid4().hex}", + batch_job_uid=batch_job_uid, + sequence_no=text_index, + part_index=part_index, + token_count=( + token_counts[text_index] + if text_index < len(token_counts) + else 0 + ), + item_status="queued", + ) + ) + return items + + +def _mark_job_completed(job: LlmBatchJob, items: list[LlmBatchItem]) -> None: + job.job_status = "completed" + job.completed_items = job.total_items + for item in items: + item.item_status = "completed" + + +def _mark_job_failed( + job: LlmBatchJob, items: list[LlmBatchItem], *, error_code: str +) -> None: + job.job_status = "failed" + job.failed_items = job.total_items + job.error_code = error_code + for item in items: + item.item_status = "failed" diff --git a/backend/services/email_import_service.py b/backend/services/email_import_service.py index c699d6182..6a604f55c 100644 --- a/backend/services/email_import_service.py +++ b/backend/services/email_import_service.py @@ -22,6 +22,7 @@ KnowledgeGraphEdgeRecord, ) from services.archive import extract_backup_async +from services.batch_embedding_service import try_batch_import_embeddings from services.content_graph import ParseResult, parse_content from services.email_dedupe_service import strong_email_fingerprint from services.email_parser import EmailData, parse_eml_bytes @@ -65,6 +66,19 @@ class EmailImportEmbeddingProvider: embedding_model: str +@dataclass(frozen=True) +class EmailImportBatchContext: + """Scope needed to route bulk import embeddings through pg-llm-batch. + + Carried alongside the embedding provider so ``_generate_import_embeddings`` + can resolve per-tenant batch settings (Fernet DB) and record batch jobs. + """ + + session: AsyncSession + user_id: str + organization_id: str | None + + @dataclass class EmailImportItemResult: filename: str @@ -224,6 +238,7 @@ async def _release_owner_import_quota_lock( async def _extract_and_generate_embeddings( parsed: EmailData, embedding_provider: EmailImportEmbeddingProvider | None, + batch_context: EmailImportBatchContext | None = None, ) -> tuple[list[dict], list[list[float]]]: attachment_payloads = list(parsed.get("attachments", [])) embedding_texts = [str(parsed.get("body") or "")] @@ -233,6 +248,7 @@ async def _extract_and_generate_embeddings( fitted_embeddings = await _generate_import_embeddings( embedding_texts, embedding_provider=embedding_provider, + batch_context=batch_context, ) return attachment_payloads, fitted_embeddings @@ -633,6 +649,7 @@ async def _import_single_eml( user_id: str, organization_id: str, embedding_provider: EmailImportEmbeddingProvider | None = None, + batch_context: EmailImportBatchContext | None = None, ) -> EmailImportItemResult: try: content, parsed = await asyncio.to_thread(_read_and_parse_eml, eml_path) @@ -670,7 +687,7 @@ async def _import_single_eml( ) attachment_payloads, fitted_embeddings = await _extract_and_generate_embeddings( - parsed, embedding_provider + parsed, embedding_provider, batch_context ) email_obj, attachment_count = _build_email_object( @@ -716,9 +733,24 @@ async def _generate_import_embeddings( texts: list[str], *, embedding_provider: EmailImportEmbeddingProvider | None, + batch_context: "EmailImportBatchContext | None" = None, ) -> list[list[float]]: if embedding_provider is None: return [_zero_embedding() for _ in texts] + if batch_context is not None and texts: + # Bulk import embeddings are latency-tolerant: try the pg-llm-batch + # engine first. A None result means batch is unconfigured/unavailable, so + # we transparently fall through to the existing per-request path below. + batched = await try_batch_import_embeddings( + batch_context.session, + texts, + embedding_provider=embedding_provider, + user_id=batch_context.user_id, + organization_id=batch_context.organization_id, + dimension=EMBEDDING_DIMENSION, + ) + if batched is not None: + return batched try: provider_embeddings = await generate_embeddings( texts, @@ -887,6 +919,11 @@ async def import_email_uploads( lock_acquired = await _acquire_owner_import_quota_lock( session, user_id=user_id, organization_id=organization_id ) + batch_context = EmailImportBatchContext( + session=session, + user_id=user_id, + organization_id=organization_id, + ) try: result = EmailImportResult() existing_email_count = await _owner_email_import_count( @@ -937,6 +974,7 @@ async def import_email_uploads( user_id=user_id, organization_id=organization_id, embedding_provider=embedding_provider, + batch_context=batch_context, ) ) finally: diff --git a/backend/tests/test_alembic_migrations.py b/backend/tests/test_alembic_migrations.py index 2503bdddf..cc3d89c25 100644 --- a/backend/tests/test_alembic_migrations.py +++ b/backend/tests/test_alembic_migrations.py @@ -217,6 +217,33 @@ def test_project_graph_projection_has_incremental_revision(): assert "if_exists=True" in revision_text +def test_llm_batch_embedding_has_incremental_revision(): + versions_dir = BACKEND_ROOT / "alembic" / "versions" + revision_path = versions_dir / "0010_llm_batch_embedding.py" + assert revision_path.exists() + revision_text = revision_path.read_text() + + assert 'revision = "0010_llm_batch_embedding"' in revision_text + assert 'down_revision = "0009_project_graph_projection"' in revision_text + assert '"llm_batch_jobs"' in revision_text + assert '"llm_batch_items"' in revision_text + assert '"batch_job_uid"' in revision_text + assert '"batch_item_uid"' in revision_text + assert '"batch_embedding_dsn"' in revision_text + assert '"batch_embedding_enabled"' in revision_text + assert "ix_llm_batch_jobs_scope_status" in revision_text + assert "ix_llm_batch_items_job_sequence" in revision_text + assert "ForeignKeyConstraint" in revision_text + assert "has_table" in revision_text + assert "has_column" in revision_text + assert "op.create_table(" in revision_text + assert "op.add_column(" in revision_text + assert "op.create_index(" in revision_text + assert "if_not_exists=True" in revision_text + assert "op.drop_index(" in revision_text + assert "if_exists=True" in revision_text + + def test_migration_runner_uses_alembic_upgrade_head_not_bootstrap_create_all(): migration_runner = BACKEND_ROOT / "scripts" / "migrate_db.py" diff --git a/backend/tests/test_batch_embedding_service.py b/backend/tests/test_batch_embedding_service.py new file mode 100644 index 000000000..af5554a79 --- /dev/null +++ b/backend/tests/test_batch_embedding_service.py @@ -0,0 +1,328 @@ +"""Tests for batch-tolerant embedding routing via the pg-llm-batch submodule. + +These are fast, fully-mocked unit tests: no live Postgres, pg_tiktoken, or the +submodule itself is required. They verify three properties the integration +promises: + +* bulk import embeddings route through the batch engine when a tenant has + enabled + configured batching; +* the path degrades gracefully (returns ``None`` so callers fall back to the + per-item path) when batching is disabled or the submodule is not importable; +* batch config (enablement + DSN) is read from the per-tenant Fernet-encrypted + ``tenant_configs`` row, never from ``os.getenv``. +""" + +import os +import types +from unittest.mock import AsyncMock + +import pytest +from cryptography.fernet import Fernet +from pydantic import SecretStr +from sqlalchemy import create_engine, text +from sqlalchemy.orm import Session + +from core.config import settings +from db.models import LlmBatchItem, LlmBatchJob, TenantConfig +from services.email_import_service import ( + EmailImportBatchContext, + EmailImportEmbeddingProvider, + _generate_import_embeddings, +) +import services.batch_embedding_service as batch_module +from services.batch_embedding_service import ( + resolve_batch_embedding_settings, + try_batch_import_embeddings, +) + + +PROVIDER = EmailImportEmbeddingProvider( + api_key="secret-provider-token", + base_url="http://gateway.internal/v1", + embedding_model="text-embedding-test", +) + + +@pytest.fixture(autouse=True) +def encryption_key(): + old_key = settings.ENCRYPTION_KEY + settings.ENCRYPTION_KEY = SecretStr(Fernet.generate_key().decode("ascii")) + yield + settings.ENCRYPTION_KEY = old_key + + +class FakeResult: + def __init__(self, tenant_config): + self._tenant_config = tenant_config + + def scalar_one_or_none(self): + return self._tenant_config + + +class FakeAsyncSession: + """Minimal async session: returns a tenant config and records .add().""" + + def __init__(self, tenant_config=None): + self.tenant_config = tenant_config + self.added: list = [] + + async def execute(self, _stmt): + return FakeResult(self.tenant_config) + + def add(self, obj): + self.added.append(obj) + + +# --- Fake pg-llm-batch engine (stands in for the submodule) ----------------- + + +class _FakeConfigStore: + def __init__(self, dsn): + self.dsn = dsn + self.closed = False + + def close(self): + self.closed = True + + +class _FakeTokenCounter: + def __init__(self, dsn, config=None): + self.dsn = dsn + self.config = config + + +class _FakeAccumulator: + """Partitions into groups of two to exercise multi-part planning.""" + + def __init__(self, counter, model): + self.model = model + self.reset() + + def reset(self): + self.record_count = 0 + + def compute_tokens(self, _system, user): + tokens = len(user) + return tokens, 0, tokens + + @staticmethod + def compute_byte_size(line): + return len(line.encode("utf-8")) + 1 + + def would_exceed(self, _tokens, _byte_size): + return self.record_count >= 2 + + def add_entry(self, _rid, _line, _tokens, _byte_size): + self.record_count += 1 + + +def _fake_engine(): + return types.SimpleNamespace( + PostgresConfigStore=_FakeConfigStore, + TokenCounter=_FakeTokenCounter, + BatchAccumulator=_FakeAccumulator, + ) + + +def _batch_tenant_config(**overrides): + config = TenantConfig(user_id="user-1", organization_id="org-acme") + config.batch_embedding_enabled = True + config.batch_embedding_dsn = "postgresql://batch-host/batch_db" + config.batch_embedding_endpoint = "primary_gateway" + config.batch_embedding_model = "text-embedding-test" + for key, value in overrides.items(): + setattr(config, key, value) + return config + + +# --- Routing ---------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_import_embeddings_route_through_batch_component(monkeypatch): + session = FakeAsyncSession(_batch_tenant_config()) + monkeypatch.setattr(batch_module, "load_batch_engine", _fake_engine) + generate = AsyncMock(side_effect=lambda texts, *a, **k: [[0.5] * 8 for _ in texts]) + monkeypatch.setattr(batch_module, "generate_embeddings", generate) + + texts = ["body", "att-1", "att-2", "att-3", "att-4"] + result = await try_batch_import_embeddings( + session, + texts, + embedding_provider=PROVIDER, + user_id="user-1", + organization_id="org-acme", + dimension=8, + ) + + assert result is not None + assert len(result) == 5 + assert all(len(vector) == 8 for vector in result) + # 5 texts partitioned two-at-a-time -> 3 batch parts -> 3 embedding calls. + assert generate.await_count == 3 + # Credentials came from the runtime provider, not os.getenv. + assert generate.await_args_list[0].args[1] == "secret-provider-token" + + jobs = [obj for obj in session.added if isinstance(obj, LlmBatchJob)] + items = [obj for obj in session.added if isinstance(obj, LlmBatchItem)] + assert len(jobs) == 1 + assert jobs[0].job_status == "completed" + assert jobs[0].total_items == 5 + assert jobs[0].part_count == 3 + assert len(items) == 5 + assert all(item.item_status == "completed" for item in items) + assert {item.part_index for item in items} == {0, 1, 2} + + +@pytest.mark.asyncio +async def test_import_embeddings_fall_back_when_batch_disabled(monkeypatch): + session = FakeAsyncSession(_batch_tenant_config(batch_embedding_enabled=False)) + monkeypatch.setattr(batch_module, "load_batch_engine", _fake_engine) + generate = AsyncMock() + monkeypatch.setattr(batch_module, "generate_embeddings", generate) + + result = await try_batch_import_embeddings( + session, + ["body"], + embedding_provider=PROVIDER, + user_id="user-1", + organization_id="org-acme", + ) + + assert result is None + generate.assert_not_awaited() + assert session.added == [] + + +@pytest.mark.asyncio +async def test_import_embeddings_fall_back_when_submodule_missing(monkeypatch): + session = FakeAsyncSession(_batch_tenant_config()) + monkeypatch.setattr(batch_module, "load_batch_engine", lambda: None) + generate = AsyncMock() + monkeypatch.setattr(batch_module, "generate_embeddings", generate) + + result = await try_batch_import_embeddings( + session, + ["body"], + embedding_provider=PROVIDER, + user_id="user-1", + organization_id="org-acme", + ) + + assert result is None + generate.assert_not_awaited() + assert session.added == [] + + +@pytest.mark.asyncio +async def test_load_batch_engine_is_import_guarded(): + # The submodule is not installed on the backend path in this environment, + # so the loader must degrade to None rather than raising ImportError. + batch_module._ENGINE_CACHE.clear() + try: + assert batch_module.load_batch_engine() is None + finally: + batch_module._ENGINE_CACHE.clear() + + +# --- email_import_service wiring -------------------------------------------- + + +@pytest.mark.asyncio +async def test_generate_import_embeddings_prefers_batch_context(monkeypatch): + context = EmailImportBatchContext( + session=FakeAsyncSession(), user_id="user-1", organization_id="org-acme" + ) + batched = [[0.1] * 1536, [0.2] * 1536] + routed = AsyncMock(return_value=batched) + monkeypatch.setattr( + "services.email_import_service.try_batch_import_embeddings", routed + ) + per_item = AsyncMock() + monkeypatch.setattr("services.email_import_service.generate_embeddings", per_item) + + result = await _generate_import_embeddings( + ["body", "attachment"], + embedding_provider=PROVIDER, + batch_context=context, + ) + + assert result == batched + routed.assert_awaited_once() + # Batch path handled it; the per-item embedding path was never touched. + per_item.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_generate_import_embeddings_falls_back_when_batch_returns_none( + monkeypatch, +): + context = EmailImportBatchContext( + session=FakeAsyncSession(), user_id="user-1", organization_id="org-acme" + ) + routed = AsyncMock(return_value=None) + monkeypatch.setattr( + "services.email_import_service.try_batch_import_embeddings", routed + ) + per_item = AsyncMock(return_value=[[0.9] * 1536, [0.9] * 1536]) + monkeypatch.setattr("services.email_import_service.generate_embeddings", per_item) + + result = await _generate_import_embeddings( + ["body", "attachment"], + embedding_provider=PROVIDER, + batch_context=context, + ) + + routed.assert_awaited_once() + # Fell through to the existing bulk/per-item path. + per_item.assert_awaited() + assert len(result) == 2 + + +# --- Config from the Fernet DB (never env) ---------------------------------- + + +@pytest.mark.asyncio +async def test_batch_config_resolves_from_fernet_db_not_env(monkeypatch): + monkeypatch.delenv("PG_LLM_BATCH_DSN", raising=False) + secret_dsn = "postgresql://batch-user:batch-pass@batch-host/batch_db" + + engine = create_engine("sqlite:///:memory:") + TenantConfig.__table__.create(engine) + try: + with Session(engine) as session: + session.add( + TenantConfig( + user_id="user-1", + organization_id="org-acme", + batch_embedding_enabled=True, + batch_embedding_dsn=secret_dsn, + batch_embedding_endpoint="primary_gateway", + batch_embedding_model="text-embedding-test", + ) + ) + session.commit() + + # Stored at rest as Fernet ciphertext, not the plaintext DSN. + raw = session.execute( + text("SELECT batch_embedding_dsn FROM tenant_configs") + ).scalar_one() + assert raw != secret_dsn + assert "batch-pass" not in raw + + reloaded = session.query(TenantConfig).one() + assert reloaded.batch_embedding_dsn == secret_dsn + finally: + engine.dispose() + + # The resolver returns the DSN decrypted from the DB row, and no env var + # supplied it. + settings_obj = await resolve_batch_embedding_settings( + FakeAsyncSession(reloaded), + user_id="user-1", + organization_id="org-acme", + ) + assert settings_obj is not None + assert settings_obj.dsn == secret_dsn + assert settings_obj.endpoint_alias == "primary_gateway" + assert "PG_LLM_BATCH_DSN" not in os.environ diff --git a/docker-compose.pg-llm-batch.yml b/docker-compose.pg-llm-batch.yml new file mode 100644 index 000000000..ce21ee47a --- /dev/null +++ b/docker-compose.pg-llm-batch.yml @@ -0,0 +1,57 @@ +# Opt-in overlay for the pg-llm-batch submodule (batch-tolerant LLM embeddings). +# +# This file is NOT part of the default naruon stack. naruon runs fine without it +# and without the submodule initialized — the batch embedding path degrades to +# the per-item embedding path when unconfigured. Bring it up only when you want +# the pg_tiktoken-enabled batch Postgres available to the backend. +# +# Prerequisites: +# git submodule update --init vendor/pg-llm-batch +# +# Usage: +# docker compose -f docker-compose.yml -f docker-compose.pg-llm-batch.yml up +# +# Then seed per-tenant batch config in naruon's Fernet DB (tenant_configs): +# batch_embedding_enabled = true +# batch_embedding_dsn = postgresql://pgllm:pgllm@pg-llm-batch-postgres:5432/pgllm +# batch_embedding_endpoint / batch_embedding_model as appropriate +# (never via os.getenv — the DSN column is Fernet-encrypted at rest). + +services: + pg-llm-batch-postgres: + # Reuses the component's pg_tiktoken + pg_cron + pgsql-http Postgres image. + build: + context: ./vendor/pg-llm-batch + dockerfile: docker/postgres/Dockerfile + args: + PG_MAJOR: "16" + ENABLE_TIKTOKEN: "1" + environment: + POSTGRES_USER: pgllm + POSTGRES_PASSWORD: pgllm + POSTGRES_DB: pgllm + ports: + - "5442:5432" + healthcheck: + test: + - CMD-SHELL + - >- + pg_isready -U pgllm -d pgllm && + psql -U pgllm -d pgllm -tAc + "SELECT bool_and(is_ready) FROM pg_llm_batch_health_check() + WHERE component IN ('database','pg_tiktoken','com_config')" | grep -q t + interval: 15s + timeout: 10s + retries: 10 + start_period: 40s + volumes: + - pg-llm-batch-data:/var/lib/postgresql/data + networks: + - naruon-network + +volumes: + pg-llm-batch-data: + +networks: + naruon-network: + external: true diff --git a/vendor/pg-llm-batch b/vendor/pg-llm-batch new file mode 160000 index 000000000..9f40192c0 --- /dev/null +++ b/vendor/pg-llm-batch @@ -0,0 +1 @@ +Subproject commit 9f40192c07d8f4883ae25c1d76643a22279ebb9c