Skip to content

Commit 48c2aa1

Browse files
Copilotowndev
andcommitted
Add version compatibility for aclose() method with fallback
Co-authored-by: owndev <69784886+owndev@users.noreply.github.com>
1 parent 8b858c1 commit 48c2aa1

1 file changed

Lines changed: 20 additions & 6 deletions

File tree

pipelines/google/google_gemini.py

Lines changed: 20 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1621,8 +1621,14 @@ async def _handle_streaming_response_with_cleanup(
16211621
finally:
16221622
# Ensure client is properly closed after streaming completes
16231623
try:
1624-
await client.aio.aclose()
1625-
self.log.debug("Streaming client closed successfully")
1624+
# Use hasattr to support both old and new SDK versions
1625+
if hasattr(client.aio, "aclose"):
1626+
await client.aio.aclose()
1627+
self.log.debug("Streaming client closed successfully")
1628+
elif hasattr(client, "close"):
1629+
# Fallback to sync close if async not available
1630+
client.close()
1631+
self.log.debug("Streaming client closed successfully (sync)")
16261632
except Exception as close_error:
16271633
self.log.warning(f"Error closing streaming client: {close_error}")
16281634

@@ -2210,10 +2216,18 @@ async def get_response():
22102216
# Ensure client is properly closed for non-streaming responses
22112217
if not stream or supports_image_generation:
22122218
try:
2213-
await client.aio.aclose()
2214-
self.log.debug(
2215-
f"Request {request_id}: Client closed successfully"
2216-
)
2219+
# Use hasattr to support both old and new SDK versions
2220+
if hasattr(client.aio, "aclose"):
2221+
await client.aio.aclose()
2222+
self.log.debug(
2223+
f"Request {request_id}: Client closed successfully"
2224+
)
2225+
elif hasattr(client, "close"):
2226+
# Fallback to sync close if async not available
2227+
client.close()
2228+
self.log.debug(
2229+
f"Request {request_id}: Client closed successfully (sync)"
2230+
)
22172231
except Exception as close_error:
22182232
self.log.warning(f"Error closing client: {close_error}")
22192233

0 commit comments

Comments
 (0)