XFE Git
XFE Studio Git
Git 首页 全局搜索
XFE 主站 文档 NuGet
公开
关注 0 Fork 0 Star 1
返回提交历史

XFEstudio/gpt4free

Improve cli argparser, Improve curl_cffi websocket support

fdc8b88d
Heiner Lohaus <hlohaus@users.noreply.github.com>
提交于

代码差异

2 个文件 +20 -18
Modified g4f/cli.py +11 -11
@@ -1,32 +1,32 @@
1 1 import argparse
2 from enum import Enum
3 2
4 import g4f
5 3 from g4f import Provider
6
7 4 from g4f.gui.run import gui_parser, run_gui_args
8 5
9 def run_gui(args):
10 print("Running GUI...")
11
12 6 def main():
13 IgnoredProviders = Enum("ignore_providers", {key: key for key in Provider.__all__})
14 7 parser = argparse.ArgumentParser(description="Run gpt4free")
15 8 subparsers = parser.add_subparsers(dest="mode", help="Mode to run the g4f in.")
16 api_parser=subparsers.add_parser("api")
9 api_parser = subparsers.add_parser("api")
17 10 api_parser.add_argument("--bind", default="0.0.0.0:1337", help="The bind string.")
18 11 api_parser.add_argument("--debug", action="store_true", help="Enable verbose logging.")
19 12 api_parser.add_argument("--workers", type=int, default=None, help="Number of workers.")
20 13 api_parser.add_argument("--disable_colors", action="store_true", help="Don't use colors.")
21 api_parser.add_argument("--ignored-providers", nargs="+", choices=[provider.name for provider in IgnoredProviders],
14 api_parser.add_argument("--ignored-providers", nargs="+", choices=[provider for provider in Provider.__map__],
22 15 default=[], help="List of providers to ignore when processing request.")
23 16 subparsers.add_parser("gui", parents=[gui_parser()], add_help=False)
24 17
25 18 args = parser.parse_args()
26 19 if args.mode == "api":
27 20 import g4f.api
28 g4f.api.api.set_list_ignored_providers(args.ignored_providers)
29 g4f.api.run_api(bind=args.bind, debug=args.debug, workers=args.workers, use_colors=not args.disable_colors)
21 g4f.api.api.set_list_ignored_providers(
22 args.ignored_providers
23 )
24 g4f.api.run_api(
25 bind=args.bind,
26 debug=args.debug,
27 workers=args.workers,
28 use_colors=not args.disable_colors
29 )
30 30 elif args.mode == "gui":
31 31 run_gui_args(args)
32 32 else:
Modified g4f/requests/curl_cffi.py +9 -7
@@ -79,10 +79,10 @@ class StreamSession(AsyncSession):
79 79 return StreamResponse(super().request(method, url, stream=True, **kwargs))
80 80
81 81 def ws_connect(self, url, *args, **kwargs):
82 return WebSocket(self, url)
82 return WebSocket(self, url, **kwargs)
83 83
84 def _ws_connect(self, url):
85 return super().ws_connect(url)
84 def _ws_connect(self, url, **kwargs):
85 return super().ws_connect(url, **kwargs)
86 86
87 87 # Defining HTTP methods as partial methods of the request method.
88 88 head = partialmethod(request, "HEAD")
@@ -102,20 +102,22 @@ else:
102 102 raise RuntimeError("CurlMimi in curl_cffi is missing | pip install -U g4f[curl_cffi]")
103 103
104 104 class WebSocket():
105 def __init__(self, session, url) -> None:
105 def __init__(self, session, url, **kwargs) -> None:
106 106 if not has_curl_ws:
107 107 raise RuntimeError("CurlWsFlag in curl_cffi is missing | pip install -U g4f[curl_cffi]")
108 108 self.session: StreamSession = session
109 109 self.url: str = url
110 del kwargs["autoping"]
111 self.options: dict = kwargs
110 112
111 113 async def __aenter__(self):
112 self.inner = await self.session._ws_connect(self.url)
114 self.inner = await self.session._ws_connect(self.url, **self.options)
113 115 return self
114 116
115 117 async def __aexit__(self, *args):
116 self.inner.aclose()
118 await self.inner.aclose()
117 119
118 async def receive_str(self) -> str:
120 async def receive_str(self, **kwargs) -> str:
119 121 bytes, _ = await self.inner.arecv()
120 122 return bytes.decode(errors="ignore")
121 123