Revert "Revert "Add support for multiple LLM Providers (#60)""

This reverts commit 43d6bc80e9.
This commit is contained in:
José Ulises Niño Rivera 2024-09-25 09:52:20 -07:00
parent 43d6bc80e9
commit d38246ceaf
12 changed files with 456 additions and 127 deletions

View file

@ -6,27 +6,31 @@ from dotenv import load_dotenv
load_dotenv()
OPEN_API_KEY=os.getenv("OPENAI_API_KEY")
OPENAI_API_KEY=os.getenv("OPENAI_API_KEY")
MISTRAL_API_KEY = os.getenv("MISTRAL_API_KEY")
CHAT_COMPLETION_ENDPOINT = os.getenv("CHAT_COMPLETION_ENDPOINT")
MODEL_NAME = os.getenv("MODEL_NAME", "gpt-3.5-turbo")
log.info("CHAT_COMPLETION_ENDPOINT: ", CHAT_COMPLETION_ENDPOINT)
client = OpenAI(api_key=OPEN_API_KEY, base_url=CHAT_COMPLETION_ENDPOINT)
client = OpenAI(api_key=OPENAI_API_KEY, base_url=CHAT_COMPLETION_ENDPOINT)
def predict(message, history):
# history_openai_format = []
# for human, assistant in history:
# history_openai_format.append({"role": "user", "content": human })
# history_openai_format.append({"role": "assistant", "content":assistant})
history.append({"role": "user", "content": message})
log.info("CHAT_COMPLETION_ENDPOINT: ", CHAT_COMPLETION_ENDPOINT)
log.info("history: ", history)
# Custom headers
custom_headers = {
'x-bolt-openai-api-key': f"{OPENAI_API_KEY}",
'x-bolt-mistral-api-key': f"{MISTRAL_API_KEY}",
}
try:
response = client.chat.completions.create(model=MODEL_NAME,
messages= history,
temperature=1.0
temperature=1.0,
headers=custom_headers
)
except Exception as e:
log.info(e)
@ -36,10 +40,6 @@ def predict(message, history):
log.info("Error calling gateway API: {}".format(e.message))
raise gr.Error("Error calling gateway API: {}".format(e.message))
# for chunk in response:
# if chunk.choices[0].delta.content is not None:
# partial_message = partial_message + chunk.choices[0].delta.content
# yield partial_message
choices = response.choices
message = choices[0].message
content = message.content