From 389f6a677c5c4cdfea99e6a20b2bcb6090661ea3 Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Wed, 2 Aug 2023 16:06:29 -0700 Subject: [PATCH] v0 --- guardrails/llm_providers.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/guardrails/llm_providers.py b/guardrails/llm_providers.py index 574fbeb92..65be37869 100644 --- a/guardrails/llm_providers.py +++ b/guardrails/llm_providers.py @@ -4,6 +4,7 @@ from typing import Any, Awaitable, Callable, Dict, List, Optional, cast import openai +from litellm import completion from pydantic import BaseModel from tenacity import retry, retry_if_exception_type, wait_exponential_jitter @@ -128,8 +129,9 @@ def openai_chat_wrapper( api_key = os.environ.get("OPENAI_API_KEY") # TODO: update this as new models are released + # Supported models here: https://litellm.readthedocs.io/en/latest/supported/ if base_model: - openai_response = openai.ChatCompletion.create( + openai_response = completion( api_key=api_key, model=model, messages=chat_prompt(text, instructions, **kwargs), @@ -140,7 +142,7 @@ def openai_chat_wrapper( ) return openai_response["choices"][0]["message"]["function_call"]["arguments"] else: - openai_response = openai.ChatCompletion.create( + openai_response = completion( api_key=api_key, model=model, messages=chat_prompt(text, instructions, **kwargs),