Cleanup litellm chat/embedding model argument assignment.

Author: dworthen

graphrag/language_model/providers/litellm/chat_model.py

@@ -64,27 +64,17 @@ def _create_base_completions(
         "frequency_penalty": model_config.frequency_penalty,
         "presence_penalty": model_config.presence_penalty,
         "drop_params": True,
+        "api_base": model_config.api_base,
+        "api_version": model_config.api_version,
+        "api_key": model_config.api_key,
+        "organization": model_config.organization,
+        "proxy": model_config.proxy,
+        "audience": model_config.audience,
+        "max_tokens": model_config.max_tokens,
+        "max_completion_tokens": model_config.max_completion_tokens,
+        "reasoning_effort": model_config.reasoning_effort,
     }
 
-    if model_config.api_base:
-        base_args["api_base"] = model_config.api_base
-    if model_config.api_version:
-        base_args["api_version"] = model_config.api_version
-    if model_config.api_key:
-        base_args["api_key"] = model_config.api_key
-    if model_config.organization:
-        base_args["organization"] = model_config.organization
-    if model_config.proxy:
-        base_args["proxy"] = model_config.proxy
-    if model_config.audience:
-        base_args["audience"] = model_config.audience
-    if model_config.max_tokens:
-        base_args["max_tokens"] = model_config.max_tokens
-    if model_config.max_completion_tokens:
-        base_args["max_completion_tokens"] = model_config.max_completion_tokens
-    if model_config.reasoning_effort:
-        base_args["reasoning_effort"] = model_config.reasoning_effort
-
     if model_config.auth_type == AuthType.AzureManagedIdentity:
         if model_config.model_provider != "azure":
             msg = "Azure Managed Identity authentication is only supported for Azure models."

graphrag/language_model/providers/litellm/embedding_model.py

@@ -53,21 +53,14 @@ def _create_base_embeddings(
         "model": f"{model_provider}/{model}",
         "timeout": model_config.request_timeout,
         "drop_params": True,
+        "api_base": model_config.api_base,
+        "api_version": model_config.api_version,
+        "api_key": model_config.api_key,
+        "organization": model_config.organization,
+        "proxy": model_config.proxy,
+        "audience": model_config.audience,
     }
 
-    if model_config.api_base:
-        base_args["api_base"] = model_config.api_base
-    if model_config.api_version:
-        base_args["api_version"] = model_config.api_version
-    if model_config.api_key:
-        base_args["api_key"] = model_config.api_key
-    if model_config.organization:
-        base_args["organization"] = model_config.organization
-    if model_config.proxy:
-        base_args["proxy"] = model_config.proxy
-    if model_config.audience:
-        base_args["audience"] = model_config.audience
-
     if model_config.auth_type == AuthType.AzureManagedIdentity:
         if model_config.model_provider != "azure":
             msg = "Azure Managed Identity authentication is only supported for Azure models."