From 91c88181f5c644a03bfedaa85d7c68ec4c293359 Mon Sep 17 00:00:00 2001
From: ishaan-jaff <ishaanjaffer0324@gmail.com>
Date: Sat, 9 Sep 2023 09:04:18 -0700
Subject: [PATCH] v0 litellm

---
 backend/config/utils.py     | 6 +++---
 backend/embeddings/utils.py | 9 +++++----
 backend/llm/summarize.py    | 7 ++++---
 backend/requirements.txt    | 1 +
 4 files changed, 13 insertions(+), 10 deletions(-)

diff --git a/backend/config/utils.py b/backend/config/utils.py
index f2aec09..915136d 100644
--- a/backend/config/utils.py
+++ b/backend/config/utils.py
@@ -1,9 +1,9 @@
-import openai
+import litellm
 
 def is_valid_api_key(openai_api_key: str):
-    openai.api_key = openai_api_key
+    litellm.api_key = openai_api_key
     try:
-        openai.ChatCompletion.create(
+        litellm.completion(
             model="gpt-3.5-turbo",
             messages=[
                 {"role": "system", "content": "You are a helpful assistant."}
diff --git a/backend/embeddings/utils.py b/backend/embeddings/utils.py
index 3aa6011..e3b1c8f 100644
--- a/backend/embeddings/utils.py
+++ b/backend/embeddings/utils.py
@@ -1,6 +1,7 @@
 from typing import List
 from gpt4all import GPT4All
 import openai
+import litellm
 import os
 
 
@@ -12,8 +13,8 @@ def openai_ask(prompt: str, context: str = None, pages: List[int] = None, questi
     print(pages)
     # TODO: make answer to same language
     if openai_model.startswith("gpt"):
-        openai.api_key = openai_api_key
-        completion = openai.ChatCompletion.create(
+        litellm.api_key = openai_api_key
+        completion = litellm.completion(
             model=f"{openai_model}",
             messages=[
                 {"role": "user", "content": f"{prompt}"}
@@ -36,8 +37,8 @@ def openai_ask(prompt: str, context: str = None, pages: List[int] = None, questi
 
 def openai_ask_no_aixplora_brain(question: str, openai_api_key: str = None, openai_model: str = "gpt-3.5-turbo"):
     if openai_model.startswith("gpt"):
-        openai.api_key = openai_api_key
-        completion = openai.ChatCompletion.create(
+        litellm.api_key = openai_api_key
+        completion = litellm.completion(
             model=f"{openai_model}",
             messages=[
                 {"role": "user", "content": f"{question}"}
diff --git a/backend/llm/summarize.py b/backend/llm/summarize.py
index 935bc3e..d315a7c 100644
--- a/backend/llm/summarize.py
+++ b/backend/llm/summarize.py
@@ -1,4 +1,5 @@
 import openai
+import litellm
 
 from schemas.question import Document
 from database.database import Database
@@ -72,7 +73,7 @@ def get_summary(self):
                     text_split = text[words_to_feed * i:words_to_feed * (i + 1)]
                     time.sleep(0.25)
                     if self.model.startswith("gpt"):
-                        response = openai.ChatCompletion.create(
+                        response = litellm.completion(
                             api_key=f"{self.openai_api_key}",
                             model=f"{self.openai_model}",
                             temperature=0.2,
@@ -100,7 +101,7 @@ def get_summary(self):
                         raise ("The summary is still too long. Please try again.")
                 count_token = num_tokens_from_string("".join(summary), "cl100k_base")
                 if self.model.startswith("gpt"):
-                    response = openai.ChatCompletion.create(
+                    response = litellm.completion(
                         api_key=f"{self.openai_api_key}",
                         model=f"{self.openai_model}",
                         temperature=0.2,
@@ -121,7 +122,7 @@ def get_summary(self):
                     return {"summary": response["choices"][0]["message"]["content"], "summary_list": "<hr>".join(summary)}
             else:
                 if self.model.startswith("gpt"):
-                    response = openai.ChatCompletion.create(
+                    response = litellm.completion(
                         api_key=f"{self.openai_api_key}",
                         model=f"{self.openai_model}",
                         temperature=0.2,
diff --git a/backend/requirements.txt b/backend/requirements.txt
index 3b39fe9..5edb6ac 100644
--- a/backend/requirements.txt
+++ b/backend/requirements.txt
@@ -34,6 +34,7 @@ idna==3.4
 iniconfig==2.0.0
 joblib==1.2.0
 langchain==0.0.178
+litellm>=0.1.574
 lxml==4.9.2
 lz4==4.3.2
 marshmallow==3.19.0