Update litellm_ai_handler.py

updates made based on review on https://github.com/qodo-ai/pr-agent/pull/1564
This commit is contained in:
chandan84
2025-02-25 14:44:03 -05:00
committed by GitHub
parent 84983f3e9d
commit 93e34703ab

View File

@ -244,12 +244,6 @@ class LiteLLMAIHandler(BaseAiHandler):
if self.repetition_penalty: if self.repetition_penalty:
kwargs["repetition_penalty"] = self.repetition_penalty kwargs["repetition_penalty"] = self.repetition_penalty
get_logger().debug("Prompts", artifact={"system": system, "user": user})
if get_settings().config.verbosity_level >= 2:
get_logger().info(f"\nSystem prompt:\n{system}")
get_logger().info(f"\nUser prompt:\n{user}")
#Added support for extra_headers while using litellm to call underlying model, via a api management gateway, would allow for passing custom headers for security and authorization #Added support for extra_headers while using litellm to call underlying model, via a api management gateway, would allow for passing custom headers for security and authorization
if get_settings().get("LITELLM.EXTRA_HEADERS", None): if get_settings().get("LITELLM.EXTRA_HEADERS", None):
try: try:
@ -260,6 +254,12 @@ class LiteLLMAIHandler(BaseAiHandler):
raise ValueError(f"LITELLM.EXTRA_HEADERS contains invalid JSON: {str(e)}") raise ValueError(f"LITELLM.EXTRA_HEADERS contains invalid JSON: {str(e)}")
kwargs["extra_headers"] = litellm_extra_headers kwargs["extra_headers"] = litellm_extra_headers
get_logger().debug("Prompts", artifact={"system": system, "user": user})
if get_settings().config.verbosity_level >= 2:
get_logger().info(f"\nSystem prompt:\n{system}")
get_logger().info(f"\nUser prompt:\n{user}")
response = await acompletion(**kwargs) response = await acompletion(**kwargs)
except (openai.APIError, openai.APITimeoutError) as e: except (openai.APIError, openai.APITimeoutError) as e:
get_logger().warning(f"Error during LLM inference: {e}") get_logger().warning(f"Error during LLM inference: {e}")