Update app.py

app.py

@@ -1,5 +1,5 @@
 from fastapi import FastAPI, Request, HTTPException
-from fastapi.responses import StreamingResponse, HTMLResponse
+from fastapi.responses import StreamingResponse, HTMLResponse, JSONResponse
 from fastapi.middleware.cors import CORSMiddleware
 import httpx
 import json
@@ -13,7 +13,7 @@ import logging
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 
 # 从环境变量或者文件中加载配置
-BASE_URL = os.environ.get("BASE_URL", "https://api.siliconflow.cn/v1/chat/completions")
+BASE_URL = os.environ.get("BASE_URL", "https://api.siliconflow.cn/v1")
 MODEL_MAP = {
     "qwen-72b": "Qwen/Qwen2.5-72B-Instruct",
     "qwen-32b": "Qwen/Qwen2.5-32B-Instruct",
@@ -47,6 +47,7 @@ LAST_UPDATE_TIME = ""
 
 app = FastAPI()
 
+
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -121,7 +122,6 @@ async def get_model_info():
         models += f"<h2>{key}————{value}</h2>"
     return models
 
-
 @app.get("/", response_class=HTMLResponse)
 async def root():
     """根路由，返回 HTML 页面，展示模型信息和更新时间"""
@@ -140,53 +140,96 @@ async def root():
     </html>
     """
 
+
 @app.get("/check")
 async def check():
     """手动触发检查 API Key 余额的路由"""
     await key_manager.check_keys_balance()
     return f"更新成功：{key_manager.last_update_time}"
 
-async def _forward_request(request: Request, api_type: str):
+async def _forward_request(request: Request, api_type: str, image_generations:bool = False):
     """转发请求到硅流 API"""
-    body = await request.json()
+    
     try:
-          key = key_manager.get_key()
+        key = key_manager.get_key()
     except HTTPException as e:
         return e
     logging.info(f"using key {key[0:4]}***{key[-4:]} to {api_type}")
     headers = {"Authorization": f"Bearer {key}"}
-    # 处理模型映射
-    if "model" in body and body["model"] in MODEL_MAP:
-        body["model"] = MODEL_MAP[body["model"]]
-
-    if api_type == "chat" and "stream" in body and body["stream"]:
-        async def generate_response():
+    if  api_type == "chat" and not image_generations :
+            body = await request.json()
+            # 处理模型映射
+            if "model" in body and body["model"] in MODEL_MAP:
+                body["model"] = MODEL_MAP[body["model"]]
+
+            if api_type == "chat" and "stream" in body and body["stream"]:
+                async def generate_response():
+                   async with httpx.AsyncClient() as client:
+                        async with client.stream("POST", f"{BASE_URL}/chat/completions" , headers=headers, json=body) as response:
+                                response.raise_for_status()  # 检查响应状态码
+                                async for chunk in response.aiter_bytes():
+                                    if chunk:
+                                       yield chunk
+                return StreamingResponse(generate_response(), media_type="text/event-stream")
+            else :
+                async with httpx.AsyncClient() as client:
+                       try:
+                            response = await client.post(f"{BASE_URL}/chat/completions", headers=headers, json=body)
+                            response.raise_for_status()
+                            return response.json()
+                       except httpx.HTTPError as exc :
+                                logging.error("httpx request error:" + str(exc))
+                                raise HTTPException(
+                                        status_code=500,
+                                        detail=f"Request failed with status: {exc.response.status_code},url:{exc.request.url},detail:{exc.response.text}"
+                                )
+    elif image_generations:
+            url =f"{BASE_URL}/images/generations"
+            body = await request.json()
+            headers["Content-Type"] = "application/json"
             async with httpx.AsyncClient() as client:
-                async with client.stream("POST", BASE_URL, headers=headers, json=body) as response:
-                        response.raise_for_status()  # 检查响应状态码
-                        async for chunk in response.aiter_bytes():
-                            if chunk:
-                                yield chunk
-        return StreamingResponse(generate_response(), media_type="text/event-stream")
-    else:
-        async with httpx.AsyncClient() as client:
-            try:
-                response = await client.post(BASE_URL, headers=headers, json=body)
-                response.raise_for_status()
-                return response.json()
-            except httpx.HTTPError as exc :
-                 logging.error("httpx request error:" + str(exc))
-                 raise HTTPException(
-                     status_code=500,
-                     detail=f"Request failed with status: {exc.response.status_code},url:{exc.request.url},detail:{exc.response.text}"
-                 )
+                
+                 try:
+                         response = await client.post(url, headers=headers,json=body)
+                         response.raise_for_status()
+                         return JSONResponse(content=response.json(), status_code=response.status_code)
+                 except httpx.HTTPError as exc :
+                          logging.error("httpx request error:" + str(exc))
+                          raise HTTPException(
+                               status_code=500,
+                                detail=f"Request failed with status: {exc.response.status_code}, url:{exc.request.url} ,detail:{exc.response.text}"
+                           )
+    elif api_type =="embedding":
+          url=f"{BASE_URL}/embeddings"
+          body=await request.json()
+          if "model" in body and body["model"] in MODEL_MAP:
+              body["model"] = MODEL_MAP[body["model"]]
+          async with httpx.AsyncClient() as client:
+              try:
+                        response = await client.post(url, json=body,headers=headers)
+                        response.raise_for_status()
+                        return response.json()
+              except  httpx.HTTPError as exc :
+                           logging.error(f"httpx request embedding error :{exc}")
+                           raise HTTPException(
+                               status_code=500,
+                                detail=f"Request failed with status: {exc.response.status_code}, url:{exc.request.url}, detail:{exc.response.text}"
+                           )
+
+
 
 @app.post("/hf/v1/chat/completions")
 async def chat_completions(request: Request):
         """转发 chat 完成请求"""
         return await _forward_request(request, "chat")
+
 @app.post("/hf/v1/embeddings")
 async def embeddings(request: Request):
     """转发 embeddings 请求"""
     return await _forward_request(request, "embedding")
 
+@app.post("/hf/v1/images/generations")
+async def image_generations(request:Request):
+     """转发图片生成请求"""
+     return await _forward_request(request, api_type="image",image_generations=True)
+