Update main.py

main.py

@@ -1,6 +1,5 @@
 import os
 from datetime import datetime
-import time
 from fastapi import FastAPI, HTTPException, Request, Depends
 import httpx
 import json
@@ -11,6 +10,8 @@ import asyncio
 import threading
 import string
 
+import logging
+
 app = FastAPI()
 app.add_middleware(
     CORSMiddleware,
@@ -44,7 +45,6 @@ async def chat_chunk(body, headers):
                     headers = await get_headers()
                     continue
                 async for line in response.aiter_lines():
-                    time.sleep(0.05)
                     if line.startswith("data: [DONE]"):
                         char_data["choices"][0]["finish_reason"] = "stop"
                         char_data["choices"][0]["delta"] = {}
@@ -150,10 +150,11 @@ async def auth(request: Request):
     token = key
     #获取Authorization头部
     auth_header = request.headers.get('Authorization')
-    # print(f'{formatted_now()} auth_chunk: {request.headers}')
+    print()
+    print(f'{formatted_now()} auth_chunk: {request.headers}')
     # 如果没有Authorization头部，返回401
-    # if not auth_header:
-        # print('没有验证')
+    if not auth_header:
+        print('没有验证')
     #     raise HTTPException(status_code=401, detail="Unauthorized")
     #     # 分割头部字符串
     # parts = auth_header.split()
@@ -213,6 +214,16 @@ async def chat_completions(req: Request, authorized: str = Depends(auth)):
             raise HTTPException(status_code=400, detail=str(e))
         except httpx.ReadTimeout:  # 捕获超时异常
             raise HTTPException(status_code=408, detail="Request Timeout")
+    elif 'embedding' in usermodel.lower():
+        try:
+
+            url = 'https://open.bigmodel.cn/api/paas/v4/embeddings'
+            headers = {'Content-Type': 'application/json',
+                       'Authorization': f'Bearer {apiToken}'
+                       }
+            return await zhipu_chat(url, headers, body)
+        except json.JSONDecodeError as e:
+            raise HTTPException(status_code=400, detail=str(e))
 
     else:
         chat_url = f"{baseurl}/duckchat/v1/chat"
@@ -263,7 +274,7 @@ async def chat_completions(req: Request, authorized: str = Depends(auth)):
         for attempt in range(max_retries):
             try:
                 async with httpx.AsyncClient() as client:
-                    print(f'{formatted_now()} duckgo_chat: {headers}')
+                    # print(f'{formatted_now()} duckgo_chat: {headers}')
                     resp = await client.post(chat_url, json=payload, headers=headers, timeout=60)
                     resp.raise_for_status()
                     if resp.status_code == 200:
@@ -450,6 +461,16 @@ async def zhipu_chat_completions(req: Request, authorized: str = Depends(auth)):
     except httpx.ReadTimeout:  # 捕获超时异常
         raise HTTPException(status_code=408, detail="Request Timeout")
 
+@app.post("/api/paas/v4/embeddings")
+async def zhipu_embeddings(req: Request, authorized: str = Depends(auth)):
+    body = await req.json()
+    url = 'https://open.bigmodel.cn/api/paas/v4/embeddings'
+    headers = {'Content-Type': 'application/json',
+               'Authorization': f'Bearer {apiToken}'
+               }
+    return await zhipu_chat(url, headers, body)
+
+
 def formatted_now():
     # 获取当前时间
     now = datetime.now()
@@ -457,7 +478,7 @@ def formatted_now():
 
 async def zhipu_chat(url, headers, body):
     while True:
-        print(f'{formatted_now()} zhipu_chat: {headers}')
+        # print(f'{formatted_now()} zhipu_chat: {headers}')
         async with httpx.AsyncClient() as client:
             response = await client.post(url, headers=headers, json=body, timeout=120)
             if response.status_code == 401:
@@ -465,7 +486,7 @@ async def zhipu_chat(url, headers, body):
                 continue
             break
 
-    print(f'{formatted_now()}: {response.json()}')
+    # print(f'{formatted_now()}: {response.json()}')
     return response.json()
 
 
@@ -513,16 +534,18 @@ def get_token():
 
 def get_external_ip():
     try:
-        response = httpx.get('https://api.ipify.org?format=json')
+        response = httpx.get('https://ipinfo.io/json')
         response.raise_for_status()  # 检查请求是否成功
         ip_info = response.json()
-        print(f"外网 IP 地址: {ip_info}")
+        ip_info2 = f"ip:{ip_info['ip']} city:{ip_info['city']} region:{ip_info['region']} country:{ip_info['country']} timezone:{ip_info['timezone']}"
+        # print(f"外网 IP 地址: {ip_info}")
+        logging.error(ip_info2)
+        logging.shutdown()
         return ip_info['ip']
     except Exception as e:
         print(f"获取外网 IP 地址时出错: {e}")
         return None
 def schedule_token_update():
-    # get_external_ip()
     get_token()
     # 创建一个新的定时器，设置为300秒后再次调用schedule_token_update函数
     threading.Timer(300, schedule_token_update).start()
@@ -552,7 +575,9 @@ if __name__ == "__main__":
             "uvicorn.access": {"handlers": ["default"], "level": "INFO", "propagate": False},
         },
     }
+    get_external_ip()
     import uvicorn
+
     schedule_token_update()
     # uvicorn.run(app, host='0.0.0.0', port=17890, log_level='warning')
     uvicorn.run(app, host='0.0.0.0', port=7860, log_config=log_config)