from fastapi import FastAPI, Request, Response, UploadFile, File from fastapi.responses import StreamingResponse, FileResponse from fastapi.staticfiles import StaticFiles import httpx import json import asyncio import time import base64 from typing import Optional, Dict, Any, List from io import BytesIO # 配置常量 QWEN_API_URL = "https://chat.qwenlm.ai/api/chat/completions" QWEN_MODELS_URL = "https://chat.qwenlm.ai/api/models" QWEN_FILES_URL = "https://chat.qwenlm.ai/api/v1/files/" MAX_RETRIES = 3 RETRY_DELAY = 1 # 1秒 # 缓存设置 cached_models = None cached_models_timestamp = 0 CACHE_TTL = 60 * 60 # 缓存1小时 app = FastAPI() app.mount("/static", StaticFiles(directory="static"), name="static") client = httpx.AsyncClient() @app.get("/") async def root(): return FileResponse("index.html") async def sleep(seconds: float): await asyncio.sleep(seconds) # 添加 base64 转换为文件的函数 async def base64_to_file(base64_str: str) -> BytesIO: try: # 去除 data:image/jpeg;base64, 这样的前缀 if ',' in base64_str: base64_str = base64_str.split(',', 1)[1] # 解码 base64 数据 image_data = base64.b64decode(base64_str) return BytesIO(image_data) except Exception as e: raise Exception(f"Failed to convert base64 to file: {str(e)}") # 添加图片上传函数 async def upload_image_to_qwen(auth_header: str, image_data: BytesIO) -> str: try: files = {'file': ('image.jpg', image_data, 'image/jpeg')} headers = { "Authorization": auth_header, "accept": "application/json" } async with httpx.AsyncClient() as client: response = await client.post( QWEN_FILES_URL, headers=headers, files=files ) if response.is_success: data = response.json() if not data.get('id'): raise Exception("File upload failed: No valid file ID returned") return data['id'] else: raise Exception(f"File upload failed with status {response.status_code}") except Exception as e: raise Exception(f"Failed to upload image: {str(e)}") # 添加消息处理函数 async def process_messages(messages: List[Dict], auth_header: str) -> List[Dict]: processed_messages = [] for message in messages: if isinstance(message.get('content'), list): new_content = [] for content in message['content']: if (content.get('type') == 'image_url' and content.get('image_url', {}).get('url', '').startswith('data:')): # 处理 base64 图片 image_data = await base64_to_file(content['image_url']['url']) image_id = await upload_image_to_qwen(auth_header, image_data) new_content.append({ 'type': 'image', 'image': image_id }) else: new_content.append(content) message['content'] = new_content processed_messages.append(message) return processed_messages async def fetch_with_retry(url: str, options: Dict, retries: int = MAX_RETRIES): last_error = None for i in range(retries): try: response = await client.request( method=options.get("method", "GET"), url=url, headers=options.get("headers", {}), json=options.get("json"), ) if response.is_success: return response content_type = response.headers.get("content-type", "") if response.status_code >= 500 or "text/html" in content_type: last_error = { "status": response.status_code, "content_type": content_type, "response_text": response.text[:1000], "headers": dict(response.headers) } if i < retries - 1: await sleep(RETRY_DELAY * (i + 1)) continue else: last_error = { "status": response.status_code, "headers": dict(response.headers) } break except Exception as error: last_error = error if i < retries - 1: await sleep(RETRY_DELAY * (i + 1)) continue raise Exception(json.dumps({ "error": True, "message": "All retry attempts failed", "last_error": str(last_error), "retries": retries })) async def process_line(line: str, previous_content: str) -> tuple[str, Optional[dict]]: try: data = json.loads(line[6:]) # 移除 "data: " 前缀 if (data.get("choices") and data["choices"][0].get("delta")): delta = data["choices"][0]["delta"] current_content = delta.get("content", "") # 计算增量内容 if previous_content and current_content: if current_content.startswith(previous_content): new_content = current_content[len(previous_content):] else: new_content = current_content else: new_content = current_content # 构造新的响应数据 new_data = { "choices": [{ "delta": { "role": delta.get("role", "assistant"), "content": new_content } }] } return current_content, new_data return previous_content, data except: return previous_content, None async def stream_generator(response: httpx.Response): buffer = "" previous_content = "" async for chunk in response.aiter_bytes(): chunk_text = chunk.decode() buffer += chunk_text lines = buffer.split("\n") buffer = lines.pop() if lines else "" for line in lines: line = line.strip() if line.startswith("data: "): previous_content, data = await process_line(line, previous_content) if data: yield f"data: {json.dumps(data)}\n\n" if buffer: previous_content, data = await process_line(buffer, previous_content) if data: yield f"data: {json.dumps(data)}\n\n" yield "data: [DONE]\n\n" @app.get("/healthz") async def health_check(): return {"status": "ok"} @app.get("/api/models") async def get_models(request: Request): global cached_models, cached_models_timestamp auth_header = request.headers.get("Authorization") if not auth_header or not auth_header.startswith("Bearer "): return Response(status_code=401, content="Unauthorized") now = time.time() if cached_models and now - cached_models_timestamp < CACHE_TTL: return Response( content=cached_models, media_type="application/json" ) try: response = await fetch_with_retry( QWEN_MODELS_URL, {"headers": {"Authorization": auth_header}} ) cached_models = response.text cached_models_timestamp = now return Response( content=cached_models, media_type="application/json" ) except Exception as error: return Response( content=json.dumps({"error": True, "message": str(error)}), status_code=500 ) @app.post("/api/chat/completions") async def chat_completions(request: Request): auth_header = request.headers.get("Authorization") if not auth_header or not auth_header.startswith("Bearer "): return Response(status_code=401, content="Unauthorized") request_data = await request.json() messages = request_data.get("messages") stream = request_data.get("stream", False) model = request_data.get("model") max_tokens = request_data.get("max_tokens") if not model: return Response( content=json.dumps({"error": True, "message": "Model parameter is required"}), status_code=400 ) try: # 处理消息中的图片 processed_messages = await process_messages(messages, auth_header) qwen_request = { "model": model, "messages": processed_messages, "stream": stream } if max_tokens is not None: qwen_request["max_tokens"] = max_tokens response = await client.post( QWEN_API_URL, headers={ "Content-Type": "application/json", "Authorization": auth_header }, json=qwen_request ) if stream: return StreamingResponse( stream_generator(response), media_type="text/event-stream" ) return Response( content=response.text, media_type="application/json" ) except Exception as error: return Response( content=json.dumps({"error": True, "message": str(error)}), status_code=500 ) if __name__ == "__main__": import uvicorn uvicorn.run(app, host="0.0.0.0", port=8000)