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:
                # chunk.choices[0].delta.content 또는 ''를 반환
                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)