Repository: microsoft/promptflow
File: src/promptflow-tracing/promptflow/tracing/_openai_utils.py (Python)

Code before:
            prompt_tokens += tokens_per_message
            for key, value in message.items():
                prompt_tokens += len(enc.encode(value))
                if key == "name":
                    prompt_tokens += tokens_per_name
        prompt_tokens += 3
        return prompt_tokens

    def _get_completion_tokens_for_chat_api(self, output, enc):
        completion_tokens = 0
        choices = output.get("choices")
        if isinstance(choices, list):
            for ch in choices:
                if isinstance(ch, dict):
                    message = ch.get("message")
                    if isinstance(message, dict):
                        content = message.get("content")
                        if isinstance(content, str):
                            completion_tokens += len(enc.encode(content))
        return completion_tokens

    def get_openai_metrics_for_completion_api(self, inputs, output):
        metrics = {}
        enc = self._get_encoding_for_completion_api(self._try_get_model(inputs, output))
        metrics["prompt_tokens"] = 0
        prompt = inputs.get("prompt")
        if isinstance(prompt, str):
            metrics["prompt_tokens"] = len(enc.encode(prompt))
        elif isinstance(prompt, list):
            for pro in prompt:
                metrics["prompt_tokens"] += len(enc.encode(pro))
        if isinstance(output, list):
            if IS_LEGACY_OPENAI:
                metrics["completion_tokens"] = len(output)
            else:
                metrics["completion_tokens"] = len(
                    [chunk for chunk in output if chunk.choices and chunk.choices[0].text]
                )
        else:
            metrics["completion_tokens"] = self._get_completion_tokens_for_completion_api(output, enc)
        metrics["total_tokens"] = metrics["prompt_tokens"] + metrics["completion_tokens"]
        return metrics

    def _get_encoding_for_completion_api(self, model):
        try:
            return tiktoken.encoding_for_model(model)
        except KeyError:
            return tiktoken.get_encoding("p50k_base")

    def _get_completion_tokens_for_completion_api(self, output, enc):
        completion_tokens = 0

Review comment (suggestion):
I have created a constant class `TokenKeys` for these 3 tokens, , consider using those instead of hard coding, no need to do it in this pr tho.

Code after:
            prompt_tokens += tokens_per_message
            for key, value in message.items():
                prompt_tokens += len(enc.encode(value))
                if key == "name":
                    prompt_tokens += tokens_per_name
        prompt_tokens += 3
        return prompt_tokens

    def _get_completion_tokens_for_chat_api(self, output, enc):
        completion_tokens = 0
        choices = output.get("choices")
        if isinstance(choices, list):
            for ch in choices:
                if isinstance(ch, dict):
                    message = ch.get("message")
                    if isinstance(message, dict):
                        content = message.get("content")
                        if isinstance(content, str):
                            completion_tokens += len(enc.encode(content))
        return completion_tokens

    def get_openai_metrics_for_completion_api(self, inputs, output):
        metrics = {}
        try:
            enc = self._get_encoding_for_completion_api(self._try_get_model(inputs, output))
            metrics["prompt_tokens"] = 0
            prompt = inputs.get("prompt")
            if isinstance(prompt, str):
                metrics["prompt_tokens"] = len(enc.encode(prompt))
            elif isinstance(prompt, list):
                for pro in prompt:
                    metrics["prompt_tokens"] += len(enc.encode(pro))
            if isinstance(output, list):
                if IS_LEGACY_OPENAI:
                    metrics["completion_tokens"] = len(output)
                else:
                    metrics["completion_tokens"] = len(
                        [chunk for chunk in output if chunk.choices and chunk.choices[0].text]
                    )
            else:
                metrics["completion_tokens"] = self._get_completion_tokens_for_completion_api(output, enc)
            metrics["total_tokens"] = metrics["prompt_tokens"] + metrics["completion_tokens"]
        except Exception as ex:
            self._log_warning(f"Failed to calculate metrics due to exception: {ex}.")
        return metrics

    def _get_encoding_for_completion_api(self, model):
        try:
            return tiktoken.encoding_for_model(model)
        except KeyError:
            return tiktoken.get_encoding("p50k_base")