From 23218c4aa398af9fec97fa44c9e7a9461251a7f5 Mon Sep 17 00:00:00 2001 From: hlohaus <983577+hlohaus@users.noreply.github.com> Date: Fri, 31 Oct 2025 14:17:41 +0100 Subject: [PATCH] Refactor Copilot and PollinationsAI classes for improved error handling and timeout adjustments; add rate limiting in API class based on user IP. --- g4f/Provider/Copilot.py | 4 +- g4f/Provider/PollinationsAI.py | 74 ++-------------------------------- g4f/api/__init__.py | 27 ++++++++++++- 3 files changed, 30 insertions(+), 75 deletions(-) diff --git a/g4f/Provider/Copilot.py b/g4f/Provider/Copilot.py index a93838e1..1e39b25f 100644 --- a/g4f/Provider/Copilot.py +++ b/g4f/Provider/Copilot.py @@ -284,7 +284,7 @@ class Copilot(AsyncAuthedProvider, ProviderModelMixin): sources = {} while not wss.closed: try: - msg_txt, _ = await asyncio.wait_for(wss.recv(), 3 if done else timeout) + msg_txt, _ = await asyncio.wait_for(wss.recv(), 1 if done else timeout) msg = json.loads(msg_txt) except: break @@ -369,7 +369,7 @@ async def get_access_token_and_cookies(url: str, proxy: str = None, needs_auth: button = await page.select("[data-testid=\"submit-button\"]") if button: await button.click() - turnstile = await page.select('#cf-turnstile', 300) + turnstile = await page.select('#cf-turnstile') if turnstile: debug.log("Found Element: 'cf-turnstile'") await asyncio.sleep(3) diff --git a/g4f/Provider/PollinationsAI.py b/g4f/Provider/PollinationsAI.py index bb205314..23ce1677 100644 --- a/g4f/Provider/PollinationsAI.py +++ b/g4f/Provider/PollinationsAI.py @@ -31,37 +31,6 @@ DEFAULT_HEADERS = { "origin": "https://pollinations.ai", } -FOLLOWUPS_TOOLS = [{ - "type": "function", - "function": { - "name": "options", - "description": "Provides options for the conversation", - "parameters": { - "properties": { - "title": { - "title": "Conversation title. Prefixed with one or more emojies", - "type": "string" - }, - "followups": { - "items": { - "type": "string" - }, - "title": "Suggested 4 Followups (only user messages)", - "type": "array" - } - }, - "title": "Conversation", - "type": "object" - } - } -}] - -FOLLOWUPS_DEVELOPER_MESSAGE = [{ - "role": "developer", - "content": "Provide conversation options.", -}] - - class PollinationsAI(AsyncGeneratorProvider, ProviderModelMixin): label = "Pollinations AI 🌸" url = "https://pollinations.ai" @@ -375,12 +344,12 @@ class PollinationsAI(AsyncGeneratorProvider, ProviderModelMixin): await raise_for_status(response) except Exception as e: responses.add(e) - debug.error(f"Error fetching image: {e}") + debug.error(f"Error fetching image:", e) if response.headers['content-type'].startswith("image/"): responses.add(ImageResponse(str(response.url), prompt, {"headers": headers})) else: t_ = await response.text() - debug.error(f"UnHandel Error fetching image: {t_}") + debug.error(f"UnHandel Error fetching image:", t_) responses.add(t_) tasks: list[asyncio.Task] = [] @@ -465,43 +434,6 @@ class PollinationsAI(AsyncGeneratorProvider, ProviderModelMixin): async with session.post(cls.openai_endpoint, json=data, headers=headers) as response: if response.status in (400, 500): debug.error(f"Error: {response.status} - Bad Request: {data}") - full_resposne = [] async for chunk in read_response(response, stream, format_media_prompt(messages), cls.get_dict(), kwargs.get("download_media", True)): - if isinstance(chunk, str): - full_resposne.append(chunk) - yield chunk - if full_resposne: - full_content = "".join(full_resposne) - if kwargs.get("action") == "next" and model != "evil": - tool_messages = [] - for message in messages: - if message.get("role") == "user": - if isinstance(message.get("content"), str): - tool_messages.append({"role": "user", "content": message.get("content")}) - elif isinstance(message.get("content"), list): - next_value = message.get("content").pop() - if isinstance(next_value, dict): - next_value = next_value.get("text") - if next_value: - tool_messages.append({"role": "user", "content": next_value}) - tool_messages.append({"role": "assistant", "content": full_content}) - data = { - "model": "openai", - "messages": tool_messages + FOLLOWUPS_DEVELOPER_MESSAGE, - "tool_choice": "required", - "tools": FOLLOWUPS_TOOLS - } - async with session.post(cls.openai_endpoint, json=data, headers=headers) as response: - try: - await raise_for_status(response) - tool_calls = (await response.json()).get("choices", [{}])[0].get("message", {}).get( - "tool_calls", []) - if tool_calls: - arguments = json.loads(tool_calls.pop().get("function", {}).get("arguments")) - if arguments.get("title"): - yield TitleGeneration(arguments.get("title")) - if arguments.get("followups"): - yield SuggestedFollowups(arguments.get("followups")) - except Exception as e: - debug.error("Error generating title and followups:", e) + yield chunk \ No newline at end of file diff --git a/g4f/api/__init__.py b/g4f/api/__init__.py index 30c8e47b..00c32f3e 100644 --- a/g4f/api/__init__.py +++ b/g4f/api/__init__.py @@ -65,7 +65,7 @@ from g4f.client.helper import filter_none from g4f.config import DEFAULT_PORT, DEFAULT_TIMEOUT, DEFAULT_STREAM_TIMEOUT from g4f.image import EXTENSIONS_MAP, is_data_an_media, process_image from g4f.image.copy_images import get_media_dir, copy_media, get_source_url -from g4f.errors import ProviderNotFoundError, ModelNotFoundError, MissingAuthError, NoValidHarFileError, MissingRequirementsError +from g4f.errors import ProviderNotFoundError, ModelNotFoundError, MissingAuthError, NoValidHarFileError, MissingRequirementsError, RateLimitError from g4f.cookies import read_cookie_files, get_cookies_dir from g4f.providers.types import ProviderType from g4f.providers.response import AudioResponse @@ -417,6 +417,8 @@ class Api: }) return ErrorResponse.from_message("The model does not exist.", HTTP_404_NOT_FOUND) + most_wanted = {} + failure_counts = {} responses = { HTTP_200_OK: {"model": ChatCompletion}, HTTP_401_UNAUTHORIZED: {"model": ErrorResponseModel}, @@ -433,8 +435,29 @@ class Api: provider: str = None, conversation_id: str = None, x_user: Annotated[str | None, Header()] = None, - cf_ipcountry: Annotated[str | None, Header()] = None + cf_ipcountry: Annotated[str | None, Header()] = None, + x_forwarded_for: Annotated[str | None, Header()] = None ): + if AppConfig.demo and x_forwarded_for is not None: + current_most_wanted = next(iter(most_wanted.values()), 0) + is_most_wanted = False + if x_forwarded_for in most_wanted: + if failure_counts.get(x_forwarded_for, 0) > 0: + failure_counts[x_forwarded_for] -= 1 + most_wanted[x_forwarded_for] += 1 + elif most_wanted[x_forwarded_for] >= current_most_wanted: + if x_forwarded_for not in failure_counts: + failure_counts[x_forwarded_for] = 0 + failure_counts[x_forwarded_for] += 1 + is_most_wanted = True + else: + most_wanted[x_forwarded_for] += 1 + else: + most_wanted[x_forwarded_for] = 1 + sorted_most_wanted = dict(sorted(most_wanted.items(), key=lambda item: item[1], reverse=True)) + debug.log(f"Most wanted IPs: {sorted_most_wanted}") + if is_most_wanted: + raise RateLimitError("You are most wanted! Please wait before making another request.") if provider is not None and provider not in Provider.__map__: if provider in model_map: config.model = provider