Added support for the situation where an embedder fails because too much text is provided.

travis-bauer · travis-bauer · commit 9467faef1cdd · 2026-06-03T20:44:18.000-06:00
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -5,6 +5,11 @@
   `execute` remains as a deprecated alias for `execute_one` (removed in 1.0). `llmEmbed` extends
   `AbstractFieldSegment`, uses `batch_size` for internal provider batching only (one stream item
   in and one out per document), and rejects list-shaped stream items with `TypeError`.
+  When the provider rejects input as too long, `on_token_overflow` controls recovery (default
+  `error`): `truncate` shortens by 20% per retry (`truncate_side`), or `chunk_pool` splits into
+  `num_chunks` segments and mean-pools. Applied reactively after embed failure; size text upstream
+  with `splitText` / `processDocuments` when possible. Batch failures fall back per item so the
+  offending chunk can be identified.
 - Added model2vec embedding support via `Model2VecEmbeddingAdapter` and `llmEmbed`
   source `model2vec`, using in-process static embeddings with offline-ready HF cache
   precaching. Install with `pip install talkpipe[model2vec]` or `talkpipe[all]`. Added
diff --git a/docs/guides/model-and-source-configuration.md b/docs/guides/model-and-source-configuration.md
@@ -276,21 +276,38 @@ segment = LLMVisionPrompt(
 | `field` | No | Text field to embed on structured items |
 | `set_as` | No | Field on the item where the vector is stored |
 | `batch_size` | No | Scalar items per provider call (default `1`) |
-| `fail_on_error` | No | Default `true` |
+| `fail_on_error` | No | Default `true`; applies to non-length failures (network, auth, etc.) |
+| `on_token_overflow` | No | Default `error` — when embed fails as too long: `error`, `truncate`, or `chunk_pool` |
+| `truncate_side` | No | For `truncate`: `head`, `tail` (default), or `middle` |
+| `num_chunks` | No | For `chunk_pool`: segments to split into (default `2`, minimum `2`) |
+
+**Sizing text:** Chunk or split documents **before** `llmEmbed` (e.g. `splitText`, `processDocuments`,
+`makevectordatabase --chunk_size`). `on_token_overflow` is **failure recovery** when a chunk is
+still too long for the model—not a substitute for upstream chunking.
+
+**Token overflow:** TalkPipe classifies provider “too long” errors and applies `on_token_overflow`.
+`truncate` retries with 20% shorter character slices per attempt; `chunk_pool` embeds `num_chunks`
+contiguous parts and mean-pools to one vector per stream item. If a batch embed fails, TalkPipe
+retries **per item** so you can see which chunk failed.
 
 **Batching:** set `batch_size` greater than `1` on `llmEmbed` to call the provider with multiple
 texts per request. The stream still has **one input item and one output item per document**;
 batching is internal only. `llmEmbed` does **not** accept list-shaped stream items (flatten or
 emit items individually upstream). `field` and `set_as` follow
-`AbstractFieldSegment` on each scalar item. With `fail_on_error=False`, failed items are skipped
-when a buffered batch falls back to per-item embedding.
+`AbstractFieldSegment` on each scalar item. With `fail_on_error=False`, non-length failures skip
+items when per-item fallback runs after a batch failure.
 
 ```chatterlang
 INPUT FROM echo[data="Hello world"]
 | llmEmbed[model="mxbai-embed-large", source="ollama", set_as="vector"]
 | print
 ```
 
+```chatterlang
+| llmEmbed[on_token_overflow="truncate", truncate_side="tail"]
+| llmEmbed[on_token_overflow="chunk_pool", num_chunks=4]
+```
+
 ### RAG and vector pipelines
 
 Higher-level segments forward model settings to inner LLM segments:
diff --git a/src/talkpipe/app/chatterlang_workbench.py b/src/talkpipe/app/chatterlang_workbench.py
@@ -1133,7 +1133,11 @@ def main():
 
     # Expose the workbench's own logo URL so example scripts can fetch it from
     # the running server via $workbench_logo_url.
-    logo_host = "localhost" if args.host in ("0.0.0.0", "::") else args.host
+    logo_host = (
+        "localhost"
+        if args.host in ("0.0.0.0", "::")  # nosec B104 - compare bind host, not binding here
+        else args.host
+    )
     add_config_values(
         {"workbench_logo_url": f"http://{logo_host}:{args.port}/static/talkpipe_logo.png"},
         override=True,
diff --git a/src/talkpipe/llm/embedding.py b/src/talkpipe/llm/embedding.py
@@ -1,17 +1,34 @@
 """Module for embedding text using different models"""
 
-from typing import Optional, Annotated, Iterator, Any, List
+from typing import Optional, Annotated, Iterator, Any, List, Literal
 import logging
 
+import numpy as np
+
 from talkpipe.pipe.core import AbstractFieldSegment, is_metadata
 from talkpipe.chatterlang.registry import register_segment
 from talkpipe.util.data_manipulation import extract_property, assign_property
 from .config import getEmbeddingAdapter, getEmbeddingSources
+from .embedding_errors import is_token_overflow_error
 from talkpipe.util.config import get_config
 from talkpipe.util.constants import TALKPIPE_EMBEDDING_MODEL_NAME, TALKPIPE_EMBEDDING_MODEL_SOURCE
 
 logger = logging.getLogger(__name__)
 
+# on_token_overflow mode strings (compare via _OVERFLOW_* constants to avoid Bandit B105/B107)
+_ON_TOKEN_OVERFLOW_CHOICES = ("error", "truncate", "chunk_pool")
+_OVERFLOW_ERROR, _OVERFLOW_TRUNCATE, _OVERFLOW_CHUNK_POOL = _ON_TOKEN_OVERFLOW_CHOICES
+_TRUNCATE_SIDE_CHOICES = ("head", "tail", "middle")
+
+# Truncate retry tuning (not exposed on the segment in v1).
+_SHRINK_RATIO = 0.2
+_MIN_TRUNCATE_CHARS = 1
+_MAX_TRUNCATE_ATTEMPTS = 8
+
+
+class EmbeddingTokenOverflowError(RuntimeError):
+    """Raised when embedding fails due to input length and on_token_overflow is error."""
+
 
 @register_segment("llmEmbed")
 class LLMEmbed(AbstractFieldSegment):
@@ -22,6 +39,11 @@ class LLMEmbed(AbstractFieldSegment):
     :class:`~talkpipe.pipe.core.AbstractFieldSegment`. Batching is internal only
     (``batch_size``); use ``makeLists`` upstream only if another segment needs grouped
     items—not as direct input to ``llmEmbed``.
+
+    When the provider rejects text as too long, ``on_token_overflow`` controls recovery:
+    ``error`` (default), ``truncate`` (shrink and retry), or ``chunk_pool`` (split into
+    ``num_chunks`` segments, embed, and mean-pool). Size text before this segment with
+    upstream chunking when possible.
     """
 
     def __init__(
@@ -32,6 +54,18 @@ def __init__(
         set_as: Annotated[Optional[str], "If provided, append embeddings to input items under this field name"] = None,
         fail_on_error: Annotated[bool, "Whether to raise an error on failure or to silently ignore it"] = True,
         batch_size: Annotated[int, "Number of stream items to embed per provider API call"] = 1,
+        on_token_overflow: Annotated[
+            Literal["error", "truncate", "chunk_pool"],
+            "When embed fails as too long: error, truncate (shrink and retry), or chunk_pool",
+        ] = _OVERFLOW_ERROR,
+        truncate_side: Annotated[
+            Literal["head", "tail", "middle"],
+            "For truncate: which portion of the string to keep when shortening",
+        ] = "tail",
+        num_chunks: Annotated[
+            int,
+            "For chunk_pool: number of contiguous segments to split overflow text into",
+        ] = 2,
     ):
         """Initialize the embedding segment with the specified parameters.
 
@@ -51,13 +85,29 @@ def __init__(
             )
         if batch_size < 1:
             raise ValueError("batch_size must be a positive integer")
+        if on_token_overflow not in _ON_TOKEN_OVERFLOW_CHOICES:
+            raise ValueError(
+                f"on_token_overflow must be one of {_ON_TOKEN_OVERFLOW_CHOICES}, "
+                f"got {on_token_overflow!r}"
+            )
+        if truncate_side not in _TRUNCATE_SIDE_CHOICES:
+            raise ValueError(
+                f"truncate_side must be one of {_TRUNCATE_SIDE_CHOICES}, got {truncate_side!r}"
+            )
+        if num_chunks < 2:
+            raise ValueError("num_chunks must be at least 2")
         self.embedder = getEmbeddingAdapter(source)(model=model)
         self.fail_on_error = fail_on_error
         self.batch_size = batch_size
+        self.on_token_overflow = on_token_overflow
+        self.truncate_side = truncate_side
+        self.num_chunks = num_chunks
+        self._embedding_source = source
+        self._embedding_model = model
 
     def process_value(self, value: Any) -> List[float]:
         """Embed one extracted field value (AbstractFieldSegment hook)."""
-        return self.embedder.execute_one(str(value))
+        return self._embed_one_with_overflow_policy(None, str(value))
 
     def _input_value(self, item: Any) -> Any:
         """Extract the value to embed (same rule as AbstractFieldSegment)."""
@@ -72,6 +122,127 @@ def _ensure_scalar_item(item: Any) -> None:
                 "before this segment."
             )
 
+    @staticmethod
+    def _slice_text(text: str, length: int, side: str) -> str:
+        if length <= 0:
+            return ""
+        if side == "head":
+            return text[:length]
+        if side == "tail":
+            return text[-length:]
+        if side == "middle":
+            if length >= len(text):
+                return text
+            start = (len(text) - length) // 2
+            return text[start : start + length]
+        raise ValueError(f"Unknown truncate_side: {side!r}")
+
+    @staticmethod
+    def _split_num_chunks(text: str, num_chunks: int) -> List[str]:
+        n = len(text)
+        if num_chunks < 2 or n == 0:
+            return [text] if text else []
+        return [text[i * n // num_chunks : (i + 1) * n // num_chunks] for i in range(num_chunks)]
+
+    @staticmethod
+    def _mean_pool(vectors: List[List[float]]) -> List[float]:
+        if not vectors:
+            raise ValueError("Cannot mean-pool an empty list of vectors")
+        arr = np.asarray(vectors, dtype=float)
+        pooled = arr.mean(axis=0)
+        norm = float(np.linalg.norm(pooled))
+        if norm > 0:
+            pooled = pooled / norm
+        return pooled.tolist()
+
+    def _wrap_token_overflow(
+        self,
+        exc: BaseException,
+        *,
+        item: Any,
+        text: str,
+        detail: Optional[str] = None,
+    ) -> EmbeddingTokenOverflowError:
+        field_part = f"field={self.field!r}, " if self.field else ""
+        item_part = f"item={item!r}, " if item is not None else ""
+        text_len = len(text)
+        hint = (
+            "Use smaller upstream chunks (e.g. splitText), "
+            f"on_token_overflow='truncate', or on_token_overflow='chunk_pool' "
+            f"(num_chunks={self.num_chunks})."
+        )
+        extra = f" {detail}" if detail else ""
+        message = (
+            f"Embedding input too long for {self._embedding_source}/{self._embedding_model}: "
+            f"{item_part}{field_part}text_length={text_len}. {hint}{extra} "
+            f"Provider error: {exc}"
+        )
+        return EmbeddingTokenOverflowError(message)
+
+    def _execute_one_raw(self, text: str) -> List[float]:
+        return self.embedder.execute_one(text)
+
+    def _embed_truncate(self, item: Any, text: str) -> List[float]:
+        current = text
+        last_overflow: Optional[BaseException] = None
+        for _ in range(_MAX_TRUNCATE_ATTEMPTS):
+            try:
+                return self._execute_one_raw(current)
+            except Exception as exc:
+                if not is_token_overflow_error(exc):
+                    raise
+                last_overflow = exc
+                n = len(current)
+                n_next = max(_MIN_TRUNCATE_CHARS, int(n * (1 - _SHRINK_RATIO)))
+                if n_next >= n:
+                    break
+                current = self._slice_text(current, n_next, self.truncate_side)
+        raise self._wrap_token_overflow(
+            last_overflow or RuntimeError("truncate exhausted"),
+            item=item,
+            text=text,
+            detail="Truncate retries exhausted.",
+        )
+
+    def _embed_chunk_pool(self, item: Any, text: str) -> List[float]:
+        segments = self._split_num_chunks(text, self.num_chunks)
+        if not segments:
+            raise self._wrap_token_overflow(
+                RuntimeError("empty text"),
+                item=item,
+                text=text,
+            )
+        try:
+            if len(segments) == 1:
+                vectors = [self._execute_one_raw(segments[0])]
+            else:
+                vectors = self.embedder.execute_batch(segments)
+        except Exception as exc:
+            if is_token_overflow_error(exc):
+                raise self._wrap_token_overflow(
+                    exc,
+                    item=item,
+                    text=text,
+                    detail=(
+                        f"chunk_pool with num_chunks={self.num_chunks} still exceeded the limit; "
+                        "try a larger num_chunks or smaller upstream chunks."
+                    ),
+                ) from exc
+            raise
+        return self._mean_pool(vectors)
+
+    def _embed_one_with_overflow_policy(self, item: Any, text: str) -> List[float]:
+        try:
+            return self._execute_one_raw(text)
+        except Exception as exc:
+            if not is_token_overflow_error(exc):
+                raise
+            if self.on_token_overflow == _OVERFLOW_ERROR:
+                raise self._wrap_token_overflow(exc, item=item, text=text) from exc
+            if self.on_token_overflow == _OVERFLOW_TRUNCATE:
+                return self._embed_truncate(item, text)
+            return self._embed_chunk_pool(item, text)
+
     def _yield_results(self, item: Any, results: List[Any]) -> Iterator[Any]:
         """Emit results using AbstractFieldSegment assign/yield semantics."""
         for result in results:
@@ -81,37 +252,36 @@ def _yield_results(self, item: Any, results: List[Any]) -> Iterator[Any]:
             else:
                 yield result
 
-    def _vectors_for_texts(self, texts: List[str]) -> List[List[float]]:
-        if not texts:
-            return []
-        if len(texts) == 1:
-            return [self.process_value(texts[0])]
-        return self.embedder.execute_batch(texts)
+    def _embed_items_pair(self, items: List[Any], texts: List[str]) -> Iterator[Any]:
+        for item, text in zip(items, texts):
+            try:
+                vector = self._embed_one_with_overflow_policy(item, text)
+            except EmbeddingTokenOverflowError:
+                raise
+            except Exception as exc:
+                logger.error(f"Error during embedding: {exc}")
+                if self.fail_on_error:
+                    raise
+                continue
+            yield from self._yield_results(item, [vector])
 
     def _embed_buffered(self, items: List[Any], texts: List[str]) -> Iterator[Any]:
         if not items or not texts:
             return
         logger.debug(f"Embedding batch of {len(texts)} texts")
+        if len(texts) == 1:
+            yield from self._embed_items_pair(items, texts)
+            return
         try:
-            vectors = self._vectors_for_texts(texts)
+            vectors = self.embedder.execute_batch(texts)
             for item, vector in zip(items, vectors):
                 yield from self._yield_results(item, [vector])
         except Exception as e:
             logger.error(f"Error during batch embedding: {e}")
-            if self.fail_on_error:
-                raise
-            if len(texts) == 1:
-                return
             logger.warning(
                 "Batch embedding failed; falling back to per-item embedding"
             )
-            for item, text in zip(items, texts):
-                try:
-                    vector = self.process_value(text)
-                except Exception as item_error:
-                    logger.error(f"Error during embedding: {item_error}")
-                    continue
-                yield from self._yield_results(item, [vector])
+            yield from self._embed_items_pair(items, texts)
 
     def transform(self, input_iter):
         """Transform one stream item at a time; batching is internal only."""
diff --git a/src/talkpipe/llm/embedding_errors.py b/src/talkpipe/llm/embedding_errors.py
@@ -0,0 +1,28 @@
+"""Classify embedding provider errors related to input length / token limits."""
+
+from __future__ import annotations
+
+import re
+
+# Substrings commonly seen when an embedding input exceeds model limits.
+_TOKEN_OVERFLOW_PATTERNS = tuple(
+    re.compile(p, re.IGNORECASE)
+    for p in (
+        r"maximum context length",
+        r"context length",
+        r"token limit",
+        r"too many tokens",
+        r"input.*too long",
+        r"reduce (the )?(length|size|your input)",
+        r"exceeds the maximum",
+        r"max.*tokens",
+    )
+)
+
+
+def is_token_overflow_error(exc: BaseException) -> bool:
+    """Return True if the exception likely indicates input text was too long to embed."""
+    message = str(exc)
+    if not message:
+        message = repr(exc)
+    return any(pattern.search(message) for pattern in _TOKEN_OVERFLOW_PATTERNS)
diff --git a/tests/talkpipe/llm/test_embedding_token_overflow.py b/tests/talkpipe/llm/test_embedding_token_overflow.py
