Thanks to visit codestin.com
Credit goes to github.com

Skip to content

🐛 fixed asyncio streaming bugs #584

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Sep 26, 2023
Merged

🐛 fixed asyncio streaming bugs #584

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
41 changes: 30 additions & 11 deletions openai/api_requestor.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,11 +6,10 @@
import threading
import time
import warnings
from contextlib import asynccontextmanager
from json import JSONDecodeError
from typing import (
AsyncContextManager,
AsyncGenerator,
AsyncIterator,
Callable,
Dict,
Iterator,
Expand Down Expand Up @@ -366,8 +365,9 @@ async def arequest(
request_id: Optional[str] = None,
request_timeout: Optional[Union[float, Tuple[float, float]]] = None,
) -> Tuple[Union[OpenAIResponse, AsyncGenerator[OpenAIResponse, None]], bool, str]:
ctx = aiohttp_session()
ctx = AioHTTPSession()
session = await ctx.__aenter__()
result = None
try:
result = await self.arequest_raw(
method.lower(),
Expand All @@ -381,6 +381,9 @@ async def arequest(
)
resp, got_stream = await self._interpret_async_response(result, stream)
except Exception:
# Close the request before exiting session context.
if result is not None:
result.release()
await ctx.__aexit__(None, None, None)
raise
if got_stream:
Expand All @@ -391,10 +394,15 @@ async def wrap_resp():
async for r in resp:
yield r
finally:
# Close the request before exiting session context. Important to do it here
# as if stream is not fully exhausted, we need to close the request nevertheless.
result.release()
await ctx.__aexit__(None, None, None)

return wrap_resp(), got_stream, self.api_key
else:
# Close the request before exiting session context.
result.release()
await ctx.__aexit__(None, None, None)
return resp, got_stream, self.api_key

Expand Down Expand Up @@ -768,11 +776,22 @@ def _interpret_response_line(
return resp


@asynccontextmanager
async def aiohttp_session() -> AsyncIterator[aiohttp.ClientSession]:
user_set_session = openai.aiosession.get()
if user_set_session:
yield user_set_session
else:
async with aiohttp.ClientSession() as session:
yield session
class AioHTTPSession(AsyncContextManager):
def __init__(self):
self._session = None
self._should_close_session = False

async def __aenter__(self):
self._session = openai.aiosession.get()
if self._session is None:
self._session = await aiohttp.ClientSession().__aenter__()
self._should_close_session = True

return self._session

async def __aexit__(self, exc_type, exc_value, traceback):
if self._session is None:
raise RuntimeError("Session is not initialized")

if self._should_close_session:
await self._session.__aexit__(exc_type, exc_value, traceback)