v1.6.0-beta1 conversations isolated by token

This commit is contained in:
lanqian528 2024-10-29 05:49:34 +08:00
parent d293d8893d
commit 59596bcb3c
7 changed files with 142 additions and 62 deletions

View File

@ -1,4 +1,5 @@
import asyncio
import json
import re
import time
import types
@ -41,7 +42,7 @@ app.add_middleware(
@app.on_event("startup")
async def app_start():
if scheduled_refresh:
scheduler.add_job(id='refresh', func=refresh_all_tokens, trigger='cron', hour=3, minute=0, day='*/4',
scheduler.add_job(id='refresh', func=refresh_all_tokens, trigger='cron', hour=3, minute=0, day='*/2',
kwargs={'force_refresh': True})
scheduler.start()
asyncio.get_event_loop().call_later(0, lambda: asyncio.create_task(refresh_all_tokens(force_refresh=False)))
@ -171,9 +172,43 @@ if enable_gateway:
return {"gizmos": []}
# @app.get("/backend-api/conversations")
# async def get_conversations():
# return {"items": [], "total": 0, "limit": 28, "offset": 0, "has_missing_conversations": False}
@app.get("/backend-api/conversations")
async def get_conversations(request: Request):
limit = request.query_params.get("limit", 28)
offset = request.query_params.get("offset", 0)
token = request.headers.get("Authorization", "").replace("Bearer ", "")
if len(token) == 45 or token.startswith("eyJhbGciOi"):
return await chatgpt_reverse_proxy(request, "backend-api/conversations")
else:
items = []
for conversation_id in globals.seed_map.get(token, {}).get("conversations", []):
conversation = globals.conversation_map.get(conversation_id, None)
if conversation:
items.append(conversation)
items = items[int(offset):int(offset) + int(limit)]
conversations = {
"items": items,
"total": len(items),
"limit": 28,
"offset": 0,
"has_missing_conversations": False
}
return conversations
@app.get("/backend-api/conversation/{conversation_id}")
async def update_conversations(request: Request, conversation_id: str):
token = request.headers.get("Authorization", "").replace("Bearer ", "")
if len(token) == 45 or token.startswith("eyJhbGciOi"):
return await chatgpt_reverse_proxy(request, f"backend-api/conversation/{conversation_id}")
else:
conversation_details_str = (await chatgpt_reverse_proxy(request, f"backend-api/conversation/{conversation_id}")).body.decode('utf-8')
conversation_details = json.loads(conversation_details_str)
if conversation_id in globals.conversation_map:
globals.conversation_map[conversation_id]["title"] = conversation_details.get("title", None)
globals.conversation_map[conversation_id]["is_archived"] = conversation_details.get("is_archived", False)
with open(globals.CONVERSATION_MAP_FILE, "w", encoding="utf-8") as f:
json.dump(globals.conversation_map, f, indent=4)
return conversation_details
# @app.patch("/backend-api/conversations")
# async def get_conversations():
@ -238,6 +273,12 @@ if enable_gateway:
@app.api_route("/{path:path}", methods=["GET", "POST", "PUT", "DELETE", "OPTIONS", "HEAD", "PATCH", "TRACE"])
async def reverse_proxy(request: Request, path: str):
if re.match("ces/v1", path):
return {"success": True}
if re.match("backend-api/edge", path):
return Response(status_code=204)
for chatgpt_path in chatgpt_paths:
if re.match(chatgpt_path, path):
return await chatgpt_html(request)

View File

@ -11,17 +11,17 @@ from chatgpt.refreshToken import rt2ac
from utils.Logger import logger
from utils.config import authorization_list, random_token
os.environ['PYTHONHASHSEED'] = '0'
random.seed(0)
def get_req_token(req_token, seed=None):
available_token_list = list(set(globals.token_list) - set(globals.error_token_list))
length = len(available_token_list)
if seed and length > 0:
req_token = globals.token_list[hash(seed) % length]
while req_token in globals.error_token_list:
req_token = random.choice(globals.token_list)
if seed not in globals.seed_map.keys():
globals.seed_map[seed] = {"token": random.choice(available_token_list), "conversations": []}
with open(globals.SEED_MAP_FILE, "w") as f:
json.dump(globals.seed_map, f, indent=4)
else:
req_token = globals.seed_map[seed]["token"]
return req_token
if req_token in authorization_list:
@ -63,7 +63,7 @@ def get_ua(req_token):
}
globals.user_agent_map[req_token] = user_agent
with open(globals.USER_AGENTS_FILE, "w", encoding="utf-8") as f:
f.write(json.dumps(globals.user_agent_map, indent=4))
json.dump(globals.user_agent_map, f, indent=4)
return user_agent
else:
return user_agent
@ -94,7 +94,7 @@ async def refresh_all_tokens(force_refresh=False):
for token in list(set(globals.token_list) - set(globals.error_token_list)):
if len(token) == 45:
try:
await asyncio.sleep(2)
await asyncio.sleep(0.5)
await rt2ac(token, force_refresh=force_refresh)
except HTTPException:
pass

View File

@ -12,6 +12,8 @@ REFRESH_MAP_FILE = os.path.join(DATA_FOLDER, "refresh_map.json")
ERROR_TOKENS_FILE = os.path.join(DATA_FOLDER, "error_token.txt")
WSS_MAP_FILE = os.path.join(DATA_FOLDER, "wss_map.json")
USER_AGENTS_FILE = os.path.join(DATA_FOLDER, "user_agents.json")
SEED_MAP_FILE = os.path.join(DATA_FOLDER, "seed_map.json")
CONVERSATION_MAP_FILE = os.path.join(DATA_FOLDER, "conversation_map.json")
count = 0
token_list = []
@ -19,6 +21,8 @@ error_token_list = []
refresh_map = {}
wss_map = {}
user_agent_map = {}
seed_map = {}
conversation_map = {}
impersonate_list = [
"chrome99",
"chrome100",
@ -38,18 +42,54 @@ if not os.path.exists(DATA_FOLDER):
os.makedirs(DATA_FOLDER)
if os.path.exists(REFRESH_MAP_FILE):
with open(REFRESH_MAP_FILE, "r") as file:
refresh_map = json.load(file)
with open(REFRESH_MAP_FILE, "r") as f:
try:
refresh_map = json.load(f)
except:
refresh_map = {}
else:
refresh_map = {}
if os.path.exists(WSS_MAP_FILE):
with open(WSS_MAP_FILE, "r") as file:
wss_map = json.load(file)
with open(WSS_MAP_FILE, "r") as f:
try:
wss_map = json.load(f)
except:
wss_map = {}
else:
wss_map = {}
if os.path.exists(USER_AGENTS_FILE):
with open(USER_AGENTS_FILE, "r", encoding="utf-8") as f:
try:
user_agent_map = json.load(f)
except:
user_agent_map = {}
else:
user_agent_map = {}
if os.path.exists(SEED_MAP_FILE):
with open(SEED_MAP_FILE, "r") as f:
try:
seed_map = json.load(f)
except:
seed_map = {}
else:
seed_map = {}
if os.path.exists(CONVERSATION_MAP_FILE):
with open(CONVERSATION_MAP_FILE, "r") as f:
try:
conversation_map = json.load(f)
except:
conversation_map = {}
else:
conversation_map = {}
if os.path.exists(TOKENS_FILE):
with open(TOKENS_FILE, "r", encoding="utf-8") as f:
for line in f:
@ -69,40 +109,6 @@ else:
with open(ERROR_TOKENS_FILE, "w", encoding="utf-8") as f:
pass
if os.path.exists(USER_AGENTS_FILE):
with open(USER_AGENTS_FILE, "r", encoding="utf-8") as f:
try:
user_agent_map = json.load(f)
except json.JSONDecodeError:
user_agent_map = {}
# token数量变化时更新ua
if len(user_agent_map.keys()) != len(token_list):
new_tokens = list(set(token_list) - user_agent_map.keys())
for token in new_tokens:
ua = ua_generator.generate(device='desktop', browser=('chrome', 'edge'), platform=('windows', 'macos'))
ua_dict = {
"user-agent": ua.text,
"sec-ch-ua-platform": ua.platform,
"sec-ch-ua": ua.ch.brands,
"sec-ch-ua-mobile": ua.ch.mobile,
"impersonate": random.choice(impersonate_list),
}
user_agent_map[token] = ua_dict
with open(USER_AGENTS_FILE, "w", encoding="utf-8") as f:
f.write(json.dumps(user_agent_map, indent=4))
else:
for token in token_list:
ua = ua_generator.generate(device='desktop', browser=('chrome', 'edge'), platform=('windows', 'macos'))
ua_dict = {
"user-agent": ua.text,
"sec-ch-ua-platform": ua.platform,
"sec-ch-ua": ua.ch.brands,
"sec-ch-ua-mobile": ua.ch.mobile,
"impersonate": random.choice(impersonate_list),
}
user_agent_map[token] = ua_dict
with open(USER_AGENTS_FILE, "w", encoding="utf-8") as f:
f.write(json.dumps(user_agent_map, indent=4))
if token_list:
logger.info(f"Token list count: {len(token_list)}, Error token list count: {len(error_token_list)}")

View File

@ -10,21 +10,17 @@ from utils.config import proxy_url_list
import chatgpt.globals as globals
def save_refresh_map(refresh_map):
with open(globals.REFRESH_MAP_FILE, "w") as file:
json.dump(refresh_map, file)
async def rt2ac(refresh_token, force_refresh=False):
if not force_refresh and (refresh_token in globals.refresh_map and int(time.time()) - globals.refresh_map.get(refresh_token, {}).get("timestamp", 0) < 5 * 24 * 60 * 60):
access_token = globals.refresh_map[refresh_token]["token"]
logger.info(f"refresh_token -> access_token from cache")
# logger.info(f"refresh_token -> access_token from cache")
return access_token
else:
try:
access_token = await chat_refresh(refresh_token)
globals.refresh_map[refresh_token] = {"token": access_token, "timestamp": int(time.time())}
save_refresh_map(globals.refresh_map)
with open(globals.REFRESH_MAP_FILE, "w") as f:
json.dump(globals.refresh_map, f, indent=4)
logger.info(f"refresh_token -> access_token with openai: {access_token}")
return access_token
except HTTPException as e:

View File

@ -6,8 +6,17 @@ from fastapi.responses import StreamingResponse, Response
from starlette.background import BackgroundTask
from chatgpt.authorization import verify_token, get_req_token, get_ua
import chatgpt.globals as globals
from utils.Client import Client
from utils.config import chatgpt_base_url_list, proxy_url_list, enable_gateway
from utils.config import chatgpt_base_url_list, proxy_url_list
from datetime import datetime, timezone
def generate_current_time():
current_time = datetime.now(timezone.utc)
formatted_time = current_time.isoformat(timespec='microseconds').replace('+00:00', 'Z')
return formatted_time
headers_reject_list = [
"x-real-ip",
@ -70,6 +79,34 @@ async def get_real_req_token(token):
return req_token
async def content_generator(r, token):
first_chunk = None
async for chunk in r.aiter_content():
if first_chunk is None and len(token) != 45 and not token.startswith("eyJhbGciOi"):
first_chunk = chunk.decode('utf-8')
conversation_id = json.loads(first_chunk[6:]).get("conversation_id")
conversation_detail = {
"id": conversation_id,
"title": "New Chat",
"update_time": generate_current_time(),
"workspace_id": None,
}
if conversation_id not in globals.conversation_map:
globals.conversation_map[conversation_id] = conversation_detail
else:
globals.conversation_map[conversation_id]["update_time"] = generate_current_time()
if conversation_id not in globals.seed_map[token]["conversations"]:
globals.seed_map[token]["conversations"].insert(0, conversation_id)
else:
globals.seed_map[token]["conversations"].remove(conversation_id)
globals.seed_map[token]["conversations"].insert(0, conversation_id)
with open(globals.CONVERSATION_MAP_FILE, "w", encoding="utf-8") as f:
json.dump(globals.conversation_map, f)
with open(globals.SEED_MAP_FILE, "w", encoding="utf-8") as f:
json.dump(globals.seed_map, f, indent=4)
yield chunk
async def chatgpt_reverse_proxy(request: Request, path: str):
try:
origin_host = request.url.netloc
@ -129,7 +166,7 @@ async def chatgpt_reverse_proxy(request: Request, path: str):
.replace("cdn.oaistatic.com", origin_host)
.replace("https", petrol)}, background=background)
elif 'stream' in r.headers.get("content-type", ""):
return StreamingResponse(r.aiter_content(), media_type=r.headers.get("content-type", ""),
return StreamingResponse(content_generator(r, token), media_type=r.headers.get("content-type", ""),
background=background)
else:
if "/backend-api/conversation" in path or "/register-websocket" in path:

View File

@ -6,8 +6,8 @@ import chatgpt.globals as globals
def save_wss_map(wss_map):
with open(globals.WSS_MAP_FILE, "w") as file:
json.dump(wss_map, file)
with open(globals.WSS_MAP_FILE, "w") as f:
json.dump(wss_map, f, indent=4)
async def token2wss(token):

View File

@ -1 +1 @@
1.5.12-beta2
1.6.0-beta1