app.py

import uuid
import json
import time
import os
import re
import random
import requests
from flask import Flask, request, Response, stream_with_context, jsonify
from flask_cors import CORS
from functools import lru_cache
from concurrent.futures import ThreadPoolExecutor

app = Flask(__name__)
CORS(app, resources={r"/*": {"origins": "*"}})

executor = ThreadPoolExecutor(max_workers=10)

# 使用LRU缓存，最大容量为10，用于缓存文件内容
@lru_cache(maxsize=10)
def read_file(filename):
    try:
        with open(filename, 'r') as f:
            return f.read().strip()
    except FileNotFoundError:
        print(f"文件 {filename} 未找到")
        return ""
    except Exception as e:
        print(f"读取文件 {filename} 时发生错误: {e}")
        return ""

# 从环境变量或文件中获取信息
def get_env_or_file(env_var, filename):
    return os.environ.get(env_var) or read_file(filename)

NOTDIAMOND_URLS = [
    'https://chat.notdiamond.ai',
    'https://chat.notdiamond.ai/mini-chat'
]

# 从预设的URL列表中随机选择一个
def get_notdiamond_url():
    return random.choice(NOTDIAMOND_URLS)

# 使用LRU缓存，缓存头信息，最大容量为1
@lru_cache(maxsize=1)
def get_notdiamond_headers():
    return {
        'accept': 'text/event-stream',
        'accept-language': 'zh-CN,zh;q=0.9',
        'content-type': 'application/json',
        'next-action': get_env_or_file('NEXT_ACTION', 'next_action.txt'),
        'user-agent': 'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/128.0.0.0 Safari/537.36',
        'cookie': get_env_or_file('COOKIES', 'cookies.txt')
    }

# 模型信息字典，包含各模型的提供商和映射
MODEL_INFO = {
    "gpt-4-turbo-2024-04-09": {
        "provider": "openai",
        "mapping": "gpt-4-turbo-2024-04-09"
    },
    "gemini-1.5-pro-exp-0801": {
        "provider": "google",
        "mapping": "models/gemini-1.5-pro-exp-0801"
    },
    "Meta-Llama-3.1-70B-Instruct-Turbo": {
        "provider": "togetherai",
        "mapping": "meta.llama3-1-70b-instruct-v1:0"
    },
    "Meta-Llama-3.1-405B-Instruct-Turbo": {
        "provider": "togetherai",
        "mapping": "meta.llama3-1-405b-instruct-v1:0"
    },
    "llama-3.1-sonar-large-128k-online": {
        "provider": "perplexity",
        "mapping": "llama-3.1-sonar-large-128k-online"
    },
    "gemini-1.5-pro-latest": {
        "provider": "google",
        "mapping": "models/gemini-1.5-pro-latest"
    },
    "claude-3-5-sonnet-20240620": {
        "provider": "anthropic",
        "mapping": "anthropic.claude-3-5-sonnet-20240620-v1:0"
    },
    "claude-3-haiku-20240307": {
        "provider": "anthropic",
        "mapping": "anthropic.claude-3-haiku-20240307-v1:0"
    },
    "gpt-4o-mini": {
        "provider": "openai",
        "mapping": "gpt-4o-mini"
    },
    "gpt-4o": {
        "provider": "openai",
        "mapping": "gpt-4o"
    },
    "mistral-large-2407": {
        "provider": "mistral",
        "mapping": "mistral.mistral-large-2407-v1:0"
    }
}

# 生成系统指纹，用于追踪和安全目的
def generate_system_fingerprint():
    return f"fp_{uuid.uuid4().hex[:10]}"

# 创建OpenAI格式的分块内容
def create_openai_chunk(content, model, finish_reason=None, usage=None):
    system_fingerprint = generate_system_fingerprint()
    chunk = {
        "id": f"chatcmpl-{uuid.uuid4()}",
        "object": "chat.completion.chunk",
        "created": int(time.time()),
        "model": model,
        "system_fingerprint": system_fingerprint,
        "choices": [
            {
                "index": 0,
                "delta": {"content": content} if content else {},
                "logprobs": None,
                "finish_reason": finish_reason
            }
        ]
    }
    if usage is not None:
        chunk["usage"] = usage
    return chunk

# 处理带美元符号的字符串
def process_dollars(s):
    return s.replace('$$', '$')

def stream_notdiamond_response(response, model):
    buffer = ""
    last_content = ""
    total_tokens = 0

    for chunk in response.iter_content(chunk_size=1024):
        if chunk:
            buffer += chunk.decode('utf-8')
            lines = buffer.split('\n')
            buffer = lines.pop()
            for line in lines:
                if line.strip():
                    match = re.match(r'^([0-9a-zA-Z]+):(.*)$', line)
                    if match:
                        try:
                            _, content = match.groups()
                            data = json.loads(content)
                            content = ''
                            if 'output' in data and 'curr' in data['output']:
                                content = process_dollars(data['output']['curr'])
                            elif 'curr' in data:
                                content = process_dollars(data['curr'])
                            elif 'diff' in data and isinstance(data['diff'], list):
                                if len(data['diff']) > 1:
                                    new_content = process_dollars(data['diff'][1])
                                    content = last_content + new_content
                                elif len(data['diff']) == 1:
                                    content = last_content

                            if content:
                                total_tokens += len(content.split()) 
                                last_content = content
                                yield create_openai_chunk(content, model)
                        except json.JSONDecodeError:
                            print(f"Error processing line: {line}")  # 可考虑替换为日志记录
    
    usage = {
        "prompt_tokens": 0,
        "completion_tokens": total_tokens,
        "total_tokens": total_tokens
    }
    yield create_openai_chunk('', model, 'stop', usage=usage)

# 处理非流式响应
def handle_non_stream_response(response, model):
    full_content = ""
    for chunk in stream_notdiamond_response(response, model):
        if chunk['choices'][0]['delta'].get('content'):
            full_content += chunk['choices'][0]['delta']['content']

    return jsonify({
        "id": f"chatcmpl-{uuid.uuid4()}",
        "object": "chat.completion",
        "created": int(time.time()),
        "model": model,
        "system_fingerprint": generate_system_fingerprint(),
        "choices": [
            {
                "index": 0,
                "message": {
                    "role": "assistant",
                    "content": full_content
                },
                "finish_reason": "stop"
            }
        ],
        "usage": {
            "prompt_tokens": len(full_content) // 4,
            "completion_tokens": len(full_content) // 4,
            "total_tokens": len(full_content) // 2
        }
    })

# 生成流式响应
def generate_stream_response(response, model):
    for chunk in stream_notdiamond_response(response, model):
        yield f"data: {json.dumps(chunk)}\n\n"
    yield "data: [DONE]\n\n"

# 获取模型列表的API
@app.route('/ai/v1/models', methods=['GET'])
def proxy_models():
    models = [
        {
            "id": model_id,
            "object": "model",
            "provider": info["provider"]
        } for model_id, info in MODEL_INFO.items()
    ]
    return jsonify({
        "object": "list",
        "data": models
    })

# 处理请求的API
@app.route('/ai/v1/chat/completions', methods=['POST'])
def handle_request():
    try:
        request_data = request.get_json()
        messages = request_data.get('messages')
        model_id = request_data.get('model', '')
        model = MODEL_INFO.get(model_id, {}).get('mapping', model_id)
        stream = request_data.get('stream', False)

        payload = {
            "messages": messages,
            "model": model,
            "stream": stream,
            "frequency_penalty": request_data.get('frequency_penalty', 0),
            "presence_penalty": request_data.get('presence_penalty', 0),
            "temperature": request_data.get('temperature', 0.8),
            "top_p": request_data.get('top_p', 1)
        }

        headers = get_notdiamond_headers()
        url = get_notdiamond_url()
        
        future = executor.submit(requests.post, url, headers=headers, json=[payload], stream=True)
        response = future.result()
        response.raise_for_status()

        if stream:
            return Response(stream_with_context(generate_stream_response(response, model)), content_type='text/event-stream')
        else:
            return handle_non_stream_response(response, model)
         
    except json.JSONDecodeError as e:
        return jsonify({
            'error': 'Invalid JSON format',
            'details': str(e)
        }), 400  # 返回400错误，因为请求数据格式不正确
    except requests.exceptions.RequestException as e:
        return jsonify({
            'error': 'Request failed',
            'details': str(e)
        }), 502  # 返回502错误，表示请求失败
    except Exception as e:
        return jsonify({
            'error': 'Internal Server Error',
            'details': str(e)
        }), 500

if __name__ == '__main__':
    port = int(os.environ.get("PORT", 8000))
    app.run(debug=True, host='0.0.0.0', port=port, threaded=True)