diff --git a/rag/llm/cv_model.py b/rag/llm/cv_model.py index f3a037998..9b083acca 100644 --- a/rag/llm/cv_model.py +++ b/rag/llm/cv_model.py @@ -298,12 +298,24 @@ class GptV4(Base): return res.choices[0].message.content.strip(), total_token_count_from_response(res) +def _resolve_azure_credentials(key): + try: + key_obj = json.loads(key) + if isinstance(key_obj, dict): + return key_obj.get("api_key", ""), key_obj.get("api_version", "2024-02-01") + logging.warning( + "Azure credential payload parsed as JSON but is not an object; using raw api_key string" + ) + except (json.JSONDecodeError, TypeError): + logging.warning("Azure credential payload is not valid JSON; using raw api_key string") + return key, "2024-02-01" + + class AzureGptV4(GptV4): _FACTORY_NAME = "Azure-OpenAI" def __init__(self, key, model_name, lang="Chinese", **kwargs): - api_key = json.loads(key).get("api_key", "") - api_version = json.loads(key).get("api_version", "2024-02-01") + api_key, api_version = _resolve_azure_credentials(key) self.client = AzureOpenAI(api_key=api_key, azure_endpoint=kwargs["base_url"], api_version=api_version) self.async_client = AsyncAzureOpenAI(api_key=api_key, azure_endpoint=kwargs["base_url"], api_version=api_version) self.model_name = model_name diff --git a/rag/llm/embedding_model.py b/rag/llm/embedding_model.py index 46c43d475..616382cf7 100644 --- a/rag/llm/embedding_model.py +++ b/rag/llm/embedding_model.py @@ -235,14 +235,26 @@ class LocalAIEmbed(Base): return np.array(embds[0]), cnt +def _resolve_azure_credentials(key): + try: + key_obj = json.loads(key) + if isinstance(key_obj, dict): + return key_obj.get("api_key", ""), key_obj.get("api_version", "2024-02-01") + logging.warning( + "Azure credential payload parsed as JSON but is not an object; using raw api_key string" + ) + except (json.JSONDecodeError, TypeError): + logging.warning("Azure credential payload is not valid JSON; using raw api_key string") + return key, "2024-02-01" + + class AzureEmbed(OpenAIEmbed): _FACTORY_NAME = "Azure-OpenAI" def __init__(self, key, model_name, **kwargs): from openai.lib.azure import AzureOpenAI - api_key = json.loads(key).get("api_key", "") - api_version = json.loads(key).get("api_version", "2024-02-01") + api_key, api_version = _resolve_azure_credentials(key) self.client = AzureOpenAI(api_key=api_key, azure_endpoint=kwargs["base_url"], api_version=api_version) self.model_name = model_name