refactor(g4f/api/__init__.py): refactor API structure and improve async handling

@@ -14,17 +14,18 @@ from starlette.status import HTTP_422_UNPROCESSABLE_ENTITY, HTTP_401_UNAUTHORIZE

14 14 from fastapi.encoders import jsonable_encoder

15 15 from fastapi.middleware.cors import CORSMiddleware

16 16 from pydantic import BaseModel

17 from typing import Union, Optional

17 from typing import Union, Optional, Iterator

18 18

19 19 import g4f

20 20 import g4f.debug

21 from g4f.client import Client

21 from g4f.client import Client, ChatCompletion, ChatCompletionChunk, ImagesResponse

22 22 from g4f.typing import Messages

23 23 from g4f.cookies import read_cookie_files

24 24

25 def create_app():

25 def create_app(g4f_api_key: str = None):

26 26 app = FastAPI()

27 api = Api(app)

27

28 # Add CORS middleware

28 29 app.add_middleware(

29 30 CORSMiddleware,

30 31 allow_origin_regex=".*",

@@ -32,18 +33,19 @@ def create_app():

32 33 allow_methods=["*"],

33 34 allow_headers=["*"],

34 35 )

36

37 api = Api(app, g4f_api_key=g4f_api_key)

35 38 api.register_routes()

36 39 api.register_authorization()

37 40 api.register_validation_exception_handler()

41

42 # Read cookie files if not ignored

38 43 if not AppConfig.ignore_cookie_files:

39 44 read_cookie_files()

40 return app

41 45

42 def create_app_debug():

43 g4f.debug.logging = True

44 return create_app()

46 return app

45 47

46 class ChatCompletionsForm(BaseModel):

48 class ChatCompletionsConfig(BaseModel):

47 49 messages: Messages

48 50 model: str

49 51 provider: Optional[str] = None

@@ -55,15 +57,12 @@ class ChatCompletionsForm(BaseModel):

55 57 web_search: Optional[bool] = None

56 58 proxy: Optional[str] = None

57 59

58 class ImagesGenerateForm(BaseModel):

59 model: Optional[str] = None

60 provider: Optional[str] = None

60 class ImageGenerationConfig(BaseModel):

61 61 prompt: str

62 response_format: Optional[str] = None

63 api_key: Optional[str] = None

64 proxy: Optional[str] = None

62 model: Optional[str] = None

63 response_format: str = "url"

65 64

66 class AppConfig():

65 class AppConfig:

67 66 ignored_providers: Optional[list[str]] = None

68 67 g4f_api_key: Optional[str] = None

69 68 ignore_cookie_files: bool = False

@@ -74,16 +73,23 @@ class AppConfig():

74 73 for key, value in data.items():

75 74 setattr(cls, key, value)

76 75

76 list_ignored_providers: list[str] = None

77

78 def set_list_ignored_providers(ignored: list[str]):

79 global list_ignored_providers

80 list_ignored_providers = ignored

81

77 82 class Api:

78 def __init__(self, app: FastAPI) -> None:

83 def __init__(self, app: FastAPI, g4f_api_key=None) -> None:

79 84 self.app = app

80 85 self.client = Client()

86 self.g4f_api_key = g4f_api_key

81 87 self.get_g4f_api_key = APIKeyHeader(name="g4f-api-key")

82 88

83 89 def register_authorization(self):

84 90 @self.app.middleware("http")

85 91 async def authorization(request: Request, call_next):

86 if AppConfig.g4f_api_key and request.url.path in ["/v1/chat/completions", "/v1/completions"]:

92

            if self.g4f_api_key and request.url.path in ["/v1/chat/completions", "/v1/completions", "/v1/images/generate"]:

87 93 try:

88 94 user_g4f_api_key = await self.get_g4f_api_key(request)

89 95 except HTTPException as e:

@@ -92,22 +98,26 @@ class Api:

92 98 status_code=HTTP_401_UNAUTHORIZED,

93 99 content=jsonable_encoder({"detail": "G4F API key required"}),

94 100 )

95 if not secrets.compare_digest(AppConfig.g4f_api_key, user_g4f_api_key):

101 if not secrets.compare_digest(self.g4f_api_key, user_g4f_api_key):

96 102 return JSONResponse(

97 103 status_code=HTTP_403_FORBIDDEN,

98 104 content=jsonable_encoder({"detail": "Invalid G4F API key"}),

99 105 )

100 return await call_next(request)

106

107 response = await call_next(request)

108 return response

101 109

102 110 def register_validation_exception_handler(self):

103 111 @self.app.exception_handler(RequestValidationError)

104 112 async def validation_exception_handler(request: Request, exc: RequestValidationError):

105 113 details = exc.errors()

106 modified_details = [{

107 "loc": error["loc"],

108 "message": error["msg"],

109 "type": error["type"],

110 } for error in details]

114 modified_details = []

115 for error in details:

116 modified_details.append({

117 "loc": error["loc"],

118 "message": error["msg"],

119 "type": error["type"],

120 })

111 121 return JSONResponse(

112 122 status_code=HTTP_422_UNPROCESSABLE_ENTITY,

113 123 content=jsonable_encoder({"detail": modified_details}),

@@ -121,25 +131,23 @@ class Api:

121 131 @self.app.get("/v1")

122 132 async def read_root_v1():

123 133 return HTMLResponse('g4f API: Go to '

124 '<a href="/v1/chat/completions">chat/completions</a> '

125 'or <a href="/v1/models">models</a>.')

134 '<a href="/v1/chat/completions">chat/completions</a>, '

135 '<a href="/v1/models">models</a>, or '

136 '<a href="/v1/images/generate">images/generate</a>.')

126 137

127 138 @self.app.get("/v1/models")

128 139 async def models():

129 model_list = {

130 model: g4f.models.ModelUtils.convert[model]

140 model_list = dict(

141 (model, g4f.models.ModelUtils.convert[model])

131 142 for model in g4f.Model.__all__()

132 }

143 )

133 144 model_list = [{

134 145 'id': model_id,

135 146 'object': 'model',

136 147 'created': 0,

137 148 'owned_by': model.base_provider

138 149 } for model_id, model in model_list.items()]

139 return JSONResponse({

140 "object": "list",

141 "data": model_list,

142 })

150 return JSONResponse(model_list)

143 151

144 152 @self.app.get("/v1/models/{model_name}")

145 153 async def model_info(model_name: str):

@@ -155,7 +163,7 @@ class Api:

155 163 return JSONResponse({"error": "The model does not exist."})

156 164

157 165 @self.app.post("/v1/chat/completions")

158 async def chat_completions(config: ChatCompletionsForm, request: Request = None, provider: str = None):

166 async def chat_completions(config: ChatCompletionsConfig, request: Request = None, provider: str = None):

159 167 try:

160 168 config.provider = provider if config.provider is None else config.provider

161 169 if config.api_key is None and request is not None:

@@ -164,17 +172,27 @@ class Api:

164 172 auth_header = auth_header.split(None, 1)[-1]

165 173 if auth_header and auth_header != "Bearer":

166 174 config.api_key = auth_header

167 # Use the asynchronous create method and await it

168 response = await self.client.chat.completions.async_create(

175

176 # Create the completion response

177 response = self.client.chat.completions.create(

169 178 **{

170 179 **AppConfig.defaults,

171 180 **config.dict(exclude_none=True),

172 181 },

173 182 ignored=AppConfig.ignored_providers

174 183 )

175 if not config.stream:

184

185 # Check if the response is synchronous or asynchronous

186 if isinstance(response, ChatCompletion):

187 # Synchronous response

176 188 return JSONResponse(response.to_json())

177 189

190 if not config.stream:

191 # If the response is an iterator but not streaming, collect the result

192 response_list = list(response) if isinstance(response, Iterator) else [response]

193 return JSONResponse(response_list[0].to_json())

194

195 # Streaming response

178 196 async def streaming():

179 197 try:

180 198 async for chunk in response:

@@ -185,41 +203,38 @@ class Api:

185 203 logging.exception(e)

186 204 yield f'data: {format_exception(e, config)}\n\n'

187 205 yield "data: [DONE]\n\n"

206

188 207 return StreamingResponse(streaming(), media_type="text/event-stream")

189 208

190 209 except Exception as e:

191 210 logging.exception(e)

192 211 return Response(content=format_exception(e, config), status_code=500, media_type="application/json")

193 212

194 @self.app.post("/v1/completions")

195 async def completions():

196 return Response(content=json.dumps({'info': 'Not working yet.'}, indent=4), media_type="application/json")

197

198 @self.app.post("/v1/images/generations")

199 async def images_generate(config: ImagesGenerateForm, request: Request = None, provider: str = None):

213 @self.app.post("/v1/images/generate")

214 async def generate_image(config: ImageGenerationConfig):

200 215 try:

201 config.provider = provider if config.provider is None else config.provider

202 if config.api_key is None and request is not None:

203 auth_header = request.headers.get("Authorization")

204 if auth_header is not None:

205 auth_header = auth_header.split(None, 1)[-1]

206 if auth_header and auth_header != "Bearer":

207 config.api_key = auth_header

208 # Use the asynchronous generate method and await it

209 response = await self.client.images.async_generate(

210 **config.dict(exclude_none=True),

216 response: ImagesResponse = await self.client.images.async_generate(

217 prompt=config.prompt,

218 model=config.model,

219 response_format=config.response_format

211 220 )

212 return JSONResponse(response.to_json())

221 # Convert Image objects to dictionaries

222 response_data = [image.to_dict() for image in response.data]

223 return JSONResponse({"data": response_data})

213 224 except Exception as e:

214 225 logging.exception(e)

215 226 return Response(content=format_exception(e, config), status_code=500, media_type="application/json")

216 227

217 def format_exception(e: Exception, config: ChatCompletionsForm) -> str:

228 @self.app.post("/v1/completions")

229 async def completions():

230 return Response(content=json.dumps({'info': 'Not working yet.'}, indent=4), media_type="application/json")

231

232 def format_exception(e: Exception, config: Union[ChatCompletionsConfig, ImageGenerationConfig]) -> str:

218 233 last_provider = g4f.get_last_provider(True)

219 234 return json.dumps({

220 235 "error": {"message": f"{e.__class__.__name__}: {e}"},

221 "model": last_provider.get("model") if last_provider else config.model,

222 "provider": last_provider.get("name") if last_provider else config.provider

236 "model": last_provider.get("model") if last_provider else getattr(config, 'model', None),

237 "provider": last_provider.get("name") if last_provider else getattr(config, 'provider', None)

223 238 })

224 239

225 240 def run_api(

@@ -228,18 +243,22 @@ def run_api(

228 243 bind: str = None,

229 244 debug: bool = False,

230 245 workers: int = None,

231 use_colors: bool = None

246 use_colors: bool = None,

247 g4f_api_key: str = None

232 248 ) -> None:

233 249 print(f'Starting server... [g4f v-{g4f.version.utils.current_version}]' + (" (debug)" if debug else ""))

234 250 if use_colors is None:

235 251 use_colors = debug

236 252 if bind is not None:

237 253 host, port = bind.split(":")

254 if debug:

255 g4f.debug.logging = True

238 256 uvicorn.run(

239 f"g4f.api:create_app{'_debug' if debug else ''}",

240 host=host, port=int(port),

241 workers=workers,

242 use_colors=use_colors,

243 factory=True,

257 "g4f.api:create_app",

258 host=host,

259 port=int(port),

260 workers=workers,

261 use_colors=use_colors,

262 factory=True,

244 263 reload=debug

245 264 )

gpt4free

refactor(g4f/api/init.py): refactor API structure and improve async handling

代码差异