Upload openai.py

openai.py

@@ -0,0 +1,423 @@
+import json
+import struct
+import gzip
+import time
+import uuid
+import os
+import re
+import asyncio
+import hashlib
+import queue
+from concurrent.futures import ThreadPoolExecutor
+from typing import List, Optional, Dict, Any
+from fastapi import FastAPI, Request, HTTPException, Depends
+from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
+from fastapi.responses import StreamingResponse
+from curl_cffi import requests
+import uvicorn
+
+app = FastAPI()
+security = HTTPBearer()
+
+# 配置项，支持环境变量覆盖
+COOKIES_PATH = os.environ.get("COOKIES_PATH", "cookies.json")
+PROXY = os.environ.get("HTTP_PROXY", None)  # 不设置则不走代理
+
+# 同步阻塞调用用的线程池
+_executor = ThreadPoolExecutor(max_workers=16)
+
+
+def _load_cookies(path: str) -> dict:
+    try:
+        with open(path, 'r', encoding='utf-8') as f:
+            cookies_list = json.load(f)
+        return {c['name']: c['value'] for c in cookies_list}
+    except Exception as e:
+        print(f"Error loading cookies: {e}")
+        return {}
+
+
+def _generate_device_id(seed: str) -> str:
+    h = hashlib.sha256(seed.encode()).hexdigest()
+    return str(int(h[:16], 16))[:19]
+
+
+def _generate_session_id(seed: str) -> str:
+    h = hashlib.sha256(("session-" + seed).encode()).hexdigest()
+    return str(int(h[:16], 16))[:19]
+
+
+def pack_connect_message(data: dict) -> bytes:
+    payload = json.dumps(data, separators=(',', ':')).encode('utf-8')
+    header = struct.pack('>BI', 0, len(payload))
+    return header + payload
+
+
+def _convert_citations(text: str) -> str:
+    """将 Kimi 的 [^N^] 引用格式转换为 [N]"""
+    return re.sub(r'\[\^(\d+)\^\]', r'[\1]', text)
+
+
+def _format_references(refs: list) -> str:
+    """将搜索引用格式化为 markdown 脚注"""
+    if not refs:
+        return ""
+    lines = ["\n\n---", "**Sources:**"]
+    for ref in refs:
+        base = ref.get("base", {})
+        title = base.get("title", "")
+        url = base.get("url", "")
+        ref_id = ref.get("id", "")
+        if title and url:
+            lines.append(f"[{ref_id}] [{title}]({url})")
+    return "\n".join(lines) + "\n"
+
+
+# ── 帧解析 ──────────────────────────────────────────────
+
+def _parse_kimi_frames(buffer: bytes):
+    """解析 connect 帧，返回 (events, remaining_buffer)。
+    event 类型:
+      - {"type": "text", "content": "..."}
+      - {"type": "tool_status", "name": "...", "status": "..."}
+      - {"type": "search_refs", "refs": [...]}
+      - {"type": "done"}
+    """
+    events = []
+    while len(buffer) >= 5:
+        flag, length = struct.unpack_from('>BI', buffer, 0)
+        if len(buffer) < 5 + length:
+            break
+        payload_bytes = buffer[5:5 + length]
+        buffer = buffer[5 + length:]
+
+        if flag == 2:
+            try:
+                payload_bytes = gzip.decompress(payload_bytes)
+            except:
+                pass
+
+        if flag not in (0, 2):
+            continue
+
+        try:
+            data = json.loads(payload_bytes.decode('utf-8'))
+        except Exception as e:
+            print(f"DEBUG: Error decoding frame JSON: {e}")
+            continue
+
+        # done 信号
+        if "done" in data:
+            events.append({"type": "done"})
+            continue
+
+        # heartbeat 跳过
+        if "heartbeat" in data:
+            continue
+
+        op = data.get("op")
+        if op not in ("set", "append"):
+            continue
+
+        # 文本内容
+        if "block" in data and "text" in data["block"]:
+            content = data["block"]["text"].get("content", "")
+            if content:
+                events.append({"type": "text", "content": content})
+
+        # message.blocks 里的文本 — 只提取 assistant 角色的，跳过 user/system 回显
+        if "message" in data and "blocks" in data.get("message", {}):
+            msg_role = data["message"].get("role", "")
+            if msg_role == "assistant":
+                content = ""
+                for block in data["message"]["blocks"]:
+                    if "text" in block:
+                        content += block["text"].get("content", "")
+                if content:
+                    events.append({"type": "text", "content": content})
+
+        # 工具调用状态
+        if "block" in data and "tool" in data["block"]:
+            tool = data["block"]["tool"]
+            name = tool.get("name", "")
+            status = tool.get("status", "")
+            if name and status:
+                events.append({"type": "tool_status", "name": name, "status": status})
+
+        # 搜索引用 (usedSearchChunks 优先)
+        msg = data.get("message", {})
+        refs = msg.get("refs", {})
+        if "usedSearchChunks" in refs:
+            events.append({"type": "search_refs", "refs": refs["usedSearchChunks"]})
+
+    return events, buffer
+
+
+# 硬编码的 API Key，匹配时使用 cookies.json 认证
+API_KEY = "sk-sseworld-kimi"
+
+# ── Kimi Bridge ─────────────────────────────────────────
+
+class KimiBridge:
+    def __init__(self):
+        self.base_url = "https://www.kimi.com"
+
+    def create_session(self, api_key: str):
+        if api_key == API_KEY:
+            cookies = _load_cookies(COOKIES_PATH)
+            auth_token = cookies.get("kimi-auth", "")
+            fingerprint_seed = "cookies-default"
+        else:
+            cookies = {}
+            auth_token = api_key
+            fingerprint_seed = api_key
+
+        device_id = _generate_device_id(fingerprint_seed)
+        session_id = _generate_session_id(fingerprint_seed)
+
+        headers = {
+            "accept": "*/*",
+            "accept-language": "zh-CN,zh;q=0.9",
+            "authorization": f"Bearer {auth_token}",
+            "content-type": "application/connect+json",
+            "connect-protocol-version": "1",
+            "origin": "https://www.kimi.com",
+            "referer": "https://www.kimi.com/",
+            "user-agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/144.0.0.0 Safari/537.36",
+            "x-language": "zh-CN",
+            "x-msh-device-id": device_id,
+            "x-msh-platform": "web",
+            "x-msh-session-id": session_id,
+            "x-msh-version": "1.0.0",
+            "x-traffic-id": f"u{device_id[:20]}",
+        }
+
+        return requests.Session(
+            headers=headers,
+            cookies=cookies,
+            impersonate="chrome124",
+            proxy=PROXY,
+        )
+
+
+bridge = KimiBridge()
+
+
+# ── OpenAI 格式化 ──────────────────────────────────────
+
+def format_openai_stream_chunk(content: str, model: str, chat_id: str, *, role: str = None, finish_reason: str = None):
+    delta = {}
+    if role:
+        delta["role"] = role
+    if content:
+        delta["content"] = content
+    chunk = {
+        "id": chat_id,
+        "object": "chat.completion.chunk",
+        "created": int(time.time()),
+        "model": model,
+        "choices": [{
+            "index": 0,
+            "delta": delta,
+            "finish_reason": finish_reason
+        }]
+    }
+    return f"data: {json.dumps(chunk, ensure_ascii=False)}\n\n"
+
+
+# ── 同步辅助函数 (线程池中执行) ────────────────────────
+
+def _sync_kimi_request(session, url, body_bytes):
+    return session.post(url, data=body_bytes, stream=True, timeout=30)
+
+
+def _sync_read_all(response):
+    """同步读取完整响应，返回 (full_text, search_refs)"""
+    full_content = ""
+    search_refs = []
+    buffer = b""
+    for chunk in response.iter_content(chunk_size=None):
+        if not chunk:
+            continue
+        buffer += chunk
+        events, buffer = _parse_kimi_frames(buffer)
+        for ev in events:
+            if ev["type"] == "text":
+                full_content += ev["content"]
+            elif ev["type"] == "search_refs":
+                search_refs = ev["refs"]
+    full_content = _convert_citations(full_content)
+    if search_refs:
+        full_content += _format_references(search_refs)
+    return full_content
+
+
+# ── 路由 ────────────────────────────────────────────────
+
+@app.middleware("http")
+async def log_requests(request: Request, call_next):
+    print(f"DEBUG: Incoming request: {request.method} {request.url}")
+    response = await call_next(request)
+    print(f"DEBUG: Response status: {response.status_code}")
+    return response
+
+
+KIMI_MODELS = {
+    "kimi-k2.5": {"scenario": "SCENARIO_K2D5", "thinking": False},
+    "kimi-k2.5-thinking": {"scenario": "SCENARIO_K2D5", "thinking": True},
+}
+DEFAULT_MODEL = "kimi-k2.5"
+
+
+@app.get("/v1/models")
+async def list_models():
+    return {
+        "object": "list",
+        "data": [
+            {"id": mid, "object": "model", "created": 0, "owned_by": "moonshot"}
+            for mid in KIMI_MODELS
+        ]
+    }
+
+
+@app.post("/v1/chat/completions")
+async def chat_completions(request: Request, credentials: HTTPAuthorizationCredentials = Depends(security)):
+    api_key = credentials.credentials
+    print(f"DEBUG: chat_completions endpoint hit, key prefix: {api_key[:6]}...")
+    print(f"DEBUG: Request headers: {dict(request.headers)}")
+
+    session = bridge.create_session(api_key)
+
+    try:
+        body = await request.json()
+    except Exception as e:
+        print(f"DEBUG: Failed to parse request JSON: {e}")
+        raise HTTPException(status_code=400, detail="Invalid JSON body")
+
+    messages = body.get("messages", [])
+    model = body.get("model", "kimi-k2.5")
+    stream = body.get("stream", False)
+    model_config = KIMI_MODELS.get(model, KIMI_MODELS[DEFAULT_MODEL])
+
+    print(f"DEBUG: Received request: model={model}, thinking={model_config['thinking']}, stream={stream}, messages_count={len(messages)}")
+
+    if not messages:
+        raise HTTPException(status_code=400, detail="Messages are required")
+
+    # 构造 Kimi 的请求
+    kimi_blocks = []
+    for msg in messages:
+        role = msg.get("role", "user")
+        content = msg.get("content", "")
+        prefix = "User: " if role == "user" else "Assistant: "
+        kimi_blocks.append({"message_id": "", "text": {"content": f"{prefix}{content}\n"}})
+
+    kimi_payload = {
+        "scenario": model_config["scenario"],
+        "tools": [{"type": "TOOL_TYPE_SEARCH", "search": {}}],
+        "message": {
+            "role": "user",
+            "blocks": kimi_blocks,
+            "scenario": model_config["scenario"]
+        },
+        "options": {"thinking": model_config["thinking"]}
+    }
+
+    print(f"DEBUG: Kimi payload size: {len(json.dumps(kimi_payload))}")
+
+    url = f"{bridge.base_url}/apiv2/kimi.gateway.chat.v1.ChatService/Chat"
+    body_bytes = pack_connect_message(kimi_payload)
+
+    print(f"DEBUG: Forwarding to Kimi: {url}")
+
+    loop = asyncio.get_event_loop()
+
+    try:
+        response = await loop.run_in_executor(_executor, _sync_kimi_request, session, url, body_bytes)
+        print(f"DEBUG: Kimi response status: {response.status_code}")
+    except Exception as e:
+        print(f"DEBUG: Request to Kimi failed: {e}")
+        session.close()
+        raise HTTPException(status_code=500, detail=f"Failed to connect to Kimi: {str(e)}")
+
+    if response.status_code != 200:
+        error_text = response.text
+        print(f"DEBUG: Kimi error: {error_text}")
+        session.close()
+        raise HTTPException(status_code=response.status_code, detail=f"Kimi API error: {error_text}")
+
+    chat_id = str(uuid.uuid4())
+
+    if stream:
+        async def generate():
+            q = queue.Queue()
+            sentinel = object()
+            sent_role = False
+
+            def _stream_worker():
+                try:
+                    buf = b""
+                    search_refs = []
+                    for chunk in response.iter_content(chunk_size=None):
+                        if not chunk:
+                            continue
+                        buf += chunk
+                        events, buf = _parse_kimi_frames(buf)
+                        for ev in events:
+                            if ev["type"] == "text":
+                                q.put(("text", _convert_citations(ev["content"])))
+                            elif ev["type"] == "tool_status" and ev["status"] == "STATUS_RUNNING":
+                                q.put(("text", "\n\n> [Searching...]\n\n"))
+                            elif ev["type"] == "search_refs":
+                                search_refs = ev["refs"]
+                    # 流结束，追加引用
+                    if search_refs:
+                        q.put(("text", _format_references(search_refs)))
+                finally:
+                    q.put(sentinel)
+                    session.close()
+
+            loop.run_in_executor(_executor, _stream_worker)
+
+            while True:
+                try:
+                    item = await loop.run_in_executor(None, q.get, True, 0.5)
+                except:
+                    continue
+                if item is sentinel:
+                    break
+                _, content = item
+                if not sent_role:
+                    yield format_openai_stream_chunk(content, model, chat_id, role="assistant")
+                    sent_role = True
+                else:
+                    yield format_openai_stream_chunk(content, model, chat_id)
+
+            # finish_reason: stop
+            yield format_openai_stream_chunk("", model, chat_id, finish_reason="stop")
+            yield "data: [DONE]\n\n"
+
+        return StreamingResponse(generate(), media_type="text/event-stream")
+    else:
+        try:
+            full_content = await loop.run_in_executor(_executor, _sync_read_all, response)
+        finally:
+            session.close()
+
+        return {
+            "id": chat_id,
+            "object": "chat.completion",
+            "created": int(time.time()),
+            "model": model,
+            "choices": [{
+                "index": 0,
+                "message": {"role": "assistant", "content": full_content},
+                "finish_reason": "stop"
+            }],
+            "usage": {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}
+        }
+
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run("openai:app", host="127.0.0.1", port=8001, reload=False, log_level="debug")