From 5785bcf2ba58d12eb7822082bf661eaef89c8378 Mon Sep 17 00:00:00 2001 From: amstrongzyf <201840057@smail.nju.edu.cn> Date: Mon, 18 Aug 2025 17:20:41 +0800 Subject: [PATCH] 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 | 25 +++- rdagent/oai/llm_conf.py | 1 + rdagent/oai/utils/embedding.py | 134 ++++++++++++++++++++++ test/oai/test_embedding_and_similarity.py | 19 ++- 4 files changed, 174 insertions(+), 5 deletions(-) create mode 100644 rdagent/oai/utils/embedding.py diff --git a/rdagent/oai/backend/base.py b/rdagent/oai/backend/base.py index 4f890231..7210ed1d 100644 --- a/rdagent/oai/backend/base.py +++ b/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")) diff --git a/rdagent/oai/llm_conf.py b/rdagent/oai/llm_conf.py index db17777b..84804541 100644 --- a/rdagent/oai/llm_conf.py +++ b/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 diff --git a/rdagent/oai/utils/embedding.py b/rdagent/oai/utils/embedding.py new file mode 100644 index 00000000..496b5e16 --- /dev/null +++ b/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 diff --git a/test/oai/test_embedding_and_similarity.py b/test/oai/test_embedding_and_similarity.py index 8c12b918..4235c466 100644 --- a/test/oai/test_embedding_and_similarity.py +++ b/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()