From b7624b75a3ed6a48733ddaa0c6ea87227922d173 Mon Sep 17 00:00:00 2001 From: Heiner Lohaus Date: Tue, 21 May 2024 20:39:31 +0200 Subject: Add streaming and conversation support to gemini --- g4f/client/async_client.py | 4 ++++ g4f/client/client.py | 4 ++++ 2 files changed, 8 insertions(+) (limited to 'g4f/client') diff --git a/g4f/client/async_client.py b/g4f/client/async_client.py index dbfa6b70..2fe4640b 100644 --- a/g4f/client/async_client.py +++ b/g4f/client/async_client.py @@ -18,6 +18,7 @@ from ..Provider import ProviderUtils from ..typing import Union, Messages, AsyncIterator, ImageType from ..errors import NoImageResponseError, ProviderNotFoundError from ..requests.aiohttp import get_connector +from ..providers.conversation import BaseConversation from ..image import ImageResponse as ImageProviderResponse, ImageDataResponse try: @@ -42,6 +43,9 @@ async def iter_response( if isinstance(chunk, FinishReason): finish_reason = chunk.reason break + elif isinstance(chunk, BaseConversation): + yield chunk + continue content += str(chunk) count += 1 if max_tokens is not None and count >= max_tokens: diff --git a/g4f/client/client.py b/g4f/client/client.py index acf53c70..63bae4fe 100644 --- a/g4f/client/client.py +++ b/g4f/client/client.py @@ -6,6 +6,7 @@ import string from ..typing import Union, Iterator, Messages, ImageType from ..providers.types import BaseProvider, ProviderType, FinishReason +from ..providers.conversation import BaseConversation from ..image import ImageResponse as ImageProviderResponse from ..errors import NoImageResponseError from .stubs import ChatCompletion, ChatCompletionChunk, Image, ImagesResponse @@ -29,6 +30,9 @@ def iter_response( if isinstance(chunk, FinishReason): finish_reason = chunk.reason break + elif isinstance(chunk, BaseConversation): + yield chunk + continue content += str(chunk) if max_tokens is not None and idx + 1 >= max_tokens: finish_reason = "length" -- cgit v1.2.3