OpenLLMetry

[[inputs.opentelemetry]]
  ## customer_tags はホワイトリストとして機能し、タグがデータセンターへ送信されるのを防ぎます。
  ## すべての . は _ に置き換えられます。例:
    customer_tags = ["llm.request.type","traceloop.entity.path","llm.is_streaming","gen_ai.openai.api_base","gen_ai.prompt.1.content","gen_ai.response.model","gen_ai.completion.0.content","gen_ai.request.model","gen_ai.request.temperature","gen_ai.system","traceloop.workflow.name"]
  ...

from traceloop.sdk import Traceloop

# OpenLit を初期化
# Traceloop.init()

Traceloop.init(app_name="kimi_openllmetry_stream_flask")

import os
import httpx
from flask import Flask, request, Response,jsonify,stream_with_context
from traceloop.sdk import Traceloop
from traceloop.sdk.decorators import workflow,task
from openai import OpenAI

from opentelemetry.instrumentation.flask import FlaskInstrumentor


app = Flask(__name__)
# FlaskInstrumentor を使用して Flask アプリケーションを自動計装する
FlaskInstrumentor().instrument_app(app)

# OpenLit を初期化
Traceloop.init(app_name="kimi_openllmetry_stream_flask")

# 環境変数から API Key を取得する
api_key = os.getenv("MOONSHOT_API_KEY")
if not api_key:
    raise ValueError("MOONSHOT_API_KEY 環境変数を設定してください")

client = OpenAI(
    api_key=api_key,
    base_url="https://api.moonshot.cn/v1",
)

def estimate_token_count(input_messages) -> int:
    """
    入力メッセージのトークン数を計算します。
    """
    try:
        header = {
            "Authorization": f"Bearer {api_key}",
        }
        data = {
            "model": "moonshot-v1-128k",
            "messages": input_messages,
        }
        with httpx.Client() as client:
            print("API を呼び出しています")
            r = client.post("https://api.moonshot.cn/v1/tokenizers/estimate-token-count", headers=header, json=data)
            r.raise_for_status()
            response_data = r.json()
            print(response_data["data"]["total_tokens"])
            return response_data["data"]["total_tokens"]
    except httpx.RequestError as e:
        print(f"リクエストに失敗しました: {e}")
        raise
    except (KeyError, ValueError) as e:
        print(f"レスポンスの解析に失敗しました: {e}")
        raise

def select_model(input_messages, max_tokens=1024) -> str:
    """
    入力されたコンテキストメッセージと想定される max_tokens 値に応じて、適切なモデルを選択します。
    """
    if not isinstance(max_tokens, int) or max_tokens <= 0:
        raise ValueError("max_tokens は正の整数でなければなりません")

    prompt_tokens = estimate_token_count(input_messages)
    total_tokens = prompt_tokens + max_tokens

    if total_tokens <= 8 * 1024:
        return "moonshot-v1-8k"
    elif total_tokens <= 32 * 1024:
        return "moonshot-v1-32k"
    elif total_tokens <= 128 * 1024:
        return "moonshot-v1-128k"
    else:
        raise ValueError("トークン数が上限を超えています 😢")

@app.route('/ask', methods=['POST'])
@workflow(name="ask_workflow")
def ask():
    data = request.json
    messages = data.get('messages')
    max_tokens = data.get('max_tokens', 2048)

    if not messages:
        return jsonify({"error": "messages フィールドは空にできません"}), 400

    try:
        model = select_model(messages, max_tokens)

        completion = client.chat.completions.create(
            model=model,
            messages=messages,
            max_tokens=max_tokens,
            temperature=0.3,
            stream=True  # ストリーミング生成を有効化
        )

        def generate():
            for chunk in completion:
                # yield chunk.choices[0].delta.content or ''
                delta = chunk.choices[0].delta
                if delta.content:
                    print(delta.content, end="")
                    yield delta.content or ''

        return Response(stream_with_context(generate()), content_type='text/event-stream')
    except Exception as e:
        return jsonify({"error": str(e)}), 500
if __name__ == '__main__':
    app.run(debug=True,port=5001)