llmstudio 0.2.2__tar.gz → 0.2.3__tar.gz

llmstudio-0.2.3/MANIFEST.in

@@ -0,0 +1 @@
+recursive-include llmstudio/ui/build *

{llmstudio-0.2.2 → llmstudio-0.2.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: llmstudio
-Version: 0.2.2
+Version: 0.2.3
 Summary: Prompt Perfection at Your Fingertips
 Home-page: https://llmstudio.ai/
 Author: TensorOps

{llmstudio-0.2.2 → llmstudio-0.2.3}/llmstudio/__init__.py

@@ -1,5 +1,5 @@
 name = "version"
-__version__ = "0.2.2"
+__version__ = "0.2.3"
 
 __requirements__ = [
     "pydantic",

llmstudio-0.2.3/llmstudio/engine/providers/__init__.py

@@ -0,0 +1,18 @@
+from llmstudio.engine.config import Provider
+from llmstudio.engine.providers.base_provider import BaseProvider
+
+
+def get_provider(provider: Provider) -> BaseProvider:
+    from llmstudio.engine.providers.bedrock import BedrockProvider
+    from llmstudio.engine.providers.openai import OpenAIProvider
+    from llmstudio.engine.providers.vertexai import VertexAIProvider
+
+    provider_to_class = {
+        Provider.OPENAI: OpenAIProvider,
+        Provider.VERTEXAI: VertexAIProvider,
+        Provider.BEDROCK: BedrockProvider,
+    }
+    if prov := provider_to_class.get(provider):
+        return prov
+
+    raise ValueError(f"Provider {provider} not found")

llmstudio-0.2.3/llmstudio/engine/providers/base_provider.py

@@ -0,0 +1,70 @@
+from abc import ABC
+from concurrent.futures import ThreadPoolExecutor
+
+
+class BaseProvider(ABC):
+    """
+    Abstract base class for LLMStudio engine providers.
+
+    This class defines the core interface for providers that integrate
+    with LLMStudio's engine. It is intended to be subclassed by specific
+    provider implementations.
+    """
+
+    def __init__(self):
+        super().__init__()
+        self.executor = ThreadPoolExecutor(max_workers=10)
+
+    async def chat(self, data) -> dict:
+        """
+        Asynchronously handle a chat request.
+
+        Parameters:
+            data: The data payload for the chat operation.
+
+        Returns:
+            dict: A dictionary containing the chat response.
+
+        Raises:
+            NotImplementedError: This method is intended to be overridden by subclasses.
+        """
+        raise NotImplementedError
+
+    async def test(test, data) -> dict:
+        """
+        Asynchronously handle a test request.
+
+        Parameters:
+            data: The data payload for the test operation.
+
+        Returns:
+            dict: A dictionary containing the test response.
+
+        Raises:
+            NotImplementedError: This method is intended to be overridden by subclasses.
+        """
+        raise NotImplementedError
+
+    def validate_model_field(self, data, model_list):
+        """
+        Validate the 'model_name' field in the request data.
+
+        Parameters:
+            data: The data payload containing the 'model_name'.
+            model_list: List of valid model names.
+
+        Raises:
+            HTTPException: If the 'model_name' is not provided or is not in the list of valid models.
+        """
+        from fastapi import HTTPException
+
+        if not data.model_name:
+            raise HTTPException(
+                status_code=422,
+                detail="The parameter 'model_name' is mandatory to be passed in the request body.",
+            )
+        if data.model_name not in model_list:
+            raise HTTPException(
+                status_code=422,
+                detail=f"The model '{data['model_name']}' does not exist.",
+            )

llmstudio-0.2.3/llmstudio/engine/providers/bedrock.py

@@ -0,0 +1,313 @@
+import asyncio
+import json
+import random
+import time
+from typing import Optional, Tuple
+
+import boto3
+from fastapi.responses import StreamingResponse
+from pydantic import BaseModel, Field, validator
+
+from llmstudio.engine.config import BedrockConfig
+from llmstudio.engine.constants import BEDROCK_MODELS, CLAUDE_MODELS, END_TOKEN, TITAN_MODELS
+from llmstudio.engine.providers.base_provider import BaseProvider
+from llmstudio.engine.utils import validate_provider_config
+
+
+class ClaudeParameters(BaseModel):
+    """
+    Model for validating and storing parameters specific to Claude model.
+
+    Attributes:
+        temperature (Optional[float]): Controls randomness in the model's output.
+        max_tokens (Optional[int]): The maximum number of tokens in the output.
+        top_p (Optional[float]): Influences the diversity of output by controlling token sampling.
+        top_k (Optional[float]): Sets the number of the most likely next tokens to filter for.
+    """
+
+    temperature: Optional[float] = Field(1, ge=0, le=1)
+    max_tokens: Optional[int] = Field(300, ge=1, le=2048)
+    top_p: Optional[float] = Field(0.999, ge=0, le=1)
+    top_k: Optional[int] = Field(250, ge=1, le=500)
+
+
+class TitanParameters(BaseModel):
+    """
+    Model for validating and storing parameters specific to Titan model.
+
+    Attributes:
+        temperature (Optional[float]): Controls randomness in the model's output.
+        max_tokens (Optional[int]): The maximum number of tokens in the output.
+        top_p (Optional[float]): Influences the diversity of output by controlling token sampling.
+    """
+
+    temperature: Optional[float] = Field(0, ge=0, le=1)
+    max_tokens: Optional[int] = Field(512, ge=1, le=4096)
+    top_p: Optional[float] = Field(0.9, ge=0.1, le=1)
+
+
+class BedrockRequest(BaseModel):
+    """
+    Represents a request to the Bedrock API.
+
+    Attributes:
+        api_key (Optional[str]): The API key for authenticating with the Bedrock API.
+        api_secret (Optional[str]): The API secret for authenticating with the Bedrock API.
+        api_region (Optional[str]): The region where the Bedrock API is hosted.
+        model_name (str): The name of the model to be used for the request.
+        chat_input (str): The input string for the chat.
+        parameters (Optional[BaseModel]): Additional parameters for the model, encapsulated in a BaseModel.
+        is_stream (Optional[bool]): Flag to indicate if the request is for streaming. Defaults to False.
+    """
+
+    api_key: Optional[str]
+    api_secret: Optional[str]
+    api_region: Optional[str]
+    model_name: str
+    chat_input: str
+    parameters: Optional[BaseModel]
+    is_stream: Optional[bool] = False
+
+    @validator("parameters", pre=True, always=True)
+    def validate_parameters_based_on_model_name(cls, parameters, values):
+        """
+        Validate and convert parameters based on the model_name.
+
+        Args:
+            parameters (Dict[str, Any]): Parameters to validate and convert.
+            values (Dict[str, Any]): Contains previously validated fields.
+
+        Returns:
+            BaseModel: An instance of `TitanParameters` or `ClaudeParameters` based on `model_name`.
+
+        Raises:
+            ValueError if model_name is invalid.
+        """
+        model_name = values.get("model_name")
+        if model_name in TITAN_MODELS:
+            return TitanParameters(**parameters)
+        if model_name in CLAUDE_MODELS:
+            return ClaudeParameters(**parameters)
+
+        raise ValueError(f"Invalid model_name: {model_name}")
+
+
+class BedrockTest(BaseModel):
+    """
+    A Pydantic model for validating Bedrock API requests.
+
+    Attributes:
+        api_key (str): The API key provided by the user for authentication with Bedrock's API.
+        api_secret (str): The API secret key provided by the user for authentication.
+        api_region (str): The API region for Bedrock API requests.
+        model_name (str): The name of the model intended for use with the Bedrock API.
+
+    Methods:
+        validate_model_name: Ensures that `model_name` is one of the allowed values.
+    """
+
+    api_key: Optional[str]
+    api_secret: Optional[str]
+    api_region: Optional[str]
+    model_name: str
+
+
+class BedrockProvider(BaseProvider):
+    """
+    BedrockProvider class to interact with the Bedrock API.
+
+    Attributes:
+        bedrock_config (BedrockConfig): Configuration for the Bedrock API.
+    """
+
+    def __init__(self, config: BedrockConfig, api_key: dict):
+        """
+        Initialize the BedrockProvider class.
+
+        Args:
+            config (BedrockConfig): Configuration for the Bedrock API.
+            api_key (dict): API key required for the Bedrock API.
+
+        Raises:
+            ValidationError: If the provided config and API key are invalid.
+        """
+        super().__init__()
+        self.bedrock_config = validate_provider_config(config, api_key)
+
+    async def chat(self, data: BedrockRequest) -> dict:
+        """
+        Endpoint to process chat input via Bedrock API and generate a model's response.
+
+        Args:
+            data (BedrockRequest): Validated API request data.
+
+        Returns:
+            Union[StreamingResponse, dict]: Streaming response if is_stream is True, otherwise a dict with chat and token data.
+        """
+        data = BedrockRequest(**data)
+        self.validate_model_field(data, BEDROCK_MODELS)
+        loop = asyncio.get_event_loop()
+        session = boto3.Session(
+            aws_access_key_id=self.bedrock_config["api_key"],
+            aws_secret_access_key=self.bedrock_config["api_secret"],
+        )
+        bedrock = session.client(
+            service_name="bedrock", region_name=self.bedrock_config["api_region"]
+        )
+
+        body, response_keys = generate_body_and_response(data)
+
+        if data.is_stream:
+            response = await loop.run_in_executor(
+                None,
+                lambda: bedrock.invoke_model_with_response_stream(
+                    body=json.dumps(body),
+                    modelId=data.model_name,
+                    accept="application/json",
+                    contentType="application/json",
+                ).get("body"),
+            )
+            return StreamingResponse(generate_stream_response(response, response_keys))
+        else:
+            response = await loop.run_in_executor(
+                None,
+                lambda: json.loads(
+                    bedrock.invoke_model(
+                        body=json.dumps(body),
+                        modelId=data.model_name,
+                        accept="application/json",
+                        contentType="application/json",
+                    )
+                    .get("body")
+                    .read()
+                ),
+            )
+
+        response = response["results"][0] if response_keys["use_results"] else response
+
+        data = {
+            "id": random.randint(0, 1000),
+            "chatInput": data.chat_input,
+            "chatOutput": response[response_keys["output_key"]],
+            "inputTokens": response.get(response_keys["input_tokens_key"], 0),
+            "outputTokens": response.get(response_keys["output_tokens_key"], 0),
+            "totalTokens": response.get(response_keys["input_tokens_key"], 0)
+            + response.get(response_keys["output_tokens_key"], 0),
+            "cost": 0,  # TODO
+            "timestamp": time.time(),
+            "modelName": data.model_name,
+            "parameters": data.parameters.dict(),
+        }
+        return data
+
+    async def test(self, data: BedrockTest) -> bool:
+        """
+        Test the validity of the Bedrock API credentials and model name.
+
+        Args:
+            data (BedrockTest): A model instance containing the Bedrock API credentials
+                            and model name to test.
+
+        Returns:
+            bool: `True` if the API credentials and model name are valid, otherwise `False`.
+        """
+        data = BedrockTest(**data)
+        try:
+            session = boto3.Session(
+                aws_access_key_id=self.bedrock_config["api_key"],
+                aws_secret_access_key=self.bedrock_config["api_secret"],
+            )
+            bedrock = session.client(
+                service_name="bedrock", region_name=self.bedrock_config["api_region"]
+            )
+            response = bedrock.list_foundation_models()
+
+            if data.model_name in [i["modelId"] for i in response["modelSummaries"]]:
+                return True
+            else:
+                return False
+        except Exception:
+            return False
+
+
+def generate_body_and_response(data: BedrockProvider) -> Tuple[dict, dict]:
+    """
+    Generate request body and response keys based on model name.
+
+    Args:
+        data (BedrockRequest): Validated API request data.
+
+    Returns:
+        Tuple[dict, dict]: Tuple of request body and response keys.
+
+    Raises:
+        ValueError if model name is invalid.
+    """
+    if data.model_name in TITAN_MODELS:
+        return {
+            "inputText": data.chat_input,
+            "textGenerationConfig": {
+                "maxTokenCount": data.parameters.max_tokens,
+                "temperature": data.parameters.temperature,
+                "topP": data.parameters.top_p,
+            },
+        }, {
+            "output_key": "outputText",
+            "input_tokens_key": "inputTextTokenCount",
+            "output_tokens_key": "tokenCount",
+            "use_results": True,
+        }
+    if data.model_name in CLAUDE_MODELS:
+        return {
+            "prompt": data.chat_input,
+            "max_tokens_to_sample": data.parameters.max_tokens,
+            "temperature": data.parameters.temperature,
+            "top_k": data.parameters.top_k,
+            "top_p": data.parameters.top_p,
+        }, {
+            "output_key": "completion",
+            "input_tokens_key": None,
+            "output_tokens_key": None,
+            "use_results": False,
+        }
+    else:
+        raise ValueError(f"Invalid model_name: {data.model_name}")
+
+
+def get_cost(input_tokens: int, output_tokens: int) -> float:
+    """
+    Calculate the cost based on input and output tokens.
+
+    Args:
+        input_tokens (int): Number of tokens in the input.
+        output_tokens (int): Number of tokens in the output.
+
+    Returns:
+        float: Cost.
+    """
+    return None
+
+
+def generate_stream_response(response, response_keys):
+    """
+    Generate streaming response based on response events and keys.
+
+    Args:
+        response (Any): Response from the Bedrock API call.
+        response_keys (Dict[str, Any]): Keys to extract relevant data from the response.
+
+    Yields:
+        str: Extracted data from response chunks.
+    """
+    chat_output = ""
+    for event in response:
+        chunk = event.get("chunk")
+        if chunk:
+            chunk_content = json.loads(chunk.get("bytes").decode())[response_keys["output_key"]]
+            chat_output += chunk_content
+            yield chunk_content
+
+    input_tokens = 0
+    output_tokens = 0
+    cost = 0
+    yield f"{END_TOKEN},{input_tokens},{output_tokens},{cost}"

llmstudio-0.2.3/llmstudio/engine/providers/openai.py

@@ -0,0 +1,233 @@
+import asyncio
+import random
+import time
+from typing import Optional
+
+import openai
+import tiktoken
+from fastapi.responses import StreamingResponse
+from pydantic import BaseModel, Field
+
+from llmstudio.engine.config import OpenAIConfig
+from llmstudio.engine.constants import END_TOKEN, OPENAI_PRICING_DICT
+from llmstudio.engine.providers.base_provider import BaseProvider
+from llmstudio.engine.utils import validate_provider_config
+
+
+class OpenAIParameters(BaseModel):
+    """
+    A Pydantic model for encapsulating parameters used in OpenAI API requests.
+
+    Attributes:
+        temperature (Optional[float]): Controls randomness in the model's output.
+        max_tokens (Optional[int]): The maximum number of tokens in the output.
+        top_p (Optional[float]): Influences the diversity of output by controlling token sampling.
+        frequency_penalty (Optional[float]): Modifies the likelihood of tokens appearing based on their frequency.
+        presence_penalty (Optional[float]): Adjusts the likelihood of new tokens appearing.
+    """
+
+    temperature: Optional[float] = Field(default=1, ge=0, le=2)
+    max_tokens: Optional[int] = Field(default=256, ge=1, le=2048)
+    top_p: Optional[float] = Field(default=1, ge=0, le=1)
+    frequency_penalty: Optional[float] = Field(default=0, ge=0, le=1)
+    presence_penalty: Optional[float] = Field(default=0, ge=0, le=1)
+
+
+class OpenAIRequest(BaseModel):
+    """
+    A Pydantic model that represents a request to an OpenAI API.
+
+    Attributes:
+        api_key (Optional[str]): The API key to use for authenticating the request.
+        model_name (str): The name of the language model to query.
+        chat_input (str): The input text to send to the model.
+        parameters (Optional[OpenAIParameters]): An optional instance of OpenAIParameters to further configure the request.
+        is_stream (Optional[bool]): Indicates if the request should be a streaming request; default is False.
+
+    """
+
+    api_key: Optional[str]
+    model_name: str
+    chat_input: str
+    parameters: Optional[OpenAIParameters] = OpenAIParameters()
+    is_stream: Optional[bool] = False
+
+
+class OpenAITest(BaseModel):
+    """
+    A Pydantic model for validating OpenAI API requests.
+
+    Attributes:
+        api_key (str): The API key provided by the user authentication with OpenAI API.
+        model_name (str): The name of the model to be used for generating text
+
+    Methods:
+        validate_model_name: Ensures that `model_name` is one of the allowed values.
+    ```
+    """
+
+    api_key: Optional[str]
+    model_name: str
+
+
+class OpenAIProvider(BaseProvider):
+    """
+    A provider class to handle interactions with the OpenAI GPT models.
+
+    Attributes:
+        openai_config (OpenAIConfig): Configuration settings for OpenAI API.
+    """
+
+    def __init__(self, config: OpenAIConfig, api_key: str):
+        """
+        Initialize the OpenAIProvider with given config and API key.
+
+        Args:
+            config (OpenAIConfig): Configuration settings for OpenAI API.
+            api_key (str): API key for authentication.
+        """
+        super().__init__()
+        if isinstance(config, OpenAIConfig):
+            self.openai_config = config
+        else:
+            self.openai_config = OpenAIConfig(**validate_provider_config(config, api_key))
+
+    async def chat(self, data: OpenAIRequest) -> dict:
+        """
+        Generate chat-based model completions using OpenAI API.
+
+        Args:
+            data (OpenAIRequest): A model instance containing chat input, model name, and additional parameters.
+
+        Returns:
+            dict: A dictionary containing chat input, chat output, tokens information, cost, and other metadata.
+
+        Raises:
+            ValueError: If the specified model field is invalid.
+        """
+
+        data = OpenAIRequest(**data)
+
+        self.validate_model_field(data, OPENAI_PRICING_DICT.keys())
+        openai.api_key = self.openai_config.api_key
+
+        # Asynchronous call, for parallelism
+        loop = asyncio.get_event_loop()
+
+        response = await loop.run_in_executor(
+            self.executor,
+            lambda: openai.ChatCompletion.create(
+                model=data.model_name,
+                messages=[
+                    {
+                        "role": "user",
+                        "content": data.chat_input,
+                    }
+                ],
+                temperature=data.parameters.temperature,
+                max_tokens=data.parameters.max_tokens,
+                top_p=data.parameters.top_p,
+                frequency_penalty=data.parameters.frequency_penalty,
+                presence_penalty=data.parameters.presence_penalty,
+                stream=data.is_stream,
+            ),
+        )
+
+        if data.is_stream:
+            return StreamingResponse(generate_stream_response(response, data))
+
+        input_tokens = get_tokens(data.chat_input, data.model_name)
+        output_tokens = get_tokens(response["choices"][0]["message"]["content"], data.model_name)
+
+        data = {
+            "id": random.randint(0, 1000),
+            "chatInput": data.chat_input,
+            "chatOutput": response["choices"][0]["message"]["content"],
+            "inputTokens": input_tokens,
+            "outputTokens": output_tokens,
+            "totalTokens": input_tokens + output_tokens,
+            "cost": get_cost(input_tokens, output_tokens, data.model_name),
+            "timestamp": time.time(),
+            "modelName": data.model_name,
+            "parameters": data.parameters.dict(),
+        }
+        return data
+
+    async def test(self, data: OpenAITest) -> bool:
+        """
+        Test the validity of the OpenAI API key.
+
+        Args:
+            data (OpenAITest): A model instance which includes the API key for OpenAI.
+
+        Returns:
+            bool: `True` if the API key is valid and initialization succeeds, otherwise `False`.
+        """
+        openai.api_key = self.openai_config.api_key
+        data = OpenAITest(**data)
+        try:
+            self.validate_model_field(data, OPENAI_PRICING_DICT.keys())
+            openai.Model.retrieve(data.model_name)
+            return True
+        except Exception:
+            return False
+
+
+def get_cost(input_tokens: int, output_tokens: int, model_name: str) -> float:
+    """
+    Calculate the cost of using the OpenAI API based on token usage and model.
+
+    Args:
+        input_tokens (int): Number of tokens in the input.
+        output_tokens (int): Number of tokens in the output.
+        model_name (str): Identifier of the model used.
+
+    Returns:
+        float: The calculated cost for the API usage.
+    """
+    return (
+        OPENAI_PRICING_DICT[model_name]["input_tokens"] * input_tokens
+        + OPENAI_PRICING_DICT[model_name]["output_tokens"] * output_tokens
+    )
+
+
+def get_tokens(chat_input: str, model_name: str) -> int:
+    """
+    Determine the number of tokens in a given input string using the specified model’s tokenizer.
+
+    Args:
+        chat_input (str): Text to be tokenized.
+        model_name (str): Identifier of the model, determines tokenizer used.
+
+    Returns:
+        int: Number of tokens in the input string.
+    """
+    tokenizer = tiktoken.encoding_for_model(model_name)
+    return len(tokenizer.encode(chat_input))
+
+
+def generate_stream_response(response: dict, data: OpenAIProvider):
+    """
+    Generate stream responses, yielding chat output or tokens and cost information at stream end.
+
+    Args:
+        response (dict): Dictionary containing chunks of responses from the OpenAI API.
+        data (OpenAIRequest): OpenAIRequest object containing necessary parameters for the API call.
+
+    Yields:
+        str: A chunk of chat output or, at stream end, tokens counts and cost information.
+    """
+    chat_output = ""
+    for chunk in response:
+        if (
+            chunk["choices"][0]["finish_reason"] != "stop"
+            and chunk["choices"][0]["finish_reason"] != "length"
+        ):
+            chunk_content = chunk["choices"][0]["delta"]["content"]
+            chat_output += chunk_content
+            yield chunk_content
+        else:
+            input_tokens = get_tokens(data.chat_input, data.model_name)
+            output_tokens = get_tokens(chat_output, data.model_name)
+            cost = get_cost(input_tokens, output_tokens, data.model_name)
+            yield f"{END_TOKEN},{input_tokens},{output_tokens},{cost}"  # json