Add support for Palm, Claude-2, Llama2, CodeLlama (100+LLMs)

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."}

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}"}

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,

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