Fix AsyncOpenAI "RuntimeError: Event loop is closed bug" when instances of AsyncOpenAI are rapidly created & destroyed

llama-index-integrations/llms/llama-index-llms-openai/llama_index/llms/openai/base.py

@@ -1,4 +1,6 @@
+import json
 from typing import (
+    TYPE_CHECKING,
     Any,
     Awaitable,
     Callable,
@@ -11,14 +13,20 @@
     cast,
     get_args,
     runtime_checkable,
-    TYPE_CHECKING,
 )
 
-from llama_index.core.llms.llm import ToolSelection
-from llama_index.core.llms.function_calling import FunctionCallingLLM
-import json
 import httpx
 import tiktoken
+from llama_index.core.base.llms.generic_utils import (
+    achat_to_completion_decorator,
+    acompletion_to_chat_decorator,
+    astream_chat_to_completion_decorator,
+    astream_completion_to_chat_decorator,
+    chat_to_completion_decorator,
+    completion_to_chat_decorator,
+    stream_chat_to_completion_decorator,
+    stream_completion_to_chat_decorator,
+)
 from llama_index.core.base.llms.types import (
     ChatMessage,
     ChatResponse,
@@ -39,23 +47,15 @@
     llm_chat_callback,
     llm_completion_callback,
 )
-from llama_index.core.base.llms.generic_utils import (
-    achat_to_completion_decorator,
-    acompletion_to_chat_decorator,
-    astream_chat_to_completion_decorator,
-    astream_completion_to_chat_decorator,
-    chat_to_completion_decorator,
-    completion_to_chat_decorator,
-    stream_chat_to_completion_decorator,
-    stream_completion_to_chat_decorator,
-)
+from llama_index.core.llms.function_calling import FunctionCallingLLM
+from llama_index.core.llms.llm import ToolSelection
 from llama_index.core.types import BaseOutputParser, PydanticProgramMode
 from llama_index.llms.openai.utils import (
     OpenAIToolCall,
     create_retry_decorator,
+    from_openai_completion_logprobs,
     from_openai_message,
     from_openai_token_logprobs,
-    from_openai_completion_logprobs,
     is_chat_model,
     is_function_calling_model,
     openai_modelname_to_contextsize,
@@ -101,7 +101,8 @@ def force_single_tool_call(response: ChatResponse) -> None:
 
 
 class OpenAI(FunctionCallingLLM):
-    """OpenAI LLM.
+    """
+    OpenAI LLM.
 
     Examples:
         `pip install llama-index-llms-openai`
@@ -359,11 +360,21 @@ def _get_model_kwargs(self, **kwargs: Any) -> Dict[str, Any]:
     def _chat(self, messages: Sequence[ChatMessage], **kwargs: Any) -> ChatResponse:
         client = self._get_client()
         message_dicts = to_openai_message_dicts(messages)
-        response = client.chat.completions.create(
-            messages=message_dicts,
-            stream=False,
-            **self._get_model_kwargs(**kwargs),
-        )
+
+        if self.reuse_client:
+            response = client.chat.completions.create(
+                messages=message_dicts,
+                stream=False,
+                **self._get_model_kwargs(**kwargs),
+            )
+        else:
+            with client:
+                response = client.chat.completions.create(
+                    messages=message_dicts,
+                    stream=False,
+                    **self._get_model_kwargs(**kwargs),
+                )
+
         openai_message = response.choices[0].message
         message = from_openai_message(openai_message)
         openai_token_logprobs = response.choices[0].logprobs
@@ -383,7 +394,8 @@ def _update_tool_calls(
         tool_calls: List[ChoiceDeltaToolCall],
         tool_calls_delta: Optional[List[ChoiceDeltaToolCall]],
     ) -> List[ChoiceDeltaToolCall]:
-        """Use the tool_calls_delta objects received from openai stream chunks
+        """
+        Use the tool_calls_delta objects received from openai stream chunks
         to update the running tool_calls object.
 
         Args:
@@ -483,11 +495,19 @@ def _complete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
         all_kwargs = self._get_model_kwargs(**kwargs)
         self._update_max_tokens(all_kwargs, prompt)
 
-        response = client.completions.create(
-            prompt=prompt,
-            stream=False,
-            **all_kwargs,
-        )
+        if self.reuse_client:
+            response = client.completions.create(
+                prompt=prompt,
+                stream=False,
+                **all_kwargs,
+            )
+        else:
+            with client:
+                response = client.completions.create(
+                    prompt=prompt,
+                    stream=False,
+                    **all_kwargs,
+                )
         text = response.choices[0].text
 
         openai_completion_logprobs = response.choices[0].logprobs
@@ -617,16 +637,30 @@ async def _achat(
     ) -> ChatResponse:
         aclient = self._get_aclient()
         message_dicts = to_openai_message_dicts(messages)
-        response = await aclient.chat.completions.create(
-            messages=message_dicts, stream=False, **self._get_model_kwargs(**kwargs)
-        )
-        message_dict = response.choices[0].message
-        message = from_openai_message(message_dict)
-        logprobs_dict = response.choices[0].logprobs
+
+        if self.reuse_client:
+            response = await aclient.chat.completions.create(
+                messages=message_dicts, stream=False, **self._get_model_kwargs(**kwargs)
+            )
+        else:
+            async with aclient:
+                response = await aclient.chat.completions.create(
+                    messages=message_dicts,
+                    stream=False,
+                    **self._get_model_kwargs(**kwargs),
+                )
+
+        openai_message = response.choices[0].message
+        message = from_openai_message(openai_message)
+        openai_token_logprobs = response.choices[0].logprobs
+        logprobs = None
+        if openai_token_logprobs and openai_token_logprobs.content:
+            logprobs = from_openai_token_logprobs(openai_token_logprobs.content)
 
         return ChatResponse(
             message=message,
             raw=response,
+            logprobs=logprobs,
             additional_kwargs=self._get_response_token_counts(response),
         )
 
@@ -700,11 +734,19 @@ async def _acomplete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
         all_kwargs = self._get_model_kwargs(**kwargs)
         self._update_max_tokens(all_kwargs, prompt)
 
-        response = await aclient.completions.create(
-            prompt=prompt,
-            stream=False,
-            **all_kwargs,
-        )
+        if self.reuse_client:
+            response = await aclient.completions.create(
+                prompt=prompt,
+                stream=False,
+                **all_kwargs,
+            )
+        else:
+            async with aclient:
+                response = await aclient.completions.create(
+                    prompt=prompt,
+                    stream=False,
+                    **all_kwargs,
+                )
 
         text = response.choices[0].text
         openai_completion_logprobs = response.choices[0].logprobs

llama-index-integrations/llms/llama-index-llms-openai/pyproject.toml

@@ -29,7 +29,7 @@ exclude = ["**/BUILD"]
 license = "MIT"
 name = "llama-index-llms-openai"
 readme = "README.md"
-version = "0.1.15"
+version = "0.1.16"
 
 [tool.poetry.dependencies]
 python = ">=3.8.1,<4.0"