llmeter 0.1.0__py3-none-any.whl

llmeter/endpoints/__init__.py

@@ -0,0 +1,13 @@
+import importlib.util
+
+from .base import Endpoint, InvocationResponse  # noqa: F401
+from .bedrock import BedrockConverse, BedrockConverseStream  # noqa: F401
+from .sagemaker import SageMakerEndpoint, SageMakerStreamEndpoint  # noqa: F401
+
+spec = importlib.util.find_spec("openai")
+if spec:
+    from .openai import OpenAIEndpoint, OpenAICompletionEndpoint  # noqa: F401
+
+spec = importlib.util.find_spec("litellm")
+if spec:
+    from .litellm import LiteLLM, LiteLLMStreaming  # noqa: F401

llmeter/endpoints/base.py

@@ -0,0 +1,232 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import importlib
+import json
+import os
+from abc import ABC, abstractmethod
+from dataclasses import asdict, dataclass
+from typing import Dict, TypeVar
+from uuid import uuid4
+
+from upath import UPath as Path
+
+Self = TypeVar(
+    "Self", bound="Endpoint"
+)  # for python >= 3.11 can be replaced with direct import of `Self`
+
+
+@dataclass
+class InvocationResponse:
+    """
+    A class representing a invocation result.
+
+    Attributes:
+        response_text (str): The invocation output.
+        id (str): A unique identifier for the invocation.
+        time_to_last_token (float): The time taken to generate the response in seconds.
+        time_to_first_token (float): The time taken to receive the first token of the response in seconds.
+        num_tokens_output (Optional[int]): The number of tokens in the response.
+        num_tokens_input (Optional[int]): The number of tokens in the invocation payload.
+        input_prompt (str): The input prompt used in the invocation.
+        time_per_output_token (float): The average time taken to generate each token in the response.
+        error (str): Any error that occurred during invocation.
+    """
+
+    response_text: str | None
+    id: str | None = None
+    input_prompt: str | None = None
+    time_to_first_token: float | None = None
+    time_to_last_token: float | None = None
+    num_tokens_input: int | None = None
+    num_tokens_output: int | None = None
+    time_per_output_token: float | None = None
+    error: str | None = None
+
+    def to_json(self, **kwargs) -> str:
+        return json.dumps(self.__dict__, **kwargs)
+
+    @staticmethod
+    def error_output(
+        input_prompt: str | None = None, error=None, id: str | None = None
+    ):
+        return InvocationResponse(
+            id=id or uuid4().hex,
+            response_text=None,
+            input_prompt=input_prompt,
+            time_to_last_token=None,
+            error="invocation failed" if error is None else str(error),
+        )
+
+    def __repr__(self):
+        return self.to_json(default=str)
+
+    def __str__(self):
+        return self.to_json(indent=4, default=str)
+
+    def to_dict(self):
+        return asdict(self)
+
+
+class Endpoint(ABC):
+    """
+    An abstract base class for endpoint implementations.
+
+    This class defines the basic structure and interface for all endpoint classes.
+    It provides abstract methods that must be implemented by subclasses.
+    """
+
+    @abstractmethod
+    def __init__(
+        self,
+        endpoint_name: str,
+        model_id: str,
+        provider: str,
+    ):
+        """
+        Initialize the BaseEndpoint.
+
+        Args:
+            endpoint_name (str): The name of the endpoint.
+            model_id (str): The identifier of the model associated with this endpoint.
+            provider (str): The provider of the endpoint.
+
+        Returns:
+            None
+        """
+        self.endpoint_name = endpoint_name
+        self.model_id = model_id
+        self.provider = provider
+
+    @abstractmethod
+    def invoke(self, payload: Dict) -> InvocationResponse:
+        """
+        Invoke the endpoint with the given payload.
+
+        This method must be implemented by subclasses to define how the endpoint
+        is invoked and how the response is processed.
+
+        Args:
+            payload (Dict): The input payload for the model.
+
+        Returns:
+            InvocationResponse: An object containing the model's response and associated metrics.
+
+        Raises:
+            NotImplementedError: If the method is not implemented by a subclass.
+        """
+        raise NotImplementedError
+
+    @staticmethod
+    def create_payload(*args, **kwargs):
+        """
+        Create a payload for the endpoint invocation.
+
+        This static method should be implemented by subclasses to define
+        how the payload is created based on the given arguments.
+
+        Args:
+            *args: Variable length argument list.
+            **kwargs: Arbitrary keyword arguments.
+
+        Returns:
+            NotImplemented: This method returns NotImplemented in the base class.
+        """
+        return NotImplemented
+
+    @classmethod
+    def __subclasshook__(cls, C):
+        """
+        Determine if a class is considered a subclass of BaseEndpoint.
+
+        This method is used to implement a custom subclass check. A class
+        is considered a subclass of BaseEndpoint if it has an 'invoke' method.
+
+        Args:
+            C: The class to check.
+
+        Returns:
+            bool or NotImplemented: True if the class is a subclass, False if it isn't,
+                                    or NotImplemented if the check is inconclusive.
+        """
+        if cls is Endpoint:
+            if any("invoke" in B.__dict__ for B in C.__mro__):
+                return True
+        return NotImplemented
+
+    def save(self, output_path: os.PathLike) -> os.PathLike:
+        """
+        Save the endpoint configuration to a JSON file.
+
+        This method serializes the endpoint's configuration (excluding private attributes)
+        to a JSON file at the specified path.
+
+        Args:
+            output_path (str | UPath): The path where the configuration file will be saved.
+
+        Returns:
+            None
+        """
+        output_path = Path(output_path)
+        output_path.parent.mkdir(parents=True, exist_ok=True)
+        with output_path.open("w") as f:
+            endpoint_conf = self.to_dict()
+            json.dump(endpoint_conf, f, indent=4, default=str)
+        return output_path
+
+    def to_dict(self) -> Dict:
+        """
+        Convert the endpoint configuration to a dictionary.
+
+        Returns:
+            Dict: A dictionary representation of the endpoint configuration.
+        """
+        endpoint_conf = {k: v for k, v in vars(self).items() if not k.startswith("_")}
+        endpoint_conf["endpoint_type"] = self.__class__.__name__
+        return endpoint_conf
+
+    @classmethod
+    def load_from_file(cls, input_path: os.PathLike) -> Self:
+        """
+        Load an endpoint configuration from a JSON file.
+
+        This class method reads a JSON file containing an endpoint configuration,
+        determines the appropriate endpoint class, and instantiates it with the
+        loaded configuration.
+
+        Args:
+            input_path (str|UPath): The path to the JSON configuration file.
+
+        Returns:
+            Endpoint: An instance of the appropriate endpoint class, initialized
+                      with the configuration from the file.
+        """
+
+        input_path = Path(input_path)
+        with input_path.open("r") as f:
+            data = json.load(f)
+        endpoint_type = data.pop("endpoint_type")
+        endpoint_module = importlib.import_module("llmeter.endpoints")
+        endpoint_class = getattr(endpoint_module, endpoint_type)
+        return endpoint_class(**data)
+
+    @classmethod
+    def load(cls, endpoint_config: Dict) -> Self:  # type: ignore
+        """
+        Load an endpoint configuration from a dictionary.
+
+        This class method reads a dictionary containing an endpoint configuration,
+        determines the appropriate endpoint class, and instantiates it with the
+        loaded configuration.
+
+        Args:
+            data (Dict): A dictionary containing the endpoint configuration.
+
+        Returns:
+            Endpoint: An instance of the appropriate endpoint class, initialized
+                      with the configuration from the dictionary.
+        """
+        endpoint_type = endpoint_config.pop("endpoint_type")
+        endpoint_module = importlib.import_module("llmeter.endpoints")
+        endpoint_class = getattr(endpoint_module, endpoint_type)
+        return endpoint_class(**endpoint_config)

llmeter/endpoints/bedrock.py

@@ -0,0 +1,192 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import logging
+import time
+from typing import Dict, Sequence
+from uuid import uuid4
+
+import boto3
+import jmespath
+from botocore.exceptions import ClientError
+
+from .base import Endpoint, InvocationResponse
+
+logger = logging.getLogger(__name__)
+
+
+class BedrockBase(Endpoint):
+    def __init__(
+        self,
+        model_id: str,
+        endpoint_name: str | None = None,
+        region: str | None = None,
+        inference_config: Dict | None = None,
+    ):
+        """
+        Base class for Amazon Bedrock endpoints.
+
+        This class provides the foundation for interacting with Amazon Bedrock services,
+        including client initialization and payload handling.
+        """
+
+        super().__init__(
+            model_id=model_id, endpoint_name=endpoint_name or "", provider="bedrock"
+        )
+        """
+        Initialize the BedrockBase instance.
+
+        Args:
+            model_id (str): The ID of the model to use.
+            region (str | None, optional): The AWS region to use. If None, uses the default session region.
+            inference_config (Dict | None, optional): Configuration for inference.
+        """
+
+        self.endpoint_name = "amazon bedrock"
+
+        self.region = region or boto3.session.Session().region_name
+        logger.info(f"Using AWS region: {self.region}")
+
+        self._bedrock_client = boto3.client("bedrock-runtime", region_name=self.region)
+        self._inference_config = inference_config
+
+    def _parse_payload(self, payload):
+        """
+        Parse the payload to extract text content.
+
+        Args:
+            payload (dict): The payload containing messages.
+
+        Returns:
+            str: Concatenated text content from the messages.
+        """
+        jpath = "[:].content[:].text"
+        messages = payload.get("messages")
+        return "\n".join([k for j in jmespath.search(jpath, messages) for k in j])
+
+    @staticmethod
+    def create_payload(
+        user_message: str | Sequence[str], max_tokens: int = 256, **kwargs
+    ):
+        """
+        Create a payload for the Bedrock Converse API request.
+
+        Args:
+            user_message (str | Sequence[str]): The user's message or a sequence of messages.
+            max_tokens (int, optional): The maximum number of tokens to generate. Defaults to 256.
+            **kwargs: Additional keyword arguments to include in the payload.
+
+        Returns:
+            dict: The formatted payload for the Bedrock API request.
+        """
+
+        if isinstance(user_message, str):
+            user_message = [user_message]
+        payload: dict = {
+            "messages": [
+                {"role": "user", "content": [{"text": k}]} for k in user_message
+            ],
+        }
+        payload.update(kwargs)
+        if payload.get("inferenceConfig") is None:
+            payload["inferenceConfig"] = {}
+
+        payload["inferenceConfig"] = {
+            **payload["inferenceConfig"],
+            "maxTokens": max_tokens,
+        }
+        return payload
+
+
+class BedrockConverse(BedrockBase):
+    def _parse_converse_response(self, response: Dict) -> InvocationResponse:
+        output_text = response["output"]["message"]["content"][0]["text"]
+        usage = response.get("usage", {})
+        # metrics = response.get("metrics", {})
+
+        return InvocationResponse(
+            id=uuid4().hex,
+            response_text=output_text,
+            num_tokens_input=usage.get("inputTokens"),
+            num_tokens_output=usage.get("outputTokens"),
+            # time_to_last_token=metrics.get("latencyMs") / 1e3,
+        )
+
+    def invoke(self, payload: Dict, **kwargs) -> InvocationResponse:
+        payload = {**kwargs, **payload}
+        if payload.get("inferenceConfig") is None:
+            payload["inferenceConfig"] = self._inference_config or {}
+
+        payload["modelId"] = self.model_id
+        try:
+            start_t = time.perf_counter()
+            client_response = self._bedrock_client.converse(**payload)
+            time_to_last_token = time.perf_counter() - start_t
+        except (ClientError, Exception) as e:
+            logger.error(e)
+            return InvocationResponse.error_output(id=uuid4().hex, error=str(e))
+        response = self._parse_converse_response(client_response)  # type: ignore
+        response.input_prompt = self._parse_payload(payload)
+        response.time_to_last_token = time_to_last_token
+        return response
+
+
+class BedrockConverseStream(BedrockConverse):
+    def invoke(self, payload: Dict, **kwargs) -> InvocationResponse:
+        payload = {**kwargs, **payload}
+        if payload.get("inferenceConfig") is None:
+            payload["inferenceConfig"] = self._inference_config or {}
+
+        payload["modelId"] = self.model_id
+        start_t = time.perf_counter()
+        try:
+            client_response = self._bedrock_client.converse_stream(**payload)
+        except (ClientError, Exception) as e:
+            logger.error(e)
+            return InvocationResponse.error_output(id=uuid4().hex, error=str(e))
+        response = self._parse_conversation_stream(client_response, start_t)  # type: ignore
+        response.input_prompt = self._parse_payload(payload)
+        return response
+
+    def _parse_conversation_stream(
+        self, client_response: Dict, start_t: float
+    ) -> InvocationResponse:
+        time_flag = True
+        time_to_first_token = None
+        output_text = ""
+        for chunk in client_response["stream"]:
+            if "contentBlockDelta" in chunk:
+                output_text += chunk["contentBlockDelta"]["delta"].get("text") or ""
+                if time_flag:
+                    time_to_first_token = time.perf_counter() - start_t
+                    time_flag = False
+
+            if "contentBlockStop" in chunk:
+                time_to_last_token = time.perf_counter() - start_t
+
+            if "metadata" in chunk:
+                metadata = chunk["metadata"]
+
+        response = InvocationResponse(
+            id=uuid4().hex,
+            response_text=output_text,
+            time_to_last_token=time_to_last_token,
+            time_to_first_token=time_to_first_token,
+        )
+
+        if metadata:
+            # time_to_last_token = metadata.get("metrics", {}).get("latencyMs")
+            usage = metadata.get("usage", {})
+            response.num_tokens_input = usage.get("inputTokens")
+            response.num_tokens_output = usage.get("outputTokens")
+            if (
+                response.num_tokens_output
+                and time_to_last_token
+                and time_to_first_token
+            ):
+                generation_time = time_to_last_token - time_to_first_token
+                response.time_per_output_token = (response.num_tokens_output - 1) and (
+                    generation_time / (response.num_tokens_output - 1)
+                )
+
+        return response

llmeter/endpoints/litellm.py

@@ -0,0 +1,165 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import json
+import logging
+import os
+import time
+from typing import Sequence
+from uuid import uuid4
+
+import litellm
+from litellm import CustomStreamWrapper, completion
+from litellm.types.utils import ModelResponse
+from litellm.utils import get_llm_provider  # type: ignore
+
+from llmeter.endpoints import Endpoint, InvocationResponse
+
+logger = logging.getLogger(__name__)
+
+litellm.json_logs = True  # type: ignore
+litellm.turn_off_message_logging = True
+litellm.suppress_debug_info = True
+
+os.environ["LITELLM_LOG"] = "CRITICAL"
+os.environ["LITELLM_DONT_SHOW_FEEDBACK_BOX"] = "true"
+
+
+class LiteLLMBase(Endpoint):
+    def __init__(
+        self,
+        litellm_model: str,
+        model_id: str | None = None,
+    ):
+        self.litellm_model = litellm_model
+        model_id_inferred, provider, _, _ = get_llm_provider(litellm_model)
+
+        logger.info(f"Using model {model_id_inferred} from provider {provider}")
+        super().__init__(
+            model_id=model_id or model_id_inferred,
+            provider=provider,
+            endpoint_name=model_id_inferred,
+        )
+
+    def _parse_payload(self, payload):
+        return json.dumps(payload.get("messages"))
+
+    @staticmethod
+    def create_payload(
+        user_message: str | Sequence[str],
+        max_tokens: int = 256,
+        system_message: str | None = None,
+        **kwargs,
+    ):
+        """
+        Create a payload for the LiteLLM `completion()` request.
+
+        Args:
+            user_message (str | Sequence[str]): The user's message or a sequence of messages.
+            max_tokens (int, optional): The maximum number of tokens to generate. Defaults to 256.
+            **kwargs: Additional keyword arguments to include in the payload.
+
+        Returns:
+            dict: The formatted payload for the Bedrock API request.
+        """
+
+        if isinstance(user_message, str):
+            user_message = [user_message]
+        payload = {
+            "messages": [{"role": "user", "content": k} for k in user_message],
+            "max_tokens": max_tokens,
+        }
+        payload.update(kwargs)
+        if system_message:
+            payload["messages"].append({"role": "system", "content": system_message})
+        return payload
+
+
+class LiteLLM(LiteLLMBase):
+    def invoke(self, payload, **kwargs):
+        try:
+            response = completion(model=self.litellm_model, **payload, **kwargs)
+            assert isinstance(response, ModelResponse)
+            response = self._parse_converse_response(response)
+            response.input_prompt = self._parse_payload(payload)
+            return response
+
+        except Exception as e:
+            logger.exception(e)
+            return InvocationResponse.error_output(
+                id=uuid4().hex, error=str(e), input_prompt=self._parse_payload(payload)
+            )
+
+    def _parse_converse_response(
+        self, client_response: ModelResponse
+    ) -> InvocationResponse:
+        response = InvocationResponse(
+            id=client_response.id,
+            response_text=client_response.choices[0].message.content,  # type: ignore
+        )
+        try:
+            usage = client_response.usage  # type: ignore
+            response.num_tokens_input = usage.prompt_tokens
+            response.num_tokens_output = usage.completion_tokens
+        except AttributeError:
+            pass
+
+        return response
+
+
+class LiteLLMStreaming(LiteLLMBase):
+    def invoke(self, payload, **kwargs):
+        if ("stream" not in kwargs) or ("stream" not in payload):
+            kwargs["stream"] = True
+
+        if ("stream_options" not in kwargs) or ("stream_options" not in payload):
+            kwargs["stream_options"] = {"include_usage": True}
+
+        try:
+            start_t = time.perf_counter()
+            response = completion(model=self.litellm_model, **payload, **kwargs)
+        except Exception as e:
+            logger.exception(e)
+            return InvocationResponse.error_output(
+                id=uuid4().hex, error=str(e), input_prompt=self._parse_payload(payload)
+            )
+
+        assert isinstance(response, CustomStreamWrapper)
+        response = self._parse_stream(response, start_t)
+        response.input_prompt = self._parse_payload(payload)
+        return response
+
+    def _parse_stream(
+        self, client_response: CustomStreamWrapper, start_t: float
+    ) -> InvocationResponse:
+        usage = None
+        time_flag = True
+        time_to_first_token = None
+        output_text = ""
+        for chunk in client_response:
+            output_text += chunk.choices[0].delta.content or ""  # type: ignore
+            if time_flag:
+                time_to_first_token = time.perf_counter() - start_t
+                time_flag = False
+                id = chunk.id
+            try:
+                usage = chunk.usage  # type: ignore
+            except AttributeError:
+                continue
+
+        time_to_last_token = time.perf_counter() - start_t
+
+        response = InvocationResponse(
+            id=id,
+            response_text=output_text,
+            num_tokens_input=usage and usage.prompt_tokens,
+            num_tokens_output=usage and usage.completion_tokens,
+            time_to_first_token=time_to_first_token,
+            time_to_last_token=time_to_last_token,
+        )
+        if response.num_tokens_output and time_to_last_token and time_to_first_token:
+            generation_time = time_to_last_token - time_to_first_token
+            response.time_per_output_token = (response.num_tokens_output - 1) and (
+                generation_time / (response.num_tokens_output - 1)
+            )
+        return response

llmeter/endpoints/openai.py

@@ -0,0 +1,78 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import logging
+from typing import Dict, Sequence
+from uuid import uuid4
+
+import jmespath
+from openai import APIConnectionError, OpenAI
+from openai.types.chat import ChatCompletion
+
+from .base import Endpoint, InvocationResponse
+
+logger = logging.getLogger(__name__)
+
+
+class OpenAIEndpoint(Endpoint):
+    def __init__(
+        self,
+        model_id: str,
+        endpoint_name: str = "openai",
+        api_key: str | None = None,
+        *args,
+        **kwargs,
+    ):
+        super().__init__(
+            endpoint_name,
+            model_id,
+            *args,
+            **kwargs,
+        )
+
+        self._client = OpenAI(api_key=api_key, **kwargs)
+
+    def _parse_payload(self, payload):
+        jpath = "[:].content"
+        messages = payload.get("messages")
+        return "\n".join([k for j in jmespath.search(jpath, messages) for k in j])
+
+    @staticmethod
+    def create_payload(
+        user_message: str | Sequence[str], max_tokens: int = 256, **kwargs
+    ):
+        if isinstance(user_message, str):
+            user_message = [user_message]
+        payload = {
+            "messages": [{"role": "user", "content": k} for k in user_message],
+            "max_tokens": max_tokens,
+        }
+        payload.update(kwargs)
+        return payload
+
+
+class OpenAICompletionEndpoint(OpenAIEndpoint):
+    def invoke(self, payload: Dict, **kwargs) -> InvocationResponse:
+        payload = {**kwargs, **payload}
+
+        payload["model"] = self.model_id
+        try:
+            client_response: ChatCompletion = self._client.chat.completions.create(
+                **payload
+            )
+        except (APIConnectionError, Exception) as e:
+            logger.error(e)
+            return InvocationResponse.error_output(id=uuid4().hex, error=str(e))
+        response = self._parse_converse_response(client_response)
+        response.input_prompt = self._parse_payload(payload)
+        return response
+
+    def _parse_converse_response(self, client_response: ChatCompletion):
+        usage = client_response.usage
+
+        return InvocationResponse(
+            id=client_response.id,
+            response_text=client_response.choices[0].message.content,
+            num_tokens_input=usage and usage.prompt_tokens,
+            num_tokens_output=usage and usage.completion_tokens,
+        )