From 102edcdcf1b38f16028c45fd66b32542fb69f9a6 Mon Sep 17 00:00:00 2001
From: Krish Dholakia <krrishdholakia@gmail.com>
Date: Thu, 3 Aug 2023 12:04:08 -0700
Subject: [PATCH] adding support for Anthropic, Cohere, Replicate, Azure

---
 pr_agent/algo/ai_handler.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/pr_agent/algo/ai_handler.py b/pr_agent/algo/ai_handler.py
index aa92f5e5..11fd56ef 100644
--- a/pr_agent/algo/ai_handler.py
+++ b/pr_agent/algo/ai_handler.py
@@ -3,7 +3,7 @@ import logging
 import openai
 from openai.error import APIError, RateLimitError, Timeout, TryAgain
 from retry import retry
-
+from litellm import acompletion
 from pr_agent.config_loader import get_settings
 
 OPENAI_RETRIES=5
@@ -57,7 +57,7 @@ class AiHandler:
             TryAgain: If there is an attribute error during OpenAI inference.
         """
         try:
-            response = await openai.ChatCompletion.acreate(
+            response = await acompletion(
                             model=model,
                             deployment_id=self.deployment_id,
                             messages=[