From 36209caf19881502a506f3d02c494845e656d3c0 Mon Sep 17 00:00:00 2001
From: ishaan-jaff <ishaanjaffer0324@gmail.com>
Date: Fri, 4 Aug 2023 15:26:50 -0700
Subject: [PATCH] chatCompleio

---
 factool/utils/claim_extractor.py | 3 ++-
 factool/utils/openai_wrapper.py  | 4 +++-
 setup.cfg                        | 1 +
 3 files changed, 6 insertions(+), 2 deletions(-)

diff --git a/factool/utils/claim_extractor.py b/factool/utils/claim_extractor.py
index 8b92ecbb..7e7f4a31 100644
--- a/factool/utils/claim_extractor.py
+++ b/factool/utils/claim_extractor.py
@@ -6,6 +6,7 @@
 import asyncio
 from tqdm import tqdm
 from factool.env_config import factool_env_config
+from litellm import completion
 
 
 # env
@@ -26,7 +27,7 @@
 # Make api calls asynchronously
 async def run_api(messages):
     async def single_run(message):
-        output = openai.ChatCompletion.create(
+        output = completion(
             model=config['model_name'],
             messages=message,
             max_tokens=config['max_tokens'],
diff --git a/factool/utils/openai_wrapper.py b/factool/utils/openai_wrapper.py
index 929f8b94..af931caa 100644
--- a/factool/utils/openai_wrapper.py
+++ b/factool/utils/openai_wrapper.py
@@ -11,6 +11,7 @@
 from typing import Any, List
 import os
 import pathlib
+from litellm import acompletion
 
 
 from factool.env_config import factool_env_config
@@ -27,6 +28,7 @@ def __init__(
             top_p=1,
             request_timeout=60,
     ):
+        # TODO: modify this to allow Anthropic, Claude, etc
         if 'gpt' not in model_name:
             openai.api_base = "http://localhost:8000/v1"
         self.config = {
@@ -64,7 +66,7 @@ async def dispatch_openai_requests(
         async def _request_with_retry(messages, retry=3):
             for _ in range(retry):
                 try:
-                    response = await openai.ChatCompletion.acreate(
+                    response = await acompletion(
                         model=self.config['model_name'],
                         messages=messages,
                         max_tokens=self.config['max_tokens'],
diff --git a/setup.cfg b/setup.cfg
index 8dc93a75..1dd1a5f5 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -21,6 +21,7 @@ packages = find:
 include_package_data = True
 install_requires =
     openai==0.27.8
+    litellm==0.1.232
     PyYAML==6.0
     asyncio==3.4.3
     numpy