gpt4free

公开

关注 0 Fork 0 Star 0

XFEstudio/gpt4free

Add websocket media streaming for OpenaiChat

Introduces the wss_media method to stream media updates via websocket in OpenaiChat, and updates logic to yield media as it becomes available. Also adds wait_media as a fallback polling method, tracks image generation tasks in Conversation, and fixes a bug in curl_cffi.py when deleting the 'autoping' key from kwargs.

e5ca0221

Ammar <ammar.alkotb@gmail.com>

提交于 2025-12-06 16:34

父提交 027d486

代码差异

2 个文件 +158 -3

Modified g4f/Provider/needs_auth/OpenaiChat.py +156 -2

@@ -10,7 +10,9 @@ import re

10 10 import time

11 11 import uuid

12 12 from copy import copy

13 from typing import AsyncIterator, Iterator, Optional, Generator, Dict, Union, List, Any

13 from typing import AsyncIterator, Iterator, Optional, Generator, Dict, Union, List, Any, AsyncGenerator, Set

15 from curl_cffi import AsyncSession

14 16

15 17 try:

16 18 import nodriver

@@ -341,7 +343,7 @@ class OpenaiChat(AsyncAuthedProvider, ProviderModelMixin):

341 343 debug.error(e)

342 344 if download_urls:

343 345 # status = None, finished_successfully

344 if is_sediment and status is None:

346 if is_sediment and status != "finished_successfully":

345 347 return ImagePreview(download_urls, prompt, {"status": status, "headers": auth_result.headers})

346 348 else:

347 349 return ImageResponse(download_urls, prompt, {"status": status, "headers": auth_result.headers})

@@ -703,8 +705,157 @@ class OpenaiChat(AsyncAuthedProvider, ProviderModelMixin):

703 705 await asyncio.sleep(5)

704 706 else:

705 707 break

708

709 if conversation.task and kwargs.get("wait_media", True):

710 async for _m in cls.wss_media(session, conversation, auth_result.headers, auth_result):

711 yield _m

712 # if kwargs.get("wait_media"):

713 # async for _m in cls.wait_media(session, conversation, headers, auth_result):

714 # yield _m

715

706 716 yield FinishReason(conversation.finish_reason)

707 717

718 @classmethod

719 async def wss_media(

720 cls,

721 _session,

722 conversation: Conversation,

723 headers: Dict[str, str],

724 auth_result: AuthResult,

725 timeout: Optional[int] = 20,

726 ):

727 seen_assets: Set[str] = set()

728 async with AsyncSession(

729 timeout=timeout,

730 impersonate="chrome",

731 headers=headers,

732 cookies=auth_result.cookies

733 ) as session:

734 response = await session.get(

735 "https://chatgpt.com/backend-api/celsius/ws/user",

736 headers=headers,

737 )

738 response.raise_for_status()

739 websocket_url = response.json().get("websocket_url")

740 started = False

741 wss = await session.ws_connect(websocket_url, timeout=3)

742 while not wss.closed:

743 try:

744 last_msg = await wss.recv_json(timeout=60 if not started else timeout)

745 except:

746 break

747 conversation_id = conversation.task.get("conversation_id")

748 message_id = conversation.task.get("message", {}).get("id")

749 if isinstance(last_msg, dict) and last_msg.get("type") == "conversation-update":

750 if last_msg.get("payload", {}).get("conversation_id") != conversation_id:

751 continue

752

753 message = last_msg.get("payload", {}).get("update_content", {}).get("message", {})

754 if message.get("id") != message_id:

755 continue

756

757 # if last_msg.get("payload", {}).get("update_type") == 'async-task-start':

758 # started = True

759 started = True

760 if last_msg.get("payload", {}).get("update_type") == 'async-task-update-message':

761

762 status = message.get("status")

763 parts = message.get("content").get("parts") or []

764 for part in parts:

765 if part.get("content_type") != "image_asset_pointer":

766 continue

767 asset = part.get("asset_pointer")

768 if not asset or asset in seen_assets:

769 continue

770 seen_assets.add(asset)

771 generated_images = await cls.get_generated_image(

772 _session,

773 auth_result,

774 asset,

775 conversation.prompt or "",

776 conversation.conversation_id,

777 status,

778 )

779 if generated_images is not None:

780 yield generated_images

781 if message.get("status") == "finished_successfully":

782 await wss.close()

783 return

784

785 @classmethod

786 async def wait_media(

787 cls,

788 session,

789 conversation,

790 headers: Dict[str, str],

791 auth_result: AuthResult,

792 poll_interval: int = 10,

793 timeout: Optional[int] = None,

794 ) -> AsyncGenerator[Any, None]:

795 start_time = asyncio.get_event_loop().time()

796 seen_assets: Set[str] = set()

797 running = True

798 has_image_task = False

799 generation_started = False

800

801 while running:

802 if timeout is not None:

803 elapsed = asyncio.get_event_loop().time() - start_time

804 if elapsed > timeout:

805 return

806 # https://chatgpt.com/backend-api/tasks

807 async with session.get(

808 f"https://chatgpt.com/backend-api/conversation/{conversation.conversation_id}",

809 headers=headers,

810 ) as response:

811 await raise_for_status(response)

812 data = await response.json()

813

814 mapping = data.get("mapping") or {}

815 if not mapping:

816 return

817

818 last_node = list(mapping.values())[-1] or {}

819 last_message = last_node.get("message") or {}

820 metadata = last_message.get("metadata") or {}

821 status = last_message.get("status")

822 image_task_id = metadata.get("image_gen_task_id")

823 if not has_image_task and not image_task_id:

824 return

825

826 if image_task_id and not has_image_task:

827 debug.log(f"OpenaiChat: Wait Task: {image_task_id}")

828 has_image_task = True

829 if status == "in_progress":

830 generation_started = True

831 elif generation_started and status == "finished_successfully":

832 running = False

833 if generation_started:

834 content = last_message.get("content") or {}

835 parts = content.get("parts") or []

836 for part in parts:

837 if part.get("content_type") != "image_asset_pointer":

838 continue

839 asset = part.get("asset_pointer")

840 if not asset or asset in seen_assets:

841 continue

842 seen_assets.add(asset)

843 generated_images = await cls.get_generated_image(

844 session,

845 auth_result,

846 asset,

847 conversation.prompt

848 or metadata.get("async_task_title")

849 or "",

850 conversation.conversation_id,

851 status,

852 )

853 if generated_images is not None:

854 yield generated_images

855 if generation_started and status == "finished_successfully":

856 return

857 await asyncio.sleep(poll_interval)

858

708 859 @classmethod

709 860 async def iter_messages_line(cls, session: StreamSession, auth_result: AuthResult, line: bytes,

710 861 fields: Conversation, sources: OpenAISources,

@@ -850,6 +1001,8 @@ class OpenaiChat(AsyncAuthedProvider, ProviderModelMixin):

850 1001 if fields.parent_message_id is None:

851 1002 fields.parent_message_id = v.get("message", {}).get("id")

852 1003 fields.message_id = v.get("message", {}).get("id")

1004 if m.get("status") == "finished_successfully" and m.get("metadata", {}).get("image_gen_task_id"):

1005 fields.task = v

853 1006 return

854 1007 if "error" in line and line.get("error"):

855 1008 raise RuntimeError(line.get("error"))

@@ -1046,6 +1199,7 @@ class Conversation(JsonConversation):

1046 1199 self.thoughts_summary = ""

1047 1200 self.prompt = None

1048 1201 self.generated_images: ImagePreview = None

1202 self.task: dict = None

1049 1203

1050 1204

1051 1205 def get_cookies(

Modified g4f/requests/curl_cffi.py +2 -1

@@ -148,7 +148,8 @@ if has_curl_cffi and has_curl_ws:

148 148 def __init__(self, session, url, **kwargs) -> None:

149 149 self.session: StreamSession = session

150 150 self.url: str = url

151 del kwargs["autoping"]

151 if "autoping" in kwargs:

152 del kwargs["autoping"]

152 153 self.options: dict = kwargs

153 154

154 155 async def __aenter__(self):