import asyncio import json import logging import random from typing import AsyncGenerator, Tuple import httpx import config from auth_manager import auth_manager import fingerprint import timing logger = logging.getLogger(__name__) def _content_to_text(content) -> str: if isinstance(content, list): return "".join( part.get("text", "") for part in content if isinstance(part, dict) and part.get("type") == "text" ) return str(content or "") def _build_prompt(messages: list) -> str: if not messages: return "Hello" lines = [] for msg in messages: role = msg.get("role", "user") text = _content_to_text(msg.get("content", "")).strip() if not text: continue lines.append(f"{role}: {text}") if not lines: return "Hello" return "\n".join(lines) def _resolve_model(model_name: str) -> Tuple[str, str]: if model_name in config.MODEL_MAP: return config.MODEL_MAP[model_name] if "__" in model_name: parts = model_name.split("__") if len(parts) >= 3: return parts[0], parts[-1] # Safe fallback for unknown names return "Anthropic", "claude-opus-4-6" def _random_referer() -> str: paths = [ "/app", "/app/chat", f"/app/chat/{random.randint(100, 9999)}", "/app/search", ] return f"{config.ONYX_BASE_URL}{random.choice(paths)}" def _headers(email: str, with_json: bool = False) -> dict: persona = fingerprint.get_persona_for_account(email) headers = fingerprint.get_base_headers(persona) headers["origin"] = config.ONYX_BASE_URL headers["referer"] = _random_referer() headers["accept"] = "application/json" if with_json: headers["content-type"] = "application/json" return headers def _build_cookies(auth_cookie: str, csrf_cookie: str) -> dict: """构建包含 auth 和 csrf 的完整 cookies 字典。""" cookies = {"fastapiusersauth": auth_manager._extract_auth_value(auth_cookie)} if csrf_cookie: cookies["fastapiusersoauthcsrf"] = csrf_cookie return cookies async def create_chat_session_with_cookie(client: httpx.AsyncClient) -> Tuple[str, str, str, str]: # 最大重试次数 = 账号总数（确保每个账号都有机会被尝试），至少 3 次 max_attempts = max(len(auth_manager.accounts), 3) last_error = "" for attempt in range(max_attempts): try: auth_cookie, csrf_cookie, email = await auth_manager.get_valid_cookie(client) except RuntimeError as e: # 所有账号都已不可用 raise RuntimeError(f"所有账号均不可用: {e}") payload = {"persona_id": config.ONYX_PERSONA_ID, "description": None, "project_id": None} response = await client.post( f"{config.ONYX_BASE_URL}/api/chat/create-chat-session", headers=_headers(email, with_json=True), json=payload, cookies=_build_cookies(auth_cookie, csrf_cookie), timeout=httpx.Timeout(config.REQUEST_TIMEOUT, connect=15.0), ) if response.status_code == 401: logger.warning("create_chat_session HTTP 401: Cookie 过期，尝试刷新... (第 %d/%d 次)", attempt + 1, max_attempts) await auth_manager.report_unauthorized(client, auth_cookie) continue elif response.status_code == 403: logger.warning("create_chat_session HTTP 403: 额度耗尽或被封禁，尝试重新登录后切换账号 (第 %d/%d 次)", attempt + 1, max_attempts) await auth_manager.report_forbidden(client, auth_cookie) last_error = "403 Forbidden" continue if response.status_code != 200: raise RuntimeError(f"Onyx create-chat-session HTTP {response.status_code}: {response.text[:300]}") data = response.json() chat_session_id = data.get("chat_session_id") or data.get("id") if not chat_session_id: raise RuntimeError(f"create-chat-session missing chat_session_id: {data}") return chat_session_id, auth_cookie, csrf_cookie, email raise RuntimeError(f"所有 {max_attempts} 个账号均已尝试失败 (最后错误: {last_error})") async def create_chat_session(client: httpx.AsyncClient) -> str: session_id, _, _, _ = await create_chat_session_with_cookie(client) return session_id async def _delete_session_bg( client: httpx.AsyncClient, chat_session_id: str, auth_cookie: str, csrf_cookie: str, email: str ) -> None: try: await timing.micro_delay("between_requests") await client.request( "DELETE", f"{config.ONYX_BASE_URL}/api/chat/delete-chat-session/{chat_session_id}", headers=_headers(email, with_json=True), cookies=_build_cookies(auth_cookie, csrf_cookie), timeout=5.0, ) except Exception as e: logger.debug("Failed to background delete chat session %s: %s", chat_session_id, e) async def stream_chat( client: httpx.AsyncClient, messages: list, model_name: str, ) -> AsyncGenerator[Tuple[str, str], None]: created_sessions = [] has_yielded_data = False # 方案 A：跟踪是否已经开始发送数据 try: # 确保 session 和后续消息使用同一个 cookie（同时携带 CSRF token） chat_session_id, auth_cookie, csrf_cookie, email = await create_chat_session_with_cookie(client) created_sessions.append((chat_session_id, auth_cookie, csrf_cookie, email)) provider, version = _resolve_model(model_name) payload = { "message": _build_prompt(messages), "chat_session_id": chat_session_id, "parent_message_id": None, "file_descriptors": [], "internal_search_filters": { "source_type": None, "document_set": None, "time_cutoff": None, "tags": [], }, "deep_research": False, "forced_tool_id": None, "llm_override": { "temperature": 0.5, "model_provider": provider, "model_version": version, }, "origin": config.ONYX_ORIGIN, } await timing.micro_delay("typing") for attempt in range(3): async with client.stream( "POST", f"{config.ONYX_BASE_URL}/api/chat/send-chat-message", headers=_headers(email, with_json=True), json=payload, cookies=_build_cookies(auth_cookie, csrf_cookie), timeout=httpx.Timeout(config.REQUEST_TIMEOUT, connect=15.0), ) as response: if response.status_code == 401: if attempt < 2 and not has_yielded_data: logger.warning("stream_chat HTTP 401: Cookie expired. Refreshing token... (第 %d/3 次)", attempt + 1) await auth_manager.report_unauthorized(client, auth_cookie) # 需要使用刷新后的 token 重建完整会话 chat_session_id, auth_cookie, csrf_cookie, email = await create_chat_session_with_cookie(client) created_sessions.append((chat_session_id, auth_cookie, csrf_cookie, email)) payload["chat_session_id"] = chat_session_id continue else: raise RuntimeError("Onyx stream_chat failed - token expired and retry failed") elif response.status_code == 403: logger.warning("stream_chat HTTP 403: Quota exhausted or account banned.") await auth_manager.report_forbidden(client, auth_cookie) if attempt < 2 and not has_yielded_data: chat_session_id, auth_cookie, csrf_cookie, email = await create_chat_session_with_cookie(client) created_sessions.append((chat_session_id, auth_cookie, csrf_cookie, email)) payload["chat_session_id"] = chat_session_id continue else: raise RuntimeError("Onyx stream_chat failed due to HTTP 403 Forbidden") if response.status_code != 200: body = await response.aread() raise RuntimeError(f"Onyx send-chat-message HTTP {response.status_code}: {body.decode(errors='replace')[:300]}") async for line in response.aiter_lines(): if not line: continue line = line.strip() if not line: continue try: item = json.loads(line) except json.JSONDecodeError: continue obj = item.get("obj", {}) item_type = obj.get("type") if item_type == "reasoning_delta": delta = obj.get("reasoning", "") if delta: has_yielded_data = True yield "thinking", delta elif item_type == "message_delta": delta = obj.get("content", "") if delta: has_yielded_data = True yield "text", delta elif item_type == "stop": break # Successfully finished without 401 break finally: for sid, auth, csrf, sid_email in created_sessions: asyncio.create_task(_delete_session_bg(client, sid, auth, csrf, sid_email)) async def full_chat(client: httpx.AsyncClient, messages: list, model_name: str) -> Tuple[str, str]: text_parts = [] thinking_parts = [] async for item_type, content in stream_chat(client, messages, model_name): if item_type == "thinking": thinking_parts.append(content) else: text_parts.append(content) return "".join(text_parts), "".join(thinking_parts)