diff --git a/litellm/llms/openai.py b/litellm/llms/openai.py index c923cbf2d..0731bd509 100644 --- a/litellm/llms/openai.py +++ b/litellm/llms/openai.py @@ -284,7 +284,7 @@ class OpenAIChatCompletion(BaseLLM): additional_args={"complete_input_dict": data}, ) return convert_to_model_response_object(response_object=json.loads(stringified_response), model_response_object=model_response) - except Exception as e: + except Exception as e: raise e def streaming(self, @@ -631,24 +631,27 @@ class OpenAITextCompletion(BaseLLM): api_key: str, model: str): async with httpx.AsyncClient() as client: - response = await client.post(api_base, json=data, headers=headers, timeout=litellm.request_timeout) - response_json = response.json() - if response.status_code != 200: - raise OpenAIError(status_code=response.status_code, message=response.text) - - ## LOGGING - logging_obj.post_call( - input=prompt, - api_key=api_key, - original_response=response, - additional_args={ - "headers": headers, - "api_base": api_base, - }, - ) + try: + response = await client.post(api_base, json=data, headers=headers, timeout=litellm.request_timeout) + response_json = response.json() + if response.status_code != 200: + raise OpenAIError(status_code=response.status_code, message=response.text) + + ## LOGGING + logging_obj.post_call( + input=prompt, + api_key=api_key, + original_response=response, + additional_args={ + "headers": headers, + "api_base": api_base, + }, + ) - ## RESPONSE OBJECT - return self.convert_to_model_response_object(response_object=response_json, model_response_object=model_response) + ## RESPONSE OBJECT + return self.convert_to_model_response_object(response_object=response_json, model_response_object=model_response) + except Exception as e: + raise e def streaming(self, logging_obj, @@ -687,9 +690,12 @@ class OpenAITextCompletion(BaseLLM): method="POST", timeout=litellm.request_timeout ) as response: - if response.status_code != 200: - raise OpenAIError(status_code=response.status_code, message=response.text) - - streamwrapper = CustomStreamWrapper(completion_stream=response.aiter_lines(), model=model, custom_llm_provider="text-completion-openai",logging_obj=logging_obj) - async for transformed_chunk in streamwrapper: - yield transformed_chunk \ No newline at end of file + try: + if response.status_code != 200: + raise OpenAIError(status_code=response.status_code, message=response.text) + + streamwrapper = CustomStreamWrapper(completion_stream=response.aiter_lines(), model=model, custom_llm_provider="text-completion-openai",logging_obj=logging_obj) + async for transformed_chunk in streamwrapper: + yield transformed_chunk + except Exception as e: + raise e \ No newline at end of file diff --git a/litellm/main.py b/litellm/main.py index 6e8931191..52d2ae5b6 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -2205,7 +2205,8 @@ def text_completion( if stream == True or kwargs.get("stream", False) == True: response = TextCompletionStreamWrapper(completion_stream=response, model=model) return response - + if kwargs.get("acompletion", False) == True: + return response transformed_logprobs = None # only supported for TGI models try: diff --git a/litellm/tests/test_text_completion.py b/litellm/tests/test_text_completion.py index 9257a07f3..f75bd2f7f 100644 --- a/litellm/tests/test_text_completion.py +++ b/litellm/tests/test_text_completion.py @@ -169,17 +169,37 @@ def test_text_completion_stream(): # test_text_completion_stream() -async def test_text_completion_async_stream(): - try: - response = await atext_completion( - model="text-completion-openai/text-davinci-003", - prompt="good morning", - stream=True, - max_tokens=10, - ) - async for chunk in response: - print(f"chunk: {chunk}") - except Exception as e: - pytest.fail(f"GOT exception for HF In streaming{e}") +# async def test_text_completion_async_stream(): +# try: +# response = await atext_completion( +# model="text-completion-openai/text-davinci-003", +# prompt="good morning", +# stream=True, +# max_tokens=10, +# ) +# async for chunk in response: +# print(f"chunk: {chunk}") +# except Exception as e: +# pytest.fail(f"GOT exception for HF In streaming{e}") -asyncio.run(test_text_completion_async_stream()) \ No newline at end of file +# asyncio.run(test_text_completion_async_stream()) + +def test_async_text_completion(): + litellm.set_verbose = True + print('test_async_text_completion') + async def test_get_response(): + try: + response = await litellm.atext_completion( + model="gpt-3.5-turbo-instruct", + prompt="good morning", + stream=False, + max_tokens=10 + ) + print(f"response: {response}") + except litellm.Timeout as e: + print(e) + except Exception as e: + print(e) + + asyncio.run(test_get_response()) +test_async_text_completion() \ No newline at end of file