fix: enable embedding truncation (#1188)

* fix: enable embedding auto truncation

* move embedding_utils to utils.embedding

* fix bug

* fix(embedding): improve text truncation with retry strategies and binary search optimization

* new way for embedding truncate

* fix ci errors

rdagent/oai/backend/base.py

@@ -22,6 +22,7 @@ from rdagent.log import LogColors
 from rdagent.log import rdagent_logger as logger
 from rdagent.log.timer import RD_Agent_TIMER_wrapper
 from rdagent.oai.llm_conf import LLM_SETTINGS
+from rdagent.oai.utils.embedding import truncate_content_list
 from rdagent.utils import md5_hash
 
 try:
@@ -466,6 +467,7 @@ class APIBackend(ABC):
         max_retry = LLM_SETTINGS.max_retry if LLM_SETTINGS.max_retry is not None else max_retry
         timeout_count = 0
         violation_count = 0
+        embedding_truncated = False  # Track if we've already tried truncation
         for i in range(max_retry):
             API_start_time = datetime.now()
             try:
@@ -479,10 +481,25 @@ class APIBackend(ABC):
                     or "\\'messages\\' must contain the word \\'json\\' in some form" in e.message
                 ):
                     kwargs["add_json_in_prompt"] = True
-                elif hasattr(e, "message") and embedding and "maximum context length" in e.message:
-                    kwargs["input_content_list"] = [
-                        content[: len(content) // 2] for content in kwargs.get("input_content_list", [])
-                    ]
+                too_long_error_message = "maximum context length" in e.message or "input must have less than" in e.message  # type: ignore[attr-defined]
+
+                if hasattr(e, "message") and embedding and too_long_error_message:
+                    if not embedding_truncated:
+                        # Handle embedding text too long error - truncate once and retry
+                        model_name = LLM_SETTINGS.embedding_model
+                        logger.warning(f"Embedding text too long for model {model_name}, truncating content")
+
+                        # Apply truncation to content list and continue to retry
+                        original_content_list = kwargs.get("input_content_list", [])
+                        kwargs["input_content_list"] = truncate_content_list(original_content_list, model_name)
+                        embedding_truncated = True  # Mark that we've tried truncation
+                        # Continue to next iteration to retry embedding with truncated content
+                    else:
+                        # Already tried truncation, raise error with guidance
+                        raise RuntimeError(
+                            f"Embedding failed even after truncation. "
+                            f"Please set LLM_SETTINGS.embedding_max_length to a smaller value."
+                        ) from e
                 else:
                     RD_Agent_TIMER_wrapper.api_fail_count += 1
                     RD_Agent_TIMER_wrapper.latest_api_fail_time = datetime.now(pytz.timezone("Asia/Shanghai"))

rdagent/oai/llm_conf.py

@@ -85,6 +85,7 @@ class LLMSettings(ExtendedBaseSettings):
     embedding_azure_api_base: str = ""
     embedding_azure_api_version: str = ""
     embedding_max_str_num: int = 50
+    embedding_max_length: int = 8192
 
     # offline llama2 related config
     use_llama2: bool = False

rdagent/oai/utils/embedding.py

@@ -0,0 +1,134 @@
+"""
+Embedding utilities for handling token limits and text truncation.
+"""
+
+from typing import Optional
+
+from litellm import decode, encode, get_max_tokens, token_counter
+
+from rdagent.log import rdagent_logger as logger
+from rdagent.oai.llm_conf import LLM_SETTINGS
+
+# Common embedding model token limits
+EMBEDDING_MODEL_LIMITS = {
+    "text-embedding-ada-002": 8191,
+    "text-embedding-3-small": 8191,
+    "text-embedding-3-large": 8191,
+    "Qwen3-Embedding-8B": 32000,
+    "Qwen3-Embedding-4B": 32000,
+    "Qwen3-Embedding-0.6B": 32000,
+    "bge-m3": 8191,
+    "bce-embedding-base_v1": 511,
+    "bge-large-zh-v1.5": 511,
+    "bge-large-en-v1.5": 511,
+}
+
+
+def get_embedding_max_tokens(model: str) -> int:
+    """
+    Get maximum token limit for embedding model.
+
+    Three-level fallback strategy:
+    1. Use litellm.get_max_tokens()
+    2. Query EMBEDDING_MODEL_LIMITS mapping
+    3. Use default value 8192
+
+    Args:
+        model: Model name
+
+    Returns:
+        Maximum token limit
+    """
+    # Remove prefix (e.g., "provider/model" -> "model")
+    model_name = model.split("/")[-1] if "/" in model else model
+
+    # Level 1: Try litellm
+    try:
+        max_tokens = get_max_tokens(model_name)
+        if max_tokens and max_tokens > 0:
+            return max_tokens
+    except Exception as e:
+        logger.warning(f"Failed to get max tokens for {model_name}: {e}")
+
+    # Level 2: Query mapping table
+    if model_name in EMBEDDING_MODEL_LIMITS:
+        return EMBEDDING_MODEL_LIMITS[model_name]
+
+    # Level 3: fallback to LLM_SETTINGS.embedding_max_length
+    default_max_tokens = LLM_SETTINGS.embedding_max_length
+    logger.warning(f"Unknown embedding model {model}, using default max_tokens={default_max_tokens}")
+    return default_max_tokens
+
+
+def trim_text_for_embedding(text: str, model: str, max_tokens: Optional[int] = None) -> str:
+    """
+    Truncate text for embedding model using encode/decode approach.
+
+    Args:
+        text: Input text
+        model: Model name
+        max_tokens: Maximum token limit, auto-detected if None. If still exceeds limit,
+                   raises error directing user to set LLM_SETTINGS.embedding_max_length
+
+    Returns:
+        Truncated text
+    """
+    if not text:
+        return ""
+
+    # Get model's maximum token limit
+    if max_tokens is None:
+        max_tokens = get_embedding_max_tokens(model)
+
+    # Apply safety margin
+    safe_max_tokens = int(max_tokens * 0.9)
+
+    # Calculate current token count
+    current_tokens = token_counter(model=model, text=text)
+
+    if current_tokens <= safe_max_tokens:
+        return text
+
+    logger.warning(
+        f"Text too long for embedding model {model}: "
+        f"{current_tokens} tokens > {safe_max_tokens} limit (with safety margin). "
+        f"Truncating using encode/decode approach."
+    )
+
+    try:
+        # Use encode/decode approach for precise truncation
+        enc_ids = encode(model=model, text=text)
+        enc_ids_trunc = enc_ids[:safe_max_tokens]
+        text_trunc = decode(model=model, tokens=enc_ids_trunc)
+        # Ensure we return a string type (mypy type safety)
+        text_trunc = str(text_trunc) if text_trunc is not None else ""
+
+        final_tokens = token_counter(model=model, text=text_trunc)
+        logger.warning(f"Truncation completed: {current_tokens} -> {final_tokens} tokens")
+
+        return text_trunc
+    except Exception as e:
+        raise RuntimeError(
+            f"Failed to truncate text for embedding model {model}. "
+            f"Please set LLM_SETTINGS.embedding_max_length to a smaller value. "
+            f"Original error: {e}"
+        ) from e
+
+
+def truncate_content_list(content_list: list[str], model: str) -> list[str]:
+    """
+    Truncate a list of content strings.
+
+    Args:
+        content_list: List of content strings to truncate
+        model: Model name
+
+    Returns:
+        List of truncated content strings
+    """
+    truncated_list = []
+    for content in content_list:
+        truncated_content = trim_text_for_embedding(content, model)
+        truncated_list.append(truncated_content)
+
+    return truncated_list

test/oai/test_embedding_and_similarity.py

@@ -24,9 +24,26 @@ class TestEmbedding(unittest.TestCase):
         assert similarity is not None
         assert isinstance(similarity, float)
         min_similarity_threshold = 0.8
-        print(f"similarity: {similarity}")
         assert similarity >= min_similarity_threshold
 
+    def test_embedding_long_text_truncation(self) -> None:
+        """Test embedding with very long text that exceeds token limits"""
+        # Create a very long text that will definitely exceed embedding token limits
+        # Using a repetitive pattern to simulate a real long document
+        long_content = (
+            """
+        This is a very long document that contains a lot of repetitive content to test the embedding truncation functionality.
+        We need to make this text long enough to exceed the typical embedding model token limits of around 8192 tokens.
+        """
+            * 1000
+        )  # This should create a text with approximately 50,000+ tokens
+        # This should trigger the gradual truncation mechanism
+        emb = APIBackend().create_embedding(long_content)
+
+        assert emb is not None
+        assert isinstance(emb, list)
+        assert len(emb) > 0
+
 
 if __name__ == "__main__":
     unittest.main()