Monitors and processes traffic to and from Azure OpenAI endpoints.
"""Several methods and classes around configuration."""

import json
import os

import yaml
from plugins.base import PowerProxyPlugin, foreach_plugin

from .dicts import QueryDict


class Configuration:
    """Configuration class."""

    def __init__(self, values_dict):
        """Constructor."""
        self.values_dict = QueryDict(values_dict)
        self.clients = [client["name"] for client in self.get("clients")]
        self.key_client_map = {client["key"]: client["name"] for client in self.get("clients")}
        self.plugin_names = [plugin["name"] for plugin in self.get("plugins")]

        # instantiate plugins
        self.plugins = [
            PowerProxyPlugin.get_plugin_instance(
                plugin_config["name"], self, QueryDict(plugin_config)
            )
            for plugin_config in self.get("plugins")
        ]
        foreach_plugin(self.plugins, "on_plugin_instantiated")

    def __getitem__(self, key):
        """Dunder method to get config value via ["..."] syntax."""
        return self.values_dict[key]

    def get(self, path, default=None):
        """Return value under given path."""
        return self.values_dict.get(path, default)

    def get_client_settings(self, client):
        """Return the value of a client's setting."""
        return next(
            (client_config for client_config in self["clients"] if client_config["name"] == client)
        )

    def print(self):
        """Print the current configuration."""
        Configuration.print_setting("Clients identified by API Key", ", ".join(self.clients))
        Configuration.print_setting(
            "Fixed client overwrite",
            f"{self['fixed_client'] if self['fixed_client'] else '(not set)'}",
        )
        if self["aoai/endpoint"]:
            Configuration.print_setting("Azure OpenAI endpoint (backend)", self["aoai/endpoint"])
        if self["aoai/mock_response"]:
            Configuration.print_setting("Azure OpenAI mock response", self["aoai/mock_response"])

        Configuration.print_setting("Plugins enabled", ", ".join(self.plugin_names))

    @staticmethod
    def print_setting(name, value, level=0):
        """Print the given setting name and value."""
        print((f"{' ' * level * 3}{name.ljust(32) if level==0 else name}: {value}"))

    @staticmethod
    def from_file(file_path):
        """Load configuration from file."""
        with open(file_path, "r", encoding="utf-8") as file:
            return Configuration(yaml.safe_load(file))

    @staticmethod
    def from_json_string(json_string):
        """Load configuration from JSON string."""
        try:
            return Configuration(json.loads(json_string))
        except ValueError as exception:
            raise ValueError(
                (f"The provided config string '{json_string}' is not a valid JSON document.")
            ) from exception

    @staticmethod
    def from_env_var(env_var_name="POWERPROXY_CONFIG_STRING", skip_no_env_var_exception=False):
        """Load configuration from environment variable."""
        if env_var_name in os.environ:
            return Configuration.from_json_string(os.environ[env_var_name])
        if not skip_no_env_var_exception:
            raise ValueError(
                f"Cannot load configuration from environment variable '{env_var_name}' because it "
                f"does not exist."
            )

    @staticmethod
    def from_args(args):
        """Load configuration from script arguments."""
        result = None
        if args.config_file:
            result = Configuration.from_file(args.config_file)
        elif args.config_env_var and args.config_env_var in os.environ:
            result = Configuration.from_env_var(args.config_env_var)
        elif args.config_env_var and args.config_env_var not in os.environ:
            raise ValueError(
                (
                    f"The specified environment variable '{args.config_env_var}', which shall "
                    "contain the configuration for PowerProxy, does not exist."
                )
            )
        elif args.config_string:
            result = Configuration.from_json_string(args.config_string)
        elif "POWERPROXY_CONFIG_STRING" in os.environ:
            result = Configuration.from_env_var(
                "POWERPROXY_CONFIG_STRING", skip_no_env_var_exception=True
            )
        else:
            raise ValueError(
                (
                    "No configuration provided. Ensure that you pass in a valid configuration "
                    "either by using argument '--config-file', '-config-env-var', or "
                    "'--config-string' or provide a valid config string in env variable named "
                    "'POWERPROXY_CONFIG_STRING' (in single-line JSON format)."
                )
            )
        return result

app/helpers/config.py
timoklimmer/powerproxy-aoai
"""Helper functions for working with dicts."""

import re


class QueryDict(dict):
    """Helper class to query and update a dict more conveniently."""

    def __getitem__(self, key):
        """Dunder method to get value at by ["..."] syntax."""
        return self.get(key)

    def get(self, path, default=None, separator="/", escape_sequence="''"):
        """
        Return the value at the given path.

        If the path cannot be found, the default value is returned.

        Keys in the path can be surrounded by the given escape sequence to treat separators as
        regular chars instead of separators. For instance, "abc/''def/ghi''" will return the value
        at abc -> def/ghi instead of abc -> def -> ghi (assuming default values for parameters
        'separator' and 'escape_sequence' are used).

        Optionally, a path can start with a separator ("/" by default). If path is only the
        separator, the entire dict is returned.
        """
        if not path:
            return default

        if path == separator:
            return dict(self)

        keys_from_path = QueryDict._get_keys_from_path(path, separator, escape_sequence)

        path_length = len(keys_from_path)
        for i, key_from_path in enumerate(keys_from_path):
            if i == 0:
                parent_element = dict.get(self, key_from_path, default)
                continue
            if i > 0 and i < path_length:
                try:
                    parent_element = parent_element[key_from_path]
                    continue
                except:
                    return default
        return parent_element

    def set(self, path, value, separator="/", escape_sequence="''"):
        """
        Set the given value at the given path.

        Path items which do not exist yet, will be added.

        Keys in the path can be surrounded by the given escape sequence to treat separators as
        regular chars instead of separators. For instance, "abc/''def/ghi''" will set the value at
        abc -> def/ghi instead of abc -> def -> ghi (assuming default values for parameters
        'separator' and 'escape_sequence' are used).

        Optionally, a path can start with a separator ("/" by default). If path is only the
        separator, the entire dict
        is set.
        """
        keys_from_path = QueryDict._get_keys_from_path(path, separator, escape_sequence)
        parent_element = self
        is_last_element_in_path = False
        for i, key_from_path in enumerate(keys_from_path):
            is_last_element_in_path = i == len(keys_from_path) - 1
            if not key_from_path in parent_element:
                parent_element[key_from_path] = {}
            if not is_last_element_in_path and not isinstance(parent_element[key_from_path], dict):
                raise ValueError(
                    (
                        "Cannot set value. All items on the way to the last element must be of "
                        "type dict to avoid that data is unintentionally overwritten."
                    )
                )
            if not is_last_element_in_path:
                parent_element = parent_element[key_from_path]
            else:
                parent_element[key_from_path] = value

    @staticmethod
    def get_last_item_from_path(path, separator="/", escape_sequence="''"):
        """Return the last item from the given path."""
        return QueryDict._get_keys_from_path(path, separator, escape_sequence)[-1]

    @staticmethod
    def _get_keys_from_path(path, separator, escape_sequence):
        """Get the different keys from the given path."""
        if path.startswith("/"):
            path = path[1:]

        escaped_escape_sequence = re.escape(escape_sequence)
        keys_from_path = [
            re.sub(
                rf"^{escaped_escape_sequence}|{escaped_escape_sequence}$",
                "",
                element.replace(chr(0), separator),
            )
            for element in re.sub(
                rf"{escaped_escape_sequence}.*?{escaped_escape_sequence}",
                lambda match: match.group().replace(separator, chr(0)),
                path,
            ).split(separator)
        ]
        return keys_from_path

app/helpers/dicts.py
timoklimmer/powerproxy-aoai
"""Several methods around printing headers."""

def print_header(header_string):
    """Print the given string as header."""
    print(len(header_string) * "-")
    print(header_string)
    print(len(header_string) * "-")

app/helpers/header.py
timoklimmer/powerproxy-aoai
"""Several methods around tokens."""

# notes:
# - for non-streaming responses, usage info is provided in Azure OpenAI's response. Try to use these
#   infos wherever possible.
# - for streaming responses, (Azure) OpenAI do not provide accurate usage infos, unfortunately.
#   hence, tokens need to be estimated. since there is no proper documentation, these estimations
#   might not be accurate but still more useful than no estimations. as of now, functions are not
#   taken into account. for exact numbers, disable the streaming mode.
# - code here is based on infos provided at
#   https://github.com/openai/openai-cookbook/blob/main/examples/How_to_count_tokens_with_tiktoken.ipynb
# - code will need update as new models come out and more documentation or usage infos are available

import io
import json

import tiktoken


def estimate_prompt_tokens_from_request_body(request_body):
    """Return the estimated number of tokes in the given request body string."""
    try:
        request_body_dict = json.load(io.BytesIO(request_body))
        return estimate_tokens_from_messages(request_body_dict["messages"])
    except:
        # eat any exception and return 0 to avoid breaking the system
        return 0


def estimate_tokens_from_string(string, encoding_name="cl100k_base"):
    """Return the estimated number of tokens in a text string."""
    encoding = tiktoken.get_encoding(encoding_name)
    num_tokens = len(encoding.encode(string))
    return num_tokens


def estimate_tokens_from_messages(messages, model="gpt-3.5-turbo-0613"):
    """Return the estimated number of tokens used by a list of messages."""
    try:
        encoding = tiktoken.encoding_for_model(model)
    except KeyError:
        encoding = tiktoken.get_encoding("cl100k_base")
    if model in {
        "gpt-3.5-turbo",
        "gpt-3.5-turbo-0613",
        "gpt-3.5-turbo-16k-0613",
        "gpt-4-0314",
        "gpt-4-32k-0314",
        "gpt-4",
        "gpt-4-0613",
        "gpt-4-32k-0613",
    }:
        tokens_per_message = 3
        tokens_per_name = 1
    elif model == "gpt-3.5-turbo-0301":
        tokens_per_message = 4  # every message follows <|start|>{role/name}\n{content}<|end|>\n
        tokens_per_name = -1  # if there's a name, the role is omitted
    else:
        raise NotImplementedError(
            f"Method is not implemented for model {model}. See "
            "https://github.com/openai/openai-python/blob/main/chatml.md for information on how "
            "messages are converted to tokens."
        )
    num_tokens = 0
    for message in messages:
        num_tokens += tokens_per_message
        for key, value in message.items():
            num_tokens += len(encoding.encode(value))
            if key == "name":
                num_tokens += tokens_per_name
    num_tokens += 3  # every reply is primed with <|start|>assistant<|message|>
    return num_tokens

app/helpers/tokens.py
timoklimmer/powerproxy-aoai
"""Declares a plugin which limits the usage rate for clients."""

import time

from plugins.base import TokenCountingPlugin, ImmediateResponseException
from fastapi import status
from fastapi.responses import Response


class LimitUsage(TokenCountingPlugin):
    """Limits the usage rate for clients."""

    budgets = {}
    max_tokens_per_minute_k_cache = {}

    def on_client_identified(self, routing_slip):
        """Run when the client has been identified."""
        super().on_client_identified(routing_slip)
        client = routing_slip["client"]

        # ensure there is a budget for the current client and minute (leaving budget as is if it
        # pre-exists for the current minute)
        current_minute = f"{time.strftime('%Y-%m-%d %H:%M')}"
        if client not in self.budgets or (
            client in self.budgets and self.budgets[client]["minute"] != current_minute
        ):
            self.budgets[client] = {
                "minute": current_minute,
                "budget": self._get_max_tokens_per_minute_k_for_client(client),
            }

        # ensure that the client has enough budget left for the current minute and return a 429
        # response if not
        if client in self.budgets:
            budget_entry = self.budgets[client]
            current_minute = f"{time.strftime('%Y-%m-%d %H:%M')}"
            if budget_entry["minute"] == current_minute and budget_entry["budget"] <= 0:
                raise ImmediateResponseException(
                    Response(
                        content=f"Too many requests for client '{client}'. Try again later.",
                        status_code=status.HTTP_429_TOO_MANY_REQUESTS,
                    )
                )

    def on_token_counts_for_request_available(self, routing_slip):
        """Is invoked when token counts are available for the request."""
        super().on_token_counts_for_request_available(routing_slip)

        # decrement the client's budget by the total tokens
        client = routing_slip["client"]
        self.budgets[client]["budget"] -= self.total_tokens

    def _get_max_tokens_per_minute_k_for_client(self, client):
        """Return the number of maximum tokens per minute in thousands for the given client."""
        if client not in self.max_tokens_per_minute_k_cache:
            client_settings = self.app_configuration.get_client_settings(client)
            if "max_tokens_per_minute_k" not in client_settings:
                raise ImmediateResponseException(
                    Response(
                        content=(
                            f"Configuration for client '{client}' misses a "
                            "'max_tokens_per_minute_k' setting. This needs to be set when the "
                            "LimitUsage plugin is enabled."
                        ),
                        status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
                    )
                )
            self.max_tokens_per_minute_k_cache[client] = (
                float(client_settings["max_tokens_per_minute_k"]) * 1000
            )
        return self.max_tokens_per_minute_k_cache[client]

app/plugins/LimitUsage/LimitUsage.py
timoklimmer/powerproxy-aoai
"""Declares the base class for a plugin that logs usage information."""

import datetime
import time
from abc import abstractmethod
from datetime import timezone

from plugins.base import TokenCountingPlugin


class LogUsageBase(TokenCountingPlugin):
    """Base class for a plugin that logs usage."""

    openai_region = None
    openai_processing_ms = None
    openai_region = None
    request_start_minute = None
    request_start_minute_utc = None

    def on_new_request_received(self, routing_slip):
        """Run when a new request is received."""
        super().on_new_request_received(routing_slip)

        self.openai_region = None
        self.openai_processing_ms = None
        self.openai_region = None
        self.request_start_minute = f"{time.strftime('%Y-%m-%d %H:%M')}"
        self.request_start_minute_utc = f"{datetime.datetime.now(timezone.utc):%Y-%m-%d %H:%M}"

    def on_headers_from_target_received(self, routing_slip):
        """Run when headers from target have been received."""
        super().on_headers_from_target_received(routing_slip)

        headers_from_target = routing_slip["headers_from_target"]
        for header, value in headers_from_target.items():
            if header == "openai-processing-ms":
                self.openai_processing_ms = float(value)
            if header == "x-ms-region":
                self.openai_region = value

    def on_body_dict_from_target_available(self, routing_slip):
        """Run when the body was received from AOAI (only for one-time, non-streaming requests)."""
        super().on_body_dict_from_target_available(routing_slip)

        self._append_line(
            request_start_minute=self.request_start_minute,
            request_start_minute_utc=self.request_start_minute_utc,
            client=routing_slip["client"],
            is_streaming=False,
            prompt_tokens=self.prompt_tokens,
            completion_tokens=self.completion_tokens,
            total_tokens=self.total_tokens,
            openai_processing_ms=self.openai_processing_ms,
            openai_region=self.openai_region,
        )

    def on_end_of_target_response_stream_reached(self, routing_slip):
        """Process the end of a stream (needs streaming requested)."""
        super().on_end_of_target_response_stream_reached(routing_slip)

        self._append_line(
            request_start_minute=self.request_start_minute,
            request_start_minute_utc=self.request_start_minute_utc,
            client=routing_slip["client"],
            is_streaming=True,
            prompt_tokens=self.prompt_tokens,
            completion_tokens=self.completion_tokens,
            total_tokens=self.total_tokens,
            openai_processing_ms=self.openai_processing_ms,
            openai_region=self.openai_region,
        )

    @abstractmethod
    def _append_line(
        self,
        request_start_minute,
        request_start_minute_utc,
        client,
        is_streaming,
        prompt_tokens,
        completion_tokens,
        total_tokens,
        openai_processing_ms,
        openai_region,
    ):
        pass

app/plugins/LogUsage/LogUsageBase.py
timoklimmer/powerproxy-aoai
"""Declares a plugin to log usage infos to console."""

from plugins.LogUsage.LogUsageBase import LogUsageBase


class LogUsageToConsole(LogUsageBase):
    """Logs Azure OpenAI usage info to console."""

    def _append_line(
        self,
        request_start_minute,
        request_start_minute_utc,
        client,
        is_streaming,
        prompt_tokens,
        completion_tokens,
        total_tokens,
        openai_processing_ms,
        openai_region,
    ):
        """Append a new line with the given infos."""
        print(
            "---\n"
            f"Request start minute           : {request_start_minute}\n"
            f"Request start minute UTC       : {request_start_minute_utc}\n"
            f"Client                         : {client}\n"
            f"Is Streaming                   : {is_streaming}\n"
            f"Prompt Tokens                  : {prompt_tokens}\n"
            f"Completion Tokens              : {completion_tokens}\n"
            f"Total Tokens                   : {total_tokens}\n"
            f"OpenAI Processing Milliseconds : {openai_processing_ms} ms\n"
            f"OpenAI Region                  : {openai_region}\n"
        )

app/plugins/LogUsage/LogUsageToConsole.py
timoklimmer/powerproxy-aoai
"""Declares a plugin to log usage infos to a CSV file."""

import os
import time

from plugins.LogUsage.LogUsageBase import LogUsageBase


class LogUsageToCsvFile(LogUsageBase):
    """Logs Azure OpenAI usage info to CSV file."""

    log_dir = None
    log_file_name = None
    log_file_path = None

    columns = [
        "request_start_minute",
        "request_start_minute_utc",
        "client",
        "is_streaming",
        "prompt_tokens",
        "completion_tokens",
        "total_tokens",
        "openai_processing_ms",
        "openai_region",
    ]

    def on_plugin_instantiated(self):
        """Run directly after the new plugin instance has been instantiated."""
        super().on_plugin_instantiated()

        # set log file etc.
        self.log_dir = "../logs"
        self.log_file_name = f"{time.strftime('%Y%m%d-%H%M%S')}.logs.csv"
        self.log_file_path = os.path.join(self.log_dir, self.log_file_name)

        # ensure we have a logs directory
        if not os.path.exists(self.log_dir):
            os.mkdir(self.log_dir)

        # create a new file containing the header for the CSV file
        with open(self.log_file_path, "w", encoding="utf-8") as log_file:
            log_file.write(",".join(self.columns))

    def _append_line(
        self,
        request_start_minute,
        request_start_minute_utc,
        client,
        is_streaming,
        prompt_tokens,
        completion_tokens,
        total_tokens,
        openai_processing_ms,
        openai_region,
    ):
        """Append a new line with the given infos."""
        with open(self.log_file_path, "a", encoding="utf-8") as log_file:
            log_file.write(
                "\n"
                f"{request_start_minute},"
                f"{request_start_minute_utc},"
                f"{client},"
                f"{1 if is_streaming else 0},"
                f"{prompt_tokens},"
                f"{completion_tokens},"
                f"{total_tokens},"
                f"{openai_processing_ms},"
                f"{openai_region}"
            )

app/plugins/LogUsage/LogUsageToCsvFile.py
timoklimmer/powerproxy-aoai
"""Declares a plugin to log usage infos to a CSV file."""

from azure.identity import ChainedTokenCredential, ClientSecretCredential, ManagedIdentityCredential
from azure.monitor.ingestion import LogsIngestionClient
from helpers.config import Configuration
from helpers.dicts import QueryDict
from plugins.LogUsage.LogUsageBase import LogUsageBase


class LogUsageToLogAnalytics(LogUsageBase):
    """Logs Azure OpenAI usage info to a Log Analytics table."""

    log_ingestion_endpoint = None
    credential_tenant_id = None
    credential_client_id = None
    credential_client_secret = None
    data_collection_rule_id = None
    stream_name = None

    log_analytics_client = None

    def __init__(self, app_configuration, plugin_configuration: QueryDict):
        """Constructor."""
        super().__init__(app_configuration, plugin_configuration)

        self.log_ingestion_endpoint = plugin_configuration.get("log_ingestion_endpoint")
        self.credential_tenant_id = plugin_configuration.get("credential_tenant_id")
        self.credential_client_id = plugin_configuration.get("credential_client_id")
        self.credential_client_secret = plugin_configuration.get("credential_client_secret")
        self.data_collection_rule_id = plugin_configuration.get("data_collection_rule_id")
        self.stream_name = plugin_configuration.get("stream_name")

    def on_plugin_instantiated(self):
        """Run directly after the new plugin instance has been instantiated."""
        super().on_plugin_instantiated()

        self.log_analytics_client = LogsIngestionClient(
            endpoint=self.log_ingestion_endpoint,
            credential=ChainedTokenCredential(
                ManagedIdentityCredential(),
                ClientSecretCredential(
                    tenant_id=self.credential_tenant_id,
                    client_id=self.credential_client_id,
                    client_secret=self.credential_client_secret,
                ),
            ),
            logging_enable=True,
        )

    def on_print_configuration(self):
        """Print plugin-specific configuration."""
        super().on_print_configuration()

        Configuration.print_setting("Log ingestion endpoint", self.log_ingestion_endpoint, 1)
        Configuration.print_setting("Credential Tenant ID", self.credential_tenant_id, 1)
        Configuration.print_setting("Credential Client ID", self.credential_client_id, 1)
        Configuration.print_setting("Data Collection Rule ID", self.data_collection_rule_id, 1)
        Configuration.print_setting("Stream Name", self.stream_name, 1)

    def _append_line(
        self,
        request_start_minute,
        request_start_minute_utc,
        client,
        is_streaming,
        prompt_tokens,
        completion_tokens,
        total_tokens,
        openai_processing_ms,
        openai_region,
    ):
        """Append a new line with the given infos."""
        # pylint: disable=no-value-for-parameter
        self.log_analytics_client.upload(
            rule_id=self.data_collection_rule_id,
            stream_name=self.stream_name,
            logs=[
                {
                    "Client": client,
                    "RequestStartMinute": request_start_minute_utc,
                    "IsStreaming": is_streaming,
                    "PromptTokens": prompt_tokens,
                    "CompletionTokens": completion_tokens,
                    "TotalTokens": total_tokens,
                    "OpenAIProcessingMS": openai_processing_ms,
                    "OpenAIRegion": openai_region,
                }
            ],
        )
        # pylint: enable=no-value-for-parameter

app/plugins/LogUsage/LogUsageToLogAnalytics.py
timoklimmer/powerproxy-aoai
"""Defines the foundation for PowerProxy plugins."""

import importlib
import re

from helpers.tokens import estimate_prompt_tokens_from_request_body


def foreach_plugin(plugins, method_name, *args):
    """Have each plugin run the method with the given name and arguments."""
    for plugin in plugins:
        if hasattr(plugin, method_name):
            getattr(plugin, method_name)(*args)
        else:
            raise ValueError(
                (
                    f"Plugin class '{plugin.__class__()}' does not have a method named "
                    f"'{method_name}'."
                )
            )


class PowerProxyPlugin:
    """A plugin for PowerProxy, doing different things at different events."""

    def __init__(self, app_configuration, plugin_configuration):
        """Constructor."""
        self.app_configuration = app_configuration
        self.plugin_configuration = plugin_configuration

    def on_plugin_instantiated(self):
        """Run directly after the new plugin instance has been instantiated."""

    def on_print_configuration(self):
        """Print plugin-specific configuration."""
        print(f"Plugin config: {self.__class__.__name__}")

    def on_new_request_received(self, routing_slip):
        """Run when a new request has been received."""

    def on_client_identified(self, routing_slip):
        """Run when the client has been identified."""

    def on_headers_from_target_received(self, routing_slip):
        """Run when the headers from the target have been received."""

    def on_body_dict_from_target_available(self, routing_slip):
        """Run when the body returned from the target could be converted to a dict."""

    def on_data_event_from_target_received(self, routing_slip):
        """Run when a data event has been received from the target (only on streaming)."""

    def on_end_of_target_response_stream_reached(self, routing_slip):
        """Run when the end of the target's response stream has been reached (only on streaming)."""

    @staticmethod
    def get_plugin_class(plugin_name):
        """Return the class for the given plugin name."""
        plugin_group = re.sub("To.+$", "", plugin_name)
        return getattr(
            importlib.import_module(f"plugins.{plugin_group}.{plugin_name}"), plugin_name
        )

    @staticmethod
    def get_plugin_instance(plugin_name, app_configuration, plugin_configuration):
        """Return an instance of the plugin with the given name."""
        plugin_class = PowerProxyPlugin.get_plugin_class(plugin_name)
        return plugin_class(app_configuration, plugin_configuration)


class TokenCountingPlugin(PowerProxyPlugin):
    """A plugin which counts tokens."""

    prompt_tokens = None
    streaming_prompt_tokens = None
    completion_tokens = None
    streaming_completion_tokens = None
    total_tokens = None

    def on_new_request_received(self, routing_slip):
        """Run when a new request is received."""
        super().on_new_request_received(routing_slip)

        self.prompt_tokens = None
        self.completion_tokens = None
        self.streaming_prompt_tokens = None
        self.streaming_completion_tokens = None
        self.total_tokens = None

    def on_body_dict_from_target_available(self, routing_slip):
        """Run when the body was received from AOAI (only for one-time, non-streaming requests)."""
        super().on_body_dict_from_target_available(routing_slip)

        usage = routing_slip["body_dict_from_target"]["usage"]
        self.completion_tokens = usage["completion_tokens"]
        self.prompt_tokens = usage["prompt_tokens"]
        self.total_tokens = usage["total_tokens"]

        self.on_token_counts_for_request_available(routing_slip)

    def on_data_event_from_target_received(self, routing_slip):
        """Run when a data event has been received by AOAI (needs streaming requested)."""
        super().on_data_event_from_target_received(routing_slip)

        self.streaming_completion_tokens = (
            self.streaming_completion_tokens + 1 if self.streaming_completion_tokens else 1
        )

    def on_end_of_target_response_stream_reached(self, routing_slip):
        """Process the end of a stream (needs streaming requested)."""
        super().on_end_of_target_response_stream_reached(routing_slip)

        self.prompt_tokens = estimate_prompt_tokens_from_request_body(
            routing_slip["incoming_request_body"]
        )
        self.completion_tokens = self.streaming_completion_tokens
        self.total_tokens = (
            self.prompt_tokens + self.completion_tokens
            if self.prompt_tokens is not None and self.completion_tokens is not None
            else None
        )
        self.on_token_counts_for_request_available(routing_slip)

    def on_token_counts_for_request_available(self, routing_slip):
        """Is invoked when token counts are available for the request."""


class ImmediateResponseException(Exception):
    """Exception thrown when a plugin needs a certain response returned immediately."""

    def __init__(self, response):
        """Constructor."""
        self.response = response

app/plugins/base.py
timoklimmer/powerproxy-aoai
"""
PowerProxy for AOAI - reverse proxy to process requests and responses to/from Azure Open AI.

- Use the "Debug powershell.py" launch configuration in VS.Code to develop and debug this script.
- Adjust the launch configuration in VS.Code as needed (esp. for plugins enabled)
"""

import argparse
import asyncio
import io
import json

import httpx
import uvicorn
from fastapi import FastAPI, Request, status
from fastapi.responses import Response, StreamingResponse
from helpers.config import Configuration
from helpers.header import print_header
from plugins.base import ImmediateResponseException, foreach_plugin
from version import VERSION

## define script arguments
parser = argparse.ArgumentParser()
# --config-file
parser.add_argument(
    "--config-file",
    type=str,
    help="Path to config file",
)
# --config-env-var
parser.add_argument(
    "--config-env-var",
    type=str,
    help="Name of the environment variable containing the configuration as JSON string.",
)
# --config-string
parser.add_argument(
    "--config-string",
    type=str,
    help="String containing the configuration as JSON string.",
)
# --port
parser.add_argument(
    "--port",
    type=int,
    default=80,
    help=(
        "Port where the proxy runs. Ports <= 1024 may need special permissions in Linux. "
        "Default: 80."
    ),
)
args, unknown = parser.parse_known_args()

## load configuration
config = Configuration.from_args(args)

## define and run proxy app
app = FastAPI()


# app startup event
@app.on_event("startup")
async def startup_event():
    """Invoked when the app is started."""
    # print header and config values
    print_header(f"PowerProxy for Azure OpenAI - v{VERSION}")
    Configuration.print_setting("Proxy port", args.port)
    config.print()
    foreach_plugin(config.plugins, "on_print_configuration")

    # instantiate HTTP client for AOAI endpoint
    if config.get("aoai/mock_response"):

        async def get_mock_response(request):
            ms_to_wait_before_return = config.get("aoai/mock_response/ms_to_wait_before_return")
            if ms_to_wait_before_return:
                ms_to_wait_before_return = float(ms_to_wait_before_return)
                await asyncio.sleep(ms_to_wait_before_return / 1_000)
            return httpx.Response(200, json=config.get("aoai/mock_response/json"))

        app.state.target_client: httpx.AsyncClient = httpx.AsyncClient(
            base_url="https://mock/",
            transport=httpx.MockTransport(get_mock_response),
        )
    else:
        app.state.target_client: httpx.AsyncClient = httpx.AsyncClient(
            base_url=config["aoai/endpoint"]
        )

    # print serve notification
    print()
    print("Serving incoming requests...")
    print()


# app shutdown
@app.on_event("shutdown")
async def shutdown_event():
    """Invoked when the app is shut down."""
    await app.state.target_client.aclose()


@app.exception_handler(ImmediateResponseException)
async def exception_callback(request: Request, exception: ImmediateResponseException):
    """Immediately return given response when an ImmediateResponseException is raised."""
    return exception.response


# liveness probe
@app.get(
    "/powerproxy/health/liveness",
    status_code=status.HTTP_204_NO_CONTENT,
    description="Liveness probe for the PowerProxy",
)
async def liveness_probe():
    """
    Return a 204/No Content if the container is live.
    Note: This is required by some hosting services to know if there are issues with the container.
    """
    return None


# all other GETs and POSTs
@app.get("/{path:path}")
@app.post("/{path:path}")
async def handle_request(request: Request, path: str):
    """Handle any incoming request."""
    # create a new routing slip and tell plugins about new request
    routing_slip = {
        "incoming_request": request,
        "incoming_request_body": await request.body(),
    }
    foreach_plugin(config.plugins, "on_new_request_received", routing_slip)

    # identify client and replace API key if needed
    # notes: - When API authentication is used, we get an API key in header 'api-key'. This
    #          would usually be the API key for Azure Open AI, but we configure and use
    #          client-specific keys here for the proxy to identify the client and replace the
    #          API key against the real AOAI key afterwards.
    #        - For Azure AD authentication, we should get no API key but an Azure AD token in
    #          header 'Authorization'. Unfortunately, we cannot interpret or modify that token,
    #          so we need another mechanism to identify clients. In that case, we need a
    #          separate instance of PowerProxy for each client, whereby each client uses a fixed
    #          client.
    #        - Some requests may neither contain an API key or an Azure AD token. In that case,
    #          we need to make sure that the proxy continues to work.
    headers = {
        key: request.headers[key]
        for key in set(request.headers.keys())
        - {"Host", "host", "Content-Length", "content-length"}
    }
    client = None
    if config["FIXED_CLIENT"]:
        client = config["FIXED_CLIENT"]
    if "api-key" in headers:
        if headers["api-key"] not in config.key_client_map:
            raise ValueError(
                (
                    "The provided API key is not a valid PowerProxy key. Ensure that the "
                    "'api-key' header contains valid API key from the PowerProxy's "
                    "configuration."
                )
            )
        client = config.key_client_map[headers["api-key"]] if client is None else client
        headers["api-key"] = config["aoai/key"] or ""
    routing_slip["client"] = client
    foreach_plugin(config.plugins, "on_client_identified", routing_slip)

    # forward request to target endpoint and get response
    aoai_response: httpx.Response = await app.state.target_client.request(
        request.method,
        path,
        params=request.query_params,
        headers=headers,
        content=routing_slip["incoming_request_body"],
    )
    routing_slip["headers_from_target"] = aoai_response.headers
    foreach_plugin(config.plugins, "on_headers_from_target_received", routing_slip)

    # determine if it's an event stream or not
    routing_slip["is_event_stream"] = (
        "content-type" in aoai_response.headers
        and aoai_response.headers["content-type"] == "text/event-stream"
    )

    # return different response types depending if it's an event stream or not
    routing_slip["response_headers_from_target"] = {
        header_item[0].decode(): header_item[1].decode()
        for header_item in aoai_response.headers.raw
    }
    match routing_slip["is_event_stream"]:
        case False:
            # non-streamed response
            body = await aoai_response.aread()
            try:
                routing_slip["body_dict_from_target"] = json.load(io.BytesIO(body))
                foreach_plugin(config.plugins, "on_body_dict_from_target_available", routing_slip)
            except:
                # eat any exception in case the response cannot be parsed
                pass
            return Response(
                content=body,
                status_code=aoai_response.status_code,
                headers=routing_slip["response_headers_from_target"],
            )
        case True:
            # event stream
            # forward and process events as they come in
            # note: see https://learn.microsoft.com/de-de/azure/ai-services/openai/reference
            async def yield_data_events():
                """Stream response while invoking plugins."""
                async for line in aoai_response.aiter_lines():
                    yield f"{line}\r\n"
                    routing_slip["data_from_target"] = None
                    if line.startswith("data: "):
                        data = line[6:]
                        if data != "[DONE]":
                            routing_slip["data_from_target"] = data
                            foreach_plugin(
                                config.plugins,
                                "on_data_event_from_target_received",
                                routing_slip,
                            )
                foreach_plugin(
                    config.plugins, "on_end_of_target_response_stream_reached", routing_slip
                )

            return StreamingResponse(
                yield_data_events(),
                status_code=aoai_response.status_code,
                headers=routing_slip["response_headers_from_target"],
            )


if __name__ == "__main__":
    # note: this applies only when the powerproxy.py script is executed directly. In the Dockerfile
    #       provided, we use a uvicorn command to run the app, so parameters might need to be
    #       modified there AS WELL.
    uvicorn.run(
        app,
        host="0.0.0.0",
        port=int(args.port),
        log_level="warning",
        server_header=False,
        date_header=False,
        proxy_headers=False,
        timeout_keep_alive=120,
        timeout_graceful_shutdown=120,
    )

app/powerproxy.py
timoklimmer/powerproxy-aoai
"""Declares the version of the app."""

# version
VERSION = "0.0.0"

app/version.py
timoklimmer/powerproxy-aoai
"""Return a one-line JSON string of the YAML file specified."""

import argparse
import json

import yaml

parser = argparse.ArgumentParser()
parser.add_argument("--yaml-file", type=str, help="Path to the YAML file to convert")
args = parser.parse_args()

with open(args.yaml_file, "r", encoding="utf-8") as yaml_file:
    dict_from_file = yaml.safe_load(yaml_file)

print(json.dumps(dict_from_file))

config/to_json_string.py
timoklimmer/powerproxy-aoai
"""Script to test the proxy's support for requests responding with a one-time response."""

import openai

openai.api_type = "azure"
openai.api_base = "http://localhost"
openai.api_version = "2023-03-15-preview"
openai.api_key = "04ae14bc78184621d37f1ce57a52eb7"

response = openai.ChatCompletion.create(
    engine="gpt-35-turbo",
    messages=[
        {
            "role": "system",
            "content": "You are an AI assistant that helps people find information.",
        },
        {"role": "user", "content": "Tell me a joke!"},
        {
            "role": "assistant",
            "content": "Why did the tomato turn red? Because it saw the salad dressing!",
        },
        {"role": "user", "content": "Yeah, that's a great one."},
    ],
    temperature=0,
    max_tokens=800,
    top_p=0.95,
    frequency_penalty=0,
    presence_penalty=0,
    stop=None,
    stream=False,
)

print(response)

print("")

test/python/test_non_streaming.py
timoklimmer/powerproxy-aoai
"""Script to test the proxy's ability to support response streaming when functions are used."""

import openai

openai.api_type = "azure"
openai.api_base = "http://localhost"
openai.api_version = "2023-07-01-preview"
openai.api_key = "04ae14bc78184621d37f1ce57a52eb7"

response = openai.ChatCompletion.create(
    engine="gpt-35-turbo-0613",
    messages=[
        {
            "role": "user",
            "content": "Find beachfront hotels in San Diego for less than $300 a month with free breakfast.",
        }
    ],
    functions=[
        {
            "name": "search_hotels",
            "description": "Retrieves hotels from the search index based on the parameters provided",
            "parameters": {
                "type": "object",
                "properties": {
                    "location": {
                        "type": "string",
                        "description": "The location of the hotel (i.e. Seattle, WA)",
                    },
                    "max_price": {
                        "type": "number",
                        "description": "The maximum price for the hotel",
                    },
                    "features": {
                        "type": "string",
                        "description": "A comma separated list of features (i.e. beachfront, free wifi, etc.)",
                    },
                },
                "required": ["location"],
            },
        }
    ],
    temperature=0,
    function_call="auto",
    stream=False,
)

print(response)

test/python/test_non_streaming_functions.py
timoklimmer/powerproxy-aoai
"""Script to test the proxy's ability to support response streaming."""

import openai

openai.api_type = "azure"
openai.api_base = "http://localhost"
openai.api_version = "2023-03-15-preview"
openai.api_key = "72bd81ef32763530b29e3da63d46ad6"

response = openai.ChatCompletion.create(
    engine="gpt-35-turbo",
    messages=[
        {
            "role": "system",
            "content": "You are an AI assistant that helps people find information.",
        },
        {"role": "user", "content": "Tell me a joke!"},
        {
            "role": "assistant",
            "content": "Why did the tomato turn red? Because it saw the salad dressing!",
        },
        {"role": "user", "content": "Yeah, that's a great one."},
    ],
    temperature=0,
    max_tokens=800,
    top_p=0.95,
    frequency_penalty=0,
    presence_penalty=0,
    stop=None,
    stream=True,
)

for item in response:
    delta = item["choices"][0]["delta"]
    if "content" in delta:
        print(delta["content"], end="")

print("")

test/python/test_streaming.py
timoklimmer/powerproxy-aoai
"""Script to test the proxy's ability to support response streaming when functions are used."""

import openai

openai.api_type = "azure"
openai.api_base = "http://localhost"
openai.api_version = "2023-07-01-preview"
openai.api_key = "72bd81ef32763530b29e3da63d46ad6"

response = openai.ChatCompletion.create(
    engine="gpt-35-turbo-0613",
    messages=[
        {
            "role": "user",
            "content": "Find beachfront hotels in San Diego for less than $300 a month with free breakfast.",
        }
    ],
    functions=[
        {
            "name": "search_hotels",
            "description": "Retrieves hotels from the search index based on the parameters provided",
            "parameters": {
                "type": "object",
                "properties": {
                    "location": {
                        "type": "string",
                        "description": "The location of the hotel (i.e. Seattle, WA)",
                    },
                    "max_price": {
                        "type": "number",
                        "description": "The maximum price for the hotel",
                    },
                    "features": {
                        "type": "string",
                        "description": "A comma separated list of features (i.e. beachfront, free wifi, etc.)",
                    },
                },
                "required": ["location"],
            },
        }
    ],
    temperature=0,
    function_call="auto",
    stream=True,
)

for item in response:
    print(item)

test/python/test_streaming_functions.py
timoklimmer/powerproxy-aoai