onyx.py

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)