fix(openai.py): fix linting issues

This commit is contained in:
Krrish Dholakia 2023-11-16 12:57:53 -08:00
parent a94c09c13c
commit d7f7694848

View file

@ -233,9 +233,9 @@ class OpenAIChatCompletion(BaseLLM):
async def acompletion(self, async def acompletion(self,
data: dict, data: dict,
model_response: ModelResponse, model_response: ModelResponse,
api_base: str, api_key: Optional[str]=None,
api_key: str): api_base: Optional[str]=None):
response = None response = None
try: try:
openai_aclient = AsyncOpenAI(api_key=api_key, base_url=api_base) openai_aclient = AsyncOpenAI(api_key=api_key, base_url=api_base)
@ -251,8 +251,8 @@ class OpenAIChatCompletion(BaseLLM):
logging_obj, logging_obj,
data: dict, data: dict,
model: str, model: str,
api_key: str, api_key: Optional[str]=None,
api_base: str api_base: Optional[str]=None
): ):
openai_client = OpenAI(api_key=api_key, base_url=api_base) openai_client = OpenAI(api_key=api_key, base_url=api_base)
response = openai_client.chat.completions.create(**data) response = openai_client.chat.completions.create(**data)
@ -264,15 +264,14 @@ class OpenAIChatCompletion(BaseLLM):
logging_obj, logging_obj,
data: dict, data: dict,
model: str, model: str,
api_key: str, api_key: Optional[str]=None,
api_base: str): api_base: Optional[str]=None):
openai_aclient = AsyncOpenAI(api_key=api_key, base_url=api_base) openai_aclient = AsyncOpenAI(api_key=api_key, base_url=api_base)
response = await openai_aclient.chat.completions.create(**data) response = await openai_aclient.chat.completions.create(**data)
streamwrapper = CustomStreamWrapper(completion_stream=response, model=model, custom_llm_provider="openai",logging_obj=logging_obj) streamwrapper = CustomStreamWrapper(completion_stream=response, model=model, custom_llm_provider="openai",logging_obj=logging_obj)
async for transformed_chunk in streamwrapper: async for transformed_chunk in streamwrapper:
yield transformed_chunk yield transformed_chunk
def embedding(self, def embedding(self,
model: str, model: str,
input: list, input: list,
@ -284,7 +283,7 @@ class OpenAIChatCompletion(BaseLLM):
super().embedding() super().embedding()
exception_mapping_worked = False exception_mapping_worked = False
try: try:
openai_client = OpenAI(api_key=api_key, api_base=api_base) openai_client = OpenAI(api_key=api_key, base_url=api_base)
model = model model = model
data = { data = {
"model": model, "model": model,