diff --git a/api.py b/api.py index b44b21b..a9fccac 100644 --- a/api.py +++ b/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( - engine=engine, - prompt=prompt, + messages=[{ "content": prompt,"role": "user"}] + completions = completion( + 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 diff --git a/requirements.txt b/requirements.txt index 6e131c8..5b74f23 100644 --- a/requirements.txt +++ b/requirements.txt @@ -6,3 +6,4 @@ tensorflow-hub openai==0.10.2 gradio langchain-serve>=0.0.19 +litellm