Merge pull request #94 from krrishdholakia/patch-1

Adding support for Claude, PaLM, Llama2, and Cohere models

api.py

@@ -4,7 +4,7 @@ import shutil
 import urllib.request
 from pathlib import Path
 from tempfile import NamedTemporaryFile
-
+from litellm import completion
 import fitz
 import numpy as np
 import openai
@@ -109,17 +109,19 @@ def load_recommender(path, start_page=1):
 
 
 def generate_text(openAI_key, prompt, engine="text-davinci-003"):
-    openai.api_key = openAI_key
+    # openai.api_key = openAI_key
     try:
-        completions = openai.Completion.create(
+        messages=[{ "content": prompt,"role": "user"}]
-            engine=engine,
+        completions = completion(
-            prompt=prompt,
+            model=engine,
+            messages=messages,
             max_tokens=512,
             n=1,
             stop=None,
             temperature=0.7,
+            api_key=openAI_key
         )
-        message = completions.choices[0].text
+        message = completions['choices'][0]['message']['content']
     except Exception as e:
         message = f'API Error: {str(e)}'
     return message 

requirements.txt

@@ -6,3 +6,4 @@ tensorflow_hub==0.13.0
 openai==0.27.4
 gradio==3.34.0
 langchain-serve>=0.0.19
+litellm