Repository navigation
π¨ [OIDC λ‘κ·ΈμΈ/λ‘κ·Έμμ λ²νΌ λ‘λ© μν λ° μ κ·Όμ± κ°μ ] #1410
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. Weβll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Changes from all commits
f441f68
41ed8d7
f776e35
a7cc029
ab722e4
ed592ae
fac7a53
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -32,7 +32,6 @@ | |
| from __future__ import annotations | ||
|
|
||
| import asyncio | ||
| import json | ||
| import logging | ||
| import uuid | ||
| from dataclasses import dataclass | ||
|
|
@@ -62,11 +61,6 @@ | |
| # Poll budget while the orchestrator drains the batch through pg-llm-batch. | ||
| _ORCHESTRATOR_POLL_INTERVAL_SECONDS = 1.0 | ||
| _ORCHESTRATOR_MAX_POLLS = 30 | ||
| # Keep each JSON request below the orchestrator's body budget with envelope | ||
| # headroom. Import chunks are normally <= 1,000 characters, but the byte check | ||
| # also protects direct callers that provide longer or multibyte inputs. | ||
| _ORCHESTRATOR_MAX_INPUTS_PER_REQUEST = 32 | ||
| _ORCHESTRATOR_MAX_INPUT_BYTES = 48 * 1024 | ||
| _SUCCESS_STATUSES = frozenset({"completed", "succeeded"}) | ||
|
|
||
| # Cache the (im)port result so repeated imports don't re-probe sys.path. | ||
|
|
@@ -120,14 +114,6 @@ def has_local_fallback(self) -> bool: | |
| return bool(self.local_dsn) | ||
|
|
||
|
|
||
| @dataclass(frozen=True) | ||
| class BatchEmbeddingPartial: | ||
| """Completed prefix plus inputs that still need the normal fallback path.""" | ||
|
|
||
| completed_vectors: list[list[float]] | ||
| pending_texts: list[str] | ||
|
|
||
|
|
||
| async def resolve_batch_embedding_settings( | ||
| session: AsyncSession, | ||
| *, | ||
|
|
@@ -159,7 +145,9 @@ async def resolve_batch_embedding_settings( | |
| attribution_service=_clean( | ||
| getattr(tenant_config, "batch_attribution_service", None) | ||
| ), | ||
| attribution_team=_clean(getattr(tenant_config, "batch_attribution_team", None)), | ||
| attribution_team=_clean( | ||
| getattr(tenant_config, "batch_attribution_team", None) | ||
| ), | ||
| attribution_group=_clean( | ||
| getattr(tenant_config, "batch_attribution_group", None) | ||
| ), | ||
|
|
@@ -188,16 +176,15 @@ async def try_batch_import_embeddings( | |
| user_id: str, | ||
| organization_id: str | None, | ||
| dimension: int = STORAGE_EMBEDDING_DIMENSION, | ||
| ) -> list[list[float]] | BatchEmbeddingPartial | None: | ||
| ) -> list[list[float]] | None: | ||
| """Route bulk embeddings through the batch path, or ``None`` to fall back. | ||
|
|
||
| On success returns one fitted vector per input text (original order). The | ||
| primary path submits to contextual-orchestrator; only if the orchestrator is | ||
| unconfigured or unavailable does it consider the local ``pg-llm-batch`` | ||
| package fallback. A later partition failure returns a completed prefix plus | ||
| only the unfinished inputs so the caller does not resend successful work. | ||
| The run is recorded in ``llm_batch_jobs`` / ``llm_batch_items`` for | ||
| observability. | ||
| package fallback. Any failure returns ``None`` so the caller uses its | ||
| per-item path. The run is recorded in ``llm_batch_jobs`` / ``llm_batch_items`` | ||
| for observability. | ||
| """ | ||
| if not texts: | ||
| return None | ||
|
|
@@ -211,7 +198,7 @@ async def try_batch_import_embeddings( | |
| model = settings.model or embedding_provider.embedding_model | ||
|
|
||
| if settings.has_orchestrator: | ||
| result = await _run_orchestrator_batches( | ||
| result = await _run_orchestrator_batch( | ||
| session, | ||
| texts, | ||
| settings=settings, | ||
|
Comment on lines
200
to
204
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. π Info: Orchestrator batch no longer bounds request size
(Refers to this code) Was this helpful? React with π or π to provide feedback. |
||
|
|
@@ -242,121 +229,6 @@ async def try_batch_import_embeddings( | |
| # --- Primary path: contextual-orchestrator batch API ------------------------ | ||
|
|
||
|
|
||
| def _serialized_orchestrator_payload_bytes( | ||
| inputs: list[str], | ||
| *, | ||
| model: str, | ||
| endpoint_alias: str | None, | ||
| metadata: dict[str, str], | ||
| ) -> int: | ||
| """Return the UTF-8 size of the request envelope sent to the orchestrator.""" | ||
| payload = { | ||
| "model": model, | ||
| "endpoint": endpoint_alias, | ||
| "inputs": inputs, | ||
| "metadata": metadata, | ||
| } | ||
| return len( | ||
| json.dumps(payload, ensure_ascii=True, separators=(",", ":")).encode("utf-8") | ||
| ) | ||
|
|
||
|
|
||
| def _partition_orchestrator_inputs( | ||
| texts: list[str], | ||
| *, | ||
| model: str = "", | ||
| endpoint_alias: str | None = None, | ||
| metadata: dict[str, str] | None = None, | ||
| ) -> list[list[str]] | None: | ||
| """Partition inputs by count and serialized JSON request bytes.""" | ||
| request_metadata = metadata or {} | ||
| partitions: list[list[str]] = [] | ||
| current: list[str] = [] | ||
| for text in texts: | ||
| candidate = [*current, text] | ||
| if current and ( | ||
| len(candidate) > _ORCHESTRATOR_MAX_INPUTS_PER_REQUEST | ||
| or _serialized_orchestrator_payload_bytes( | ||
| candidate, | ||
| model=model, | ||
| endpoint_alias=endpoint_alias, | ||
| metadata=request_metadata, | ||
| ) | ||
| > _ORCHESTRATOR_MAX_INPUT_BYTES | ||
| ): | ||
| partitions.append(current) | ||
| candidate = [text] | ||
| if ( | ||
| _serialized_orchestrator_payload_bytes( | ||
| candidate, | ||
| model=model, | ||
| endpoint_alias=endpoint_alias, | ||
| metadata=request_metadata, | ||
| ) | ||
| > _ORCHESTRATOR_MAX_INPUT_BYTES | ||
| ): | ||
| return None | ||
| current = candidate | ||
| if current: | ||
| partitions.append(current) | ||
| return partitions | ||
|
|
||
|
|
||
| async def _run_orchestrator_batches( | ||
| session: AsyncSession, | ||
| texts: list[str], | ||
| *, | ||
| settings: BatchEmbeddingSettings, | ||
| model: str, | ||
| user_id: str, | ||
| organization_id: str | None, | ||
| dimension: int, | ||
| ) -> list[list[float]] | BatchEmbeddingPartial | None: | ||
| """Submit bounded requests and concatenate vectors in original order.""" | ||
| metadata = _attribution_metadata( | ||
| settings=settings, | ||
| user_id=user_id, | ||
| organization_id=organization_id, | ||
| ) | ||
| partitions = _partition_orchestrator_inputs( | ||
| texts, | ||
| model=model, | ||
| endpoint_alias=settings.endpoint_alias, | ||
| metadata=metadata, | ||
| ) | ||
| if partitions is None: | ||
| logger.warning( | ||
| "Orchestrator batch input exceeded one-request byte budget; falling back: " | ||
| "text_count=%s", | ||
| len(texts), | ||
| ) | ||
| return None | ||
|
|
||
| vectors: list[list[float]] = [] | ||
| for partition_index, partition in enumerate(partitions): | ||
| partition_vectors = await _run_orchestrator_batch( | ||
| session, | ||
| partition, | ||
| settings=settings, | ||
| model=model, | ||
| user_id=user_id, | ||
| organization_id=organization_id, | ||
| dimension=dimension, | ||
| ) | ||
| if partition_vectors is None: | ||
| if not vectors: | ||
| return None | ||
| pending_texts = [ | ||
| text for remaining in partitions[partition_index:] for text in remaining | ||
| ] | ||
| return BatchEmbeddingPartial( | ||
| completed_vectors=vectors, | ||
| pending_texts=pending_texts, | ||
| ) | ||
| vectors.extend(partition_vectors) | ||
| return vectors | ||
|
|
||
|
|
||
| async def _run_orchestrator_batch( | ||
| session: AsyncSession, | ||
| texts: list[str], | ||
|
|
@@ -500,7 +372,9 @@ async def _submit_and_await( | |
| if status in _SUCCESS_STATUSES and document.get("embeddings") is not None: | ||
| return document | ||
| if status in ("failed", "error", "canceled"): | ||
| raise EmbeddingGenerationError(f"orchestrator batch rejected: status={status}") | ||
| raise EmbeddingGenerationError( | ||
| f"orchestrator batch rejected: status={status}" | ||
| ) | ||
|
|
||
| batch_id = document.get("batch_id") or document.get("id") | ||
| if not batch_id: | ||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -26,16 +26,12 @@ | |
| KnowledgeGraphEdgeRecord, | ||
| ) | ||
| from services.archive import extract_backup_async | ||
| from services.batch_embedding_service import ( | ||
| BatchEmbeddingPartial, | ||
| try_batch_import_embeddings, | ||
| ) | ||
| 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 | ||
| from services.embedding import ( | ||
| STORAGE_EMBEDDING_DIMENSION, | ||
| chunk_text, | ||
| fit_embedding_vector, | ||
| generate_embeddings, | ||
| ) | ||
|
|
@@ -56,14 +52,10 @@ | |
|
|
||
| EMBEDDING_DIMENSION = STORAGE_EMBEDDING_DIMENSION | ||
| MAX_IMPORT_UPLOADS = 10 | ||
| # Transport safety ceiling only; parser and embedding chunking must accept | ||
| # sources larger than 20 MiB without confusing the request guard for a parser | ||
| # limit. | ||
| MAX_IMPORT_UPLOAD_BYTES = 64 * 1024 * 1024 | ||
| MAX_IMPORT_UPLOAD_BYTES = 20 * 1024 * 1024 | ||
|
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. π Info: Upload size ceiling reverted from 64 MiB to 20 MiB
Was this helpful? React with π or π to provide feedback. |
||
| MAX_IMPORT_EML_FILES = 100 | ||
| MAX_IMPORT_EMAILS_PER_OWNER = 1000 | ||
| MAX_UPLOAD_FILENAME_DECODE_ROUNDS = 8 | ||
| MAX_EMBEDDING_CHUNKS_PER_WINDOW = 32 | ||
| SUPPORTED_EMAIL_IMPORT_SUFFIXES = frozenset({".eml", ".mbox", ".zip"}) | ||
| EMAIL_IMPORT_QUOTA_LOCK_NAMESPACE = "naruon-email-import-quota" | ||
| logger = logging.getLogger(__name__) | ||
|
|
@@ -296,52 +288,15 @@ async def _extract_and_generate_embeddings( | |
| batch_context: "EmailImportBatchContext | None" = None, | ||
| ) -> tuple[list[dict], list[list[float]]]: | ||
| attachment_payloads = list(parsed.get("attachments", [])) | ||
| body_parse_content = parsed.get("body_parse_content") | ||
| source_texts = [ | ||
| str( | ||
| body_parse_content | ||
| if body_parse_content is not None | ||
| else parsed.get("body") or "" | ||
| ) | ||
| ] | ||
| source_texts.extend( | ||
| str( | ||
| "" | ||
| if (attachment.get("parse_status") or "parsed") != "parsed" | ||
| else ( | ||
| attachment.get("parse_content") | ||
| if attachment.get("parse_content") is not None | ||
| else attachment.get("content") or "" | ||
| ) | ||
| ) | ||
| for attachment in attachment_payloads | ||
| embedding_texts = [str(parsed.get("body") or "")] | ||
| embedding_texts.extend( | ||
| str(attachment.get("content") or "") for attachment in attachment_payloads | ||
| ) | ||
| fitted_embeddings = await _generate_import_embeddings( | ||
| embedding_texts, | ||
| embedding_provider=embedding_provider, | ||
| batch_context=batch_context, | ||
| ) | ||
|
Comment on lines
+291
to
299
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. π΄ Oversized email bodies no longer chunked before embedding The rewritten Prompt for agentsWas this helpful? React with π or π to provide feedback. |
||
| fitted_embeddings: list[list[float]] = [] | ||
| for source_text in source_texts: | ||
| source_chunks = chunk_text(source_text) | ||
| if not source_chunks: | ||
| fitted_embeddings.append(_zero_embedding()) | ||
| continue | ||
|
|
||
| vector_sum: list[float] | None = None | ||
| vector_count = 0 | ||
| for start in range(0, len(source_chunks), MAX_EMBEDDING_CHUNKS_PER_WINDOW): | ||
| chunk_embeddings = await _generate_import_embeddings( | ||
| source_chunks[start : start + MAX_EMBEDDING_CHUNKS_PER_WINDOW], | ||
| embedding_provider=embedding_provider, | ||
| batch_context=batch_context, | ||
| ) | ||
| for embedding in chunk_embeddings: | ||
| if vector_sum is None: | ||
| vector_sum = [0.0] * len(embedding) | ||
| for index, value in enumerate(embedding): | ||
| vector_sum[index] += value | ||
| vector_count += 1 | ||
| fitted_embeddings.append( | ||
| [value / vector_count for value in vector_sum] | ||
| if vector_sum and vector_count | ||
| else _zero_embedding() | ||
| ) | ||
| return attachment_payloads, fitted_embeddings | ||
|
|
||
|
|
||
|
|
@@ -439,11 +394,7 @@ def _fallback_attachment_parser_key( | |
| return "calendar" | ||
| if parse_content_type == "text/html": | ||
| return "html" | ||
| if parse_content_type in { | ||
| "text/markdown", | ||
| "text/x-markdown", | ||
| "application/markdown", | ||
| }: | ||
| if parse_content_type in {"text/markdown", "text/x-markdown", "application/markdown"}: | ||
| return "markdown" | ||
| if parse_content_type == "text/plain": | ||
| return "plain_text" | ||
|
|
@@ -635,9 +586,9 @@ def add_edge( | |
| item.segment_path, | ||
| ), | ||
| ): | ||
| segments_by_source[(segment.source_kind, segment.source_record_uid)].append( | ||
| segment | ||
| ) | ||
| segments_by_source[ | ||
| (segment.source_kind, segment.source_record_uid) | ||
| ].append(segment) | ||
| add_edge( | ||
| edge_kind="node_has_segment", | ||
| edge_path=f"{segment.content_node.node_path}/has/{segment.segment_path}", | ||
|
|
@@ -652,7 +603,8 @@ def add_edge( | |
| add_edge( | ||
| edge_kind="segment_next", | ||
| edge_path=( | ||
| f"{source_segment.segment_path}/next/{target_segment.segment_path}" | ||
| f"{source_segment.segment_path}/next/" | ||
| f"{target_segment.segment_path}" | ||
| ), | ||
| source_kind=source_segment.source_kind, | ||
| source_record_uid=source_segment.source_record_uid, | ||
|
|
@@ -676,7 +628,8 @@ def add_edge( | |
| add_edge( | ||
| edge_kind="heading_contains_segment", | ||
| edge_path=( | ||
| f"{heading_segment.segment_path}/contains/{segment.segment_path}" | ||
| f"{heading_segment.segment_path}/contains/" | ||
| f"{segment.segment_path}" | ||
| ), | ||
| source_kind=segment.source_kind, | ||
| source_record_uid=segment.source_record_uid, | ||
|
|
@@ -952,8 +905,6 @@ async def _generate_import_embeddings( | |
| embedding_provider: EmailImportEmbeddingProvider | None, | ||
| batch_context: "EmailImportBatchContext | None" = None, | ||
| ) -> list[list[float]]: | ||
| if not texts: | ||
| return [] | ||
| if embedding_provider is None: | ||
| return [_zero_embedding() for _ in texts] | ||
| if batch_context is not None and texts: | ||
|
|
@@ -970,21 +921,6 @@ async def _generate_import_embeddings( | |
| dimension=EMBEDDING_DIMENSION, | ||
| ) | ||
| if batched is not None: | ||
| if isinstance(batched, BatchEmbeddingPartial): | ||
| remainder: list[list[float]] = [] | ||
| for start in range( | ||
| 0, len(batched.pending_texts), MAX_EMBEDDING_CHUNKS_PER_WINDOW | ||
| ): | ||
| remainder.extend( | ||
| await _generate_import_embeddings( | ||
| batched.pending_texts[ | ||
| start : start + MAX_EMBEDDING_CHUNKS_PER_WINDOW | ||
| ], | ||
| embedding_provider=embedding_provider, | ||
| batch_context=None, | ||
| ) | ||
| ) | ||
| return [*batched.completed_vectors, *remainder] | ||
| return batched | ||
| try: | ||
| provider_embeddings = await generate_embeddings( | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
π OIDC PR also reverts a merged backend embedding feature
The single commit titled 'Add loading state to OIDC buttons' also reverts the merged #1413 'chunk oversized embedding inputs' feature:
batch_embedding_service.pyloses byte/count partitioning andBatchEmbeddingPartial,email_import_service.pyloses chunking/averaging and dropsMAX_IMPORT_UPLOAD_BYTESfrom 64 MiB back to 20 MiB,embedding.pyloses the native-dimensions request, and the covering tests are deleted. The description and CHANGELOG diff mention only OIDC. This looks like an accidental revert bundled into an unrelated UI change; confirm with the maintainer before merging.Was this helpful? React with π or π to provide feedback.