mirror of
https://github.com/qodo-ai/pr-agent.git
synced 2025-07-11 00:00:38 +08:00
Merge pull request #567 from zmeir/zmeir/enhance/support_azure_in_langchain_ai_handler
Add support for Azure OpenAI in LangChainOpenAIHandler
This commit is contained in:
@ -1,5 +1,5 @@
|
||||
try:
|
||||
from langchain.chat_models import ChatOpenAI
|
||||
from langchain.chat_models import ChatOpenAI, AzureChatOpenAI
|
||||
from langchain.schema import SystemMessage, HumanMessage
|
||||
except: # we don't enforce langchain as a dependency, so if it's not installed, just move on
|
||||
pass
|
||||
@ -10,21 +10,39 @@ from pr_agent.log import get_logger
|
||||
|
||||
from openai.error import APIError, RateLimitError, Timeout, TryAgain
|
||||
from retry import retry
|
||||
import functools
|
||||
|
||||
OPENAI_RETRIES = 5
|
||||
|
||||
class LangChainOpenAIHandler(BaseAiHandler):
|
||||
def __init__(self):
|
||||
# Initialize OpenAIHandler specific attributes here
|
||||
try:
|
||||
super().__init__()
|
||||
self.azure = get_settings().get("OPENAI.API_TYPE", "").lower() == "azure"
|
||||
try:
|
||||
if self.azure:
|
||||
# using a partial function so we can set the deployment_id later to support fallback_deployments
|
||||
# but still need to access the other settings now so we can raise a proper exception if they're missing
|
||||
self._chat = functools.partial(
|
||||
lambda **kwargs: AzureChatOpenAI(**kwargs),
|
||||
openai_api_key=get_settings().openai.key,
|
||||
openai_api_base=get_settings().openai.api_base,
|
||||
openai_api_version=get_settings().openai.api_version,
|
||||
)
|
||||
else:
|
||||
self._chat = ChatOpenAI(openai_api_key=get_settings().openai.key)
|
||||
|
||||
except AttributeError as e:
|
||||
raise ValueError("OpenAI key is required") from e
|
||||
if getattr(e, "name"):
|
||||
raise ValueError(f"OpenAI {e.name} is required") from e
|
||||
else:
|
||||
raise e
|
||||
|
||||
@property
|
||||
def chat(self):
|
||||
if self.azure:
|
||||
# we must set the deployment_id only here (instead of the __init__ method) to support fallback_deployments
|
||||
return self._chat(deployment_name=self.deployment_id)
|
||||
else:
|
||||
return self._chat
|
||||
|
||||
@property
|
||||
|
@ -101,11 +101,6 @@ class LiteLLMAIHandler(BaseAiHandler):
|
||||
"""
|
||||
try:
|
||||
deployment_id = self.deployment_id
|
||||
if get_settings().config.verbosity_level >= 2:
|
||||
get_logger().debug(
|
||||
f"Generating completion with {model}"
|
||||
f"{(' from deployment ' + deployment_id) if deployment_id else ''}"
|
||||
)
|
||||
if self.azure:
|
||||
model = 'azure/' + model
|
||||
messages = [{"role": "system", "content": system}, {"role": "user", "content": user}]
|
||||
|
@ -226,6 +226,11 @@ async def retry_with_fallback_models(f: Callable):
|
||||
# try each (model, deployment_id) pair until one is successful, otherwise raise exception
|
||||
for i, (model, deployment_id) in enumerate(zip(all_models, all_deployments)):
|
||||
try:
|
||||
if get_settings().config.verbosity_level >= 2:
|
||||
get_logger().debug(
|
||||
f"Generating prediction with {model}"
|
||||
f"{(' from deployment ' + deployment_id) if deployment_id else ''}"
|
||||
)
|
||||
get_settings().set("openai.deployment_id", deployment_id)
|
||||
return await f(model)
|
||||
except Exception as e:
|
||||
|
Reference in New Issue
Block a user