Fix streaming request handling for GoogleProviderHandler

parent 8a701a57
...@@ -153,21 +153,29 @@ class GoogleProviderHandler(BaseProviderHandler): ...@@ -153,21 +153,29 @@ class GoogleProviderHandler(BaseProviderHandler):
# Handle streaming request # Handle streaming request
if stream: if stream:
logging.info(f"GoogleProviderHandler: Using streaming API") logging.info(f"GoogleProviderHandler: Using streaming API")
# Create a new client instance for each streaming request to ensure it remains open # Create a new client instance for streaming to ensure it stays open
# This prevents "Cannot send a request, as the client has been closed" errors
from google import genai from google import genai
stream_client = genai.Client(api_key=self.api_key) stream_client = genai.Client(api_key=self.api_key)
response = stream_client.models.generate_content_stream(
# We need to iterate over the streaming response immediately without yielding control
# to ensure the client stays alive
chunks = []
for chunk in stream_client.models.generate_content_stream(
model=model, model=model,
contents=content, contents=content,
config=config config=config
) ):
logging.info(f"GoogleProviderHandler: Streaming response received") chunks.append(chunk)
logging.info(f"GoogleProviderHandler: Streaming response received (total chunks: {len(chunks)})")
self.record_success() self.record_success()
# Return the synchronous iterator directly # Now yield chunks asynchronously
# The handler will iterate over it and convert to OpenAI format async def async_generator():
return response for chunk in chunks:
yield chunk
return async_generator()
else: else:
# Non-streaming request # Non-streaming request
# Generate content using the google-genai client # Generate content using the google-genai client
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment