AI butler Jarvis

agent_jarvis/jarvis/gpt/ai_function.py

@@ -0,0 +1,41 @@
+from typing import List
+
+from jarvis import CFG
+from jarvis.gpt import gpt
+from jarvis.gpt.message import Message
+from jarvis.logger import logger
+
+
+async def acall_ai_function(function: str, args: list, description: str, model: str | None = None) -> str:
+    """Call an AI function
+
+    This is a magic function that can do anything with no-code. See
+    https://github.com/Torantulino/AI-Functions for more info.
+
+    Args:
+        function (str): The function to call
+        args (list): The arguments to pass to the function
+        description (str): The description of the function
+        model (str, optional): The model to use. Defaults to None.
+
+    Returns:
+        str: The response from the function
+    """
+    if model is None:
+        model = CFG.small_llm_model
+    # For each arg, if any are None, convert to "None":
+    args = [str(arg) if arg is not None else "None" for arg in args]
+    # parse args to comma separated string
+    args: str = ", ".join(args)
+    messages: List[Message] = [
+        {
+            "role": "system",
+            "content": f"You are now the following python function: ```# {description}"
+                       f"\n{function}```\n\nOnly respond with your `return` value.",
+        },
+        {"role": "user", "content": args},
+    ]
+
+    logger.debug(str(messages))
+
+    return await gpt.acreate_chat_completion(model=model, messages=messages, temperature=0)

agent_jarvis/jarvis/gpt/gpt.py

@@ -0,0 +1,134 @@
+import asyncio
+
+import openai
+from openai.error import RateLimitError, APIError, Timeout
+
+from jarvis import CFG
+from jarvis.gpt.message import Message
+from jarvis.logger import logger
+from typing import Callable
+
+openai.api_key = CFG.openai_api_key
+if CFG.openai_url_base is not None:
+    openai.api_base = CFG.openai_url_base
+
+print_total_cost = CFG.debug_mode
+
+
+async def acreate_chat_completion_once(
+        messages: list,  # type: ignore
+        model: str | None = None,
+        temperature: float = CFG.temperature,
+        max_tokens: int | None = None,
+        deployment_id=None,
+        request_timeout=40,
+) -> str:
+    """
+    Create a chat completion and update the cost.
+    Args:
+    messages (list): The list of messages to send to the API.
+    model (str): The model to use for the API call.
+    temperature (float): The temperature to use for the API call.
+    max_tokens (int): The maximum number of tokens for the API call.
+    Returns:
+    str: The AI's response.
+    """
+    if deployment_id is not None:
+        response = await openai.ChatCompletion.acreate(
+            deployment_id=deployment_id,
+            model=model,
+            messages=messages,
+            temperature=temperature,
+            max_tokens=max_tokens,
+            request_timeout=request_timeout
+        )
+    else:
+        response = await openai.ChatCompletion.acreate(
+            model=model,
+            messages=messages,
+            temperature=temperature,
+            max_tokens=max_tokens,
+            request_timeout=request_timeout
+        )
+    if CFG.debug_mode:
+        logger.debug(f"Response: {response}")
+    # prompt_tokens = response.usage.prompt_tokens
+    # completion_tokens = response.usage.completion_tokens
+    return response
+
+
+# Overly simple abstraction until we create something better
+# simple retry mechanism when getting a rate error or a bad gateway
+async def acreate_chat_completion(
+        messages: list[Message],  # type: ignore
+        model: str = None,
+        temperature: float = CFG.temperature,
+        max_tokens: int = None,
+        request_timeout: int = 40,
+        num_retries=3,
+        on_single_request_timeout: Callable = None
+):
+    """Create a chat completion using the OpenAI API
+
+    Args:
+        messages (List[Message]): The messages to send to the chat completion
+        model (str, optional): The model to use. Defaults to None.
+        temperature (float, optional): The temperature to use. Defaults to 0.9.
+        max_tokens (int, optional): The max tokens to use. Defaults to None.
+        request_timeout (int, optional): The request_timeout of a single openai request.
+        num_retries (int, optional): The max retries.
+        on_single_request_timeout (Callable, optional): This function will be called each time a single openai request
+            timeout, must be an async function, the last timeout will not emit callback.
+
+    Returns:
+        str: The response from the chat completion
+    """
+    if CFG.debug_mode:
+        logger.debug(
+            f"Creating chat completion with model {model}, temperature {temperature}, max_tokens {max_tokens}"
+        )
+
+    response = None
+
+    for attempt in range(num_retries):
+        backoff = min(2 ** (attempt + 2), 8)
+        try:
+            if CFG.use_azure:
+                response = await acreate_chat_completion_once(
+                    deployment_id=CFG.get_azure_deployment_id_for_model(model),
+                    model=model,
+                    messages=messages,
+                    temperature=temperature,
+                    max_tokens=max_tokens,
+                    request_timeout=request_timeout,
+                )
+            else:
+                response = await acreate_chat_completion_once(
+                    model=model,
+                    messages=messages,
+                    temperature=temperature,
+                    max_tokens=max_tokens,
+                    request_timeout=request_timeout,
+                )
+            break
+        except RateLimitError:
+            if CFG.debug_mode:
+                logger.debug(f"Error: Reached rate limit, passing...")
+        except (APIError, Timeout) as e:
+            if isinstance(e, Timeout):
+                if on_single_request_timeout:
+                    await on_single_request_timeout(num_retries < num_retries - 1)
+            if e.http_status != 502:
+                raise
+            if attempt == num_retries - 1:
+                raise
+        if CFG.debug_mode:
+            logger.debug(
+                f"Error: API Bad gateway. Waiting {backoff} seconds..."
+            )
+        await asyncio.sleep(backoff)
+    if response is None:
+        logger.error(f"Failed to get response from GPT after {num_retries} retries")
+        raise RuntimeError(f"Failed to get response after {num_retries} retries")
+    resp = response.choices[0].message["content"]
+    return resp

agent_jarvis/jarvis/gpt/message.py

@@ -0,0 +1,9 @@
+"""Type helpers for working with the OpenAI library"""
+from typing import TypedDict
+
+
+class Message(TypedDict):
+    """OpenAI Message object containing a role and the message content"""
+
+    role: str
+    content: str

agent_jarvis/jarvis/gpt/token_counter.py

@@ -0,0 +1,75 @@
+"""Functions for counting the number of tokens in a message or string."""
+from __future__ import annotations
+
+from typing import List
+
+import tiktoken_async
+
+from jarvis.gpt.message import Message
+
+
+async def count_message_tokens(
+        messages: List[Message], model: str = "gpt-3.5-turbo-0301"
+) -> int:
+    """
+    Returns the number of tokens used by a list of messages.
+
+    Args:
+        messages (list): A list of messages, each of which is a dictionary
+            containing the role and content of the message.
+        model (str): The name of the model to use for tokenization.
+            Defaults to "gpt-3.5-turbo-0301".
+
+    Returns:
+        int: The number of tokens used by the list of messages.
+    """
+    try:
+        encoding = await tiktoken_async.encoding_for_model(model)
+    except KeyError:
+        print("Warning: model not found. Using cl100k_base encoding.")
+        encoding = await tiktoken_async.get_encoding("cl100k_base")
+    if model == "gpt-3.5-turbo":
+        # !Note: gpt-3.5-turbo may change over time.
+        # Returning num tokens assuming Mgpt-3.5-turbo-0301.")
+        return await count_message_tokens(messages, model="gpt-3.5-turbo-0301")
+    elif model == "gpt-4":
+        # !Note: gpt-4 may change over time. Returning num tokens assuming gpt-4-0314.")
+        return await count_message_tokens(messages, model="gpt-4-0314")
+    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
+    elif model == "gpt-4-0314":
+        tokens_per_message = 3
+        tokens_per_name = 1
+    else:
+        raise NotImplementedError(
+            f"num_tokens_from_messages() is not implemented for model {model}.\n"
+            " 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
+
+
+async def count_string_tokens(string: str, model_name: str) -> int:
+    """
+    Returns the number of tokens in a text string.
+
+    Args:
+        string (str): The text string.
+        model_name (str): The name of the encoding to use. (e.g., "gpt-3.5-turbo")
+
+    Returns:
+        int: The number of tokens in the text string.
+    """
+    encoding = await tiktoken_async.encoding_for_model(model_name)
+    return len(encoding.encode(string))