监控 Microsoft Azure OpenAI

import os

from flask import Flask
from openai import AzureOpenAI
from opentelemetry import trace

from monitor import count_completion_requests_and_tokens


# Initialize Flask app
app = Flask(__name__)

# Set OpenAI API key
client = AzureOpenAI(
    api_key=os.getenv("AZURE_OPENAI_API_KEY"),
    api_version=os.getenv("OPENAI_API_VERSION"),
    azure_endpoint=os.getenv("AZURE_OPENAI_ENDPOINT"),
)

# Monkey-patch the openai.Completion.create function
client.chat.completions.create = count_completion_requests_and_tokens(
    client.chat.completions.create
)

tracer = trace.get_tracer("counter")


@app.route("/completion")
@tracer.start_as_current_span("completion")
def completion():
    response = client.chat.completions.create(
        model="gpt-4",
        messages=[
            {
                "role": "user",
                "content": "How do I send my APM data to Elastic Observability?",
            }
        ],
        max_tokens=20,
        temperature=0,
    )
    return response.choices[0].message.content.strip()

def count_completion_requests_and_tokens(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        counters["completion_count"] += 1
        response = func(*args, **kwargs)

        token_count = response.usage.total_tokens
        prompt_tokens = response.usage.prompt_tokens
        completion_tokens = response.usage.completion_tokens
        cost = calculate_cost(response)
        strResponse = json.dumps(response, default=str)

        # Set OpenTelemetry attributes
        span = trace.get_current_span()
        if span:
            span.set_attribute("completion_count", counters["completion_count"])
            span.set_attribute("token_count", token_count)
            span.set_attribute("prompt_tokens", prompt_tokens)
            span.set_attribute("completion_tokens", completion_tokens)
            span.set_attribute("model", response.model)
            span.set_attribute("cost", cost)
            span.set_attribute("response", strResponse)
        return response

    return wrapper

def calculate_cost(response):
    if response.model in ["gpt-4", "gpt-4-0314"]:
        cost = (
            response.usage.prompt_tokens * 0.03
            + response.usage.completion_tokens * 0.06
        ) / 1000
    elif response.model in ["gpt-4-32k", "gpt-4-32k-0314"]:
        cost = (
            response.usage.prompt_tokens * 0.06
            + response.usage.completion_tokens * 0.12
        ) / 1000
    elif "gpt-3.5-turbo" in response.model:
        cost = response.usage.total_tokens * 0.002 / 1000
    elif "davinci" in response.model:
        cost = response.usage.total_tokens * 0.02 / 1000
    elif "curie" in response.model:
        cost = response.usage.total_tokens * 0.002 / 1000
    elif "babbage" in response.model:
        cost = response.usage.total_tokens * 0.0005 / 1000
    elif "ada" in response.model:
        cost = response.usage.total_tokens * 0.0004 / 1000
    else:
        cost = 0
    return cost