diff --git a/pr_agent/algo/ai_handlers/openai_ai_handler.py b/pr_agent/algo/ai_handlers/openai_ai_handler.py index 999f3d3f..8e755b46 100644 --- a/pr_agent/algo/ai_handlers/openai_ai_handler.py +++ b/pr_agent/algo/ai_handlers/openai_ai_handler.py @@ -1,6 +1,7 @@ +from os import environ from pr_agent.algo.ai_handlers.base_ai_handler import BaseAiHandler import openai -from openai.error import APIError, RateLimitError, Timeout, TryAgain +from openai import APIError, AsyncOpenAI, RateLimitError, Timeout from retry import retry from pr_agent.config_loader import get_settings @@ -14,7 +15,7 @@ class OpenAIHandler(BaseAiHandler): # Initialize OpenAIHandler specific attributes here try: super().__init__() - openai.api_key = get_settings().openai.key + environ["OPENAI_API_KEY"] = get_settings().openai.key if get_settings().get("OPENAI.ORG", None): openai.organization = get_settings().openai.org if get_settings().get("OPENAI.API_TYPE", None): @@ -24,7 +25,7 @@ class OpenAIHandler(BaseAiHandler): if get_settings().get("OPENAI.API_VERSION", None): openai.api_version = get_settings().openai.api_version if get_settings().get("OPENAI.API_BASE", None): - openai.api_base = get_settings().openai.api_base + environ["OPENAI_BASE_URL"] = get_settings().openai.api_base except AttributeError as e: raise ValueError("OpenAI key is required") from e @@ -36,7 +37,7 @@ class OpenAIHandler(BaseAiHandler): """ return get_settings().get("OPENAI.DEPLOYMENT_ID", None) - @retry(exceptions=(APIError, Timeout, TryAgain, AttributeError, RateLimitError), + @retry(exceptions=(APIError, Timeout, AttributeError, RateLimitError), tries=OPENAI_RETRIES, delay=2, backoff=2, jitter=(1, 3)) async def chat_completion(self, model: str, system: str, user: str, temperature: float = 0.2): try: @@ -44,20 +45,19 @@ class OpenAIHandler(BaseAiHandler): get_logger().info("System: ", system) get_logger().info("User: ", user) messages = [{"role": "system", "content": system}, {"role": "user", "content": user}] - - chat_completion = await openai.ChatCompletion.acreate( + client = AsyncOpenAI() + chat_completion = await client.chat.completions.create( model=model, - deployment_id=deployment_id, messages=messages, temperature=temperature, ) - resp = chat_completion["choices"][0]['message']['content'] - finish_reason = chat_completion["choices"][0]["finish_reason"] - usage = chat_completion.get("usage") + resp = chat_completion.choices[0].message.content + finish_reason = chat_completion.choices[0].finish_reason + usage = chat_completion.usage get_logger().info("AI response", response=resp, messages=messages, finish_reason=finish_reason, model=model, usage=usage) return resp, finish_reason - except (APIError, Timeout, TryAgain) as e: + except (APIError, Timeout) as e: get_logger().error("Error during OpenAI inference: ", e) raise except (RateLimitError) as e: @@ -65,4 +65,4 @@ class OpenAIHandler(BaseAiHandler): raise except (Exception) as e: get_logger().error("Unknown error during OpenAI inference: ", e) - raise TryAgain from e + raise \ No newline at end of file