app.py

import os
import asyncio
import json
import time
import uuid

from flask import Flask, request, Response, jsonify
from flask_cors import CORS
import websockets

CONFIG = {
    "WS_URI": "wss://api.inkeep.com/graphql",
    "AUTH_TOKEN": f"Bearer {os.getenv('AUTH_TOKEN', 'xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx')}",
    "SUBSCRIBE_ID": str(uuid.uuid4()),
    "ORG_ID": "org_xxxxxxxxxxxxxxx",
    "INTEGRATION_ID": "xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx",
    "DEFAULT_MESSAGE": "Hello.",
    "MODEL": "claude",
}

# 是否启用上下文
ENABLE_CONTEXT = True

app = Flask(__name__)
CORS(app)


def process_messages(messages):
    """
    处理消息历史，将历史对话和当前消息组合成一个完整的消息
    Args:
        messages: 消息列表，每个消息包含 role 和 content
    Returns:
        str: 处理后的消息内容
    """
    if not messages:
        return CONFIG["DEFAULT_MESSAGE"]
    
    # 添加引导语
    instruction = "我们接下来的讨论是基于我们在使用Claude时遇到的一些问题。你只需要回答问题就行了，回答内容不要引用文档，明白吗？\n\n"
    
    # 构建完整的对话历史
    full_context = instruction
    
    # 处理历史消息（除了最后一条）
    for msg in messages[:-1]:
        role = msg["role"]
        content = msg["content"].strip()
        full_context += f"{role}: {content}\n\n"
    
    # 添加最新的消息
    current_message = messages[-1]["content"]
    full_message = f"{full_context}{current_message}"
    
    return full_message




async def perform_handshake(websocket, message_input):
    init_msg = {
        "type": "connection_init",
        "payload": {"headers": {"Authorization": CONFIG["AUTH_TOKEN"]}},
    }
    await websocket.send(json.dumps(init_msg))
    while True:
        resp = json.loads(await websocket.recv())
        if resp.get("type") == "connection_ack":
            break
    subscribe_msg = {
        "id": CONFIG["SUBSCRIBE_ID"],
        "type": "subscribe",
        "payload": {
            "variables": {
                "messageInput": message_input,
                "messageContext": None,
                "organizationId": CONFIG["ORG_ID"],
                "integrationId": CONFIG["INTEGRATION_ID"],
                "chatMode": "AUTO",
                "messageAttributes": {},
                "includeAIAnnotations": False,
                "environment": "production",
            },
            "extensions": {},
            "operationName": "OnNewSessionChatResult",
            "query": (
                "subscription OnNewSessionChatResult($messageInput: String!, $messageContext: String, $organizationId: ID!, "
                "$integrationId: ID, $chatMode: ChatMode, $filters: ChatFiltersInput, $messageAttributes: JSON, $tags: [String!], "
                "$workflowId: String, $context: String, $guidance: String, $includeAIAnnotations: Boolean!, $environment: String) {"
                "  newSessionChatResult(input: {messageInput: $messageInput, messageContext: $messageContext, organizationId: $organizationId, "
                "integrationId: $integrationId, chatMode: $chatMode, messageAttributes: $messageAttributes, environment: $environment}) {"
                "    isEnd sessionId message { id content __typename }"
                "  }"
                "}"
            ),
        },
    }
    await websocket.send(json.dumps(subscribe_msg))


async def openai_compatible_stream(message_input):
    async with websockets.connect(
        CONFIG["WS_URI"], subprotocols=["graphql-transport-ws"]
    ) as websocket:
        await perform_handshake(websocket, message_input)
        created = int(time.time())
        unique_id = f"chatcmpl-{uuid.uuid4()}"
        last_content = ""
        while True:
            raw = await websocket.recv()
            message = json.loads(raw)
            if message.get("type") == "next":
                content = message["payload"]["data"]["newSessionChatResult"]["message"][
                    "content"
                ]
                delta = content[len(last_content) :]
                if delta:
                    chunk = {
                        "id": unique_id,
                        "object": "chat.completion.chunk",
                        "created": created,
                        "model": CONFIG["MODEL"],
                        "choices": [
                            {
                                "delta": {"content": delta},
                                "index": 0,
                                "finish_reason": None,
                            }
                        ],
                    }
                    yield f"data: {json.dumps(chunk)}\n\n"
                    last_content = content
                if message["payload"]["data"]["newSessionChatResult"].get("isEnd"):
                    final_chunk = {
                        "id": unique_id,
                        "object": "chat.completion.chunk",
                        "created": created,
                        "model": CONFIG["MODEL"],
                        "choices": [{"delta": {}, "index": 0, "finish_reason": "stop"}],
                    }
                    yield f"data: {json.dumps(final_chunk)}\n\n"
                    break
        yield "data: [DONE]\n\n"


async def openai_compatible_complete(model_name, message_input):
    async with websockets.connect(
        CONFIG["WS_URI"], subprotocols=["graphql-transport-ws"]
    ) as websocket:
        await perform_handshake(websocket, message_input)
        created = int(time.time())
        unique_id = f"chatcmpl-{uuid.uuid4()}"
        last_content = ""
        while True:
            raw = await websocket.recv()
            message = json.loads(raw)
            if message.get("type") == "next":
                content = message["payload"]["data"]["newSessionChatResult"]["message"][
                    "content"
                ]
                last_content = content
                if message["payload"]["data"]["newSessionChatResult"].get("isEnd"):
                    break
        return {
            "id": unique_id,
            "object": "chat.completion",
            "created": created,
            "model": model_name,
            "choices": [
                {
                    "finish_reason": "stop",
                    "index": 0,
                    "logprobs": None,
                    "message": {
                        "role": "assistant",
                        "content": last_content,
                        "refusal": None,
                    },
                }
            ],
            "system_fingerprint": None,
        }


def sync_openai_stream(message_input):
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)
    gen = openai_compatible_stream(message_input)
    try:
        while True:
            try:
                yield loop.run_until_complete(gen.__anext__())
            except StopAsyncIteration:
                break
    finally:
        loop.close()


def sync_openai_complete(model_name, message_input):
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)
    try:
        return loop.run_until_complete(
            openai_compatible_complete(model_name, message_input)
        )
    finally:
        loop.close()


@app.route("/hf/v1/chat/completions", methods=["POST"])
def chat_completions():
    req = request.get_json(silent=True) or {}

    if req.get("model") != CONFIG["MODEL"]:
        return jsonify({"error": "Unsupported model."}), 400

    messages = req.get("messages", [])
    if ENABLE_CONTEXT:
        message_input = process_messages(messages)
    else:
        message_input = (
            messages[-1]["content"] if messages else CONFIG["DEFAULT_MESSAGE"]
        )
    if req.get("stream"):
        return Response(sync_openai_stream(message_input), mimetype="text/event-stream")
    result = sync_openai_complete(CONFIG["MODEL"], message_input)
    return jsonify(result)

if __name__ == "__main__":
    app.run(host="0.0.0.0", port=7860)