core/utils.py

import asyncio
import codecs
import hashlib
import os
from datetime import datetime
import json
import uuid
# from http.client import HTTPException
from typing import Dict, Any, Optional
from fastapi import HTTPException
import ssl
import httpx
from httpx import ConnectError, TransportError
from starlette import status

from core.config import get_settings
from core.logger import setup_logger
from core.models import ChatRequest
from core.sqlite_store import SQLiteConnectionPool, DatabaseManager

settings = get_settings()
logger = setup_logger(__name__)

def decode_unicode_escape(s):
    # 检查输入是否为字典类型
    if isinstance(s, dict):
        return s
    # 如果需要，将输入转换为字符串
    if not isinstance(s, (str, bytes)):
        s = str(s)
    # 如果是字符串，转换为字节
    if isinstance(s, str):
        s = s.encode('utf-8')
    return codecs.decode(s, 'unicode_escape')

FIREBASE_API_KEY = settings.FIREBASE_API_KEY
async def refresh_token_via_rest(refresh_token):
    # Firebase Auth REST API endpoint
    url = f"https://securetoken.googleapis.com/v1/token?key={FIREBASE_API_KEY}"

    payload = {
        'grant_type': 'refresh_token',
        'refresh_token': refresh_token
    }

    try:
        async with httpx.AsyncClient() as client:
            response = await client.post(url, json=payload)
            if response.status_code == 200:
                data = response.json()
                print(json.dumps(data, indent=2))
                # return {
                #     'id_token': data['id_token'],
                #     'refresh_token': data.get('refresh_token'),
                #     'expires_in': data['expires_in']
                # }
                return data['id_token']
            else:
                print(f"刷新失败: {response.text}")
                return None
    except Exception as e:
        print(f"请求异常: {e}")
        return None


async def sign_in_with_idp():
    url = "https://identitytoolkit.googleapis.com/v1/accounts:signInWithIdp"

    # 查询参数
    params = {
        "key": FIREBASE_API_KEY
    }

    # 请求头
    headers = {
        "X-Client-Version": "Node/JsCore/10.5.2/FirebaseCore-web",
        "X-Firebase-gmpid": "1:252179682924:web:9c80c6a32cb4682cbfaa49",
        "Content-Type": "application/json",
        "User-Agent": "node-fetch/1.0 (+https://github.com/bitinn/node-fetch)"
    }

    # 请求体
    data = {
        "requestUri": "http://localhost",
        "returnSecureToken": True,
        "postBody": f"&id_token={settings.AUTHORIZATION_TOKEN}&providerId=google.com"
    }
    print("Request Headers:", json.dumps(headers, indent=2))  # 格式化打印
    print("Request Body:", json.dumps(data, indent=2))  # 格式化打印
    print("Request params:", json.dumps(params, indent=2))  # 格式化打印

    async with httpx.AsyncClient() as client:
        response = await client.post(
            url,
            params=params,
            headers=headers,
            json=data
        )
        # 检查状态码
        if response.status_code == 200:
            return response.json()
        else:
            raise Exception(f"Request failed with status code: {response.status_code}")

async def handle_firebase_response(response) -> str:
    try:
        # 如果响应是字典（已经解析的 JSON）
        if isinstance(response, dict):
            print(json.dumps(response, indent=2))
            if response.get('error', {}).get('code') == 400:
                print("Invalid id_token in IdP response")
            # 保存refresh_token到配置中
            if 'refreshToken' in response:
                os.environ["REFRESH_TOKEN"] = response['refreshToken']
            if 'idToken' in response:
                return response['idToken']
            else:
                raise ValueError("dict case Response does not contain idToken")

        # 如果响应是 Response 对象
        elif hasattr(response, 'status_code'):
            if response.status_code == 200:
                data = response.json()
                print(data)
                # 保存refresh_token到配置中
                if 'refreshToken' in data:
                    os.environ["REFRESH_TOKEN"] = data['refreshToken']
                if 'idToken' in data:
                    return data['idToken']
                else:
                    raise ValueError("response case Response does not contain idToken")

            # 处理其他状态码
            elif response.status_code == 400:
                error_data = response.json()
                raise ValueError(f"Bad Request: {error_data.get('error', {}).get('message', 'Unknown error')}")
            elif response.status_code == 401:
                raise ValueError("Unauthorized: Invalid credentials")
            elif response.status_code == 403:
                raise ValueError("Forbidden: Insufficient permissions")
            elif response.status_code == 404:
                raise ValueError("Not Found: Resource doesn't exist")
            else:
                raise ValueError(f"Unexpected status code: {response.status_code}")

        else:
            raise ValueError(f"Unexpected response type: {type(response)}")

    except json.JSONDecodeError:
        raise ValueError("Invalid JSON response")
    except Exception as e:
        raise ValueError(f"Error processing response: {str(e)}")

# SHA-256
def _sha256_hash(text):
    sha256 = hashlib.sha256()
    sha256.update(text.encode('utf-8'))
    return sha256.hexdigest()

# 处理字典列表
def sha256_hash_messages(messages):
    # 只提取 role 为 "user" 的消息的 content 字段
    message_data = [str(msg['content']) for msg in messages if msg['role'] == "user"]
    print("Filtered contents:", message_data)  # 调试用
    json_str = json.dumps(message_data, sort_keys=True)
    print("JSON string:", json_str)  # 调试用
    return hashlib.sha256(json_str.encode('utf-8')).hexdigest()


def create_chat_completion_data(
        content: str, model: str, timestamp: int, finish_reason: Optional[str] = None
) -> Dict[str, Any]:
    return {
        "id": f"chatcmpl-{uuid.uuid4()}",
        "object": "chat.completion.chunk",
        "created": timestamp,
        "model": model,
        "choices": [
            {
                "index": 0,
                "delta": {"content": content, "role": "assistant"},
                "finish_reason": finish_reason,
            }
        ],
        "usage": None,
    }
pool = SQLiteConnectionPool('~/tmp/merlin-sqlite.db', max_connections=5)
db = DatabaseManager(pool)
async def process_streaming_response(request: ChatRequest, app_secret: str):
    # 创建自定义 SSL 上下文
    ssl_context = ssl.create_default_context()
    ssl_context.check_hostname = True
    ssl_context.verify_mode = ssl.CERT_REQUIRED
    # 1. 获取消息列表并转换为字典列表
    # previous_messages = request['messages'][:-1]
    # previous_messages = [msg.model_dump() for msg in request.messages[:-1]]
    # previous_messages = [dict(msg) for msg in request.messages[:-1]]
    # 如果用户有系统提示词，则不包含在上下文中 merlin chat
    previous_messages = [
        msg.model_dump(include={'role', 'content'})
        for msg in request.messages[:-1]
        if msg.role == "user"
    ]
    # 对话上下文处理
    # sqlite数据存储结构
    # api_key chat_id parent_id sha256_hash # 用于快速查找
    message = [dict(msg) for msg in request.messages]
    if not previous_messages:
        chat_id = str(uuid.uuid4())
        child_id = str(uuid.uuid4())
        parent_id = "root"
        # messages = [msg.model_dump() for msg in request.messages]
        db.insert_context_record(app_secret, chat_id, child_id, sha256_hash_messages(message))
        print("Insert new context record: hash=", sha256_hash_messages(message))
    else:
        sha256_hash = sha256_hash_messages(previous_messages)
        print("func get_context_record_by_sha256_hash SHA-256 hash:", sha256_hash)
        context_record = db.get_context_record_by_sha256_hash(sha256_hash)
        chat_id = context_record['chat_id']
        child_id = str(uuid.uuid4())
        parent_id = context_record['parent_id']
    json_data = {
        "attachments": [],
        "chatId": chat_id,
        "language": "AUTO",
        "message": {
            "childId": child_id,
            "content": request.messages[-1].content,
            "context": "",
            "id": str(uuid.uuid4()),
            "parentId": parent_id
        },
        "metadata": {
            "largeContext": False,
            "merlinMagic": False,
            "proFinderMode": False,
            "webAccess": False
        },
        "mode": "UNIFIED_CHAT",
        # "model": "claude-3.5-sonnet"
        "model": request.model
    }
    async with httpx.AsyncClient(
            verify=ssl_context,
            # timeout=30.0,  # 增加超时时间
            # http2=True    # 启用 HTTP/2
    ) as client:
        try:
            request_headers = {**settings.HEADERS, 'authorization': f"Bearer {os.getenv('TOKEN', '')}"} # 从环境变量中获取新的TOKEN

            print("Request Headers:", json.dumps(request_headers, indent=2))  # 格式化打印
            print("Request Body:", json.dumps(json_data, indent=2))  # 格式化打印
            async with client.stream(
                    "POST",
                    f"https://arcane.getmerlin.in/v1/thread/unified",
                    headers=request_headers,
                    json=json_data,
                    timeout=100,
            ) as response:
                response.raise_for_status()
                timestamp = int(datetime.now().timestamp())
                async for line in response.aiter_lines():
                    # print(f"{type(line)}: {line}")
                    if line and line.startswith("data: "):
                        try:
                            data_str = line[6:]  # 去掉 'data: ' 前缀
                            # 解析JSON
                            json_data = json.loads(data_str)
                            # print(json_data)
                            # 解码Unicode转义序列
                            # human_readable = decode_unicode_escape(decoded_chunk)
                            # print(human_readable)
                            # 每日用量超限时 返回的数据处理
                            # data: {"message":"Daily pro model limit exhausted. Switch to any base model or Go Pro!","type":"USAGE_LIMIT_REACHED_FREE"}
                            if 'type' in json_data and json_data['type'] == 'USAGE_LIMIT_REACHED_FREE':
                                content = json_data['message']
                                print(content, end='', flush=True)  # 使用 end='' 来避免额外换行
                                yield f"data: {json.dumps(create_chat_completion_data(content, request.model, timestamp))}\n\n"
                            # data: {"status":"system","data":{"content":" ","eventType":"DONE"}}
                            if 'data' in json_data and 'eventType' in json_data['data'] and json_data['data']['eventType'] == 'DONE':
                                await response.aclose()
                                break
                            #     yield f"data: {json.dumps(create_chat_completion_data('', 'gpt-4o-mini', timestamp, 'stop'))}\n\n"
                            #     yield "data: [DONE]\n\n"
                            #     await asyncio.sleep(0)
                                # return
                            # 提取并解码content内容
                            # if 'data' in human_readable and 'content' in human_readable['data']:
                            # data: {"data":{"content":"xx"}}
                            if 'data' in json_data and 'content' in json_data['data'] and json_data['data']['content'] != '':
                                content = json_data['data']['content']
                                # human_readable_content = decode_unicode_escape(content)
                                print(content, end='', flush=True)  # 使用 end='' 来避免额外换行
                                yield f"data: {json.dumps(create_chat_completion_data(content, request.model, timestamp))}\n\n"

                        except json.JSONDecodeError as e:
                            print(f"JSON解析错误: {e}")
                            print(f"原始数据: {line}")
                            continue

                yield f"data: {json.dumps(create_chat_completion_data('', request.model, timestamp, 'stop'))}\n\n"
                yield "data: [DONE]\n\n"
                # 更新sha256_hash
                db.update_context_record_by_chat_id(app_secret, chat_id, child_id, sha256_hash_messages(message))
        except ConnectError as e:
            logger.error(f"Connection error details: {str(e)}")
            raise HTTPException(
                status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
                detail="Service temporarily unavailable. Please try again later."
            )
        except TransportError as e:
            logger.error(f"Transport error details: {str(e)}")
            raise HTTPException(
                status_code=status.HTTP_502_BAD_GATEWAY,
                detail="Network transport error occurred."
            )
        except httpx.HTTPStatusError as e:
            # 这里需要处理401错误
            logger.error(f"HTTP error occurred: {e}")
            raise HTTPException(status_code=e.response.status_code, detail=str(e))
        except httpx.RequestError as e:
            logger.error(f"Error occurred during request: {e}")
            raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=str(e))
        finally:
            await response.aclose()
        # except ConnectError as e:
        #     logger.error(f"Connection error: {str(e)}")
        #     raise HTTPException(
        #         status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
        #         detail="Service temporarily unavailable. Please try again later."
        #     )
        # except asyncio.CancelledError:
        #     logger.info("Request was cancelled by client")
        #     raise  # Let the framework handle cancellation
        # except Exception as e:
        #     logger.error(f"Unexpected error: {str(e)}", exc_info=True)
        #     raise HTTPException(
        #         status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
        #         detail="An unexpected error occurred"
        #     )
        # finally:
        #     # 确保资源被正确清理
        #     if 'client' in locals():
        #         await client.aclose()

# async def process_non_streaming_response(request: ChatRequest):
#     json_data = {
#         "messages": [message_to_dict(msg) for msg in request.messages],
#         "previewToken": None,
#         "userId": None,
#         "codeModelMode": True,
#         "agentMode": {},
#         "trendingAgentMode": {},
#         "isMicMode": False,
#         "userSystemPrompt": None,
#         "maxTokens": request.max_tokens,
#         "playgroundTopP": request.top_p,
#         "playgroundTemperature": request.temperature,
#         "isChromeExt": False,
#         "githubToken": None,
#         "clickedAnswer2": False,
#         "clickedAnswer3": False,
#         "clickedForceWebSearch": False,
#         "visitFromDelta": False,
#         "mobileClient": False,
#         "userSelectedModel": MODEL_MAPPING.get(request.model),
#         "validated": validate.getVid(),
#         "webSearchModePrompt": True if request.model.endswith("-search") else False
#     }
#     full_response = ""
#     async with httpx.AsyncClient() as client:
#         async with client.stream(
#                 method="POST", url=f"{BASE_URL}/api/chat", headers=settings.HEADERS, json=json_data
#         ) as response:
#             async for chunk in response.aiter_text():
#                 full_response += chunk
#     if "https://www.blackbox.ai" in full_response:
#         validate.getVid(True)
#         full_response = "vid已刷新，重新对话即可"
#     if full_response.startswith("$@$v=undefined-rv1$@$"):
#         full_response = full_response[21:]
#     return {
#         "id": f"chatcmpl-{uuid.uuid4()}",
#         "object": "chat.completion",
#         "created": int(datetime.now().timestamp()),
#         "model": request.model,
#         "choices": [
#             {
#                 "index": 0,
#                 "message": {"role": "assistant", "content": full_response},
#                 "finish_reason": "stop",
#             }
#         ],
#         "usage": None,
#     }

# POST /v1/accounts:signInWithIdp?key=AIzaSyCMMynYm5VRHj1NOwkfWinX-HYsFArdUbk HTTP/1.1
# X-Client-Version: Node/JsCore/10.5.2/FirebaseCore-web
# X-Firebase-gmpid: 1:252179682924:web:9c80c6a32cb4682cbfaa49
# Content-Type: application/json
# Accept: */*
# Content-Length: 1269
# User-Agent: node-fetch/1.0 (+https://github.com/bitinn/node-fetch)
# Accept-Encoding: gzip,deflate
# Connection: close
# Host: identitytoolkit.googleapis.com
#
# {"requestUri":"http://localhost","returnSecureToken":true,"postBody":"&id_token=aaa.eyJpc3MiOiJodHRwczovL2FjY291bnRzLmdvb2dsZS5jb20iLCJhenAiOiIyNTIxNzk2ODI5MjQtYWhmcTh2amtyOGEwOGFrMjMxMGFiajZqM29jZ2I5azIuYXBwcy5nb29nbGV1c2VyY29udGVudC5jb20iLCJhdWQiOiIyNTIxNzk2ODI5MjQtYWhmcTh2amtyOGEwOGFrMjMxMGFiajZqM29jZ2I5azIuYXBwcy5nb29nbGV1c2VyY29udGVudC5jb20iLCJzdWIiOiIxMDI3Mzc4NDA0NjA0MDg3ODQ5NTEiLCJoZCI6ImNoYXRncHQubnljLm1uIiwiZW1haWwiOiJrYmluQGNoYXRncHQubnljLm1uIiwiZW1haWxfdmVyaWZpZWQiOnRydWUsImF0X2hhc2giOiJhQ3JLcVlzZ0JzT1JIQnFqQ2k1N0VBIiwibmFtZSI6ImtCaW4gbGVlIiwicGljdHVyZSI6Imh0dHBzOi8vbGgzLmdvb2dsZXVzZXJjb250ZW50LmNvbS9hL0FDZzhvY0ktUWJ1WHQzd3UzaUJsMUNkMzVIN1JrQ29xRW5NUmxPLWxwMEcxUVBsYm5ZTGd1dz1zOTYtYyIsImdpdmVuX25hbWUiOiJrQmluIiwiZmFtaWx5X25hbWUiOiJsZWUiLCJpYXQiOjE3MzYwNDAyODMsImV4cCI6MTczNjA0Mzg4M30.hC14X8oJhKzXOU_STfmmgXItPGOT_RYQKNtO_KpisXA6d8NRuCArAV6YOy10pvHVYElTk-Avpe2ymQCl58K2Itw05xpuGQ1EhF-8u_rUFmiBs0w1wtFPt57tnMdzbqyLs_OK7ndUYl-myVoTi-2JVM6P2rMSYqCLkM0jkAeGsiOLFjXw3wQpiOAmUVbiJ8mjztdK2jBMqK91C18vJ9BROkPog_rX5hQYPbTNvrdKPplNJb94NVssXXDuAUShO6ADOCJZ6EXj4IxyD5vyUP_0sHX9tdQs6wJXcmEgNBwgEPJ7DjdcBaEzeG9o-F3v9uMHjCbLCH4ben5KcolAqHVzmw&providerId=google.com"}
#
#
#
# async def sign_in_with_idp():
#     url = "https://identitytoolkit.googleapis.com/v1/accounts:signInWithIdp"
#
#     # 查询参数
#     params = {
#         "key": "AIzaSyCMMynYm5VRHj1NOwkfWinX-HYsFArdUbk"
#     }
#
#     # 请求头
#     headers = {
#         "X-Client-Version": "Node/JsCore/10.5.2/FirebaseCore-web",
#         "X-Firebase-gmpid": "1:252179682924:web:9c80c6a32cb4682cbfaa49",
#         "Content-Type": "application/json",
#         "User-Agent": "node-fetch/1.0 (+https://github.com/bitinn/node-fetch)"
#     }
#
#     # 请求体
#     data = {
#         "requestUri": "http://localhost",
#         "returnSecureToken": True,
#         "postBody": f"&id_token={settings.TOKEN}&providerId=google.com"
#     }
#
#     async with httpx.AsyncClient() as client:
#         response = await client.post(
#             url,
#             params=params,
#             headers=headers,
#             json=data
#         )
#         return response.json()
#
# async def main():
#     result = await sign_in_with_idp()
#     print(result)
#
# if __name__ == "__main__":
#     asyncio.run(main())