diff --git a/rdagent/oai/llm_utils.py b/rdagent/oai/llm_utils.py index b27268ab..d0cb8abe 100644 --- a/rdagent/oai/llm_utils.py +++ b/rdagent/oai/llm_utils.py @@ -327,7 +327,7 @@ class APIBackend: self.chat_model = LLM_SETTINGS.chat_model if chat_model is None else chat_model self.chat_model_map = json.loads(LLM_SETTINGS.chat_model_map) - self.encoder = tiktoken.encoding_for_model(self.chat_model) + self.encoder = self._get_encoder() self.chat_api_base = LLM_SETTINGS.chat_azure_api_base if chat_api_base is None else chat_api_base self.chat_api_version = ( LLM_SETTINGS.chat_azure_api_version if chat_api_version is None else chat_api_version @@ -399,6 +399,33 @@ class APIBackend: self.use_gcr_endpoint = LLM_SETTINGS.use_gcr_endpoint self.retry_wait_seconds = LLM_SETTINGS.retry_wait_seconds + def _get_encoder(self): + """ + tiktoken.encoding_for_model(self.chat_model) does not cover all cases it should consider. + + This function attempts to handle several edge cases. + """ + + # 1) cases + def _azure_patch(model: str) -> str: + """ + When using Azure API, self.chat_model is the deployment name that can be any string. + For example, it may be `gpt-4o_2024-08-06`. But tiktoken.encoding_for_model can't handle this. + """ + return model.replace("_", "-") + + model = self.chat_model + try: + return tiktoken.encoding_for_model(model) + except KeyError: + logger.warning(f"Failed to get encoder. Trying to patch the model name") + for patch_func in [_azure_patch]: + try: + return tiktoken.encoding_for_model(patch_func(model)) + except KeyError: + logger.error(f"Failed to get encoder even after patching with {patch_func.__name__}") + raise + def build_chat_session( self, conversation_id: str | None = None,